ATLAS-3700: added option to append value to array, map type attributes

intg/src/main/java/org/apache/atlas/model/instance/EntityMutationResponse.java

@@ -209,6 +209,11 @@ public class EntityMutationResponse {
     }
 
     @JsonIgnore
+    public AtlasEntityHeader getFirstPartialUpdatedEntityByTypeName(String typeName) {
+        return getFirstEntityByType(getEntitiesByOperation(EntityOperation.PARTIAL_UPDATE), typeName);
+    }
+
+    @JsonIgnore
     public void addEntity(EntityOperation op, AtlasEntityHeader header) {
         // if an entity is already included in CREATE, ignore subsequent UPDATE, PARTIAL_UPDATE
         if (op == EntityOperation.UPDATE || op == EntityOperation.PARTIAL_UPDATE) {

intg/src/main/java/org/apache/atlas/model/typedef/AtlasStructDef.java

@@ -263,10 +263,11 @@ public class AtlasStructDef extends AtlasBaseTypeDef implements Serializable {
         public static final int       DEFAULT_SEARCHWEIGHT          = -1;
 
 
-        public static final String    SEARCH_WEIGHT_ATTR_NAME       = "searchWeight";
-        public static final String    INDEX_TYPE_ATTR_NAME          = "indexType";
-        public static final String    ATTRDEF_OPTION_SOFT_REFERENCE = "isSoftReference";
-        private final String          STRING_TRUE                   = "true";
+        public static final String    SEARCH_WEIGHT_ATTR_NAME                 = "searchWeight";
+        public static final String    INDEX_TYPE_ATTR_NAME                    = "indexType";
+        public static final String    ATTRDEF_OPTION_SOFT_REFERENCE           = "isSoftReference";
+        public static final String    ATTRDEF_OPTION_APPEND_ON_PARTIAL_UPDATE = "isAppendOnPartialUpdate";
+        private final String          STRING_TRUE                             = "true";
 
         /**
          * single-valued attribute or multi-valued attribute.
@@ -520,6 +521,13 @@ public class AtlasStructDef extends AtlasBaseTypeDef implements Serializable {
         }
 
         @JsonIgnore
+        public boolean isAppendOnPartialUpdate() {
+            String val = getOption(AtlasAttributeDef.ATTRDEF_OPTION_APPEND_ON_PARTIAL_UPDATE);
+
+            return val != null && Boolean.valueOf(val);
+        }
+
+        @JsonIgnore
         public void setOption(String name, String value) {
             if (this.options == null) {
                 this.options = new HashMap<>();

intg/src/test/java/org/apache/atlas/TestUtilsV2.java

@@ -627,6 +627,13 @@ public final class TestUtilsV2 {
     }
 
     public static AtlasTypesDef defineSimpleAttrType() {
+        AtlasAttributeDef attrPuArray = new AtlasAttributeDef("puArray", "array<string>", true, SINGLE, 1, 1, false, false, false, null);
+        AtlasAttributeDef attrPuMap   = new AtlasAttributeDef("puMap", "map<string,string>",  true, SINGLE, 1,1, false, false, false, null);
+
+        attrPuArray.setOption(AtlasAttributeDef.ATTRDEF_OPTION_APPEND_ON_PARTIAL_UPDATE, "true");
+        attrPuMap.setOption(AtlasAttributeDef.ATTRDEF_OPTION_APPEND_ON_PARTIAL_UPDATE, "true");
+
+
         AtlasEntityDef simpleAttributesEntityType =
             createClassTypeDef(ENTITY_TYPE_WITH_SIMPLE_ATTR, ENTITY_TYPE_WITH_SIMPLE_ATTR + "_description", null,
                 createUniqueRequiredAttrDef("name", "string"),
@@ -641,7 +648,11 @@ public final class TestUtilsV2 {
                     false, false, false, null),
 
                 new AtlasAttributeDef("mapOfStrings", "map<string,string>",
-                    true, SINGLE, 1,1, false, false, false, null)
+                    true, SINGLE, 1,1, false, false, false, null),
+
+                attrPuArray,
+
+                attrPuMap
             );
 
         AtlasTypesDef ret = AtlasTypeUtil.getTypesDef(Collections.<AtlasEnumDef>emptyList(),
@@ -659,6 +670,8 @@ public final class TestUtilsV2 {
         entity.setAttribute("stringAtrr", "DummyThree");
         entity.setAttribute("arrayOfStrings", Arrays.asList("DummyOne", "DummyTwo"));
         entity.setAttribute("mapOfStrings", Collections.singletonMap("one", "DummyString"));
+        entity.setAttribute("puArray", Arrays.asList("DummyOne", "DummyTwo"));
+        entity.setAttribute("puMap", Collections.singletonMap("one", "DummyString"));
 
         return new AtlasEntityWithExtInfo(entity);
     }

repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphMapper.java

@@ -284,6 +284,8 @@ public class EntityGraphMapper {
             }
         }
 
+        EntityOperation updateType = isPartialUpdate ? PARTIAL_UPDATE : UPDATE;
+
         if (CollectionUtils.isNotEmpty(updatedEntities)) {
             for (AtlasEntity updatedEntity : updatedEntities) {
                 String          guid       = updatedEntity.getGuid();
@@ -292,14 +294,10 @@ public class EntityGraphMapper {
 
                 mapRelationshipAttributes(updatedEntity, entityType, vertex, UPDATE, context);
 
-                mapAttributes(updatedEntity, entityType, vertex, UPDATE, context);
+                mapAttributes(updatedEntity, entityType, vertex, updateType, context);
                 setCustomAttributes(vertex,updatedEntity);
 
-                if (isPartialUpdate) {
-                    resp.addEntity(PARTIAL_UPDATE, constructHeader(updatedEntity, entityType, vertex));
-                } else {
-                    resp.addEntity(UPDATE, constructHeader(updatedEntity, entityType, vertex));
-                }
+                resp.addEntity(updateType, constructHeader(updatedEntity, entityType, vertex));
 
                 if (replaceClassifications) {
                     deleteClassifications(guid);
@@ -325,12 +323,7 @@ public class EntityGraphMapper {
         }
 
         for (AtlasEntityHeader entity : req.getUpdatedEntities()) {
-            if (isPartialUpdate) {
-                resp.addEntity(PARTIAL_UPDATE, entity);
-            }
-            else {
-                resp.addEntity(UPDATE, entity);
-            }
+            resp.addEntity(updateType, entity);
         }
 
         RequestContext.get().endMetricRecord(metric);
@@ -623,7 +616,7 @@ public class EntityGraphMapper {
                     mapAttribute(attribute, attrValue, vertex, op, context);
                 }
 
-            } else if (op.equals(UPDATE)) {
+            } else if (op.equals(UPDATE) || op.equals(PARTIAL_UPDATE)) {
                 for (String attrName : struct.getAttributes().keySet()) {
                     AtlasAttribute attribute = structType.getAttribute(attrName);
 
@@ -665,7 +658,7 @@ public class EntityGraphMapper {
                     mapAttribute(attribute, attrValue, vertex, op, context);
                 }
 
-            } else if (op.equals(UPDATE)) {
+            } else if (op.equals(UPDATE) || op.equals(PARTIAL_UPDATE)) {
                 // relationship attributes mapping
                 for (String attrName : entityType.getRelationshipAttributes().keySet()) {
                     if (entity.hasRelationshipAttribute(attrName)) {
@@ -1251,6 +1244,22 @@ public class EntityGraphMapper {
         boolean             isReference = isReference(mapType.getValueType());
         boolean             isSoftReference = ctx.getAttribute().getAttributeDef().isSoftReferenced();
 
+        if (PARTIAL_UPDATE.equals(ctx.getOp()) && attribute.getAttributeDef().isAppendOnPartialUpdate() && MapUtils.isNotEmpty(currentMap)) {
+            if (MapUtils.isEmpty(newVal)) {
+                newVal = new HashMap<>(currentMap);
+            } else {
+                Map<Object, Object> mergedVal = new HashMap<>(currentMap);
+
+                for (Map.Entry<Object, Object> entry : newVal.entrySet()) {
+                    String newKey = entry.getKey().toString();
+
+                    mergedVal.put(newKey, entry.getValue());
+                }
+
+                newVal = mergedVal;
+            }
+        }
+
         boolean isNewValNull = newVal == null;
 
         if (isNewValNull) {
@@ -1334,7 +1343,6 @@ public class EntityGraphMapper {
         Cardinality    cardinality         = attribute.getAttributeDef().getCardinality();
         List<Object>   newElementsCreated  = new ArrayList<>();
         List<Object>   currentElements;
-        boolean isNewElementsNull          = newElements == null;
 
         if (isReference && !isSoftReference) {
             currentElements = (List) getCollectionElementsUsingRelationship(ctx.getReferringVertex(), attribute);
@@ -1342,6 +1350,20 @@ public class EntityGraphMapper {
             currentElements = (List) getArrayElementsProperty(elementType, isSoftReference, ctx.getReferringVertex(), ctx.getVertexProperty());
         }
 
+        if (PARTIAL_UPDATE.equals(ctx.getOp()) && attribute.getAttributeDef().isAppendOnPartialUpdate() && CollectionUtils.isNotEmpty(currentElements)) {
+            if (CollectionUtils.isEmpty(newElements)) {
+                newElements = new ArrayList<>(currentElements);
+            } else {
+                List<Object> mergedVal = new ArrayList<>(currentElements);
+
+                mergedVal.addAll(newElements);
+
+                newElements = mergedVal;
+            }
+        }
+
+        boolean isNewElementsNull = newElements == null;
+
         if (isNewElementsNull) {
             newElements = new ArrayList();
         }

repository/src/test/java/org/apache/atlas/repository/store/graph/v2/AtlasComplexAttributesTest.java

@@ -32,6 +32,7 @@ import org.apache.atlas.repository.graphdb.AtlasEdgeDirection;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.atlas.type.AtlasEntityType;
 import org.apache.commons.lang.time.DateUtils;
+import org.locationtech.jts.util.Assert;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Guice;
 import org.testng.annotations.Test;
@@ -212,6 +213,8 @@ public class AtlasComplexAttributesTest extends AtlasEntityTestBase {
         createdSimpleEntity.setAttribute("stringAtrr", null);
         createdSimpleEntity.setAttribute("mapOfStrings", Collections.emptyMap());
         createdSimpleEntity.setAttribute("arrayOfStrings", Collections.emptyList());
+        createdSimpleEntity.setAttribute("puArray", Collections.emptyList());
+        createdSimpleEntity.setAttribute("puMap", Collections.emptyMap());
         EntityMutationResponse responseUpdated      = entityStore.createOrUpdate(new AtlasEntityStream(createdSimpleEntity), false);
         AtlasEntityHeader simpleEntityUpdatedHeader = responseUpdated.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_SIMPLE_ATTR);
         AtlasEntity updatedSimpleEntity             = getEntityFromStore(simpleEntityUpdatedHeader);
@@ -219,10 +222,14 @@ public class AtlasComplexAttributesTest extends AtlasEntityTestBase {
         assertNull(updatedSimpleEntity.getAttribute("stringAtrr"));
         assertEquals(updatedSimpleEntity.getAttribute("mapOfStrings"), Collections.emptyMap());
         assertEquals(updatedSimpleEntity.getAttribute("arrayOfStrings"), Collections.emptyList());
+        assertEquals(updatedSimpleEntity.getAttribute("puArray"), Collections.emptyList());
+        assertEquals(updatedSimpleEntity.getAttribute("puMap"), Collections.emptyMap());
 
         updatedSimpleEntity.setAttribute("stringAtrr", "");
         updatedSimpleEntity.setAttribute("mapOfStrings", null);
         updatedSimpleEntity.setAttribute("arrayOfStrings", null);
+        updatedSimpleEntity.setAttribute("puArray", null);
+        updatedSimpleEntity.setAttribute("puMap", null);
         EntityMutationResponse responseUpdatedAgain      = entityStore.createOrUpdate(new AtlasEntityStream(updatedSimpleEntity), false);
         AtlasEntityHeader simpleEntityUpdatedAgainHeader = responseUpdatedAgain.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_SIMPLE_ATTR);
         AtlasEntity updatedAgainSimpleEntity             = getEntityFromStore(simpleEntityUpdatedAgainHeader);
@@ -230,10 +237,14 @@ public class AtlasComplexAttributesTest extends AtlasEntityTestBase {
         assertEquals(updatedAgainSimpleEntity.getAttribute("stringAtrr"), "");
         assertNull(updatedAgainSimpleEntity.getAttribute("arrayOfStrings"));
         assertNull(updatedAgainSimpleEntity.getAttribute("mapOfStrings"));
+        assertNull(updatedAgainSimpleEntity.getAttribute("puArray"));
+        assertNull(updatedAgainSimpleEntity.getAttribute("puMap"));
 
         updatedAgainSimpleEntity.setAttribute("stringAtrr", "Dummy String Test 3");
         updatedAgainSimpleEntity.setAttribute("mapOfStrings", Collections.singletonMap("key1", "val1"));
         updatedAgainSimpleEntity.setAttribute("arrayOfStrings", Arrays.asList("DummyTest3", "DummyTest4"));
+        updatedAgainSimpleEntity.setAttribute("puArray", Arrays.asList("1"));
+        updatedAgainSimpleEntity.setAttribute("puMap", Collections.singletonMap("1", "1"));
         EntityMutationResponse updateRes   = entityStore.createOrUpdate(new AtlasEntityStream(updatedAgainSimpleEntity), false);
         AtlasEntityHeader updateHeader = updateRes.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_SIMPLE_ATTR);
         AtlasEntity updateEntity = getEntityFromStore(updateHeader);
@@ -241,6 +252,94 @@ public class AtlasComplexAttributesTest extends AtlasEntityTestBase {
         assertEquals(updateEntity.getAttribute("stringAtrr"), "Dummy String Test 3");
         assertEquals(updateEntity.getAttribute("arrayOfStrings"), Arrays.asList("DummyTest3", "DummyTest4"));
         assertEquals(updateEntity.getAttribute("mapOfStrings"), Collections.singletonMap("key1", "val1"));
+        assertEquals(updateEntity.getAttribute("puArray"), Arrays.asList("1"));
+        assertEquals(updateEntity.getAttribute("puMap"), Collections.singletonMap("1", "1"));
+
+        // full-update puArray and puMap; existing values should be replaced
+        updatedAgainSimpleEntity.setAttribute("stringAtrr", "Dummy String Test 3");
+        updatedAgainSimpleEntity.setAttribute("mapOfStrings", Collections.singletonMap("key1", "val1"));
+        updatedAgainSimpleEntity.setAttribute("arrayOfStrings", Arrays.asList("DummyTest3", "DummyTest4"));
+        updatedAgainSimpleEntity.setAttribute("puArray", Arrays.asList("10"));
+        updatedAgainSimpleEntity.setAttribute("puMap", Collections.singletonMap("10", "10"));
+        updateRes    = entityStore.createOrUpdate(new AtlasEntityStream(updatedAgainSimpleEntity), false);
+        updateHeader = updateRes.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_SIMPLE_ATTR);
+        updateEntity = getEntityFromStore(updateHeader);
+
+        assertEquals(updateEntity.getAttribute("stringAtrr"), "Dummy String Test 3");
+        assertEquals(updateEntity.getAttribute("arrayOfStrings"), Arrays.asList("DummyTest3", "DummyTest4"));
+        assertEquals(updateEntity.getAttribute("mapOfStrings"), Collections.singletonMap("key1", "val1"));
+        assertEquals(updateEntity.getAttribute("puArray"), Arrays.asList("10"));
+        assertEquals(updateEntity.getAttribute("puMap"), Collections.singletonMap("10", "10"));
+
+        // partial-update tests
+        // set puArray and puMap to null
+        updatedAgainSimpleEntity.setAttribute("stringAtrr", "Dummy String Test 3");
+        updatedAgainSimpleEntity.setAttribute("mapOfStrings", Collections.singletonMap("key1", "val1"));
+        updatedAgainSimpleEntity.setAttribute("arrayOfStrings", Arrays.asList("DummyTest3", "DummyTest4"));
+        updatedAgainSimpleEntity.setAttribute("puArray", null);
+        updatedAgainSimpleEntity.setAttribute("puMap", null);
+        updateRes    = entityStore.createOrUpdate(new AtlasEntityStream(updatedAgainSimpleEntity), false);
+        updateHeader = updateRes.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_SIMPLE_ATTR);
+        updateEntity = getEntityFromStore(updateHeader);
+
+        assertEquals(updateEntity.getAttribute("stringAtrr"), "Dummy String Test 3");
+        assertEquals(updateEntity.getAttribute("arrayOfStrings"), Arrays.asList("DummyTest3", "DummyTest4"));
+        assertEquals(updateEntity.getAttribute("mapOfStrings"), Collections.singletonMap("key1", "val1"));
+        assertNull(updateEntity.getAttribute("puArray"));
+        assertNull(updateEntity.getAttribute("puMap"));
+
+        List<String>        puArray = new ArrayList<>();
+        Map<String, String> puMap   = new HashMap<>();
+
+        // partial-update: current value as null
+        updatedAgainSimpleEntity.getAttributes().clear();
+        updatedAgainSimpleEntity.setAttribute("puArray", Collections.singletonList("1"));
+        updatedAgainSimpleEntity.setAttribute("puMap", Collections.singletonMap("1", "1"));
+        updateRes    = entityStore.createOrUpdate(new AtlasEntityStream(updatedAgainSimpleEntity), true);
+        updateHeader = updateRes.getFirstPartialUpdatedEntityByTypeName(ENTITY_TYPE_WITH_SIMPLE_ATTR);
+        updateEntity = getEntityFromStore(updateHeader);
+
+        puArray.addAll(Collections.singletonList("1"));
+        puMap.putAll(Collections.singletonMap("1", "1"));
+
+        Assert.equals(updateEntity.getAttribute("puArray"), puArray);
+        Assert.equals(updateEntity.getAttribute("puMap"), puMap);
+
+        // partial-update: append to existing value
+        updatedAgainSimpleEntity.getAttributes().clear();
+        updatedAgainSimpleEntity.setAttribute("puArray", Collections.singletonList("2"));
+        updatedAgainSimpleEntity.setAttribute("puMap", Collections.singletonMap("2", "2"));
+        updateRes    = entityStore.createOrUpdate(new AtlasEntityStream(updatedAgainSimpleEntity), true);
+        updateHeader = updateRes.getFirstPartialUpdatedEntityByTypeName(ENTITY_TYPE_WITH_SIMPLE_ATTR);
+        updateEntity = getEntityFromStore(updateHeader);
+
+        puArray.addAll(Collections.singletonList("2"));
+        puMap.putAll(Collections.singletonMap("2", "2"));
+
+        Assert.equals(updateEntity.getAttribute("puArray"), puArray);
+        Assert.equals(updateEntity.getAttribute("puMap"), puMap);
+
+        //  partial-update: with null value; existing value should be retained
+        updatedAgainSimpleEntity.getAttributes().clear();
+        updatedAgainSimpleEntity.setAttribute("puArray", null);
+        updatedAgainSimpleEntity.setAttribute("puMap", null);
+        updateRes    = entityStore.createOrUpdate(new AtlasEntityStream(updatedAgainSimpleEntity), true);
+        updateHeader = updateRes.getFirstPartialUpdatedEntityByTypeName(ENTITY_TYPE_WITH_SIMPLE_ATTR);
+        updateEntity = getEntityFromStore(updateHeader);
+
+        Assert.equals(updateEntity.getAttribute("puArray"), puArray);
+        Assert.equals(updateEntity.getAttribute("puMap"), puMap);
+
+        //  partial-update: with empty value; existing value should be retained
+        updatedAgainSimpleEntity.getAttributes().clear();
+        updatedAgainSimpleEntity.setAttribute("puArray", Collections.emptyList());
+        updatedAgainSimpleEntity.setAttribute("puMap", Collections.emptyMap());
+        updateRes    = entityStore.createOrUpdate(new AtlasEntityStream(updatedAgainSimpleEntity), true);
+        updateHeader = updateRes.getFirstPartialUpdatedEntityByTypeName(ENTITY_TYPE_WITH_SIMPLE_ATTR);
+        updateEntity = getEntityFromStore(updateHeader);
+
+        Assert.equals(updateEntity.getAttribute("puArray"), puArray);
+        Assert.equals(updateEntity.getAttribute("puMap"), puMap);
     }
 
     @Test(dependsOnMethods = "testCreateComplexAttributeEntity")
@@ -570,4 +669,4 @@ public class AtlasComplexAttributesTest extends AtlasEntityTestBase {
             assertEquals(entity.getStatus(), AtlasEntity.Status.DELETED);
         }
     }
-}
+}
\ No newline at end of file