This is an automated email from the ASF dual-hosted git repository.
ashb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new a27f8e74da3 Let XCom backends tell loop iterations apart through
store_value (#74858)
a27f8e74da3 is described below
commit a27f8e74da303caff84ba3e747f667067b4bc4dc
Author: Ash Berlin-Taylor <[email protected]>
AuthorDate: Sun Oct 18 17:23:53 2026 +0100
Let XCom backends tell loop iterations apart through store_value (#74858)
Every iteration of a task inside a loop is its own task instance, but an
XCom
backend's `serialize_value` only receives dag_id, run_id, task_id, map_index
and key, which is shared by different iterations of the same task. A backend
might naievly derives its storage path from those values would every
iteration
to the same location, so every read would give back the last iteration's
value
n times. It only shows when the XComs are read from outside the loop, which
makes it easy to miss.
To work around this, without breaking old XCom backends, I have introduced
a new method on the backend, `store_value`, and instead of 4 (or 5 or 6)
arguments it takes a single arg of type `XComCoordinate` to prevent arg
sprawl
and to make this new fn future proof.
`XCom.set` on another task (i.e. passign a different `task_id`) cannot know
which task instance the server will resolve its coordinates to, so there the
id is unknown and the location has to be unique to the write. Leaving that
distinction to each backend would make the collision easy to reintroduce, so
`XComCoordinate.build_path` builds the location for both cases and shortens
the readable parts only when the path would exceed the backend-provided key
limit. `serialize_value` already combined converting a value with writing
it somewhere, and one more keyword would have kept growing its signature.
For each old and new hook fn, the override closest to the class wins, so a
subclass of a store_value backend that overrides serialize_value keeps its
override. Airflow can't tell whether a legacy backend's locations are
unique,
so implementing store_value is what marks a backend as safe in a loop: a Dag
that uses a loop fails at startup under a serialize_value-only backend
instead
of overwriting results silently.
Writes and deletes of another task's XCom did not follow the rule reads
use: they resolved to the caller's own task instance or the single live
one, and failed with 409 when several iterations of a loop task were live.
Resolving them like reads, from the calling task instance, lets a task in
a loop write its own iteration of a sibling and gives a caller outside the
loop the same 400 a read gets.
The backend's store and load talk to external storage, so the async paths,
which share an event loop with the triggerer's other triggers, run them in
a thread.
The trigger runner already carried its task instance id in a context
variable so the API could resolve a trigger's requests. It moved into the
Task SDK and holds the task instance, because the SDK has to recognise when
a trigger writes its own coordinates and then pass the id and region on.
The common.io ObjectStore backend now builds its object paths with
build_path, so
they now sit under `ti-<id>` or `write-<rand_uuid>`. Rows already written
keep
working because they store full paths in the metadata, and read just reads
the
path it is given.
Overwritten and failed writes still leave their old objects behind, as they
do today; cleaning those up is separate work. The configured backend also
applies only to the XComs that Python tasks and triggers write and read.
Tasks written in other languages still store theirs directly in the
database, and fixing that needs the supervisor to apply the backend for
them. Again, this doesn't change from today.
---
.../docs/authoring-and-scheduling/loops.rst | 1 +
airflow-core/docs/core-concepts/xcoms.rst | 21 +-
.../task-loops-custom-backends.significant.rst | 28 ++
.../api_fastapi/execution_api/routes/xcoms.py | 110 +++---
.../src/airflow/dag_processing/processor.py | 9 +-
.../src/airflow/jobs/triggerer_job_runner.py | 16 +-
.../src/airflow/models/task_coordinates.py | 4 +-
.../execution_api/versions/head/test_xcoms.py | 100 +++++-
.../tests/unit/dag_processing/test_processor.py | 23 +-
airflow-core/tests/unit/jobs/test_triggerer_job.py | 33 +-
devel-common/src/tests_common/pytest_plugin.py | 95 +++--
.../src/tests_common/test_utils/xcom_backends.py | 28 ++
generated/known_sdk_imports_in_core.txt | 2 +-
providers/common/io/docs/changelog.rst | 10 +
.../airflow/providers/common/io/xcom/backend.py | 75 ++--
.../io/tests/unit/common/io/xcom/test_backend.py | 55 +++
task-sdk/docs/api.rst | 2 +
task-sdk/src/airflow/sdk/bases/xcom.py | 383 ++++++++++++++++++---
.../airflow/sdk/execution_time/lazy_sequence.py | 4 +-
.../src/airflow/sdk/execution_time/supervisor.py | 8 +-
.../src/airflow/sdk/execution_time/task_runner.py | 45 ++-
task-sdk/src/airflow/sdk/types.py | 2 +
task-sdk/tests/task_sdk/bases/test_xcom.py | 379 +++++++++++++++++++-
.../task_sdk/execution_time/test_lazy_sequence.py | 27 ++
.../task_sdk/execution_time/test_supervisor.py | 30 +-
.../task_sdk/execution_time/test_task_runner.py | 83 ++++-
26 files changed, 1333 insertions(+), 240 deletions(-)
diff --git a/airflow-core/docs/authoring-and-scheduling/loops.rst
b/airflow-core/docs/authoring-and-scheduling/loops.rst
index 161567a6935..b0ea5b97575 100644
--- a/airflow-core/docs/authoring-and-scheduling/loops.rst
+++ b/airflow-core/docs/authoring-and-scheduling/loops.rst
@@ -200,6 +200,7 @@ task's latest iteration.
Limitations:
* ``include_prior_dates=True`` cannot select a loop iteration from another Dag
run. Push the result to XCom through a task outside the loop if later Dag runs
need to retrieve it with an ordinary XCom pull; see :ref:`loops-outside-result`.
+* A Dag that uses a loop needs an XCom backend that supports loops. The
default backend does. If your backend does not, for example a custom one
written before loops existed, the Dag fails with an error when it is parsed and
again when a task starts. See Custom XCom Backends in
:doc:`/core-concepts/xcoms`.
* The experimental DagRun wait API also requires an outside-loop result task.
Marking a loop member as the Dag result raises ``ValueError`` while the Dag is
parsed, and the wait API rejects results selected directly from a loop member.
Mark a task outside the loop as the result instead, for example one that
returns the last element of a loop task's results. See :ref:`dag-result`.
Connect a loop to other tasks
diff --git a/airflow-core/docs/core-concepts/xcoms.rst
b/airflow-core/docs/core-concepts/xcoms.rst
index d1d8fe75ee3..c2205d7853a 100644
--- a/airflow-core/docs/core-concepts/xcoms.rst
+++ b/airflow-core/docs/core-concepts/xcoms.rst
@@ -106,11 +106,22 @@ The default XCom backend, BaseXCom, stores XComs in the
Airflow database, which
Custom XCom Backends
--------------------
-The XCom system has interchangeable backends, and you can set which backend is
being used via the ``xcom_backend`` configuration option.
-
-If you want to implement your own backend, you should subclass
:class:`~airflow.sdk.bases.xcom.BaseXCom`, and override the ``serialize_value``
and ``deserialize_value`` methods.
-
-You can override the ``purge`` method in the ``BaseXCom`` class to have
control over purging the XCom data from the custom backend. This will be called
as part of ``delete``.
+The XCom system has interchangeable backends, and you can set which backend is
being used via the ``xcom_backend`` configuration option. A backend is a
subclass of :class:`~airflow.sdk.bases.xcom.BaseXCom`.
+
+A backend that keeps values outside the metadata database overrides
``store_value``, which receives an
+:class:`~airflow.sdk.bases.xcom.XComCoordinate`, stores the value and returns
what the XCom row holds, and
+``load_value``, which turns that row value back into the value.
``XComCoordinate.build_path`` builds a storage
+location that is unique per task instance, or per write when the writer cannot
know the task instance.
+The key is chosen by the Dag author and can look like a file extension, so
pass a ``suffix`` that names the
+stored format if your backend reads the format back from the object name.
+
+Override ``purge`` to remove the stored data when the Task SDK deletes an
XCom, either through ``XCom.delete`` or
+when a task instance clears its XComs as it starts again. ``purge`` is not
called when a row is removed through
+the REST API, the UI or ``airflow db clean``, so those leave the stored data
behind.
+
+A backend that only overrides ``serialize_value`` and ``deserialize_value``
keeps working, but cannot be used
+with Dags that use task loops: every iteration of a task in a loop is a
separate task instance, and only
+``store_value`` can tell them apart.
Verifying Custom XCom Backend usage in Containers
-------------------------------------------------
diff --git
a/airflow-core/newsfragments/task-loops-custom-backends.significant.rst
b/airflow-core/newsfragments/task-loops-custom-backends.significant.rst
new file mode 100644
index 00000000000..23de6f1b0e5
--- /dev/null
+++ b/airflow-core/newsfragments/task-loops-custom-backends.significant.rst
@@ -0,0 +1,28 @@
+Custom XCom backends that override ``serialize_value`` must implement
``store_value`` to support task loops
+
+Every iteration of a task inside a loop is its own task instance, but
``serialize_value``
+receives the same ``dag_id``, ``run_id``, ``task_id``, ``map_index`` and
``key`` for all of
+them, so a backend that builds its storage location from those makes every
iteration's XCom
+read back as the last iteration's value. Backends should override
``store_value`` and
+``load_value`` instead, and build the storage location with
``XComCoordinate.build_path``,
+which is unique per task instance, or per write when the write cannot know the
task instance.
+
+A Dag that uses a loop fails on startup when the configured XCom backend still
stores values
+through ``serialize_value``, even if it already writes to a unique location.
This includes
+``XComObjectStorageBackend`` from ``apache-airflow-providers-common-io``
1.10.1 and earlier;
+upgrade the provider before using task loops with it.
+
+Custom Task State Store backends will likewise need to include
+``scope.region_id`` in their storage key to support mapped tasks inside a loop
+and loop iterations regenerated by clearing.
+
+* Types of change
+
+ * [ ] Dag changes
+ * [ ] Config changes
+ * [x] API changes
+ * [ ] CLI changes
+ * [x] Behaviour changes
+ * [ ] Plugin changes
+ * [ ] Dependency changes
+ * [x] Code interface changes
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 ee4a1edb0ae..93cc67fe7b9 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
@@ -62,6 +62,7 @@ from airflow.models.task_coordinates import (
TaskCoordinateResolver,
build_coordinate_filters,
enclosing_loop,
+ get_public_map_index,
)
from airflow.models.taskinstance import TaskInstance
from airflow.models.trigger import Trigger
@@ -682,7 +683,8 @@ def get_xcom(
[
(
status.HTTP_400_BAD_REQUEST,
- "The key is empty or reserved, the value is too large to map,
or is unserializable",
+ "The key is empty or reserved, the value is too large to map
or is unserializable, "
+ "or the target is in a loop the caller is not in",
),
(
status.HTTP_409_CONFLICT,
@@ -697,6 +699,7 @@ def set_xcom(
task_id: str,
key: Annotated[str, Path(min_length=1)],
session: SessionDep,
+ dag_bag: DagBagDep,
value: Annotated[
JsonValue,
Body(
@@ -765,19 +768,26 @@ def set_xcom(
# TODO: Can/should we check if a client _hasn't_ provided this for an
upstream of a mapped task? That
# means loading the serialized dag and that seems like a relatively costly
operation for minimal benefit
# (the mapped task would fail in a moment as it can't be expanded anyway.)
+ owner = _find_writer_id(
+ token,
+ dag_id=dag_id,
+ run_id=run_id,
+ task_id=task_id,
+ map_index=map_index,
+ session=session,
+ dag_bag=dag_bag,
+ )
+ if owner is None:
+ raise HTTPException(
+ status.HTTP_404_NOT_FOUND,
+ detail={"reason": "not_found", "message": f"Task Instance
{task_id!r} not found in {run_id!r}"},
+ )
try:
# We expect serialised value from the caller - sdk, do not serialise
in here
XComModel.set_for_attempt(
key=key,
value=value,
- task_instance_id=_get_writer_id(
- token.id,
- dag_id=dag_id,
- run_id=run_id,
- task_id=task_id,
- map_index=map_index,
- session=session,
- ),
+ task_instance_id=owner,
serialize=False,
dag_result=dag_result,
mapped_length=mapped_length,
@@ -801,7 +811,10 @@ def set_xcom(
"/{dag_id}/{run_id}/{task_id}/{key:path}",
responses={
**create_openapi_http_exception_doc(
- [(status.HTTP_409_CONFLICT, "Several live task instances match the
producer")]
+ [
+ (status.HTTP_400_BAD_REQUEST, "The target is in a loop the
caller is not in"),
+ (status.HTTP_409_CONFLICT, "Several live task instances match
the producer"),
+ ]
),
status.HTTP_404_NOT_FOUND: {"description": "XCom not found"},
},
@@ -809,6 +822,7 @@ def set_xcom(
)
def delete_xcom(
session: SessionDep,
+ dag_bag: DagBagDep,
dag_id: str,
run_id: str,
task_id: str,
@@ -818,12 +832,13 @@ def delete_xcom(
):
"""Delete a single XCom Value."""
owner = _find_writer_id(
- token.id,
+ token,
dag_id=dag_id,
run_id=run_id,
task_id=task_id,
map_index=map_index,
session=session,
+ dag_bag=dag_bag,
)
if owner is not None:
XComModel.delete_for_attempts(
@@ -835,58 +850,39 @@ def delete_xcom(
def _find_writer_id(
- attempt_id: UUID,
+ token: TIToken,
*,
dag_id: str,
run_id: str,
task_id: str,
map_index: int,
- session: SessionDep,
+ session: Session,
+ dag_bag: DBDagBag,
) -> UUID | None:
- """Resolve the attempt a write targets: the caller's own, or the live
attempt of another task."""
- coordinates = (
- TaskInstance.dag_id == dag_id,
- TaskInstance.run_id == run_id,
- TaskInstance.task_id == task_id,
- *build_coordinate_filters(TaskInstance, map_index=map_index),
- )
- own = session.scalar(
- select(TaskInstance.id)
- .where(TaskInstance.id == attempt_id, *coordinates)
- .execution_options(include_all_attempts=True)
- )
- if own is not None:
- return own
- live =
session.scalars(select(TaskInstance.id).where(*coordinates).limit(2)).all()
- if len(live) > 1:
+ """Resolve the task instance a write targets with the rule reads use,
working from the caller."""
+ caller = session.get(TaskInstance, token.id,
execution_options={"include_all_attempts": True})
+ # A write admitted for an attempt that has since been archived must stay
with that attempt.
+ if (
+ caller is not None
+ and (caller.dag_id, caller.run_id, caller.task_id) == (dag_id, run_id,
task_id)
+ and get_public_map_index(caller, session=session) == map_index
+ ):
+ return caller.id
+ try:
+ ids = session.scalars(
+ TaskCoordinateResolver(dag_bag, session)
+ .select_producer_ids(
+ dag_id=dag_id, run_id=run_id, task_id=task_id, caller=caller,
map_indexes=map_index
+ )
+ .limit(2)
+ ).all()
+ except AmbiguousProducerError as error:
+ raise HTTPException(status.HTTP_409_CONFLICT, str(error)) from error
+ except ValueError as error:
+ raise HTTPException(status.HTTP_400_BAD_REQUEST, str(error)) from error
+ if len(ids) > 1:
raise HTTPException(
status.HTTP_409_CONFLICT,
- f"More than one live task instance matches {dag_id}.{task_id} in
run {run_id}; "
- "a task inside a loop can only write XComs of its own attempt",
- )
- return live[0] if live else None
-
-
-def _get_writer_id(
- attempt_id: UUID,
- *,
- dag_id: str,
- run_id: str,
- task_id: str,
- map_index: int,
- session: SessionDep,
-) -> UUID:
- owner = _find_writer_id(
- attempt_id,
- dag_id=dag_id,
- run_id=run_id,
- task_id=task_id,
- map_index=map_index,
- session=session,
- )
- if owner is None:
- raise HTTPException(
- status.HTTP_404_NOT_FOUND,
- detail={"reason": "not_found", "message": f"Task Instance
{task_id!r} not found in {run_id!r}"},
+ f"More than one live task instance matches {dag_id}.{task_id} in
run {run_id}",
)
- return owner
+ return ids[0] if ids else None
diff --git a/airflow-core/src/airflow/dag_processing/processor.py
b/airflow-core/src/airflow/dag_processing/processor.py
index 9530d7be8c9..4be5623d71d 100644
--- a/airflow-core/src/airflow/dag_processing/processor.py
+++ b/airflow-core/src/airflow/dag_processing/processor.py
@@ -73,7 +73,11 @@ from airflow.sdk.execution_time.comms import (
XComSequenceSliceResult,
)
from airflow.sdk.execution_time.supervisor import WatchedSubprocess,
register_request_method
-from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance,
_send_error_email_notification
+from airflow.sdk.execution_time.task_runner import (
+ RuntimeTaskInstance,
+ _send_error_email_notification,
+ check_xcom_backend_supports_loops,
+)
from airflow.sdk.importers import DagSourceCode # noqa: TC001
from airflow.serialization.serialized_objects import DagSerialization,
LazyDeserializedDAG
from airflow.utils.dag_version_inflation_checker import
check_dag_file_stability
@@ -283,6 +287,9 @@ def _serialize_dags(
serialized_dags = []
for dag in bag.dags.values():
try:
+ # Only possible because importing the SDK resolves the XCom
backend here; the worker check is the
+ # one that must stay, so drop this one rather than keep that
import working for its sake.
+ check_xcom_backend_supports_loops(dag)
data = DagSerialization.to_dict(dag)
serialized_dags.append(LazyDeserializedDAG(data=data,
last_loaded=dag.last_loaded))
except Exception:
diff --git a/airflow-core/src/airflow/jobs/triggerer_job_runner.py
b/airflow-core/src/airflow/jobs/triggerer_job_runner.py
index b29cf66a42e..326eae953e4 100644
--- a/airflow-core/src/airflow/jobs/triggerer_job_runner.py
+++ b/airflow-core/src/airflow/jobs/triggerer_job_runner.py
@@ -63,6 +63,7 @@ from airflow.models.task_coordinates import
TaskCoordinateResolver
from airflow.models.trigger import Trigger
from airflow.observability.metrics import stats_utils
from airflow.sdk.api.datamodels._generated import HITLDetailResponse
+from airflow.sdk.bases.xcom import _current_trigger_ti
from airflow.sdk.definitions.asset import Asset
from airflow.sdk.execution_time import supervisor
from airflow.sdk.execution_time.comms import (
@@ -164,12 +165,6 @@ _REASSIGNED_CANCEL_MSG = "__airflow_reassigned__"
_ON_CANCEL_TIMEOUT: int = conf.getint("triggerer", "on_kill_timeout",
fallback=30)
-# Task instance of the trigger whose code is running, so its SDK requests act
as that task instance; tasks
-# serving several triggers must not inherit it.
-_current_trigger_ti_id: contextvars.ContextVar[UUID | None] =
contextvars.ContextVar(
- "airflow_current_trigger_ti_id", default=None
-)
-
# Task instance of the trigger whose request the supervisor is handling, sent
to the API as
# ``TRIGGER_TI_ID_HEADER``.
_handling_trigger_ti_id: contextvars.ContextVar[UUID | None] =
contextvars.ContextVar(
@@ -1261,8 +1256,11 @@ class TriggerCommsDecoder(CommsDecoder[ToTriggerRunner,
ToTriggerSupervisor]):
# non-blocking for the foreign loop.
cf = asyncio.run_coroutine_threadsafe(self.asend(msg), self._loop)
return await asyncio.wrap_future(cf)
+ trigger_ti = _current_trigger_ti.get()
frame = _RequestFrame(
- id=next(self.id_counter), body=msg.model_dump(),
ti_id=_current_trigger_ti_id.get()
+ id=next(self.id_counter),
+ body=msg.model_dump(),
+ ti_id=trigger_ti.id if trigger_ti else None,
)
future: asyncio.Future = current_loop.create_future()
self._pending[frame.id] = future
@@ -1738,7 +1736,7 @@ class TriggerRunner:
ti = trigger.task_instance
if ti is not None:
- _current_trigger_ti_id.set(ti.id)
+ _current_trigger_ti.set(ti)
bind_log_contextvars(
trigger_id=trigger_id,
ti_id=str(ti.id) if ti else None,
@@ -1783,7 +1781,7 @@ class TriggerRunner:
# subscribe() may start the group's poll task, which
copies the current context but
# serves several triggers.
subscribe_ctx = contextvars.copy_context()
- subscribe_ctx.run(_current_trigger_ti_id.set, None)
+ subscribe_ctx.run(_current_trigger_ti.set, None)
shared_stream = subscribe_ctx.run(
self._shared_streams.subscribe,
trigger_id=trigger_id,
diff --git a/airflow-core/src/airflow/models/task_coordinates.py
b/airflow-core/src/airflow/models/task_coordinates.py
index 3d97ef3eade..73c2f87ed8d 100644
--- a/airflow-core/src/airflow/models/task_coordinates.py
+++ b/airflow-core/src/airflow/models/task_coordinates.py
@@ -114,7 +114,9 @@ def get_public_map_index(task_instance: TaskInstance, *,
session: Session) -> in
if task_instance.region_id == SENTINEL_REGION_ID:
return task_instance.region_index
return session.execute(
-
select(public_map_index_expression(TaskInstance)).where(TaskInstance.id ==
task_instance.id)
+ select(public_map_index_expression(TaskInstance))
+ .where(TaskInstance.id == task_instance.id)
+ .execution_options(include_all_attempts=True)
).scalar_one()
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 560cbe0f462..7768fe75336 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
@@ -38,7 +38,7 @@ from airflow.api_fastapi.execution_api.security import
_jwt_bearer, require_auth
from airflow.exceptions import TaskNotFound
from airflow.models.dag import DagModel
from airflow.models.dagrun import DagRun
-from airflow.models.dynamic_region import SENTINEL_REGION_ID, DynamicRegion
+from airflow.models.dynamic_region import SENTINEL_REGION_ID,
AmbiguousProducerError, DynamicRegion
from airflow.models.task_coordinates import TaskCoordinateResolver
from airflow.models.taskinstance import LegacyTaskDataOwner, TaskInstance
from airflow.models.trigger import Trigger
@@ -782,12 +782,102 @@ def
test_loop_prior_dates_for_unknown_target_run_falls_through_to_plain_read(cli
assert response.status_code == 404
-def test_loop_xcom_write_and_delete_for_several_live_passes_conflict(client,
loop_xcoms):
- dr, producer, _, _ = loop_xcoms
[email protected]
+def single_pass_loop(dag_maker, session, authenticate_as):
+ @task_group
+ def body():
+ EmptyOperator(task_id="step")
+
+ with dag_maker(serialized=True):
+ create_loop(body, max_iterations=3)
+ EmptyOperator(task_id="outside")
+ PythonOperator.partial(task_id="mapped",
python_callable=str).expand(op_args=[[1], [2]])
+ dr = dag_maker.create_dagrun()
+ session.commit()
+ tis = {(ti.task_id, ti.map_index): ti for ti in dr.task_instances}
+ return dr, tis, authenticate_as
+
+
[email protected]("method", ["post", "delete"])
+def test_xcom_change_to_a_loop_task_from_outside_the_loop_is_rejected(client,
single_pass_loop, method):
+ dr, tis, authenticate_as = single_pass_loop
+ authenticate_as(tis["outside", -1])
+ url = get_loop_url(dr, "body.step")
+
+ response = client.post(url, json="value") if method == "post" else
client.delete(url)
+
+ assert response.status_code == 400
+ assert "consumer inside the loop" in response.json()["detail"]
+
+
[email protected](
+ ("caller", "target", "params"),
+ [
+ pytest.param(("body.step", 0), "body.step", {},
id="own-xcom-inside-loop"),
+ pytest.param(("body.step", 0), "outside", {},
id="from-inside-loop-to-outside"),
+ pytest.param(("outside", -1), "mapped", {"map_index": 1},
id="mapped-task-outside-loops"),
+ ],
+)
+def test_xcom_writes_outside_another_tasks_loop_are_allowed(client,
single_pass_loop, caller, target, params):
+ dr, tis, authenticate_as = single_pass_loop
+ authenticate_as(tis[caller])
+
+ response = client.post(get_loop_url(dr, target), json="value",
params=params)
+
+ assert response.status_code == 201
+
+
[email protected]("method", ["post", "delete"])
[email protected]("several", ["ambiguous-producer",
"several-task-instances"])
+def test_write_and_delete_conflict_when_several_task_instances_match(
+ client, single_pass_loop, mocker, method, several
+):
+ dr, tis, authenticate_as = single_pass_loop
+ authenticate_as(tis["outside", -1])
+ select_producer_ids = mocker.patch.object(TaskCoordinateResolver,
"select_producer_ids", autospec=True)
+ if several == "ambiguous-producer":
+ select_producer_ids.side_effect = AmbiguousProducerError("several live
gates")
+ else:
+ select_producer_ids.return_value = select(TaskInstance.id).where(
+ TaskInstance.dag_id == dr.dag_id, TaskInstance.run_id == dr.run_id
+ )
+ url = get_loop_url(dr, "mapped")
+
+ response = (
+ client.post(url, json="value", params={"map_index": 1})
+ if method == "post"
+ else client.delete(url, params={"map_index": 1})
+ )
+
+ assert response.status_code == 409
+
+
+def test_archived_attempt_in_a_loop_keeps_writing_its_own_xcom(client,
single_pass_loop, session):
+ dr, tis, authenticate_as = single_pass_loop
+ step = tis["body.step", 0]
+ authenticate_as(step)
+ step.archive(reason="test", session=session)
+ session.commit()
+
+ response = client.post(get_loop_url(dr, "body.step"), json="late")
+
+ assert response.status_code == 201
+ assert XComModelV2.get_for_attempt(step.id, "key", session=session).value
== "late"
+
+
+def test_loop_xcom_write_and_delete_resolve_to_the_callers_iteration(client,
loop_xcoms, session):
+ dr, producer, _, previous = loop_xcoms
url = f"/execution/xcoms/{dr.dag_id}/{dr.run_id}/{producer.task_id}/key"
- assert client.post(url, json="value").status_code == 409
- assert client.delete(url).status_code == 409
+ assert client.post(url, json="rewritten").status_code == 201
+ session.expire_all()
+ assert XComModelV2.get_for_attempt(producer.id, "key",
session=session).value == "rewritten"
+ assert XComModelV2.get_for_attempt(previous.id, "key",
session=session).value == "previous"
+
+ assert client.delete(url).status_code == 200
+ session.expire_all()
+ assert XComModelV2.get_for_attempt(producer.id, "key", session=session) is
None
+ assert XComModelV2.get_for_attempt(previous.id, "key",
session=session).value == "previous"
@pytest.mark.parametrize("current_regional", [False, True])
diff --git a/airflow-core/tests/unit/dag_processing/test_processor.py
b/airflow-core/tests/unit/dag_processing/test_processor.py
index 629cb6b1c49..0271a5e07fb 100644
--- a/airflow-core/tests/unit/dag_processing/test_processor.py
+++ b/airflow-core/tests/unit/dag_processing/test_processor.py
@@ -67,12 +67,14 @@ from airflow.dag_processing.processor import (
_parse_file,
_parse_file_entrypoint,
_pre_import_airflow_modules,
+ _serialize_dags,
)
from airflow.models import DagRun
from airflow.models.dagwarning import DagWarning
-from airflow.sdk import DAG, BaseOperator
+from airflow.sdk import DAG, BaseOperator, task_group
from airflow.sdk.api.client import Client
from airflow.sdk.api.datamodels._generated import ConnectionResponse,
DagRunState, VariableResponse
+from airflow.sdk.definitions._internal.loop import create_loop
from airflow.sdk.execution_time import comms, supervisor
from airflow.sdk.execution_time.comms import (
GetConnection,
@@ -95,6 +97,7 @@ from airflow.utils.session import create_session
from airflow.utils.state import TaskInstanceState
from tests_common.test_utils.config import conf_vars, env_vars
+from tests_common.test_utils.xcom_backends import SerializeValueOnlyXCom
if TYPE_CHECKING:
from kgb import SpyAgency
@@ -567,6 +570,24 @@ def disable_capturing():
sys.stdin, sys.stdout, sys.stderr = old_in, old_out, old_err
+def
test_serialize_dags_reports_loop_dag_with_legacy_xcom_backend_as_import_error(monkeypatch):
+ @task_group
+ def body():
+ BaseOperator(task_id="work")
+
+ with DAG("loop_dag", schedule=None) as dag:
+ create_loop(body, max_iterations=2)
+ dag.relative_fileloc = "loop_dag.py"
+ bag = MagicMock(spec=DagBag)
+ bag.dags = {dag.dag_id: dag}
+ monkeypatch.setattr("airflow.sdk.execution_time.task_runner.XCom",
SerializeValueOnlyXCom)
+
+ serialized_dags, import_errors = _serialize_dags(bag,
structlog.get_logger())
+
+ assert serialized_dags == []
+ assert "uses a loop" in import_errors["loop_dag.py"]
+
+
@pytest.mark.usefixtures("testing_dag_bundle")
@pytest.mark.usefixtures("disable_capturing")
def test_parse_file_entrypoint_parses_dag_callbacks(mocker):
diff --git a/airflow-core/tests/unit/jobs/test_triggerer_job.py
b/airflow-core/tests/unit/jobs/test_triggerer_job.py
index ab6a9e6f4f1..0573e0d3d6d 100644
--- a/airflow-core/tests/unit/jobs/test_triggerer_job.py
+++ b/airflow-core/tests/unit/jobs/test_triggerer_job.py
@@ -67,7 +67,6 @@ from airflow.jobs.triggerer_job_runner import (
TriggerLoggingFactory,
TriggerRunner,
TriggerRunnerSupervisor,
- _current_trigger_ti_id,
_make_trigger_span,
messages,
)
@@ -86,6 +85,7 @@ from airflow.providers.standard.triggers.temporal import
DateTimeTrigger, TimeDe
from airflow.sdk import DAG, Asset, BaseHook, BaseOperator, task_group
from airflow.sdk.api.client import Client
from airflow.sdk.api.datamodels._generated import AssetStateStoreResponse
+from airflow.sdk.bases.xcom import _current_trigger_ti
from airflow.sdk.definitions._internal.loop import create_loop
from airflow.sdk.exceptions import ErrorType
from airflow.sdk.execution_time import supervisor
@@ -3534,15 +3534,15 @@ async def
test_asend_frame_carries_the_current_trigger_ti_id(decoder_pair):
thread = threading.Thread(target=serve, daemon=True)
thread.start()
- ti_id = uuid.uuid4()
- token = _current_trigger_ti_id.set(ti_id)
+ ti = _make_ti_dto()
+ token = _current_trigger_ti.set(ti)
try:
await decoder.asend(messages.TriggerStateChanges(events=None,
finished=[1], failures=None))
finally:
- _current_trigger_ti_id.reset(token)
+ _current_trigger_ti.reset(token)
thread.join(timeout=5)
- assert [frame.ti_id for frame in frames] == [ti_id]
+ assert [frame.ti_id for frame in frames] == [ti.id]
def test_trigger_state_messages_round_trip_with_seqs():
@@ -3754,16 +3754,17 @@ def
test_run_trigger_appends_none_seq_for_non_shared_trigger():
@pytest.mark.asyncio
@pytest.mark.parametrize("has_task_instance", [True, False])
-async def
test_run_trigger_sets_current_trigger_ti_id_only_for_task_instance_triggers(has_task_instance):
- ti = _make_ti_dto()
- seen: list[uuid.UUID | None] = []
+async def
test_run_trigger_sets_current_trigger_ti_only_for_task_instance_triggers(has_task_instance):
+ ti = _make_ti_dto(map_index=3)
+ ti.region_id, ti.region_index = uuid.uuid4(), 2
+ seen: list[TaskInstanceDTO | None] = []
class _RecordingTrigger(BaseTrigger):
def serialize(self):
return (f"{type(self).__module__}.{type(self).__qualname__}", {})
async def run(self):
- seen.append(_current_trigger_ti_id.get())
+ seen.append(_current_trigger_ti.get())
yield TriggerEvent(True)
trigger_runner = TriggerRunner()
@@ -3776,13 +3777,13 @@ async def
test_run_trigger_sets_current_trigger_ti_id_only_for_task_instance_tri
await asyncio.create_task(trigger_runner.run_trigger(5, trigger))
- assert seen == [ti.id if has_task_instance else None]
+ assert seen == [ti if has_task_instance else None]
@pytest.mark.asyncio
-async def
test_shared_stream_group_tasks_do_not_inherit_the_subscribing_trigger_ti_id():
- ti = _make_ti_dto()
- seen: dict[str, uuid.UUID | None] = {}
+async def
test_shared_stream_group_tasks_do_not_inherit_the_subscribing_trigger_ti():
+ ti = _make_ti_dto(map_index=3)
+ seen: dict[str, TaskInstanceDTO | None] = {}
class _SharedTrigger(BaseEventTrigger):
def serialize(self):
@@ -3793,12 +3794,12 @@ async def
test_shared_stream_group_tasks_do_not_inherit_the_subscribing_trigger_
@classmethod
async def open_shared_stream(cls, kwargs):
- seen["group"] = _current_trigger_ti_id.get()
+ seen["group"] = _current_trigger_ti.get()
yield {}
await asyncio.Event().wait()
async def filter_shared_stream(self, shared_stream):
- seen["trigger"] = _current_trigger_ti_id.get()
+ seen["trigger"] = _current_trigger_ti.get()
async for raw in shared_stream:
yield TriggerEvent(raw)
@@ -3821,7 +3822,7 @@ async def
test_shared_stream_group_tasks_do_not_inherit_the_subscribing_trigger_
with contextlib.suppress(asyncio.CancelledError):
await run_task
- assert seen == {"group": None, "trigger": ti.id}
+ assert seen == {"group": None, "trigger": ti}
@pytest.mark.asyncio
diff --git a/devel-common/src/tests_common/pytest_plugin.py
b/devel-common/src/tests_common/pytest_plugin.py
index c82b2644551..da328441d4c 100644
--- a/devel-common/src/tests_common/pytest_plugin.py
+++ b/devel-common/src/tests_common/pytest_plugin.py
@@ -2986,7 +2986,7 @@ def create_runtime_ti(mocked_parse):
@pytest.fixture
-def run_task(create_runtime_ti, mock_supervisor_comms, spy_agency) ->
RunTaskCallable:
+def run_task(create_runtime_ti, mock_supervisor_comms, spy_agency) ->
Generator[RunTaskCallable, None, None]:
"""
Fixture to run a task without defining a dag file.
@@ -3018,35 +3018,61 @@ def run_task(create_runtime_ti, mock_supervisor_comms,
spy_agency) -> RunTaskCal
timezone = _import_timezone()
+ # kgb leaves a classmethod it unspied bound to the class it spied, which
breaks subclasses in later tests.
+ # A task's own pushes go through _push/_apush on newer Task SDK versions
and through set/aset on older ones.
+ push_methods = [name for name in ("set", "_push", "_apush") if
hasattr(XCom, name)]
+ spied_methods = [*push_methods, "get_one"]
+ classmethods = {
+ name: next((owner, vars(owner)[name]) for owner in XCom.__mro__ if
name in vars(owner))
+ for name in spied_methods
+ }
+
# Set up spies once at fixture level
- if hasattr(XCom.set, "spy"):
- spy_agency.unspy(XCom.set)
- if hasattr(XCom.get_one, "spy"):
- spy_agency.unspy(XCom.get_one)
- spy_agency.spy_on(XCom.set, call_original=True)
+ for name in spied_methods:
+ if hasattr(getattr(XCom, name), "spy"):
+ spy_agency.unspy(getattr(XCom, name))
+ for name in push_methods:
+ spy_agency.spy_on(getattr(XCom, name), call_original=True)
spy_agency.spy_on(
XCom.get_one, call_fake=lambda cls, *args, **kwargs:
_get_one_from_set_calls(*args, **kwargs)
)
+ def _recorded_pushes() -> list[tuple[str, Any, dict[str, Any]]]:
+ """List the pushed ``(key, value, coordinates and options)`` of every
spied push method."""
+ pushes = []
+ for name in push_methods:
+ for call in getattr(XCom, name).calls:
+ if name == "set":
+ if len(call.args) >= 2:
+ pushes.append((call.args[0], call.args[1],
dict(call.kwargs)))
+ else:
+ ti, key, value = call.args[:3]
+ coordinates = {
+ "task_id": ti.task_id,
+ "dag_id": ti.dag_id,
+ "run_id": ti.run_id,
+ "map_index": ti.map_index,
+ }
+ pushes.append((key, value, {**coordinates, **call.kwargs}))
+ return pushes
+
def _get_one_from_set_calls(*args, **kwargs) -> Any | None:
- """Get the most recent value from XCom.set calls that matches the
criteria."""
+ """Get the last pushed value that matches the criteria, per push
method."""
key = kwargs.get("key")
task_id = kwargs.get("task_id")
dag_id = kwargs.get("dag_id")
run_id = kwargs.get("run_id")
map_index = kwargs.get("map_index") or -1
- for call in reversed(XCom.set.calls):
+ for call_key, value, call_kwargs in reversed(_recorded_pushes()):
if (
- call.kwargs.get("task_id") == task_id
- and call.kwargs.get("dag_id") == dag_id
- and call.kwargs.get("run_id") == run_id
- and call.kwargs.get("map_index") == map_index
+ call_kwargs.get("task_id") == task_id
+ and call_kwargs.get("dag_id") == dag_id
+ and call_kwargs.get("run_id") == run_id
+ and call_kwargs.get("map_index") == map_index
+ and call_key == key
):
- if call.args and len(call.args) >= 2:
- call_key, value = call.args
- if call_key == key:
- return value
+ return value
return None
class XComHelper:
@@ -3091,23 +3117,30 @@ def run_task(create_runtime_ti, mock_supervisor_comms,
spy_agency) -> RunTaskCal
run_id = run_id or self._ti.run_id
map_index = map_index if map_index is not None else
self._ti.map_index
- spy_agency.assert_spy_called_with(
- XCom.set,
- key,
- value,
- task_id=task_id,
- dag_id=dag_id,
- run_id=run_id,
- map_index=map_index,
+ expected = {
+ "task_id": task_id,
+ "dag_id": dag_id,
+ "run_id": run_id,
+ "map_index": map_index,
**kwargs,
+ }
+ pushes = _recorded_pushes()
+ for call_key, call_value, call_kwargs in pushes:
+ if (
+ call_key == key
+ and call_value == value
+ and all(call_kwargs.get(name) == wanted for name, wanted
in expected.items())
+ ):
+ return
+ raise AssertionError(
+ f"No XCom push matches key={key!r} value={value!r} {expected};
recorded pushes: {pushes}"
)
def clear(self):
"""Clear all XCom calls."""
- if hasattr(XCom.set, "spy"):
- spy_agency.unspy(XCom.set)
- if hasattr(XCom.get_one, "spy"):
- spy_agency.unspy(XCom.get_one)
+ for name in spied_methods:
+ if hasattr(getattr(XCom, name), "spy"):
+ spy_agency.unspy(getattr(XCom, name))
class RunTaskWithXCom(RunTaskCallable):
def __init__(self, create_runtime_ti):
@@ -3200,7 +3233,11 @@ def run_task(create_runtime_ti, mock_supervisor_comms,
spy_agency) -> RunTaskCal
return state, msg, error
- return RunTaskWithXCom(create_runtime_ti)
+ yield RunTaskWithXCom(create_runtime_ti)
+
+ spy_agency.unspy_all()
+ for name, (owner, original) in classmethods.items():
+ setattr(owner, name, original)
@pytest.fixture
diff --git a/devel-common/src/tests_common/test_utils/xcom_backends.py
b/devel-common/src/tests_common/test_utils/xcom_backends.py
new file mode 100644
index 00000000000..bc5a280b12c
--- /dev/null
+++ b/devel-common/src/tests_common/test_utils/xcom_backends.py
@@ -0,0 +1,28 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from airflow.sdk.bases.xcom import BaseXCom
+
+
+class SerializeValueOnlyXCom(BaseXCom):
+ """A legacy backend: it stores values through ``serialize_value`` and has
no ``store_value``."""
+
+ @staticmethod
+ def serialize_value(value, *, key=None, task_id=None, dag_id=None,
run_id=None, map_index=None):
+ return value
diff --git a/generated/known_sdk_imports_in_core.txt
b/generated/known_sdk_imports_in_core.txt
index 91d224dbcce..01a4bee78a5 100644
--- a/generated/known_sdk_imports_in_core.txt
+++ b/generated/known_sdk_imports_in_core.txt
@@ -11,7 +11,7 @@ airflow-core/src/airflow/dag_processing/manager.py::5
airflow-core/src/airflow/dag_processing/processor.py::14
airflow-core/src/airflow/exceptions.py::1
airflow-core/src/airflow/executors/base_executor.py::3
-airflow-core/src/airflow/jobs/triggerer_job_runner.py::18
+airflow-core/src/airflow/jobs/triggerer_job_runner.py::19
airflow-core/src/airflow/models/__init__.py::5
airflow-core/src/airflow/models/asset.py::2
airflow-core/src/airflow/models/connection.py::5
diff --git a/providers/common/io/docs/changelog.rst
b/providers/common/io/docs/changelog.rst
index 56b05e3a9cd..d113222d853 100644
--- a/providers/common/io/docs/changelog.rst
+++ b/providers/common/io/docs/changelog.rst
@@ -25,6 +25,16 @@
Changelog
---------
+.. note::
+ On Airflow 3.4 and later, ``XComObjectStorageBackend`` implements
``store_value``, which Dags that use
+ task loops require, and writes objects to
+ ``<dag_id>/<run_id>/<task_id>/ti-<task instance id>/<quoted key>.json``
instead of a random file name, or to
+ ``.../write-<uuid>/<quoted key>.json`` for an XCom written for another
task. ``dag_id``, ``run_id`` and
+ ``task_id`` are shortened only when the path would exceed 1,024 bytes. The
quoted key is replaced by its
+ SHA-256 digest when it is longer than 200 characters, and a compression
suffix is appended when
+ compression is configured. A subclass that overrides ``serialize_value``
keeps it, but cannot be used
+ with task loops until it implements ``store_value`` as well.
+
1.10.1
......
diff --git
a/providers/common/io/src/airflow/providers/common/io/xcom/backend.py
b/providers/common/io/src/airflow/providers/common/io/xcom/backend.py
index 43fbfc5b119..3328e341d8b 100644
--- a/providers/common/io/src/airflow/providers/common/io/xcom/backend.py
+++ b/providers/common/io/src/airflow/providers/common/io/xcom/backend.py
@@ -19,6 +19,7 @@ from __future__ import annotations
import contextlib
import json
import uuid
+from collections.abc import Callable
from functools import cache
from typing import TYPE_CHECKING, Any, TypeVar
from urllib.parse import urlsplit
@@ -32,6 +33,7 @@ from airflow.utils.json import XComDecoder, XComEncoder
if TYPE_CHECKING:
from sqlalchemy.orm import Session
+ from airflow.sdk.bases.xcom import XComCoordinate
from airflow.sdk.execution_time.comms import XComResult
if AIRFLOW_V_3_0_PLUS:
@@ -45,6 +47,9 @@ T = TypeVar("T")
SECTION = "common.io"
+# S3 limits an object key to 1,024 bytes.
+_MAX_PATH_LENGTH = 1024
+
def _get_compression_suffix(compression: str) -> str:
"""
@@ -107,6 +112,22 @@ class XComObjectStorageBackend(BaseXCom):
raise ValueError(f"Not a valid url: {data}")
+ @classmethod
+ def store_value(cls, coordinate: XComCoordinate, key: str, value: Any) ->
str:
+ base = _get_base_path()
+ path = cls._write_large_value(
+ value,
+ lambda suffix: base.joinpath(
+ coordinate.build_path(
+ "",
+ key,
+ suffix=f".json{suffix}",
+ max_length=_MAX_PATH_LENGTH - len(str(base).encode()) - 1,
+ )
+ ),
+ )
+ return BaseXCom.serialize_value(value if path is None else str(path))
+
@staticmethod
def serialize_value( # type: ignore[override]
value: T,
@@ -117,38 +138,46 @@ class XComObjectStorageBackend(BaseXCom):
run_id: str | None = None,
map_index: int | None = None,
) -> bytes | str:
- # We will use this serialized value to write to the object store.
- s_val = json.dumps(value, cls=XComEncoder)
- s_val_encoded = s_val.encode("utf-8")
+ def get_unused_path(suffix: str) -> ObjectStoragePath:
+ base_path = _get_base_path()
+ while True: # Safeguard against collisions.
+ p = base_path.joinpath(
+ dag_id or "NO_DAG_ID",
+ run_id or "NO_RUN_ID",
+ task_id or "NO_TASK_ID",
+ f"{uuid.uuid4()}{suffix}",
+ )
+ if not p.exists():
+ return p
+
+ if (path := XComObjectStorageBackend._write_large_value(value,
get_unused_path)) is not None:
+ return BaseXCom.serialize_value(str(path))
+ if AIRFLOW_V_3_0_PLUS:
+ return BaseXCom.serialize_value(value)
+ # TODO: Remove this branch once we drop support for Airflow 2
+ # This is for Airflow 2.10 where the value is expected to be bytes
+ return json.dumps(value, cls=XComEncoder).encode("utf-8")
+
+ @staticmethod
+ def _write_large_value(
+ value: Any, get_path: Callable[[str], ObjectStoragePath]
+ ) -> ObjectStoragePath | None:
+ """Write ``value`` to object storage if it reaches the threshold,
returning where it was written."""
+ s_val_encoded = json.dumps(value, cls=XComEncoder).encode("utf-8")
+ threshold = _get_threshold()
+ if threshold < 0 or len(s_val_encoded) < threshold: # Either no
threshold or value is small enough.
+ return None
if compression := _get_compression():
suffix = f".{_get_compression_suffix(compression)}"
else:
suffix = ""
-
- threshold = _get_threshold()
- if threshold < 0 or len(s_val_encoded) < threshold: # Either no
threshold or value is small enough.
- if AIRFLOW_V_3_0_PLUS:
- return BaseXCom.serialize_value(value)
- # TODO: Remove this branch once we drop support for Airflow 2
- # This is for Airflow 2.10 where the value is expected to be bytes
- return s_val_encoded
-
- base_path = _get_base_path()
- while True: # Safeguard against collisions.
- p = base_path.joinpath(
- dag_id or "NO_DAG_ID",
- run_id or "NO_RUN_ID",
- task_id or "NO_TASK_ID",
- f"{uuid.uuid4()}{suffix}",
- )
- if not p.exists():
- break
+ p = get_path(suffix)
p.parent.mkdir(parents=True, exist_ok=True)
with p.open(mode="wb", compression=compression) as f:
f.write(s_val_encoded)
- return BaseXCom.serialize_value(str(p))
+ return p
@staticmethod
def deserialize_value(result) -> Any:
diff --git a/providers/common/io/tests/unit/common/io/xcom/test_backend.py
b/providers/common/io/tests/unit/common/io/xcom/test_backend.py
index fe4fe5a7e03..9fd26f2187c 100644
--- a/providers/common/io/tests/unit/common/io/xcom/test_backend.py
+++ b/providers/common/io/tests/unit/common/io/xcom/test_backend.py
@@ -17,11 +17,13 @@
# under the License.
from __future__ import annotations
+import uuid
from unittest.mock import MagicMock, patch
import pytest
import airflow.models.xcom
+from airflow.providers.common.io.xcom import backend as xcom_backend
from airflow.providers.common.io.xcom.backend import XComObjectStorageBackend
from airflow.providers.standard.operators.empty import EmptyOperator
@@ -39,6 +41,8 @@ pytestmark = [pytest.mark.db_test]
if AIRFLOW_V_3_1_PLUS:
from airflow.models.xcom import XComModel
+if AIRFLOW_V_3_4_PLUS:
+ from airflow.sdk.bases.xcom import XComCoordinate
if AIRFLOW_V_3_0_PLUS:
from airflow.models.xcom import XComModel
from airflow.sdk import ObjectStoragePath
@@ -391,6 +395,57 @@ class TestXComObjectStorageBackend:
assert p.exists() is False
assert not value
+ @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="store_value is called
from Airflow 3.4")
+ @conf_vars({("common.io", "xcom_objectstorage_compression"): ""})
+ @pytest.mark.parametrize("key", ["a/key", "k" * 300, ".", "report.gz",
"data.zip"])
+ @pytest.mark.parametrize("ti_id", [uuid.uuid4(), None], ids=["own-push",
"another-tasks-xcom"])
+ def
test_store_value_writes_a_path_for_the_task_instance_or_for_the_write(self,
ti_id, key):
+ value = {"key": "bigvaluebigvaluebigvalue" * 100}
+ coordinate = XComCoordinate("dag", "run", "task", ti_id)
+
+ stored = XComObjectStorageBackend.store_value(coordinate, key, value)
+ again = XComObjectStorageBackend.store_value(coordinate, key, value)
+
+ assert stored.startswith(f"{self.path}/dag/run/task/{'ti' if ti_id
else 'write'}-")
+ assert (stored == again) == (ti_id is not None)
+ assert XComObjectStorageBackend.load_value(stored) == value
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="store_value is called
from Airflow 3.4")
+ @conf_vars({("common.io", "xcom_objectstorage_compression"): ""})
+ def test_store_value_keeps_a_small_value_in_the_database(self, tmp_path):
+ coordinate = XComCoordinate("dag", "run", "task", uuid.uuid4())
+
+ stored = XComObjectStorageBackend.store_value(coordinate, "key",
"small")
+
+ assert stored == "small"
+ assert XComObjectStorageBackend.load_value(stored) == "small"
+ assert not list((tmp_path / "xcom").iterdir())
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="store_value is called
from Airflow 3.4")
+ @conf_vars({("common.io", "xcom_objectstorage_compression"): "gzip"})
+ def test_store_value_appends_the_compression_suffix(self):
+ value = {"key": "bigvaluebigvaluebigvalue" * 100}
+ coordinate = XComCoordinate("dag", "run", "task", uuid.uuid4())
+
+ stored = XComObjectStorageBackend.store_value(coordinate, "key", value)
+
+ assert stored.endswith("/key.json.gz")
+ assert XComObjectStorageBackend.load_value(stored) == value
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="store_value is called
from Airflow 3.4")
+ @conf_vars({("common.io", "xcom_objectstorage_compression"): ""})
+ def test_store_value_shortens_the_path_to_fit_the_key_limit(self,
monkeypatch):
+ value = {"key": "bigvaluebigvaluebigvalue" * 100}
+ base = str(xcom_backend._get_base_path())
+ monkeypatch.setattr(xcom_backend, "_MAX_PATH_LENGTH",
len(base.encode()) + 1 + 200)
+ coordinate = XComCoordinate("d" * 250, "r" * 250, "t" * 250,
uuid.uuid4())
+
+ stored = XComObjectStorageBackend.store_value(coordinate, "key", value)
+
+ assert len(stored.removeprefix(f"{base}/").encode()) <= 200
+ assert stored.endswith(f"/ti-{coordinate.ti_id}/key.json")
+ assert XComObjectStorageBackend.load_value(stored) == value
+
@conf_vars({("common.io", "xcom_objectstorage_compression"): "gzip"})
def test_compression(self, task_instance, session, mock_supervisor_comms):
session.add(task_instance)
diff --git a/task-sdk/docs/api.rst b/task-sdk/docs/api.rst
index e6750cdf4bc..824fde2638e 100644
--- a/task-sdk/docs/api.rst
+++ b/task-sdk/docs/api.rst
@@ -94,6 +94,8 @@ Bases
.. autoapiclass:: airflow.sdk.BaseXCom
+.. autoapiclass:: airflow.sdk.bases.xcom.XComCoordinate
+
.. autoapiclass:: airflow.sdk.BranchMixIn
.. autoapiclass:: airflow.sdk.PokeReturnValue
diff --git a/task-sdk/src/airflow/sdk/bases/xcom.py
b/task-sdk/src/airflow/sdk/bases/xcom.py
index 1cbe0882845..f1f6e949373 100644
--- a/task-sdk/src/airflow/sdk/bases/xcom.py
+++ b/task-sdk/src/airflow/sdk/bases/xcom.py
@@ -17,9 +17,15 @@
from __future__ import annotations
+import asyncio
import collections
-from typing import Any, Protocol
+import contextvars
+import hashlib
+from typing import TYPE_CHECKING, Any, Protocol
+from urllib.parse import quote
+from uuid import UUID, uuid4
+import attrs
import structlog
from airflow.sdk.execution_time.comms import (
@@ -31,12 +37,112 @@ from airflow.sdk.execution_time.comms import (
XComSequenceSliceResult,
)
+if TYPE_CHECKING:
+ from pydantic import JsonValue
+
+ from airflow.sdk.types import TaskInstance
+
# Lightweight wrapper for XCom values
_XComValueWrapper = collections.namedtuple("_XComValueWrapper", "value")
log = structlog.get_logger(logger_name="task")
[email protected]
+class XComCoordinate:
+ """
+ The task instance an XCom is written for.
+
+ This is the type passed to :meth:`BaseXCom.store_value`, used when
implementing a custom XCom backend.
+
+ ``dag_id``, ``run_id`` and ``task_id`` do not identify a task instance:
every iteration of a task inside a
+ loop has the same three, and so does every map index of a mapped task.
``ti_id`` does. It is set when the
+ write is made by the task instance itself, or by its trigger while it is
deferred, and ``None`` when one
+ task instance writes another's XCom through coordinates, since only the
server can tell which task
+ instance those coordinates resolve to.
+
+ ``region_id`` and ``region_index`` are the stored position of the task
instance: ``region_index`` is its
+ place within the region ``region_id``, which is the map index of a mapped
task and the loop iteration of
+ a task directly inside a loop. They are known only together with
``ti_id``, and are ``None`` otherwise and
+ for a task instance outside any region.
+ """
+
+ dag_id: str
+ run_id: str
+ task_id: str
+ ti_id: UUID | None = None
+ region_id: UUID | None = None
+ region_index: int | None = None
+
+ def build_path(self, prefix: str, key: str, *, suffix: str = "",
max_length: int = 1024) -> str:
+ """
+ Build
``<prefix>/<dag_id>/<run_id>/<task_id>/<instance>/<key><suffix>`` as the
storage path for ``key``.
+
+ ``<instance>`` is ``ti-<ti_id>`` when ``ti_id`` is set, so that a task
instance rewriting a key reuses
+ its path, and ``write-<uuid>`` otherwise, so that no two writes share
a path. A backend that builds
+ its paths here needs no case of its own for either.
+
+ ``dag_id``, ``run_id`` and ``task_id`` are shortened, longest first,
only when the path would exceed
+ ``max_length`` bytes. The key is quoted, or replaced by its SHA-256
digest when it is too long to
+ quote, and ``suffix`` follows it directly. The key is chosen by the
Dag author and can look like a
+ file extension, so pass a ``suffix`` that names the stored format if
the backend reads the format
+ back from the object name.
+
+ :raises ValueError: if the path cannot be made to fit in
``max_length`` bytes
+ """
+ instance = f"ti-{self.ti_id}" if self.ti_id is not None else
f"write-{uuid4()}"
+ name = f"{_build_object_name(key)}{suffix}"
+ readable = [self.dag_id, self.run_id, self.task_id]
+ head = prefix.removesuffix("/")
+
+ def join() -> str:
+ return "/".join(part for part in (head, *readable, instance, name)
if part)
+
+ while len(join().encode()) > max_length:
+ longest = max(range(len(readable)), key=lambda i:
len(readable[i].encode()))
+ encoded = readable[longest].encode()
+ if len(encoded) <= _MIN_READABLE_BYTES:
+ raise ValueError(
+ f"Cannot fit an XCom path for key {key!r} in {max_length}
bytes under {prefix!r}"
+ )
+ readable[longest] = encoded[:-1].decode(errors="ignore")
+ return join()
+
+
+_MAX_QUOTED_KEY_LENGTH = 200
+_MIN_READABLE_BYTES = 8
+
+
+def _build_object_name(key: str) -> str:
+ quoted = quote(key, safe="")
+ # A name made only of dots would be the directory itself.
+ if not quoted.strip("."):
+ quoted = quoted.replace(".", "%2E")
+ # Keys are up to 512 characters and quoting expands them, while file
systems cap a name at 255 bytes.
+ if len(quoted) <= _MAX_QUOTED_KEY_LENGTH:
+ return quoted
+ return hashlib.sha256(key.encode()).hexdigest()
+
+
+class _OwnTaskInstance(Protocol):
+ """The task instance a write can come from: a running task, or the task a
trigger is deferred for."""
+
+ id: UUID
+ dag_id: str
+ run_id: str
+ task_id: str
+ map_index: int | None
+ region_id: UUID | None
+ region_index: int | None
+
+
+# Task instance of the trigger whose code is running, so its XCom writes and
SDK requests act as that task
+# instance; tasks serving several triggers must not inherit it.
+_current_trigger_ti: contextvars.ContextVar[_OwnTaskInstance | None] =
contextvars.ContextVar(
+ "airflow_current_trigger_ti", default=None
+)
+
+
class TIKeyProtocol(Protocol):
dag_id: str
task_id: str
@@ -60,7 +166,6 @@ class BaseXCom:
run_id: str,
map_index: int = -1,
dag_result: bool = False,
- _mapped_length: int | None = None,
) -> None:
"""
Store an XCom value.
@@ -73,29 +178,13 @@ class BaseXCom:
:param map_index: Optional map index to assign XCom for a mapped task.
The default is ``-1`` (set for a non-mapped task).
"""
- from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
-
- value = cls.serialize_value(
- value=value,
- key=key,
- task_id=task_id,
- dag_id=dag_id,
- run_id=run_id,
- map_index=map_index,
+ coordinate = _find_own_coordinate(dag_id, run_id, task_id, map_index)
or XComCoordinate(
+ dag_id, run_id, task_id
)
-
- SUPERVISOR_COMMS.send(
- SetXCom(
- key=key,
- value=value,
- dag_id=dag_id,
- task_id=task_id,
- run_id=run_id,
- map_index=map_index,
- dag_result=dag_result,
- mapped_length=_mapped_length,
- ),
+ msg = cls._build_set_message(
+ coordinate, map_index, key, value, dag_result=dag_result,
mapped_length=None
)
+ cls._send_set(msg)
@classmethod
async def aset(
@@ -108,7 +197,6 @@ class BaseXCom:
run_id: str,
map_index: int = -1,
dag_result: bool = False,
- _mapped_length: int | None = None,
) -> None:
"""
Store an XCom value asynchronously.
@@ -121,30 +209,134 @@ class BaseXCom:
:param map_index: Optional map index to assign XCom for a mapped task.
The default is ``-1`` (set for a non-mapped task).
"""
- from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+ coordinate = _find_own_coordinate(dag_id, run_id, task_id, map_index)
or XComCoordinate(
+ dag_id, run_id, task_id
+ )
+ msg = await cls._abuild_set_message(
+ coordinate, map_index, key, value, dag_result=dag_result,
mapped_length=None
+ )
+ await cls._asend_set(msg)
+
+ @classmethod
+ def _push(
+ cls,
+ ti: TaskInstance,
+ key: str,
+ value: Any,
+ *,
+ dag_result: bool = False,
+ mapped_length: int | None = None,
+ ) -> None:
+ msg = cls._build_set_message(
+ _own_coordinate(ti),
+ _public_map_index(ti),
+ key,
+ value,
+ dag_result=dag_result,
+ mapped_length=mapped_length,
+ )
+ cls._send_set(msg)
+
+ @classmethod
+ async def _apush(
+ cls,
+ ti: TaskInstance,
+ key: str,
+ value: Any,
+ *,
+ dag_result: bool = False,
+ mapped_length: int | None = None,
+ ) -> None:
+ msg = await cls._abuild_set_message(
+ _own_coordinate(ti),
+ _public_map_index(ti),
+ key,
+ value,
+ dag_result=dag_result,
+ mapped_length=mapped_length,
+ )
+ await cls._asend_set(msg)
- value = cls.serialize_value(
+ @classmethod
+ def _store(cls, coordinate: XComCoordinate, map_index: int, key: str,
value: Any) -> JsonValue:
+ """Store ``value`` through the backend, returning what the XCom row
holds."""
+ if _is_preferred_over(cls, "store_value", "serialize_value"):
+ return cls.store_value(coordinate, key, value)
+ return cls.serialize_value(
value=value,
key=key,
- task_id=task_id,
- dag_id=dag_id,
- run_id=run_id,
+ task_id=coordinate.task_id,
+ dag_id=coordinate.dag_id,
+ run_id=coordinate.run_id,
map_index=map_index,
)
- await SUPERVISOR_COMMS.asend(
- SetXCom(
- key=key,
- value=value,
- dag_id=dag_id,
- task_id=task_id,
- run_id=run_id,
- map_index=map_index,
- dag_result=dag_result,
- mapped_length=_mapped_length,
- ),
+ @classmethod
+ def _build_set_message(
+ cls,
+ coordinate: XComCoordinate,
+ map_index: int,
+ key: str,
+ value: Any,
+ *,
+ dag_result: bool,
+ mapped_length: int | None,
+ ) -> SetXCom:
+ stored = cls._store(coordinate, map_index, key, value)
+ return cls._make_set_message(
+ coordinate, map_index, key, stored, dag_result=dag_result,
mapped_length=mapped_length
+ )
+
+ @classmethod
+ async def _abuild_set_message(
+ cls,
+ coordinate: XComCoordinate,
+ map_index: int,
+ key: str,
+ value: Any,
+ *,
+ dag_result: bool,
+ mapped_length: int | None,
+ ) -> SetXCom:
+ # The backend writes to external storage, and the event loop also
serves the triggerer's other triggers.
+ stored = await asyncio.to_thread(cls._store, coordinate, map_index,
key, value)
+ return cls._make_set_message(
+ coordinate, map_index, key, stored, dag_result=dag_result,
mapped_length=mapped_length
+ )
+
+ @staticmethod
+ def _make_set_message(
+ coordinate: XComCoordinate,
+ map_index: int,
+ key: str,
+ stored: JsonValue,
+ *,
+ dag_result: bool,
+ mapped_length: int | None,
+ ) -> SetXCom:
+ return SetXCom(
+ key=key,
+ value=stored,
+ dag_id=coordinate.dag_id,
+ task_id=coordinate.task_id,
+ run_id=coordinate.run_id,
+ map_index=map_index,
+ dag_result=dag_result,
+ mapped_length=mapped_length,
)
+ @classmethod
+ def _send_set(cls, msg: SetXCom) -> None:
+ from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+
+ SUPERVISOR_COMMS.send(msg)
+
+ @classmethod
+ async def _asend_set(cls, msg: SetXCom) -> None:
+ from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+
+ await SUPERVISOR_COMMS.asend(msg)
+
@classmethod
def _set_xcom_in_db(
cls,
@@ -334,7 +526,7 @@ class BaseXCom:
raise TypeError(f"Expected XComResult, received: {type(msg)}
{msg}")
if msg.value is not None:
- return cls.deserialize_value(msg)
+ return cls._load(msg)
log.debug(
"No XCom value found; defaulting to None.",
key=key,
@@ -396,7 +588,7 @@ class BaseXCom:
raise TypeError(f"Expected XComResult, received: {type(msg)}
{msg}")
if msg.value is not None:
- return cls.deserialize_value(msg)
+ return await asyncio.to_thread(cls._load, msg)
log.debug(
"No XCom value found; defaulting to None.",
key=key,
@@ -456,7 +648,7 @@ class BaseXCom:
if not msg.root:
return None
- return [cls.deserialize_value(_XComValueWrapper(value)) for value in
msg.root]
+ return cls._load_many(msg.root)
@classmethod
async def aget_all(
@@ -507,7 +699,52 @@ class BaseXCom:
if not msg.root:
return None
- return [cls.deserialize_value(_XComValueWrapper(value)) for value in
msg.root]
+ return await asyncio.to_thread(cls._load_many, msg.root)
+
+ @classmethod
+ def store_value(cls, coordinate: XComCoordinate, key: str, value: Any) ->
JsonValue:
+ """
+ Store an XCom value for ``coordinate`` and return what its XCom row in
the metadata database holds.
+
+ Override this to keep values outside the metadata database, in any
format. The return value must be
+ JSON-serializable and must locate what was stored, since
:meth:`load_value` and :meth:`purge` only
+ receive the row. Every iteration of a task inside a loop is a separate
task instance with the same
+ ``dag_id``, ``run_id``, ``task_id``, ``map_index`` and ``key``, so the
storage location must tell
+ task instances apart.
+
+ :meth:`XComCoordinate.build_path` builds such a location. When a task
instance, or its trigger,
+ writes its own XCom, ``coordinate.ti_id`` is that task instance's id,
and a location unique per
+ ``(coordinate.ti_id, key)`` is enough. When one task writes another's
XCom, ``coordinate.ti_id`` is
+ ``None``, because only the server knows which task instance the
coordinates resolve to, so the
+ location must be unique to each write. A backend that does not use
``build_path`` has to handle both
+ cases itself, or call ``super()``.
+
+ The default stores the value in the metadata database, so an override
can call ``super()`` for the
+ values it does not keep elsewhere.
+ """
+ return BaseXCom.serialize_value(value)
+
+ @classmethod
+ def load_value(cls, stored: JsonValue) -> Any:
+ """
+ Return the value for an XCom row, loading it from storage if
:meth:`store_value` put it there.
+
+ Rows can also hold values the backend did not store, such as those
written through the REST API or
+ under an earlier backend, so pass those to ``super()``. The default
delegates to
+ :meth:`deserialize_value`.
+ """
+ return cls.deserialize_value(_XComValueWrapper(stored))
+
+ @classmethod
+ def _load_many(cls, values: list[JsonValue]) -> list[Any]:
+ return [cls._load(_XComValueWrapper(value)) for value in values]
+
+ @classmethod
+ def _load(cls, result: Any) -> Any:
+ # Legacy backends keep receiving the result object deserialize_value
has always been given.
+ if _is_preferred_over(cls, "load_value", "deserialize_value"):
+ return cls.load_value(result.value)
+ return cls.deserialize_value(result)
@staticmethod
def serialize_value(
@@ -566,3 +803,61 @@ class BaseXCom:
),
)
cls.purge(xcom_result)
+
+
+def _is_preferred_over(xcom: type[BaseXCom], new_hook: str, old_hook: str) ->
bool:
+ # The override closest to the class wins, so a subclass that overrides
only the old hook (e.g. of a
+ # backend that implements store_value) keeps its override.
+ mro = xcom.__mro__
+ return next(i for i, k in enumerate(mro) if new_hook in vars(k)) <= next(
+ i for i, k in enumerate(mro) if old_hook in vars(k)
+ )
+
+
+def _is_legacy_xcom_backend(xcom: type[BaseXCom]) -> bool:
+ """
+ Whether the XCom backend stores task XComs through ``serialize_value``
rather than ``store_value``.
+
+ Such a backend does not receive the task instance id, so it cannot keep
values of different task
+ instances with the same ``dag_id``, ``run_id``, ``task_id``, ``map_index``
and ``key`` apart.
+ """
+ return not _is_preferred_over(xcom, "store_value", "serialize_value")
+
+
+def _public_map_index(ti: _OwnTaskInstance) -> int:
+ return ti.map_index if ti.map_index is not None else -1
+
+
+def _own_coordinate(ti: _OwnTaskInstance) -> XComCoordinate:
+ # The top-level region and a missing one both mean the task instance is
outside any region.
+ in_region = ti.region_id not in (None, UUID(int=0))
+ return XComCoordinate(
+ ti.dag_id,
+ ti.run_id,
+ ti.task_id,
+ ti.id,
+ ti.region_id if in_region else None,
+ ti.region_index if in_region else None,
+ )
+
+
+def _find_own_coordinate(dag_id: str, run_id: str, task_id: str, map_index:
int) -> XComCoordinate | None:
+ """Return the task instance making the write as a coordinate when the
coordinates are its own."""
+ # The Dag definitions import this module, so importing them at the top
level is circular.
+ from airflow.sdk.definitions._internal.contextmanager import
_CURRENT_CONTEXT
+
+ owner: _OwnTaskInstance
+ if (trigger_ti := _current_trigger_ti.get()) is not None:
+ owner = trigger_ti
+ elif _CURRENT_CONTEXT and (running := _CURRENT_CONTEXT[-1].get("ti")) is
not None:
+ owner = running
+ else:
+ return None
+ if (owner.dag_id, owner.run_id, owner.task_id, _public_map_index(owner))
== (
+ dag_id,
+ run_id,
+ task_id,
+ map_index,
+ ):
+ return _own_coordinate(owner)
+ return None
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 4ce198d14f5..743f9ab7db7 100644
--- a/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py
+++ b/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py
@@ -145,7 +145,7 @@ class LazyXComSequence(Sequence[T]):
)
if not isinstance(msg, XComSequenceSliceResult):
raise TypeError(f"Got unexpected response to
GetXComSequenceSlice: {msg!r}")
- return [XCom.deserialize_value(_XComWrapper(value)) for value in
msg.root]
+ return [XCom._load(_XComWrapper(value)) for value in msg.root]
if not isinstance(key, int):
if (index := getattr(key, "__index__", None)) is not None:
@@ -167,7 +167,7 @@ class LazyXComSequence(Sequence[T]):
raise IndexError(key)
if not isinstance(msg, XComSequenceIndexResult):
raise TypeError(f"Got unexpected response to GetXComSequenceItem:
{msg!r}")
- return XCom.deserialize_value(_XComWrapper(msg.root))
+ return XCom._load(_XComWrapper(msg.root))
def _coerce_slice_index(value: Any) -> int | None:
diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
index 8a552a3c6e5..8b23068ee50 100644
--- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
+++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
@@ -2612,8 +2612,14 @@ class InProcessTestSupervisor(ActivitySubprocess):
**kwargs,
)
- from airflow.sdk.execution_time.task_runner import
RuntimeTaskInstance, finalize, run
+ from airflow.sdk.execution_time.task_runner import (
+ RuntimeTaskInstance,
+ check_xcom_backend_supports_loops,
+ finalize,
+ run,
+ )
+ check_xcom_backend_supports_loops(task.dag)
supervisor.comms = InProcessSupervisorComms(supervisor=supervisor)
with set_supervisor_comms(supervisor.comms):
supervisor.ti = what # type: ignore[assignment]
diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py
b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
index bfa6e493434..92a8b1ed0cb 100644
--- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py
+++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
@@ -60,11 +60,11 @@ from airflow.sdk.api.datamodels._generated import (
)
from airflow.sdk.bases.operator import BaseOperator, ExecutorSafeguard
from airflow.sdk.bases.skipmixin import XCOM_SKIPMIXIN_KEY
-from airflow.sdk.bases.xcom import BaseXCom
+from airflow.sdk.bases.xcom import BaseXCom, _is_legacy_xcom_backend
from airflow.sdk.configuration import conf
from airflow.sdk.coordinators._dag_importer import find_claiming_importer
from airflow.sdk.definitions._internal.dag_parsing_context import
_airflow_parsing_context_manager
-from airflow.sdk.definitions._internal.loop import
is_loop_task_read_from_outside
+from airflow.sdk.definitions._internal.loop import LoopTaskGroup,
is_loop_task_read_from_outside
from airflow.sdk.definitions._internal.types import NOTSET, ArgNotSet,
is_arg_set
from airflow.sdk.definitions.asset import (
Asset,
@@ -962,7 +962,7 @@ def _xcom_push(
*,
mapped_length: int | None = None,
) -> None:
- """Push a XCom through XCom.set, which pushes to XCom Backend if
configured."""
+ """Push a XCom for ``ti`` through the configured XCom backend."""
# Private function, as we don't want to expose the ability to manually set
`mapped_length` to SDK
# consumers
@@ -974,16 +974,7 @@ def _xcom_push(
_xcom_push_to_db(ti, key, value)
return
- XCom.set(
- key=key,
- value=value,
- dag_id=ti.dag_id,
- task_id=ti.task_id,
- run_id=ti.run_id,
- map_index=ti.map_index,
- dag_result=ti.task.returns_dag_result,
- _mapped_length=mapped_length,
- )
+ XCom._push(ti, key, value, dag_result=ti.task.returns_dag_result,
mapped_length=mapped_length)
async def _axcom_push(
@@ -993,20 +984,11 @@ async def _axcom_push(
*,
mapped_length: int | None = None,
) -> None:
- """Push a XCom through XCom.aset, which pushes to XCom Backend if
configured."""
+ """Push a XCom for ``ti`` through the configured XCom backend."""
# Private function, as we don't want to expose the ability to manually set
`mapped_length` to SDK
# consumers
- await XCom.aset(
- key=key,
- value=value,
- dag_id=ti.dag_id,
- task_id=ti.task_id,
- run_id=ti.run_id,
- map_index=ti.map_index,
- dag_result=ti.task.returns_dag_result,
- _mapped_length=mapped_length,
- )
+ await XCom._apush(ti, key, value, dag_result=ti.task.returns_dag_result,
mapped_length=mapped_length)
def _xcom_push_to_db(ti: RuntimeTaskInstance, key: str, value: Any) -> None:
@@ -1113,6 +1095,19 @@ def _fail_lang_sdk_task(what: StartupDetails,
coordinator_classpath: str, log: L
sys.exit(0)
+def check_xcom_backend_supports_loops(dag) -> None:
+ if not _is_legacy_xcom_backend(XCom):
+ return
+ if not any(isinstance(group, LoopTaskGroup) for group in
dag.task_group.get_task_group_dict().values()):
+ return
+ raise TypeError(
+ f"Dag {dag.dag_id!r} uses a loop, but XCom backend
{XCom.__module__}.{XCom.__qualname__} stores "
+ "values through serialize_value rather than store_value, so every
iteration of a task could store "
+ "its XCom in the same place. Implement store_value and include ti.id
in the storage location; for "
+ "XComObjectStorageBackend, upgrade apache-airflow-providers-common-io."
+ )
+
+
@detail_span("parse")
def parse(what: StartupDetails, log: Logger) -> RuntimeTaskInstance:
# TODO: Task-SDK:
@@ -1169,6 +1164,8 @@ def parse(what: StartupDetails, log: Logger) ->
RuntimeTaskInstance:
f"task is of the wrong type, got {type(task)}, wanted
{BaseOperator} or {MappedOperator}"
)
+ check_xcom_backend_supports_loops(dag)
+
# Register operator-declared XCom model classes (e.g. ``output_type``) for
# the whole DAG so this task can deserialize structured output produced by
# any other task -- including mapped producers and DAGs loaded from a cache
diff --git a/task-sdk/src/airflow/sdk/types.py
b/task-sdk/src/airflow/sdk/types.py
index 583e33da9bf..eccee7a26a1 100644
--- a/task-sdk/src/airflow/sdk/types.py
+++ b/task-sdk/src/airflow/sdk/types.py
@@ -138,6 +138,8 @@ class RuntimeTaskInstanceProtocol(Protocol):
run_id: str
try_number: int
map_index: int | None
+ region_id: uuid.UUID | None
+ region_index: int | None
max_tries: int
hostname: str | None = None
start_date: AwareDatetime
diff --git a/task-sdk/tests/task_sdk/bases/test_xcom.py
b/task-sdk/tests/task_sdk/bases/test_xcom.py
index caff4df089a..5ce759815d5 100644
--- a/task-sdk/tests/task_sdk/bases/test_xcom.py
+++ b/task-sdk/tests/task_sdk/bases/test_xcom.py
@@ -17,15 +17,26 @@
from __future__ import annotations
+import hashlib
+import threading
+import uuid
+from typing import Any
from unittest import mock
import pytest
-from airflow.sdk.bases.xcom import BaseXCom
+from airflow.sdk import BaseOperator
+from airflow.sdk.bases.xcom import (
+ BaseXCom,
+ XComCoordinate,
+ _current_trigger_ti,
+ _is_legacy_xcom_backend,
+)
from airflow.sdk.execution_time.comms import (
DeleteXCom,
GetXCom,
GetXComSequenceSlice,
+ SetXCom,
XComResult,
XComSequenceSliceResult,
)
@@ -288,3 +299,369 @@ class TestBaseXCom:
include_prior_dates=False,
)
)
+
+
+class LegacyXCom(BaseXCom):
+ @staticmethod
+ def serialize_value(value, *, key=None, task_id=None, dag_id=None,
run_id=None, map_index=None):
+ return f"legacy://{dag_id}/{run_id}/{task_id}/{map_index}/{key}"
+
+ @staticmethod
+ def deserialize_value(result):
+ return f"deserialized {result.value}"
+
+
+class StoreXCom(BaseXCom):
+ @classmethod
+ def store_value(cls, coordinate, key, value):
+ return f"store://{coordinate.ti_id}/{key}"
+
+ @classmethod
+ def load_value(cls, stored):
+ return f"loaded {stored}"
+
+
+class CompatXCom(BaseXCom):
+ store_value = StoreXCom.__dict__["store_value"]
+ load_value = StoreXCom.__dict__["load_value"]
+ serialize_value = LegacyXCom.__dict__["serialize_value"]
+ deserialize_value = LegacyXCom.__dict__["deserialize_value"]
+
+
+class SubclassOverridingOldHooks(StoreXCom):
+ serialize_value = LegacyXCom.__dict__["serialize_value"]
+ deserialize_value = LegacyXCom.__dict__["deserialize_value"]
+
+
+class StoreMixin:
+ @classmethod
+ def store_value(cls, coordinate, key, value):
+ return "mixin"
+
+
[email protected](
+ ("backend", "expected"),
+ [
+ pytest.param(BaseXCom, False, id="base"),
+ pytest.param(
+ type("ReadOnly", (BaseXCom,), {"deserialize_value":
staticmethod(str)}), False, id="read-only"
+ ),
+ pytest.param(LegacyXCom, True, id="serialize-only"),
+ pytest.param(type("LegacySubclass", (LegacyXCom,), {}), True,
id="inherits-serialize-only"),
+ pytest.param(StoreXCom, False, id="store"),
+ pytest.param(CompatXCom, False, id="store-and-serialize-in-one-class"),
+ pytest.param(SubclassOverridingOldHooks, True,
id="serialize-overridden-below-store"),
+ pytest.param(
+ type("StoreBelowLegacy", (LegacyXCom,), {"store_value":
StoreXCom.__dict__["store_value"]}),
+ False,
+ id="store-overridden-below-serialize",
+ ),
+ pytest.param(type("MixinStore", (StoreMixin, LegacyXCom), {}), False,
id="store-from-mixin"),
+ ],
+)
+def test_is_legacy_xcom_backend(backend, expected):
+ assert _is_legacy_xcom_backend(backend) is expected
+
+
+class TestStoreAndLoadValue:
+ @pytest.mark.parametrize(
+ ("backend", "expected_value"),
+ [
+ pytest.param(StoreXCom, "store://{ti_id}/key", id="store"),
+ pytest.param(CompatXCom, "store://{ti_id}/key",
id="store-and-serialize-in-one-class"),
+ pytest.param(LegacyXCom, "legacy://test_dag/test_run/task/-1/key",
id="legacy"),
+ pytest.param(
+ SubclassOverridingOldHooks,
+ "legacy://test_dag/test_run/task/-1/key",
+ id="serialize-overridden-below-store",
+ ),
+ ],
+ )
+ @pytest.mark.parametrize("use_async", [False, True])
+ @pytest.mark.asyncio
+ async def test_push_stores_value_for_the_task_instance(
+ self, backend, expected_value, use_async, create_runtime_ti,
mock_supervisor_comms
+ ):
+ ti = create_runtime_ti(task=BaseOperator(task_id="task"))
+ mock_supervisor_comms.asend = mock.AsyncMock()
+
+ if use_async:
+ await backend._apush(ti, "key", 1, dag_result=True)
+ sent = mock_supervisor_comms.asend.call_args.args[0]
+ else:
+ backend._push(ti, "key", 1, dag_result=True)
+ sent = mock_supervisor_comms.send.call_args.args[0]
+
+ assert sent == SetXCom(
+ key="key",
+ value=expected_value.format(ti_id=ti.id),
+ dag_id="test_dag",
+ task_id="task",
+ run_id="test_run",
+ map_index=-1,
+ dag_result=True,
+ )
+
+ @pytest.mark.parametrize(
+ ("region_id", "region_index"),
+ [
+ pytest.param(None, None, id="no-region"),
+ pytest.param(uuid.UUID(int=0), 2, id="top-level-region"),
+ pytest.param(uuid.uuid4(), 3, id="loop-region"),
+ ],
+ )
+ @pytest.mark.parametrize("use_async", [False, True])
+ @pytest.mark.asyncio
+ async def test_push_passes_the_task_instance_and_its_region_to_store_value(
+ self, region_id, region_index, use_async, create_runtime_ti,
mock_supervisor_comms
+ ):
+ ti = create_runtime_ti(task=BaseOperator(task_id="task"), map_index=2)
+ ti.region_id, ti.region_index = region_id, region_index
+ mock_supervisor_comms.asend = mock.AsyncMock()
+
+ with mock.patch.object(StoreXCom, "store_value", autospec=True,
return_value="stored") as store_value:
+ if use_async:
+ await StoreXCom._apush(ti, "key", 1)
+ else:
+ StoreXCom._push(ti, "key", 1)
+
+ in_region = region_id not in (None, uuid.UUID(int=0))
+ store_value.assert_called_once_with(
+ XComCoordinate(
+ "test_dag",
+ "test_run",
+ "task",
+ ti.id,
+ region_id if in_region else None,
+ region_index if in_region else None,
+ ),
+ "key",
+ 1,
+ )
+ sent = (mock_supervisor_comms.asend if use_async else
mock_supervisor_comms.send).call_args.args[0]
+ assert sent.map_index == 2
+
+ @pytest.mark.parametrize("operation", ["aset", "apush", "aget_one",
"aget_all"])
+ @pytest.mark.asyncio
+ async def test_async_paths_call_the_backend_hooks_off_the_calling_thread(
+ self, operation, create_runtime_ti, mock_supervisor_comms
+ ):
+ threads: list[int] = []
+
+ class RecordingXCom(BaseXCom):
+ @classmethod
+ def store_value(cls, coordinate, key, value):
+ threads.append(threading.get_ident())
+ return "stored"
+
+ @classmethod
+ def load_value(cls, stored):
+ threads.append(threading.get_ident())
+ return stored
+
+ ti = create_runtime_ti(task=BaseOperator(task_id="task"))
+ coordinates = {"key": "key", "dag_id": "dag", "task_id": "task",
"run_id": "run"}
+ reply = (
+ XComSequenceSliceResult(root=["ref"])
+ if operation == "aget_all"
+ else XComResult(key="key", value="ref")
+ )
+ mock_supervisor_comms.asend = mock.AsyncMock(return_value=reply)
+
+ match operation:
+ case "aset":
+ await RecordingXCom.aset(value=1, **coordinates)
+ case "apush":
+ await RecordingXCom._apush(ti, "key", 1)
+ case "aget_one":
+ await RecordingXCom.aget_one(**coordinates)
+ case "aget_all":
+ await RecordingXCom.aget_all(**coordinates)
+
+ assert threads
+ assert threading.get_ident() not in threads
+
+ def test_super_store_value_keeps_the_value_in_the_metadata_database(self):
+ coordinate = XComCoordinate("dag", "run", "task")
+
+ assert BaseXCom.store_value(coordinate, "key", {"a": 1}) == {"a": 1}
+
+ @pytest.mark.parametrize(
+ ("backend", "expected_value"),
+ [
+ pytest.param(StoreXCom, "store://None/key", id="store_value"),
+ pytest.param(CompatXCom, "store://None/key",
id="store-and-serialize-in-one-class"),
+ pytest.param(LegacyXCom, "legacy://dag/run/task/3/key",
id="serialize_value"),
+ ],
+ )
+ @pytest.mark.parametrize("use_async", [False, True])
+ @pytest.mark.asyncio
+ async def
test_set_by_coordinates_goes_through_the_backend_without_a_task_instance_id(
+ self, backend, expected_value, use_async, mock_supervisor_comms
+ ):
+ mock_supervisor_comms.asend = mock.AsyncMock()
+ coordinates = dict(key="key", value=1, dag_id="dag", task_id="task",
run_id="run", map_index=3)
+
+ if use_async:
+ await backend.aset(**coordinates)
+ sent = mock_supervisor_comms.asend.call_args.args[0]
+ else:
+ backend.set(**coordinates)
+ sent = mock_supervisor_comms.send.call_args.args[0]
+
+ assert sent == SetXCom(
+ key="key", value=expected_value, dag_id="dag", task_id="task",
run_id="run", map_index=3
+ )
+
+ @pytest.mark.parametrize(
+ ("owner", "overrides", "expected_id"),
+ [
+ pytest.param("task", {}, "{ti_id}",
id="running-task-own-coordinates"),
+ pytest.param("task", {"task_id": "other"}, "None",
id="running-task-another-task"),
+ pytest.param("task", {"map_index": 2}, "None",
id="running-task-another-map-index"),
+ pytest.param("task", {"dag_id": "other_dag"}, "None",
id="running-task-another-dag"),
+ pytest.param("task", {"run_id": "other_run"}, "None",
id="running-task-another-run"),
+ pytest.param("trigger", {}, "{ti_id}",
id="trigger-own-coordinates"),
+ pytest.param("trigger", {"task_id": "other"}, "None",
id="trigger-another-task"),
+ pytest.param(None, {}, "None", id="no-owner"),
+ ],
+ )
+ @pytest.mark.parametrize("use_async", [False, True])
+ @pytest.mark.asyncio
+ async def
test_set_passes_the_task_instance_id_only_for_the_writers_own_coordinates(
+ self, owner, overrides, expected_id, use_async, create_runtime_ti,
mock_supervisor_comms
+ ):
+ ti = create_runtime_ti(task=BaseOperator(task_id="task"))
+ mock_supervisor_comms.asend = mock.AsyncMock()
+ context = [{"ti": ti}] if owner == "task" else []
+ coordinates: dict[str, Any] = {
+ "dag_id": "test_dag",
+ "run_id": "test_run",
+ "task_id": "task",
+ "map_index": -1,
+ **overrides,
+ }
+ token = _current_trigger_ti.set(ti if owner == "trigger" else None)
+ try:
+ with
mock.patch("airflow.sdk.definitions._internal.contextmanager._CURRENT_CONTEXT",
context):
+ if use_async:
+ await StoreXCom.aset(key="key", value=1, **coordinates)
+ sent = mock_supervisor_comms.asend.call_args.args[0]
+ else:
+ StoreXCom.set(key="key", value=1, **coordinates)
+ sent = mock_supervisor_comms.send.call_args.args[0]
+ finally:
+ _current_trigger_ti.reset(token)
+
+ assert sent.value == f"store://{expected_id.format(ti_id=ti.id)}/key"
+
+ @pytest.mark.parametrize(
+ ("backend", "prefix"),
+ [
+ pytest.param(StoreXCom, "loaded", id="load_value"),
+ pytest.param(SubclassOverridingOldHooks, "deserialized",
id="deserialize-overridden-below-load"),
+ ],
+ )
+ @pytest.mark.parametrize("use_async", [False, True])
+ @pytest.mark.asyncio
+ async def test_get_one_loads_the_stored_value(self, backend, prefix,
use_async, mock_supervisor_comms):
+ response = XComResult(key="key", value="ref")
+ coordinates = dict(key="key", dag_id="dag", task_id="task",
run_id="run")
+
+ if use_async:
+ mock_supervisor_comms.asend = mock.AsyncMock(return_value=response)
+ result = await backend.aget_one(**coordinates)
+ else:
+ mock_supervisor_comms.send.return_value = response
+ result = backend.get_one(**coordinates)
+
+ assert result == f"{prefix} ref"
+
+ @pytest.mark.parametrize(
+ ("backend", "prefix"),
+ [
+ pytest.param(StoreXCom, "loaded", id="load_value"),
+ pytest.param(SubclassOverridingOldHooks, "deserialized",
id="deserialize-overridden-below-load"),
+ ],
+ )
+ @pytest.mark.parametrize("use_async", [False, True])
+ @pytest.mark.asyncio
+ async def test_get_all_loads_each_stored_value(self, backend, prefix,
use_async, mock_supervisor_comms):
+ response = XComSequenceSliceResult(root=["a", "b"])
+ coordinates = dict(key="key", dag_id="dag", task_id="task",
run_id="run")
+
+ if use_async:
+ mock_supervisor_comms.asend = mock.AsyncMock(return_value=response)
+ result = await backend.aget_all(**coordinates)
+ else:
+ mock_supervisor_comms.send.return_value = response
+ result = backend.get_all(**coordinates)
+
+ assert result == [f"{prefix} a", f"{prefix} b"]
+
+
+class TestXComCoordinate:
+ def test_build_path_for_a_task_instance_reuses_its_path_for_a_key(self):
+ ti_id = uuid.uuid4()
+ coordinate = XComCoordinate("dag", "run", "task", ti_id)
+
+ path = coordinate.build_path("s3://bucket/xcom", "key", suffix=".gz")
+
+ assert path == f"s3://bucket/xcom/dag/run/task/ti-{ti_id}/key.gz"
+ assert coordinate.build_path("s3://bucket/xcom/", "key", suffix=".gz")
== path
+
+ def test_build_path_without_a_task_instance_is_unique_per_write(self):
+ coordinate = XComCoordinate("dag", "run", "task")
+
+ first = coordinate.build_path("p", "key")
+ second = coordinate.build_path("p", "key")
+
+ head, instance, name = first.rsplit("/", 2)
+ assert (head, name) == ("p/dag/run/task", "key")
+ assert uuid.UUID(instance.removeprefix("write-"))
+ assert first != second
+
+ @pytest.mark.parametrize(
+ ("key", "expected"),
+ [
+ pytest.param("a/key", "a%2Fkey", id="quoted"),
+ pytest.param(".", "%2E", id="dot"),
+ pytest.param("..", "%2E%2E", id="dots"),
+ pytest.param("k" * 300, hashlib.sha256(b"k" * 300).hexdigest(),
id="too-long-to-quote"),
+ pytest.param("/" * 100, hashlib.sha256(b"/" * 100).hexdigest(),
id="quoting-makes-it-too-long"),
+ pytest.param("a" * 200, "a" * 200, id="at-the-limit"),
+ pytest.param("a" * 201, hashlib.sha256(b"a" * 201).hexdigest(),
id="just-over-the-limit"),
+ ],
+ )
+ def test_build_path_names_the_key_safely(self, key, expected):
+ path = XComCoordinate("dag", "run", "task",
uuid.uuid4()).build_path("p", key)
+
+ assert path.rsplit("/", 1)[1] == expected
+
+ @pytest.mark.parametrize("dag_id", ["d" * 100, "é" * 100])
+ def test_build_path_shortens_the_readable_parts_only_when_it_must(self,
dag_id):
+ ti_id = uuid.uuid4()
+ coordinate = XComCoordinate(dag_id, "r" * 100, "t" * 100, ti_id)
+ whole = coordinate.build_path("prefix", "key")
+
+ shortened = coordinate.build_path("prefix", "key",
max_length=len(whole.encode()) - 20)
+
+ assert len(shortened.encode()) <= len(whole.encode()) - 20
+ assert shortened.startswith("prefix/")
+ assert shortened.endswith(f"/ti-{ti_id}/key")
+ assert coordinate.build_path("prefix", "key",
max_length=len(whole.encode())) == whole
+
+ def test_build_path_shortens_the_longest_part_first(self):
+ ti_id = uuid.uuid4()
+ coordinate = XComCoordinate("d" * 10, "r" * 100, "t" * 10, ti_id)
+ whole = coordinate.build_path("p", "key")
+
+ shortened = coordinate.build_path("p", "key", max_length=len(whole) -
20)
+
+ assert shortened == f"p/{'d' * 10}/{'r' * 80}/{'t' *
10}/ti-{ti_id}/key"
+
+ def test_build_path_raises_when_the_fixed_parts_do_not_fit(self):
+ coordinate = XComCoordinate("dag", "run", "task", uuid.uuid4())
+
+ with pytest.raises(ValueError, match="Cannot fit"):
+ coordinate.build_path("p" * 300, "key", max_length=100)
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 16149ae9e26..73ab4ff225b 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,33 @@ def test_getitem_calls_correct_deserialise(monkeypatch,
mock_supervisor_comms, l
)
+class LoadValueXCom(BaseXCom):
+ @classmethod
+ def load_value(cls, stored):
+ return f"loaded {stored}"
+
+
[email protected](
+ ("index", "response", "expected"),
+ [
+ pytest.param(4, XComSequenceIndexResult(root="ref-4"), "loaded ref-4",
id="index"),
+ pytest.param(
+ slice(0, 2),
+ XComSequenceSliceResult(root=["ref-0", "ref-1"]),
+ ["loaded ref-0", "loaded ref-1"],
+ id="slice",
+ ),
+ ],
+)
+def test_getitem_loads_through_load_value(
+ index, response, expected, monkeypatch, mock_supervisor_comms,
lazy_sequence
+):
+ mock_supervisor_comms.send.return_value = response
+ monkeypatch.setattr(airflow.sdk.execution_time.xcom, "XCom", LoadValueXCom)
+
+ assert lazy_sequence[index] == expected
+
+
def test_getitem_indexerror(mock_supervisor_comms, lazy_sequence):
mock_supervisor_comms.send.return_value = ErrorResponse(
error=ErrorType.XCOM_NOT_FOUND,
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 c210c8f7292..ae5c2af97d7 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
@@ -56,7 +56,7 @@ from task_sdk.execution_time import exec_probe_target
from uuid6 import uuid7
from airflow.executors.workloads import BundleInfo
-from airflow.sdk import DAG, BaseOperator, timezone
+from airflow.sdk import DAG, BaseOperator, task_group, timezone
from airflow.sdk._shared.secrets_masker import mask_secret
from airflow.sdk.api import client as sdk_client
from airflow.sdk.api.client import ServerResponseError
@@ -73,6 +73,7 @@ from airflow.sdk.api.datamodels._generated import (
TaskInstance,
TaskInstanceState,
)
+from airflow.sdk.definitions._internal.loop import create_loop
from airflow.sdk.exceptions import AirflowRuntimeError, ErrorType,
TaskAlreadyRunningError
from airflow.sdk.execution_time import supervisor, task_runner
from airflow.sdk.execution_time.comms import (
@@ -187,6 +188,7 @@ from airflow.sdk.execution_time.supervisor import (
from airflow.sdk.execution_time.task_runner import run
from tests_common.test_utils.config import conf_vars
+from tests_common.test_utils.xcom_backends import SerializeValueOnlyXCom
if TYPE_CHECKING:
from collections.abc import Callable
@@ -4901,6 +4903,32 @@ class TestSetSupervisorComms:
class TestInProcessTestSupervisor:
+ def test_rejects_loop_dag_with_legacy_xcom_backend(self, mocker,
monkeypatch):
+ @task_group
+ def body():
+ BaseOperator(task_id="work")
+
+ with DAG(dag_id="loop_dag") as dag:
+ create_loop(body, max_iterations=2)
+ task = dag.task_dict["body.work"]
+ ti = TaskInstance(
+ id=uuid7(),
+ dag_version_id=uuid7(),
+ dag_id="loop_dag",
+ task_id=task.task_id,
+ run_id="test_run",
+ try_number=1,
+ queue="default",
+ )
+ client = mocker.Mock(spec=sdk_client.Client)
+ client.task_instances =
mocker.create_autospec(sdk_client.TaskInstanceOperations, instance=True)
+ monkeypatch.setattr(task_runner, "XCom", SerializeValueOnlyXCom)
+
+ with pytest.raises(TypeError, match="uses a loop"):
+ InProcessTestSupervisor.start(what=ti, task=task, client=client)
+
+ client.task_instances.start.assert_not_called()
+
@pytest.mark.parametrize("callback_error", [None, RuntimeError])
def test_retry_callback_finishes_before_retry_report(self,
make_ti_context, mocker, callback_error):
client = mocker.Mock(spec=sdk_client.Client)
diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
index bf64c087a3e..14e1ec4801e 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
@@ -183,6 +183,7 @@ from airflow.sdk.execution_time.lazy_sequence import
LazyXComSequence
from airflow.sdk.execution_time.task_runner import (
RuntimeTaskInstance,
TaskRunnerMarker,
+ _axcom_push,
_defer_task,
_execute_task,
_find_native_dag_importer,
@@ -192,6 +193,7 @@ from airflow.sdk.execution_time.task_runner import (
_run_execute_callable,
_serialize_outlet_events,
_xcom_push,
+ check_xcom_backend_supports_loops,
finalize,
get_startup_details,
parse,
@@ -207,6 +209,7 @@ from airflow.triggers.testing import SuccessTrigger
from tests_common.test_utils.config import conf_vars
from tests_common.test_utils.mock_operators import AirflowLink
+from tests_common.test_utils.xcom_backends import SerializeValueOnlyXCom
if TYPE_CHECKING:
from kgb import SpyAgency
@@ -290,8 +293,9 @@ def test_parse(test_dags_dir: Path, make_ti_context,
coordinates, map_index):
assert ti.map_index == map_index
[email protected]("airflow.sdk.execution_time.task_runner.check_xcom_backend_supports_loops",
autospec=True)
@mock.patch("airflow.dag_processing.dagbag.BundleDagBag")
-def test_parse_dag_bag(mock_dagbag, test_dags_dir: Path, make_ti_context):
+def test_parse_dag_bag(mock_dagbag, mock_check_xcom_backend, test_dags_dir:
Path, make_ti_context):
"""Test that checks that the BundleDagBag is constructed as expected
during parsing"""
mock_bag_instance = mock.Mock()
mock_dagbag.return_value = mock_bag_instance
@@ -339,6 +343,7 @@ def test_parse_dag_bag(mock_dagbag, test_dags_dir: Path,
make_ti_context):
):
parse(what, mock.Mock())
+ mock_check_xcom_backend.assert_called_once_with(mock_dag)
mock_dagbag.assert_called_once_with(
dag_folder=mock.ANY,
safe_mode=False,
@@ -1097,6 +1102,45 @@ def
test_parse_fails_a_task_of_a_native_dag_without_retries(
)
+class _StoreXCom(SerializeValueOnlyXCom):
+ @classmethod
+ def store_value(cls, coordinate, key, value):
+ return value
+
+
+def _make_dag(with_loop: bool) -> DAG:
+ @task_group
+ def body():
+ BaseOperator(task_id="work")
+
+ with DAG("xcom_backend_check", schedule=None) as dag:
+ if with_loop:
+ create_loop(body, max_iterations=2)
+ else:
+ body()
+ return dag
+
+
[email protected](
+ ("backend", "with_loop"),
+ [
+ pytest.param(SerializeValueOnlyXCom, False, id="legacy-without-loop"),
+ pytest.param(_StoreXCom, True, id="store-value-with-loop"),
+ ],
+)
+def test_check_xcom_backend_supports_loops_accepts(backend, with_loop,
monkeypatch):
+ monkeypatch.setattr("airflow.sdk.execution_time.task_runner.XCom", backend)
+
+ check_xcom_backend_supports_loops(_make_dag(with_loop))
+
+
+def
test_check_xcom_backend_supports_loops_rejects_legacy_backend_in_loop_dag(monkeypatch):
+ monkeypatch.setattr("airflow.sdk.execution_time.task_runner.XCom",
SerializeValueOnlyXCom)
+
+ with pytest.raises(TypeError, match="uses a loop"):
+ check_xcom_backend_supports_loops(_make_dag(with_loop=True))
+
+
@patch("airflow.dag_processing.dagbag.BundleDagBag", autospec=True)
def test_parse_exits_non_zero_when_the_failure_cannot_be_reported(
mock_bag, native_dag_startup, mock_supervisor_comms
@@ -4639,19 +4683,23 @@ class TestXComAfterTaskExecution:
runtime_ti = create_runtime_ti(task=task)
- with mock.patch.object(XCom, "set") as mock_xcom_set:
+ with mock.patch.object(XCom, "_push", autospec=True) as mock_xcom_push:
_xcom_push(runtime_ti, BaseXCom.XCOM_RETURN_KEY, result,
mapped_length=7)
- mock_xcom_set.assert_called_once_with(
- key=BaseXCom.XCOM_RETURN_KEY,
- value=result,
- dag_id=runtime_ti.dag_id,
- task_id=runtime_ti.task_id,
- run_id=runtime_ti.run_id,
- map_index=runtime_ti.map_index,
- dag_result=False,
- _mapped_length=7,
+ mock_xcom_push.assert_called_once_with(
+ runtime_ti, BaseXCom.XCOM_RETURN_KEY, result,
dag_result=False, mapped_length=7
)
+ @pytest.mark.asyncio
+ async def test_axcom_push_passes_the_task_instance(self,
create_runtime_ti):
+ runtime_ti = create_runtime_ti(task=BaseOperator(task_id="push_task"))
+
+ with mock.patch.object(XCom, "_apush", autospec=True) as
mock_xcom_apush:
+ await _axcom_push(runtime_ti, "key", "value")
+
+ mock_xcom_apush.assert_awaited_once_with(
+ runtime_ti, "key", "value", dag_result=False, mapped_length=None
+ )
+
def test_xcom_with_multiple_outputs_and_no_mapping_result(self,
create_runtime_ti, spy_agency):
"""Test that error is raised when multiple outputs are returned
without mapping."""
result = "value1"
@@ -4710,15 +4758,12 @@ class TestXComAfterTaskExecution:
run(runtime_ti, context=runtime_ti.get_template_context(),
log=mock.MagicMock())
- mock_xcom_backend.set.assert_called_once_with(
- key=BaseXCom.XCOM_RETURN_KEY,
- value="pushing to xcom backend!",
- dag_id="test_dag",
- task_id="pull_task",
- run_id="test_run",
- map_index=-1,
+ mock_xcom_backend._push.assert_called_once_with(
+ runtime_ti,
+ BaseXCom.XCOM_RETURN_KEY,
+ "pushing to xcom backend!",
dag_result=False,
- _mapped_length=None,
+ mapped_length=None,
)
# assert that we didn't call the API when XCom backend is configured
