Merge pull request #3646 from atlanhq/DG-1720

DG-1720 Support Product lineage with current lineage APIs
atlanhq · Oct 17, 2024 · 6cffb84 · 6cffb84
2 parents 0e9c5e4 + 853d5bb
commit 6cffb84
Show file tree

Hide file tree

Showing 7 changed files with 240 additions and 88 deletions.
diff --git a/common/src/main/java/org/apache/atlas/repository/Constants.java b/common/src/main/java/org/apache/atlas/repository/Constants.java
@@ -206,6 +206,7 @@ public final class Constants {
     public static final String PROCESS_INPUTS = "__Process.inputs";
 
     public static String[] PROCESS_EDGE_LABELS = {PROCESS_OUTPUTS, PROCESS_INPUTS};
+    public static final String CONNECTION_PROCESS_ENTITY_TYPE = "ConnectionProcess";
 
     /**
      * The homeId field is used when saving into Atlas a copy of an object that is being imported from another

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
@@ -13,6 +13,8 @@
 @JsonAutoDetect(getterVisibility = PUBLIC_ONLY, setterVisibility = PUBLIC_ONLY, fieldVisibility = NONE)
 @JsonSerialize(include = JsonSerialize.Inclusion.NON_NULL)
 public class LineageListRequest {
+    public static final String LINEAGE_TYPE_DATASET_PROCESS_LINEAGE = "DatasetProcessLineage";
+    public static final String LINEAGE_TYPE_PRODUCT_ASSET_LINEAGE = "ProductAssetLineage";
     private String                          guid;
     private Integer                         size;
     private Integer                         from;
@@ -35,6 +37,9 @@ public void setImmediateNeighbours(Boolean immediateNeighbours) {
         this.immediateNeighbours = immediateNeighbours;
     }
 
+
+    private String                         lineageType = LINEAGE_TYPE_DATASET_PROCESS_LINEAGE;
+
     public enum LineageDirection {INPUT, OUTPUT}
 
     public LineageListRequest() {
@@ -90,6 +95,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
@@ -11,6 +11,7 @@
 
 import static com.fasterxml.jackson.annotation.JsonAutoDetect.Visibility.NONE;
 import static com.fasterxml.jackson.annotation.JsonAutoDetect.Visibility.PUBLIC_ONLY;
+import static org.apache.atlas.model.lineage.LineageListRequest.LINEAGE_TYPE_DATASET_PROCESS_LINEAGE;
 
 @JsonAutoDetect(getterVisibility = PUBLIC_ONLY, setterVisibility = PUBLIC_ONLY, fieldVisibility = NONE)
 @JsonSerialize(include = JsonSerialize.Inclusion.NON_NULL)
@@ -22,6 +23,7 @@ public class LineageOnDemandRequest {
     private Set<String>                             attributes;
     private Set<String>                             relationAttributes;
     private LineageOnDemandBaseParams               defaultParams;
+    private String                                  lineageType = LINEAGE_TYPE_DATASET_PROCESS_LINEAGE;
 
     public LineageOnDemandRequest() {
         this.attributes = new HashSet<>();
@@ -64,6 +66,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
@@ -8,6 +8,8 @@
 import org.apache.commons.collections.Predicate;
 import java.util.Set;
 
+import static org.apache.atlas.model.lineage.LineageListRequest.LINEAGE_TYPE_DATASET_PROCESS_LINEAGE;
+
 public final class AtlasLineageListContext {
     private String                              guid;
     private int                                 size;
@@ -23,6 +25,7 @@ public final class AtlasLineageListContext {
     private int                                 currentEntityCounter;
     private boolean                             depthLimitReached;
     private boolean                             hasMoreUpdated;
+    private String                              lineageType = LINEAGE_TYPE_DATASET_PROCESS_LINEAGE;
     private Boolean                             immediateNeighbours;
 
     public AtlasLineageListContext(LineageListRequest lineageListRequest, AtlasTypeRegistry typeRegistry) {
@@ -35,6 +38,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();
         this.immediateNeighbours = lineageListRequest.getImmediateNeighbours();
     }
@@ -131,6 +135,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
@@ -14,6 +14,9 @@
 import java.util.Map;
 import java.util.Set;
 
+import static org.apache.atlas.model.lineage.LineageListRequest.LINEAGE_TYPE_DATASET_PROCESS_LINEAGE;
+
+
 public class AtlasLineageOnDemandContext {
     private static final Logger LOG = LoggerFactory.getLogger(AtlasLineageContext.class);
 
@@ -24,13 +27,17 @@ public class AtlasLineageOnDemandContext {
     private Set<String>                             relationAttributes;
     private LineageOnDemandBaseParams               defaultParams;
 
+
+    private String                                  lineageType = LINEAGE_TYPE_DATASET_PROCESS_LINEAGE;
+
     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 +63,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;