Fixes #27150: Bulk-fetch TaskInstances per DAG to eliminate N+1 in yield_pipeline_status

Author: RajdeepKushwaha5

ingestion/src/metadata/ingestion/source/pipeline/airflow/metadata.py

@@ -320,18 +320,15 @@ def get_pipeline_status(self, dag_id: str) -> List[DagRun]:
             return []
 
     def get_task_instances(
-        self, dag_id: str, run_id: str, serialized_tasks: List[AirflowTask]
-    ) -> List[OMTaskInstance]:
+        self, dag_id: str, run_ids: List[str], serialized_tasks: List[AirflowTask]
+    ) -> Dict[str, List[OMTaskInstance]]:
         """
-        We are building our own scoped TaskInstance
-        class to only focus on core properties required
-        by the metadata ingestion.
-
-        This makes the versioning more flexible on which Airflow
-        sources we support.
+        Fetch all TaskInstances for the given DAG and run IDs in a single query,
+        returning a dict keyed by run_id. This avoids an N+1 pattern where a
+        separate query was previously fired for each DagRun.
         """
-        task_instance_list = None
         serialized_tasks_ids = {task.task_id for task in serialized_tasks}
+        result: Dict[str, List[OMTaskInstance]] = defaultdict(list)
 
         try:
             task_instance_list = (
@@ -344,50 +341,59 @@ class to only focus on core properties required
                 )
                 .filter(
                     TaskInstance.dag_id == dag_id,
-                    TaskInstance.run_id == run_id,
+                    TaskInstance.run_id.in_(run_ids),
                     # updating old runs flag deleted tasks as `removed`
                     TaskInstance.state != AirflowTaskStatus.REMOVED.value,
                 )
                 .all()
             )
+            for elem in task_instance_list:
+                row = elem._asdict()
+                if row.get("task_id") in serialized_tasks_ids:
+                    result[row["run_id"]].append(
+                        OMTaskInstance(
+                            task_id=row.get("task_id"),
+                            state=row.get("state"),
+                            start_date=row.get("start_date"),
+                            end_date=row.get("end_date"),
+                        )
+                    )
         except Exception as exc:
             logger.debug(traceback.format_exc())
             logger.warning(
-                f"Tried to get TaskInstances with run_id. It might not be available in older Airflow versions - {exc}."
+                f"Tried to get TaskInstances for run_ids. The run_id column might not be available in older Airflow DB schemas - {exc}."
             )
 
-        task_instance_dict = (
-            [elem._asdict() for elem in task_instance_list]
-            if task_instance_list
-            else []
-        )
-
-        return [
-            OMTaskInstance(
-                task_id=elem.get("task_id"),
-                state=elem.get("state"),
-                start_date=elem.get("start_date"),
-                end_date=elem.get("end_date"),
-            )
-            for elem in task_instance_dict
-            if elem.get("task_id") in serialized_tasks_ids
-        ]
+        return result
 
     def yield_pipeline_status(
         self, pipeline_details: AirflowDagDetails
     ) -> Iterable[Either[OMetaPipelineStatus]]:
         try:
             dag_run_list = self.get_pipeline_status(pipeline_details.dag_id)
 
+            # Collect all run_ids up front so we can fetch all TaskInstances in
+            # one query instead of one per DagRun (N+1 avoidance).
+            run_ids = [
+                dag_run.run_id
+                for dag_run in dag_run_list
+                if dag_run.run_id and self.context.get().task_names
+            ]
+            tasks_by_run_id = (
+                self.get_task_instances(
+                    dag_id=pipeline_details.dag_id,
+                    run_ids=run_ids,
+                    serialized_tasks=pipeline_details.tasks,
+                )
+                if run_ids
+                else {}
+            )
+
             for dag_run in dag_run_list:
                 if (
                     dag_run.run_id and self.context.get().task_names
                 ):  # Airflow dags can have old task which are turned off/commented out in code
-                    tasks = self.get_task_instances(
-                        dag_id=dag_run.dag_id,
-                        run_id=dag_run.run_id,
-                        serialized_tasks=pipeline_details.tasks,
-                    )
+                    tasks = tasks_by_run_id.get(dag_run.run_id, [])
 
                     task_statuses = [
                         TaskStatus(

ingestion/tests/unit/topology/pipeline/test_airflow.py

@@ -31,6 +31,7 @@
 from metadata.ingestion.source.pipeline.airflow.models import (
     AirflowDag,
     AirflowDagDetails,
+    AirflowTask,
 )
 from metadata.ingestion.source.pipeline.airflow.utils import get_schedule_interval
 
@@ -998,3 +999,72 @@ def test_task_source_url_airflow_3x_format(self):
         assert f"/dags/{quote(dag_id)}/tasks/{quote(task_id)}" in task_url
         assert "/taskinstance/list/" not in task_url
         assert "_flt_3_dag_id=" not in task_url
+
+    def test_get_task_instances_bulk_query(self):
+        """
+        Verify that get_task_instances fires a single DB query for all run_ids
+        (no N+1 per DagRun) and groups the returned rows by run_id.
+        Tasks not present in serialized_tasks are excluded from the result.
+        """
+        from unittest.mock import MagicMock
+
+        serialized_tasks = [
+            AirflowTask(task_id="task_a"),
+            AirflowTask(task_id="task_b"),
+        ]
+
+        row_run1 = MagicMock()
+        row_run1._asdict.return_value = {
+            "task_id": "task_a",
+            "state": "success",
+            "start_date": None,
+            "end_date": None,
+            "run_id": "run_1",
+        }
+        row_run2 = MagicMock()
+        row_run2._asdict.return_value = {
+            "task_id": "task_b",
+            "state": "failed",
+            "start_date": None,
+            "end_date": None,
+            "run_id": "run_2",
+        }
+        unknown_task_row = MagicMock()
+        unknown_task_row._asdict.return_value = {
+            "task_id": "task_unknown",
+            "state": "success",
+            "start_date": None,
+            "end_date": None,
+            "run_id": "run_1",
+        }
+
+        mock_query = MagicMock()
+        mock_query.filter.return_value = mock_query
+        mock_query.all.return_value = [row_run1, row_run2, unknown_task_row]
+        mock_session = MagicMock()
+        mock_session.query.return_value = mock_query
+
+        original_session = getattr(self.airflow, "_session", None)
+        self.airflow._session = mock_session
+        try:
+            result = self.airflow.get_task_instances(
+                "my_dag", ["run_1", "run_2"], serialized_tasks
+            )
+        finally:
+            self.airflow._session = original_session
+
+        # Single DB query — not one per run_id
+        mock_session.query.assert_called_once()
+
+        # Results grouped correctly by run_id
+        self.assertIn("run_1", result)
+        self.assertIn("run_2", result)
+
+        # task_unknown is not in serialized_tasks so it must be excluded
+        self.assertEqual(len(result["run_1"]), 1)
+        self.assertEqual(result["run_1"][0].task_id, "task_a")
+        self.assertEqual(result["run_1"][0].state, "success")
+
+        self.assertEqual(len(result["run_2"]), 1)
+        self.assertEqual(result["run_2"][0].task_id, "task_b")
+        self.assertEqual(result["run_2"][0].state, "failed")