From 02002693df1c78322f71f51995db6db7b8eb6c3b Mon Sep 17 00:00:00 2001 From: suraj5077 Date: Mon, 21 Aug 2023 15:07:58 +0530 Subject: [PATCH] fix: add base entity to guidEntityMap in lineage response --- .../atlas/discovery/EntityLineageService.java | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/repository/src/main/java/org/apache/atlas/discovery/EntityLineageService.java b/repository/src/main/java/org/apache/atlas/discovery/EntityLineageService.java index 2fd2cc350a..a1fa489129 100644 --- a/repository/src/main/java/org/apache/atlas/discovery/EntityLineageService.java +++ b/repository/src/main/java/org/apache/atlas/discovery/EntityLineageService.java @@ -284,26 +284,23 @@ private AtlasLineageOnDemandInfo getLineageInfoOnDemand(String guid, AtlasLineag LineageOnDemandConstraints lineageConstraintsByGuid = getAndValidateLineageConstraintsByGuid(guid, atlasLineageOnDemandContext); AtlasLineageOnDemandInfo.LineageDirection direction = lineageConstraintsByGuid.getDirection(); int depth = lineageConstraintsByGuid.getDepth(); - AtlasLineageOnDemandInfo ret = initializeLineageOnDemandInfo(guid); - if (depth == 0) { + if (depth == 0) depth = -1; - } - - if (!ret.getRelationsOnDemand().containsKey(guid)) { + if (!ret.getRelationsOnDemand().containsKey(guid)) ret.getRelationsOnDemand().put(guid, new LineageInfoOnDemand(lineageConstraintsByGuid)); - } + AtomicInteger inputEntitiesTraversed = new AtomicInteger(0); AtomicInteger outputEntitiesTraversed = new AtomicInteger(0); if (isDataSet) { AtlasVertex datasetVertex = AtlasGraphUtilsV2.findByGuid(this.graph, guid); - if (direction == AtlasLineageOnDemandInfo.LineageDirection.INPUT || direction == AtlasLineageOnDemandInfo.LineageDirection.BOTH) { + if (direction == AtlasLineageOnDemandInfo.LineageDirection.INPUT || direction == AtlasLineageOnDemandInfo.LineageDirection.BOTH) traverseEdgesOnDemand(datasetVertex, true, depth, new HashSet<>(), atlasLineageOnDemandContext, ret, guid, inputEntitiesTraversed); - } - if (direction == AtlasLineageOnDemandInfo.LineageDirection.OUTPUT || direction == AtlasLineageOnDemandInfo.LineageDirection.BOTH) { + if (direction == AtlasLineageOnDemandInfo.LineageDirection.OUTPUT || direction == AtlasLineageOnDemandInfo.LineageDirection.BOTH) traverseEdgesOnDemand(datasetVertex, false, depth, new HashSet<>(), atlasLineageOnDemandContext, ret, guid, outputEntitiesTraversed); - } + AtlasEntityHeader baseEntityHeader = entityRetriever.toAtlasEntityHeader(datasetVertex, atlasLineageOnDemandContext.getAttributes()); + ret.getGuidEntityMap().put(guid, baseEntityHeader); } else { AtlasVertex processVertex = AtlasGraphUtilsV2.findByGuid(this.graph, guid); // make one hop to the next dataset vertices from process vertex and traverse with 'depth = depth - 1' @@ -317,7 +314,6 @@ private AtlasLineageOnDemandInfo getLineageInfoOnDemand(String guid, AtlasLineag } } RequestContext.get().endMetricRecord(metricRecorder); - return ret; }