From 4adc04708e7f61de94955800d76a40a044495a30 Mon Sep 17 00:00:00 2001 From: ColtenOuO Date: Wed, 26 Aug 2026 14:57:00 +0000 Subject: [PATCH 1/4] Add a batch XCom lookup endpoint to the Execution API Lays the groundwork for resolving multiple XComs in one round trip instead of one per lookup. Kept separate from the caller that will actually use it so this additive, backward-compatible piece can be reviewed on its own before the behavior change that depends on it. --- .../execution_api/datamodels/xcom.py | 31 +++++ .../execution_api/routes/__init__.py | 1 + .../api_fastapi/execution_api/routes/xcoms.py | 124 ++++++++++++++---- .../execution_api/versions/__init__.py | 7 +- .../execution_api/versions/v2026_10_30.py | 12 ++ .../execution_api/versions/head/test_xcoms.py | 119 +++++++++++++++++ task-sdk/src/airflow/sdk/api/client.py | 21 +++ .../airflow/sdk/api/datamodels/_generated.py | 38 ++++++ task-sdk/src/airflow/sdk/exceptions.py | 3 + .../src/airflow/sdk/execution_time/comms.py | 40 ++++++ .../sdk/execution_time/request_handlers.py | 16 +++ .../sdk/execution_time/schema/schema.json | 112 ++++++++++++++++ .../schema/versions/__init__.py | 3 +- .../schema/versions/v2026_10_30.py | 15 +++ .../airflow/sdk/execution_time/supervisor.py | 4 + task-sdk/tests/task_sdk/api/test_client.py | 61 +++++++++ .../execution_time/test_supervisor.py | 40 ++++++ 17 files changed, 618 insertions(+), 29 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/xcom.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/xcom.py index ea62e8de25649..434e943864516 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/xcom.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/xcom.py @@ -40,3 +40,34 @@ class XComSequenceSliceResponse(RootModel): """XCom schema with minimal structure for slice-based access.""" root: list[JsonValue] + + +class XComBatchItemRequest(BaseModel): + """One XCom lookup within a batch request.""" + + task_id: str + key: str + map_index: int = -1 + + +class XComBatchRequestBody(BaseModel): + """Body for a batch XCom lookup, scoped to a single dag_id/run_id.""" + + items: list[XComBatchItemRequest] + + +class XComBatchItemResponse(BaseModel): + """One XCom lookup result within a batch response.""" + + task_id: str + key: str + map_index: int + found: bool + value: JsonValue = None + """The returned XCom value in a JSON-compatible format. Meaningless when ``found`` is False.""" + + +class XComBatchResponse(BaseModel): + """Batch XCom lookup response, ordered the same as the request's items.""" + + items: list[XComBatchItemResponse] diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/__init__.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/__init__.py index 7b19f3ddd3055..87fed3436e227 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/__init__.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/__init__.py @@ -61,6 +61,7 @@ ) authenticated_router.include_router(variables.router, prefix="/variables", tags=["Variables"]) authenticated_router.include_router(xcoms.router, prefix="/xcoms", tags=["XComs"]) +authenticated_router.include_router(xcoms.batch_router, prefix="/xcoms", tags=["XComs"]) authenticated_router.include_router(hitl.router, prefix="/hitlDetails", tags=["Human in the Loop"]) authenticated_router.include_router(task_state_store.router, prefix="/store/ti", tags=["Task State Store"]) authenticated_router.include_router( diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py index 0cb6ccb23ce54..b1d00c39e2462 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py @@ -22,12 +22,15 @@ from fastapi import APIRouter, Body, Depends, HTTPException, Path, Query, Request, Response, status from pydantic import JsonValue -from sqlalchemy import delete +from sqlalchemy import delete, select, tuple_ from sqlalchemy.sql.selectable import Select from airflow.api_fastapi.common.db.common import SessionDep from airflow.api_fastapi.core_api.base import BaseModel from airflow.api_fastapi.execution_api.datamodels.xcom import ( + XComBatchItemResponse, + XComBatchRequestBody, + XComBatchResponse, XComResponse, XComSequenceIndexResponse, XComSequenceSliceResponse, @@ -38,23 +41,14 @@ from airflow.utils.db import get_query_count -def has_xcom_access( - dag_id: str, - run_id: str, - task_id: str, - xcom_key: Annotated[str, Path(alias="key", min_length=1)], - request: Request, - session: SessionDep, - token=CurrentTIToken, -) -> bool: +def _check_dag_team_access(dag_id: str, *, write: bool, session: SessionDep, token) -> None: """ - Check whether the requesting task may access the XCom for ``dag_id``. + Raise 403 unless the requesting task's team may access XComs for ``dag_id``. In multi-team mode, XCom access is scoped by team ownership (resolved via the ``dag -> bundle -> team`` chain). There is no cross-team XCom sharing: - * reads (``GET``/``HEAD``) are allowed for the requester's own team or for - global (teamless) dags; + * reads are allowed for the requester's own team or for global (teamless) dags; * writes and deletes are allowed only for the requester's own team; a team task may not mutate a global dag's XCom, mirroring how team-scoped Variables and Connections behave. @@ -67,18 +61,8 @@ def has_xcom_access( """ from airflow.configuration import conf - write = request.method not in {"GET", "HEAD", "OPTIONS"} - - log.debug( - "Checking %s XCom access for task instance '%s' to XCom '%s' on dag '%s'", - "write" if write else "read", - token.id, - xcom_key, - dag_id, - ) - if not conf.getboolean("core", "multi_team"): - return True + return from airflow.api_fastapi.execution_api.security import ( _team_name_for_dag_stmt, @@ -90,10 +74,10 @@ def has_xcom_access( # Same team (including a teamless task accessing a global, teamless dag) is always allowed. if target_team == requester_team: - return True + return # Reads may additionally reach global (teamless) dags; writes and deletes may not. if not write and target_team is None: - return True + return raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, @@ -104,6 +88,40 @@ def has_xcom_access( ) +def has_xcom_access( + dag_id: str, + run_id: str, + task_id: str, + xcom_key: Annotated[str, Path(alias="key", min_length=1)], + request: Request, + session: SessionDep, + token=CurrentTIToken, +) -> bool: + """Check whether the requesting task may access the XCom for ``dag_id``. See ``_check_dag_team_access``.""" + write = request.method not in {"GET", "HEAD", "OPTIONS"} + log.debug( + "Checking %s XCom access for task instance '%s' to XCom '%s' on dag '%s'", + "write" if write else "read", + token.id, + xcom_key, + dag_id, + ) + _check_dag_team_access(dag_id, write=write, session=session, token=token) + return True + + +def has_xcom_batch_access( + dag_id: str, + run_id: str, + session: SessionDep, + token=CurrentTIToken, +) -> bool: + """Check whether the requesting task may batch-read XComs for ``dag_id``. See ``_check_dag_team_access``.""" + log.debug("Checking read XCom batch access for task instance '%s' on dag '%s'", token.id, dag_id) + _check_dag_team_access(dag_id, write=False, session=session, token=token) + return True + + router = APIRouter( responses={ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"}, @@ -482,3 +500,57 @@ def delete_xcom( ) session.execute(query) return {"message": f"XCom with key: {key} successfully deleted."} + + +# A separate router (registered directly by routes/__init__.py, not merged via +# router.include_router) because ``router`` above carries ``has_xcom_access`` as a +# constructor-level dependency, which needs ``task_id``/``key`` path params this +# route doesn't have. Merging would leak that dependency onto this route's OpenAPI +# operation as an unresolvable path parameter. +batch_router = APIRouter(dependencies=[Depends(has_xcom_batch_access)]) + + +@batch_router.post( + "/{dag_id}/{run_id}/batch", + description="Look up multiple XCom values in one request, scoped to a single dag run", +) +def get_xcom_batch( + dag_id: str, + run_id: str, + body: XComBatchRequestBody, + session: SessionDep, +) -> XComBatchResponse: + """ + Batch-fetch XComs from the database - not other XCom Backends. + + Collapses what would otherwise be one Execution API round trip per requested + XCom (e.g. one per XComArg kwarg of a mapped task's ``.expand()``) into a + single request. Missing items are reported via ``found=False`` rather than + failing the whole batch, since a partial miss (e.g. an upstream that hasn't + pushed yet) is an expected outcome for a batch, not an error. + """ + if not body.items: + return XComBatchResponse(items=[]) + + requested = [(item.task_id, item.key, item.map_index) for item in body.items] + query = select(XComModel.task_id, XComModel.key, XComModel.map_index, XComModel.value).where( + XComModel.dag_id == dag_id, + XComModel.run_id == run_id, + tuple_(XComModel.task_id, XComModel.key, XComModel.map_index).in_(requested), + ) + found_values = { + (task_id, key, map_index): value for task_id, key, map_index, value in session.execute(query) + } + + return XComBatchResponse( + items=[ + XComBatchItemResponse( + task_id=item.task_id, + key=item.key, + map_index=item.map_index, + found=(item.task_id, item.key, item.map_index) in found_values, + value=found_values.get((item.task_id, item.key, item.map_index)), + ) + for item in body.items + ] + ) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py index d56ec735c8f13..9517b9dd7559c 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py @@ -51,11 +51,14 @@ AddTeamNameField, AddVariableKeysEndpoint, ) -from airflow.api_fastapi.execution_api.versions.v2026_10_30 import AddArgBindingsToTIRunContext +from airflow.api_fastapi.execution_api.versions.v2026_10_30 import ( + AddArgBindingsToTIRunContext, + AddXComBatchEndpoint, +) bundle = VersionBundle( HeadVersion(), - Version("2026-10-30", AddArgBindingsToTIRunContext), + Version("2026-10-30", AddArgBindingsToTIRunContext, AddXComBatchEndpoint), Version( "2026-06-30", AddVariableKeysEndpoint, diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py index 1c85aed252c06..3e0e40de3f503 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py @@ -19,14 +19,26 @@ from cadwyn import ( ResponseInfo, + VersionChange, VersionChangeWithSideEffects, convert_response_to_previous_version_for, + endpoint, schema, ) from airflow.api_fastapi.execution_api.datamodels.taskinstance import TIRunContext +class AddXComBatchEndpoint(VersionChange): + """Add a batch XCom lookup endpoint that resolves multiple XComs in one request.""" + + description = __doc__ + + instructions_to_migrate_to_previous_version = ( + endpoint("/xcoms/{dag_id}/{run_id}/batch", ["POST"]).didnt_exist, + ) + + class AddArgBindingsToTIRunContext(VersionChangeWithSideEffects): """Add the ``arg_bindings`` argument-binding spec for stub (foreign-runtime) tasks.""" diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py index ae476d0344ca0..5e95afcb06455 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py @@ -618,6 +618,92 @@ def test_xcom_delete_endpoint(self, client, create_task_instance, session): assert xcom_ti is not None +class TestXComsBatchEndpoint: + def test_batch_multiple_items_found(self, client, create_task_instance, session): + """A batch request returns each requested item, ordered the same as the request.""" + ti = create_task_instance() + session.add_all( + [ + XComModel( + key="a", + value="value_a", + dag_run_id=ti.dag_run.id, + run_id=ti.run_id, + task_id=ti.task_id, + dag_id=ti.dag_id, + ), + XComModel( + key="b", + value="value_b", + dag_run_id=ti.dag_run.id, + run_id=ti.run_id, + task_id=ti.task_id, + dag_id=ti.dag_id, + ), + ] + ) + session.commit() + + response = client.post( + f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/batch", + json={ + "items": [ + {"task_id": ti.task_id, "key": "a"}, + {"task_id": ti.task_id, "key": "b"}, + ] + }, + ) + + assert response.status_code == 200, response.json() + assert response.json() == { + "items": [ + {"task_id": ti.task_id, "key": "a", "map_index": -1, "found": True, "value": "value_a"}, + {"task_id": ti.task_id, "key": "b", "map_index": -1, "found": True, "value": "value_b"}, + ] + } + + def test_batch_partial_miss_reported_not_found(self, client, create_task_instance, session): + """A missing item is reported as found=False instead of failing the whole batch.""" + ti = create_task_instance() + session.add( + XComModel( + key="a", + value="value_a", + dag_run_id=ti.dag_run.id, + run_id=ti.run_id, + task_id=ti.task_id, + dag_id=ti.dag_id, + ) + ) + session.commit() + + response = client.post( + f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/batch", + json={ + "items": [ + {"task_id": ti.task_id, "key": "a"}, + {"task_id": ti.task_id, "key": "missing"}, + ] + }, + ) + + assert response.status_code == 200, response.json() + assert response.json() == { + "items": [ + {"task_id": ti.task_id, "key": "a", "map_index": -1, "found": True, "value": "value_a"}, + {"task_id": ti.task_id, "key": "missing", "map_index": -1, "found": False, "value": None}, + ] + } + + def test_batch_empty_request(self, client, create_task_instance): + ti = create_task_instance() + + response = client.post(f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/batch", json={"items": []}) + + assert response.status_code == 200, response.json() + assert response.json() == {"items": []} + + class TestXComTeamAccess: """Multi-team isolation for the Execution API XCom routes (no cross-team sharing).""" @@ -766,3 +852,36 @@ def test_teamless_requester_scoping(self, client, session, dag_maker): assert forbidden.status_code == 403, forbidden.json() assert allowed.status_code == 200, allowed.json() + + def test_batch_cross_team_access_forbidden(self, client, exec_app, session, dag_maker): + """A task cannot batch-read another team's XComs, scoped by the batch's dag_id.""" + _, requester_ti = self._make_dag(session, dag_maker, f"req_{uuid4().hex}", "team_a") + target_dag = f"tgt_{uuid4().hex}" + target_dr, _ = self._make_dag(session, dag_maker, target_dag, "team_b") + self._insert_xcom(session, target_dr, target_dag) + self._authenticate_as(exec_app, requester_ti.id) + + with conf_vars({("core", "multi_team"): "True"}): + response = client.post( + f"/execution/xcoms/{target_dag}/run1/batch", json={"items": [{"task_id": "task", "key": "k"}]} + ) + + assert response.status_code == 403, response.json() + assert response.json()["detail"]["reason"] == "access_denied" + + def test_batch_same_team_read_allowed(self, client, exec_app, session, dag_maker): + """A task may batch-read XComs within its own team.""" + dag_id = f"dag_{uuid4().hex}" + dag_run, ti = self._make_dag(session, dag_maker, dag_id, "team_a") + self._insert_xcom(session, dag_run, dag_id, key="k", value="v") + self._authenticate_as(exec_app, ti.id) + + with conf_vars({("core", "multi_team"): "True"}): + response = client.post( + f"/execution/xcoms/{dag_id}/run1/batch", json={"items": [{"task_id": "task", "key": "k"}]} + ) + + assert response.status_code == 200, response.json() + assert response.json() == { + "items": [{"task_id": "task", "key": "k", "map_index": -1, "found": True, "value": "v"}] + } diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index a0c0a30de088f..8acb69cc5278e 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -85,6 +85,9 @@ VariableKeysResponse, VariablePostBody, VariableResponse, + XComBatchItemRequest, + XComBatchRequestBody, + XComBatchResponse, XComResponse, XComSequenceIndexResponse, XComSequenceSliceResponse, @@ -716,6 +719,24 @@ def get_sequence_slice( resp = self.client.get(f"xcoms/{dag_id}/{run_id}/{task_id}/{key}/slice", params=params) return XComSequenceSliceResponse.model_validate_json(resp.read()) + def get_batch( + self, + dag_id: str, + run_id: str, + items: list[XComBatchItemRequest], + ) -> XComBatchResponse | ErrorResponse: + """Look up multiple XComs, scoped to one dag run, in a single API call.""" + try: + resp = self.client.post( + f"xcoms/{dag_id}/{run_id}/batch", + content=XComBatchRequestBody(items=items).model_dump_json(), + ) + except ServerResponseError as e: + if e.response.status_code == HTTPStatus.NOT_FOUND: + return ErrorResponse(error=ErrorType.XCOM_BATCH_NOT_SUPPORTED, detail={}) + raise + return XComBatchResponse.model_validate_json(resp.read()) + class TaskStateStoreOperations: __slots__ = ("client",) diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index f31794560ef5c..24244aeac23b0 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -513,6 +513,44 @@ class VariableResponse(BaseModel): value: Annotated[str | None, Field(title="Value")] +class XComBatchItemRequest(BaseModel): + """ + One XCom lookup within a batch request. + """ + + task_id: Annotated[str, Field(title="Task Id")] + key: Annotated[str, Field(title="Key")] + map_index: Annotated[int | None, Field(title="Map Index")] = -1 + + +class XComBatchItemResponse(BaseModel): + """ + One XCom lookup result within a batch response. + """ + + task_id: Annotated[str, Field(title="Task Id")] + key: Annotated[str, Field(title="Key")] + map_index: Annotated[int, Field(title="Map Index")] + found: Annotated[bool, Field(title="Found")] + value: JsonValue | None = None + + +class XComBatchRequestBody(BaseModel): + """ + Body for a batch XCom lookup, scoped to a single dag_id/run_id. + """ + + items: Annotated[list[XComBatchItemRequest], Field(title="Items")] + + +class XComBatchResponse(BaseModel): + """ + Batch XCom lookup response, ordered the same as the request's items. + """ + + items: Annotated[list[XComBatchItemResponse], Field(title="Items")] + + class XComResponse(BaseModel): """ XCom schema for responses with fields that are needed for Runtime. diff --git a/task-sdk/src/airflow/sdk/exceptions.py b/task-sdk/src/airflow/sdk/exceptions.py index 6f43d5421ecf2..4a8d164604068 100644 --- a/task-sdk/src/airflow/sdk/exceptions.py +++ b/task-sdk/src/airflow/sdk/exceptions.py @@ -90,6 +90,9 @@ class ErrorType(enum.Enum): CONNECTION_NOT_FOUND = "CONNECTION_NOT_FOUND" VARIABLE_NOT_FOUND = "VARIABLE_NOT_FOUND" XCOM_NOT_FOUND = "XCOM_NOT_FOUND" + # The batch XCom lookup endpoint doesn't exist on this API server (older server, + # newer Task SDK). Callers fall back to resolving XComs one at a time. + XCOM_BATCH_NOT_SUPPORTED = "XCOM_BATCH_NOT_SUPPORTED" ASSET_NOT_FOUND = "ASSET_NOT_FOUND" TASK_STORE_NOT_FOUND = "TASK_STORE_NOT_FOUND" ASSET_STORE_NOT_FOUND = "ASSET_STORE_NOT_FOUND" diff --git a/task-sdk/src/airflow/sdk/execution_time/comms.py b/task-sdk/src/airflow/sdk/execution_time/comms.py index 4a26f56e297dc..c885d3976144b 100644 --- a/task-sdk/src/airflow/sdk/execution_time/comms.py +++ b/task-sdk/src/airflow/sdk/execution_time/comms.py @@ -96,6 +96,7 @@ TriggerDAGRunPayload, UpdateHITLDetailPayload, VariableResponse, + XComBatchResponse, XComResponse, XComSequenceIndexResponse, XComSequenceSliceResponse, @@ -561,6 +562,28 @@ class XComCountResponse(BaseModel): type: Literal["XComCountResponse"] = "XComCountResponse" +class XComBatchResultItem(BaseModel): + task_id: str + key: str + map_index: int + found: bool + value: JsonValue = None + + +class XComBatchResult(BaseModel): + """Response to GetXComBatch request.""" + + items: list[XComBatchResultItem] + type: Literal["XComBatchResult"] = "XComBatchResult" + + @classmethod + def from_response(cls, response: XComBatchResponse) -> XComBatchResult: + return cls( + items=[XComBatchResultItem(**item.model_dump()) for item in response.items], + type="XComBatchResult", + ) + + class XComSequenceIndexResult(BaseModel): root: JsonValue type: Literal["XComSequenceIndexResult"] = "XComSequenceIndexResult" @@ -841,6 +864,7 @@ def from_api_response(cls, dag_response: DagResponse) -> DagResult: | VariableKeysResult | XComCountResponse | XComResult + | XComBatchResult | XComSequenceIndexResult | XComSequenceSliceResult | InactiveAssetsResult @@ -927,6 +951,21 @@ class GetXComCount(BaseModel): type: Literal["GetXComCount"] = "GetXComCount" +class GetXComBatchItem(BaseModel): + task_id: str + key: str + map_index: int = -1 + + +class GetXComBatch(BaseModel): + """Look up multiple XComs, scoped to a single dag run, in one request.""" + + dag_id: str + run_id: str + items: list[GetXComBatchItem] + type: Literal["GetXComBatch"] = "GetXComBatch" + + class GetXComSequenceItem(BaseModel): key: str dag_id: str @@ -1282,6 +1321,7 @@ class GetDag(BaseModel): | GetVariable | GetVariableKeys | GetXCom + | GetXComBatch | GetXComCount | GetXComSequenceItem | GetXComSequenceSlice diff --git a/task-sdk/src/airflow/sdk/execution_time/request_handlers.py b/task-sdk/src/airflow/sdk/execution_time/request_handlers.py index a31596b3d2b82..a3ee11c2980fe 100644 --- a/task-sdk/src/airflow/sdk/execution_time/request_handlers.py +++ b/task-sdk/src/airflow/sdk/execution_time/request_handlers.py @@ -36,6 +36,8 @@ DagRunStateResponse, TaskStatesResponse, VariableResponse, + XComBatchItemRequest, + XComBatchResponse, XComResponse, XComSequenceIndexResponse, XComSequenceSliceResponse, @@ -62,6 +64,7 @@ GetVariable, GetVariableKeys, GetXCom, + GetXComBatch, GetXComCount, GetXComSequenceItem, GetXComSequenceSlice, @@ -74,6 +77,7 @@ TaskStatesResult, VariableKeysResult, VariableResult, + XComBatchResult, XComResult, XComSequenceIndexResult, XComSequenceSliceResult, @@ -284,6 +288,18 @@ def handle_get_xcom(client: Client, msg: GetXCom) -> tuple[BaseModel | None, dic return xcom, {} +def handle_get_xcom_batch(client: Client, msg: GetXComBatch) -> tuple[BaseModel | None, dict[str, bool]]: + """Fetch multiple XComs in one request and normalize the result for supervisor response handling.""" + items = [ + XComBatchItemRequest(task_id=item.task_id, key=item.key, map_index=item.map_index) + for item in msg.items + ] + result = client.xcoms.get_batch(msg.dag_id, msg.run_id, items) + if isinstance(result, XComBatchResponse): + return XComBatchResult.from_response(result), {} + return result, {} + + def handle_get_asset_state_store_by_name( client: Client, msg: GetAssetStateStoreByName ) -> tuple[BaseModel | None, dict[str, bool]]: diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json index 85f8b7cfa4e2b..51d2c26b9222e 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -1861,6 +1861,7 @@ "CONNECTION_NOT_FOUND", "VARIABLE_NOT_FOUND", "XCOM_NOT_FOUND", + "XCOM_BATCH_NOT_SUPPORTED", "ASSET_NOT_FOUND", "TASK_STORE_NOT_FOUND", "ASSET_STORE_NOT_FOUND", @@ -2852,6 +2853,62 @@ "title": "GetXCom", "type": "object" }, + "GetXComBatch": { + "description": "Look up multiple XComs, scoped to a single dag run, in one request.", + "properties": { + "dag_id": { + "title": "Dag Id", + "type": "string" + }, + "run_id": { + "title": "Run Id", + "type": "string" + }, + "items": { + "items": { + "$ref": "#/$defs/GetXComBatchItem" + }, + "title": "Items", + "type": "array" + }, + "type": { + "const": "GetXComBatch", + "default": "GetXComBatch", + "title": "Type", + "type": "string" + } + }, + "required": [ + "dag_id", + "run_id", + "items" + ], + "title": "GetXComBatch", + "type": "object" + }, + "GetXComBatchItem": { + "properties": { + "task_id": { + "title": "Task Id", + "type": "string" + }, + "key": { + "title": "Key", + "type": "string" + }, + "map_index": { + "default": -1, + "title": "Map Index", + "type": "integer" + } + }, + "required": [ + "task_id", + "key" + ], + "title": "GetXComBatchItem", + "type": "object" + }, "GetXComCount": { "description": "Get the number of (mapped) XCom values available.", "properties": { @@ -4507,6 +4564,61 @@ "title": "VariableResult", "type": "object" }, + "XComBatchResult": { + "description": "Response to GetXComBatch request.", + "properties": { + "items": { + "items": { + "$ref": "#/$defs/XComBatchResultItem" + }, + "title": "Items", + "type": "array" + }, + "type": { + "const": "XComBatchResult", + "default": "XComBatchResult", + "title": "Type", + "type": "string" + } + }, + "required": [ + "items" + ], + "title": "XComBatchResult", + "type": "object" + }, + "XComBatchResultItem": { + "properties": { + "task_id": { + "title": "Task Id", + "type": "string" + }, + "key": { + "title": "Key", + "type": "string" + }, + "map_index": { + "title": "Map Index", + "type": "integer" + }, + "found": { + "title": "Found", + "type": "boolean" + }, + "value": { + "$ref": "#/$defs/JsonValue", + "default": null + } + }, + "required": [ + "task_id", + "key", + "map_index", + "found" + ], + "title": "XComBatchResultItem", + "type": "object" + }, "XComCountResponse": { "properties": { "len": { diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py index 7e5ce93f86bdc..ae8632232b06f 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py +++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py @@ -39,11 +39,12 @@ def get_bundle() -> VersionBundle: from airflow.sdk.execution_time.schema.versions.v2026_10_30 import ( AddArgBindingsToSupervisorTIRunContext, + AddXComBatchMessages, ) return VersionBundle( HeadVersion(), - Version("2026-10-30", AddArgBindingsToSupervisorTIRunContext), + Version("2026-10-30", AddArgBindingsToSupervisorTIRunContext, AddXComBatchMessages), Version("2026-06-16"), ) diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py index e6b93f5dea805..c307282a808dd 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py +++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py @@ -22,6 +22,21 @@ from airflow.sdk.api.datamodels._generated import TIRunContext +class AddXComBatchMessages(VersionChange): + """ + Add the ``GetXComBatch``/``XComBatchResult`` message pair to the task-execution channel. + + A wholly new discriminated-union member, not a field change on an existing one, so + there is nothing for older schema consumers to migrate away from -- the head shape + already is the schema for this body (see schema/AGENTS.md). This entry exists only + to satisfy the per-commit ``check-supervisor-schemas-versions`` snapshot check. + """ + + description = __doc__ + + instructions_to_migrate_to_previous_version = () + + class AddArgBindingsToSupervisorTIRunContext(VersionChange): """ Add the ``arg_bindings`` argument-binding spec for stub (foreign-runtime) tasks. diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 4c3ae38d21fab..9de5e904855cd 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -102,6 +102,7 @@ GetVariable, GetVariableKeys, GetXCom, + GetXComBatch, GetXComCount, GetXComSequenceItem, GetXComSequenceSlice, @@ -147,6 +148,7 @@ handle_get_variable, handle_get_variable_keys, handle_get_xcom, + handle_get_xcom_batch, handle_get_xcom_count, handle_get_xcom_sequence_item, handle_get_xcom_sequence_slice, @@ -1779,6 +1781,8 @@ def _handle_request(self, msg: ToSupervisor, log: FilteringBoundLogger, req_id: resp, dump_opts = handle_get_variable_keys(self.client, msg) elif isinstance(msg, GetXCom): resp, dump_opts = handle_get_xcom(self.client, msg) + elif isinstance(msg, GetXComBatch): + resp, dump_opts = handle_get_xcom_batch(self.client, msg) elif isinstance(msg, GetXComSequenceItem): resp, dump_opts = handle_get_xcom_sequence_item(self.client, msg) elif isinstance(msg, GetXComSequenceSlice): diff --git a/task-sdk/tests/task_sdk/api/test_client.py b/task-sdk/tests/task_sdk/api/test_client.py index 8bd558b461b1b..fe0f017fe11b7 100644 --- a/task-sdk/tests/task_sdk/api/test_client.py +++ b/task-sdk/tests/task_sdk/api/test_client.py @@ -48,6 +48,8 @@ TaskStateStoreResponse, TerminalTIState, VariableResponse, + XComBatchItemRequest, + XComBatchResponse, XComResponse, ) from airflow.sdk.exceptions import ErrorType, TaskAlreadyRunningError @@ -1127,6 +1129,65 @@ def handle_request(request: httpx.Request) -> httpx.Response: ) assert result == OKResponse(ok=True) + def test_xcom_get_batch_success(self): + def handle_request(request: httpx.Request) -> httpx.Response: + if request.url.path == "/xcoms/dag_id/run_id/batch": + assert json.loads(request.read()) == { + "items": [ + {"task_id": "task_a", "key": "return_value", "map_index": -1}, + {"task_id": "task_b", "key": "return_value", "map_index": -1}, + ] + } + return httpx.Response( + status_code=200, + json={ + "items": [ + { + "task_id": "task_a", + "key": "return_value", + "map_index": -1, + "found": True, + "value": "value_a", + }, + { + "task_id": "task_b", + "key": "return_value", + "map_index": -1, + "found": False, + "value": None, + }, + ] + }, + ) + return httpx.Response(status_code=400, json={"detail": "Bad Request"}) + + client = make_client(transport=httpx.MockTransport(handle_request)) + result = client.xcoms.get_batch( + dag_id="dag_id", + run_id="run_id", + items=[ + XComBatchItemRequest(task_id="task_a", key="return_value"), + XComBatchItemRequest(task_id="task_b", key="return_value"), + ], + ) + assert isinstance(result, XComBatchResponse) + assert result.items[0].found is True + assert result.items[0].value == "value_a" + assert result.items[1].found is False + + def test_xcom_get_batch_not_supported_by_server(self): + def handle_request(request: httpx.Request) -> httpx.Response: + return httpx.Response(status_code=404, json={"detail": "Not Found"}) + + client = make_client(transport=httpx.MockTransport(handle_request)) + result = client.xcoms.get_batch( + dag_id="dag_id", + run_id="run_id", + items=[XComBatchItemRequest(task_id="task_a", key="return_value")], + ) + assert isinstance(result, ErrorResponse) + assert result.error == ErrorType.XCOM_BATCH_NOT_SUPPORTED + class TestConnectionOperations: """ diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py index c9a977c9793c8..164b7c118a97e 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -67,6 +67,7 @@ PreviousTIResponse, TaskInstance, TaskInstanceState, + XComBatchItemRequest, ) from airflow.sdk.exceptions import AirflowRuntimeError, ErrorType, TaskAlreadyRunningError from airflow.sdk.execution_time import supervisor, task_runner @@ -117,6 +118,8 @@ GetVariable, GetVariableKeys, GetXCom, + GetXComBatch, + GetXComBatchItem, GetXComCount, GetXComSequenceItem, GetXComSequenceSlice, @@ -152,6 +155,8 @@ ValidateInletsAndOutlets, VariableKeysResult, VariableResult, + XComBatchResult, + XComBatchResultItem, XComCountResponse, XComResult, XComSequenceIndexResult, @@ -1914,6 +1919,41 @@ class RequestTestCase: ), expected_body={"key": "test_key", "value": None, "type": "XComResult"}, ), + RequestTestCase( + message=GetXComBatch( + dag_id="test_dag", + run_id="test_run", + items=[GetXComBatchItem(task_id="test_task", key="test_key")], + ), + test_id="get_xcom_batch", + client_mock=ClientMock( + method_path="xcoms.get_batch", + args=( + "test_dag", + "test_run", + [XComBatchItemRequest(task_id="test_task", key="test_key", map_index=-1)], + ), + response=XComBatchResult( + items=[ + XComBatchResultItem( + task_id="test_task", key="test_key", map_index=-1, found=True, value="test_value" + ) + ] + ), + ), + expected_body={ + "items": [ + { + "task_id": "test_task", + "key": "test_key", + "map_index": -1, + "found": True, + "value": "test_value", + } + ], + "type": "XComBatchResult", + }, + ), RequestTestCase( message=SetXCom( dag_id="test_dag", From 1477c464238da921c17492fcccd16e31c7050876 Mon Sep 17 00:00:00 2001 From: ColtenOuO Date: Wed, 26 Aug 2026 14:57:52 +0000 Subject: [PATCH 2/4] Resolve a mapped task's plain expand() XCom arguments in one batched call A mapped task with N plain, non-mapped-upstream XComArg kwargs made N sequential Execution API round trips to resolve them, growing linearly with the number of kwargs. Batching collapses that into one round trip for the common case, falling back to the old per-item behavior for custom XCom backends, older API servers, and anything not eligible for batching (mapped upstreams, task groups, composite XComArgs). --- .../sdk/definitions/_internal/expandinput.py | 68 +++++- .../src/airflow/sdk/definitions/xcom_arg.py | 25 ++- .../airflow/sdk/execution_time/task_runner.py | 18 ++ .../definitions/test_mappedoperator.py | 25 +++ .../task_sdk/definitions/test_xcom_arg.py | 198 +++++++++++++++++- 5 files changed, 324 insertions(+), 10 deletions(-) diff --git a/task-sdk/src/airflow/sdk/definitions/_internal/expandinput.py b/task-sdk/src/airflow/sdk/definitions/_internal/expandinput.py index b6ffbd2214253..afdcadaacb66e 100644 --- a/task-sdk/src/airflow/sdk/definitions/_internal/expandinput.py +++ b/task-sdk/src/airflow/sdk/definitions/_internal/expandinput.py @@ -17,12 +17,14 @@ # under the License. from __future__ import annotations +import collections from collections.abc import Iterable, Mapping, Sequence, Sized from typing import TYPE_CHECKING, Any, ClassVar, Union import attrs from airflow.sdk.definitions._internal.mixins import ResolveMixin +from airflow.sdk.definitions._internal.types import NOTSET if TYPE_CHECKING: from typing import TypeGuard @@ -30,6 +32,11 @@ from airflow.sdk.definitions.xcom_arg import XComArg from airflow.sdk.types import Operator +# Wraps a raw batch-response value so it can be passed to BaseXCom.deserialize_value, +# which expects an object with a ``.value`` attribute (mirrors lazy_sequence.py's +# private helper of the same shape; kept separate to avoid reaching into that module). +_XComWrapper = collections.namedtuple("_XComWrapper", "value") + ExpandInput = Union["DictOfListsExpandInput", "ListOfDictsExpandInput"] # Each keyword argument to expand() can be an XComArg, sequence, or dict (not @@ -132,8 +139,11 @@ def _get_map_lengths( they will not be present in the dict. """ - # TODO: This initiates one API call for each XComArg. Would it be - # more efficient to do one single call and unpack the value here? + # NOTE: For plain, non-mapped-upstream XComArgs this operates on values already + # resolved by resolve() below (batched where possible, see + # _batch_resolve_plain_xcom_args). For a mapped-upstream XComArg, resolved_vals[k] + # is a LazyXComSequence and len() on it still triggers its own API call + # (GetXComCount) -- that case isn't batched yet. def _get_length(k: str, v: OperatorExpandArgument) -> int | None: from airflow.sdk.definitions.xcom_arg import XComArg, get_task_map_length @@ -184,6 +194,52 @@ def iter_references(self) -> Iterable[tuple[Operator, str]]: if isinstance(x, XComArg): yield from x.iter_references() + def _batch_resolve_plain_xcom_args(self, context: Mapping[str, Any]) -> dict[str, Any]: + """ + Batch-resolve the plain, non-mapped-upstream XComArg kwargs in one Execution API call. + + Returns resolved values keyed by kwarg name, for whichever entries were actually + batched. Any key not in the returned dict falls through to the normal per-item + ``.resolve()`` path in the caller -- this happens when: fewer than two kwargs are + eligible (a single XComArg has nothing to batch with), a custom XCom backend is + configured (its ``get_one``/``get_all`` may bypass the batch endpoint's semantics + entirely, so batching is skipped rather than risking a mismatch), or the API server + doesn't support the batch endpoint yet (old server, newer Task SDK). + """ + from airflow.sdk.definitions.xcom_arg import PlainXComArg + from airflow.sdk.execution_time.comms import ErrorResponse, GetXComBatchItem, XComBatchResult + from airflow.sdk.execution_time.xcom import BaseXCom as ResolvedBaseXCom, XCom + + if XCom is not ResolvedBaseXCom: + return {} + + batchable = {k: v for k, v in self.value.items() if isinstance(v, PlainXComArg) and v.is_batchable} + if len(batchable) < 2: + return {} + + ti = context["ti"] + items = [GetXComBatchItem(task_id=v.operator.task_id, key=v.key) for v in batchable.values()] + response = ti.xcom_pull_batch(items) + if isinstance(response, ErrorResponse): + return {} + if not isinstance(response, XComBatchResult): + raise TypeError( + f"Expected XComBatchResult or ErrorResponse, received: {type(response)} {response!r}" + ) + + by_task_key = {(item.task_id, item.key): item for item in response.items} + results: dict[str, Any] = {} + for k, xcom_arg in batchable.items(): + item = by_task_key[(xcom_arg.operator.task_id, xcom_arg.key)] + if not item.found: + raw_value = NOTSET + elif item.value is None: + raw_value = None + else: + raw_value = ResolvedBaseXCom.deserialize_value(_XComWrapper(item.value)) + results[k] = xcom_arg._finalize_pulled_value(raw_value, dag_id=ti.dag_id) + return results + def resolve(self, context: Mapping[str, Any]) -> tuple[Mapping[str, Any], set[int]]: map_index: int | None = context["ti"].map_index if map_index is None or map_index < 0: @@ -193,11 +249,13 @@ def resolve(self, context: Mapping[str, Any]) -> tuple[Mapping[str, Any], set[in # When empty, individual XComArgs will compute their map_indexes lazily in xcom_arg.py. upstream_map_indexes = getattr(context["ti"], "_upstream_map_indexes", None) or {} - # TODO: This initiates one API call for each XComArg. Would it be - # more efficient to do one single call and unpack the value here? + batch_results = self._batch_resolve_plain_xcom_args(context) resolved = { - k: v.resolve(context) if _needs_run_time_resolution(v) else v for k, v in self.value.items() + k: batch_results[k] + if k in batch_results + else (v.resolve(context) if _needs_run_time_resolution(v) else v) + for k, v in self.value.items() } sized_resolved = {k: v for k, v in resolved.items() if isinstance(v, Sized)} diff --git a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py index a38e47a466bcf..003ba3d33942e 100644 --- a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py +++ b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py @@ -366,8 +366,27 @@ def resolve(self, context: Mapping[str, Any]) -> Any: default=NOTSET, map_indexes=map_indexes, ) - if is_arg_set(result): - return result + return self._finalize_pulled_value(result, dag_id=ti.dag_id) + + @property + def is_batchable(self) -> bool: + """ + Whether this XComArg can be resolved via a single-value batch XCom lookup. + + True for the common, non-mapped-upstream case: the same condition that makes + ``resolve()`` above pull with ``map_indexes=None`` instead of returning a + ``LazyXComSequence`` or resolving per-map-index upstream indexes. + """ + return not self.operator.is_mapped and self.operator.get_closest_mapped_task_group() is None + + def _finalize_pulled_value(self, raw_value: Any, *, dag_id: str) -> Any: + """ + Apply the not-found/default fallback shared by the per-item and batched pull paths. + + ``raw_value`` is ``NOTSET`` (or unset) when no matching XCom was found. + """ + if is_arg_set(raw_value): + return raw_value if self.key == BaseXCom.XCOM_RETURN_KEY: return None if getattr(self.operator, "multiple_outputs", False): @@ -377,7 +396,7 @@ def resolve(self, context: Mapping[str, Any]) -> Any: # different names than the predefined "XCOM_RETURN_KEY" and won't be found. # Therefore, it's better to return "None" like we did above where self.key==XCOM_RETURN_KEY. return None - raise XComNotFound(ti.dag_id, task_id, self.key) + raise XComNotFound(dag_id, self.operator.task_id, self.key) def _get_callable_name(f: Callable | str) -> str: diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py index eeb2788062248..fdab1294e9979 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -103,6 +103,8 @@ GetTaskRescheduleStartDate, GetTaskStates, GetTICount, + GetXComBatch, + GetXComBatchItem, InactiveAssetsResult, PreviousDagRunResult, PreviousTIResult, @@ -124,6 +126,7 @@ ToTask, TriggerDagRun, ValidateInletsAndOutlets, + XComBatchResult, ) from airflow.sdk.execution_time.context import ( AssetStateStoreAccessors, @@ -651,6 +654,21 @@ async def axcom_push(self, key: str, value: Any): """ await _axcom_push(self, key, value) + def xcom_pull_batch(self, items: list[GetXComBatchItem]) -> XComBatchResult | ErrorResponse: + """ + Pull multiple XComs in a single Execution API round trip. + + Used by ``DictOfListsExpandInput.resolve()`` to batch the plain, non-mapped- + upstream XComArg kwargs of a mapped task's ``.expand()`` call instead of + pulling each one individually. Returns an ``ErrorResponse`` with + ``ErrorType.XCOM_BATCH_NOT_SUPPORTED`` when talking to an API server that + predates this endpoint; callers fall back to per-item ``xcom_pull()``. + """ + response = SUPERVISOR_COMMS.send(GetXComBatch(dag_id=self.dag_id, run_id=self.run_id, items=items)) + if TYPE_CHECKING: + assert isinstance(response, (XComBatchResult, ErrorResponse)) + return response + def get_relevant_upstream_map_indexes( self, upstream: BaseOperator, ti_count: int | None, session: Any ) -> int | range | None: diff --git a/task-sdk/tests/task_sdk/definitions/test_mappedoperator.py b/task-sdk/tests/task_sdk/definitions/test_mappedoperator.py index 4aadf923e9c85..f48ae32d7e061 100644 --- a/task-sdk/tests/task_sdk/definitions/test_mappedoperator.py +++ b/task-sdk/tests/task_sdk/definitions/test_mappedoperator.py @@ -35,10 +35,13 @@ ErrorResponse, GetTICount, GetXCom, + GetXComBatch, GetXComSequenceItem, GetXComSequenceSlice, SetXCom, TICount, + XComBatchResult, + XComBatchResultItem, XComResult, XComSequenceIndexResult, XComSequenceSliceResult, @@ -52,6 +55,19 @@ DEFAULT_DATE = datetime(2016, 1, 1) +def _xcom_batch_result(msg: GetXComBatch, value_for_task_id) -> XComBatchResult: + """Build a found=True XComBatchResult for every item in msg, mirroring a GetXCom mock's value.""" + get_value = value_for_task_id if callable(value_for_task_id) else (lambda _: value_for_task_id) + return XComBatchResult( + items=[ + XComBatchResultItem( + task_id=item.task_id, key=item.key, map_index=-1, found=True, value=get_value(item.task_id) + ) + for item in msg.items + ] + ) + + def test_task_mapping_with_dag(): with DAG("test-dag") as dag: task1 = BaseOperator(task_id="op1") @@ -266,6 +282,8 @@ def execute(self, context): def mock_comms(msg): if isinstance(msg, GetXCom): return XComResult(key=BaseXCom.XCOM_RETURN_KEY, value=["{{ ds }}"]) + if isinstance(msg, GetXComBatch): + return _xcom_batch_result(msg, ["{{ ds }}"]) if isinstance(msg, GetXComSequenceSlice): return XComSequenceSliceResult(root=["{{ ds }}"]) if isinstance(msg, GetTICount): @@ -460,12 +478,17 @@ def show(number, letter): numbers = [1, 2] letters = {"a": "x", "b": "y", "c": "z"} + def _value_for(task_id): + return numbers if task_id == "emit_numbers" else letters + def mock_comms(msg): if isinstance(msg, GetXCom): if msg.task_id == "emit_numbers": return XComResult(key=BaseXCom.XCOM_RETURN_KEY, value=numbers) if msg.task_id == "emit_letters": return XComResult(key=BaseXCom.XCOM_RETURN_KEY, value=letters) + elif isinstance(msg, GetXComBatch): + return _xcom_batch_result(msg, _value_for) elif isinstance(msg, GetXComSequenceSlice): if msg.task_id == "emit_numbers": return XComSequenceSliceResult(root=numbers) @@ -516,6 +539,8 @@ def mock_comms(msg): if isinstance(msg, GetXCom): if msg.task_id == "emit_numbers": return XComResult(key=BaseXCom.XCOM_RETURN_KEY, value=numbers) + elif isinstance(msg, GetXComBatch): + return _xcom_batch_result(msg, numbers) elif isinstance(msg, GetXComSequenceSlice): if msg.task_id == "emit_numbers": return XComSequenceSliceResult(root=numbers) diff --git a/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py b/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py index 2d1f8b2a4fa52..64864307868c9 100644 --- a/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py +++ b/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py @@ -25,10 +25,19 @@ from airflow.sdk import TaskInstanceState from airflow.sdk.bases.xcom import BaseXCom +from airflow.sdk.definitions._internal.expandinput import DictOfListsExpandInput from airflow.sdk.definitions.dag import DAG from airflow.sdk.definitions.xcom_arg import PlainXComArg -from airflow.sdk.exceptions import AirflowSkipException -from airflow.sdk.execution_time.comms import GetXCom, XComResult, XComSequenceSliceResult +from airflow.sdk.exceptions import AirflowSkipException, ErrorType +from airflow.sdk.execution_time.comms import ( + ErrorResponse, + GetXCom, + GetXComBatch, + XComBatchResult, + XComBatchResultItem, + XComResult, + XComSequenceSliceResult, +) from airflow.sdk.execution_time.lazy_sequence import LazyXComSequence from airflow.sdk.serde import deserialize, serialize @@ -416,3 +425,188 @@ def test_resolve_uses_xcom_pull_for_specific_index(self): assert resolved == "value-0" ti.xcom_pull.assert_called_once() assert ti.xcom_pull.call_args.kwargs["map_indexes"] == 0 + + +def test_expand_batches_plain_xcom_args_into_one_call(run_ti: RunTI, mock_supervisor_comms): + """Multiple plain, non-mapped-upstream XComArg kwargs resolve via a single GetXComBatch.""" + results = [] + + with DAG("test") as dag: + + @dag.task + def push_a(): + return ["a"] + + @dag.task + def push_b(): + return ["b"] + + @dag.task + def push_c(): + return ["c"] + + @dag.task + def pull(x, y, z): + results.append((x, y, z)) + + pull.expand(x=push_a(), y=push_b(), z=push_c()) + + calls = {"GetXComBatch": 0, "GetXCom": 0} + values = {"push_a": ["a"], "push_b": ["b"], "push_c": ["c"]} + + def comms(msg): + if isinstance(msg, GetXComBatch): + calls["GetXComBatch"] += 1 + return XComBatchResult( + items=[ + XComBatchResultItem( + task_id=item.task_id, + key=item.key, + map_index=-1, + found=True, + value=values[item.task_id], + ) + for item in msg.items + ] + ) + if isinstance(msg, GetXCom): + calls["GetXCom"] += 1 + return mock.DEFAULT + + mock_supervisor_comms.send.side_effect = comms + + assert run_ti(dag, "pull", 0) == TaskInstanceState.SUCCESS + assert calls == {"GetXComBatch": 1, "GetXCom": 0} + assert results == [("a", "b", "c")] + + +def test_expand_batch_falls_back_on_old_server(run_ti: RunTI, mock_supervisor_comms): + """When the API server doesn't support batch lookups, kwargs resolve individually instead.""" + results = [] + + with DAG("test") as dag: + + @dag.task + def push_a(): + return ["a"] + + @dag.task + def push_b(): + return ["b"] + + @dag.task + def pull(x, y): + results.append((x, y)) + + pull.expand(x=push_a(), y=push_b()) + + calls = {"GetXComBatch": 0, "GetXCom": 0} + values = {"push_a": ["a"], "push_b": ["b"]} + + def comms(msg): + if isinstance(msg, GetXComBatch): + calls["GetXComBatch"] += 1 + return ErrorResponse(error=ErrorType.XCOM_BATCH_NOT_SUPPORTED, detail={}) + if isinstance(msg, GetXCom): + calls["GetXCom"] += 1 + return XComResult(key=BaseXCom.XCOM_RETURN_KEY, value=values[msg.task_id]) + return mock.DEFAULT + + mock_supervisor_comms.send.side_effect = comms + + assert run_ti(dag, "pull", 0) == TaskInstanceState.SUCCESS + assert calls == {"GetXComBatch": 1, "GetXCom": 2} + assert results == [("a", "b")] + + +def test_expand_batch_skipped_for_custom_xcom_backend(run_ti: RunTI, mock_supervisor_comms): + """A custom XCom backend may bypass the batch endpoint's semantics, so batching is skipped.""" + import airflow.sdk.execution_time.xcom as xcom_module + + results = [] + + with DAG("test") as dag: + + @dag.task + def push_a(): + return ["a"] + + @dag.task + def push_b(): + return ["b"] + + @dag.task + def pull(x, y): + results.append((x, y)) + + pull.expand(x=push_a(), y=push_b()) + + calls = {"GetXComBatch": 0, "GetXCom": 0} + values = {"push_a": ["a"], "push_b": ["b"]} + + def comms(msg): + if isinstance(msg, GetXComBatch): + calls["GetXComBatch"] += 1 + if isinstance(msg, GetXCom): + calls["GetXCom"] += 1 + return XComResult(key=BaseXCom.XCOM_RETURN_KEY, value=values[msg.task_id]) + return mock.DEFAULT + + mock_supervisor_comms.send.side_effect = comms + + class CustomXCom(BaseXCom): + pass + + with mock.patch.object(xcom_module, "XCom", CustomXCom): + assert run_ti(dag, "pull", 0) == TaskInstanceState.SUCCESS + + assert calls == {"GetXComBatch": 0, "GetXCom": 2} + assert results == [("a", "b")] + + +class TestDictOfListsExpandInputBatching: + """Unit-level tests for _batch_resolve_plain_xcom_args, below the full task-run scaffold.""" + + @staticmethod + def _make_plain_arg(task_id: str, *, mapped_task_group: bool = False) -> PlainXComArg: + operator = mock.MagicMock() + operator.is_mapped = False + operator.task_id = task_id + operator.dag_id = "test_dag" + operator.get_closest_mapped_task_group.return_value = mock.MagicMock() if mapped_task_group else None + return PlainXComArg(operator=operator, key=BaseXCom.XCOM_RETURN_KEY) + + def test_single_eligible_arg_is_not_batched(self): + """Nothing to batch with, so the caller's normal per-item path handles it instead.""" + expand_input = DictOfListsExpandInput({"x": self._make_plain_arg("push_a")}) + ti = mock.MagicMock(dag_id="test_dag") + ti.xcom_pull_batch.side_effect = AssertionError("should not batch a single kwarg") + + assert expand_input._batch_resolve_plain_xcom_args({"ti": ti}) == {} + + def test_mixed_eligibility_only_batches_plain_args(self): + """A mapped-task-group XComArg alongside plain ones: only the plain ones batch.""" + expand_input = DictOfListsExpandInput( + { + "x": self._make_plain_arg("push_a"), + "y": self._make_plain_arg("push_b"), + "z": self._make_plain_arg("push_c", mapped_task_group=True), + } + ) + ti = mock.MagicMock(dag_id="test_dag") + ti.xcom_pull_batch.return_value = XComBatchResult( + items=[ + XComBatchResultItem( + task_id="push_a", key=BaseXCom.XCOM_RETURN_KEY, map_index=-1, found=True, value="a" + ), + XComBatchResultItem( + task_id="push_b", key=BaseXCom.XCOM_RETURN_KEY, map_index=-1, found=True, value="b" + ), + ] + ) + + results = expand_input._batch_resolve_plain_xcom_args({"ti": ti}) + + assert results == {"x": "a", "y": "b"} + ti.xcom_pull_batch.assert_called_once() + assert len(ti.xcom_pull_batch.call_args.args[0]) == 2 From 299628cc40bdddfa6d65ae29b7ff7b011229b59b Mon Sep 17 00:00:00 2001 From: ColtenOuO Date: Wed, 26 Aug 2026 15:57:48 +0000 Subject: [PATCH 3/4] Address review feedback on the XCom batch endpoint Bounds the batch endpoint to a fixed max item count so an authenticated task can't force an oversized bind-parameter list or unbounded work on the API server. Also replaces unspecced MagicMocks in the new batch tests with autospec'd ones so a signature or attribute drift on the mocked classes would actually fail the test, and adds coverage for the found=False path (a missing return_value XCom resolves to None; a missing non-return_value key raises, matching the per-item xcom_pull path) and for a batch item with an explicit, non-default map_index, which the original tests didn't exercise. --- .../execution_api/datamodels/xcom.py | 9 ++- .../execution_api/versions/v2026_10_30.py | 2 +- .../execution_api/versions/head/test_xcoms.py | 47 ++++++++++++ .../task_sdk/definitions/test_xcom_arg.py | 73 +++++++++++++++++-- 4 files changed, 121 insertions(+), 10 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/xcom.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/xcom.py index 434e943864516..086b5c3390cbd 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/xcom.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/xcom.py @@ -17,10 +17,15 @@ from __future__ import annotations -from pydantic import JsonValue, RootModel +from pydantic import Field, JsonValue, RootModel from airflow.api_fastapi.core_api.base import BaseModel +# Each item expands to 3 SQL bind parameters (task_id, key, map_index); this bounds a +# single batch request to at most 3000 bind parameters, safely under Postgres/MySQL/ +# modern-SQLite limits regardless of how many kwargs a caller's .expand() call has. +MAX_XCOM_BATCH_ITEMS = 1000 + class XComResponse(BaseModel): """XCom schema for responses with fields that are needed for Runtime.""" @@ -53,7 +58,7 @@ class XComBatchItemRequest(BaseModel): class XComBatchRequestBody(BaseModel): """Body for a batch XCom lookup, scoped to a single dag_id/run_id.""" - items: list[XComBatchItemRequest] + items: list[XComBatchItemRequest] = Field(max_length=MAX_XCOM_BATCH_ITEMS) class XComBatchItemResponse(BaseModel): diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py index 3e0e40de3f503..a4408a081099f 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py @@ -30,7 +30,7 @@ class AddXComBatchEndpoint(VersionChange): - """Add a batch XCom lookup endpoint that resolves multiple XComs in one request.""" + """Add a batch XCom lookup endpoint that resolves multiple XComs in one request, capped at MAX_XCOM_BATCH_ITEMS items per request.""" description = __doc__ diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py index 5e95afcb06455..0ce9ba6af5bc0 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py @@ -703,6 +703,53 @@ def test_batch_empty_request(self, client, create_task_instance): assert response.status_code == 200, response.json() assert response.json() == {"items": []} + def test_batch_rejects_too_many_items(self, client, create_task_instance): + """A request over MAX_XCOM_BATCH_ITEMS is rejected before touching the database.""" + from airflow.api_fastapi.execution_api.datamodels.xcom import MAX_XCOM_BATCH_ITEMS + + ti = create_task_instance() + items = [{"task_id": ti.task_id, "key": f"k{i}"} for i in range(MAX_XCOM_BATCH_ITEMS + 1)] + + response = client.post(f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/batch", json={"items": items}) + + assert response.status_code == 422, response.json() + + def test_batch_non_default_map_index(self, client, dag_maker, session): + """A batch item with an explicit map_index resolves the matching mapped-task-instance XCom.""" + + class MyOperator(EmptyOperator): + def __init__(self, *, x, **kwargs): + super().__init__(**kwargs) + self.x = x + + with dag_maker(dag_id="dag"): + MyOperator.partial(task_id="task").expand(x=[1, 2, 3]) + dag_run = dag_maker.create_dagrun(run_id="runid") + + for map_index in (0, 1, 2): + session.add( + XComModel( + key="k", + value=f"value_{map_index}", + dag_run_id=dag_run.id, + run_id=dag_run.run_id, + task_id="task", + dag_id="dag", + map_index=map_index, + ) + ) + session.commit() + + response = client.post( + "/execution/xcoms/dag/runid/batch", + json={"items": [{"task_id": "task", "key": "k", "map_index": 1}]}, + ) + + assert response.status_code == 200, response.json() + assert response.json() == { + "items": [{"task_id": "task", "key": "k", "map_index": 1, "found": True, "value": "value_1"}] + } + class TestXComTeamAccess: """Multi-team isolation for the Execution API XCom routes (no cross-team sharing).""" diff --git a/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py b/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py index 64864307868c9..6039b6eb33065 100644 --- a/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py +++ b/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py @@ -24,11 +24,12 @@ import structlog from airflow.sdk import TaskInstanceState +from airflow.sdk.bases.operator import BaseOperator from airflow.sdk.bases.xcom import BaseXCom from airflow.sdk.definitions._internal.expandinput import DictOfListsExpandInput from airflow.sdk.definitions.dag import DAG from airflow.sdk.definitions.xcom_arg import PlainXComArg -from airflow.sdk.exceptions import AirflowSkipException, ErrorType +from airflow.sdk.exceptions import AirflowSkipException, ErrorType, XComNotFound from airflow.sdk.execution_time.comms import ( ErrorResponse, GetXCom, @@ -39,6 +40,7 @@ XComSequenceSliceResult, ) from airflow.sdk.execution_time.lazy_sequence import LazyXComSequence +from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance from airflow.sdk.serde import deserialize, serialize log = structlog.get_logger(__name__) @@ -568,18 +570,28 @@ class TestDictOfListsExpandInputBatching: """Unit-level tests for _batch_resolve_plain_xcom_args, below the full task-run scaffold.""" @staticmethod - def _make_plain_arg(task_id: str, *, mapped_task_group: bool = False) -> PlainXComArg: - operator = mock.MagicMock() + def _make_plain_arg( + task_id: str, *, key: str = BaseXCom.XCOM_RETURN_KEY, mapped_task_group: bool = False + ) -> PlainXComArg: + operator = mock.create_autospec(BaseOperator, instance=True) operator.is_mapped = False operator.task_id = task_id operator.dag_id = "test_dag" - operator.get_closest_mapped_task_group.return_value = mock.MagicMock() if mapped_task_group else None - return PlainXComArg(operator=operator, key=BaseXCom.XCOM_RETURN_KEY) + operator.multiple_outputs = False + # Only the None-ness matters to is_batchable, so a plain sentinel object is enough. + operator.get_closest_mapped_task_group.return_value = object() if mapped_task_group else None + return PlainXComArg(operator=operator, key=key) + + @staticmethod + def _make_ti() -> mock.MagicMock: + ti = mock.create_autospec(RuntimeTaskInstance, instance=True) + ti.dag_id = "test_dag" + return ti def test_single_eligible_arg_is_not_batched(self): """Nothing to batch with, so the caller's normal per-item path handles it instead.""" expand_input = DictOfListsExpandInput({"x": self._make_plain_arg("push_a")}) - ti = mock.MagicMock(dag_id="test_dag") + ti = self._make_ti() ti.xcom_pull_batch.side_effect = AssertionError("should not batch a single kwarg") assert expand_input._batch_resolve_plain_xcom_args({"ti": ti}) == {} @@ -593,7 +605,7 @@ def test_mixed_eligibility_only_batches_plain_args(self): "z": self._make_plain_arg("push_c", mapped_task_group=True), } ) - ti = mock.MagicMock(dag_id="test_dag") + ti = self._make_ti() ti.xcom_pull_batch.return_value = XComBatchResult( items=[ XComBatchResultItem( @@ -610,3 +622,50 @@ def test_mixed_eligibility_only_batches_plain_args(self): assert results == {"x": "a", "y": "b"} ti.xcom_pull_batch.assert_called_once() assert len(ti.xcom_pull_batch.call_args.args[0]) == 2 + + def test_not_found_return_value_key_resolves_to_none(self): + """A missing return_value XCom resolves to None, matching the per-item xcom_pull path.""" + expand_input = DictOfListsExpandInput( + { + "x": self._make_plain_arg("push_a"), + "y": self._make_plain_arg("push_b"), + } + ) + ti = self._make_ti() + ti.xcom_pull_batch.return_value = XComBatchResult( + items=[ + XComBatchResultItem( + task_id="push_a", key=BaseXCom.XCOM_RETURN_KEY, map_index=-1, found=True, value="a" + ), + XComBatchResultItem( + task_id="push_b", key=BaseXCom.XCOM_RETURN_KEY, map_index=-1, found=False, value=None + ), + ] + ) + + results = expand_input._batch_resolve_plain_xcom_args({"ti": ti}) + + assert results == {"x": "a", "y": None} + + def test_not_found_other_key_raises_xcom_not_found(self): + """A missing non-return_value XCom raises, matching the per-item xcom_pull path.""" + expand_input = DictOfListsExpandInput( + { + "x": self._make_plain_arg("push_a"), + "y": self._make_plain_arg("push_b", key="custom_key"), + } + ) + ti = self._make_ti() + ti.xcom_pull_batch.return_value = XComBatchResult( + items=[ + XComBatchResultItem( + task_id="push_a", key=BaseXCom.XCOM_RETURN_KEY, map_index=-1, found=True, value="a" + ), + XComBatchResultItem( + task_id="push_b", key="custom_key", map_index=-1, found=False, value=None + ), + ] + ) + + with pytest.raises(XComNotFound): + expand_input._batch_resolve_plain_xcom_args({"ti": ti}) From f57b7ac94b14ef4f33f071b08fd2105402d6d2d3 Mon Sep 17 00:00:00 2001 From: ColtenOuO Date: Wed, 26 Aug 2026 17:22:23 +0000 Subject: [PATCH 4/4] Fix CI failures from the XCom batch endpoint change The regenerated Task SDK client was stale after adding the item-count cap to XComBatchRequestBody, missing the max_length constraint in the committed _generated.py. The DAG File Processor already forwards every other XCom read message (GetXCom, GetXComCount, GetXComSequenceItem, GetXComSequenceSlice) to the API server the same way the task-runner supervisor does, so GetXComBatch needs the same wiring for the message-type completeness checks to hold. The Triggerer has no equivalent need for it (deferred triggers don't resolve expand() kwargs), so it's excluded there the same way the existing sequence/count XCom messages already are. --- .../src/airflow/dag_processing/processor.py | 7 +++ .../unit/dag_processing/test_processor.py | 51 ++++++++++++++++++- .../tests/unit/jobs/test_triggerer_job.py | 2 + .../airflow/sdk/api/datamodels/_generated.py | 2 +- .../schema/versions/v2026_10_30.py | 6 +++ 5 files changed, 66 insertions(+), 2 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/processor.py b/airflow-core/src/airflow/dag_processing/processor.py index 3805de9fdf55b..9af78696b842f 100644 --- a/airflow-core/src/airflow/dag_processing/processor.py +++ b/airflow-core/src/airflow/dag_processing/processor.py @@ -54,6 +54,7 @@ GetVariable, GetVariableKeys, GetXCom, + GetXComBatch, GetXComCount, GetXComSequenceItem, GetXComSequenceSlice, @@ -66,6 +67,7 @@ TaskStatesResult, VariableKeysResult, VariableResult, + XComBatchResult, XComCountResponse, XComResult, XComSequenceIndexResult, @@ -80,6 +82,7 @@ handle_get_ti_count, handle_get_variable_keys, handle_get_xcom, + handle_get_xcom_batch, handle_get_xcom_count, handle_get_xcom_sequence_item, handle_get_xcom_sequence_slice, @@ -158,6 +161,7 @@ class DagFileParsingResult(BaseModel): | GetPreviousDagRun | GetPreviousTI | GetXCom + | GetXComBatch | GetXComCount | GetXComSequenceItem | GetXComSequenceSlice @@ -176,6 +180,7 @@ class DagFileParsingResult(BaseModel): | PrevSuccessfulDagRunResult | ErrorResponse | OKResponse + | XComBatchResult | XComCountResponse | XComResult | XComSequenceIndexResult @@ -718,6 +723,8 @@ def _handle_request(self, msg: ToManager, log: FilteringBoundLogger, req_id: int resp, dump_opts = handle_get_prev_successful_dag_run(self.client, self.id) elif isinstance(msg, GetXCom): resp, dump_opts = handle_get_xcom(self.client, msg) + elif isinstance(msg, GetXComBatch): + resp, dump_opts = handle_get_xcom_batch(self.client, msg) elif isinstance(msg, GetXComCount): resp, dump_opts = handle_get_xcom_count(self.client, msg) elif isinstance(msg, GetXComSequenceItem): diff --git a/airflow-core/tests/unit/dag_processing/test_processor.py b/airflow-core/tests/unit/dag_processing/test_processor.py index f54b82fef8291..4f33af6728592 100644 --- a/airflow-core/tests/unit/dag_processing/test_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_processor.py @@ -68,7 +68,13 @@ from airflow.models import DagRun from airflow.sdk import DAG, BaseOperator from airflow.sdk.api.client import Client -from airflow.sdk.api.datamodels._generated import ConnectionResponse, DagRunState, VariableResponse +from airflow.sdk.api.datamodels._generated import ( + ConnectionResponse, + DagRunState, + VariableResponse, + XComBatchItemResponse, + XComBatchResponse, +) from airflow.sdk.execution_time import comms, supervisor from airflow.sdk.execution_time.comms import ( GetConnection, @@ -76,6 +82,8 @@ GetTICount, GetVariable, GetXCom, + GetXComBatch, + GetXComBatchItem, GetXComSequenceSlice, TaskStatesResult, TICount, @@ -2428,6 +2436,47 @@ def test_handle_request_get_variable_masks_value_with_key(self, proc): "type": "VariableResult", } + def test_handle_request_get_xcom_batch(self, proc): + proc.client.xcoms.get_batch.return_value = XComBatchResponse( + items=[ + XComBatchItemResponse( + task_id="test_task", key="return_value", map_index=-1, found=True, value="test_value" + ) + ] + ) + + with patch.object(DagFileProcessorProcess, "send_msg", autospec=True) as mock_send_msg: + proc._handle_request( + GetXComBatch( + dag_id="test_dag", + run_id="test_run", + items=[GetXComBatchItem(task_id="test_task", key="return_value")], + ), + structlog.get_logger(), + req_id=789, + ) + + proc.client.xcoms.get_batch.assert_called_once() + + mock_send_msg.assert_called_once() + _, args, kwargs = mock_send_msg.mock_calls[0] + assert args[0] is proc + msg = args[1] + assert kwargs["request_id"] == 789 + assert kwargs["error"] is None + assert msg.model_dump() == { + "items": [ + { + "task_id": "test_task", + "key": "return_value", + "map_index": -1, + "found": True, + "value": "test_value", + } + ], + "type": "XComBatchResult", + } + class TestMultiTeamCallbackMetrics: """Tests for team_name tag on dag.callback_exceptions in multi-team mode.""" diff --git a/airflow-core/tests/unit/jobs/test_triggerer_job.py b/airflow-core/tests/unit/jobs/test_triggerer_job.py index 07f861f4f9198..339b727445014 100644 --- a/airflow-core/tests/unit/jobs/test_triggerer_job.py +++ b/airflow-core/tests/unit/jobs/test_triggerer_job.py @@ -2726,6 +2726,7 @@ def get_type_names(union_type): "GetPreviousDagRun", "GetTaskBreadcrumbs", "GetTaskRescheduleStartDate", + "GetXComBatch", "GetXComCount", "GetXComSequenceItem", "GetXComSequenceSlice", @@ -2760,6 +2761,7 @@ def get_type_names(union_type): "InactiveAssetsResult", "CreateHITLDetailPayload", "PrevSuccessfulDagRunResult", + "XComBatchResult", "XComCountResponse", "XComSequenceIndexResult", "XComSequenceSliceResult", diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 24244aeac23b0..42e6f08165d6c 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -540,7 +540,7 @@ class XComBatchRequestBody(BaseModel): Body for a batch XCom lookup, scoped to a single dag_id/run_id. """ - items: Annotated[list[XComBatchItemRequest], Field(title="Items")] + items: Annotated[list[XComBatchItemRequest], Field(max_length=1000, title="Items")] class XComBatchResponse(BaseModel): diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py index c307282a808dd..1583972ba6ada 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py +++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py @@ -26,6 +26,12 @@ class AddXComBatchMessages(VersionChange): """ Add the ``GetXComBatch``/``XComBatchResult`` message pair to the task-execution channel. + Also wired into the Dag-processing channel (``ToManager``/``ToDagProcessor`` in + ``airflow.dag_processing.processor``), the same way the DAG File Processor already + forwards ``GetXCom``/``GetXComCount``/``GetXComSequenceItem``/``GetXComSequenceSlice``. + Not wired into the Triggerer channel -- deferred triggers never resolve ``expand()`` + kwargs. + A wholly new discriminated-union member, not a field change on an existing one, so there is nothing for older schema consumers to migrate away from -- the head shape already is the schema for this body (see schema/AGENTS.md). This entry exists only