You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by as...@apache.org on 2026/10/10 22:50:02 UTC
(airflow) 07/12: Run loop gates and advance iterations atomically (#74348)
This is an automated email from the ASF dual-hosted git repository.
ashb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 551d54bdaeefeb992839a39ff5ee7fc185f2686b
Author: Ash Berlin-Taylor <as...@astronomer.io>
AuthorDate: Sat Oct 10 23:49:04 2026 +0100
Run loop gates and advance iterations atomically (#74348)
The scheduler, the API and the UI must never see every task of a loop finished
while its next pass does not exist yet. If we recorded a continuing gate
as successful and created the next pass afterwards, a crash between the two
would leave recovery unable to tell an unfinished advancement from a completed
decision. Gate completion and next-pass creation therefore happen in one
transaction under the DagRun lock, which also stops a duplicate completion from
advancing the loop twice.
The continue-or-stop decision is a reserved XCom stored under the gate's attempt
UUID. A retried gate starts without a decision, and a retired attempt cannot
advance the loop. An invalid decision moves the gate attempt to retry or failed
within the same request; rejecting the request instead would leave the attempt
running with nobody to finish it.
Gates use normal task dependencies and an explicit loop context. A barrier
that waited for every body task would override trigger rules that allow early
progress. Workers and dag.test supply the same coordinates, so decisions and
previous-iteration reads reach the intended gate and data.
Skipping downstream of a task from outside a loop skips every live pass.
Skipping only one would let later passes run work the author meant to skip.
---
.../docs/authoring-and-scheduling/loops.rst | 8 +-
.../execution_api/datamodels/taskinstance.py | 12 +
.../execution_api/routes/task_instances.py | 122 +++-
.../api_fastapi/execution_api/routes/xcoms.py | 179 ++++--
.../api_fastapi/execution_api/versions/__init__.py | 6 +
.../execution_api/versions/v2026_10_30.py | 11 +
.../versions/v2026_10_30_xcom_params.py | 35 ++
.../airflow/example_dags}/example_task_loops.py | 0
.../src/airflow/jobs/scheduler_job_runner.py | 7 +-
airflow-core/src/airflow/models/dagrun.py | 144 ++++-
airflow-core/src/airflow/models/dynamic_region.py | 349 +++++++++--
.../src/airflow/models/task_coordinates.py | 415 +++++++++++--
.../airflow/serialization/definitions/xcom_arg.py | 23 +-
airflow-core/src/airflow/ti_deps/dep_context.py | 7 +-
.../versions/head/test_task_instances.py | 642 ++++++++++++++++++++-
.../execution_api/versions/head/test_xcoms.py | 475 ++++++++++++++-
.../versions/v2026_10_30/test_task_instances.py | 2 +
airflow-core/tests/unit/jobs/test_scheduler_job.py | 50 ++
airflow-core/tests/unit/models/test_dagrun.py | 259 ++++++++-
.../tests/unit/models/test_task_coordinates.py | 241 +++++++-
airflow-core/tests/unit/models/test_xcom_arg.py | 45 ++
.../unit/serialization/test_dag_serialization.py | 24 +
.../unit/ti_deps/deps/test_trigger_rule_dep.py | 74 +++
.../src/tests_common/test_utils/asserts.py | 20 +
go-sdk/pkg/execution/genmodels/models.gen.go | 59 +-
go-sdk/schema/supervisor-schema.json | 67 +++
.../unit/edge3/worker_api/routes/test_logs.py | 14 +-
.../check_template_context_variable_in_sync.py | 2 +
task-sdk/src/airflow/sdk/api/client.py | 25 +-
.../src/airflow/sdk/api/datamodels/_generated.py | 13 +
task-sdk/src/airflow/sdk/bases/decorator.py | 24 +-
task-sdk/src/airflow/sdk/bases/xcom.py | 2 +
.../src/airflow/sdk/definitions/_internal/loop.py | 65 +++
.../src/airflow/sdk/definitions/_internal/node.py | 15 +-
task-sdk/src/airflow/sdk/definitions/context.py | 4 +-
task-sdk/src/airflow/sdk/definitions/dag.py | 9 +-
.../sdk/definitions/decorators/task_group.py | 24 +-
task-sdk/src/airflow/sdk/definitions/xcom_arg.py | 3 +-
task-sdk/src/airflow/sdk/execution_time/comms.py | 4 +
.../airflow/sdk/execution_time/lazy_sequence.py | 4 +
task-sdk/src/airflow/sdk/execution_time/loop.py | 104 ++++
.../airflow/sdk/execution_time/request_handlers.py | 11 +-
.../airflow/sdk/execution_time/schema/schema.json | 67 +++
.../sdk/execution_time/schema/versions/__init__.py | 2 +
.../execution_time/schema/versions/v2026_10_30.py | 17 +
.../src/airflow/sdk/execution_time/task_runner.py | 45 +-
.../task_sdk/definitions/_internal/test_loop.py | 231 +++++++-
.../definitions/decorators/test_task_group.py | 87 ++-
.../tests/task_sdk/definitions/test_xcom_arg.py | 51 +-
.../execution_time/schema/test_migrator.py | 39 +-
.../tests/task_sdk/execution_time/test_loop.py | 245 ++++++++
.../execution_time/test_request_handlers.py | 44 ++
.../task_sdk/execution_time/test_supervisor.py | 28 +-
.../task_sdk/execution_time/test_task_runner.py | 64 ++
ts-sdk/schema/supervisor-schema.json | 67 +++
ts-sdk/src/generated/supervisor.ts | 27 +
56 files changed, 4370 insertions(+), 243 deletions(-)
diff --git a/airflow-core/docs/authoring-and-scheduling/loops.rst b/airflow-core/docs/authoring-and-scheduling/loops.rst
index 8cd0c76f06d..b588fdbe421 100644
--- a/airflow-core/docs/authoring-and-scheduling/loops.rst
+++ b/airflow-core/docs/authoring-and-scheduling/loops.rst
@@ -68,7 +68,7 @@ iterations.
This example improves an estimate of the square root of two until the error
is small enough:
-.. exampleinclude:: /authoring-and-scheduling/examples/example_task_loops.py
+.. exampleinclude:: /../src/airflow/example_dags/example_task_loops.py
:start-after: [START refine_estimate]
:end-before: [END refine_estimate]
@@ -98,7 +98,7 @@ For a fixed-count loop, omit ``until``. The definition ``accumulate.loop(max_ite
three iterations, carrying results between them. Reaching the cap completes
a fixed-count loop successfully. Its gate is named ``__loop_gate`` within the group.
-.. exampleinclude:: /authoring-and-scheduling/examples/example_task_loops.py
+.. exampleinclude:: /../src/airflow/example_dags/example_task_loops.py
:start-after: [START fixed_loop]
:end-before: [END fixed_loop]
@@ -106,7 +106,7 @@ For a task-group function with arguments, supply them with ``.partial()`` before
calling ``.loop()``. Use ``.override()`` to configure the group, for example to
give another loop a different ``group_id``:
-.. exampleinclude:: /authoring-and-scheduling/examples/example_task_loops.py
+.. exampleinclude:: /../src/airflow/example_dags/example_task_loops.py
:start-after: [START partial_override_loop]
:end-before: [END partial_override_loop]
@@ -258,7 +258,7 @@ created.
In this example, ``choose_items`` returns the values to map over: ``[1, 2]`` in the
first iteration, then each previous result plus one.
-.. exampleinclude:: /authoring-and-scheduling/examples/example_task_loops.py
+.. exampleinclude:: /../src/airflow/example_dags/example_task_loops.py
:start-after: [START mapped_loop]
:end-before: [END mapped_loop]
diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py
index 446aea95b01..2d158476ada 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py
@@ -427,12 +427,24 @@ class DagRun(StrictBaseModel):
return values
+class LoopContext(BaseModel):
+ """Pinned loop definition and enclosing iteration for a task execution."""
+
+ node_id: str
+ index: int = Field(ge=0)
+ max_iterations: int = Field(gt=0)
+ terminal_task_id: str
+ terminal_is_mapped: bool
+
+
class TIRunContext(BaseModel):
"""Response schema for TaskInstance run context."""
dag_run: DagRun
"""DAG run information for the task instance."""
+ loop: LoopContext | None = None
+
task_reschedule_count: int = 0
"""How many times the task has been rescheduled."""
diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
index 6c67e87ead9..99cd52a3e97 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
@@ -35,7 +35,7 @@ from opentelemetry import trace
from opentelemetry.trace import StatusCode
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
from pydantic import JsonValue, ValidationError
-from sqlalchemy import and_, func, or_, tuple_, update
+from sqlalchemy import and_, func, or_, tuple_, union, update
from sqlalchemy.engine import CursorResult
from sqlalchemy.exc import DataError, NoResultFound, SQLAlchemyError
from sqlalchemy.orm import Session, contains_eager, joinedload
@@ -57,6 +57,7 @@ from airflow.api_fastapi.execution_api.datamodels.task_arg_binding import get_ar
from airflow.api_fastapi.execution_api.datamodels.taskinstance import (
DagRunNoteUpdatePayload,
InactiveAssetsResponse,
+ LoopContext,
PreviousTIResponse,
PrevSuccessfulDagRunResponse,
TaskBreadcrumbsResponse,
@@ -95,11 +96,12 @@ from airflow.models.asset import AssetActive
from airflow.models.base import ID_LEN
from airflow.models.dag import DagModel
from airflow.models.dagbag import DBDagBag
-from airflow.models.dagrun import DagRun as DR
+from airflow.models.dagrun import DagRun as DR, InvalidLoopDecision
from airflow.models.dynamic_region import SENTINEL_REGION_ID, AmbiguousProducerError
from airflow.models.hitl import HITLDetail
from airflow.models.log import Log
from airflow.models.task_coordinates import (
+ LOOP_GATE_OPERATOR,
TaskCoordinateResolver,
build_coordinate_filters,
get_public_region,
@@ -113,7 +115,7 @@ from airflow.serialization.definitions.assets import SerializedAsset, Serialized
from airflow.state import get_state_backend
from airflow.triggers.base import TriggerEvent
from airflow.utils.sqlalchemy import get_dialect_name
-from airflow.utils.state import DagRunState, TaskInstanceState, TerminalTIState
+from airflow.utils.state import DagRunState, IntermediateTIState, TaskInstanceState, TerminalTIState
router = VersionedAPIRouter()
@@ -177,6 +179,7 @@ def ti_run(
TI.dag_id,
TI.run_id,
TI.task_id,
+ TI.region_id,
TI.region_index,
TI.try_number,
TI.max_tries,
@@ -336,6 +339,19 @@ def ti_run(
should_retry=_is_eligible_to_retry(previous_state, ti.try_number, ti.max_tries),
multi_team=conf.getboolean("core", "multi_team"),
)
+ resolver = TaskCoordinateResolver(dag_bag, session)
+ if loop_context := resolver.loop_context(ti):
+ group, index = loop_context
+ terminal = resolver.get_task(
+ ti.dag_id, ti.run_id, group.terminal_task_id, dag_version_id=ti.dag_version_id
+ )
+ context.loop = LoopContext(
+ node_id=group.node_id,
+ index=index,
+ max_iterations=group.max_iterations,
+ terminal_task_id=group.terminal_task_id,
+ terminal_is_mapped=terminal.get_needs_expansion(),
+ )
# Only set for lang-SDK (foreign-runtime) tasks with a captured TaskFlow arg
# spec; the route excludes unset fields, keeping regular responses lean.
@@ -435,6 +451,46 @@ def ti_update_state(
raise HTTPException(status_code=409, detail={"reason": "invalid_state"})
return Response(status_code=status.HTTP_204_NO_CONTENT)
+ loop_group = None
+ loop_gate = None
+ if isinstance(ti_patch_payload, (TISuccessStatePayload, TITerminalStatePayload)):
+ gate_run = session.execute(
+ select(TI.dag_id, TI.run_id).where(
+ TI.id == task_instance_id,
+ TI.working_set.is_(True),
+ TI.operator == LOOP_GATE_OPERATOR,
+ )
+ ).one_or_none()
+ if gate_run is not None:
+ session.execute(
+ select(DR).where(DR.dag_id == gate_run.dag_id, DR.run_id == gate_run.run_id).with_for_update()
+ ).scalar_one()
+ loop_gate = session.scalar(
+ select(TI)
+ .where(TI.id == task_instance_id, TI.working_set.is_(True))
+ .with_for_update(of=TI)
+ .execution_options(populate_existing=True)
+ )
+ if loop_gate is not None:
+ try:
+ loop_context = TaskCoordinateResolver(dag_bag, session).loop_context(loop_gate)
+ except TaskNotFound:
+ # A gate renamed mid-run on an unversioned bundle leaves this attempt without a loop.
+ log.warning("Loop gate is missing from its pinned Dag version; failing it")
+ if isinstance(ti_patch_payload, TISuccessStatePayload):
+ ti_patch_payload = _build_rejected_gate_payload(
+ ti_patch_payload,
+ reason="Loop gate is missing from its pinned Dag version",
+ retry=False,
+ )
+ loop_context = None
+ if loop_context is None:
+ loop_gate = None
+ elif loop_context[0].gate_task_id != loop_gate.task_id:
+ raise HTTPException(status_code=409, detail={"reason": "invalid_loop_gate"})
+ else:
+ loop_group = loop_context[0]
+
old = (
select(
TI.state,
@@ -522,6 +578,25 @@ def ti_update_state(
detail={"reason": "invalid_partition_key", "message": str(e)},
) from e
+ gate_completed = False
+ if (
+ loop_gate is not None
+ and loop_group is not None
+ and isinstance(ti_patch_payload, TISuccessStatePayload)
+ ):
+ try:
+ loop_gate.dag_run.complete_loop_gate(
+ loop_gate, loop_group, TaskInstanceState.SUCCESS, session=session
+ )
+ gate_completed = True
+ except InvalidLoopDecision as error:
+ log.warning("Loop gate success rejected", error=str(error))
+ ti_patch_payload = _build_rejected_gate_payload(
+ ti_patch_payload,
+ reason=f"Loop gate success rejected: {error}",
+ retry=_is_eligible_to_retry(previous_state, try_number, max_tries),
+ )
+
# We exclude_unset to avoid updating fields that are not set in the payload
data = ti_patch_payload.model_dump(
exclude={"task_outlets", "outlet_events", "retry_delay_seconds", "retry_reason"},
@@ -545,6 +620,8 @@ def ti_update_state(
# Let DataErrorHandler return a 422 instead of silently marking the TI FAILED below.
raise
except Exception:
+ if loop_gate is not None:
+ raise
# Set a task to failed in case any unexpected exception happened during task state update
log.exception(
"Error updating Task Instance state. Setting the task to failed.",
@@ -595,6 +672,9 @@ def ti_update_state(
# Defer to app-level SQLAlchemyError handler (returns HTTP 500).
raise
+ if loop_gate is not None and loop_group is not None and not gate_completed:
+ loop_gate.dag_run.complete_loop_gate(loop_gate, loop_group, updated_state, session=session)
+
if updated_state == TaskInstanceState.SUCCESS:
if conf.getboolean("state_store", "clear_on_success"):
scope = TaskScope(
@@ -628,6 +708,15 @@ def ti_update_state(
callback()
+def _build_rejected_gate_payload(
+ payload: TISuccessStatePayload, *, reason: str, retry: bool
+) -> TIRetryStatePayload | TITerminalStatePayload:
+ carried = payload.model_dump(include={"end_date", "rendered_map_index"}, exclude_unset=True)
+ if retry:
+ return TIRetryStatePayload(state=IntermediateTIState.UP_FOR_RETRY, retry_reason=reason, **carried)
+ return TITerminalStatePayload(state=TerminalStateNonSuccess.FAILED, retry_reason=reason, **carried)
+
+
def _emit_task_span(ti, state, *, resolver: TaskCoordinateResolver):
# just to be safe
if not ti.dag_run:
@@ -941,29 +1030,27 @@ def ti_skip_downstream(
for task_id, index in targets
if not resolver.has_regions(dag_id, run_id, task_id)
]
- selected_ids = {
- ti.id
+ selects = [
+ resolver.select_skip_target_ids(caller=caller, task_id=task_id, map_indexes=index)
for task_id, index in targets
if (task_id, index) not in plain
- for ti in resolver.resolve(
- dag_id=dag_id, run_id=run_id, task_id=task_id, caller=caller, map_indexes=index
- )
- }
+ ]
if plain:
- selected_ids.update(
- session.scalars(
- resolver.select_legacy_task_ids(
- dag_id=dag_id,
- run_id=run_id,
- task_ids=[task_id for task_id, index in plain if index is None],
- slots=[(task_id, index) for task_id, index in plain if index is not None],
- )
+ selects.append(
+ resolver.select_legacy_task_ids(
+ dag_id=dag_id,
+ run_id=run_id,
+ task_ids=[task_id for task_id, index in plain if index is None],
+ slots=[(task_id, index) for task_id, index in plain if index is not None],
)
)
except AmbiguousProducerError as error:
raise HTTPException(status.HTTP_409_CONFLICT, str(error)) from error
except ValueError as error:
raise HTTPException(status.HTTP_400_BAD_REQUEST, str(error)) from error
+ if not selects:
+ return
+ selected_ids = set(session.scalars(union(*selects)))
# Don't overwrite tasks that are already executing or finished.
# See: https://github.com/apache/airflow/issues/59378
@@ -981,6 +1068,7 @@ def ti_skip_downstream(
query = (
update(TI)
.where(
+ TI.working_set.is_(True),
TI.id.in_(selected_ids),
skippable_state_clause,
)
diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py
index c8d96e5c5d7..02fb6eceb6b 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py
@@ -23,6 +23,7 @@ from functools import partial
from typing import TYPE_CHECKING, Annotated
from uuid import UUID
+import attrs
from fastapi import APIRouter, Body, Depends, HTTPException, Path, Query, Request, Response, status
from pydantic import JsonValue
from sqlalchemy import select
@@ -40,8 +41,14 @@ from airflow.api_fastapi.execution_api.datamodels.xcom import (
from airflow.api_fastapi.execution_api.security import CurrentTIToken
from airflow.api_fastapi.execution_api.versions.v2026_10_30 import IdentifyArchivedTaskStateUpdates
from airflow.exceptions import TaskNotFound
-from airflow.models.dagrun import DagRun
-from airflow.models.dynamic_region import SENTINEL_REGION_ID, AmbiguousProducerError
+from airflow.models.dagrun import DagRun, InvalidLoopDecision, validate_loop_decision
+from airflow.models.dynamic_region import (
+ LOOP_DECISION_KEY,
+ LOOP_XCOM_PREFIX,
+ SENTINEL_REGION_ID,
+ AmbiguousProducerError,
+ build_loop_sequence_order,
+)
from airflow.models.task_coordinates import (
TaskCoordinateResolver,
build_coordinate_filters,
@@ -50,9 +57,12 @@ from airflow.models.task_coordinates import (
from airflow.models.taskinstance import TaskInstance
from airflow.models.xcom import XCOM_RETURN_KEY, XComModel, xcom_entity
from airflow.utils.db import get_query_count
+from airflow.utils.state import TaskInstanceState
if TYPE_CHECKING:
from sqlalchemy.orm import Session
+ from sqlalchemy.orm.util import AliasedClass
+ from sqlalchemy.sql.elements import ColumnElement
from airflow.api_fastapi.execution_api.datamodels.token import TIToken
from airflow.models.dagbag import DBDagBag
@@ -136,6 +146,23 @@ router = APIRouter(
log = logging.getLogger(__name__)
+@attrs.frozen
+class _XComRead:
+ """The XCom rows a read selects, and whether they are the executions of a loop task seen from outside it."""
+
+ statement: Select
+ loop_sequence: bool = False
+
+ def get_order(self, entity: AliasedClass, *, descending: bool = False) -> tuple[ColumnElement, ...]:
+ """Order the rows by their place in the sequence the read returns, or reverse that order."""
+ columns = (
+ build_loop_sequence_order(entity.region_id, entity.region_index)
+ if self.loop_sequence
+ else (entity.map_index,)
+ )
+ return tuple(column.desc() if descending else column.asc() for column in columns)
+
+
def _build_xcom_read(
*,
dag_id: str,
@@ -147,8 +174,15 @@ def _build_xcom_read(
token: TIToken,
map_index: int | None = None,
include_prior_dates: bool = False,
-) -> Select:
- """Select the XCom rows of the producers visible to the calling task instance."""
+ previous_iteration: bool = False,
+ sequence: bool = False,
+) -> _XComRead:
+ """
+ Select the XCom rows of the producers visible to the calling task instance.
+
+ A ``sequence`` read returns several values in order. When it targets a loop task from outside that
+ loop, it covers every iteration of the task rather than the one a caller inside the loop would see.
+ """
resolver = TaskCoordinateResolver(dag_bag, session)
read = partial(
XComModel.get_many,
@@ -159,8 +193,12 @@ def _build_xcom_read(
include_prior_dates=include_prior_dates,
)
try:
- if not resolver.has_regions(dag_id, None if include_prior_dates else run_id, task_id):
- return read(region_id=SENTINEL_REGION_ID, map_indexes=map_index)
+ if previous_iteration and include_prior_dates:
+ raise ValueError("previous_iteration cannot be combined with prior dates")
+ if not previous_iteration and not resolver.has_regions(
+ dag_id, None if include_prior_dates else run_id, task_id
+ ):
+ return _XComRead(read(region_id=SENTINEL_REGION_ID, map_indexes=map_index))
caller = session.get(TaskInstance, token.id)
if include_prior_dates:
task = None
@@ -176,18 +214,27 @@ def _build_xcom_read(
)
if task is not None and enclosing_loop(task) is not None:
raise ValueError("Prior-date lookups are not supported for tasks inside a loop")
- return read(region_id=SENTINEL_REGION_ID, include_node_regions=True, map_indexes=map_index)
- producers = resolver.resolve(
- dag_id=dag_id,
- run_id=run_id,
- task_id=task_id,
- caller=caller,
- map_indexes=map_index,
+ return _XComRead(
+ read(region_id=SENTINEL_REGION_ID, include_node_regions=True, map_indexes=map_index)
+ )
+ loop_sequence = sequence and resolver.is_loop_task_read_from_outside(
+ dag_id=dag_id, run_id=run_id, task_id=task_id, caller=caller
)
- return XComModel.get_many(
- run_id=run_id,
- key=key,
- producer_ids=select(TaskInstance.id).where(TaskInstance.id.in_([ti.id for ti in producers])),
+ return _XComRead(
+ XComModel.get_many(
+ run_id=run_id,
+ key=key,
+ producer_ids=resolver.select_producer_ids(
+ dag_id=dag_id,
+ run_id=run_id,
+ task_id=task_id,
+ caller=caller,
+ map_indexes=map_index,
+ previous_iteration=previous_iteration,
+ all_iterations=loop_sequence,
+ ),
+ ),
+ loop_sequence=loop_sequence,
)
except AmbiguousProducerError as error:
raise HTTPException(status.HTTP_409_CONFLICT, str(error)) from error
@@ -203,6 +250,7 @@ def xcom_query(
session: SessionDep,
dag_bag: DagBagDep,
map_index: Annotated[int | None, Query()] = None,
+ previous_iteration: bool = False,
token=CurrentTIToken,
) -> Select:
return _build_xcom_read(
@@ -214,7 +262,9 @@ def xcom_query(
dag_bag=dag_bag,
token=token,
map_index=map_index,
- )
+ previous_iteration=previous_iteration,
+ sequence=True,
+ ).statement
@router.get(
@@ -235,6 +285,7 @@ def get_mapped_xcom_by_index(
offset: int,
session: SessionDep,
dag_bag: DagBagDep,
+ previous_iteration: bool = False,
token=CurrentTIToken,
) -> XComSequenceIndexResponse:
xcom_read = _build_xcom_read(
@@ -245,14 +296,15 @@ def get_mapped_xcom_by_index(
session=session,
dag_bag=dag_bag,
token=token,
+ previous_iteration=previous_iteration,
+ sequence=True,
)
- entity = xcom_entity(xcom_read)
- xcom_query = xcom_read
- xcom_query = xcom_query.order_by(None)
+ entity = xcom_entity(xcom_read.statement)
+ xcom_query = xcom_read.statement.order_by(None)
if offset >= 0:
- xcom_query = xcom_query.order_by(entity.map_index.asc()).offset(offset)
+ xcom_query = xcom_query.order_by(*xcom_read.get_order(entity)).offset(offset)
else:
- xcom_query = xcom_query.order_by(entity.map_index.desc()).offset(-1 - offset)
+ xcom_query = xcom_query.order_by(*xcom_read.get_order(entity, descending=True)).offset(-1 - offset)
result: tuple[XComModel] | None
if (result := session.scalars(xcom_query.limit(1)).first()) is None:
@@ -273,6 +325,7 @@ class GetXComSliceFilterParams(BaseModel):
stop: int | None = None
step: int | None = None
include_prior_dates: bool = False
+ previous_iteration: bool = False
def _get_sliced_query_or_empty(query: Select, low: int, high: int) -> Select:
@@ -310,10 +363,12 @@ def get_mapped_xcom_by_slice(
dag_bag=dag_bag,
token=token,
include_prior_dates=params.include_prior_dates,
+ previous_iteration=params.previous_iteration,
+ sequence=True,
)
- entity = xcom_entity(xcom_read)
- query = xcom_read
- query = query.order_by(None)
+ entity = xcom_entity(xcom_read.statement)
+ query = xcom_read.statement.order_by(None)
+ ascending, descending = xcom_read.get_order(entity), xcom_read.get_order(entity, descending=True)
step = params.step or 1
@@ -323,25 +378,25 @@ def get_mapped_xcom_by_slice(
if (start := params.start) is None:
if (stop := params.stop) is None:
if step >= 0:
- query = query.order_by(entity.map_index.asc())
+ query = query.order_by(*ascending)
else:
- query = query.order_by(entity.map_index.desc())
+ query = query.order_by(*descending)
step = -step
elif stop >= 0:
- query = query.order_by(entity.map_index.asc())
+ query = query.order_by(*ascending)
if step >= 0:
query = query.limit(stop)
else:
query = query.offset(stop + 1)
else:
- query = query.order_by(entity.map_index.desc())
+ query = query.order_by(*descending)
step = -step
if step > 0:
query = query.limit(-stop - 1)
else:
query = query.offset(-stop)
elif start >= 0:
- query = query.order_by(entity.map_index.asc())
+ query = query.order_by(*ascending)
if (stop := params.stop) is None:
if step >= 0:
query = query.offset(start)
@@ -355,7 +410,7 @@ def get_mapped_xcom_by_slice(
else:
query = _get_sliced_query_or_empty(query, stop + 1, start + 1)
else:
- query = query.order_by(entity.map_index.desc())
+ query = query.order_by(*descending)
step = -step
if (stop := params.stop) is None:
if step > 0:
@@ -419,6 +474,7 @@ class GetXcomFilterParams(BaseModel):
map_index: int = -1
include_prior_dates: bool = False
offset: int | None = None
+ previous_iteration: bool = False
@router.get(
@@ -482,15 +538,19 @@ def get_xcom(
token=token,
map_index=params.map_index if params.offset is None else None,
include_prior_dates=params.include_prior_dates,
+ previous_iteration=params.previous_iteration,
+ sequence=params.offset is not None,
)
- entity = xcom_entity(xcom_read)
- xcom_query = xcom_read
+ entity = xcom_entity(xcom_read.statement)
+ xcom_query = xcom_read.statement
if params.offset is not None:
xcom_query = xcom_query.where(entity.value.is_not(None)).order_by(None)
if params.offset >= 0:
- xcom_query = xcom_query.order_by(entity.map_index.asc()).offset(params.offset)
+ xcom_query = xcom_query.order_by(*xcom_read.get_order(entity)).offset(params.offset)
else:
- xcom_query = xcom_query.order_by(entity.map_index.desc()).offset(-1 - params.offset)
+ xcom_query = xcom_query.order_by(*xcom_read.get_order(entity, descending=True)).offset(
+ -1 - params.offset
+ )
# We use `BaseXCom.get_many` to fetch XComs directly from the database, bypassing the XCom Backend.
# This avoids deserialization via the backend (e.g., from a remote storage like S3) and instead
@@ -528,7 +588,10 @@ def get_xcom(
status.HTTP_400_BAD_REQUEST,
"The key is empty, the value is too large to map, or is unserializable",
),
- (status.HTTP_409_CONFLICT, "Several live task instances match the producer"),
+ (
+ status.HTTP_409_CONFLICT,
+ "Several live task instances match the producer, or the loop gate is no longer running",
+ ),
]
),
)
@@ -538,6 +601,7 @@ def set_xcom(
task_id: str,
key: Annotated[str, Path(min_length=1)],
session: SessionDep,
+ dag_bag: DagBagDep,
value: Annotated[
JsonValue,
Body(
@@ -568,6 +632,45 @@ def set_xcom(
"""Set an Airflow XCom."""
from airflow.configuration import conf
+ region_id: UUID | None = None
+ region_index: int | None = None
+ if key == LOOP_DECISION_KEY:
+ caller = session.get(TaskInstance, token.id)
+ if (
+ value not in ("continue", "stop")
+ or mapped_length is not None
+ or dag_result
+ or map_index != -1
+ or caller is None
+ or (dag_id, run_id, task_id) != (caller.dag_id, caller.run_id, caller.task_id)
+ ):
+ raise HTTPException(status.HTTP_400_BAD_REQUEST, "Invalid loop decision write")
+ session.execute(
+ select(DagRun)
+ .where(DagRun.dag_id == caller.dag_id, DagRun.run_id == caller.run_id)
+ .with_for_update()
+ ).scalar_one()
+ caller = session.scalar(
+ select(TaskInstance)
+ .where(TaskInstance.id == token.id)
+ .with_for_update()
+ .execution_options(populate_existing=True)
+ )
+ if caller is None:
+ raise HTTPException(status.HTTP_409_CONFLICT, "Loop gate is no longer running")
+ loop = TaskCoordinateResolver(dag_bag, session).loop_context(caller)
+ if loop is None or loop[0].gate_task_id != caller.task_id:
+ raise HTTPException(status.HTTP_400_BAD_REQUEST, "Loop decisions require the pinned loop gate")
+ if caller.state != TaskInstanceState.RUNNING or caller.working_set is not True:
+ raise HTTPException(status.HTTP_409_CONFLICT, "Loop gate is no longer running")
+ try:
+ validate_loop_decision(value, caller.region_index, loop[0])
+ except InvalidLoopDecision as error:
+ raise HTTPException(status.HTTP_400_BAD_REQUEST, str(error)) from error
+ region_id, region_index = caller.region_id, caller.region_index
+ elif key.startswith(LOOP_XCOM_PREFIX):
+ raise HTTPException(status.HTTP_400_BAD_REQUEST, "Reserved loop XCom key")
+
# Validate that the provided key is not empty
# XCom keys must be non-empty strings to ensure proper data retrieval and avoid ambiguity.
if not key:
@@ -614,6 +717,8 @@ def set_xcom(
run_id=run_id,
task_id=task_id,
map_index=map_index,
+ region_id=region_id,
+ region_index=region_index,
session=session,
),
serialize=False,
diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
index b87c99b7c27..f23972db004 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
@@ -56,12 +56,16 @@ from airflow.api_fastapi.execution_api.versions.v2026_10_30 import (
AddArgBindingsToTIRunContext,
AddCallbackRunEndpoint,
AddDagRunNoteUpdateEndpoint,
+ AddLoopContext,
AddMultiTeamToTIRunContext,
AddStoppedTaskReport,
AddTaskInstanceRegionCoordinates,
AddTerminalStateRetryReasonField,
IdentifyArchivedTaskStateUpdates,
)
+from airflow.api_fastapi.execution_api.versions.v2026_10_30_xcom_params import (
+ AddPreviousIterationToXComFilterParams,
+)
bundle = VersionBundle(
HeadVersion(),
@@ -70,8 +74,10 @@ bundle = VersionBundle(
AddArgBindingsToTIRunContext,
AddCallbackRunEndpoint,
AddDagRunNoteUpdateEndpoint,
+ AddLoopContext,
AddTerminalStateRetryReasonField,
AddMultiTeamToTIRunContext,
+ AddPreviousIterationToXComFilterParams,
AddStoppedTaskReport,
AddTaskInstanceRegionCoordinates,
IdentifyArchivedTaskStateUpdates,
diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py
index fd8ca698ed3..a75c7591982 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py
@@ -36,6 +36,17 @@ from airflow.api_fastapi.execution_api.datamodels.taskinstance import (
)
+class AddLoopContext(VersionChange):
+ """Provide the pinned loop definition and enclosing iteration at task startup."""
+
+ description = __doc__
+ instructions_to_migrate_to_previous_version = (schema(TIRunContext).field("loop").didnt_exist,)
+
+ @convert_response_to_previous_version_for(TIRunContext) # type: ignore[arg-type]
+ def remove_loop_context(response: ResponseInfo) -> None: # type: ignore[misc]
+ response.body.pop("loop", None)
+
+
class AddTaskInstanceRegionCoordinates(VersionChange):
"""Carry task execution coordinates independently of public mapped position, omitting them outside a region."""
diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30_xcom_params.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30_xcom_params.py
new file mode 100644
index 00000000000..4a97ff9af91
--- /dev/null
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30_xcom_params.py
@@ -0,0 +1,35 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from cadwyn import VersionChange, schema
+
+# These models live next to the routes that use them, and the routes import this version package, so the
+# changes that touch them cannot sit in ``v2026_10_30`` without a circular import.
+from airflow.api_fastapi.execution_api.routes.xcoms import GetXcomFilterParams, GetXComSliceFilterParams
+
+
+class AddPreviousIterationToXComFilterParams(VersionChange):
+ """Add the `previous_iteration` field to GetXComSliceFilterParams and GetXcomFilterParams."""
+
+ description = __doc__
+
+ instructions_to_migrate_to_previous_version = (
+ schema(GetXComSliceFilterParams).field("previous_iteration").didnt_exist,
+ schema(GetXcomFilterParams).field("previous_iteration").didnt_exist,
+ )
diff --git a/airflow-core/docs/authoring-and-scheduling/examples/example_task_loops.py b/airflow-core/src/airflow/example_dags/example_task_loops.py
similarity index 100%
rename from airflow-core/docs/authoring-and-scheduling/examples/example_task_loops.py
rename to airflow-core/src/airflow/example_dags/example_task_loops.py
diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
index ca0604af6b6..5c5857ca8ec 100644
--- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
+++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
@@ -114,7 +114,7 @@ from airflow.models.dynamic_region import SENTINEL_REGION_ID
from airflow.models.log import resolve_team_name
from airflow.models.pool import normalize_pool_name_for_stats
from airflow.models.serialized_dag import SerializedDagModel
-from airflow.models.task_coordinates import TaskCoordinateResolver
+from airflow.models.task_coordinates import LOOP_GATE_OPERATOR, TaskCoordinateResolver
from airflow.models.taskinstance import TaskInstance
from airflow.models.team import Team
from airflow.models.trigger import TRIGGER_FAIL_REPR, Trigger, TriggerFailureReason, handle_event_submit
@@ -1685,6 +1685,11 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
task = dag.get_task(ti.task_id)
except Exception:
+ if state == TaskInstanceState.SUCCESS and ti.operator == LOOP_GATE_OPERATOR:
+ cls.logger().exception("Failing loop gate %s: %s", ti, msg)
+ ti.task = None
+ ti.handle_failure(error=msg, session=session)
+ continue
cls.logger().exception("Marking task instance %s as %s", ti, state)
ti.set_state(state)
continue
diff --git a/airflow-core/src/airflow/models/dagrun.py b/airflow-core/src/airflow/models/dagrun.py
index 8d4ebcc91d3..0643e613bab 100644
--- a/airflow-core/src/airflow/models/dagrun.py
+++ b/airflow-core/src/airflow/models/dagrun.py
@@ -88,7 +88,12 @@ from airflow.models import Deadline, Log
from airflow.models.backfill import Backfill
from airflow.models.base import Base, StringID
from airflow.models.deadline_alert import DeadlineAlert as DeadlineAlertModel
-from airflow.models.dynamic_region import SENTINEL_REGION_ID, DynamicRegion, public_region_filter
+from airflow.models.dynamic_region import (
+ LOOP_DECISION_KEY,
+ SENTINEL_REGION_ID,
+ DynamicRegion,
+ public_region_filter,
+)
from airflow.models.taskinstance import TaskInstance as TI, _add_and_prime_mapped_ti, clear_task_instances
from airflow.models.tasklog import LogTemplate
from airflow.serialization.definitions.deadline import SerializedReferenceModels
@@ -127,6 +132,7 @@ if TYPE_CHECKING:
from airflow.sdk import DAG as SDKDAG
from airflow.serialization.definitions.dag import SerializedDAG
from airflow.serialization.definitions.mappedoperator import Operator
+ from airflow.serialization.definitions.taskgroup import SerializedLoopTaskGroup
from airflow.timetables.base import Timetable
CreatedTasks = TypeVar("CreatedTasks", Iterator["dict[str, Any]"], Iterator[TI])
@@ -141,6 +147,21 @@ log = structlog.get_logger(__name__)
tracer = trace.get_tracer(__name__)
+class InvalidLoopDecision(ValueError):
+ """A gate cannot complete successfully with the available decision."""
+
+
+def validate_loop_decision(decision: str | None, region_index: int, group: SerializedLoopTaskGroup) -> None:
+ """Raise :class:`InvalidLoopDecision` unless ``decision`` is valid for the gate of pass ``region_index``."""
+ at_limit = region_index + 1 >= group.max_iterations
+ if decision not in ("continue", "stop"):
+ raise InvalidLoopDecision("Successful loop gate requires a decision")
+ if decision == "continue" and at_limit:
+ raise InvalidLoopDecision("Loop cannot continue beyond its iteration limit")
+ if not group.has_until and (decision == "stop") != at_limit:
+ raise InvalidLoopDecision("Fixed-count loop decision does not match its iteration limit")
+
+
class TISchedulingDecision(NamedTuple):
"""Type of return for DagRun.task_instance_scheduling_decisions."""
@@ -1931,11 +1952,10 @@ class DagRun(Base, LoggingMixin):
# Create the missing tasks, including mapped tasks
missing_tasks = [task for task in dag.task_dict.values() if task_filter(task)]
tis_to_create = list(
- self._create_tasks(
+ self._create_initial_tasks(
missing_tasks,
task_creator,
session=session,
- expand_literals=True,
)
)
self._create_task_instances(self.dag_id, tis_to_create, created_counts, hook_is_noop, session=session)
@@ -2205,6 +2225,121 @@ class DagRun(Base, LoggingMixin):
creator = create_ti
return creator
+ def complete_loop_gate(
+ self,
+ gate: TI,
+ group: SerializedLoopTaskGroup,
+ state: TaskInstanceState,
+ *,
+ session: Session,
+ ) -> None:
+ """Consume a gate decision while the caller holds the DagRun and TI locks."""
+ from airflow.models.xcom import XComModelV2
+ from airflow.settings import task_instance_mutation_hook
+
+ if gate.dag_version_id is None:
+ raise InvalidLoopDecision("Loop gate requires a pinned DAG version")
+ signal = XComModelV2.get_for_attempt(gate.id, LOOP_DECISION_KEY, session=session)
+ later_gate = session.scalar(
+ select(TI.id)
+ .where(
+ TI.working_set.is_(True),
+ TI.dag_id == self.dag_id,
+ TI.run_id == self.run_id,
+ TI.task_id == gate.task_id,
+ TI.region_id == gate.region_id,
+ TI.region_index > gate.region_index,
+ )
+ .limit(1)
+ )
+ if state != TaskInstanceState.SUCCESS or later_gate:
+ if signal is not None:
+ session.delete(signal)
+ return
+ decision = signal.value if signal is not None else None
+ validate_loop_decision(decision, gate.region_index, group)
+ session.delete(signal)
+ if decision == "stop":
+ return
+ created_counts: dict[str, int] = defaultdict(int)
+ hook_is_noop: Literal[True, False] = getattr(task_instance_mutation_hook, "is_noop", False)
+ creator = self._get_task_creator(
+ created_counts, task_instance_mutation_hook, hook_is_noop, gate.dag_version_id
+ )
+ tasks = list(
+ self._create_tasks(
+ group.iter_tasks(),
+ creator,
+ session=session,
+ parent_region=(gate.region_id, gate.region_index + 1),
+ )
+ )
+ self._create_task_instances(
+ self.dag_id, tasks, created_counts, hook_is_noop, session=session, propagate_errors=True
+ )
+
+ def _create_initial_tasks(
+ self,
+ tasks: Iterable[Operator],
+ task_creator: Callable[[Operator, Iterable[int], UUID], CreatedTasks],
+ *,
+ session: Session,
+ ) -> CreatedTasks:
+ from airflow.models.task_coordinates import enclosing_loop
+
+ ordinary_tasks = []
+ loop_tasks = defaultdict(list)
+ loops = {}
+ for task in tasks:
+ if loop := enclosing_loop(task):
+ loops[loop.group_id] = loop
+ loop_tasks[loop.group_id].append(task)
+ else:
+ ordinary_tasks.append(task)
+ yield from self._create_tasks(ordinary_tasks, task_creator, session=session, expand_literals=True)
+ for group_id, members in loop_tasks.items():
+ coordinates = (
+ session.execute(
+ select(TI.region_id, TI.region_index).where(
+ TI.working_set.is_(True),
+ TI.dag_id == self.dag_id,
+ TI.run_id == self.run_id,
+ TI.task_id == loops[group_id].gate_task_id,
+ )
+ )
+ .tuples()
+ .all()
+ )
+ if not coordinates:
+ if session.scalar(
+ select(DynamicRegion.id)
+ .where(
+ DynamicRegion.dag_id == self.dag_id,
+ DynamicRegion.run_id == self.run_id,
+ DynamicRegion.node_id == group_id,
+ )
+ .limit(1)
+ ):
+ log.warning(
+ "skipping loop that has regions but no live gate",
+ dag_id=self.dag_id,
+ run_id=self.run_id,
+ loop=group_id,
+ )
+ continue
+ region = DynamicRegion.get_or_create(
+ dag_id=self.dag_id, run_id=self.run_id, node_id=loops[group_id].node_id, session=session
+ )
+ coordinates = [(region.id, 0)]
+ for region_id, region_index in coordinates:
+ yield from self._create_tasks(
+ members,
+ task_creator,
+ session=session,
+ parent_region=(region_id, region_index),
+ expand_literals=True,
+ )
+
def _create_tasks(
self,
tasks: Iterable[Operator],
@@ -2251,6 +2386,7 @@ class DagRun(Base, LoggingMixin):
hook_is_noop: bool,
*,
session: Session,
+ propagate_errors: bool = False,
) -> None:
"""
Create the necessary task instances from the given tasks.
@@ -2280,6 +2416,8 @@ class DagRun(Base, LoggingMixin):
)
session.flush()
except (IntegrityError, StaleDataError) as exc:
+ if propagate_errors:
+ raise
self.log.info(
"Hit %s while creating the TIs for %s- %s",
type(exc).__name__,
diff --git a/airflow-core/src/airflow/models/dynamic_region.py b/airflow-core/src/airflow/models/dynamic_region.py
index e9073047bf3..fbef2054312 100644
--- a/airflow-core/src/airflow/models/dynamic_region.py
+++ b/airflow-core/src/airflow/models/dynamic_region.py
@@ -20,7 +20,8 @@ import hashlib
import struct
from collections.abc import Collection, Iterable
from datetime import datetime
-from typing import TYPE_CHECKING
+from functools import partial
+from typing import TYPE_CHECKING, Any
from uuid import UUID
import attrs
@@ -35,6 +36,8 @@ from sqlalchemy import (
UniqueConstraint,
and_,
event,
+ false,
+ func,
literal,
or_,
select,
@@ -48,8 +51,11 @@ from airflow.models.base import Base, StringID
from airflow.utils.sqlalchemy import CompactUUID, UtcDateTime
SENTINEL_REGION_ID = UUID(int=0)
+LOOP_XCOM_PREFIX = "_airflow_loop_"
+LOOP_DECISION_KEY = f"{LOOP_XCOM_PREFIX}decision"
if TYPE_CHECKING:
+ from sqlalchemy import Select
from sqlalchemy.engine import Connection
from sqlalchemy.orm import Mapper, Session
from sqlalchemy.sql.elements import ColumnElement
@@ -466,29 +472,88 @@ def _match_any(column, value: str | Collection[str]) -> ColumnElement[bool]:
return column == value if isinstance(value, str) else column.in_(value)
-def resolve_current_producers(
+def _filter_producers(
+ query: Select,
*,
dag_id: str,
run_id: str,
task_id: str,
is_mapped: bool,
- context: ProducerContext | None = None,
- map_indexes: int | Collection[int] | None = None,
- region_id: UUID | None = None,
- region_index: int | None = None,
- session: Session,
-) -> tuple[TaskInstance, ...]:
- """Resolve the live producer task instances whose data the caller reads by task instance UUID."""
+ map_indexes: int | Collection[int] | None,
+ region_id: UUID | None,
+ region_index: int | None,
+ top_level_only: bool,
+ loop_pass: _LoopPassRegions | None = None,
+) -> Select:
+ """Push every coordinate predicate into SQL."""
from airflow.models.taskinstance import TaskInstance
+ query = query.where(
+ TaskInstance.dag_id == dag_id,
+ TaskInstance.run_id == run_id,
+ TaskInstance.task_id == task_id,
+ TaskInstance.working_set.is_(True),
+ )
+ if region_id is not None:
+ query = query.where(TaskInstance.region_id == region_id)
+ if region_index is not None:
+ query = query.where(TaskInstance.region_index == region_index)
+ if loop_pass is not None:
+ query = query.where(_build_loop_pass_filter(loop_pass))
+ if top_level_only:
+ if is_mapped:
+ top_level_region = (
+ select(DynamicRegion.id)
+ .where(DynamicRegion.id == TaskInstance.region_id, DynamicRegion.parent_region_id.is_(None))
+ .correlate(TaskInstance)
+ .exists()
+ )
+ query = query.where(or_(TaskInstance.region_id == SENTINEL_REGION_ID, top_level_region))
+ else:
+ query = query.where(TaskInstance.region_id == SENTINEL_REGION_ID, TaskInstance.region_index == -1)
+ if map_indexes is None:
+ return query
+ if not is_mapped:
+ wanted = map_indexes == -1 if isinstance(map_indexes, int) else -1 in map_indexes
+ return query if wanted else query.where(false())
+ if isinstance(map_indexes, int):
+ return query.where(TaskInstance.region_index == map_indexes)
+ if isinstance(map_indexes, range) and map_indexes.step == 1:
+ return query.where(
+ TaskInstance.region_index >= map_indexes.start, TaskInstance.region_index < map_indexes.stop
+ )
+ return query.where(TaskInstance.region_index.in_(list(map_indexes)))
+
+
+def _validate_producer_request(
+ *, context: ProducerContext | None, region_id: UUID | None, region_index: int | None
+) -> None:
if region_index is not None and region_id is None:
raise ValueError("region_index requires an explicit producer region_id")
if context and context.previous_iteration and context.loop_node_id is None:
raise ValueError("Previous-iteration lookup requires a loop context")
+
+
+def _build_candidate_query(
+ query: Select[Any],
+ *,
+ dag_id: str,
+ run_id: str,
+ task_id: str,
+ is_mapped: bool,
+ context: ProducerContext | None,
+ map_indexes: int | Collection[int] | None,
+ region_id: UUID | None,
+ region_index: int | None,
+ session: Session,
+) -> Select[Any] | None:
+ """Narrow ``query`` to the candidate producers of the caller's scope, or ``None`` if there are none."""
regions: dict[UUID, DynamicRegion] = {}
position = None
if context and region_id is None:
- regions = load_region_ancestry({context.region_id}, dag_id=dag_id, run_id=run_id, session=session)
+ regions = load_region_ancestry(
+ {context.region_id} - {SENTINEL_REGION_ID}, dag_id=dag_id, run_id=run_id, session=session
+ )
if context.loop_node_id is not None:
position = loop_position(regions, context.region_id, context.region_index, context.loop_node_id)
if position is None:
@@ -496,47 +561,239 @@ def resolve_current_producers(
if context.previous_iteration:
position = position[0], position[1] - 1
if position[1] < 0:
- return ()
- query = select(TaskInstance).where(
- TaskInstance.dag_id == dag_id,
- TaskInstance.run_id == run_id,
- TaskInstance.task_id == task_id,
- TaskInstance.working_set.is_(True),
+ return None
+
+ return _filter_producers(
+ query,
+ dag_id=dag_id,
+ run_id=run_id,
+ task_id=task_id,
+ is_mapped=is_mapped,
+ map_indexes=map_indexes,
+ region_id=region_id,
+ region_index=region_index,
+ top_level_only=region_id is None and position is None,
+ loop_pass=_load_loop_pass_regions(position, session=session) if position is not None else None,
)
- if region_id is not None:
- query = query.where(TaskInstance.region_id == region_id)
- if region_index is not None:
- query = query.where(TaskInstance.region_index == region_index)
- if position is not None:
- query = query.where(_build_loop_pass_filter(_load_loop_pass_regions(position, session=session)))
+
+
+def _check_producers_unambiguous(
+ region_indexes: Iterable[int], *, dag_id: str, run_id: str, task_id: str, is_mapped: bool
+) -> None:
+ seen: set[int] = set()
+ for region_index in region_indexes:
+ public_index = region_index if is_mapped else -1
+ if public_index in seen:
+ raise AmbiguousProducerError(
+ f"Multiple live producers for {dag_id}/{run_id}/{task_id} index {public_index}"
+ )
+ seen.add(public_index)
+
+
+def resolve_current_producers(
+ *,
+ dag_id: str,
+ run_id: str,
+ task_id: str,
+ is_mapped: bool,
+ context: ProducerContext | None = None,
+ map_indexes: int | Collection[int] | None = None,
+ region_id: UUID | None = None,
+ region_index: int | None = None,
+ session: Session,
+) -> tuple[TaskInstance, ...]:
+ """Resolve the live producer task instances whose data the caller reads by task instance UUID."""
+ from airflow.models.taskinstance import TaskInstance
+
+ _validate_producer_request(context=context, region_id=region_id, region_index=region_index)
+ query = _build_candidate_query(
+ select(TaskInstance),
+ dag_id=dag_id,
+ run_id=run_id,
+ task_id=task_id,
+ is_mapped=is_mapped,
+ context=context,
+ map_indexes=map_indexes,
+ region_id=region_id,
+ region_index=region_index,
+ session=session,
+ )
+ if query is None:
+ return ()
candidates = session.scalars(query).all()
- if region_id is None and position is None:
- regions.update(
- load_region_ancestry(
- {ti.region_id for ti in candidates} - set(regions),
+ _check_producers_unambiguous(
+ (ti.region_index for ti in candidates),
+ dag_id=dag_id,
+ run_id=run_id,
+ task_id=task_id,
+ is_mapped=is_mapped,
+ )
+ return tuple(sorted(candidates, key=lambda ti: ti.region_index if is_mapped else -1))
+
+
+def select_current_producer_ids(
+ *,
+ dag_id: str,
+ run_id: str,
+ task_id: str,
+ is_mapped: bool,
+ context: ProducerContext | None = None,
+ map_indexes: int | Collection[int] | None = None,
+ region_id: UUID | None = None,
+ region_index: int | None = None,
+ session: Session,
+) -> Select[tuple[UUID]]:
+ """
+ Select the ids of the live producers :func:`resolve_current_producers` would return.
+
+ A task whose live executions all share one region cannot have two producers at one index, so
+ that common case stays a pure SQL selection whose cost does not depend on the number of
+ mapped instances. Everything else resolves the rows and pins their ids.
+ """
+ from airflow.models.taskinstance import TaskInstance
+
+ _validate_producer_request(context=context, region_id=region_id, region_index=region_index)
+ filter_producers = partial(
+ _filter_producers,
+ dag_id=dag_id,
+ run_id=run_id,
+ task_id=task_id,
+ is_mapped=is_mapped,
+ map_indexes=map_indexes,
+ region_id=region_id,
+ region_index=region_index,
+ top_level_only=region_id is None,
+ )
+ query = filter_producers(select(TaskInstance.id))
+ if region_id is None:
+ if context is not None:
+ single_region = False
+ else:
+ first_region = filter_producers(select(TaskInstance.region_id)).limit(1)
+ other_region = filter_producers(select(TaskInstance.id)).where(
+ TaskInstance.region_id != first_region.scalar_subquery()
+ )
+ single_region = not session.scalar(select(other_region.exists()))
+ if not single_region:
+ candidate_query = _build_candidate_query(
+ select(TaskInstance.id, TaskInstance.region_index),
dag_id=dag_id,
run_id=run_id,
+ task_id=task_id,
+ is_mapped=is_mapped,
+ context=context,
+ map_indexes=map_indexes,
+ region_id=None,
+ region_index=None,
session=session,
)
+ candidates = session.execute(candidate_query).all() if candidate_query is not None else []
+ _check_producers_unambiguous(
+ (candidate.region_index for candidate in candidates),
+ dag_id=dag_id,
+ run_id=run_id,
+ task_id=task_id,
+ is_mapped=is_mapped,
+ )
+ return select(TaskInstance.id).where(
+ TaskInstance.id.in_([candidate.id for candidate in candidates])
+ )
+ return query
+
+
+def _load_loop_pass_owners(
+ *, dag_id: str, run_id: str, loop_node_id: str, session: Session
+) -> list[tuple[UUID, int | None]]:
+ """
+ Return each region that holds passes of a loop with the first pass it no longer owns.
+
+ A fork takes over every pass from its ``resumes_from_index``, so a region owns only the passes
+ before the earliest fork that follows it. The newest region owns every pass from its own start.
+ """
+ regions = session.scalars(
+ select(DynamicRegion).where(
+ DynamicRegion.dag_id == dag_id,
+ DynamicRegion.run_id == run_id,
+ DynamicRegion.node_id == loop_node_id,
+ )
+ ).all()
+ originals = [region for region in regions if region.forked_from_region_id is None]
+ if len(originals) > 1:
+ raise ValueError(f"Loop {loop_node_id!r} has several executions in Dag run {run_id!r}")
+ successors = {region.forked_from_region_id: region for region in regions}
+ chain = originals[:]
+ while chain and chain[-1].id in successors:
+ chain.append(successors[chain[-1].id])
+ owners: list[tuple[UUID, int | None]] = []
+ cut: int | None = None
+ for region in reversed(chain):
+ owners.append((region.id, cut))
+ cut = region.resumes_from_index if cut is None else min(cut, region.resumes_from_index)
+ return owners
+
+
+def _build_pass_ownership_filter(
+ region_column, index_column, owners: Iterable[tuple[UUID, int | None]]
+) -> ColumnElement[bool]:
+ return or_(
+ *(
+ region_column == region_id
+ if cut is None
+ else and_(region_column == region_id, index_column < cut)
+ for region_id, cut in owners
)
+ )
- selected: dict[int, TaskInstance] = {}
- for ti in candidates:
- if region_id is None and position is None:
- if ti.region_id != SENTINEL_REGION_ID:
- if not is_mapped or regions[ti.region_id].parent_region_id is not None:
- continue
- elif not is_mapped and ti.region_index != -1:
- continue
- public_index = ti.region_index if is_mapped else -1
- if isinstance(map_indexes, int):
- if public_index != map_indexes:
- continue
- elif map_indexes is not None and public_index not in map_indexes:
- continue
- if public_index in selected:
- raise AmbiguousProducerError(
- f"Multiple live producers for {dag_id}/{run_id}/{task_id} index {public_index}"
- )
- selected[public_index] = ti
- return tuple(selected[index] for index in sorted(selected))
+
+def select_loop_producer_ids(
+ *, dag_id: str, run_id: str, loop_node_id: str, task_id: str, is_mapped: bool, session: Session
+) -> Select[tuple[UUID]]:
+ """
+ Select every live execution of ``task_id`` across all passes of the loop ``loop_node_id``.
+
+ That is what a task outside the loop sees of a loop task. A pass that a fork replaced does not count
+ even while its task instances are still live, because the fork's task instances stand in for it.
+ Order the executions with :func:`build_loop_sequence_order`.
+ """
+ from airflow.models.taskinstance import TaskInstance
+
+ query = select(TaskInstance.id).where(
+ TaskInstance.dag_id == dag_id,
+ TaskInstance.run_id == run_id,
+ TaskInstance.task_id == task_id,
+ TaskInstance.working_set.is_(True),
+ )
+ owners = _load_loop_pass_owners(dag_id=dag_id, run_id=run_id, loop_node_id=loop_node_id, session=session)
+ if not owners:
+ return query.where(false())
+ if not is_mapped:
+ return query.where(
+ _build_pass_ownership_filter(TaskInstance.region_id, TaskInstance.region_index, owners)
+ )
+ pass_region = (
+ select(DynamicRegion.id)
+ .where(
+ DynamicRegion.id == TaskInstance.region_id,
+ _build_pass_ownership_filter(
+ DynamicRegion.parent_region_id, DynamicRegion.parent_region_index, owners
+ ),
+ )
+ .correlate(TaskInstance)
+ .exists()
+ )
+ return query.where(pass_region)
+
+
+def build_loop_sequence_order(region_id_column, region_index_column) -> tuple[ColumnElement[int], ...]:
+ """
+ Order the rows :func:`select_loop_producer_ids` selects depth-first: by pass, then by map index.
+
+ An unmapped task sits in the loop's own region, whose ``region_index`` is the pass. A mapped task sits
+ in a region nested under the pass, whose parent index is the pass and whose rows carry the map index.
+ """
+ nested_pass = (
+ select(DynamicRegion.parent_region_index)
+ .where(DynamicRegion.id == region_id_column)
+ .scalar_subquery()
+ )
+ return func.coalesce(nested_pass, region_index_column), region_index_column
diff --git a/airflow-core/src/airflow/models/task_coordinates.py b/airflow-core/src/airflow/models/task_coordinates.py
index 7f7930192f1..40bdcc89079 100644
--- a/airflow-core/src/airflow/models/task_coordinates.py
+++ b/airflow-core/src/airflow/models/task_coordinates.py
@@ -34,6 +34,8 @@ from airflow.models.dynamic_region import (
load_region_ancestry,
loop_position,
resolve_current_producers,
+ select_current_producer_ids,
+ select_loop_producer_ids,
)
from airflow.models.taskinstance import TaskInstance
from airflow.serialization.definitions.mappedoperator import is_mapped
@@ -49,6 +51,23 @@ if TYPE_CHECKING:
from airflow.serialization.definitions.dag import SerializedDAG, SerializedOperator
+LOOP_GATE_OPERATOR = "LoopGateOperator"
+"""Operator class name the Task SDK gives the gate task of a loop."""
+
+
+@attrs.frozen(kw_only=True)
+class _ProducerRequest:
+ dag_id: str
+ run_id: str
+ task_id: str
+ is_mapped: bool
+ context: ProducerContext | None
+ map_indexes: int | Collection[int] | None
+ region_id: UUID | None
+ region_index: int | None
+ loop_node_id: str | None = None
+
+
class TaskCoordinate(Protocol):
"""Stored task identity shared by live and retained task data."""
@@ -57,6 +76,7 @@ class TaskCoordinate(Protocol):
task_id: str
region_id: UUID
region_index: int
+ dag_version_id: UUID | None
def enclosing_loop(task: SerializedOperator) -> SerializedLoopTaskGroup | None:
@@ -156,6 +176,21 @@ class TaskCoordinateResolver:
if dag is not None and dag.dag_version_id is not None:
self._dags.setdefault(dag.dag_version_id, dag)
+ def loop_context(self, ti: TaskCoordinate) -> tuple[SerializedLoopTaskGroup, int] | None:
+ if ti.region_id == SENTINEL_REGION_ID:
+ return None
+ task = self.get_task(ti.dag_id, ti.run_id, ti.task_id, dag_version_id=ti.dag_version_id)
+ group = enclosing_loop(task)
+ if group is None:
+ return None
+ regions = load_region_ancestry(
+ [ti.region_id], dag_id=ti.dag_id, run_id=ti.run_id, session=self.session
+ )
+ position = loop_position(regions, ti.region_id, ti.region_index, group.node_id)
+ if position is None:
+ raise ValueError("Task coordinates do not belong to the pinned loop")
+ return group, position[1]
+
def get_task(
self, dag_id: str, run_id: str, task_id: str, *, dag_version_id: UUID | None = None
) -> SerializedOperator:
@@ -371,39 +406,137 @@ class TaskCoordinateResolver:
query = query.where(TaskInstance.run_id == run_id)
return bool(self.session.scalar(select(query.exists())))
- def resolve(
+ def resolve_dependency(self, caller: TaskInstance, task_id: str) -> tuple[TaskInstance, ...]:
+ try:
+ producer = self.get_task(
+ caller.dag_id, caller.run_id, task_id, dag_version_id=caller.dag_version_id
+ )
+ except TaskNotFound:
+ return self.resolve(dag_id=caller.dag_id, run_id=caller.run_id, task_id=task_id, caller=caller)
+ loop = enclosing_loop(producer)
+ caller_task = self.get_task(
+ caller.dag_id, caller.run_id, caller.task_id, dag_version_id=caller.dag_version_id
+ )
+ caller_loop = enclosing_loop(caller_task)
+ if (
+ loop is not None
+ and loop.gate_task_id == task_id
+ and (caller_loop is None or caller_loop.group_id != loop.group_id)
+ ):
+ gates = self.session.scalars(
+ select(TaskInstance)
+ .join(DynamicRegion, DynamicRegion.id == TaskInstance.region_id)
+ .where(
+ TaskInstance.working_set.is_(True),
+ TaskInstance.dag_id == caller.dag_id,
+ TaskInstance.run_id == caller.run_id,
+ TaskInstance.task_id == task_id,
+ DynamicRegion.node_id == loop.group_id,
+ )
+ .order_by(TaskInstance.region_index.desc())
+ .limit(2)
+ ).all()
+ if len(gates) == 2 and gates[0].region_index == gates[1].region_index:
+ raise AmbiguousProducerError(f"Multiple live loop gates at pass {gates[0].region_index}")
+ return tuple(gates[:1])
+ return self.resolve(
+ dag_id=caller.dag_id, run_id=caller.run_id, task_id=task_id, caller=caller, all_iterations=True
+ )
+
+ @staticmethod
+ def _filter_region_indexes(
+ query: Select, region_index: int | None, map_indexes: int | Collection[int] | None
+ ) -> Select:
+ if region_index is not None:
+ query = query.where(TaskInstance.region_index == region_index)
+ if isinstance(map_indexes, int):
+ query = query.where(TaskInstance.region_index == map_indexes)
+ elif map_indexes is not None:
+ query = query.where(TaskInstance.region_index.in_(map_indexes))
+ return query
+
+ def _filter_legacy_producers(
self,
+ query: Select,
*,
dag_id: str,
run_id: str,
task_id: str,
- caller: TaskInstance | None = None,
- region_id: UUID | None = None,
- region_index: int | None = None,
- map_indexes: int | Collection[int] | None = None,
- previous_iteration: bool = False,
- ) -> tuple[TaskInstance, ...]:
- if region_index is not None and region_id is None:
- raise ValueError("region_index requires an explicit producer region_id")
- if not previous_iteration and (
+ region_index: int | None,
+ map_indexes: int | Collection[int] | None,
+ ) -> Select:
+ query = query.where(
+ TaskInstance.working_set.is_(True),
+ TaskInstance.dag_id == dag_id,
+ TaskInstance.run_id == run_id,
+ TaskInstance.task_id == task_id,
+ TaskInstance.region_id == SENTINEL_REGION_ID,
+ )
+ return self._filter_region_indexes(query, region_index, map_indexes)
+
+ def _filter_removed_task_producers(
+ self,
+ query: Select,
+ *,
+ dag_id: str,
+ run_id: str,
+ task_id: str,
+ region_id: UUID | None,
+ region_index: int | None,
+ map_indexes: int | Collection[int] | None,
+ ) -> Select:
+ """Find live rows of a task whose definition is gone, using only its own expansion regions."""
+ query = query.join(DynamicRegion, DynamicRegion.id == TaskInstance.region_id).where(
+ TaskInstance.working_set.is_(True),
+ TaskInstance.dag_id == dag_id,
+ TaskInstance.run_id == run_id,
+ TaskInstance.task_id == task_id,
+ DynamicRegion.node_id == task_id,
+ )
+ if region_id is not None:
+ query = query.where(TaskInstance.region_id == region_id)
+ return self._filter_region_indexes(query, region_index, map_indexes)
+
+ def _is_legacy_lookup(
+ self, dag_id: str, run_id: str, task_id: str, region_id: UUID | None, previous_iteration: bool
+ ) -> bool:
+ return not previous_iteration and (
region_id == SENTINEL_REGION_ID
or (region_id is None and not self.has_regions(dag_id, run_id, task_id))
- ):
- query = select(TaskInstance).where(
- TaskInstance.working_set.is_(True),
- TaskInstance.dag_id == dag_id,
- TaskInstance.run_id == run_id,
- TaskInstance.task_id == task_id,
- TaskInstance.region_id == SENTINEL_REGION_ID,
- )
- if region_index is not None:
- query = query.where(TaskInstance.region_index == region_index)
- if isinstance(map_indexes, int):
- query = query.where(TaskInstance.region_index == map_indexes)
- elif map_indexes is not None:
- query = query.where(TaskInstance.region_index.in_(map_indexes))
- return tuple(self.session.scalars(query.order_by(TaskInstance.region_index)))
+ )
+
+ def _is_outside_loop(self, loop: SerializedLoopTaskGroup, caller: TaskInstance) -> bool:
+ caller_loop = enclosing_loop(
+ self.get_task(caller.dag_id, caller.run_id, caller.task_id, dag_version_id=caller.dag_version_id)
+ )
+ return caller_loop is None or caller_loop.group_id != loop.group_id
+
+ def is_loop_task_read_from_outside(
+ self, *, dag_id: str, run_id: str, task_id: str, caller: TaskInstance | None
+ ) -> bool:
+ """Tell whether ``caller``, of the same Dag run and outside the loop, reads a task of that loop."""
+ if caller is None or (caller.dag_id, caller.run_id) != (dag_id, run_id):
+ return False
+ try:
+ task = self.get_task(dag_id, run_id, task_id, dag_version_id=caller.dag_version_id)
+ except TaskNotFound:
+ return False
+ loop = enclosing_loop(task)
+ return loop is not None and self._is_outside_loop(loop, caller)
+ def _build_producer_request(
+ self,
+ *,
+ dag_id: str,
+ run_id: str,
+ task_id: str,
+ caller: TaskInstance | None,
+ region_id: UUID | None,
+ region_index: int | None,
+ map_indexes: int | Collection[int] | None,
+ previous_iteration: bool,
+ all_iterations: bool = False,
+ ) -> _ProducerRequest | None:
shared_run = caller is not None and (caller.dag_id, caller.run_id) == (dag_id, run_id)
try:
task = (
@@ -414,30 +547,36 @@ class TaskCoordinateResolver:
)
)
except TaskNotFound:
- return self._resolve_removed_task(
- dag_id=dag_id,
- run_id=run_id,
- task_id=task_id,
- region_id=region_id,
- region_index=region_index,
- map_indexes=map_indexes,
- )
+ return None
loop = enclosing_loop(task)
context = None
if region_id is None and loop is not None:
if caller is None or (caller.dag_id, caller.run_id) != (dag_id, run_id):
raise ValueError("A loop producer requires an explicit scope or a consumer inside the loop")
- caller_loop = enclosing_loop(
- self.get_task(dag_id, run_id, caller.task_id, dag_version_id=caller.dag_version_id)
- )
- if caller_loop is None or caller_loop.group_id != loop.group_id:
- raise ValueError("A loop producer requires an explicit scope or a consumer inside the loop")
+ if self._is_outside_loop(loop, caller):
+ if not all_iterations:
+ raise ValueError(
+ "A loop producer requires an explicit scope or a consumer inside the loop"
+ )
+ if previous_iteration:
+ raise ValueError("Previous-iteration lookup requires a shared loop scope")
+ return _ProducerRequest(
+ dag_id=dag_id,
+ run_id=run_id,
+ task_id=task_id,
+ is_mapped=task.get_needs_expansion(),
+ context=None,
+ map_indexes=map_indexes,
+ region_id=None,
+ region_index=None,
+ loop_node_id=loop.node_id,
+ )
context = ProducerContext(
caller.region_id, caller.region_index, loop.group_id, previous_iteration
)
if previous_iteration and context is None:
raise ValueError("Previous-iteration lookup requires a shared loop scope")
- return resolve_current_producers(
+ return _ProducerRequest(
dag_id=dag_id,
run_id=run_id,
task_id=task_id,
@@ -446,37 +585,187 @@ class TaskCoordinateResolver:
map_indexes=map_indexes,
region_id=region_id,
region_index=region_index,
- session=self.session,
)
- def _resolve_removed_task(
+ def resolve(
self,
*,
dag_id: str,
run_id: str,
task_id: str,
- region_id: UUID | None,
- region_index: int | None,
- map_indexes: int | Collection[int] | None,
+ caller: TaskInstance | None = None,
+ region_id: UUID | None = None,
+ region_index: int | None = None,
+ map_indexes: int | Collection[int] | None = None,
+ previous_iteration: bool = False,
+ all_iterations: bool = False,
) -> tuple[TaskInstance, ...]:
- """Find live rows of a task whose definition is gone, using only its own expansion regions."""
- query = (
- select(TaskInstance)
- .join(DynamicRegion, DynamicRegion.id == TaskInstance.region_id)
- .where(
- TaskInstance.working_set.is_(True),
- TaskInstance.dag_id == dag_id,
- TaskInstance.run_id == run_id,
- TaskInstance.task_id == task_id,
- DynamicRegion.node_id == task_id,
+ """
+ Resolve the live task instances of ``task_id`` that ``caller`` reads.
+
+ With ``all_iterations``, a caller of the same Dag run that sits outside the loop of ``task_id``
+ reads every live iteration of it instead of being refused.
+ """
+ if region_index is not None and region_id is None:
+ raise ValueError("region_index requires an explicit producer region_id")
+ if self._is_legacy_lookup(dag_id, run_id, task_id, region_id, previous_iteration):
+ query = self._filter_legacy_producers(
+ select(TaskInstance),
+ dag_id=dag_id,
+ run_id=run_id,
+ task_id=task_id,
+ region_index=region_index,
+ map_indexes=map_indexes,
)
+ return tuple(self.session.scalars(query.order_by(TaskInstance.region_index)))
+ request = self._build_producer_request(
+ dag_id=dag_id,
+ run_id=run_id,
+ task_id=task_id,
+ caller=caller,
+ region_id=region_id,
+ region_index=region_index,
+ map_indexes=map_indexes,
+ previous_iteration=previous_iteration,
+ all_iterations=all_iterations,
+ )
+ if request is not None and request.loop_node_id is not None:
+ return tuple(
+ self.session.scalars(
+ select(TaskInstance).where(TaskInstance.id.in_(self._select_loop_producer_ids(request)))
+ )
+ )
+ if request is None:
+ return tuple(
+ self.session.scalars(
+ self._filter_removed_task_producers(
+ select(TaskInstance),
+ dag_id=dag_id,
+ run_id=run_id,
+ task_id=task_id,
+ region_id=region_id,
+ region_index=region_index,
+ map_indexes=map_indexes,
+ ).order_by(TaskInstance.region_index)
+ )
+ )
+ return resolve_current_producers(
+ dag_id=request.dag_id,
+ run_id=request.run_id,
+ task_id=request.task_id,
+ is_mapped=request.is_mapped,
+ context=request.context,
+ map_indexes=request.map_indexes,
+ region_id=request.region_id,
+ region_index=request.region_index,
+ session=self.session,
+ )
+
+ def select_skip_target_ids(
+ self, *, caller: TaskInstance, task_id: str, map_indexes: int | None = None
+ ) -> Select[tuple[UUID]]:
+ """
+ Select the task instance ids a skip request from ``caller`` reaches for ``task_id``.
+
+ A caller outside the loop that encloses ``task_id`` reaches every live pass and region of it.
+ """
+ try:
+ task = self.get_task(caller.dag_id, caller.run_id, task_id, dag_version_id=caller.dag_version_id)
+ except TaskNotFound:
+ task = None
+ loop = enclosing_loop(task) if task is not None else None
+ if loop is not None:
+ caller_loop = enclosing_loop(
+ self.get_task(
+ caller.dag_id, caller.run_id, caller.task_id, dag_version_id=caller.dag_version_id
+ )
+ )
+ if caller_loop is None or caller_loop.group_id != loop.group_id:
+ query = select(TaskInstance.id).where(
+ TaskInstance.working_set.is_(True),
+ TaskInstance.dag_id == caller.dag_id,
+ TaskInstance.run_id == caller.run_id,
+ TaskInstance.task_id == task_id,
+ )
+ if map_indexes is not None:
+ query = query.where(public_map_index_expression(TaskInstance) == map_indexes)
+ return query
+ return self.select_producer_ids(
+ dag_id=caller.dag_id,
+ run_id=caller.run_id,
+ task_id=task_id,
+ caller=caller,
+ map_indexes=map_indexes,
+ )
+
+ def _select_loop_producer_ids(self, request: _ProducerRequest) -> Select[tuple[UUID]]:
+ if TYPE_CHECKING:
+ assert request.loop_node_id is not None
+ return select_loop_producer_ids(
+ dag_id=request.dag_id,
+ run_id=request.run_id,
+ loop_node_id=request.loop_node_id,
+ task_id=request.task_id,
+ is_mapped=request.is_mapped,
+ session=self.session,
+ )
+
+ def select_producer_ids(
+ self,
+ *,
+ dag_id: str,
+ run_id: str,
+ task_id: str,
+ caller: TaskInstance | None = None,
+ region_id: UUID | None = None,
+ region_index: int | None = None,
+ map_indexes: int | Collection[int] | None = None,
+ previous_iteration: bool = False,
+ all_iterations: bool = False,
+ ) -> Select[tuple[UUID]]:
+ """Select the task instance ids :meth:`resolve` returns, without loading the producers."""
+ if region_index is not None and region_id is None:
+ raise ValueError("region_index requires an explicit producer region_id")
+ if self._is_legacy_lookup(dag_id, run_id, task_id, region_id, previous_iteration):
+ return self._filter_legacy_producers(
+ select(TaskInstance.id),
+ dag_id=dag_id,
+ run_id=run_id,
+ task_id=task_id,
+ region_index=region_index,
+ map_indexes=map_indexes,
+ )
+ request = self._build_producer_request(
+ dag_id=dag_id,
+ run_id=run_id,
+ task_id=task_id,
+ caller=caller,
+ region_id=region_id,
+ region_index=region_index,
+ map_indexes=map_indexes,
+ previous_iteration=previous_iteration,
+ all_iterations=all_iterations,
+ )
+ if request is not None and request.loop_node_id is not None:
+ return self._select_loop_producer_ids(request)
+ if request is None:
+ return self._filter_removed_task_producers(
+ select(TaskInstance.id),
+ dag_id=dag_id,
+ run_id=run_id,
+ task_id=task_id,
+ region_id=region_id,
+ region_index=region_index,
+ map_indexes=map_indexes,
+ )
+ return select_current_producer_ids(
+ dag_id=request.dag_id,
+ run_id=request.run_id,
+ task_id=request.task_id,
+ is_mapped=request.is_mapped,
+ context=request.context,
+ map_indexes=request.map_indexes,
+ region_id=request.region_id,
+ region_index=request.region_index,
+ session=self.session,
)
- if region_id is not None:
- query = query.where(TaskInstance.region_id == region_id)
- if region_index is not None:
- query = query.where(TaskInstance.region_index == region_index)
- if isinstance(map_indexes, int):
- query = query.where(TaskInstance.region_index == map_indexes)
- elif map_indexes is not None:
- query = query.where(TaskInstance.region_index.in_(map_indexes))
- return tuple(self.session.scalars(query.order_by(TaskInstance.region_index)))
diff --git a/airflow-core/src/airflow/serialization/definitions/xcom_arg.py b/airflow-core/src/airflow/serialization/definitions/xcom_arg.py
index 4891448ebaf..93327772a65 100644
--- a/airflow-core/src/airflow/serialization/definitions/xcom_arg.py
+++ b/airflow-core/src/airflow/serialization/definitions/xcom_arg.py
@@ -22,13 +22,14 @@ from functools import singledispatch
from typing import TYPE_CHECKING, Any
import attrs
-from sqlalchemy import func, select
+from sqlalchemy import func, or_
from sqlalchemy.orm import Session
-from airflow.models.dynamic_region import ProducerContext, resolve_current_producers
+from airflow.models.dynamic_region import ProducerContext, select_current_producer_ids
from airflow.models.referencemixin import ReferenceMixin
from airflow.models.xcom import XCOM_RETURN_KEY
from airflow.serialization.definitions.notset import NOTSET, is_arg_set
+from airflow.utils.db import exists_query
from airflow.utils.state import State
__all__ = ["SchedulerXComArg", "deserialize_xcom_arg", "get_task_map_length"]
@@ -175,7 +176,7 @@ def _(
if not mapped and xcom_arg.operator.get_needs_expansion():
return None
- producers = resolve_current_producers(
+ producer_ids = select_current_producer_ids(
dag_id=dag_id,
run_id=run_id,
task_id=task_id,
@@ -183,14 +184,20 @@ def _(
context=(producer_contexts or {}).get(task_id),
session=session,
)
- if not producers:
- return None
- if mapped and any(ti.state in State.unfinished for ti in producers):
- return None
+ if mapped and exists_query(
+ TaskInstance.id.in_(producer_ids),
+ # 'state' can be NULL, which "NOT IN" would silently drop.
+ or_(
+ TaskInstance.state.is_(None),
+ TaskInstance.state.in_(s.value for s in State.unfinished if s is not None),
+ ),
+ session=session,
+ ):
+ return None # Not all of the expanded tis are done yet.
read = XComModel.get_many(
run_id=run_id,
key=XCOM_RETURN_KEY,
- producer_ids=select(TaskInstance.id).where(TaskInstance.id.in_([ti.id for ti in producers])),
+ producer_ids=producer_ids,
)
entity = xcom_entity(read)
diff --git a/airflow-core/src/airflow/ti_deps/dep_context.py b/airflow-core/src/airflow/ti_deps/dep_context.py
index 43a9aee4607..4083793c067 100644
--- a/airflow-core/src/airflow/ti_deps/dep_context.py
+++ b/airflow-core/src/airflow/ti_deps/dep_context.py
@@ -169,12 +169,7 @@ class DepContext:
task_id,
)
if key not in self.producer_tis:
- self.producer_tis[key] = resolver.resolve(
- dag_id=ti.dag_id,
- run_id=ti.run_id,
- task_id=task_id,
- caller=ti,
- )
+ self.producer_tis[key] = resolver.resolve_dependency(ti, task_id)
return self.producer_tis[key]
mapped_group_skip_decisions: dict[
diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
index 7c488791b93..2bb39e8692e 100644
--- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
+++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
@@ -18,8 +18,11 @@
from __future__ import annotations
import re
+import time
+from concurrent.futures import ThreadPoolExecutor
from contextlib import nullcontext
from datetime import datetime
+from threading import Barrier, Event, local
from types import SimpleNamespace
from typing import TYPE_CHECKING
from unittest import mock
@@ -35,31 +38,36 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanE
from opentelemetry.trace import StatusCode
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
from pydantic import ValidationError
-from sqlalchemy import select, update
+from sqlalchemy import delete, event, select, update
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import Session
+from sqlalchemy.orm.exc import StaleDataError
from airflow._shared.observability.traces import OverrideableRandomIdGenerator
from airflow._shared.state import TaskScope
from airflow._shared.timezones import timezone
from airflow.api_fastapi.auth.tokens import JWTGenerator, JWTValidator
from airflow.api_fastapi.execution_api.app import lifespan
+from airflow.api_fastapi.execution_api.datamodels.taskinstance import TISuccessStatePayload
from airflow.api_fastapi.execution_api.datamodels.token import TIClaims, TIToken
-from airflow.api_fastapi.execution_api.routes.task_instances import _emit_task_span
+from airflow.api_fastapi.execution_api.routes import task_instances as task_instances_route
+from airflow.api_fastapi.execution_api.routes.task_instances import _emit_task_span, ti_update_state
+from airflow.api_fastapi.execution_api.routes.xcoms import set_xcom
from airflow.api_fastapi.execution_api.security import require_auth
from airflow.exceptions import AirflowSkipException, TaskNotFound
from airflow.jobs.scheduler_job_runner import SchedulerJobRunner
-from airflow.models import RenderedTaskInstanceFields, TaskReschedule, Trigger
+from airflow.models import DagRun, RenderedTaskInstanceFields, TaskReschedule, Trigger
from airflow.models.asset import AssetActive, AssetAliasModel, AssetEvent, AssetModel
from airflow.models.dag import DagModel
+from airflow.models.dag_version import DagVersion
from airflow.models.dagbag import DBDagBag
from airflow.models.dynamic_region import DynamicRegion
from airflow.models.log import Log
-from airflow.models.task_coordinates import TaskCoordinateResolver
+from airflow.models.task_coordinates import LOOP_GATE_OPERATOR, TaskCoordinateResolver
from airflow.models.task_state_store import TaskStateStoreModel
from airflow.models.taskinstance import TaskInstance, clear_task_instances
-from airflow.models.xcom import XComModel
+from airflow.models.xcom import XComModel, XComModelV2
from airflow.providers.standard.operators.empty import EmptyOperator
from airflow.providers.standard.operators.python import PythonOperator
from airflow.sdk import Asset, TaskGroup, TriggerRule, task, task_group
@@ -399,6 +407,59 @@ def test_execution_breadcrumbs_keep_regional_identity_separate_from_map_index(cl
class TestTIRunState:
+ @pytest.mark.parametrize("mapped", [False, True])
+ @pytest.mark.parametrize("version", ["2026-06-30", "2026-10-30"])
+ def test_startup_loop_context_uses_enclosing_iteration(self, client, session, dag_maker, mapped, version):
+ @task_group
+ def body():
+ if mapped:
+ PythonOperator.partial(task_id="terminal", python_callable=str).expand(op_args=[[1], [2]])
+ else:
+ EmptyOperator(task_id="terminal")
+
+ with dag_maker(serialized=True):
+ loop = create_loop(body, max_iterations=4)
+ dr = dag_maker.create_dagrun(state=DagRunState.RUNNING)
+ ti = next(ti for ti in dr.task_instances if ti.task_id == loop.terminal_task_id)
+ parent = DynamicRegion.get_or_create(
+ dag_id=dr.dag_id, run_id=dr.run_id, node_id=loop.group_id, session=session
+ )
+ session.add(parent)
+ session.flush()
+ ti.region_id, ti.region_index = parent.id, 2
+ if mapped:
+ child = DynamicRegion.get_or_create(
+ dag_id=dr.dag_id,
+ run_id=dr.run_id,
+ node_id=ti.task_id,
+ parent_region_id=parent.id,
+ parent_region_index=2,
+ session=session,
+ )
+ session.add(child)
+ session.flush()
+ ti.region_id, ti.region_index = child.id, 0
+ ti.state = State.QUEUED
+ session.commit()
+
+ response = client.patch(
+ f"/execution/task-instances/{ti.id}/run",
+ json=self.RUN_PAYLOAD,
+ headers={"Airflow-API-Version": version},
+ )
+
+ assert response.status_code == 200
+ if version == "2026-06-30":
+ assert "loop" not in response.json()
+ return
+ assert response.json()["loop"] == {
+ "node_id": loop.group_id,
+ "index": 2,
+ "max_iterations": 4,
+ "terminal_task_id": loop.terminal_task_id,
+ "terminal_is_mapped": mapped,
+ }
+
@pytest.mark.parametrize("regional", [False, True])
def test_startup_xcom_cleanup_keys_belong_to_the_starting_try(
self, client, session, create_task_instance, regional
@@ -1604,6 +1665,545 @@ class TestTIRunState:
class TestTIUpdateState:
+ def test_loop_continues_while_earlier_body_branch_is_running(self, client, session, dag_maker):
+ @task_group
+ def body():
+ left = PythonOperator(task_id="left", python_callable=list)
+ right = PythonOperator(task_id="right", python_callable=list)
+ terminal = PythonOperator(
+ task_id="terminal", python_callable=list, trigger_rule=TriggerRule.ONE_SUCCESS
+ )
+ [left, right] >> terminal
+
+ with dag_maker(serialized=True, session=session):
+ loop = create_loop(body, max_iterations=2)
+ dr = dag_maker.create_dagrun()
+ dr.dag = dag_maker.serialized_dag
+ tis = {ti.task_id: ti for ti in dr.task_instances}
+ tis["body.left"].state = State.SUCCESS
+ right = tis["body.right"]
+ right.state = State.RUNNING
+ right_id = right.id
+ terminal = tis["body.terminal"]
+ gate = tis[loop.gate_task_id]
+ session.flush()
+ assert terminal in dr.task_instance_scheduling_decisions(session=session).schedulable_tis
+ terminal.state, terminal.start_date = State.RUNNING, DEFAULT_START_DATE
+ session.commit()
+ assert (
+ client.patch(
+ f"/execution/task-instances/{terminal.id}/state",
+ json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat()},
+ ).status_code
+ == 204
+ )
+ session.expire_all()
+ assert gate in dr.task_instance_scheduling_decisions(session=session).schedulable_tis
+ gate.state, gate.start_date = State.RUNNING, DEFAULT_START_DATE
+ session.commit()
+ exec_app = client.app.routes[-1].app
+ exec_app.dependency_overrides[require_auth] = lambda: TIToken(id=gate.id, claims=TIClaims())
+ assert (
+ client.post(
+ f"/execution/xcoms/{dr.dag_id}/{dr.run_id}/{gate.task_id}/_airflow_loop_decision",
+ json="continue",
+ ).status_code
+ == 201
+ )
+
+ response = client.patch(
+ f"/execution/task-instances/{gate.id}/state",
+ json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat()},
+ )
+
+ assert response.status_code == 204, response.text
+ session.expire_all()
+ assert session.get(TaskInstance, right_id).state == State.RUNNING
+ ready = dr.task_instance_scheduling_decisions(session=session).schedulable_tis
+ assert {(ti.task_id, ti.region_index) for ti in ready} == {("body.left", 1), ("body.right", 1)}
+
+ @conf_vars({("state_store", "clear_on_success"): "true"})
+ @pytest.mark.parametrize("mapped", [False, True])
+ @pytest.mark.parametrize(
+ ("decision", "state", "max_iterations", "conditional", "rejected", "passes"),
+ [
+ ("continue", State.SUCCESS, 3, False, False, [0, 1]),
+ ("stop", State.SUCCESS, 3, False, True, [0]),
+ (None, State.SUCCESS, 3, False, True, [0]),
+ ("continue", State.FAILED, 3, False, False, [0]),
+ ("continue", State.SKIPPED, 3, False, False, [0]),
+ ("stop", State.SUCCESS, 1, False, False, [0]),
+ ("continue", State.SUCCESS, 1, False, True, [0]),
+ ("stop", State.SUCCESS, 3, True, False, [0]),
+ ("stop", State.SUCCESS, 1, True, False, [0]),
+ ("continue", State.SUCCESS, 1, True, True, [0]),
+ (None, State.FAILED, 1, True, False, [0]),
+ ],
+ )
+ def test_loop_gate_completion_consumes_decision_atomically(
+ self,
+ client,
+ session,
+ dag_maker,
+ decision,
+ state,
+ max_iterations,
+ conditional,
+ mapped,
+ rejected,
+ passes,
+ mocker,
+ ):
+ backend = mocker.create_autospec(MetastoreBackend, instance=True)
+ mocker.patch(
+ "airflow.api_fastapi.execution_api.routes.task_instances.get_state_backend",
+ autospec=True,
+ return_value=backend,
+ )
+
+ @task_group
+ def body():
+ if mapped:
+ PythonOperator.partial(task_id="terminal", python_callable=list).expand(op_kwargs=[{}, {}])
+ else:
+ PythonOperator(task_id="terminal", python_callable=list)
+
+ with dag_maker(serialized=True, session=session):
+ loop = create_loop(
+ body, max_iterations=max_iterations, until=(lambda loop: True) if conditional else None
+ )
+ dr = dag_maker.create_dagrun()
+ gate = next(ti for ti in dr.task_instances if ti.task_id == loop.gate_task_id)
+ gate.state = State.RUNNING
+ gate.start_date = DEFAULT_START_DATE
+ session.commit()
+ if decision and rejected:
+ XComModel.set_for_attempt(
+ task_instance_id=gate.id,
+ key="_airflow_loop_decision",
+ value=decision,
+ serialize=False,
+ session=session,
+ )
+ session.commit()
+ elif decision:
+ exec_app = client.app.routes[-1].app
+ exec_app.dependency_overrides[require_auth] = lambda: TIToken(id=gate.id, claims=TIClaims())
+ response = client.post(
+ f"/execution/xcoms/{dr.dag_id}/{dr.run_id}/{gate.task_id}/_airflow_loop_decision",
+ json=decision,
+ )
+ assert response.status_code == 201
+
+ response = client.patch(
+ f"/execution/task-instances/{gate.id}/state",
+ json={"state": state, "end_date": DEFAULT_END_DATE.isoformat()},
+ )
+
+ assert response.status_code == 204, response.text
+ assert backend.clear.call_count == (state == State.SUCCESS and not rejected)
+ session.expire_all()
+ assert gate.state == (State.FAILED if rejected else state)
+ current = dr.get_task_instances(session=session)
+ assert sorted(ti.region_index for ti in current if ti.task_id == gate.task_id) == passes
+ if mapped:
+ children = session.scalars(
+ select(DynamicRegion).where(DynamicRegion.parent_region_id == gate.region_id)
+ ).all()
+ assert sorted(region.parent_region_index for region in children) == passes
+ region_ids = {gate.region_id, *(region.id for region in children)}
+ assert all(ti.region_id in region_ids for ti in current)
+ else:
+ assert all(ti.region_id == gate.region_id for ti in current)
+ signal = session.scalar(select(XComModel).where(XComModel.task_id == gate.task_id))
+ assert signal is None
+ if not rejected:
+ duplicate = client.patch(
+ f"/execution/task-instances/{gate.id}/state",
+ json={"state": state, "end_date": DEFAULT_END_DATE.isoformat()},
+ )
+ assert duplicate.status_code == 200
+ assert (
+ sorted(
+ ti.region_index
+ for ti in dr.get_task_instances(session=session)
+ if ti.task_id == gate.task_id
+ )
+ == passes
+ )
+ if mapped and passes == [0, 1]:
+ successor_region = next(region for region in children if region.parent_region_index == 1)
+ assert [ti.region_index for ti in current if ti.region_id == successor_region.id] == [-1]
+ dr.dag = dag_maker.serialized_dag
+
+ dr.task_instance_scheduling_decisions(session=session)
+
+ assert sorted(
+ ti.region_index
+ for ti in dr.get_task_instances(session=session)
+ if ti.region_id == successor_region.id
+ ) == [0, 1]
+
+ @pytest.fixture
+ def running_loop_gate(self, session, dag_maker):
+ @task_group
+ def body():
+ PythonOperator(task_id="terminal", python_callable=list)
+
+ with dag_maker(serialized=True, session=session):
+ loop = create_loop(body, max_iterations=3)
+ dr = dag_maker.create_dagrun()
+ gate = next(ti for ti in dr.task_instances if ti.task_id == loop.gate_task_id)
+ gate.state, gate.start_date = State.RUNNING, DEFAULT_START_DATE
+ XComModel.set_for_attempt(
+ task_instance_id=gate.id,
+ key="_airflow_loop_decision",
+ value="continue",
+ serialize=False,
+ session=session,
+ )
+ session.commit()
+ return gate
+
+ def test_loop_gate_materialization_error_rolls_back_state_signal_and_successor(
+ self, client, session, running_loop_gate, mocker
+ ):
+ gate = running_loop_gate
+ dr = gate.dag_run
+ original = DagRun._create_task_instances
+ observed = []
+
+ def fail_after_insert(*args, **kwargs):
+ original(*args, **kwargs)
+ with Session(bind=session.get_bind()) as observer:
+ observed.append(observer.get(TaskInstance, gate.id).state)
+ observed.append(
+ observer.scalars(
+ select(TaskInstance.region_index).where(
+ TaskInstance.dag_id == dr.dag_id, TaskInstance.task_id == gate.task_id
+ )
+ ).all()
+ )
+ raise StaleDataError("injected materialization failure")
+
+ mocker.patch.object(DagRun, "_create_task_instances", autospec=True, side_effect=fail_after_insert)
+ response = client.patch(
+ f"/execution/task-instances/{gate.id}/state",
+ json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat()},
+ )
+
+ assert response.status_code == 500
+ session.expire_all()
+ assert observed == [State.RUNNING, [0]]
+ assert gate.state == State.RUNNING
+ assert len(dr.get_task_instances(session=session)) == 2
+ assert session.scalar(select(XComModel.value).where(XComModel.task_id == gate.task_id)) == "continue"
+
+ @pytest.mark.parametrize(
+ ("max_tries", "expected_state"),
+ [(0, State.FAILED), (1, State.UP_FOR_RETRY)],
+ )
+ def test_loop_gate_with_invalid_decision_leaves_running_in_same_request(
+ self, client, session, running_loop_gate, max_tries, expected_state
+ ):
+ gate = running_loop_gate
+ dr = gate.dag_run
+ gate_id = gate.id
+ gate.max_tries = max_tries
+ session.execute(delete(XComModelV2).where(XComModelV2.task_instance_id == gate.id))
+ session.commit()
+
+ response = client.patch(
+ f"/execution/task-instances/{gate.id}/state",
+ json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat()},
+ )
+
+ assert response.status_code == 204
+ session.expire_all()
+ current = dr.get_task_instances(session=session)
+ assert sorted(ti.state for ti in current if ti.task_id == gate.task_id) == [expected_state]
+ assert len(current) == 2
+ assert "requires a decision" in session.get(TaskInstance, gate_id).retry_reason
+
+ @pytest.mark.backend("mysql", "postgres")
+ def test_concurrent_loop_gate_completions_create_one_successor(self, session, running_loop_gate):
+ gate = running_loop_gate
+ dr = gate.dag_run
+ gate_id = gate.id
+ bind = session.get_bind()
+ barrier = Barrier(2)
+
+ def complete():
+ with Session(bind=bind) as request_session:
+ barrier.wait(timeout=10)
+ result = ti_update_state(
+ task_instance_id=gate_id,
+ ti_patch_payload=TISuccessStatePayload(state="success", end_date=DEFAULT_END_DATE),
+ session=request_session,
+ dag_bag=DBDagBag(),
+ )
+ return result.status_code if result is not None else 204
+
+ with ThreadPoolExecutor(max_workers=2) as pool:
+ futures = [pool.submit(complete) for _ in range(2)]
+ assert sorted(future.result(timeout=20) for future in futures) == [200, 204]
+ session.expire_all()
+ assert gate.state == State.SUCCESS
+ assert sorted(
+ ti.region_index for ti in dr.get_task_instances(session=session) if ti.task_id == gate.task_id
+ ) == [0, 1]
+ assert session.scalar(select(XComModel).where(XComModel.task_id == gate.task_id)) is None
+
+ @pytest.mark.backend("mysql", "postgres")
+ def test_loop_decision_rewrite_and_completion_use_same_lock_order(
+ self, session, running_loop_gate, request
+ ):
+ gate = running_loop_gate
+ gate_id, dag_id, run_id, task_id = gate.id, gate.dag_id, gate.run_id, gate.task_id
+ XComModel.set_for_attempt(
+ task_instance_id=gate_id,
+ key="_airflow_loop_decision",
+ value="stop",
+ serialize=False,
+ session=session,
+ )
+ session.commit()
+ bind = session.get_bind()
+ role = local()
+ writer_locked = Event()
+ completion_at_lock = Event()
+
+ def note_completion_lock(conn, cursor, statement, parameters, context, executemany):
+ if getattr(role, "name", None) == "completion" and "FOR UPDATE" in statement:
+ completion_at_lock.set()
+
+ def pause_writer_after_first_lock(conn, cursor, statement, parameters, context, executemany):
+ if (
+ getattr(role, "name", None) == "writer"
+ and "FOR UPDATE" in statement
+ and not writer_locked.is_set()
+ ):
+ writer_locked.set()
+ assert completion_at_lock.wait(timeout=10)
+ # Give completion time to take its first lock, so a writer that locks in a different
+ # order deadlocks against it.
+ time.sleep(0.5)
+
+ event.listen(bind, "before_cursor_execute", note_completion_lock)
+ event.listen(bind, "after_cursor_execute", pause_writer_after_first_lock)
+ request.addfinalizer(lambda: event.remove(bind, "before_cursor_execute", note_completion_lock))
+ request.addfinalizer(
+ lambda: event.remove(bind, "after_cursor_execute", pause_writer_after_first_lock)
+ )
+
+ def rewrite():
+ role.name = "writer"
+ with Session(bind=bind) as request_session:
+ set_xcom(
+ dag_id=dag_id,
+ run_id=run_id,
+ task_id=task_id,
+ key="_airflow_loop_decision",
+ session=request_session,
+ dag_bag=DBDagBag(),
+ value="continue",
+ token=TIToken(id=gate_id, claims=TIClaims()),
+ )
+ request_session.commit()
+
+ def complete():
+ role.name = "completion"
+ assert writer_locked.wait(timeout=10)
+ with Session(bind=bind) as request_session:
+ ti_update_state(
+ task_instance_id=gate_id,
+ ti_patch_payload=TISuccessStatePayload(state="success", end_date=DEFAULT_END_DATE),
+ session=request_session,
+ dag_bag=DBDagBag(),
+ )
+
+ with ThreadPoolExecutor(max_workers=2) as pool:
+ futures = [pool.submit(rewrite), pool.submit(complete)]
+ for future in futures:
+ future.result(timeout=20)
+ session.expire_all()
+ assert gate.state == State.SUCCESS
+ assert session.scalar(select(XComModel).where(XComModel.task_id == gate.task_id)) is None
+ assert sorted(
+ ti.region_index
+ for ti in gate.dag_run.get_task_instances(session=session)
+ if ti.task_id == gate.task_id
+ ) == [0, 1]
+
+ def test_loop_gate_completion_uses_its_pinned_limit(self, client, session, running_loop_gate, dag_maker):
+ gate = running_loop_gate
+ original_version = gate.dag_version_id
+
+ @task_group
+ def body():
+ PythonOperator(task_id="terminal", python_callable=list)
+
+ with dag_maker(serialized=True, session=session):
+ create_loop(body, max_iterations=1)
+
+ response = client.patch(
+ f"/execution/task-instances/{gate.id}/state",
+ json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat()},
+ )
+
+ assert response.status_code == 204, response.text
+ session.expire_all()
+ successors = [ti for ti in gate.dag_run.get_task_instances(session=session) if ti.region_index == 1]
+ assert len(successors) == 2
+ assert all(ti.dag_version_id == original_version for ti in successors)
+
+ @pytest.mark.parametrize("state", ["success", "failed"])
+ def test_loop_gate_missing_from_its_pinned_version_fails(
+ self, client, session, running_loop_gate, dag_maker, state
+ ):
+ gate = running_loop_gate
+
+ def renamed_gate(loop):
+ return True
+
+ @task_group
+ def body():
+ PythonOperator(task_id="terminal", python_callable=list)
+
+ with dag_maker(serialized=True, session=session):
+ create_loop(body, max_iterations=3, until=renamed_gate)
+ gate.dag_version_id = DagVersion.get_latest_version(gate.dag_id, session=session).id
+ session.commit()
+
+ response = client.patch(
+ f"/execution/task-instances/{gate.id}/state",
+ json={"state": state, "end_date": DEFAULT_END_DATE.isoformat()},
+ )
+
+ assert response.status_code == 204, response.text
+ session.expire_all()
+ assert gate.state == State.FAILED
+ assert [ti.region_index for ti in gate.dag_run.get_task_instances(session=session)] == [0, 0]
+
+ def test_late_loop_decision_cannot_recreate_consumed_signal(self, client, session, running_loop_gate):
+ gate = running_loop_gate
+ response = client.patch(
+ f"/execution/task-instances/{gate.id}/state",
+ json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat()},
+ )
+ assert response.status_code == 204
+ exec_app = client.app.routes[-1].app
+ exec_app.dependency_overrides[require_auth] = lambda: TIToken(id=gate.id, claims=TIClaims())
+
+ response = client.post(
+ f"/execution/xcoms/{gate.dag_id}/{gate.run_id}/{gate.task_id}/_airflow_loop_decision",
+ json="continue",
+ )
+
+ assert response.status_code == 409
+ assert session.scalar(select(XComModel).where(XComModel.task_id == gate.task_id)) is None
+
+ def test_loop_gate_state_update_error_rolls_back_next_pass_and_keeps_decision(
+ self, client, session, running_loop_gate, mocker
+ ):
+ gate = running_loop_gate
+ mocker.patch.object(
+ task_instances_route,
+ "_create_ti_state_update_query_and_update_state",
+ autospec=True,
+ side_effect=StaleDataError("injected state update failure"),
+ )
+
+ response = client.patch(
+ f"/execution/task-instances/{gate.id}/state",
+ json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat()},
+ )
+
+ assert response.status_code == 500
+ session.expire_all()
+ assert gate.state == State.RUNNING
+ assert len(gate.dag_run.get_task_instances(session=session)) == 2
+ assert session.scalar(select(XComModel.value).where(XComModel.task_id == gate.task_id)) == "continue"
+
+ @pytest.mark.parametrize("state", ["success", "failed"])
+ def test_non_loop_operator_named_loop_gate_completes_normally(
+ self, client, session, create_task_instance, state
+ ):
+ ti = create_task_instance(task_id="user_gate", start_date=DEFAULT_START_DATE, state=State.RUNNING)
+ ti.operator = LOOP_GATE_OPERATOR
+ session.commit()
+
+ response = client.patch(
+ f"/execution/task-instances/{ti.id}/state",
+ json={"state": state, "end_date": DEFAULT_END_DATE.isoformat()},
+ )
+
+ assert response.status_code == 204
+ session.expire_all()
+ assert ti.state == state
+
+ def test_loop_gate_inner_insert_error_is_not_swallowed(self, client, session, running_loop_gate, mocker):
+ gate = running_loop_gate
+ bulk_insert = mocker.patch.object(
+ Session,
+ "bulk_insert_mappings",
+ autospec=True,
+ side_effect=StaleDataError("injected insert error"),
+ )
+
+ response = client.patch(
+ f"/execution/task-instances/{gate.id}/state",
+ json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat()},
+ )
+
+ assert response.status_code == 500
+ bulk_insert.assert_called_once()
+ session.expire_all()
+ assert gate.state == State.RUNNING
+ assert len(gate.dag_run.get_task_instances(session=session)) == 2
+ assert session.scalar(select(XComModel.value).where(XComModel.task_id == gate.task_id)) == "continue"
+
+ @pytest.mark.parametrize("state", [State.SUCCESS, State.FAILED, State.SKIPPED])
+ def test_rerun_gate_preserves_existing_later_pass(self, client, session, running_loop_gate, state):
+ gate = running_loop_gate
+ response = client.patch(
+ f"/execution/task-instances/{gate.id}/state",
+ json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat()},
+ )
+ assert response.status_code == 204
+ session.refresh(gate)
+ successor_ids = {
+ ti.id for ti in gate.dag_run.get_task_instances(session=session) if ti.region_index == 1
+ }
+ archived_id = gate.id
+ dag_run = gate.dag_run
+ clear_task_instances([gate], session=session)
+ session.expire_all()
+ gate = session.scalar(
+ select(TaskInstance).where(
+ TaskInstance.dag_id == dag_run.dag_id,
+ TaskInstance.run_id == dag_run.run_id,
+ TaskInstance.task_id == gate.task_id,
+ TaskInstance.region_index == 0,
+ TaskInstance.working_set.is_(True),
+ )
+ )
+ assert gate.id != archived_id
+ gate.state = State.RUNNING
+ session.commit()
+
+ response = client.patch(
+ f"/execution/task-instances/{gate.id}/state",
+ json={"state": state, "end_date": DEFAULT_END_DATE.isoformat()},
+ )
+
+ assert response.status_code == 204, response.text
+ session.expire_all()
+ assert gate.state == state
+ assert {
+ ti.id for ti in gate.dag_run.get_task_instances(session=session) if ti.region_index == 1
+ } == successor_ids
+
@pytest.mark.parametrize(
("interruption", "expected_state", "expected_status"),
[("clear", State.RESTARTING, 409), ("failed", State.FAILED, 409)],
@@ -3459,6 +4059,38 @@ class TestTISkipDownstream:
assert target.state == State.SKIPPED
assert previous.state is None
+ @pytest.mark.parametrize("tasks", [["body.target"], [["body.target", -1]]])
+ def test_skip_downstream_from_outside_loop_skips_every_live_pass(self, client, session, dag_maker, tasks):
+ @task_group
+ def body():
+ EmptyOperator(task_id="target")
+
+ with dag_maker(serialized=True) as dag:
+ EmptyOperator(task_id="branch") >> create_loop(body, max_iterations=3)
+ dr = dag_maker.create_dagrun()
+ region = DynamicRegion.get_or_create(
+ dag_id=dr.dag_id, run_id=dr.run_id, node_id="body", session=session
+ )
+ session.add(region)
+ session.flush()
+ tis = {ti.task_id: ti for ti in dr.task_instances}
+ caller, target = tis["branch"], tis["body.target"]
+ other_pass = TaskInstance(
+ task=dag.get_task(target.task_id), run_id=dr.run_id, dag_version_id=target.dag_version_id
+ )
+ other_pass.region_id, other_pass.region_index = region.id, 1
+ session.add(other_pass)
+ session.commit()
+
+ response = client.patch(
+ f"/execution/task-instances/{caller.id}/skip-downstream", json={"tasks": tasks}
+ )
+
+ assert response.status_code == 204
+ session.expire_all()
+ assert target.state == State.SKIPPED
+ assert other_pass.state == State.SKIPPED
+
def setup_method(self):
clear_db_runs()
diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py
index a3e95dca3a5..5c8d49c3f1e 100644
--- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py
+++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py
@@ -25,7 +25,8 @@ from uuid import uuid4
import pytest
from fastapi import FastAPI, HTTPException, Path, Request, status
-from sqlalchemy import delete, select, update
+from sqlalchemy import delete, event, select, update
+from sqlalchemy.orm import Session
from airflow._shared.timezones import timezone
from airflow.api_fastapi.auth.tokens import JWTValidator
@@ -129,6 +130,478 @@ def loop_xcoms(dag_maker, session, authenticate_as):
return dr, producer, consumer, previous
+@pytest.mark.parametrize("suffix", ["", "/item/0", "/item/-1", "/slice"])
+def test_previous_iteration_xcom_uses_retained_predecessor(client, loop_xcoms, suffix):
+ dr, producer, _, _ = loop_xcoms
+ url = f"/execution/xcoms/{dr.dag_id}/{dr.run_id}/{producer.task_id}/key"
+ params = {"previous_iteration": True}
+
+ response = client.get(url + suffix, params=params)
+
+ assert response.status_code == 200
+ expected = {"": {"key": "key", "value": "previous"}, "/slice": ["previous"]}
+ assert response.json() == expected.get(suffix, "previous")
+ assert client.head(url, params=params).headers["Content-Range"] == "map_indexes 1"
+
+
+def test_previous_iteration_rejects_prior_dates(client, loop_xcoms):
+ dr, producer, _, _ = loop_xcoms
+
+ response = client.get(
+ f"/execution/xcoms/{dr.dag_id}/{dr.run_id}/{producer.task_id}/key",
+ params={"previous_iteration": True, "include_prior_dates": True},
+ )
+
+ assert response.status_code == 400
+
+
+def test_previous_mapped_iteration_reads_retained_slots_before_count_and_slice(
+ client, dag_maker, session, authenticate_as
+):
+ @task_group
+ def body():
+ PythonOperator.partial(task_id="mapped", python_callable=str).expand(
+ op_args=[[1], [2]]
+ ) >> EmptyOperator(task_id="consumer")
+
+ with dag_maker(serialized=True) as dag:
+ create_loop(body, max_iterations=3)
+ dr = dag_maker.create_dagrun()
+ parent = DynamicRegion.get_or_create(dag_id=dr.dag_id, run_id=dr.run_id, node_id="body", session=session)
+ session.add(parent)
+ session.flush()
+ previous = DynamicRegion.get_or_create(
+ dag_id=dr.dag_id,
+ run_id=dr.run_id,
+ node_id="body.mapped",
+ parent_region_id=parent.id,
+ parent_region_index=0,
+ session=session,
+ )
+ current = DynamicRegion.get_or_create(
+ dag_id=dr.dag_id,
+ run_id=dr.run_id,
+ node_id="body.mapped",
+ parent_region_id=parent.id,
+ parent_region_index=1,
+ session=session,
+ )
+ session.add_all([previous, current])
+ session.flush()
+ replacement = DynamicRegion(
+ dag_id=dr.dag_id,
+ run_id=dr.run_id,
+ node_id="body.mapped",
+ parent_region_id=parent.id,
+ parent_region_index=0,
+ forked_from_region_id=previous.id,
+ )
+ session.add(replacement)
+ session.flush()
+ consumer = next(ti for ti in dr.task_instances if ti.task_id == "body.consumer")
+ consumer.region_id, consumer.region_index = parent.id, 1
+ for ti in dr.task_instances:
+ if ti.task_id == "body.mapped":
+ ti.region_id = current.id
+ producers = {}
+ for region, index in [(previous, 0), (previous, 1), (replacement, 1)]:
+ ti = TaskInstance(
+ task=dag.get_task("body.mapped"), run_id=dr.run_id, dag_version_id=consumer.dag_version_id
+ )
+ ti.region_id, ti.region_index = region.id, index
+ session.add(ti)
+ producers[region.id, index] = ti
+ session.flush()
+ producers[previous.id, 1].archive(reason="cleared", session=session)
+ for (region, index), value in [
+ ((previous, 0), "retained"),
+ ((previous, 1), "archived"),
+ ((replacement, 1), "replacement"),
+ ]:
+ XComModel.set_for_attempt(
+ task_instance_id=producers[region.id, index].id,
+ key="key",
+ value=value,
+ serialize=False,
+ session=session,
+ )
+ session.commit()
+ authenticate_as(consumer)
+ url = f"/execution/xcoms/{dr.dag_id}/{dr.run_id}/body.mapped/key"
+ params = {"previous_iteration": True}
+ assert client.head(url, params=params).headers["Content-Range"] == "map_indexes 2"
+ assert client.get(url + "/slice", params=params).json() == ["retained", "replacement"]
+ assert client.get(url + "/item/-1", params=params).json() == "replacement"
+
+
+LOOP_SEQUENCES = {
+ "unmapped": ("body.step", ["step-0", "step-1", "step-2"]),
+ "mapped": ("body.fanout", ["0-0", "0-1", "1-0", "1-1", "2-0", "2-1"]),
+}
+
+
+@pytest.fixture
+def outside_loop_run(dag_maker, session, authenticate_as):
+ @task_group
+ def body():
+ EmptyOperator(task_id="step") >> PythonOperator.partial(task_id="fanout", python_callable=str).expand(
+ op_args=[[1], [2]]
+ )
+
+ with dag_maker(serialized=True) as dag:
+ loop = create_loop(body, max_iterations=3)
+ PythonOperator(
+ task_id="outside",
+ python_callable=list,
+ op_args=[loop["step"].output, loop["fanout"].output],
+ )
+ dr = dag_maker.create_dagrun()
+ root = session.scalar(
+ select(DynamicRegion).where(
+ DynamicRegion.dag_id == dr.dag_id,
+ DynamicRegion.run_id == dr.run_id,
+ DynamicRegion.node_id == "body",
+ )
+ )
+ tis = {(ti.task_id, ti.region_id, ti.region_index): ti for ti in dr.task_instances}
+ steps = {0: tis["body.step", root.id, 0]}
+ fanouts = {(0, ti.region_index): ti for ti in dr.task_instances if ti.task_id == "body.fanout"}
+ for iteration in (1, 2):
+ steps[iteration] = TaskInstance(
+ dag.get_task("body.step"),
+ dr.created_dag_version_id,
+ run_id=dr.run_id,
+ region_id=root.id,
+ region_index=iteration,
+ )
+ nested = DynamicRegion.get_or_create(
+ dag_id=dr.dag_id,
+ run_id=dr.run_id,
+ node_id="body.fanout",
+ parent_region_id=root.id,
+ parent_region_index=iteration,
+ session=session,
+ )
+ session.add_all([steps[iteration], nested])
+ session.flush()
+ for slot in range(2):
+ fanouts[iteration, slot] = TaskInstance(
+ dag.get_task("body.fanout"),
+ dr.created_dag_version_id,
+ run_id=dr.run_id,
+ region_id=nested.id,
+ region_index=slot,
+ )
+ session.add(fanouts[iteration, slot])
+ session.flush()
+ for iteration in (2, 1, 0):
+ XComModel.set_for_attempt(
+ task_instance_id=steps[iteration].id,
+ key="key",
+ value=f"step-{iteration}",
+ serialize=False,
+ session=session,
+ )
+ XComModel.set_for_attempt(
+ task_instance_id=steps[iteration].id,
+ key=XCOM_RETURN_KEY,
+ value=f"returned-{iteration}",
+ serialize=False,
+ session=session,
+ )
+ for slot in (1, 0):
+ XComModel.set_for_attempt(
+ task_instance_id=fanouts[iteration, slot].id,
+ key="key",
+ value=f"{iteration}-{slot}",
+ serialize=False,
+ session=session,
+ )
+ session.commit()
+ outside = next(ti for ti in dr.task_instances if ti.task_id == "outside")
+ authenticate_as(outside)
+ return dr, dag, root, steps, authenticate_as
+
+
+def get_loop_url(dr, task_id):
+ return f"/execution/xcoms/{dr.dag_id}/{dr.run_id}/{task_id}/key"
+
+
+@pytest.mark.parametrize("shape", LOOP_SEQUENCES)
+def test_outside_caller_reads_every_iteration_of_a_loop_task(client, outside_loop_run, shape):
+ dr = outside_loop_run[0]
+ task_id, expected = LOOP_SEQUENCES[shape]
+ url = get_loop_url(dr, task_id)
+
+ assert client.head(url).headers["Content-Range"] == f"map_indexes {len(expected)}"
+ assert client.get(f"{url}/slice").json() == expected
+ for offset, value in enumerate(expected):
+ assert client.get(f"{url}/item/{offset}").json() == value
+ assert client.get(f"{url}/item/{offset - len(expected)}").json() == value
+ assert client.get(f"{url}/item/{len(expected)}").status_code == 404
+ assert client.get(url, params={"offset": -1}).json() == {"key": "key", "value": expected[-1]}
+
+
+def test_outside_task_that_uses_loop_outputs_waits_for_the_gate_and_reads_the_loop_task_xcom(
+ client, dag_maker, outside_loop_run
+):
+ dr = outside_loop_run[0]
+ consumer = dag_maker.serialized_dag.get_task("outside")
+
+ assert consumer.upstream_task_ids == {"body.__loop_gate"}
+ url = f"/execution/xcoms/{dr.dag_id}/{dr.run_id}/body.step/{XCOM_RETURN_KEY}"
+ assert client.head(url).headers["Content-Range"] == "map_indexes 3"
+ assert client.get(f"{url}/slice").json() == ["returned-0", "returned-1", "returned-2"]
+ assert client.get(f"{url}/item/-1").json() == "returned-2"
+
+
+@pytest.mark.parametrize("shape", LOOP_SEQUENCES)
+@pytest.mark.parametrize(
+ "params",
+ [
+ {"start": 1},
+ {"stop": 3},
+ {"stop": -1},
+ {"start": -2},
+ {"start": 1, "stop": -1},
+ {"start": -3, "stop": 2},
+ {"step": -1},
+ {"step": 2},
+ ],
+)
+def test_outside_caller_slices_a_loop_task_depth_first(client, outside_loop_run, shape, params):
+ dr = outside_loop_run[0]
+ task_id, expected = LOOP_SEQUENCES[shape]
+
+ response = client.get(f"{get_loop_url(dr, task_id)}/slice", params=params)
+
+ assert response.json() == expected[slice(params.get("start"), params.get("stop"), params.get("step"))]
+
+
+@pytest.mark.parametrize("shape", LOOP_SEQUENCES)
+def test_outside_caller_reads_a_forked_loop_from_the_region_that_owns_each_pass(
+ client, session, outside_loop_run, shape
+):
+ dr, dag, root, steps, _ = outside_loop_run
+ fork = DynamicRegion(
+ dag_id=dr.dag_id,
+ run_id=dr.run_id,
+ node_id="body",
+ forked_from_region_id=root.id,
+ resumes_from_index=2,
+ )
+ nested = DynamicRegion(
+ dag_id=dr.dag_id,
+ run_id=dr.run_id,
+ node_id="body.fanout",
+ parent_region_id=None,
+ parent_region_index=None,
+ )
+ session.add(fork)
+ session.flush()
+ nested.parent_region_id, nested.parent_region_index = fork.id, 2
+ session.add(nested)
+ session.flush()
+ replacements = [
+ TaskInstance(
+ dag.get_task("body.step"),
+ dr.created_dag_version_id,
+ run_id=dr.run_id,
+ region_id=fork.id,
+ region_index=2,
+ ),
+ *(
+ TaskInstance(
+ dag.get_task("body.fanout"),
+ dr.created_dag_version_id,
+ run_id=dr.run_id,
+ region_id=nested.id,
+ region_index=slot,
+ )
+ for slot in range(2)
+ ),
+ ]
+ session.add_all(replacements)
+ session.flush()
+ for ti, value in zip(replacements, ["step-2-again", "2-0-again", "2-1-again"]):
+ XComModel.set_for_attempt(
+ task_instance_id=ti.id, key="key", value=value, serialize=False, session=session
+ )
+ session.commit()
+ task_id = LOOP_SEQUENCES[shape][0]
+ expected = {
+ "unmapped": ["step-0", "step-1", "step-2-again"],
+ "mapped": ["0-0", "0-1", "1-0", "1-1", "2-0-again", "2-1-again"],
+ }[shape]
+ url = get_loop_url(dr, task_id)
+
+ assert client.head(url).headers["Content-Range"] == f"map_indexes {len(expected)}"
+ assert client.get(f"{url}/slice").json() == expected
+ assert client.get(f"{url}/item/-1").json() == expected[-1]
+
+
+def test_outside_caller_cannot_read_a_single_value_of_a_loop_task(client, outside_loop_run):
+ dr = outside_loop_run[0]
+
+ assert client.get(get_loop_url(dr, "body.step")).status_code == 400
+
+
+def test_outside_caller_cannot_read_the_previous_iteration_of_a_loop_task(client, outside_loop_run):
+ dr = outside_loop_run[0]
+
+ response = client.get(f"{get_loop_url(dr, 'body.step')}/slice", params={"previous_iteration": True})
+
+ assert response.status_code == 400
+
+
+def test_outside_caller_cannot_read_a_loop_with_several_executions(client, session, outside_loop_run):
+ dr, _, root, _, _ = outside_loop_run
+ session.add(
+ DynamicRegion(
+ dag_id=dr.dag_id,
+ run_id=dr.run_id,
+ node_id="body",
+ parent_region_id=root.id,
+ parent_region_index=0,
+ )
+ )
+ session.commit()
+
+ response = client.get(f"{get_loop_url(dr, 'body.step')}/slice")
+
+ assert response.status_code == 400
+ assert "several executions" in response.json()["detail"]
+
+
+def test_loop_task_reads_from_inside_the_loop_stay_on_one_pass(client, outside_loop_run):
+ dr, _, _, steps, authenticate_as = outside_loop_run
+ authenticate_as(steps[1])
+ step_url = get_loop_url(dr, "body.step")
+ fanout_url = get_loop_url(dr, "body.fanout")
+
+ assert client.head(step_url).headers["Content-Range"] == "map_indexes 1"
+ assert client.get(f"{step_url}/slice").json() == ["step-1"]
+ assert client.get(f"{step_url}/slice", params={"previous_iteration": True}).json() == ["step-0"]
+ assert client.get(f"{fanout_url}/slice").json() == ["1-0", "1-1"]
+ assert client.get(f"{fanout_url}/item/-1").json() == "1-1"
+
+
+@pytest.mark.parametrize("key", ["_airflow_loop_decision", "_airflow_loop_other"])
+def test_normal_xcom_push_rejects_loop_reserved_prefix(client, create_task_instance, key):
+ ti = create_task_instance()
+ response = client.post(f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/{key}", json="continue")
+ assert response.status_code == 400
+
+
+@pytest.mark.parametrize(
+ ("value", "region_index", "expected_status"),
+ [
+ ("continue", 0, 201),
+ ("stop", 2, 201),
+ ("continue", 2, 400),
+ ("stop", 0, 400),
+ (True, 2, 400),
+ ("invalid", 2, 400),
+ ({"decision": "stop"}, 2, 400),
+ ],
+)
+def test_internal_loop_decision_uses_token_gate_coordinates(
+ client, loop_xcoms, session, authenticate_as, value, region_index, expected_status
+):
+ dr, _, consumer, _ = loop_xcoms
+ gate = session.scalar(
+ select(TaskInstance).where(
+ TaskInstance.dag_id == dr.dag_id, TaskInstance.task_id == "body.__loop_gate"
+ )
+ )
+ gate.region_id, gate.region_index = consumer.region_id, region_index
+ gate.state = "running"
+ session.commit()
+ authenticate_as(gate)
+ url = f"/execution/xcoms/{dr.dag_id}/{dr.run_id}/{gate.task_id}/_airflow_loop_decision"
+
+ response = client.post(url, json=value)
+
+ assert response.status_code == expected_status
+ stored = session.scalar(select(XComModel).where(XComModel.task_id == gate.task_id))
+ if expected_status == 400:
+ assert stored is None
+ return
+ assert (stored.region_id, stored.region_index, stored.value) == (gate.region_id, region_index, value)
+ authenticate_as(consumer)
+ assert client.post(url, json="stop").status_code == 400
+ assert (
+ client.post(
+ f"/execution/xcoms/{dr.dag_id}/{dr.run_id}/{consumer.task_id}/_airflow_loop_decision",
+ json="stop",
+ ).status_code
+ == 400
+ )
+
+
+def test_loop_result_item_read_loads_no_producer_rows(dag_maker, session, client, authenticate_as):
+ @task_group
+ def body():
+ PythonOperator.partial(task_id="mapped", python_callable=str).expand(op_args=[[1], [2], [3]]) >> (
+ EmptyOperator(task_id="consumer")
+ )
+
+ with dag_maker(serialized=True) as dag:
+ create_loop(body, max_iterations=3)
+ dr = dag_maker.create_dagrun()
+ tis = {(ti.task_id, ti.map_index): ti for ti in dr.task_instances}
+ consumer = tis["body.consumer", 0]
+ loop_region = DynamicRegion.get_or_create(
+ dag_id=dr.dag_id, run_id=dr.run_id, node_id="body", session=session
+ )
+ session.add(loop_region)
+ session.flush()
+ consumer.region_id, consumer.region_index = loop_region.id, 1
+ mapped_task = dag.get_task("body.mapped")
+ for iteration in range(3):
+ region = DynamicRegion.get_or_create(
+ dag_id=dr.dag_id,
+ run_id=dr.run_id,
+ node_id="body.mapped",
+ parent_region_id=loop_region.id,
+ parent_region_index=iteration,
+ session=session,
+ )
+ session.add(region)
+ session.flush()
+ for map_index in range(3):
+ ti = tis.get(("body.mapped", map_index)) if iteration == 0 else None
+ if ti is None:
+ ti = TaskInstance(mapped_task, consumer.dag_version_id, run_id=dr.run_id, map_index=map_index)
+ session.add(ti)
+ ti.region_id = region.id
+ session.flush()
+ XComModel.set_for_attempt(
+ task_instance_id=ti.id,
+ key="return_value",
+ value=f"{iteration}-{map_index}",
+ serialize=False,
+ session=session,
+ )
+ session.commit()
+ authenticate_as(consumer)
+ session.expunge_all()
+ loaded = []
+
+ def record(_, instance):
+ if isinstance(instance, TaskInstance) and instance.task_id == "body.mapped":
+ loaded.append(instance.id)
+
+ event.listen(Session, "loaded_as_persistent", record)
+ try:
+ response = client.get(f"/execution/xcoms/{dr.dag_id}/{dr.run_id}/body.mapped/return_value/item/0")
+ finally:
+ event.remove(Session, "loaded_as_persistent", record)
+
+ assert response.json() == "1-0"
+ assert loaded == []
+
+
def test_loop_xcom_omission_uses_consumer_pass_and_retained_region(client, loop_xcoms):
dr, producer, consumer, _ = loop_xcoms
base = f"/execution/xcoms/{dr.dag_id}/{dr.run_id}"
diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py
index 2f4182505cb..55280830e14 100644
--- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py
+++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py
@@ -49,6 +49,8 @@ def test_task_coordinate_schema_by_version(client, version):
previous = response.json()["components"]["schemas"]["PreviousTIResponse"]["properties"]
assert ("region_id" in previous) == (version == "2026-10-30")
assert ("region_index" in previous) == (version == "2026-10-30")
+ context = response.json()["components"]["schemas"]["TIRunContext"]["properties"]
+ assert ("loop" in context) == (version == "2026-10-30")
@pytest.mark.parametrize(
diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py
index a66e3f6c5ca..3203bd23dcd 100644
--- a/airflow-core/tests/unit/jobs/test_scheduler_job.py
+++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py
@@ -147,7 +147,9 @@ from airflow.sdk import (
StartOfDayMapper,
StartOfHourMapper,
task,
+ task_group,
)
+from airflow.sdk.definitions._internal.loop import create_loop
from airflow.sdk.definitions.callback import AsyncCallback, SyncCallback
from airflow.sdk.definitions.timetables.assets import PartitionedAssetTimetable
from airflow.serialization.definitions.dag import SerializedDAG
@@ -1180,6 +1182,54 @@ class TestSchedulerJob:
with pytest.raises(TypeError, match="Unknown workload key type in event buffer"):
self.job_runner._process_executor_events(executor=executor, session=session)
+ @pytest.mark.parametrize("retries", [0, 2])
+ def test_missing_gate_definition_executor_success_uses_failure_policy(
+ self, dag_maker, session, mocker, retries
+ ):
+ @task_group
+ def body():
+ EmptyOperator(task_id="terminal")
+
+ with dag_maker(serialized=True, default_args={"retries": retries}, session=session):
+ loop = create_loop(body, max_iterations=3)
+ dr = dag_maker.create_dagrun()
+ gate = next(ti for ti in dr.task_instances if ti.task_id == loop.gate_task_id)
+ executor = MockExecutor(do_update=False)
+ scheduler_job = Job()
+ session.add(scheduler_job)
+ session.flush()
+ runner = SchedulerJobRunner(scheduler_job, executors=[executor])
+ mocker.patch.object(runner.scheduler_dag_bag, "get_dag_for_run", autospec=True, return_value=None)
+ gate.state = State.RUNNING
+ gate.try_number = 1
+ gate.queued_by_job_id = scheduler_job.id
+ session.commit()
+ original_id = gate.id
+ callback = mocker.patch.object(executor, "send_callback", autospec=True)
+ executor.event_buffer[TaskInstanceUuid(gate.id)] = State.SUCCESS, None
+
+ runner._process_executor_events(executor=executor, session=session)
+
+ gate = session.scalar(
+ select(TaskInstance).where(
+ TaskInstance.dag_id == dr.dag_id,
+ TaskInstance.run_id == dr.run_id,
+ TaskInstance.task_id == loop.gate_task_id,
+ TaskInstance.working_set.is_(True),
+ )
+ )
+ assert gate.state == (State.UP_FOR_RETRY if retries else State.FAILED)
+ assert (gate.id != original_id) is bool(retries)
+ callback.assert_not_called()
+ if retries:
+ archived = session.scalar(
+ select(TaskInstance)
+ .where(TaskInstance.id == original_id)
+ .execution_options(include_all_attempts=True)
+ )
+ assert archived.working_set is None
+ assert len(dr.get_task_instances(session=session)) == 2
+
@mock.patch("airflow.jobs.scheduler_job_runner.TaskCallbackRequest")
@mock.patch("airflow._shared.observability.metrics.stats._get_backend")
def test_process_executor_event_missing_dag(
diff --git a/airflow-core/tests/unit/models/test_dagrun.py b/airflow-core/tests/unit/models/test_dagrun.py
index 4b769b4d214..47ad1c63955 100644
--- a/airflow-core/tests/unit/models/test_dagrun.py
+++ b/airflow-core/tests/unit/models/test_dagrun.py
@@ -55,6 +55,7 @@ from airflow._shared.observability.traces import (
from airflow._shared.timezones import timezone
from airflow.api_fastapi.execution_api.datamodels.taskinstance import TaskInstance as TIDataModel
from airflow.callbacks.callback_requests import DagCallbackRequest, DagRunContext
+from airflow.dag_processing.dagbag import DagBag
from airflow.models.dag import DagModel, infer_automated_data_interval
from airflow.models.dag_version import DagVersion
from airflow.models.dagbag import DBDagBag
@@ -85,12 +86,15 @@ from airflow.sdk import (
task_group,
teardown,
)
+from airflow.sdk.definitions._internal.loop import create_loop
from airflow.sdk.definitions.callback import AsyncCallback
from airflow.sdk.definitions.deadline import DeadlineAlert, DeadlineReference, VariableInterval
from airflow.serialization.definitions.deadline import SerializedReferenceModels
from airflow.serialization.serialized_objects import DagSerialization, LazyDeserializedDAG
from airflow.settings import get_policy_plugin_manager
from airflow.task.trigger_rule import TriggerRule
+from airflow.ti_deps.dep_context import DepContext
+from airflow.ti_deps.deps.trigger_rule_dep import TriggerRuleDep
from airflow.triggers.base import StartTriggerArgs
from airflow.utils.session import create_session
from airflow.utils.sqlalchemy import prohibit_commit
@@ -101,9 +105,10 @@ from tests_common.test_utils import db
from tests_common.test_utils.asserts import assert_queries_count
from tests_common.test_utils.config import conf_vars
from tests_common.test_utils.dag import sync_dag_to_db
-from tests_common.test_utils.db import clear_db_dags, clear_db_runs
+from tests_common.test_utils.db import clear_db_dags, clear_db_runs, clear_db_serialized_dags, clear_db_xcom
from tests_common.test_utils.mapping import expand_mapped_task, push_mapped_length
from tests_common.test_utils.mock_operators import MockOperator
+from tests_common.test_utils.paths import AIRFLOW_CORE_SOURCES_PATH
from tests_common.test_utils.taskinstance import create_task_instance, run_task_instance
from unit.models import DEFAULT_DATE as _DEFAULT_DATE
from unit.plugins.priority_weight_strategy import DecreasingPriorityStrategy, TestPriorityWeightStrategyPlugin
@@ -123,6 +128,258 @@ def get_callback_ti(ti):
)
+@pytest.mark.parametrize("values", [[], [{}, {}]])
+def test_loop_first_pass_uses_one_region_and_mapped_child(dag_maker, session, values):
+ @task_group
+ def body():
+ plain = PythonOperator(task_id="plain", python_callable=list)
+ mapped = PythonOperator.partial(task_id="mapped", python_callable=list).expand(op_kwargs=values)
+ plain >> mapped
+
+ with dag_maker(serialized=True):
+ outside = EmptyOperator(task_id="outside")
+ loop = create_loop(body, max_iterations=3)
+ outside >> loop
+ dr = dag_maker.create_dagrun()
+ regions = session.scalars(
+ select(DynamicRegion).where(DynamicRegion.dag_id == dr.dag_id, DynamicRegion.run_id == dr.run_id)
+ ).all()
+ assert {region.node_id for region in regions} == {loop.group_id, "body.mapped"}
+ parent = next(region for region in regions if region.node_id == loop.group_id)
+ mapped_region = next(region for region in regions if region.node_id == "body.mapped")
+ assert (mapped_region.parent_region_id, mapped_region.parent_region_index) == (parent.id, 0)
+ original_ids = {ti.id for ti in dr.task_instances}
+ for ti in dr.task_instances:
+ if ti.task_id == "outside":
+ assert (ti.region_id, ti.region_index) == (SENTINEL_REGION_ID, -1)
+ elif ti.task_id == "body.mapped":
+ assert ti.region_id == mapped_region.id
+ else:
+ assert (ti.region_id, ti.region_index) == (parent.id, 0)
+
+ dr.verify_integrity(session=session, dag_version_id=dr.created_dag_version_id)
+
+ assert {ti.id for ti in dr.get_task_instances(session=session)} == original_ids
+ assert (
+ session.scalar(
+ select(func.count()).select_from(DynamicRegion).where(DynamicRegion.dag_id == dr.dag_id)
+ )
+ == 2
+ )
+
+
+@pytest.mark.parametrize(
+ ("trigger_rule", "left_state", "expected"),
+ [
+ (TriggerRule.ALWAYS, State.RUNNING, True),
+ (TriggerRule.ONE_SUCCESS, State.SUCCESS, True),
+ (TriggerRule.ALL_SUCCESS, State.SUCCESS, False),
+ (TriggerRule.ALL_DONE, State.FAILED, False),
+ (TriggerRule.ALL_SUCCESS, State.SKIPPED, False),
+ ],
+)
+def test_loop_entry_uses_normal_trigger_rule_before_all_upstreams_finish(
+ dag_maker, session, trigger_rule, left_state, expected
+):
+ @task_group
+ def body():
+ PythonOperator(task_id="entry", python_callable=list, trigger_rule=trigger_rule)
+
+ with dag_maker(serialized=True):
+ left = PythonOperator(task_id="left", python_callable=list)
+ right = PythonOperator(task_id="right", python_callable=list)
+ loop = create_loop(body, max_iterations=3)
+ [left, right] >> loop
+ dr = dag_maker.create_dagrun()
+ tis = {ti.task_id: ti for ti in dr.task_instances}
+ tis["left"].state = left_state
+ tis["right"].state = State.RUNNING
+ entry = tis["body.entry"]
+ entry.task = dag_maker.serialized_dag.get_task(entry.task_id)
+ assert entry.region_id != SENTINEL_REGION_ID
+ session.flush()
+
+ assert TriggerRuleDep().is_met(ti=entry, dep_context=DepContext(), session=session) is expected
+
+
+def test_loop_new_member_joins_live_pass_after_reserialization(dag_maker, session):
+ @task_group
+ def body(add_member=False):
+ terminal = PythonOperator(task_id="terminal", python_callable=list)
+ if add_member:
+ PythonOperator(task_id="new", python_callable=list) >> terminal
+
+ with dag_maker(serialized=True, session=session):
+ loop = create_loop(body, max_iterations=3)
+ dr = dag_maker.create_dagrun()
+ existing = dr.get_task_instances(session=session)
+ original_ids = {ti.id for ti in existing}
+ gate = next(ti for ti in existing if ti.task_id == loop.gate_task_id)
+ gate.state = State.SUCCESS
+ session.flush()
+
+ with dag_maker(serialized=True, session=session):
+ body.partial(add_member=True).loop(max_iterations=3)
+ dr.dag = dag_maker.serialized_dag
+ version_id = DagVersion.get_latest_version(dr.dag_id, session=session).id
+
+ dr.verify_integrity(session=session, dag_version_id=version_id)
+
+ current = dr.get_task_instances(session=session)
+ assert {ti.id for ti in current if ti.task_id != "body.new"} == original_ids
+ added = next(ti for ti in current if ti.task_id == "body.new")
+ assert (added.region_id, added.region_index, added.dag_version_id) == (gate.region_id, 0, version_id)
+ assert session.get(TI, gate.id).state == State.SUCCESS
+
+
+def test_loop_integrity_skips_loop_with_regions_but_no_live_gate(dag_maker, session, caplog):
+ @task_group
+ def body():
+ PythonOperator(task_id="terminal", python_callable=list)
+
+ with dag_maker(serialized=True, session=session):
+ EmptyOperator(task_id="outside")
+ loop = create_loop(body, max_iterations=3)
+ dr = dag_maker.create_dagrun()
+ for ti in dr.get_task_instances(session=session):
+ session.delete(ti)
+ session.flush()
+
+ with caplog.at_level("WARNING"):
+ dr.verify_integrity(session=session, dag_version_id=dr.created_dag_version_id)
+
+ assert [ti.task_id for ti in dr.get_task_instances(session=session)] == ["outside"]
+ assert {
+ "event": "skipping loop that has regions but no live gate",
+ "dag_id": dr.dag_id,
+ "run_id": dr.run_id,
+ "loop": loop.group_id,
+ "log_level": "warning",
+ } in caplog
+ assert (
+ session.scalar(
+ select(func.count()).select_from(DynamicRegion).where(DynamicRegion.dag_id == dr.dag_id)
+ )
+ == 1
+ )
+
+
+def test_loop_integrity_reuses_region_created_by_a_concurrent_run(dag_maker, session):
+ @task_group
+ def body():
+ PythonOperator(task_id="terminal", python_callable=list)
+
+ with dag_maker(serialized=True, session=session):
+ loop = create_loop(body, max_iterations=3)
+ dr = dag_maker.create_dagrun()
+ region_id = session.scalar(select(DynamicRegion.id).where(DynamicRegion.dag_id == dr.dag_id))
+ for ti in dr.get_task_instances(session=session):
+ session.delete(ti)
+ session.flush()
+
+ with mock.patch.object(session, "scalar", autospec=True, return_value=None):
+ dr.verify_integrity(session=session, dag_version_id=dr.created_dag_version_id)
+
+ gate = next(ti for ti in dr.get_task_instances(session=session) if ti.task_id == loop.gate_task_id)
+ assert gate.region_id == region_id
+ assert (
+ session.scalar(
+ select(func.count()).select_from(DynamicRegion).where(DynamicRegion.dag_id == dr.dag_id)
+ )
+ == 1
+ )
+
+
+@pytest.mark.parametrize(
+ "mode", ["fixed", "mapped", "empty", "until_true", "until_false", "handled_failure", "unhandled_failure"]
+)
+def test_loop_runs_through_scheduler_sdk_and_api(dag_maker, session, mode):
+ @task
+ def terminal(value, *, ti, loop):
+ assert ti.region_id != SENTINEL_REGION_ID
+ if mode == "mapped":
+ assert ti.map_index == value - 1
+ else:
+ assert ti.map_index == -1
+ assert ti.region_index == loop.index
+ return value
+
+ @task
+ def fail():
+ raise ValueError("body branch failed")
+
+ @task_group
+ def body():
+ if mode in {"mapped", "empty"}:
+ terminal.expand(value=[1, 2] if mode == "mapped" else [])
+ elif mode == "handled_failure":
+ fail() >> terminal.override(trigger_rule=TriggerRule.ALL_DONE)(1)
+ elif mode == "unhandled_failure":
+ fail() >> terminal(1)
+ else:
+ terminal(1)
+
+ def converged(*, loop):
+ return mode == "until_true" and loop.index == 1
+
+ with dag_maker(serialized=False) as dag:
+ loop = body.loop(max_iterations=2, until=converged if mode.startswith("until") else None)
+ loop >> PythonOperator(task_id="outside", python_callable=list, trigger_rule=TriggerRule.ALL_DONE)
+
+ dr = dag.test()
+
+ gates = sorted(
+ (ti.region_index, ti.state)
+ for ti in dr.get_task_instances(session=session)
+ if ti.task_id == loop.gate_task_id
+ )
+ if mode == "empty":
+ assert gates == [(0, State.SKIPPED)]
+ elif mode == "until_false":
+ assert gates == [(0, State.SUCCESS), (1, State.FAILED)]
+ elif mode == "unhandled_failure":
+ assert gates == [(0, State.UPSTREAM_FAILED)]
+ else:
+ assert gates == [(0, State.SUCCESS), (1, State.SUCCESS)]
+ outside = dr.get_task_instance(task_id="outside", session=session)
+ assert outside.state == State.SUCCESS
+
+
+@pytest.fixture
+def clear_loop_example_runs():
+ yield
+ clear_db_runs()
+ clear_db_xcom()
+ clear_db_dags()
+ clear_db_serialized_dags()
+
+
+@pytest.mark.usefixtures("clear_loop_example_runs")
+@pytest.mark.parametrize(
+ ("dag_id", "expected_iterations"),
+ [
+ ("refine_estimate", 4),
+ ("fixed_task_loop", 3),
+ ("mapped_task_loop", 2),
+ ("partial_override_loop", 3),
+ ],
+)
+def test_documented_task_loop_examples_complete(session, dag_id, expected_iterations):
+ bag = DagBag(
+ dag_folder=str(AIRFLOW_CORE_SOURCES_PATH / "airflow" / "example_dags" / "example_task_loops.py"),
+ )
+ assert bag.import_errors == {}
+ dag = bag.get_dag(dag_id)
+
+ dr = dag.test()
+
+ assert dr.state == DagRunState.SUCCESS
+ gates = [ti for ti in dr.get_task_instances(session=session) if ti.operator == "LoopGateOperator"]
+ assert sorted((ti.region_index, ti.state) for ti in gates) == [
+ (index, State.SUCCESS) for index in range(expected_iterations)
+ ]
+
+
async def empty_callback_for_deadline():
"""Used in a number of tests to confirm that Deadlines and DeadlineAlerts function correctly."""
pass
diff --git a/airflow-core/tests/unit/models/test_task_coordinates.py b/airflow-core/tests/unit/models/test_task_coordinates.py
index 30131fe2277..c872b653b4a 100644
--- a/airflow-core/tests/unit/models/test_task_coordinates.py
+++ b/airflow-core/tests/unit/models/test_task_coordinates.py
@@ -19,6 +19,8 @@ from __future__ import annotations
import pytest
from sqlalchemy import select
+from airflow.api_fastapi.execution_api.datamodels.token import TIClaims, TIToken
+from airflow.api_fastapi.execution_api.routes.xcoms import _build_xcom_read
from airflow.executors.workloads.base import BundleInfo
from airflow.executors.workloads.task import ExecuteTask
from airflow.models.dag_version import DagVersion
@@ -37,7 +39,11 @@ from airflow.sdk import task_group
from airflow.sdk.definitions._internal.loop import create_loop
from airflow.utils.log.task_log_address import prepare_task_log_contexts
-from tests_common.test_utils.asserts import assert_queries_count, capture_orm_selects
+from tests_common.test_utils.asserts import (
+ assert_queries_count,
+ capture_orm_selects,
+ count_loaded_task_instances,
+)
pytestmark = pytest.mark.db_test
@@ -141,6 +147,113 @@ def test_shared_loop_uses_pinned_graph_and_retained_producer(loop_coordinates, s
)
+@pytest.fixture
+def forked_loop_coordinates(loop_coordinates, session):
+ replacement = session.scalars(
+ select(DynamicRegion).where(DynamicRegion.forked_from_region_id.is_not(None))
+ ).one()
+ replacement.resumes_from_index = 2
+ session.flush()
+ return loop_coordinates
+
+
+def test_outside_caller_reads_the_live_iterations_each_region_owns(forked_loop_coordinates, session):
+ dr, producer, consumer, previous, outside = forked_loop_coordinates
+ resolver = TaskCoordinateResolver(DBDagBag(), session)
+ arguments = dict(
+ dag_id=dr.dag_id, run_id=dr.run_id, task_id=producer.task_id, caller=outside, all_iterations=True
+ )
+
+ selected = resolver.resolve(**arguments)
+ selected_ids = set(session.scalars(resolver.select_producer_ids(**arguments)))
+
+ assert selected == (previous,)
+ assert selected_ids == {previous.id}
+
+
+@pytest.mark.parametrize("by", ["resolve", "select_producer_ids"])
+def test_outside_caller_is_refused_a_loop_task_unless_it_reads_every_iteration(loop_coordinates, session, by):
+ dr, producer, consumer, previous, outside = loop_coordinates
+ resolver = TaskCoordinateResolver(DBDagBag(), session)
+
+ with pytest.raises(ValueError, match="explicit scope or a consumer inside the loop"):
+ getattr(resolver, by)(dag_id=dr.dag_id, run_id=dr.run_id, task_id=producer.task_id, caller=outside)
+
+
+def test_reading_every_iteration_cannot_be_combined_with_the_previous_one(loop_coordinates, session):
+ dr, producer, consumer, previous, outside = loop_coordinates
+ resolver = TaskCoordinateResolver(DBDagBag(), session)
+
+ with pytest.raises(ValueError, match="requires a shared loop scope"):
+ resolver.resolve(
+ dag_id=dr.dag_id,
+ run_id=dr.run_id,
+ task_id=producer.task_id,
+ caller=outside,
+ all_iterations=True,
+ previous_iteration=True,
+ )
+
+
+def test_all_iterations_does_not_widen_the_read_of_a_caller_inside_the_loop(loop_coordinates, session):
+ dr, producer, consumer, previous, outside = loop_coordinates
+ resolver = TaskCoordinateResolver(DBDagBag(), session)
+
+ selected = resolver.resolve(
+ dag_id=dr.dag_id, run_id=dr.run_id, task_id=producer.task_id, caller=consumer, all_iterations=True
+ )
+
+ assert selected == (producer,)
+
+
+def test_dependency_of_an_outside_task_on_a_loop_task_is_every_live_iteration(loop_coordinates, session):
+ dr, producer, consumer, previous, outside = loop_coordinates
+ replacement = session.scalars(
+ select(DynamicRegion).where(DynamicRegion.forked_from_region_id.is_not(None))
+ ).one()
+ replacement.resumes_from_index = 3
+ first = TaskInstance(
+ task=producer.task,
+ run_id=dr.run_id,
+ dag_version_id=producer.dag_version_id,
+ region_id=producer.region_id,
+ )
+ first.region_index = 0
+ session.add(first)
+ session.flush()
+ resolver = TaskCoordinateResolver(DBDagBag(), session)
+
+ assert set(resolver.resolve_dependency(outside, producer.task_id)) == {first, previous, producer}
+
+
+@pytest.mark.parametrize(
+ ("caller", "task_id", "run_id", "expected"),
+ [
+ pytest.param("outside", "body.producer", None, True, id="outside caller of a loop task"),
+ pytest.param("consumer", "body.producer", None, False, id="caller inside the loop"),
+ pytest.param("outside", "outside", None, False, id="task outside any loop"),
+ pytest.param("outside", "removed", None, False, id="task missing from the pinned dag"),
+ pytest.param("outside", "body.producer", "other_run", False, id="caller of another run"),
+ pytest.param(None, "body.producer", None, False, id="no caller"),
+ ],
+)
+def test_loop_task_read_from_outside_needs_a_caller_of_the_same_run_outside_the_loop(
+ loop_coordinates, session, caller, task_id, run_id, expected
+):
+ dr, producer, consumer, previous, outside = loop_coordinates
+ resolver = TaskCoordinateResolver(DBDagBag(), session)
+
+ assert (
+ resolver.is_loop_task_read_from_outside(
+ dag_id=dr.dag_id,
+ run_id=run_id or dr.run_id,
+ task_id=task_id,
+ caller={"consumer": consumer, "outside": outside, None: None}[caller],
+ )
+ is expected
+ )
+
+
def test_loop_passes_load_the_region_ancestry_once_for_all_task_instances(loop_coordinates, session):
_, producer, consumer, previous, outside = loop_coordinates
resolver = TaskCoordinateResolver(DBDagBag(), session)
@@ -152,6 +265,31 @@ def test_loop_passes_load_the_region_ancestry_once_for_all_task_instances(loop_c
assert len(statements) == 1
+def test_dependency_rejects_duplicate_latest_gate_coordinates(loop_coordinates, session):
+ dr, _, consumer, _, outside = loop_coordinates
+ resolver = TaskCoordinateResolver(DBDagBag(), session)
+ gate = next(ti for ti in dr.task_instances if ti.task_id == "body.__loop_gate")
+ duplicate = TaskInstance(
+ task=resolver.get_task(dr.dag_id, dr.run_id, gate.task_id, dag_version_id=gate.dag_version_id),
+ run_id=dr.run_id,
+ dag_version_id=gate.dag_version_id,
+ region_id=consumer.region_id,
+ region_index=gate.region_index,
+ )
+ session.add(duplicate)
+ session.flush()
+
+ with pytest.raises(AmbiguousProducerError, match="Multiple live loop gates"):
+ resolver.resolve_dependency(outside, gate.task_id)
+
+
+def test_dependency_on_task_missing_from_pinned_dag_falls_back_to_plain_resolution(loop_coordinates, session):
+ _, _, consumer, _, _ = loop_coordinates
+ resolver = TaskCoordinateResolver(DBDagBag(), session)
+
+ assert resolver.resolve_dependency(consumer, "removed_upstream") == ()
+
+
def test_loop_without_context_requires_explicit_scope(loop_coordinates, session):
dr, producer, _, _, outside = loop_coordinates
resolver = TaskCoordinateResolver(DBDagBag(), session)
@@ -274,3 +412,104 @@ def test_unversioned_run_keeps_task_definition_after_latest_graph_changes(
region_id=mapped.region_id,
region_index=0,
) == (mapped,)
+
+
+@pytest.fixture
+def mapped_run(dag_maker, session):
+ def create(mapped_count: int):
+ with dag_maker(serialized=True):
+ mapped = PythonOperator.partial(task_id="mapped", python_callable=str).expand(
+ op_args=[[i] for i in range(mapped_count)]
+ )
+ mapped >> EmptyOperator(task_id="reduce")
+ dr = dag_maker.create_dagrun()
+ caller = session.scalars(
+ select(TaskInstance).where(TaskInstance.run_id == dr.run_id, TaskInstance.task_id == "reduce")
+ ).one()
+ session.expire_all()
+ return dr, caller
+
+ return create
+
+
+@pytest.mark.parametrize("mapped_count", [3, 60])
+def test_xcom_read_of_one_mapped_slot_loads_no_producer_rows(mapped_run, session, mapped_count):
+ dr, caller = mapped_run(mapped_count)
+ token = TIToken(id=caller.id, claims=TIClaims())
+
+ with count_loaded_task_instances("mapped") as loaded, assert_queries_count(8):
+ read = _build_xcom_read(
+ dag_id=dr.dag_id,
+ run_id=dr.run_id,
+ task_id="mapped",
+ key="return_value",
+ session=session,
+ dag_bag=DBDagBag(),
+ token=token,
+ map_index=2,
+ )
+ session.scalars(read.statement).all()
+
+ assert loaded == []
+
+
+def test_loaded_task_instance_collector_keeps_only_the_requested_task(mapped_run, session):
+ dr, _ = mapped_run(3)
+
+ with (
+ count_loaded_task_instances("mapped") as loaded,
+ count_loaded_task_instances("other") as other_loaded,
+ ):
+ rows = session.scalars(
+ select(TaskInstance).where(TaskInstance.run_id == dr.run_id, TaskInstance.task_id == "mapped")
+ ).all()
+
+ assert len(rows) == 3
+ assert sorted(ti.id for ti in loaded) == sorted(ti.id for ti in rows)
+ assert other_loaded == []
+
+
+@pytest.mark.parametrize("mapped_count", [3, 60])
+def test_loop_read_of_one_mapped_slot_loads_no_producer_rows(dag_maker, session, mapped_count):
+ @task_group
+ def body():
+ PythonOperator.partial(task_id="mapped", python_callable=str).expand(
+ op_args=[[i] for i in range(mapped_count)]
+ ) >> EmptyOperator(task_id="reduce")
+
+ with dag_maker(serialized=True):
+ create_loop(body, max_iterations=3)
+ dr = dag_maker.create_dagrun()
+ caller = session.scalars(
+ select(TaskInstance).where(TaskInstance.run_id == dr.run_id, TaskInstance.task_id == "body.reduce")
+ ).one()
+ session.expire_all()
+ token = TIToken(id=caller.id, claims=TIClaims())
+
+ with count_loaded_task_instances("body.mapped") as loaded:
+ read = _build_xcom_read(
+ dag_id=dr.dag_id,
+ run_id=dr.run_id,
+ task_id="body.mapped",
+ key="return_value",
+ session=session,
+ dag_bag=DBDagBag(),
+ token=token,
+ map_index=2,
+ )
+ session.scalars(read.statement).all()
+
+ assert loaded == []
+
+
+def test_selected_mapped_producers_match_resolved_producers(mapped_run, session):
+ dr, caller = mapped_run(6)
+ resolver = TaskCoordinateResolver(DBDagBag(), session)
+
+ for map_indexes in (None, 4, range(1, 3), [0, 5]):
+ arguments = dict(
+ dag_id=dr.dag_id, run_id=dr.run_id, task_id="mapped", caller=caller, map_indexes=map_indexes
+ )
+ resolved = {ti.id for ti in resolver.resolve(**arguments)}
+ assert set(session.scalars(resolver.select_producer_ids(**arguments))) == resolved
+ assert resolved
diff --git a/airflow-core/tests/unit/models/test_xcom_arg.py b/airflow-core/tests/unit/models/test_xcom_arg.py
index d2887c0d027..c7c3a3136a2 100644
--- a/airflow-core/tests/unit/models/test_xcom_arg.py
+++ b/airflow-core/tests/unit/models/test_xcom_arg.py
@@ -17,6 +17,7 @@
from __future__ import annotations
import pytest
+from sqlalchemy import delete
from airflow.models.dynamic_region import DynamicRegion, ProducerContext
from airflow.models.expandinput import NotFullyPopulated
@@ -37,6 +38,7 @@ from airflow.serialization.definitions.xcom_arg import (
)
from airflow.utils.state import TaskInstanceState
+from tests_common.test_utils.asserts import assert_queries_count, count_loaded_task_instances
from tests_common.test_utils.db import clear_db_dags, clear_db_runs
pytestmark = pytest.mark.db_test
@@ -463,3 +465,46 @@ def test_member_of_mapped_task_group_has_no_map_length(dag_maker, session):
)
is None
)
+
+
+def test_mapped_producer_without_live_instances_has_zero_length(dag_maker, session):
+ with dag_maker(session=session, serialized=True) as dag:
+
+ @dag.task
+ def source(value):
+ return value
+
+ source.expand(value=[1, 2])
+
+ dr = dag_maker.create_dagrun()
+ session.execute(delete(TaskInstance).where(TaskInstance.run_id == dr.run_id))
+ session.flush()
+ argument = SchedulerPlainXComArg(dag_maker.serialized_dag.get_task("source"), XCOM_RETURN_KEY)
+
+ assert get_task_map_length(argument, dr.run_id, session=session) == 0
+
+
+@pytest.mark.parametrize("mapped_count", [3, 60])
+def test_mapped_producer_length_loads_no_producer_rows(dag_maker, session, mapped_count):
+ with dag_maker(session=session, serialized=True) as dag:
+
+ @dag.task
+ def source(value):
+ return value
+
+ source.expand(value=list(range(mapped_count)))
+
+ dr = dag_maker.create_dagrun()
+ for ti in dr.task_instances:
+ ti.state = TaskInstanceState.SUCCESS
+ XComModel.set_for_attempt(
+ task_instance_id=ti.id, key=XCOM_RETURN_KEY, value=ti.map_index, session=session
+ )
+ session.flush()
+ session.expire_all()
+ argument = SchedulerPlainXComArg(dag_maker.serialized_dag.get_task("source"), XCOM_RETURN_KEY)
+ with count_loaded_task_instances("source") as loaded, assert_queries_count(4):
+ length = get_task_map_length(argument, dr.run_id, session=session)
+
+ assert length == mapped_count
+ assert loaded == []
diff --git a/airflow-core/tests/unit/serialization/test_dag_serialization.py b/airflow-core/tests/unit/serialization/test_dag_serialization.py
index 85af2ca9dfd..d40f36ebf6d 100644
--- a/airflow-core/tests/unit/serialization/test_dag_serialization.py
+++ b/airflow-core/tests/unit/serialization/test_dag_serialization.py
@@ -1861,6 +1861,30 @@ class TestStringifiedDAGs:
assert restored_loop.dag is restored
assert restored_loop.parent_group is restored.task_group
+ def test_task_using_a_loop_task_serializes_with_the_gate_as_its_upstream(self):
+ @task_decorator
+ def report(value):
+ return value
+
+ @task_group(group_id="body")
+ def body():
+ BaseOperator(task_id="first") >> BaseOperator(task_id="terminal")
+
+ with DAG("loop_outside_edge", schedule=None) as dag:
+ loop = create_loop(body, max_iterations=3)
+ report(loop["first"].output)
+
+ encoded = DagSerialization.to_dict(dag)
+ DagSerialization.validate_schema(encoded)
+ restored = DagSerialization.from_json(DagSerialization.to_json(dag))
+
+ assert restored.get_task("report").upstream_task_ids == {"body.__loop_gate"}
+ assert restored.get_task("body.__loop_gate").downstream_task_ids == {"report"}
+ assert restored.get_task("body.first").downstream_task_ids == {"body.terminal"}
+ assert [arg.operator.task_id for arg in dag.get_task("report").op_args] == ["body.first"]
+ assert "body.first" in str(restored.get_task("report").op_args[0])
+ assert "body.__loop_gate" not in str(restored.get_task("report").op_args[0])
+
def test_task_group_serialization(self):
"""
Test TaskGroup serialization/deserialization.
diff --git a/airflow-core/tests/unit/ti_deps/deps/test_trigger_rule_dep.py b/airflow-core/tests/unit/ti_deps/deps/test_trigger_rule_dep.py
index 4850dde9426..53d33165d0e 100644
--- a/airflow-core/tests/unit/ti_deps/deps/test_trigger_rule_dep.py
+++ b/airflow-core/tests/unit/ti_deps/deps/test_trigger_rule_dep.py
@@ -56,6 +56,80 @@ SUCCESS = TaskInstanceState.SUCCESS
FAILED = TaskInstanceState.FAILED
+@pytest.mark.parametrize("latest_state", [None, SUCCESS, FAILED, SKIPPED, TaskInstanceState.RESTARTING])
+@pytest.mark.parametrize("rule", [TriggerRule.ALL_SUCCESS, TriggerRule.ONE_SUCCESS, TriggerRule.ALL_DONE])
+def test_outside_dependency_observes_only_latest_live_loop_gate(dag_maker, session, latest_state, rule):
+ @task_group
+ def body():
+ PythonOperator(task_id="terminal", python_callable=list)
+
+ with dag_maker(serialized=True):
+ loop = create_loop(body, max_iterations=3)
+ loop >> PythonOperator(task_id="outside", python_callable=list, trigger_rule=rule)
+ dr = dag_maker.create_dagrun()
+ previous = next(ti for ti in dr.task_instances if ti.task_id == loop.gate_task_id)
+ previous.state = SUCCESS
+ latest = TaskInstance(
+ task=dag_maker.serialized_dag.get_task(previous.task_id),
+ run_id=dr.run_id,
+ dag_version_id=previous.dag_version_id,
+ region_id=previous.region_id,
+ region_index=1,
+ state=latest_state,
+ )
+ session.add(latest)
+ outside = next(ti for ti in dr.task_instances if ti.task_id == "outside")
+ outside.task = dag_maker.serialized_dag.get_task(outside.task_id)
+ session.flush()
+ context = DepContext()
+
+ actual = TriggerRuleDep().is_met(ti=outside, dep_context=context, session=session)
+
+ expected = latest_state == SUCCESS or (rule == TriggerRule.ALL_DONE and latest_state in {FAILED, SKIPPED})
+ assert actual is expected
+ assert context.upstream_tis(outside, previous.task_id, session=session) == (latest,)
+ with pytest.raises(ValueError, match="loop producer requires"):
+ context.coordinate_resolver(outside, session=session).resolve(
+ dag_id=dr.dag_id, run_id=dr.run_id, task_id=previous.task_id, caller=outside
+ )
+
+
+GATE_RULE_MATRIX = {
+ TriggerRule.ALL_SUCCESS: {SUCCESS},
+ TriggerRule.ONE_SUCCESS: {SUCCESS},
+ TriggerRule.ONE_FAILED: {FAILED, UPSTREAM_FAILED},
+ TriggerRule.NONE_FAILED: {SUCCESS, SKIPPED},
+ TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS: {SUCCESS},
+ TriggerRule.ALL_DONE: {SUCCESS, FAILED, UPSTREAM_FAILED, SKIPPED},
+ TriggerRule.ALWAYS: {SUCCESS, FAILED, UPSTREAM_FAILED, SKIPPED, None},
+}
+
+
+@pytest.mark.parametrize("gate_state", [None, SUCCESS, FAILED, UPSTREAM_FAILED, SKIPPED])
+@pytest.mark.parametrize("rule", list(GATE_RULE_MATRIX))
+def test_task_using_a_loop_task_follows_the_loop_gate_not_the_loop_task(dag_maker, session, rule, gate_state):
+ @task_group
+ def body():
+ EmptyOperator(task_id="first") >> EmptyOperator(task_id="terminal")
+
+ with dag_maker(serialized=True):
+ loop = create_loop(body, max_iterations=3)
+ loop["first"] >> PythonOperator(task_id="outside", python_callable=list, trigger_rule=rule)
+ dr = dag_maker.create_dagrun()
+ tis = {ti.task_id: ti for ti in dr.task_instances}
+ tis["body.first"].state = SUCCESS
+ tis["body.terminal"].state = SUCCESS if gate_state in (SUCCESS, FAILED, SKIPPED) else FAILED
+ tis[loop.gate_task_id].state = gate_state
+ outside = tis["outside"]
+ outside.task = dag_maker.serialized_dag.get_task("outside")
+ session.flush()
+
+ met = TriggerRuleDep().is_met(ti=outside, dep_context=DepContext(), session=session)
+
+ assert outside.task.upstream_task_ids == {loop.gate_task_id}
+ assert met is (gate_state in GATE_RULE_MATRIX[rule])
+
+
@pytest.mark.parametrize("current_state", [SUCCESS, FAILED])
def test_trigger_rule_ignores_other_loop_passes(dag_maker, session, current_state):
@task_group
diff --git a/devel-common/src/tests_common/test_utils/asserts.py b/devel-common/src/tests_common/test_utils/asserts.py
index e938bb0abbf..ab79ad768a0 100644
--- a/devel-common/src/tests_common/test_utils/asserts.py
+++ b/devel-common/src/tests_common/test_utils/asserts.py
@@ -41,6 +41,8 @@ if TYPE_CHECKING:
from sqlalchemy.orm import ORMExecuteState
+ from airflow.models.taskinstance import TaskInstance
+
log = logging.getLogger(__name__)
@@ -282,3 +284,21 @@ def assert_no_cartesian_products() -> Generator[list[Select], None, None]:
for element in joined - set(linter.froms)
)
assert not problems, "Cartesian product in generated SQL:\n" + "\n".join(sorted(set(problems)))
+
+
+@contextmanager
+def count_loaded_task_instances(task_id: str) -> Generator[list[TaskInstance], None, None]:
+ """Collect the ``task_id`` task instance rows the ORM loads inside the block."""
+ from airflow.models.taskinstance import TaskInstance
+
+ loaded: list[TaskInstance] = []
+
+ def listener(target: TaskInstance, _context: object) -> None:
+ if target.task_id == task_id:
+ loaded.append(target)
+
+ event.listen(TaskInstance, "load", listener)
+ try:
+ yield loaded
+ finally:
+ event.remove(TaskInstance, "load", listener)
diff --git a/go-sdk/pkg/execution/genmodels/models.gen.go b/go-sdk/pkg/execution/genmodels/models.gen.go
index 49a2f603a37..d7d779b15cf 100644
--- a/go-sdk/pkg/execution/genmodels/models.gen.go
+++ b/go-sdk/pkg/execution/genmodels/models.gen.go
@@ -1142,6 +1142,9 @@ type GetXCom struct {
// MapIndex corresponds to the JSON schema field "map_index".
MapIndex interface{} `msgpack:"map_index,omitempty"`
+ // PreviousIteration corresponds to the JSON schema field "previous_iteration".
+ PreviousIteration bool `msgpack:"previous_iteration,omitempty"`
+
// RunID corresponds to the JSON schema field "run_id".
RunID string `msgpack:"run_id"`
@@ -1160,6 +1163,9 @@ type GetXComCount struct {
// Key corresponds to the JSON schema field "key".
Key string `msgpack:"key"`
+ // PreviousIteration corresponds to the JSON schema field "previous_iteration".
+ PreviousIteration bool `msgpack:"previous_iteration,omitempty"`
+
// RunID corresponds to the JSON schema field "run_id".
RunID string `msgpack:"run_id"`
@@ -1180,6 +1186,9 @@ type GetXComSequenceItem struct {
// Offset corresponds to the JSON schema field "offset".
Offset int `msgpack:"offset"`
+ // PreviousIteration corresponds to the JSON schema field "previous_iteration".
+ PreviousIteration bool `msgpack:"previous_iteration,omitempty"`
+
// RunID corresponds to the JSON schema field "run_id".
RunID string `msgpack:"run_id"`
@@ -1200,6 +1209,9 @@ type GetXComSequenceSlice struct {
// Key corresponds to the JSON schema field "key".
Key string `msgpack:"key"`
+ // PreviousIteration corresponds to the JSON schema field "previous_iteration".
+ PreviousIteration bool `msgpack:"previous_iteration,omitempty"`
+
// RunID corresponds to the JSON schema field "run_id".
RunID string `msgpack:"run_id"`
@@ -1306,6 +1318,24 @@ type LiteralArgBinding struct {
type LogicalDates []time.Time
+// Pinned loop definition and enclosing iteration for a task execution.
+type LoopContext struct {
+ // Index corresponds to the JSON schema field "index".
+ Index int `msgpack:"index"`
+
+ // MaxIterations corresponds to the JSON schema field "max_iterations".
+ MaxIterations int `msgpack:"max_iterations"`
+
+ // NodeID corresponds to the JSON schema field "node_id".
+ NodeID string `msgpack:"node_id"`
+
+ // TerminalIsMapped corresponds to the JSON schema field "terminal_is_mapped".
+ TerminalIsMapped bool `msgpack:"terminal_is_mapped"`
+
+ // TerminalTaskID corresponds to the JSON schema field "terminal_task_id".
+ TerminalTaskID string `msgpack:"terminal_task_id"`
+}
+
// Add a new value to be redacted in task logs.
type MaskSecret struct {
// Name corresponds to the JSON schema field "name".
@@ -1654,6 +1684,9 @@ type TIRunContext struct {
// DagRun corresponds to the JSON schema field "dag_run".
DagRun DagRun `msgpack:"dag_run"`
+ // Loop corresponds to the JSON schema field "loop".
+ Loop *LoopContext `msgpack:"loop,omitempty"`
+
// MaxTries corresponds to the JSON schema field "max_tries".
MaxTries int `msgpack:"max_tries"`
@@ -1847,6 +1880,19 @@ type TaskStatesResult struct {
Type string `msgpack:"type,omitempty"`
}
+// Variable schema for responses with fields that are needed for Runtime.
+type VariableResponse struct {
+ // Key corresponds to the JSON schema field "key".
+ Key string `msgpack:"key"`
+
+ // Value corresponds to the JSON schema field "value".
+ Value interface{} `msgpack:"value"`
+}
+
+type Warnings []interface{}
+
+type TriggerKwargs map[string]interface{}
+
type TriggerDagRun struct {
// Conf corresponds to the JSON schema field "conf".
Conf *Conf `msgpack:"conf,omitempty"`
@@ -1876,8 +1922,6 @@ type TriggerDagRun struct {
Type string `msgpack:"type,omitempty"`
}
-type TriggerKwargs map[string]interface{}
-
type UpdateDagRunNote struct {
// Note corresponds to the JSON schema field "note".
Note interface{} `msgpack:"note"`
@@ -1889,19 +1933,8 @@ type UpdateDagRunNote struct {
Type string `msgpack:"type,omitempty"`
}
-// Variable schema for responses with fields that are needed for Runtime.
-type VariableResponse struct {
- // Key corresponds to the JSON schema field "key".
- Key string `msgpack:"key"`
-
- // Value corresponds to the JSON schema field "value".
- Value interface{} `msgpack:"value"`
-}
-
type VersionData map[string]interface{}
-type Warnings []interface{}
-
// Update the response content part of an existing Human-in-the-loop response.
type UpdateHITLDetail struct {
// ChosenOptions corresponds to the JSON schema field "chosen_options".
diff --git a/go-sdk/schema/supervisor-schema.json b/go-sdk/schema/supervisor-schema.json
index b0d5797b788..7e09a736de7 100644
--- a/go-sdk/schema/supervisor-schema.json
+++ b/go-sdk/schema/supervisor-schema.json
@@ -2836,6 +2836,11 @@
},
"GetXCom": {
"properties": {
+ "previous_iteration": {
+ "default": false,
+ "title": "Previous Iteration",
+ "type": "boolean"
+ },
"key": {
"title": "Key",
"type": "string"
@@ -2888,6 +2893,11 @@
"GetXComCount": {
"description": "Get the number of (mapped) XCom values available.",
"properties": {
+ "previous_iteration": {
+ "default": false,
+ "title": "Previous Iteration",
+ "type": "boolean"
+ },
"key": {
"title": "Key",
"type": "string"
@@ -2922,6 +2932,11 @@
},
"GetXComSequenceItem": {
"properties": {
+ "previous_iteration": {
+ "default": false,
+ "title": "Previous Iteration",
+ "type": "boolean"
+ },
"key": {
"title": "Key",
"type": "string"
@@ -2961,6 +2976,11 @@
},
"GetXComSequenceSlice": {
"properties": {
+ "previous_iteration": {
+ "default": false,
+ "title": "Previous Iteration",
+ "type": "boolean"
+ },
"key": {
"title": "Key",
"type": "string"
@@ -3222,6 +3242,42 @@
"title": "LazyDeserializedDAG",
"type": "object"
},
+ "LoopContext": {
+ "description": "Pinned loop definition and enclosing iteration for a task execution.",
+ "properties": {
+ "node_id": {
+ "title": "Node Id",
+ "type": "string"
+ },
+ "index": {
+ "minimum": 0,
+ "title": "Index",
+ "type": "integer"
+ },
+ "max_iterations": {
+ "exclusiveMinimum": 0,
+ "title": "Max Iterations",
+ "type": "integer"
+ },
+ "terminal_task_id": {
+ "title": "Terminal Task Id",
+ "type": "string"
+ },
+ "terminal_is_mapped": {
+ "title": "Terminal Is Mapped",
+ "type": "boolean"
+ }
+ },
+ "required": [
+ "node_id",
+ "index",
+ "max_iterations",
+ "terminal_task_id",
+ "terminal_is_mapped"
+ ],
+ "title": "LoopContext",
+ "type": "object"
+ },
"MaskSecret": {
"description": "Add a new value to be redacted in task logs.",
"properties": {
@@ -5154,6 +5210,17 @@
"dag_run": {
"$ref": "#/$defs/DagRun"
},
+ "loop": {
+ "anyOf": [
+ {
+ "$ref": "#/$defs/LoopContext"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null
+ },
"task_reschedule_count": {
"default": 0,
"title": "Task Reschedule Count",
diff --git a/providers/edge3/tests/unit/edge3/worker_api/routes/test_logs.py b/providers/edge3/tests/unit/edge3/worker_api/routes/test_logs.py
index ea9245f96ec..d69dc99b1c8 100644
--- a/providers/edge3/tests/unit/edge3/worker_api/routes/test_logs.py
+++ b/providers/edge3/tests/unit/edge3/worker_api/routes/test_logs.py
@@ -75,8 +75,9 @@ class TestLogsApiRoutes:
create_loop(body, max_iterations=3)
dr = dag_maker.create_dagrun(run_id="loop_run")
ti = next(ti for ti in dr.task_instances if ti.task_id == "body.work")
- region = DynamicRegion(dag_id=dr.dag_id, run_id=dr.run_id, node_id="body")
- session.add(region)
+ region = DynamicRegion.get_or_create(
+ dag_id=dr.dag_id, run_id=dr.run_id, node_id="body", session=session
+ )
session.flush()
ti.region_id, ti.region_index = region.id, 2
session.add(
@@ -118,19 +119,20 @@ class TestLogsApiRoutes:
create_loop(body, max_iterations=3)
dr = dag_maker.create_dagrun(run_id="mapped_loop_run")
first = next(ti for ti in dr.task_instances if ti.task_id == "body.mapped")
- loop_region = DynamicRegion(dag_id=dr.dag_id, run_id=dr.run_id, node_id="body")
- session.add(loop_region)
+ loop_region = DynamicRegion.get_or_create(
+ dag_id=dr.dag_id, run_id=dr.run_id, node_id="body", session=session
+ )
session.flush()
tis = []
for iteration in range(2):
- region = DynamicRegion(
+ region = DynamicRegion.get_or_create(
dag_id=dr.dag_id,
run_id=dr.run_id,
node_id="body.mapped",
parent_region_id=loop_region.id,
parent_region_index=iteration,
+ session=session,
)
- session.add(region)
session.flush()
ti = (
first
diff --git a/scripts/ci/prek/check_template_context_variable_in_sync.py b/scripts/ci/prek/check_template_context_variable_in_sync.py
index 57459a863f3..88bc3cbed3a 100755
--- a/scripts/ci/prek/check_template_context_variable_in_sync.py
+++ b/scripts/ci/prek/check_template_context_variable_in_sync.py
@@ -51,6 +51,8 @@ IGNORE = {
# AIP-103: task_state_store/asset_state_store aren't documented in templates-ref yet. Will be done in a later PR.
"task_state_store",
"asset_state_store",
+ # Only available inside private loop groups.
+ "loop",
}
diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py
index 872c11cca13..3d6b7db7e10 100644
--- a/task-sdk/src/airflow/sdk/api/client.py
+++ b/task-sdk/src/airflow/sdk/api/client.py
@@ -606,9 +606,14 @@ class XComOperations:
run_id: str,
task_id: str,
key: str,
+ *,
+ previous_iteration: bool = False,
) -> XComCountResponse:
"""Get the number of mapped XCom values."""
- resp = self.client.head(f"xcoms/{dag_id}/{run_id}/{task_id}/{key}")
+ params: dict[str, str | int] = {}
+ if previous_iteration:
+ params["previous_iteration"] = True
+ resp = self.client.head(f"xcoms/{dag_id}/{run_id}/{task_id}/{key}", params=params)
# content_range: str | None
if not (content_range := resp.headers["Content-Range"]) or not content_range.startswith(
@@ -625,9 +630,13 @@ class XComOperations:
key: str,
map_index: int | None = None,
include_prior_dates: bool = False,
+ *,
+ previous_iteration: bool = False,
) -> XComResponse:
"""Get a XCom value from the API server."""
params: dict[str, str | int] = {}
+ if previous_iteration:
+ params["previous_iteration"] = True
if map_index is not None and map_index >= 0:
params.update({"map_index": map_index})
if include_prior_dates:
@@ -704,9 +713,17 @@ class XComOperations:
task_id: str,
key: str,
offset: int,
+ *,
+ previous_iteration: bool = False,
) -> XComSequenceIndexResponse | ErrorResponse:
+ params: dict[str, str | int] = {}
+ if previous_iteration:
+ params["previous_iteration"] = True
try:
- resp = self.client.get(f"xcoms/{dag_id}/{run_id}/{task_id}/{key}/item/{offset}")
+ resp = self.client.get(
+ f"xcoms/{dag_id}/{run_id}/{task_id}/{key}/item/{offset}",
+ params=params,
+ )
except ServerResponseError as e:
if e.response.status_code == HTTPStatus.NOT_FOUND:
log.error(
@@ -742,8 +759,12 @@ class XComOperations:
stop: int | None,
step: int | None,
include_prior_dates: bool = False,
+ *,
+ previous_iteration: bool = False,
) -> XComSequenceSliceResponse:
params: dict[str, str | int] = {}
+ if previous_iteration:
+ params["previous_iteration"] = True
if start is not None:
params["start"] = start
if stop is not None:
diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
index bde0a9560d7..d6e9e45f0d2 100644
--- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
+++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
@@ -233,6 +233,18 @@ class IntermediateTIState(str, Enum):
AWAITING_INPUT = "awaiting_input"
+class LoopContext(BaseModel):
+ """
+ Pinned loop definition and enclosing iteration for a task execution.
+ """
+
+ node_id: Annotated[str, Field(title="Node Id")]
+ index: Annotated[int, Field(ge=0, title="Index")]
+ max_iterations: Annotated[int, Field(gt=0, title="Max Iterations")]
+ terminal_task_id: Annotated[str, Field(title="Terminal Task Id")]
+ terminal_is_mapped: Annotated[bool, Field(title="Terminal Is Mapped")]
+
+
class PrevSuccessfulDagRunResponse(BaseModel):
"""
Schema for response with previous successful DagRun information for Task Template Context.
@@ -842,6 +854,7 @@ class TIRunContext(BaseModel):
"""
dag_run: DagRun
+ loop: LoopContext | None = None
task_reschedule_count: Annotated[int | None, Field(title="Task Reschedule Count")] = 0
max_tries: Annotated[int, Field(title="Max Tries")]
variables: Annotated[list[VariableResponse] | None, Field(title="Variables")] = None
diff --git a/task-sdk/src/airflow/sdk/bases/decorator.py b/task-sdk/src/airflow/sdk/bases/decorator.py
index baaee4a9894..896bedb3abc 100644
--- a/task-sdk/src/airflow/sdk/bases/decorator.py
+++ b/task-sdk/src/airflow/sdk/bases/decorator.py
@@ -45,7 +45,7 @@ from airflow.sdk.definitions._internal.expandinput import (
ListOfDictsExpandInput,
is_mappable,
)
-from airflow.sdk.definitions._internal.loop import check_dag_result_outside_loop
+from airflow.sdk.definitions._internal.loop import LoopTaskGroup, check_dag_result_outside_loop
from airflow.sdk.definitions._internal.types import NOTSET
from airflow.sdk.definitions.asset import Asset
from airflow.sdk.definitions.context import KNOWN_CONTEXT_KEYS
@@ -68,6 +68,14 @@ if TYPE_CHECKING:
from airflow.sdk.definitions.taskgroup import TaskGroup
+def _context_keys_for_task_group(task_group: TaskGroup | None) -> set[str]:
+ while task_group is not None:
+ if isinstance(task_group, LoopTaskGroup):
+ return KNOWN_CONTEXT_KEYS | {"loop"}
+ task_group = task_group.parent_group
+ return KNOWN_CONTEXT_KEYS
+
+
class ExpandableFactory(Protocol):
"""
Protocol providing inspection against wrapped function.
@@ -332,19 +340,24 @@ class DecoratedOperator(BaseOperator):
# since values for those will be provided when the task is run. Since
# we're not actually running the function, None is good enough here.
signature = inspect.signature(python_callable)
+ task_group = kwargs.get("task_group") or TaskGroupContext.get_current(kwargs.get("dag"))
+ context_keys = _context_keys_for_task_group(task_group)
+ if "loop" in context_keys and (parameter := signature.parameters.get("loop")) is not None:
+ if parameter.kind != inspect.Parameter.KEYWORD_ONLY:
+ raise ValueError("The loop context parameter must be keyword-only")
# Don't allow context argument defaults other than None to avoid ambiguities.
faulty_parameters = [
param.name
for param in signature.parameters.values()
- if param.name in KNOWN_CONTEXT_KEYS and param.default not in (None, inspect.Parameter.empty)
+ if param.name in context_keys and param.default not in (None, inspect.Parameter.empty)
]
if faulty_parameters:
message = f"Context key parameter {faulty_parameters[0]} can't have a default other than None"
raise ValueError(message)
parameters = [
- param.replace(default=None) if param.name in KNOWN_CONTEXT_KEYS else param
+ param.replace(default=None) if param.name in context_keys else param
for param in signature.parameters.values()
]
@@ -368,7 +381,7 @@ class DecoratedOperator(BaseOperator):
and param.default == inspect.Parameter.empty
):
new_parameters.append(param.replace(default=None))
- if param.name not in KNOWN_CONTEXT_KEYS:
+ if param.name not in context_keys:
injected_for_ordering.add(param.name)
else:
new_parameters.append(param)
@@ -576,7 +589,8 @@ class _TaskDecorator(ExpandableFactory, Generic[FParams, FReturn, OperatorSubcla
def _validate_arg_names(self, func: ValidationSource, kwargs: dict[str, Any]):
# Ensure that context variables are not shadowed.
- context_keys_being_mapped = KNOWN_CONTEXT_KEYS.intersection(kwargs)
+ task_group = self.kwargs.get("task_group") or TaskGroupContext.get_current(self.kwargs.get("dag"))
+ context_keys_being_mapped = _context_keys_for_task_group(task_group).intersection(kwargs)
if len(context_keys_being_mapped) == 1:
(name,) = context_keys_being_mapped
raise ValueError(f"cannot call {func}() on task context variable {name!r}")
diff --git a/task-sdk/src/airflow/sdk/bases/xcom.py b/task-sdk/src/airflow/sdk/bases/xcom.py
index 50e54c458b0..1cbe0882845 100644
--- a/task-sdk/src/airflow/sdk/bases/xcom.py
+++ b/task-sdk/src/airflow/sdk/bases/xcom.py
@@ -291,6 +291,7 @@ class BaseXCom:
run_id: str,
map_index: int | None = None,
include_prior_dates: bool = False,
+ _previous_iteration: bool = False,
) -> Any | None:
"""
Retrieve an XCom value, optionally meeting certain criteria.
@@ -325,6 +326,7 @@ class BaseXCom:
run_id=run_id,
map_index=map_index,
include_prior_dates=include_prior_dates,
+ previous_iteration=_previous_iteration,
),
)
diff --git a/task-sdk/src/airflow/sdk/definitions/_internal/loop.py b/task-sdk/src/airflow/sdk/definitions/_internal/loop.py
index da5f6bca753..d4c3767ccaa 100644
--- a/task-sdk/src/airflow/sdk/definitions/_internal/loop.py
+++ b/task-sdk/src/airflow/sdk/definitions/_internal/loop.py
@@ -16,6 +16,7 @@
# under the License.
from __future__ import annotations
+import inspect
from collections.abc import Callable
from typing import TYPE_CHECKING, Any
@@ -25,8 +26,12 @@ from airflow.sdk.bases.operator import BaseOperator
from airflow.sdk.definitions.taskgroup import MappedTaskGroup, TaskGroup
if TYPE_CHECKING:
+ from airflow.sdk.definitions.context import Context
from airflow.sdk.definitions.decorators.task_group import _TaskGroupFactory
+LOOP_XCOM_PREFIX = "_airflow_loop_"
+LOOP_DECISION_KEY = f"{LOOP_XCOM_PREFIX}decision"
+
def _validate_max_iterations(instance, attribute, value: int) -> None:
if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
@@ -53,6 +58,55 @@ class LoopTaskGroup(TaskGroup):
super().__attrs_post_init__()
+def get_enclosing_loop(operator: Any) -> LoopTaskGroup | None:
+ group = operator.task_group
+ while isinstance(group, TaskGroup):
+ if isinstance(group, LoopTaskGroup):
+ return group
+ group = group.parent_group
+ return None
+
+
+def is_loop_task_read_from_outside(producer: Any, consumer: Any) -> bool:
+ """Tell whether ``consumer`` reads ``producer``, a task of a loop, without being in that loop itself."""
+ loop = get_enclosing_loop(producer)
+ if loop is None:
+ return False
+ consumer_loop = get_enclosing_loop(consumer)
+ return consumer_loop is None or consumer_loop.group_id != loop.group_id
+
+
+def get_edge_source(source: Any, target: Any) -> Any:
+ """
+ Return the task that ``target`` depends on when it is made downstream of ``source``.
+
+ A task outside a loop waits for the loop's gate rather than for one of its tasks: the gate is what
+ creates the next iteration, so a member alone finishes before the loop has produced its later
+ iterations and a consumer would see only some of them.
+ """
+ loop = get_enclosing_loop(source)
+ if loop is None or (gate_task_id := getattr(loop, "gate_task_id", None)) is None:
+ return source
+ target_loop = get_enclosing_loop(target)
+ if target_loop is not None and target_loop.group_id == loop.group_id:
+ return source
+ return source.dag.get_task(gate_task_id)
+
+
+def redirect_outside_edges_to_gate(group: LoopTaskGroup, gate: BaseOperator) -> None:
+ """Make tasks outside the loop that the body already pointed at depend on the gate instead."""
+ for task in list(group):
+ for node_id in list(task.downstream_task_ids):
+ downstream = gate.dag.get_task(node_id)
+ inside = get_enclosing_loop(downstream)
+ if inside is not None and inside.group_id == group.group_id:
+ continue
+ task.downstream_task_ids.discard(node_id)
+ downstream.upstream_task_ids.discard(task.node_id)
+ gate.downstream_task_ids.add(node_id)
+ downstream.upstream_task_ids.add(gate.node_id)
+
+
def check_dag_result_outside_loop(operator: Any) -> None:
"""Raise if the operator sits in a loop group, whose tasks cannot be the Dag result."""
group = operator.task_group
@@ -106,6 +160,12 @@ class LoopGateOperator(BaseOperator):
super().__init__(**kwargs)
self.until = until
+ def execute(self, context: Context) -> None:
+ # execution_time.loop imports LOOP_DECISION_KEY from this module.
+ from airflow.sdk.execution_time.loop import execute_loop_gate
+
+ execute_loop_gate(self, context)
+
def create_loop(
factory: _TaskGroupFactory,
@@ -119,6 +179,10 @@ def create_loop(
if until is not None and not callable(until):
raise TypeError("until must be callable")
+ if until is not None and (
+ inspect.iscoroutinefunction(until) or inspect.iscoroutinefunction(type(until).__call__)
+ ):
+ raise TypeError("until must be synchronous")
group = LoopTaskGroup(
add_suffix_on_collision=True,
**factory.tg_kwargs,
@@ -142,5 +206,6 @@ def create_loop(
terminal >> gate
group.terminal_task_id = terminal.task_id
group.gate_task_id = gate.task_id
+ redirect_outside_edges_to_gate(group, gate)
factory._task_group_created = True
return group
diff --git a/task-sdk/src/airflow/sdk/definitions/_internal/node.py b/task-sdk/src/airflow/sdk/definitions/_internal/node.py
index 803ca837825..ac0a61d6277 100644
--- a/task-sdk/src/airflow/sdk/definitions/_internal/node.py
+++ b/task-sdk/src/airflow/sdk/definitions/_internal/node.py
@@ -92,6 +92,7 @@ class DAGNode(GenericDAGNode["DAG", "Operator", "TaskGroup"], DependencyMixin, m
) -> None:
"""Set relatives for the task or task list."""
from airflow.sdk.bases.operator import BaseOperator
+ from airflow.sdk.definitions._internal.loop import get_edge_source
from airflow.sdk.definitions.mappedoperator import MappedOperator
if not isinstance(task_or_task_list, Sequence):
@@ -131,15 +132,17 @@ class DAGNode(GenericDAGNode["DAG", "Operator", "TaskGroup"], DependencyMixin, m
# If the other task does not yet have a Dag, add it to the same Dag as this task and
dag.add_task(task) # type: ignore[arg-type]
if upstream:
- task.downstream_task_ids.add(self.node_id)
- self.upstream_task_ids.add(task.node_id)
+ source = get_edge_source(task, self)
+ source.downstream_task_ids.add(self.node_id)
+ self.upstream_task_ids.add(source.node_id)
if edge_modifier:
- edge_modifier.add_edge_info(dag, task.node_id, self.node_id)
+ edge_modifier.add_edge_info(dag, source.node_id, self.node_id)
else:
- self.downstream_task_ids.add(task.node_id)
- task.upstream_task_ids.add(self.node_id)
+ source = get_edge_source(self, task)
+ source.downstream_task_ids.add(task.node_id)
+ task.upstream_task_ids.add(source.node_id)
if edge_modifier:
- edge_modifier.add_edge_info(dag, self.node_id, task.node_id)
+ edge_modifier.add_edge_info(dag, source.node_id, task.node_id)
def set_downstream(
self,
diff --git a/task-sdk/src/airflow/sdk/definitions/context.py b/task-sdk/src/airflow/sdk/definitions/context.py
index 98c8051a247..6f85470cc68 100644
--- a/task-sdk/src/airflow/sdk/definitions/context.py
+++ b/task-sdk/src/airflow/sdk/definitions/context.py
@@ -33,6 +33,7 @@ if TYPE_CHECKING:
InletEventsAccessors,
TaskStateStoreAccessor,
)
+ from airflow.sdk.execution_time.loop import LoopContextAccessor
from airflow.sdk.types import (
DagRunProtocol,
Operator,
@@ -57,6 +58,7 @@ class Context(TypedDict, total=False):
inlets: list
inlet_events: InletEventsAccessors
logical_date: DateTime
+ loop: LoopContextAccessor
macros: Any
map_index_template: NotRequired[str | None]
outlets: list
@@ -90,7 +92,7 @@ class Context(TypedDict, total=False):
var: Any
-KNOWN_CONTEXT_KEYS: set[str] = set(Context.__annotations__.keys())
+KNOWN_CONTEXT_KEYS: set[str] = set(Context.__annotations__.keys()) - {"loop"}
def context_merge(context: Context, *args: Any, **kwargs: Any) -> None:
diff --git a/task-sdk/src/airflow/sdk/definitions/dag.py b/task-sdk/src/airflow/sdk/definitions/dag.py
index be581d32105..14ddc937dda 100644
--- a/task-sdk/src/airflow/sdk/definitions/dag.py
+++ b/task-sdk/src/airflow/sdk/definitions/dag.py
@@ -1488,7 +1488,8 @@ def _run_task(
from airflow.utils.session import create_session
taskrun_result: TaskRunResult | None
- log.info("[DAG TEST] starting task_id=%s map_index=%s", ti.task_id, ti.map_index)
+ map_index = ti.region_index if task.get_needs_expansion() else -1
+ log.info("[DAG TEST] starting task_id=%s map_index=%s", ti.task_id, map_index)
while True:
try:
log.info("[DAG TEST] running task %s", ti)
@@ -1507,7 +1508,9 @@ def _run_task(
dag_id=ti.dag_id,
run_id=ti.run_id,
try_number=ti.try_number,
- map_index=ti.map_index,
+ map_index=map_index,
+ region_id=ti.region_id,
+ region_index=ti.region_index,
dag_version_id=UUID(str(ti.dag_version_id)),
)
@@ -1559,7 +1562,7 @@ def _run_task(
break
raise
- log.info("[DAG TEST] end task task_id=%s map_index=%s", ti.task_id, ti.map_index)
+ log.info("[DAG TEST] end task task_id=%s map_index=%s", ti.task_id, map_index)
return taskrun_result
diff --git a/task-sdk/src/airflow/sdk/definitions/decorators/task_group.py b/task-sdk/src/airflow/sdk/definitions/decorators/task_group.py
index 34943dd12a0..b0148278fb4 100644
--- a/task-sdk/src/airflow/sdk/definitions/decorators/task_group.py
+++ b/task-sdk/src/airflow/sdk/definitions/decorators/task_group.py
@@ -39,6 +39,7 @@ from airflow.sdk.definitions._internal.expandinput import (
ListOfDictsExpandInput,
MappedArgument,
)
+from airflow.sdk.definitions._internal.loop import create_loop
from airflow.sdk.definitions._internal.node import DAGNode
from airflow.sdk.definitions.mappedoperator import ensure_xcomarg_return_value, prevent_duplicates
from airflow.sdk.definitions.taskgroup import MappedTaskGroup, TaskGroup
@@ -126,14 +127,33 @@ class _TaskGroupFactory(ExpandableFactory, Generic[FParams, FReturn]):
def override(self, **kwargs: Any) -> _TaskGroupFactory[FParams, FReturn]:
# TODO: FIXME when mypy gets compatible with new attrs
- return attr.evolve(self, tg_kwargs={**self.tg_kwargs, **kwargs}) # type: ignore[arg-type]
+ evolved = attr.evolve(self, tg_kwargs={**self.tg_kwargs, **kwargs}) # type: ignore[arg-type]
+ # The evolved copy owns the "never mapped" warning, so the intermediate factory must not repeat it.
+ self._task_group_created = True
+ return evolved
+
+ def loop(self, *, max_iterations: int, until: Callable[..., bool] | None = None) -> TaskGroup:
+ """
+ Repeat this group's tasks, creating each iteration when the previous gate continues.
+
+ The body must have one terminal task definition, which may be mapped.
+ Use ``partial()`` to supply body arguments and ``override()`` to configure the group.
+
+ :param max_iterations: Positive upper bound on the number of iterations.
+ :param until: Synchronous callable receiving task context keyword arguments.
+ True stops the loop. False at the iteration limit fails the gate.
+ Omit for a fixed-count loop that succeeds at its limit.
+ """
+ return create_loop(self, max_iterations=max_iterations, until=until)
def partial(self, **kwargs: Any) -> _TaskGroupFactory[FParams, FReturn]:
self._validate_arg_names("partial", kwargs)
prevent_duplicates(self.partial_kwargs, kwargs, fail_reason="duplicate partial")
kwargs.update(self.partial_kwargs)
# TODO: FIXME when mypy gets compatible with new attrs
- return attr.evolve(self, partial_kwargs=kwargs) # type: ignore[arg-type]
+ evolved = attr.evolve(self, partial_kwargs=kwargs) # type: ignore[arg-type]
+ self._task_group_created = True
+ return evolved
def expand(self, **kwargs: OperatorExpandArgument) -> DAGNode:
if not kwargs:
diff --git a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
index a38e47a466b..36b393c165f 100644
--- a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
+++ b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
@@ -28,6 +28,7 @@ import attrs
from airflow.sdk import TriggerRule
from airflow.sdk.definitions._internal.abstractoperator import AbstractOperator
+from airflow.sdk.definitions._internal.loop import is_loop_task_read_from_outside
from airflow.sdk.definitions._internal.mixins import DependencyMixin, ResolveMixin
from airflow.sdk.definitions._internal.setup_teardown import SetupTeardownContext
from airflow.sdk.definitions._internal.types import NOTSET, is_arg_set
@@ -333,7 +334,7 @@ class PlainXComArg(XComArg):
ti = context["ti"]
task_id = self.operator.task_id
- if self.operator.is_mapped:
+ if self.operator.is_mapped or is_loop_task_read_from_outside(self.operator, ti.task):
return LazyXComSequence(xcom_arg=self, ti=ti)
tg = self.operator.get_closest_mapped_task_group()
if tg is None:
diff --git a/task-sdk/src/airflow/sdk/execution_time/comms.py b/task-sdk/src/airflow/sdk/execution_time/comms.py
index 6472863fbd8..7c957027ba7 100644
--- a/task-sdk/src/airflow/sdk/execution_time/comms.py
+++ b/task-sdk/src/airflow/sdk/execution_time/comms.py
@@ -941,6 +941,7 @@ class SkipDownstreamTasks(TISkippedDownstreamTasksStatePayload):
class GetXCom(BaseModel):
+ previous_iteration: bool = False
key: str
dag_id: str
run_id: str
@@ -953,6 +954,7 @@ class GetXCom(BaseModel):
class GetXComCount(BaseModel):
"""Get the number of (mapped) XCom values available."""
+ previous_iteration: bool = False
key: str
dag_id: str
run_id: str
@@ -961,6 +963,7 @@ class GetXComCount(BaseModel):
class GetXComSequenceItem(BaseModel):
+ previous_iteration: bool = False
key: str
dag_id: str
run_id: str
@@ -970,6 +973,7 @@ class GetXComSequenceItem(BaseModel):
class GetXComSequenceSlice(BaseModel):
+ previous_iteration: bool = False
key: str
dag_id: str
run_id: str
diff --git a/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py b/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py
index 4efb0b71368..4ce198d14f5 100644
--- a/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py
+++ b/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py
@@ -65,6 +65,7 @@ class LazyXComSequence(Sequence[T]):
_len: int | None = attrs.field(init=False, default=None)
_xcom_arg: PlainXComArg = attrs.field(alias="xcom_arg")
_ti: RuntimeTaskInstance = attrs.field(alias="ti")
+ _previous_iteration: bool = attrs.field(default=False, alias="previous_iteration", kw_only=True)
def __repr__(self) -> str:
if self._len is not None:
@@ -96,6 +97,7 @@ class LazyXComSequence(Sequence[T]):
msg = SUPERVISOR_COMMS.send(
GetXComCount(
+ previous_iteration=self._previous_iteration,
key=self._xcom_arg.key,
dag_id=task.dag_id,
run_id=self._ti.run_id,
@@ -131,6 +133,7 @@ class LazyXComSequence(Sequence[T]):
source = (xcom_arg := self._xcom_arg).operator
msg = SUPERVISOR_COMMS.send(
GetXComSequenceSlice(
+ previous_iteration=self._previous_iteration,
key=xcom_arg.key,
dag_id=source.dag_id,
task_id=source.task_id,
@@ -152,6 +155,7 @@ class LazyXComSequence(Sequence[T]):
source = (xcom_arg := self._xcom_arg).operator
msg = SUPERVISOR_COMMS.send(
GetXComSequenceItem(
+ previous_iteration=self._previous_iteration,
key=xcom_arg.key,
dag_id=source.dag_id,
task_id=source.task_id,
diff --git a/task-sdk/src/airflow/sdk/execution_time/loop.py b/task-sdk/src/airflow/sdk/execution_time/loop.py
new file mode 100644
index 00000000000..a3b52282b1f
--- /dev/null
+++ b/task-sdk/src/airflow/sdk/execution_time/loop.py
@@ -0,0 +1,104 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import attrs
+
+from airflow.sdk.bases.decorator import determine_kwargs
+from airflow.sdk.definitions._internal.loop import LOOP_DECISION_KEY
+from airflow.sdk.definitions.xcom_arg import PlainXComArg
+from airflow.sdk.exceptions import AirflowFailException
+from airflow.sdk.execution_time.comms import SetXCom
+from airflow.sdk.execution_time.lazy_sequence import LazyXComSequence
+from airflow.sdk.execution_time.xcom import XCom
+
+if TYPE_CHECKING:
+ from airflow.sdk.api.datamodels._generated import LoopContext
+ from airflow.sdk.definitions._internal.loop import LoopGateOperator
+ from airflow.sdk.definitions.context import Context
+ from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
+
+
+class LoopMaxIterationsExceeded(AirflowFailException):
+ """The loop condition remained false at its iteration cap."""
+
+
+@attrs.define(frozen=True)
+class LoopContextAccessor:
+ _descriptor: LoopContext
+ _ti: RuntimeTaskInstance
+
+ @property
+ def index(self) -> int:
+ return self._descriptor.index
+
+ @property
+ def max_iterations(self) -> int:
+ return self._descriptor.max_iterations
+
+ @property
+ def previous(self) -> Any:
+ return None if self.index == 0 else self._result(previous_iteration=True)
+
+ @property
+ def result(self) -> Any:
+ return self._result(previous_iteration=False)
+
+ def _result(self, *, previous_iteration: bool) -> Any:
+ task = self._ti.task.dag.get_task(self._descriptor.terminal_task_id)
+ if self._descriptor.terminal_is_mapped:
+ return LazyXComSequence(
+ xcom_arg=PlainXComArg(task), ti=self._ti, previous_iteration=previous_iteration
+ )
+ return XCom.get_one(
+ key=XCom.XCOM_RETURN_KEY,
+ dag_id=self._ti.dag_id,
+ run_id=self._ti.run_id,
+ task_id=task.task_id,
+ _previous_iteration=previous_iteration,
+ )
+
+
+def execute_loop_gate(gate: LoopGateOperator, context: Context) -> None:
+ # task_runner imports LoopContextAccessor from this module.
+ from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+
+ loop = context["loop"]
+ if gate.until is None:
+ stop = loop.index + 1 >= loop.max_iterations
+ else:
+ stop = gate.until(**determine_kwargs(gate.until, (), context))
+ if not isinstance(stop, bool):
+ raise AirflowFailException(f"until must return a bool, got {type(stop).__name__}")
+ if not stop and loop.index + 1 >= loop.max_iterations:
+ raise LoopMaxIterationsExceeded(
+ f"Loop condition remained false at max_iterations={loop.max_iterations}"
+ )
+ decision = "stop" if stop else "continue"
+ gate.log.info("Loop iteration %s: %s (max_iterations=%s)", loop.index, decision, loop.max_iterations)
+ ti = context["ti"]
+ SUPERVISOR_COMMS.send(
+ SetXCom(
+ dag_id=ti.dag_id,
+ run_id=ti.run_id,
+ task_id=ti.task_id,
+ key=LOOP_DECISION_KEY,
+ value=decision,
+ )
+ )
diff --git a/task-sdk/src/airflow/sdk/execution_time/request_handlers.py b/task-sdk/src/airflow/sdk/execution_time/request_handlers.py
index 594a871f3a8..2a2d16833df 100644
--- a/task-sdk/src/airflow/sdk/execution_time/request_handlers.py
+++ b/task-sdk/src/airflow/sdk/execution_time/request_handlers.py
@@ -257,7 +257,13 @@ def handle_get_prev_successful_dag_run(
def handle_get_xcom_count(client: Client, msg: GetXComCount) -> tuple[BaseModel | None, dict[str, bool]]:
"""Fetch XCom count metadata."""
- resp = client.xcoms.head(msg.dag_id, msg.run_id, msg.task_id, msg.key)
+ resp = client.xcoms.head(
+ msg.dag_id,
+ msg.run_id,
+ msg.task_id,
+ msg.key,
+ previous_iteration=msg.previous_iteration,
+ )
return resp, {}
@@ -271,6 +277,7 @@ def handle_get_xcom_sequence_item(
msg.task_id,
msg.key,
msg.offset,
+ previous_iteration=msg.previous_iteration,
)
if isinstance(xcom, XComSequenceIndexResponse):
return XComSequenceIndexResult.from_response(xcom), {}
@@ -290,6 +297,7 @@ def handle_get_xcom_sequence_slice(
msg.stop,
msg.step,
msg.include_prior_dates,
+ previous_iteration=msg.previous_iteration,
)
if isinstance(xcoms, XComSequenceSliceResponse):
return XComSequenceSliceResult.from_response(xcoms), {}
@@ -305,6 +313,7 @@ def handle_get_xcom(client: Client, msg: GetXCom) -> tuple[BaseModel | None, dic
msg.key,
msg.map_index,
msg.include_prior_dates,
+ previous_iteration=msg.previous_iteration,
)
if isinstance(xcom, XComResponse):
xcom_result = XComResult.from_xcom_response(xcom)
diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json
index b0d5797b788..7e09a736de7 100644
--- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json
+++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json
@@ -2836,6 +2836,11 @@
},
"GetXCom": {
"properties": {
+ "previous_iteration": {
+ "default": false,
+ "title": "Previous Iteration",
+ "type": "boolean"
+ },
"key": {
"title": "Key",
"type": "string"
@@ -2888,6 +2893,11 @@
"GetXComCount": {
"description": "Get the number of (mapped) XCom values available.",
"properties": {
+ "previous_iteration": {
+ "default": false,
+ "title": "Previous Iteration",
+ "type": "boolean"
+ },
"key": {
"title": "Key",
"type": "string"
@@ -2922,6 +2932,11 @@
},
"GetXComSequenceItem": {
"properties": {
+ "previous_iteration": {
+ "default": false,
+ "title": "Previous Iteration",
+ "type": "boolean"
+ },
"key": {
"title": "Key",
"type": "string"
@@ -2961,6 +2976,11 @@
},
"GetXComSequenceSlice": {
"properties": {
+ "previous_iteration": {
+ "default": false,
+ "title": "Previous Iteration",
+ "type": "boolean"
+ },
"key": {
"title": "Key",
"type": "string"
@@ -3222,6 +3242,42 @@
"title": "LazyDeserializedDAG",
"type": "object"
},
+ "LoopContext": {
+ "description": "Pinned loop definition and enclosing iteration for a task execution.",
+ "properties": {
+ "node_id": {
+ "title": "Node Id",
+ "type": "string"
+ },
+ "index": {
+ "minimum": 0,
+ "title": "Index",
+ "type": "integer"
+ },
+ "max_iterations": {
+ "exclusiveMinimum": 0,
+ "title": "Max Iterations",
+ "type": "integer"
+ },
+ "terminal_task_id": {
+ "title": "Terminal Task Id",
+ "type": "string"
+ },
+ "terminal_is_mapped": {
+ "title": "Terminal Is Mapped",
+ "type": "boolean"
+ }
+ },
+ "required": [
+ "node_id",
+ "index",
+ "max_iterations",
+ "terminal_task_id",
+ "terminal_is_mapped"
+ ],
+ "title": "LoopContext",
+ "type": "object"
+ },
"MaskSecret": {
"description": "Add a new value to be redacted in task logs.",
"properties": {
@@ -5154,6 +5210,17 @@
"dag_run": {
"$ref": "#/$defs/DagRun"
},
+ "loop": {
+ "anyOf": [
+ {
+ "$ref": "#/$defs/LoopContext"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null
+ },
"task_reschedule_count": {
"default": 0,
"title": "Task Reschedule Count",
diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py
index bdbb825c2d1..647e5dc9193 100644
--- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py
+++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py
@@ -40,6 +40,7 @@ def get_bundle() -> VersionBundle:
from airflow.sdk.execution_time.schema.versions.v2026_10_30 import (
AddArgBindingsToSupervisorTIRunContext,
AddDagDefinitionsToDagFileParsingResult,
+ AddLoopRuntime,
AddRegionCoordinatesToSupervisorTaskInstance,
AddRetryReasonToTaskState,
)
@@ -49,6 +50,7 @@ def get_bundle() -> VersionBundle:
Version(
"2026-10-30",
AddArgBindingsToSupervisorTIRunContext,
+ AddLoopRuntime,
AddRetryReasonToTaskState,
AddDagDefinitionsToDagFileParsingResult,
AddRegionCoordinatesToSupervisorTaskInstance,
diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py
index e093ecd07ec..baffc8c9d0f 100644
--- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py
+++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py
@@ -31,10 +31,27 @@ from cadwyn import VersionChange, schema
from airflow.dag_processing.processor import DagFileParsingResult # noqa: SDK002
from airflow.sdk.api.datamodels._generated import PreviousTIResponse, TaskInstance, TIRunContext
from airflow.sdk.execution_time.comms import (
+ GetXCom,
+ GetXComCount,
+ GetXComSequenceItem,
+ GetXComSequenceSlice,
TaskState,
)
+class AddLoopRuntime(VersionChange):
+ """Carry loop context and previous-iteration reads."""
+
+ description = __doc__
+ instructions_to_migrate_to_previous_version = (
+ schema(TIRunContext).field("loop").didnt_exist,
+ schema(GetXCom).field("previous_iteration").didnt_exist,
+ schema(GetXComCount).field("previous_iteration").didnt_exist,
+ schema(GetXComSequenceItem).field("previous_iteration").didnt_exist,
+ schema(GetXComSequenceSlice).field("previous_iteration").didnt_exist,
+ )
+
+
class AddRegionCoordinatesToSupervisorTaskInstance(VersionChange):
"""Carry task coordinates in supervisor startup and previous-instance messages."""
diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
index aa6aa477c0c..d4fbec7c25a 100644
--- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py
+++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
@@ -63,6 +63,7 @@ from airflow.sdk.bases.xcom import BaseXCom
from airflow.sdk.configuration import conf
from airflow.sdk.coordinators._dag_importer import find_claiming_importer
from airflow.sdk.definitions._internal.dag_parsing_context import _airflow_parsing_context_manager
+from airflow.sdk.definitions._internal.loop import LOOP_XCOM_PREFIX, is_loop_task_read_from_outside
from airflow.sdk.definitions._internal.types import NOTSET, ArgNotSet, is_arg_set
from airflow.sdk.definitions.asset import (
Asset,
@@ -73,6 +74,7 @@ from airflow.sdk.definitions.asset import (
)
from airflow.sdk.definitions.mappedoperator import MappedOperator
from airflow.sdk.definitions.param import process_params
+from airflow.sdk.definitions.xcom_arg import PlainXComArg
from airflow.sdk.exceptions import (
AirflowException,
AirflowFailException,
@@ -85,6 +87,7 @@ from airflow.sdk.exceptions import (
ErrorType,
TaskAwaitingInput,
TaskDeferred,
+ TaskNotFound,
)
from airflow.sdk.execution_time.bundles import initialize_ti_bundle
from airflow.sdk.execution_time.callback_runner import create_executable_runner
@@ -150,6 +153,8 @@ from airflow.sdk.execution_time.email_backend import (
_ErrorEmailNotifier,
_LegacyEmailBackendNotifier,
)
+from airflow.sdk.execution_time.lazy_sequence import LazyXComSequence
+from airflow.sdk.execution_time.loop import LoopContextAccessor
from airflow.sdk.execution_time.sentry import Sentry
from airflow.sdk.execution_time.tracing import detail_span
from airflow.sdk.execution_time.xcom import XCom
@@ -350,6 +355,8 @@ class RuntimeTaskInstance(TaskInstance):
if TYPE_CHECKING:
assert self._cached_template_context is not None
if from_server:
+ if from_server.loop is not None:
+ self._cached_template_context["loop"] = LoopContextAccessor(from_server.loop, self)
dag_run = from_server.dag_run
context_from_server: Context = {
# TODO: Assess if we need to pass these through timezone.coerce_datetime
@@ -493,6 +500,20 @@ class RuntimeTaskInstance(TaskInstance):
map_indexes_iterable,
)
+ def _get_loop_task_sequence(
+ self, task_id: str, *, key: str, dag_id: str, run_id: str, include_prior_dates: bool
+ ) -> LazyXComSequence | None:
+ """Return every iteration of a loop task as a sequence when this task instance is outside that loop."""
+ if include_prior_dates or (dag_id, run_id) != (self.dag_id, self.run_id):
+ return None
+ try:
+ producer = self.task.dag.get_task(task_id)
+ except TaskNotFound:
+ return None
+ if not is_loop_task_read_from_outside(producer, self.task):
+ return None
+ return LazyXComSequence(xcom_arg=PlainXComArg(producer, key), ti=self)
+
def xcom_pull(
self,
task_ids: str | Iterable[str] | None = None,
@@ -551,6 +572,16 @@ class RuntimeTaskInstance(TaskInstance):
xcoms: list[Any] = []
if not is_arg_set(map_indexes_iterable):
+ if single_task_requested:
+ sequence = self._get_loop_task_sequence(
+ task_ids[0],
+ key=key,
+ dag_id=dag_id,
+ run_id=run_id,
+ include_prior_dates=include_prior_dates,
+ )
+ if sequence is not None:
+ return sequence
# map_indexes was not specified — fetch all map indexes for each task
for t_id in task_ids:
values = XCom.get_all(
@@ -600,6 +631,17 @@ class RuntimeTaskInstance(TaskInstance):
xcoms: list[Any] = []
if not is_arg_set(map_indexes_iterable):
+ if single_task_requested and (
+ self._get_loop_task_sequence(
+ task_ids[0],
+ key=key,
+ dag_id=dag_id,
+ run_id=run_id,
+ include_prior_dates=include_prior_dates,
+ )
+ is not None
+ ):
+ return await XCom.aget_all(run_id=run_id, key=key, task_id=task_ids[0], dag_id=dag_id) or []
# map_indexes was not specified — fetch all map indexes for each task
for t_id in task_ids:
values = await XCom.aget_all(
@@ -1575,7 +1617,8 @@ def _run_task_and_map_outcome(
if ti._ti_context_from_server and (keys_to_delete := ti._ti_context_from_server.xcom_keys_to_clear):
for x in keys_to_delete:
log.debug("Clearing XCom with key", key=x)
- XCom.delete(
+ backend = BaseXCom if x.startswith(LOOP_XCOM_PREFIX) else XCom
+ backend.delete(
key=x,
dag_id=ti.dag_id,
task_id=ti.task_id,
diff --git a/task-sdk/tests/task_sdk/definitions/_internal/test_loop.py b/task-sdk/tests/task_sdk/definitions/_internal/test_loop.py
index 28bb2eda92e..d2b8db90654 100644
--- a/task-sdk/tests/task_sdk/definitions/_internal/test_loop.py
+++ b/task-sdk/tests/task_sdk/definitions/_internal/test_loop.py
@@ -20,10 +20,68 @@ from functools import partial
import pytest
-from airflow.sdk import DAG, BaseOperator, TaskGroup, TriggerRule, task, task_group
+from airflow.sdk import (
+ DAG,
+ BaseOperator,
+ Label,
+ TaskGroup,
+ TriggerRule,
+ chain,
+ cross_downstream,
+ task,
+ task_group,
+)
from airflow.sdk.definitions._internal.loop import LoopGateOperator, create_loop
+@pytest.mark.parametrize("mapped", [False, True])
+def test_loop_body_accepts_omitted_keyword_only_loop_context(mapped):
+ @task
+ def terminal(value, *, loop):
+ return value + loop.index
+
+ @task_group
+ def body():
+ with TaskGroup("nested"):
+ if mapped:
+ terminal.expand(value=[1, 2])
+ else:
+ terminal(1)
+
+ with DAG("loop_context_argument", schedule=None) as dag:
+ group = create_loop(body, max_iterations=2)
+
+ assert group.terminal_task_id == "body.nested.terminal"
+ assert dag.get_task(group.gate_task_id).upstream_task_ids == {group.terminal_task_id}
+
+
+def test_ordinary_task_loop_argument_remains_required():
+ @task
+ def terminal(*, loop):
+ return loop
+
+ with DAG("ordinary_loop_argument", schedule=None):
+ with pytest.raises(TypeError, match="missing a required (keyword-only )?argument: 'loop'"):
+ terminal()
+ result = terminal(loop="user value")
+
+ assert result.operator.op_kwargs == {"loop": "user value"}
+
+
+def test_loop_context_cannot_be_replaced_by_mapped_input():
+ @task
+ def terminal(*, loop):
+ return loop.index
+
+ @task_group
+ def body():
+ terminal.expand(loop=["user value"])
+
+ with DAG("mapped_loop_context", schedule=None):
+ with pytest.raises(ValueError, match="task context variable 'loop'"):
+ create_loop(body, max_iterations=2)
+
+
def test_loop_dependencies_follow_terminal_instead_of_body_return():
@task_group
def body():
@@ -144,6 +202,25 @@ def test_noncallable_condition_is_rejected(until):
create_loop(body, max_iterations=2, until=until)
+@pytest.mark.parametrize("kind", ["function", "partial", "object"])
+def test_async_condition_is_rejected(kind):
+ @task_group
+ def body():
+ BaseOperator(task_id="terminal")
+
+ async def until(*, loop):
+ return True
+
+ class AsyncCondition:
+ async def __call__(self, *, loop):
+ return True
+
+ condition = {"function": until, "partial": partial(until), "object": AsyncCondition()}[kind]
+ with DAG("async_condition", schedule=None):
+ with pytest.raises(TypeError, match="synchronous"):
+ create_loop(body, max_iterations=2, until=condition)
+
+
@pytest.mark.parametrize("terminal_count", [0, 2])
def test_loop_requires_one_terminal_task_definition(terminal_count):
@task_group
@@ -200,3 +277,155 @@ def test_mapping_whole_loop_is_rejected():
with DAG("mapped_loop", schedule=None):
with pytest.raises(NotImplementedError, match="mapped task group"):
mapped.expand(value=[1, 2])
+
+
+def make_edge_dag():
+ @task_group
+ def body():
+ BaseOperator(task_id="first") >> BaseOperator(task_id="second")
+
+ with DAG("loop_edges", schedule=None) as dag:
+ loop = create_loop(body, max_iterations=3)
+ BaseOperator(task_id="before")
+ BaseOperator(task_id="after")
+ return dag, loop
+
+
+def link_with_rshift(member, consumer):
+ member >> consumer
+
+
+def link_with_lshift(member, consumer):
+ consumer << member
+
+
+def link_with_set_downstream(member, consumer):
+ member.set_downstream(consumer)
+
+
+def link_with_set_upstream(member, consumer):
+ consumer.set_upstream(member)
+
+
+def link_with_list(member, consumer):
+ [member] >> consumer
+
+
+def link_with_label(member, consumer):
+ member >> Label("loop result") >> consumer
+
+
+def link_with_chain(member, consumer):
+ chain(member, consumer)
+
+
+def link_with_cross_downstream(member, consumer):
+ cross_downstream([member], [consumer])
+
+
+def link_with_xcom_arg(member, consumer):
+ consumer.set_upstream(member.output)
+
+
+@pytest.mark.parametrize(
+ "link",
+ [
+ link_with_rshift,
+ link_with_lshift,
+ link_with_set_downstream,
+ link_with_set_upstream,
+ link_with_list,
+ link_with_label,
+ link_with_chain,
+ link_with_cross_downstream,
+ link_with_xcom_arg,
+ ],
+)
+def test_outside_task_depends_on_the_gate_not_on_the_loop_task(link):
+ dag, loop = make_edge_dag()
+
+ link(dag.get_task("body.first"), dag.get_task("after"))
+
+ assert dag.get_task("after").upstream_task_ids == {loop.gate_task_id}
+ assert "after" in dag.get_task(loop.gate_task_id).downstream_task_ids
+ assert dag.get_task("body.first").downstream_task_ids == {"body.second"}
+
+
+def test_task_group_edge_to_an_outside_task_uses_the_gate():
+ dag, loop = make_edge_dag()
+
+ loop >> dag.get_task("after")
+
+ assert dag.get_task("after").upstream_task_ids == {loop.gate_task_id}
+
+
+def test_edges_that_stay_inside_or_enter_the_loop_are_unchanged():
+ dag, loop = make_edge_dag()
+
+ with dag, loop:
+ extra = BaseOperator(task_id="extra")
+ dag.get_task("body.first") >> extra
+ dag.get_task("before") >> dag.get_task("body.first")
+ dag.get_task("before") >> loop
+
+ assert extra.upstream_task_ids == {"body.first"}
+ assert dag.get_task("body.first").upstream_task_ids == {"before"}
+ assert dag.get_task("body.second").upstream_task_ids == {"body.first"}
+ assert dag.get_task(loop.gate_task_id).upstream_task_ids == {"body.second"}
+ assert dag.get_task("before").downstream_task_ids == {"body.first"}
+
+
+def test_xcom_argument_of_a_loop_task_keeps_reading_the_task_but_waits_for_the_gate():
+ @task
+ def produce():
+ return 1
+
+ @task
+ def report(value):
+ return value
+
+ @task_group
+ def body():
+ produce()
+
+ with DAG("loop_xcom_edge", schedule=None):
+ loop = create_loop(body, max_iterations=2)
+ consumer = report(loop["produce"].output)
+
+ assert [arg.operator.task_id for arg in consumer.operator.op_args] == ["body.produce"]
+ assert consumer.operator.upstream_task_ids == {loop.gate_task_id}
+
+
+def test_mapped_consumer_expanded_over_a_loop_task_waits_for_the_gate():
+ @task
+ def produce():
+ return [1, 2]
+
+ @task
+ def use(value):
+ return value
+
+ @task_group
+ def body():
+ produce()
+
+ with DAG("loop_expand_edge", schedule=None):
+ loop = create_loop(body, max_iterations=2)
+ consumer = use.expand(value=loop["produce"].output)
+
+ assert consumer.operator.upstream_task_ids == {loop.gate_task_id}
+
+
+def test_edge_to_an_outside_task_made_while_the_body_is_built_waits_for_the_gate():
+ with DAG("loop_early_edge", schedule=None) as dag:
+ outside = BaseOperator(task_id="outside")
+
+ @task_group
+ def body():
+ BaseOperator(task_id="first") >> outside
+
+ loop = create_loop(body, max_iterations=2)
+
+ assert outside.upstream_task_ids == {loop.gate_task_id}
+ assert dag.get_task("body.first").downstream_task_ids == {loop.gate_task_id}
+ assert outside.task_id in dag.get_task(loop.gate_task_id).downstream_task_ids
diff --git a/task-sdk/tests/task_sdk/definitions/decorators/test_task_group.py b/task-sdk/tests/task_sdk/definitions/decorators/test_task_group.py
index 7a777b1ea19..a3519946f84 100644
--- a/task-sdk/tests/task_sdk/definitions/decorators/test_task_group.py
+++ b/task-sdk/tests/task_sdk/definitions/decorators/test_task_group.py
@@ -17,13 +17,15 @@
# under the License.
from __future__ import annotations
+import gc
+import warnings
from datetime import timedelta
import pendulum
import pytest
from airflow.providers.standard.operators.empty import EmptyOperator
-from airflow.sdk import dag, task, task_group
+from airflow.sdk import DAG, dag, task, task_group
from airflow.sdk.definitions._internal.expandinput import (
DictOfListsExpandInput,
ListOfDictsExpandInput,
@@ -32,6 +34,89 @@ from airflow.sdk.definitions._internal.expandinput import (
from airflow.sdk.definitions.taskgroup import MappedTaskGroup
+@pytest.mark.parametrize("mapped", [False, True])
+@pytest.mark.parametrize("conditional", [False, True])
+def test_loop_preserves_group_configuration_and_body_arguments(mapped, conditional):
+ @task
+ def terminal(value, *, loop):
+ return value + loop.index
+
+ @task_group
+ def body(value):
+ """Refine the input."""
+ if mapped:
+ return terminal.expand(value=[value, value + 1])
+ return terminal(value)
+
+ def converged(*, loop):
+ return loop.index == 1
+
+ with DAG("public_loop", schedule=None) as pipeline:
+ start = EmptyOperator(task_id="start")
+ finish = EmptyOperator(task_id="finish")
+ group = (
+ body.override(group_id="refine")
+ .partial(value=2)
+ .loop(max_iterations=3, until=converged if conditional else None)
+ )
+ start >> group >> finish
+
+ gate = pipeline.get_task("refine.converged" if conditional else "refine.__loop_gate")
+ assert group.doc_md == "Refine the input."
+ assert gate.until is (converged if conditional else None)
+ assert gate.upstream_task_ids == {"refine.terminal"}
+ assert gate.downstream_task_ids == {"finish"}
+ result = pipeline.get_task("refine.terminal")
+ assert result.upstream_task_ids == {"start"}
+ if mapped:
+ assert result.op_kwargs_expand_input.value == {"value": [2, 3]}
+ else:
+ assert result.op_args == (2,)
+
+
+@pytest.mark.parametrize("max_iterations", [0, True])
+def test_loop_rejects_invalid_iteration_limit(max_iterations):
+ @task_group
+ def body():
+ EmptyOperator(task_id="terminal")
+
+ with DAG("invalid_public_loop", schedule=None):
+ with pytest.raises(ValueError, match="positive integer"):
+ body.loop(max_iterations=max_iterations)
+
+
+def test_loop_rejects_multiple_terminal_definitions():
+ @task_group
+ def body():
+ EmptyOperator(task_id="first")
+ EmptyOperator(task_id="second")
+
+ with DAG("ambiguous_public_loop", schedule=None):
+ with pytest.raises(ValueError, match="exactly one terminal"):
+ body.loop(max_iterations=2)
+
+
+@pytest.mark.parametrize("partial_first", [False, True])
+def test_loop_after_partial_and_override_in_either_order_does_not_warn(partial_first):
+ @task_group
+ def body(value):
+ EmptyOperator(task_id="terminal")
+
+ with DAG("ordered_public_loop", schedule=None) as pipeline:
+ with warnings.catch_warnings(record=True) as recorded:
+ warnings.simplefilter("always")
+ if partial_first:
+ factory = body.partial(value=2).override(group_id="refine")
+ else:
+ factory = body.override(group_id="refine").partial(value=2)
+ factory.loop(max_iterations=2)
+ del factory
+ gc.collect()
+
+ assert pipeline.get_task("refine.terminal")
+ assert [str(warning.message) for warning in recorded if "never mapped" in str(warning.message)] == []
+
+
def test_task_group_with_overridden_kwargs():
@task_group(
default_args={
diff --git a/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py b/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py
index 2d1f8b2a4fa..53b1139dacb 100644
--- a/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py
+++ b/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py
@@ -23,12 +23,19 @@ from unittest import mock
import pytest
import structlog
-from airflow.sdk import TaskInstanceState
+from airflow.sdk import BaseOperator, TaskInstanceState, task_group
from airflow.sdk.bases.xcom import BaseXCom
+from airflow.sdk.definitions._internal.loop import create_loop
from airflow.sdk.definitions.dag import DAG
from airflow.sdk.definitions.xcom_arg import PlainXComArg
from airflow.sdk.exceptions import AirflowSkipException
-from airflow.sdk.execution_time.comms import GetXCom, XComResult, XComSequenceSliceResult
+from airflow.sdk.execution_time.comms import (
+ GetXCom,
+ GetXComCount,
+ XComCountResponse,
+ XComResult,
+ XComSequenceSliceResult,
+)
from airflow.sdk.execution_time.lazy_sequence import LazyXComSequence
from airflow.sdk.serde import deserialize, serialize
@@ -416,3 +423,43 @@ class TestPlainXComArgResolveMappedGroup:
assert resolved == "value-0"
ti.xcom_pull.assert_called_once()
assert ti.xcom_pull.call_args.kwargs["map_indexes"] == 0
+
+
+class TestPlainXComArgResolveLoopTask:
+ @pytest.fixture
+ def dag(self):
+ @task_group
+ def body():
+ BaseOperator(task_id="first") >> BaseOperator(task_id="second")
+
+ with DAG("loop_dag", schedule=None) as dag:
+ create_loop(body, max_iterations=3) >> BaseOperator(task_id="after")
+ return dag
+
+ def test_resolve_from_outside_the_loop_returns_a_sequence_of_every_iteration(
+ self, dag, create_runtime_ti, mock_supervisor_comms
+ ):
+ ti = create_runtime_ti(task=dag.get_task("after"))
+ mock_supervisor_comms.send.return_value = XComCountResponse(len=3)
+
+ resolved = PlainXComArg(dag.get_task("body.first")).resolve({"ti": ti})
+
+ assert isinstance(resolved, LazyXComSequence)
+ mock_supervisor_comms.send.assert_not_called()
+ assert len(resolved) == 3
+ sent = mock_supervisor_comms.send.call_args.args[0]
+ assert sent == GetXComCount(
+ key="return_value", dag_id=ti.dag_id, run_id=ti.run_id, task_id="body.first"
+ )
+ assert set(sent.model_dump()) == {"type", "previous_iteration", "key", "dag_id", "run_id", "task_id"}
+
+ def test_resolve_from_inside_the_loop_pulls_the_value_of_the_same_iteration(
+ self, dag, create_runtime_ti, mock_supervisor_comms
+ ):
+ ti = create_runtime_ti(task=dag.get_task("body.second"))
+ mock_supervisor_comms.send.return_value = XComResult(key="return_value", value="value")
+
+ resolved = PlainXComArg(dag.get_task("body.first")).resolve({"ti": ti})
+
+ assert resolved == "value"
+ assert isinstance(mock_supervisor_comms.send.call_args.args[0], GetXCom)
diff --git a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py
index da3f120f063..a5ac1de22d2 100644
--- a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py
+++ b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py
@@ -50,8 +50,12 @@ from task_sdk.execution_time.schema._mock_version_bundle import (
)
from airflow.sdk import TaskInstanceState
-from airflow.sdk.api.datamodels._generated import PreviousTIResponse
+from airflow.sdk.api.datamodels._generated import LoopContext, PreviousTIResponse
from airflow.sdk.execution_time.comms import (
+ GetXCom,
+ GetXComCount,
+ GetXComSequenceItem,
+ GetXComSequenceSlice,
PreviousTIResult,
TaskState,
)
@@ -94,6 +98,23 @@ def test_previous_ti_response_coordinates_follow_supervisor_version(version):
assert result["region_index"] == 3
+@pytest.mark.parametrize(
+ ("model", "extra", "field"),
+ [
+ (GetXCom, {}, "previous_iteration"),
+ (GetXComCount, {}, "previous_iteration"),
+ (GetXComSequenceItem, {"offset": 0}, "previous_iteration"),
+ (GetXComSequenceSlice, {"start": None, "stop": None, "step": None}, "previous_iteration"),
+ ],
+)
+def test_loop_selectors_follow_supervisor_version(model, extra, field):
+ message = model(dag_id="dag", run_id="run", task_id="task", key="key", **extra, **{field: True})
+ migrator = get_schema_version_migrator()
+
+ assert field not in migrator.downgrade(message, "2026-06-16").model_dump()
+ assert migrator.downgrade(message, "2026-10-30").model_dump()[field] is True
+
+
class _IntroduceQueueCapacity(VersionChange):
"""3026-04-17: introduce queue_capacity."""
@@ -485,6 +506,22 @@ class TestRealBundleArgBindingsDowngrade:
out = real_migrator.downgrade(startup_details, "2026-06-16").model_dump()
assert "arg_bindings" not in out["ti_context"]
+ @pytest.mark.parametrize("version", ["2026-06-16", "2026-10-30"])
+ def test_loop_context_follows_supervisor_version(self, real_migrator, startup_details, version):
+ startup_details.ti_context.loop = LoopContext(
+ node_id="body",
+ index=2,
+ max_iterations=4,
+ terminal_task_id="body.terminal",
+ terminal_is_mapped=True,
+ )
+ context = real_migrator.downgrade(startup_details, version).model_dump()["ti_context"]
+
+ if version == "2026-06-16":
+ assert "loop" not in context
+ else:
+ assert context["loop"] == startup_details.ti_context.loop.model_dump()
+
@pytest.mark.parametrize("version", ["2026-06-16", "2026-10-30"])
def test_task_coordinates_follow_supervisor_version(self, real_migrator, startup_details, version):
ti = real_migrator.downgrade(startup_details, version).model_dump()["ti"]
diff --git a/task-sdk/tests/task_sdk/execution_time/test_loop.py b/task-sdk/tests/task_sdk/execution_time/test_loop.py
new file mode 100644
index 00000000000..a2485584670
--- /dev/null
+++ b/task-sdk/tests/task_sdk/execution_time/test_loop.py
@@ -0,0 +1,245 @@
+# 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 pytest
+
+from airflow.sdk import DAG, BaseOperator, TaskInstanceState, task, task_group
+from airflow.sdk.api.datamodels._generated import LoopContext
+from airflow.sdk.bases.xcom import BaseXCom
+from airflow.sdk.definitions._internal.loop import create_loop
+from airflow.sdk.exceptions import AirflowFailException
+from airflow.sdk.execution_time import task_runner
+from airflow.sdk.execution_time.comms import (
+ DeleteXCom,
+ GetXCom,
+ GetXComCount,
+ GetXComSequenceItem,
+ GetXComSequenceSlice,
+ SetXCom,
+ XComCountResponse,
+ XComResult,
+ XComSequenceIndexResult,
+ XComSequenceSliceResult,
+)
+from airflow.sdk.execution_time.lazy_sequence import LazyXComSequence
+from airflow.sdk.execution_time.loop import LoopMaxIterationsExceeded
+
+
+@pytest.fixture
+def loop_ti(create_runtime_ti):
+ def make(*, index=0, max_iterations=3, until=None, map_index=-1):
+ @task_group
+ def body():
+ BaseOperator(task_id="terminal")
+
+ with DAG("loop_runtime", schedule=None) as dag:
+ group = create_loop(body, max_iterations=max_iterations, until=until)
+ ti = create_runtime_ti(task=dag.get_task(group.gate_task_id), map_index=map_index)
+ ti._ti_context_from_server.loop = LoopContext(
+ node_id=group.group_id,
+ index=index,
+ max_iterations=max_iterations,
+ terminal_task_id=group.terminal_task_id,
+ terminal_is_mapped=False,
+ )
+ return ti
+
+ return make
+
+
+def test_first_iteration_previous_does_not_read_xcom(loop_ti, mock_supervisor_comms):
+ ti = loop_ti(map_index=7)
+ loop = ti.get_template_context()["loop"]
+
+ assert loop.index == 0
+ assert loop.max_iterations == 3
+ assert loop.previous is None
+ assert ti.map_index == 7
+ mock_supervisor_comms.send.assert_not_called()
+
+
+@pytest.mark.parametrize("mapped", [False, True])
+def test_body_callable_receives_loop_context_after_unmapping(create_runtime_ti, mapped):
+ @task
+ def terminal(value, *, loop, ti):
+ return value, loop.index, ti.map_index
+
+ @task_group
+ def body():
+ if mapped:
+ terminal.expand(value=[5])
+ else:
+ terminal(5)
+
+ with DAG("loop_body_runtime", schedule=None) as dag:
+ group = create_loop(body, max_iterations=4)
+ operator = dag.get_task(group.terminal_task_id)
+ if mapped:
+ operator = operator.unmap({"op_kwargs": {"value": 5}})
+ ti = create_runtime_ti(task=operator, map_index=0 if mapped else -1)
+ ti._ti_context_from_server.loop = LoopContext(
+ node_id=group.group_id,
+ index=2,
+ max_iterations=4,
+ terminal_task_id=group.terminal_task_id,
+ terminal_is_mapped=mapped,
+ )
+
+ assert operator.execute(ti.get_template_context()) == (5, 2, 0 if mapped else -1)
+
+
+@pytest.mark.parametrize("value", [0, False, [], ""])
+def test_previous_and_current_result_preserve_falsey_values(loop_ti, mock_supervisor_comms, value):
+ ti = loop_ti(index=1)
+ mock_supervisor_comms.send.return_value = XComResult(key="return_value", value=value)
+ loop = ti.get_template_context()["loop"]
+
+ assert loop.previous == value
+ previous = mock_supervisor_comms.send.call_args.args[0]
+ assert isinstance(previous, GetXCom)
+ assert previous.task_id == "body.terminal"
+ assert previous.previous_iteration is True
+ assert loop.result == value
+ assert mock_supervisor_comms.send.call_args.args[0].previous_iteration is False
+
+
+@pytest.mark.parametrize(
+ ("index", "condition", "decision"),
+ [(0, None, "continue"), (2, None, "stop"), (0, False, "continue"), (0, True, "stop"), (2, True, "stop")],
+)
+def test_gate_publishes_successful_decision(loop_ti, mock_supervisor_comms, index, condition, decision):
+ def until(*, loop):
+ assert loop.index == index
+ return condition
+
+ ti = loop_ti(index=index, until=until if condition is not None else None)
+ context = ti.get_template_context()
+
+ ti.task.execute(context)
+
+ message = mock_supervisor_comms.send.call_args.args[0]
+ assert isinstance(message, SetXCom)
+ assert message.key == "_airflow_loop_decision"
+ assert message.value == decision
+
+
+@pytest.mark.parametrize(
+ ("raises", "error", "message"),
+ [(False, LoopMaxIterationsExceeded, "max_iterations"), (True, ValueError, "condition failed")],
+)
+def test_unsuccessful_gate_does_not_publish_decision(loop_ti, mock_supervisor_comms, raises, error, message):
+ def until(*, loop):
+ if raises:
+ raise ValueError("condition failed")
+ return False
+
+ ti = loop_ti(index=2, until=until)
+ with pytest.raises(error, match=message):
+ ti.task.execute(ti.get_template_context())
+
+ mock_supervisor_comms.send.assert_not_called()
+
+
+def test_gate_fails_without_retrying_when_the_condition_is_false_at_the_cap(loop_ti, mock_supervisor_comms):
+ ti = loop_ti(index=2, until=lambda: False)
+
+ with pytest.raises(LoopMaxIterationsExceeded, match="max_iterations=3") as exc_info:
+ ti.task.execute(ti.get_template_context())
+
+ assert isinstance(exc_info.value, AirflowFailException)
+ mock_supervisor_comms.send.assert_not_called()
+
+
+@pytest.mark.parametrize("result", [None, 0, 1, [], "yes"])
+def test_gate_fails_when_until_does_not_return_a_bool(loop_ti, mock_supervisor_comms, result):
+ ti = loop_ti(index=0, until=lambda: result)
+
+ with pytest.raises(AirflowFailException, match=f"got {type(result).__name__}"):
+ ti.task.execute(ti.get_template_context())
+
+ mock_supervisor_comms.send.assert_not_called()
+
+
+@pytest.mark.parametrize("previous_iteration", [False, True])
+@pytest.mark.parametrize("values", [[], [0], [0, False, ""]])
+def test_mapped_terminal_keeps_iteration_across_lazy_reads(
+ loop_ti, mock_supervisor_comms, previous_iteration, values
+):
+ ti = loop_ti(index=1)
+ ti._ti_context_from_server.loop.terminal_is_mapped = True
+
+ def respond(message):
+ assert message.previous_iteration is previous_iteration
+ if isinstance(message, GetXComCount):
+ return XComCountResponse(len=len(values))
+ if isinstance(message, GetXComSequenceItem):
+ return XComSequenceIndexResult(root=values[message.offset])
+ assert isinstance(message, GetXComSequenceSlice)
+ return XComSequenceSliceResult(root=values[slice(message.start, message.stop, message.step)])
+
+ mock_supervisor_comms.send.side_effect = respond
+ context = ti.get_template_context()["loop"]
+ result = context.previous if previous_iteration else context.result
+
+ assert isinstance(result, LazyXComSequence)
+ assert len(result) == len(values)
+ assert result[:] == values
+ assert result[::-1] == values[::-1]
+ if values:
+ assert result[-1] == values[-1]
+
+
+def test_loop_decision_bypasses_custom_backend(loop_ti, mock_supervisor_comms, mocker):
+ serialize = mocker.patch("airflow.sdk.execution_time.xcom.XCom.serialize_value", autospec=True)
+ ti = loop_ti()
+
+ ti.task.execute(ti.get_template_context())
+
+ serialize.assert_not_called()
+ assert mock_supervisor_comms.send.call_args.args[0].value == "continue"
+
+
+def test_gate_retry_clears_old_signal_without_custom_backend(loop_ti, mock_supervisor_comms, monkeypatch):
+ class CustomXCom(BaseXCom):
+ @classmethod
+ def purge(cls, xcom, *args):
+ raise AssertionError("Control metadata must not reach custom backend")
+
+ ti = loop_ti()
+ ti._ti_context_from_server.xcom_keys_to_clear = ["_airflow_loop_decision"]
+ monkeypatch.setattr(task_runner, "XCom", CustomXCom)
+ mock_supervisor_comms.send.side_effect = lambda message: (
+ XComResult(key="_airflow_loop_decision", value="stop") if isinstance(message, GetXCom) else None
+ )
+
+ state, _, error = task_runner.run(ti, context=ti.get_template_context(), log=ti.task.log)
+
+ assert error is None
+ assert state == TaskInstanceState.SUCCESS
+ decisions = [
+ call.args[0]
+ for call in mock_supervisor_comms.send.call_args_list
+ if call.args and isinstance(call.args[0], SetXCom)
+ ]
+ assert [decision.value for decision in decisions] == ["continue"]
+ deletes = [
+ call.args[0]
+ for call in mock_supervisor_comms.send.call_args_list
+ if call.args and isinstance(call.args[0], DeleteXCom)
+ ]
+ assert [delete.key for delete in deletes] == ["_airflow_loop_decision"]
diff --git a/task-sdk/tests/task_sdk/execution_time/test_request_handlers.py b/task-sdk/tests/task_sdk/execution_time/test_request_handlers.py
index 12a31532462..9678dd4e09a 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_request_handlers.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_request_handlers.py
@@ -16,9 +16,12 @@
# under the License.
from __future__ import annotations
+import json
from unittest.mock import MagicMock
+import httpx
import pytest
+from task_sdk import make_client
from airflow.sdk.api import client as sdk_client
from airflow.sdk.api.datamodels._generated import AssetStateStoreResponse
@@ -32,6 +35,10 @@ from airflow.sdk.execution_time.comms import (
ErrorResponse,
GetAssetStateStoreByName,
GetAssetStateStoreByUri,
+ GetXCom,
+ GetXComCount,
+ GetXComSequenceItem,
+ GetXComSequenceSlice,
SetAssetStateStoreByName,
SetAssetStateStoreByUri,
)
@@ -42,6 +49,10 @@ from airflow.sdk.execution_time.request_handlers import (
handle_delete_asset_state_store_by_uri,
handle_get_asset_state_store_by_name,
handle_get_asset_state_store_by_uri,
+ handle_get_xcom,
+ handle_get_xcom_count,
+ handle_get_xcom_sequence_item,
+ handle_get_xcom_sequence_slice,
handle_set_asset_state_store_by_name,
handle_set_asset_state_store_by_uri,
)
@@ -59,6 +70,39 @@ def client_ssl_cache():
sdk_client.Client._get_ssl_context_cached.cache_clear()
+@pytest.mark.parametrize(
+ ("message", "handler", "extra", "response", "field"),
+ [
+ (GetXCom, handle_get_xcom, {}, {"key": "key", "value": 0}, "previous_iteration"),
+ (GetXComCount, handle_get_xcom_count, {}, None, "previous_iteration"),
+ (GetXComSequenceItem, handle_get_xcom_sequence_item, {"offset": -1}, 0, "previous_iteration"),
+ (
+ GetXComSequenceSlice,
+ handle_get_xcom_sequence_slice,
+ {"start": None, "stop": None, "step": -1},
+ [],
+ "previous_iteration",
+ ),
+ ],
+)
+def test_loop_selectors_survive_comms_and_http_transport(
+ message, handler, extra, response, field, client_ssl_cache
+):
+ msg = message(dag_id="dag", run_id="run", task_id="task", key="key", **extra, **{field: True})
+ requests = []
+
+ def respond(request):
+ requests.append(request)
+ return httpx.Response(200, content=json.dumps(response), headers={"Content-Range": "map_indexes 1"})
+
+ handler(
+ make_client(transport=httpx.MockTransport(respond)),
+ message.model_validate_json(msg.model_dump_json()),
+ )
+
+ assert requests[0].url.params[field] == "true"
+
+
def test_get_asset_state_store_by_name_wraps_response_as_result(client):
client.asset_state_store.get.return_value = AssetStateStoreResponse(value="2026-01-01")
diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
index d109dd4ffd7..0f507a87f8e 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
@@ -2441,6 +2441,7 @@ REQUEST_TEST_CASES = [
client_mock=ClientMock(
method_path="xcoms.get",
args=("test_dag", "test_run", "test_task", "test_key", None, False),
+ kwargs={"previous_iteration": False},
response=XComResult(key="test_key", value="test_value"),
),
expected_body={"key": "test_key", "value": "test_value", "type": "XComResult"},
@@ -2453,6 +2454,7 @@ REQUEST_TEST_CASES = [
client_mock=ClientMock(
method_path="xcoms.get",
args=("test_dag", "test_run", "test_task", "test_key", 2, False),
+ kwargs={"previous_iteration": False},
response=XComResult(key="test_key", value="test_value"),
),
expected_body={"key": "test_key", "value": "test_value", "type": "XComResult"},
@@ -2463,6 +2465,7 @@ REQUEST_TEST_CASES = [
client_mock=ClientMock(
method_path="xcoms.get",
args=("test_dag", "test_run", "test_task", "test_key", None, False),
+ kwargs={"previous_iteration": False},
response=XComResult(key="test_key", value=None, type="XComResult"),
),
expected_body={"key": "test_key", "value": None, "type": "XComResult"},
@@ -2479,6 +2482,7 @@ REQUEST_TEST_CASES = [
client_mock=ClientMock(
method_path="xcoms.get",
args=("test_dag", "test_run", "test_task", "test_key", None, True),
+ kwargs={"previous_iteration": False},
response=XComResult(key="test_key", value=None, type="XComResult"),
),
expected_body={"key": "test_key", "value": None, "type": "XComResult"},
@@ -2501,7 +2505,10 @@ REQUEST_TEST_CASES = [
'{"key": "test_key", "value": {"key2": "value2"}}',
None,
),
- kwargs={"dag_result": False, "mapped_length": None},
+ kwargs={
+ "dag_result": False,
+ "mapped_length": None,
+ },
response=OKResponse(ok=True),
),
test_id="set_xcom",
@@ -2525,7 +2532,10 @@ REQUEST_TEST_CASES = [
'{"key": "test_key", "value": {"key2": "value2"}}',
2,
),
- kwargs={"dag_result": False, "mapped_length": None},
+ kwargs={
+ "dag_result": False,
+ "mapped_length": None,
+ },
response=OKResponse(ok=True),
),
test_id="set_xcom_with_map_index",
@@ -2550,7 +2560,10 @@ REQUEST_TEST_CASES = [
'{"key": "test_key", "value": {"key2": "value2"}}',
2,
),
- kwargs={"dag_result": False, "mapped_length": 3},
+ kwargs={
+ "dag_result": False,
+ "mapped_length": 3,
+ },
response=OKResponse(ok=True),
),
test_id="set_xcom_with_map_index_and_mapped_length",
@@ -2574,7 +2587,10 @@ REQUEST_TEST_CASES = [
'{"key": "test_key", "value": {"key2": "value2"}}',
None,
),
- kwargs={"dag_result": True, "mapped_length": None},
+ kwargs={
+ "dag_result": True,
+ "mapped_length": None,
+ },
response=OKResponse(ok=True),
),
test_id="set_xcom_with_dag_result",
@@ -3570,6 +3586,7 @@ REQUEST_TEST_CASES = [
client_mock=ClientMock(
method_path="xcoms.get_sequence_item",
args=("test_dag", "test_run", "test_task", "test_key", 0),
+ kwargs={"previous_iteration": False},
response=XComSequenceIndexResult(root="test_value"),
),
test_id="get_xcom_seq_item",
@@ -3586,6 +3603,7 @@ REQUEST_TEST_CASES = [
client_mock=ClientMock(
method_path="xcoms.get_sequence_item",
args=("test_dag", "test_run", "test_task", "test_key", 2),
+ kwargs={"previous_iteration": False},
response=ErrorResponse(error=ErrorType.XCOM_NOT_FOUND),
),
test_id="get_xcom_seq_item_not_found",
@@ -3605,6 +3623,7 @@ REQUEST_TEST_CASES = [
client_mock=ClientMock(
method_path="xcoms.get_sequence_slice",
args=("test_dag", "test_run", "test_task", "test_key", None, None, None, False),
+ kwargs={"previous_iteration": False},
response=XComSequenceSliceResult(root=["foo", "bar"]),
),
test_id="get_xcom_seq_slice",
@@ -3667,6 +3686,7 @@ REQUEST_TEST_CASES = [
client_mock=ClientMock(
method_path="xcoms.head",
args=("test_dag", "test_run", "test_task", "test_key"),
+ kwargs={"previous_iteration": False},
response=XComCountResponse(len=5),
),
test_id="get_xcom_count",
diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
index 66fa2411082..8d4ade1a340 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
@@ -64,6 +64,7 @@ from airflow.sdk import (
dag as dag_decorator,
get_current_context,
task as task_decorator,
+ task_group,
timezone,
)
from airflow.sdk._shared.observability.metrics.base_stats_logger import StatsLogger
@@ -83,6 +84,7 @@ from airflow.sdk.bases.operator import ExecutorSafeguard
from airflow.sdk.bases.xcom import BaseXCom
from airflow.sdk.coordinators._dag_importer import CoordinatorDagImporter
from airflow.sdk.coordinators._subprocess import SubprocessCoordinator
+from airflow.sdk.definitions._internal.loop import create_loop
from airflow.sdk.definitions._internal.types import NOTSET, SET_DURING_EXECUTION, is_arg_set
from airflow.sdk.definitions.asset import Asset, AssetAlias, AssetUniqueKey, AssetUriRef, Dataset, Model
from airflow.sdk.definitions.param import DagParam
@@ -177,6 +179,7 @@ from airflow.sdk.execution_time.context import (
VariableAccessor,
_wrap_external_ref,
)
+from airflow.sdk.execution_time.lazy_sequence import LazyXComSequence
from airflow.sdk.execution_time.task_runner import (
RuntimeTaskInstance,
TaskRunnerMarker,
@@ -3246,6 +3249,67 @@ class TestRuntimeTaskInstance:
),
)
+ @staticmethod
+ def _make_loop_dag():
+ @task_group
+ def body():
+ BaseOperator(task_id="first") >> BaseOperator(task_id="second")
+
+ with DAG("loop_dag", schedule=None) as dag:
+ create_loop(body, max_iterations=3) >> BaseOperator(task_id="after")
+ return dag
+
+ def test_xcom_pull_of_a_loop_task_from_outside_the_loop_returns_every_iteration(
+ self, create_runtime_ti, mock_supervisor_comms
+ ):
+ dag = self._make_loop_dag()
+ runtime_ti = create_runtime_ti(task=dag.get_task("after"))
+ mock_supervisor_comms.send.return_value = XComSequenceSliceResult(root=[1, 2, 3])
+
+ result = runtime_ti.xcom_pull(key="test_key", task_ids="body.first")
+
+ assert isinstance(result, LazyXComSequence)
+ mock_supervisor_comms.send.assert_not_called()
+ assert result[:] == [1, 2, 3]
+ sent = mock_supervisor_comms.send.call_args.args[0]
+ assert (sent.task_id, sent.key, sent.previous_iteration) == ("body.first", "test_key", False)
+
+ @pytest.mark.parametrize(
+ ("task_id", "pull_kwargs"),
+ [
+ pytest.param("body.second", {"task_ids": "body.first"}, id="caller inside the loop"),
+ pytest.param("after", {"task_ids": "body.first", "include_prior_dates": True}, id="prior dates"),
+ pytest.param("after", {"task_ids": "body.first", "run_id": "other_run"}, id="other run"),
+ pytest.param("after", {"task_ids": "body.first", "map_indexes": -1}, id="explicit map index"),
+ pytest.param("after", {"task_ids": ["body.first"]}, id="several task ids"),
+ pytest.param("after", {"task_ids": "removed"}, id="task missing from the dag"),
+ pytest.param("after", {"task_ids": "after"}, id="task outside any loop"),
+ ],
+ )
+ def test_xcom_pull_keeps_pulling_values_unless_an_outside_caller_reads_a_loop_task(
+ self, create_runtime_ti, mock_supervisor_comms, task_id, pull_kwargs
+ ):
+ dag = self._make_loop_dag()
+ runtime_ti = create_runtime_ti(task=dag.get_task(task_id))
+ mock_supervisor_comms.send.return_value = XComSequenceSliceResult(root=[1, 2, 3])
+
+ with patch.object(XCom, "get_one", return_value=7):
+ result = runtime_ti.xcom_pull(key="test_key", **pull_kwargs)
+
+ assert not isinstance(result, LazyXComSequence)
+
+ @pytest.mark.asyncio
+ async def test_axcom_pull_of_a_loop_task_from_outside_the_loop_returns_a_list_for_one_iteration(
+ self, create_runtime_ti, mock_supervisor_comms
+ ):
+ dag = self._make_loop_dag()
+ runtime_ti = create_runtime_ti(task=dag.get_task("after"))
+ mock_supervisor_comms.asend.return_value = XComSequenceSliceResult(root=["only"])
+
+ result = await runtime_ti.axcom_pull(key="test_key", task_ids="body.first")
+
+ assert result == ["only"]
+
def test_get_param_from_context(
self, mocked_parse, make_ti_context, mock_supervisor_comms, create_runtime_ti
):
diff --git a/ts-sdk/schema/supervisor-schema.json b/ts-sdk/schema/supervisor-schema.json
index b0d5797b788..7e09a736de7 100644
--- a/ts-sdk/schema/supervisor-schema.json
+++ b/ts-sdk/schema/supervisor-schema.json
@@ -2836,6 +2836,11 @@
},
"GetXCom": {
"properties": {
+ "previous_iteration": {
+ "default": false,
+ "title": "Previous Iteration",
+ "type": "boolean"
+ },
"key": {
"title": "Key",
"type": "string"
@@ -2888,6 +2893,11 @@
"GetXComCount": {
"description": "Get the number of (mapped) XCom values available.",
"properties": {
+ "previous_iteration": {
+ "default": false,
+ "title": "Previous Iteration",
+ "type": "boolean"
+ },
"key": {
"title": "Key",
"type": "string"
@@ -2922,6 +2932,11 @@
},
"GetXComSequenceItem": {
"properties": {
+ "previous_iteration": {
+ "default": false,
+ "title": "Previous Iteration",
+ "type": "boolean"
+ },
"key": {
"title": "Key",
"type": "string"
@@ -2961,6 +2976,11 @@
},
"GetXComSequenceSlice": {
"properties": {
+ "previous_iteration": {
+ "default": false,
+ "title": "Previous Iteration",
+ "type": "boolean"
+ },
"key": {
"title": "Key",
"type": "string"
@@ -3222,6 +3242,42 @@
"title": "LazyDeserializedDAG",
"type": "object"
},
+ "LoopContext": {
+ "description": "Pinned loop definition and enclosing iteration for a task execution.",
+ "properties": {
+ "node_id": {
+ "title": "Node Id",
+ "type": "string"
+ },
+ "index": {
+ "minimum": 0,
+ "title": "Index",
+ "type": "integer"
+ },
+ "max_iterations": {
+ "exclusiveMinimum": 0,
+ "title": "Max Iterations",
+ "type": "integer"
+ },
+ "terminal_task_id": {
+ "title": "Terminal Task Id",
+ "type": "string"
+ },
+ "terminal_is_mapped": {
+ "title": "Terminal Is Mapped",
+ "type": "boolean"
+ }
+ },
+ "required": [
+ "node_id",
+ "index",
+ "max_iterations",
+ "terminal_task_id",
+ "terminal_is_mapped"
+ ],
+ "title": "LoopContext",
+ "type": "object"
+ },
"MaskSecret": {
"description": "Add a new value to be redacted in task logs.",
"properties": {
@@ -5154,6 +5210,17 @@
"dag_run": {
"$ref": "#/$defs/DagRun"
},
+ "loop": {
+ "anyOf": [
+ {
+ "$ref": "#/$defs/LoopContext"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null
+ },
"task_reschedule_count": {
"default": 0,
"title": "Task Reschedule Count",
diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts
index 1d785b4d476..31196f868a5 100644
--- a/ts-sdk/src/generated/supervisor.ts
+++ b/ts-sdk/src/generated/supervisor.ts
@@ -231,6 +231,11 @@ export type TaskInstanceState =
| "skipped"
| "deferred"
| "awaiting_input";
+export type NodeId = string;
+export type Index = number;
+export type MaxIterations = number;
+export type TerminalTaskId = string;
+export type TerminalIsMapped = boolean;
export type TaskRescheduleCount = number;
export type MaxTries = number;
export type Key = string;
@@ -473,6 +478,7 @@ export type Prefix = string | null;
export type Limit2 = number;
export type Offset = number;
export type Type49 = "GetVariableKeys";
+export type PreviousIteration = boolean;
export type Key10 = string;
export type DagId16 = string;
export type RunId9 = string;
@@ -480,17 +486,20 @@ export type TaskId4 = string;
export type MapIndex5 = number | null;
export type IncludePriorDates = boolean;
export type Type50 = "GetXCom";
+export type PreviousIteration1 = boolean;
export type Key11 = string;
export type DagId17 = string;
export type RunId10 = string;
export type TaskId5 = string;
export type Type51 = "GetXComCount";
+export type PreviousIteration2 = boolean;
export type Key12 = string;
export type DagId18 = string;
export type RunId11 = string;
export type TaskId6 = string;
export type Offset1 = number;
export type Type52 = "GetXComSequenceItem";
+export type PreviousIteration3 = boolean;
export type Key13 = string;
export type DagId19 = string;
export type RunId12 = string;
@@ -1038,6 +1047,7 @@ export interface TaskCallbackRequest {
*/
export interface TIRunContext {
dag_run: DagRun;
+ loop?: LoopContext | null;
task_reschedule_count?: TaskRescheduleCount;
max_tries: MaxTries;
variables?: Variables;
@@ -1050,6 +1060,19 @@ export interface TIRunContext {
arg_bindings?: ArgBindings;
multi_team?: MultiTeam;
}
+/**
+ * Pinned loop definition and enclosing iteration for a task execution.
+ *
+ * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema
+ * via the `definition` "LoopContext".
+ */
+export interface LoopContext {
+ node_id: NodeId;
+ index: Index;
+ max_iterations: MaxIterations;
+ terminal_task_id: TerminalTaskId;
+ terminal_is_mapped: TerminalIsMapped;
+}
/**
* Variable schema for responses with fields that are needed for Runtime.
*
@@ -1514,6 +1537,7 @@ export interface GetVariableKeys {
* via the `definition` "GetXCom".
*/
export interface GetXCom {
+ previous_iteration?: PreviousIteration;
key: Key10;
dag_id: DagId16;
run_id: RunId9;
@@ -1529,6 +1553,7 @@ export interface GetXCom {
* via the `definition` "GetXComCount".
*/
export interface GetXComCount {
+ previous_iteration?: PreviousIteration1;
key: Key11;
dag_id: DagId17;
run_id: RunId10;
@@ -1540,6 +1565,7 @@ export interface GetXComCount {
* via the `definition` "GetXComSequenceItem".
*/
export interface GetXComSequenceItem {
+ previous_iteration?: PreviousIteration2;
key: Key12;
dag_id: DagId18;
run_id: RunId11;
@@ -1552,6 +1578,7 @@ export interface GetXComSequenceItem {
* via the `definition` "GetXComSequenceSlice".
*/
export interface GetXComSequenceSlice {
+ previous_iteration?: PreviousIteration3;
key: Key13;
dag_id: DagId19;
run_id: RunId12;