From 28b2c96615d168f83e30dbbc69bd7a56875fe0b4 Mon Sep 17 00:00:00 2001
From: Henry Chen <henryhenry0512@gmail.com>
Date: Fri, 3 Apr 2026 01:58:44 +0800
Subject: [PATCH 1/3] improve grid/ti_summaries and grid/runs (#64034)

* improve grid/ti_summaries and grid/runs

* remove serdag
---
 .../api_fastapi/core_api/routes/ui/grid.py    | 100 +++++++----
 .../api_fastapi/core_api/services/ui/grid.py  | 170 ++++++++++++------
 .../core_api/routes/ui/test_grid.py           |  23 +++
 3 files changed, 196 insertions(+), 97 deletions(-)

diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/grid.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/grid.py
index 69d5f60785e9b..ebb13bd2be704 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/grid.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/grid.py
@@ -17,8 +17,7 @@
 
 from __future__ import annotations
 
-import collections
-from collections.abc import Generator, Sequence
+from collections.abc import Generator, Iterable
 from typing import TYPE_CHECKING, Annotated, Any
 from uuid import UUID
 
@@ -26,11 +25,12 @@
 from fastapi import Depends, HTTPException, Query, status
 from fastapi.responses import StreamingResponse
 from sqlalchemy import exists, select
-from sqlalchemy.orm import Session, joinedload, load_only, selectinload
+from sqlalchemy.orm import Session, joinedload, load_only
 
 from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity
 from airflow.api_fastapi.common.dagbag import DagBagDep
 from airflow.api_fastapi.common.db.common import SessionDep, paginated_select
+from airflow.api_fastapi.common.db.dag_runs import attach_dag_versions_to_runs
 from airflow.api_fastapi.common.parameters import (
     QueryDagRunRunTypesFilter,
     QueryDagRunStateFilter,
@@ -55,6 +55,7 @@
 from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
 from airflow.api_fastapi.core_api.security import requires_access_dag
 from airflow.api_fastapi.core_api.services.ui.grid import (
+    GridNodeAgg,
     _find_aggregates,
     _get_aggs_for_node,
     _merge_node_dicts,
@@ -63,13 +64,11 @@
     get_task_group_children_getter,
     task_group_to_dict_grid,
 )
-from airflow.models.dag import DagModel
 from airflow.models.dag_version import DagVersion
 from airflow.models.dagrun import DagRun
 from airflow.models.deadline import Deadline
 from airflow.models.serialized_dag import SerializedDagModel
 from airflow.models.taskinstance import TaskInstance
-from airflow.models.taskinstancehistory import TaskInstanceHistory
 from airflow.utils.session import create_session
 
 if TYPE_CHECKING:
@@ -307,16 +306,8 @@ def get_grid_runs(
                 DagRun.run_type,
                 DagRun.bundle_version,
             ),
-            joinedload(DagRun.dag_model).load_only(DagModel._dag_display_property_value),
             joinedload(DagRun.created_dag_version).joinedload(DagVersion.bundle),
-            selectinload(DagRun.task_instances)
-            .load_only(TaskInstance.dag_version_id)
-            .joinedload(TaskInstance.dag_version)
-            .joinedload(DagVersion.bundle),
-            selectinload(DagRun.task_instances_histories)
-            .load_only(TaskInstanceHistory.dag_version_id)
-            .joinedload(TaskInstanceHistory.dag_version)
-            .joinedload(DagVersion.bundle),
+            joinedload(DagRun.created_dag_version).joinedload(DagVersion.dag_model),
         )
     )
 
@@ -338,43 +329,67 @@ def get_grid_runs(
         return_total_entries=False,
     )
     results = session.execute(dag_runs_select_filter).unique().all()
+    dag_runs = [run for run, _ in results]
+    attach_dag_versions_to_runs(dag_runs, session=session)
     grid_runs = []
     for run, has_missed in results:
