AIP-65 Add DagVersion to TaskInstanceResponse Serializer (apache#46320)

* AIP-65 Add DagVersion to TaskInstanceResponse Serializer

* Update following code review

* Fix test

Author: pierrejeambrun

airflow/api_fastapi/common/parameters.py

@@ -48,6 +48,7 @@
     TaskOutletAssetReference,
 )
 from airflow.models.dag import DagModel, DagTag
+from airflow.models.dag_version import DagVersion
 from airflow.models.dagrun import DagRun
 from airflow.models.pool import Pool
 from airflow.models.taskinstance import TaskInstance
@@ -635,6 +636,17 @@ def _transform_ti_states(states: list[str] | None) -> list[TaskInstanceState | N
 QueryTITaskDisplayNamePatternSearch = Annotated[
     _SearchParam, Depends(search_param_factory(TaskInstance.task_display_name, "task_display_name_pattern"))
 ]
+QueryTIDagVersionFilter = Annotated[
+    FilterParam[list[int]],
+    Depends(
+        filter_param_factory(
+            DagVersion.version_number,
+            list[int],
+            FilterOptionEnum.ANY_EQUAL,
+            default_factory=list,
+        )
+    ),
+]
 
 # Assets
 QueryAssetNamePatternSearch = Annotated[

airflow/api_fastapi/core_api/datamodels/dag_versions.py

@@ -0,0 +1,33 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datetime import datetime
+from uuid import UUID
+
+from airflow.api_fastapi.core_api.base import BaseModel
+
+
+class DagVersionResponse(BaseModel):
+    """Dag Version serializer for responses."""
+
+    id: UUID
+    version_number: int
+    dag_id: str
+    bundle_name: str
+    bundle_version: str | None
+    created_at: datetime

airflow/api_fastapi/core_api/datamodels/task_instances.py

@@ -33,6 +33,7 @@
 )
 
 from airflow.api_fastapi.core_api.base import BaseModel, StrictBaseModel
+from airflow.api_fastapi.core_api.datamodels.dag_versions import DagVersionResponse
 from airflow.api_fastapi.core_api.datamodels.job import JobResponse
 from airflow.api_fastapi.core_api.datamodels.trigger import TriggerResponse
 from airflow.utils.state import TaskInstanceState
@@ -76,6 +77,7 @@ class TaskInstanceResponse(BaseModel):
     )
     trigger: TriggerResponse | None
     queued_by_job: JobResponse | None = Field(alias="triggerer_job")
+    dag_version: DagVersionResponse | None
 
 
 class TaskInstanceCollectionResponse(BaseModel):

airflow/api_fastapi/core_api/openapi/v1-generated.yaml

@@ -4842,6 +4842,14 @@ paths:
           items:
             type: string
           title: Executor
+      - name: version_number
+        in: query
+        required: false
+        schema:
+          type: array
+          items:
+            type: integer
+          title: Version Number
       - name: limit
         in: query
         required: false
@@ -5464,6 +5472,14 @@ paths:
           items:
             type: string
           title: Executor
+      - name: version_number
+        in: query
+        required: false
+        schema:
+          type: array
+          items:
+            type: integer
+          title: Version Number
       - name: limit
         in: query
         required: false
@@ -8649,6 +8665,40 @@ components:
       - dag_id
       title: DagTagResponse
       description: DAG Tag serializer for responses.
+    DagVersionResponse:
+      properties:
+        id:
+          type: string
+          format: uuid
+          title: Id
+        version_number:
+          type: integer
+          title: Version Number
+        dag_id:
+          type: string
+          title: Dag Id
+        bundle_name:
+          type: string
+          title: Bundle Name
+        bundle_version:
+          anyOf:
+          - type: string
+          - type: 'null'
+          title: Bundle Version
+        created_at:
+          type: string
+          format: date-time
+          title: Created At
+      type: object
+      required:
+      - id
+      - version_number
+      - dag_id
+      - bundle_name
+      - bundle_version
+      - created_at
+      title: DagVersionResponse
+      description: Dag Version serializer for responses.
     DagWarningType:
       type: string
       enum:
@@ -9900,6 +9950,10 @@ components:
           anyOf:
           - $ref: '#/components/schemas/JobResponse'
           - type: 'null'
