From 13132ee34b56d74437a59d93f0766e24882261cb Mon Sep 17 00:00:00 2001
From: Pierre Jeambrun <pierrejbrun@gmail.com>
Date: Mon, 27 Apr 2026 17:54:17 +0200
Subject: [PATCH] UI: Serve grid TI summaries from shared cached DagBag
 (#65775)

The grid TI summaries streaming endpoint used a per-request dict to
dedupe serialized Dag loads within a single response. Switch to the
app-wide DBDagBag so the serialized Dag is cached across requests (via
the LRU+TTL configured by [api] dag_cache_size / dag_cache_ttl) and
the in-memory representation is the lightweight SerializedDAG instead
of the SerializedDagModel ORM row holding the full JSON blob.

(cherry picked from commit be41cb4031e38050314c9c9605885edee5dd1df0)
---
 .../core_api/openapi/_private_ui.yaml         |  6 +-
 .../api_fastapi/core_api/routes/ui/grid.py    | 85 ++++++++++---------
 .../ui/openapi-gen/queries/ensureQueryData.ts |  5 +-
 .../ui/openapi-gen/queries/prefetch.ts        |  5 +-
 .../airflow/ui/openapi-gen/queries/queries.ts |  5 +-
 .../ui/openapi-gen/queries/suspense.ts        |  5 +-
 .../ui/openapi-gen/requests/services.gen.ts   |  5 +-
 7 files changed, 64 insertions(+), 52 deletions(-)

diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml
index dc5c6d101dcad..e6f1854f6eff7 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml
+++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml
@@ -1145,9 +1145,11 @@ paths:
         progressively without waiting for all runs to complete.
 
 
-        The serialized Dag structure is loaded once and reused for all runs that
+        The serialized Dag structure is served from the app-wide ``DBDagBag`` cache
 
-        share the same ``dag_version_id``, avoiding repeated deserialization.'
+        (keyed by ``dag_version_id``), which avoids repeated deserialization across
+
+        runs of the same version *and* across requests.'
       operationId: get_grid_ti_summaries_stream
       security:
       - OAuth2PasswordBearer: []
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 ecc4f8e4655ac..71a02273d2d40 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
@@ -20,14 +20,16 @@
 import collections
 from collections.abc import Generator, Sequence
 from typing import TYPE_CHECKING, Annotated, Any
+from uuid import UUID
 
 import structlog
 from fastapi import Depends, HTTPException, Query, status
 from fastapi.responses import StreamingResponse
 from sqlalchemy import exists, select
-from sqlalchemy.orm import joinedload, load_only, selectinload
+from sqlalchemy.orm import Session, joinedload, load_only, selectinload
 
 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.parameters import (
     QueryDagRunRunTypesFilter,
@@ -69,6 +71,10 @@
 from airflow.models.taskinstancehistory import TaskInstanceHistory
 from airflow.utils.session import create_session
 
+if TYPE_CHECKING:
+    from airflow.models.dagbag import DBDagBag
+    from airflow.serialization.definitions.dag import SerializedDAG
+
 log = structlog.get_logger(logger_name=__name__)
 grid_router = AirflowRouter(prefix="/grid", tags=["Grid"])
 
@@ -90,32 +96,28 @@ def _get_latest_serdag(dag_id, session):
     return serdag
 
 
-def _get_serdag(dag_id, dag_version_id, session) -> SerializedDagModel | None:
-    # this is a simplification - we account for structure based on the first task
-    version = session.scalar(
-        select(DagVersion)
-        .where(DagVersion.id == dag_version_id)
-        .options(joinedload(DagVersion.serialized_dag))
+def _get_serdag(
+    dag_bag: DBDagBag,
+    dag_id: str,
+    dag_version_id: UUID | None,
+    session: Session,
+) -> SerializedDAG | None:
+    """Resolve the serialized Dag for a grid TI summary via the shared (cached) ``DBDagBag``."""
+    if dag_version_id is not None:
+        serdag = dag_bag._get_dag(dag_version_id, session=session)
+        if serdag is None:
+            log.error("No serialized dag found", dag_id=dag_id, version_id=dag_version_id)
+        return serdag
+
+    # Fallback: pre-3.0 upgrade — pick the oldest DagVersion for this dag_id.
+    oldest_version_id = session.scalar(
+        select(DagVersion.id).where(DagVersion.dag_id == dag_id).order_by(DagVersion.id).limit(1)
     )
-    if not version:
-        version = session.scalar(
-            select(DagVersion)
-            .where(
-                DagVersion.dag_id == dag_id,
-            )
-            .options(joinedload(DagVersion.serialized_dag))
-            .order_by(DagVersion.id)  # ascending cus this is mostly for pre-3.0 upgrade
-            .limit(1)
-        )
-    if not version:
+    if oldest_version_id is None:
         return None
-    if not (serdag := version.serialized_dag):
-        log.error(
-            "No serialized dag found",
-            dag_id=dag_id,
-            version_id=version.id,
-            version_number=version.version_number,
-        )
+    serdag = dag_bag._get_dag(oldest_version_id, session=session)
+    if serdag is None:
+        log.error("No serialized dag found", dag_id=dag_id, version_id=oldest_version_id)
     return serdag
 
 
@@ -341,7 +343,12 @@ def get_grid_runs(
 
 
 def _build_ti_summaries(
-    dag_id: str, run_id: str, task_instances: Sequence, session, serdag: SerializedDagModel | None = None
+    dag_id: str,
+    run_id: str,
+    task_instances: Sequence,
+    session: Session,
+    *,
+    dag_bag: DBDagBag,
 ) -> dict:
     ti_details: dict = collections.defaultdict(list)
     for ti in task_instances:
@@ -353,19 +360,19 @@ def _build_ti_summaries(
                 "dag_version_number": getattr(ti, "version_number", None),
             }
         )
-    if serdag is None:
-        serdag = _get_serdag(
-            dag_id=dag_id,
-            dag_version_id=task_instances[0].dag_version_id,
-            session=session,
-        )
+    serdag = _get_serdag(
+        dag_bag,
+        dag_id,
+        task_instances[0].dag_version_id,
+        session,
+    )
     if TYPE_CHECKING:
         assert serdag
 
     def get_node_summaries():
         yielded_task_ids: set[str] = set()
         for node in _find_aggregates(
-            node=serdag.dag.task_group,
+            node=serdag.task_group,
             parent_node=None,
             ti_details=ti_details,
         ):
@@ -427,6 +434,7 @@ def get_node_summaries():
 )
 def get_grid_ti_summaries_stream(
     dag_id: str,
+    dag_bag: DagBagDep,
     run_ids: Annotated[list[str] | None, Query()] = None,
 ) -> StreamingResponse:
     """
@@ -436,8 +444,9 @@ def get_grid_ti_summaries_stream(
     run's task instances have been processed, so the client can render columns
     progressively without waiting for all runs to complete.
 
-    The serialized Dag structure is loaded once and reused for all runs that
-    share the same ``dag_version_id``, avoiding repeated deserialization.
+    The serialized Dag structure is served from the app-wide ``DBDagBag`` cache
+    (keyed by ``dag_version_id``), which avoids repeated deserialization across
+    runs of the same version *and* across requests.
     """
 
     def _generate() -> Generator[str, None, None]:
@@ -446,7 +455,6 @@ def _generate() -> Generator[str, None, None]:
         # released between yields.  This prevents a slow client from holding a
         # database connection open for the entire stream duration.
         # See https://github.com/apache/airflow/issues/65010.
-        serdag_cache: dict = {}
         for run_id in run_ids or []:
             with create_session(scoped=False) as session:
                 tis = session.execute(
@@ -465,10 +473,7 @@ def _generate() -> Generator[str, None, None]:
                 ).all()
                 if not tis:
                     continue
-                version_id = tis[0].dag_version_id
-                if version_id not in serdag_cache:
-                    serdag_cache[version_id] = _get_serdag(dag_id, version_id, session)
-                summary = _build_ti_summaries(dag_id, run_id, tis, session, serdag=serdag_cache[version_id])
+                summary = _build_ti_summaries(dag_id, run_id, tis, session, dag_bag=dag_bag)
             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/ui/openapi-gen/queries/ensureQueryData.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
index 724d8f8b3ca8a..2c9eda8c18154 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
@@ -1705,8 +1705,9 @@ export const ensureUseGridServiceGetGridRunsData = (queryClient: QueryClient, {
 * run's task instances have been processed, so the client can render columns
 * progressively without waiting for all runs to complete.
 *
-* The serialized Dag structure is loaded once and reused for all runs that
-* share the same ``dag_version_id``, avoiding repeated deserialization.
+* The serialized Dag structure is served from the app-wide ``DBDagBag`` cache
+* (keyed by ``dag_version_id``), which avoids repeated deserialization across
+* runs of the same version *and* across requests.
 * @param data The data for the request.
 * @param data.dagId
 * @param data.runIds
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
index 449c99ef94152..67ee45a44aca4 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
@@ -1705,8 +1705,9 @@ export const prefetchUseGridServiceGetGridRuns = (queryClient: QueryClient, { da
 * run's task instances have been processed, so the client can render columns
 * progressively without waiting for all runs to complete.
 *
-* The serialized Dag structure is loaded once and reused for all runs that
-* share the same ``dag_version_id``, avoiding repeated deserialization.
+* The serialized Dag structure is served from the app-wide ``DBDagBag`` cache
+* (keyed by ``dag_version_id``), which avoids repeated deserialization across
+* runs of the same version *and* across requests.
 * @param data The data for the request.
 * @param data.dagId
 * @param data.runIds
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
index 805700964ce11..6d13b63815c65 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
@@ -1705,8 +1705,9 @@ export const useGridServiceGetGridRuns = <TData = Common.GridServiceGetGridRunsD
 * run's task instances have been processed, so the client can render columns
 * progressively without waiting for all runs to complete.
 *
-* The serialized Dag structure is loaded once and reused for all runs that
-* share the same ``dag_version_id``, avoiding repeated deserialization.
+* The serialized Dag structure is served from the app-wide ``DBDagBag`` cache
+* (keyed by ``dag_version_id``), which avoids repeated deserialization across
+* runs of the same version *and* across requests.
 * @param data The data for the request.
 * @param data.dagId
 * @param data.runIds
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
index 109d7f87676f8..0ffb43cdee825 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
@@ -1705,8 +1705,9 @@ export const useGridServiceGetGridRunsSuspense = <TData = Common.GridServiceGetG
 * run's task instances have been processed, so the client can render columns
 * progressively without waiting for all runs to complete.
 *
-* The serialized Dag structure is loaded once and reused for all runs that
-* share the same ``dag_version_id``, avoiding repeated deserialization.
+* The serialized Dag structure is served from the app-wide ``DBDagBag`` cache
+* (keyed by ``dag_version_id``), which avoids repeated deserialization across
+* runs of the same version *and* across requests.
 * @param data The data for the request.
 * @param data.dagId
 * @param data.runIds
diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
index 56a884f1bab4e..a196f1f47faf6 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
@@ -4199,8 +4199,9 @@ export class GridService {
      * run's task instances have been processed, so the client can render columns
      * progressively without waiting for all runs to complete.
      *
-     * The serialized Dag structure is loaded once and reused for all runs that
-     * share the same ``dag_version_id``, avoiding repeated deserialization.
+     * The serialized Dag structure is served from the app-wide ``DBDagBag`` cache
+     * (keyed by ``dag_version_id``), which avoids repeated deserialization across
+     * runs of the same version *and* across requests.
      * @param data The data for the request.
      * @param data.dagId
      * @param data.runIds
