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..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.""" @@ -40,3 +45,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] = Field(max_length=MAX_XCOM_BATCH_ITEMS) + + +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..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 @@ -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, capped at MAX_XCOM_BATCH_ITEMS items per 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/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/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..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 @@ -618,6 +618,139 @@ 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": []} + + 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).""" @@ -766,3 +899,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/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/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..42e6f08165d6c 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(max_length=1000, 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/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/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..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 @@ -22,6 +22,27 @@ from airflow.sdk.api.datamodels._generated import TIRunContext +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 + 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/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/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/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..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,12 +24,23 @@ 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 -from airflow.sdk.execution_time.comms import GetXCom, XComResult, XComSequenceSliceResult +from airflow.sdk.exceptions import AirflowSkipException, ErrorType, XComNotFound +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.execution_time.task_runner import RuntimeTaskInstance from airflow.sdk.serde import deserialize, serialize log = structlog.get_logger(__name__) @@ -416,3 +427,245 @@ 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, *, 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.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 = 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}) == {} + + 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 = 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=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 + + 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}) 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",