ATLAS-3762: Improve edge creation using genuine iterator. Part 2

graphdb/api/src/main/java/org/apache/atlas/repository/graphdb/AtlasGraph.java

@@ -48,6 +48,15 @@ public interface AtlasGraph<V, E> {
     AtlasEdge<V, E> addEdge(AtlasVertex<V, E> outVertex, AtlasVertex<V, E> inVertex, String label);
 
     /**
+     * Fetch edges between two vertices using relationshipLabel
+     * @param fromVertex
+     * @param toVertex
+     * @param relationshipLabel
+     * @return
+     */
+    AtlasEdge<V, E> getEdgeBetweenVertices(AtlasVertex fromVertex, AtlasVertex toVertex, String relationshipLabel);
+
+        /**
          * Adds a vertex to the graph.
          *
          * @return

graphdb/janus/src/main/java/org/apache/atlas/repository/graphdb/janus/AtlasJanusGraph.java

@@ -172,6 +172,17 @@ public class AtlasJanusGraph implements AtlasGraph<AtlasJanusVertex, AtlasJanusE
     }
 
     @Override
+    public AtlasEdge getEdgeBetweenVertices(AtlasVertex fromVertex, AtlasVertex toVertex, String edgeLabel) {
+        GraphTraversal gt = V(fromVertex.getId()).outE(edgeLabel).where(__.otherV().hasId(toVertex.getId()));
+        Object o = gt.hasNext() ? gt.next() : null;
+        if (o == null) {
+            return null;
+        }
+
+        return GraphDbObjectFactory.createEdge(this, (Edge) o);
+    }
+
+    @Override
     public AtlasEdge<AtlasJanusVertex, AtlasJanusEdge> getEdge(String edgeId) {
         Iterator<Edge> it = getGraph().edges(edgeId);
         Edge           e  = getSingleElement(it, edgeId);

repository/src/main/java/org/apache/atlas/repository/graph/GraphHelper.java

@@ -739,6 +739,7 @@ public final class GraphHelper {
 
     public static List<AtlasVertex> getPropagationEnabledClassificationVertices(AtlasVertex entityVertex) {
         List<AtlasVertex> ret   = new ArrayList<>();
+        if (entityVertex.hasEdges(AtlasEdgeDirection.OUT, CLASSIFICATION_LABEL)) {
             Iterable edges = entityVertex.query().direction(AtlasEdgeDirection.OUT).label(CLASSIFICATION_LABEL).edges();
 
             if (edges != null) {
@@ -756,6 +757,7 @@ public final class GraphHelper {
                     }
                 }
             }
+        }
 
         return ret;
     }

repository/src/main/java/org/apache/atlas/repository/store/graph/v2/AtlasRelationshipStoreV2.java

@@ -784,36 +784,13 @@ public class AtlasRelationshipStoreV2 implements AtlasRelationshipStore {
         AtlasEdge ret = null;
 
         if (toVertex.hasEdges(AtlasEdgeDirection.IN, relationshipLabel) && fromVertex.hasEdges(AtlasEdgeDirection.OUT, relationshipLabel)) {
-            long fromVertexOutgoingEdgeCount = graphHelper.getOutGoingEdgesCountByLabel(fromVertex, relationshipLabel);
-            long toVertexIncomingEdgeCount = graphHelper.getInComingEdgesCountByLabel(toVertex, relationshipLabel);
-            if (toVertexIncomingEdgeCount < fromVertexOutgoingEdgeCount) {
-                Iterator<AtlasEdge> edgesIteratorIn = graphHelper.getIncomingEdgesByLabel(toVertex, relationshipLabel);
-                ret = getActiveEdgeFromList(edgesIteratorIn, fromVertex.getId(), e -> e.getOutVertex().getId());
-            } else {
-                Iterator<AtlasEdge> edgesIteratorOut = graphHelper.getOutGoingEdgesByLabel(fromVertex, relationshipLabel);
-                ret = getActiveEdgeFromList(edgesIteratorOut, toVertex.getId(), e -> e.getInVertex().getId());
-            }
+            ret = graph.getEdgeBetweenVertices(fromVertex, toVertex, relationshipLabel);
         }
 
         RequestContext.get().endMetricRecord(metric);
         return ret;
     }
 
-    private AtlasEdge getActiveEdgeFromList(Iterator<AtlasEdge> edgesIterator, Object vertexIdToCompare, Function<AtlasEdge, Object> edgeIdFn) {
-        while (edgesIterator != null && edgesIterator.hasNext()) {
-            AtlasEdge edge = edgesIterator.next();
-            if (edge != null) {
-                Status status = graphHelper.getStatus(edge);
-
-                if ((status == null || status == ACTIVE) && edgeIdFn.apply(edge).equals(vertexIdToCompare)) {
-                    return edge;
-                }
-            }
-        }
-
-        return null;
-    }
-
     private Long getRelationshipVersion(AtlasRelationship relationship) {
         Long ret = relationship != null ? relationship.getVersion() : null;