From 90bd05240921e8879ddc9a3854506978cb1793f6 Mon Sep 17 00:00:00 2001
From: aarshi <aarshi.arora@atlan.com>
Date: Mon, 16 Dec 2024 20:18:22 +0530
Subject: [PATCH] fix property keys

---
 .../store/graph/v2/EntityGraphRetriever.java  | 21 +++++++++++--------
 1 file changed, 12 insertions(+), 9 deletions(-)

diff --git a/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java b/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java
index b7732f5610..a0a094222e 100644
--- a/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java
+++ b/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java
@@ -50,6 +50,7 @@
 import org.apache.atlas.repository.graphdb.AtlasGraph;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.atlas.repository.graphdb.janus.AtlasJanusEdge;
+import org.apache.atlas.repository.graphdb.janus.AtlasJanusPropertyKey;
 import org.apache.atlas.repository.graphdb.janus.AtlasJanusVertex;
 import org.apache.atlas.repository.util.AccessControlUtils;
 import org.apache.atlas.type.AtlasArrayType;
@@ -1065,23 +1066,25 @@ private Map<String, String> fetchEdgeNames(AtlasEntityType entityType){
     private void retrieveEdgeLabels(AtlasVertex entityVertex, AtlasEdgeDirection edgeDirection, Set<String> attributes, Map<String, String> relationshipsLookup,Map<String, Object> propertiesMap) throws AtlasBaseException {
         Iterator<AtlasJanusEdge> edges = GraphHelper.getOnlyActiveEdges(entityVertex, edgeDirection);
 
+        List<String> edgeLabelsDebug = new ArrayList<>();
+        Map<String, String> edgesTypeName = new HashMap();
 
-
-        List<AtlasEdge> edgesDebug= new ArrayList<>();
         while (edges.hasNext()) {
             AtlasJanusEdge edge = edges.next();
-            edgesDebug.add(edge);
+            String label = edge.getLabel();
+            edgeLabelsDebug.add(label);
+            edgesTypeName.putIfAbsent(label, edge.getProperty(TYPE_NAME_PROPERTY_KEY, String.class));
         }
 
-        Set<String> edgeLabels= new HashSet<>();
-
-        edgesDebug.stream().filter(Objects::nonNull).forEach(edge -> attributes.forEach(attribute->{
-               if (edge.getLabel().contains(attribute)){
+        Set<String> edgeLabels = new HashSet<>();
+        edgeLabelsDebug.stream().filter(Objects::nonNull).forEach(edgeLabel -> attributes.forEach(attribute->{
+               if (edgeLabel.contains(attribute)){
                    edgeLabels.add(attribute);
                    return;
                }
-            String edgeName = edge.getProperty(TYPE_NAME_PROPERTY_KEY, String.class);
-            if (relationshipsLookup.containsKey(edgeName) && attribute.equals(relationshipsLookup.get(edgeName))) {
+
+            String edgeTypeName = edgesTypeName.get(edgeLabel);
+            if (relationshipsLookup.containsKey(edgeTypeName) && attribute.equals(relationshipsLookup.get(edgeTypeName))) {
                 edgeLabels.add(attribute);
             }
         }));