Merge pull request #3378 from atlanhq/LIN-1022

LIN-1022 : [add] add relationship attributes search in lineage list API
atlanhq · Aug 2, 2024 · 94c6f5e · 94c6f5e
2 parents e154fe8 + fcafe07
commit 94c6f5e
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 1 deletion.
diff --git a/intg/src/main/java/org/apache/atlas/model/lineage/LineageListRequest.java b/intg/src/main/java/org/apache/atlas/model/lineage/LineageListRequest.java
@@ -22,18 +22,20 @@ public class LineageListRequest {
     private SearchParameters.FilterCriteria entityTraversalFilters;
     private SearchParameters.FilterCriteria relationshipTraversalFilters;
     private Set<String>                     attributes;
+    private Set<String>                     relationAttributes;
     private Boolean                         excludeMeanings;
     private Boolean                         excludeClassifications;
 
     public enum LineageDirection {INPUT, OUTPUT}
 
     public LineageListRequest() {
         this.attributes = new HashSet<>();
+        this.relationAttributes = new HashSet<>();
     }
 
     public LineageListRequest(String guid, Integer size, Integer from, Integer depth, LineageDirection direction, SearchParameters.FilterCriteria entityFilters,
                               SearchParameters.FilterCriteria entityTraversalFilters, SearchParameters.FilterCriteria relationshipTraversalFilters,
-                              Set<String> attributes, boolean excludeMeanings, boolean excludeClassifications) {
+                              Set<String> attributes, boolean excludeMeanings, boolean excludeClassifications, Set<String> relationAttributes) {
         this.guid = guid;
         this.size = size;
         this.from = from;
@@ -45,6 +47,7 @@ public LineageListRequest(String guid, Integer size, Integer from, Integer depth
         this.attributes = attributes;
         this.excludeMeanings = excludeMeanings;
         this.excludeClassifications = excludeClassifications;
+        this.relationAttributes = relationAttributes;
     }
 
     public String getGuid() {
@@ -119,6 +122,14 @@ public void setAttributes(Set<String> attributes) {
         this.attributes = attributes;
     }
 
+    public Set<String> getRelationAttributes() {
+        return relationAttributes;
+    }
+
+    public void setRelationAttributes(Set<String> relationAttributes) {
+        this.relationAttributes = relationAttributes;
+    }
+
     public Boolean isExcludeMeanings() {
         return excludeMeanings;
     }

diff --git a/repository/src/main/java/org/apache/atlas/discovery/AtlasLineageListContext.java b/repository/src/main/java/org/apache/atlas/discovery/AtlasLineageListContext.java
@@ -18,6 +18,7 @@ public final class AtlasLineageListContext {
     private Predicate                           vertexTraversalPredicate;
     private Predicate                           edgeTraversalPredicate;
     private Set<String>                         attributes;
+    private Set<String>                         relationAttributes;
     private int                                 currentFromCounter;
     private int                                 currentEntityCounter;
     private boolean                             depthLimitReached;
@@ -33,6 +34,7 @@ public AtlasLineageListContext(LineageListRequest lineageListRequest, AtlasTypeR
         this.vertexTraversalPredicate = constructInMemoryPredicate(typeRegistry, lineageListRequest.getEntityTraversalFilters());
         this.edgeTraversalPredicate = constructInMemoryPredicate(typeRegistry, lineageListRequest.getRelationshipTraversalFilters());
         this.attributes = lineageListRequest.getAttributes();
+        this.relationAttributes = lineageListRequest.getRelationAttributes();
     }
 
     public String getGuid() {
@@ -111,6 +113,14 @@ public void setAttributes(Set<String> attributes) {
         this.attributes = attributes;
     }
 
+    public Set<String> getRelationAttributes() {
+        return relationAttributes;
+    }
+
+    public void setRelationAttributes(Set<String> relationAttributes) {
+        this.relationAttributes = relationAttributes;
+    }
+
     public int getCurrentFromCounter() {
         return currentFromCounter;
     }

diff --git a/repository/src/main/java/org/apache/atlas/discovery/EntityLineageService.java b/repository/src/main/java/org/apache/atlas/discovery/EntityLineageService.java
@@ -193,6 +193,8 @@ public AtlasLineageListInfo getLineageListInfoOnDemand(String guid, LineageListR
         AtlasPerfMetrics.MetricRecorder metricRecorder = RequestContext.get().startMetricRecord("getLineageListInfoOnDemand");
 
         AtlasLineageListInfo ret = new AtlasLineageListInfo(new ArrayList<>());
+        RequestContext.get().setRelationAttrsForSearch(lineageListRequest.getRelationAttributes());
+
         traverseEdgesUsingBFS(guid, new AtlasLineageListContext(lineageListRequest, atlasTypeRegistry), ret);
         ret.setSearchParameters(lineageListRequest);