Merge pull request #3130 from atlanhq/feat/lineage-out-time

LIN-749: New finishTime attribute in lineage graph API
atlanhq · May 20, 2024 · 5aa3ba4 · 5aa3ba4
2 parents c6a4833 + 62d0dd4
commit 5aa3ba4
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 2 deletions.
diff --git a/intg/src/main/java/org/apache/atlas/model/instance/AtlasEntityHeader.java b/intg/src/main/java/org/apache/atlas/model/instance/AtlasEntityHeader.java
@@ -73,6 +73,7 @@ public class AtlasEntityHeader extends AtlasStruct implements Serializable {
     private String                          deleteHandler       = null;
     private Integer                         depth               = null;
     private Integer                         traversalOrder      = null;
+    private Integer                         finishTime          = null;
 
     private Map<String, AtlasSearchResult>  collapse    = null;
 
@@ -162,6 +163,14 @@ public AtlasEntityHeader(AtlasEntity entity) {
 
     public void setTraversalOrder(Integer traversalOrder) { this.traversalOrder = traversalOrder; }
 
+    public Integer getFinishTime() {
+        return finishTime;
+    }
+
+    public void setFinishTime(Integer finishTime) {
+        this.finishTime = finishTime;
+    }
+
     public AtlasEntity.Status getStatus() {
         return status;
     }

diff --git a/repository/src/main/java/org/apache/atlas/discovery/EntityLineageService.java b/repository/src/main/java/org/apache/atlas/discovery/EntityLineageService.java
@@ -302,8 +302,7 @@ private AtlasLineageOnDemandInfo getLineageInfoOnDemand(String guid, AtlasLineag
             if (direction == AtlasLineageOnDemandInfo.LineageDirection.OUTPUT || direction == AtlasLineageOnDemandInfo.LineageDirection.BOTH)
                 traverseEdgesOnDemand(datasetVertex, false, depth, level, new HashSet<>(), atlasLineageOnDemandContext, ret, guid, outputEntitiesTraversed, traversalOrder);
             AtlasEntityHeader baseEntityHeader = entityRetriever.toAtlasEntityHeader(datasetVertex, atlasLineageOnDemandContext.getAttributes());
-            baseEntityHeader.setDepth(level);
-            baseEntityHeader.setTraversalOrder(0);
+            setGraphTraversalMetadata(level, traversalOrder, baseEntityHeader);
             ret.getGuidEntityMap().put(guid, baseEntityHeader);
         } else  {
             AtlasVertex processVertex = AtlasGraphUtilsV2.findByGuid(this.graph, guid);
@@ -321,6 +320,12 @@ private AtlasLineageOnDemandInfo getLineageInfoOnDemand(String guid, AtlasLineag
         return ret;
     }
 
+    private static void setGraphTraversalMetadata(int level, AtomicInteger traversalOrder, AtlasEntityHeader baseEntityHeader) {
+        baseEntityHeader.setDepth(level);
+        baseEntityHeader.setTraversalOrder(0);
+        baseEntityHeader.setFinishTime(traversalOrder.get());
+    }
+
     private void traverseEdgesOnDemand(Iterator<AtlasEdge> processEdges, boolean isInput, int depth, int level, AtlasLineageOnDemandContext atlasLineageOnDemandContext, AtlasLineageOnDemandInfo ret, AtlasVertex processVertex, String baseGuid, AtomicInteger entitiesTraversed, AtomicInteger traversalOrder) throws AtlasBaseException {
         AtlasLineageOnDemandInfo.LineageDirection direction = isInput ? AtlasLineageOnDemandInfo.LineageDirection.INPUT : AtlasLineageOnDemandInfo.LineageDirection.OUTPUT;
         int nextLevel = isInput ? level - 1: level + 1;
@@ -427,6 +432,8 @@ private void traverseEdgesOnDemand(AtlasVertex datasetVertex, boolean isInput, i
                     }
                     if (entityVertex != null && !visitedVertices.contains(getId(entityVertex))) {
                         traverseEdgesOnDemand(entityVertex, isInput, depth - 1, nextLevel, visitedVertices, atlasLineageOnDemandContext, ret, baseGuid, entitiesTraversed, traversalOrder); // execute inner depth
+                        AtlasEntityHeader traversedEntity = ret.getGuidEntityMap().get(AtlasGraphUtilsV2.getIdFromVertex(entityVertex));
+                        traversedEntity.setFinishTime(traversalOrder.get());
                     }
                 }
             }