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

LIN-1022 : [add] add relationship attributes search in lineage list API

Author: akshaysw

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;
     }

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;
     }

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);