From a078cceb2af3ae6886889c18ea2c3766a950d0cd Mon Sep 17 00:00:00 2001
From: Pierre Jeambrun <pierrejbrun@gmail.com>
Date: Sat, 1 Feb 2025 16:55:11 +0100
Subject: [PATCH] AIP-84 Fix version_number in grid endpoint (#46323)

---
 airflow/api_fastapi/core_api/datamodels/ui/grid.py    |  3 +--
 .../api_fastapi/core_api/openapi/v1-generated.yaml    |  3 +--
 airflow/api_fastapi/core_api/routes/ui/grid.py        | 11 +++++++----
 airflow/ui/openapi-gen/requests/schemas.gen.ts        |  3 +--
 airflow/ui/openapi-gen/requests/types.gen.ts          |  2 +-
 5 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/airflow/api_fastapi/core_api/datamodels/ui/grid.py b/airflow/api_fastapi/core_api/datamodels/ui/grid.py
index 00d6b496b55b5..cb05c53aaa36a 100644
--- a/airflow/api_fastapi/core_api/datamodels/ui/grid.py
+++ b/airflow/api_fastapi/core_api/datamodels/ui/grid.py
@@ -18,7 +18,6 @@
 from __future__ import annotations
 
 from datetime import datetime
-from uuid import UUID
 
 from pydantic import BaseModel, Field
 
@@ -51,7 +50,7 @@ class GridDAGRunwithTIs(BaseModel):
     run_type: DagRunType
     data_interval_start: datetime | None
     data_interval_end: datetime | None
-    version_number: UUID | None
+    version_number: int | None
     note: str | None
     task_instances: list[GridTaskInstanceSummary]
 
diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml
index 1911c401be002..d06d09428a5d6 100644
--- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml
+++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml
@@ -8786,8 +8786,7 @@ components:
           title: Data Interval End
         version_number:
           anyOf:
-          - type: string
-            format: uuid
+          - type: integer
           - type: 'null'
           title: Version Number
         note:
diff --git a/airflow/api_fastapi/core_api/routes/ui/grid.py b/airflow/api_fastapi/core_api/routes/ui/grid.py
index ca8fd01684094..fdfbc6e746853 100644
--- a/airflow/api_fastapi/core_api/routes/ui/grid.py
+++ b/airflow/api_fastapi/core_api/routes/ui/grid.py
@@ -23,6 +23,7 @@
 
 from fastapi import Depends, HTTPException, Request, status
 from sqlalchemy import select
+from sqlalchemy.orm import joinedload
 
 from airflow import DAG
 from airflow.api_fastapi.common.db.common import SessionDep, paginated_select
@@ -93,7 +94,9 @@ def grid_data(
     base_query = (
         select(DagRun)
         .join(DagRun.dag_run_note, isouter=True)
+        .join(DagRun.dag_version, isouter=True)
         .select_from(DagRun)
+        .options(joinedload(DagRun.dag_version))
         .where(DagRun.dag_id == dag.dag_id)
     )
 
@@ -212,11 +215,11 @@ def grid_data(
             run_type=dag_run.run_type,
             data_interval_start=dag_run.data_interval_start,
             data_interval_end=dag_run.data_interval_end,
-            version_number=dag_run.dag_version_id,
+            version_number=dag_run.dag_version.version_number if dag_run.dag_version else None,
             note=dag_run.note,
-            task_instances=task_instance_summaries[dag_run.run_id]
-            if dag_run.run_id in task_instance_summaries
-            else [],
+            task_instances=(
+                task_instance_summaries[dag_run.run_id] if dag_run.run_id in task_instance_summaries else []
+            ),
         )
         for dag_run in dag_runs
     ]
diff --git a/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow/ui/openapi-gen/requests/schemas.gen.ts
index d30aacd6be03e..bbb7a033c3167 100644
--- a/airflow/ui/openapi-gen/requests/schemas.gen.ts
+++ b/airflow/ui/openapi-gen/requests/schemas.gen.ts
@@ -3401,8 +3401,7 @@ export const $GridDAGRunwithTIs = {
     version_number: {
       anyOf: [
         {
-          type: "string",
-          format: "uuid",
+          type: "integer",
         },
         {
           type: "null",
diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts
index 9746699f1b521..e687dacaa0188 100644
--- a/airflow/ui/openapi-gen/requests/types.gen.ts
+++ b/airflow/ui/openapi-gen/requests/types.gen.ts
@@ -889,7 +889,7 @@ export type GridDAGRunwithTIs = {
   run_type: DagRunType;
   data_interval_start: string | null;
   data_interval_end: string | null;
-  version_number: string | null;
+  version_number: number | null;
   note: string | null;
   task_instances: Array<GridTaskInstanceSummary>;
 };