[v3-1-test] Fix n+1 queries in ui/dags endpoint and use a compact response model. (#57425)

airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/dag_runs.py

@@ -18,6 +18,8 @@
 
 from datetime import datetime
 
+from pydantic import computed_field
+
 from airflow.api_fastapi.core_api.base import BaseModel
 from airflow.utils.state import DagRunState
 
@@ -33,3 +35,9 @@ class DAGRunLightResponse(BaseModel):
     start_date: datetime | None
     end_date: datetime | None
     state: DagRunState
+
+    @computed_field
+    def duration(self) -> float | None:
+        if self.end_date and self.start_date:
+            return (self.end_date - self.start_date).total_seconds()
+        return None

airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/dags.py

@@ -18,16 +18,16 @@
 from __future__ import annotations
 
 from airflow.api_fastapi.core_api.base import BaseModel
-from airflow.api_fastapi.core_api.datamodels.dag_run import DAGRunResponse
 from airflow.api_fastapi.core_api.datamodels.dags import DAGResponse
 from airflow.api_fastapi.core_api.datamodels.hitl import HITLDetail
+from airflow.api_fastapi.core_api.datamodels.ui.dag_runs import DAGRunLightResponse
 
 
 class DAGWithLatestDagRunsResponse(DAGResponse):
     """DAG with latest dag runs response serializer."""
 
     asset_expression: dict | None
-    latest_dag_runs: list[DAGRunResponse]
+    latest_dag_runs: list[DAGRunLightResponse]
     pending_actions: list[HITLDetail]
     is_favorite: bool
 

airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml

@@ -1347,137 +1347,24 @@ components:
           title: End Date
         state:
           $ref: '#/components/schemas/DagRunState'
-      type: object
-      required:
-      - id
-      - dag_id
-      - run_id
-      - logical_date
-      - run_after
-      - start_date
-      - end_date
-      - state
-      title: DAGRunLightResponse
-      description: DAG Run serializer for responses.
-    DAGRunResponse:
-      properties:
-        dag_run_id:
-          type: string
-          title: Dag Run Id
-        dag_id:
-          type: string
-          title: Dag Id
-        logical_date:
-          anyOf:
-          - type: string
-            format: date-time
-          - type: 'null'
-          title: Logical Date
-        queued_at:
-          anyOf:
-          - type: string
-            format: date-time
-          - type: 'null'
-          title: Queued At
-        start_date:
-          anyOf:
-          - type: string
-            format: date-time
-          - type: 'null'
-          title: Start Date
-        end_date:
-          anyOf:
-          - type: string
-            format: date-time
-          - type: 'null'
-          title: End Date
         duration:
           anyOf:
           - type: number
           - type: 'null'
           title: Duration
-        data_interval_start:
-          anyOf:
-          - type: string
-            format: date-time
-          - type: 'null'
-          title: Data Interval Start
-        data_interval_end:
-          anyOf:
-          - type: string
-            format: date-time
-          - type: 'null'
-          title: Data Interval End
-        run_after:
-          type: string
-          format: date-time
-          title: Run After
-        last_scheduling_decision:
-          anyOf:
-          - type: string
-            format: date-time
-          - type: 'null'
-          title: Last Scheduling Decision
-        run_type:
-          $ref: '#/components/schemas/DagRunType'
-        state:
-          $ref: '#/components/schemas/DagRunState'
-        triggered_by:
-          anyOf:
-          - $ref: '#/components/schemas/DagRunTriggeredByType'
-          - type: 'null'
-        triggering_user_name:
-          anyOf:
-          - type: string
-          - type: 'null'
-          title: Triggering User Name
-        conf:
-          anyOf:
-          - additionalProperties: true
-            type: object
-          - type: 'null'
-          title: Conf
-        note:
-          anyOf:
-          - type: string
-          - type: 'null'
-          title: Note
-        dag_versions:
-          items:
-            $ref: '#/components/schemas/DagVersionResponse'
-          type: array
-          title: Dag Versions
-        bundle_version:
-          anyOf:
-          - type: string
-          - type: 'null'
-          title: Bundle Version
-        dag_display_name:
-          type: string
-          title: Dag Display Name
+          readOnly: true
       type: object
       required:
-      - dag_run_id
+      - id
       - dag_id
+      - run_id
       - logical_date
-      - queued_at
+      - run_after
       - start_date
       - end_date
-      - duration
-      - data_interval_start
-      - data_interval_end
-      - run_after
-      - last_scheduling_decision
-      - run_type
       - state
-      - triggered_by
-      - triggering_user_name
-      - conf
-      - note
-      - dag_versions
-      - bundle_version
-      - dag_display_name
-      title: DAGRunResponse
+      - duration
+      title: DAGRunLightResponse
       description: DAG Run serializer for responses.
     DAGRunStates:
       properties:
@@ -1662,7 +1549,7 @@ components:
           title: Asset Expression
         latest_dag_runs:
           items:
-            $ref: '#/components/schemas/DAGRunResponse'
+            $ref: '#/components/schemas/DAGRunLightResponse'
           type: array
           title: Latest Dag Runs
         pending_actions:
@@ -1728,19 +1615,6 @@ components:
         so please ensure that their values always match the ones with the
 
         same name in TaskInstanceState.'
-    DagRunTriggeredByType:
-      type: string
-      enum:
-      - cli
-      - operator
-      - rest_api
-      - ui
-      - test
-      - timetable
-      - asset
-      - backfill
-      title: DagRunTriggeredByType
-      description: Class with TriggeredBy types for DagRun.
     DagRunType:
       type: string
       enum:

airflow-core/src/airflow/api_fastapi/core_api/routes/ui/dags.py

@@ -52,7 +52,6 @@
     filter_param_factory,
 )
 from airflow.api_fastapi.common.router import AirflowRouter
