fix marker logic for all relationships

atlanhq · Dec 16, 2024 · e2433e2 · e2433e2
1 parent aaa26a0
commit e2433e2
Showing 1 changed file with 39 additions and 15 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
@@ -86,6 +86,7 @@
 import java.util.*;
 import java.util.concurrent.*;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import static org.apache.atlas.glossary.GlossaryUtils.TERM_ASSIGNMENT_ATTR_CONFIDENCE;
 import static org.apache.atlas.glossary.GlossaryUtils.TERM_ASSIGNMENT_ATTR_CREATED_BY;
@@ -1007,8 +1008,14 @@ private Map<String, Object> preloadProperties(AtlasVertex entityVertex, AtlasEnt
 
         // Execute the traversal to fetch properties
         Iterator<VertexProperty<Object>> traversal = ((AtlasJanusVertex)entityVertex).getWrappedElement().properties();
+
+        //  retrieve all the valid relationships for this entityType
+        Map<String, String> relationshipsLookup = fetchEdgeNames(entityType);
+
         // Fetch edges in both directions
-        retrieveEdgeLabels(entityVertex, AtlasEdgeDirection.BOTH, attributes, propertiesMap);
+        retrieveEdgeLabels(entityVertex, AtlasEdgeDirection.BOTH, attributes, relationshipsLookup, propertiesMap);
+
+
 
         // Iterate through the resulting VertexProperty objects
         while (traversal.hasNext()) {
@@ -1041,25 +1048,43 @@ private Map<String, Object> preloadProperties(AtlasVertex entityVertex, AtlasEnt
         return propertiesMap;
     }
 
-    private void retrieveEdgeLabels(AtlasVertex entityVertex, AtlasEdgeDirection edgeDirection, Set<String> attributes, Map<String, Object> propertiesMap) throws AtlasBaseException {
+    private Map<String, String> fetchEdgeNames(AtlasEntityType entityType){
+        Map<String, Map<String, AtlasAttribute>> relationships = entityType.getRelationshipAttributes();
+        Map<String, String> edgeNames = new HashMap<>();
+        relationships.forEach((k,v) -> {
+            v.forEach((k1,v1) -> {
+                edgeNames.put(k1, k);
+            });
+        });
+        return edgeNames;
+    }
+
+    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<>();
+
+        List<AtlasEdge> edgesDebug= new ArrayList<>();
         while (edges.hasNext()) {
             AtlasJanusEdge edge = edges.next();
-            edgeLabelsDebug.add(edge.getLabel());
+            edgesDebug.add(edge);
         }
 
-        Set<String> edgeLabels =
-                edgeLabelsDebug.stream()
-                        .map(edgeLabel -> {
-                            Optional<String> matchingAttrOpt = attributes.stream().filter(ele -> (edgeLabel.contains(ele))).findFirst();
-                            return matchingAttrOpt.orElse(null);
-                        }).filter(Objects::nonNull)
-                        .collect(Collectors.toSet());
+        Set<String> edgeLabels= new HashSet<>();
+
+        edgesDebug.stream().filter(Objects::nonNull).forEach(edge -> attributes.forEach(attribute->{
+               if (edge.getLabel().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))) {
+                edgeLabels.add(attribute);
+            }
+        }));
 
         edgeLabels.stream().forEach(e -> propertiesMap.put(e, StringUtils.SPACE));
+
     }
     private void updateAttrValue( Map<String, Object> propertiesMap, VertexProperty<Object> property){
         Object value = propertiesMap.get(property.key());
@@ -1923,7 +1948,7 @@ public Object getVertexAttributePreFetchCache(AtlasVertex vertex, AtlasAttribute
         }
 
         // if value is empty && element is array and not inward relation, return empty list
-        if (properties.get(attribute.getName()) == null && typeCategory.equals(TypeCategory.ARRAY) && !AtlasRelationshipEdgeDirection.IN.equals(attribute.getRelationshipEdgeDirection())) {
+        if (properties.get(attribute.getName()) == null && typeCategory.equals(TypeCategory.ARRAY)) {
             return new ArrayList<>();
         }
 
@@ -1932,9 +1957,8 @@ public Object getVertexAttributePreFetchCache(AtlasVertex vertex, AtlasAttribute
             return null;
         }
 
-        // value is present as marker or is inward/outward relation, fetch the value from the vertex
-        if (properties.get(attribute.getName()) == StringUtils.SPACE || AtlasRelationshipEdgeDirection.IN.equals(attribute.getRelationshipEdgeDirection())
-                ||  AtlasRelationshipEdgeDirection.OUT.equals(attribute.getRelationshipEdgeDirection())) {
+        // value is present as marker , fetch the value from the vertex
+        if (properties.get(attribute.getName()) == StringUtils.SPACE) {
             return mapVertexToAttribute(vertex, attribute, null, false);
         }