Repository: atlas Updated Branches: refs/heads/master 5c2f7a0c8 -> 4c2f7bf20
ATLAS-2372: Export-Import: Support for Relationship Attributes Missing Signed-off-by: Sarath Subramanian <[email protected]> Project: http://git-wip-us.apache.org/repos/asf/atlas/repo Commit: http://git-wip-us.apache.org/repos/asf/atlas/commit/4c2f7bf2 Tree: http://git-wip-us.apache.org/repos/asf/atlas/tree/4c2f7bf2 Diff: http://git-wip-us.apache.org/repos/asf/atlas/diff/4c2f7bf2 Branch: refs/heads/master Commit: 4c2f7bf203d2e688e4ef80570657ce5a979cf93d Parents: 5c2f7a0 Author: Ashutosh Mestry <[email protected]> Authored: Fri Jan 19 15:10:23 2018 -0800 Committer: Sarath Subramanian <[email protected]> Committed: Fri Jan 19 15:10:23 2018 -0800 ---------------------------------------------------------------------- .../atlas/repository/impexp/ExportService.java | 14 ----------- .../graph/v1/AtlasEntityGraphDiscoveryV1.java | 26 ++++++++++++-------- .../store/graph/v1/EntityGraphMapper.java | 7 +++++- .../store/graph/v1/EntityMutationContext.java | 3 +++ .../repository/impexp/ImportServiceTest.java | 15 +++++++++-- 5 files changed, 38 insertions(+), 27 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/atlas/blob/4c2f7bf2/repository/src/main/java/org/apache/atlas/repository/impexp/ExportService.java ---------------------------------------------------------------------- diff --git a/repository/src/main/java/org/apache/atlas/repository/impexp/ExportService.java b/repository/src/main/java/org/apache/atlas/repository/impexp/ExportService.java index c5e6534..ca8a106 100644 --- a/repository/src/main/java/org/apache/atlas/repository/impexp/ExportService.java +++ b/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); http://git-wip-us.apache.org/repos/asf/atlas/blob/4c2f7bf2/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityGraphDiscoveryV1.java ---------------------------------------------------------------------- diff --git a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityGraphDiscoveryV1.java b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityGraphDiscoveryV1.java index 450fa87..f092d18 100644 --- a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityGraphDiscoveryV1.java +++ b/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(); http://git-wip-us.apache.org/repos/asf/atlas/blob/4c2f7bf2/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 7474a8a..b961e20 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 @@ -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) { http://git-wip-us.apache.org/repos/asf/atlas/blob/4c2f7bf2/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 8a6a0e3..df5bd09 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 @@ -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) { http://git-wip-us.apache.org/repos/asf/atlas/blob/4c2f7bf2/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 ab25faa..136fe8a 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 @@ -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);
