Add Patch Task Instance Summary

Author: pierrejeambrun

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

@@ -6002,6 +6002,99 @@ paths:
             application/json:
               schema:
                 $ref: '#/components/schemas/HTTPValidationError'
+  /api/v2/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/summary:
+    patch:
+      tags:
+      - Task Instance
+      summary: Patch Task Instance Summary
+      description: 'Update a task instance summary.
+
+
+        If the task is unmapped this is similar to the patch task instance endpoint.
+        For mapped task
+
+        updates all the mapped task instances at once.'
+      operationId: patch_task_instance_summary
+      security:
+      - OAuth2PasswordBearer: []
+      parameters:
+      - name: dag_id
+        in: path
+        required: true
+        schema:
+          type: string
+          title: Dag Id
+      - name: dag_run_id
+        in: path
+        required: true
+        schema:
+          type: string
+          title: Dag Run Id
+      - name: task_id
+        in: path
+        required: true
+        schema:
+          type: string
+          title: Task Id
+      - name: update_mask
+        in: query
+        required: false
+        schema:
+          anyOf:
+          - type: array
+            items:
+              type: string
+          - type: 'null'
+          title: Update Mask
+      requestBody:
+        required: true
+        content:
+          application/json:
+            schema:
+              $ref: '#/components/schemas/PatchTaskInstanceBody'
+      responses:
+        '200':
+          description: Successful Response
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/TaskInstanceCollectionResponse'
+        '401':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Unauthorized
+        '403':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Forbidden
+        '400':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Bad Request
+        '404':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Not Found
+        '409':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Conflict
+        '422':
+          description: Validation Error
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPValidationError'
   /api/v2/dags/{dag_id}/tasks:
     get:
       tags:

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

@@ -24,7 +24,6 @@
 from fastapi.exceptions import RequestValidationError
 from pydantic import ValidationError
 from sqlalchemy import or_, select
-from sqlalchemy.exc import MultipleResultsFound
 from sqlalchemy.orm import joinedload
 from sqlalchemy.sql.selectable import Select
 
@@ -736,9 +735,9 @@ def _patch_ti_validate_request(
     dag_bag: DagBagDep,
     body: PatchTaskInstanceBody,
     session: SessionDep,
-    map_index: int = -1,
+    map_index: int | None = -1,
     update_mask: list[str] | None = Query(None),
-) -> tuple[DAG, TI, dict]:
+) -> tuple[DAG, list[TI], dict]:
     dag = dag_bag.get_dag(dag_id)
     if not dag:
         raise HTTPException(status.HTTP_404_NOT_FOUND, f"DAG {dag_id} not found")
@@ -752,20 +751,15 @@ def _patch_ti_validate_request(
         .join(TI.dag_run)
         .options(joinedload(TI.rendered_task_instance_fields))
     )
-    query = query.where(TI.map_index == map_index)
+    if map_index is not None:
+        query = query.where(TI.map_index == map_index)
 
-    try:
-        ti = session.scalar(query)
-    except MultipleResultsFound:
-        raise HTTPException(
-            status.HTTP_400_BAD_REQUEST,
-            "Multiple task instances found. As the TI is mapped, add the map_index value to the URL",
-        )
+    tis = session.scalars(query).all()
 
     err_msg_404 = (
         f"The Task Instance with dag_id: `{dag_id}`, run_id: `{dag_run_id}`, task_id: `{task_id}` and map_index: `{map_index}` was not found",
     )
-    if ti is None:
+    if len(tis) == 0:
         raise HTTPException(status.HTTP_404_NOT_FOUND, err_msg_404)
 
     fields_to_update = body.model_fields_set
