ATLAS-1584 Fix issues with owned map reference and add tests (sumasai)

intg/src/main/java/org/apache/atlas/type/AtlasStructType.java

@@ -559,6 +559,10 @@ public class AtlasStructType extends AtlasType {
             type = ((AtlasArrayType)type).getElementType();
         }
 
+        if (type instanceof AtlasMapType) {
+            type = ((AtlasMapType)type).getValueType();
+        }
+
         return type instanceof AtlasEntityType ? (AtlasEntityType)type : null;
     }
 

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

@@ -713,7 +713,10 @@ public final class TestUtilsV2 {
                                 true,
                                 AtlasAttributeDef.Cardinality.SINGLE, 0, 1,
                                 false, false,
-                                Collections.<AtlasStructDef.AtlasConstraintDef>emptyList()
+                                 new ArrayList<AtlasStructDef.AtlasConstraintDef>() {{
+                                     add(new AtlasStructDef.AtlasConstraintDef(
+                                             AtlasConstraintDef.CONSTRAINT_TYPE_OWNED_REF));
+                                     }}
                                 ),
                         //map of structs
                         new AtlasAttributeDef("partitionsMap",

release-log.txt

@@ -9,7 +9,8 @@ ATLAS-1060 Add composite indexes for exact match performance improvements for al
 ATLAS-1127 Modify creation and modification timestamps to Date instead of Long(sumasai)
 
 ALL CHANGES:
-ATLAS_1589 DSL queries return wrong object when filter traverses an edge (jnhagelberg)
+ATLAS-1584 Fix issues with owned map reference and add tests (sumasai)
+ATLAS-1589 DSL queries return wrong object when filter traverses an edge (jnhagelberg)
 ATLAS-1590 UI : Edit Button is enabled for Deleted entities. (kevalbhatt)
 ATLAS-1487 Create Entity in UI : types list doesn't list fs_permissions (struct type) hence no entity could be created for it. (kevalbhatt)
 ATLAS-1573 Full text mapping for Entity store V2

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasGraphUtilsV1.java

@@ -225,8 +225,7 @@ public class AtlasGraphUtilsV1 {
 
     public static AtlasVertex findByGuid(String guid) {
         AtlasGraphQuery query = AtlasGraphProvider.getGraphInstance().query()
-                                                  .has(Constants.GUID_PROPERTY_KEY, guid)
-                                                  .has(Constants.STATE_PROPERTY_KEY, AtlasEntity.Status.ACTIVE.name());
+                                                  .has(Constants.GUID_PROPERTY_KEY, guid);
 
         Iterator<AtlasVertex> results = query.vertices().iterator();
 

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

@@ -59,6 +59,7 @@ import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -68,6 +69,7 @@ import static org.apache.atlas.model.instance.EntityMutations.EntityOperation.CR
 import static org.apache.atlas.model.instance.EntityMutations.EntityOperation.PARTIAL_UPDATE;
 import static org.apache.atlas.model.instance.EntityMutations.EntityOperation.UPDATE;
 import static org.apache.atlas.model.instance.EntityMutations.EntityOperation.DELETE;
+import static org.apache.atlas.repository.Constants.STATE_PROPERTY_KEY;
 import static org.apache.atlas.repository.graph.GraphHelper.string;
 
 
@@ -437,7 +439,11 @@ public class EntityGraphMapper {
 
             Map<String, Object> finalMap = removeUnusedMapEntries(attribute, ctx.getReferringVertex(), ctx.getVertexProperty(), currentMap, newMap);
 
-            Set<String> newKeys = new HashSet<>(newMap.keySet());
+            for (Object newEntry : newMap.values()) {
+                updateInConsistentOwnedMapVertices(ctx, mapType, newEntry);
+            }
+
+            Set<String> newKeys = new LinkedHashSet<>(newMap.keySet());
             newKeys.addAll(finalMap.keySet());
 
             // for dereference on way out
@@ -534,21 +540,21 @@ public class EntityGraphMapper {
 
     private Object mapCollectionElementsToVertex(AttributeMutationContext ctx, EntityMutationContext context) throws AtlasBaseException {
         switch(ctx.getAttrType().getTypeCategory()) {
-            case PRIMITIVE:
-            case ENUM:
-                return ctx.getValue();
+        case PRIMITIVE:
+        case ENUM:
+            return ctx.getValue();
 
-            case STRUCT:
-                return mapStructValue(ctx, context);
+        case STRUCT:
+            return mapStructValue(ctx, context);
 
-            case OBJECT_ID_TYPE:
-                AtlasEntityType instanceType = getInstanceType(ctx.getValue());
-                ctx.setElementType(instanceType);
-                return mapObjectIdValue(ctx, context);
+        case OBJECT_ID_TYPE:
+            AtlasEntityType instanceType = getInstanceType(ctx.getValue());
+            ctx.setElementType(instanceType);
+            return mapObjectIdValue(ctx, context);
 
-            case MAP:
-            case ARRAY:
-            default:
+        case MAP:
+        case ARRAY:
+        default:
                 throw new AtlasBaseException(AtlasErrorCode.TYPE_CATEGORY_INVALID, ctx.getAttrType().getTypeCategory().name());
         }
     }
@@ -776,6 +782,19 @@ public class EntityGraphMapper {
         return new AtlasEntityHeader(id.getTypeName(), id.getGuid(), id.getUniqueAttributes());
     }
 
+    private void updateInConsistentOwnedMapVertices(AttributeMutationContext ctx, AtlasMapType mapType, Object val) {
+        if (mapType.getValueType().getTypeCategory() == TypeCategory.OBJECT_ID_TYPE) {
+            AtlasEdge edge = (AtlasEdge) val;
+            if (ctx.getAttribute().isOwnedRef() &&
+                GraphHelper.getStatus(edge) == AtlasEntity.Status.DELETED &&
+                GraphHelper.getStatus(edge.getInVertex()) == AtlasEntity.Status.DELETED) {
+                //Resurrect the vertex and edge to ACTIVE state
+                GraphHelper.setProperty(edge, STATE_PROPERTY_KEY, AtlasEntity.Status.ACTIVE.name());
+                GraphHelper.setProperty(edge.getInVertex(), STATE_PROPERTY_KEY, AtlasEntity.Status.ACTIVE.name());
+            }
+        }
+    }
+
     public void addClassifications(final EntityMutationContext context, String guid, List<AtlasClassification> classifications)
         throws AtlasBaseException {
 

repository/src/test/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityStoreV1Test.java

@@ -261,7 +261,6 @@ public class AtlasEntityStoreV1Test {
         partsMap.put("part0", new AtlasStruct(TestUtils.PARTITION_STRUCT_TYPE, TestUtilsV2.NAME, "test"));
 
         tableEntity.setAttribute("partitionsMap", partsMap);
-        entitiesInfo.addReferredEntity(tableEntity);
 
         init();
         EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo), false);
@@ -365,10 +364,14 @@ public class AtlasEntityStoreV1Test {
         AtlasEntityHeader tableDefinition6 = response.getFirstUpdatedEntityByTypeName(TABLE_TYPE);
         validateEntity(entitiesInfo, getEntityFromStore(tableDefinition6));
 
+        Assert.assertEquals(entityStore.getById(col0.getGuid()).getEntity().getStatus(), AtlasEntity.Status.ACTIVE);
+        Assert.assertEquals(entityStore.getById(col1.getGuid()).getEntity().getStatus(), AtlasEntity.Status.ACTIVE);
+
         //Drop the first key and change the class type as well to col0
         columnsMap.clear();
         columnsMap.put("col0", AtlasTypeUtil.getAtlasObjectId(col0));
         init();
+
         response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo), false);
         AtlasEntityHeader tableDefinition7 = response.getFirstUpdatedEntityByTypeName(TABLE_TYPE);
         validateEntity(entitiesInfo, getEntityFromStore(tableDefinition7));
@@ -494,25 +497,6 @@ public class AtlasEntityStoreV1Test {
         validateEntity(entitiesInfo, getEntityFromStore(updatedTable));
     }
 
-//    private AtlasEntity clearSubOrdinates(List<AtlasObjectId> employees, int index) {
-//
-//        AtlasEntity ret = null;
-//        AtlasObjectId employee = employees.get(index);
-//        AtlasEntity subOrdClone = new ArrayList<>(subOrdinates);
-//        ret = subOrdClone.remove(index);
-//
-//        employees.get(index).setAttribute("subordinates", subOrdClone);
-//        return ret;
-//    }
-//
-//    private int addSubordinate(AtlasEntity manager, AtlasEntity employee) {
-//        List<AtlasEntity> subOrdinates = (List<AtlasEntity>) manager.getAttribute("subordinates");
-//        subOrdinates.add(employee);
-//
-//        manager.setAttribute("subordinates", subOrdinates);
-//        return subOrdinates.size() - 1;
-//    }
-
     @Test(dependsOnMethods = "testCreate")
     public void testClassUpdate() throws Exception {
 

repository/src/test/java/org/apache/atlas/repository/store/graph/v1/HardDeleteHandlerV1Test.java

@@ -26,6 +26,7 @@ import org.apache.atlas.model.instance.AtlasEntity;
 import org.apache.atlas.model.instance.AtlasEntityHeader;
 import org.apache.atlas.model.instance.AtlasObjectId;
 import org.apache.atlas.repository.Constants;
+import org.apache.atlas.repository.graph.GraphHelper;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.atlas.type.AtlasTypeRegistry;
 import org.apache.atlas.typesystem.IStruct;
@@ -146,6 +147,23 @@ public class HardDeleteHandlerV1Test extends AtlasDeleteHandlerV1Test {
         }
     }
 
+    protected void assertTestDisconnectMapReferenceFromClassType(final String mapOwnerGuid) throws Exception {
+        // Verify map references from mapOwner were disconnected.
+        AtlasEntity.AtlasEntityWithExtInfo mapOwnerInstance = entityStore.getById(mapOwnerGuid);
+        Map<String, AtlasObjectId> map =
+            (Map<String, AtlasObjectId>) mapOwnerInstance.getEntity().getAttribute("map");
+        Assert.assertNull(map);
+        Map<String, AtlasObjectId> biMap =
+            (Map<String, AtlasObjectId>) mapOwnerInstance.getEntity().getAttribute("biMap");
+        Assert.assertNull(biMap);
+
+        AtlasVertex mapOwnerVertex = GraphHelper.getInstance().getVertexForGUID(mapOwnerGuid);
+        Object object = mapOwnerVertex.getProperty("MapOwner.map.value1", String.class);
+        assertNull(object);
+        object = mapOwnerVertex.getProperty("MapOwner.biMap.value1", String.class);
+        assertNull(object);
+    }
+
     @Override
     protected void assertTestDisconnectUnidirectionalArrayReferenceFromStructAndTraitTypes(String structContainerGuid)
         throws Exception {

repository/src/test/java/org/apache/atlas/repository/store/graph/v1/SoftDeleteHandlerV1Test.java

@@ -183,6 +183,19 @@ public class SoftDeleteHandlerV1Test extends AtlasDeleteHandlerV1Test {
     }
 
     @Override
+    protected void assertTestDisconnectMapReferenceFromClassType(final String mapOwnerGuid) throws Exception {
+        AtlasEntity.AtlasEntityWithExtInfo mapOwnerInstance = entityStore.getById(mapOwnerGuid);
+        Map<String, AtlasObjectId> map =
+            (Map<String, AtlasObjectId>) mapOwnerInstance.getEntity().getAttribute("map");
+        assertNotNull(map);
+        assertEquals(map.size(), 1);
+        Map<String, AtlasObjectId> biMap =
+            (Map<String, AtlasObjectId>) mapOwnerInstance.getEntity().getAttribute("biMap");
+        assertNotNull(biMap);
+        assertEquals(biMap.size(), 1);
+    }
+
+    @Override
     protected void assertTestDisconnectUnidirectionalArrayReferenceFromStructAndTraitTypes(final String structContainerGuid) throws Exception {
         // Verify that the unidirectional references from the struct and trait instances
         // to the deleted entities were not disconnected.