From 1ee5a1d1133e402828b20ad5315a59823fec2d9e Mon Sep 17 00:00:00 2001 From: ColtenOuO Date: Thu, 20 Aug 2026 18:56:48 +0000 Subject: [PATCH 1/2] Reduce per-task DB round trips in the asset dependency graph API get_data_dependencies() already batches Dag entry-point resolution into one query per BFS round, but still issued one inlet/outlet TaskInletAssetReference/TaskOutletAssetReference query per individual task. An asset shared by many Dags/tasks (common in multi-team deployments) turned opening its dependency graph into hundreds of sequential round trips within a single request. Collect the task keys needing an inlet or outlet lookup across the whole round and resolve each with one batched query, the same tuple_(...).in_(...) pattern already used for entry-point resolution and elsewhere in the API layer. --- .../core_api/services/ui/dependencies.py | 56 ++++++++++--------- .../core_api/routes/ui/test_dependencies.py | 34 +++++++++++ 2 files changed, 63 insertions(+), 27 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/core_api/services/ui/dependencies.py b/airflow-core/src/airflow/api_fastapi/core_api/services/ui/dependencies.py index b519a61def3b7..b48a248050284 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/services/ui/dependencies.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/services/ui/dependencies.py @@ -394,7 +394,7 @@ def get_data_dependencies( asset_id: int, session: Session, readable_dag_ids: set[str] | None = None ) -> dict[str, list[dict]]: """Get full task dependencies for an asset.""" - from sqlalchemy import select, union_all + from sqlalchemy import select, tuple_, union_all from sqlalchemy.orm import selectinload from airflow.models.asset import ( @@ -458,6 +458,8 @@ def get_data_dependencies( next_frontier: set[int] = set() pending_dag_ids: set[str] = set() triggering_asset_node_ids_by_dag: dict[str, set[str]] = defaultdict(set) + pending_inlet_keys: set[tuple[str, str]] = set() + pending_outlet_keys: set[tuple[str, str]] = set() for asset in assets: asset_node_id = f"asset:{asset.id}" @@ -491,15 +493,7 @@ def get_data_dependencies( # Find other assets this task consumes (inlets) to trace upstream if task_key not in processed_tasks: processed_tasks.add(task_key) - inlet_refs = session.scalars( - select(TaskInletAssetReference).where( - TaskInletAssetReference.dag_id == ref.dag_id, - TaskInletAssetReference.task_id == ref.task_id, - ) - ).all() - for inlet_ref in inlet_refs: - if inlet_ref.asset_id not in processed_assets: - next_frontier.add(inlet_ref.asset_id) + pending_inlet_keys.add(task_key) # Process consuming tasks (tasks that input this asset) for ref in asset.consuming_tasks: @@ -523,15 +517,7 @@ def get_data_dependencies( # Find other assets this task produces (outlets) to trace downstream if task_key not in processed_tasks: processed_tasks.add(task_key) - outlet_refs = session.scalars( - select(TaskOutletAssetReference).where( - TaskOutletAssetReference.dag_id == ref.dag_id, - TaskOutletAssetReference.task_id == ref.task_id, - ) - ).all() - for outlet_ref in outlet_refs: - if outlet_ref.asset_id not in processed_assets: - next_frontier.add(outlet_ref.asset_id) + pending_outlet_keys.add(task_key) # Process Dags scheduled by this asset at the Dag level (`schedule=...`). These have # no task-level inlet reference and would otherwise be a dead end in this graph -- @@ -599,15 +585,31 @@ def get_data_dependencies( # Find other assets this entry task produces (outlets) to trace downstream. if task_key not in processed_tasks: processed_tasks.add(task_key) - outlet_refs = session.scalars( - select(TaskOutletAssetReference).where( - TaskOutletAssetReference.dag_id == dag_id, - TaskOutletAssetReference.task_id == entry_task_id, + pending_outlet_keys.add(task_key) + + if pending_inlet_keys: + inlet_refs = session.scalars( + select(TaskInletAssetReference).where( + tuple_(TaskInletAssetReference.dag_id, TaskInletAssetReference.task_id).in_( + pending_inlet_keys + ) + ) + ).all() + for inlet_ref in inlet_refs: + if inlet_ref.asset_id not in processed_assets: + next_frontier.add(inlet_ref.asset_id) + + if pending_outlet_keys: + outlet_refs = session.scalars( + select(TaskOutletAssetReference).where( + tuple_(TaskOutletAssetReference.dag_id, TaskOutletAssetReference.task_id).in_( + pending_outlet_keys ) - ).all() - for outlet_ref in outlet_refs: - if outlet_ref.asset_id not in processed_assets: - next_frontier.add(outlet_ref.asset_id) + ) + ).all() + for outlet_ref in outlet_refs: + if outlet_ref.asset_id not in processed_assets: + next_frontier.add(outlet_ref.asset_id) frontier = next_frontier diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_dependencies.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_dependencies.py index 3c266f856d9a1..05dee90101aeb 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_dependencies.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_dependencies.py @@ -868,6 +868,40 @@ def test_data_dependencies_batches_entry_point_resolution_across_scheduled_dags( entry_task_node_id = f"task:{dag_id}__SEPARATOR__entry_task" assert nodes_by_id[entry_task_node_id]["type"] == "task" + def test_data_dependencies_batches_inlet_outlet_resolution_across_producing_tasks( + self, dag_maker, test_client, session + ): + asset_a = Asset(uri="s3://data-dep-bucket/inlet-batch-a", name="data_dep_inlet_batch_asset_a") + dag_ids = [f"data_dep_inlet_batch_producer_{i}" for i in range(3)] + + for dag_id in dag_ids: + with dag_maker( + dag_id=dag_id, + schedule=None, + serialized=True, + session=session, + ): + EmptyOperator(task_id="produce", outlets=[asset_a]) + + dag_maker.sync_dagbag_to_db() + + asset_a_id = session.scalar( + select(AssetModel.id).where(AssetModel.name == "data_dep_inlet_batch_asset_a") + ) + + expected_query_count = 9 + with assert_queries_count(expected_query_count): + response = test_client.get( + "/dependencies", params={"node_id": f"asset:{asset_a_id}", "dependency_type": "data"} + ) + assert response.status_code == 200 + + result = response.json() + nodes_by_id = {node["id"]: node for node in result["nodes"]} + for dag_id in dag_ids: + task_node_id = f"task:{dag_id}__SEPARATOR__produce" + assert nodes_by_id[task_node_id]["type"] == "task" + def test_data_dependencies_attributes_alias_produced_asset_to_source_task( self, dag_maker, test_client, session ): From 34a8c84e9a8835004a136df8da635a7c1a5546c4 Mon Sep 17 00:00:00 2001 From: ColtenOuO Date: Sat, 29 Aug 2026 19:40:55 +0000 Subject: [PATCH 2/2] Add outlet-resolution batching test for the asset dependency graph The new regression test only exercised the inlet-lookup batching path, leaving the outlet-lookup path (traced from consuming tasks) uncovered. --- .../core_api/routes/ui/test_dependencies.py | 40 ++++++++++++++++++- 1 file changed, 39 insertions(+), 1 deletion(-) diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_dependencies.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_dependencies.py index 05dee90101aeb..426c11134d011 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_dependencies.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_dependencies.py @@ -868,9 +868,11 @@ def test_data_dependencies_batches_entry_point_resolution_across_scheduled_dags( entry_task_node_id = f"task:{dag_id}__SEPARATOR__entry_task" assert nodes_by_id[entry_task_node_id]["type"] == "task" - def test_data_dependencies_batches_inlet_outlet_resolution_across_producing_tasks( + def test_data_dependencies_batches_inlet_resolution_across_producing_tasks( self, dag_maker, test_client, session ): + """Producing tasks trace upstream through their own inlets; that inlet lookup must be + batched across all producing tasks in a round, not issued once per task.""" asset_a = Asset(uri="s3://data-dep-bucket/inlet-batch-a", name="data_dep_inlet_batch_asset_a") dag_ids = [f"data_dep_inlet_batch_producer_{i}" for i in range(3)] @@ -902,6 +904,42 @@ def test_data_dependencies_batches_inlet_outlet_resolution_across_producing_task task_node_id = f"task:{dag_id}__SEPARATOR__produce" assert nodes_by_id[task_node_id]["type"] == "task" + def test_data_dependencies_batches_outlet_resolution_across_consuming_tasks( + self, dag_maker, test_client, session + ): + """Consuming tasks trace downstream through their own outlets; that outlet lookup must be + batched across all consuming tasks in a round, not issued once per task.""" + asset_a = Asset(uri="s3://data-dep-bucket/outlet-batch-a", name="data_dep_outlet_batch_asset_a") + dag_ids = [f"data_dep_outlet_batch_consumer_{i}" for i in range(3)] + + for dag_id in dag_ids: + with dag_maker( + dag_id=dag_id, + schedule=None, + serialized=True, + session=session, + ): + EmptyOperator(task_id="consume", inlets=[asset_a]) + + dag_maker.sync_dagbag_to_db() + + asset_a_id = session.scalar( + select(AssetModel.id).where(AssetModel.name == "data_dep_outlet_batch_asset_a") + ) + + expected_query_count = 9 + with assert_queries_count(expected_query_count): + response = test_client.get( + "/dependencies", params={"node_id": f"asset:{asset_a_id}", "dependency_type": "data"} + ) + assert response.status_code == 200 + + result = response.json() + nodes_by_id = {node["id"]: node for node in result["nodes"]} + for dag_id in dag_ids: + task_node_id = f"task:{dag_id}__SEPARATOR__consume" + assert nodes_by_id[task_node_id]["type"] == "task" + def test_data_dependencies_attributes_alias_produced_asset_to_source_task( self, dag_maker, test_client, session ):