ATLAS-2610: Relationship id during import. Tag propgation set to false. Inverse…

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

@@ -76,7 +76,6 @@ import static org.apache.atlas.model.instance.EntityMutations.EntityOperation.PA
 import static org.apache.atlas.model.instance.EntityMutations.EntityOperation.UPDATE;
 import static org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef.Cardinality.SET;
 import static org.apache.atlas.repository.Constants.CLASSIFICATION_EDGE_STATE_PROPERTY_KEY;
-import static org.apache.atlas.repository.Constants.CLASSIFICATION_ENTITY_GUID;
 import static org.apache.atlas.repository.Constants.CLASSIFICATION_LABEL;
 import static org.apache.atlas.repository.Constants.STATE_PROPERTY_KEY;
 import static org.apache.atlas.repository.Constants.TRAIT_NAMES_PROPERTY_KEY;
@@ -426,7 +425,7 @@ public class EntityGraphMapper {
                     // legacy case update inverse attribute
                     if (ctx.getAttribute().getInverseRefAttribute() != null) {
                         // Update the inverse reference using relationship on the target entity
-                        addInverseReference(ctx.getAttribute().getInverseRefAttribute(), newEdge, getRelationshipAttributes(ctx.getValue()));
+                        addInverseReference(context, ctx.getAttribute().getInverseRefAttribute(), newEdge, getRelationshipAttributes(ctx.getValue()));
                     }
                 }
 
@@ -473,7 +472,7 @@ public class EntityGraphMapper {
         }
     }
 
