Fetch several XCom values of a Taskinstance in one Execution API request

.github/CODEOWNERS

@@ -163,6 +163,8 @@ Dockerfile.ci @potiuk @ashb @gopidesupavan @amoghrajesh @jscheffl @bugraoz93 @ja
 # AIP-72 - Task SDK
 # Python SDK
 /task-sdk/ @ashb @amoghrajesh
+/task-sdk/src/airflow/sdk/execution_time/executor.py @ashb @amoghrajesh @dabla
+/task-sdk/src/airflow/sdk/definitions/iterableoperator.py @ashb @amoghrajesh @dabla
 
 # AIP-108 - Coordinators
 /task-sdk/src/airflow/sdk/coordinators/ @jason810496 @uranusjr

airflow-core/newsfragments/62922.feature.rst

@@ -0,0 +1 @@
+Task Iteration (AIP-104): ``.iterate()`` and ``.iterate_kwargs()`` on operators and ``@task`` run the mapped inputs inside one task instance, ``task_concurrency`` of them at a time, instead of one task instance per item as ``.expand()`` does. Each iteration's result is checkpointed in the task state store, so a retry resumes where the previous attempt stopped, and the task's return value is an ``XComIterable``, a lazy sequence downstream tasks index, iterate or ``.expand()`` over. The Task SDK docs page "Mapped tasks vs iterable tasks" says when to use which.

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]

airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py

@@ -30,6 +30,7 @@
 from airflow.api_fastapi.core_api.base import BaseModel
 from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
 from airflow.api_fastapi.execution_api.datamodels.xcom import (
+    XComKeysRequest,
     XComResponse,
     XComSequenceIndexResponse,
     XComSequenceSliceResponse,
@@ -45,7 +46,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,
@@ -71,6 +71,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'",
@@ -116,6 +117,7 @@ def has_xcom_access(
     dependencies=[Depends(has_xcom_access)],
 )
 
+
 log = logging.getLogger(__name__)
 
 
@@ -405,6 +407,41 @@ 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([])
+
+    xcom_read = XComModel.get_many(
+        run_id=run_id,
+        task_ids=task_id,
+        dag_ids=dag_id,
+        map_indexes=map_index,
+    )
+    entity = xcom_entity(xcom_read)
+    query = (
+        xcom_read.with_only_columns(entity.key, entity.value).where(entity.key.in_(key_list)).order_by(None)
+    )
+    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(

airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py

@@ -59,6 +59,7 @@
     AddMultiTeamToTIRunContext,
     AddStoppedTaskReport,
     AddTerminalStateRetryReasonField,
+    AddXComKeysEndpoint,
     IdentifyArchivedTaskStateUpdates,
 )
 
@@ -73,6 +74,7 @@
         AddMultiTeamToTIRunContext,
         AddStoppedTaskReport,
         IdentifyArchivedTaskStateUpdates,
+        AddXComKeysEndpoint,
     ),
     Version(
         "2026-06-30",

airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py

@@ -108,3 +108,13 @@ class AddMultiTeamToTIRunContext(VersionChange):
     def remove_multi_team_field(response: ResponseInfo) -> None:  # type: ignore[misc]
         """Strip ``multi_team`` from the run context for older clients."""
         response.body.pop("multi_team", None)
+
+
+class AddXComKeysEndpoint(VersionChange):
+    """Add the ``xcoms/{dag_id}/{run_id}/{task_id}/keys`` endpoint that fetches several XCom values by key in one query."""
+
+    description = __doc__
+
+    instructions_to_migrate_to_previous_version = (
+        endpoint("/xcoms/{dag_id}/{run_id}/{task_id}/keys", ["POST"]).didnt_exist,
+    )

airflow-core/src/airflow/dag_processing/processor.py

@@ -55,6 +55,7 @@
     GetVariable,
     GetVariableKeys,
     GetXCom,
+    GetXComByKeys,
     GetXComCount,
     GetXComSequenceItem,
     GetXComSequenceSlice,
@@ -150,6 +151,7 @@ class DagFileParsingResult(BaseModel):
     | GetXCom
     | GetXComCount
     | GetXComSequenceItem
+    | GetXComByKeys
     | GetXComSequenceSlice
     | MaskSecret,
     Field(discriminator="type"),
@@ -636,6 +638,7 @@ def _handle_parsing_result(
         GetVariable,
         GetVariableKeys,
         GetXCom,
+        GetXComByKeys,
         GetXComCount,
         GetXComSequenceItem,
         GetXComSequenceSlice,

airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py

@@ -79,11 +79,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={
@@ -875,3 +874,109 @@ 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(XComModelV2(key=key, value=value, task_instance_id=ti.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, dag_maker, session):
+        """Only XCom rows matching the requested map_index are returned."""
+
+        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(XComModelV2(key="return_value_0", value="for_index_0", task_instance_id=tis[0].id))
+        session.add(XComModelV2(key="return_value_0", value="for_index_1", task_instance_id=tis[1].id))
+        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/{dag_id}/{dag_run.run_id}/{task_id}/keys",
+            json={"keys": ["return_value_0"]},
+        )
+        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() == [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."""
+        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]

airflow-core/tests/unit/jobs/test_triggerer_job.py

@@ -3054,6 +3054,7 @@ def get_type_names(union_type):
             "GetPreviousDagRun",
             "GetTaskBreadcrumbs",
             "GetTaskRescheduleStartDate",
+            "GetXComByKeys",
             "GetXComCount",
             "GetXComSequenceItem",
             "GetXComSequenceSlice",