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..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 @@ -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,40 @@ 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: + any_pushed = XComModel.get_many( + run_id=run_id, + key=key, + task_ids=task_id, + dag_ids=dag_id, + ).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, key=key, @@ -157,15 +213,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,14 +240,33 @@ def get_mapped_xcom_by_slice( params: Annotated[GetXComSliceFilterParams, Query()], session: SessionDep, ) -> XComSequenceSliceResponse: - query = XComModel.get_many( + mapped_length = ( + None if params.include_prior_dates else _get_mapped_length(dag_id, run_id, task_id, session=session) + ) + 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, - ) - query = query.order_by(None) + ).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 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)) + } + values = [by_map_index.get(i) for i in range(mapped_length)] + return XComSequenceSliceResponse(values[params.start : params.stop : params.step]) + + query = base_query step = params.step or 1 @@ -276,6 +345,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 +360,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..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 @@ -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,68 @@ 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" + + 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", [ @@ -278,7 +342,112 @@ 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"] + + 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"), 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,