Add a POST /xcoms/{dag_id}/{run_id}/{task_id}/keys execution API endpoint

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

@@ -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(

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):

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={
@@ -766,3 +765,138 @@ 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, 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(
+            XComModel(
+                key="return_value_0",
+                value="for_index_0",
+                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_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/{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

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

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",)

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.

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