From 017df90fd80cd5f5fe2e14ff956689ccf67ad784 Mon Sep 17 00:00:00 2001
From: aarshi <aarshi.arora@atlan.com>
Date: Thu, 19 Dec 2024 19:52:10 +0530
Subject: [PATCH] add attributes as set

---
 .../store/graph/v2/EntityGraphRetriever.java        | 13 +++++++------
 1 file changed, 7 insertions(+), 6 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 2da180cf62..a4482c88cb 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
@@ -1014,7 +1014,7 @@ private Map<String, Object> preloadProperties(AtlasVertex entityVertex, AtlasEnt
         Iterator<VertexProperty<Object>> traversal = ((AtlasJanusVertex)entityVertex).getWrappedElement().properties();
 
         //  retrieve all the valid relationships for this entityType
-        Map<String, String> relationshipsLookup = fetchEdgeNames(entityType);
+        Map<String, Set<String>> relationshipsLookup = fetchEdgeNames(entityType);
 
         // Fetch edges in both directions
         retrieveEdgeLabels(entityVertex, AtlasEdgeDirection.BOTH, attributes, relationshipsLookup, propertiesMap);
@@ -1052,12 +1052,13 @@ private Map<String, Object> preloadProperties(AtlasVertex entityVertex, AtlasEnt
         return propertiesMap;
     }
 
-    private Map<String, String> fetchEdgeNames(AtlasEntityType entityType){
+    private Map<String, Set<String>> fetchEdgeNames(AtlasEntityType entityType){
         Map<String, Map<String, AtlasAttribute>> relationships = entityType.getRelationshipAttributes();
-        Map<String, String> edgeNames = new HashMap<>();
+        Map<String, Set<String>> edgeNames = new HashMap<>();
         relationships.forEach((k,v) -> {
             v.forEach((k1,v1) -> {
-                edgeNames.put(k1, k);
+                edgeNames.putIfAbsent(k1, new HashSet<>());
+                edgeNames.get(k1).add(k);
             });
         });
 
@@ -1065,7 +1066,7 @@ private Map<String, String> fetchEdgeNames(AtlasEntityType entityType){
         return edgeNames;
     }
 
-    private void retrieveEdgeLabels(AtlasVertex entityVertex, AtlasEdgeDirection edgeDirection, Set<String> attributes, Map<String, String> relationshipsLookup,Map<String, Object> propertiesMap) throws AtlasBaseException {
+    private void retrieveEdgeLabels(AtlasVertex entityVertex, AtlasEdgeDirection edgeDirection, Set<String> attributes, Map<String, Set<String>> relationshipsLookup,Map<String, Object> propertiesMap) throws AtlasBaseException {
         Iterator<AtlasJanusEdge> edges = GraphHelper.getOnlyActiveEdges(entityVertex, edgeDirection);
 
         List<String> edgeLabelsDebug = new ArrayList<>();
@@ -1091,7 +1092,7 @@ private void retrieveEdgeLabels(AtlasVertex entityVertex, AtlasEdgeDirection edg
             String edgeTypeName = edgesTypeName.get(edgeLabel);
 
              LOG.warn("edgeTypeName for attribute: {} :{}", edgeTypeName, relationshipsLookup.containsKey(edgeTypeName));
-            if (MapUtils.isNotEmpty(relationshipsLookup) && relationshipsLookup.containsKey(edgeTypeName) && attribute.equals(relationshipsLookup.get(edgeTypeName))) {
+            if (MapUtils.isNotEmpty(relationshipsLookup) && relationshipsLookup.containsKey(edgeTypeName) && relationshipsLookup.get(edgeTypeName).contains(attribute)) {
                 edgeLabels.add(attribute);
                 LOG.warn("attribute matched by edgeTypeName");
             }