Add update XCom endpoint in RestAPI

airflow/api_fastapi/core_api/datamodels/xcom.py

@@ -65,3 +65,10 @@ class XComCreateBody(StrictBaseModel):
     key: str
     value: Any
     map_index: int = -1
+
+
+class XComUpdateBody(StrictBaseModel):
+    """Payload serializer for updating an XCom entry."""
+
+    value: Any
+    map_index: int = -1

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

@@ -4380,6 +4380,80 @@ paths:
             application/json:
               schema:
                 $ref: '#/components/schemas/HTTPValidationError'
+    patch:
+      tags:
+      - XCom
+      summary: Update Xcom Entry
+      description: Update an existing XCom entry.
+      operationId: update_xcom_entry
+      parameters:
+      - name: dag_id
+        in: path
+        required: true
+        schema:
+          type: string
+          title: Dag Id
+      - name: task_id
+        in: path
+        required: true
+        schema:
+          type: string
+          title: Task Id
+      - name: dag_run_id
+        in: path
+        required: true
+        schema:
+          type: string
+          title: Dag Run Id
+      - name: xcom_key
+        in: path
+        required: true
+        schema:
+          type: string
+          title: Xcom Key
+      requestBody:
+        required: true
+        content:
+          application/json:
+            schema:
+              $ref: '#/components/schemas/XComUpdateBody'
+      responses:
+        '200':
+          description: Successful Response
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/XComResponseNative'
+        '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
+        '422':
+          description: Validation Error
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPValidationError'
   /public/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/xcomEntries:
     get:
       tags:
@@ -10747,3 +10821,17 @@ components:
       - value
       title: XComResponseString
       description: XCom response serializer with string return type.
+    XComUpdateBody:
+      properties:
+        value:
+          title: Value
+        map_index:
+          type: integer
+          title: Map Index
+          default: -1
+      additionalProperties: false
+      type: object
+      required:
+      - value
+      title: XComUpdateBody
+      description: Payload serializer for updating an XCom entry.

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

@@ -30,6 +30,7 @@
     XComCreateBody,
     XComResponseNative,
     XComResponseString,
+    XComUpdateBody,
 )
 from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
 from airflow.exceptions import TaskNotFound
@@ -222,3 +223,48 @@ def create_xcom_entry(
     )
 
     return XComResponseNative.model_validate(xcom)
+
+
+@xcom_router.patch(
+    "/{xcom_key}",
+    status_code=status.HTTP_200_OK,
+    responses=create_openapi_http_exception_doc(
+        [
+            status.HTTP_400_BAD_REQUEST,
+            status.HTTP_404_NOT_FOUND,
+        ]
+    ),
+)
+def update_xcom_entry(
+    dag_id: str,
+    task_id: str,
+    dag_run_id: str,
+    xcom_key: str,
+    patch_body: XComUpdateBody,
+    session: SessionDep,
+) -> XComResponseNative:
+    """Update an existing XCom entry."""
+    # Check if XCom entry exists
+    xcom_new_value = XCom.serialize_value(patch_body.value)
+    xcom_entry = session.scalar(
+        select(XCom)
+        .where(
+            XCom.dag_id == dag_id,
+            XCom.task_id == task_id,
+            XCom.run_id == dag_run_id,
+            XCom.key == xcom_key,
+            XCom.map_index == patch_body.map_index,
+        )
+        .limit(1)
+    )
+
+    if not xcom_entry:
+        raise HTTPException(
+            status.HTTP_404_NOT_FOUND,
+            f"The XCom with key: `{xcom_key}` with mentioned task instance doesn't exist.",
+        )
+
+    # Update XCom entry
+    xcom_entry.value = XCom.serialize_value(xcom_new_value)
+
+    return XComResponseNative.model_validate(xcom_entry)

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

@@ -1715,6 +1715,9 @@ export type TaskInstanceServicePatchTaskInstanceDryRun1MutationResult = Awaited<
 >;
 export type PoolServicePatchPoolMutationResult = Awaited<ReturnType<typeof PoolService.patchPool>>;
 export type PoolServiceBulkPoolsMutationResult = Awaited<ReturnType<typeof PoolService.bulkPools>>;
+export type XcomServiceUpdateXcomEntryMutationResult = Awaited<
+  ReturnType<typeof XcomService.updateXcomEntry>
+>;
 export type VariableServicePatchVariableMutationResult = Awaited<
   ReturnType<typeof VariableService.patchVariable>
 >;

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

@@ -52,6 +52,7 @@ import {
   TriggerDAGRunPostBody,
   VariableBody,
   XComCreateBody,
+  XComUpdateBody,
 } from "../requests/types.gen";
 import * as Common from "./common";
 
@@ -3982,6 +3983,61 @@ export const usePoolServiceBulkPools = <
     mutationFn: ({ requestBody }) => PoolService.bulkPools({ requestBody }) as unknown as Promise<TData>,
     ...options,
   });
