ATLAS-2372: Export-Import: Support for Relationship Attributes Missing

repository/src/main/java/org/apache/atlas/repository/impexp/ExportService.java

@@ -442,7 +442,6 @@ public class ExportService {
             return;
         }
 
-        removeRelationshipAttributes(entity);
         context.sink.add(entity);
 
         context.result.incrementMeticsCounter(String.format("entity:%s", entity.getEntity().getTypeName()));
@@ -456,19 +455,6 @@ public class ExportService {
         context.reportProgress();
     }
 
-    private void removeRelationshipAttributes(AtlasEntityWithExtInfo entity) {
-        if (entity.getEntity().getRelationshipAttributes() != null) {
-            entity.getEntity().getRelationshipAttributes().clear();
-        }
-        if (entity.getReferredEntities() != null) {
-            for (AtlasEntity e : entity.getReferredEntities().values()) {
-                if (e.getRelationshipAttributes() != null) {
-                    e.getRelationshipAttributes().clear();
-                }
-            }
-        }
-    }
-
     private void addTypes(AtlasEntity entity, ExportContext context) {
         addEntityType(entity.getTypeName(), context);
 

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

@@ -305,16 +305,8 @@ public class AtlasEntityGraphDiscoveryV1 implements EntityGraphDiscovery {
         List<String> visitedAttributes = new ArrayList<>();
 
         // visit relationship attributes
-        for (AtlasAttribute attribute : entityType.getRelationshipAttributes().values()) {
-            AtlasType attrType = attribute.getAttributeType();
-            String    attrName = attribute.getName();
-            Object    attrVal  = entity.getRelationshipAttribute(attrName);
-
-            if (entity.hasRelationshipAttribute(attrName)) {
-                visitAttribute(attrType, attrVal);
-
-                visitedAttributes.add(attrName);
-            }
+        if(!(this.discoveryContext.getEntityStream() instanceof EntityImportStream)) {
+            visitRelationships(entityType, entity, visitedAttributes);
         }
 
         // visit struct attributes
@@ -329,6 +321,20 @@ public class AtlasEntityGraphDiscoveryV1 implements EntityGraphDiscovery {
         }
     }
 
+    private void visitRelationships(AtlasEntityType entityType, AtlasEntity entity, List<String> visitedAttributes) throws AtlasBaseException {
+        for (AtlasAttribute attribute : entityType.getRelationshipAttributes().values()) {
+            AtlasType attrType = attribute.getAttributeType();
+            String attrName = attribute.getName();
+            Object attrVal = entity.getRelationshipAttribute(attrName);
+
+            if (entity.hasRelationshipAttribute(attrName)) {
+                visitAttribute(attrType, attrVal);
+
+                visitedAttributes.add(attrName);
+            }
+        }
+    }
+
     void visitStruct(AtlasStructType structType, AtlasStruct struct) throws AtlasBaseException {
         for (AtlasAttribute attribute : structType.getAllAttributes().values()) {
             AtlasType attrType = attribute.getAttributeType();

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

@@ -377,7 +377,7 @@ public class EntityGraphMapper {
                 AtlasEdge currentEdge;
 
                 // if relationshipGuid is assigned in AtlasRelatedObjectId use it to fetch existing AtlasEdge
-                if (StringUtils.isNotEmpty(relationshipGuid)) {
+                if (StringUtils.isNotEmpty(relationshipGuid) && !context.isImport()) {
                     currentEdge = graphHelper.getEdgeForGUID(relationshipGuid);
                 } else {
                     currentEdge = graphHelper.getEdgeForLabel(ctx.getReferringVertex(), edgeLabel, edgeDirection);
@@ -699,6 +699,11 @@ public class EntityGraphMapper {
 
                     ret = getOrCreateRelationship(fromVertex, toVertex, relationshipName, relationshipAttributes);
 
+                    // for import use the relationship guid provided
+                    if (context.isImport()) {
+                        AtlasGraphUtilsV1.setProperty(ret, Constants.GUID_PROPERTY_KEY, getRelationshipGuid(ctx.getValue()));
+                    }
+
                     // if relationship did not exist before and new relationship was created
                     // record entity update on both relationship vertices
                     if (!relationshipExists) {

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

@@ -87,6 +87,9 @@ public class EntityMutationContext {
 
     public AtlasVertex getVertex(String guid) { return entityVsVertex.get(guid); }
 
+    public boolean isImport() {
+        return (context != null) && context.getEntityStream() instanceof EntityImportStream;
+    }
 
     @Override
     public boolean equals(final Object o) {

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

@@ -135,7 +135,6 @@ public class ImportServiceTest {
         return getZipSource("salesNewTypeAttrs-next.zip");
     }
 
-
     @Test(dataProvider = "salesNewTypeAttrs-next", dependsOnMethods = "importDB4")
     public void importDB5(ZipSource zipSource) throws AtlasBaseException, IOException {
         final String newEnumDefName = "database_action";
@@ -186,7 +185,6 @@ public class ImportServiceTest {
         return getZipSource("hdfs_path1.zip");
     }
 
-
     @Test(dataProvider = "hdfs_path1", expectedExceptions = AtlasBaseException.class)
     public void importHdfs_path1(ZipSource zipSource) throws IOException, AtlasBaseException {
         loadBaseModel();
@@ -204,6 +202,19 @@ public class ImportServiceTest {
         }
     }
 
+    @DataProvider(name = "relationship")
+    public static Object[][] getImportWithRelationships(ITestContext context) throws IOException {
+        return getZipSource("stocks-rel-2.zip");
+    }
+
+    @Test(dataProvider = "relationship")
+    public void importDB7(ZipSource zipSource) throws AtlasBaseException, IOException {
+        loadBaseModel();
+        loadHiveModel();
+        AtlasImportRequest request = getDefaultImportRequest();
+        runImportWithParameters(importService, request, zipSource);
+    }
+
     @Test
     public void importServiceProcessesIOException() {
         ImportService importService = new ImportService(typeDefStore, typeRegistry, null);