add attributes as set

atlanhq · Dec 19, 2024 · 017df90 · 017df90
1 parent e9683e5
commit 017df90
Showing 1 changed file with 7 additions and 6 deletions.
diff --git a/...sitory/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java b/...sitory/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,20 +1052,21 @@ 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);
             });
         });
 
         LOG.warn("get all edges: {}", edgeNames);
         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");
             }