Merge pull request #3608 from atlanhq/DG-1791

DG-1791 | Retrieve only the active edges
atlanhq · Oct 9, 2024 · b33e745 · b33e745
2 parents 756d64d + 0eafe33
commit b33e745
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 6 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
@@ -1967,6 +1967,31 @@ public static Iterator<AtlasVertex> getActiveChildrenVertices(AtlasVertex vertex
         return getActiveVertices(vertex, childrenEdgeLabel, AtlasEdgeDirection.OUT);
     }
 
+    /**
+     * 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> getActiveVertices(AtlasVertex vertex, String childrenEdgeLabel, AtlasEdgeDirection direction) throws AtlasBaseException {
         AtlasPerfMetrics.MetricRecorder metricRecorder = RequestContext.get().startMetricRecord("CategoryPreProcessor.getEdges");
 
@@ -1985,7 +2010,6 @@ public static Iterator<AtlasVertex> getActiveVertices(AtlasVertex vertex, String
             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
@@ -1447,7 +1447,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) {
@@ -1477,7 +1477,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) {
@@ -1494,12 +1494,13 @@ 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 = GraphHelper.getActiveEdges(processVertex, edgeLabel, AtlasEdgeDirection.BOTH);
+
                 boolean activeEdgeFound = false;
 
                 while (edgeIterator.hasNext()) {
                     AtlasEdge edge = edgeIterator.next();
-                    if (getStatus(edge) == ACTIVE && !removedEdges.contains(edge)) {
+                    if (!removedEdges.contains(edge)) {
                         AtlasVertex relatedAssetVertex = edge.getInVertex();
 
                         if (getStatus(relatedAssetVertex) == ACTIVE) {
@@ -1514,7 +1515,7 @@ 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 = GraphHelper.getActiveEdges(processVertex, oppositeEdgeLabel, AtlasEdgeDirection.BOTH);
 
                     while (processEdgeIterator.hasNext()) {
                         AtlasEdge edge = processEdgeIterator.next();