+        dag_version:
+          anyOf:
+          - $ref: '#/components/schemas/DagVersionResponse'
+          - type: 'null'
       type: object
       required:
       - id
@@ -9931,6 +9985,7 @@ components:
       - rendered_map_index
       - trigger
       - triggerer_job
+      - dag_version
       title: TaskInstanceResponse
       description: TaskInstance serializer for responses.
     TaskInstanceState:

airflow/api_fastapi/core_api/routes/public/task_instances.py

@@ -35,6 +35,7 @@
     OffsetFilter,
     QueryLimit,
     QueryOffset,
+    QueryTIDagVersionFilter,
     QueryTIExecutorFilter,
     QueryTIPoolFilter,
     QueryTIQueueFilter,
@@ -87,6 +88,7 @@ def get_task_instance(
         .where(TI.dag_id == dag_id, TI.run_id == dag_run_id, TI.task_id == task_id)
         .join(TI.dag_run)
         .options(joinedload(TI.rendered_task_instance_fields))
+        .options(joinedload(TI.dag_version))
     )
     task_instance = session.scalar(query)
 
@@ -121,6 +123,7 @@ def get_mapped_task_instances(
     pool: QueryTIPoolFilter,
     queue: QueryTIQueueFilter,
     executor: QueryTIExecutorFilter,
+    version_number: QueryTIDagVersionFilter,
     limit: QueryLimit,
     offset: QueryOffset,
     order_by: Annotated[
@@ -156,6 +159,7 @@ def get_mapped_task_instances(
         select(TI)
         .where(TI.dag_id == dag_id, TI.run_id == dag_run_id, TI.task_id == task_id, TI.map_index >= 0)
         .join(TI.dag_run)
+        .options(joinedload(TI.dag_version))
     )
     # 0 can mean a mapped TI that expanded to an empty list, so it is not an automatic 404
     unfiltered_total_count = get_query_count(query, session=session)
@@ -185,6 +189,7 @@ def get_mapped_task_instances(
             pool,
             queue,
             executor,
+            version_number,
         ],
         order_by=order_by,
         offset=offset,
@@ -330,6 +335,7 @@ def get_mapped_task_instance(
         .where(TI.dag_id == dag_id, TI.run_id == dag_run_id, TI.task_id == task_id, TI.map_index == map_index)
         .join(TI.dag_run)
         .options(joinedload(TI.rendered_task_instance_fields))
+        .options(joinedload(TI.dag_version))
     )
     task_instance = session.scalar(query)
 
@@ -361,6 +367,7 @@ def get_task_instances(
     pool: QueryTIPoolFilter,
     queue: QueryTIQueueFilter,
     executor: QueryTIExecutorFilter,
+    version_number: QueryTIDagVersionFilter,
     limit: QueryLimit,
     offset: QueryOffset,
     order_by: Annotated[
@@ -397,7 +404,7 @@ def get_task_instances(
     This endpoint allows specifying `~` as the dag_id, dag_run_id to retrieve Task Instances for all DAGs
     and DAG runs.
     """
-    query = select(TI).join(TI.dag_run)
+    query = select(TI).join(TI.dag_run).outerjoin(TI.dag_version).options(joinedload(TI.dag_version))
 
     if dag_id != "~":
         dag = request.app.state.dag_bag.get_dag(dag_id)
@@ -428,6 +435,7 @@ def get_task_instances(
             executor,
             task_id,
             task_display_name_pattern,
+            version_number,
         ],
         order_by=order_by,
         offset=offset,

airflow/api_fastapi/core_api/routes/ui/structure.py

@@ -89,7 +89,7 @@ def structure_data(
         start_edges: list[dict] = []
         end_edges: list[dict] = []
 
-        for dependency_dag_id, dependencies in SerializedDagModel.get_dag_dependencies().items():
+        for dependency_dag_id, dependencies in sorted(SerializedDagModel.get_dag_dependencies().items()):
             for dependency in dependencies:
                 # Dependencies not related to `dag_id` are ignored
                 if dependency_dag_id != dag_id and dependency.target != dag_id:

airflow/ui/openapi-gen/queries/common.ts

@@ -982,6 +982,7 @@ export const UseTaskInstanceServiceGetMappedTaskInstancesKeyFn = (
     taskId,
     updatedAtGte,
     updatedAtLte,
+    versionNumber,
   }: {
     dagId: string;
     dagRunId: string;
@@ -1003,6 +1004,7 @@ export const UseTaskInstanceServiceGetMappedTaskInstancesKeyFn = (
     taskId: string;
     updatedAtGte?: string;
     updatedAtLte?: string;
+    versionNumber?: number[];
   },
   queryKey?: Array<unknown>,
 ) => [
@@ -1029,6 +1031,7 @@ export const UseTaskInstanceServiceGetMappedTaskInstancesKeyFn = (
       taskId,
       updatedAtGte,
       updatedAtLte,
+      versionNumber,
     },
   ]),
 ];
@@ -1191,6 +1194,7 @@ export const UseTaskInstanceServiceGetTaskInstancesKeyFn = (
     taskId,
     updatedAtGte,
     updatedAtLte,
+    versionNumber,
   }: {
     dagId: string;
     dagRunId: string;
@@ -1213,6 +1217,7 @@ export const UseTaskInstanceServiceGetTaskInstancesKeyFn = (
     taskId?: string;
     updatedAtGte?: string;
     updatedAtLte?: string;
+    versionNumber?: number[];
   },
   queryKey?: Array<unknown>,
 ) => [
@@ -1240,6 +1245,7 @@ export const UseTaskInstanceServiceGetTaskInstancesKeyFn = (
       taskId,
       updatedAtGte,
       updatedAtLte,
+      versionNumber,
     },
   ]),
 ];

airflow/ui/openapi-gen/queries/prefetch.ts

@@ -1328,6 +1328,7 @@ export const prefetchUseTaskInstanceServiceGetTaskInstance = (
  * @param data.pool
  * @param data.queue
  * @param data.executor
+ * @param data.versionNumber
  * @param data.limit
  * @param data.offset
  * @param data.orderBy
@@ -1357,6 +1358,7 @@ export const prefetchUseTaskInstanceServiceGetMappedTaskInstances = (
     taskId,
     updatedAtGte,
     updatedAtLte,
+    versionNumber,
   }: {
     dagId: string;
     dagRunId: string;
@@ -1378,6 +1380,7 @@ export const prefetchUseTaskInstanceServiceGetMappedTaskInstances = (
     taskId: string;
     updatedAtGte?: string;
     updatedAtLte?: string;
+    versionNumber?: number[];
   },
 ) =>
   queryClient.prefetchQuery({
@@ -1402,6 +1405,7 @@ export const prefetchUseTaskInstanceServiceGetMappedTaskInstances = (
       taskId,
       updatedAtGte,
       updatedAtLte,
+      versionNumber,
     }),
     queryFn: () =>
       TaskInstanceService.getMappedTaskInstances({
@@ -1425,6 +1429,7 @@ export const prefetchUseTaskInstanceServiceGetMappedTaskInstances = (
         taskId,
         updatedAtGte,
         updatedAtLte,
+        versionNumber,
       }),
   });
 /**
@@ -1611,6 +1616,7 @@ export const prefetchUseTaskInstanceServiceGetMappedTaskInstance = (
  * @param data.pool
  * @param data.queue
  * @param data.executor
+ * @param data.versionNumber
  * @param data.limit
  * @param data.offset
  * @param data.orderBy
@@ -1641,6 +1647,7 @@ export const prefetchUseTaskInstanceServiceGetTaskInstances = (
     taskId,
     updatedAtGte,
     updatedAtLte,
+    versionNumber,
   }: {
     dagId: string;
     dagRunId: string;
@@ -1663,6 +1670,7 @@ export const prefetchUseTaskInstanceServiceGetTaskInstances = (
     taskId?: string;
     updatedAtGte?: string;
     updatedAtLte?: string;
+    versionNumber?: number[];
   },
 ) =>
   queryClient.prefetchQuery({
@@ -1688,6 +1696,7 @@ export const prefetchUseTaskInstanceServiceGetTaskInstances = (
       taskId,
       updatedAtGte,
       updatedAtLte,
+      versionNumber,
     }),
     queryFn: () =>
       TaskInstanceService.getTaskInstances({
@@ -1712,6 +1721,7 @@ export const prefetchUseTaskInstanceServiceGetTaskInstances = (
         taskId,
         updatedAtGte,
         updatedAtLte,
+        versionNumber,
       }),
   });
 /**