Repository: atlas
Updated Branches:
  refs/heads/master 65f7da17e -> 30dd2f5b8


ATLAS-2610: Relationship id during import. Tag propgation set to false. Inverse 
relationship mapped correctly.

Signed-off-by: Ashutosh Mestry <[email protected]>


Project: http://git-wip-us.apache.org/repos/asf/atlas/repo
Commit: http://git-wip-us.apache.org/repos/asf/atlas/commit/30dd2f5b
Tree: http://git-wip-us.apache.org/repos/asf/atlas/tree/30dd2f5b
Diff: http://git-wip-us.apache.org/repos/asf/atlas/diff/30dd2f5b

Branch: refs/heads/master
Commit: 30dd2f5b80cab252a217d486558d9384ba3bfeb6
Parents: 65f7da1
Author: Ashutosh Mestry <[email protected]>
Authored: Wed Apr 25 10:47:38 2018 -0700
Committer: Ashutosh Mestry <[email protected]>
Committed: Wed Apr 25 16:43:54 2018 -0700

----------------------------------------------------------------------
 .../store/graph/v1/EntityGraphMapper.java       | 50 +++++++++++++++----
 .../store/graph/v1/EntityMutationContext.java   | 29 ++++++++++-
 .../repository/impexp/ImportServiceTest.java    | 51 ++++++++++++++++++++
 3 files changed, 119 insertions(+), 11 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/atlas/blob/30dd2f5b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityGraphMapper.java
----------------------------------------------------------------------
diff --git 
a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityGraphMapper.java
 
b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityGraphMapper.java
index 4225a80..43e67df 100644
--- 
a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityGraphMapper.java
+++ 
b/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

http://git-wip-us.apache.org/repos/asf/atlas/blob/30dd2f5b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityMutationContext.java
----------------------------------------------------------------------
diff --git 
a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityMutationContext.java
 
b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityMutationContext.java
index 0a55d55..e3f7121 100644
--- 
a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityMutationContext.java
+++ 
b/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;
+    }
 }

http://git-wip-us.apache.org/repos/asf/atlas/blob/30dd2f5b/repository/src/test/java/org/apache/atlas/repository/impexp/ImportServiceTest.java
----------------------------------------------------------------------
diff --git 
a/repository/src/test/java/org/apache/atlas/repository/impexp/ImportServiceTest.java
 
b/repository/src/test/java/org/apache/atlas/repository/impexp/ImportServiceTest.java
index 5089070..68b57ed 100644
--- 
a/repository/src/test/java/org/apache/atlas/repository/impexp/ImportServiceTest.java
+++ 
b/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());
+    }
 }

Reply via email to