Improved Code visuals etc.

atlanhq · Aug 21, 2024 · b040cc4 · b040cc4
1 parent 09c35ea
commit b040cc4
Show file tree

Hide file tree

Showing 7 changed files with 128 additions and 75 deletions.
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
@@ -26,6 +26,9 @@ public class LineageListRequest {
     private Boolean                         excludeMeanings;
     private Boolean                         excludeClassifications;
 
+
+    private String                         lineageType = "DatasetProcessLineage";
+
     public enum LineageDirection {INPUT, OUTPUT}
 
     public LineageListRequest() {
@@ -81,6 +84,13 @@ public Integer getDepth() {
     public void setDepth(Integer depth) {
         this.depth = depth;
     }
+    public String getLineageType() {
+        return lineageType;
+    }
+
+    public void setLineageType(String lineageType) {
+        this.lineageType = lineageType;
+    }
 
     public LineageDirection getDirection() {
         return direction;

diff --git a/intg/src/main/java/org/apache/atlas/model/lineage/LineageOnDemandRequest.java b/intg/src/main/java/org/apache/atlas/model/lineage/LineageOnDemandRequest.java
@@ -22,6 +22,7 @@ public class LineageOnDemandRequest {
     private Set<String>                             attributes;
     private Set<String>                             relationAttributes;
     private LineageOnDemandBaseParams               defaultParams;
+    private String                                  lineageType = "DatasetProcessLineage";
 
     public LineageOnDemandRequest() {
         this.attributes = new HashSet<>();
@@ -64,6 +65,13 @@ public void setRelationshipTraversalFilters(SearchParameters.FilterCriteria rela
         this.relationshipTraversalFilters = relationshipTraversalFilters;
     }
 
+    public String getLineageType() {
+        return lineageType;
+    }
+
+    public void setLineageType(String lineageType) {
+        this.lineageType = lineageType;
+    }
     public Set<String> getAttributes() {
         return attributes;
     }

diff --git a/repository/src/main/java/org/apache/atlas/discovery/AtlasLineageListContext.java b/repository/src/main/java/org/apache/atlas/discovery/AtlasLineageListContext.java
@@ -23,6 +23,7 @@ public final class AtlasLineageListContext {
     private int                                 currentEntityCounter;
     private boolean                             depthLimitReached;
     private boolean                             hasMoreUpdated;
+    private String                              lineageType = "DatasetProcessLineage";
 
     public AtlasLineageListContext(LineageListRequest lineageListRequest, AtlasTypeRegistry typeRegistry) {
         this.guid = lineageListRequest.getGuid();
@@ -34,6 +35,7 @@ public AtlasLineageListContext(LineageListRequest lineageListRequest, AtlasTypeR
         this.vertexTraversalPredicate = constructInMemoryPredicate(typeRegistry, lineageListRequest.getEntityTraversalFilters());
         this.edgeTraversalPredicate = constructInMemoryPredicate(typeRegistry, lineageListRequest.getRelationshipTraversalFilters());
         this.attributes = lineageListRequest.getAttributes();
+        this.lineageType = lineageListRequest.getLineageType();
         this.relationAttributes = lineageListRequest.getRelationAttributes();
     }
 
@@ -129,6 +131,14 @@ public void setCurrentFromCounter(int currentFromCounter) {
         this.currentFromCounter = currentFromCounter;
     }
 
+    public String getLineageType() {
+        return lineageType;
+    }
+
+    public void setLineageType(String lineageType) {
+        this.lineageType = lineageType;
+    }
+
     public int getCurrentEntityCounter() {
         return currentEntityCounter;
     }

diff --git a/repository/src/main/java/org/apache/atlas/discovery/AtlasLineageOnDemandContext.java b/repository/src/main/java/org/apache/atlas/discovery/AtlasLineageOnDemandContext.java
@@ -24,13 +24,17 @@ public class AtlasLineageOnDemandContext {
     private Set<String>                             relationAttributes;
     private LineageOnDemandBaseParams               defaultParams;
 
+
+    private String                                  lineageType = "DatasetProcessLineage";
+
     public AtlasLineageOnDemandContext(LineageOnDemandRequest lineageOnDemandRequest, AtlasTypeRegistry typeRegistry) {
         this.constraints        = lineageOnDemandRequest.getConstraints();
         this.attributes         = lineageOnDemandRequest.getAttributes();
         this.relationAttributes = lineageOnDemandRequest.getRelationAttributes();
         this.defaultParams      = lineageOnDemandRequest.getDefaultParams();
         this.vertexPredicate    = constructInMemoryPredicate(typeRegistry, lineageOnDemandRequest.getEntityTraversalFilters());
         this.edgePredicate      = constructInMemoryPredicate(typeRegistry, lineageOnDemandRequest.getRelationshipTraversalFilters());
+        this.lineageType        = lineageOnDemandRequest.getLineageType();
     }
 
     public Map<String, LineageOnDemandConstraints> getConstraints() {
@@ -56,6 +60,13 @@ public Predicate getEdgePredicate() {
     public void setEdgePredicate(Predicate edgePredicate) {
         this.edgePredicate = edgePredicate;
     }
+    public String getLineageType() {
+        return lineageType;
+    }
+
+    public void setLineageType(String lineageType) {
+        this.lineageType = lineageType;
+    }
 
     public Set<String> getAttributes() {
         return attributes;