From 3bd5f61f9081c105052a7ab3522f54c2cc06f21d Mon Sep 17 00:00:00 2001
From: mohitdeuex <mohit.y@deuexsolutions.com>
Date: Wed, 18 Dec 2024 19:05:38 +0530
Subject: [PATCH] Cleanup lineage on pipeline and store procedure removal

---
 .../openmetadata/service/jdbi3/CollectionDAO.java  |  5 +----
 .../service/jdbi3/LineageRepository.java           |  2 +-
 .../service/jdbi3/PipelineRepository.java          | 14 ++++++++++++++
 .../service/jdbi3/StoredProcedureRepository.java   | 13 +++++++++++++
 4 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/CollectionDAO.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/CollectionDAO.java
index a423db901cd3..5926e5413af3 100644
--- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/CollectionDAO.java
+++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/CollectionDAO.java
@@ -1180,10 +1180,7 @@ void deleteLineageBySource(
                 + "AND json->>'source' = :source",
         connectionType = POSTGRES)
     void deleteLineageBySourcePipeline(
-        @BindUUID("toId") UUID toId,
-        @Bind("toEntity") String toEntity,
-        @Bind("source") String source,
-        @Bind("relation") int relation);
+        @BindUUID("toId") UUID toId, @Bind("source") String source, @Bind("relation") int relation);
 
     class FromRelationshipMapper implements RowMapper<EntityRelationshipRecord> {
       @Override
diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/LineageRepository.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/LineageRepository.java
index 134fe998ec15..9026c4f36daf 100644
--- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/LineageRepository.java
+++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/LineageRepository.java
@@ -633,7 +633,7 @@ public void deleteLineageBySource(UUID toId, String toEntity, String source) {
               .findLineageBySourcePipeline(toId, toEntity, source, Relationship.UPSTREAM.ordinal());
       // Finally, delete lineage relationship
       dao.relationshipDAO()
-          .deleteLineageBySourcePipeline(toId, toEntity, source, Relationship.UPSTREAM.ordinal());
+          .deleteLineageBySourcePipeline(toId, toEntity, Relationship.UPSTREAM.ordinal());
     } else {
       relations =
           dao.relationshipDAO()
diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/PipelineRepository.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/PipelineRepository.java
index 8aaf4db08d8f..bab495ad27d0 100644
--- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/PipelineRepository.java
+++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/PipelineRepository.java
@@ -42,6 +42,8 @@
 import org.openmetadata.schema.type.EntityReference;
 import org.openmetadata.schema.type.FieldChange;
 import org.openmetadata.schema.type.Include;
+import org.openmetadata.schema.type.LineageDetails;
+import org.openmetadata.schema.type.Relationship;
 import org.openmetadata.schema.type.Status;
 import org.openmetadata.schema.type.TagLabel;
 import org.openmetadata.schema.type.Task;
@@ -296,6 +298,18 @@ public void storeEntity(Pipeline pipeline, boolean update) {
     pipeline.withService(service).withTasks(taskWithTagsAndOwners);
   }
 
+  @Override
+  protected void cleanup(Pipeline pipeline) {
+    // When a pipeline is removed , the linege needs to be removed
+    daoCollection
+        .relationshipDAO()
+        .deleteLineageBySourcePipeline(
+            pipeline.getId(),
+            LineageDetails.Source.PIPELINE_LINEAGE.value(),
+            Relationship.UPSTREAM.ordinal());
+    super.cleanup(pipeline);
+  }
+
   @Override
   public void storeRelationships(Pipeline pipeline) {
     addServiceRelationship(pipeline, pipeline.getService());
diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/StoredProcedureRepository.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/StoredProcedureRepository.java
index 60ad7db78d23..11624a475808 100644
--- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/StoredProcedureRepository.java
+++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/StoredProcedureRepository.java
@@ -11,6 +11,7 @@
 import org.openmetadata.schema.entity.data.StoredProcedure;
 import org.openmetadata.schema.type.EntityReference;
 import org.openmetadata.schema.type.Include;
+import org.openmetadata.schema.type.LineageDetails;
 import org.openmetadata.schema.type.Relationship;
 import org.openmetadata.service.Entity;
 import org.openmetadata.service.resources.databases.StoredProcedureResource;
@@ -69,6 +70,18 @@ public void storeRelationships(StoredProcedure storedProcedure) {
         Relationship.CONTAINS);
   }
 
+  @Override
+  protected void cleanup(StoredProcedure storedProcedure) {
+    // When a pipeline is removed , the linege needs to be removed
+    daoCollection
+        .relationshipDAO()
+        .deleteLineageBySourcePipeline(
+            storedProcedure.getId(),
+            LineageDetails.Source.QUERY_LINEAGE.value(),
+            Relationship.UPSTREAM.ordinal());
+    super.cleanup(storedProcedure);
+  }
+
   @Override
   public void setInheritedFields(StoredProcedure storedProcedure, EntityUtil.Fields fields) {
     DatabaseSchema schema =