-    private void addInverseReference(AtlasAttribute inverseAttribute, AtlasEdge edge, Map<String, Object> relationshipAttributes) throws AtlasBaseException {
+    private void addInverseReference(EntityMutationContext context, AtlasAttribute inverseAttribute, AtlasEdge edge, Map<String, Object> relationshipAttributes) throws AtlasBaseException {
         AtlasStructType inverseType      = inverseAttribute.getDefinedInType();
         AtlasVertex     inverseVertex    = edge.getInVertex();
         String          inverseEdgeLabel = inverseAttribute.getRelationshipEdgeLabel();
@@ -481,7 +480,7 @@ public class EntityGraphMapper {
         String          propertyName     = AtlasGraphUtilsV1.getQualifiedAttributePropertyKey(inverseType, inverseAttribute.getName());
 
         // create new inverse reference
-        AtlasEdge newEdge = createInverseReferenceUsingRelationship(inverseAttribute, edge, relationshipAttributes);
+        AtlasEdge newEdge = createInverseReferenceUsingRelationship(context, inverseAttribute, edge, relationshipAttributes);
 
         boolean inverseUpdated = true;
         switch (inverseAttribute.getAttributeType().getTypeCategory()) {
@@ -532,7 +531,7 @@ public class EntityGraphMapper {
         }
     }
 
-    private AtlasEdge createInverseReferenceUsingRelationship(AtlasAttribute inverseAttribute, AtlasEdge edge, Map<String, Object> relationshipAttributes) throws AtlasBaseException {
+    private AtlasEdge createInverseReferenceUsingRelationship(EntityMutationContext context, AtlasAttribute inverseAttribute, AtlasEdge edge, Map<String, Object> relationshipAttributes) throws AtlasBaseException {
         if (LOG.isDebugEnabled()) {
             LOG.debug("==> createInverseReferenceUsingRelationship()");
         }
@@ -568,9 +567,34 @@ public class EntityGraphMapper {
             LOG.debug("<== createInverseReferenceUsingRelationship()");
         }
 
+        updateRelationshipGuidForImport(context, inverseAttributeName, inverseVertex, ret);
+
         return ret;
     }
 
+    private void updateRelationshipGuidForImport(EntityMutationContext context, String inverseAttributeName, AtlasVertex inverseVertex, AtlasEdge edge) throws AtlasBaseException {
+        if (!context.isImport()) {
+            return;
+        }
+
+        String parentGuid = GraphHelper.getGuid(inverseVertex);
+        if(StringUtils.isEmpty(parentGuid)) {
+            return;
+        }
+
+        AtlasEntity entity = context.getCreatedOrUpdatedEntity(parentGuid);
+        if(entity == null) {
+            return;
+        }
+
+        String parentRelationshipGuid = getRelationshipGuid(entity.getRelationshipAttribute(inverseAttributeName));
+        if(StringUtils.isEmpty(parentRelationshipGuid)) {
+            return;
+        }
+
+        AtlasGraphUtilsV1.setProperty(edge, Constants.RELATIONSHIP_GUID_PROPERTY_KEY, parentRelationshipGuid);
+    }
+
     // legacy method to create edges for inverse reference
     private AtlasEdge createInverseReference(AtlasAttribute inverseAttribute, AtlasStructType inverseAttributeType,
                                              AtlasVertex inverseVertex, AtlasVertex vertex) throws AtlasBaseException {
@@ -737,7 +761,10 @@ public class EntityGraphMapper {
 
                     // for import use the relationship guid provided
                     if (context.isImport()) {
-                        AtlasGraphUtilsV1.setProperty(ret, Constants.RELATIONSHIP_GUID_PROPERTY_KEY, getRelationshipGuid(ctx.getValue()));
+                        String relationshipGuid = getRelationshipGuid(ctx.getValue());
+                        if(!StringUtils.isEmpty(relationshipGuid)) {
+                            AtlasGraphUtilsV1.setProperty(ret, Constants.RELATIONSHIP_GUID_PROPERTY_KEY, relationshipGuid);
+                        }
                     }
 
                     // if relationship did not exist before and new relationship was created
@@ -813,7 +840,7 @@ public class EntityGraphMapper {
                     if (isReference && newEntry instanceof AtlasEdge && inverseRefAttribute != null) {
                         AtlasEdge newEdge = (AtlasEdge) newEntry;
 
-                        addInverseReference(inverseRefAttribute, newEdge, getRelationshipAttributes(ctx.getValue()));
+                        addInverseReference(context, inverseRefAttribute, newEdge, getRelationshipAttributes(ctx.getValue()));
                     }
                 }
             }
@@ -877,7 +904,7 @@ public class EntityGraphMapper {
                     // Update the inverse reference value.
                     AtlasEdge newEdge = (AtlasEdge) newEntry;
 
-                    addInverseReference(inverseRefAttribute, newEdge, getRelationshipAttributes(ctx.getValue()));
+                    addInverseReference(context, inverseRefAttribute, newEdge, getRelationshipAttributes(ctx.getValue()));
                 }
 
                 if(newEntry != null) {
@@ -1331,7 +1358,12 @@ public class EntityGraphMapper {
                 Boolean propagateTags      = classification.isPropagate();
 
                 if (propagateTags == null) {
-                    propagateTags = true;
+                    if(context.isImport()) {
+                        propagateTags = false;
+                        classification.setPropagate(propagateTags);
+                    } else {
+                        propagateTags = true;
+                    }
                 }
 
                 // set associated entity id to classification

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

@@ -19,11 +19,9 @@ package org.apache.atlas.repository.store.graph.v1;
 
 import org.apache.atlas.model.instance.AtlasEntity;
 
-import org.apache.atlas.model.instance.AtlasObjectId;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.atlas.repository.store.graph.EntityGraphDiscoveryContext;
 import org.apache.atlas.type.AtlasEntityType;
-import org.apache.atlas.type.AtlasType;
 import org.apache.commons.lang.StringUtils;
 
 import java.util.*;
@@ -127,4 +125,31 @@ public class EntityMutationContext {
             ", entityVsVertex=" + entityVsVertex +
             '}';
     }
+
+    public AtlasEntity getCreatedEntity(String parentGuid) {
+        return getFromCollection(parentGuid, getCreatedEntities());
+    }
+
+    public AtlasEntity getUpdatedEntity(String parentGuid) {
+        return getFromCollection(parentGuid, getUpdatedEntities());
+    }
+
+    private AtlasEntity getFromCollection(String parentGuid, Collection<AtlasEntity> coll) {
+        for (AtlasEntity e : coll) {
+            if(e.getGuid().equalsIgnoreCase(parentGuid)) {
+                return e;
+            }
+        }
+
+        return null;
+    }
+
+    public AtlasEntity getCreatedOrUpdatedEntity(String parentGuid) {
+        AtlasEntity e = getCreatedEntity(parentGuid);
+        if(e == null) {
+            return getUpdatedEntity(parentGuid);
+        }
+
+        return e;
+    }
 }

repository/src/test/java/org/apache/atlas/repository/impexp/ImportServiceTest.java

@@ -24,9 +24,13 @@ import org.apache.atlas.TestModules;
 import org.apache.atlas.TestUtilsV2;
 import org.apache.atlas.discovery.EntityDiscoveryService;
 import org.apache.atlas.exception.AtlasBaseException;
+import org.apache.atlas.model.discovery.AtlasSearchResult;
 import org.apache.atlas.model.impexp.AtlasImportRequest;
+import org.apache.atlas.model.instance.AtlasEntity;
 import org.apache.atlas.model.instance.AtlasEntityHeader;
+import org.apache.atlas.model.instance.AtlasRelatedObjectId;
 import org.apache.atlas.repository.graph.AtlasGraphProvider;
+import org.apache.atlas.repository.store.graph.AtlasEntityStore;
 import org.apache.atlas.runner.LocalSolrRunner;
 import org.apache.atlas.store.AtlasTypeDefStore;
 import org.apache.atlas.type.AtlasClassificationType;
@@ -49,10 +53,12 @@ import java.util.List;
 import java.util.Map;
 
 import static org.apache.atlas.graph.GraphSandboxUtil.useLocalSolr;
+import static org.apache.atlas.repository.Constants.RELATIONSHIP_GUID_PROPERTY_KEY;
 import static org.apache.atlas.repository.impexp.ZipFileResourceTestUtils.*;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
 import static org.testng.Assert.assertNotNull;
 
 @Guice(modules = TestModules.TestOnlyModule.class)
@@ -71,6 +77,9 @@ public class ImportServiceTest {
     private EntityDiscoveryService discoveryService;
 
     @Inject
+    AtlasEntityStore entityStore;
+
+    @Inject
     public ImportServiceTest(ImportService importService) {
         this.importService = importService;
     }
@@ -99,6 +108,19 @@ public class ImportServiceTest {
     public void importDB1(ZipSource zipSource) throws AtlasBaseException, IOException {
         loadBaseModel();
         runAndVerifyQuickStart_v1_Import(importService, zipSource);
+
+        assertEntityCount("DB_v1", "bfe88eb8-7556-403c-8210-647013f44a44", 1);
+
+
+        List<AtlasEntityHeader> entityHeader = getEntitiesFromDB("Table_v1", "fe91bf93-eb0c-4638-8361-15937390c810");
+        assertEquals(entityHeader.size(), 1);
+
+        AtlasEntity.AtlasEntityWithExtInfo entityWithExtInfo = getEntity(entityHeader.get(0));
+        assertNotNull(entityWithExtInfo);
+
+        AtlasEntity entity = entityWithExtInfo.getEntity();
+        assertEquals(entity.getClassifications().size(), 1);
+        assertFalse(entity.getClassifications().get(0).isPropagate(), "Default propagate should be false");
     }
 
     @DataProvider(name = "reporting")
@@ -184,6 +206,31 @@ public class ImportServiceTest {
         runImportWithNoParameters(importService, zipSource);
     }
 
+    @DataProvider(name = "stocks-legacy")
+    public static Object[][] getDataFromLegacyStocks(ITestContext context) throws IOException {
+        return getZipSource("stocks.zip");
+    }
+
+    @Test(dataProvider = "stocks-legacy")
+    public void importLegacy(ZipSource zipSource) throws IOException, AtlasBaseException {
+        loadBaseModel();
+        loadFsModel();
+        loadHiveModel();
+
+        runImportWithNoParameters(importService, zipSource);
+        List<AtlasEntityHeader> result = getEntitiesFromDB("hive_db", "886c5e9c-3ac6-40be-8201-fb0cebb64783");
+        assertEquals(result.size(), 1);
+
+        AtlasEntity.AtlasEntityWithExtInfo entityWithExtInfo = getEntity(result.get(0));
+        Map<String, Object> relationshipAttributes = entityWithExtInfo.getEntity().getRelationshipAttributes();
+        assertNotNull(relationshipAttributes);
+        assertNotNull(relationshipAttributes.get("tables"));
+
+        List<AtlasRelatedObjectId> relatedList = (List<AtlasRelatedObjectId>) relationshipAttributes.get("tables");
+        AtlasRelatedObjectId relatedObjectId = relatedList.get(0);
+        assertNotNull(relatedObjectId.getRelationshipGuid());
+    }
+
     @DataProvider(name = "stocks-glossary")
     public static Object[][] getDataFromGlossary(ITestContext context) throws IOException {
         return getZipSource("stocks-glossary.zip");
@@ -303,4 +350,8 @@ public class ImportServiceTest {
     private void loadGlossary() throws IOException, AtlasBaseException {
         loadModelFromJson("0000-Area0/0011-glossary_model.json", typeDefStore, typeRegistry);
     }
+
+    private AtlasEntity.AtlasEntityWithExtInfo getEntity(AtlasEntityHeader header) throws AtlasBaseException {
+        return entityStore.getById(header.getGuid());
+    }
 }