From 59eae6831f2afcf15ffe8e79ebddc53fce53ad2f Mon Sep 17 00:00:00 2001 From: Eason09053360 <185830721+Eason09053360@users.noreply.github.com> Date: Sat, 12 Sep 2026 22:36:43 +0800 Subject: [PATCH] Speed up airflow dags list on deployments with many Dags Listing Dags issued two extra queries for every Dag printed: one to load the DagModel row, and another when the response model reached the lazily loaded tags relationship. The command therefore got slower in proportion to the number of Dags, which is exactly the case it matters most for. The batched lookup is chunked rather than issued as a single IN clause so that installations with tens of thousands of Dags stay below the backend's bind parameter cap, where one oversized statement would fail outright. --- .../src/airflow/cli/commands/dag_command.py | 42 ++++++++++++------ .../unit/cli/commands/test_dag_command.py | 44 +++++++++++++++++++ 2 files changed, 73 insertions(+), 13 deletions(-) diff --git a/airflow-core/src/airflow/cli/commands/dag_command.py b/airflow-core/src/airflow/cli/commands/dag_command.py index bd42dadf925fc..fbb5b24b7b040 100644 --- a/airflow-core/src/airflow/cli/commands/dag_command.py +++ b/airflow-core/src/airflow/cli/commands/dag_command.py @@ -31,6 +31,7 @@ from typing import TYPE_CHECKING, cast from sqlalchemy import func, or_, select +from sqlalchemy.orm import selectinload from airflow._shared.timezones import timezone from airflow.api.client import get_current_api_client @@ -79,6 +80,9 @@ # Chunk size for bulk delete. _RUN_CHUNK_SIZE = 500 +# Chunk size for looking up Dags by id. +_DAG_CHUNK_SIZE = 500 + @deprecated_for_airflowctl("airflowctl dags trigger") @cli_utils.action_cli @@ -599,15 +603,6 @@ def dag_list_dags(args, *, session: Session = NEW_SESSION) -> None: file=sys.stderr, ) - def get_dag_detail(dag: DAG) -> dict: - if dag_model := DagModel.get_dagmodel(dag.dag_id, session=session): - dag_detail = DAGResponse.model_validate(dag_model, from_attributes=True).model_dump() - else: - dag_detail = _get_dagbag_dag_details(dag) - if not cols: - return dag_detail - return {col: dag_detail[col] for col in cols if col in DAG_DETAIL_FIELDS} - def filter_dags_by_bundle(dags: Iterable[DAG], bundle_names: list[str] | None) -> Iterable[DAG]: """Filter DAGs based on the specified bundle name, if provided.""" if not bundle_names: @@ -619,11 +614,32 @@ def filter_dags_by_bundle(dags: Iterable[DAG], bundle_names: list[str] | None) - ) return (dag for dag in dags if dag.dag_id in selected_dag_ids) + dags_to_show = sorted( + filter_dags_by_bundle(dags_list, args.bundle_name if not args.local else None), + key=operator.attrgetter("dag_id"), + ) + # Chunked so the IN clause stays under the backend's bind parameter cap on deployments + # with many Dags (SQLite allows 32766, PostgreSQL 65535). + dag_models: dict[str, DagModel] = {} + for dag_id_chunk in chunks([dag.dag_id for dag in dags_to_show], _DAG_CHUNK_SIZE): + dag_models.update( + (dag_model.dag_id, dag_model) + for dag_model in session.scalars( + select(DagModel).options(selectinload(DagModel.tags)).where(DagModel.dag_id.in_(dag_id_chunk)) + ) + ) + + def get_dag_detail(dag: DAG) -> dict: + if dag_model := dag_models.get(dag.dag_id): + dag_detail = DAGResponse.model_validate(dag_model, from_attributes=True).model_dump() + else: + dag_detail = _get_dagbag_dag_details(dag) + if not cols: + return dag_detail + return {col: dag_detail[col] for col in cols if col in DAG_DETAIL_FIELDS} + AirflowConsole().print_as( - data=sorted( - filter_dags_by_bundle(dags_list, args.bundle_name if not args.local else None), - key=operator.attrgetter("dag_id"), - ), + data=dags_to_show, output=args.output, mapper=get_dag_detail, ) diff --git a/airflow-core/tests/unit/cli/commands/test_dag_command.py b/airflow-core/tests/unit/cli/commands/test_dag_command.py index fcf819910e612..1c91beee7578d 100644 --- a/airflow-core/tests/unit/cli/commands/test_dag_command.py +++ b/airflow-core/tests/unit/cli/commands/test_dag_command.py @@ -56,6 +56,7 @@ from airflow.utils.state import DagRunState, DagSchedulingState, TaskInstanceState from airflow.utils.types import DagRunType +from tests_common.test_utils.asserts import count_queries from tests_common.test_utils.config import conf_vars from tests_common.test_utils.dag import sync_dag_to_db from tests_common.test_utils.db import ( @@ -64,6 +65,7 @@ clear_db_runs, parse_and_sync_to_db, ) +from tests_common.test_utils.stream_capture_manager import StdoutCaptureManager from unit.models import TEST_DAGS_FOLDER DEFAULT_DATE = timezone.make_aware(datetime(2015, 1, 1), timezone=timezone.utc) @@ -384,6 +386,48 @@ def test_cli_list_dags_with_multiple_dag_versions(self, dag_maker, stdout_captur # Rebuild Test DB for other tests self.setup_class() + def test_cli_list_dags_does_not_query_per_dag(self): + def count_list_dags_queries(dag_count: int) -> int: + clear_db_dags() + for i in range(dag_count): + with DAG( + f"test_query_count_{i}", schedule=None, start_date=DEFAULT_DATE, tags=["a", "b"] + ) as dag: + EmptyOperator(task_id="task1") + sync_dag_to_db(dag) + + args = self.parser.parse_args(["dags", "list", "--output", "json"]) + with count_queries() as result, StdoutCaptureManager() as temp_stdout: + dag_command.dag_list_dags(args) + assert len(json.loads(temp_stdout.getvalue())) == dag_count + return sum(result.values()) + + assert count_list_dags_queries(2) == count_list_dags_queries(6) + + # Rebuild Test DB for other tests + self.setup_class() + + @mock.patch("airflow.cli.commands.dag_command._DAG_CHUNK_SIZE", 2) + def test_cli_list_dags_lists_every_dag_when_lookup_is_chunked(self, stdout_capture): + clear_db_dags() + for i in range(5): + with DAG(f"test_chunked_{i}", schedule=None, start_date=DEFAULT_DATE) as dag: + EmptyOperator(task_id="task1") + sync_dag_to_db(dag) + + args = self.parser.parse_args(["dags", "list", "--columns", "dag_id,is_paused", "--output", "json"]) + with stdout_capture as temp_stdout: + dag_command.dag_list_dags(args) + + rows = json.loads(temp_stdout.getvalue()) + assert [row["dag_id"] for row in rows] == [f"test_chunked_{i}" for i in range(5)] + # Only DagModel carries is_paused; the dagbag fallback leaves it None, so this catches a + # chunked lookup that drops Dags from the prefetch map. + assert all(row["is_paused"] is not None for row in rows) + + # Rebuild Test DB for other tests + self.setup_class() + def test_cli_list_local_dags(self, stdout_capture): # Clear the database clear_db_dags()