Fix patch_task_instance endpoint (#50550)

Author: pierrejeambrun

airflow-core/src/airflow/api_fastapi/core_api/openapi/v1-rest-api-generated.yaml

@@ -4543,8 +4543,9 @@ paths:
         in: query
         required: false
         schema:
-          type: integer
-          default: -1
+          anyOf:
+          - type: integer
+          - type: 'null'
           title: Map Index
       - name: update_mask
         in: query
@@ -4568,7 +4569,7 @@ paths:
           content:
             application/json:
               schema:
-                $ref: '#/components/schemas/TaskInstanceResponse'
+                $ref: '#/components/schemas/TaskInstanceCollectionResponse'
         '401':
           content:
             application/json:
@@ -5256,7 +5257,9 @@ paths:
         in: path
         required: true
         schema:
-          type: integer
+          anyOf:
+          - type: integer
+          - type: 'null'
           title: Map Index
       - name: update_mask
         in: query
@@ -5280,7 +5283,7 @@ paths:
           content:
             application/json:
               schema:
-                $ref: '#/components/schemas/TaskInstanceResponse'
+                $ref: '#/components/schemas/TaskInstanceCollectionResponse'
         '401':
           content:
             application/json:
@@ -5853,7 +5856,9 @@ paths:
         in: path
         required: true
         schema:
-          type: integer
+          anyOf:
+          - type: integer
+          - type: 'null'
           title: Map Index
       - name: update_mask
         in: query
@@ -5940,8 +5945,9 @@ paths:
         in: query
         required: false
         schema:
-          type: integer
-          default: -1
+          anyOf:
+          - type: integer
+          - type: 'null'
           title: Map Index
       - name: update_mask
         in: query

airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py

@@ -803,10 +803,13 @@ def patch_task_instance_dry_run(
     dag_bag: DagBagDep,
     body: PatchTaskInstanceBody,
     session: SessionDep,
-    map_index: int = -1,
+    map_index: int | None = None,
     update_mask: list[str] | None = Query(None),
 ) -> TaskInstanceCollectionResponse:
     """Update a task instance dry_run mode."""
+    if map_index is None:
+        map_index = -1
+
     dag, ti, data = _patch_ti_validate_request(
         dag_id, dag_run_id, task_id, dag_bag, body, session, map_index, update_mask
     )
@@ -874,10 +877,13 @@ def patch_task_instance(
     body: PatchTaskInstanceBody,
     user: GetUserDep,
     session: SessionDep,
-    map_index: int = -1,
+    map_index: int | None = None,
     update_mask: list[str] | None = Query(None),
-) -> TaskInstanceResponse:
+) -> TaskInstanceCollectionResponse:
     """Update a task instance."""
+    if map_index is None:
+        map_index = -1
+
     dag, ti, data = _patch_ti_validate_request(
         dag_id, dag_run_id, task_id, dag_bag, body, session, map_index, update_mask
     )
@@ -924,7 +930,14 @@ def patch_task_instance(
                     ti.task_instance_note.user_id = user.get_id()
                 session.commit()
 
-    return TaskInstanceResponse.model_validate(ti)
+    return TaskInstanceCollectionResponse(
+        task_instances=[
+            TaskInstanceResponse.model_validate(
+                ti,
+            )
+        ],
+        total_entries=1,
+    )
 
 
 @task_instances_router.delete(

airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts

@@ -4066,7 +4066,7 @@ export const useDagServicePatchDag = <
  * @param data.requestBody
  * @param data.mapIndex
  * @param data.updateMask
- * @returns TaskInstanceResponse Successful Response
+ * @returns TaskInstanceCollectionResponse Successful Response
  * @throws ApiError
  */
 export const useTaskInstanceServicePatchTaskInstance = <
@@ -4125,7 +4125,7 @@ export const useTaskInstanceServicePatchTaskInstance = <
  * @param data.mapIndex
  * @param data.requestBody
  * @param data.updateMask
- * @returns TaskInstanceResponse Successful Response
+ * @returns TaskInstanceCollectionResponse Successful Response
  * @throws ApiError
  */
 export const useTaskInstanceServicePatchTaskInstanceByMapIndex = <

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

@@ -1923,7 +1923,7 @@ export class TaskInstanceService {
    * @param data.requestBody
    * @param data.mapIndex
    * @param data.updateMask
-   * @returns TaskInstanceResponse Successful Response
+   * @returns TaskInstanceCollectionResponse Successful Response
    * @throws ApiError
    */
   public static patchTaskInstance(data: PatchTaskInstanceData): CancelablePromise<PatchTaskInstanceResponse> {
@@ -2231,7 +2231,7 @@ export class TaskInstanceService {
    * @param data.mapIndex
    * @param data.requestBody
    * @param data.updateMask
-   * @returns TaskInstanceResponse Successful Response
+   * @returns TaskInstanceCollectionResponse Successful Response
    * @throws ApiError
    */
   public static patchTaskInstanceByMapIndex(

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

@@ -2270,13 +2270,13 @@ export type GetTaskInstanceResponse = TaskInstanceResponse;
 export type PatchTaskInstanceData = {
   dagId: string;
   dagRunId: string;
-  mapIndex?: number;
+  mapIndex?: number | null;
   requestBody: PatchTaskInstanceBody;
   taskId: string;
   updateMask?: Array<string> | null;
 };
 
-export type PatchTaskInstanceResponse = TaskInstanceResponse;
+export type PatchTaskInstanceResponse = TaskInstanceCollectionResponse;
 
 export type DeleteTaskInstanceData = {
   dagId: string;
@@ -2363,13 +2363,13 @@ export type GetMappedTaskInstanceResponse = TaskInstanceResponse;
 export type PatchTaskInstanceByMapIndexData = {
   dagId: string;
   dagRunId: string;
-  mapIndex: number;
+  mapIndex: number | null;
   requestBody: PatchTaskInstanceBody;
   taskId: string;
   updateMask?: Array<string> | null;
 };
 
-export type PatchTaskInstanceByMapIndexResponse = TaskInstanceResponse;
+export type PatchTaskInstanceByMapIndexResponse = TaskInstanceCollectionResponse;
 
 export type GetTaskInstancesData = {
   dagId: string;
@@ -2441,7 +2441,7 @@ export type PostClearTaskInstancesResponse = TaskInstanceCollectionResponse;
 export type PatchTaskInstanceDryRunByMapIndexData = {
   dagId: string;
   dagRunId: string;
-  mapIndex: number;
+  mapIndex: number | null;
   requestBody: PatchTaskInstanceBody;
   taskId: string;
   updateMask?: Array<string> | null;
@@ -2452,7 +2452,7 @@ export type PatchTaskInstanceDryRunByMapIndexResponse = TaskInstanceCollectionRe
 export type PatchTaskInstanceDryRunData = {
   dagId: string;
   dagRunId: string;
-  mapIndex?: number;
+  mapIndex?: number | null;
   requestBody: PatchTaskInstanceBody;
   taskId: string;
   updateMask?: Array<string> | null;
@@ -4307,7 +4307,7 @@ export type $OpenApiTs = {
         /**
          * Successful Response
          */
-        200: TaskInstanceResponse;
+        200: TaskInstanceCollectionResponse;
         /**
          * Bad Request
          */
@@ -4527,7 +4527,7 @@ export type $OpenApiTs = {
         /**
          * Successful Response
          */
-        200: TaskInstanceResponse;
+        200: TaskInstanceCollectionResponse;
         /**
          * Bad Request
          */