From 219749dbdddfb2348580102e5aba0117456ecf2b Mon Sep 17 00:00:00 2001 From: ColtenOuO Date: Sun, 23 Aug 2026 17:03:59 +0000 Subject: [PATCH 1/3] Fix mismatched XCom values when a dynamically mapped task has a gap Reading XCom values for a mapped task or task group by sequence index or slice treated the position among existing XCom rows as the map_index. When a mapped task instance hasn't pushed an XCom yet (still up_for_reschedule, or skipped), the gap silently shifted every later map_index's value down by one, so downstream tasks could read the wrong value or hit an IndexError even though the upstream task instances existed. --- .../api_fastapi/execution_api/routes/xcoms.py | 88 ++++++++++++-- .../execution_api/versions/head/test_xcoms.py | 115 ++++++++++++++++-- .../execution_time/test_lazy_sequence.py | 9 ++ 3 files changed, 190 insertions(+), 22 deletions(-) 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..2bad84c444540 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,7 +22,7 @@ from fastapi import APIRouter, Body, Depends, HTTPException, Path, Query, Request, Response, status from pydantic import JsonValue -from sqlalchemy import delete +from sqlalchemy import delete, func, select from sqlalchemy.sql.selectable import Select from airflow.api_fastapi.common.db.common import SessionDep @@ -33,6 +33,7 @@ XComSequenceSliceResponse, ) from airflow.api_fastapi.execution_api.security import CurrentTIToken +from airflow.models.taskinstance import TaskInstance from airflow.models.taskmap import TaskMap from airflow.models.xcom import XComModel from airflow.utils.db import get_query_count @@ -133,6 +134,27 @@ async def xcom_query( return query +def _get_mapped_length(dag_id: str, run_id: str, task_id: str, *, session: SessionDep) -> int | None: + """ + Get the number of mapped task instances ``task_id`` expands into for this run. + + Returns *None* if ``task_id`` is not a mapped task in this run (no task instance + with ``map_index >= 0`` exists), in which case callers should fall back to treating + the XCom rows as an unmapped/non-sparse sequence. + """ + max_map_index = session.scalar( + select(func.max(TaskInstance.map_index)).where( + TaskInstance.dag_id == dag_id, + TaskInstance.run_id == run_id, + TaskInstance.task_id == task_id, + TaskInstance.map_index >= 0, + ) + ) + if max_map_index is None: + return None + return max_map_index + 1 + + @router.get( "/{dag_id}/{run_id}/{task_id}/{key:path}/item/{offset}", description="Get a single XCom value from a mapped task by sequence index", @@ -145,6 +167,33 @@ def get_mapped_xcom_by_index( offset: int, session: SessionDep, ) -> XComSequenceIndexResponse: + mapped_length = _get_mapped_length(dag_id, run_id, task_id, session=session) + + def not_found() -> HTTPException: + message = ( + f"XCom with {key=} {offset=} not found for task {task_id!r} in DAG run {run_id!r} of {dag_id!r}" + ) + return HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail={"reason": "not_found", "message": message}, + ) + + if mapped_length is not None: + map_index = offset if offset >= 0 else mapped_length + offset + if not (0 <= map_index < mapped_length): + raise not_found() + xcom_query = XComModel.get_many( + run_id=run_id, + key=key, + task_ids=task_id, + dag_ids=dag_id, + map_indexes=map_index, + ) + result: tuple[XComModel] | None = session.scalars(xcom_query.order_by(None)).first() + if result is None: + return XComSequenceIndexResponse(None) + return XComSequenceIndexResponse((result[0] if isinstance(result, tuple) else result).value) + xcom_query = XComModel.get_many( run_id=run_id, key=key, @@ -157,15 +206,9 @@ def get_mapped_xcom_by_index( else: xcom_query = xcom_query.order_by(XComModel.map_index.desc()).offset(-1 - offset) - result: tuple[XComModel] | None - if (result := session.scalars(xcom_query).first()) is None: - message = ( - f"XCom with {key=} {offset=} not found for task {task_id!r} in DAG run {run_id!r} of {dag_id!r}" - ) - raise HTTPException( - status_code=status.HTTP_404_NOT_FOUND, - detail={"reason": "not_found", "message": message}, - ) + result = session.scalars(xcom_query).first() + if result is None: + raise not_found() return XComSequenceIndexResponse((result[0] if isinstance(result, tuple) else result).value) @@ -190,6 +233,23 @@ def get_mapped_xcom_by_slice( params: Annotated[GetXComSliceFilterParams, Query()], session: SessionDep, ) -> XComSequenceSliceResponse: + mapped_length = ( + None if params.include_prior_dates else _get_mapped_length(dag_id, run_id, task_id, session=session) + ) + if mapped_length is not None: + query = XComModel.get_many( + run_id=run_id, + key=key, + task_ids=task_id, + dag_ids=dag_id, + ).order_by(None) + by_map_index = { + row.map_index: row.value + for row in session.execute(query.with_only_columns(XComModel.map_index, XComModel.value)) + } + values = [by_map_index.get(i) for i in range(mapped_length)] + return XComSequenceSliceResponse(values[params.start : params.stop : params.step]) + query = XComModel.get_many( run_id=run_id, key=key, @@ -276,6 +336,9 @@ def get_mapped_xcom_by_slice( description="Returns the count of mapped XCom values found in the `Content-Range` response header", ) def head_xcom( + dag_id: str, + run_id: str, + task_id: str, response: Response, session: SessionDep, xcom_query: Annotated[Select, Depends(xcom_query)], @@ -288,7 +351,10 @@ def head_xcom( detail={"reason": "invalid_request", "message": "Cannot specify map_index in a HEAD request"}, ) - count = get_query_count(xcom_query, session=session) + if (mapped_length := _get_mapped_length(dag_id, run_id, task_id, session=session)) is not None: + count = mapped_length + else: + count = get_query_count(xcom_query, session=session) # Tell the caller how many items in this query. We define a custom range unit (HTTP spec only defines # "bytes" but we can add our own) response.headers["Content-Range"] = f"map_indexes {count}" 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..ef51b7ab03f32 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 @@ -35,7 +35,7 @@ from airflow.providers.standard.operators.empty import EmptyOperator from airflow.sdk.serde import deserialize, serialize from airflow.utils.session import create_session -from airflow.utils.state import DagRunState +from airflow.utils.state import DagRunState, TaskInstanceState from tests_common.test_utils.config import conf_vars @@ -142,38 +142,40 @@ def test_xcom_access_denied(self, client, caplog): ("offset", "expected_status", "expected_json"), [ pytest.param( - -4, + -5, 404, { "detail": { "reason": "not_found", "message": ( - "XCom with key='xcom_1' offset=-4 not found " + "XCom with key='xcom_1' offset=-5 not found " "for task 'task' in DAG run 'runid' of 'dag'" ), }, }, - id="-4", + id="-5", ), - pytest.param(-3, 200, "f", id="-3"), + pytest.param(-4, 200, "f", id="-4"), + pytest.param(-3, 200, None, id="-3"), pytest.param(-2, 200, "o", id="-2"), pytest.param(-1, 200, "b", id="-1"), pytest.param(0, 200, "f", id="0"), - pytest.param(1, 200, "o", id="1"), - pytest.param(2, 200, "b", id="2"), + pytest.param(1, 200, None, id="1"), + pytest.param(2, 200, "o", id="2"), + pytest.param(3, 200, "b", id="3"), pytest.param( - 3, + 4, 404, { "detail": { "reason": "not_found", "message": ( - "XCom with key='xcom_1' offset=3 not found " + "XCom with key='xcom_1' offset=4 not found " "for task 'task' in DAG run 'runid' of 'dag'" ), }, }, - id="3", + id="4", ), ], ) @@ -218,6 +220,30 @@ def __init__(self, *, x, **kwargs): assert response.status_code == expected_status assert response.json() == expected_json + def test_xcom_get_with_offset_unmapped_task_falls_back_to_position( + self, client, create_task_instance, session + ): + """ + An unmapped task (map_index=-1) has no logical map_index range to key a sparse + array off of, so offset access falls back to indexing by position among existing + XCom rows, same as before this endpoint gained sparse-index resolution. + """ + ti = create_task_instance() + x = XComModel( + key="xcom_1", + value="f", + dag_run_id=ti.dag_run.id, + run_id=ti.run_id, + task_id=ti.task_id, + dag_id=ti.dag_id, + ) + session.add(x) + session.commit() + + response = client.get(f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/xcom_1/item/0") + assert response.status_code == 200 + assert response.json() == "f" + @pytest.mark.parametrize( "key", [ @@ -278,7 +304,74 @@ def __init__(self, *, x, **kwargs): response = client.get(f"/execution/xcoms/dag/runid/task/xcom_1/slice?{urllib.parse.urlencode(qs)}") assert response.status_code == 200 - assert response.json() == ["f", "o", "b"][key] + assert response.json() == xcom_values[key] + + def test_xcom_get_with_slice_unmapped_task_falls_back_to_compaction( + self, client, create_task_instance, session + ): + """ + An unmapped task (map_index=-1) has no logical map_index range to key a sparse + array off of, so slicing falls back to the existing XCom rows in map_index order, + same as before this endpoint gained sparse-index resolution. + """ + ti = create_task_instance() + x = XComModel( + key="xcom_1", + value="f", + dag_run_id=ti.dag_run.id, + run_id=ti.run_id, + task_id=ti.task_id, + dag_id=ti.dag_id, + ) + session.add(x) + session.commit() + + response = client.get(f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/xcom_1/slice") + assert response.status_code == 200 + assert response.json() == ["f"] + + def test_xcom_get_with_slice_and_count_unfinished_mapped_task(self, client, dag_maker, session): + """ + Regression test for a mapped task instance that hasn't pushed any XCom yet (e.g. + still up_for_reschedule), as opposed to one that pushed a None value. + + Both the slice and count (HEAD) endpoints must reflect the full logical + map_index range, not just the map indexes that have an XCom row. + """ + + 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=["f", "o", "o", "b"]) + dag_run = dag_maker.create_dagrun(run_id="runid") + tis = {ti.map_index: ti for ti in dag_run.task_instances} + + tis[1].state = TaskInstanceState.UP_FOR_RESCHEDULE + for map_index, value in {0: "f", 2: "o", 3: "b"}.items(): + ti = tis[map_index] + session.add( + XComModel( + key="xcom_1", + value=value, + dag_run_id=ti.dag_run.id, + run_id=ti.run_id, + task_id=ti.task_id, + dag_id=ti.dag_id, + map_index=map_index, + ) + ) + session.commit() + + head_response = client.head("/execution/xcoms/dag/runid/task/xcom_1") + assert head_response.status_code == 200 + assert head_response.headers["Content-Range"] == "map_indexes 4" + + slice_response = client.get("/execution/xcoms/dag/runid/task/xcom_1/slice") + assert slice_response.status_code == 200 + assert slice_response.json() == ["f", None, "o", "b"] @pytest.mark.parametrize( ("include_prior_dates", "expected_xcoms"), 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..708da3b8b6ff9 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 @@ -139,6 +139,15 @@ def test_getitem_calls_correct_deserialise(monkeypatch, mock_supervisor_comms, l ) +def test_getitem_index_sparse_gap_returns_none(mock_supervisor_comms, lazy_sequence): + """ + A mapped task instance that hasn't pushed an XCom yet (still within the logical + map_index range, e.g. up_for_reschedule or skipped) resolves to None, not IndexError. + """ + mock_supervisor_comms.send.return_value = XComSequenceIndexResult(root=None) + assert lazy_sequence[4] is None + + def test_getitem_indexerror(mock_supervisor_comms, lazy_sequence): mock_supervisor_comms.send.return_value = ErrorResponse( error=ErrorType.XCOM_NOT_FOUND, From e38c3f2a7eb0d372015f7ea018d2fbe3ee2a14ad Mon Sep 17 00:00:00 2001 From: ColtenOuO Date: Sun, 23 Aug 2026 17:31:49 +0000 Subject: [PATCH 2/3] Avoid loading every XCom row when a mapped task has no gap The sparse-index fix for slice reads unconditionally fetched all of a mapped task's XCom rows into memory to place each value at its true map_index. That regressed the common case (no gap) versus the previous SQL-side pagination, and XCom payloads have no size cap -- only the number of mapped instances does. Only fall into the full-fetch path once a gap is actually detected via a COUNT query. --- .../api_fastapi/execution_api/routes/xcoms.py | 30 +++++++-------- .../execution_api/versions/head/test_xcoms.py | 38 +++++++++++++++++++ 2 files changed, 52 insertions(+), 16 deletions(-) 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 2bad84c444540..6030101d69813 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 @@ -236,28 +236,26 @@ def get_mapped_xcom_by_slice( mapped_length = ( None if params.include_prior_dates else _get_mapped_length(dag_id, run_id, task_id, session=session) ) - if mapped_length is not None: - query = XComModel.get_many( - run_id=run_id, - key=key, - task_ids=task_id, - dag_ids=dag_id, - ).order_by(None) + base_query = XComModel.get_many( + run_id=run_id, + key=key, + task_ids=task_id, + dag_ids=dag_id, + include_prior_dates=params.include_prior_dates, + ).order_by(None) + + # Only fall back to fetching every XCom row (needed to place each value at its true + # map_index) when there's an actual gap. The common case -- every mapped instance has + # pushed its XCom -- keeps the cheaper SQL-side OFFSET/LIMIT/slice() below. + if mapped_length is not None and get_query_count(base_query, session=session) < mapped_length: by_map_index = { row.map_index: row.value - for row in session.execute(query.with_only_columns(XComModel.map_index, XComModel.value)) + for row in session.execute(base_query.with_only_columns(XComModel.map_index, XComModel.value)) } values = [by_map_index.get(i) for i in range(mapped_length)] return XComSequenceSliceResponse(values[params.start : params.stop : params.step]) - query = XComModel.get_many( - run_id=run_id, - key=key, - task_ids=task_id, - dag_ids=dag_id, - include_prior_dates=params.include_prior_dates, - ) - query = query.order_by(None) + query = base_query step = params.step or 1 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 ef51b7ab03f32..cdf016623a1ff 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 @@ -373,6 +373,44 @@ def __init__(self, *, x, **kwargs): assert slice_response.status_code == 200 assert slice_response.json() == ["f", None, "o", "b"] + def test_xcom_get_with_slice_mapped_task_without_gap_uses_sql_pagination( + self, client, dag_maker, session + ): + """ + When every mapped instance has already pushed an XCom, the slice endpoint must + not fall into the gap-filling branch that loads every row into memory -- it should + keep using SQL-side OFFSET/LIMIT/slice() pagination. + """ + + 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=["f", "o", "o", "b"]) + dag_run = dag_maker.create_dagrun(run_id="runid") + tis = {ti.map_index: ti for ti in dag_run.task_instances} + + for map_index, value in enumerate(["f", "o", "o", "b"]): + ti = tis[map_index] + session.add( + XComModel( + key="xcom_1", + value=value, + dag_run_id=ti.dag_run.id, + run_id=ti.run_id, + task_id=ti.task_id, + dag_id=ti.dag_id, + map_index=map_index, + ) + ) + session.commit() + + response = client.get("/execution/xcoms/dag/runid/task/xcom_1/slice") + assert response.status_code == 200 + assert response.json() == ["f", "o", "o", "b"] + @pytest.mark.parametrize( ("include_prior_dates", "expected_xcoms"), [[True, ["earlier_value", "later_value"]], [False, ["later_value"]]], From 562d9b5cf98974af5bab57237644dd2bf82b33cc Mon Sep 17 00:00:00 2001 From: ColtenOuO Date: Sun, 23 Aug 2026 19:17:01 +0000 Subject: [PATCH 3/3] Fix 404 for a nonexistent XCom key on a mapped task's sequence endpoints The sparse-index handling only checked whether the task was mapped, not whether the requested key had ever been pushed at all. Querying a wrong or nonexistent key against a mapped task incorrectly took the sparse path and returned every offset as null / the full-length array filled with null, instead of 404 / an empty list. --- .../api_fastapi/execution_api/routes/xcoms.py | 37 +++++++++++------- .../execution_api/versions/head/test_xcoms.py | 38 +++++++++++++++++++ 2 files changed, 62 insertions(+), 13 deletions(-) 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 6030101d69813..c6686024ace91 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 @@ -179,20 +179,27 @@ def not_found() -> HTTPException: ) if mapped_length is not None: - map_index = offset if offset >= 0 else mapped_length + offset - if not (0 <= map_index < mapped_length): - raise not_found() - xcom_query = XComModel.get_many( + any_pushed = XComModel.get_many( run_id=run_id, key=key, task_ids=task_id, dag_ids=dag_id, - map_indexes=map_index, - ) - result: tuple[XComModel] | None = session.scalars(xcom_query.order_by(None)).first() - if result is None: - return XComSequenceIndexResponse(None) - return XComSequenceIndexResponse((result[0] if isinstance(result, tuple) else result).value) + ).order_by(None) + if session.execute(any_pushed.limit(1)).first() is not None: + map_index = offset if offset >= 0 else mapped_length + offset + if not (0 <= map_index < mapped_length): + raise not_found() + xcom_query = XComModel.get_many( + run_id=run_id, + key=key, + task_ids=task_id, + dag_ids=dag_id, + map_indexes=map_index, + ) + result: tuple[XComModel] | None = session.scalars(xcom_query.order_by(None)).first() + if result is None: + return XComSequenceIndexResponse(None) + return XComSequenceIndexResponse((result[0] if isinstance(result, tuple) else result).value) xcom_query = XComModel.get_many( run_id=run_id, @@ -245,9 +252,13 @@ def get_mapped_xcom_by_slice( ).order_by(None) # Only fall back to fetching every XCom row (needed to place each value at its true - # map_index) when there's an actual gap. The common case -- every mapped instance has - # pushed its XCom -- keeps the cheaper SQL-side OFFSET/LIMIT/slice() below. - if mapped_length is not None and get_query_count(base_query, session=session) < mapped_length: + # map_index) when there's an actual gap for this key -- i.e. some, but not all, of the + # mapped instances have pushed it. If none have (e.g. the key doesn't exist at all), + # that's not a gap to fill with None; it's an empty/not-found result, handled below by + # the unchanged SQL-side path. The common case -- every instance has pushed -- also + # keeps that cheaper OFFSET/LIMIT/slice() path. + existing_count = get_query_count(base_query, session=session) if mapped_length is not None else 0 + if mapped_length is not None and 0 < existing_count < mapped_length: by_map_index = { row.map_index: row.value for row in session.execute(base_query.with_only_columns(XComModel.map_index, XComModel.value)) 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 cdf016623a1ff..0cb362dcb6d68 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 @@ -244,6 +244,44 @@ def test_xcom_get_with_offset_unmapped_task_falls_back_to_position( assert response.status_code == 200 assert response.json() == "f" + def test_xcom_get_with_offset_mapped_task_wrong_key_not_found(self, client, dag_maker, session): + """ + A mapped task with no XCom rows at all for the requested key (e.g. the key is + wrong, or nothing has been pushed yet) must 404, not be treated as a sparse + sequence and resolve every offset to None. + """ + + 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=["f", "o", "o", "b"]) + dag_maker.create_dagrun(run_id="runid") + + response = client.get("/execution/xcoms/dag/runid/task/non_existent_key/item/0") + assert response.status_code == 404 + + def test_xcom_get_with_slice_mapped_task_wrong_key_not_found(self, client, dag_maker, session): + """ + Same as above, but for the slice endpoint: no XCom rows at all for the key means + an empty result, not a full-length list of None. + """ + + 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=["f", "o", "o", "b"]) + dag_maker.create_dagrun(run_id="runid") + + response = client.get("/execution/xcoms/dag/runid/task/non_existent_key/slice") + assert response.status_code == 200 + assert response.json() == [] + @pytest.mark.parametrize( "key", [