dabla commented on code in PR #62922: URL: https://github.com/apache/airflow/pull/62922#discussion_r4248483468
########## task-sdk/tests/task_sdk/definitions/test_iterableoperator.py: ########## @@ -0,0 +1,4036 @@ +# +# 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 + +import asyncio +import threading +import time +from collections.abc import Iterator +from contextlib import contextmanager +from datetime import UTC, datetime, timedelta +from types import SimpleNamespace +from typing import TYPE_CHECKING, Any +from unittest.mock import AsyncMock, MagicMock, Mock, call, create_autospec, patch + +try: + # Python 3.11+ + BaseExceptionGroup +except NameError: + from exceptiongroup import BaseExceptionGroup + +import jinja2 +import pytest +from task_sdk.definitions.conftest import make_xcom_arg + +from airflow.sdk import ( + DAG, + Asset, + BaseAsyncOperator, + BaseOperator, + BaseXCom, + TaskInstanceState, + get_current_context, +) +from airflow.sdk.bases.operator import event_loop +from airflow.sdk.definitions._internal.abstractoperator import DEFAULT_RETRIES +from airflow.sdk.definitions._internal.expandinput import ( + DictOfListsExpandInput, + ExpandInput, + ListOfDictsExpandInput, + Resolved, +) +from airflow.sdk.definitions.context import clone_context +from airflow.sdk.definitions.iterableoperator import ( + Checkpoints, + IndexedTaskInstanceNotStarted, + IndexedTaskOutcomes, + IterableOperator, + IterationState, +) +from airflow.sdk.exceptions import ( + AirflowFailException, + AirflowRescheduleException, + AirflowSkipException, + AirflowTaskTerminated, + AirflowTaskTimeout, + DagRunTriggerException, + DownstreamTasksSkipped, + TaskDeferred, + UnmappableXComTypePushed, +) +from airflow.sdk.execution_time.comms import DeadlockImminentError +from airflow.sdk.execution_time.context import InletEventsAccessors, OutletEventAccessors +from airflow.sdk.execution_time.executor import AsyncAwareExecutor +from airflow.sdk.execution_time.task_runner import ( + IndexedTaskInstance, + IndexedTaskRunner, + IndexedTaskState, + RuntimeTaskInstance, +) +from airflow.sdk.execution_time.xcom import XCom + +from tests_common.test_utils.mock_context import mock_context as _mock_context_base + +if TYPE_CHECKING: + from airflow.sdk.definitions._internal.expandinput import ExpandInput + from airflow.sdk.definitions.mappedoperator import MappedOperator + + from tests_common.test_utils.compat import Context + + +class MockTaskStateStoreAccessor: + """Minimal in-memory stand-in for ``TaskStateStoreAccessor``, exposing only the async + ``aget``/``aset`` methods used by ``IterableOperator`` to checkpoint per-index sub-task + progress (see ``IterableOperator._run_task``), plus the sync ``get``/``set``/``delete`` used for + the completion marker (see ``IterableOperator._run_tasks``).""" + + def __init__(self): + self._data: dict[str, Any] = {} + + async def aget(self, key: str, default: Any = None) -> Any: + return self._data.get(key, default) + + async def aset(self, key: str, value: Any, **kwargs) -> None: + self._data[key] = value + + def get(self, key: str, default: Any = None) -> Any: + return self._data.get(key, default) + + def set(self, key: str, value: Any, **kwargs) -> None: + self._data[key] = value + + def delete(self, key: str) -> None: + self._data.pop(key, None) + + def __contains__(self, key: str) -> bool: + return key in self._data + + def __getitem__(self, key: str) -> Any: + return self._data[key] + + +@contextmanager +def mock_context(task, run_id: str | None = None) -> Iterator[Context]: + """Create a mock context for IterableOperator tests. + + The context includes the task state store, asset event accessors, DAG/run + information, and a mocked XCom backend. + """ + task_state_store = MockTaskStateStoreAccessor() + context = _mock_context_base(task=task, run_id=run_id) + context["dag"] = task.dag # type: ignore[typeddict-item] + context["dag_run"] = SimpleNamespace(conf={}) # type: ignore[typeddict-item] + context["task_state_store"] = task_state_store # type: ignore[typeddict-item] + context["outlet_events"] = OutletEventAccessors() + context["inlet_events"] = InletEventsAccessors(inlets=[]) + + def _set( + cls, + key, + value, + *, + dag_id, + task_id, + run_id, + map_index=-1, + **kwargs, + ): + context["ti"].xcom_push(key=key, value=value) + + async def _aset( + cls, + key, + value, + *, + dag_id, + task_id, + run_id, + map_index=-1, + **kwargs, + ): + context["ti"].xcom_push(key=key, value=value) + + def _get_one( + cls, + *, + key, + dag_id, + task_id, + run_id, + map_index=None, + include_prior_dates=False, + **kwargs, + ): + return context["ti"].xcom_pull( + task_ids=task_id, + dag_id=dag_id, + key=key, + ) + + def _get_all(cls, *, key, dag_id, task_id, run_id, include_prior_dates=False, **kwargs): + """One map index in these tests, so the values of all of them are the one value, or none.""" + value = context["ti"].xcom_pull(task_ids=task_id, dag_id=dag_id, key=key) + return None if value is None else [value] + + with ( + patch.object(XCom, "set", classmethod(_set)), + patch.object(XCom, "aset", classmethod(_aset)), + patch.object(XCom, "get_one", classmethod(_get_one)), + patch.object(XCom, "get_all", classmethod(_get_all)), + patch.object(RuntimeTaskInstance, "task_state_store", property(lambda self: task_state_store)), + ): + yield context + + +class MockOperator(BaseOperator): + """Mock operator for testing IterableOperator expansion.""" + + template_fields = ("arg1", "arg2", "arg3") + + def __init__( + self, + arg1=None, + arg2=None, + arg3=None, + fail_on_first_attempt=False, + raise_exception: BaseException | None = None, + **kwargs, + ): + super().__init__(**kwargs) + self.arg1 = arg1 + self.arg2 = arg2 + self.arg3 = arg3 + self.fail_on_first_attempt = fail_on_first_attempt + self.raise_exception = raise_exception + + def execute(self, context): + """Execute the operator and return passed arguments as tuple if do_xcom_push is True.""" + expected = clone_context(context) + + if self.raise_exception is not None: + raise self.raise_exception + if self.fail_on_first_attempt: + self.fail_on_first_attempt = False + raise RuntimeError + if not self.do_xcom_push: + return None + + assert context == expected, "Context was unexpectedly mutated during task execution" + return self.arg1, self.arg2, self.arg3 + + +class MockOperatorWithCustomName(MockOperator): + """MockOperator subclass with a custom display name, mimicking a @task-decorated callable, + used to verify IterableOperator.operator_name forwards the wrapped operator's own + operator_name rather than falling back to this wrapper's task_type.""" + + custom_operator_name = "@mock_task" + + +class MockOutletEventOperator(BaseOperator): + """Operator that records an outlet asset event on execute, used to test that + IterableOperator merges/replays per-sub-task outlet events (see ``_run_task``).""" + + template_fields = () + + def __init__(self, extra_value: str = "v", **kwargs): + super().__init__(**kwargs) + self.extra_value = extra_value + + def execute(self, context): + context["outlet_events"][Asset(name="a", uri="s3://bucket/a")].extra["value"] = self.extra_value + return "done" + + +class MockOnKillOperator(BaseOperator): + """Operator that records whether ``on_kill()`` was called on it, used to test that + IterableOperator.on_kill() propagates to currently in-flight sub-tasks (see ``on_kill``).""" + + template_fields = () + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.killed = False + + def execute(self, context): + return "done" + + def on_kill(self): + self.killed = True + + +class MockDeferredOperator(BaseOperator): + """Operator that immediately defers on execute, simulating a deferrable operator.""" + + template_fields = () + + def execute(self, context): + raise TaskDeferred(trigger=None, method_name="execute_complete") # type: ignore[arg-type] + + +class MockRescheduleSensor(BaseOperator): + """Operator that raises AirflowRescheduleException on execute, simulating a reschedule-mode sensor.""" + + template_fields = () + + def execute(self, context): + from datetime import timedelta + + from airflow.sdk import timezone + + raise AirflowRescheduleException(timezone.utcnow() + timedelta(seconds=60)) + + +class MockStateStoreOperator(BaseOperator): + """Sync operator that keeps its own state in the task state store, as a paginated fetch would.""" + + template_fields = ("offset",) + + def __init__(self, offset=None, **kwargs): + super().__init__(**kwargs) + self.offset = offset + + def execute(self, context): + store = context["task_state_store"] + assert store is context["ti"].task_state_store + store.set("last_offset", self.offset) + return store.get("last_offset") + + +class MockAsyncStateStoreOperator(BaseAsyncOperator): + """Async twin of MockStateStoreOperator.""" + + template_fields = ("offset",) + + def __init__(self, offset=None, **kwargs): + super().__init__(**kwargs) + self.offset = offset + + async def aexecute(self, context): + store = context["task_state_store"] + await store.aset("last_offset", self.offset) + return await store.aget("last_offset") + + +class MockAttemptOperator(BaseAsyncOperator): + """ + Async operator whose result tells which attempt produced it. + + For the values in ``times_out_on`` the parent's execution timeout strikes instead, which ends + the whole iteration there and leaves the items after it unreached, as a crash would. Async, + because that timeout lands on the main thread, where async operators run. + """ + + template_fields = ("arg1",) + times_out_on: set = set() + + def __init__(self, arg1=None, **kwargs): + super().__init__(**kwargs) + self.arg1 = arg1 + + async def aexecute(self, context): + if self.arg1 in self.times_out_on: + raise AirflowTaskTimeout("the iteration ran out of time") + return f"{self.arg1}@attempt{context['ti'].try_number}" + + +FIRED_CALLBACKS: list = [] + + +class MockCallbackTimeoutOperator(BaseAsyncOperator): + """ + Async operator in which the parent's execution timeout strikes; records which callback fired. + + Async, because the parent's limit lands on the main thread, where async operators run; a sync + operator raising ``AirflowTaskTimeout`` in its worker thread raised its own. + """ + + template_fields = () + + def __init__(self, **kwargs): + kwargs["on_failure_callback"] = lambda context: FIRED_CALLBACKS.append("failure") + kwargs["on_retry_callback"] = lambda context: FIRED_CALLBACKS.append("retry") + super().__init__(**kwargs) + + async def aexecute(self, context): + raise AirflowTaskTimeout("the task ran out of time") + + +class MockCallbackAsyncOperator(BaseAsyncOperator): + """Async operator that defers for ``arg1="defer"`` and sleeps otherwise; records its callbacks.""" + + template_fields = ("arg1",) + + def __init__(self, arg1=None, **kwargs): + kwargs["on_failure_callback"] = lambda context: FIRED_CALLBACKS.append(("failure", self.arg1)) + kwargs["on_retry_callback"] = lambda context: FIRED_CALLBACKS.append(("retry", self.arg1)) + super().__init__(**kwargs) + self.arg1 = arg1 + + async def aexecute(self, context): + if self.arg1 == "defer": + raise TaskDeferred(trigger=None, method_name="execute_complete") # type: ignore[arg-type] + await asyncio.sleep(60) + + +class MockPushingOperator(BaseOperator): + """Operator that pushes an extra XCom next to its return value, and fails for ``fail=True``.""" + + template_fields = ("arg1",) + executed: list = [] + + def __init__(self, arg1=None, fail=False, **kwargs): + super().__init__(**kwargs) + self.arg1 = arg1 + self.fail = fail + + def execute(self, context): + type(self).executed.append(self.arg1) + context["ti"].xcom_push(key="foo", value=f"foo-of-{self.arg1}") + if self.fail: + raise RuntimeError("sibling failed") + return self.arg1 + + +class MockClearingStateStoreOperator(BaseOperator): + template_fields = () + + def execute(self, context): + context["task_state_store"].clear() + + +def wait_until(condition, timeout: float = 5.0) -> bool: + """Poll ``condition`` until it holds; on_kill() kills in a thread of its own.""" + deadline = time.monotonic() + timeout + while not condition(): + if time.monotonic() > deadline: + return False + time.sleep(0.01) + return True + + +def create_mapped_operator( + dag: DAG, + expand_input: ExpandInput, + task_id: str = "my_task", + retries: int = DEFAULT_RETRIES, + do_xcom_push: bool = True, + task_concurrency: int | None = None, + execution_timeout: timedelta | None = None, + operator_class: type[BaseOperator] = MockOperator, +) -> MappedOperator: + """ + Create a MappedOperator and assign it to a DAG. + + :param expand_input: The input to expand + :param dag: The DAG to assign the operator to + :param task_id: Task ID for the operator + :param do_xcom_push: Whether to push XCom (default True) + :param operator_class: Operator class to wrap (default MockOperator) + """ + return operator_class.partial( + task_id=task_id, + dag=dag, + retries=retries, + task_concurrency=task_concurrency, + do_xcom_push=do_xcom_push, + execution_timeout=execution_timeout, + )._expand( + expand_input, + strict=True, + register_with_dag=False, + ) + + +def create_iterable_operator( + dag: DAG, + expand_input: ExpandInput, + task_id: str = "my_task", + task_concurrency: int | None = None, + retries: int = DEFAULT_RETRIES, + do_xcom_push: bool = True, + operator_class: type[BaseOperator] = MockOperator, +) -> IterableOperator: + """Create an IterableOperator with a MappedOperator and ExpandInput.""" + mapped_op = create_mapped_operator( + dag=dag, + expand_input=expand_input, + task_id=task_id, + retries=retries, + do_xcom_push=do_xcom_push, + task_concurrency=task_concurrency, + operator_class=operator_class, + ) + return IterableOperator( + operator=mapped_op, + expand_input=expand_input, + dag=dag, + ) + + +def _items(expand_input: ExpandInput) -> list: + """Every item of the input in index order, the way IterableOperator.execute reads it.""" + + async def read(): + length, aget = await expand_input.aresolve({}) + return [await aget(index) for index in range(length)] + + return asyncio.run(read()) + + +class TestIterableOperator: + @pytest.mark.parametrize( + ("actual", "expected"), + [ + ([{"a": 1}, {"a": 2}], [{"a": 1}, {"a": 2}]), + ([{"a": 1, "b": 2}], [{"a": 1, "b": 2}]), + ([], []), + ], + ) + def test_list_of_dicts_expand_input_aresolve(self, actual, expected): + """Test IterableOperator with ListOfDictsExpandInput expand_input.""" + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput(actual) + iterable_op = create_iterable_operator(dag, expand_input) + + assert _items(iterable_op.expand_input) == expected + + @pytest.mark.parametrize( + ("actual", "expected"), + [ + ({"a": [1, 2, 3]}, [{"a": 1}, {"a": 2}, {"a": 3}]), + ( + {"a": [1, 2], "b": [10, 20]}, + [{"a": 1, "b": 10}, {"a": 1, "b": 20}, {"a": 2, "b": 10}, {"a": 2, "b": 20}], + ), + ({"a": [1, 2]}, [{"a": 1}, {"a": 2}]), + ( + {"a": {"x": 1, "y": 2}}, + [{"a": ("x", 1)}, {"a": ("y", 2)}], + ), + ], + ) + def test_dict_of_lists_expand_input_aresolve(self, actual, expected): + """Test IterableOperator with DictOfListsExpandInput expand_input. + + A dict value expands to its (key, value) pairs (not just its keys), matching + the classic .expand() resolve() path's handling of dict values. + """ + with DAG("test_dag") as dag: + expand_input = DictOfListsExpandInput(actual) + iterable_op = create_iterable_operator(dag, expand_input) + + assert _items(iterable_op.expand_input) == expected + + def test_task_type(self): + """ + IterableOperator reports its own class as its type, so whatever resolves a class by it + (OpenLineage's extractors, the task's class reference) finds the class that runs, while + the name it is shown under is the wrapped operator's. + """ + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{"a": 1}]) + iterable_op = create_iterable_operator(dag, expand_input) + + assert isinstance(iterable_op, IterableOperator) + assert iterable_op.task_type == "IterableOperator" + assert iterable_op.operator_name == "MockOperator" + + def test_operator_name(self): + """Test that IterableOperator forwards the wrapped operator's operator_name (e.g. a + @task-decorated callable's custom_operator_name), not just its own task_type.""" + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{"a": 1}]) + iterable_op = create_iterable_operator( + dag, expand_input, operator_class=MockOperatorWithCustomName + ) + + assert isinstance(iterable_op, IterableOperator) + assert iterable_op.task_type == "IterableOperator" + assert iterable_op.operator_name == "@mock_task" + + def test_forwards_params_weight_rule_and_retry_policy(self): + """Test that IterableOperator forwards params, weight_rule, and retry_policy from the + wrapped operator onto its own DAG node, not just onto the generated sub-tasks.""" + from airflow.sdk import WeightRule + from airflow.sdk.definitions.retry_policy import ExceptionRetryPolicy + + retry_policy = ExceptionRetryPolicy(rules=[]) + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{"a": 1}]) + mapped_op = MockOperator.partial( + task_id="my_task", + dag=dag, + params={"p": 1}, + weight_rule=WeightRule.UPSTREAM, + retry_policy=retry_policy, + )._expand(expand_input, strict=True, register_with_dag=False) + iterable_op = IterableOperator(operator=mapped_op, expand_input=expand_input, dag=dag) + + assert iterable_op.params["p"] == 1 + assert iterable_op.weight_rule == WeightRule.UPSTREAM + assert iterable_op.retry_policy is retry_policy + + def test_default_args_callbacks_and_hooks_stay_with_the_items(self): + """ + Callbacks and execute hooks in ``default_args`` reach the items, not the iterated task. + + ``_apply_defaults`` fills every ``BaseOperator.__init__`` parameter the call leaves out + from ``default_args``; the IterableOperator passes them as ``None``, so a DAG-level + ``on_failure_callback`` runs once per failed item, as through ``.partial()``, and not once + more for the task with the parent's context. + """ + callback = Mock() + with DAG( + "test_dag", + default_args={ + "on_failure_callback": callback, + "on_success_callback": callback, + "on_execute_callback": callback, + "pre_execute": callback, + "post_execute": callback, + }, + ): + iterated = MockOperator.partial(task_id="op").iterate(arg1=["a", "b"]) + + assert iterated.on_failure_callback == [] + assert iterated.on_success_callback == [] + assert iterated.on_execute_callback == [] + assert iterated._pre_execute_hook is None + assert iterated._post_execute_hook is None + item = iterated._operator.unmap({"arg1": "a"}) + assert item.on_failure_callback == [callback] + assert item.on_success_callback == [callback] + assert item._pre_execute_hook is callback + + def test_forwards_do_xcom_push(self): + """Test that IterableOperator forwards do_xcom_push from the wrapped operator's + partial_kwargs onto its own DAG node.""" + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{"a": 1}]) + iterable_op = create_iterable_operator(dag, expand_input, do_xcom_push=False) + + assert iterable_op.do_xcom_push is False + + def test_forwards_is_setup_is_teardown_and_on_failure_fail_dagrun(self): + """Test that IterableOperator forwards is_setup/is_teardown/on_failure_fail_dagrun from + the wrapped operator's partial_kwargs, mirroring what unmap() reads (via + ``_get_unmap_kwargs``) to apply the same flags to each generated sub-task.""" + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{"a": 1}]) + mapped_op = create_mapped_operator(dag, expand_input) + # unmap() only ever applies these three flags to sub-tasks by reading them off + # partial_kwargs (see MappedOperator._get_unmap_kwargs), so that's the only place + # IterableOperator needs to source them from too. + mapped_op.partial_kwargs["is_teardown"] = True + mapped_op.partial_kwargs["on_failure_fail_dagrun"] = True + iterable_op = IterableOperator(operator=mapped_op, expand_input=expand_input, dag=dag) + + assert iterable_op.is_setup is False + assert iterable_op.is_teardown is True + assert iterable_op.on_failure_fail_dagrun is True + + def test_applies_upstream_relationship_for_partial_kwargs_template_fields(self): + """Test that an XComArg passed via a partial kwarg matching the wrapped operator's own + template_fields is wired as an upstream dependency of the IterableOperator, mirroring what + MappedOperator.__attrs_post_init__ does for a normal mapped task.""" + with DAG("test_dag") as dag: + upstream = MockOperator(task_id="upstream", dag=dag) + expand_input = ListOfDictsExpandInput([{"arg1": 1}]) + mapped_op = MockOperator.partial( + task_id="my_task", + dag=dag, + arg1=upstream.output, + )._expand(expand_input, strict=True, register_with_dag=False) + iterable_op = IterableOperator(operator=mapped_op, expand_input=expand_input, dag=dag) + + assert upstream.task_id in iterable_op.upstream_task_ids + + def test_iterate_inside_mapped_task_group_raises_not_implemented_error(self): + """Test that wrapping an operator expanded inside a mapped task group with IterableOperator + raises NotImplementedError, since operator expansion in an expanded task group is not + supported (mirrors MappedOperator's own guard for the analogous .expand() case).""" + from airflow.decorators import task_group + + with DAG("test_dag") as dag: + + @task_group + def tg(va): + expand_input = ListOfDictsExpandInput([{"arg1": 1}]) + mapped_op = MockOperator.partial( + task_id="my_task", + dag=dag, + )._expand(expand_input, strict=True, register_with_dag=False) + + with pytest.raises(NotImplementedError, match="operator expansion in an expanded task group"): + IterableOperator(operator=mapped_op, expand_input=expand_input, dag=dag) + + tg.expand(va=[["a", "b"], [4]]) + + def test_task_retries(self): + """Test that IterableOperator inherits retries from the wrapped operator, since + the whole IterableOperator is now retried via Airflow's standard retry mechanism.""" + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{"a": 1}]) + iterable_op = create_iterable_operator(dag, expand_input, retries=3) + + assert isinstance(iterable_op, IterableOperator) + assert iterable_op.retries == 3 + assert iterable_op.task_retries == 3 + + def test_task_id(self): + """Test that IterableOperator inherits task_id from operator.""" + with DAG("test_dag") as dag: + task_id = "my_task" + expand_input = ListOfDictsExpandInput([{"a": 1}]) + iterable_op = create_iterable_operator(dag, expand_input, task_id=task_id) + + assert iterable_op.task_id == task_id + + def test_with_task_concurrency(self): + """Test that IterableOperator respects task_concurrency parameter.""" + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{"a": 1}]) + iterable_op = create_iterable_operator(dag, expand_input, task_concurrency=4) + + assert iterable_op.max_workers == 4 + + def test_direct_instantiation_rejects_task_concurrency(self): + """A directly instantiated operator can never reach IterableOperator, + so task_concurrency must be rejected instead of silently accepted as a dead value.""" + with pytest.raises(TypeError, match="which is now max_active_tis_per_dag"): + MockOperator(task_id="my_task", task_concurrency=4) + + def test_expand_rejects_task_concurrency(self): + """.expand() produces a plain MappedOperator, never an IterableOperator, so + task_concurrency (only meaningful for .iterate()/.iterate_kwargs()) must be rejected.""" + with DAG("test_dag"): + with pytest.raises(TypeError, match="which is now max_active_tis_per_dag"): + MockOperator.partial(task_id="my_task", task_concurrency=4).expand(arg1=[1, 2, 3]) + + def test_expand_kwargs_rejects_task_concurrency(self): + """.expand_kwargs() produces a plain MappedOperator, never an IterableOperator, so + task_concurrency (only meaningful for .iterate()/.iterate_kwargs()) must be rejected.""" + with DAG("test_dag"): + with pytest.raises(TypeError, match="which is now max_active_tis_per_dag"): + MockOperator.partial(task_id="my_task", task_concurrency=4).expand_kwargs( + [{"arg1": 1}, {"arg1": 2}] + ) + + def test_iterate_accepts_task_concurrency(self): + """.iterate() is the one public entry point where task_concurrency is meaningful: it + produces an IterableOperator, which reads task_concurrency out of partial_kwargs as + max_workers rather than forwarding it to BaseOperator.__init__.""" + with DAG("test_dag"): + iterable_op = MockOperator.partial(task_id="my_task", task_concurrency=4).iterate(arg1=[1, 2, 3]) + + assert isinstance(iterable_op, IterableOperator) + assert iterable_op.max_workers == 4 + + def test_iterate_kwargs_accepts_task_concurrency(self): + """.iterate_kwargs() is the list-of-dicts counterpart to .iterate() and must accept + task_concurrency the same way.""" + with DAG("test_dag"): + iterable_op = MockOperator.partial(task_id="my_task", task_concurrency=4).iterate_kwargs( + [{"arg1": 1}, {"arg1": 2}] + ) + + assert isinstance(iterable_op, IterableOperator) + assert iterable_op.max_workers == 4 + + @pytest.mark.parametrize("invalid_value", [0, -1, -10]) + def test_task_concurrency_validation_rejects_non_positive_values(self, invalid_value): + """Test that IterableOperator raises ValueError for task_concurrency < 1.""" + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{"a": 1}]) + with pytest.raises(ValueError, match=f"task_concurrency must be at least 1, got {invalid_value}"): + create_iterable_operator(dag, expand_input, task_concurrency=invalid_value) + + def test_partial_kwargs_not_mutated(self): + """Test that creating IterableOperator does not mutate the original MappedOperator's partial_kwargs.""" + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{"a": 1}]) + mapped_op = create_mapped_operator(dag, expand_input, task_concurrency=4) + original_partial_kwargs = mapped_op.partial_kwargs.copy() + + IterableOperator(operator=mapped_op, expand_input=expand_input, dag=dag) + + # Verify that mapped_op.partial_kwargs was not mutated + assert mapped_op.partial_kwargs == original_partial_kwargs + assert "task_concurrency" in mapped_op.partial_kwargs + + def test_expand_input_stored(self): + """Test that IterableOperator stores expand_input correctly.""" + with DAG("test_dag") as dag: + expand_input_data = ListOfDictsExpandInput([{"a": 1}, {"a": 2}]) + iterable_op = create_iterable_operator(dag, expand_input_data) + + assert iterable_op.expand_input is expand_input_data + assert isinstance(iterable_op.expand_input, (ListOfDictsExpandInput, DictOfListsExpandInput)) + + def test_partial_kwargs_stored(self): + """Test that IterableOperator stores partial_kwargs from operator.""" + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{"a": 1}]) + iterable_op = create_iterable_operator(dag, expand_input) + + assert hasattr(iterable_op, "partial_kwargs") + assert isinstance(iterable_op.partial_kwargs, dict) + + def test_xcom_push_delegates_to_task(self): + """_xcom_push awaits task.axcom_push with the default XCom return key.""" + from unittest import mock + + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{}]) + iterable_op = create_iterable_operator(dag, expand_input) + + task = mock.MagicMock() + task.axcom_push = mock.AsyncMock() + task.task_id = "my_task" + task.index = 0 + + asyncio.run(iterable_op.axcom_push(task=task, value="result_value")) + + task.axcom_push.assert_awaited_once_with(key=BaseXCom.XCOM_RETURN_KEY, value="result_value") + + def test_execute_list_of_dicts(self): + """Test executing IterableOperator with ListOfDictsExpandInput.""" + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{"arg1": 1}, {"arg1": 2}]) + iterable_op = create_iterable_operator(dag, expand_input, task_id="exec_list_of_dicts") + + with mock_context(task=iterable_op) as context: + result = iterable_op.execute(context=context) + materialized = list(result) + assert materialized == [(1, None, None), (2, None, None)] + + @pytest.mark.parametrize( + "operator_class", + [MockStateStoreOperator, MockAsyncStateStoreOperator], + ids=["sync", "async"], + ) + def test_execute_gives_each_iteration_its_own_task_state_store_keys(self, operator_class): + """Iterations share one task instance, so a key an iteration stores carries its index.""" + with DAG("test_dag") as dag: + expand_input = DictOfListsExpandInput({"offset": [10, 20, 30]}) + iterable_op = create_iterable_operator( + dag, expand_input, task_id="exec_state_store", operator_class=operator_class + ) + + with mock_context(task=iterable_op) as context: + results = list(iterable_op.execute(context=context)) + store = context["task_state_store"] + + assert results == [10, 20, 30] + assert {key: store[key] for key in ("last_offset_0", "last_offset_1", "last_offset_2")} == { + "last_offset_0": 10, + "last_offset_1": 20, + "last_offset_2": 30, + } + assert "last_offset" not in store + # the operator's own checkpoints keep their unsuffixed keys in the same store + assert all(f"_iterable_{index}" in store for index in range(3)) + + def test_task_state_store_clear_is_refused_inside_an_iteration(self): + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{}]) + iterable_op = create_iterable_operator( + dag, expand_input, task_id="exec_clear", operator_class=MockClearingStateStoreOperator + ) + + with mock_context(task=iterable_op) as context: + with pytest.raises(RuntimeError, match="not available inside an iterated task"): + iterable_op.execute(context=context) + + def test_execute_marks_iteration_completed_once_every_index_succeeds(self): + """ + Once every sub-task index has succeeded, ``_run_tasks`` writes a single completion marker + instead of deleting the per-index checkpoints one by one, and leaves any state written by + user code inside the iterated task untouched. + """ + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{"arg1": 1}, {"arg1": 2}]) + iterable_op = create_iterable_operator(dag, expand_input, task_id="exec_completed") + + with mock_context(task=iterable_op) as context: + store = context["task_state_store"] + store._data["last_offset"] = 42 + + list(iterable_op.execute(context=context)) + + assert store["last_offset"] == 42 + assert store["_iterable_completed"]["completed"] is True + assert store["_iterable_0"]["status"] == "success" + assert store["_iterable_1"]["status"] == "success" + + def test_execute_reruns_every_index_after_a_clear_that_follows_success(self): + """ + A manual clear does not reset the parent TI's ``try_number`` (see ``clear_task_instances``), + so the next attempt looks like a retry. The completion marker left by the previous fully + successful run tells ``_run_tasks`` to ignore the stale ``SUCCESS`` checkpoints and run every + index again, recording first where the rerun starts so a crash mid-rerun resumes from the + new checkpoints only. + """ + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{"arg1": 1}, {"arg1": 2}]) + iterable_op = create_iterable_operator(dag, expand_input, task_id="exec_rerun") + + with mock_context(task=iterable_op) as context: + context["ti"].try_number = 2 + store = context["task_state_store"] + store._data["_iterable_completed"] = {"completed": True, "try_number": 1} + for index in (0, 1): + store._data[f"_iterable_{index}"] = IndexedTaskState( + status=TaskInstanceState.SUCCESS, result="stale" + ).serialize() + + materialized = list(iterable_op.execute(context=context)) + + assert materialized == [(1, None, None), (2, None, None)] + assert IndexedTaskState.deserialize(store["_iterable_0"]).result == (1, None, None) + assert store["_iterable_completed"] == {"completed": True, "try_number": 2} + + def test_execute_after_a_crashed_rerun_does_not_replay_results_from_before_the_clear(self, monkeypatch): + """ + A task that succeeded is cleared and its rerun stops part-way. The next attempt resumes + what the rerun finished and runs the rest again: an index the rerun did not get to still + holds its checkpoint from before the clear, and replaying that one is not what clearing + the task asked for. + """ + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{"arg1": 1}, {"arg1": 2}, {"arg1": 3}]) + iterable_op = create_iterable_operator( + dag, + expand_input, + task_id="exec_crashed_rerun", + operator_class=MockAttemptOperator, + task_concurrency=1, + ) + + with mock_context(task=iterable_op) as context: + context["ti"].try_number = 1 + first_run = list(iterable_op.execute(context=context)) + assert first_run == ["1@attempt1", "2@attempt1", "3@attempt1"] + + # Cleared; the rerun finishes the first item and stops at the second. + context["ti"].try_number = 2 + monkeypatch.setattr(MockAttemptOperator, "times_out_on", {2}) + with pytest.raises(AirflowTaskTimeout): + iterable_op.execute(context=context) + + context["ti"].try_number = 3 + monkeypatch.setattr(MockAttemptOperator, "times_out_on", set()) + materialized = list(iterable_op.execute(context=context)) + + assert materialized == ["1@attempt2", "2@attempt3", "3@attempt3"] + assert context["task_state_store"]["_iterable_completed"] == { + "completed": True, + "try_number": 3, + } + + def test_execute_resumes_from_checkpoints_on_retry_after_failure(self): + """ + On a genuine retry (no completion marker) an index checkpointed as ``SUCCESS`` is skipped and + its checkpointed result is pushed again, while an index left ``UP_FOR_RETRY`` runs, and the + returned iterable yields both. + """ + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{"arg1": 1}, {"arg1": 2}]) + iterable_op = create_iterable_operator(dag, expand_input, task_id="exec_resume") + + with mock_context(task=iterable_op) as context: + context["ti"].try_number = 2 + store = context["task_state_store"] + store._data["_iterable_0"] = IndexedTaskState( + status=TaskInstanceState.SUCCESS, + result="from_checkpoint", + fingerprint=IterableOperator._fingerprint({"arg1": 1}), + ).serialize() + store._data["_iterable_1"] = IndexedTaskState( + status=TaskInstanceState.UP_FOR_RETRY + ).serialize() + + materialized = list(iterable_op.execute(context=context)) + + assert materialized == ["from_checkpoint", (2, None, None)] + assert store["_iterable_1"]["status"] == "success" + assert store["_iterable_completed"]["completed"] is True + + def test_execute_runs_an_index_again_when_its_input_changed(self): + """ + A retry can run on another input than the attempt that wrote the checkpoints: the upstream + was cleared together with this task and produced other items. The index whose item changed + runs again instead of replaying a result computed from the old item; the other is resumed. + """ + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{"arg1": "new"}, {"arg1": 2}]) + iterable_op = create_iterable_operator(dag, expand_input, task_id="exec_input_changed") + + with mock_context(task=iterable_op) as context: + context["ti"].try_number = 2 + store = context["task_state_store"] + store._data["_iterable_0"] = IndexedTaskState( + status=TaskInstanceState.SUCCESS, + result="from_the_old_item", + fingerprint=IterableOperator._fingerprint({"arg1": "old"}), + ).serialize() + store._data["_iterable_1"] = IndexedTaskState( + status=TaskInstanceState.SUCCESS, + result="from_checkpoint", + fingerprint=IterableOperator._fingerprint({"arg1": 2}), + ).serialize() + + materialized = list(iterable_op.execute(context=context)) + + assert materialized == [("new", None, None), "from_checkpoint"] + assert store["_iterable_0"]["fingerprint"] == IterableOperator._fingerprint({"arg1": "new"}) + + def test_fingerprint(self): + """The digest follows the content, not the key order, and gives up on what serde cannot serialize.""" + assert IterableOperator._fingerprint({"a": 1, "b": [1, 2]}) == IterableOperator._fingerprint( + {"b": [1, 2], "a": 1} + ) + assert IterableOperator._fingerprint({"a": 1}) != IterableOperator._fingerprint({"a": 2}) + assert IterableOperator._fingerprint({"a": object()}) is None + + def test_execute_does_not_leak_unmapped_operator_into_parent_context(self): + """ + Regression test: unmapping a sub-task must not mutate the parent context's own `ti`. + + ``context_update_for_unmapped`` sets ``context["ti"].task = task`` in place. Since + ``context.copy()`` is only a shallow copy, ``context["ti"]`` in the copy is the *same* + object as the parent's. If ``_create_task`` rendered the unmapped sub-operator against a + context still carrying the parent's `ti`, the parent's `ti.task` would end up pointing at + whichever sub-task was unmapped last, corrupting anything the runner reads off `ti.task` + after `execute()` returns (e.g. `do_xcom_push`/`multiple_outputs` in `_push_xcom_if_needed`). + """ + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{"arg1": 1}, {"arg1": 2}, {"arg1": 3}]) + iterable_op = create_iterable_operator(dag, expand_input, task_id="exec_no_leak") + + with mock_context(task=iterable_op) as context: + parent_ti_task = context["ti"].task + + list(iterable_op.execute(context=context)) + + assert context["ti"].task is parent_ti_task + assert context["task"] is iterable_op + + def test_an_item_reads_back_the_xcom_it_pushed(self): + """ + ``ti.xcom_push("progress", v)`` in an item lands under ``progress_<index>``; the item's own + ``ti.xcom_pull(key="progress")`` must read that back, not the parent's unsuffixed key. + """ + seen: list = [] + + class PushThenPull(BaseOperator): + def __init__(self, arg1=None, **kwargs): + super().__init__(**kwargs) + self.arg1 = arg1 + + def execute(self, context): + ti = context["ti"] + ti.xcom_push(key="progress", value=self.arg1 * 10) + seen.append( + ( + self.arg1, + ti.xcom_pull(key="progress"), + ti.xcom_pull(task_ids=ti.task_id, key="progress"), + ) + ) + + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{"arg1": 1}, {"arg1": 2}]) + mapped_op = PushThenPull.partial(task_id="push_pull", dag=dag, task_concurrency=1)._expand( + expand_input, strict=True, register_with_dag=False + ) + iterable_op = IterableOperator(operator=mapped_op, expand_input=expand_input, dag=dag) + with mock_context(task=iterable_op) as context: + iterable_op.execute(context=context) + + assert sorted(seen) == [(1, 10, 10), (2, 20, 20)] + + def test_an_item_sees_its_own_operator_as_the_contexts_task(self): + """ + ``context["task"]`` inside an item, and in its callbacks, is the item's unmapped operator. + + ``.expand()`` gives each mapped task instance its own operator under ``context["task"]`` + (``context_update_for_unmapped`` sets it next to ``ti.task``); the item's view of the context + must swap that key as well, not only ``ti``, or user code reads the IterableOperator there. + """ + seen: list[tuple[str, object, object]] = [] + + class MockContextTaskOperator(BaseOperator): + def execute(self, context): + seen.append(("execute", context["task"], context["ti"].task)) + + def on_success(context): + seen.append(("on_success_callback", context["task"], context["ti"].task)) + + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{}, {}]) + mapped_op = MockContextTaskOperator.partial( + task_id="context_task", dag=dag, on_success_callback=on_success + )._expand(expand_input, strict=True, register_with_dag=False) + iterable_op = IterableOperator(operator=mapped_op, expand_input=expand_input, dag=dag) + + with mock_context(task=iterable_op) as context: + iterable_op.execute(context=context) + + assert len(seen) == 4 + for where, task, ti_task in seen: + assert task is ti_task, where + assert isinstance(task, MockContextTaskOperator), where + assert task is not iterable_op, where + assert context["task"] is iterable_op + + def test_execute_resolves_and_reads_the_expand_input_on_the_running_loop(self): + """ + Regression test for the frozen IterableOperator: sub-task inputs were pulled from the main + thread between two runs of the event loop, where a blocking supervisor call deadlocked with + the ``asend`` of a sub-task parked mid-call. The input must be resolved with ``aresolve`` + and every index read while the loop runs; the synchronous ``resolve`` must stay untouched. + """ + loop_running: list[bool] = [] + original_aresolve = ListOfDictsExpandInput.aresolve + + async def aresolve(self, context): + loop_running.append(asyncio.get_running_loop().is_running()) + length, aget = await original_aresolve(self, context) + + async def aget_on_loop(index): + loop_running.append(asyncio.get_running_loop().is_running()) + return await aget(index) + + return Resolved(length, aget_on_loop) + + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{"arg1": 1}, {"arg1": 2}, {"arg1": 3}]) + iterable_op = create_iterable_operator(dag, expand_input, task_id="exec_async_input") + + with ( + mock_context(task=iterable_op) as context, + patch.object(ListOfDictsExpandInput, "aresolve", aresolve), + patch.object( + ListOfDictsExpandInput, "resolve", side_effect=AssertionError("sync resolve used") + ), + ): + materialized = sorted(iterable_op.execute(context=context)) + + assert materialized == [(1, None, None), (2, None, None), (3, None, None)] + assert loop_running == [True, True, True, True] + + def test_execute_pulls_xcom_arg_inputs_through_aresolve(self): + """An XComArg input is resolved with ``aresolve`` (``ti.axcom_pull``), never with blocking ``resolve``.""" + with DAG("test_dag") as dag: + xcom_arg = make_xcom_arg([{"arg1": 1}, {"arg1": 2}]) + xcom_arg.resolve = lambda *a, **kw: pytest.fail("synchronous resolve() used on the loop") + expand_input = ListOfDictsExpandInput(xcom_arg) + iterable_op = create_iterable_operator(dag, expand_input, task_id="exec_xcom_arg_input") + + with mock_context(task=iterable_op) as context: + materialized = sorted(iterable_op.execute(context=context)) + + assert materialized == [(1, None, None), (2, None, None)] + + def test_a_map_callable_skipping_one_element_skips_only_that_item(self): + """ + A ``.map()`` callable that raises ``AirflowSkipException`` filters out its element, as it + skips only that mapped task instance under ``.expand()``: the other items run and the task + succeeds without it. + """ + + def skip_b(value): + if value == "b": + raise AirflowSkipException("not this one") + return {"arg1": value} + + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput(make_xcom_arg(["a", "b", "c"]).map(skip_b)) + iterable_op = create_iterable_operator(dag, expand_input, task_id="map_skips_one") + + with mock_context(task=iterable_op) as context: + materialized = sorted(iterable_op.execute(context=context)) + + assert materialized == [("a", None, None), ("c", None, None)] + + def test_a_map_callable_failing_on_one_element_fails_only_that_item(self): + """ + A ``.map()`` callable that raises for one element fails that item alone: the others run and + are checkpointed, so a retry replays them and only that element is read again. + """ + + def fail_on_b(value): + if value == "b": + raise ValueError("cannot map b") + return {"arg1": value} + + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput(make_xcom_arg(["a", "b", "c"]).map(fail_on_b)) + iterable_op = create_iterable_operator( + dag, expand_input, task_id="map_fails_one", task_concurrency=1 + ) + + with mock_context(task=iterable_op) as context: + store = context["task_state_store"] + with pytest.raises(ValueError, match="cannot map b"): + iterable_op.execute(context=context) + + assert store["_iterable_0"]["status"] == "success" + assert "_iterable_1" not in store + assert store["_iterable_2"]["status"] == "success" + + def test_a_template_failing_for_one_item_fails_only_that_item(self): + """Rendering is part of reading an item's input: a template that fails for one item fails it alone.""" + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput( + [{"arg1": "a"}, {"arg1": "{{ no_such_name }}"}, {"arg1": "c"}] + ) + iterable_op = create_iterable_operator( + dag, expand_input, task_id="render_fails_one", task_concurrency=1 + ) + + with mock_context(task=iterable_op) as context: + store = context["task_state_store"] + with pytest.raises(jinja2.UndefinedError): + iterable_op.execute(context=context) + + assert store["_iterable_0"]["status"] == "success" + assert "_iterable_1" not in store + assert store["_iterable_2"]["status"] == "success" + + def test_execute_refuses_an_unmappable_upstream_value(self): + """ + An upstream returning a JSON string fails the task instead of being iterated as one item. + + ``.expand()`` gets this from the upstream's push (``UnmappableXComTypePushed`` in + ``_push_xcom_if_needed``), which only fires for a ``MappedOperator`` dependant, so the + iterated path checks the resolved value itself. + """ + executed: list = [] + + class RecordingOperator(MockOperator): + def execute(self, context): + executed.append(self.arg1) + return super().execute(context) + + with DAG("test_dag") as dag: + expand_input = DictOfListsExpandInput({"arg1": make_xcom_arg('{"a": 1}')}) + iterable_op = create_iterable_operator( + dag, expand_input, task_id="exec_unmappable", operator_class=RecordingOperator + ) + with mock_context(task=iterable_op) as context: + with pytest.raises(UnmappableXComTypePushed, match="str"): + iterable_op.execute(context=context) + + assert executed == [] + + def test_a_template_reads_the_indexed_tasks_own_state_store(self): + """A template field reads the same suffixed store as ``execute`` does, not the parent's key.""" + template = "{{ task_state_store.get('offset') }}" + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{"offset": template}, {"offset": template}]) + iterable_op = create_iterable_operator( + dag, expand_input, task_id="template_store", operator_class=MockStateStoreOperator + ) + + with mock_context(task=iterable_op) as context: + context["task_state_store"].set("offset", "parents") + context["task_state_store"].set("offset_1", 7) + task = iterable_op._create_task( + context=context, + index=1, + mapped_kwargs={"offset": template}, + jinja_env=iterable_op.get_template_env(dag=dag), + ) + + assert task.task.offset == "7" + + def test_a_template_reads_the_items_own_index(self): + """ + ``{{ ti.index }}`` in a partial kwarg renders per item, so an operator that finds its remote + work by labels (a ``KubernetesPodOperator`` that reattaches) can be told the item apart. + """ + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{"arg1": "a"}, {"arg1": "b"}]) + mapped_op = MockOperator.partial(task_id="indexed_label", dag=dag, arg2="{{ ti.index }}")._expand( + expand_input, strict=True, register_with_dag=False + ) + iterable_op = IterableOperator(operator=mapped_op, expand_input=expand_input, dag=dag) + + with mock_context(task=iterable_op) as context: + results = sorted(iterable_op.execute(context=context)) + + assert results == [("a", "0", None), ("b", "1", None)] + + def test_execute_renders_template_fields_off_the_loop_thread(self): + """ + Rendering a sub-task's template fields may call the supervisor synchronously: an XComArg in + a partial kwarg resolves with ``resolve``, ``{{ var.value.x }}`` with ``Variable.get``. On + the loop thread that call raises ``DeadlockImminentError`` whenever a sibling's ``asend`` is + in flight, so the rendering has to happen in a worker thread. + """ + on_running_loop: list[bool] = [] + + def resolve(context): + try: + asyncio.get_running_loop() + except RuntimeError: + on_running_loop.append(False) + else: + on_running_loop.append(True) + return "pulled" + + with DAG("test_dag") as dag: + xcom_arg = make_xcom_arg(None) + xcom_arg.resolve = resolve + expand_input = ListOfDictsExpandInput([{"arg1": 1}, {"arg1": 2}]) + mapped_op = MockOperator.partial(task_id="render_off_loop", dag=dag, arg2=xcom_arg)._expand( + expand_input, strict=True, register_with_dag=False + ) + iterable_op = IterableOperator(operator=mapped_op, expand_input=expand_input, dag=dag) + + with mock_context(task=iterable_op) as context: + materialized = sorted(iterable_op.execute(context=context)) + + assert materialized == [(1, "pulled", None), (2, "pulled", None)] + assert on_running_loop == [False, False] + + def test_execute_dict_of_lists(self): + """Test executing IterableOperator with DictOfListsExpandInput.""" + with DAG("test_dag") as dag: + expand_input = DictOfListsExpandInput({"arg1": [1, 2, 3]}) + iterable_op = create_iterable_operator(dag, expand_input, task_id="exec_dict_of_lists") + + with mock_context(task=iterable_op) as context: + result = iterable_op.execute(context=context) + materialized = list(result) + assert materialized == [(1, None, None), (2, None, None), (3, None, None)] + + def test_execute_multiple_key_dict_of_lists(self): + """Test executing IterableOperator with multiple keys in DictOfListsExpandInput.""" + with DAG("test_dag") as dag: + expand_input = DictOfListsExpandInput({"arg1": [1, 2], "arg2": [10, 20], "arg3": ["x", "y"]}) + iterable_op = create_iterable_operator(dag, expand_input, task_id="exec_multi_key") + + with mock_context(task=iterable_op) as context: + result = iterable_op.execute(context=context) + materialized = list(result) + # Cartesian product expected order: + # (1,10,'x'), (1,10,'y'), (1,20,'x'), (1,20,'y'), + # (2,10,'x'), (2,10,'y'), (2,20,'x'), (2,20,'y') + assert materialized == [ + (1, 10, "x"), + (1, 10, "y"), + (1, 20, "x"), + (1, 20, "y"), + (2, 10, "x"), + (2, 10, "y"), + (2, 20, "x"), + (2, 20, "y"), + ] + + def test_execute_with_task_concurrency_setting(self): + """Test executing IterableOperator with task_concurrency parameter.""" + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{"arg1": 1}, {"arg1": 2}, {"arg1": 3}]) + iterable_op = create_iterable_operator( + dag, expand_input, task_id="exec_concurrency", task_concurrency=2 + ) + + with mock_context(task=iterable_op) as context: + result = iterable_op.execute(context=context) + materialized = list(result) + assert materialized == [(1, None, None), (2, None, None), (3, None, None)] + assert iterable_op.max_workers == 2 + + def test_execute_all_parameters(self): + """Test executing IterableOperator with all arg1, arg2, arg3 parameters.""" + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput( + [ + {"arg1": 1, "arg2": 10, "arg3": 100}, + {"arg1": 2, "arg2": 20, "arg3": 200}, + ] + ) + iterable_op = create_iterable_operator(dag, expand_input, task_id="exec_all_args") + + with mock_context(task=iterable_op) as context: + result = iterable_op.execute(context=context) + materialized = list(result) + assert materialized == [(1, 10, 100), (2, 20, 200)] + + def test_execute_with_do_xcom_push_false(self): + """With do_xcom_push=False no return_value_<index> XCom is pushed for any sub-task.""" + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{"arg1": 1}, {"arg1": 2}]) + iterable_op = create_iterable_operator( + dag, expand_input, task_id="no_xcom_push", do_xcom_push=False + ) + + with ( + mock_context(task=iterable_op) as context, + patch.object( + IterableOperator, "axcom_push", new=AsyncMock(spec=IterableOperator.axcom_push) + ) as axcom_push, + ): + result = iterable_op.execute(context=context) + + assert result is None + axcom_push.assert_not_awaited() + + def test_execute_does_not_push_xcom_for_none_results(self): + """A sub-task returning None has nothing to push, even when do_xcom_push is True.""" + + class NoneOperator(BaseOperator): + def __init__(self, arg1=None, **kwargs): + super().__init__(**kwargs) + self.arg1 = arg1 + + def execute(self, context): + return None + + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{"arg1": 1}, {"arg1": 2}]) + iterable_op = create_iterable_operator( + dag, expand_input, task_id="none_results", operator_class=NoneOperator + ) + + with ( + mock_context(task=iterable_op) as context, + patch.object( + IterableOperator, "axcom_push", new=AsyncMock(spec=IterableOperator.axcom_push) + ) as axcom_push, + ): + iterable_op.execute(context=context) + + axcom_push.assert_not_awaited() + + def test_execute_with_failed_tasks_raises_regardless_of_retries(self): + """ + Test executing IterableOperator where a sub-task fails. + + This test verifies that: + 1. Tasks with fail_on_first_attempt=True raise an exception on first attempt + 2. IterableOperator no longer retries failed sub-tasks in-process — retries (if any) are + handled by Airflow retrying the whole IterableOperator task instance + 3. The failing sub-task's own exception is raised, regardless of whether the wrapped operator + has retries configured (one failure is not wrapped in a group, see _failure_for_the_runner) + """ + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput( + [ + {"arg1": 1, "arg2": 10}, + {"arg1": 2, "arg2": 20, "fail_on_first_attempt": True}, + {"arg1": 3, "arg2": 30}, + ] + ) + iterable_op = create_iterable_operator( + dag, + expand_input, + task_id="exec_with_failures", + retries=1, + ) + + with mock_context(task=iterable_op) as context: + with pytest.raises(RuntimeError): + iterable_op.execute(context=context) + + def test_execute_all_sub_tasks_skipped_raises_single_skip_exception(self): + """When every sub-task raises AirflowSkipException, IterableOperator must re-raise a single + AirflowSkipException so the runner marks it SKIPPED instead of a retryable failure.""" + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput( + [ + {"raise_exception": AirflowSkipException("skip 1")}, + {"raise_exception": AirflowSkipException("skip 2")}, + ] + ) + iterable_op = create_iterable_operator(dag, expand_input, task_id="all_skipped") + + with mock_context(task=iterable_op) as context: + with pytest.raises(AirflowSkipException): + iterable_op.execute(context=context) + + def test_execute_over_an_empty_input_skips(self): + """An empty input skips the task, as ``.expand()`` over nothing does, rather than returning nothing.""" + with DAG("test_dag") as dag: + iterable_op = create_iterable_operator(dag, ListOfDictsExpandInput([]), task_id="empty_input") + + with mock_context(task=iterable_op) as context: + with pytest.raises(AirflowSkipException, match="empty"): + iterable_op.execute(context=context) + + def test_execute_skip_next_to_a_failure_raises_only_the_failure(self): + """A skipped sub-task is not a failure, so the group holds what failed and nothing else.""" + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput( + [ + {"raise_exception": AirflowSkipException("skip 1")}, + {"raise_exception": RuntimeError("boom")}, + ] + ) + iterable_op = create_iterable_operator(dag, expand_input, task_id="partial_skip") + + with mock_context(task=iterable_op) as context: + with pytest.raises(RuntimeError, match="boom") as raised: + iterable_op.execute(context=context) + + # A skipped sub-task is not a failure: the only failure is raised on its own. + assert raised.value.__cause__ is None + + def test_execute_skip_next_to_successes_succeeds(self): + """ + A sub-task that skips does not fail the task, as a skipped mapped task instance would not: + the others' results are pushed, the skipped index is left out, and the run counts as complete. + """ + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput( + [{"raise_exception": AirflowSkipException("nothing to do")}, {"arg1": 2}] + ) + iterable_op = create_iterable_operator(dag, expand_input, task_id="skip_and_success") + + with mock_context(task=iterable_op) as context: + materialized = list(iterable_op.execute(context=context)) + store = context["task_state_store"] + + assert materialized == [(2, None, None)] + assert store["_iterable_0"]["status"] == "skipped" + assert store["_iterable_1"]["status"] == "success" + assert store["_iterable_completed"]["completed"] is True + + def test_an_iteration_returning_none_keeps_its_position_and_reads_as_none(self): + """ + Only a skip takes an index out of the sequence. An iteration that returned ``None`` pushed + nothing, as a mapped task instance would, but is counted and reads as ``None``; ``.expand()`` + would not count it, which the docs page says. + """ + + class ValueOrNothing(BaseOperator): + def __init__(self, value=None, skip=False, **kwargs): + super().__init__(**kwargs) + self.value = value + self.skip = skip + + def execute(self, context): + if self.skip: + raise AirflowSkipException("nothing to do") + return self.value + + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput( + [{"value": 1}, {"value": None}, {"skip": True}, {"value": 4}] + ) + iterable_op = create_iterable_operator( + dag, expand_input, task_id="none_and_skip", operator_class=ValueOrNothing + ) + + with mock_context(task=iterable_op) as context: + result = iterable_op.execute(context=context) + store = context["task_state_store"] + + assert result.skipped == [2] + assert len(result) == 3 + assert list(result) == [1, None, 4] + assert store["_iterable_1"]["status"] == "success" + assert "result" not in store["_iterable_1"] + + def test_skipped_iteration_is_left_out_of_the_result(self): + """As a skipped mapped task instance, a skipped iteration is not among the values downstream reads.""" + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput( + [{"arg1": 1}, {"raise_exception": AirflowSkipException("nothing to do")}, {"arg1": 3}] + ) + iterable_op = create_iterable_operator(dag, expand_input, task_id="skip_in_the_middle") + + with mock_context(task=iterable_op) as context: + result = iterable_op.execute(context=context) + + assert result.skipped == [1] + assert len(result) == 2 + assert [value[0] for value in result] == [1, 3] + + @pytest.mark.parametrize( + ("trigger_rule", "skipped"), + [ + ("all_success", True), + ("none_skipped", True), + ("all_done_min_one_success", True), + ("none_failed", False), + ("all_done", False), + ("one_success", False), + ], + ) + def test_partial_skip_skips_the_downstream_tasks_a_skipped_mapped_instance_would( + self, trigger_rule, skipped + ): + """ + A downstream task whose trigger rule skips it when an upstream task instance skipped is skipped, + after the result is pushed for the others; any other downstream task is left to run. + """ + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput( + [{"arg1": 1}, {"raise_exception": AirflowSkipException("nothing to do")}] + ) + iterable_op = create_iterable_operator(dag, expand_input, task_id="partial_skip") + iterable_op >> BaseOperator(task_id="downstream", trigger_rule=trigger_rule) + + with ( + mock_context(task=iterable_op) as context, + patch("airflow.sdk.definitions.iterableoperator._push_xcom_if_needed", autospec=True) as push, + ): + if skipped: + with pytest.raises(DownstreamTasksSkipped) as raised: + iterable_op.execute(context=context) + assert raised.value.tasks == ["downstream"] + (pushed, ti, _), _ = push.call_args + assert ti is context["ti"] + assert [value[0] for value in pushed] == [1] + else: + assert len(iterable_op.execute(context=context)) == 1 + push.assert_not_called() + + def test_no_downstream_task_is_skipped_without_a_skipped_iteration(self): + with DAG("test_dag") as dag: + iterable_op = create_iterable_operator( + dag, ListOfDictsExpandInput([{"arg1": 1}, {"arg1": 2}]), task_id="no_skip" + ) + iterable_op >> BaseOperator(task_id="downstream") + + with mock_context(task=iterable_op) as context: + assert len(iterable_op.execute(context=context)) == 2 + + def test_clear_after_every_iteration_skipped_runs_every_iteration_again(self): + """ + Regression test: a task whose iterations all skipped is ``SKIPPED``, a final state. Clearing it + must run the iterations again, as clearing a skipped mapped task instance does, instead of + replaying the skips recorded before the clear. + """ + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput( + [ + {"raise_exception": AirflowSkipException("nothing to do yet")}, + {"raise_exception": AirflowSkipException("nothing to do yet")}, + ] + ) + iterable_op = create_iterable_operator(dag, expand_input, task_id="all_skipped_then_cleared") + + with mock_context(task=iterable_op) as context: + context["ti"].try_number = 1 + with pytest.raises(AirflowSkipException): + iterable_op.execute(context=context) + + # Cleared: the next attempt runs on the same try_number sequence, now with work to do. + context["ti"].try_number = 2 + with patch.object(MockOperator, "execute", autospec=True, return_value="done") as execute: + result = iterable_op.execute(context=context) + + assert execute.call_count == 2 + assert len(result) == 2 + + @pytest.mark.asyncio + async def test_run_task_does_not_rerun_a_sub_task_skipped_on_a_previous_attempt(self): + """On a retry a ``SKIPPED`` checkpoint is honoured like a ``SUCCESS`` one: the sub-task stays skipped.""" + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{"arg1": 1}]) + iterable_op = create_iterable_operator(dag, expand_input, task_id="skipped_before") + + with mock_context(task=iterable_op) as context: + task = iterable_op._create_task( + context=context, + index=0, + mapped_kwargs={"raise_exception": RuntimeError("must not run again")}, + jinja_env=iterable_op.get_template_env(dag=dag), + ) + await context["task_state_store"].aset( + task.state_key, + IndexedTaskState( + status=TaskInstanceState.SKIPPED, fingerprint=task.input_fingerprint + ).serialize(), + ) + + with AsyncAwareExecutor(loop=asyncio.get_running_loop(), max_workers=1) as executor: + _, result, raised = await iterable_op._run_task( + executor, + context, + task, + trust_checkpoints=True, + outcomes=IndexedTaskOutcomes(iterable_op, IterationState(), context), + ) + + assert result is None + assert isinstance(raised, AirflowSkipException) + + def test_failed_iteration_checkpoint_records_its_input_and_attempt(self): + """ + A failed iteration's checkpoint carries the input digest and the attempt, like a successful + one, so the next attempt does not take a plain retry for a clear or a changed input. + """ + with DAG("test_dag") as dag: + failing = {"arg1": 1, "fail_on_first_attempt": True} # serializable, so it has a digest + iterable_op = create_iterable_operator( + dag, ListOfDictsExpandInput([failing]), task_id="failed_checkpoint" + ) + + with mock_context(task=iterable_op) as context: + context["ti"].try_number = 3 + with pytest.raises(RuntimeError): + iterable_op.execute(context=context) + + checkpoint = context["task_state_store"]["_iterable_0"] + + assert checkpoint["status"] == "up_for_retry" + assert IterableOperator._fingerprint(failing) is not None + assert checkpoint["fingerprint"] == IterableOperator._fingerprint(failing) + assert checkpoint["try_number"] == 3 + + def test_parent_timeout_runs_the_retry_callback_while_retries_are_left(self): + """The task is retried for the timeout, so its iteration reports a retry, not a failure.""" + FIRED_CALLBACKS.clear() + with DAG("test_dag") as dag: + iterable_op = create_iterable_operator( + dag, + ListOfDictsExpandInput([{}]), + task_id="timed_out", + retries=3, + operator_class=MockCallbackTimeoutOperator, + ) + + with mock_context(task=iterable_op) as context: + context["ti"].try_number = 1 + context["ti"].max_tries = 3 + with pytest.raises(AirflowTaskTimeout): + iterable_op.execute(context=context) + + assert FIRED_CALLBACKS == ["retry"] + + def test_a_sync_items_own_timeout_is_its_failure(self): + """ + A sync operator may raise ``AirflowTaskTimeout`` itself, as a hook that gave up waiting does. + In its worker thread the parent's limit never strikes, so that is the item's own failure, as + it is the mapped task instance's under ``.expand()``: the sibling runs, nothing is killed, and + the task fails with it as with any failed item. + """ + killed: list = [] + + class GivingUpOperator(MockOperator): + def on_kill(self): + killed.append(self.arg1) + + with DAG("test_dag") as dag: + iterable_op = create_iterable_operator( + dag, + ListOfDictsExpandInput( + [{"arg1": 1, "raise_exception": AirflowTaskTimeout("gave up waiting")}, {"arg1": 2}] + ), + task_id="own_timeout", + task_concurrency=1, + operator_class=GivingUpOperator, + ) + + with mock_context(task=iterable_op) as context: + store = context["task_state_store"] + with pytest.raises(AirflowTaskTimeout, match="gave up waiting"): + iterable_op.execute(context=context) + + assert store["_iterable_0"]["status"] == "up_for_retry" + assert store["_iterable_1"]["status"] == "success" + assert killed == [] + + def test_an_own_timeout_next_to_another_failure_reports_a_retry(self): + """ + An item's own ``AirflowTaskTimeout`` and another item's ``ValueError`` reach the runner as a + ``BaseExceptionGroup``, which it retries like any failure while retries are left. Both + items' callbacks report that retry, not a failure. + """ + FIRED_CALLBACKS.clear() + + class RecordingOperator(MockOperator): + def __init__(self, **kwargs): + kwargs["on_failure_callback"] = lambda context: FIRED_CALLBACKS.append(("failure", self.arg1)) + kwargs["on_retry_callback"] = lambda context: FIRED_CALLBACKS.append(("retry", self.arg1)) + super().__init__(**kwargs) + + with DAG("test_dag") as dag: + iterable_op = create_iterable_operator( + dag, + ListOfDictsExpandInput( + [ + {"arg1": 1, "raise_exception": AirflowTaskTimeout("gave up waiting")}, + {"arg1": 2, "raise_exception": ValueError("bad item")}, + ] + ), + task_id="timeout_and_failure", + retries=1, + operator_class=RecordingOperator, + ) + + with mock_context(task=iterable_op) as context: + context["ti"].try_number = 1 + context["ti"].max_tries = 1 + with pytest.raises(BaseExceptionGroup): + iterable_op.execute(context=context) + + assert sorted(FIRED_CALLBACKS) == [("retry", 1), ("retry", 2)] + + def test_iteration_cancelled_by_a_sibling_runs_no_callback(self): + """One iteration stops the task; the sibling cancelled on the way did not fail.""" + FIRED_CALLBACKS.clear() + with DAG("test_dag") as dag: + iterable_op = create_iterable_operator( + dag, + ListOfDictsExpandInput([{"arg1": "defer"}, {"arg1": "sleeper"}]), + task_id="cancelled_sibling", + retries=3, + task_concurrency=2, + operator_class=MockCallbackAsyncOperator, + ) + + with mock_context(task=iterable_op) as context: + context["ti"].try_number = 1 + context["ti"].max_tries = 3 + with pytest.raises(AirflowFailException, match="attempted to defer"): + iterable_op.execute(context=context) + + assert ("failure", "sleeper") not in FIRED_CALLBACKS + assert ("retry", "sleeper") not in FIRED_CALLBACKS + + def test_failed_publish_keeps_the_success_checkpoint_and_the_retry_only_republishes(self): + """ + When pushing an item's result fails after its SUCCESS checkpoint was written, the checkpoint + stays, and the retry replays the result instead of running the operator again, which would + repeat whatever it did outside Airflow. + """ + runs = [] + original_execute = MockOperator.execute + + def counting_execute(self, context): + runs.append(self.arg1) + return original_execute(self, context) + + with DAG("test_dag") as dag: + iterable_op = create_iterable_operator( + dag, ListOfDictsExpandInput([{"arg1": 1}]), task_id="failed_publish" + ) + + with ( + mock_context(task=iterable_op) as context, + patch.object(MockOperator, "execute", counting_execute), + ): + context["ti"].try_number = 1 + with patch.object( + IterableOperator, "axcom_push", side_effect=RuntimeError("xcom backend down") + ): + with pytest.raises(RuntimeError, match="xcom backend down"): + iterable_op.execute(context=context) + checkpoint_after_failure = context["task_state_store"]["_iterable_0"]["status"] + + context["ti"].try_number = 2 + result = iterable_op.execute(context=context) + pushed = list(result) + + assert checkpoint_after_failure == "success" + assert runs == [1] + assert pushed == [(1, None, None)] + + def test_the_success_callback_waits_for_the_checkpoint(self): + """ + An item's success callback fires once its SUCCESS checkpoint is written, not when execute returns. + + With the checkpoint write failing, the item announces nothing: the attempt fails, the retry + runs the item again and reports it then, once. A plain task whose result cannot be pushed + fires no success callback either. + """ + CALLBACKS.clear() + runs = [] + original_execute = MockCallbackOperator.execute + + def counting_execute(self, context): + runs.append(self.arg1) + return original_execute(self, context) + + with DAG("test_dag") as dag: + iterable_op = create_iterable_operator( + dag, + ListOfDictsExpandInput([{"arg1": "a"}]), + task_id="checkpoint_fails", + retries=2, + operator_class=MockCallbackOperator, + ) + + with ( + mock_context(task=iterable_op) as context, + patch.object(MockCallbackOperator, "execute", counting_execute), + ): + store = context["task_state_store"] + original_aset = store.aset + failed: list[str] = [] + + async def aset_failing_the_first_success(key, value, **kwargs): + if value.get("status") == "success" and not failed: + failed.append(key) + raise RuntimeError("state store down") + await original_aset(key, value, **kwargs) + + store.aset = aset_failing_the_first_success + context["ti"].try_number = 1 + with pytest.raises(RuntimeError, match="state store down"): + iterable_op.execute(context=context) + fired_after_the_failed_write = list(CALLBACKS) + + context["ti"].try_number = 2 + iterable_op.execute(context=context) + + assert fired_after_the_failed_write == [] + assert runs == ["a", "a"] + assert CALLBACKS == [("success", "a")] + + def test_a_failed_publish_fires_the_success_callback_once(self): + """ + The checkpoint is written before the result is pushed, so a push that fails leaves work the + retry replays rather than runs again: the success callback fired with the checkpoint and + does not fire again on the replay. + """ + CALLBACKS.clear() + with DAG("test_dag") as dag: + iterable_op = create_iterable_operator( + dag, + ListOfDictsExpandInput([{"arg1": "a"}]), + task_id="publish_fails", + retries=2, + operator_class=MockCallbackOperator, + ) + + with mock_context(task=iterable_op) as context: + context["ti"].try_number = 1 + with patch.object( + IterableOperator, "axcom_push", side_effect=RuntimeError("xcom backend down") + ): + with pytest.raises(RuntimeError, match="xcom backend down"): + iterable_op.execute(context=context) + fired_after_the_failed_push = list(CALLBACKS) + + context["ti"].try_number = 2 + result = iterable_op.execute(context=context) + pushed = list(result) + + assert fired_after_the_failed_push == [("success", "a")] + assert CALLBACKS == [("success", "a")] + assert pushed == ["a"] + + def test_a_success_callback_that_raises_keeps_the_checkpoint(self): + """ + The success callback runs after the checkpoint is written and may raise what nothing catches: + a ``DeadlockImminentError`` from a synchronous SDK call in an async item's callback. The item's + work is done, so its SUCCESS checkpoint stays and the retry replays the result instead of + running the item again; the callback fired once. + """ + runs: list[str] = [] + fired: list[str] = [] + + class CallbackRaisesOperator(BaseAsyncOperator): + def __init__(self, arg1=None, **kwargs): + kwargs["on_success_callback"] = self._callback + super().__init__(**kwargs) + self.arg1 = arg1 + + def _callback(self, context): + fired.append(self.arg1) + raise DeadlockImminentError("Variable.get on the event loop thread") + + async def aexecute(self, context): + runs.append(self.arg1) + return self.arg1 + + with DAG("test_dag") as dag: + iterable_op = create_iterable_operator( + dag, + ListOfDictsExpandInput([{"arg1": "a"}]), + task_id="callback_raises", + retries=2, + operator_class=CallbackRaisesOperator, + ) + + with mock_context(task=iterable_op) as context: + store = context["task_state_store"] + context["ti"].try_number = 1 + with pytest.raises(AirflowFailException, match="synchronous SDK call"): + iterable_op.execute(context=context) + status_after_the_failed_callback = store["_iterable_0"]["status"] + + context["ti"].try_number = 2 + result = iterable_op.execute(context=context) + pushed = list(result) + + assert status_after_the_failed_callback == "success" + assert runs == ["a"] + assert fired == ["a"] + assert pushed == ["a"] + + def test_extra_xcoms_are_checkpointed_once_and_pushed_again_when_a_retry_skips_the_item(self): + """ + The runner deletes every XCom before a retry. An item skipped because it already succeeded + gets its other pushed keys back from its checkpoint, written once with its SUCCESS state. + """ + MockPushingOperator.executed = [] + with DAG("test_dag") as dag: + iterable_op = create_iterable_operator( + dag, + ListOfDictsExpandInput([{"arg1": "a"}, {"arg1": "b", "fail": True}]), + task_id="extra_xcoms", + task_concurrency=1, + operator_class=MockPushingOperator, + ) + + with mock_context(task=iterable_op) as context: + store = context["task_state_store"] + context["ti"].try_number = 1 + with pytest.raises(RuntimeError, match="sibling failed"): + iterable_op.execute(context=context) + checkpoint = store["_iterable_0"] + + # The retry: the failing item succeeds now; the finished one must not run again. + pushed = [] + + async def recording_aset(cls, key, value, **kwargs): + pushed.append((key, value)) + + iterable_op.expand_input = ListOfDictsExpandInput([{"arg1": "a"}, {"arg1": "b"}]) + context["ti"].try_number = 2 + with patch.object(XCom, "aset", classmethod(recording_aset)): + iterable_op.execute(context=context) + + assert checkpoint["status"] == "success" + assert checkpoint["xcoms"] == {"foo": "foo-of-a"} + assert "return_value" not in checkpoint["xcoms"] + assert MockPushingOperator.executed == ["a", "b", "b"] + assert ("foo_0", "foo-of-a") in pushed + assert ("return_value_0", "a") in pushed + + def test_an_item_pushing_many_keys_writes_one_checkpoint(self): + """The extra pushes are kept in memory and written with the one SUCCESS checkpoint, not per push.""" + + class ManyKeysOperator(BaseOperator): + template_fields = () + + def execute(self, context): + for key in range(5): + context["ti"].xcom_push(key=f"key{key}", value=key) + return "done" + + with DAG("test_dag") as dag: + iterable_op = create_iterable_operator( + dag, ListOfDictsExpandInput([{}]), task_id="many_keys", operator_class=ManyKeysOperator + ) + + with mock_context(task=iterable_op) as context: + store = context["task_state_store"] + writes = [] + original_aset = store.aset + + async def counting_aset(key, value, **kwargs): + writes.append(key) + await original_aset(key, value, **kwargs) + + store.aset = counting_aset + iterable_op.execute(context=context) + + assert writes == ["_iterable_0"] + assert store["_iterable_0"]["xcoms"] == {f"key{key}": key for key in range(5)} + + def test_an_async_items_extra_xcoms_are_pushed_again_on_retry(self): + """The async path pushes through axcom_push; its keys are recorded and replayed the same way.""" + + class AsyncPushingOperator(BaseAsyncOperator): + template_fields = ("arg1",) + executed: list = [] + + def __init__(self, arg1=None, fail=False, **kwargs): + super().__init__(**kwargs) + self.arg1 = arg1 + self.fail = fail + + async def aexecute(self, context): + type(self).executed.append(self.arg1) + await context["ti"].axcom_push(key="bar", value=f"bar-of-{self.arg1}") + if self.fail: + raise RuntimeError("sibling failed") + return self.arg1 + + with DAG("test_dag") as dag: + iterable_op = create_iterable_operator( + dag, + ListOfDictsExpandInput([{"arg1": "a"}, {"arg1": "b", "fail": True}]), + task_id="async_extra_xcoms", + task_concurrency=1, + operator_class=AsyncPushingOperator, + ) + + with mock_context(task=iterable_op) as context: + context["ti"].try_number = 1 + with pytest.raises(RuntimeError, match="sibling failed"): + iterable_op.execute(context=context) + + pushed = [] + + async def recording_aset(cls, key, value, **kwargs): + pushed.append((key, value)) + + iterable_op.expand_input = ListOfDictsExpandInput([{"arg1": "a"}, {"arg1": "b"}]) + context["ti"].try_number = 2 + with patch.object(XCom, "aset", classmethod(recording_aset)): + iterable_op.execute(context=context) + + assert AsyncPushingOperator.executed == ["a", "b", "b"] + assert ("bar_0", "bar-of-a") in pushed + + def test_execute_failed_attempt_leaves_no_completion_marker_so_retry_resumes(self): + """ + Regression test: an attempt that fails must not write the completion marker. + + The marker tells a retry apart from a rerun after a manual clear: when it is present, every + checkpoint is ignored and every index runs again. Raising the collected sub-task failures + outside the ``Checkpoints`` block would let the block exit cleanly and write the marker on a + failed attempt, so the retry would re-run the succeeded indices too instead of resuming. + """ + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{"arg1": 1}, {"raise_exception": ValueError("boom")}]) + iterable_op = create_iterable_operator(dag, expand_input, task_id="exec_failed_no_marker") + + with mock_context(task=iterable_op) as context: + store = context["task_state_store"] + + with pytest.raises(ValueError, match="boom"): + iterable_op.execute(context=context) + + assert "_iterable_completed" not in store + assert store["_iterable_0"]["status"] == "success" + assert store["_iterable_1"]["status"] == "up_for_retry" + + # The retry: the failing item succeeds now; index 0 must be replayed, not re-run. + context["ti"].try_number = 2 + iterable_op.expand_input = ListOfDictsExpandInput([{"arg1": 1}, {"arg1": 2}]) + with patch.object( + IndexedTaskRunner, "run", autospec=True, side_effect=IndexedTaskRunner.run + ) as run: + materialized = list(iterable_op.execute(context=context)) + + assert materialized == [(1, None, None), (2, None, None)] + assert run.call_count == 1 + assert run.call_args.args[0].task_index == 1 # only the failed index runs again + assert store["_iterable_completed"] == {"completed": True, "try_number": 2} + + def test_execute_fail_exception_re_raised_directly_without_retry(self): + """A sub-task that raises AirflowFailException must be re-raised directly (not wrapped in a + BaseExceptionGroup) so the IterableOperator fails without being retried.""" + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput( + [ + {"arg1": 1}, + {"raise_exception": AirflowFailException("boom")}, + ] + ) + iterable_op = create_iterable_operator(dag, expand_input, task_id="fail_exception", retries=3) + + with mock_context(task=iterable_op) as context: + with pytest.raises(AirflowFailException, match="boom"): + iterable_op.execute(context=context) + + @pytest.mark.parametrize( + "raised_exception", + [ + DagRunTriggerException( + trigger_dag_id="triggered_dag", + dag_run_id="triggered_run", + conf={}, + reset_dag_run=False, + skip_when_already_exists=False, + wait_for_completion=False, + allowed_states=["success"], + failed_states=["failed"], + poke_interval=1, + deferrable=False, + ), + DownstreamTasksSkipped(tasks=["downstream_task"]), + ], + ids=["DagRunTriggerException", "DownstreamTasksSkipped"], + ) + def test_execute_rejects_trigger_and_downstream_skip_exceptions(self, raised_exception): + """TriggerDagRunOperator (DagRunTriggerException) and downstream-skip operators like + ShortCircuitOperator (DownstreamTasksSkipped) are not supported inside IterableOperator: a + sub-task index has no DAG run or downstream tasks of its own for the trigger/skip to apply + to, so this must fail the whole IterableOperator immediately with a clear error.""" + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{"raise_exception": raised_exception}]) + iterable_op = create_iterable_operator(dag, expand_input, task_id="unsupported_exception") + + with mock_context(task=iterable_op) as context: + with pytest.raises(AirflowFailException, match="not supported inside IterableOperator"): + iterable_op.execute(context=context) + + def test_execute_rejects_reschedule_exception(self): + """A reschedule-mode sensor (AirflowRescheduleException) is not supported inside + IterableOperator: the sub-task's index has no task instance of its own to reschedule, so this + must fail the whole IterableOperator immediately with a clear error rather than being silently + aggregated into a retryable BaseExceptionGroup.""" + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput( + [{"raise_exception": AirflowRescheduleException(reschedule_date=None)}] + ) + iterable_op = create_iterable_operator(dag, expand_input, task_id="reschedule_exception") + + with mock_context(task=iterable_op) as context: + with pytest.raises(AirflowFailException, match="not supported inside IterableOperator"): + iterable_op.execute(context=context) + + @pytest.mark.asyncio + async def test_run_task_skips_sub_task_already_checkpointed_as_succeeded(self): + """ + When the task_state_store already records a sub-task index as succeeded (e.g. because + Airflow retried the whole IterableOperator after a previous partial failure), ``_run_task`` + must skip re-executing that sub-task entirely. + """ + from unittest import mock + + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{"arg1": 1, "arg2": 10}]) + iterable_op = create_iterable_operator(dag, expand_input, task_id="checkpoint_skip") + + with mock_context(task=iterable_op) as context: + jinja_env = iterable_op.get_template_env(dag=dag) + task = iterable_op._create_task( + context=context, index=0, mapped_kwargs={"arg1": 1, "arg2": 10}, jinja_env=jinja_env + ) + task.try_number = 2 # checkpoint is only consulted from the second attempt onwards + await context["task_state_store"].aset( + task.state_key, + IndexedTaskState( + status=TaskInstanceState.SUCCESS, fingerprint=task.input_fingerprint + ).serialize(), + ) + + executor = mock.MagicMock() + _, result, raised = await iterable_op._run_task( + executor, + context, + task, + trust_checkpoints=True, + outcomes=IndexedTaskOutcomes(iterable_op, IterationState(), context), + ) + + assert result is None + assert raised is None + executor.run_sync.assert_not_called() + + @pytest.mark.asyncio + async def test_run_task_reruns_and_checkpoints_success_after_up_for_retry_state(self): + """ + A sub-task whose checkpoint records ``UP_FOR_RETRY`` (left behind by a previous failed or + crashed attempt) is re-executed rather than skipped — only a ``SUCCESS`` checkpoint causes + ``_run_task`` to skip re-execution — and a new ``SUCCESS`` checkpoint recording its result is + stored once it completes. + """ + from airflow.sdk.execution_time.executor import AsyncAwareExecutor + + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{"arg1": 1, "arg2": 10}]) + iterable_op = create_iterable_operator(dag, expand_input, task_id="checkpoint_pending") + + with mock_context(task=iterable_op) as context: + jinja_env = iterable_op.get_template_env(dag=dag) + task = iterable_op._create_task( + context=context, index=0, mapped_kwargs={"arg1": 1, "arg2": 10}, jinja_env=jinja_env + ) + task.try_number = 2 # checkpoint is only consulted from the second attempt onwards + await context["task_state_store"].aset( + task.state_key, + IndexedTaskState(status=TaskInstanceState.UP_FOR_RETRY).serialize(), + ) + + with AsyncAwareExecutor(loop=asyncio.get_running_loop(), max_workers=1) as executor: + result_task, result, raised = await iterable_op._run_task( + executor, + context, + task, + trust_checkpoints=True, + outcomes=IndexedTaskOutcomes(iterable_op, IterationState(), context), + ) + + assert raised is None + assert result == (1, 10, None) + assert ( + result_task.try_number == 2 + ) # try_number is inherited from the parent TI, never mutated + store = context["task_state_store"] + assert ( + store[task.state_key] + == IndexedTaskState( + status=TaskInstanceState.SUCCESS, + result=result, + fingerprint=task.input_fingerprint, + try_number=2, + ).serialize() + ) + + @pytest.mark.asyncio + async def test_run_task_merges_outlet_events_into_shared_context_on_success(self): + """ + A succeeded item's outlet asset events are merged into the task's ``context["outlet_events"]``. + + The task instance reports its outlet events once, from that accessor, when the whole + iteration finishes, so an item's events only reach the server through it. + """ + from airflow.sdk.execution_time.executor import AsyncAwareExecutor + + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{}]) + iterable_op = create_iterable_operator( + dag, expand_input, task_id="outlet_merge", operator_class=MockOutletEventOperator + ) + + with mock_context(task=iterable_op) as context: + jinja_env = iterable_op.get_template_env(dag=dag) + task = iterable_op._create_task( + context=context, index=0, mapped_kwargs={}, jinja_env=jinja_env + ) + + with AsyncAwareExecutor(loop=asyncio.get_running_loop(), max_workers=1) as executor: + _, result, raised = await iterable_op._run_task( + executor, + context, + task, + trust_checkpoints=True, + outcomes=IndexedTaskOutcomes(iterable_op, IterationState(), context), + ) + + assert raised is None + assert result == "done" + accessor = context["outlet_events"][Asset(name="a", uri="s3://bucket/a")] + assert accessor.extra == {"value": "v"} + store = context["task_state_store"] + checkpoint = IndexedTaskState.deserialize(store[task.state_key]) + assert checkpoint.outlet_events == [ + { + "kind": "asset", + "name": "a", + "uri": "s3://bucket/a", + "extra": {"value": "v"}, + "partition_keys": [], + } + ] + + @pytest.mark.asyncio + async def test_run_task_replays_outlet_events_when_skipping_already_succeeded_sub_task(self): + """A sub-task skipped on retry (because its checkpoint already records SUCCESS) never + re-executes, so it would otherwise never re-populate the fresh ``outlet_events`` accessor + created for the new attempt; ``_run_task`` must replay the events it recorded on its + earlier successful attempt instead of silently losing them.""" + from unittest import mock + + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{}]) + iterable_op = create_iterable_operator( + dag, expand_input, task_id="outlet_replay", operator_class=MockOutletEventOperator + ) + + with mock_context(task=iterable_op) as context: + jinja_env = iterable_op.get_template_env(dag=dag) + task = iterable_op._create_task( + context=context, index=0, mapped_kwargs={}, jinja_env=jinja_env + ) + task.try_number = 2 # checkpoint is only consulted from the second attempt onwards + await context["task_state_store"].aset( + task.state_key, + IndexedTaskState( + status=TaskInstanceState.SUCCESS, + fingerprint=task.input_fingerprint, + outlet_events=[ + { + "kind": "asset", + "name": "a", + "uri": "s3://bucket/a", + "extra": {"value": "v"}, + "partition_keys": [], + } + ], + ).serialize(), + ) + + executor = mock.MagicMock() + _, result, raised = await iterable_op._run_task( + executor, + context, + task, + trust_checkpoints=True, + outcomes=IndexedTaskOutcomes(iterable_op, IterationState(), context), + ) + + assert raised is None + assert result is None + executor.run_sync.assert_not_called() + accessor = context["outlet_events"][Asset(name="a", uri="s3://bucket/a")] + assert accessor.extra == {"value": "v"} + + def test_on_kill_propagates_to_active_sub_operators(self): + """IterableOperator.on_kill() (SIGTERM or execution_timeout) must propagate the kill + signal to every sub-task currently in flight, since the default BaseOperator.on_kill() + no-op would otherwise leave running sub-tasks completely unaware of the kill.""" + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{}]) + iterable_op = create_iterable_operator( + dag, expand_input, task_id="on_kill_test", operator_class=MockOnKillOperator + ) + + active_operator = MockOnKillOperator(task_id="active_sub_task") + iterable_op._state.register(active_operator) + + iterable_op.on_kill() + + assert wait_until(lambda: active_operator.killed) + + def test_on_kill_reaches_sub_operators_that_compare_equal_and_kills_each_once(self): + """ + The sub-operators of one iterated task compare equal (same task_id), so the register is keyed + by identity; and the runner's second on_kill() after a timeout does not kill them again. + """ + with DAG("test_dag") as dag: + iterable_op = create_iterable_operator( + dag, ListOfDictsExpandInput([{}]), task_id="on_kill_equal", operator_class=MockOnKillOperator + ) + + first, second = MockOnKillOperator(task_id="same"), MockOnKillOperator(task_id="same") + assert first == second + kills = [] + first.on_kill = lambda: kills.append("first") # type: ignore[method-assign] + second.on_kill = lambda: kills.append("second") # type: ignore[method-assign] + iterable_op._state.register(first) + iterable_op._state.register(second) + + iterable_op.on_kill() + iterable_op.on_kill() + + assert wait_until(lambda: len(kills) == 2) + time.sleep(0.05) # a second kill would arrive here + assert sorted(kills) == ["first", "second"] + + def test_on_kill_reaches_every_sub_operator_when_one_raises_a_base_exception(self): + """A sub-operator's on_kill may raise DeadlockImminentError, a BaseException; the rest are still killed.""" + + class RaisingOnKill(MockOnKillOperator): + def on_kill(self): + raise DeadlockImminentError("sync SDK call on the loop thread") + + with DAG("test_dag") as dag: + iterable_op = create_iterable_operator( + dag, ListOfDictsExpandInput([{}]), task_id="on_kill_raises", operator_class=MockOnKillOperator + ) + first, second = RaisingOnKill(task_id="first"), MockOnKillOperator(task_id="second") + iterable_op._state.register(first) + iterable_op._state.register(second) + + iterable_op.on_kill() + + assert wait_until(lambda: second.killed) + + def test_on_kill_with_the_loop_paused_kills_off_the_main_thread(self): + """ + Between two ``run_until_complete`` calls the loop is paused and a parked ``asend`` may hold + the comms lock: a sub-operator's sync SDK call in ``on_kill`` on the main thread would wait + for it forever, since only the paused loop can release it. The kill runs in a thread. + """ + seen: dict[str, bool] = {} + done = threading.Event() + main = threading.get_ident() + + class Op(MockOnKillOperator): + def on_kill(self): + seen["on_main"] = threading.get_ident() == main + done.set() + + with DAG("test_dag") as dag: + iterable_op = create_iterable_operator( + dag, ListOfDictsExpandInput([{}]), task_id="on_kill_paused", operator_class=Op + ) + op = Op(task_id="active") + iterable_op._state.register(op) + + iterable_op.on_kill() # no loop running on this thread, as when the loop is paused + + assert done.wait(5) + assert seen == {"on_main": False} + assert iterable_op._state.stop_requested() + + def test_on_kill_inside_a_running_loop_kills_off_the_loop_thread(self): + """ + The runner's SIGTERM handler calls on_kill on the main thread, usually while the loop runs + there; the sub-operators' on_kill may make a sync SDK call, which must not run on the loop. + """ + seen: dict[str, bool] = {} + done = threading.Event() + main = threading.get_ident() + + class Op(MockOnKillOperator): + def on_kill(self): + seen["on_running_loop"] = asyncio._get_running_loop() is not None + seen["on_main"] = threading.get_ident() == main + done.set() + + with DAG("test_dag") as dag: + iterable_op = create_iterable_operator( + dag, ListOfDictsExpandInput([{}]), task_id="on_kill_in_loop", operator_class=Op + ) + op = Op(task_id="active") + iterable_op._state.register(op) + + async def kill_from_the_loop(): + iterable_op.on_kill() + await asyncio.wait_for(asyncio.to_thread(done.wait, 5), 6) + + with event_loop() as loop: + loop.run_until_complete(kill_from_the_loop()) + + assert done.is_set() + assert seen == {"on_running_loop": False, "on_main": False} + + def test_on_kill_takes_no_lock_of_the_iteration_state(self): + """ + The runner's SIGTERM handler runs ``on_kill()`` on the main thread, between two bytecodes of + the loop thread, which may be inside ``register``/``unregister`` and hold the state's lock: + ``on_kill()`` returns at once while the lock is held and the kill follows once it is free. + """ + with DAG("test_dag") as dag: + iterable_op = create_iterable_operator( + dag, ListOfDictsExpandInput([{}]), task_id="on_kill_locked", operator_class=MockOnKillOperator + ) + active_operator = MockOnKillOperator(task_id="active_sub_task") + iterable_op._state.register(active_operator) + + with iterable_op._state._lock: + iterable_op.on_kill() + assert iterable_op._state.stop_requested() + assert not active_operator.killed + + assert wait_until(lambda: active_operator.killed) + + def test_on_kill_is_noop_when_no_sub_operators_are_active(self): + """on_kill() must not raise when called with no in-flight sub-tasks (e.g. the + IterableOperator is killed before any sub-task has started, or after all finished).""" + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{}]) + iterable_op = create_iterable_operator( + dag, expand_input, task_id="on_kill_noop", operator_class=MockOnKillOperator + ) + + iterable_op.on_kill() # should not raise + + def test_the_kill_thread_is_waited_for_before_the_run_concludes(self): + """ + ``on_kill()`` kills in a thread of its own and returns at once; the run waits for that + thread once the killed indexed tasks came back, so the task does not end, and the process + with it, while a sub-operator's ``on_kill`` is still cleaning up. + """ + executing = threading.Event() + killed = threading.Event() + cleanup_may_finish = threading.Event() + + class SlowToKill(BaseOperator): + def __init__(self, arg1=None, **kwargs): + super().__init__(**kwargs) + self.arg1 = arg1 + + def execute(self, context): + executing.set() + killed.wait(5) + raise RuntimeError("killed") + + def on_kill(self): + killed.set() + cleanup_may_finish.wait(5) + + outcome: list[BaseException] = [] + with DAG("test_dag") as dag: + iterable_op = create_iterable_operator( + dag, ListOfDictsExpandInput([{"arg1": 1}]), task_id="joined_kill", operator_class=SlowToKill + ) + + with mock_context(task=iterable_op) as context: + + def run(): + try: + iterable_op.execute(context=context) + except BaseException as e: + outcome.append(e) + + run_thread = threading.Thread(target=run) + run_thread.start() + assert executing.wait(5) + iterable_op.on_kill() + assert killed.wait(5) + + run_thread.join(0.5) + assert run_thread.is_alive(), "the run concluded while on_kill was still cleaning up" + cleanup_may_finish.set() + run_thread.join(5) + + assert not run_thread.is_alive() + assert isinstance(outcome[0], AirflowTaskTerminated) + + def test_the_wait_for_the_kill_thread_is_bounded(self): + """A sub-operator whose ``on_kill`` never returns holds the run for the shutdown timeout, not forever.""" + executing = threading.Event() + killed = threading.Event() + released = threading.Event() + + class StuckKill(BaseOperator): + def __init__(self, arg1=None, **kwargs): + super().__init__(**kwargs) + self.arg1 = arg1 + + def execute(self, context): + executing.set() + killed.wait(5) + raise RuntimeError("killed") + + def on_kill(self): + killed.set() + released.wait(30) + + outcome: list[BaseException] = [] + with DAG("test_dag") as dag: + iterable_op = create_iterable_operator( + dag, ListOfDictsExpandInput([{"arg1": 1}]), task_id="stuck_kill", operator_class=StuckKill + ) + + with ( + mock_context(task=iterable_op) as context, + patch.object(IterableOperator, "_SHUTDOWN_TIMEOUT", 0.2), + ): + + def run(): + try: + iterable_op.execute(context=context) + except BaseException as e: + outcome.append(e) + + run_thread = threading.Thread(target=run) + run_thread.start() + assert executing.wait(5) + iterable_op.on_kill() + run_thread.join(5) + + released.set() + assert not run_thread.is_alive() + assert isinstance(outcome[0], AirflowTaskTerminated) + + def test_run_task_tracks_active_sub_operator_during_execution(self, monkeypatch: pytest.MonkeyPatch): + """``_run_task`` must register the sub-task's unmapped operator in the iteration state + only for the duration of its execution, so ``on_kill()`` propagates only to sub-tasks that + are actually running.""" + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{}]) + iterable_op = create_iterable_operator( + dag, expand_input, task_id="on_kill_tracking", operator_class=MockOnKillOperator + ) + + with mock_context(task=iterable_op) as context: + jinja_env = iterable_op.get_template_env(dag=dag) + task = iterable_op._create_task( + context=context, index=0, mapped_kwargs={}, jinja_env=jinja_env + ) + + seen_active_during_run = [] + + def tracking_execute(context, ti, log): + seen_active_during_run.append(task.task in iterable_op._state) + + monkeypatch.setattr("airflow.sdk.execution_time.task_runner._execute_task", tracking_execute) + + with event_loop() as loop, AsyncAwareExecutor(loop=loop, max_workers=1) as executor: + _, _, raised = loop.run_until_complete( + iterable_op._run_task( + executor, + context, + task, + trust_checkpoints=False, + outcomes=IndexedTaskOutcomes(iterable_op, IterationState(), context), + ) + ) + + assert raised is None + assert seen_active_during_run == [True] + assert task.task not in iterable_op._state + + def test_multiple_outputs_is_ignored(self): + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{"arg1": 1}]) + mapped_op = create_mapped_operator(dag=dag, expand_input=expand_input, task_id="multi") + + iterable_op = IterableOperator( + operator=mapped_op, expand_input=expand_input, dag=dag, multiple_outputs=True + ) + + assert iterable_op.multiple_outputs is False + + def test_iterable_execution_timeout_caps_whole_iteration_and_wrapped_operator_retains_it(self): + """IterableOperator keeps execution_timeout as the wall-clock cap on the whole iteration, which + the runner enforces on the outer TI; the wrapped operator keeps the value too, as any of its + own parameters, without a limit of its own per indexed task.""" + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{"arg1": 1}]) + execution_timeout = timedelta(seconds=7) + mapped_op = create_mapped_operator( + dag, expand_input, task_id="timeout_task", execution_timeout=execution_timeout + ) + + iterable_op = IterableOperator(operator=mapped_op, expand_input=expand_input, dag=dag) + + assert iterable_op._operator.execution_timeout == execution_timeout + assert iterable_op.execution_timeout == execution_timeout + + @pytest.mark.parametrize( + "base_exception", + [ + SystemExit(1), + KeyboardInterrupt(), + ], + ids=["SystemExit", "KeyboardInterrupt"], + ) + def test_base_exception_not_retried_raises_airflow_fail_exception(self, base_exception): + """ + BaseException subclasses (e.g., SystemExit, KeyboardInterrupt) must never + be retried—they signal conditions where continuing iteration is meaningless. + They should raise AirflowFailException immediately. + """ + with DAG("test_dag") as dag: + # Create a mapped operator that raises a BaseException + expand_input = ListOfDictsExpandInput([{"raise_exception": base_exception}]) + mapped_op = MockOperator.partial( + task_id="base_exception_task", + dag=dag, + retries=3, # Has retries available, but should NOT use them + )._expand( + expand_input, + strict=True, + register_with_dag=False, + ) + iterable_op = IterableOperator(operator=mapped_op, expand_input=expand_input, dag=dag) + + with mock_context(task=iterable_op) as context: + with pytest.raises(AirflowFailException): + iterable_op.execute(context=context) + + def test_parent_timeout_landing_in_a_sub_task_stays_a_timeout(self): + """ + The parent's ``execution_timeout`` is raised by a signal handler on the main thread, so it + can surface inside whichever async sub-task is running there. It is not that sub-task's + outcome: it has to reach the runner as ``AirflowTaskTimeout``, which retries the task, and + not be turned into a non-retryable ``AirflowFailException`` blamed on the sub-task. + """ + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{}]) + iterable_op = create_iterable_operator( + dag, + expand_input, + task_id="timeout_task", + retries=3, + operator_class=MockCallbackTimeoutOperator, + ) + + with mock_context(task=iterable_op) as context: + with pytest.raises(AirflowTaskTimeout): + iterable_op.execute(context=context) + + assert "_iterable_0" not in context["task_state_store"] + + @pytest.mark.asyncio + async def test_run_task_lets_a_cancellation_through(self): + """A cancelled sub-task stays cancelled: no checkpoint write, no result handed back.""" + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{"arg1": 1}]) + iterable_op = create_iterable_operator(dag, expand_input, task_id="cancelled_task") + + with mock_context(task=iterable_op) as context: + task = iterable_op._create_task( + context=context, + index=0, + mapped_kwargs={"raise_exception": asyncio.CancelledError()}, + jinja_env=iterable_op.get_template_env(dag=dag), + ) + + with AsyncAwareExecutor(loop=asyncio.get_running_loop(), max_workers=1) as executor: + with pytest.raises(asyncio.CancelledError): + await iterable_op._run_task( + executor, + context, + task, + trust_checkpoints=False, + outcomes=IndexedTaskOutcomes(iterable_op, IterationState(), context), + ) + + assert task.state_key not in context["task_state_store"] + + def test_deadlock_imminent_error_raises_actionable_airflow_fail_exception(self): + """A sub-task that raises DeadlockImminentError (a sync SDK call made from an async + sub-task) must never be retried and must surface an actionable error pointing at async-safe + SDK alternatives, rather than the generic non-Exception BaseException message.""" + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput( + [{"raise_exception": DeadlockImminentError("simulated sync SDK call")}] + ) + mapped_op = MockOperator.partial( + task_id="deadlock_task", + dag=dag, + retries=3, # Has retries available, but should NOT use them + )._expand( + expand_input, + strict=True, + register_with_dag=False, + ) + iterable_op = IterableOperator(operator=mapped_op, expand_input=expand_input, dag=dag) + + with mock_context(task=iterable_op) as context: + with pytest.raises(AirflowFailException, match="synchronous SDK call"): + iterable_op.execute(context=context) + + def test_deferred_operator_raises_airflow_fail_exception(self): + """A sub-task that raises TaskDeferred must cause IterableOperator to raise AirflowFailException.""" + with DAG("test_dag") as dag: + expand_input = ListOfDictsExpandInput([{}, {}]) + mapped_op = MockDeferredOperator.partial(task_id="deferred_task", dag=dag)._expand( + expand_input, + strict=True, + register_with_dag=False, + ) + iterable_op = IterableOperator(operator=mapped_op, expand_input=expand_input, dag=dag) + + with mock_context(task=iterable_op) as context: + with pytest.raises(AirflowFailException, match="attempted to defer"): + iterable_op.execute(context=context) + + def test_reschedule_mode_sensor_raises_base_exception_group(self): Review Comment: Right, that text was left from before the fail-fast change. Renamed to `test_reschedule_mode_sensor_fails_the_task_without_retry`, and the docstring now says what it asserts: the reschedule fails the whole task through the `AirflowFailException` from `_fail_fast_for` (8b88691d50). Drafted-by: Claude Opus 5.5; reviewed by @dabla before posting -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