@@ -777,7 +771,7 @@ def _patch_ti_validate_request(
         except ValidationError as e:
             raise RequestValidationError(errors=e.errors())
 
-    return dag, ti, body.model_dump(include=fields_to_update, by_alias=True)
+    return dag, list(tis), body.model_dump(include=fields_to_update, by_alias=True)
 
 
 @task_instances_router.patch(
@@ -807,21 +801,16 @@ def patch_task_instance_dry_run(
     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, tis, data = _patch_ti_validate_request(
         dag_id, dag_run_id, task_id, dag_bag, body, session, map_index, update_mask
     )
 
-    tis: list[TI] = []
-
     if data.get("new_state"):
         tis = (
             dag.set_task_instance_state(
                 task_id=task_id,
                 run_id=dag_run_id,
-                map_indexes=[map_index],
+                map_indexes=[map_index] if map_index is not None else None,
                 state=data["new_state"],
                 upstream=body.include_upstream,
                 downstream=body.include_downstream,
@@ -833,9 +822,6 @@ def patch_task_instance_dry_run(
             or []
         )
 
-    elif "note" in data:
-        tis = [ti]
-
     return TaskInstanceCollectionResponse(
         task_instances=[
             TaskInstanceResponse.model_validate(
@@ -881,19 +867,16 @@ def patch_task_instance(
     update_mask: list[str] | None = Query(None),
 ) -> TaskInstanceCollectionResponse:
     """Update a task instance."""
-    if map_index is None:
-        map_index = -1
-
-    dag, ti, data = _patch_ti_validate_request(
+    dag, tis, data = _patch_ti_validate_request(
         dag_id, dag_run_id, task_id, dag_bag, body, session, map_index, update_mask
     )
 
     for key, _ in data.items():
         if key == "new_state":
-            tis: list[TI] = dag.set_task_instance_state(
+            tis = dag.set_task_instance_state(
                 task_id=task_id,
                 run_id=dag_run_id,
-                map_indexes=[map_index],
+                map_indexes=[map_index] if map_index is not None else None,
                 state=data["new_state"],
                 upstream=body.include_upstream,
                 downstream=body.include_downstream,
@@ -906,37 +889,39 @@ def patch_task_instance(
                 raise HTTPException(
                     status.HTTP_409_CONFLICT, f"Task id {task_id} is already in {data['new_state']} state"
                 )
-            ti = tis[0] if isinstance(tis, list) else tis
-            try:
-                if data["new_state"] == TaskInstanceState.SUCCESS:
-                    get_listener_manager().hook.on_task_instance_success(
-                        previous_state=None, task_instance=ti
-                    )
-                elif data["new_state"] == TaskInstanceState.FAILED:
-                    get_listener_manager().hook.on_task_instance_failed(
-                        previous_state=None,
-                        task_instance=ti,
-                        error=f"TaskInstance's state was manually set to `{TaskInstanceState.FAILED}`.",
-                    )
-            except Exception:
-                log.exception("error calling listener")
+
+            for ti in tis:
+                try:
+                    if data["new_state"] == TaskInstanceState.SUCCESS:
+                        get_listener_manager().hook.on_task_instance_success(
+                            previous_state=None, task_instance=ti
+                        )
+                    elif data["new_state"] == TaskInstanceState.FAILED:
+                        get_listener_manager().hook.on_task_instance_failed(
+                            previous_state=None,
+                            task_instance=ti,
+                            error=f"TaskInstance's state was manually set to `{TaskInstanceState.FAILED}`.",
+                        )
+                except Exception:
+                    log.exception("error calling listener")
 
         elif key == "note":
-            if update_mask or body.note is not None:
-                if ti.task_instance_note is None:
-                    ti.note = (body.note, user.get_id())
-                else:
-                    ti.task_instance_note.content = body.note
-                    ti.task_instance_note.user_id = user.get_id()
-                session.commit()
+            for ti in tis:
+                if update_mask or body.note is not None:
+                    if ti.task_instance_note is None:
+                        ti.note = (body.note, user.get_id())
+                    else:
+                        ti.task_instance_note.content = body.note
+                        ti.task_instance_note.user_id = user.get_id()
 
     return TaskInstanceCollectionResponse(
         task_instances=[
             TaskInstanceResponse.model_validate(
                 ti,
             )
+            for ti in tis
         ],
-        total_entries=1,
+        total_entries=len(tis),
     )
 
 

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

@@ -1918,6 +1918,9 @@ export type TaskInstanceServicePatchTaskInstanceDryRunByMapIndexMutationResult =
 export type TaskInstanceServicePatchTaskInstanceDryRunMutationResult = Awaited<
   ReturnType<typeof TaskInstanceService.patchTaskInstanceDryRun>
 >;
+export type TaskInstanceServicePatchTaskInstanceSummaryMutationResult = Awaited<
+  ReturnType<typeof TaskInstanceService.patchTaskInstanceSummary>
+>;
 export type PoolServicePatchPoolMutationResult = Awaited<ReturnType<typeof PoolService.patchPool>>;
 export type PoolServiceBulkPoolsMutationResult = Awaited<ReturnType<typeof PoolService.bulkPools>>;
 export type XcomServiceUpdateXcomEntryMutationResult = Awaited<

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

@@ -4292,6 +4292,64 @@ export const useTaskInstanceServicePatchTaskInstanceDryRun = <
       }) as unknown as Promise<TData>,
     ...options,
   });
+/**
+ * Patch Task Instance Summary
+ * Update a task instance summary.
+ *
+ * If the task is unmapped this is similar to the patch task instance endpoint. For mapped task
+ * updates all the mapped task instances at once.
+ * @param data The data for the request.
+ * @param data.dagId
+ * @param data.dagRunId
+ * @param data.taskId
+ * @param data.requestBody
+ * @param data.updateMask
+ * @returns TaskInstanceCollectionResponse Successful Response
+ * @throws ApiError
+ */
+export const useTaskInstanceServicePatchTaskInstanceSummary = <
+  TData = Common.TaskInstanceServicePatchTaskInstanceSummaryMutationResult,
+  TError = unknown,
+  TContext = unknown,
+>(
+  options?: Omit<
+    UseMutationOptions<
+      TData,
+      TError,
+      {
+        dagId: string;
+        dagRunId: string;
+        requestBody: PatchTaskInstanceBody;
+        taskId: string;
+        updateMask?: string[];
+      },
+      TContext
+    >,
+    "mutationFn"
+  >,
+) =>
+  useMutation<
+    TData,
+    TError,
+    {
+      dagId: string;
+      dagRunId: string;
+      requestBody: PatchTaskInstanceBody;
+      taskId: string;
+      updateMask?: string[];
+    },
+    TContext
+  >({
+    mutationFn: ({ dagId, dagRunId, requestBody, taskId, updateMask }) =>
+      TaskInstanceService.patchTaskInstanceSummary({
+        dagId,
+        dagRunId,
+        requestBody,
+        taskId,
+        updateMask,
+      }) as unknown as Promise<TData>,
+    ...options,
+  });
 /**
  * Patch Pool
  * Update a Pool.