From 3dba9d333881825b657703f6f14f460fcef1d85c Mon Sep 17 00:00:00 2001 From: Mandar Ambawane Date: Wed, 3 Sep 2025 23:59:34 +0530 Subject: [PATCH] ATLAS-5097 : Support for fetching multiple types which are referenced by same relationship attribute name in export operation --- distro/src/conf/atlas-application.properties | 6 +- .../RelationshipAttributesExtractor.java | 101 ++++++++++++++++++ 2 files changed, 106 insertions(+), 1 deletion(-) diff --git a/distro/src/conf/atlas-application.properties b/distro/src/conf/atlas-application.properties index b5734d7a8de..bbb418bca41 100755 --- a/distro/src/conf/atlas-application.properties +++ b/distro/src/conf/atlas-application.properties @@ -281,4 +281,8 @@ atlas.search.gremlin.enable=false ######### Skip check for the same attribute name in Parent type and Child type ######### -#atlas.skip.check.for.parent.child.attribute.name=true \ No newline at end of file +#atlas.skip.check.for.parent.child.attribute.name=true + +#atlas.relationship.types.for.= +#atlas.relationship.attr.for.= +#atlas.relationship.edge.label.between..= \ No newline at end of file diff --git a/repository/src/main/java/org/apache/atlas/repository/impexp/RelationshipAttributesExtractor.java b/repository/src/main/java/org/apache/atlas/repository/impexp/RelationshipAttributesExtractor.java index 6f42e1b85f9..a52b8627347 100644 --- a/repository/src/main/java/org/apache/atlas/repository/impexp/RelationshipAttributesExtractor.java +++ b/repository/src/main/java/org/apache/atlas/repository/impexp/RelationshipAttributesExtractor.java @@ -17,18 +17,28 @@ */ package org.apache.atlas.repository.impexp; +import org.apache.atlas.ApplicationProperties; +import org.apache.atlas.AtlasException; import org.apache.atlas.model.instance.AtlasEntity; import org.apache.atlas.model.instance.AtlasRelatedObjectId; +import org.apache.atlas.model.instance.AtlasRelationship; +import org.apache.atlas.model.instance.AtlasStruct; import org.apache.atlas.model.typedef.AtlasBaseTypeDef; import org.apache.atlas.model.typedef.AtlasEntityDef; +import org.apache.atlas.repository.graph.GraphHelper; +import org.apache.atlas.repository.graphdb.AtlasEdge; +import org.apache.atlas.repository.graphdb.AtlasVertex; import org.apache.atlas.repository.impexp.ExportService.TraversalDirection; +import org.apache.atlas.repository.store.graph.v2.AtlasGraphUtilsV2; import org.apache.atlas.type.AtlasTypeRegistry; import org.apache.atlas.type.AtlasTypeUtil; +import org.apache.commons.collections.CollectionUtils; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import java.util.ArrayList; import java.util.Collection; +import java.util.Iterator; import java.util.List; import static org.apache.atlas.repository.impexp.ExportService.ExportContext; @@ -40,6 +50,10 @@ public class RelationshipAttributesExtractor implements ExtractStrategy { private static final Logger LOG = LoggerFactory.getLogger(RelationshipAttributesExtractor.class); + private static final String DISPLAY_NAME_PROPERTY = "Asset.__s_name"; + private static final String RELATIONSHIP_GUID_PROPERTY = "_r__guid"; + private static final String QUALIFIED_NAME_PROPERTY = "Referenceable.qualifiedName"; + private final AtlasTypeRegistry typeRegistry; public RelationshipAttributesExtractor(AtlasTypeRegistry typeRegistry) { @@ -148,6 +162,93 @@ private List getRelatedObjectIds(AtlasEntity entity) { } } + String typeName = entity.getTypeName(); + List relatedTypes = checkForRelationshipTypes(typeName); + if (CollectionUtils.isNotEmpty(relatedTypes)) { + String alreadyPopulatedType = checkAlreadyPopulatedRelationshipType(entity); + relatedTypes.remove(alreadyPopulatedType); + } + String entityGuid = entity.getGuid(); + List relatedObjectIdList = populateRelatedObjectIds(typeName, entityGuid, relatedTypes); + relatedObjectIds.addAll(relatedObjectIdList); + return relatedObjectIds; + } + + private String checkAlreadyPopulatedRelationshipType(AtlasEntity entity) { + String typeName = entity.getTypeName(); + String relationshipAttr; + try { + relationshipAttr = ApplicationProperties.get().getString("atlas.relationship.attr.for." + typeName); + } catch (AtlasException e) { + throw new RuntimeException(e); + } + Object o = entity.getRelationshipAttribute(relationshipAttr); + String relatedTypeName = null; + if (o instanceof AtlasRelatedObjectId) { + AtlasRelatedObjectId relatedObjectId = (AtlasRelatedObjectId) o; + relatedTypeName = relatedObjectId.getTypeName(); + } else if (o instanceof Collection) { + List relatedObjectIds = (List) o; + if (!CollectionUtils.isEmpty(relatedObjectIds)) { + relatedTypeName = relatedObjectIds.get(0).getTypeName(); + } + } + return relatedTypeName; + } + + private List checkForRelationshipTypes(String typeName) { + List relatedTypeList = null; + try { + String[] relatedTypes = ApplicationProperties.get().getStringArray("atlas.relationship.types.for." + typeName); + if (relatedTypes != null && relatedTypes.length > 0) { + relatedTypeList = new ArrayList<>(); + for (String relatedType : relatedTypes) { + relatedTypeList.add(relatedType); + } + } + } catch (AtlasException e) { + throw new RuntimeException(e); + } + return relatedTypeList; + } + + private List populateRelatedObjectIds(String typeName, String entityGuid, List relatedTypes) { + List relatedObjectIds = new ArrayList<>(); + AtlasVertex vertex = AtlasGraphUtilsV2.findByGuid(entityGuid); + if (CollectionUtils.isEmpty(relatedTypes)) { + return relatedObjectIds; + } + for (String relatedType : relatedTypes) { + String edgeLabel; + try { + edgeLabel = ApplicationProperties.get().getString("atlas.relationship.edge.label.between." + typeName + "." + relatedType); + } catch (AtlasException e) { + throw new RuntimeException(e); + } + Iterator edges = GraphHelper.getIncomingEdgesByLabel(vertex, edgeLabel); + while (edges.hasNext()) { + AtlasEdge relationshipEdge = edges.next(); + AtlasVertex relationshipVertex = relationshipEdge.getOutVertex(); + AtlasRelatedObjectId relatedObjectId = new AtlasRelatedObjectId(); + AtlasEntity.Status vertexStatus = GraphHelper.getStatus(relationshipVertex); + relatedObjectId.setEntityStatus(vertexStatus); + String displayText = relationshipVertex.getProperty(DISPLAY_NAME_PROPERTY, String.class); + relatedObjectId.setDisplayText(displayText); + relatedObjectId.setRelationshipType(GraphHelper.getTypeName(relationshipEdge)); + String relationshipGuid = relationshipEdge.getProperty(RELATIONSHIP_GUID_PROPERTY, String.class); + relatedObjectId.setRelationshipGuid(relationshipGuid); + AtlasRelationship.Status edgeStatus = GraphHelper.getEdgeStatus(relationshipEdge); + relatedObjectId.setRelationshipStatus(edgeStatus); + AtlasStruct relationshipAttributes = new AtlasStruct(); + relationshipAttributes.setTypeName(GraphHelper.getTypeName(relationshipEdge)); + relatedObjectId.setRelationshipAttributes(relationshipAttributes); + String qualifiedName = relationshipVertex.getProperty(QUALIFIED_NAME_PROPERTY, String.class); + relatedObjectId.setQualifiedName(qualifiedName); + relatedObjectId.setGuid(GraphHelper.getGuid(relationshipVertex)); + relatedObjectId.setTypeName(GraphHelper.getTypeName(relationshipVertex)); + relatedObjectIds.add(relatedObjectId); + } + } return relatedObjectIds; } }