Merge pull request #3371 from atlanhq/beta-DG-1720-helper

Improved Code visuals etc.
atlanhq · Jul 30, 2024 · 0b51ebb · 0b51ebb
2 parents 4f3bdc6 + 53f8576
commit 0b51ebb
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
@@ -25,6 +25,9 @@ public class LineageListRequest {
     private Boolean                         excludeMeanings;
     private Boolean                         excludeClassifications;
 
+
+    private String                         lineageType = "DatasetProcessLineage";
+
     public enum LineageDirection {INPUT, OUTPUT}
 
     public LineageListRequest() {
@@ -78,6 +81,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
@@ -22,6 +22,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();
@@ -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.lineageType = lineageListRequest.getLineageType();
     }
 
     public String getGuid() {
@@ -119,6 +121,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;