Commit 4c2f7bf2 by Ashutosh Mestry Committed by Sarath Subramanian

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

parent 5c2f7a0c
......@@ -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);
......
......@@ -305,26 +305,32 @@ public class AtlasEntityGraphDiscoveryV1 implements EntityGraphDiscovery {
List<String> visitedAttributes = new ArrayList<>();
// visit relationship attributes
for (AtlasAttribute attribute : entityType.getRelationshipAttributes().values()) {
if(!(this.discoveryContext.getEntityStream() instanceof EntityImportStream)) {
visitRelationships(entityType, entity, visitedAttributes);
}
// visit struct attributes
for (AtlasAttribute attribute : entityType.getAllAttributes().values()) {
AtlasType attrType = attribute.getAttributeType();
String attrName = attribute.getName();
Object attrVal = entity.getRelationshipAttribute(attrName);
Object attrVal = entity.getAttribute(attrName);
if (entity.hasRelationshipAttribute(attrName)) {
if (entity.hasAttribute(attrName) && !visitedAttributes.contains(attrName)) {
visitAttribute(attrType, attrVal);
visitedAttributes.add(attrName);
}
}
}
// visit struct attributes
for (AtlasAttribute attribute : entityType.getAllAttributes().values()) {
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.getAttribute(attrName);
Object attrVal = entity.getRelationshipAttribute(attrName);
if (entity.hasAttribute(attrName) && !visitedAttributes.contains(attrName)) {
if (entity.hasRelationshipAttribute(attrName)) {
visitAttribute(attrType, attrVal);
visitedAttributes.add(attrName);
}
}
}
......
......@@ -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) {
......
......@@ -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) {
......
......@@ -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);
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment