Merge pull request #3377 from atlanhq/LIN-1022-beta

LIN-1022 : [add] add relationship attributes search in lineage list API
atlanhq · Aug 1, 2024 · 5ce9e98 · 5ce9e98
2 parents 2e04501 + 6a3e4b6
commit 5ce9e98
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,6 +22,7 @@ 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;
 
@@ -32,11 +33,12 @@ 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;
@@ -48,6 +50,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() {
@@ -129,6 +132,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;
@@ -35,6 +36,7 @@ public AtlasLineageListContext(LineageListRequest lineageListRequest, AtlasTypeR
         this.edgeTraversalPredicate = constructInMemoryPredicate(typeRegistry, lineageListRequest.getRelationshipTraversalFilters());
         this.attributes = lineageListRequest.getAttributes();
         this.lineageType = lineageListRequest.getLineageType();
+        this.relationAttributes = lineageListRequest.getRelationAttributes();
     }
 
     public String getGuid() {
@@ -113,6 +115,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
@@ -209,6 +209,8 @@ public AtlasLineageListInfo getLineageListInfoOnDemand(String guid, LineageListR
         RequestContext.get().setLineageInputLabel(LINEAGE_MAP.get(lineageListRequest.getLineageType())[0]);
         RequestContext.get().setLineageOutputLabel(LINEAGE_MAP.get(lineageListRequest.getLineageType())[1]);
         AtlasLineageListInfo ret = new AtlasLineageListInfo(new ArrayList<>());
+        RequestContext.get().setRelationAttrsForSearch(lineageListRequest.getRelationAttributes());
+
         traverseEdgesUsingBFS(guid, new AtlasLineageListContext(lineageListRequest, atlasTypeRegistry), ret);
         ret.setSearchParameters(lineageListRequest);