Merge pull request #3609 from atlanhq/dev/DG-1791

DG-1791 | Retrieve only the active edges
atlanhq · Oct 8, 2024 · 25feb1d · 25feb1d
2 parents 3f43169 + 9aac7e5
commit 25feb1d
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 5 deletions.
diff --git a/repository/src/main/java/org/apache/atlas/repository/graph/GraphHelper.java b/repository/src/main/java/org/apache/atlas/repository/graph/GraphHelper.java
@@ -2006,6 +2006,31 @@ public static Iterator<AtlasVertex> getActiveVertices(AtlasVertex vertex, AtlasE
         }
     }
 
+    /**
+     * Get all the active edges
+     * @param vertex entity vertex
+     * @param childrenEdgeLabel Edge label of children
+     * @return Iterator of children edges
+     */
+    public static Iterator<AtlasEdge> getActiveEdges(AtlasVertex vertex, String childrenEdgeLabel, AtlasEdgeDirection direction) throws AtlasBaseException {
+        AtlasPerfMetrics.MetricRecorder metricRecorder = RequestContext.get().startMetricRecord("GraphHelper.getActiveEdges");
+
+        try {
+            return vertex.query()
+                    .direction(direction)
+                    .label(childrenEdgeLabel)
+                    .has(STATE_PROPERTY_KEY, ACTIVE_STATE_VALUE)
+                    .edges()
+                    .iterator();
+        } catch (Exception e) {
+            LOG.error("Error while getting active edges of vertex for edge label " + childrenEdgeLabel, e);
+            throw new AtlasBaseException(AtlasErrorCode.INTERNAL_ERROR, e);
+        }
+        finally {
+            RequestContext.get().endMetricRecord(metricRecorder);
+        }
+    }
+
     public static Iterator<AtlasVertex> getAllChildrenVertices(AtlasVertex vertex, String childrenEdgeLabel) throws AtlasBaseException {
         return getAllVertices(vertex, childrenEdgeLabel, AtlasEdgeDirection.OUT);
     }

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
@@ -1465,7 +1465,7 @@ private boolean skipClassificationTaskCreation(String classificationId) throws A
     }
 
 
-    public void removeHasLineageOnDelete(Collection<AtlasVertex> vertices) {
+    public void removeHasLineageOnDelete(Collection<AtlasVertex> vertices) throws AtlasBaseException {
         AtlasPerfMetrics.MetricRecorder metricRecorder = RequestContext.get().startMetricRecord("removeHasLineageOnDelete");
 
         for (AtlasVertex vertexToBeDeleted : vertices) {
@@ -1495,7 +1495,7 @@ public void removeHasLineageOnDelete(Collection<AtlasVertex> vertices) {
     }
 
 
-    public void resetHasLineageOnInputOutputDelete(Collection<AtlasEdge> removedEdges, AtlasVertex deletedVertex) {
+    public void resetHasLineageOnInputOutputDelete(Collection<AtlasEdge> removedEdges, AtlasVertex deletedVertex) throws AtlasBaseException {
         AtlasPerfMetrics.MetricRecorder metricRecorder = RequestContext.get().startMetricRecord("resetHasLineageOnInputOutputDelete");
 
         for (AtlasEdge atlasEdge : removedEdges) {
@@ -1513,13 +1513,15 @@ public void resetHasLineageOnInputOutputDelete(Collection<AtlasEdge> removedEdge
             if (getStatus(processVertex) == ACTIVE && !processVertex.equals(deletedVertex)) {
                 String edgeLabel = isOutputEdge ? PROCESS_OUTPUTS : PROCESS_INPUTS;
 
-                Iterator<AtlasEdge> edgeIterator = processVertex.getEdges(AtlasEdgeDirection.BOTH, edgeLabel).iterator();
+              //  Iterator<AtlasEdge> edgeIterator = processVertex.getEdges(AtlasEdgeDirection.BOTH, edgeLabel).iterator();
+                Iterator<AtlasEdge> edgeIterator = GraphHelper.getActiveEdges(processVertex, edgeLabel, AtlasEdgeDirection.BOTH);
+
                 boolean activeEdgeFound = false;
 
                 while (edgeIterator.hasNext()) {
                     AtlasPerfMetrics.MetricRecorder edgeIteratorMetricRecorder1 = RequestContext.get().startMetricRecord("resetHasLineageOnInputOutputDelete_edgeIterator1");
                     AtlasEdge edge = edgeIterator.next();
-                    if (getStatus(edge) == ACTIVE && !removedEdges.contains(edge)) {
+                    if (!removedEdges.contains(edge)) {
                         AtlasVertex relatedAssetVertex = edge.getInVertex();
 
                         if (getStatus(relatedAssetVertex) == ACTIVE) {
@@ -1535,7 +1537,8 @@ public void resetHasLineageOnInputOutputDelete(Collection<AtlasEdge> removedEdge
 
                     String oppositeEdgeLabel = isOutputEdge ? PROCESS_INPUTS : PROCESS_OUTPUTS;
 
-                    Iterator<AtlasEdge> processEdgeIterator = processVertex.getEdges(AtlasEdgeDirection.BOTH, oppositeEdgeLabel).iterator();
+                    // Iterator<AtlasEdge> processEdgeIterator = processVertex.getEdges(AtlasEdgeDirection.BOTH, oppositeEdgeLabel).iterator();
+                    Iterator<AtlasEdge> processEdgeIterator = GraphHelper.getActiveEdges(processVertex, oppositeEdgeLabel, AtlasEdgeDirection.BOTH);
 
                     while (processEdgeIterator.hasNext()) {
                         AtlasPerfMetrics.MetricRecorder edgeIteratorMetricRecorder2 = RequestContext.get().startMetricRecord("resetHasLineageOnInputOutputDelete_edgeIterator2");