Speed up airflow dags list on deployments with many Dags

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,
     )

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()