Merge pull request #2395 from atlanhq/nb/hasLineage

[stag] DG-128 Fix __hasLineage filter issue
atlanhq · Oct 3, 2023 · a339e74 · a339e74
2 parents e8a528c + 89462d8
commit a339e74
Show file tree

Hide file tree

Showing 4 changed files with 72 additions and 16 deletions.
diff --git a/intg/src/main/java/org/apache/atlas/model/instance/AtlasHasLineageRequest.java b/intg/src/main/java/org/apache/atlas/model/instance/AtlasHasLineageRequest.java
@@ -40,6 +40,7 @@ public class AtlasHasLineageRequest implements Serializable {
     public AtlasHasLineageRequest() {
     }
 
+    private String assetGuid;
     private String processGuid;
     private String endGuid;
     private String label;
@@ -68,9 +69,18 @@ public void setLabel(String label) {
         this.label = label;
     }
 
+    public String getAssetGuid() {
+        return assetGuid;
+    }
+
+    public void setAssetGuid(String assetGuid) {
+        this.assetGuid = assetGuid;
+    }
+
     @Override
     public String toString() {
         return "AtlasHasLineageRequest{" +
+                "assetGuid='" + assetGuid + '\'' +
                 "processGuid='" + processGuid + '\'' +
                 ", endGuid='" + endGuid + '\'' +
                 ", label='" + label + '\'' +

diff --git a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/DeleteHandlerV1.java b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/DeleteHandlerV1.java
@@ -1527,6 +1527,8 @@ public void resetHasLineageOnInputOutputDelete(Collection<AtlasEdge> removedEdge
     private void updateAssetHasLineageStatus(AtlasVertex assetVertex, AtlasEdge currentEdge, Collection<AtlasEdge> removedEdges) {
         AtlasPerfMetrics.MetricRecorder metricRecorder = RequestContext.get().startMetricRecord("updateAssetHasLineageStatus");
 
+        removedEdges.forEach(edge -> RequestContext.get().addToDeletedEdgesIdsForResetHasLineage(edge.getIdForDisplay()));
+
         Iterator<AtlasEdge> edgeIterator = assetVertex.query()
                 .direction(AtlasEdgeDirection.BOTH)
                 .label(PROCESS_EDGE_LABELS)
@@ -1537,7 +1539,7 @@ private void updateAssetHasLineageStatus(AtlasVertex assetVertex, AtlasEdge curr
         int processHasLineageCount = 0;
         while (edgeIterator.hasNext()) {
             AtlasEdge edge = edgeIterator.next();
-            if (!removedEdges.contains(edge) && !currentEdge.equals(edge)) {
+            if (!RequestContext.get().getDeletedEdgesIdsForResetHasLineage().contains(edge.getIdForDisplay()) && !currentEdge.equals(edge)) {
                 AtlasVertex relatedProcessVertex = edge.getOutVertex();
                 boolean processHasLineage = getEntityHasLineage(relatedProcessVertex);
                 if (processHasLineage) {

diff --git a/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/AtlasEntityStoreV2.java b/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/AtlasEntityStoreV2.java
@@ -2443,28 +2443,61 @@ public void repairHasLineage(AtlasHasLineageRequests requests) throws AtlasBaseE
         Set<AtlasEdge> inputOutputEdges = new HashSet<>();
 
         for (AtlasHasLineageRequest request : requests.getRequest()) {
-            AtlasVertex processVertex = AtlasGraphUtilsV2.findByGuid(this.graph, request.getProcessGuid());
-            AtlasVertex assetVertex = AtlasGraphUtilsV2.findByGuid(this.graph, request.getEndGuid());
-            AtlasEdge edge = null;
-            try {
-                if (processVertex != null && assetVertex != null) {
-                    edge = graphHelper.getEdge(processVertex, assetVertex, request.getLabel());
-                } else {
-                    LOG.warn("Skipping since vertex is null for processGuid {} and asset Guid {}"
-                            ,request.getProcessGuid(),request.getEndGuid()  );
+            if (StringUtils.isNotEmpty(request.getAssetGuid())) {
+                //only supports repairing scenario mentioned here - https://atlanhq.atlassian.net/browse/DG-128?focusedCommentId=20652
+                repairHasLineageForAsset(request);
+
+            } else {
+                AtlasVertex processVertex = AtlasGraphUtilsV2.findByGuid(this.graph, request.getProcessGuid());
+                AtlasVertex assetVertex = AtlasGraphUtilsV2.findByGuid(this.graph, request.getEndGuid());
+                AtlasEdge edge = null;
+                try {
+                    if (processVertex != null && assetVertex != null) {
+                        edge = graphHelper.getEdge(processVertex, assetVertex, request.getLabel());
+                    } else {
+                        LOG.warn("Skipping since vertex is null for processGuid {} and asset Guid {}"
+                                ,request.getProcessGuid(),request.getEndGuid()  );
+                    }
+                } catch (RepositoryException re) {
+                    throw new AtlasBaseException(AtlasErrorCode.HAS_LINEAGE_GET_EDGE_FAILED, re);
                 }
-            } catch (RepositoryException re) {
-                throw new AtlasBaseException(AtlasErrorCode.HAS_LINEAGE_GET_EDGE_FAILED, re);
-            }
 
-            if (edge != null) {
-                inputOutputEdges.add(edge);
+                if (edge != null) {
+                    inputOutputEdges.add(edge);
+                }
             }
         }
-        repairHasLineageWithAtlasEdges(inputOutputEdges);
+
+        if (CollectionUtils.isNotEmpty(inputOutputEdges)) {
+            repairHasLineageWithAtlasEdges(inputOutputEdges);
+        }
+
         RequestContext.get().endMetricRecord(metricRecorder);
     }
 
+    private void repairHasLineageForAsset(AtlasHasLineageRequest request) {
+        //only supports repairing scenario mentioned here - https://atlanhq.atlassian.net/browse/DG-128?focusedCommentId=20652
+
+        AtlasVertex assetVertex = AtlasGraphUtilsV2.findByGuid(this.graph, request.getAssetGuid());
+
+        if (getEntityHasLineage(assetVertex)) {
+            Iterator<AtlasEdge> lineageEdges = assetVertex.getEdges(AtlasEdgeDirection.BOTH, PROCESS_EDGE_LABELS).iterator();
+            boolean foundActiveRel = false;
+
+            while (lineageEdges.hasNext()) {
+                AtlasEdge edge = lineageEdges.next();
+                if (getStatus(edge) == ACTIVE) {
+                    foundActiveRel = true;
+                    break;
+                }
+            }
+
+            if (!foundActiveRel) {
+                AtlasGraphUtilsV2.setEncodedProperty(assetVertex, HAS_LINEAGE, false);
+            }
+        }
+    }
+
     public void repairHasLineageWithAtlasEdges(Set<AtlasEdge> inputOutputEdges) {
         AtlasPerfMetrics.MetricRecorder metricRecorder = RequestContext.get().startMetricRecord("repairHasLineageWithAtlasEdges");
 

diff --git a/server-api/src/main/java/org/apache/atlas/RequestContext.java b/server-api/src/main/java/org/apache/atlas/RequestContext.java
@@ -94,6 +94,8 @@ public class RequestContext {
     private boolean     allowDuplicateDisplayName;
     private MetricsRegistry metricsRegistry;
     private boolean skipAuthorizationCheck = false;
+    private Set<String> deletedEdgesIdsForResetHasLineage = new HashSet<>(0);
+
 
     private RequestContext() {
     }
@@ -147,6 +149,7 @@ public void clearCache() {
         this.removedElementsMap.clear();
         this.deletedEdgesIds.clear();
         this.processGuidIds.clear();
+        this.deletedEdgesIdsForResetHasLineage.clear();
         this.requestContextHeaders.clear();
         this.relationshipEndToVertexIdMap.clear();
         this.relationshipMutationMap.clear();
@@ -389,6 +392,14 @@ public Set<String> getDeletedEdgesIds() {
         return deletedEdgesIds;
     }
 
+    public void addToDeletedEdgesIdsForResetHasLineage(String edgeId) {
+        deletedEdgesIdsForResetHasLineage.add(edgeId);
+    }
+
+    public Set<String> getDeletedEdgesIdsForResetHasLineage() {
+        return deletedEdgesIdsForResetHasLineage;
+    }
+
     public Set<String> getProcessGuidIds() {
         return processGuidIds;
     }