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/06 14:58:19 UTC
(airflow) 01/01: Give every mapped expansion a region and store its index as region_index
This is an automated email from the ASF dual-hosted git repository.
ashb pushed a commit to branch task-loops-stack-5
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 8c43b9103f237678bdeece41aa1d31d353763ae2
Author: Ash Berlin-Taylor <as...@apache.org>
AuthorDate: Sun Oct 4 08:49:24 2026 +0100
Give every mapped expansion a region and store its index as region_index
Mapped tasks were the one kind of dynamic work still outside the region
model: their instances sat in the sentinel region and were told apart only by
map_index. A loop has to be able to contain an expansion, one per pass, and that
only works if the expansion owns a region of its own.
The coordinate column is renamed in place so large deployments pay for metadata
changes, not a table rewrite. No table is rebuilt because SQLite's rebuild would
cascade-delete the rows that reference task_instance.id. legacy_task_data_owner,
xcom_v1 and rtif_v1 keep map_index because they record the coordinates their
rows had before regions existed. The ORM keeps map_index as an alias: Jinja
log-filename templates and provider code still read ti.map_index from ORM rows,
and removing it needs its own audit.
Mapping keeps its existing behaviour. A task's region is born with its
placeholder, re-expansion grows or shrinks it in place, and a map never forks.
Promoting the placeholder to index 0 stays in place too. Attempts retired before
the promotion stay at index -1: they own their XCom, rendered fields and log
files by UUID, and moving them could collide with retired attempts already
at index 0. A pre-region placeholder never ran, so it has no history to keep;
it moves into a freshly minted region and is promoted like any other, instead
of being retired as a phantom failed try. Pre-region expansions stay in the
sentinel region and keep revising in place, so the upgrade does not rewrite
existing rows.
Callers that address a task instance by public coordinates now find a task's own
region as well as the sentinel. The DagRun wait endpoint reads current producers
only and rejects loop results, because which pass's result to return would be
ambiguous.
---
airflow-core/docs/migrations-ref.rst | 4 +-
.../airflow/api_fastapi/common/parameters/misc.py | 2 +-
.../api_fastapi/common/parameters/task_instance.py | 6 +-
.../core_api/openapi/v2-rest-api-generated.yaml | 6 +
.../api_fastapi/core_api/routes/public/dag_run.py | 7 +-
.../core_api/routes/public/task_state_store.py | 32 +-
.../api_fastapi/core_api/routes/ui/gantt.py | 2 +-
.../core_api/services/public/dag_run.py | 109 ++++--
.../execution_api/routes/asset_state_store.py | 2 +-
.../execution_api/routes/task_instances.py | 14 +-
.../airflow/cli/commands/state_store_command.py | 6 +-
.../src/airflow/cli/commands/task_command.py | 86 ++++-
.../src/airflow/jobs/scheduler_job_runner.py | 51 +--
.../versions/0144_3_4_0_rename_stored_map_index.py | 61 +++
airflow-core/src/airflow/models/asset.py | 2 +-
airflow-core/src/airflow/models/dagbag.py | 21 ++
airflow-core/src/airflow/models/dagrun.py | 196 ++++++----
airflow-core/src/airflow/models/dynamic_region.py | 62 +++-
airflow-core/src/airflow/models/expandinput.py | 31 +-
airflow-core/src/airflow/models/log.py | 2 +-
.../src/airflow/models/renderedtifields.py | 2 +-
airflow-core/src/airflow/models/serialized_dag.py | 4 +-
.../src/airflow/models/task_coordinates.py | 64 +++-
.../src/airflow/models/task_state_store.py | 11 +-
airflow-core/src/airflow/models/taskinstance.py | 184 +++++++---
airflow-core/src/airflow/models/xcom.py | 44 +--
.../serialization/definitions/baseoperator.py | 2 +
.../src/airflow/serialization/definitions/dag.py | 98 +----
.../serialization/definitions/mappedoperator.py | 51 ++-
.../airflow/serialization/definitions/xcom_arg.py | 69 ++--
airflow-core/src/airflow/serialization/schema.json | 1 +
airflow-core/src/airflow/state/metastore.py | 22 +-
airflow-core/src/airflow/ti_deps/dep_context.py | 76 +++-
.../airflow/ti_deps/deps/mapped_task_expanded.py | 2 +-
.../ti_deps/deps/mapped_task_upstream_dep.py | 37 +-
.../ti_deps/deps/not_previously_skipped_dep.py | 8 +-
.../airflow/ti_deps/deps/ready_to_reschedule.py | 9 +-
.../src/airflow/ti_deps/deps/trigger_rule_dep.py | 98 +++--
.../ui/openapi-gen/requests/services.gen.ts | 2 +
.../airflow/ui/openapi-gen/requests/types.gen.ts | 4 +
airflow-core/src/airflow/utils/db.py | 2 +-
.../src/airflow/utils/log/task_log_address.py | 27 +-
.../core_api/routes/public/test_dag_run.py | 176 ++++++++-
.../core_api/routes/public/test_task_instances.py | 2 +-
.../routes/public/test_task_state_store.py | 30 +-
.../versions/head/test_task_instances.py | 2 +-
.../versions/head/test_task_state_store.py | 4 +-
.../execution_api/versions/head/test_xcoms.py | 28 +-
.../unit/cli/commands/test_state_store_command.py | 10 +-
.../tests/unit/cli/commands/test_task_command.py | 244 +++++++++++-
airflow-core/tests/unit/jobs/test_scheduler_job.py | 38 +-
.../test_0143_add_dynamic_region_storage.py | 30 +-
...age.py => test_0144_rename_stored_map_index.py} | 106 +++---
airflow-core/tests/unit/models/test_dag.py | 59 +++
airflow-core/tests/unit/models/test_dagrun.py | 115 +++++-
.../tests/unit/models/test_dynamic_region.py | 66 +++-
.../tests/unit/models/test_mappedoperator.py | 408 ++++++++++++++++++++-
airflow-core/tests/unit/models/test_task_data.py | 22 +-
.../tests/unit/models/test_task_state_store.py | 18 +-
.../tests/unit/models/test_taskinstance.py | 114 +++---
.../unit/serialization/test_dag_serialization.py | 26 ++
airflow-core/tests/unit/state/test_metastore.py | 10 +-
.../ti_deps/deps/test_mapped_task_upstream_dep.py | 31 +-
.../ti_deps/deps/test_ready_to_reschedule_dep.py | 26 ++
.../unit/ti_deps/deps/test_trigger_rule_dep.py | 212 ++++++++++-
.../tests/unit/utils/log/test_task_log_address.py | 24 +-
airflow-core/tests/unit/utils/test_db_cleanup.py | 8 +-
.../tests_common/test_utils/attempt_ownership.py | 10 +-
go-sdk/airflow/spec.gen.go | 3 +
go-sdk/schema/dag-schema.json | 1 +
java-sdk/sdk/schema/dag-schema.json | 1 +
shared/state/src/airflow_shared/state/__init__.py | 41 ++-
shared/state/tests/state/test_state.py | 13 +
task-sdk/src/airflow/sdk/bases/operator.py | 1 +
.../src/airflow/sdk/definitions/mappedoperator.py | 1 -
ts-sdk/schema/dag-schema.json | 1 +
ts-sdk/src/generated/dag-schema-fields.ts | 3 +
77 files changed, 2721 insertions(+), 682 deletions(-)
diff --git a/airflow-core/docs/migrations-ref.rst b/airflow-core/docs/migrations-ref.rst
index aa991bfc2b7..df781acc85d 100644
--- a/airflow-core/docs/migrations-ref.rst
+++ b/airflow-core/docs/migrations-ref.rst
@@ -39,7 +39,9 @@ Here's the list of all the Database Migrations that are executed via when you ru
+-------------------------+------------------+-------------------+--------------------------------------------------------------+
| Revision ID | Revises ID | Airflow Version | Description |
+=========================+==================+===================+==============================================================+
-| ``54a27b6f9d01`` (head) | ``e7c2a91bd540`` | ``3.4.0`` | Add dynamic region storage. |
+| ``7f8c9a2d410e`` (head) | ``54a27b6f9d01`` | ``3.4.0`` | Rename stored map indexes to region indexes. |
++-------------------------+------------------+-------------------+--------------------------------------------------------------+
+| ``54a27b6f9d01`` | ``e7c2a91bd540`` | ``3.4.0`` | Add dynamic region storage. |
+-------------------------+------------------+-------------------+--------------------------------------------------------------+
| ``e7c2a91bd540`` | ``90e4d18ccadf`` | ``3.4.0`` | Unify task attempt ownership without rewriting legacy XCom |
| | | | data. |
diff --git a/airflow-core/src/airflow/api_fastapi/common/parameters/misc.py b/airflow-core/src/airflow/api_fastapi/common/parameters/misc.py
index 6a5b2a44bdc..01770746d03 100644
--- a/airflow-core/src/airflow/api_fastapi/common/parameters/misc.py
+++ b/airflow-core/src/airflow/api_fastapi/common/parameters/misc.py
@@ -245,7 +245,7 @@ QueryHITLDetailMapIndexFilter = Annotated[
FilterParam[int | None],
Depends(
filter_param_factory(
- TaskInstance.map_index,
+ TaskInstance.region_index,
int | None,
filter_name="map_index",
)
diff --git a/airflow-core/src/airflow/api_fastapi/common/parameters/task_instance.py b/airflow-core/src/airflow/api_fastapi/common/parameters/task_instance.py
index d1fdb868656..2a0dc93c2e0 100644
--- a/airflow-core/src/airflow/api_fastapi/common/parameters/task_instance.py
+++ b/airflow-core/src/airflow/api_fastapi/common/parameters/task_instance.py
@@ -237,7 +237,11 @@ QueryTIMapIndexFilter = Annotated[
FilterParam[list[int]],
Depends(
filter_param_factory(
- TaskInstance.map_index, list[int], FilterOptionEnum.ANY_EQUAL, default_factory=list
+ TaskInstance.region_index,
+ list[int],
+ FilterOptionEnum.ANY_EQUAL,
+ filter_name="map_index",
+ default_factory=list,
)
),
]
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
index 1952f25f707..6f2c1b0d089 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
+++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
@@ -3083,6 +3083,12 @@ paths:
schema:
$ref: '#/components/schemas/HTTPExceptionResponse'
description: Not Found
+ '409':
+ content:
+ application/json:
+ schema:
+ $ref: '#/components/schemas/HTTPExceptionResponse'
+ description: Conflict
'422':
description: Validation Error
content:
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py
index 1ce68758710..9eef9661799 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py
@@ -863,7 +863,7 @@ def trigger_dag_run(
summary="Experimental: Wait for a dag run to complete, and return task results if requested.",
description="🚧 This is an experimental endpoint and may change or be removed without notice.Successful response are streamed as newline-delimited JSON (NDJSON). Each line is a JSON object representing the Dag run state.",
responses={
- **create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]),
+ **create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND, status.HTTP_409_CONFLICT]),
status.HTTP_200_OK: {
"description": "Successful Response",
"content": {
@@ -886,6 +886,7 @@ def trigger_dag_run(
def wait_dag_run_until_finished(
dag_id: str,
dag_run_id: str,
+ dag_bag: DagBagDep,
session: SessionDep,
user: GetUserDep,
interval: Annotated[float, Query(gt=0.0, description="Seconds to wait between dag run state checks")],
@@ -917,7 +918,8 @@ def wait_dag_run_until_finished(
"User is not authorized to read XCom data for this Dag",
)
result_task_ids = [] # Explicitly not returning any XCom results.
- if not session.scalar(select(1).where(DagRun.dag_id == dag_id, DagRun.run_id == dag_run_id)):
+ dag_run = session.scalar(select(DagRun).where(DagRun.dag_id == dag_id, DagRun.run_id == dag_run_id))
+ if dag_run is None:
raise HTTPException(
status.HTTP_404_NOT_FOUND,
f"The DagRun with dag_id: `{dag_id}` and run_id: `{dag_run_id}` was not found",
@@ -928,6 +930,7 @@ def wait_dag_run_until_finished(
interval=interval,
result_task_ids=result_task_ids,
)
+ waiter.validate_result_selection(dag_run, dag_bag, session=session)
return StreamingResponse(waiter.wait())
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_state_store.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_state_store.py
index 8c0e5be5ef3..c92e77ba16c 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_state_store.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_state_store.py
@@ -42,6 +42,8 @@ from airflow.models.taskinstance import TaskInstance as TI
from airflow.state.metastore import _get_db_backend
if TYPE_CHECKING:
+ from uuid import UUID
+
from sqlalchemy.orm import Session
task_state_store_router = AirflowRouter(
@@ -56,16 +58,17 @@ def _require_task_instance(
task_id: str,
map_index: int | None,
session: Session,
-) -> None:
- """Raise 404 unless the task instance exists. ``map_index=None`` matches any map index."""
- statement = select(TI.task_id).where(
+) -> UUID:
+ """Return the region of the addressed task instance, raising 404 when it does not exist."""
+ statement = select(TI.region_id).where(
TI.dag_id == dag_id,
TI.run_id == dag_run_id,
TI.task_id == task_id,
)
if map_index is not None:
- statement = statement.where(TI.map_index == map_index)
- if session.scalar(statement.limit(1)) is None:
+ statement = statement.where(TI.region_index == map_index)
+ region_id = session.scalar(statement.order_by(TI.region_id).limit(1))
+ if region_id is None:
addressed_by = "all_map_indices=True" if map_index is None else f"map_index={map_index}"
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
@@ -74,6 +77,7 @@ def _require_task_instance(
f"task_id={task_id!r}, {addressed_by}"
),
)
+ return region_id
def _resolve_scope(
@@ -91,8 +95,10 @@ TaskScopeDep = Annotated[TaskScope, Depends(_resolve_scope)]
def _validate_scope(scope: TaskScopeDep, session: SessionDep) -> TaskScope:
"""Resolve the scope, 404ing when the task instance it addresses does not exist."""
- _require_task_instance(scope.dag_id, scope.run_id, scope.task_id, scope.map_index, session)
- return scope
+ region_id = _require_task_instance(scope.dag_id, scope.run_id, scope.task_id, scope.map_index, session)
+ return TaskScope(
+ scope.dag_id, scope.run_id, scope.task_id, region_index=scope.region_index, region_id=region_id
+ )
ValidatedTaskScopeDep = Annotated[TaskScope, Depends(_validate_scope)]
@@ -109,10 +115,12 @@ def _validate_clear_scope(
``all_map_indices`` addresses the task across every index, so it is validated against any
instance -- an expanded mapped task has no ``map_index=-1`` instance to check.
"""
- _require_task_instance(
+ region_id = _require_task_instance(
scope.dag_id, scope.run_id, scope.task_id, None if all_map_indices else scope.map_index, session
)
- return scope
+ return TaskScope(
+ scope.dag_id, scope.run_id, scope.task_id, region_index=scope.region_index, region_id=region_id
+ )
ValidatedClearTaskScopeDep = Annotated[TaskScope, Depends(_validate_clear_scope)]
@@ -161,7 +169,7 @@ def list_task_state_store(
TaskStateStoreModel.dag_id == scope.dag_id,
TaskStateStoreModel.run_id == scope.run_id,
TaskStateStoreModel.task_id == scope.task_id,
- TaskStateStoreModel.map_index == scope.map_index,
+ TaskStateStoreModel.region_index == scope.region_index,
)
.order_by(TaskStateStoreModel.key.asc())
)
@@ -204,7 +212,7 @@ def get_task_state_store(
TaskStateStoreModel.dag_id == scope.dag_id,
TaskStateStoreModel.run_id == scope.run_id,
TaskStateStoreModel.task_id == scope.task_id,
- TaskStateStoreModel.map_index == scope.map_index,
+ TaskStateStoreModel.region_index == scope.region_index,
TaskStateStoreModel.key == key,
)
).one_or_none()
@@ -256,7 +264,7 @@ def patch_task_state_store(
TaskStateStoreModel.dag_id == scope.dag_id,
TaskStateStoreModel.run_id == scope.run_id,
TaskStateStoreModel.task_id == scope.task_id,
- TaskStateStoreModel.map_index == scope.map_index,
+ TaskStateStoreModel.region_index == scope.region_index,
TaskStateStoreModel.key == key,
)
).one_or_none()
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/gantt.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/gantt.py
index ba4778b4a9c..660bb8d3568 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/gantt.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/gantt.py
@@ -73,7 +73,7 @@ def get_gantt_data(
).where(
TaskInstance.dag_id == dag_id,
TaskInstance.run_id == run_id,
- TaskInstance.map_index == -1,
+ TaskInstance.region_index == -1,
or_(TaskInstance.state != TaskInstanceState.UP_FOR_RETRY, TaskInstance.state.is_(None)),
)
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/services/public/dag_run.py b/airflow-core/src/airflow/api_fastapi/core_api/services/public/dag_run.py
index d91fe06a077..9ef6e15f3a0 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/services/public/dag_run.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/services/public/dag_run.py
@@ -20,7 +20,6 @@ from __future__ import annotations
import asyncio
import itertools
import json
-import operator
from typing import TYPE_CHECKING, Any
import attrs
@@ -61,13 +60,15 @@ from airflow.api_fastapi.core_api.services.public.task_instances import _emit_st
from airflow.listeners.listener import get_listener_manager
from airflow.models.dagrun import DagRun, clear_partition_runs
from airflow.models.renderedtifields import load_legacy_rendered_fields
+from airflow.models.task_coordinates import enclosing_loop, public_map_index_expression
from airflow.models.taskinstance import TaskInstance
from airflow.models.xcom import XCOM_RETURN_KEY, XComModel, xcom_entity
from airflow.utils.session import create_session_async
from airflow.utils.state import State, TaskInstanceState
if TYPE_CHECKING:
- from collections.abc import AsyncGenerator, Iterator
+ from collections.abc import AsyncGenerator
+ from uuid import UUID
from airflow.serialization.definitions.dag import SerializedDAG
@@ -256,44 +257,90 @@ class DagRunWaiter:
interval: float
result_task_ids: list[str] | None
+ def validate_result_selection(self, dag_run: DagRun, dag_bag: DagBagDep, *, session: Session) -> None:
+ if self.result_task_ids == []:
+ return
+ query = select(TaskInstance.task_id, TaskInstance.dag_version_id).where(
+ TaskInstance.dag_id == self.dag_id, TaskInstance.run_id == self.run_id
+ )
+ if self.result_task_ids is not None:
+ query = query.where(TaskInstance.task_id.in_(self.result_task_ids))
+ versions: dict[UUID | None, set[str]] = {}
+ for task_id, version in session.execute(query.distinct()):
+ versions.setdefault(version or dag_run.created_dag_version_id, set()).add(task_id)
+ current_task_ids = {task_id for task_ids in versions.values() for task_id in task_ids}
+ created_dag = (
+ dag_bag.get_dag(dag_run.created_dag_version_id, session=session)
+ if dag_run.created_dag_version_id is not None
+ else None
+ )
+ if created_dag is not None:
+ requested = (
+ set(self.result_task_ids) if self.result_task_ids is not None else set(created_dag.task_ids)
+ )
+ versions.setdefault(dag_run.created_dag_version_id, set()).update(requested - current_task_ids)
+ for version, task_ids in versions.items():
+ if version is None:
+ continue
+ dag = (
+ created_dag
+ if version == dag_run.created_dag_version_id
+ else dag_bag.get_dag(version, session=session)
+ )
+ if dag is None:
+ continue
+ for task_id in task_ids:
+ if not dag.has_task(task_id):
+ continue
+ task = dag.get_task(task_id)
+ if self.result_task_ids is None and not task.returns_dag_result:
+ continue
+ if enclosing_loop(task) is not None:
+ raise HTTPException(
+ status.HTTP_409_CONFLICT,
+ f"Cannot collect loop task {task_id!r} as a DagRun result; "
+ "publish the result through a task outside the loop.",
+ )
+
async def _get_dag_run(self) -> DagRun:
async with create_session_async() as session:
return await session.scalar(select(DagRun).filter_by(dag_id=self.dag_id, run_id=self.run_id))
async def _serialize_xcoms(self) -> dict[str, Any]:
+ xcom_query = XComModel.get_many(
+ run_id=self.run_id,
+ key=XCOM_RETURN_KEY,
+ dag_ids=self.dag_id,
+ task_ids=self.result_task_ids,
+ region_id=None,
+ )
+ entity = xcom_entity(xcom_query)
if self.result_task_ids is None: # Return dag-author-specified results.
- xcom_read = XComModel.get_many(
- run_id=self.run_id,
- key=XCOM_RETURN_KEY,
- dag_ids=self.dag_id,
- )
- entity = xcom_entity(xcom_read)
- xcom_query = xcom_read.where(entity.dag_result.is_(True))
- else: # Explicitly API user-specified results.
- xcom_read = XComModel.get_many(
- run_id=self.run_id,
- key=XCOM_RETURN_KEY,
- task_ids=self.result_task_ids,
- dag_ids=self.dag_id,
- )
- xcom_query = xcom_read
- entity = xcom_entity(xcom_read)
+ xcom_query = xcom_query.where(entity.dag_result.is_(True))
# XComModel.get_many() orders XCom by timestamp. Reset this to make
# mapped task results stable since execution order is not guaranteed.
- xcom_query = xcom_query.order_by(None).order_by(entity.task_id, entity.map_index)
+ xcom_query = (
+ xcom_query.add_columns(public_map_index_expression(entity))
+ .order_by(None)
+ .order_by(entity.task_id, entity.region_index)
+ )
async with create_session_async() as session:
- xcom_results = (await session.scalars(xcom_query)).all()
-
- def _group_xcoms(g: Iterator[XComModel | tuple[XComModel]]) -> Any:
- entries = [row[0] if isinstance(row, tuple) else row for row in g]
- if len(entries) == 1 and entries[0].map_index < 0: # Unpack non-mapped task xcom.
- return entries[0].value
- return [entry.value for entry in entries] # Task is mapped; return all xcoms in a list.
-
- return {
- task_id: _group_xcoms(g)
- for task_id, g in itertools.groupby(xcom_results, key=operator.attrgetter("task_id"))
- }
+ xcom_results = (await session.execute(xcom_query)).all()
+
+ results = {}
+ for task_id, group in itertools.groupby(xcom_results, key=lambda row: row[0].task_id):
+ entries = list(group)
+ if len({entry.region_id for entry, _ in entries}) > 1 or (
+ len(entries) > 1 and any(map_index < 0 for _, map_index in entries)
+ ):
+ raise ValueError(
+ f"Ambiguous current producers for task {task_id!r} in DAG run {self.run_id!r}"
+ )
+ if len(entries) == 1 and entries[0][1] < 0:
+ results[task_id] = entries[0][0].value
+ else:
+ results[task_id] = [entry.value for entry, _ in entries]
+ return results
async def _serialize_response(self, dag_run: DagRun) -> str:
resp = {"state": dag_run.state}
diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_state_store.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_state_store.py
index 4424e4ad3f1..0d701a5b48f 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_state_store.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_state_store.py
@@ -60,7 +60,7 @@ def _fetch_ti_writer_fields(token: TIToken, session: SessionDep) -> _TIWriterFie
TaskInstance.dag_id,
TaskInstance.run_id,
TaskInstance.task_id,
- TaskInstance.map_index,
+ TaskInstance.region_index.label("map_index"),
).where(TaskInstance.id == token.id)
).one_or_none()
if row is None:
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 ce09f4b4bb4..89088b245a5 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
@@ -171,7 +171,7 @@ def ti_run(
TI.dag_id,
TI.run_id,
TI.task_id,
- TI.map_index,
+ TI.region_index,
TI.region_id,
TI.try_number,
TI.max_tries,
@@ -259,7 +259,7 @@ def ti_run(
task_id=ti.task_id,
dag_id=ti.dag_id,
run_id=ti.run_id,
- map_index=ti.map_index,
+ map_index=ti.region_index,
try_number=ti.try_number,
logical_date=ti.logical_date,
owner=ti.owners,
@@ -438,7 +438,7 @@ def ti_update_state(
TI.dag_id,
TI.task_id,
TI.run_id,
- TI.map_index,
+ TI.region_index,
TI.region_id,
TI.hostname,
DR.logical_date,
@@ -460,7 +460,7 @@ def ti_update_state(
dag_id,
task_id,
run_id,
- map_index,
+ region_index,
region_id,
hostname,
logical_date,
@@ -576,7 +576,7 @@ def ti_update_state(
task_id=task_id,
dag_id=dag_id,
run_id=run_id,
- map_index=map_index,
+ map_index=region_index,
try_number=try_number,
logical_date=logical_date,
owner=owners,
@@ -596,7 +596,7 @@ def ti_update_state(
dag_id=dag_id,
run_id=run_id,
task_id=task_id,
- map_index=map_index if map_index is not None else -1,
+ region_index=region_index if region_index is not None else -1,
region_id=region_id,
)
try:
@@ -606,7 +606,7 @@ def ti_update_state(
dag_id=dag_id,
run_id=run_id,
task_id=task_id,
- map_index=map_index,
+ region_index=region_index,
)
except Exception:
log.warning(
diff --git a/airflow-core/src/airflow/cli/commands/state_store_command.py b/airflow-core/src/airflow/cli/commands/state_store_command.py
index c5139522d61..aca074102b8 100644
--- a/airflow-core/src/airflow/cli/commands/state_store_command.py
+++ b/airflow-core/src/airflow/cli/commands/state_store_command.py
@@ -52,8 +52,10 @@ def clean_state_store(args) -> None:
print("Nothing to delete.")
return
print(f"Would delete {len(expired)} task state store row(s):\n")
- for dag_id, run_id, task_id, map_index, key in expired:
- print(f" Dag {dag_id!r}, run {run_id!r}, task {task_id!r}, map_index {map_index!r}, key {key!r}")
+ for dag_id, run_id, task_id, region_index, key in expired:
+ print(
+ f" Dag {dag_id!r}, run {run_id!r}, task {task_id!r}, region_index {region_index!r}, key {key!r}"
+ )
return
log.info("Running task state store cleanup")
diff --git a/airflow-core/src/airflow/cli/commands/task_command.py b/airflow-core/src/airflow/cli/commands/task_command.py
index 32aa98d083a..d330cbcdfb2 100644
--- a/airflow-core/src/airflow/cli/commands/task_command.py
+++ b/airflow-core/src/airflow/cli/commands/task_command.py
@@ -27,6 +27,9 @@ import textwrap
from contextlib import redirect_stdout
from typing import TYPE_CHECKING, Protocol, cast
+import uuid6
+from sqlalchemy import select
+
from airflow import settings
from airflow._shared.timezones import timezone
from airflow.cli.simple_table import AirflowConsole
@@ -34,10 +37,16 @@ from airflow.cli.utils import deprecated_for_airflowctl, fetch_dag_run_from_run_
from airflow.exceptions import AirflowConfigException, DagRunNotFound, NotMapped, TaskInstanceNotFound
from airflow.models import TaskInstance
from airflow.models.dag_version import DagVersion
+from airflow.models.dagbag import DBDagBag
from airflow.models.dagrun import DagRun, get_or_create_dagrun
-from airflow.models.dynamic_region import SENTINEL_REGION_ID
+from airflow.models.dynamic_region import SENTINEL_REGION_ID, DynamicRegion
from airflow.models.expandinput import NotFullyPopulated
from airflow.models.serialized_dag import SerializedDagModel
+from airflow.models.task_coordinates import (
+ TaskCoordinateResolver,
+ enclosing_loop,
+ public_map_index_expression,
+)
from airflow.sdk.definitions.dag import DAG, _run_task
from airflow.sdk.definitions.param import ParamsDict
from airflow.serialization.definitions.dag import SerializedDAG
@@ -185,6 +194,8 @@ def _get_ti(
if not logical_date_or_run_id and not create_if_necessary:
raise ValueError("Must provide `logical_date_or_run_id` if not `create_if_necessary`.")
+ if enclosing_loop(task) is not None:
+ raise ValueError("A loop task requires an explicit region scope")
if task.get_needs_expansion():
if map_index < 0:
raise RuntimeError("No map_index passed to mapped task")
@@ -197,7 +208,13 @@ def _get_ti(
session=session,
)
- ti_or_none = dag_run.get_task_instance(task.task_id, map_index=map_index, session=session)
+ producers = TaskCoordinateResolver(DBDagBag(), session).resolve(
+ dag_id=dag_run.dag_id,
+ run_id=dag_run.run_id,
+ task_id=task.task_id,
+ map_indexes=map_index,
+ )
+ ti_or_none = producers[0] if producers else None
ti: TaskInstance
if ti_or_none is None:
if not create_if_necessary:
@@ -223,7 +240,48 @@ def _get_ti(
raise ValueError(
f"Cannot create TaskInstance for {dag.dag_id} because the Dag is not serialized."
)
- ti = TaskInstance(task, run_id=dag_run.run_id, map_index=map_index, dag_version_id=dag_version.id)
+ region_id = SENTINEL_REGION_ID
+ if map_index >= 0:
+ existing_region_id = session.scalar(
+ select(DynamicRegion.id)
+ .where(
+ DynamicRegion.dag_id == dag.dag_id,
+ DynamicRegion.run_id == dag_run.run_id,
+ DynamicRegion.node_id == task.task_id,
+ )
+ .order_by(DynamicRegion.id.desc())
+ .limit(1)
+ )
+ has_legacy_instances = existing_region_id is None and session.scalar(
+ select(
+ select(TaskInstance.id)
+ .where(
+ TaskInstance.working_set.is_(True),
+ TaskInstance.dag_id == dag.dag_id,
+ TaskInstance.run_id == dag_run.run_id,
+ TaskInstance.task_id == task.task_id,
+ )
+ .exists()
+ )
+ )
+ if existing_region_id is not None:
+ region_id = existing_region_id
+ elif not has_legacy_instances:
+ region_id = uuid6.uuid7()
+ if dag_run in session:
+ session.add(
+ DynamicRegion(
+ id=region_id, dag_id=dag.dag_id, run_id=dag_run.run_id, node_id=task.task_id
+ )
+ )
+ session.flush()
+ ti = TaskInstance(
+ task,
+ run_id=dag_run.run_id,
+ region_id=region_id,
+ region_index=map_index,
+ dag_version_id=dag_version.id,
+ )
if dag_run in session:
session.add(ti)
ti.dag_run = dag_run
@@ -389,7 +447,19 @@ def task_states_for_dag_run(args, *, session: Session = NEW_SESSION) -> None:
"not found"
)
- has_mapped_instances = any(ti.map_index >= 0 for ti in dag_run.task_instances)
+ rows = session.execute(
+ select(TaskInstance, public_map_index_expression(TaskInstance)).where(
+ TaskInstance.working_set.is_(True),
+ TaskInstance.dag_id == dag_run.dag_id,
+ TaskInstance.run_id == dag_run.run_id,
+ )
+ ).all()
+ task_instances = [ti for ti, _ in rows]
+ map_indexes = {ti.id: map_index for ti, map_index in rows}
+ has_mapped_instances = any(map_index >= 0 for map_index in map_indexes.values())
+ has_loop_instances = any(
+ ti.region_id != SENTINEL_REGION_ID and map_indexes[ti.id] < 0 for ti in task_instances
+ )
def format_task_instance(ti: TaskInstance) -> dict[str, str]:
data = {
@@ -401,10 +471,14 @@ def task_states_for_dag_run(args, *, session: Session = NEW_SESSION) -> None:
"end_date": ti.end_date.isoformat() if ti.end_date else "",
}
if has_mapped_instances:
- data["map_index"] = str(ti.map_index) if ti.map_index is not None and ti.map_index >= 0 else ""
+ map_index = map_indexes[ti.id]
+ data["map_index"] = str(map_index) if map_index >= 0 else ""
+ if has_loop_instances:
+ data["region_id"] = str(ti.region_id)
+ data["region_index"] = str(ti.region_index)
return data
- AirflowConsole().print_as(data=dag_run.task_instances, output=args.output, mapper=format_task_instance)
+ AirflowConsole().print_as(data=task_instances, output=args.output, mapper=format_task_instance)
@cli_utils.action_cli(check_db=False)
diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
index c7eb7458aa0..c22753b6a4c 100644
--- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
+++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
@@ -110,6 +110,7 @@ from airflow.models.dagbag import CachedDBDagBag, DBDagBag
from airflow.models.dagbundle import DagBundleModel
from airflow.models.dagrun import DagRun
from airflow.models.dagwarning import DagWarning, DagWarningType
+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
@@ -1041,7 +1042,7 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
isouter=True,
)
.where(func.coalesce(dr_task_concurrency_subquery.c.task_per_dr_count, 0) < DM.max_active_tasks)
- .order_by(-TI.priority_weight, DR.logical_date, TI.map_index)
+ .order_by(-TI.priority_weight, DR.logical_date, TI.region_index)
)
# Starvation filters should be applied before computing the row_num based on the
@@ -1069,14 +1070,14 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
func.row_number()
.over(
partition_by=[TI.dag_id, TI.run_id],
- order_by=[-TI.priority_weight, DR.logical_date, TI.map_index],
+ order_by=[-TI.priority_weight, DR.logical_date, TI.region_index],
)
.label("row_num"),
DM.max_active_tasks.label("dr_max_active_tasks"),
# Create columns for the order_by checks here for sqlite.
TI.priority_weight.label("priority_weight_for_ordering"),
DR.logical_date.label("logical_date_for_ordering"),
- TI.map_index.label("map_index_for_ordering"),
+ TI.region_index.label("region_index_for_ordering"),
)
).subquery()
@@ -1086,17 +1087,14 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
.select_from(ranked_query)
.join(
TI,
- (TI.dag_id == ranked_query.c.dag_id)
- & (TI.task_id == ranked_query.c.task_id)
- & (TI.run_id == ranked_query.c.run_id)
- & (TI.map_index == ranked_query.c.map_index),
+ TI.id == ranked_query.c.id,
)
.where(ranked_query.c.row_num <= ranked_query.c.dr_max_active_tasks)
# Add the order_by columns from the ranked query for sqlite.
.order_by(
-ranked_query.c.priority_weight_for_ordering,
ranked_query.c.logical_date_for_ordering,
- ranked_query.c.map_index_for_ordering,
+ ranked_query.c.region_index_for_ordering,
)
.options(selectinload(TI.dag_model))
# Eager-load the run's pinned DagVersion (dag_run.created_dag_version): TIs become
@@ -1575,7 +1573,7 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
continue
msg = (
- "TaskInstance Finished: dag_id=%s, task_id=%s, run_id=%s, map_index=%s, ti_id=%s, "
+ "TaskInstance Finished: dag_id=%s, task_id=%s, run_id=%s, %s, ti_id=%s, "
"run_start_date=%s, run_end_date=%s, "
"run_duration=%s, state=%s, executor=%s, executor_state=%s, try_number=%s, max_tries=%s, "
"pool=%s, queue=%s, priority_weight=%d, operator=%s, queued_dttm=%s, scheduled_dttm=%s,"
@@ -1586,7 +1584,9 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
ti.dag_id,
ti.task_id,
ti.run_id,
- ti.map_index,
+ f"map_index={ti.region_index}"
+ if ti.region_id == SENTINEL_REGION_ID
+ else f"region_id={ti.region_id}, region_index={ti.region_index}",
ti.id,
ti.start_date,
ti.end_date,
@@ -3462,7 +3462,7 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
Log.dag_id == ti.dag_id,
Log.task_id == ti.task_id,
Log.run_id == ti.run_id,
- Log.map_index == ti.map_index,
+ Log.map_index == ti.region_index,
Log.try_number == ti.try_number,
Log.event == "running",
)
@@ -3477,7 +3477,7 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
Log.task_id == ti.task_id,
Log.dag_id == ti.dag_id,
Log.run_id == ti.run_id,
- Log.map_index == ti.map_index,
+ Log.map_index == ti.region_index,
Log.try_number == ti.try_number,
Log.event == TASK_STUCK_IN_QUEUED_RESCHEDULE_EVENT,
)
@@ -3652,7 +3652,9 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
TI.dag_id,
TI.task_id,
TI.run_id,
- TI.map_index,
+ TI.region_id,
+ TI.region_index,
+ TI.dag_version_id,
TI.try_number,
TI.state,
TI.external_executor_id,
@@ -3971,7 +3973,15 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
has_callback_version = _ensure_ti_has_dag_version_id(ti, session, self.log)
if not has_callback_version and ti.state != TaskInstanceState.RESTARTING:
continue
+ callback_map_index = None
if has_callback_version:
+ try:
+ callback_map_index = coordinates.public_map_index(ti)
+ except (TaskNotFound, ValueError):
+ self.log.warning(
+ "Skipping heartbeat callback for %s: pinned task definition is missing", ti
+ )
+ if callback_map_index is not None:
context_from_server = TIRunContext(
dag_run=DRDataModel.model_validate(ti.dag_run, from_attributes=True),
max_tries=ti.max_tries,
@@ -3984,9 +3994,7 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
bundle_name=bundle_name,
bundle_version=bundle_version,
version_data=version_data,
- ti=task_instance_to_runtime(
- ti, model=TIDataModel, map_index=coordinates.public_map_index(ti)
- ),
+ ti=task_instance_to_runtime(ti, model=TIDataModel, map_index=callback_map_index),
msg=msg,
task_callback_type=task_callback_type,
context_from_server=context_from_server,
@@ -4003,9 +4011,7 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
bundle_name=bundle_name,
bundle_version=bundle_version,
version_data=version_data,
- ti=task_instance_to_runtime(
- ti, model=TIDataModel, map_index=coordinates.public_map_index(ti)
- ),
+ ti=task_instance_to_runtime(ti, model=TIDataModel, map_index=callback_map_index),
msg=msg,
email_type=(
"retry" if task_callback_type == TaskInstanceState.UP_FOR_RETRY else "failure"
@@ -4055,8 +4061,11 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
"Run Id": ti.run_id,
}
- if ti.map_index != -1:
- task_instance_heartbeat_timeout_message_details["Map Index"] = ti.map_index
+ if ti.region_id != SENTINEL_REGION_ID:
+ task_instance_heartbeat_timeout_message_details["Region Id"] = str(ti.region_id)
+ task_instance_heartbeat_timeout_message_details["Region Index"] = ti.region_index
+ elif ti.region_index != -1:
+ task_instance_heartbeat_timeout_message_details["Map Index"] = ti.region_index
if ti.hostname:
task_instance_heartbeat_timeout_message_details["Hostname"] = ti.hostname
if ti.external_executor_id:
diff --git a/airflow-core/src/airflow/migrations/versions/0144_3_4_0_rename_stored_map_index.py b/airflow-core/src/airflow/migrations/versions/0144_3_4_0_rename_stored_map_index.py
new file mode 100644
index 00000000000..65b0491aee7
--- /dev/null
+++ b/airflow-core/src/airflow/migrations/versions/0144_3_4_0_rename_stored_map_index.py
@@ -0,0 +1,61 @@
+# 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.
+
+"""
+Rename stored map indexes to region indexes.
+
+Revision ID: 7f8c9a2d410e
+Revises: 54a27b6f9d01
+Create Date: 2026-09-27 23:00:00.000000
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+
+revision = "7f8c9a2d410e"
+down_revision = "54a27b6f9d01"
+branch_labels = None
+depends_on = None
+airflow_version = "3.4.0"
+
+# ``legacy_task_data_owner``, ``xcom_v1`` and ``rtif_v1`` keep ``map_index``: they record
+# the coordinates their rows had before regions existed and are never rewritten.
+_TABLES = ("task_instance", "task_state_store")
+
+
+def _rename(old: str, new: str) -> None:
+ for table in _TABLES:
+ op.alter_column(
+ table,
+ old,
+ new_column_name=new,
+ existing_type=sa.Integer(),
+ existing_nullable=False,
+ existing_server_default="-1",
+ )
+
+
+def upgrade():
+ """Rename the stored coordinate without rewriting rows, keys or indexes."""
+ _rename("map_index", "region_index")
+
+
+def downgrade():
+ """Restore the prior column name without changing coordinate identity."""
+ _rename("region_index", "map_index")
diff --git a/airflow-core/src/airflow/models/asset.py b/airflow-core/src/airflow/models/asset.py
index ddac396f182..28cf3107564 100644
--- a/airflow-core/src/airflow/models/asset.py
+++ b/airflow-core/src/airflow/models/asset.py
@@ -883,7 +883,7 @@ class AssetEvent(Base):
AssetEvent.source_dag_id == foreign(TaskInstance.dag_id),
AssetEvent.source_run_id == foreign(TaskInstance.run_id),
AssetEvent.source_task_id == foreign(TaskInstance.task_id),
- AssetEvent.source_map_index == foreign(TaskInstance.map_index),
+ AssetEvent.source_map_index == foreign(TaskInstance.region_index),
TaskInstance.working_set.is_(True),
)""",
viewonly=True,
diff --git a/airflow-core/src/airflow/models/dagbag.py b/airflow-core/src/airflow/models/dagbag.py
index 1665699bec1..4415d8b0b2b 100644
--- a/airflow-core/src/airflow/models/dagbag.py
+++ b/airflow-core/src/airflow/models/dagbag.py
@@ -22,6 +22,7 @@ import math
import time
from collections.abc import MutableMapping
from contextlib import nullcontext
+from functools import cache
from threading import RLock
from typing import TYPE_CHECKING, Any, NamedTuple
from uuid import UUID
@@ -45,6 +46,9 @@ if TYPE_CHECKING:
from airflow.serialization.definitions.dag import SerializedDAG
+SHARED_DAG_CACHE_SIZE = 64
+
+
class _CacheEntry(NamedTuple):
"""A cached deserialized DAG plus the metadata needed to detect staleness on lookup."""
@@ -289,6 +293,23 @@ class CachedDBDagBag(DBDagBag):
stats.gauge(f"{self._stats_prefix}.cache_size", size, rate=rate)
+@cache
+def get_shared_dag_bag() -> DBDagBag:
+ """
+ Return the process-wide Dag cache for structural lookups on hot paths.
+
+ Scheduling, log addressing and task coordinate resolution only need a Dag's shape, and a
+ version's definition is immutable, so one revalidated cache is reused rather than
+ deserializing the serialized Dag for every pass or request.
+ """
+ return CachedDBDagBag(
+ load_op_links=False,
+ cache_size=SHARED_DAG_CACHE_SIZE,
+ cache_ttl=0,
+ stats_prefix="shared.dag_bag",
+ )
+
+
def generate_md5_hash(context):
bundle_name = context.get_current_parameters()["bundle_name"]
relative_fileloc = context.get_current_parameters()["relative_fileloc"]
diff --git a/airflow-core/src/airflow/models/dagrun.py b/airflow-core/src/airflow/models/dagrun.py
index a5b880dd82a..40b4efe3941 100644
--- a/airflow-core/src/airflow/models/dagrun.py
+++ b/airflow-core/src/airflow/models/dagrun.py
@@ -88,7 +88,7 @@ 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
+from airflow.models.dynamic_region import 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
@@ -973,7 +973,7 @@ class DagRun(Base, LoggingMixin):
TI.dag_id == dag_id,
TI.run_id == run_id,
)
- .order_by(TI.task_id, TI.map_index)
+ .order_by(TI.task_id, TI.region_id, TI.region_index)
)
if state:
@@ -1063,13 +1063,14 @@ class DagRun(Base, LoggingMixin):
task_id: str,
*,
map_index: int = -1,
- region_id: UUID = SENTINEL_REGION_ID,
+ region_id: UUID | None = None,
session: Session = NEW_SESSION,
) -> TI | None:
"""
Return the task instance specified by task_id for this dag run.
:param task_id: the task id
+ :param region_id: the exact region; by default the task instance addressed by its public ``map_index``
:param session: Sqlalchemy ORM Session
"""
return DagRun.fetch_task_instance(
@@ -1089,7 +1090,7 @@ class DagRun(Base, LoggingMixin):
task_id: str,
*,
map_index: int = -1,
- region_id: UUID = SENTINEL_REGION_ID,
+ region_id: UUID | None = None,
session: Session = NEW_SESSION,
) -> TI | None:
"""
@@ -1100,11 +1101,15 @@ class DagRun(Base, LoggingMixin):
:param task_id: the task id
:param session: Sqlalchemy ORM Session
"""
- return session.scalars(
- select(TI).filter_by(
- dag_id=dag_id, run_id=dag_run_id, task_id=task_id, map_index=map_index, region_id=region_id
- )
- ).one_or_none()
+ query = select(TI).filter_by(
+ dag_id=dag_id, run_id=dag_run_id, task_id=task_id, region_index=map_index
+ )
+ query = query.where(
+ public_region_filter(TI, dag_ids=dag_id, run_ids=dag_run_id, task_ids=task_id)
+ if region_id is None
+ else TI.region_id == region_id
+ )
+ return session.scalars(query).one_or_none()
def get_dag(self) -> SerializedDAG:
"""
@@ -1532,10 +1537,9 @@ class DagRun(Base, LoggingMixin):
task_instance_to_runtime,
)
from airflow.models.dag_version import DagVersion
- from airflow.models.dagbag import DBDagBag
from airflow.models.task_coordinates import TaskCoordinateResolver
- coordinates = TaskCoordinateResolver(DBDagBag(), session)
+ coordinates = TaskCoordinateResolver.for_dag(self.dag, session)
if relevant_ti.dag_version_id is not None:
return task_instance_to_runtime(
@@ -1717,10 +1721,7 @@ class DagRun(Base, LoggingMixin):
"""
Try to expand the ti, if needed.
- If the ti needs expansion, newly created task instances are
- returned as well as the original ti.
- The original ti is also modified in-place and assigned the
- ``map_index`` of 0.
+ Expansion replaces the placeholder with newly created task instances.
If the ti does not need expansion, either because the task is not
mapped, or has already been expanded, *None* is returned.
@@ -1730,7 +1731,7 @@ class DagRun(Base, LoggingMixin):
if TYPE_CHECKING:
assert ti.task
- if ti.map_index >= 0: # Already expanded, we're good.
+ if ti.region_index >= 0: # Already expanded, we're good.
return None
if is_mapped(ti.task):
@@ -1764,7 +1765,7 @@ class DagRun(Base, LoggingMixin):
# expanded before executed. Also see _revise_map_indexes_if_mapped
# docstring for additional information.
new_tis = None
- if schedulable.map_index < 0:
+ if schedulable.region_index < 0:
new_tis = _expand_mapped_task_if_needed(schedulable)
if new_tis is not None:
additional_tis.extend(new_tis)
@@ -1779,9 +1780,7 @@ class DagRun(Base, LoggingMixin):
expansion_key = (schedulable.task.task_id, schedulable.region_id)
if expansion_key not in revised_map_index_task_ids:
revised_tis = self._revise_map_indexes_if_mapped(
- schedulable.task,
- dag_version_id=schedulable.dag_version_id,
- region_id=schedulable.region_id,
+ schedulable,
session=session,
)
ready_tis.extend(revised_tis)
@@ -1930,12 +1929,42 @@ class DagRun(Base, LoggingMixin):
)
# Create the missing tasks, including mapped tasks
- tis_to_create = self._create_tasks(
- (task for task in dag.task_dict.values() if task_filter(task)),
- task_creator,
- session=session,
+ missing_tasks = [task for task in dag.task_dict.values() if task_filter(task)]
+ tis_to_create = list(
+ self._create_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)
+ empty_literal_tasks = [
+ task.task_id for task in missing_tasks if self._get_literal_expansion_count(task) == 0
+ ]
+ if empty_literal_tasks:
+ session.execute(
+ update(TI)
+ .where(
+ TI.working_set.is_(True),
+ TI.dag_id == self.dag_id,
+ TI.run_id == self.run_id,
+ TI.task_id.in_(empty_literal_tasks),
+ TI.region_index == -1,
+ )
+ .values(state=TaskInstanceState.SKIPPED)
+ .execution_options(synchronize_session=False)
+ )
+
+ @staticmethod
+ def _get_literal_expansion_count(task: Operator) -> int | None:
+ """Return how many instances a task whose inputs are all literals expands to, else None."""
+ from airflow.models.expandinput import NotFullyPopulated
+
+ try:
+ return task.get_parse_time_mapped_ti_count()
+ except (NotMapped, NotFullyPopulated):
+ return None
def _check_for_removed_or_restored_tasks(
self, dag: SerializedDAG, ti_mutation_hook, *, session: Session
@@ -1951,9 +1980,11 @@ class DagRun(Base, LoggingMixin):
"""
from airflow.models.expandinput import NotFullyPopulated
+ from airflow.models.task_coordinates import TaskCoordinateResolver
from airflow.serialization.definitions.mappedoperator import get_mapped_ti_count
tis = self.get_task_instances(session=session)
+ coordinates = TaskCoordinateResolver.for_dag(dag, session)
# check for removed or restored tasks
task_ids = set()
@@ -1981,10 +2012,15 @@ class DagRun(Base, LoggingMixin):
except NotFullyPopulated:
# What if it is _now_ dynamically mapped, but wasn't before?
try:
- total_length = get_mapped_ti_count(task, self.run_id, session=session)
+ total_length = get_mapped_ti_count(
+ task,
+ self.run_id,
+ session=session,
+ producer_contexts=coordinates.producer_contexts(ti),
+ )
except NotFullyPopulated:
# Not all upstreams finished, so we can't tell what should be here. Remove everything.
- if ti.map_index >= 0:
+ if ti.region_index >= 0:
self.log.debug(
"Removing the unmapped TI '%s' as the mapping can't be resolved yet", ti
)
@@ -1992,7 +2028,7 @@ class DagRun(Base, LoggingMixin):
continue
else:
# Upstreams finished, check there aren't any extras
- if ti.map_index >= total_length:
+ if ti.region_index >= total_length:
self.log.debug(
"Removing task '%s' as the map_index is longer than the resolved mapping list (%d)",
ti,
@@ -2002,7 +2038,7 @@ class DagRun(Base, LoggingMixin):
continue
else:
# Check if the number of mapped literals has changed, and we need to mark this TI as removed.
- if ti.map_index >= num_mapped_tis:
+ if ti.region_index >= num_mapped_tis:
self.log.debug(
"Removing task '%s' as the map_index is longer than the literal mapping list (%s)",
ti,
@@ -2010,9 +2046,7 @@ class DagRun(Base, LoggingMixin):
)
ti.state = TaskInstanceState.REMOVED
continue
- if ti.map_index < 0:
- self.log.debug("Removing the unmapped TI '%s' as the mapping can now be performed", ti)
- ti.state = TaskInstanceState.REMOVED
+ if ti.region_index < 0:
continue
if ti.state == TaskInstanceState.REMOVED:
@@ -2033,7 +2067,7 @@ class DagRun(Base, LoggingMixin):
ti_mutation_hook: Callable,
hook_is_noop: Literal[True],
dag_version_id: UUID,
- ) -> Callable[[Operator, Iterable[int]], Iterator[dict[str, Any]]]: ...
+ ) -> Callable[[Operator, Iterable[int], UUID], Iterator[dict[str, Any]]]: ...
@overload
def _get_task_creator(
@@ -2042,7 +2076,7 @@ class DagRun(Base, LoggingMixin):
ti_mutation_hook: Callable,
hook_is_noop: Literal[False],
dag_version_id: UUID,
- ) -> Callable[[Operator, Iterable[int]], Iterator[TI]]: ...
+ ) -> Callable[[Operator, Iterable[int], UUID], Iterator[TI]]: ...
def _get_task_creator(
self,
@@ -2050,7 +2084,7 @@ class DagRun(Base, LoggingMixin):
ti_mutation_hook: Callable,
hook_is_noop: Literal[True, False],
dag_version_id: UUID,
- ) -> Callable[[Operator, Iterable[int]], Iterator[dict[str, Any]] | Iterator[TI]]:
+ ) -> Callable[[Operator, Iterable[int], UUID], Iterator[dict[str, Any]] | Iterator[TI]]:
"""
Get the task creator function.
@@ -2063,13 +2097,16 @@ class DagRun(Base, LoggingMixin):
"""
if hook_is_noop:
- def create_ti_mapping(task: Operator, indexes: Iterable[int]) -> Iterator[dict[str, Any]]:
+ def create_ti_mapping(
+ task: Operator, indexes: Iterable[int], region_id: UUID
+ ) -> Iterator[dict[str, Any]]:
created_counts[task.task_type] += 1
- for map_index in indexes:
+ for region_index in indexes:
yield TI.insert_mapping(
self.run_id,
task,
- map_index=map_index,
+ region_index=region_index,
+ region_id=region_id,
dag_version_id=dag_version_id,
dag_run=self,
)
@@ -2078,9 +2115,15 @@ class DagRun(Base, LoggingMixin):
else:
- def create_ti(task: Operator, indexes: Iterable[int]) -> Iterator[TI]:
- for map_index in indexes:
- ti = TI(task, run_id=self.run_id, map_index=map_index, dag_version_id=dag_version_id)
+ def create_ti(task: Operator, indexes: Iterable[int], region_id: UUID) -> Iterator[TI]:
+ for region_index in indexes:
+ ti = TI(
+ task,
+ run_id=self.run_id,
+ region_index=region_index,
+ region_id=region_id,
+ dag_version_id=dag_version_id,
+ )
ti_mutation_hook(ti, dag_run=self)
if ti.operator:
created_counts[ti.operator] += 1
@@ -2092,38 +2135,50 @@ class DagRun(Base, LoggingMixin):
def _create_tasks(
self,
tasks: Iterable[Operator],
- task_creator: Callable[[Operator, Iterable[int]], CreatedTasks],
+ task_creator: Callable[[Operator, Iterable[int], UUID], CreatedTasks],
*,
session: Session,
+ parent_region: tuple[UUID, int] | None = None,
+ expand_literals: bool = False,
) -> CreatedTasks:
"""
- Create missing tasks -- and expand any MappedOperator that _only_ have literals as input.
+ Create ordinary tasks and a region-owned placeholder for each mapped task.
:param tasks: Tasks to create jobs for in the DAG run
:param task_creator: Function to create task instances
+ :param expand_literals: Create the slots of a task whose inputs are all literals right away,
+ so its placeholder is never expanded one task at a time.
"""
- from airflow.models.expandinput import NotFullyPopulated
- from airflow.serialization.definitions.mappedoperator import get_mapped_ti_count
-
- map_indexes: Iterable[int]
+ tasks = list(tasks)
+ regions = {
+ task.task_id: DynamicRegion(
+ dag_id=self.dag_id,
+ run_id=self.run_id,
+ node_id=task.task_id,
+ parent_region_id=parent_region[0] if parent_region else None,
+ parent_region_index=parent_region[1] if parent_region else None,
+ )
+ for task in tasks
+ if task.get_needs_expansion()
+ }
+ session.add_all(regions.values())
+ session.flush()
for task in tasks:
- try:
- count = get_mapped_ti_count(task, self.run_id, session=session)
- except (NotMapped, NotFullyPopulated):
- map_indexes = (-1,)
+ indexes: Iterable[int]
+ if task.task_id in regions:
+ region_id, indexes = regions[task.task_id].id, (-1,)
+ if expand_literals and (count := self._get_literal_expansion_count(task)):
+ indexes = range(count)
+ elif parent_region:
+ region_id, indexes = parent_region[0], (parent_region[1],)
else:
- if count:
- map_indexes = range(count)
- else:
- # Make sure to always create at least one ti; this will be
- # marked as REMOVED later at runtime.
- map_indexes = (-1,)
- yield from task_creator(task, map_indexes)
+ region_id, indexes = SENTINEL_REGION_ID, (-1,)
+ yield from task_creator(task, indexes, region_id)
def _create_task_instances(
self,
dag_id: str,
- tasks: Iterator[dict[str, Any]] | Iterator[TI],
+ tasks: Iterable[dict[str, Any]] | Iterable[TI],
created_counts: dict[str, int],
hook_is_noop: bool,
*,
@@ -2172,11 +2227,9 @@ class DagRun(Base, LoggingMixin):
def _revise_map_indexes_if_mapped(
self,
- task: Operator,
+ ti: TI,
*,
- dag_version_id: UUID | None,
session: Session,
- region_id: UUID = SENTINEL_REGION_ID,
) -> list[TI]:
"""
Check if task increased or reduced in length and handle appropriately.
@@ -2187,17 +2240,30 @@ class DagRun(Base, LoggingMixin):
for more details.
"""
from airflow.models.expandinput import NotFullyPopulated
+ from airflow.models.task_coordinates import TaskCoordinateResolver
from airflow.serialization.definitions.mappedoperator import get_mapped_ti_count
+ task = ti.task
+ if task is None:
+ raise ValueError(f"Task definition is required to revise {ti}")
+ if not task.get_needs_expansion():
+ return []
+ region_id = ti.region_id
+ dag_version_id = ti.dag_version_id
try:
- total_length = get_mapped_ti_count(task, self.run_id, session=session)
+ total_length = get_mapped_ti_count(
+ task,
+ self.run_id,
+ session=session,
+ producer_contexts=TaskCoordinateResolver.for_dag(self.dag, session).producer_contexts(ti),
+ )
except NotMapped:
return [] # Not a mapped task, don't need to do anything.
except NotFullyPopulated:
return [] # Upstreams not ready, don't need to revise this yet.
query = session.scalars(
- select(TI.map_index).where(
+ select(TI.region_index).where(
TI.dag_id == self.dag_id,
TI.task_id == task.task_id,
TI.run_id == self.run_id,
@@ -2214,8 +2280,8 @@ class DagRun(Base, LoggingMixin):
TI.dag_id == self.dag_id,
TI.task_id == task.task_id,
TI.run_id == self.run_id,
+ TI.region_index.in_(removed_indexes),
TI.region_id == region_id,
- TI.map_index.in_(removed_indexes),
)
.values(state=TaskInstanceState.REMOVED)
)
@@ -2238,7 +2304,7 @@ class DagRun(Base, LoggingMixin):
ti = TI(
task,
run_id=self.run_id,
- map_index=index,
+ region_index=index,
region_id=region_id,
state=None,
dag_version_id=dag_version_id,
diff --git a/airflow-core/src/airflow/models/dynamic_region.py b/airflow-core/src/airflow/models/dynamic_region.py
index 72e02449fc4..2324ba6a4e7 100644
--- a/airflow-core/src/airflow/models/dynamic_region.py
+++ b/airflow-core/src/airflow/models/dynamic_region.py
@@ -23,7 +23,19 @@ from uuid import UUID
import attrs
import uuid6
-from sqlalchemy import CheckConstraint, ForeignKeyConstraint, Index, Integer, UniqueConstraint, Uuid, select
+from sqlalchemy import (
+ CheckConstraint,
+ ForeignKeyConstraint,
+ Index,
+ Integer,
+ UniqueConstraint,
+ Uuid,
+ and_,
+ literal,
+ or_,
+ select,
+ union_all,
+)
from sqlalchemy.orm import Mapped, mapped_column
from airflow._shared.timezones import timezone
@@ -34,6 +46,7 @@ SENTINEL_REGION_ID = UUID(int=0)
if TYPE_CHECKING:
from sqlalchemy.orm import Session
+ from sqlalchemy.sql.elements import ColumnElement
from airflow.models.taskinstance import TaskInstance
@@ -102,6 +115,53 @@ class DynamicRegion(Base):
)
+def public_region_filter(
+ model,
+ *,
+ dag_ids: str | Collection[str] | None = None,
+ run_ids: str | Collection[str] | None = None,
+ task_ids: str | Collection[str] | None = None,
+) -> ColumnElement[bool]:
+ """
+ Match rows addressed by public coordinates: legacy rows and a task's own top-level region.
+
+ Pass the dag, run and task the caller is looking up so the regions that can match are
+ resolved once from ``dynamic_region`` instead of being checked per row. The row filter is then an
+ equality-or-IN on ``region_id`` that the unique key serves, where a per-row check has to read
+ every row of the task.
+ """
+ own_region = (
+ select(DynamicRegion.id)
+ .where(
+ DynamicRegion.id == model.region_id,
+ DynamicRegion.node_id == model.task_id,
+ DynamicRegion.parent_region_id.is_(None),
+ )
+ .correlate(model)
+ .exists()
+ )
+ exact = or_(model.region_id == SENTINEL_REGION_ID, own_region)
+ if dag_ids is None or run_ids is None or task_ids is None:
+ return exact
+ candidates = union_all(
+ select(literal(SENTINEL_REGION_ID, Uuid()).label("id")),
+ select(DynamicRegion.id.label("id")).where(
+ _match_any(DynamicRegion.dag_id, dag_ids),
+ _match_any(DynamicRegion.run_id, run_ids),
+ _match_any(DynamicRegion.node_id, task_ids),
+ DynamicRegion.parent_region_id.is_(None),
+ ),
+ ).subquery()
+ narrowed = model.region_id.in_(select(candidates.c.id))
+ if isinstance(dag_ids, str) and isinstance(run_ids, str) and isinstance(task_ids, str):
+ return narrowed
+ return and_(narrowed, exact)
+
+
+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(
*,
dag_id: str,
diff --git a/airflow-core/src/airflow/models/expandinput.py b/airflow-core/src/airflow/models/expandinput.py
index 0363bae9262..b0d482dc0b7 100644
--- a/airflow-core/src/airflow/models/expandinput.py
+++ b/airflow-core/src/airflow/models/expandinput.py
@@ -31,6 +31,7 @@ if TYPE_CHECKING:
from sqlalchemy.orm import Session
from typing_extensions import TypeIs
+ from airflow.models.dynamic_region import ProducerContext
from airflow.serialization.definitions.mappedoperator import Operator
from airflow.serialization.definitions.xcom_arg import SchedulerXComArg
@@ -117,7 +118,13 @@ class SchedulerDictOfListsExpandInput:
raise NotFullyPopulated(set(self.value).difference(literal_keys))
return functools.reduce(operator.mul, literal_values, 1)
- def _get_map_lengths(self, run_id: str, *, session: Session) -> dict[str, int]:
+ def _get_map_lengths(
+ self,
+ run_id: str,
+ *,
+ session: Session,
+ producer_contexts: Mapping[str, ProducerContext] | None = None,
+ ) -> dict[str, int]:
"""
Return dict of argument name to map length.
@@ -130,7 +137,7 @@ class SchedulerDictOfListsExpandInput:
# more efficient to do one single db call and unpack the value here?
def _get_length(v: ExpandArgument) -> int | None:
if isinstance(v, SchedulerXComArg):
- return get_task_map_length(v, run_id, session=session)
+ return get_task_map_length(v, run_id, session=session, producer_contexts=producer_contexts)
# Unfortunately a user-defined TypeGuard cannot apply negative type
# narrowing. https://github.com/python/typing/discussions/1013
@@ -145,10 +152,16 @@ class SchedulerDictOfListsExpandInput:
raise NotFullyPopulated(set(self.value).difference(map_lengths))
return map_lengths
- def get_total_map_length(self, run_id: str, *, session: Session) -> int:
+ def get_total_map_length(
+ self,
+ run_id: str,
+ *,
+ session: Session,
+ producer_contexts: Mapping[str, ProducerContext] | None = None,
+ ) -> int:
if not self.value:
return 0
- lengths = self._get_map_lengths(run_id, session=session)
+ lengths = self._get_map_lengths(run_id, session=session, producer_contexts=producer_contexts)
return functools.reduce(operator.mul, (lengths[name] for name in self.value), 1)
def iter_references(self) -> Iterable[tuple[Operator, str]]:
@@ -177,12 +190,18 @@ class SchedulerListOfDictsExpandInput:
return len(self.value)
raise NotFullyPopulated({"expand_kwargs() argument"})
- def get_total_map_length(self, run_id: str, *, session: Session) -> int:
+ def get_total_map_length(
+ self,
+ run_id: str,
+ *,
+ session: Session,
+ producer_contexts: Mapping[str, ProducerContext] | None = None,
+ ) -> int:
from airflow.serialization.definitions.xcom_arg import get_task_map_length
if isinstance(self.value, Sized):
return len(self.value)
- length = get_task_map_length(self.value, run_id, session=session)
+ length = get_task_map_length(self.value, run_id, session=session, producer_contexts=producer_contexts)
if length is None:
raise NotFullyPopulated({"expand_kwargs() argument"})
return length
diff --git a/airflow-core/src/airflow/models/log.py b/airflow-core/src/airflow/models/log.py
index 220ab2b0b91..7f9e03d16d8 100644
--- a/airflow-core/src/airflow/models/log.py
+++ b/airflow-core/src/airflow/models/log.py
@@ -74,7 +74,7 @@ class Log(Base):
"TaskInstance",
viewonly=True,
foreign_keys=[dag_id, task_id, run_id, map_index],
- primaryjoin="and_(Log.dag_id == TaskInstance.dag_id, Log.task_id == TaskInstance.task_id, Log.run_id == TaskInstance.run_id, Log.map_index == TaskInstance.map_index, TaskInstance.working_set.is_(True))",
+ primaryjoin="and_(Log.dag_id == TaskInstance.dag_id, Log.task_id == TaskInstance.task_id, Log.run_id == TaskInstance.run_id, Log.map_index == TaskInstance.region_index, TaskInstance.working_set.is_(True))",
lazy="raise",
)
diff --git a/airflow-core/src/airflow/models/renderedtifields.py b/airflow-core/src/airflow/models/renderedtifields.py
index cea27c1b446..00152a33126 100644
--- a/airflow-core/src/airflow/models/renderedtifields.py
+++ b/airflow-core/src/airflow/models/renderedtifields.py
@@ -227,7 +227,7 @@ class RenderedTaskInstanceFields(Base):
TaskInstance.dag_id == ti.dag_id,
TaskInstance.task_id == ti.task_id,
TaskInstance.run_id == ti.run_id,
- TaskInstance.map_index == ti.map_index,
+ TaskInstance.region_index == ti.map_index,
TaskInstance.try_number == ti.try_number,
)
.execution_options(include_all_attempts=True)
diff --git a/airflow-core/src/airflow/models/serialized_dag.py b/airflow-core/src/airflow/models/serialized_dag.py
index b933194470d..c8fcea2d634 100644
--- a/airflow-core/src/airflow/models/serialized_dag.py
+++ b/airflow-core/src/airflow/models/serialized_dag.py
@@ -939,7 +939,9 @@ class SerializedDagModel(Base):
data = json.loads(self.data)
else:
raise ValueError("invalid or missing serialized DAG data")
- return DagSerialization.from_dict(data)
+ dag = DagSerialization.from_dict(data)
+ dag.dag_version_id = self.dag_version_id
+ return dag
@classmethod
@provide_session
diff --git a/airflow-core/src/airflow/models/task_coordinates.py b/airflow-core/src/airflow/models/task_coordinates.py
index 262948424ad..3be481ddc52 100644
--- a/airflow-core/src/airflow/models/task_coordinates.py
+++ b/airflow-core/src/airflow/models/task_coordinates.py
@@ -23,6 +23,7 @@ import attrs
from sqlalchemy import case, or_, select
from airflow.exceptions import TaskNotFound
+from airflow.models.dagbag import DBDagBag
from airflow.models.dagrun import DagRun
from airflow.models.dynamic_region import (
SENTINEL_REGION_ID,
@@ -32,7 +33,8 @@ from airflow.models.dynamic_region import (
resolve_current_producers,
)
from airflow.models.taskinstance import TaskInstance
-from airflow.serialization.definitions.taskgroup import SerializedLoopTaskGroup
+from airflow.serialization.definitions.mappedoperator import is_mapped
+from airflow.serialization.definitions.taskgroup import SerializedLoopTaskGroup, SerializedMappedTaskGroup
if TYPE_CHECKING:
from collections.abc import Collection
@@ -40,7 +42,6 @@ if TYPE_CHECKING:
from sqlalchemy.orm import Session
from sqlalchemy.sql.elements import ColumnElement
- from airflow.models.dagbag import DBDagBag
from airflow.serialization.definitions.dag import SerializedDAG, SerializedOperator
@@ -85,17 +86,33 @@ class TaskCoordinateResolver:
_dags: dict[UUID, SerializedDAG] = attrs.field(factory=dict, init=False)
_region_nodes: dict[UUID, str | None] = attrs.field(factory=dict, init=False)
+ @classmethod
+ def for_dag(cls, dag: SerializedDAG | None, session: Session) -> TaskCoordinateResolver:
+ """Build a resolver that reads from the Dag the caller already holds, loading others on demand."""
+ resolver = cls(DBDagBag(load_op_links=False), session)
+ resolver.adopt_dag(dag)
+ return resolver
+
+ def adopt_dag(self, dag: SerializedDAG | None) -> None:
+ if dag is not None and dag.dag_version_id is not None:
+ self._dags.setdefault(dag.dag_version_id, dag)
+
def get_task(
self, dag_id: str, run_id: str, task_id: str, *, dag_version_id: UUID | None = None
) -> SerializedOperator:
if dag_version_id is None:
return self._producer_task(dag_id, run_id, task_id)
+ dag = self.get_dag(dag_version_id)
+ if dag is None or dag.dag_id != dag_id:
+ raise ValueError(f"Pinned Dag for {dag_id}/{run_id} not found")
+ return dag.get_task(task_id)
+
+ def get_dag(self, dag_version_id: UUID) -> SerializedDAG | None:
if dag_version_id not in self._dags:
- dag = self.dag_bag.get_dag(dag_version_id, session=self.session)
- if dag is None or dag.dag_id != dag_id:
- raise ValueError(f"Pinned Dag for {dag_id}/{run_id} not found")
+ if (dag := self.dag_bag.get_dag(dag_version_id, session=self.session)) is None:
+ return None
self._dags[dag_version_id] = dag
- return self._dags[dag_version_id].get_task(task_id)
+ return self._dags[dag_version_id]
def _producer_task(
self,
@@ -158,6 +175,41 @@ class TaskCoordinateResolver:
)
return self._region_nodes[ti.region_id] == ti.task_id
+ def producer_contexts(
+ self,
+ caller: TaskInstance,
+ producer_task_ids: Collection[str] | None = None,
+ ) -> dict[str, ProducerContext]:
+ task = caller.task or self.get_task(
+ caller.dag_id, caller.run_id, caller.task_id, dag_version_id=caller.dag_version_id
+ )
+ if producer_task_ids is None:
+ producer_task_ids = (
+ {op.task_id for op in task.iter_mapped_dependencies()} if is_mapped(task) else set()
+ )
+ group = task.task_group
+ while group is not None:
+ if isinstance(group, SerializedMappedTaskGroup):
+ producer_task_ids.update(op.task_id for op in group.iter_mapped_dependencies())
+ group = group.parent_group
+ caller_loop = enclosing_loop(task)
+ contexts = {}
+ for task_id in producer_task_ids:
+ producer = (
+ task.dag.get_task(task_id)
+ if caller.task is not None and task.dag is not None
+ else self.get_task(
+ caller.dag_id, caller.run_id, task_id, dag_version_id=caller.dag_version_id
+ )
+ )
+ loop = enclosing_loop(producer)
+ if loop is None:
+ continue
+ if caller_loop is None or caller_loop.group_id != loop.group_id:
+ raise ValueError("A loop producer requires a consumer inside the loop")
+ contexts[task_id] = ProducerContext(caller.region_id, caller.region_index, loop.group_id)
+ return contexts
+
def has_regions(self, dag_id: str, run_id: str | None, task_id: str) -> bool:
query = select(TaskInstance.id).where(
TaskInstance.working_set.is_(True),
diff --git a/airflow-core/src/airflow/models/task_state_store.py b/airflow-core/src/airflow/models/task_state_store.py
index 9a4d443aa81..c0b3dc66d1f 100644
--- a/airflow-core/src/airflow/models/task_state_store.py
+++ b/airflow-core/src/airflow/models/task_state_store.py
@@ -21,7 +21,7 @@ from uuid import UUID
from sqlalchemy import ForeignKeyConstraint, Index, Integer, String, Text, UniqueConstraint, Uuid
from sqlalchemy.dialects.mysql import MEDIUMTEXT
-from sqlalchemy.orm import Mapped, mapped_column, synonym
+from sqlalchemy.orm import Mapped, mapped_column
from airflow._shared.timezones import timezone
from airflow.models.base import COLLATION_ARGS, Base, StringID
@@ -32,7 +32,7 @@ class TaskStateStoreModel(Base):
"""
Persists key/value state for a task within a single DAG run.
- Scoped to (dag_run_id, task_id, region_id, map_index). Retries of the same task share
+ Scoped to (dag_run_id, task_id, region_id, region_index). Retries of the same task share
the same rows — that is the point. Different DAG runs have different dag_run_id
values so they get independent namespaces automatically.
"""
@@ -43,8 +43,7 @@ class TaskStateStoreModel(Base):
dag_run_id: Mapped[int] = mapped_column(Integer, nullable=False)
task_id: Mapped[str] = mapped_column(StringID(), nullable=False)
- map_index: Mapped[int] = mapped_column(Integer, nullable=False, server_default="-1")
- region_index = synonym("map_index")
+ region_index: Mapped[int] = mapped_column(Integer, nullable=False, server_default="-1")
region_id: Mapped[UUID] = mapped_column(
Uuid(), nullable=False, default=UUID(int=0), server_default="00000000000000000000000000000000"
)
@@ -63,7 +62,7 @@ class TaskStateStoreModel(Base):
__table_args__ = (
UniqueConstraint(
- "dag_run_id", "task_id", "region_id", "map_index", "key", name="task_state_store_uq"
+ "dag_run_id", "task_id", "region_id", "region_index", "key", name="task_state_store_uq"
),
ForeignKeyConstraint(
["dag_run_id"],
@@ -71,6 +70,6 @@ class TaskStateStoreModel(Base):
name="task_state_store_dag_run_fkey",
ondelete="CASCADE",
),
- Index("idx_task_state_store_lookup", "dag_id", "run_id", "task_id", "map_index"),
+ Index("idx_task_state_store_lookup", "dag_id", "run_id", "task_id", "region_index"),
Index("idx_task_state_store_expires_at", "expires_at"),
)
diff --git a/airflow-core/src/airflow/models/taskinstance.py b/airflow-core/src/airflow/models/taskinstance.py
index 7a4ce0d3755..003a6e06065 100644
--- a/airflow-core/src/airflow/models/taskinstance.py
+++ b/airflow-core/src/airflow/models/taskinstance.py
@@ -24,7 +24,7 @@ import logging
import math
import warnings
from collections import defaultdict
-from collections.abc import Callable, Collection, Iterable, Sequence
+from collections.abc import Callable, Collection, Iterable, Mapping, Sequence
from datetime import datetime, timedelta
from typing import TYPE_CHECKING, Any, NamedTuple
from urllib.parse import quote
@@ -101,7 +101,12 @@ from airflow.models.base import Base, StringID
from airflow.models.dag_version import DagVersion
from airflow.models.deadline import Deadline, ReferenceModels
from airflow.models.deadline_alert import DeadlineAlert as DeadlineAlertModel
-from airflow.models.dynamic_region import SENTINEL_REGION_ID
+from airflow.models.dynamic_region import (
+ SENTINEL_REGION_ID,
+ DynamicRegion,
+ ProducerContext,
+ public_region_filter,
+)
# Import HITLDetail at runtime so SQLAlchemy can resolve the relationship
from airflow.models.hitl import HITLDetail # noqa: F401
@@ -594,8 +599,11 @@ def _log_state(*, task_instance: TaskInstance, lead_msg: str = "") -> None:
task_instance.run_id,
]
message = "%sMarking task as %s. dag_id=%s, task_id=%s, run_id=%s, "
- if task_instance.map_index >= 0:
- params.append(task_instance.map_index)
+ if task_instance.region_id != SENTINEL_REGION_ID:
+ params.extend((str(task_instance.region_id), task_instance.region_index))
+ message += "region_id=%s, region_index=%d, "
+ elif task_instance.region_index >= 0:
+ params.append(task_instance.region_index)
message += "map_index=%d, "
message += "logical_date=%s, start_date=%s, end_date=%s"
log.info(
@@ -621,6 +629,14 @@ def _date_or_empty(*, task_instance: TaskInstance, attr: str) -> str:
return result.strftime("%Y%m%dT%H%M%S") if result else ""
+def _resolve_stored_index(map_index: int | None, region_index: int | None) -> int:
+ if map_index is not None and region_index is not None and map_index != region_index:
+ raise ValueError("map_index and region_index must identify the same stored index")
+ if region_index is not None:
+ return region_index
+ return map_index if map_index is not None else -1
+
+
def uuid7() -> UUID:
"""Generate a new UUID7."""
return uuid6.uuid7()
@@ -696,8 +712,8 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
task_id: Mapped[str] = mapped_column(StringID(), nullable=False)
dag_id: Mapped[str] = mapped_column(StringID(), nullable=False)
run_id: Mapped[str] = mapped_column(StringID(), nullable=False)
- map_index: Mapped[int] = mapped_column(Integer, nullable=False, server_default="-1")
- region_index = synonym("map_index")
+ region_index: Mapped[int] = mapped_column(Integer, nullable=False, server_default="-1")
+ map_index = synonym("region_index")
region_id: Mapped[UUID] = mapped_column(
Uuid(),
nullable=False,
@@ -797,7 +813,7 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
"task_id",
"run_id",
"region_id",
- "map_index",
+ "region_index",
"working_set",
name="task_instance_current_key",
),
@@ -806,7 +822,7 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
"task_id",
"run_id",
"region_id",
- "map_index",
+ "region_index",
"try_number",
name="task_instance_try_key",
),
@@ -873,14 +889,15 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
dag_version_id: UUID | None,
run_id: str | None = None,
state: str | None = None,
- map_index: int = -1,
+ map_index: int | None = None,
*,
region_id: UUID = SENTINEL_REGION_ID,
+ region_index: int | None = None,
):
super().__init__()
self.dag_id = task.dag_id
self.task_id = task.task_id
- self.map_index = map_index
+ self.region_index = _resolve_stored_index(map_index, region_index)
self.region_id = region_id
if run_id is not None:
self.run_id = run_id
@@ -907,7 +924,7 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
self.dag_version_id = dag_version_id
def __hash__(self):
- return hash((self.task_id, self.dag_id, self.run_id, self.map_index))
+ return hash((self.task_id, self.dag_id, self.run_id, self.region_index))
@property
def stats_tags(self) -> dict[str, str]:
@@ -919,17 +936,19 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
def insert_mapping(
run_id: str,
task: Operator,
- map_index: int,
+ map_index: int | None = None,
*,
dag_version_id: UUID | None,
dag_run: DagRun,
region_id: UUID = SENTINEL_REGION_ID,
+ region_index: int | None = None,
) -> dict[str, Any]:
"""
Insert mapping.
:meta private:
"""
+ region_index = _resolve_stored_index(map_index, region_index)
weight_rule = task.weight_rule
if not hasattr(weight_rule, "get_weight"):
weight_rule = validate_and_load_priority_weight_strategy(weight_rule)
@@ -937,7 +956,7 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
TaskInstance(
task=task,
run_id=run_id,
- map_index=map_index,
+ region_index=region_index,
region_id=region_id,
dag_version_id=dag_version_id,
)
@@ -961,7 +980,7 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
"executor_config": task.executor_config,
"operator": task.task_type,
"custom_operator_name": getattr(task, "operator_name", None),
- "map_index": map_index,
+ "region_index": region_index,
"region_id": region_id,
"_task_display_property_value": task.task_display_name,
"dag_version_id": dag_version_id,
@@ -990,15 +1009,15 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
def rendered_map_index(self) -> str | None:
if self._rendered_map_index is not None:
return self._rendered_map_index
- if self.map_index >= 0:
- return str(self.map_index)
+ if self.region_index >= 0:
+ return str(self.region_index)
return None
@rendered_map_index.expression # type: ignore[no-redef]
def rendered_map_index(cls):
return case(
(cls._rendered_map_index.isnot(None), cls._rendered_map_index),
- (cls.map_index >= 0, cast(cls.map_index, String)),
+ (cls.region_index >= 0, cast(cls.region_index, String)),
else_=None,
)
@@ -1007,7 +1026,7 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
"""Log URL for TaskInstance."""
run_id = quote(self.run_id)
base_url = conf.get("api", "base_url", fallback="http://localhost:8080/")
- map_index = f"/mapped/{self.map_index}" if self.map_index >= 0 else ""
+ map_index = f"/mapped/{self.region_index}" if self.region_index >= 0 else ""
try_number = f"?try_number={self.try_number}" if self.try_number > 0 else ""
_log_uri = f"{base_url.rstrip('/')}/dags/{self.dag_id}/runs/{run_id}/tasks/{self.task_id}{map_index}{try_number}"
@@ -1040,7 +1059,7 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
map_index: int,
lock_for_update: bool = False,
*,
- region_id: UUID = SENTINEL_REGION_ID,
+ region_id: UUID | None = None,
session: Session = NEW_SESSION,
) -> TaskInstance | None:
query = (
@@ -1050,8 +1069,12 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
dag_id=dag_id,
run_id=run_id,
task_id=task_id,
- map_index=map_index,
- region_id=region_id,
+ region_index=map_index,
+ )
+ .where(
+ public_region_filter(TaskInstance, dag_ids=dag_id, run_ids=run_id, task_ids=task_id)
+ if region_id is None
+ else TaskInstance.region_id == region_id
)
)
@@ -1148,7 +1171,7 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
@property
def key(self) -> TaskInstanceKey:
"""Returns a tuple that identifies the task instance uniquely."""
- return TaskInstanceKey(self.dag_id, self.task_id, self.run_id, self.try_number, self.map_index)
+ return TaskInstanceKey(self.dag_id, self.task_id, self.run_id, self.try_number, self.region_index)
def get_dag_id(self) -> str:
"""Return the DAG ID for scheduler routing."""
@@ -1241,18 +1264,18 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
if map_indexes is not None and not map_indexes:
return {}
statement = (
- select(cls.map_index, func.max(cls.try_number))
+ select(cls.region_index, func.max(cls.try_number))
.where(
cls.dag_id == dag_id,
cls.task_id == task_id,
cls.run_id == run_id,
cls.region_id == region_id,
)
- .group_by(cls.map_index)
+ .group_by(cls.region_index)
.execution_options(include_all_attempts=True)
)
if map_indexes is not None:
- statement = statement.where(cls.map_index.in_(map_indexes))
+ statement = statement.where(cls.region_index.in_(map_indexes))
return {map_index: last_try for map_index, last_try in session.execute(statement)}
def prepare_db_for_next_try(self, session: Session) -> TaskInstance:
@@ -1463,9 +1486,12 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
return "<deferred>"
prefix = f"<TaskInstance: {field('dag_id')}.{field('task_id')} {field('run_id')} "
- map_index = field("map_index")
- if map_index != -1:
- prefix += f"map_index={map_index} "
+ region_id = field("region_id")
+ region_index = field("region_index")
+ if region_id != SENTINEL_REGION_ID:
+ prefix += f"region_id={region_id} region_index={region_index} "
+ elif region_index != -1:
+ prefix += f"map_index={region_index} "
return prefix + f"[{field('state')}] ti_id={field('id')}>"
def next_retry_datetime(self):
@@ -2409,14 +2435,24 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
@staticmethod
def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) -> ColumnElement[bool] | None:
"""Return SQLAlchemy filter to query selected task instances."""
- by_region: dict[UUID, list[TaskInstance | TaskInstanceKey]] = defaultdict(list)
+ by_region: dict[UUID | None, list[TaskInstance | TaskInstanceKey]] = defaultdict(list)
for ti in tis:
- by_region[SENTINEL_REGION_ID if isinstance(ti, TaskInstanceKey) else ti.region_id].append(ti)
+ by_region[None if isinstance(ti, TaskInstanceKey) else ti.region_id].append(ti)
if not by_region:
return None
return or_(
*(
- and_(TaskInstance.region_id == region_id, TaskInstance._filter_for_tis(group))
+ and_(
+ public_region_filter(
+ TaskInstance,
+ dag_ids={ti.dag_id for ti in group},
+ run_ids={ti.run_id for ti in group},
+ task_ids={ti.task_id for ti in group},
+ )
+ if region_id is None
+ else TaskInstance.region_id == region_id,
+ TaskInstance._filter_for_tis(group),
+ )
for region_id, group in by_region.items()
)
)
@@ -2446,21 +2482,21 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
return and_(
TaskInstance.dag_id == dag_id,
TaskInstance.run_id == run_id,
- TaskInstance.map_index == map_index,
+ TaskInstance.region_index == map_index,
TaskInstance.task_id.in_(task_ids),
)
if dag_ids == {dag_id} and task_ids == {first_task_id} and map_indices == {map_index}:
return and_(
TaskInstance.dag_id == dag_id,
TaskInstance.run_id.in_(run_ids),
- TaskInstance.map_index == map_index,
+ TaskInstance.region_index == map_index,
TaskInstance.task_id == first_task_id,
)
if dag_ids == {dag_id} and run_ids == {run_id} and task_ids == {first_task_id}:
return and_(
TaskInstance.dag_id == dag_id,
TaskInstance.run_id == run_id,
- TaskInstance.map_index.in_(map_indices),
+ TaskInstance.region_index.in_(map_indices),
TaskInstance.task_id == first_task_id,
)
@@ -2470,8 +2506,9 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
task_id_groups: dict[tuple, dict[Any, list[Any]]] = defaultdict(lambda: defaultdict(list))
map_index_groups: dict[tuple, dict[Any, list[Any]]] = defaultdict(lambda: defaultdict(list))
for t in tis:
- task_id_groups[(t.dag_id, t.run_id)][t.task_id].append(t.map_index)
- map_index_groups[(t.dag_id, t.run_id)][t.map_index].append(t.task_id)
+ index = t.map_index
+ task_id_groups[(t.dag_id, t.run_id)][t.task_id].append(index)
+ map_index_groups[(t.dag_id, t.run_id)][index].append(t.task_id)
# this assumes that most dags have dag_id as the largest grouping, followed by run_id. even
# if its not, this is still a significant optimization over querying for every single tuple key
@@ -2487,7 +2524,7 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
TaskInstance.dag_id == cur_dag_id,
TaskInstance.run_id == cur_run_id,
TaskInstance.task_id == cur_task_id,
- TaskInstance.map_index.in_(cur_map_indices),
+ TaskInstance.region_index.in_(cur_map_indices),
)
)
else:
@@ -2497,7 +2534,7 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
TaskInstance.dag_id == cur_dag_id,
TaskInstance.run_id == cur_run_id,
TaskInstance.task_id.in_(cur_task_ids),
- TaskInstance.map_index == cur_map_index,
+ TaskInstance.region_index == cur_map_index,
)
)
@@ -2519,7 +2556,7 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
if task_id_only:
filters.append(cls.task_id.in_(task_id_only))
if with_map_index:
- filters.append(tuple_(cls.task_id, cls.map_index).in_(with_map_index))
+ filters.append(tuple_(cls.task_id, cls.region_index).in_(with_map_index))
if not filters:
return false()
@@ -2533,16 +2570,18 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
ti_count: int | None,
*,
session: Session,
+ producer_contexts: Mapping[str, ProducerContext] | None = None,
) -> int | range | None:
if TYPE_CHECKING:
assert self.task
return _get_relevant_map_indexes(
run_id=self.run_id,
- map_index=self.map_index,
+ map_index=self.region_index,
ti_count=ti_count,
task=self.task,
relative=upstream,
session=session,
+ producer_contexts=producer_contexts,
)
def clear_db_references(self, session: Session):
@@ -2567,11 +2606,16 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
"""
Create the mapped task instances for mapped task.
+ A mapped task's region is born with its placeholder and keeps it through re-expansion. The
+ placeholder of a pre-region expansion lives in the sentinel region; it moves, with its
+ identity, into a newly minted region.
+
:raise NotMapped: If this task does not need expansion.
:return: The newly created mapped task instances (if any) in ascending
order by map index, and the maximum map index value.
"""
from airflow.models.expandinput import NotFullyPopulated
+ from airflow.models.task_coordinates import TaskCoordinateResolver
from airflow.serialization.definitions.baseoperator import SerializedBaseOperator
from airflow.serialization.definitions.mappedoperator import (
SerializedMappedOperator,
@@ -2587,7 +2631,14 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
)
try:
- total_length: int | None = get_mapped_ti_count(task, run_id, session=session)
+ total_length: int | None = get_mapped_ti_count(
+ task,
+ run_id,
+ session=session,
+ producer_contexts=TaskCoordinateResolver.for_dag(
+ getattr(task, "dag", None), session
+ ).producer_contexts(self),
+ )
except NotFullyPopulated as e:
if not task.dag or not task.dag.partial:
task.log.error(
@@ -2598,28 +2649,39 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
)
total_length = None
+ region_id = self.region_id
state: str | None = None
unmapped_ti: TaskInstance | None = session.scalars(
select(TaskInstance).where(
TaskInstance.dag_id == task.dag_id,
TaskInstance.task_id == task.task_id,
TaskInstance.run_id == run_id,
- TaskInstance.map_index == -1,
- TaskInstance.region_id == self.region_id,
+ TaskInstance.region_index == -1,
+ TaskInstance.region_id == region_id,
or_(TaskInstance.state.in_(State.unfinished), TaskInstance.state.is_(None)),
)
).one_or_none()
+ all_expanded_tis: list[TaskInstance] = []
+
+ if unmapped_ti and total_length is not None and region_id == SENTINEL_REGION_ID:
+ region = DynamicRegion(dag_id=task.dag_id, run_id=run_id, node_id=task.task_id)
+ session.add(region)
+ session.flush()
+ region_id = region.id
+ unmapped_ti.region_id = region_id
+ session.flush()
+ if total_length < 1:
+ all_expanded_tis.append(unmapped_ti)
+
last_tries = TaskInstance.get_last_try_numbers(
dag_id=task.dag_id,
task_id=task.task_id,
run_id=run_id,
- region_id=self.region_id,
+ region_id=region_id,
session=session,
)
- all_expanded_tis: list[TaskInstance] = []
-
if unmapped_ti:
if TYPE_CHECKING:
assert task.dag is None
@@ -2648,7 +2710,8 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
TaskInstance.dag_id == task.dag_id,
TaskInstance.task_id == task.task_id,
TaskInstance.run_id == run_id,
- TaskInstance.map_index == 0,
+ TaskInstance.region_index == 0,
+ TaskInstance.region_id == region_id,
session=session,
)
if not zero_index_ti_exists:
@@ -2659,7 +2722,7 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
(unmapped_ti.max_tries or 0) + promoted_try - unmapped_ti.try_number
)
unmapped_ti.try_number = promoted_try
- unmapped_ti.map_index = 0
+ unmapped_ti.region_index = 0
task.log.debug("Updated in place to become %s", unmapped_ti)
all_expanded_tis.append(unmapped_ti)
# execute hook for task instance map index 0
@@ -2676,11 +2739,11 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
# Only create "missing" ones.
current_max_mapping = (
session.scalar(
- select(func.max(TaskInstance.map_index)).where(
+ select(func.max(TaskInstance.region_index)).where(
TaskInstance.dag_id == task.dag_id,
TaskInstance.task_id == task.task_id,
TaskInstance.run_id == run_id,
- TaskInstance.region_id == self.region_id,
+ TaskInstance.region_id == region_id,
)
)
or 0
@@ -2709,8 +2772,8 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
ti = TaskInstance(
task,
run_id=run_id,
- map_index=index,
- region_id=self.region_id,
+ region_index=index,
+ region_id=region_id,
state=state,
dag_version_id=dag_version_id,
)
@@ -2735,8 +2798,8 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
TaskInstance.dag_id == task.dag_id,
TaskInstance.task_id == task.task_id,
TaskInstance.run_id == run_id,
- TaskInstance.map_index >= total_expanded_ti_count,
- TaskInstance.region_id == self.region_id,
+ TaskInstance.region_index >= total_expanded_ti_count,
+ TaskInstance.region_id == region_id,
)
to_update = session.scalars(with_row_locks(query, of=TaskInstance, session=session, skip_locked=True))
for ti in to_update:
@@ -2846,6 +2909,7 @@ def _get_relevant_map_indexes(
relative: Operator,
ti_count: int | None,
session: Session,
+ producer_contexts: Mapping[str, ProducerContext] | None = None,
) -> int | range | None:
"""
Infer the map indexes of a relative that's "relevant" to this ti.
@@ -2914,7 +2978,9 @@ def _get_relevant_map_indexes(
# should use a "partial" value. Let's break down the mapped ti count
# between the ancestor and further expansion happened inside it.
- ancestor_ti_count = get_mapped_ti_count(common_ancestor, run_id, session=session)
+ ancestor_ti_count = get_mapped_ti_count(
+ common_ancestor, run_id, session=session, producer_contexts=producer_contexts
+ )
ancestor_map_index = map_index * ancestor_ti_count // ti_count
# If the task is NOT further expanded inside the common ancestor, we
@@ -3032,8 +3098,12 @@ class TaskInstanceNote(Base):
def __repr__(self):
prefix = f"<{self.__class__.__name__}: {self.task_instance.dag_id}.{self.task_instance.task_id} {self.task_instance.run_id}"
- if self.task_instance.map_index != -1:
- prefix += f" map_index={self.task_instance.map_index}"
+ if self.task_instance.region_id != SENTINEL_REGION_ID:
+ prefix += (
+ f" region_id={self.task_instance.region_id} region_index={self.task_instance.region_index}"
+ )
+ elif self.task_instance.region_index != -1:
+ prefix += f" map_index={self.task_instance.region_index}"
return prefix + f" TI ID: {self.ti_id}>"
diff --git a/airflow-core/src/airflow/models/xcom.py b/airflow-core/src/airflow/models/xcom.py
index 0b0b940623c..b33575c0caf 100644
--- a/airflow-core/src/airflow/models/xcom.py
+++ b/airflow-core/src/airflow/models/xcom.py
@@ -41,7 +41,6 @@ from sqlalchemy import (
delete,
event,
func,
- or_,
select,
union_all,
)
@@ -52,7 +51,7 @@ from sqlalchemy.sql.visitors import cloned_traverse
from airflow._shared.timezones import timezone
from airflow.models.base import COLLATION_ARGS, ID_LEN, Base, TaskInstanceDependencies
-from airflow.models.dynamic_region import SENTINEL_REGION_ID, DynamicRegion
+from airflow.models.dynamic_region import SENTINEL_REGION_ID, public_region_filter
from airflow.utils.db import LazySelectSequence
from airflow.utils.helpers import is_container
from airflow.utils.json import XComDecoder, XComEncoder
@@ -318,7 +317,7 @@ class _XComOperations:
dag_ids: str | Iterable[str] | None = None,
map_indexes: int | Iterable[int] | None = None,
region_id: UUID | None = SENTINEL_REGION_ID,
- include_node_regions: bool = False,
+ include_node_regions: bool = True,
producer_ids: Select | None = None,
include_prior_dates: bool = False,
limit: int | None = None,
@@ -331,9 +330,9 @@ class _XComOperations:
just want one stored value, use :meth:`get_one` instead.
``region_id`` is the exact producer region (the legacy sentinel by default); pass ``None`` to
- enumerate across regions. ``include_node_regions`` widens it to the top-level region each
- producer task owns, which is where its mapped instances live. ``producer_ids`` replaces the
- coordinate filters with attempts already resolved by
+ enumerate across regions. For the sentinel, ``include_node_regions`` also admits the top-level
+ region each producer task owns, which is where its mapped instances live. ``producer_ids``
+ replaces the coordinate filters with task instances already resolved by
:func:`~airflow.models.dynamic_region.resolve_current_producers`.
Use :func:`xcom_entity` for columns added to the returned statement.
@@ -448,7 +447,12 @@ def _rows():
data = ("key", "value", "timestamp", "dag_result", "mapped_length")
run_join = and_(DagRun.dag_id == TaskInstance.dag_id, DagRun.run_id == TaskInstance.run_id)
context = [
- *(getattr(TaskInstance, name) for name in coordinates),
+ TaskInstance.dag_id,
+ TaskInstance.task_id,
+ TaskInstance.run_id,
+ TaskInstance.region_index.label("map_index"),
+ TaskInstance.region_id,
+ TaskInstance.region_index,
DagRun.id.label("dag_run_id"),
DagRun.logical_date,
DagRun.run_after,
@@ -595,26 +599,24 @@ def select_producers(
task_ids=None,
map_indexes=None,
region_id=SENTINEL_REGION_ID,
- include_node_regions=False,
+ include_node_regions=True,
include_prior_dates=False,
try_number=None,
):
from airflow.models.dagrun import DagRun
from airflow.models.taskinstance import TaskInstance
+ dag_ids = list(dag_ids) if is_container(dag_ids) else dag_ids
+ task_ids = list(task_ids) if is_container(task_ids) else task_ids
query = select(TaskInstance.id)
if region_id is not None:
region_filter = TaskInstance.region_id == region_id
- if include_node_regions:
- region_filter = or_(
- region_filter,
- select(DynamicRegion.id)
- .where(
- DynamicRegion.id == TaskInstance.region_id,
- DynamicRegion.node_id == TaskInstance.task_id,
- DynamicRegion.parent_region_id.is_(None),
- )
- .exists(),
+ if include_node_regions and region_id == SENTINEL_REGION_ID:
+ region_filter = public_region_filter(
+ TaskInstance,
+ dag_ids=dag_ids,
+ run_ids=None if include_prior_dates else run_id,
+ task_ids=task_ids,
)
query = query.where(region_filter)
if try_number is not None:
@@ -626,12 +628,12 @@ def select_producers(
query = query.where(column == value)
if isinstance(map_indexes, range) and map_indexes.step == 1:
query = query.where(
- TaskInstance.map_index >= map_indexes.start, TaskInstance.map_index < map_indexes.stop
+ TaskInstance.region_index >= map_indexes.start, TaskInstance.region_index < map_indexes.stop
)
elif is_container(map_indexes):
- query = query.where(TaskInstance.map_index.in_(map_indexes))
+ query = query.where(TaskInstance.region_index.in_(map_indexes))
elif map_indexes is not None:
- query = query.where(TaskInstance.map_index == map_indexes)
+ query = query.where(TaskInstance.region_index == map_indexes)
if include_prior_dates:
requested_run = aliased(DagRun)
cutoff = (
diff --git a/airflow-core/src/airflow/serialization/definitions/baseoperator.py b/airflow-core/src/airflow/serialization/definitions/baseoperator.py
index 7c8915470e7..8c28b17c377 100644
--- a/airflow-core/src/airflow/serialization/definitions/baseoperator.py
+++ b/airflow-core/src/airflow/serialization/definitions/baseoperator.py
@@ -77,6 +77,7 @@ class SerializedBaseOperator(DAGNode):
dag: SerializedDAG | None = None
depends_on_past: bool = False
do_xcom_push: bool = True
+ returns_dag_result: bool = False
doc: str | None = None
doc_md: str | None = None
doc_json: str | None = None
@@ -181,6 +182,7 @@ class SerializedBaseOperator(DAGNode):
"allow_nested_operators",
"depends_on_past",
"do_xcom_push",
+ "returns_dag_result",
"doc",
"doc_json",
"doc_md",
diff --git a/airflow-core/src/airflow/serialization/definitions/dag.py b/airflow-core/src/airflow/serialization/definitions/dag.py
index 0acb5a2faa3..041cc88a358 100644
--- a/airflow-core/src/airflow/serialization/definitions/dag.py
+++ b/airflow-core/src/airflow/serialization/definitions/dag.py
@@ -48,7 +48,7 @@ from airflow.models.dagbundle import DagBundleModel
from airflow.models.dagrun import DagRun
from airflow.models.deadline import Deadline
from airflow.models.deadline_alert import DeadlineAlert as DeadlineAlertModel
-from airflow.models.taskinstancekey import TaskInstanceKey
+from airflow.models.task_coordinates import public_map_index_expression
from airflow.models.tasklog import LogTemplate
from airflow.serialization.decoders import decode_deadline_alert
from airflow.serialization.definitions.deadline import (
@@ -68,6 +68,7 @@ if TYPE_CHECKING:
import datetime
from collections.abc import Collection, Iterable, Sequence
from typing import Any, Literal
+ from uuid import UUID
from pendulum.tz.timezone import FixedTimezone, Timezone
from pydantic import NonNegativeInt
@@ -152,6 +153,8 @@ class SerializedDAG:
last_loaded: datetime.datetime = attrs.field(init=False)
fileloc: str = attrs.field(init=False)
relative_fileloc: str | None = attrs.field(init=False)
+ # Set when the Dag is deserialized from a stored version, so holders can tell which version they have.
+ dag_version_id: UUID | None = attrs.field(init=False, default=None)
# Determine the relative fileloc based only on the serialize dag.
_processor_dags_folder: str = attrs.field(init=False)
@@ -1008,40 +1011,10 @@ class SerializedDAG:
return altered
- @overload
- def _get_task_instances(
- self,
- *,
- task_ids: Collection[str | tuple[str, int]] | None,
- start_date: datetime.datetime | None,
- end_date: datetime.datetime | None,
- run_id: str | None,
- state: TaskInstanceState | Sequence[TaskInstanceState],
- exclude_task_ids: Collection[str | tuple[str, int]] | None,
- exclude_run_ids: frozenset[str] | None,
- session: Session,
- ) -> Iterable[TaskInstance]: ... # pragma: no cover
-
- @overload
- def _get_task_instances(
- self,
- *,
- task_ids: Collection[str | tuple[str, int]] | None,
- as_pk_tuple: Literal[True],
- start_date: datetime.datetime | None,
- end_date: datetime.datetime | None,
- run_id: str | None,
- state: TaskInstanceState | Sequence[TaskInstanceState],
- exclude_task_ids: Collection[str | tuple[str, int]] | None,
- exclude_run_ids: frozenset[str] | None,
- session: Session,
- ) -> set[TaskInstanceKey]: ... # pragma: no cover
-
def _get_task_instances(
self,
*,
task_ids: Collection[str | tuple[str, int]] | None,
- as_pk_tuple: Literal[True, None] = None,
start_date: datetime.datetime | None,
end_date: datetime.datetime | None,
run_id: str | None,
@@ -1049,29 +1022,10 @@ class SerializedDAG:
exclude_task_ids: Collection[str | tuple[str, int]] | None,
exclude_run_ids: frozenset[str] | None,
session: Session,
- ) -> Iterable[TaskInstance] | set[TaskInstanceKey]:
+ ) -> Iterable[TaskInstance]:
from airflow.models.taskinstance import TaskInstance
- # If we are looking at dependent dags we want to avoid UNION calls
- # in SQL (it doesn't play nice with fields that have no equality operator,
- # like JSON types), we instead build our result set separately.
- #
- # This will be empty if we are only looking at one dag, in which case
- # we can return the filtered TI query object directly.
- result: set[TaskInstanceKey] = set()
-
- # Do we want full objects, or just the primary columns?
- if as_pk_tuple:
- tis_pk = select(
- TaskInstance.dag_id,
- TaskInstance.task_id,
- TaskInstance.run_id,
- TaskInstance.map_index,
- )
- tis_pk = tis_pk.join(TaskInstance.dag_run)
- else:
- tis_full = select(TaskInstance)
- tis_full = tis_full.join(TaskInstance.dag_run)
+ tis_full = select(TaskInstance).join(TaskInstance.dag_run)
# Apply common filters
def apply_filters(query):
@@ -1092,10 +1046,7 @@ class SerializedDAG:
query = query.where(DagRun.logical_date <= end_date)
return query
- if as_pk_tuple:
- tis_pk = apply_filters(tis_pk)
- else:
- tis_full = apply_filters(tis_full)
+ tis_full = apply_filters(tis_full)
def apply_state_filter(query):
if state:
@@ -1120,42 +1071,17 @@ class SerializedDAG:
query = query.where(TaskInstance.run_id.not_in(exclude_run_ids))
return query
- if as_pk_tuple:
- tis_pk = apply_state_filter(tis_pk)
- else:
- tis_full = apply_state_filter(tis_full)
-
- if result or as_pk_tuple:
- # Only execute the `ti` query if we have also collected some other results
- if as_pk_tuple:
- tis_query = session.execute(tis_pk).all()
- result.update(TaskInstanceKey(**cols._mapping) for cols in tis_query)
- else:
- result.update(ti.key for ti in session.scalars(tis_full))
-
- if exclude_task_ids is not None:
- result = {
- task
- for task in result
- if task.task_id not in exclude_task_ids
- and (task.task_id, task.map_index) not in exclude_task_ids
- }
+ tis_full = apply_state_filter(tis_full)
- if as_pk_tuple:
- return result
- if result:
- # We've been asked for objects, lets combine it all back in to a result set
- ti_filters = TaskInstance.filter_for_tis(result)
- if ti_filters is not None:
- tis_final = select(TaskInstance).where(ti_filters)
- return session.scalars(tis_final)
- elif exclude_task_ids is None:
+ if exclude_task_ids is None:
pass # Disable filter if not set.
elif isinstance(next(iter(exclude_task_ids), None), str):
tis_full = tis_full.where(TaskInstance.task_id.notin_(exclude_task_ids))
else:
tis_full = tis_full.where(
- tuple_(TaskInstance.task_id, TaskInstance.map_index).not_in(exclude_task_ids)
+ tuple_(TaskInstance.task_id, public_map_index_expression(TaskInstance)).not_in(
+ exclude_task_ids
+ )
)
return session.scalars(tis_full)
diff --git a/airflow-core/src/airflow/serialization/definitions/mappedoperator.py b/airflow-core/src/airflow/serialization/definitions/mappedoperator.py
index 8762c6734be..9f3029e2657 100644
--- a/airflow-core/src/airflow/serialization/definitions/mappedoperator.py
+++ b/airflow-core/src/airflow/serialization/definitions/mappedoperator.py
@@ -20,7 +20,7 @@ from __future__ import annotations
import functools
import operator
-from collections.abc import Iterable
+from collections.abc import Iterable, Mapping
from typing import TYPE_CHECKING, Any, ClassVar, TypeAlias, TypeGuard, overload
import attrs
@@ -29,6 +29,7 @@ import structlog
from sqlalchemy.orm import Session
from airflow.exceptions import NotMapped
+from airflow.models.dynamic_region import ProducerContext
from airflow.sdk import BaseOperator as TaskSDKBaseOperator
from airflow.sdk.definitions.mappedoperator import MappedOperator as TaskSDKMappedOperator
from airflow.serialization.definitions.baseoperator import DEFAULT_OPERATOR_DEPS, SerializedBaseOperator
@@ -113,6 +114,7 @@ class SerializedMappedOperator(DAGNode):
start_trigger_args: StartTriggerArgs | None = None
start_from_trigger: bool = False
_needs_expansion: bool = True
+ returns_dag_result: bool = attrs.field(init=False, default=False)
doc: str | None = attrs.field(init=False)
doc_json: str | None = attrs.field(init=False)
@@ -361,6 +363,7 @@ class SerializedMappedOperator(DAGNode):
"operator_extra_links",
"params",
"partial_kwargs",
+ "returns_dag_result",
"start_date",
"start_from_trigger",
"start_trigger_args",
@@ -504,7 +507,13 @@ class SerializedMappedOperator(DAGNode):
@functools.singledispatch
-def get_mapped_ti_count(task: DAGNode | TaskSDKDAGNode, run_id: str, *, session: Session) -> int:
+def get_mapped_ti_count(
+ task: DAGNode | TaskSDKDAGNode,
+ run_id: str,
+ *,
+ session: Session,
+ producer_contexts: Mapping[str, ProducerContext] | None = None,
+) -> int:
raise NotImplementedError(f"Not implemented for {type(task)}")
@@ -512,18 +521,30 @@ def get_mapped_ti_count(task: DAGNode | TaskSDKDAGNode, run_id: str, *, session:
# TODO (GH-52141): Rewrite tests so we can drop SDK references at some point.
@get_mapped_ti_count.register(SerializedBaseOperator)
@get_mapped_ti_count.register(TaskSDKBaseOperator)
-def _(task: SerializedBaseOperator | TaskSDKBaseOperator, run_id: str, *, session: Session) -> int:
+def _(
+ task: SerializedBaseOperator | TaskSDKBaseOperator,
+ run_id: str,
+ *,
+ session: Session,
+ producer_contexts: Mapping[str, ProducerContext] | None = None,
+) -> int:
group = task.get_closest_mapped_task_group()
if group is None:
raise NotMapped()
- return get_mapped_ti_count(group, run_id, session=session)
+ return get_mapped_ti_count(group, run_id, session=session, producer_contexts=producer_contexts)
# Still accept TaskSDKMappedOperator because some tests don't go through serialization.
# TODO (GH-52141): Rewrite tests so we can drop SDK references at some point.
@get_mapped_ti_count.register(SerializedMappedOperator)
@get_mapped_ti_count.register(TaskSDKMappedOperator)
-def _(task: SerializedMappedOperator | TaskSDKMappedOperator, run_id: str, *, session: Session) -> int:
+def _(
+ task: SerializedMappedOperator | TaskSDKMappedOperator,
+ run_id: str,
+ *,
+ session: Session,
+ producer_contexts: Mapping[str, ProducerContext] | None = None,
+) -> int:
from airflow.serialization.serialized_objects import BaseSerialization, _ExpandInputRef
exp_input = task._get_specified_expand_input()
@@ -537,20 +558,28 @@ def _(task: SerializedMappedOperator | TaskSDKMappedOperator, run_id: str, *, se
BaseSerialization.deserialize(BaseSerialization.serialize(exp_input.value)),
)
.deref(task.dag)
- .get_total_map_length(run_id, session=session)
+ .get_total_map_length(run_id, session=session, producer_contexts=producer_contexts)
)
else:
- current_count = exp_input.get_total_map_length(run_id, session=session)
+ current_count = exp_input.get_total_map_length(
+ run_id, session=session, producer_contexts=producer_contexts
+ )
group = task.get_closest_mapped_task_group()
if group is None:
return current_count
- parent_count = get_mapped_ti_count(group, run_id, session=session)
+ parent_count = get_mapped_ti_count(group, run_id, session=session, producer_contexts=producer_contexts)
return parent_count * current_count
@get_mapped_ti_count.register
-def _(group: SerializedTaskGroup, run_id: str, *, session: Session) -> int:
+def _(
+ group: SerializedTaskGroup,
+ run_id: str,
+ *,
+ session: Session,
+ producer_contexts: Mapping[str, ProducerContext] | None = None,
+) -> int:
"""
Return the number of instances a task in this group should be mapped to at run time.
@@ -579,7 +608,9 @@ def _(group: SerializedTaskGroup, run_id: str, *, session: Session) -> int:
exp_input.EXPAND_INPUT_TYPE,
BaseSerialization.deserialize(BaseSerialization.serialize(exp_input.value)),
).deref(group.dag)
- yield exp_input.get_total_map_length(run_id, session=session)
+ yield exp_input.get_total_map_length(
+ run_id, session=session, producer_contexts=producer_contexts
+ )
group = group.parent_group
return functools.reduce(operator.mul, iter_mapped_task_group_lengths(group))
diff --git a/airflow-core/src/airflow/serialization/definitions/xcom_arg.py b/airflow-core/src/airflow/serialization/definitions/xcom_arg.py
index 991ab48618b..866027ad802 100644
--- a/airflow-core/src/airflow/serialization/definitions/xcom_arg.py
+++ b/airflow-core/src/airflow/serialization/definitions/xcom_arg.py
@@ -22,14 +22,13 @@ from functools import singledispatch
from typing import TYPE_CHECKING, Any
import attrs
-from sqlalchemy import func, or_, select
+from sqlalchemy import func, select
from sqlalchemy.orm import Session
-from airflow.models.dynamic_region import SENTINEL_REGION_ID, ProducerContext, resolve_current_producers
+from airflow.models.dynamic_region import ProducerContext, resolve_current_producers
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"]
@@ -174,49 +173,25 @@ def _(
task_id = xcom_arg.operator.task_id
mapped = is_mapped(xcom_arg.operator)
- if producer_contexts is not None:
- producers = resolve_current_producers(
- dag_id=dag_id,
- run_id=run_id,
- task_id=task_id,
- is_mapped=mapped,
- context=producer_contexts.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
- read = XComModel.get_many(
- dag_ids=dag_id,
- run_id=run_id,
- task_ids=task_id,
- key=XCOM_RETURN_KEY,
- producer_ids=select(TaskInstance.id).where(TaskInstance.id.in_([ti.id for ti in producers])),
- )
- elif mapped:
- unfinished_ti_exists = exists_query(
- 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,
- # Special NULL treatment is needed because 'state' can be NULL.
- # The "IN" part would produce "NULL NOT IN ..." and eventually
- # "NULl = NULL", which is a big no-no in SQL.
- or_(
- TaskInstance.state.is_(None),
- TaskInstance.state.in_(s.value for s in State.unfinished if s is not None),
- ),
- session=session,
- )
- if unfinished_ti_exists:
- return None # Not all of the expanded tis are done yet.
- read = XComModel.get_many(dag_ids=dag_id, run_id=run_id, task_ids=task_id, key=XCOM_RETURN_KEY)
- else:
- read = XComModel.get_many(
- dag_ids=dag_id, run_id=run_id, task_ids=task_id, map_indexes=-1, key=XCOM_RETURN_KEY
- )
+ producers = resolve_current_producers(
+ dag_id=dag_id,
+ run_id=run_id,
+ task_id=task_id,
+ is_mapped=mapped,
+ 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
+ read = XComModel.get_many(
+ dag_ids=dag_id,
+ run_id=run_id,
+ task_ids=task_id,
+ key=XCOM_RETURN_KEY,
+ producer_ids=select(TaskInstance.id).where(TaskInstance.id.in_([ti.id for ti in producers])),
+ )
entity = xcom_entity(read)
if mapped:
@@ -224,8 +199,6 @@ def _(
read.order_by(None).where(entity.map_index >= 0).with_only_columns(func.count(entity.map_index))
)
# Not xcom_arg.key: the SDK records the length of the whole return value, never per key.
- if producer_contexts is None:
- read = read.where(entity.map_index == -1)
return session.scalar(read.order_by(None).with_only_columns(entity.mapped_length))
diff --git a/airflow-core/src/airflow/serialization/schema.json b/airflow-core/src/airflow/serialization/schema.json
index 4a29deda53c..ed2251ce3da 100644
--- a/airflow-core/src/airflow/serialization/schema.json
+++ b/airflow-core/src/airflow/serialization/schema.json
@@ -331,6 +331,7 @@
"executor": { "type": "string" },
"executor_config": { "$ref": "#/definitions/dict" },
"do_xcom_push": { "type": "boolean", "default": true },
+ "returns_dag_result": { "type": "boolean", "default": false },
"email_on_failure": { "type": "boolean", "default": true },
"email_on_retry": { "type": "boolean", "default": true },
"ui_color": { "type": "string", "default": "#fff" },
diff --git a/airflow-core/src/airflow/state/metastore.py b/airflow-core/src/airflow/state/metastore.py
index 7847a8f0643..33d92d41e78 100644
--- a/airflow-core/src/airflow/state/metastore.py
+++ b/airflow-core/src/airflow/state/metastore.py
@@ -234,7 +234,7 @@ class MetastoreBackend(BaseStoreBackend):
TaskStateStoreModel.run_id == scope.run_id,
TaskStateStoreModel.task_id == scope.task_id,
TaskStateStoreModel.region_id == scope.region_id,
- TaskStateStoreModel.map_index == scope.map_index,
+ TaskStateStoreModel.region_index == scope.region_index,
TaskStateStoreModel.key == key,
)
)
@@ -264,7 +264,7 @@ class MetastoreBackend(BaseStoreBackend):
run_id=scope.run_id,
task_id=scope.task_id,
region_id=scope.region_id,
- map_index=scope.map_index,
+ region_index=scope.region_index,
key=key,
value=value,
updated_at=now,
@@ -273,7 +273,7 @@ class MetastoreBackend(BaseStoreBackend):
stmt = _build_upsert_stmt(
get_dialect_name(session),
TaskStateStoreModel,
- ["dag_run_id", "task_id", "region_id", "map_index", "key"],
+ ["dag_run_id", "task_id", "region_id", "region_index", "key"],
values,
dict(value=value, updated_at=now, expires_at=expires_at),
)
@@ -286,7 +286,7 @@ class MetastoreBackend(BaseStoreBackend):
TaskStateStoreModel.run_id == scope.run_id,
TaskStateStoreModel.task_id == scope.task_id,
TaskStateStoreModel.region_id == scope.region_id,
- TaskStateStoreModel.map_index == scope.map_index,
+ TaskStateStoreModel.region_index == scope.region_index,
TaskStateStoreModel.key == key,
)
)
@@ -301,7 +301,7 @@ class MetastoreBackend(BaseStoreBackend):
TaskStateStoreModel.region_id == scope.region_id,
]
if not all_map_indices:
- conditions.append(TaskStateStoreModel.map_index == scope.map_index)
+ conditions.append(TaskStateStoreModel.region_index == scope.region_index)
session.execute(delete(TaskStateStoreModel).where(*conditions))
def _get_asset_state_store(self, scope: AssetScope, key: str, *, session: Session) -> str | None:
@@ -423,7 +423,7 @@ class MetastoreBackend(BaseStoreBackend):
TaskStateStoreModel.dag_id,
TaskStateStoreModel.run_id,
TaskStateStoreModel.task_id,
- TaskStateStoreModel.map_index,
+ TaskStateStoreModel.region_index,
TaskStateStoreModel.key,
)
with create_session() as session:
@@ -439,7 +439,7 @@ class MetastoreBackend(BaseStoreBackend):
TaskStateStoreModel.run_id == scope.run_id,
TaskStateStoreModel.task_id == scope.task_id,
TaskStateStoreModel.region_id == scope.region_id,
- TaskStateStoreModel.map_index == scope.map_index,
+ TaskStateStoreModel.region_index == scope.region_index,
TaskStateStoreModel.key == key,
)
)
@@ -469,7 +469,7 @@ class MetastoreBackend(BaseStoreBackend):
run_id=scope.run_id,
task_id=scope.task_id,
region_id=scope.region_id,
- map_index=scope.map_index,
+ region_index=scope.region_index,
key=key,
value=value,
updated_at=now,
@@ -479,7 +479,7 @@ class MetastoreBackend(BaseStoreBackend):
stmt = _build_upsert_stmt(
get_dialect_name(session.sync_session),
TaskStateStoreModel,
- ["dag_run_id", "task_id", "region_id", "map_index", "key"],
+ ["dag_run_id", "task_id", "region_id", "region_index", "key"],
values,
dict(value=value, updated_at=now, expires_at=expires_at),
)
@@ -492,7 +492,7 @@ class MetastoreBackend(BaseStoreBackend):
TaskStateStoreModel.run_id == scope.run_id,
TaskStateStoreModel.task_id == scope.task_id,
TaskStateStoreModel.region_id == scope.region_id,
- TaskStateStoreModel.map_index == scope.map_index,
+ TaskStateStoreModel.region_index == scope.region_index,
TaskStateStoreModel.key == key,
)
)
@@ -507,7 +507,7 @@ class MetastoreBackend(BaseStoreBackend):
TaskStateStoreModel.region_id == scope.region_id,
]
if not all_map_indices:
- conditions.append(TaskStateStoreModel.map_index == scope.map_index)
+ conditions.append(TaskStateStoreModel.region_index == scope.region_index)
await session.execute(delete(TaskStateStoreModel).where(*conditions))
async def _aget_asset_state_store(
diff --git a/airflow-core/src/airflow/ti_deps/dep_context.py b/airflow-core/src/airflow/ti_deps/dep_context.py
index 3f0d6e8b46f..70a197f9128 100644
--- a/airflow-core/src/airflow/ti_deps/dep_context.py
+++ b/airflow-core/src/airflow/ti_deps/dep_context.py
@@ -19,8 +19,10 @@ from __future__ import annotations
import contextlib
from typing import TYPE_CHECKING
+from uuid import UUID
import attr
+from sqlalchemy import select
from airflow.exceptions import TaskNotFound
from airflow.utils.state import State
@@ -29,6 +31,7 @@ if TYPE_CHECKING:
from sqlalchemy.orm.session import Session
from airflow.models.dagrun import DagRun
+ from airflow.models.task_coordinates import TaskCoordinateResolver
from airflow.models.taskinstance import TaskInstance
@@ -85,12 +88,12 @@ class DepContext:
have_changed_ti_states: bool = False
"""Have any of the TIs state's been changed as a result of evaluating dependencies"""
- upstream_task_id_counts: dict[tuple[str, str, frozenset[str]], list[tuple[str, int]]] = attr.ib(
+ upstream_task_id_counts: dict[tuple[str, str, frozenset[str | UUID]], list[tuple[str, int]]] = attr.ib(
factory=dict, repr=False
)
"""
Per-pass memo of the trigger-rule upstream task-instance counts, keyed by
- ``(dag_id, run_id, frozenset of direct-upstream task_ids)``.
+ ``(dag_id, run_id, frozenset of direct-upstream task ids or region ids)``.
Only populated for the "simple" case where the count-query predicate is exactly
``task_id IN (upstream_ids)`` and is therefore identical for every downstream sharing the same
@@ -106,6 +109,73 @@ class DepContext:
every ``UP_FOR_RESCHEDULE`` task instance. With ``init=False`` those instances would each get a
fresh empty dict, so they would neither read the memo nor warm it for anything else.
"""
+ regional_runs: dict[tuple[str, str], bool] = attr.ib(factory=dict, repr=False)
+ producer_tis: dict[tuple[str, str, UUID | None, UUID, int, str], tuple[TaskInstance, ...]] = attr.ib(
+ factory=dict, repr=False
+ )
+ coordinate_resolvers: dict[Session, TaskCoordinateResolver] = attr.ib(factory=dict, repr=False)
+ dynamic_dags: dict[int, bool] = attr.ib(factory=dict, repr=False)
+
+ def has_regions(self, ti: TaskInstance, *, session: Session) -> bool:
+ from airflow.models.dynamic_region import DynamicRegion
+
+ key = ti.dag_id, ti.run_id
+ if key not in self.regional_runs:
+ self.regional_runs[key] = self._dag_can_have_regions(ti) and bool(
+ session.scalar(
+ select(
+ select(DynamicRegion.id)
+ .where(
+ DynamicRegion.dag_id == ti.dag_id,
+ DynamicRegion.run_id == ti.run_id,
+ )
+ .exists()
+ )
+ )
+ )
+ return self.regional_runs[key]
+
+ def _dag_can_have_regions(self, ti: TaskInstance) -> bool:
+ from airflow.models.task_coordinates import enclosing_loop
+
+ dag = getattr(ti.task, "dag", None)
+ if dag is None:
+ return True
+ if id(dag) not in self.dynamic_dags:
+ self.dynamic_dags[id(dag)] = any(
+ task.get_needs_expansion() or enclosing_loop(task) is not None
+ for task in dag.task_dict.values()
+ )
+ return self.dynamic_dags[id(dag)]
+
+ def coordinate_resolver(self, ti: TaskInstance, *, session: Session) -> TaskCoordinateResolver:
+ from airflow.models.task_coordinates import TaskCoordinateResolver
+
+ if session not in self.coordinate_resolvers:
+ self.coordinate_resolvers[session] = TaskCoordinateResolver.for_dag(None, session)
+ resolver = self.coordinate_resolvers[session]
+ resolver.adopt_dag(getattr(ti.task, "dag", None))
+ return resolver
+
+ def upstream_tis(self, ti: TaskInstance, task_id: str, *, session: Session) -> tuple[TaskInstance, ...]:
+ resolver = self.coordinate_resolver(ti, session=session)
+ task = resolver.get_task(ti.dag_id, ti.run_id, ti.task_id, dag_version_id=ti.dag_version_id)
+ key = (
+ ti.dag_id,
+ ti.run_id,
+ ti.dag_version_id,
+ ti.region_id,
+ -1 if task.get_needs_expansion() else ti.region_index,
+ 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,
+ )
+ return self.producer_tis[key]
def ensure_finished_tis(self, dag_run: DagRun, session: Session) -> list[TaskInstance]:
"""
@@ -134,3 +204,5 @@ class DepContext:
later in the same pass recomputes the count instead of reading a stale one.
"""
self.upstream_task_id_counts.clear()
+ self.producer_tis.clear()
+ self.regional_runs.clear()
diff --git a/airflow-core/src/airflow/ti_deps/deps/mapped_task_expanded.py b/airflow-core/src/airflow/ti_deps/deps/mapped_task_expanded.py
index f3437376162..8871b50b63e 100644
--- a/airflow-core/src/airflow/ti_deps/deps/mapped_task_expanded.py
+++ b/airflow-core/src/airflow/ti_deps/deps/mapped_task_expanded.py
@@ -30,7 +30,7 @@ class MappedTaskIsExpanded(BaseTIDep):
def _get_dep_statuses(self, ti, dep_context, *, session):
if dep_context.ignore_unmapped_tasks:
return
- if ti.map_index == -1:
+ if ti.region_index == -1:
yield self._failing_status(reason="The task has yet to be mapped!")
return
yield self._passing_status(reason="The task has been mapped")
diff --git a/airflow-core/src/airflow/ti_deps/deps/mapped_task_upstream_dep.py b/airflow-core/src/airflow/ti_deps/deps/mapped_task_upstream_dep.py
index e1957b9a342..198e493bf9f 100644
--- a/airflow-core/src/airflow/ti_deps/deps/mapped_task_upstream_dep.py
+++ b/airflow-core/src/airflow/ti_deps/deps/mapped_task_upstream_dep.py
@@ -27,6 +27,8 @@ from airflow.ti_deps.deps.base_ti_dep import BaseTIDep
from airflow.utils.state import State, TaskInstanceState
if TYPE_CHECKING:
+ from collections.abc import Sequence
+
from sqlalchemy.orm import Session
from airflow.models.taskinstance import TaskInstance
@@ -69,18 +71,29 @@ class MappedTaskUpstreamDep(BaseTIDep):
# only interested in it if it hasn't been expanded yet, i.e., we filter by map_index=-1. This is
# because if it has been expanded, it did not fail and was not skipped outright which is all we need
# to know for the purposes of this check.
- mapped_dependency_tis = (
- session.scalars(
- select(TaskInstance).where(
- TaskInstance.task_id.in_(operator.task_id for operator in mapped_dependencies),
- TaskInstance.dag_id == ti.dag_id,
- TaskInstance.run_id == ti.run_id,
- TaskInstance.map_index == -1,
- )
- ).all()
- if mapped_dependencies
- else []
- )
+ dependency_ids = {operator.task_id for operator in mapped_dependencies}
+ mapped_dependency_tis: Sequence[TaskInstance]
+ if dep_context.has_regions(ti, session=session):
+ resolver = dep_context.coordinate_resolver(ti, session=session)
+ mapped_dependency_tis = [
+ upstream
+ for task_id in dependency_ids
+ for upstream in dep_context.upstream_tis(ti, task_id, session=session)
+ if resolver.public_map_index(upstream) < 0
+ ]
+ else:
+ mapped_dependency_tis = (
+ session.scalars(
+ select(TaskInstance).where(
+ TaskInstance.task_id.in_(dependency_ids),
+ TaskInstance.dag_id == ti.dag_id,
+ TaskInstance.run_id == ti.run_id,
+ TaskInstance.region_index == -1,
+ )
+ ).all()
+ if dependency_ids
+ else []
+ )
if not mapped_dependency_tis:
yield self._passing_status(reason="There are no (unexpanded) mapped dependencies!")
return
diff --git a/airflow-core/src/airflow/ti_deps/deps/not_previously_skipped_dep.py b/airflow-core/src/airflow/ti_deps/deps/not_previously_skipped_dep.py
index 94417b252fc..8c4449d2731 100644
--- a/airflow-core/src/airflow/ti_deps/deps/not_previously_skipped_dep.py
+++ b/airflow-core/src/airflow/ti_deps/deps/not_previously_skipped_dep.py
@@ -63,7 +63,13 @@ class NotPreviouslySkippedDep(BaseTIDep):
# query with -1 instead of the child's map_index. A parent inside a
# mapped task group is expanded like a mapped task, so it writes XCom
# per map index even though ``is_mapped`` is False.
- xcom_map_index = ti.map_index if parent.get_needs_expansion() else -1
+ xcom_map_index = -1
+ if parent.get_needs_expansion():
+ xcom_map_index = (
+ dep_context.coordinate_resolver(ti, session=session).public_map_index(ti)
+ if dep_context.has_regions(ti, session=session)
+ else ti.region_index
+ )
prev_result = ti.xcom_pull(
task_ids=parent.task_id,
key=XCOM_SKIPMIXIN_KEY,
diff --git a/airflow-core/src/airflow/ti_deps/deps/ready_to_reschedule.py b/airflow-core/src/airflow/ti_deps/deps/ready_to_reschedule.py
index 384ee0bff48..bfe14eb55c8 100644
--- a/airflow-core/src/airflow/ti_deps/deps/ready_to_reschedule.py
+++ b/airflow-core/src/airflow/ti_deps/deps/ready_to_reschedule.py
@@ -63,7 +63,12 @@ class ReadyToRescheduleDep(BaseTIDep):
# UP_FOR_RESCHEDULE we *always* check TaskReschedule regardless of operator type
# (e.g. startup/DAG-load rescheduling). For NONE-state tasks, only reschedule-mode
# sensors (and mapped tasks whose reschedule attr is unknown) need the DB query.
- if ti.state is None and ti.map_index < 0 and not getattr(ti.task, "reschedule", False):
+ map_index = (
+ dep_context.coordinate_resolver(ti, session=session).public_map_index(ti)
+ if dep_context.has_regions(ti, session=session)
+ else ti.region_index
+ )
+ if ti.state is None and map_index < 0 and not getattr(ti.task, "reschedule", False):
yield self._passing_status(reason="Task is not in reschedule mode.")
return
@@ -87,7 +92,7 @@ class ReadyToRescheduleDep(BaseTIDep):
if not next_reschedule_date:
# Because mapped sensors don't have the reschedule property, here's the last resort
# and we need a slightly different passing reason
- if ti.map_index >= 0:
+ if map_index >= 0:
yield self._passing_status(reason="The task is mapped and not in reschedule mode")
return
yield self._passing_status(reason="There is no reschedule request for this task instance.")
diff --git a/airflow-core/src/airflow/ti_deps/deps/trigger_rule_dep.py b/airflow-core/src/airflow/ti_deps/deps/trigger_rule_dep.py
index f49b71f7ca9..3b52d4db366 100644
--- a/airflow-core/src/airflow/ti_deps/deps/trigger_rule_dep.py
+++ b/airflow-core/src/airflow/ti_deps/deps/trigger_rule_dep.py
@@ -22,6 +22,7 @@ import functools
from collections import Counter
from collections.abc import Iterator, KeysView, Mapping
from typing import TYPE_CHECKING, NamedTuple
+from uuid import UUID
from sqlalchemy import and_, func, or_, select
@@ -138,6 +139,19 @@ class TriggerRuleDep(BaseTIDep):
if TYPE_CHECKING:
assert task
+ regional = dep_context.has_regions(ti, session=session)
+ map_index = (
+ dep_context.coordinate_resolver(ti, session=session).public_map_index(ti)
+ if regional
+ else ti.region_index
+ )
+
+ @functools.lru_cache
+ def _selected_upstream_ids(upstream_id: str) -> frozenset[UUID]:
+ return frozenset(
+ upstream.id for upstream in dep_context.upstream_tis(ti, upstream_id, session=session)
+ )
+
@functools.lru_cache
def _get_expanded_ti_count() -> int:
"""
@@ -148,7 +162,12 @@ class TriggerRuleDep(BaseTIDep):
"""
from airflow.serialization.definitions.mappedoperator import get_mapped_ti_count
- return get_mapped_ti_count(task, ti.run_id, session=session)
+ contexts = (
+ dep_context.coordinate_resolver(ti, session=session).producer_contexts(ti)
+ if regional
+ else None
+ )
+ return get_mapped_ti_count(task, ti.run_id, session=session, producer_contexts=contexts)
def _iter_expansion_dependencies(task_group: SerializedTaskGroup | None) -> Iterator[str]:
if is_mapped(task):
@@ -184,7 +203,7 @@ class TriggerRuleDep(BaseTIDep):
# expanded instance (see #50210). The task may be nested in plain task
# groups inside the mapped one (see #39801), so check the closest mapped
# ancestor rather than only the immediate parent group.
- if ti.map_index < 0 and task.get_closest_mapped_task_group() is not None:
+ if map_index < 0 and task.get_closest_mapped_task_group() is not None:
is_fast_triggered = task.trigger_rule in (
TR.ONE_SUCCESS,
TR.ONE_FAILED,
@@ -204,6 +223,11 @@ class TriggerRuleDep(BaseTIDep):
upstream=task.dag.task_dict[upstream_id],
ti_count=expanded_ti_count,
session=session,
+ producer_contexts=(
+ dep_context.coordinate_resolver(ti, session=session).producer_contexts(ti)
+ if regional
+ else None
+ ),
)
def _is_relevant_upstream(upstream: TaskInstance, relevant_ids: set[str] | KeysView[str]) -> bool:
@@ -219,22 +243,29 @@ class TriggerRuleDep(BaseTIDep):
# Not actually an upstream task.
if upstream.task_id not in relevant_ids:
return False
+ if regional and upstream.id not in _selected_upstream_ids(upstream.task_id):
+ return False
# The current task is not in a mapped task group. All tis from an
# upstream task are relevant.
if task.get_closest_mapped_task_group() is None:
return True
# The upstream ti is not expanded. The upstream may be mapped or
# not, but the ti is relevant either way.
- if upstream.map_index < 0:
+ upstream_map_index = (
+ dep_context.coordinate_resolver(ti, session=session).public_map_index(upstream)
+ if regional
+ else upstream.region_index
+ )
+ if upstream_map_index < 0:
return True
# Now we need to perform fine-grained check on whether this specific
# upstream ti's map index is relevant.
relevant = _get_relevant_upstream_map_indexes(upstream_id=upstream.task_id)
if relevant is None:
return True
- if relevant == upstream.map_index:
+ if relevant == upstream_map_index:
return True
- if isinstance(relevant, collections.abc.Container) and upstream.map_index in relevant:
+ if isinstance(relevant, collections.abc.Container) and upstream_map_index in relevant:
return True
return False
@@ -243,6 +274,16 @@ class TriggerRuleDep(BaseTIDep):
# it depends on all upstream task instances.
from airflow.models.taskinstance import TaskInstance
+ if regional:
+ for upstream_id in relevant_tasks:
+ selected = dep_context.upstream_tis(ti, upstream_id, session=session)
+ yield TaskInstance.id.in_(
+ upstream.id
+ for upstream in selected
+ if _is_relevant_upstream(upstream, relevant_tasks.keys())
+ )
+ return
+
if task.get_closest_mapped_task_group() is None:
yield TaskInstance.task_id.in_(relevant_tasks.keys())
return
@@ -257,17 +298,19 @@ class TriggerRuleDep(BaseTIDep):
# of this upstream task. Since the upstream may not have been
# expanded at this point, we also depend on the non-expanded ti
# to ensure at least one ti is included for the task.
- yield and_(TaskInstance.task_id == upstream_id, TaskInstance.map_index < 0)
+ yield and_(TaskInstance.task_id == upstream_id, TaskInstance.region_index < 0)
if isinstance(map_indexes, range) and map_indexes.step == 1:
yield and_(
TaskInstance.task_id == upstream_id,
- TaskInstance.map_index >= map_indexes.start,
- TaskInstance.map_index < map_indexes.stop,
+ TaskInstance.region_index >= map_indexes.start,
+ TaskInstance.region_index < map_indexes.stop,
)
elif isinstance(map_indexes, collections.abc.Container):
- yield and_(TaskInstance.task_id == upstream_id, TaskInstance.map_index.in_(map_indexes))
+ yield and_(
+ TaskInstance.task_id == upstream_id, TaskInstance.region_index.in_(map_indexes)
+ )
else:
- yield and_(TaskInstance.task_id == upstream_id, TaskInstance.map_index == map_indexes)
+ yield and_(TaskInstance.task_id == upstream_id, TaskInstance.region_index == map_indexes)
def _evaluate_setup_constraint(
*, relevant_setups: Mapping[str, Operator]
@@ -318,8 +361,8 @@ class TriggerRuleDep(BaseTIDep):
new_state = TaskInstanceState.UPSTREAM_FAILED
elif skipped:
new_state = TaskInstanceState.SKIPPED
- elif removed and success and ti.map_index > -1:
- if ti.map_index >= success:
+ elif removed and success and map_index > -1:
+ if map_index >= success:
new_state = TaskInstanceState.REMOVED
if new_state is not None:
@@ -344,7 +387,7 @@ class TriggerRuleDep(BaseTIDep):
dep_context.have_changed_ti_states = True
non_successes = upstream - success
- if ti.map_index > -1:
+ if map_index > -1:
non_successes -= removed
if non_successes > 0:
yield (
@@ -388,22 +431,19 @@ class TriggerRuleDep(BaseTIDep):
upstream = len(upstream_tasks)
upstream_setup = sum(1 for x in upstream_tasks.values() if x.is_setup)
else:
- # In the simple case, `_iter_upstream_conditions` emits exactly
- # `task_id IN (upstream_task_ids)` (the matching `get_closest_mapped_task_group()
- # is None` branch). That predicate, and therefore the resulting counts, are
- # identical for every downstream that shares the same set of direct upstreams, so
- # we memoize them on the DepContext and run the query once per pass instead of
- # once per downstream. The mapped-task-group case uses per-ti map-index predicates
- # and is left un-memoized. The cache lives for one scheduling pass, and
- # DagRun._get_ready_tis invalidates it for both of the things that change a mapped
- # task's instance count mid-pass: expanding an unexpanded task, and
- # _revise_map_indexes_if_mapped growing an already-expanded one. State-only changes
- # (a task finishing, instances marked REMOVED) leave the row count alone, so they
- # need no invalidation.
- cache_key: tuple[str, str, frozenset[str]] | None = None
+ # Regional selections share counts only for identical producer UUID sets;
+ # legacy runs can share by task IDs. Expansion invalidates both forms.
+ cache_key: tuple[str, str, frozenset[str | UUID]] | None = None
task_id_counts: list[tuple[str, int]] | None = None
if task.get_closest_mapped_task_group() is None:
- cache_key = (ti.dag_id, ti.run_id, frozenset(upstream_tasks))
+ scope: frozenset[str | UUID] = (
+ frozenset(
+ ti_id for task_id in upstream_tasks for ti_id in _selected_upstream_ids(task_id)
+ )
+ if regional
+ else frozenset(upstream_tasks)
+ )
+ cache_key = (ti.dag_id, ti.run_id, scope)
task_id_counts = dep_context.upstream_task_id_counts.get(cache_key)
if task_id_counts is None:
task_id_counts = [
@@ -432,8 +472,8 @@ class TriggerRuleDep(BaseTIDep):
new_state = TaskInstanceState.UPSTREAM_FAILED
elif skipped:
new_state = TaskInstanceState.SKIPPED
- elif removed and success and ti.map_index > -1:
- if ti.map_index >= success:
+ elif removed and success and map_index > -1:
+ if map_index >= success:
new_state = TaskInstanceState.REMOVED
elif trigger_rule == TR.ALL_FAILED:
if success or skipped:
diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
index 12161426709..7cd9702385b 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
@@ -1383,6 +1383,7 @@ export class DagRunService {
401: 'Unauthorized',
403: 'Forbidden',
404: 'Not Found',
+ 409: 'Conflict',
422: 'Validation Error'
}
});
@@ -1525,6 +1526,7 @@ export class ExperimentalService {
401: 'Unauthorized',
403: 'Forbidden',
404: 'Not Found',
+ 409: 'Conflict',
422: 'Validation Error'
}
});
diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
index ddadbd19f28..ae5fc838cd2 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
@@ -6468,6 +6468,10 @@ export type $OpenApiTs = {
* Not Found
*/
404: HTTPExceptionResponse;
+ /**
+ * Conflict
+ */
+ 409: HTTPExceptionResponse;
/**
* Validation Error
*/
diff --git a/airflow-core/src/airflow/utils/db.py b/airflow-core/src/airflow/utils/db.py
index 6d1b912abd7..0f92dddde2b 100644
--- a/airflow-core/src/airflow/utils/db.py
+++ b/airflow-core/src/airflow/utils/db.py
@@ -117,7 +117,7 @@ _REVISION_HEADS_MAP: dict[str, str] = {
"3.1.8": "509b94a1042d",
"3.2.0": "1d6611b6ab7c",
"3.3.0": "d2f4e1b3c5a7",
- "3.4.0": "54a27b6f9d01",
+ "3.4.0": "7f8c9a2d410e",
}
# Prefix used to identify tables holding data moved during migration.
diff --git a/airflow-core/src/airflow/utils/log/task_log_address.py b/airflow-core/src/airflow/utils/log/task_log_address.py
index 37bb1d0e784..bd82b0497ec 100644
--- a/airflow-core/src/airflow/utils/log/task_log_address.py
+++ b/airflow-core/src/airflow/utils/log/task_log_address.py
@@ -28,9 +28,9 @@ import jinja2
from jinja2.meta import find_undeclared_variables
from sqlalchemy import select, tuple_
+from airflow.models.dagbag import DBDagBag
from airflow.models.dagrun import DagRun
from airflow.models.dynamic_region import SENTINEL_REGION_ID, DynamicRegion
-from airflow.models.serialized_dag import SerializedDagModel
from airflow.models.tasklog import LogTemplate
from airflow.serialization.definitions.taskgroup import SerializedLoopTaskGroup
from airflow.utils.helpers import render_template
@@ -175,15 +175,22 @@ def prepare_task_log_contexts(
if ancestor_id is not None and ancestor_id not in regions
}
version_ids = {
- ti.dag_version_id or runs[ti.dag_id, ti.run_id].created_dag_version_id for ti in regional
+ version_id
+ for ti in regional
+ if (version_id := ti.dag_version_id or runs[ti.dag_id, ti.run_id].created_dag_version_id)
}
- for row in session.scalars(
- select(SerializedDagModel).where(SerializedDagModel.dag_version_id.in_(version_ids))
- ):
- row.load_op_links = False
- dag = row.dag
- dags[row.dag_version_id] = dag
- node_kinds[row.dag_version_id] = {
+ attached = {
+ dag.dag_version_id: dag
+ for ti in regional
+ if (dag := getattr(ti.task, "dag", None)) is not None and dag.dag_version_id is not None
+ }
+ dag_bag = DBDagBag(load_op_links=False)
+ for version_id in version_ids:
+ dag = attached.get(version_id) or dag_bag.get_dag(version_id, session=session)
+ if dag is None:
+ raise ValueError(f"Pinned Dag version {version_id} not found")
+ dags[version_id] = dag
+ node_kinds[version_id] = {
group_id: "loop"
for group_id, group in dag.task_group.get_task_group_dict().items()
if group_id is not None and isinstance(group, SerializedLoopTaskGroup)
@@ -198,7 +205,7 @@ def prepare_task_log_contexts(
"Please make sure you set up the metadatabase correctly."
)
position = ""
- map_index = ti.map_index
+ map_index = ti.region_index
if ti.region_id != SENTINEL_REGION_ID:
version_id = ti.dag_version_id or run.created_dag_version_id
if version_id is None:
diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py
index 76f7ed723aa..68af746f035 100644
--- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py
+++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py
@@ -24,7 +24,8 @@ from unittest import mock
import pytest
import time_machine
-from sqlalchemy import delete, func, select, update
+import uuid6
+from sqlalchemy import delete, func, insert, select, update
from airflow import plugins_manager
from airflow._shared.module_loading import qualname
@@ -33,15 +34,18 @@ from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessE
from airflow.api_fastapi.auth.managers.simple.user import SimpleAuthManagerUser
from airflow.api_fastapi.common.dagbag import resolve_run_on_latest_version
from airflow.api_fastapi.core_api.datamodels.dag_versions import DagVersionResponse
+from airflow.api_fastapi.core_api.services.public.dag_run import DagRunWaiter
from airflow.exceptions import ParamValidationError
from airflow.models import DagModel, DagRun, DagTag, Log
from airflow.models.asset import AssetEvent, AssetModel
from airflow.models.dagbundle import DagBundleModel
+from airflow.models.dynamic_region import DynamicRegion
from airflow.models.taskinstance import TaskInstance
from airflow.models.team import Team
from airflow.models.xcom import XComModel
from airflow.providers.standard.operators.empty import EmptyOperator
-from airflow.sdk import Asset, Param, result, task
+from airflow.sdk import Asset, Param, result, task, task_group
+from airflow.sdk.definitions._internal.loop import create_loop
from airflow.timetables.interval import CronDataIntervalTimetable
from airflow.timetables.simple import PartitionedAssetTimetable, PartitionedAtRuntime
from airflow.timetables.trigger import CronPartitionTimetable
@@ -4479,6 +4483,147 @@ class TestResolveRunOnLatestVersion:
class TestWaitDagRun:
+ @pytest.mark.parametrize("state", [DagRunState.RUNNING, DagRunState.SUCCESS])
+ @pytest.mark.parametrize("explicit", [False, True])
+ @pytest.mark.parametrize("mapped", [False, True])
+ @pytest.mark.parametrize("has_live_tis", [False, True])
+ def test_rejects_loop_results_before_streaming(
+ self, test_client, dag_maker, session, state, explicit, mapped, has_live_tis
+ ):
+ with dag_maker("wait_loop_result", session=session) as dag:
+
+ @task
+ def value(v=1):
+ return v
+
+ @task_group
+ def body():
+ output = value.expand(v=[1, 2]) if mapped else value()
+ if not explicit:
+ dag.add_result(output)
+
+ create_loop(body, max_iterations=2)
+
+ run = dag_maker.create_dagrun(state=state, session=session)
+ if not has_live_tis:
+ session.execute(delete(TaskInstance).where(TaskInstance.dag_id == run.dag_id))
+ session.commit()
+ with mock.patch.object(DagRunWaiter, "wait", autospec=True, return_value=iter(())) as stream:
+ response = test_client.get(
+ f"/dags/{run.dag_id}/dagRuns/{run.run_id}/wait",
+ params={"interval": 1, **({"result": "body.value"} if explicit else {})},
+ )
+
+ assert response.status_code == 409
+ assert "outside the loop" in response.json()["detail"]
+ stream.assert_not_called()
+
+ @pytest.mark.parametrize("explicit", [False, True])
+ def test_loop_result_preflight_uses_pinned_definition(self, test_client, dag_maker, session, explicit):
+ with dag_maker("wait_pinned_loop", session=session) as dag:
+
+ @task_group
+ def body():
+ output = EmptyOperator(task_id="value")
+ dag.add_result(output.output)
+
+ create_loop(body, max_iterations=2)
+ run = dag_maker.create_dagrun(state=DagRunState.SUCCESS, session=session)
+ session.commit()
+
+ with dag_maker("wait_pinned_loop", session=session):
+ EmptyOperator(task_id="body.value")
+
+ response = test_client.get(
+ f"/dags/{run.dag_id}/dagRuns/{run.run_id}/wait",
+ params={"interval": 1, **({"result": "body.value"} if explicit else {})},
+ )
+ assert response.status_code == 409
+
+ def test_loop_run_can_return_an_outside_result(self, test_client, dag_maker, session):
+ with dag_maker("wait_outside_result", session=session) as dag:
+
+ @task_group
+ def body():
+ EmptyOperator(task_id="value")
+
+ loop = create_loop(body, max_iterations=2)
+ outside = EmptyOperator(task_id="outside")
+ loop >> outside
+ dag.add_result(outside.output)
+ run = dag_maker.create_dagrun(state=DagRunState.SUCCESS, session=session)
+ XComModel.set(
+ key="return_value",
+ value=42,
+ dag_id=run.dag_id,
+ run_id=run.run_id,
+ task_id="outside",
+ dag_result=True,
+ session=session,
+ )
+ session.commit()
+ response = test_client.get(
+ f"/dags/{run.dag_id}/dagRuns/{run.run_id}/wait",
+ params={"interval": 1},
+ )
+ assert response.status_code == 200
+ assert response.json() == {"state": "success", "results": {"outside": "42"}}
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize("mapped", [False, True])
+ async def test_result_rejects_ambiguous_loop_producers(self, dag_maker, session, mapped):
+ with dag_maker("ambiguous_loop_result") as dag:
+
+ @result
+ @task
+ def value(v=0):
+ return v
+
+ if mapped:
+ value.expand(v=[0])
+ else:
+ value()
+ run = dag_maker.create_dagrun(state=DagRunState.SUCCESS)
+ loop = DynamicRegion(dag_id=run.dag_id, run_id=run.run_id, node_id="loop")
+ session.add(loop)
+ session.flush()
+ original = run.task_instances[0]
+ for index in range(2):
+ region = loop
+ if mapped:
+ region = DynamicRegion(
+ dag_id=run.dag_id,
+ run_id=run.run_id,
+ node_id="value",
+ parent_region_id=loop.id,
+ parent_region_index=index,
+ )
+ session.add(region)
+ session.flush()
+ ti = (
+ original
+ if index == 0
+ else TaskInstance(
+ dag.get_task("value"), run_id=run.run_id, dag_version_id=original.dag_version_id
+ )
+ )
+ ti.region_id, ti.region_index = region.id, 0 if mapped else index
+ session.add(ti)
+ session.flush()
+ XComModel.set_for_attempt(
+ task_instance_id=ti.id,
+ key="return_value",
+ value=index,
+ serialize=False,
+ dag_result=True,
+ session=session,
+ )
+ session.commit()
+
+ waiter = DagRunWaiter(run.dag_id, run.run_id, 1, None)
+ with pytest.raises(ValueError, match="Ambiguous current producers for task 'value'"):
+ await waiter._serialize_xcoms()
+
def test_should_respond_401(self, unauthenticated_test_client):
response = unauthenticated_test_client.get(
f"/dags/{DAG1_ID}/dagRuns/{DAG1_RUN1_ID}/wait",
@@ -4607,7 +4752,8 @@ class TestWaitDagRun:
data = response.json()
assert data == {"state": DagRunState.SUCCESS}
- def test_collect_mapped_task_dag_result(self, test_client, dag_maker, session):
+ @pytest.mark.parametrize("explicit", [False, True])
+ def test_collect_mapped_task_dag_result(self, test_client, dag_maker, session, explicit):
"""XComs from a mapped @result task are aggregated into a list ordered by map_index."""
with dag_maker("dag_mapped_result"):
@@ -4629,11 +4775,33 @@ class TestWaitDagRun:
)
for ti in dag_run.task_instances:
run_task_instance(ti, mapped_op, session=session)
+ first = session.scalars(
+ select(TaskInstance).where(
+ TaskInstance.dag_id == dag_run.dag_id,
+ TaskInstance.run_id == dag_run.run_id,
+ TaskInstance.region_index == 0,
+ )
+ ).one()
+ archived = {column.name: getattr(first, column.key) for column in TaskInstance.__table__.columns} | {
+ "id": uuid6.uuid7(),
+ "working_set": None,
+ "archived_reason": "retry",
+ "try_number": 7,
+ }
+ session.execute(insert(TaskInstance.__table__).values(archived))
+ XComModel.set_for_attempt(
+ task_instance_id=archived["id"],
+ key="return_value",
+ value=-999,
+ serialize=False,
+ dag_result=True,
+ session=session,
+ )
session.commit()
response = test_client.get(
f"/dags/dag_mapped_result/dagRuns/{dag_run.run_id}/wait",
- params={"interval": "1"},
+ params={"interval": "1", **({"result": "a"} if explicit else {})},
)
assert response.status_code == 200
assert response.json() == {"state": DagRunState.SUCCESS, "results": {"a": [2, 4]}}
diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
index 8768ad41e9d..be5f88268fc 100644
--- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
+++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
@@ -4411,7 +4411,7 @@ class TestPostClearTaskInstances(TestTaskInstanceEndpoint):
if task_id is not None:
stmt = stmt.where(TaskStateStoreModel.task_id == task_id)
if map_index is not None:
- stmt = stmt.where(TaskStateStoreModel.map_index == map_index)
+ stmt = stmt.where(TaskStateStoreModel.region_index == map_index)
return session.scalars(stmt).all()
@pytest.mark.db_test
diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_state_store.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_state_store.py
index 95f3285542c..e074ec4ba38 100644
--- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_state_store.py
+++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_state_store.py
@@ -29,7 +29,9 @@ from airflow.api_fastapi.core_api.datamodels.task_state_store import (
TaskStateStorePatchBody,
)
from airflow.models.dagrun import DagRun
+from airflow.models.dynamic_region import DynamicRegion
from airflow.models.task_state_store import TaskStateStoreModel
+from airflow.models.taskinstance import TaskInstance
from airflow.providers.standard.operators.empty import EmptyOperator
from airflow.utils.types import DagRunType
@@ -61,7 +63,7 @@ def _create_task_state_store_row(session, key: str, value: str, dag_run: DagRun)
dag_id=DAG_ID,
run_id=RUN_ID,
task_id=TASK_ID,
- map_index=-1,
+ region_index=-1,
key=key,
value=json.dumps(value),
)
@@ -120,7 +122,7 @@ class TestListTaskState(TestTaskStateEndpoint):
dag_id=DAG_ID,
run_id=RUN_ID,
task_id=TASK_ID,
- map_index=0,
+ region_index=0,
key="job_id",
value=json.dumps("mapped_app"),
)
@@ -461,7 +463,7 @@ class TestClearTaskState(TestTaskStateEndpoint):
dag_id=DAG_ID,
run_id=RUN_ID,
task_id=TASK_ID,
- map_index=map_index,
+ region_index=map_index,
key="job_id",
value=json.dumps(f"app_{map_index}"),
)
@@ -569,3 +571,25 @@ class TestUnknownTaskInstance(TestTaskStateEndpoint):
response = test_client.get(f"{BASE_URL}/no_such_key")
assert response.status_code == 404
assert response.json()["detail"] == "Task state store key 'no_such_key' not found"
+
+
+class TestRegionalTaskState(TestTaskStateEndpoint):
+ def test_writes_and_clears_address_the_region_of_the_instance(self, test_client):
+ ti = self._session.scalar(select(TaskInstance).where(TaskInstance.dag_id == DAG_ID))
+ region = DynamicRegion(dag_id=DAG_ID, run_id=RUN_ID, node_id=TASK_ID)
+ self._session.add(region)
+ self._session.flush()
+ ti.region_id, ti.region_index = region.id, 0
+ self._session.commit()
+
+ assert test_client.put(f"{BASE_URL}/key?map_index=0", json={"value": "v"}).status_code == 204
+
+ row = self._session.scalars(
+ select(TaskStateStoreModel).where(TaskStateStoreModel.dag_id == DAG_ID)
+ ).one()
+ assert (row.region_id, row.region_index) == (region.id, 0)
+ assert test_client.get(f"{BASE_URL}/key?map_index=0").json()["value"] == "v"
+
+ assert test_client.delete(f"{BASE_URL}?all_map_indices=true").status_code == 204
+ self._session.expire_all()
+ assert self._session.scalars(select(TaskStateStoreModel)).all() == []
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 6c7641c5ac7..7108b450492 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
@@ -2188,7 +2188,7 @@ class TestTIUpdateState:
dag_id="dag",
run_id="run",
task_id="task",
- map_index=-1,
+ region_index=-1,
try_number=1,
max_tries=0,
start_date=None,
diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_state_store.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_state_store.py
index 837d848b359..042ae4ee653 100644
--- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_state_store.py
+++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_state_store.py
@@ -334,7 +334,7 @@ class TestClearTaskState:
dag_id=ti.dag_id,
run_id=ti.run_id,
task_id=ti.task_id,
- map_index=idx,
+ region_index=idx,
key="job_id",
value=f"app_{idx}",
updated_at=now,
@@ -355,7 +355,7 @@ class TestClearTaskState:
with create_session() as session:
remaining_indices = sorted(
session.scalars(
- select(TaskStateStoreModel.map_index).where(
+ select(TaskStateStoreModel.region_index).where(
TaskStateStoreModel.dag_id == ti.dag_id,
TaskStateStoreModel.task_id == ti.task_id,
)
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 ff57be4d4f0..2f970653adf 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
@@ -216,6 +216,30 @@ def test_prior_dates_resolves_mapped_region_separately_for_each_run(
assert response.json() == {"key": "key", "value": "old"}
+@pytest.mark.parametrize("archived_run", ["old", "current"])
+def test_prior_dates_ignores_values_of_archived_expansion_task_instances(
+ client, dag_maker, session, archived_run
+):
+ with dag_maker(serialized=True):
+ PythonOperator.partial(task_id="mapped", python_callable=str).expand(op_args=[[1]])
+ old = dag_maker.create_dagrun(run_id="old", logical_date=timezone.datetime(2026, 1, 1))
+ current = dag_maker.create_dagrun(run_id="current", logical_date=timezone.datetime(2026, 1, 2))
+ dr = old if archived_run == "old" else current
+ archived = dr.task_instances[0]
+ XComModel.set_for_attempt(
+ task_instance_id=archived.id, key="key", value=99, serialize=False, session=session
+ )
+ archived.archive(reason="retry", session=session)
+ session.commit()
+
+ response = client.get(
+ f"/execution/xcoms/{current.dag_id}/{current.run_id}/mapped/key",
+ params={"map_index": 0, "include_prior_dates": True},
+ )
+
+ assert response.status_code == 404
+
+
@pytest.fixture(autouse=True)
def reset_db():
"""Reset XCom entries."""
@@ -378,7 +402,7 @@ class TestXComsGetEndpoint:
with dag_maker(dag_id="dag"):
MyOperator.partial(task_id="task").expand(x=xcom_values)
dag_run = dag_maker.create_dagrun(run_id="runid")
- tis = {ti.map_index: ti for ti in dag_run.task_instances}
+ tis = {ti.region_index: ti for ti in dag_run.task_instances}
for map_index, db_value in enumerate(xcom_values):
if db_value is None: # We don't put None to XCom.
@@ -432,7 +456,7 @@ class TestXComsGetEndpoint:
with dag_maker(dag_id="dag"):
MyOperator.partial(task_id="task").expand(x=xcom_values)
dag_run = dag_maker.create_dagrun(run_id="runid")
- tis = {ti.map_index: ti for ti in dag_run.task_instances}
+ tis = {ti.region_index: ti for ti in dag_run.task_instances}
for map_index, db_value in enumerate(xcom_values):
if db_value is None: # We don't put None to XCom.
diff --git a/airflow-core/tests/unit/cli/commands/test_state_store_command.py b/airflow-core/tests/unit/cli/commands/test_state_store_command.py
index d777d41afd7..2d91125f6e4 100644
--- a/airflow-core/tests/unit/cli/commands/test_state_store_command.py
+++ b/airflow-core/tests/unit/cli/commands/test_state_store_command.py
@@ -45,7 +45,11 @@ class TestStateStoreCleanCommand:
backend.cleanup.assert_called_once_with()
- def test_dry_run_does_not_call_backend(self, capsys):
+ @pytest.mark.parametrize(
+ ("expired", "expected"),
+ [([], "Nothing to delete"), ([("dag", "run", "task", 4, "key")], "region_index 4")],
+ )
+ def test_dry_run_does_not_call_backend(self, capsys, expired, expected):
args = Namespace(dry_run=True, verbose=False)
backend = MetastoreBackend()
with (
@@ -54,12 +58,12 @@ class TestStateStoreCleanCommand:
return_value=backend,
autospec=True,
),
- patch.object(backend, "_summary_dry_run", return_value={"expired": []}, autospec=True),
+ patch.object(backend, "_summary_dry_run", return_value={"expired": expired}, autospec=True),
):
clean_state_store(args)
captured = capsys.readouterr()
- assert "Nothing to delete" in captured.out
+ assert expected in captured.out
def test_custom_backend_is_skipped(self, capsys):
# Custom (non-metastore) backends are intentionally skipped; cleanup() must not run.
diff --git a/airflow-core/tests/unit/cli/commands/test_task_command.py b/airflow-core/tests/unit/cli/commands/test_task_command.py
index 62c14eace22..fe190d637ed 100644
--- a/airflow-core/tests/unit/cli/commands/test_task_command.py
+++ b/airflow-core/tests/unit/cli/commands/test_task_command.py
@@ -29,7 +29,7 @@ from typing import TYPE_CHECKING
from unittest import mock
import pytest
-from sqlalchemy import delete
+from sqlalchemy import delete, select
from airflow._shared.timezones import timezone
from airflow.cli import cli_parser
@@ -43,9 +43,12 @@ from airflow.exceptions import DagRunNotFound
from airflow.models import DagModel, DagRun, TaskInstance
from airflow.models.dag_version import DagVersion
from airflow.models.dagbag import DBDagBag
+from airflow.models.dynamic_region import SENTINEL_REGION_ID, DynamicRegion
from airflow.models.serialized_dag import SerializedDagModel
from airflow.providers.standard.operators.bash import BashOperator
-from airflow.sdk import DAG as SdkDAG, BaseOperator
+from airflow.providers.standard.operators.python import PythonOperator
+from airflow.sdk import DAG as SdkDAG, BaseOperator, task_group
+from airflow.sdk.definitions._internal.loop import create_loop
from airflow.serialization.serialized_objects import DagSerialization, LazyDeserializedDAG
from airflow.utils.session import create_session
from airflow.utils.state import State
@@ -74,6 +77,243 @@ def reset(dag_id):
session.execute(delete(SerializedDagModel).where(SerializedDagModel.dag_id == dag_id))
+def test_task_states_for_dag_run_projects_map_index_and_identifies_regions(dag_maker, session):
+ with dag_maker() as dag:
+ BashOperator(task_id="work", bash_command="echo work")
+ BashOperator.partial(task_id="mapped").expand(bash_command=["echo mapped"])
+ run = dag_maker.create_dagrun()
+ loop = DynamicRegion(dag_id=run.dag_id, run_id=run.run_id, node_id="loop")
+ successor = DynamicRegion(dag_id=run.dag_id, run_id=run.run_id, node_id="loop")
+ mapped = DynamicRegion(dag_id=run.dag_id, run_id=run.run_id, node_id="mapped")
+ session.add_all([loop, successor, mapped])
+ session.flush()
+ work_ti = next(ti for ti in run.task_instances if ti.task_id == "work")
+ mapped_ti = next(ti for ti in run.task_instances if ti.task_id == "mapped")
+ work_ti.region_id, work_ti.region_index = loop.id, 4
+ mapped_ti.region_id, mapped_ti.region_index = mapped.id, 0
+ sibling = TaskInstance(
+ dag.get_task("work"),
+ run_id=run.run_id,
+ dag_version_id=work_ti.dag_version_id,
+ region_id=successor.id,
+ region_index=4,
+ )
+ session.add(sibling)
+ session.flush()
+
+ with redirect_stdout(io.StringIO()) as stdout:
+ task_command.task_states_for_dag_run(
+ cli_parser.get_parser().parse_args(
+ ["tasks", "states-for-dag-run", run.dag_id, run.run_id, "--output", "json"]
+ ),
+ session=session,
+ )
+
+ rows = json.loads(stdout.getvalue())
+ assert len(rows) == 3
+ assert {(row["region_id"], row["region_index"], row["map_index"]) for row in rows} == {
+ (str(loop.id), "4", ""),
+ (str(successor.id), "4", ""),
+ (str(mapped.id), "0", "0"),
+ }
+
+
+def test_task_states_for_dag_run_lists_only_live_rows_without_region_columns_for_mapped_dag(
+ dag_maker, session
+):
+ with dag_maker():
+ BashOperator(task_id="work", bash_command="echo work")
+ BashOperator.partial(task_id="mapped").expand(bash_command=["echo a", "echo b"])
+ run = dag_maker.create_dagrun()
+ work_ti = next(ti for ti in run.task_instances if ti.task_id == "work")
+ work_ti.state = State.SUCCESS
+ work_ti.archive(reason="cleared", session=session)
+ session.flush()
+ session.expire(run, ["task_instances"])
+
+ with redirect_stdout(io.StringIO()) as stdout:
+ task_command.task_states_for_dag_run(
+ cli_parser.get_parser().parse_args(
+ ["tasks", "states-for-dag-run", run.dag_id, run.run_id, "--output", "json"]
+ ),
+ session=session,
+ )
+
+ rows = json.loads(stdout.getvalue())
+ assert sorted((row["task_id"], row["map_index"]) for row in rows) == [("mapped", "0"), ("mapped", "1")]
+ assert all(
+ set(row) == {"dag_id", "logical_date", "task_id", "state", "start_date", "end_date", "map_index"}
+ for row in rows
+ )
+
+
+@pytest.mark.parametrize("command", ["state", "failed-deps", "test", "render"])
+def test_cli_commands_reuse_existing_regional_mapped_task(dag_maker, session, mocker, capsys, command):
+ with dag_maker(dag_id="regional_cli", serialized=True) as dag:
+ PythonOperator.partial(task_id="mapped", python_callable=str).expand(op_args=[[1], [2]])
+ dr = dag_maker.create_dagrun(run_id="regional")
+ region = DynamicRegion(dag_id=dr.dag_id, run_id=dr.run_id, node_id="mapped")
+ session.add(region)
+ session.flush()
+ for ti in dr.task_instances:
+ ti.region_id = region.id
+ ti.state = State.SUCCESS
+ selected = next(ti for ti in dr.task_instances if ti.region_index == 0)
+ selected_id = selected.id
+ session.commit()
+ serialized = dag_maker.serialized_dag
+ mocker.patch.object(task_command, "get_db_dag", autospec=True, return_value=serialized)
+ mocker.patch.object(task_command, "get_bagged_dag", autospec=True, return_value=dag_maker.dag)
+ run_task = mocker.patch.object(task_command, "_run_task", autospec=True)
+ lookup = mocker.spy(task_command, "_get_ti")
+ args = cli_parser.get_parser().parse_args(
+ ["tasks", command, dag.dag_id, "mapped", dr.run_id, "--map-index", "0"]
+ )
+
+ getattr(task_command, f"task_{command.replace('-', '_')}")(args)
+
+ assert lookup.spy_return[0].id == selected_id
+ if command == "test":
+ assert run_task.call_args.kwargs["ti"].id == selected_id
+ elif command == "state":
+ assert "success" in capsys.readouterr().out
+ session.expire_all()
+ rows = session.scalars(select(TaskInstance).where(TaskInstance.dag_id == dag.dag_id)).all()
+ assert len(rows) == 2
+ assert all(ti.region_id == region.id for ti in rows)
+
+
+@pytest.mark.parametrize("create_if_necessary", [False, "db", "memory"])
+@pytest.mark.parametrize("live_rows", [False, True])
+def test_cli_lookup_rejects_loop_scope_instead_of_creating_sentinel(
+ dag_maker, session, create_if_necessary, live_rows
+):
+ @task_group
+ def body():
+ BashOperator(task_id="terminal", bash_command="true")
+
+ with dag_maker(serialized=True):
+ group = create_loop(body, max_iterations=2)
+ dr = dag_maker.create_dagrun()
+ region = DynamicRegion(dag_id=dr.dag_id, run_id=dr.run_id, node_id=group.group_id)
+ session.add(region)
+ session.flush()
+ for ti in dr.task_instances:
+ ti.region_id, ti.region_index = region.id, 0
+ if not live_rows:
+ session.execute(delete(TaskInstance).where(TaskInstance.dag_id == dr.dag_id))
+ session.commit()
+ serialized = dag_maker.serialized_dag
+
+ with pytest.raises(ValueError, match="loop.*scope"):
+ task_command._get_ti(
+ serialized.get_task(group.terminal_task_id),
+ -1,
+ logical_date_or_run_id=dr.run_id,
+ create_if_necessary=create_if_necessary,
+ session=session,
+ )
+
+ assert not session.scalars(
+ select(TaskInstance).where(TaskInstance.dag_id == dr.dag_id, TaskInstance.region_id != region.id)
+ ).all()
+
+
+@pytest.mark.parametrize("create_if_necessary", ["db", "memory"])
+@pytest.mark.parametrize("regional_expansion", [False, True])
+def test_cli_missing_mapped_slot_is_created_ad_hoc_in_existing_expansion(
+ dag_maker, session, create_if_necessary, regional_expansion
+):
+ with dag_maker(serialized=True):
+ PythonOperator.partial(task_id="mapped", python_callable=str).expand(op_args=[[1], [2]])
+ dr = dag_maker.create_dagrun()
+ if regional_expansion:
+ region = session.scalars(select(DynamicRegion).where(DynamicRegion.dag_id == dr.dag_id)).one()
+ expected_region_id = region.id
+ session.execute(delete(TaskInstance).where(TaskInstance.dag_id == dr.dag_id))
+ else:
+ expected_region_id = SENTINEL_REGION_ID
+ for ti in dr.task_instances:
+ ti.region_id = SENTINEL_REGION_ID
+ session.flush()
+ session.execute(delete(DynamicRegion).where(DynamicRegion.dag_id == dr.dag_id))
+ session.execute(
+ delete(TaskInstance).where(TaskInstance.dag_id == dr.dag_id, TaskInstance.region_index == 0)
+ )
+ session.commit()
+
+ ti, _ = task_command._get_ti(
+ dag_maker.serialized_dag.get_task("mapped"),
+ 0,
+ logical_date_or_run_id=dr.run_id,
+ create_if_necessary=create_if_necessary,
+ session=session,
+ )
+
+ assert (ti.region_id, ti.region_index) == (expected_region_id, 0)
+ assert session.scalars(select(DynamicRegion).where(DynamicRegion.dag_id == dr.dag_id)).all() == (
+ [region] if regional_expansion else []
+ )
+
+
+@pytest.mark.parametrize("create_if_necessary", ["db", "memory"])
+def test_cli_xcom_mapped_slot_is_created_ad_hoc_beside_placeholder(dag_maker, session, create_if_necessary):
+ with dag_maker(serialized=True):
+ upstream = PythonOperator(task_id="upstream", python_callable=lambda: [1, 2, 3])
+ PythonOperator.partial(task_id="mapped", python_callable=str).expand(op_args=upstream.output)
+ dr = dag_maker.create_dagrun()
+ placeholder = session.scalars(
+ select(TaskInstance).where(TaskInstance.dag_id == dr.dag_id, TaskInstance.task_id == "mapped")
+ ).all()
+ assert placeholder
+
+ ti, _ = task_command._get_ti(
+ dag_maker.serialized_dag.get_task("mapped"),
+ 2,
+ logical_date_or_run_id=dr.run_id,
+ create_if_necessary=create_if_necessary,
+ session=session,
+ )
+
+ assert ti.region_index == 2
+ assert ti.task_id == "mapped"
+
+
+@pytest.mark.parametrize("command", ["test", "render"])
+def test_ad_hoc_mapped_task_gets_root_region(dag_maker, session, mocker, command):
+ with dag_maker(dag_id="ad_hoc_mapping", serialized=True):
+ PythonOperator.partial(task_id="mapped", python_callable=str).expand(op_args=[[1], [2]])
+ run_id = "ad_hoc"
+ if command == "test":
+ dr = dag_maker.create_dagrun(run_id=run_id)
+ session.execute(delete(TaskInstance).where(TaskInstance.dag_id == dr.dag_id))
+ session.execute(delete(DynamicRegion).where(DynamicRegion.dag_id == dr.dag_id))
+ session.add(DynamicRegion(dag_id=dr.dag_id, run_id=run_id, node_id="other_task"))
+ session.commit()
+ mocker.patch.object(task_command, "get_db_dag", autospec=True, return_value=dag_maker.serialized_dag)
+ mocker.patch.object(task_command, "get_bagged_dag", autospec=True, return_value=dag_maker.dag)
+
+ def run_task(*, ti, task, run_triggerer):
+ region = session.get(DynamicRegion, ti.region_id)
+ assert region is not None
+ assert (region.node_id, region.parent_region_id) == ("mapped", None)
+
+ mocker.patch.object(task_command, "_run_task", autospec=True, side_effect=run_task)
+ lookup = mocker.spy(task_command, "_get_ti")
+ args = cli_parser.get_parser().parse_args(
+ ["tasks", command, "ad_hoc_mapping", "mapped", run_id, "--map-index", "0"]
+ )
+
+ getattr(task_command, f"task_{command}")(args)
+
+ ti = lookup.spy_return[0]
+ assert ti.region_id.int != 0
+ assert ti.region_index == 0
+ if command == "render":
+ assert session.get(DynamicRegion, ti.region_id) is None
+ assert not session.scalars(select(DagRun).where(DagRun.dag_id == "ad_hoc_mapping")).all()
+
+
class TestCliTasks:
run_id = "TEST_RUN_ID"
dag_id = "example_python_operator"
diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py
index ae66019b561..63290a9d252 100644
--- a/airflow-core/tests/unit/jobs/test_scheduler_job.py
+++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py
@@ -95,6 +95,7 @@ from airflow.models.dagwarning import DagWarning
from airflow.models.db_callback_request import DbCallbackRequest
from airflow.models.deadline import Deadline
from airflow.models.deadline_alert import DeadlineAlert
+from airflow.models.dynamic_region import DynamicRegion
from airflow.models.hitl import HITLDetail
from airflow.models.log import Log, resolve_team_name
from airflow.models.pool import Pool, PoolStats
@@ -411,6 +412,28 @@ class TestSchedulerJob:
loader_mock.return_value = mock_executors
yield default_executor
+ def test_ranked_admission_joins_exact_task_uuid(self, dag_maker, session):
+ with dag_maker(max_active_tasks=1):
+ task = EmptyOperator(task_id="task")
+ dr = dag_maker.create_dagrun(state=DagRunState.RUNNING)
+ scheduled = dr.task_instances[0]
+ scheduled.state = TaskInstanceState.SCHEDULED
+ region = DynamicRegion(dag_id=dr.dag_id, run_id=dr.run_id, node_id=task.task_id)
+ session.add(region)
+ session.flush()
+ sibling = TaskInstance(
+ task=task, run_id=dr.run_id, dag_version_id=scheduled.dag_version_id, region_id=region.id
+ )
+ sibling.state = TaskInstanceState.SUCCESS
+ session.add(sibling)
+ scheduled.dag_model.is_paused = False
+ session.flush()
+ runner = SchedulerJobRunner(Job())
+
+ selected = session.scalars(runner._build_schedulable_tis_query(set(), set(), set(), set(), 10)).all()
+
+ assert [ti.id for ti in selected] == [scheduled.id]
+
def test_is_alive(self):
scheduler_job = Job(heartrate=10, state=State.RUNNING)
self.job_runner = SchedulerJobRunner(scheduler_job)
@@ -3997,11 +4020,11 @@ class TestSchedulerJob:
def refuse_adoption(tis):
assert len(tis) == 1
- # ``repr(ti)`` in the reset path reads both ``state`` and ``map_index``; the query
+ # ``repr(ti)`` in the reset path reads both ``state`` and ``region_index``; the query
# must load both so the reset log stays accurate (and never lazy-loads on detach).
unloaded = inspect(tis[0]).unloaded
assert "state" not in unloaded
- assert "map_index" not in unloaded
+ assert "region_index" not in unloaded
# repr must render the real state, not the ``<deferred>`` fallback.
assert "queued" in repr(tis[0])
return tis
@@ -10211,6 +10234,7 @@ class TestSchedulerJob:
@pytest.mark.parametrize("missing_definition", ["dag", "task"])
@pytest.mark.parametrize(("max_tries", "expected_max_tries"), [(0, 3), (7, 7)])
@pytest.mark.parametrize("callback_version_available", [True, False])
+ @pytest.mark.parametrize("regional", [False, True])
def test_heartbeat_timeout_completes_clear_without_definition(
self,
dag_maker,
@@ -10220,11 +10244,17 @@ class TestSchedulerJob:
max_tries,
expected_max_tries,
callback_version_available,
+ regional,
):
with dag_maker(dag_id="heartbeat_clear_missing_definition", session=session):
EmptyOperator(task_id="task")
dr = dag_maker.create_dagrun(state=DagRunState.RUNNING)
ti = dr.get_task_instance("task", session=session)
+ if regional:
+ region = DynamicRegion(dag_id=dr.dag_id, run_id=dr.run_id, node_id="loop")
+ session.add(region)
+ session.flush()
+ ti.region_id, ti.region_index = region.id, 2
ti.state = State.RESTARTING
ti.try_number = 3
ti.max_tries = max_tries
@@ -10252,7 +10282,9 @@ class TestSchedulerJob:
session.refresh(ti)
assert ti.id == old_id
assert ti.working_set is None
- current = dr.get_task_instance("task", session=session)
+ current = dr.get_task_instance(
+ "task", map_index=ti.region_index, region_id=ti.region_id, session=session
+ )
assert current.id != old_id
assert (current.try_number, current.state, current.max_tries) == (4, None, expected_max_tries)
assert current.external_executor_id is None
diff --git a/airflow-core/tests/unit/migrations/test_0143_add_dynamic_region_storage.py b/airflow-core/tests/unit/migrations/test_0143_add_dynamic_region_storage.py
index eff1e73d6e7..78ac957d936 100644
--- a/airflow-core/tests/unit/migrations/test_0143_add_dynamic_region_storage.py
+++ b/airflow-core/tests/unit/migrations/test_0143_add_dynamic_region_storage.py
@@ -38,6 +38,7 @@ from tests_common.test_utils.db import clear_db_runs
pytestmark = pytest.mark.db_test
_migration = import_module("airflow.migrations.versions.0143_3_4_0_add_dynamic_region_storage")
+_index_migration = import_module("airflow.migrations.versions.0144_3_4_0_rename_stored_map_index")
MODELS = (TaskInstance, TaskStateStoreModel)
@@ -60,7 +61,7 @@ def stored_rows(dag_maker, session):
dag_id=ti.dag_id,
run_id=ti.run_id,
task_id=ti.task_id,
- map_index=-1,
+ region_index=-1,
key="state",
value="value",
)
@@ -92,22 +93,25 @@ def test_downgrade_then_upgrade_keeps_existing_rows_in_the_sentinel_region(
with Operations.context(context), context.begin_transaction(_per_migration=True):
if connection.dialect.name == "sqlite":
connection.exec_driver_sql(f"PRAGMA foreign_keys={int(sqlite_foreign_keys)}")
+ _index_migration.downgrade()
_migration.downgrade()
_migration.upgrade()
+ _index_migration.upgrade()
if connection.dialect.name == "sqlite":
assert connection.exec_driver_sql("PRAGMA foreign_keys").scalar_one() == sqlite_foreign_keys
session.expire_all()
for model in MODELS:
row = session.scalars(select(model)).one()
assert row.region_id == UUID(int=0)
- assert row.map_index == -1
+ assert row.region_index == -1
@pytest.mark.parametrize("model", MODELS)
def test_downgrade_refuses_each_old_coordinate_collision_before_ddl(stored_rows, session, model):
connection = session.connection()
table = model.__table__
- duplicate = dict(connection.execute(select(table)).mappings().one())
+ row = connection.execute(select(table)).mappings().one()
+ duplicate = {column.key: row[column.name] for column in table.columns}
duplicate["region_id"] = uuid4()
if model is TaskInstance:
duplicate["id"] = uuid4()
@@ -126,14 +130,18 @@ def test_downgrade_refuses_each_old_coordinate_collision_before_ddl(stored_rows,
connection.execute(table.insert().values(**duplicate))
try:
with Operations.context(MigrationContext.configure(connection)):
- # SQLite is checked in Python; PostgreSQL and MySQL raise from the SQL guard. PostgreSQL
- # aborts the transaction on error, but MySQL's guard procedure is DDL, which would commit
- # away a savepoint.
- savepoint = (
- connection.begin_nested() if connection.dialect.name == "postgresql" else nullcontext()
- )
- with pytest.raises((RuntimeError, DBAPIError), match=model.__tablename__), savepoint:
- _migration.downgrade()
+ _index_migration.downgrade()
+ try:
+ # SQLite is checked in Python; PostgreSQL and MySQL raise from the SQL guard. PostgreSQL
+ # aborts the transaction on error, but MySQL's guard procedure is DDL, which would commit
+ # away a savepoint.
+ savepoint = (
+ connection.begin_nested() if connection.dialect.name == "postgresql" else nullcontext()
+ )
+ with pytest.raises((RuntimeError, DBAPIError), match=model.__tablename__), savepoint:
+ _migration.downgrade()
+ finally:
+ _index_migration.upgrade()
assert inspect(connection).has_table("dynamic_region")
for retained in MODELS:
assert "region_id" in {c["name"] for c in inspect(connection).get_columns(retained.__tablename__)}
diff --git a/airflow-core/tests/unit/migrations/test_0143_add_dynamic_region_storage.py b/airflow-core/tests/unit/migrations/test_0144_rename_stored_map_index.py
similarity index 50%
copy from airflow-core/tests/unit/migrations/test_0143_add_dynamic_region_storage.py
copy to airflow-core/tests/unit/migrations/test_0144_rename_stored_map_index.py
index eff1e73d6e7..8758999ab05 100644
--- a/airflow-core/tests/unit/migrations/test_0143_add_dynamic_region_storage.py
+++ b/airflow-core/tests/unit/migrations/test_0144_rename_stored_map_index.py
@@ -16,29 +16,27 @@
# under the License.
from __future__ import annotations
-from contextlib import nullcontext
from importlib import import_module
-from uuid import UUID, uuid4
+from uuid import uuid4
import pytest
from alembic.migration import MigrationContext
from alembic.operations import Operations
-from sqlalchemy import inspect, select
-from sqlalchemy.exc import DBAPIError
+from sqlalchemy import MetaData, Table, func, select
-from airflow._shared.timezones import timezone
-from airflow.models.dagrun import DagRun
-from airflow.models.dynamic_region import DynamicRegion
from airflow.models.task_state_store import TaskStateStoreModel
-from airflow.models.taskinstance import TaskInstance
+from airflow.models.taskinstance import LegacyTaskDataOwner, TaskInstance
+from airflow.models.taskreschedule import TaskReschedule
+from airflow.models.xcom import XComModelV2
from airflow.providers.standard.operators.empty import EmptyOperator
from tests_common.test_utils.db import clear_db_runs
pytestmark = pytest.mark.db_test
-_migration = import_module("airflow.migrations.versions.0143_3_4_0_add_dynamic_region_storage")
-MODELS = (TaskInstance, TaskStateStoreModel)
+_migration = import_module("airflow.migrations.versions.0144_3_4_0_rename_stored_map_index")
+TABLES = ("task_instance", "task_state_store")
+CHILDREN = (TaskReschedule, XComModelV2, LegacyTaskDataOwner)
@pytest.fixture(autouse=True)
@@ -48,30 +46,52 @@ def clean_db():
clear_db_runs()
+def _stored(connection, table_name, index_column):
+ table = Table(table_name, MetaData(), autoload_with=connection)
+ rows = connection.execute(select(table).order_by(table.c.region_id, table.c[index_column])).mappings()
+ return [{("region_index" if k == index_column else k): v for k, v in row.items()} for row in rows]
+
+
@pytest.fixture
def stored_rows(dag_maker, session):
with dag_maker(serialized=True):
EmptyOperator(task_id="task")
dr = dag_maker.create_dagrun()
ti = dr.task_instances[0]
+ now = ti.dag_run.start_date or ti.dag_run.queued_at
+ connection = session.connection()
+ ti_table = TaskInstance.__table__
+ current = dict(connection.execute(select(ti_table)).mappings().one())
+ archived = {**current, "id": uuid4(), "working_set": None, "archived_reason": "retry", "try_number": 3}
+ regional = {**current, "id": uuid4(), "region_id": uuid4(), "region_index": 4}
+ for values in (archived, regional):
+ connection.execute(ti_table.insert().values(values))
+ session.add(TaskReschedule(ti_id=regional["id"], start_date=now, end_date=now, reschedule_date=now))
+ session.add(
+ LegacyTaskDataOwner(
+ dag_id=ti.dag_id, task_id=ti.task_id, run_id=ti.run_id, map_index=-1, task_instance_id=ti.id
+ )
+ )
+ session.add(XComModelV2(task_instance_id=regional["id"], key="key", value="value"))
session.add(
TaskStateStoreModel(
dag_run_id=dr.id,
dag_id=ti.dag_id,
run_id=ti.run_id,
task_id=ti.task_id,
- map_index=-1,
+ region_id=regional["region_id"],
+ region_index=4,
key="state",
value="value",
)
)
session.flush()
+ session.commit()
return ti
@pytest.fixture
def migration_connection(stored_rows, session):
- session.commit()
with session.get_bind().connect() as connection:
try:
yield connection
@@ -82,61 +102,25 @@ def migration_connection(stored_rows, session):
@pytest.mark.parametrize("sqlite_foreign_keys", [True, False])
-def test_downgrade_then_upgrade_keeps_existing_rows_in_the_sentinel_region(
- stored_rows, session, migration_connection, sqlite_foreign_keys
+def test_rename_roundtrip_preserves_rows_and_their_children(
+ stored_rows, migration_connection, sqlite_foreign_keys
):
connection = migration_connection
if not sqlite_foreign_keys and connection.dialect.name != "sqlite":
pytest.skip("SQLite foreign-key setting")
+ upgraded = {name: _stored(connection, name, "region_index") for name in TABLES}
+ child_counts = {model: connection.scalar(select(func.count()).select_from(model)) for model in CHILDREN}
context = MigrationContext.configure(connection, opts={"transaction_per_migration": True})
with Operations.context(context), context.begin_transaction(_per_migration=True):
if connection.dialect.name == "sqlite":
connection.exec_driver_sql(f"PRAGMA foreign_keys={int(sqlite_foreign_keys)}")
_migration.downgrade()
- _migration.upgrade()
- if connection.dialect.name == "sqlite":
- assert connection.exec_driver_sql("PRAGMA foreign_keys").scalar_one() == sqlite_foreign_keys
- session.expire_all()
- for model in MODELS:
- row = session.scalars(select(model)).one()
- assert row.region_id == UUID(int=0)
- assert row.map_index == -1
-
-
-@pytest.mark.parametrize("model", MODELS)
-def test_downgrade_refuses_each_old_coordinate_collision_before_ddl(stored_rows, session, model):
- connection = session.connection()
- table = model.__table__
- duplicate = dict(connection.execute(select(table)).mappings().one())
- duplicate["region_id"] = uuid4()
- if model is TaskInstance:
- duplicate["id"] = uuid4()
- else:
- duplicate.pop("id")
- region = DynamicRegion.__table__
- connection.execute(
- region.insert().values(
- id=duplicate["region_id"],
- dag_id=session.scalar(select(DagRun.dag_id)),
- run_id=session.scalar(select(DagRun.run_id)),
- node_id="task",
- created_at=timezone.utcnow(),
- )
- )
- connection.execute(table.insert().values(**duplicate))
- try:
- with Operations.context(MigrationContext.configure(connection)):
- # SQLite is checked in Python; PostgreSQL and MySQL raise from the SQL guard. PostgreSQL
- # aborts the transaction on error, but MySQL's guard procedure is DDL, which would commit
- # away a savepoint.
- savepoint = (
- connection.begin_nested() if connection.dialect.name == "postgresql" else nullcontext()
- )
- with pytest.raises((RuntimeError, DBAPIError), match=model.__tablename__), savepoint:
- _migration.downgrade()
- assert inspect(connection).has_table("dynamic_region")
- for retained in MODELS:
- assert "region_id" in {c["name"] for c in inspect(connection).get_columns(retained.__tablename__)}
- finally:
- connection.execute(table.delete().where(table.c.region_id == duplicate["region_id"]))
- connection.execute(region.delete().where(region.c.id == duplicate["region_id"]))
+ try:
+ assert {name: _stored(connection, name, "map_index") for name in TABLES} == upgraded
+ finally:
+ _migration.upgrade()
+
+ assert {name: _stored(connection, name, "region_index") for name in TABLES} == upgraded
+ assert {
+ model: connection.scalar(select(func.count()).select_from(model)) for model in CHILDREN
+ } == child_counts
diff --git a/airflow-core/tests/unit/models/test_dag.py b/airflow-core/tests/unit/models/test_dag.py
index 7d6ebf3419c..f6d871a2655 100644
--- a/airflow-core/tests/unit/models/test_dag.py
+++ b/airflow-core/tests/unit/models/test_dag.py
@@ -63,6 +63,7 @@ from airflow.models.dagbag import DBDagBag
from airflow.models.dagbundle import DagBundleModel
from airflow.models.dagrun import DagRun
from airflow.models.deadline_alert import DeadlineAlert as DeadlineAlertModel
+from airflow.models.dynamic_region import DynamicRegion
from airflow.models.hitl import HITLDetail
from airflow.models.serialized_dag import SerializedDagModel
from airflow.models.taskinstance import TaskInstance as TI
@@ -78,9 +79,11 @@ from airflow.sdk import (
TaskGroup,
setup,
task as task_decorator,
+ task_group,
teardown,
)
from airflow.sdk.definitions._internal.contextmanager import TaskGroupContext
+from airflow.sdk.definitions._internal.loop import create_loop
from airflow.sdk.definitions._internal.templater import NativeEnvironment, SandboxedEnvironment
from airflow.sdk.definitions.asset import Asset, AssetAlias, AssetAll, AssetAny
from airflow.sdk.definitions.callback import AsyncCallback
@@ -1721,6 +1724,62 @@ class TestDag:
assert task_group.get_child_by_label("task_with_task_group") == task_with_task_group
assert dag.get_task("task_group.task_with_task_group") == task_with_task_group
+ @pytest.mark.parametrize(
+ ("excluded", "remaining"),
+ [
+ (None, {"outside", "body.work", "body.__loop_gate", "mapped"}),
+ (set(), {"outside", "body.work", "body.__loop_gate", "mapped"}),
+ ({("outside", -1)}, {"body.work", "body.__loop_gate", "mapped"}),
+ ({"body.work"}, {"outside", "body.__loop_gate", "mapped"}),
+ ({("body.work", -1)}, {"outside", "body.__loop_gate", "mapped"}),
+ ({("mapped", 0)}, {"outside", "body.work", "body.__loop_gate", "mapped"}),
+ ],
+ )
+ @pytest.mark.need_serialized_dag
+ def test_clear_selection_excludes_public_map_indexes(self, dag_maker, session, excluded, remaining):
+ @task_group
+ def body():
+ EmptyOperator(task_id="work")
+
+ with dag_maker(serialized=True) as dag:
+ create_loop(body, max_iterations=3)
+ EmptyOperator(task_id="outside")
+ PythonOperator.partial(task_id="mapped", python_callable=list).expand(op_kwargs=[{}, {}])
+ dr = dag_maker.create_dagrun(session=session)
+ root = DynamicRegion(dag_id=dr.dag_id, run_id=dr.run_id, node_id="body")
+ session.add(root)
+ session.flush()
+ first_pass = next(ti for ti in dr.task_instances if ti.task_id == "body.work")
+ first_pass.region_id, first_pass.region_index = root.id, 0
+ for index in (1, 2):
+ session.add(
+ TI(
+ task=dag.get_task("body.work"),
+ run_id=dr.run_id,
+ dag_version_id=dr.created_dag_version_id,
+ region_id=root.id,
+ region_index=index,
+ )
+ )
+ session.flush()
+ selected = list(
+ dag._get_task_instances(
+ task_ids=None,
+ start_date=None,
+ end_date=None,
+ run_id=dr.run_id,
+ state=None,
+ exclude_task_ids=excluded,
+ exclude_run_ids=None,
+ session=session,
+ )
+ )
+ assert {ti.task_id for ti in selected} == remaining
+ if excluded == {("mapped", 0)}:
+ assert [ti.region_index for ti in selected if ti.task_id == "mapped"] == [1]
+ if "body.work" in remaining:
+ assert {ti.region_index for ti in selected if ti.task_id == "body.work"} == {0, 1, 2}
+
@pytest.mark.parametrize("dag_run_state", [DagRunState.QUEUED, DagRunState.RUNNING])
@pytest.mark.need_serialized_dag
def test_clear_set_dagrun_state(self, dag_run_state, dag_maker, session):
diff --git a/airflow-core/tests/unit/models/test_dagrun.py b/airflow-core/tests/unit/models/test_dagrun.py
index ded8127dd11..b696b5b74b2 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.timezones import timezone
from airflow.callbacks.callback_requests import DagCallbackRequest, DagRunContext
from airflow.models.dag import DagModel, infer_automated_data_interval
from airflow.models.dag_version import DagVersion
+from airflow.models.dagbag import DBDagBag
from airflow.models.dagrun import DagRun, DagRunNote, clear_partition_runs, get_or_create_dagrun
from airflow.models.deadline import Deadline
from airflow.models.deadline_alert import DeadlineAlert as DeadlineAlertModel
@@ -80,7 +81,7 @@ from airflow.sdk import (
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 LazyDeserializedDAG
+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.triggers.base import StartTriggerArgs
@@ -90,6 +91,7 @@ from airflow.utils.state import DagRunState, DagSchedulingState, State, TaskInst
from airflow.utils.types import DagRunTriggeredByType, DagRunType
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
@@ -196,7 +198,7 @@ class TestDagRun:
if task_states is not None:
for task_id, task_state in task_states.items():
- ti = dag_run.get_task_instance(task_id)
+ ti = dag_run.get_task_instance(task_id, session=session)
if TYPE_CHECKING:
assert ti
ti.set_state(task_state, session=session)
@@ -1730,10 +1732,10 @@ def test_verify_integrity_task_start_and_end_date(
@pytest.mark.parametrize("is_noop", [True, False])
-def test_expand_mapped_task_instance_at_create(is_noop, dag_maker, session):
- with mock.patch("airflow.settings.task_instance_mutation_hook") as mock_mut:
+@pytest.mark.parametrize("literal", [[], [1, 2, 3, 4]])
+def test_expand_mapped_task_instance_at_create(is_noop, literal, dag_maker, session):
+ with mock.patch("airflow.settings.task_instance_mutation_hook", autospec=True) as mock_mut:
mock_mut.is_noop = is_noop
- literal = [1, 2, 3, 4]
with dag_maker(session=session, dag_id="test_dag"):
mapped = MockOperator.partial(task_id="task_2").expand(arg2=literal)
@@ -1743,7 +1745,56 @@ def test_expand_mapped_task_instance_at_create(is_noop, dag_maker, session):
.where(TI.task_id == mapped.task_id, TI.dag_id == mapped.dag_id, TI.run_id == dr.run_id)
.order_by(TI.map_index)
).all()
- assert indices == [0, 1, 2, 3]
+ assert indices == (list(range(len(literal))) if literal else [-1])
+ region = session.scalars(select(DynamicRegion).where(DynamicRegion.dag_id == dr.dag_id)).one()
+ assert region.node_id == mapped.task_id
+ assert region.parent_region_id is None
+ assert region.forked_from_region_id is None
+ tis = session.scalars(select(TI).where(TI.dag_id == dr.dag_id, TI.run_id == dr.run_id)).all()
+ assert {ti.region_id for ti in tis} == {region.id}
+ if not literal:
+ assert tis[0].state == TaskInstanceState.SKIPPED
+ original_states = {ti.id: ti.state for ti in tis}
+
+ dr.verify_integrity(dag_version_id=dr.created_dag_version_id, session=session)
+
+ assert {ti.id: ti.state for ti in dr.get_task_instances(session=session)} == original_states
+ assert session.scalars(select(DynamicRegion).where(DynamicRegion.dag_id == dr.dag_id)).all() == [
+ region
+ ]
+
+
+@pytest.mark.parametrize("task_count", [2, 40])
+def test_creating_literal_mapped_tasks_costs_constant_queries(task_count, dag_maker, session):
+ with dag_maker(session=session, serialized=True):
+ for index in range(task_count):
+ MockOperator.partial(task_id=f"mapped_{index}").expand(arg2=[1, 2, 3])
+
+ with assert_queries_count(12):
+ dr = dag_maker.create_dagrun()
+
+ live = session.scalars(select(TI).where(TI.dag_id == dr.dag_id, TI.run_id == dr.run_id)).all()
+ assert len(live) == 3 * task_count
+
+
+def test_unresolved_mapping_has_region_before_expansion(dag_maker, session):
+ with dag_maker(serialized=True):
+ producer = BaseOperator(task_id="producer")
+ MockOperator.partial(task_id="mapped").expand(arg2=producer.output)
+ dr = dag_maker.create_dagrun()
+ tis = {ti.task_id: ti for ti in dr.task_instances}
+ mapped = tis["mapped"]
+ assert mapped.region_id.int != 0
+ assert mapped.region_index == -1
+ assert mapped.state is None
+ assert tis["producer"].region_id.int == 0
+ region = session.get(DynamicRegion, mapped.region_id)
+ assert region.node_id == "mapped"
+
+ dr.verify_integrity(dag_version_id=dr.created_dag_version_id, session=session)
+
+ assert session.scalars(select(DynamicRegion).where(DynamicRegion.dag_id == dr.dag_id)).all() == [region]
+ assert session.get(TI, mapped.id) is mapped
@pytest.mark.parametrize("is_noop", [True, False])
@@ -2161,7 +2212,7 @@ def test_restoring_removed_task_allocates_attempt_once(dag_maker, session, try_n
else:
BashOperator(task_id="task", bash_command="true")
dr = dag_maker.create_dagrun()
- ti = dr.get_task_instance("task", map_index=0 if mapped else -1, session=session)
+ ti = dr.task_instances[0]
ti.state = TaskInstanceState.REMOVED
ti.try_number = try_number
old_id = ti.id
@@ -2199,7 +2250,7 @@ def test_verifying_removed_map_index_does_not_allocate_attempt(dag_maker, sessio
with dag_maker(session=session):
BashOperator.partial(task_id="task").expand(bash_command=["true", "true"])
dr = dag_maker.create_dagrun()
- ti = dr.get_task_instance("task", map_index=1, session=session)
+ ti = next(ti for ti in dr.task_instances if ti.region_index == 1)
ti.try_number = 2
ti.state = TaskInstanceState.SUCCESS
old_id = ti.id
@@ -2398,9 +2449,11 @@ def test_mapped_literal_faulty_state_in_db(dag_maker, session):
assert len(decision.schedulable_tis) == 2
# We insert a faulty record
- session.add(
- create_task_instance(task=dag.get_task("task_2"), run_id=dr.run_id, dag_version_id=ti.dag_version_id)
+ placeholder = create_task_instance(
+ task=dag.get_task("task_2"), run_id=dr.run_id, dag_version_id=ti.dag_version_id
)
+ placeholder.region_id = decision.schedulable_tis[0].region_id
+ session.add(placeholder)
session.flush()
decision = dr.task_instance_scheduling_decisions()
@@ -2557,20 +2610,27 @@ def test_mapped_task_group_empty_operator(dag_maker, session):
dr = dag_maker.create_dagrun()
t2_task = dag.get_task("tg.t2")
- t2_0 = dr.get_task_instance(task_id="tg.t2", map_index=0)
+ region_id = session.scalar(
+ select(DynamicRegion.id).where(
+ DynamicRegion.dag_id == dr.dag_id,
+ DynamicRegion.run_id == dr.run_id,
+ DynamicRegion.node_id == "tg.t2",
+ )
+ )
+ t2_0 = dr.get_task_instance(task_id="tg.t2", map_index=0, region_id=region_id)
t2_0.refresh_from_task(t2_task)
assert t2_0.state is None
- t2_1 = dr.get_task_instance(task_id="tg.t2", map_index=1)
+ t2_1 = dr.get_task_instance(task_id="tg.t2", map_index=1, region_id=region_id)
t2_1.refresh_from_task(t2_task)
assert t2_1.state is None
dr.schedule_tis([t2_0])
- t2_0 = dr.get_task_instance(task_id="tg.t2", map_index=0)
+ t2_0 = dr.get_task_instance(task_id="tg.t2", map_index=0, region_id=region_id)
assert t2_0.state == TaskInstanceState.SUCCESS
- t2_1 = dr.get_task_instance(task_id="tg.t2", map_index=1)
+ t2_1 = dr.get_task_instance(task_id="tg.t2", map_index=1, region_id=region_id)
assert t2_1.state is None
@@ -3377,6 +3437,7 @@ def test_schedulable_task_exist_when_rerun_removed_upstream_mapped_task(session,
map_index=map_index,
dag_version_id=ti.dag_version_id,
)
+ ti_new.region_id = ti.region_id
session.add(ti_new)
ti_new.dag_run = dr
else:
@@ -3861,7 +3922,14 @@ def test_clearing_task_and_moving_from_non_mapped_to_mapped(dag_maker, session):
# Purposely omitted RenderedTaskInstanceFields because the ti need
# to be expanded but here we are mimicking and made it map_index -1
session.add(tr)
- XComModel.set(key="test", value="value", task_id=ti.task_id, dag_id=dag.dag_id, run_id=ti.run_id)
+ XComModel.set(
+ key="test",
+ value="value",
+ task_id=ti.task_id,
+ dag_id=dag.dag_id,
+ run_id=ti.run_id,
+ region_id=ti.region_id,
+ )
session.commit()
for table in [TaskInstanceNote, TaskReschedule, XComModel]:
assert session.scalar(select(func.count()).select_from(table)) == 1
@@ -5537,3 +5605,20 @@ class TestGetOrCreateDagrun:
select(func.count()).select_from(TaskInstance).where(TaskInstance.run_id == "manual__existing")
)
assert existing_ti_count == 1
+
+
+def test_scheduling_decisions_do_not_redeserialize_the_dag_per_pass(dag_maker, session):
+ with dag_maker(session=session, serialized=True):
+ producer = EmptyOperator(task_id="producer")
+ mapped = MockOperator.partial(task_id="mapped").expand(arg2=[1, 2])
+ producer >> mapped >> EmptyOperator(task_id="reduce")
+ dr = dag_maker.create_dagrun()
+ dr.dag = DBDagBag().get_dag_for_run(dr, session=session)
+
+ with mock.patch.object(
+ DagSerialization, "from_dict", autospec=True, side_effect=DagSerialization.from_dict
+ ) as read:
+ for _ in range(3):
+ dr.task_instance_scheduling_decisions(session=session)
+
+ assert read.call_count == 0
diff --git a/airflow-core/tests/unit/models/test_dynamic_region.py b/airflow-core/tests/unit/models/test_dynamic_region.py
index 2859592944a..df32cc23538 100644
--- a/airflow-core/tests/unit/models/test_dynamic_region.py
+++ b/airflow-core/tests/unit/models/test_dynamic_region.py
@@ -21,21 +21,24 @@ from typing import TYPE_CHECKING
from uuid import uuid4
import pytest
-from sqlalchemy import select
+from sqlalchemy import select, text
from airflow._shared.timezones import timezone
from airflow.models.dynamic_region import (
+ SENTINEL_REGION_ID,
AmbiguousProducerError,
DynamicRegion,
ProducerContext,
resolve_current_producers,
)
from airflow.models.taskinstance import TaskInstance
+from airflow.models.taskinstancekey import TaskInstanceKey
from airflow.models.xcom import XComModel
from airflow.providers.standard.operators.empty import EmptyOperator
from airflow.providers.standard.operators.python import PythonOperator
from airflow.utils.state import TaskInstanceState
+from tests_common.test_utils.asserts import capture_orm_selects
from tests_common.test_utils.db import clear_db_runs
if TYPE_CHECKING:
@@ -322,9 +325,7 @@ def test_mapping_revision_only_changes_selected_expansion(dag_maker, session):
regional = TaskInstance(task, version, run_id=dr.run_id, map_index=3, region_id=uuid4())
session.add_all([ordinary, regional])
session.flush()
- added = dr._revise_map_indexes_if_mapped(
- task, dag_version_id=version, region_id=regional.region_id, session=session
- )
+ added = dr._revise_map_indexes_if_mapped(regional, session=session)
assert [(ti.region_id, ti.region_index) for ti in added] == [
(regional.region_id, 0),
(regional.region_id, 1),
@@ -376,3 +377,60 @@ def test_xcom_reads_are_scoped_to_the_producer_region(regional_tis, session):
assert read() == {first.id}
assert read(region_id=second.region_id) == {second.id}
assert read(region_id=None) == {first.id, second.id}
+
+
+def test_public_task_instance_lookup_addresses_a_tasks_own_region_only(dag_maker, session):
+ with dag_maker(serialized=True):
+ PythonOperator.partial(task_id="mapped", python_callable=str).expand(op_args=[[1], [2]])
+ EmptyOperator(task_id="plain")
+ dr = dag_maker.create_dagrun()
+ mapped = dr.get_task_instance("mapped", map_index=1, session=session)
+ plain = dr.get_task_instance("plain", session=session)
+ assert mapped is not None
+ assert mapped.region_id != SENTINEL_REGION_ID
+ assert TaskInstance.get_task_instance(dr.dag_id, dr.run_id, "mapped", 1, session=session) == mapped
+ assert dr.get_task_instance("mapped", map_index=1, region_id=SENTINEL_REGION_ID, session=session) is None
+ assert dr.get_task_instance("mapped", map_index=1, region_id=mapped.region_id, session=session) == mapped
+
+ loop = make_region(dr)
+ session.add(loop)
+ session.flush()
+ plain.region_id, plain.region_index = loop.id, 2
+ session.flush()
+
+ assert dr.get_task_instance("plain", map_index=2, session=session) is None
+ assert dr.get_task_instance("plain", map_index=2, region_id=loop.id, session=session) == plain
+
+
+def _ti_search_plans(session, statements):
+ plans = []
+ for statement in statements:
+ rows = session.execute(text(f"EXPLAIN QUERY PLAN {statement}")).all()
+ plans.extend(row[-1] for row in rows if "task_instance" in row[-1] and "SEARCH" in row[-1])
+ return plans
+
+
+def test_public_lookups_use_the_task_instance_unique_key(dag_maker, session):
+ if session.get_bind().dialect.name != "sqlite":
+ pytest.skip("Reads the SQLite query plan")
+ with dag_maker(serialized=True):
+ PythonOperator.partial(task_id="mapped", python_callable=str).expand(op_args=[[i] for i in range(40)])
+ dr = dag_maker.create_dagrun()
+ key = TaskInstanceKey(dr.dag_id, "mapped", dr.run_id, map_index=7)
+ lookups = {
+ "dagrun": lambda: dr.get_task_instance("mapped", map_index=7, session=session),
+ "ti": lambda: TaskInstance.get_task_instance(dr.dag_id, dr.run_id, "mapped", 7, session=session),
+ "keys": lambda: session.scalars(select(TaskInstance).where(TaskInstance.filter_for_tis([key]))).all(),
+ "xcom": lambda: session.scalars(
+ XComModel.get_many(
+ run_id=dr.run_id, dag_ids=dr.dag_id, task_ids="mapped", map_indexes=7, key="return_value"
+ )
+ ).all(),
+ }
+
+ for name, lookup in lookups.items():
+ with capture_orm_selects("task_instance") as statements:
+ lookup()
+ plans = _ti_search_plans(session, statements)
+ assert any("region_id=? AND region_index=?" in plan for plan in plans), (name, plans)
+ assert not any("ANY(" in plan for plan in plans), (name, plans)
diff --git a/airflow-core/tests/unit/models/test_mappedoperator.py b/airflow-core/tests/unit/models/test_mappedoperator.py
index 6b48d92d854..185d58ea3aa 100644
--- a/airflow-core/tests/unit/models/test_mappedoperator.py
+++ b/airflow-core/tests/unit/models/test_mappedoperator.py
@@ -23,6 +23,7 @@ from datetime import timedelta
from typing import TYPE_CHECKING
from unittest import mock
from unittest.mock import patch
+from uuid import UUID
import pytest
from sqlalchemy import delete, event, select
@@ -32,11 +33,16 @@ from airflow.executors.base_executor import BaseExecutor
from airflow.jobs.job import Job
from airflow.jobs.scheduler_job_runner import SchedulerJobRunner
from airflow.models.dag_version import DagVersion
-from airflow.models.taskinstance import LegacyTaskDataOwner, TaskInstance
-from airflow.models.xcom import XComModelV1, build_xcom_read_query
+from airflow.models.dagbag import DBDagBag
+from airflow.models.dynamic_region import DynamicRegion
+from airflow.models.task_coordinates import TaskCoordinateResolver
+from airflow.models.taskinstance import LegacyTaskDataOwner, TaskInstance, clear_task_instances
+from airflow.models.xcom import XComModel, XComModelV1, build_xcom_read_query
from airflow.providers.standard.operators.python import PythonOperator
from airflow.sdk import DAG, BaseOperator, TaskGroup, setup, task, task_group, teardown
+from airflow.sdk.definitions._internal.loop import create_loop
from airflow.serialization.definitions.baseoperator import SerializedBaseOperator
+from airflow.serialization.definitions.mappedoperator import get_mapped_ti_count
from airflow.task.trigger_rule import TriggerRule
from airflow.utils.state import TaskInstanceState
@@ -57,6 +63,97 @@ if TYPE_CHECKING:
from airflow.sdk.definitions.context import Context
+@pytest.mark.parametrize("mapping", ["dict", "list", "group"])
+def test_mapped_count_uses_retained_producer_in_callers_loop_pass(dag_maker, session, mapping):
+ @task_group
+ def body():
+ source = PythonOperator(task_id="source", python_callable=list)
+ if mapping == "dict":
+ PythonOperator.partial(task_id="consumer", python_callable=list).expand(op_kwargs=source.output)
+ elif mapping == "list":
+ PythonOperator.partial(task_id="consumer", python_callable=list).expand_kwargs(source.output)
+ else:
+
+ @task_group
+ def mapped_group(value):
+ PythonOperator(task_id="consumer", python_callable=list)
+
+ mapped_group.expand(value=source.output)
+
+ with dag_maker(serialized=True):
+ loop = create_loop(body, max_iterations=3)
+ dr = dag_maker.create_dagrun()
+ original = DynamicRegion(dag_id=dr.dag_id, run_id=dr.run_id, node_id=loop.group_id)
+ session.add(original)
+ session.flush()
+ fork = DynamicRegion(
+ dag_id=dr.dag_id,
+ run_id=dr.run_id,
+ node_id=loop.group_id,
+ forked_from_region_id=original.id,
+ resumes_from_index=1,
+ )
+ session.add(fork)
+ session.flush()
+ source = next(ti for ti in dr.task_instances if ti.task_id == "body.source")
+ consumer = next(ti for ti in dr.task_instances if ti.task_id.endswith("consumer"))
+ source.region_id, source.region_index = original.id, 1
+ source.state = TaskInstanceState.SUCCESS
+ expansion = DynamicRegion(
+ dag_id=dr.dag_id,
+ run_id=dr.run_id,
+ node_id=consumer.task_id,
+ parent_region_id=fork.id,
+ parent_region_index=1,
+ )
+ session.add(expansion)
+ session.flush()
+ consumer.region_id, consumer.region_index = expansion.id, 0
+ earlier = TaskInstance(
+ dag_maker.serialized_dag.get_task(source.task_id),
+ dag_version_id=source.dag_version_id,
+ run_id=dr.run_id,
+ region_id=original.id,
+ region_index=0,
+ state=TaskInstanceState.SUCCESS,
+ )
+ session.add(earlier)
+ session.flush()
+ for producer, length in [(earlier, 99), (source, 2)]:
+ XComModel.set_for_attempt(
+ task_instance_id=producer.id,
+ key="return_value",
+ value=[{}] * length,
+ serialize=False,
+ mapped_length=length,
+ session=session,
+ )
+ session.flush()
+ resolver = TaskCoordinateResolver(DBDagBag(), session)
+ contexts = resolver.producer_contexts(consumer)
+
+ assert (
+ get_mapped_ti_count(
+ dag_maker.serialized_dag.get_task(consumer.task_id),
+ dr.run_id,
+ producer_contexts=contexts,
+ session=session,
+ )
+ == 2
+ )
+ if mapping == "group":
+ consumer.task = dag_maker.serialized_dag.get_task(consumer.task_id)
+ assert (
+ consumer.get_relevant_upstream_map_indexes(
+ consumer.task,
+ 2,
+ producer_contexts=contexts,
+ session=session,
+ )
+ == 0
+ )
+
+
@patch("airflow.sdk.definitions._internal.abstractoperator.AbstractOperator.render_template")
@pytest.mark.usefixtures("testing_dag_bundle")
def test_task_mapping_with_dag_and_list_of_pandas_dataframe(mock_render_template, caplog):
@@ -174,16 +271,18 @@ def test_expand_mapped_task_failed_state_in_db(dag_maker, session):
dr = dag_maker.create_dagrun()
mapped_deser = dag.task_dict[mapped.task_id]
- placeholder_id = dr.get_task_instance(mapped.task_id, session=session).id
+ placeholder = next(ti for ti in dr.task_instances if ti.task_id == mapped.task_id)
+ placeholder_id = placeholder.id
push_mapped_length(dr.get_task_instance(task1.task_id, session=session), literal, session=session)
- dag_version = DagVersion.get_latest_version(dr.dag_id)
+ dag_version = DagVersion.get_latest_version(dr.dag_id, session=session)
for index in range(2):
# Give the existing TIs a state to make sure we don't change them
ti = TaskInstance(
mapped_deser,
run_id=dr.run_id,
- map_index=index,
+ region_id=placeholder.region_id,
+ region_index=index,
state=TaskInstanceState.SUCCESS,
dag_version_id=dag_version.id,
)
@@ -241,10 +340,10 @@ def test_stable_mapped_indexes_do_not_query_historical_max_try(dag_maker, sessio
statements.append(statement)
event.listen(session.bind, "before_cursor_execute", capture_sql)
+ live = run.get_task_instance(mapped.task_id, map_index=0, session=session)
+ live.task = dag.task_dict[mapped.task_id]
try:
- new_tis = run._revise_map_indexes_if_mapped(
- dag.task_dict[mapped.task_id], dag_version_id=run.created_dag_version_id, session=session
- )
+ new_tis = run._revise_map_indexes_if_mapped(live, session=session)
finally:
event.remove(session.bind, "before_cursor_execute", capture_sql)
@@ -262,10 +361,10 @@ def test_missing_mapped_index_uses_retained_max_try(dag_maker, session):
removed = run.get_task_instance(mapped.task_id, map_index=1, session=session)
removed.try_number = 3
removed.archive(reason="retry", session=session)
+ live = run.get_task_instance(mapped.task_id, map_index=0, session=session)
+ live.task = dag.task_dict[mapped.task_id]
- new_tis = run._revise_map_indexes_if_mapped(
- dag.task_dict[mapped.task_id], dag_version_id=run.created_dag_version_id, session=session
- )
+ new_tis = run._revise_map_indexes_if_mapped(live, session=session)
assert len(new_tis) == 1
assert new_tis[0].map_index == 1
@@ -1623,10 +1722,11 @@ def test_mapped_tasks_in_mapped_task_group_waits_for_upstreams_to_complete(dag_m
tg1.expand(a=t)
dr = dag_maker.create_dagrun()
- ti = dr.get_task_instance(task_id="t1")
+ region_id = next(ti.region_id for ti in dr.task_instances if ti.task_id == "tg1.t3")
+ ti = dr.get_task_instance(task_id="t1", session=session)
run_task_instance(ti, dag.get_task(ti.task_id))
- dr.task_instance_scheduling_decisions()
- ti3 = dr.get_task_instance(task_id="tg1.t3")
+ dr.task_instance_scheduling_decisions(session=session)
+ ti3 = dr.get_task_instance(task_id="tg1.t3", region_id=region_id, session=session)
assert not ti3.state
@@ -1870,7 +1970,8 @@ def test_placeholder_promotion_keeps_legacy_owner_and_avoids_historical_try_coll
historical = TaskInstance(
serialized,
run_id=run.run_id,
- map_index=0,
+ region_id=placeholder.region_id,
+ region_index=0,
state=TaskInstanceState.FAILED,
dag_version_id=placeholder.dag_version_id,
)
@@ -1920,3 +2021,280 @@ def test_placeholder_promotion_keeps_legacy_owner_and_avoids_historical_try_coll
value = session.scalars(read).one()
assert value.map_index == 0
assert value.value == {"owner": "placeholder"}
+
+
+def _live_tis(session, dr, task_id):
+ return session.scalars(
+ select(TaskInstance)
+ .where(
+ TaskInstance.working_set.is_(True),
+ TaskInstance.dag_id == dr.dag_id,
+ TaskInstance.run_id == dr.run_id,
+ TaskInstance.task_id == task_id,
+ )
+ .order_by(TaskInstance.region_index)
+ ).all()
+
+
+def _node_regions(session, dr, task_id):
+ return session.scalars(
+ select(DynamicRegion).where(
+ DynamicRegion.dag_id == dr.dag_id,
+ DynamicRegion.run_id == dr.run_id,
+ DynamicRegion.node_id == task_id,
+ )
+ ).all()
+
+
+def _make_legacy_placeholder(session, dr, serialized):
+ """Replace the regional placeholder with one that predates regions."""
+ current = _live_tis(session, dr, serialized.task_id)[0]
+ version = current.dag_version_id
+ session.delete(current)
+ session.flush()
+ session.execute(delete(DynamicRegion).where(DynamicRegion.node_id == serialized.task_id))
+ placeholder = TaskInstance(serialized, dag_version_id=version, run_id=dr.run_id)
+ session.add(placeholder)
+ session.flush()
+ return placeholder
+
+
+@pytest.fixture
+def mapped_dag(dag_maker, session):
+ with dag_maker(session=session, serialized=True) as dag:
+ upstream = BaseOperator(task_id="upstream")
+ mapped = MockOperator.partial(task_id="mapped").expand(arg2=upstream.output)
+ dr = dag_maker.create_dagrun()
+ return dr, dag.task_dict[mapped.task_id], upstream.task_id
+
+
+def test_mapped_task_region_is_born_with_its_placeholder(mapped_dag, session):
+ dr, serialized, _ = mapped_dag
+
+ (placeholder,) = _live_tis(session, dr, serialized.task_id)
+ (region,) = _node_regions(session, dr, serialized.task_id)
+
+ assert placeholder.region_id == region.id
+ assert placeholder.region_index == -1
+ assert placeholder.map_index == -1
+ assert region.parent_region_id is None
+ assert region.forked_from_region_id is None
+
+
+def test_expansion_grows_and_shrinks_within_one_region(mapped_dag, session):
+ dr, serialized, upstream_id = mapped_dag
+ placeholder_id = _live_tis(session, dr, serialized.task_id)[0].id
+ (region,) = _node_regions(session, dr, serialized.task_id)
+
+ expand_mapped_task(serialized, dr.run_id, upstream_id, 3, session)
+ first = _live_tis(session, dr, serialized.task_id)
+ assert [ti.region_index for ti in first] == [0, 1, 2]
+ assert {ti.region_id for ti in first} == {region.id}
+ assert first[0].id == placeholder_id
+
+ push_mapped_length(dr.get_task_instance(upstream_id, session=session), list(range(5)), session=session)
+ first[0].task = serialized
+ assert [ti.region_index for ti in dr._revise_map_indexes_if_mapped(first[0], session=session)] == [3, 4]
+
+ push_mapped_length(dr.get_task_instance(upstream_id, session=session), list(range(2)), session=session)
+ assert dr._revise_map_indexes_if_mapped(first[0], session=session) == []
+ session.expire_all()
+ states = {ti.region_index: ti.state for ti in _live_tis(session, dr, serialized.task_id)}
+ assert states[0] is None
+ assert states[2] == states[3] == states[4] == TaskInstanceState.REMOVED
+ assert {ti.region_id for ti in _live_tis(session, dr, serialized.task_id)} == {region.id}
+ assert _node_regions(session, dr, serialized.task_id) == [region]
+
+
+def test_zero_length_expansion_leaves_a_region_holding_only_its_skipped_placeholder(mapped_dag, session):
+ dr, serialized, upstream_id = mapped_dag
+ (region,) = _node_regions(session, dr, serialized.task_id)
+
+ expand_mapped_task(serialized, dr.run_id, upstream_id, 0, session)
+
+ (placeholder,) = _live_tis(session, dr, serialized.task_id)
+ assert (placeholder.region_id, placeholder.region_index) == (region.id, -1)
+ assert placeholder.state == TaskInstanceState.SKIPPED
+
+
+def test_promotion_leaves_earlier_placeholder_task_instances_at_their_coordinates(mapped_dag, session):
+ dr, serialized, upstream_id = mapped_dag
+ placeholder = _live_tis(session, dr, serialized.task_id)[0]
+ region_id = placeholder.region_id
+ placeholder.task = serialized
+ placeholder.try_number = 1
+ placeholder.state = TaskInstanceState.FAILED
+ archived_id = placeholder.id
+ (current,) = clear_task_instances([placeholder], session)
+ session.flush()
+ assert current.id != archived_id
+
+ expand_mapped_task(serialized, dr.run_id, upstream_id, 2, session)
+
+ archived = session.get(TaskInstance, archived_id)
+ assert (archived.region_id, archived.region_index, archived.try_number) == (region_id, -1, 1)
+ assert archived.working_set is None
+ live = _live_tis(session, dr, serialized.task_id)
+ assert [ti.region_index for ti in live] == [0, 1]
+ assert live[0].id == current.id
+ assert live[0].try_number == current.try_number
+
+
+@pytest.fixture
+def twin_mapped_dag(dag_maker, session):
+ with dag_maker(session=session, serialized=True) as dag:
+ upstream = BaseOperator(task_id="upstream")
+ MockOperator.partial(task_id="fresh").expand(arg2=upstream.output)
+ MockOperator.partial(task_id="legacy").expand(arg2=upstream.output)
+ dr = dag_maker.create_dagrun()
+ return dr, dag, upstream.task_id
+
+
+def _expansion_outcome(session, dr, task_id, placeholder_id):
+ live = _live_tis(session, dr, task_id)
+ archived = session.scalars(
+ select(TaskInstance)
+ .where(
+ TaskInstance.working_set.is_(None),
+ TaskInstance.dag_id == dr.dag_id,
+ TaskInstance.run_id == dr.run_id,
+ TaskInstance.task_id == task_id,
+ )
+ .execution_options(include_all_attempts=True)
+ ).all()
+ return (
+ [(ti.region_index, ti.state, ti.try_number, ti.id == placeholder_id) for ti in live],
+ len({ti.region_id for ti in live}),
+ [(ti.state, ti.archived_reason) for ti in archived],
+ )
+
+
+@pytest.mark.parametrize("length", [0, 3])
+def test_legacy_placeholder_moves_into_its_new_region_without_being_archived(
+ twin_mapped_dag, session, length
+):
+ dr, dag, upstream_id = twin_mapped_dag
+ fresh = dag.task_dict["fresh"]
+ serialized = dag.task_dict["legacy"]
+ fresh_placeholder_id = _live_tis(session, dr, "fresh")[0].id
+ legacy = _make_legacy_placeholder(session, dr, serialized)
+ legacy_id = legacy.id
+
+ push_mapped_length(
+ dr.get_task_instance(upstream_id, session=session), list(range(length)), session=session
+ )
+ expand_mapped_task_instances(fresh, dr.run_id, session=session)
+ expand_mapped_task_instances(serialized, dr.run_id, session=session)
+
+ (region,) = _node_regions(session, dr, serialized.task_id)
+ live = _live_tis(session, dr, serialized.task_id)
+ assert {ti.region_id for ti in live} == {region.id}
+ assert live[0].id == legacy_id
+ assert live[0].region_index == (0 if length else -1)
+ assert (
+ session.scalar(select(TaskInstance.id).where(TaskInstance.archived_reason == "mapped_placeholder"))
+ is None
+ )
+ assert _expansion_outcome(session, dr, "legacy", legacy_id) == _expansion_outcome(
+ session, dr, "fresh", fresh_placeholder_id
+ )
+ assert _expansion_outcome(session, dr, "legacy", legacy_id)[2] == []
+
+
+def test_legacy_placeholder_keeps_its_earlier_tries_and_is_promoted_in_place(twin_mapped_dag, session):
+ dr, dag, upstream_id = twin_mapped_dag
+ serialized = dag.task_dict["legacy"]
+ earlier = _make_legacy_placeholder(session, dr, serialized)
+ earlier.try_number = 1
+ earlier.state = TaskInstanceState.UPSTREAM_FAILED
+ earlier.archive(reason="retry", session=session)
+ legacy = TaskInstance(serialized, dag_version_id=earlier.dag_version_id, run_id=dr.run_id)
+ legacy.try_number = 2
+ session.add(legacy)
+ session.flush()
+ legacy_id = legacy.id
+
+ push_mapped_length(dr.get_task_instance(upstream_id, session=session), [1, 2], session=session)
+ expand_mapped_task_instances(serialized, dr.run_id, session=session)
+
+ (region,) = _node_regions(session, dr, serialized.task_id)
+ live = _live_tis(session, dr, serialized.task_id)
+ assert [(ti.region_id, ti.region_index) for ti in live] == [(region.id, 0), (region.id, 1)]
+ assert (live[0].id, live[0].try_number) == (legacy_id, 2)
+ assert (earlier.region_id, earlier.region_index, earlier.try_number) == (UUID(int=0), -1, 1)
+ assert earlier.archived_reason == "retry"
+
+
+def test_clearing_a_legacy_expansion_keeps_it_in_the_sentinel_region(mapped_dag, session):
+ dr, serialized, _ = mapped_dag
+ _make_legacy_placeholder(session, dr, serialized)
+ legacy = [
+ TaskInstance(
+ serialized,
+ dag_version_id=dr.created_dag_version_id,
+ run_id=dr.run_id,
+ region_index=index,
+ state=TaskInstanceState.SUCCESS,
+ )
+ for index in range(2)
+ ]
+ session.add_all(legacy)
+ session.flush()
+ for ti in legacy:
+ ti.task = serialized
+
+ (cleared,) = clear_task_instances([legacy[0]], session)
+ session.flush()
+
+ assert cleared.region_id == UUID(int=0)
+ assert (cleared.region_index, cleared.try_number) == (0, 1)
+ assert cleared.state is None
+ assert legacy[1].state == TaskInstanceState.SUCCESS
+ assert _node_regions(session, dr, serialized.task_id) == []
+ archived = session.scalars(
+ select(TaskInstance)
+ .where(TaskInstance.working_set.is_(None), TaskInstance.task_id == serialized.task_id)
+ .execution_options(include_all_attempts=True)
+ ).one()
+ assert (archived.region_id, archived.region_index, archived.try_number) == (UUID(int=0), 0, 0)
+
+
+def test_expansion_nested_under_a_loop_pass_is_a_child_region_that_revises_in_place(dag_maker, session):
+ with dag_maker(session=session, serialized=True):
+ mapped = MockOperator.partial(task_id="mapped").expand(arg2=[1, 2, 3])
+ dr = dag_maker.create_dagrun()
+ serialized = dag_maker.serialized_dag.get_task(mapped.task_id)
+ session.execute(delete(TaskInstance).where(TaskInstance.task_id == mapped.task_id))
+ session.execute(delete(DynamicRegion))
+ loop_region = DynamicRegion(dag_id=dr.dag_id, run_id=dr.run_id, node_id="loop")
+ session.add(loop_region)
+ session.flush()
+ created: list[TaskInstance] = []
+
+ def creator(task, indexes, region_id):
+ for index in indexes:
+ ti = TaskInstance(
+ task,
+ dag_version_id=dr.created_dag_version_id,
+ run_id=dr.run_id,
+ region_id=region_id,
+ region_index=index,
+ )
+ created.append(ti)
+ yield ti
+
+ session.add_all(
+ list(dr._create_tasks([serialized], creator, session=session, parent_region=(loop_region.id, 2)))
+ )
+ session.flush()
+ (child,) = _node_regions(session, dr, mapped.task_id)
+ assert (child.parent_region_id, child.parent_region_index) == (loop_region.id, 2)
+ (placeholder,) = created
+ assert (placeholder.region_id, placeholder.region_index) == (child.id, -1)
+
+ placeholder.expand_mapped_task(session=session)
+
+ live = _live_tis(session, dr, mapped.task_id)
+ assert [ti.region_index for ti in live] == [0, 1, 2]
+ assert {ti.region_id for ti in live} == {child.id}
+ assert _node_regions(session, dr, mapped.task_id) == [child]
diff --git a/airflow-core/tests/unit/models/test_task_data.py b/airflow-core/tests/unit/models/test_task_data.py
index f647ad88585..a46f2c07898 100644
--- a/airflow-core/tests/unit/models/test_task_data.py
+++ b/airflow-core/tests/unit/models/test_task_data.py
@@ -49,6 +49,11 @@ from tests_common.test_utils.attempt_ownership import (
pytestmark = [pytest.mark.db_test, pytest.mark.execution_timeout(10)]
+def ti_coordinates(coordinates):
+ """Rename the legacy ``map_index`` key to the live ``region_index`` column."""
+ return {("region_index" if key == "map_index" else key): value for key, value in coordinates.items()}
+
+
def producers(session, *ids):
ti = table(session.connection(), "task_instance", "id")
return sa.select(ti.c.id).where(ti.c.id.in_(ids))
@@ -70,7 +75,7 @@ def seed_many_legacy_attempts(session):
session.execute(
ti.insert(),
[
- dict(**coord, id=attempt_id, try_number=1, pool="default_pool", pool_slots=1)
+ dict(**ti_coordinates(coord), id=attempt_id, try_number=1, pool="default_pool", pool_slots=1)
for coord, attempt_id in zip(coordinates, ids)
],
)
@@ -158,7 +163,12 @@ def add_other_task_with_legacy_xcom(session, run_id):
coordinates = {**COORDINATES, "task_id": "other", "run_id": run_id}
session.execute(
TaskInstance.__table__.insert().values(
- **coordinates, id=OTHER_ID, try_number=1, pool="default_pool", pool_slots=1, working_set=True
+ **ti_coordinates(coordinates),
+ id=OTHER_ID,
+ try_number=1,
+ pool="default_pool",
+ pool_slots=1,
+ working_set=True,
)
)
session.execute(LegacyTaskDataOwner.__table__.insert().values(**coordinates, task_instance_id=OTHER_ID))
@@ -314,7 +324,7 @@ def test_delete_removes_both_copies_without_touching_a_successor(ownership_sessi
dag_id="ownership",
task_id="task",
run_id="manual",
- map_index=-1,
+ region_index=-1,
try_number=3,
pool="default_pool",
pool_slots=1,
@@ -355,7 +365,7 @@ def test_read_projection_cannot_be_flushed_as_a_physical_row(ownership_session,
def test_legacy_lookup_uses_fixed_owner_coordinates_after_mapping_promotion(ownership_session):
session = ownership_session
ti = table(session.connection(), "task_instance", "id")
- session.execute(ti.update().where(ti.c.id == CURRENT_ID).values(map_index=0))
+ session.execute(ti.update().where(ti.c.id == CURRENT_ID).values(region_index=0))
rows = read(session, CURRENT_ID)
assert len(rows) == 1
assert rows[0].map_index == 0
@@ -527,7 +537,7 @@ def test_rendered_retention_deletes_both_stores_and_keeps_a_whole_mapped_run(own
dag_id="ownership",
task_id="task",
run_id="next",
- map_index=map_index,
+ region_index=map_index,
try_number=1,
pool="default_pool",
pool_slots=1,
@@ -560,7 +570,7 @@ def test_rendered_retention_deletes_both_stores_and_keeps_a_whole_mapped_run(own
session.execute(
ti.insert().values(
id=attempt_id,
- **coordinates,
+ **ti_coordinates(coordinates),
try_number=1,
pool="default_pool",
pool_slots=1,
diff --git a/airflow-core/tests/unit/models/test_task_state_store.py b/airflow-core/tests/unit/models/test_task_state_store.py
index dfec5bbbee8..c4032c38874 100644
--- a/airflow-core/tests/unit/models/test_task_state_store.py
+++ b/airflow-core/tests/unit/models/test_task_state_store.py
@@ -67,7 +67,7 @@ class TestTaskStateStoreModel:
row = TaskStateStoreModel(
dag_run_id=dag_run.id,
task_id=TASK_ID,
- map_index=-1,
+ region_index=-1,
key="remote_job_id",
dag_id=DAG_ID,
run_id=RUN_ID,
@@ -80,7 +80,7 @@ class TestTaskStateStoreModel:
select(TaskStateStoreModel).where(
TaskStateStoreModel.dag_run_id == dag_run.id,
TaskStateStoreModel.task_id == TASK_ID,
- TaskStateStoreModel.map_index == -1,
+ TaskStateStoreModel.region_index == -1,
TaskStateStoreModel.key == "remote_job_id",
)
)
@@ -95,7 +95,7 @@ class TestTaskStateStoreModel:
TaskStateStoreModel(
dag_run_id=dag_run.id,
task_id=TASK_ID,
- map_index=-1,
+ region_index=-1,
key="remote_job_id",
dag_id=DAG_ID,
run_id=RUN_ID,
@@ -108,7 +108,7 @@ class TestTaskStateStoreModel:
TaskStateStoreModel(
dag_run_id=dag_run.id,
task_id=TASK_ID,
- map_index=-1,
+ region_index=-1,
key="remote_job_id",
dag_id=DAG_ID,
run_id=RUN_ID,
@@ -123,7 +123,7 @@ class TestTaskStateStoreModel:
row = TaskStateStoreModel(
dag_run_id=dag_run.id,
task_id=TASK_ID,
- map_index=-1,
+ region_index=-1,
key="remote_job_id",
dag_id=DAG_ID,
run_id=RUN_ID,
@@ -152,7 +152,7 @@ class TestTaskStateStoreModel:
TaskStateStoreModel(
dag_run_id=dag_run.id,
task_id=TASK_ID,
- map_index=idx,
+ region_index=idx,
key="checkpoint",
dag_id=DAG_ID,
run_id=RUN_ID,
@@ -166,7 +166,7 @@ class TestTaskStateStoreModel:
select(TaskStateStoreModel).where(
TaskStateStoreModel.dag_run_id == dag_run.id,
TaskStateStoreModel.task_id == TASK_ID,
- TaskStateStoreModel.map_index == idx,
+ TaskStateStoreModel.region_index == idx,
TaskStateStoreModel.key == "checkpoint",
)
)
@@ -179,7 +179,7 @@ class TestTaskStateStoreModel:
TaskStateStoreModel(
dag_run_id=dag_run.id,
task_id=TASK_ID,
- map_index=-1,
+ region_index=-1,
key="remote_job_id",
dag_id=DAG_ID,
run_id=RUN_ID,
@@ -212,7 +212,7 @@ class TestTaskStateStoreModel:
TaskStateStoreModel(
dag_run_id=dag_run.id,
task_id=TASK_ID,
- map_index=-1,
+ region_index=-1,
key="watermark",
dag_id=DAG_ID,
run_id=RUN_ID,
diff --git a/airflow-core/tests/unit/models/test_taskinstance.py b/airflow-core/tests/unit/models/test_taskinstance.py
index 54cf9121301..f582ddc3a10 100644
--- a/airflow-core/tests/unit/models/test_taskinstance.py
+++ b/airflow-core/tests/unit/models/test_taskinstance.py
@@ -73,7 +73,7 @@ from airflow.models.taskinstance import (
find_relevant_relatives,
)
from airflow.models.taskreschedule import TaskReschedule
-from airflow.models.xcom import XComModel, XComModelV2, build_xcom_read_query
+from airflow.models.xcom import LazyXComSelectSequence, XComModel, XComModelV2, build_xcom_read_query
from airflow.providers.standard.operators.bash import BashOperator
from airflow.providers.standard.operators.empty import EmptyOperator
from airflow.providers.standard.operators.python import PythonOperator
@@ -188,6 +188,49 @@ def clean_db():
db.clear_db_assets()
+@pytest.mark.parametrize("bulk", [False, True])
+@pytest.mark.parametrize(
+ "index_fields", [{"region_index": 4}, {"map_index": 4}, {"map_index": 4, "region_index": 4}]
+)
+def test_canonical_region_index_is_persisted_for_new_task_instances(dag_maker, session, bulk, index_fields):
+ with dag_maker(serialized=True) as dag:
+ EmptyOperator(task_id="canonical")
+ dr = dag_maker.create_dagrun()
+ original = dr.task_instances[0]
+ task = dag.get_task("canonical")
+ kwargs = dict(task=task, run_id=dr.run_id, dag_version_id=original.dag_version_id, **index_fields)
+ if bulk:
+ row = TaskInstance.insert_mapping(**kwargs, dag_run=dr)
+ session.bulk_insert_mappings(TaskInstance, [row])
+ else:
+ session.add(TaskInstance(**kwargs))
+ session.flush()
+
+ ti = session.scalar(select(TaskInstance).where(TaskInstance.region_index == 4))
+ assert ti is not None
+ assert ti.id != original.id
+ assert ti.region_index == 4
+
+
+@pytest.mark.parametrize("bulk", [False, True])
+def test_conflicting_task_instance_index_names_are_rejected(dag_maker, bulk):
+ with dag_maker(serialized=True) as dag:
+ EmptyOperator(task_id="canonical")
+ dr = dag_maker.create_dagrun()
+ kwargs = dict(
+ task=dag.get_task("canonical"),
+ run_id=dr.run_id,
+ dag_version_id=dr.task_instances[0].dag_version_id,
+ map_index=2,
+ region_index=3,
+ )
+ constructor = TaskInstance.insert_mapping if bulk else TaskInstance
+ if bulk:
+ kwargs["dag_run"] = dr
+ with pytest.raises(ValueError, match="map_index.*region_index"):
+ constructor(**kwargs)
+
+
class TestTaskInstance:
def setup_method(self):
# We don't want to store any code for (test) dags created in this file
@@ -1382,20 +1425,33 @@ class TestTaskInstance:
dr = dag_maker.create_dagrun()
dag_maker.session.commit()
monkeypatch.setattr(_UpstreamTIStates, "calculate", lambda *_: upstream_states)
- ti = dr.get_task_instance("do_something_else", session=session)
+ ti = session.scalars(
+ select(TaskInstance).where(
+ TaskInstance.dag_id == dr.dag_id,
+ TaskInstance.run_id == dr.run_id,
+ TaskInstance.task_id == "do_something_else",
+ )
+ ).one()
+ region_id = ti.region_id
ti.map_index = 0
base_task = ti.task
for map_index in range(1, 5):
ti = TaskInstance(
- base_task, run_id=dr.run_id, map_index=map_index, dag_version_id=ti.dag_version_id
+ base_task,
+ run_id=dr.run_id,
+ map_index=map_index,
+ region_id=region_id,
+ dag_version_id=ti.dag_version_id,
)
session.add(ti)
ti.dag_run = dr
session.flush()
session.commit()
downstream = ti.task
- ti = dr.get_task_instance(task_id="do_something_else", map_index=3, session=session)
+ ti = dr.get_task_instance(
+ task_id="do_something_else", map_index=3, region_id=region_id, session=session
+ )
ti.task = downstream
dep_results = TriggerRuleDep()._evaluate_trigger_rule(
ti=ti,
@@ -1403,7 +1459,9 @@ class TestTaskInstance:
session=session,
)
completed = all(dep.passed for dep in dep_results)
- ti = dr.get_task_instance(task_id="do_something_else", map_index=3, session=session)
+ ti = dr.get_task_instance(
+ task_id="do_something_else", map_index=3, region_id=region_id, session=session
+ )
assert completed == expect_completed
assert ti.state == expect_state
@@ -2616,7 +2674,7 @@ class TestTaskInstance:
"dag_id": "test_refresh_from_db_dag",
"run_id": "test",
"region_id": uuid6.UUID(int=0),
- "map_index": -1,
+ "region_index": -1,
"start_date": run_date + datetime.timedelta(days=1),
"end_date": run_date + datetime.timedelta(days=1, seconds=1, milliseconds=234),
"duration": 1.234,
@@ -2910,7 +2968,7 @@ class TestTaskInstance:
dag_id=attempt.dag_id,
task_id=attempt.task_id,
run_id=attempt.run_id,
- map_index=map_index,
+ region_index=map_index,
try_number=number,
working_set=working_set,
state=TaskInstanceState.FAILED,
@@ -3760,21 +3818,12 @@ class TestMappedTaskInstanceReceiveValue:
dag_maker.run_ti(ti.task_id, dag_run=dag_run, map_index=ti.map_index, session=session)
assert outputs == expected_outputs
- def test_map_xcom_wide_batched_expand(self, dag_maker, session):
- """Wide XCom-driven expand goes through the batched add_all()/flush() path.
-
- Exercises ``TaskInstance.expand_mapped_task`` over a 20-element upstream XCom and
- asserts the batched expansion creates exactly N mapped TIs with contiguous
- ``map_index`` 0..N-1, the expected ``None`` (schedulable) state, and that the
- returned instances are usable: they keep their ``.task`` (no merge() that drops
- it), are attached to the session, and have ``dag_run`` primed so a later
- ``ti.get_dagrun()`` -- as dependency evaluation makes per index -- is a cache hit
- rather than an N+1 SELECT. Also pins the expansion call's query count.
- """
+ @pytest.mark.parametrize("width", [20, 200])
+ def test_map_xcom_wide_batched_expand(self, dag_maker, session, width):
+ """Expansion and subsequent DagRun access use bounded queries at either width."""
from sqlalchemy import event
from sqlalchemy.orm.base import NO_VALUE
- width = 20
upstream_return = list(range(width))
with dag_maker(dag_id="xcom_wide", session=session, serialized=True) as dag:
@@ -3795,34 +3844,21 @@ class TestMappedTaskInstanceReceiveValue:
dag_maker.run_ti(emit_ti.task_id, dag_run=dag_run, session=session)
show_task = dag_maker.serialized_dag.get_task("show")
- # Pins the query count so a regression back to per-index session.merge() -- which
- # would issue a merge-load + reload SELECT per index -- fails this test, not just
- # a slower one. Measured at 7 for this fixture; margin allows for minor backend
- # differences while staying far below what a per-index merge() would cost.
- with assert_queries_count(7, margin=2):
+ with assert_queries_count(11, margin=2):
mapped_tis, max_map_index = expand_mapped_task_instances(
show_task, dag_run.run_id, session=session
)
- # Correct count + contiguous indexes 0..N-1.
assert len(mapped_tis) == width
assert max_map_index + 1 == width
assert sorted(ti.map_index for ti in mapped_tis) == list(range(width))
- # Freshly-expanded mapped TIs are schedulable (state None) and are attached to
- # the session. The batched path (indexes >= 1) additionally keeps its ``.task``
- # because we never merge(); index 0 is the repurposed unmapped TI (loaded from
- # the DB, so ``.task`` is None) which is stock behaviour untouched by this change.
for ti in mapped_tis:
assert ti.state is None
assert ti in session
- batched_tis = [ti for ti in mapped_tis if ti.map_index >= 1]
- assert len(batched_tis) == width - 1
- for ti in batched_tis:
assert ti.task is not None
assert sa_inspect(ti).attrs.dag_run.loaded_value is not NO_VALUE
- # And they are actually persisted as rows.
persisted = session.scalars(
select(TI)
.where(TI.task_id == "show", TI.dag_id == dag_run.dag_id, TI.run_id == dag_run.run_id)
@@ -3840,7 +3876,7 @@ class TestMappedTaskInstanceReceiveValue:
event.listen(session.bind, "after_cursor_execute", _on_cursor_execute)
try:
- for ti in batched_tis:
+ for ti in mapped_tis:
returned = ti.get_dagrun(session=session)
assert returned.run_id == dag_run.run_id
finally:
@@ -3894,14 +3930,11 @@ class TestMappedTaskInstanceReceiveValue:
)
session.flush()
- # Pins the query count so a regression back to per-index session.merge() -- which
- # would issue a merge-load + reload SELECT per index -- fails this test. Measured
- # at 3 for this fixture (2 new indexes); margin allows for minor backend
- # differences while staying below what a per-index merge() would cost.
- with assert_queries_count(3, margin=1):
+ with assert_queries_count(5, margin=1):
new_tis = list(
dag_run._revise_map_indexes_if_mapped(
- show_task, dag_version_id=mapped_tis[0].dag_version_id, session=session
+ mapped_tis[0],
+ session=session,
)
)
# Only the two brand-new indexes (3, 4) are created, contiguous and usable.
@@ -4051,7 +4084,6 @@ class TestMappedTaskInstanceReceiveValue:
Test that xcom_pull returns LazyXComSelectSequence when XComs are mapped (map_index >= 0)
and map_indexes is not specified.
"""
- from airflow.models.xcom import LazyXComSelectSequence
with dag_maker(dag_id="test_xcom_mapped_values", session=session):
diff --git a/airflow-core/tests/unit/serialization/test_dag_serialization.py b/airflow-core/tests/unit/serialization/test_dag_serialization.py
index 8de02d0fb9f..3cd501ac2e0 100644
--- a/airflow-core/tests/unit/serialization/test_dag_serialization.py
+++ b/airflow-core/tests/unit/serialization/test_dag_serialization.py
@@ -1677,6 +1677,7 @@ class TestStringifiedDAGs:
"retries": 0,
"retry_delay": timedelta(0, 300),
"retry_exponential_backoff": 0,
+ "returns_dag_result": False,
"run_as_user": None,
"start_date": None,
"start_from_trigger": False,
@@ -1796,6 +1797,31 @@ class TestStringifiedDAGs:
== "<callable unit.serialization.test_dag_serialization.TestStringifiedDAGs.test_template_field_via_callable_serialization.<locals>.fn_returns_callable.<locals>.get_arg>"
)
+ @pytest.mark.parametrize("mapped", [False, True])
+ @pytest.mark.parametrize("is_result", [False, True])
+ def test_dag_result_declaration_survives_serialization(self, mapped, is_result):
+ @task_decorator
+ def value(v):
+ return v
+
+ with DAG("serialized_dag_result", schedule=None) as dag:
+ output = value.expand(v=[1, 2]) if mapped else value(1)
+ if is_result:
+ dag.add_result(output)
+
+ encoded = DagSerialization.to_dict(dag)
+ DagSerialization.validate_schema(encoded)
+ restored = DagSerialization.from_json(DagSerialization.to_json(dag))
+ assert restored.get_task("value").returns_dag_result is is_result
+
+ operator = (
+ OperatorSerialization.serialize_mapped_operator(output.operator)
+ if mapped
+ else OperatorSerialization.serialize_operator(output.operator)
+ )
+ operator.pop("returns_dag_result", None)
+ assert OperatorSerialization.deserialize_operator(operator).returns_dag_result is False
+
@pytest.mark.parametrize("conditional", [False, True])
@pytest.mark.parametrize("mapped_terminal", [False, True])
def test_loop_group_serialization(self, conditional, mapped_terminal):
diff --git a/airflow-core/tests/unit/state/test_metastore.py b/airflow-core/tests/unit/state/test_metastore.py
index 27048c70a00..da8bea41976 100644
--- a/airflow-core/tests/unit/state/test_metastore.py
+++ b/airflow-core/tests/unit/state/test_metastore.py
@@ -123,7 +123,10 @@ class TestMetastoreBackendTaskScope:
def test_operations_isolate_colliding_regions(self, session, backend, dag_run, operation):
regions = [UUID(int=0), uuid4(), uuid4()]
expected = {
- (TaskScope(DAG_ID, RUN_ID, TASK_ID, index, region_id=region), key): f"{region}:{index}:{key}"
+ (
+ TaskScope(DAG_ID, RUN_ID, TASK_ID, region_index=index, region_id=region),
+ key,
+ ): f"{region}:{index}:{key}"
for region in regions
for index in (0, 1)
for key in ("value", "other")
@@ -621,7 +624,10 @@ class TestMetastoreBackendAsync:
async def test_operations_isolate_colliding_regions(self, backend, dag_run_committed, operation):
regions = [UUID(int=0), uuid4(), uuid4()]
expected = {
- (TaskScope(DAG_ID, RUN_ID, TASK_ID, index, region_id=region), key): f"{region}:{index}:{key}"
+ (
+ TaskScope(DAG_ID, RUN_ID, TASK_ID, region_index=index, region_id=region),
+ key,
+ ): f"{region}:{index}:{key}"
for region in regions
for index in (0, 1)
for key in ("value", "other")
diff --git a/airflow-core/tests/unit/ti_deps/deps/test_mapped_task_upstream_dep.py b/airflow-core/tests/unit/ti_deps/deps/test_mapped_task_upstream_dep.py
index d4d7ad7ee2a..ef8bfbe6cf0 100644
--- a/airflow-core/tests/unit/ti_deps/deps/test_mapped_task_upstream_dep.py
+++ b/airflow-core/tests/unit/ti_deps/deps/test_mapped_task_upstream_dep.py
@@ -21,8 +21,12 @@ from typing import TYPE_CHECKING
import pytest
+from airflow.models.dynamic_region import DynamicRegion
from airflow.models.xcom import XCOM_RETURN_KEY
from airflow.providers.standard.operators.empty import EmptyOperator
+from airflow.providers.standard.operators.python import PythonOperator
+from airflow.sdk import task_group
+from airflow.sdk.definitions._internal.loop import create_loop
from airflow.sdk.exceptions import AirflowFailException, AirflowSkipException
from airflow.ti_deps.dep_context import DepContext
from airflow.ti_deps.deps.base_ti_dep import TIDepStatus
@@ -46,6 +50,31 @@ SUCCESS = TaskInstanceState.SUCCESS
UPSTREAM_FAILED = TaskInstanceState.UPSTREAM_FAILED
+def test_mapping_dependency_checks_unmapped_loop_producer_at_positive_pass(dag_maker, session):
+ @task_group
+ def body():
+ source = PythonOperator(task_id="source", python_callable=list)
+ PythonOperator.partial(task_id="consumer", python_callable=list).expand(op_kwargs=source.output)
+
+ with dag_maker(serialized=True):
+ loop = create_loop(body, max_iterations=3)
+ dr = dag_maker.create_dagrun()
+ region = DynamicRegion(dag_id=dr.dag_id, run_id=dr.run_id, node_id=loop.group_id)
+ session.add(region)
+ session.flush()
+ source = next(ti for ti in dr.task_instances if ti.task_id == "body.source")
+ consumer = next(ti for ti in dr.task_instances if ti.task_id == "body.consumer")
+ source.region_id, source.region_index, source.state = region.id, 1, FAILED
+ consumer.region_id, consumer.region_index = region.id, 1
+ consumer.task = dag_maker.serialized_dag.get_task(consumer.task_id)
+ session.flush()
+
+ statuses = list(MappedTaskUpstreamDep()._get_dep_statuses(consumer, DepContext(), session=session))
+
+ assert any(not status.passed for status in statuses)
+ assert consumer.state == UPSTREAM_FAILED
+
+
@pytest.mark.parametrize(
("task_state", "upstream_states", "expected_state", "expect_failed_dep"),
[
@@ -151,7 +180,7 @@ def test_mapped_task_upstream_dep(
]
)
assert get_dep_statuses(dr, mapped_task, session) == expected_statuses
- ti = dr.get_task_instance(session=session, task_id=mapped_task)
+ ti = dr.get_task_instance(session=session, task_id=mapped_task, region_id=tis[mapped_task].region_id)
assert ti is not None
assert ti.state == expected_state
diff --git a/airflow-core/tests/unit/ti_deps/deps/test_ready_to_reschedule_dep.py b/airflow-core/tests/unit/ti_deps/deps/test_ready_to_reschedule_dep.py
index 04f13388a5f..8c7ce7d78ac 100644
--- a/airflow-core/tests/unit/ti_deps/deps/test_ready_to_reschedule_dep.py
+++ b/airflow-core/tests/unit/ti_deps/deps/test_ready_to_reschedule_dep.py
@@ -25,7 +25,11 @@ import time_machine
from slugify import slugify
from airflow._shared.timezones import timezone
+from airflow.models.dynamic_region import DynamicRegion
from airflow.models.taskreschedule import TaskReschedule
+from airflow.providers.standard.operators.empty import EmptyOperator
+from airflow.sdk import task_group
+from airflow.sdk.definitions._internal.loop import create_loop
from airflow.ti_deps.dep_context import DepContext
from airflow.ti_deps.deps.ready_to_reschedule import ReadyToRescheduleDep
from airflow.utils.session import create_session
@@ -39,6 +43,28 @@ pytestmark = pytest.mark.db_test
DEFAULT_DATE = timezone.datetime(2016, 1, 1)
+def test_ordinary_loop_pass_skips_reschedule_query(dag_maker, session, mocker):
+ @task_group
+ def body():
+ EmptyOperator(task_id="ordinary")
+
+ with dag_maker(serialized=True):
+ loop = create_loop(body, max_iterations=3)
+ dr = dag_maker.create_dagrun()
+ region = DynamicRegion(dag_id=dr.dag_id, run_id=dr.run_id, node_id=loop.group_id)
+ session.add(region)
+ session.flush()
+ ti = dr.task_instances[0]
+ ti.region_id, ti.region_index = region.id, 2
+ ti.task = dag_maker.serialized_dag.get_task(ti.task_id)
+ session.flush()
+ query = mocker.patch.object(TaskReschedule, "stmt_for_task_instance", autospec=True)
+
+ assert ReadyToRescheduleDep().is_met(ti=ti, dep_context=DepContext(), session=session)
+
+ query.assert_not_called()
+
+
@pytest.fixture
def not_expected_tr_db_call():
def side_effect(*args, **kwargs):
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 1c4b2141c0b..fb95c8ea2ce 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
@@ -26,14 +26,17 @@ from unittest.mock import Mock
import attrs
import pytest
-from sqlalchemy import event
+from sqlalchemy import delete, event
import airflow.settings
from airflow.models.dag_version import DagVersion
+from airflow.models.dynamic_region import SENTINEL_REGION_ID, DynamicRegion
from airflow.models.taskinstance import TaskInstance
from airflow.providers.standard.operators.empty import EmptyOperator
+from airflow.providers.standard.operators.python import PythonOperator
from airflow.sdk import task, task_group
from airflow.sdk.bases.operator import BaseOperator
+from airflow.sdk.definitions._internal.loop import create_loop
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, _UpstreamTIStates
@@ -53,6 +56,117 @@ SUCCESS = TaskInstanceState.SUCCESS
FAILED = TaskInstanceState.FAILED
+@pytest.mark.parametrize("current_state", [SUCCESS, FAILED])
+def test_trigger_rule_ignores_other_loop_passes(dag_maker, session, current_state):
+ @task_group
+ def body():
+ EmptyOperator(task_id="source") >> EmptyOperator(task_id="consumer")
+
+ with dag_maker(serialized=True):
+ loop = create_loop(body, max_iterations=3)
+ dr = dag_maker.create_dagrun()
+ region = DynamicRegion(dag_id=dr.dag_id, run_id=dr.run_id, node_id=loop.group_id)
+ session.add(region)
+ session.flush()
+ source, consumer = (
+ next(ti for ti in dr.task_instances if ti.task_id == f"body.{name}")
+ for name in ("source", "consumer")
+ )
+ source.region_id, source.region_index, source.state = region.id, 1, current_state
+ consumer.region_id, consumer.region_index = region.id, 1
+ consumer.task = dag_maker.serialized_dag.get_task(consumer.task_id)
+ previous = TaskInstance(
+ task=dag_maker.serialized_dag.get_task(source.task_id),
+ run_id=dr.run_id,
+ dag_version_id=source.dag_version_id,
+ region_id=region.id,
+ map_index=0,
+ )
+ previous.state = FAILED if current_state == SUCCESS else SUCCESS
+ session.add(previous)
+ session.flush()
+ statuses = list(
+ TriggerRuleDep()._get_dep_statuses(
+ consumer,
+ DepContext(flag_upstream_failed=True),
+ session=session,
+ )
+ )
+
+ assert all(status.passed for status in statuses) is (current_state == SUCCESS)
+ assert consumer.state == (None if current_state == SUCCESS else UPSTREAM_FAILED)
+
+
+@pytest.mark.parametrize("tail_state", [TaskInstanceState.RUNNING, REMOVED])
+def test_trigger_count_cache_separates_expansions_in_different_loop_passes(dag_maker, session, tail_state):
+ @task_group
+ def body():
+ source = PythonOperator.partial(task_id="source", python_callable=list).expand(op_kwargs=[{}, {}])
+ source >> EmptyOperator(task_id="consumer")
+
+ with dag_maker(serialized=True):
+ loop = create_loop(body, max_iterations=3)
+ dr = dag_maker.create_dagrun()
+ parent = DynamicRegion(dag_id=dr.dag_id, run_id=dr.run_id, node_id=loop.group_id)
+ for ti in dr.task_instances:
+ ti.region_id = SENTINEL_REGION_ID
+ session.flush()
+ session.execute(delete(DynamicRegion).where(DynamicRegion.dag_id == dr.dag_id))
+ context = DepContext(flag_upstream_failed=True)
+ assert not context.has_regions(dr.task_instances[0], session=session)
+ session.add(parent)
+ session.flush()
+ expansions = [
+ DynamicRegion(
+ dag_id=dr.dag_id,
+ run_id=dr.run_id,
+ node_id="body.source",
+ parent_region_id=parent.id,
+ parent_region_index=index,
+ )
+ for index in (0, 1)
+ ]
+ session.add_all(expansions)
+ session.flush()
+ sources = [ti for ti in dr.task_instances if ti.task_id == "body.source"]
+ for source, expansion in zip(sources, expansions):
+ source.region_id, source.region_index, source.state = expansion.id, 0, SUCCESS
+ source_task = dag_maker.serialized_dag.get_task("body.source")
+ pending = TaskInstance(
+ task=source_task,
+ run_id=dr.run_id,
+ dag_version_id=sources[0].dag_version_id,
+ region_id=expansions[1].id,
+ map_index=1,
+ state=tail_state,
+ )
+ consumer = next(ti for ti in dr.task_instances if ti.task_id == "body.consumer")
+ consumer.region_id, consumer.region_index = parent.id, 0
+ consumer.task = dag_maker.serialized_dag.get_task(consumer.task_id)
+ following = TaskInstance(
+ task=consumer.task,
+ run_id=dr.run_id,
+ dag_version_id=consumer.dag_version_id,
+ region_id=parent.id,
+ map_index=1,
+ )
+ session.add_all([pending, following])
+ session.flush()
+ context.invalidate_upstream_task_id_counts()
+ with _count_upstream_count_queries() as counter:
+ first_statuses = list(
+ TriggerRuleDep()._evaluate_trigger_rule(ti=consumer, dep_context=context, session=session)
+ )
+ next_statuses = list(
+ TriggerRuleDep()._evaluate_trigger_rule(ti=following, dep_context=context, session=session)
+ )
+
+ assert first_statuses == []
+ assert any(not status.passed for status in next_statuses) is (tail_state == TaskInstanceState.RUNNING)
+ assert following.state is None
+ assert counter["n"] == 2
+
+
@pytest.fixture
def get_task_instance(monkeypatch, session, dag_maker):
def _get_task_instance(
@@ -139,14 +253,20 @@ def get_mapped_task_dagrun(session, dag_maker):
dr = dag_maker.create_dagrun()
def _expand_tasks(task_instance: str, upstream: str) -> BaseOperator | None:
- ti = dr.get_task_instance(task_instance, session=session)
- ti.map_index = 0
- dag_version = DagVersion.get_latest_version(dag.dag_id)
+ ti = next(ti for ti in dr.task_instances if ti.task_id == task_instance)
+ ti.region_index = 0
+ expanded_task = ti.task
+ region_id = ti.region_id
+ dag_version = DagVersion.get_latest_version(dag.dag_id, session=session)
if TYPE_CHECKING:
assert dag_version
for map_index in range(1, 5):
ti = TaskInstance(
- ti.task, run_id=dr.run_id, map_index=map_index, dag_version_id=dag_version.id
+ expanded_task,
+ run_id=dr.run_id,
+ region_id=region_id,
+ region_index=map_index,
+ dag_version_id=dag_version.id,
)
session.add(ti)
ti.dag_run = dr
@@ -154,12 +274,12 @@ def get_mapped_task_dagrun(session, dag_maker):
tis = dr.get_task_instances(session=session)
for ti in tis:
if ti.task_id == upstream:
- if ti.map_index > 2:
+ if ti.region_index > 2:
ti.state = REMOVED
else:
ti.state = state
session.merge(ti)
- return ti.task
+ return expanded_task
do_task = _expand_tasks("do_something_else", "do_something")
if add_setup_tasks:
@@ -1430,7 +1550,11 @@ class TestTriggerRuleDep:
dr, task, _ = get_mapped_task_dagrun()
# ti with removed upstream ti
- ti = dr.get_task_instance(task_id="do_something_else", map_index=3, session=session)
+ ti = next(
+ ti
+ for ti in dr.get_task_instances(session=session)
+ if ti.task_id == "do_something_else" and ti.region_index == 3
+ )
ti.task = task
upstream_states = _UpstreamTIStates(
@@ -1467,7 +1591,11 @@ class TestTriggerRuleDep:
dr, task, _ = get_mapped_task_dagrun(trigger_rule=TriggerRule.ALL_FAILED, state=FAILED)
# ti with removed upstream ti
- ti = dr.get_task_instance(task_id="do_something_else", map_index=3, session=session)
+ ti = next(
+ ti
+ for ti in dr.get_task_instances(session=session)
+ if ti.task_id == "do_something_else" and ti.region_index == 3
+ )
ti.task = task
upstream_states = _UpstreamTIStates(
@@ -1503,7 +1631,11 @@ class TestTriggerRuleDep:
dr, task, _ = get_mapped_task_dagrun(trigger_rule=trigger_rule)
# ti with removed upstream ti
- ti = dr.get_task_instance(task_id="do_something_else", map_index=3, session=session)
+ ti = next(
+ ti
+ for ti in dr.get_task_instances(session=session)
+ if ti.task_id == "do_something_else" and ti.region_index == 3
+ )
ti.task = task
upstream_states = _UpstreamTIStates(
@@ -1542,7 +1674,11 @@ class TestTriggerRuleDep:
)
# ti with removed upstream ti
- ti = dr.get_task_instance(task_id="do_something_else", map_index=3, session=session)
+ ti = next(
+ ti
+ for ti in dr.get_task_instances(session=session)
+ if ti.task_id == "do_something_else" and ti.region_index == 3
+ )
ti.task = task
upstream_states = _UpstreamTIStates(
@@ -2003,7 +2139,11 @@ def test_setup_constraint_mapped_task_upstream_removed_and_success(
"""
dr, _, setup_task = get_mapped_task_dagrun(add_setup_tasks=True)
- ti = dr.get_task_instance(task_id="setup_3", map_index=map_index, session=session)
+ ti = next(
+ ti
+ for ti in dr.get_task_instances(session=session)
+ if ti.task_id == "setup_3" and ti.region_index == map_index
+ )
ti.task = setup_task
_test_trigger_rule(
@@ -2205,6 +2345,41 @@ def _expand_mapped_task(dr, dag, task_id, states, session):
class TestTriggerRuleUpstreamCountMemo:
"""The upstream-count query is memoized per scheduling pass (one DepContext) in the simple case."""
+ def test_mapped_siblings_share_resolved_producers(self, dag_maker, session):
+ with dag_maker(serialized=True):
+ source = PythonOperator.partial(task_id="source", python_callable=list).expand(op_kwargs=[{}] * 4)
+ consumer = PythonOperator.partial(task_id="consumer", python_callable=list).expand(
+ op_kwargs=[{}] * 4
+ )
+ source >> consumer
+ dr = dag_maker.create_dagrun()
+ for task_id in ("source", "consumer"):
+ region = DynamicRegion(dag_id=dr.dag_id, run_id=dr.run_id, node_id=task_id)
+ session.add(region)
+ session.flush()
+ for ti in dr.task_instances:
+ if ti.task_id == task_id:
+ ti.region_id = region.id
+ session.flush()
+ consumers = [ti for ti in dr.task_instances if ti.task_id == "consumer"]
+ expected_ids = {ti.id for ti in dr.task_instances if ti.task_id == "source"}
+ context = DepContext()
+ assert {ti.id for ti in context.upstream_tis(consumers[0], "source", session=session)} == expected_ids
+ statements = []
+
+ def record_statement(conn, cursor, statement, parameters, context, executemany):
+ statements.append(statement)
+
+ event.listen(session.bind, "before_cursor_execute", record_statement)
+ try:
+ for consumer_ti in consumers[1:]:
+ assert {
+ ti.id for ti in context.upstream_tis(consumer_ti, "source", session=session)
+ } == expected_ids
+ finally:
+ event.remove(session.bind, "before_cursor_execute", record_statement)
+ assert statements == []
+
def _make_dag(
self, dag_maker, session, *, n_downstreams, src_states, trigger_rule=TriggerRule.ALL_SUCCESS
):
@@ -2226,9 +2401,18 @@ class TestTriggerRuleUpstreamCountMemo:
session.commit()
return dr
- def test_memoized_across_downstreams_sharing_upstream(self, dag_maker, session):
+ @pytest.mark.parametrize("regional", [False, True])
+ def test_memoized_across_downstreams_sharing_upstream(self, dag_maker, session, regional):
"""N plain downstreams of the same mapped upstream issue the count query once per pass."""
dr = self._make_dag(dag_maker, session, n_downstreams=4, src_states=[SUCCESS, SUCCESS, SUCCESS])
+ if regional:
+ region = DynamicRegion(dag_id=dr.dag_id, run_id=dr.run_id, node_id="src")
+ session.add(region)
+ session.flush()
+ for ti in dr.get_task_instances(session=session):
+ if ti.task_id == "src":
+ ti.region_id = region.id
+ session.flush()
dep_context = DepContext()
with _count_upstream_count_queries() as counter:
for k in range(4):
@@ -2336,7 +2520,7 @@ class TestTriggerRuleUpstreamCountMemo:
ser = dag_maker.serialized_dag
d1 = dr.get_task_instance("d1", session=session)
d2 = dr.get_task_instance("d2", session=session)
- src0 = dr.get_task_instance("src", map_index=0, session=session)
+ src0 = next(ti for ti in dr.task_instances if ti.task_id == "src" and ti.region_index == 0)
d1.task = ser.get_task("d1")
d2.task = ser.get_task("d2")
src0.task = ser.get_task("src")
diff --git a/airflow-core/tests/unit/utils/log/test_task_log_address.py b/airflow-core/tests/unit/utils/log/test_task_log_address.py
index ea4df810eb9..db0eed149e0 100644
--- a/airflow-core/tests/unit/utils/log/test_task_log_address.py
+++ b/airflow-core/tests/unit/utils/log/test_task_log_address.py
@@ -17,18 +17,21 @@
from __future__ import annotations
from types import SimpleNamespace
+from unittest import mock
from uuid import UUID, uuid4
import pytest
from sqlalchemy import event
from airflow.configuration import conf
+from airflow.models.dagbag import DBDagBag
from airflow.models.dynamic_region import DynamicRegion
from airflow.models.taskinstance import TaskInstance
from airflow.providers.standard.operators.empty import EmptyOperator
from airflow.providers.standard.operators.python import PythonOperator
from airflow.sdk import task_group
from airflow.sdk.definitions._internal.loop import create_loop
+from airflow.serialization.serialized_objects import DagSerialization
from airflow.utils.log.file_task_handler import FileTaskHandler
from airflow.utils.log.task_log_address import (
TaskLogContext,
@@ -45,7 +48,7 @@ from tests_common.test_utils.asserts import assert_queries_count
[(-1, "", "attempt=3.log"), (0, "", "map_index=0/attempt=3.log"), (-1, "pass=2", "pass=2/attempt=3.log")],
)
def test_default_template_preserves_exact_filename(map_index, token, suffix):
- ti = SimpleNamespace(dag_id="dag", run_id="run", task_id="task", map_index=map_index, try_number=3)
+ ti = SimpleNamespace(dag_id="dag", run_id="run", task_id="task", region_index=map_index, try_number=3)
context = TaskLogContext(conf.get("logging", "log_filename_template"), "", "", "", token, map_index)
assert render_task_log_filename(ti, 3, context=context) == f"dag_id=dag/run_id=run/task_id=task/{suffix}"
@@ -242,3 +245,22 @@ def test_mapped_log_context_projects_expansion_position(dag_maker, session, insi
assert path.endswith(f"map_index=0/{token}/attempt=1.log")
else:
assert path.endswith("task_id=body.mapped/map_index=0/attempt=1.log")
+
+
+@pytest.mark.db_test
+@pytest.mark.parametrize(("attach_dag", "expected_reads"), [(True, 0), (False, 1)])
+def test_mapped_log_contexts_deserialize_the_dag_at_most_once(dag_maker, session, attach_dag, expected_reads):
+ with dag_maker(serialized=True):
+ PythonOperator.partial(task_id="mapped", python_callable=str).expand(op_args=[[1], [2]])
+ run = dag_maker.create_dagrun()
+ tis = run.task_instances
+ dag = DBDagBag().get_dag_for_run(run, session=session)
+ for ti in tis:
+ ti.task = dag.get_task(ti.task_id) if attach_dag else None
+
+ with mock.patch.object(
+ DagSerialization, "from_dict", autospec=True, side_effect=DagSerialization.from_dict
+ ) as read:
+ prepare_task_log_contexts(tis, session=session)
+
+ assert read.call_count == expected_reads
diff --git a/airflow-core/tests/unit/utils/test_db_cleanup.py b/airflow-core/tests/unit/utils/test_db_cleanup.py
index 8f833c98088..f8e6006a6aa 100644
--- a/airflow-core/tests/unit/utils/test_db_cleanup.py
+++ b/airflow-core/tests/unit/utils/test_db_cleanup.py
@@ -1733,7 +1733,7 @@ class TestDBCleanup:
dag_id="other",
task_id="task",
run_id="manual",
- map_index=-1,
+ region_index=-1,
state="success",
try_number=1,
working_set=True,
@@ -1851,7 +1851,7 @@ class TestTaskStoreCleanup:
expired = TaskStateStoreModel(
dag_run_id=dag_run.id,
task_id="t1",
- map_index=-1,
+ region_index=-1,
key="job_id",
dag_id="ts_test_dag",
run_id="ts_test_run",
@@ -1862,7 +1862,7 @@ class TestTaskStoreCleanup:
never_expire = TaskStateStoreModel(
dag_run_id=dag_run.id,
task_id="t1",
- map_index=-1,
+ region_index=-1,
key="result",
dag_id="ts_test_dag",
run_id="ts_test_run",
@@ -1873,7 +1873,7 @@ class TestTaskStoreCleanup:
not_yet_expired = TaskStateStoreModel(
dag_run_id=dag_run.id,
task_id="t1",
- map_index=-1,
+ region_index=-1,
key="future_key",
dag_id="ts_test_dag",
run_id="ts_test_run",
diff --git a/devel-common/src/tests_common/test_utils/attempt_ownership.py b/devel-common/src/tests_common/test_utils/attempt_ownership.py
index add1dd9dac8..a2ac4fbe978 100644
--- a/devel-common/src/tests_common/test_utils/attempt_ownership.py
+++ b/devel-common/src/tests_common/test_utils/attempt_ownership.py
@@ -31,6 +31,12 @@ HISTORY_ID = UUID("01960000-0000-7000-8000-000000000001")
DANGLING_VERSION = UUID("01960000-0000-7000-8000-000000000099")
NOW = datetime(2026, 9, 29, tzinfo=timezone.utc)
COORDINATES = {"dag_id": "ownership", "task_id": "task", "run_id": "manual", "map_index": -1}
+PHYSICAL_COORDINATES = {
+ "dag_id": "ownership",
+ "task_id": "task",
+ "run_id": "manual",
+ "region_index": -1,
+}
def table(connection, name, *uuid_columns):
@@ -92,7 +98,7 @@ def ownership_session(session):
attempts = TaskInstance.__table__
session.execute(
attempts.insert().values(
- **COORDINATES,
+ **PHYSICAL_COORDINATES,
id=CURRENT_ID,
try_number=2,
pool="default_pool",
@@ -104,7 +110,7 @@ def ownership_session(session):
)
session.execute(
attempts.insert().values(
- **COORDINATES,
+ **PHYSICAL_COORDINATES,
id=HISTORY_ID,
try_number=1,
pool="default_pool",
diff --git a/go-sdk/airflow/spec.gen.go b/go-sdk/airflow/spec.gen.go
index fcbd00393c2..aba65cc807c 100644
--- a/go-sdk/airflow/spec.gen.go
+++ b/go-sdk/airflow/spec.gen.go
@@ -189,6 +189,9 @@ type TaskSpec struct {
// "retry_exponential_backoff".
RetryExponentialBackoff float64
+ // ReturnsDagResult corresponds to the JSON schema field "returns_dag_result".
+ ReturnsDagResult bool
+
// StartDate corresponds to the JSON schema field "start_date".
StartDate time.Time
diff --git a/go-sdk/schema/dag-schema.json b/go-sdk/schema/dag-schema.json
index 4a29deda53c..ed2251ce3da 100644
--- a/go-sdk/schema/dag-schema.json
+++ b/go-sdk/schema/dag-schema.json
@@ -331,6 +331,7 @@
"executor": { "type": "string" },
"executor_config": { "$ref": "#/definitions/dict" },
"do_xcom_push": { "type": "boolean", "default": true },
+ "returns_dag_result": { "type": "boolean", "default": false },
"email_on_failure": { "type": "boolean", "default": true },
"email_on_retry": { "type": "boolean", "default": true },
"ui_color": { "type": "string", "default": "#fff" },
diff --git a/java-sdk/sdk/schema/dag-schema.json b/java-sdk/sdk/schema/dag-schema.json
index 4a29deda53c..ed2251ce3da 100644
--- a/java-sdk/sdk/schema/dag-schema.json
+++ b/java-sdk/sdk/schema/dag-schema.json
@@ -331,6 +331,7 @@
"executor": { "type": "string" },
"executor_config": { "$ref": "#/definitions/dict" },
"do_xcom_push": { "type": "boolean", "default": true },
+ "returns_dag_result": { "type": "boolean", "default": false },
"email_on_failure": { "type": "boolean", "default": true },
"email_on_retry": { "type": "boolean", "default": true },
"ui_color": { "type": "string", "default": "#fff" },
diff --git a/shared/state/src/airflow_shared/state/__init__.py b/shared/state/src/airflow_shared/state/__init__.py
index 524c5fd5f41..5c0ed914b59 100644
--- a/shared/state/src/airflow_shared/state/__init__.py
+++ b/shared/state/src/airflow_shared/state/__init__.py
@@ -31,26 +31,49 @@ if TYPE_CHECKING:
from sqlalchemy.orm import Session
-@dataclass(frozen=True)
+_SENTINEL_REGION_ID = UUID(int=0)
+
+
+@dataclass(frozen=True, init=False)
class TaskScope:
"""
Identifies the state namespace for a single task instance (or retry thereof).
- ``map_index`` defaults to ``-1`` for non-mapped tasks. For mapped tasks,
- set it to the actual mapped index. ``get``/``set``/``delete`` always match
- on ``(dag_id, run_id, task_id, region_id, map_index)`` exactly. The zero
- region UUID preserves the namespace of legacy callers.
+ Storage matches ``(dag_id, run_id, task_id, region_id, region_index)`` exactly.
+ The zero region UUID preserves legacy namespaces. ``map_index`` remains a
+ compatibility name for the stored index, not a public mapped position.
"""
dag_id: str
run_id: str
task_id: str
- map_index: int = -1
- region_id: UUID = UUID(int=0)
+ region_index: int = -1
+ region_id: UUID = _SENTINEL_REGION_ID
+
+ def __init__(
+ self,
+ dag_id: str,
+ run_id: str,
+ task_id: str,
+ map_index: int | None = None,
+ region_id: UUID = _SENTINEL_REGION_ID,
+ *,
+ region_index: int | None = None,
+ ) -> None:
+ if map_index is not None and region_index is not None and map_index != region_index:
+ raise ValueError("map_index and region_index must identify the same stored index")
+ if region_index is None:
+ region_index = map_index if map_index is not None else -1
+ object.__setattr__(self, "dag_id", dag_id)
+ object.__setattr__(self, "run_id", run_id)
+ object.__setattr__(self, "task_id", task_id)
+ object.__setattr__(self, "region_id", region_id)
+ object.__setattr__(self, "region_index", region_index)
@property
- def region_index(self) -> int:
- return self.map_index
+ def map_index(self) -> int:
+ """Return the stored index for legacy state backend callers."""
+ return self.region_index
@dataclass(frozen=True)
diff --git a/shared/state/tests/state/test_state.py b/shared/state/tests/state/test_state.py
index 149127725f7..578ae302212 100644
--- a/shared/state/tests/state/test_state.py
+++ b/shared/state/tests/state/test_state.py
@@ -24,6 +24,19 @@ from airflow_shared.state import AssetScope, BaseStoreBackend, StoreScope, TaskS
class TestTaskScope:
+ def test_canonical_coordinates_preserve_legacy_scope_identity(self):
+ region_id = uuid4()
+ canonical = TaskScope("dag", "run", "task", region_id=region_id, region_index=3)
+ legacy = TaskScope("dag", "run", "task", 3, region_id)
+
+ assert canonical == legacy
+ assert hash(canonical) == hash(legacy)
+ assert canonical.map_index == 3
+
+ def test_conflicting_coordinate_names_are_rejected(self):
+ with pytest.raises(ValueError, match="map_index.*region_index"):
+ TaskScope("dag", "run", "task", map_index=2, region_index=3)
+
def test_legacy_scope_uses_sentinel_region(self):
scope = TaskScope("dag", "run", "task", 3)
assert scope.region_id == UUID(int=0)
diff --git a/task-sdk/src/airflow/sdk/bases/operator.py b/task-sdk/src/airflow/sdk/bases/operator.py
index 7634443b119..1e9216ce967 100644
--- a/task-sdk/src/airflow/sdk/bases/operator.py
+++ b/task-sdk/src/airflow/sdk/bases/operator.py
@@ -1635,6 +1635,7 @@ class BaseOperator(AbstractOperator, metaclass=BaseOperatorMeta):
"has_on_retry_callback",
"has_on_skipped_callback",
"has_retry_policy",
+ "returns_dag_result",
}
)
DagContext.pop()
diff --git a/task-sdk/src/airflow/sdk/definitions/mappedoperator.py b/task-sdk/src/airflow/sdk/definitions/mappedoperator.py
index 3a68ecadfc2..43d14f75aa3 100644
--- a/task-sdk/src/airflow/sdk/definitions/mappedoperator.py
+++ b/task-sdk/src/airflow/sdk/definitions/mappedoperator.py
@@ -379,7 +379,6 @@ class MappedOperator(AbstractOperator):
"_needs_expansion",
"partial_kwargs",
"operator_extra_links",
- "returns_dag_result",
"is_stub",
}
diff --git a/ts-sdk/schema/dag-schema.json b/ts-sdk/schema/dag-schema.json
index 4a29deda53c..ed2251ce3da 100644
--- a/ts-sdk/schema/dag-schema.json
+++ b/ts-sdk/schema/dag-schema.json
@@ -331,6 +331,7 @@
"executor": { "type": "string" },
"executor_config": { "$ref": "#/definitions/dict" },
"do_xcom_push": { "type": "boolean", "default": true },
+ "returns_dag_result": { "type": "boolean", "default": false },
"email_on_failure": { "type": "boolean", "default": true },
"email_on_retry": { "type": "boolean", "default": true },
"ui_color": { "type": "string", "default": "#fff" },
diff --git a/ts-sdk/src/generated/dag-schema-fields.ts b/ts-sdk/src/generated/dag-schema-fields.ts
index de155670fda..a3c6f767820 100644
--- a/ts-sdk/src/generated/dag-schema-fields.ts
+++ b/ts-sdk/src/generated/dag-schema-fields.ts
@@ -113,6 +113,8 @@ export interface GeneratedTaskFields {
readonly executor?: string;
/** Maps to the schema key `do_xcom_push` (schema default `true`). */
readonly doXcomPush?: boolean;
+ /** Maps to the schema key `returns_dag_result` (schema default `false`). */
+ readonly returnsDagResult?: boolean;
/** Maps to the schema key `email_on_failure` (schema default `true`). */
readonly emailOnFailure?: boolean;
/** Maps to the schema key `email_on_retry` (schema default `true`). */
@@ -201,6 +203,7 @@ export const TASK_SCHEMA_FIELDS = {
weightRule: { key: "weight_rule", type: "string", default: "downstream" },
executor: { key: "executor", type: "string" },
doXcomPush: { key: "do_xcom_push", type: "boolean", default: true },
+ returnsDagResult: { key: "returns_dag_result", type: "boolean", default: false },
emailOnFailure: { key: "email_on_failure", type: "boolean", default: true },
emailOnRetry: { key: "email_on_retry", type: "boolean", default: true },
docMd: { key: "doc_md", type: "string" },