From 1b9c256a7607b40cef4010286ce97691c07eade7 Mon Sep 17 00:00:00 2001 From: David Blain Date: Wed, 22 Jul 2026 10:36:01 +0200 Subject: [PATCH 1/6] Add POST /xcoms/{dag_id}/{run_id}/{task_id}/keys endpoint for batched XCom fetch Fetching XCom results from XComIterable (which stores per-index values as distinct keys return_value_0, return_value_1, ... under the same map_index) required one API round-trip per element. The existing GetXComSequenceSlice endpoint cannot be reused because it ranges over map_index for a single key -- the inverse structure. Add a dedicated POST route that accepts a list of keys and returns all values in one database query, preserving key order in the response. The full wiring -- comms message, client method, supervisor request handler, and route registration -- is included so the endpoint is usable end-to-end once XComIterable is updated to call it. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../execution_api/datamodels/xcom.py | 6 + .../api_fastapi/execution_api/routes/xcoms.py | 38 +++++- .../src/airflow/dag_processing/processor.py | 4 + .../execution_api/versions/head/test_xcoms.py | 113 ++++++++++++++++++ task-sdk/src/airflow/sdk/api/client.py | 16 +++ .../src/airflow/sdk/execution_time/comms.py | 12 ++ .../sdk/execution_time/request_handlers.py | 9 ++ 7 files changed, 196 insertions(+), 2 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..93f8f58195653 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,9 @@ class XComSequenceSliceResponse(RootModel): """XCom schema with minimal structure for slice-based access.""" root: list[JsonValue] + + +class XComKeysRequest(BaseModel): + """Request body for fetching multiple XCom values by key list in a single query.""" + + keys: list[str] 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..7332196a37ebb 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,13 @@ 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 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 ( + XComKeysRequest, XComResponse, XComSequenceIndexResponse, XComSequenceSliceResponse, @@ -42,7 +43,6 @@ 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, @@ -68,6 +68,7 @@ def has_xcom_access( from airflow.configuration import conf write = request.method not in {"GET", "HEAD", "OPTIONS"} + xcom_key = request.path_params.get("key") log.debug( "Checking %s XCom access for task instance '%s' to XCom '%s' on dag '%s'", @@ -113,6 +114,7 @@ def has_xcom_access( dependencies=[Depends(has_xcom_access)], ) + log = logging.getLogger(__name__) @@ -356,6 +358,38 @@ def get_xcom( return XComResponse(key=key, value=(result[0] if isinstance(result, tuple) else result).value) +@router.post( + "/{dag_id}/{run_id}/{task_id}/keys", + summary="Get multiple XCom values by keys", + description=( + "Fetch multiple XCom values by key list in a single database query. " + "Optimised for XComIterable iteration, reducing N round-trips to one." + ), +) +def get_xcom_by_keys( + dag_id: str, + run_id: str, + task_id: str, + request_body: XComKeysRequest, + session: SessionDep, + map_index: Annotated[int, Query()] = -1, +) -> XComSequenceSliceResponse: + """Fetch multiple XCom values by different keys in a single database query.""" + key_list = request_body.keys + if not key_list: + return XComSequenceSliceResponse([]) + + query = select(XComModel.key, XComModel.value).where( + XComModel.dag_id == dag_id, + XComModel.run_id == run_id, + XComModel.task_id == task_id, + XComModel.map_index == map_index, + XComModel.key.in_(key_list), + ) + rows = {row.key: row.value for row in session.execute(query)} + return XComSequenceSliceResponse([rows.get(key) for key in key_list]) + + # TODO: once we have JWT tokens, then remove dag_id/run_id/task_id from the URL and just use the info in # the token @router.post( diff --git a/airflow-core/src/airflow/dag_processing/processor.py b/airflow-core/src/airflow/dag_processing/processor.py index ca7539131f60e..6a7fa9909777b 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, + GetXComByKeys, GetXComCount, GetXComSequenceItem, GetXComSequenceSlice, @@ -80,6 +81,7 @@ handle_get_ti_count, handle_get_variable_keys, handle_get_xcom, + handle_get_xcom_by_keys, handle_get_xcom_count, handle_get_xcom_sequence_item, handle_get_xcom_sequence_slice, @@ -700,6 +702,8 @@ def _handle_request(self, msg: ToManager, log: FilteringBoundLogger, req_id: int resp, dump_opts = handle_get_xcom_sequence_item(self.client, msg) elif isinstance(msg, GetXComSequenceSlice): resp, dump_opts = handle_get_xcom_sequence_slice(self.client, msg) + elif isinstance(msg, GetXComByKeys): + resp, dump_opts = handle_get_xcom_by_keys(self.client, msg) elif isinstance(msg, MaskSecret): handle_mask_secret(msg) elif isinstance(msg, GetTICount): 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..4d63ad05e78e9 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 @@ -766,3 +766,116 @@ def test_teamless_requester_scoping(self, client, session, dag_maker): assert forbidden.status_code == 403, forbidden.json() assert allowed.status_code == 200, allowed.json() + + +class TestGetXComByKeys: + """Tests for the POST /{dag_id}/{run_id}/{task_id}/keys batched-fetch endpoint.""" + + def _insert_xcom(self, session, ti, key, value): + session.add( + XComModel( + key=key, + value=value, + dag_run_id=ti.dag_run.id, + run_id=ti.run_id, + task_id=ti.task_id, + dag_id=ti.dag_id, + ) + ) + + def test_returns_values_in_request_key_order(self, client, create_task_instance, session): + """Values are returned in the same order as the requested keys.""" + ti = create_task_instance() + self._insert_xcom(session, ti, "return_value_0", "alpha") + self._insert_xcom(session, ti, "return_value_1", "beta") + self._insert_xcom(session, ti, "return_value_2", "gamma") + session.commit() + + response = client.post( + f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/keys", + json={"keys": ["return_value_2", "return_value_0", "return_value_1"]}, + ) + + assert response.status_code == 200 + assert response.json() == ["gamma", "alpha", "beta"] + + def test_missing_key_returned_as_none(self, client, create_task_instance, session): + """Keys not found in the database are returned as null (None).""" + ti = create_task_instance() + self._insert_xcom(session, ti, "return_value_0", "exists") + session.commit() + + response = client.post( + f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/keys", + json={"keys": ["return_value_0", "return_value_1"]}, + ) + + assert response.status_code == 200 + assert response.json() == ["exists", None] + + def test_empty_key_list_returns_empty_list(self, client, create_task_instance, session): + ti = create_task_instance() + session.commit() + + response = client.post( + f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/keys", + json={"keys": []}, + ) + + assert response.status_code == 200 + assert response.json() == [] + + def test_map_index_filter(self, client, create_task_instance, session): + """Only XCom rows matching the requested map_index are returned.""" + ti = create_task_instance() + session.add( + XComModel( + key="return_value_0", + value="for_index_0", + dag_run_id=ti.dag_run.id, + run_id=ti.run_id, + task_id=ti.task_id, + dag_id=ti.dag_id, + map_index=0, + ) + ) + session.add( + XComModel( + key="return_value_0", + value="for_default", + dag_run_id=ti.dag_run.id, + run_id=ti.run_id, + task_id=ti.task_id, + dag_id=ti.dag_id, + map_index=-1, + ) + ) + session.commit() + + response_default = client.post( + f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/keys", + json={"keys": ["return_value_0"]}, + ) + response_indexed = client.post( + f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/keys", + params={"map_index": 0}, + json={"keys": ["return_value_0"]}, + ) + + assert response_default.json() == ["for_default"] + assert response_indexed.json() == ["for_index_0"] + + def test_does_not_return_other_tasks_xcoms(self, client, create_task_instance, session): + """Keys from a different task_id are not returned.""" + ti = create_task_instance() + other_ti = create_task_instance(dag_id="other_dag", task_id="other_task") + self._insert_xcom(session, other_ti, "return_value_0", "not_mine") + session.commit() + + response = client.post( + f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/keys", + json={"keys": ["return_value_0"]}, + ) + + assert response.status_code == 200 + assert response.json() == [None] diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index 7e681c1114956..998da5b96a7c5 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -710,6 +710,22 @@ 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_by_keys( + self, + dag_id: str, + run_id: str, + task_id: str, + keys: list[str], + map_index: int = -1, + ) -> XComSequenceSliceResponse: + """Fetch multiple XCom values by key list in a single round-trip.""" + resp = self.client.post( + f"xcoms/{dag_id}/{run_id}/{task_id}/keys", + params={"map_index": map_index} if map_index >= 0 else {}, + json={"keys": keys}, + ) + return XComSequenceSliceResponse.model_validate_json(resp.read()) + class TaskStateStoreOperations: __slots__ = ("client",) diff --git a/task-sdk/src/airflow/sdk/execution_time/comms.py b/task-sdk/src/airflow/sdk/execution_time/comms.py index 5c4ae14dff1cf..42227a26b1af9 100644 --- a/task-sdk/src/airflow/sdk/execution_time/comms.py +++ b/task-sdk/src/airflow/sdk/execution_time/comms.py @@ -940,6 +940,17 @@ class GetXComSequenceSlice(BaseModel): type: Literal["GetXComSequenceSlice"] = "GetXComSequenceSlice" +class GetXComByKeys(BaseModel): + """Fetch multiple XCom values by key list in a single round-trip.""" + + keys: list[str] + dag_id: str + run_id: str + task_id: str + map_index: int = -1 + type: Literal["GetXComByKeys"] = "GetXComByKeys" + + class SetXCom(BaseModel): key: str value: JsonValue @@ -1274,6 +1285,7 @@ class GetDag(BaseModel): | GetVariable | GetVariableKeys | GetXCom + | GetXComByKeys | 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..5db0bc0e92549 100644 --- a/task-sdk/src/airflow/sdk/execution_time/request_handlers.py +++ b/task-sdk/src/airflow/sdk/execution_time/request_handlers.py @@ -62,6 +62,7 @@ GetVariable, GetVariableKeys, GetXCom, + GetXComByKeys, GetXComCount, GetXComSequenceItem, GetXComSequenceSlice, @@ -273,6 +274,14 @@ def handle_get_xcom_sequence_slice( return xcoms, {} +def handle_get_xcom_by_keys(client: Client, msg: GetXComByKeys) -> tuple[BaseModel | None, dict[str, bool]]: + """Fetch multiple XCom values by key list in a single round-trip.""" + xcoms = client.xcoms.get_by_keys(msg.dag_id, msg.run_id, msg.task_id, msg.keys, msg.map_index) + if isinstance(xcoms, XComSequenceSliceResponse): + return XComSequenceSliceResult.from_response(xcoms), {} + return xcoms, {} + + def handle_get_xcom(client: Client, msg: GetXCom) -> tuple[BaseModel | None, dict[str, bool]]: """Fetch an XCom and normalize it for supervisor response handling.""" xcom = client.xcoms.get( From ec3a6d54284250ea25c5594778adeb045b68aead Mon Sep 17 00:00:00 2001 From: David Blain Date: Wed, 22 Jul 2026 11:22:09 +0200 Subject: [PATCH 2/6] Use batched XCom fetch in XComIterable to reduce N round-trips to one --- .../sdk/execution_time/lazy_sequence.py | 94 ++++++++++++++++++ .../execution_time/test_lazy_sequence.py | 95 ++++++++++++++++++- 2 files changed, 188 insertions(+), 1 deletion(-) diff --git a/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py b/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py index 4efb0b71368ca..52f86346f56d9 100644 --- a/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py +++ b/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py @@ -25,6 +25,9 @@ import attrs import structlog +from airflow.sdk import BaseXCom +from airflow.sdk.execution_time.xcom import XCom + if TYPE_CHECKING: from airflow.sdk.definitions.xcom_arg import PlainXComArg from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance @@ -166,6 +169,97 @@ def __getitem__(self, key: int | slice) -> T | Sequence[T]: return XCom.deserialize_value(_XComWrapper(msg.root)) +class XComIterable(Sequence): + """An iterable that lazily fetches XCom values one by one instead of loading all at once.""" + + def __init__(self, task_id: str, dag_id: str, run_id: str, length: int, map_index: int | None = None): + self.task_id = task_id + self.dag_id = dag_id + self.run_id = run_id + self.length = length + self.map_index = map_index + + def __iter__(self) -> Iterator: + """Fetch all XCom values in a single round-trip, deserializing lazily.""" + if not self.length: + return iter(()) + + from airflow.sdk.execution_time.comms import GetXComByKeys, XComSequenceSliceResult + from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS + + keys = [f"{BaseXCom.XCOM_RETURN_KEY}_{i}" for i in range(self.length)] + msg = SUPERVISOR_COMMS.send( + GetXComByKeys( + keys=keys, + dag_id=self.dag_id, + run_id=self.run_id, + task_id=self.task_id, + map_index=self.map_index if self.map_index is not None else -1, + ) + ) + if not isinstance(msg, XComSequenceSliceResult): + raise TypeError(f"Got unexpected response to GetXComByKeys: {msg!r}") + return iter(XCom.deserialize_value(_XComWrapper(v)) for v in msg.root) + + def __len__(self) -> int: + return self.length + + @overload + def __getitem__(self, key: int) -> Any: ... + + @overload + def __getitem__(self, key: slice) -> Sequence[Any]: ... + + def __getitem__(self, key: int | slice) -> Any | Sequence[Any]: + """Allow direct indexing so this works like a sequence.""" + if isinstance(key, slice): + from airflow.sdk.execution_time.comms import GetXComByKeys, XComSequenceSliceResult + from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS + + indices = range(*key.indices(len(self))) + if not indices: + return [] + keys = [f"{BaseXCom.XCOM_RETURN_KEY}_{i}" for i in indices] + msg = SUPERVISOR_COMMS.send( + GetXComByKeys( + keys=keys, + dag_id=self.dag_id, + run_id=self.run_id, + task_id=self.task_id, + map_index=self.map_index if self.map_index is not None else -1, + ) + ) + if not isinstance(msg, XComSequenceSliceResult): + raise TypeError(f"Got unexpected response to GetXComByKeys: {msg!r}") + return [XCom.deserialize_value(_XComWrapper(v)) for v in msg.root] + + if not (0 <= key < self.length): + raise IndexError(key) + + return XCom.get_one( + key=f"{BaseXCom.XCOM_RETURN_KEY}_{key}", + dag_id=self.dag_id, + task_id=self.task_id, + run_id=self.run_id, + map_index=self.map_index, + ) + + def serialize(self) -> dict: + """Ensure the object is JSON serializable.""" + return { + "task_id": self.task_id, + "dag_id": self.dag_id, + "run_id": self.run_id, + "length": self.length, + "map_index": self.map_index, + } + + @classmethod + def deserialize(cls, data: dict, version: int): + """Ensure the object is JSON deserializable.""" + return XComIterable(**data) + + def _coerce_slice_index(value: Any) -> int | None: """ Check slice attribute's type and convert it to int. diff --git a/task-sdk/tests/task_sdk/execution_time/test_lazy_sequence.py b/task-sdk/tests/task_sdk/execution_time/test_lazy_sequence.py index 16149ae9e26c3..dfef9cde39392 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_lazy_sequence.py +++ b/task-sdk/tests/task_sdk/execution_time/test_lazy_sequence.py @@ -26,6 +26,7 @@ from airflow.sdk.exceptions import ErrorType from airflow.sdk.execution_time.comms import ( ErrorResponse, + GetXComByKeys, GetXComCount, GetXComSequenceItem, GetXComSequenceSlice, @@ -33,7 +34,7 @@ XComSequenceIndexResult, XComSequenceSliceResult, ) -from airflow.sdk.execution_time.lazy_sequence import LazyXComSequence +from airflow.sdk.execution_time.lazy_sequence import LazyXComSequence, XComIterable from airflow.sdk.execution_time.xcom import resolve_xcom_backend from tests_common.test_utils.config import conf_vars @@ -172,3 +173,95 @@ def test_getitem_slice(mock_supervisor_comms, lazy_sequence): step=None, ), ) + + +class TestXComIterable: + @pytest.fixture + def iterable(self): + return XComIterable(task_id="task", dag_id="dag", run_id="run", length=3) + + def test_iter_sends_single_batch_request(self, mock_supervisor_comms, iterable): + mock_supervisor_comms.send.return_value = XComSequenceSliceResult(root=["a", "b", "c"]) + assert list(iterable) == ["a", "b", "c"] + mock_supervisor_comms.send.assert_called_once_with( + GetXComByKeys( + keys=["return_value_0", "return_value_1", "return_value_2"], + dag_id="dag", + run_id="run", + task_id="task", + map_index=-1, + ) + ) + + def test_iter_empty_length_skips_supervisor_call(self, mock_supervisor_comms): + iterable = XComIterable(task_id="task", dag_id="dag", run_id="run", length=0) + assert list(iterable) == [] + mock_supervisor_comms.send.assert_not_called() + + def test_iter_is_lazy(self, mock_supervisor_comms, iterable): + """Deserialization happens on demand - breaking early skips remaining items.""" + mock_supervisor_comms.send.return_value = XComSequenceSliceResult(root=["a", "b", "c"]) + it = iter(iterable) + assert next(it) == "a" + # Only one send was made for the whole batch, but only one item consumed. + mock_supervisor_comms.send.assert_called_once() + + def test_iter_uses_map_index(self, mock_supervisor_comms): + iterable = XComIterable(task_id="task", dag_id="dag", run_id="run", length=2, map_index=5) + mock_supervisor_comms.send.return_value = XComSequenceSliceResult(root=["x", "y"]) + list(iterable) + mock_supervisor_comms.send.assert_called_once_with( + GetXComByKeys( + keys=["return_value_0", "return_value_1"], + dag_id="dag", + run_id="run", + task_id="task", + map_index=5, + ) + ) + + def test_iter_unexpected_response_raises(self, mock_supervisor_comms, iterable): + mock_supervisor_comms.send.return_value = XComSequenceIndexResult(root="oops") + with pytest.raises(TypeError, match="Got unexpected response to GetXComByKeys"): + list(iterable) + + def test_getitem_integer(self, mock_supervisor_comms, iterable): + """Single-index access still uses XCom.get_one (one-item fetch).""" + from unittest.mock import patch + + with patch("airflow.sdk.execution_time.lazy_sequence.XCom") as mock_xcom: + mock_xcom.get_one.return_value = "val" + assert iterable[1] == "val" + mock_xcom.get_one.assert_called_once_with( + key="return_value_1", + dag_id="dag", + task_id="task", + run_id="run", + map_index=None, + ) + mock_supervisor_comms.send.assert_not_called() + + def test_getitem_integer_out_of_range(self, mock_supervisor_comms, iterable): + with pytest.raises(IndexError): + iterable[10] + mock_supervisor_comms.send.assert_not_called() + + def test_getitem_slice_sends_batch(self, mock_supervisor_comms, iterable): + mock_supervisor_comms.send.return_value = XComSequenceSliceResult(root=["a", "c"]) + assert iterable[::2] == ["a", "c"] + mock_supervisor_comms.send.assert_called_once_with( + GetXComByKeys( + keys=["return_value_0", "return_value_2"], + dag_id="dag", + run_id="run", + task_id="task", + map_index=-1, + ) + ) + + def test_getitem_empty_slice_skips_supervisor_call(self, mock_supervisor_comms, iterable): + assert iterable[5:2] == [] + mock_supervisor_comms.send.assert_not_called() + + def test_len(self, iterable): + assert len(iterable) == 3 From 25f685061b15ef30745579cedf71f8d9606223fb Mon Sep 17 00:00:00 2001 From: David Blain Date: Wed, 22 Jul 2026 12:49:10 +0200 Subject: [PATCH 3/6] Updated generated API datamodels and schema.json --- .../airflow/sdk/api/datamodels/_generated.py | 8 ++++ .../sdk/execution_time/schema/schema.json | 43 +++++++++++++++++++ 2 files changed, 51 insertions(+) diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 72aebeab3e49b..087e1fd90a2a3 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -513,6 +513,14 @@ class VariableResponse(BaseModel): value: Annotated[str | None, Field(title="Value")] = None +class XComKeysRequest(BaseModel): + """ + Request body for fetching multiple XCom values by key list in a single query. + """ + + keys: Annotated[list[str], Field(title="Keys")] + + class XComResponse(BaseModel): """ XCom schema for responses with fields that are needed for Runtime. 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 0ec8fe4e49a1c..749eb5e017008 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -2747,6 +2747,49 @@ "title": "GetXCom", "type": "object" }, + "GetXComByKeys": { + "description": "Fetch multiple XCom values by key list in a single round-trip.", + "properties": { + "keys": { + "items": { + "type": "string" + }, + "title": "Keys", + "type": "array" + }, + "dag_id": { + "title": "Dag Id", + "type": "string" + }, + "run_id": { + "title": "Run Id", + "type": "string" + }, + "task_id": { + "title": "Task Id", + "type": "string" + }, + "map_index": { + "default": -1, + "title": "Map Index", + "type": "integer" + }, + "type": { + "const": "GetXComByKeys", + "default": "GetXComByKeys", + "title": "Type", + "type": "string" + } + }, + "required": [ + "keys", + "dag_id", + "run_id", + "task_id" + ], + "title": "GetXComByKeys", + "type": "object" + }, "GetXComCount": { "description": "Get the number of (mapped) XCom values available.", "properties": { From dbf25d534d39425dbf5fb407412c43599307118e Mon Sep 17 00:00:00 2001 From: David Blain Date: Wed, 22 Jul 2026 13:31:46 +0200 Subject: [PATCH 4/6] Fix unit tests for TestGetXComByKeys --- .../execution_api/versions/head/test_xcoms.py | 59 +++++++++++++------ 1 file changed, 40 insertions(+), 19 deletions(-) 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 4d63ad05e78e9..77f51556ad613 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 @@ -64,11 +64,10 @@ def _( dag_id: str = Path(), run_id: str = Path(), task_id: str = Path(), - xcom_key: str = Path(alias="key"), token=CurrentTIToken, ): with create_session() as session: - has_xcom_access(dag_id, run_id, task_id, xcom_key, request, session, token) + has_xcom_access(dag_id, run_id, task_id, request, session, token) raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail={ @@ -825,45 +824,67 @@ def test_empty_key_list_returns_empty_list(self, client, create_task_instance, s assert response.status_code == 200 assert response.json() == [] - def test_map_index_filter(self, client, create_task_instance, session): + def test_map_index_filter(self, client, dag_maker, session): """Only XCom rows matching the requested map_index are returned.""" - ti = create_task_instance() + + 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=["a", "b"]) + dag_run = dag_maker.create_dagrun(run_id="run") + tis = {ti.map_index: ti for ti in dag_run.task_instances} + + # Insert XComs for map_index=0 and map_index=1 (both have real TI rows). + # map_index=-1 has no TI row for a mapped task, so it is not inserted. session.add( XComModel( key="return_value_0", value="for_index_0", - dag_run_id=ti.dag_run.id, - run_id=ti.run_id, - task_id=ti.task_id, - dag_id=ti.dag_id, + dag_run_id=dag_run.id, + run_id=dag_run.run_id, + task_id=tis[0].task_id, + dag_id=tis[0].dag_id, map_index=0, ) ) session.add( XComModel( key="return_value_0", - value="for_default", - dag_run_id=ti.dag_run.id, - run_id=ti.run_id, - task_id=ti.task_id, - dag_id=ti.dag_id, - map_index=-1, + value="for_index_1", + dag_run_id=dag_run.id, + run_id=dag_run.run_id, + task_id=tis[1].task_id, + dag_id=tis[1].dag_id, + map_index=1, ) ) session.commit() + dag_id = tis[0].dag_id + task_id = tis[0].task_id + + # Default map_index=-1: no rows stored there, so null is returned. response_default = client.post( - f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/keys", + f"/execution/xcoms/{dag_id}/{dag_run.run_id}/{task_id}/keys", json={"keys": ["return_value_0"]}, ) - response_indexed = client.post( - f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/keys", + response_index_0 = client.post( + f"/execution/xcoms/{dag_id}/{dag_run.run_id}/{task_id}/keys", params={"map_index": 0}, json={"keys": ["return_value_0"]}, ) + response_index_1 = client.post( + f"/execution/xcoms/{dag_id}/{dag_run.run_id}/{task_id}/keys", + params={"map_index": 1}, + json={"keys": ["return_value_0"]}, + ) - assert response_default.json() == ["for_default"] - assert response_indexed.json() == ["for_index_0"] + assert response_default.json() == [None] + assert response_index_0.json() == ["for_index_0"] + assert response_index_1.json() == ["for_index_1"] def test_does_not_return_other_tasks_xcoms(self, client, create_task_instance, session): """Keys from a different task_id are not returned.""" From 3da5abf95d6567f2885dc2d22a98edf9f680f46a Mon Sep 17 00:00:00 2001 From: David Blain Date: Wed, 22 Jul 2026 16:59:00 +0200 Subject: [PATCH 5/6] Add GetXComByKeys supervisor message test coverage --- .../task_sdk/execution_time/test_supervisor.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) 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 c39ceb92b0e3b..47e42d9df6dd8 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -117,6 +117,7 @@ GetVariable, GetVariableKeys, GetXCom, + GetXComByKeys, GetXComCount, GetXComSequenceItem, GetXComSequenceSlice, @@ -2823,6 +2824,22 @@ class RequestTestCase: ), test_id="get_xcom_seq_slice", ), + RequestTestCase( + message=GetXComByKeys( + keys=["return_value_0", "return_value_1"], + dag_id="test_dag", + run_id="test_run", + task_id="test_task", + map_index=-1, + ), + expected_body={"root": ["foo", "bar"], "type": "XComSequenceSliceResult"}, + client_mock=ClientMock( + method_path="xcoms.get_by_keys", + args=("test_dag", "test_run", "test_task", ["return_value_0", "return_value_1"], -1), + response=XComSequenceSliceResult(root=["foo", "bar"]), + ), + test_id="get_xcom_by_keys", + ), RequestTestCase( message=TaskState(state=TaskInstanceState.SKIPPED, end_date=timezone.parse("2024-10-31T12:00:00Z")), test_id="patch_task_instance_to_skipped", From c9e2e4292f7b8cb6a023ef6c11c8d85472f07b0e Mon Sep 17 00:00:00 2001 From: David Blain Date: Wed, 22 Jul 2026 20:11:04 +0200 Subject: [PATCH 6/6] Exclude GetXComByKeys from triggerer supervisor message check --- airflow-core/tests/unit/jobs/test_triggerer_job.py | 1 + 1 file changed, 1 insertion(+) diff --git a/airflow-core/tests/unit/jobs/test_triggerer_job.py b/airflow-core/tests/unit/jobs/test_triggerer_job.py index fa4d68f5fff6a..e53f8aa61a6e6 100644 --- a/airflow-core/tests/unit/jobs/test_triggerer_job.py +++ b/airflow-core/tests/unit/jobs/test_triggerer_job.py @@ -2596,6 +2596,7 @@ def get_type_names(union_type): "GetPreviousDagRun", "GetTaskBreadcrumbs", "GetTaskRescheduleStartDate", + "GetXComByKeys", "GetXComCount", "GetXComSequenceItem", "GetXComSequenceSlice",