fix(ingestion/airflow-plugin): airflow remove old tasks

metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_listener.py

@@ -1,3 +1,4 @@
+import asyncio
 import copy
 import functools
 import logging
@@ -8,6 +9,7 @@
 
 import airflow
 import datahub.emitter.mce_builder as builder
+from airflow.models.serialized_dag import SerializedDagModel
 from datahub.api.entities.datajob import DataJob
 from datahub.api.entities.dataprocess.dataprocess_instance import InstanceRunResult
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
@@ -68,6 +70,7 @@ def hookimpl(f: _F) -> _F:  # type: ignore[misc] # noqa: F811
     "1",
 )
 _RUN_IN_THREAD_TIMEOUT = 30
+_DATAHUB_CLEANUP_DAG = "Datahub_Cleanup"
 
 
 def get_airflow_plugin_listener() -> Optional["DataHubListener"]:
@@ -541,6 +544,48 @@ def on_dag_start(self, dag_run: "DagRun") -> None:
 
             self.emitter.emit(event)
 
+        if dag.dag_id == _DATAHUB_CLEANUP_DAG:
+            assert self.graph
+
+            logger.debug("Initiating the cleanup of obsselete data from datahub")
+
+            ingested_dataflow_urns = list(
+                self.graph.get_urns_by_filter(
+                    platform="airflow", entity_types=["dataFlow"]
+                )
+            )
+            ingested_datajob_urns = list(
+                self.graph.get_urns_by_filter(
+                    platform="airflow", entity_types=["dataJob"]
+                )
+            )
+
+            all_airflow_dags = SerializedDagModel.read_all_dags().values()
+
+            airflow_flow_urns: List = []
+            airflow_job_urns: List = []
+
+            for dag in all_airflow_dags:
+                flow_urn = builder.make_data_flow_urn(
+                    orchestrator="airflow", flow_id=dag.dag_id
+                )
+                airflow_flow_urns.append(flow_urn)
+
+                for task in dag.tasks:
+                    airflow_job_urns.append(
+                        builder.make_data_job_urn_with_flow(str(flow_urn), task.task_id)
+                    )
+
+            obsolete_pipelines = set(ingested_dataflow_urns) - set(airflow_flow_urns)
+            obsolete_tasks = set(ingested_datajob_urns) - set(airflow_job_urns)
+
+            obsolete_urns = obsolete_pipelines.union(obsolete_tasks)
+
+            asyncio.run(self._soft_delete_obsolete_urns(obsolete_urns=obsolete_urns))
+
+            logger.debug(f"total pipelines removed = {len(obsolete_pipelines)}")
+            logger.debug(f"total tasks removed = {len(obsolete_tasks)}")
+
     if HAS_AIRFLOW_DAG_LISTENER_API:
 
         @hookimpl
@@ -577,3 +622,13 @@ def on_dataset_changed(self, dataset: "Dataset") -> None:
             logger.debug(
                 f"DataHub listener got notification about dataset change for {dataset}"
             )
+
+    async def _soft_delete_obsolete_urns(self, obsolete_urns):
+        delete_tasks = [self._delete_obsolete_data(urn) for urn in obsolete_urns]
+        await asyncio.gather(*delete_tasks)
+
+    async def _delete_obsolete_data(self, obsolete_urn):
+        assert self.graph
+
+        if self.graph.exists(str(obsolete_urn)):
+            self.graph.soft_delete_entity(str(obsolete_urn))