diff --git a/airflow/api_fastapi/core_api/datamodels/xcom.py b/airflow/api_fastapi/core_api/datamodels/xcom.py index f874f8bdeed5a..1acccb702efaf 100644 --- a/airflow/api_fastapi/core_api/datamodels/xcom.py +++ b/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 diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index 4c3ebd13496da..486c371682fdb 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/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. diff --git a/airflow/api_fastapi/core_api/routes/public/xcom.py b/airflow/api_fastapi/core_api/routes/public/xcom.py index e1ef40685d39f..3da163f3e4033 100644 --- a/airflow/api_fastapi/core_api/routes/public/xcom.py +++ b/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) diff --git a/airflow/ui/openapi-gen/queries/common.ts b/airflow/ui/openapi-gen/queries/common.ts index 7a6d7eb780b24..dcb8ec827d072 100644 --- a/airflow/ui/openapi-gen/queries/common.ts +++ b/airflow/ui/openapi-gen/queries/common.ts @@ -1715,6 +1715,9 @@ export type TaskInstanceServicePatchTaskInstanceDryRun1MutationResult = Awaited< >; export type PoolServicePatchPoolMutationResult = Awaited>; export type PoolServiceBulkPoolsMutationResult = Awaited>; +export type XcomServiceUpdateXcomEntryMutationResult = Awaited< + ReturnType +>; export type VariableServicePatchVariableMutationResult = Awaited< ReturnType >; diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index d838c9ccf5e5d..5a71695254bfa 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/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, ...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, + ...options, + }); /** * Patch Variable * Update a variable by key. diff --git a/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow/ui/openapi-gen/requests/schemas.gen.ts index 9ea975f97f4fc..1d688c3442d27 100644 --- a/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/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; diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts index 69ab77337a1f9..585796aa2b2db 100644 --- a/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/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 { + 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. diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index bfe935523169b..a63ec17ffadcf 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/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: { diff --git a/tests/api_fastapi/core_api/routes/public/test_xcom.py b/tests/api_fastapi/core_api/routes/public/test_xcom.py index 7c6ce2c71e6ff..e3d9b3641a91b 100644 --- a/tests/api_fastapi/core_api/routes/public/test_xcom.py +++ b/tests/api_fastapi/core_api/routes/public/test_xcom.py @@ -579,3 +579,44 @@ def test_create_xcom_entry( assert current_data["task_id"] == task_id assert current_data["run_id"] == dag_run_id assert current_data["map_index"] == request_body.map_index + + +class TestPatchXComEntry(TestXComEndpoint): + @pytest.mark.parametrize( + "key, patch_body, expected_status, expected_detail", + [ + # Test case: Valid update, should return 200 OK + pytest.param( + TEST_XCOM_KEY, + {"value": "new_value"}, + 200, + None, + id="valid-xcom-update", + ), + # Test case: XCom entry does not exist, should return 404 + pytest.param( + TEST_XCOM_KEY, + {"value": "new_value", "map_index": -1}, + 404, + f"The XCom with key: `{TEST_XCOM_KEY}` with mentioned task instance doesn't exist.", + id="xcom-not-found", + ), + ], + ) + def test_patch_xcom_entry(self, key, patch_body, expected_status, expected_detail, test_client): + # Ensure the XCom entry exists before updating + if expected_status != 404: + self._create_xcom(TEST_XCOM_KEY, TEST_XCOM_VALUE) + new_value = XCom.serialize_value(patch_body["value"]) + + response = test_client.patch( + f"/public/dags/{TEST_DAG_ID}/dagRuns/{run_id}/taskInstances/{TEST_TASK_ID}/xcomEntries/{key}", + json=patch_body, + ) + + assert response.status_code == expected_status + + if expected_status == 200: + assert response.json()["value"] == XCom.serialize_value(new_value) + else: + assert response.json()["detail"] == expected_detail