+/**
+ * Update Xcom Entry
+ * Update an existing XCom entry.
+ * @param data The data for the request.
+ * @param data.dagId
+ * @param data.taskId
+ * @param data.dagRunId
+ * @param data.xcomKey
+ * @param data.requestBody
+ * @returns XComResponseNative Successful Response
+ * @throws ApiError
+ */
+export const useXcomServiceUpdateXcomEntry = <
+  TData = Common.XcomServiceUpdateXcomEntryMutationResult,
+  TError = unknown,
+  TContext = unknown,
+>(
+  options?: Omit<
+    UseMutationOptions<
+      TData,
+      TError,
+      {
+        dagId: string;
+        dagRunId: string;
+        requestBody: XComUpdateBody;
+        taskId: string;
+        xcomKey: string;
+      },
+      TContext
+    >,
+    "mutationFn"
+  >,
+) =>
+  useMutation<
+    TData,
+    TError,
+    {
+      dagId: string;
+      dagRunId: string;
+      requestBody: XComUpdateBody;
+      taskId: string;
+      xcomKey: string;
+    },
+    TContext
+  >({
+    mutationFn: ({ dagId, dagRunId, requestBody, taskId, xcomKey }) =>
+      XcomService.updateXcomEntry({
+        dagId,
+        dagRunId,
+        requestBody,
+        taskId,
+        xcomKey,
+      }) as unknown as Promise<TData>,
+    ...options,
+  });
 /**
  * Patch Variable
  * Update a variable by key.

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

@@ -6238,3 +6238,21 @@ export const $XComResponseString = {
   title: "XComResponseString",
   description: "XCom response serializer with string return type.",
 } as const;
+
+export const $XComUpdateBody = {
+  properties: {
+    value: {
+      title: "Value",
+    },
+    map_index: {
+      type: "integer",
+      title: "Map Index",
+      default: -1,
+    },
+  },
+  additionalProperties: false,
+  type: "object",
+  required: ["value"],
+  title: "XComUpdateBody",
+  description: "Payload serializer for updating an XCom entry.",
+} as const;

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

@@ -175,6 +175,8 @@ import type {
   GetProvidersResponse,
   GetXcomEntryData,
   GetXcomEntryResponse,
+  UpdateXcomEntryData,
+  UpdateXcomEntryResponse,
   GetXcomEntriesData,
   GetXcomEntriesResponse,
   CreateXcomEntryData,
@@ -2981,6 +2983,40 @@ export class XcomService {
     });
   }
 
+  /**
+   * Update Xcom Entry
+   * Update an existing XCom entry.
+   * @param data The data for the request.
+   * @param data.dagId
+   * @param data.taskId
+   * @param data.dagRunId
+   * @param data.xcomKey
+   * @param data.requestBody
+   * @returns XComResponseNative Successful Response
+   * @throws ApiError
+   */
+  public static updateXcomEntry(data: UpdateXcomEntryData): CancelablePromise<UpdateXcomEntryResponse> {
+    return __request(OpenAPI, {
+      method: "PATCH",
+      url: "/public/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/xcomEntries/{xcom_key}",
+      path: {
+        dag_id: data.dagId,
+        task_id: data.taskId,
+        dag_run_id: data.dagRunId,
+        xcom_key: data.xcomKey,
+      },
+      body: data.requestBody,
+      mediaType: "application/json",
+      errors: {
+        400: "Bad Request",
+        401: "Unauthorized",
+        403: "Forbidden",
+        404: "Not Found",
+        422: "Validation Error",
+      },
+    });
+  }
+
   /**
    * Get Xcom Entries
    * Get all XCom entries.

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

@@ -1542,6 +1542,14 @@ export type XComResponseString = {
   value: string | null;
 };
 
+/**
+ * Payload serializer for updating an XCom entry.
+ */
+export type XComUpdateBody = {
+  value: unknown;
+  map_index?: number;
+};
+
 export type NextRunAssetsData = {
   dagId: string;
 };
@@ -2320,6 +2328,16 @@ export type GetXcomEntryData = {
 
 export type GetXcomEntryResponse = XComResponseNative | XComResponseString;
 
+export type UpdateXcomEntryData = {
+  dagId: string;
+  dagRunId: string;
+  requestBody: XComUpdateBody;
+  taskId: string;
+  xcomKey: string;
+};
+
+export type UpdateXcomEntryResponse = XComResponseNative;
+
 export type GetXcomEntriesData = {
   dagId: string;
   dagRunId: string;
@@ -4721,6 +4739,35 @@ export type $OpenApiTs = {
         422: HTTPValidationError;
       };
     };
+    patch: {
+      req: UpdateXcomEntryData;
+      res: {
+        /**
+         * Successful Response
+         */
+        200: XComResponseNative;
+        /**
+         * Bad Request
+         */
+        400: HTTPExceptionResponse;
+        /**
+         * Unauthorized
+         */
+        401: HTTPExceptionResponse;
+        /**
+         * Forbidden
+         */
+        403: HTTPExceptionResponse;
+        /**
+         * Not Found
+         */
+        404: HTTPExceptionResponse;
+        /**
+         * Validation Error
+         */
+        422: HTTPValidationError;
+      };
+    };
   };
   "/public/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/xcomEntries": {
     get: {