-        run.has_missed_deadline = has_missed
-        grid_runs.append(GridRunsResponse.model_validate(run, from_attributes=True))
+        grid_runs.append(
+            GridRunsResponse.model_validate(
+                {
+                    "dag_id": run.dag_id,
+                    "run_id": run.run_id,
+                    "queued_at": run.queued_at,
+                    "start_date": run.start_date,
+                    "end_date": run.end_date,
+                    "run_after": run.run_after,
+                    "state": run.state,
+                    "run_type": run.run_type,
+                    "dag_versions": run.dag_versions,
+                    "has_missed_deadline": has_missed,
+                }
+            )
+        )
     return grid_runs
 
 
 def _build_ti_summaries(
     dag_id: str,
     run_id: str,
-    task_instances: Sequence,
+    task_instances: Iterable[Any],
     session: Session,
     *,
     dag_bag: DBDagBag,
-) -> dict:
-    ti_details: dict = collections.defaultdict(list)
+    serdag: SerializedDAG | None = None,
+    serdag_cache: dict[Any, SerializedDAG | None] | None = None,
+) -> dict[str, Any] | None:
+    ti_details: dict[str, GridNodeAgg] = {}
+    dag_version_id = None
     for ti in task_instances:
-        ti_details[ti.task_id].append(
-            {
-                "state": ti.state,
-                "start_date": ti.start_date,
-                "end_date": ti.end_date,
-                "dag_version_number": getattr(ti, "version_number", None),
-            }
+        dag_version_id = dag_version_id or ti.dag_version_id
+        summary = ti_details.get(ti.task_id)
+        if summary is None:
+            summary = ti_details[ti.task_id] = GridNodeAgg()
+        summary.add_ti(
+            state=ti.state,
+            start_date=ti.start_date,
+            end_date=ti.end_date,
+            dag_version_number=getattr(ti, "version_number", None),
         )
-    serdag = _get_serdag(
-        dag_bag,
-        dag_id,
-        task_instances[0].dag_version_id,
-        session,
-    )
+    if not ti_details:
+        return None
+    if serdag is None:
+        if serdag_cache is not None:
+            if dag_version_id not in serdag_cache:
+                serdag_cache[dag_version_id] = _get_serdag(dag_bag, dag_id, dag_version_id, session)
+            serdag = serdag_cache[dag_version_id]
+        else:
+            serdag = _get_serdag(dag_bag, dag_id, dag_version_id, session)
     if TYPE_CHECKING:
         assert serdag
 
-    def get_node_summaries():
+    def get_node_summaries() -> Iterable[dict[str, Any]]:
         yielded_task_ids: set[str] = set()
-        for node in _find_aggregates(
+        for node, _ in _find_aggregates(
             node=serdag.task_group,
             parent_node=None,
             ti_details=ti_details,
@@ -453,6 +468,7 @@ def get_grid_ti_summaries_stream(
     """
 
     def _generate() -> Generator[str, None, None]:
+        serdag_cache: dict[Any, SerializedDAG | None] = {}
 
         # Each iteration opens and closes its own DB session so the connection is
         # released between yields.  This prevents a slow client from holding a
@@ -473,10 +489,18 @@ def _generate() -> Generator[str, None, None]:
                     .where(TaskInstance.dag_id == dag_id)
                     .where(TaskInstance.run_id == run_id)
                     .order_by(TaskInstance.task_id)
-                ).all()
-                if not tis:
-                    continue
-                summary = _build_ti_summaries(dag_id, run_id, tis, session, dag_bag=dag_bag)
+                    .execution_options(yield_per=1000)
+                )
+                summary = _build_ti_summaries(
+                    dag_id,
+                    run_id,
+                    tis,
+                    session,
+                    dag_bag=dag_bag,
+                    serdag_cache=serdag_cache,
+                )
+            if summary is None:
+                continue
             yield GridTISummaries.model_validate(summary).model_dump_json() + "\n"
 
     return StreamingResponse(content=_generate(), media_type="application/x-ndjson")
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/services/ui/grid.py b/airflow-core/src/airflow/api_fastapi/core_api/services/ui/grid.py
index 838818ad66bb7..3ea7c5b21a4a2 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/services/ui/grid.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/services/ui/grid.py
@@ -18,7 +18,9 @@
 from __future__ import annotations
 
 from collections import Counter
-from collections.abc import Iterable
+from collections.abc import Iterable, Mapping
+from dataclasses import dataclass, field
+from datetime import datetime
 from typing import Any
 
 import structlog
@@ -33,6 +35,64 @@
 log = structlog.get_logger(logger_name=__name__)
 
 
+@dataclass
+class GridNodeAgg:
+    """Compact task instance summary used to aggregate grid state without keeping TI details."""
+
+    child_states: Counter[Any] = field(default_factory=Counter)
+    min_start_date: datetime | None = None
+    max_end_date: datetime | None = None
+    dag_version_number: int | None = None
+
+    def add_ti(
+        self,
+        *,
+        state: Any,
+        start_date: datetime | None,
+        end_date: datetime | None,
+        dag_version_number: int | None,
+    ) -> None:
+        """Merge one task instance row into the summary."""
+        self.child_states[state] += 1
+        if start_date is not None and (self.min_start_date is None or start_date < self.min_start_date):
+            self.min_start_date = start_date
+        if end_date is not None and (self.max_end_date is None or end_date > self.max_end_date):
+            self.max_end_date = end_date
+        if dag_version_number is not None and (
+            self.dag_version_number is None or dag_version_number > self.dag_version_number
+        ):
+            self.dag_version_number = dag_version_number
+
+    def merge(self, other: GridNodeAgg) -> None:
+        """Merge another summary into this one."""
+        self.child_states.update(other.child_states)
+        if other.min_start_date is not None and (
+            self.min_start_date is None or other.min_start_date < self.min_start_date
+        ):
+            self.min_start_date = other.min_start_date
+        if other.max_end_date is not None and (
+            self.max_end_date is None or other.max_end_date > self.max_end_date
+        ):
+            self.max_end_date = other.max_end_date
+        if other.dag_version_number is not None and (
+            self.dag_version_number is None or other.dag_version_number > self.dag_version_number
+        ):
+            self.dag_version_number = other.dag_version_number
+
+    def with_placeholder_state(self) -> GridNodeAgg:
+        """Represent mapped tasks without rows as a single no-status square in the grid."""
+        if self.child_states:
+            return self
+        placeholder = GridNodeAgg(dag_version_number=self.dag_version_number)
+        placeholder.add_ti(
+            state=None,
+            start_date=None,
+            end_date=None,
+            dag_version_number=self.dag_version_number,
+        )
+        return placeholder
+
+
 def _merge_node_dicts(current: list[dict[str, Any]], new: list[dict[str, Any]] | None) -> None:
     """Merge node dictionaries from different Dag versions, handling structure changes."""
     # Handle None case - can occur when merging old Dag versions
@@ -55,90 +115,82 @@ def _merge_node_dicts(current: list[dict[str, Any]], new: list[dict[str, Any]] |
 
 
 def agg_state(states):
-    states = Counter(states)
+    state_counts = states if isinstance(states, Counter) else Counter(states)
     for state in state_priority:
-        if state in states:
+        if state in state_counts:
             return state
     return None
 
 
-def _get_aggs_for_node(detail):
-    states = [x["state"] for x in detail]
-    try:
-        min_start_date = min(x["start_date"] for x in detail if x["start_date"])
-    except ValueError:
-        min_start_date = None
-    try:
-        max_end_date = max(x["end_date"] for x in detail if x["end_date"])
-    except ValueError:
-        max_end_date = None
-
-    dag_version_numbers = [
-        x.get("dag_version_number") for x in detail if x.get("dag_version_number") is not None
-    ]
-    dag_version_number = max(dag_version_numbers) if dag_version_numbers else None
-
+def _get_aggs_for_node(summary: GridNodeAgg) -> dict[str, Any]:
     return {
-        "state": agg_state(states),
-        "min_start_date": min_start_date,
-        "max_end_date": max_end_date,
-        "child_states": dict(Counter(states)),
-        "dag_version_number": dag_version_number,
+        "state": agg_state(summary.child_states),
+        "min_start_date": summary.min_start_date,
+        "max_end_date": summary.max_end_date,
+        "child_states": dict(summary.child_states),
+        "dag_version_number": summary.dag_version_number,
     }
 
 
 def _find_aggregates(
     node: SerializedTaskGroup | SerializedBaseOperator | TaskMap,
     parent_node: SerializedTaskGroup | SerializedBaseOperator | TaskMap | None,
-    ti_details: dict[str, list],
-) -> Iterable[dict]:
+    ti_details: Mapping[str, GridNodeAgg],
+) -> Iterable[tuple[dict[str, Any], GridNodeAgg]]:
     """Recursively fill the Task Group Map."""
     node_id = node.node_id
     parent_id = parent_node.node_id if parent_node else None
     # Do not mutate ti_details by accidental key creation
-    details = ti_details.get(node_id, [])
+    summary = ti_details.get(node_id)
+    if summary is None:
+        summary = GridNodeAgg()
 
     if node is None:
         return
     if isinstance(node, SerializedMappedOperator):
-        # For unmapped tasks, reflect a single None state so UI shows one square
-        mapped_details = details or [{"state": None, "start_date": None, "end_date": None}]
-        yield {
-            "task_id": node_id,
-            "task_display_name": node.task_display_name,
-            "type": "mapped_task",
-            "parent_id": parent_id,
-            **_get_aggs_for_node(mapped_details),
-            "details": mapped_details,
-        }
+        mapped_summary = summary.with_placeholder_state()
+        yield (
+            {
+                "task_id": node_id,
+                "task_display_name": node.task_display_name,
+                "type": "mapped_task",
+                "parent_id": parent_id,
+                **_get_aggs_for_node(mapped_summary),
+            },
+            mapped_summary,
+        )
 
         return
     if isinstance(node, SerializedTaskGroup):
-        children_details = []
+        children_summary = GridNodeAgg()
         for child in get_task_group_children_getter()(node):
-            for child_node in _find_aggregates(node=child, parent_node=node, ti_details=ti_details):
+            for child_node, child_summary in _find_aggregates(
+                node=child, parent_node=node, ti_details=ti_details
+            ):
                 if child_node["parent_id"] == node_id:
-                    # Collect detailed task instance data from all children
-                    if child_node.get("details"):
-                        children_details.extend(child_node["details"])
-                yield child_node
+                    children_summary.merge(child_summary)
+                yield child_node, child_summary
         if node_id:
-            yield {
-                "task_id": node_id,
-                "task_display_name": node_id,
-                "type": "group",
-                "parent_id": parent_id,
-                **_get_aggs_for_node(children_details),
-                "details": children_details,
-            }
+            yield (
+                {
+                    "task_id": node_id,
+                    "task_display_name": node_id,
+                    "type": "group",
+                    "parent_id": parent_id,
+                    **_get_aggs_for_node(children_summary),
+                },
+                children_summary,
+            )
         return
     if isinstance(node, SerializedBaseOperator):
-        yield {
-            "task_id": node_id,
-            "task_display_name": node.task_display_name,
-            "type": "task",
-            "parent_id": parent_id,
-            **_get_aggs_for_node(details),
-            "details": details,
-        }
+        yield (
+            {
+                "task_id": node_id,
+                "task_display_name": node.task_display_name,
+                "type": "task",
+                "parent_id": parent_id,
+                **_get_aggs_for_node(summary),
+            },
+            summary,
+        )
         return
diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_grid.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_grid.py
index dbcf87149e588..b0dab9012a566 100644
--- a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_grid.py
+++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_grid.py
@@ -641,6 +641,29 @@ def test_get_grid_runs(self, session, test_client):
         assert response.status_code == 200
         assert _strip_dag_version_ids(response.json()) == [GRID_RUN_1, GRID_RUN_2]
 
+    def test_get_grid_runs_multiple_dag_versions(self, session, test_client):
+        latest_dag_version = session.scalar(select(DagModel).where(DagModel.dag_id == DAG_ID_5)).dag_versions[
+            -1
+        ]
+        latest_task_instance = session.scalar(
+            select(TaskInstance)
+            .where(TaskInstance.dag_id == DAG_ID_5, TaskInstance.run_id == "run_5_2")
+            .limit(1)
+        )
+        latest_task_instance.dag_version = latest_dag_version
+        session.commit()
+
+        response = test_client.get(f"/grid/runs/{DAG_ID_5}?limit=5")
+        assert response.status_code == 200
+        dag_versions_by_run_id = {
+            run["run_id"]: [dag_version["version_number"] for dag_version in run["dag_versions"]]
+            for run in response.json()
+        }
+        assert dag_versions_by_run_id == {
+            "run_5_1": [1],
+            "run_5_2": [1, 2],
+        }
+
     @pytest.mark.parametrize(
         ("endpoint", "run_type", "expected"),
         [

From a3babfa01ce237ae6ac2f0d477a07ff23185f30d Mon Sep 17 00:00:00 2001
From: Brent Bovenzi <brent.bovenzi@gmail.com>
Date: Fri, 15 May 2026 20:15:08 -0400
Subject: [PATCH 2/3] Fix tests

---
 airflow-core/tests/unit/migrations/test_migration_utils.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/airflow-core/tests/unit/migrations/test_migration_utils.py b/airflow-core/tests/unit/migrations/test_migration_utils.py
index 1ab309421fee9..25dce90f44322 100644
--- a/airflow-core/tests/unit/migrations/test_migration_utils.py
+++ b/airflow-core/tests/unit/migrations/test_migration_utils.py
@@ -136,6 +136,7 @@ def _get_revisions_in_order() -> list[str]:
 
 
 @pytest.fixture(scope="module")
+@pytest.mark.backend("sqlite", "postgres")
 def stairway_db():
     """
     Bring the DB to the stairway base revision before the test runs, and
@@ -153,6 +154,7 @@ def stairway_db():
     upgradedb()
 
 
+@pytest.mark.backend("sqlite", "postgres")
 def test_migration_stairway(stairway_db) -> None:
     """
     Walk every incremental migration step since 3.0.0: upgrade → downgrade → re-upgrade.

From c6b6d71b35521e075d409be37ba33f2427d47cd0 Mon Sep 17 00:00:00 2001
From: Brent Bovenzi <brent.bovenzi@gmail.com>
Date: Tue, 19 May 2026 11:36:36 -0400
Subject: [PATCH 3/3] Undo migration file change

---
 airflow-core/tests/unit/migrations/test_migration_utils.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/airflow-core/tests/unit/migrations/test_migration_utils.py b/airflow-core/tests/unit/migrations/test_migration_utils.py
index 25dce90f44322..1ab309421fee9 100644
--- a/airflow-core/tests/unit/migrations/test_migration_utils.py
+++ b/airflow-core/tests/unit/migrations/test_migration_utils.py
@@ -136,7 +136,6 @@ def _get_revisions_in_order() -> list[str]:
 
 
 @pytest.fixture(scope="module")
-@pytest.mark.backend("sqlite", "postgres")
 def stairway_db():
     """
     Bring the DB to the stairway base revision before the test runs, and
@@ -154,7 +153,6 @@ def stairway_db():
     upgradedb()
 
 
-@pytest.mark.backend("sqlite", "postgres")
 def test_migration_stairway(stairway_db) -> None:
     """
     Walk every incremental migration step since 3.0.0: upgrade → downgrade → re-upgrade.