-from airflow.api_fastapi.core_api.datamodels.dag_run import DAGRunResponse
 from airflow.api_fastapi.core_api.datamodels.dags import DAGResponse
 from airflow.api_fastapi.core_api.datamodels.ui.dag_runs import DAGRunLightResponse
 from airflow.api_fastapi.core_api.datamodels.ui.dags import (
@@ -183,7 +182,15 @@ def get_dags(
     recent_dag_runs_select = (
         select(
             recent_runs_subquery.c.run_after,
-            DagRun,
+            DagRun.id,
+            DagRun.dag_id,
+            DagRun.run_id,
+            DagRun.end_date,
+            DagRun.logical_date,
+            DagRun.run_after,
+            DagRun.start_date,
+            DagRun.state,
+            DagRun.duration,
         )
         .join(
             DagRun,
@@ -205,7 +212,7 @@ def get_dags(
 
     # Fetch pending HITL actions for each Dag if we are not certain whether some of the Dag might contain HITL actions
     pending_actions_by_dag_id: dict[str, list[HITLDetail]] = {dag.dag_id: [] for dag in dags}
-    if has_pending_actions.value is not False:
+    if has_pending_actions.value:
         pending_actions_select = (
             select(
                 TaskInstance.dag_id,
@@ -241,9 +248,8 @@ def get_dags(
     }
 
     for row in recent_dag_runs:
-        _, dag_run = row
-        dag_id = dag_run.dag_id
-        dag_run_response = DAGRunResponse.model_validate(dag_run)
+        dag_run_response = DAGRunLightResponse.model_validate(row)
+        dag_id = dag_run_response.dag_id
         dag_runs_by_dag_id[dag_id].latest_dag_runs.append(dag_run_response)
 
     return DAGWithLatestDagRunsCollectionResponse(

airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts

@@ -6878,10 +6878,22 @@ export const $DAGRunLightResponse = {
         },
         state: {
             '$ref': '#/components/schemas/DagRunState'
+        },
+        duration: {
+            anyOf: [
+                {
+                    type: 'number'
+                },
+                {
+                    type: 'null'
+                }
+            ],
+            title: 'Duration',
+            readOnly: true
         }
     },
     type: 'object',
-    required: ['id', 'dag_id', 'run_id', 'logical_date', 'run_after', 'start_date', 'end_date', 'state'],
+    required: ['id', 'dag_id', 'run_id', 'logical_date', 'run_after', 'start_date', 'end_date', 'state', 'duration'],
     title: 'DAGRunLightResponse',
     description: 'DAG Run serializer for responses.'
 } as const;
@@ -7182,7 +7194,7 @@ export const $DAGWithLatestDagRunsResponse = {
         },
         latest_dag_runs: {
             items: {
-                '$ref': '#/components/schemas/DAGRunResponse'
+                '$ref': '#/components/schemas/DAGRunLightResponse'
             },
             type: 'array',
             title: 'Latest Dag Runs'

airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts

@@ -1742,6 +1742,7 @@ export type DAGRunLightResponse = {
     start_date: string | null;
     end_date: string | null;
     state: DagRunState;
+    readonly duration: number | null;
 };
 
 /**
@@ -1804,7 +1805,7 @@ export type DAGWithLatestDagRunsResponse = {
     asset_expression: {
     [key: string]: unknown;
 } | null;
-    latest_dag_runs: Array<DAGRunResponse>;
+    latest_dag_runs: Array<DAGRunLightResponse>;
     pending_actions: Array<HITLDetail>;
     is_favorite: boolean;
     /**

airflow-core/src/airflow/ui/src/mocks/handlers/dags.ts

@@ -38,20 +38,13 @@ export const handlers: Array<HttpHandler> = [
       last_parsed_time: "2025-01-13T07:34:01.593459Z",
       latest_dag_runs: [
         {
-          conf: {},
           dag_id: "tutorial_taskflow_api",
-          dag_run_id: "manual__2025-01-13T04:33:58.387988+00:00",
-          data_interval_end: "2025-01-13T04:33:58.396323Z",
-          data_interval_start: "2025-01-13T04:33:58.396323Z",
           end_date: "2025-01-13T04:34:12.143831Z",
-          external_trigger: true,
-          last_scheduling_decision: "2025-01-13T04:34:12.137382Z",
+          id: 1,
           logical_date: "2025-01-13T04:33:58.396323Z",
-          queued_at: "2025-01-13T04:33:58.404628Z",
-          run_type: "manual",
+          run_id: "manual__2025-01-13T04:33:58.387988+00:00",
           start_date: "2025-01-13T04:33:58.496197Z",
           state: "success",
-          triggered_by: "rest_api",
         },
       ],
       max_active_runs: 16,
@@ -76,20 +69,13 @@ export const handlers: Array<HttpHandler> = [
       last_parsed_time: "2025-01-13T07:34:01.593459Z",
       latest_dag_runs: [
         {
-          conf: {},
           dag_id: "tutorial_taskflow_api",
-          dag_run_id: "manual__2025-01-13T04:33:58.387988+00:00",
-          data_interval_end: "2025-01-13T04:33:58.396323Z",
-          data_interval_start: "2025-01-13T04:33:58.396323Z",
           end_date: "2025-01-13T04:34:12.143831Z",
-          external_trigger: true,
-          last_scheduling_decision: "2025-01-13T04:34:12.137382Z",
+          id: 2,
           logical_date: "2025-01-13T04:33:58.396323Z",
-          queued_at: "2025-01-13T04:33:58.404628Z",
-          run_type: "manual",
+          run_id: "manual__2025-01-13T04:33:58.387988+00:00",
           start_date: "2025-01-13T04:33:58.496197Z",
           state: "success",
-          triggered_by: "rest_api",
         },
       ],
       max_active_runs: 16,