This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 6a925a8fbe6 Resolve loop iterations through family and path columns on
dynamic regions (#74892)
6a925a8fbe6 is described below
commit 6a925a8fbe690444a8e685fdd824e79a5864833d
Author: Ash Berlin-Taylor <[email protected]>
AuthorDate: Sat Oct 17 21:44:00 2026 +0100
Resolve loop iterations through family and path columns on dynamic regions
(#74892)
First off: why is this needed, wasn't loops landed and working? "Yes, but."
Part of the goal with the data-modelling changes for task loops (AIP-111)
is to support dynamic graphs (AIP-113) and other dynamic things, such as
choosing to allow loops inside loops (currently a parse error), with no
future
migrations or data modelling changes needed. This change lets us enable
these
features without expensive recursive queries in TIDep evaulation etc. And
why
now? By getting this change in before 3.4.0 enters the RC stage, while we
can
still change the migration "at zero-cost". And adding it as a migration in
3.5
would involve a (possibly costly) migration.
So the meat of the change:
Finding the task instances that sit inside one iteration of a loop currently
means finding every region of the loop's fork chain (a fork is when a loop
is
cleared part way through and some iterations run a-fresh), plus everything
nested under that iteration, however deep. The previous lookup resolved
those
region ids with two recursive CTEs and then passed them to the task_instance
query as literal lists. The nested list holds every region under the
iteration, whichever task is being read, so once something under an
iteration
can fan out again (nested loops, or dynamic graphs inside a loop) its size
grows with the product of those widths. Every lookup hits this path: the
scheduler's trigger-rule check of a loop task repeats it on each pass until
the task is ready, and every XCom read from inside a loop pays it on the API
server. Measured with 5,000 regions under one iteration, a lookup that
returns
one row took about 30 ms on PostgreSQL and MySQL, almost all of it building,
sending and parsing the list; with 50,000 it took about 550 ms.
Two columns on dynamic_region lets us get the same result but without
recursion. `family_root_id` is the id of the first region of a fork chain,
so
a chain is one equality column comparisson. `path` contains, for each
ancestor
from the top, its family root and the index the region sits at inside it, so
everything under one iteration of a loop shares a byte prefix and is one
range
of an index. The lookup becomes a UNION ALL of two joins from task_instance
to dynamic_region, each driven by its own index; a single OR across the two
legs gives the planner no such choice. Its cost no longer depends on how
many
regions sit under the iteration. The hook that fills the columns rejects a
fork whose parent differs from its source's, because both the family
equality
and the prefix range rely on every member of a chain sharing one position.
Region inserts are rare next to task instance reads and writes, so the hook
and the two extra indexes cost comparitavely little.
The columns are NOT NULL so that raw inserts which skip the ORM fail loudly
instead of writing rows the lookup would silently miss. They are added to
the
migration that creates dynamic_region, which has only shipped in a beta, so
no
backfill is needed. A 3.4.0b1 database already ran that migration without
them
so users in that situation will need to downgrade to 3.3 and upgrade again,
or reset the database. Nesting is capped at ten ancestors, which sets the
200-byte path width on MySQL. The cap is a chosen bound rather than an index
limit, and authoring already rejects nested loops, so current real depth is
far lower.
In order to allow more easily resizing this column on MySQL in the future if
we want or need to, the length needs to be over 255 bytes (as InnoDB stores
VARBINARY with shorter values with a 1 byte length prefix.) So we have
chosen
300, allowing 15 levels, such that we can change it in future with a
meta-data
only storage change. Regardless of max length of column, all DBs only use as
much storage as needed (+ size counter).
---
.../0143_3_4_0_add_dynamic_region_storage.py | 6 +
airflow-core/src/airflow/models/dynamic_region.py | 212 +++++++++++++------
.../core_api/routes/public/test_extra_links.py | 7 +-
.../test_0143_add_dynamic_region_storage.py | 10 +-
.../test_0144_rename_stored_map_index.py | 12 +-
.../tests/unit/models/test_dynamic_region.py | 233 ++++++++++++++++++++-
airflow-core/tests/unit/utils/test_db_cleanup.py | 2 +
7 files changed, 418 insertions(+), 64 deletions(-)
diff --git
a/airflow-core/src/airflow/migrations/versions/0143_3_4_0_add_dynamic_region_storage.py
b/airflow-core/src/airflow/migrations/versions/0143_3_4_0_add_dynamic_region_storage.py
index 80d6c5e1aad..8c99d502386 100644
---
a/airflow-core/src/airflow/migrations/versions/0143_3_4_0_add_dynamic_region_storage.py
+++
b/airflow-core/src/airflow/migrations/versions/0143_3_4_0_add_dynamic_region_storage.py
@@ -157,6 +157,10 @@ def upgrade():
sa.Column("parent_region_id", CompactUUID(), nullable=True),
sa.Column("parent_region_index", sa.Integer(), nullable=True),
sa.Column("forked_from_region_id", CompactUUID(), nullable=True),
+ sa.Column("family_root_id", CompactUUID(), nullable=False),
+ sa.Column(
+ "path", sa.LargeBinary().with_variant(sa.VARBINARY(300),
"mysql", "mariadb"), nullable=False
+ ),
sa.Column("resumes_from_index", sa.Integer(), nullable=False,
server_default="0"),
sa.Column("created_at", UtcDateTime(), nullable=False),
sa.PrimaryKeyConstraint("id", name="dynamic_region_pkey"),
@@ -187,6 +191,8 @@ def upgrade():
["dag_id", "run_id", "node_id", "parent_region_id",
"parent_region_index"],
)
op.create_index("idx_dynamic_region_parent_region_id",
"dynamic_region", ["parent_region_id"])
+ op.create_index("idx_dynamic_region_family_root_id", "dynamic_region",
["family_root_id"])
+ op.create_index("idx_dynamic_region_path", "dynamic_region", ["path"])
for table_name in ("task_instance", "task_state_store"):
op.add_column(
table_name,
diff --git a/airflow-core/src/airflow/models/dynamic_region.py
b/airflow-core/src/airflow/models/dynamic_region.py
index 65470ea6bc5..494cd09a926 100644
--- a/airflow-core/src/airflow/models/dynamic_region.py
+++ b/airflow-core/src/airflow/models/dynamic_region.py
@@ -28,7 +28,9 @@ import attrs
import uuid6
from sqlalchemy import (
BINARY,
+ VARBINARY,
CheckConstraint,
+ CompoundSelect,
ForeignKeyConstraint,
Index,
Integer,
@@ -44,7 +46,7 @@ from sqlalchemy import (
union_all,
)
from sqlalchemy.exc import IntegrityError
-from sqlalchemy.orm import Mapped, aliased, mapped_column
+from sqlalchemy.orm import Mapped, Session, mapped_column, object_session
from airflow._shared.timezones import timezone
from airflow.models.base import Base, StringID
@@ -54,11 +56,19 @@ SENTINEL_REGION_ID = UUID(int=0)
LOOP_XCOM_PREFIX = "_airflow_loop_"
LOOP_DECISION_REASON_KEY = f"{LOOP_XCOM_PREFIX}decision_reason"
LOOP_DECISION_REASON_MAX_BYTES = 512
+MAX_REGION_DEPTH = 10 # Max depth we allow at runtime
+_PATH_SEGMENT_BYTES = 20
+_MAX_REGION_INDEX = 2**31 - 1
+_COORDINATES_CACHE_KEY = "dynamic_region_coordinates"
+# Over 255 bytes so MySQL can widen the column later as a metadata-only
change; 15 levels fit already.
+_PATH_TYPE = LargeBinary().with_variant(VARBINARY(300), "mysql", "mariadb")
+_pack_index = struct.Struct(">I").pack
if TYPE_CHECKING:
from sqlalchemy import Select
from sqlalchemy.engine import Connection
- from sqlalchemy.orm import Mapper, Session
+ from sqlalchemy.orm import Mapper
+ from sqlalchemy.orm.session import SessionTransaction
from sqlalchemy.sql.elements import ColumnElement
from airflow.models.taskinstance import TaskInstance
@@ -142,6 +152,34 @@ class DynamicRegion(Base):
rows from before the column existed or a raw insert that bypasses the ORM,
have a NULL key and are
not guarded: :meth:`DynamicRegion.get_or_create` finds those by their slot
columns, and a raw insert
must supply the key itself.
+
+ A *family* is an original region and the chain of forks of it, and
``family_root_id`` is the id of
+ its first region. Everything that sits inside one iteration of a loop,
however deep, has the loop
+ family and that iteration in its ``path``: for each ancestor from the top,
16 bytes of its
+ ``family_root_id`` and the 4-byte big-endian index this region sits at
within it. That makes "every
+ region of a loop family, and everything nested under it" an equality on
``family_root_id`` plus a
+ range scan of ``path``, see :func:`build_loop_family_legs`, instead of a
walk of the ancestry. Both
+ columns are set by the ``before_insert`` hook from
``forked_from_region_id`` and the parent
+ columns, so a raw insert that bypasses the ORM must supply them. A fork
must sit under the same
+ parent as its source, which the hook enforces, and nesting is limited to
:data:`MAX_REGION_DEPTH`
+ ancestors.
+
+ The same encoding in SQL, for a region ``r`` with parent ``p``, a loop
family member ``m`` and an
+ iteration ``k``. ``family_root_id`` is ``BINARY(16)`` on MySQL, so it
concatenates as is, and 32 hex
+ characters on SQLite, where ``||`` always yields text, so the path is
built in hex and decoded once::
+
+ -- r.path on PostgreSQL / MySQL / SQLite 3.41+
+ p.path || uuid_send(p.family_root_id) ||
int4send(r.parent_region_index)
+ CONCAT(p.path, p.family_root_id,
UNHEX(LPAD(HEX(r.parent_region_index), 8, '0')))
+ unhex(hex(p.path) || p.family_root_id || printf('%08X',
r.parent_region_index))
+
+ -- everything nested under iteration k of m's family, on PostgreSQL
+ path >= m.path || uuid_send(m.family_root_id) || int4send(k)
+ AND path < m.path || uuid_send(m.family_root_id) || int4send(k + 1)
+
+ -- the iteration of m's family that a nested region sits in, on
PostgreSQL / MySQL
+ ('x' || encode(substring(path from length(m.path) + 17 for 4),
'hex'))::bit(32)::integer
+ CONV(HEX(SUBSTRING(path, LENGTH(m.path) + 17, 4)), 16, 10)
"""
__tablename__ = "dynamic_region"
@@ -154,6 +192,8 @@ class DynamicRegion(Base):
parent_region_id: Mapped[UUID | None] = mapped_column(CompactUUID(),
nullable=True)
parent_region_index: Mapped[int | None] = mapped_column(Integer,
nullable=True)
forked_from_region_id: Mapped[UUID | None] = mapped_column(CompactUUID(),
nullable=True)
+ family_root_id: Mapped[UUID] = mapped_column(CompactUUID(), nullable=False)
+ path: Mapped[bytes] = mapped_column(_PATH_TYPE, nullable=False)
resumes_from_index: Mapped[int] = mapped_column(Integer, nullable=False,
default=0, server_default="0")
# Hash of the slot, set on original regions and NULL on forks; see the
class docstring.
slot_key: Mapped[bytes | None] = mapped_column(
@@ -184,6 +224,8 @@ class DynamicRegion(Base):
CheckConstraint("resumes_from_index >= 0",
name="resumes_from_index_nonnegative"),
Index("idx_dynamic_region_slot", dag_id, run_id, node_id,
parent_region_id, parent_region_index),
Index("idx_dynamic_region_parent_region_id", parent_region_id),
+ Index("idx_dynamic_region_family_root_id", family_root_id),
+ Index("idx_dynamic_region_path", path),
)
@classmethod
@@ -404,6 +446,72 @@ def _set_slot_key(mapper: Mapper, connection: Connection,
region: DynamicRegion)
)
+def _get_coordinates_cache(region: DynamicRegion) -> dict[UUID, tuple[UUID,
bytes]]:
+ """
+ Remember region coordinates up until the unit-of-work is flushed to the DB.
+
+ That is enough for siblings flushed together to read their parent once,
and for a region to find a
+ parent added in the same flush: every ``before_insert`` hook runs before
the first ``INSERT``, so the
+ parent is not in the database yet.
+ """
+ session = object_session(region)
+ if TYPE_CHECKING:
+ assert session is not None
+ return session.info.setdefault(_COORDINATES_CACHE_KEY, {})
+
+
+def _get_region_coordinates(
+ connection: Connection, region: DynamicRegion, region_id: UUID, role: str
+) -> tuple[UUID, bytes]:
+ cache = _get_coordinates_cache(region)
+ if region_id not in cache:
+ row = connection.execute(
+ select(DynamicRegion.family_root_id,
DynamicRegion.path).where(DynamicRegion.id == region_id)
+ ).one_or_none()
+ if row is None:
+ raise ValueError(f"The {role} region {region_id} does not exist")
+ cache[region_id] = (row.family_root_id, row.path)
+ return cache[region_id]
+
+
+def _build_region_path(connection: Connection, region: DynamicRegion) -> bytes:
+ if region.parent_region_id is None and region.parent_region_index is None:
+ return b""
+ if region.parent_region_id is None or region.parent_region_index is None:
+ raise ValueError("parent_region_id and parent_region_index must be set
together")
+ if region.parent_region_index < 0:
+ raise ValueError("parent_region_index must not be negative")
+ parent_root, parent_path = _get_region_coordinates(connection, region,
region.parent_region_id, "parent")
+ path = parent_path + parent_root.bytes +
_pack_index(region.parent_region_index)
+ if len(path) > MAX_REGION_DEPTH * _PATH_SEGMENT_BYTES:
+ raise ValueError(f"Dynamic regions can nest at most {MAX_REGION_DEPTH}
levels deep")
+ return path
+
+
[email protected]_for(DynamicRegion, "before_insert")
+def _set_region_coordinates(mapper: Mapper, connection: Connection, region:
DynamicRegion) -> None:
+ # The id default runs after this hook, and a region that starts a family
is named by its own id.
+ if region.id is None:
+ region.id = uuid6.uuid7()
+ path = _build_region_path(connection, region)
+ if region.forked_from_region_id is None:
+ region.family_root_id = region.id
+ else:
+ region.family_root_id, source_path = _get_region_coordinates(
+ connection, region, region.forked_from_region_id, "fork source"
+ )
+ if source_path != path:
+ raise ValueError("A fork must sit under the same parent as the
region it forks")
+ region.path = path
+ _get_coordinates_cache(region)[region.id] = (region.family_root_id,
region.path)
+
+
+# Every flush runs in a subtransaction of its own, so this fires after each
flush, not only at commit.
[email protected]_for(Session, "after_transaction_end")
+def _forget_region_coordinates(session: Session, transaction:
SessionTransaction) -> None:
+ session.info.pop(_COORDINATES_CACHE_KEY, None)
+
+
def load_region_ancestry(
region_ids: Collection[UUID], *, dag_id: str, run_id: str, session: Session
) -> dict[UUID, DynamicRegion]:
@@ -457,59 +565,39 @@ def loop_position(
return None
[email protected](frozen=True)
-class _LoopPassRegions:
- """The regions that hold one iteration of a loop: its fork family and
everything nested under it."""
-
- family_ids: tuple[UUID, ...]
- nested_ids: tuple[UUID, ...]
- iteration: int
-
-
-def _load_loop_pass_regions(position: tuple[UUID, int], *, session: Session)
-> _LoopPassRegions:
+def build_loop_family_legs(
+ query: Select[Any], family_member: DynamicRegion, *, iteration: int | None
= None
+) -> tuple[Select[Any], Select[Any]]:
"""
- Resolve, in one statement, the regions that hold the loop iteration
``position`` names.
-
- Both recursive CTEs run against ``dynamic_region`` only. The
``task_instance`` query then takes the
- ids as literals, which every backend turns into an index lookup on
``task_instance_current_key``;
- a subquery in that position leaves MySQL and Postgres scanning every row
of the task.
+ Split ``query``, a selection from ``task_instance``, into the rows that
belong to one loop family.
+
+ ``family_member`` is any region of the loop's family; the family's regions
share ``family_root_id`` and
+ ``path``. The first leg holds the rows that sit directly in a region of
the family, where
+ ``region_index`` is the iteration. The second holds the rows nested under
the family, however deep,
+ whose ``path`` starts with the family and the iteration. Pass
``iteration`` to keep one iteration, or
+ leave it out for all of them.
+
+ Each leg is joined to :class:`DynamicRegion`, so a caller can add columns
to both before combining
+ them with ``UNION ALL``. The iteration of a row in a region directly under
the family is
+ ``DynamicRegion.parent_region_index``; deeper down, it is the 4-byte index
that follows the family's
+ 16-byte root in ``path``, as decoded in the :class:`DynamicRegion`
docstring. The
+ legs never overlap. They are separate statements because the planner can
then drive each from its
+ own index, which an ``OR`` across them does not allow.
"""
- family_root, iteration = position
- family = (
- select(DynamicRegion.id).where(DynamicRegion.id ==
family_root).cte("loop_family", recursive=True)
- )
- forked = aliased(DynamicRegion)
- family = family.union(select(forked.id).where(forked.forked_from_region_id
== family.c.id))
- nested = (
- select(DynamicRegion.id)
- .where(
- DynamicRegion.parent_region_id.in_(select(family.c.id)),
- DynamicRegion.parent_region_index == iteration,
- )
- .cte("loop_iteration_regions", recursive=True)
- )
- child = aliased(DynamicRegion)
- nested = nested.union(select(child.id).where(child.parent_region_id ==
nested.c.id))
- rows = session.execute(
- union_all(
- select(family.c.id, literal(0).label("is_nested")),
- select(nested.c.id, literal(1).label("is_nested")),
- )
- ).all()
- return _LoopPassRegions(
- family_ids=tuple(region_id for region_id, is_nested in rows if not
is_nested),
- nested_ids=tuple(region_id for region_id, is_nested in rows if
is_nested),
- iteration=iteration,
- )
-
-
-def _build_loop_pass_filter(regions: _LoopPassRegions) -> ColumnElement[bool]:
from airflow.models.taskinstance import TaskInstance
- return or_(
- and_(TaskInstance.region_id.in_(regions.family_ids),
TaskInstance.region_index == regions.iteration),
- TaskInstance.region_id.in_(regions.nested_ids),
+ first, end = (0, _MAX_REGION_INDEX + 1) if iteration is None else
(iteration, iteration + 1)
+ prefix = family_member.path + family_member.family_root_id.bytes
+ direct = query.join(DynamicRegion, TaskInstance.region_id ==
DynamicRegion.id).where(
+ DynamicRegion.family_root_id == family_member.family_root_id
)
+ if iteration is not None:
+ direct = direct.where(TaskInstance.region_index == iteration)
+ nested = query.join(DynamicRegion, TaskInstance.region_id ==
DynamicRegion.id).where(
+ DynamicRegion.path >= prefix + _pack_index(first),
+ DynamicRegion.path < prefix + _pack_index(end),
+ )
+ return direct, nested
def public_region_filter(
@@ -570,9 +658,8 @@ def _filter_producers(
region_id: UUID | None,
region_index: int | None,
top_level_only: bool,
- loop_pass: _LoopPassRegions | None = None,
) -> Select:
- """Push every coordinate predicate into SQL."""
+ """Push every coordinate predicate that does not need loop ancestry into
SQL."""
from airflow.models.taskinstance import TaskInstance
query = query.where(
@@ -585,8 +672,6 @@ def _filter_producers(
query = query.where(TaskInstance.region_id == region_id)
if region_index is not None:
query = query.where(TaskInstance.region_index == region_index)
- if loop_pass is not None:
- query = query.where(_build_loop_pass_filter(loop_pass))
if top_level_only:
if is_mapped:
top_level_region = (
@@ -633,8 +718,13 @@ def _build_candidate_query(
region_id: UUID | None,
region_index: int | None,
session: Session,
-) -> Select[Any] | None:
- """Narrow ``query`` to the candidate producers of the caller's scope, or
``None`` if there are none."""
+) -> Select[Any] | CompoundSelect | None:
+ """
+ Narrow ``query`` to the candidate producers of the caller's scope, or
``None`` if there are none.
+
+ Inside a loop the result is the ``UNION ALL`` of the two legs of
:func:`build_loop_family_legs`, so
+ load ORM entities from it with ``from_statement``.
+ """
regions: dict[UUID, DynamicRegion] = {}
position = None
if context and region_id is None:
@@ -650,7 +740,7 @@ def _build_candidate_query(
if position[1] < 0:
return None
- return _filter_producers(
+ query = _filter_producers(
query,
dag_id=dag_id,
run_id=run_id,
@@ -660,8 +750,11 @@ def _build_candidate_query(
region_id=region_id,
region_index=region_index,
top_level_only=region_id is None and position is None,
- loop_pass=_load_loop_pass_regions(position, session=session) if
position is not None else None,
)
+ if position is None:
+ return query
+ family_root, iteration = position
+ return union_all(*build_loop_family_legs(query, regions[family_root],
iteration=iteration))
def _check_producers_unambiguous(
@@ -707,7 +800,8 @@ def resolve_current_producers(
)
if query is None:
return ()
- candidates = session.scalars(query).all()
+ statement = select(TaskInstance).from_statement(query) if
isinstance(query, CompoundSelect) else query
+ candidates = session.scalars(statement).all()
_check_producers_unambiguous(
(ti.region_index for ti in candidates),
dag_id=dag_id,
diff --git
a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_extra_links.py
b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_extra_links.py
index 1bc12c93260..0bb45fb62a8 100644
---
a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_extra_links.py
+++
b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_extra_links.py
@@ -140,13 +140,16 @@ def
test_serialized_links_read_xcom_from_selected_region(test_client, dag_maker,
CustomOperator(task_id="work", bash_command="echo test")
with dag_maker(serialized=True) as dag:
- loop = create_loop(body, max_iterations=4)
+ create_loop(body, max_iterations=4)
run = dag_maker.create_dagrun()
sibling = next(ti for ti in run.task_instances if ti.task_id ==
"body.work" and ti.region_index == 0)
+ source_region = session.get(DynamicRegion, sibling.region_id)
selected_region = DynamicRegion(
dag_id=run.dag_id,
run_id=run.run_id,
- node_id=loop.group_id,
+ node_id=source_region.node_id,
+ parent_region_id=source_region.parent_region_id,
+ parent_region_index=source_region.parent_region_index,
forked_from_region_id=sibling.region_id,
resumes_from_index=sibling.region_index,
)
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 16fc5206149..27e489cc007 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
@@ -90,6 +90,8 @@ dynamic_region = sa.Table(
sa.Column("dag_id", sa.String(250)),
sa.Column("run_id", sa.String(250)),
sa.Column("node_id", sa.String(250)),
+ sa.Column("family_root_id", CompactUUID),
+ sa.Column("path", sa.LargeBinary),
sa.Column("created_at", UtcDateTime),
)
TABLES = (task_instance, task_state_store)
@@ -166,7 +168,13 @@ def _add_region_row(connection, table, **overrides):
region_id = uuid4()
connection.execute(
dynamic_region.insert().values(
- id=region_id, dag_id=DAG_ID, run_id=RUN_ID, node_id="task",
created_at=timezone.utcnow()
+ id=region_id,
+ dag_id=DAG_ID,
+ run_id=RUN_ID,
+ node_id="task",
+ family_root_id=region_id,
+ path=b"",
+ created_at=timezone.utcnow(),
)
)
values = {name: value for name, value in row.items() if name != "id" or
table is task_instance}
diff --git
a/airflow-core/tests/unit/migrations/test_0144_rename_stored_map_index.py
b/airflow-core/tests/unit/migrations/test_0144_rename_stored_map_index.py
index 16ded6715be..e28fbbc5af6 100644
--- a/airflow-core/tests/unit/migrations/test_0144_rename_stored_map_index.py
+++ b/airflow-core/tests/unit/migrations/test_0144_rename_stored_map_index.py
@@ -93,6 +93,8 @@ dynamic_region = sa.Table(
sa.Column("node_id", sa.String(250)),
sa.Column("parent_region_id", CompactUUID),
sa.Column("parent_region_index", sa.Integer),
+ sa.Column("family_root_id", CompactUUID),
+ sa.Column("path", sa.LargeBinary),
sa.Column("created_at", UtcDateTime),
)
task_reschedule = sa.Table(
@@ -183,7 +185,13 @@ def stored_rows(database_at_revision):
).inserted_primary_key[0]
connection.execute(
dynamic_region.insert().values(
- id=REGION_ID, dag_id=DAG_ID, run_id=RUN_ID, node_id="task",
created_at=now
+ id=REGION_ID,
+ dag_id=DAG_ID,
+ run_id=RUN_ID,
+ node_id="task",
+ family_root_id=REGION_ID,
+ path=b"",
+ created_at=now,
)
)
connection.execute(
@@ -194,6 +202,8 @@ def stored_rows(database_at_revision):
node_id="task",
parent_region_id=REGION_ID,
parent_region_index=0,
+ family_root_id=CHILD_REGION_ID,
+ path=REGION_ID.bytes + b"\x00\x00\x00\x00",
created_at=now,
)
)
diff --git a/airflow-core/tests/unit/models/test_dynamic_region.py
b/airflow-core/tests/unit/models/test_dynamic_region.py
index 82f107abf13..54c17bf9b2e 100644
--- a/airflow-core/tests/unit/models/test_dynamic_region.py
+++ b/airflow-core/tests/unit/models/test_dynamic_region.py
@@ -17,20 +17,24 @@
from __future__ import annotations
+import struct
from typing import TYPE_CHECKING
from uuid import uuid4
import pytest
-from sqlalchemy import event, func, insert, select, text
+from sqlalchemy import event, func, insert, select, text, union_all
from sqlalchemy.exc import IntegrityError
from airflow._shared.timezones import timezone
from airflow.models.dagrun import DagRun
from airflow.models.dynamic_region import (
+ _COORDINATES_CACHE_KEY,
+ MAX_REGION_DEPTH,
SENTINEL_REGION_ID,
AmbiguousProducerError,
DynamicRegion,
ProducerContext,
+ build_loop_family_legs,
build_loop_sequence_order,
resolve_current_producers,
select_loop_producer_ids,
@@ -626,6 +630,231 @@ def
test_fork_is_exempt_but_a_source_can_only_be_forked_once(regional_tis, sessi
add_region(session, run, forked_from_region_id=fork.id,
resumes_from_index=2)
+def pack_index(value: int) -> bytes:
+ return struct.pack(">i", value)
+
+
+def test_top_level_region_is_its_own_family_with_an_empty_path(regional_tis,
session):
+ run, _ = regional_tis
+ region = add_region(session, run)
+
+ assert region.family_root_id == region.id
+ assert region.path == b""
+
+
+def
test_nested_regions_extend_the_path_by_the_parents_family_and_the_index(regional_tis,
session):
+ run, _ = regional_tis
+ outer = add_region(session, run)
+ inner = add_region(session, run, parent_region_id=outer.id,
parent_region_index=300)
+ innermost = add_region(session, run, parent_region_id=inner.id,
parent_region_index=1)
+
+ assert inner.path == outer.id.bytes + pack_index(300)
+ assert innermost.path == inner.path + inner.id.bytes + pack_index(1)
+
+
+def
test_forks_share_the_family_and_regions_nested_under_a_fork_use_its_family_root(regional_tis,
session):
+ run, _ = regional_tis
+ original = add_region(session, run)
+ fork = add_region(session, run, forked_from_region_id=original.id)
+ forked_again = add_region(session, run, forked_from_region_id=fork.id)
+ child = add_region(session, run, parent_region_id=forked_again.id,
parent_region_index=2)
+ grandchild = add_region(session, run, parent_region_id=child.id,
parent_region_index=0)
+
+ assert fork.family_root_id == forked_again.family_root_id == original.id
+ assert child.family_root_id == child.id
+ assert child.path == original.id.bytes + pack_index(2)
+ assert grandchild.path == child.path + child.id.bytes + pack_index(0)
+
+
[email protected]("moved", ["parent", "index"])
+def test_fork_under_another_parent_than_its_source_is_rejected(regional_tis,
session, moved):
+ run, _ = regional_tis
+ outer = add_region(session, run)
+ other = add_region(session, run, node_id="other")
+ original = add_region(session, run, node_id="body",
parent_region_id=outer.id, parent_region_index=0)
+ fork_parent = {
+ "parent": {"parent_region_id": other.id, "parent_region_index": 0},
+ "index": {"parent_region_id": outer.id, "parent_region_index": 1},
+ }[moved]
+
+ with pytest.raises(ValueError, match="same parent"),
session.begin_nested():
+ add_region(session, run, node_id="body",
forked_from_region_id=original.id, **fork_parent)
+
+
+def test_regions_flushed_together_resolve_a_pending_parent(regional_tis,
session):
+ run, _ = regional_tis
+ parent = make_region(run, id=uuid4())
+ child = make_region(run, node_id="child", parent_region_id=parent.id,
parent_region_index=4)
+ session.add_all([parent, child])
+ session.flush()
+
+ assert child.path == parent.id.bytes + pack_index(4)
+
+
+def test_siblings_under_one_parent_read_the_parent_once(regional_tis, session):
+ run, _ = regional_tis
+ parent_id = uuid4()
+ session.execute(
+ insert(DynamicRegion).values(
+ id=parent_id,
+ dag_id=run.dag_id,
+ run_id=run.run_id,
+ node_id="parent",
+ family_root_id=parent_id,
+ path=b"",
+ resumes_from_index=0,
+ created_at=timezone.utcnow(),
+ )
+ )
+ statements = []
+
+ def record(conn, cursor, statement, parameters, context, executemany):
+ if statement.lstrip().upper().startswith("SELECT") and
"dynamic_region" in statement:
+ statements.append(statement)
+
+ event.listen(session.get_bind(), "before_cursor_execute", record)
+ try:
+ session.add_all(
+ make_region(run, node_id=f"child{index}",
parent_region_id=parent_id, parent_region_index=index)
+ for index in range(5)
+ )
+ session.flush()
+ finally:
+ event.remove(session.get_bind(), "before_cursor_execute", record)
+
+ assert len(statements) == 1
+
+
[email protected](
+ ("kwargs", "role"),
+ [
+ pytest.param({"parent_region_id": uuid4(), "parent_region_index": 0},
"parent", id="parent"),
+ pytest.param({"forked_from_region_id": uuid4()}, "fork source",
id="fork source"),
+ ],
+)
+def test_unknown_parent_or_fork_source_is_rejected(regional_tis, session,
kwargs, role):
+ run, _ = regional_tis
+
+ with pytest.raises(ValueError, match=f"The {role} region .* does not
exist"), session.begin_nested():
+ add_region(session, run, **kwargs)
+
+
[email protected](
+ ("with_parent", "index", "message"),
+ [
+ pytest.param(True, None, "set together", id="parent without index"),
+ pytest.param(False, 3, "set together", id="index without parent"),
+ pytest.param(True, -1, "must not be negative", id="negative index"),
+ ],
+)
+def test_invalid_parent_coordinates_are_rejected(regional_tis, session,
with_parent, index, message):
+ run, _ = regional_tis
+ outer = add_region(session, run)
+ parent = {"parent_region_id": outer.id} if with_parent else {}
+
+ with pytest.raises(ValueError, match=message), session.begin_nested():
+ add_region(session, run, node_id="body", parent_region_index=index,
**parent)
+
+
+def test_regions_nest_at_most_the_maximum_depth(regional_tis, session):
+ run, _ = regional_tis
+ region = add_region(session, run)
+ for _ in range(MAX_REGION_DEPTH):
+ region = add_region(session, run, parent_region_id=region.id,
parent_region_index=0)
+
+ with pytest.raises(ValueError, match="nest at most"),
session.begin_nested():
+ add_region(session, run, parent_region_id=region.id,
parent_region_index=0)
+
+
+def test_coordinates_of_a_rolled_back_region_are_forgotten(regional_tis,
session):
+ run, _ = regional_tis
+ with session.begin_nested() as savepoint:
+ region_id = add_region(session, run).id
+ savepoint.rollback()
+
+ with pytest.raises(ValueError, match="fork source"),
session.begin_nested():
+ add_region(session, run, forked_from_region_id=region_id)
+
+
+def test_coordinates_are_forgotten_after_the_flush(regional_tis, session):
+ run, _ = regional_tis
+ add_region(session, run)
+
+ assert _COORDINATES_CACHE_KEY not in session.info
+
+
[email protected]
+def family_rows(producer_tis, session):
+ """
+ Task instances of ``mapped`` in the loop family of ``producer_tis``, in a
loop of another node and outside.
+
+ The family is three regions. Each has a row at the iteration it is named
for, and the last has a second
+ one. Nested regions hold rows at iterations 0 to 3, one of them two levels
deep.
+ """
+ tis, regions, _ = producer_tis
+ mapped = tis["mapped"]
+ other_loop = add_region(session, mapped.dag_run, node_id="other_loop")
+
+ def add_row(region: DynamicRegion, region_index: int = 0) -> TaskInstance:
+ ti = TaskInstance(
+ mapped.task,
+ mapped.dag_version_id,
+ run_id=mapped.run_id,
+ map_index=region_index,
+ region_id=region.id,
+ )
+ session.add(ti)
+ return ti
+
+ def add_nested(parent: DynamicRegion, iteration: int, node_id: str =
"mapped") -> DynamicRegion:
+ return add_region(
+ session,
+ mapped.dag_run,
+ node_id=node_id,
+ parent_region_id=parent.id,
+ parent_region_index=iteration,
+ )
+
+ nested_in_first_pass = add_nested(regions[1], 0)
+ rows = {
+ "direct 0": add_row(regions[0], 0),
+ "direct 1": add_row(regions[1], 1),
+ "direct 2": add_row(regions[2], 2),
+ "direct 3": add_row(regions[2], 3),
+ "nested 0": add_row(nested_in_first_pass),
+ "nested 2": add_row(add_nested(regions[0], 2)),
+ "nested 3": add_row(add_nested(regions[2], 3)),
+ "deep 0": add_row(add_nested(nested_in_first_pass, 0, "deep")),
+ "other loop": add_row(add_nested(other_loop, 2)),
+ "outside": add_row(add_region(session, mapped.dag_run,
node_id="outside")),
+ }
+ session.flush()
+ return regions, {name: ti.id for name, ti in rows.items()}
+
+
[email protected](
+ ("iteration", "expected"),
+ [
+ pytest.param(
+ None,
+ ["direct 0", "direct 1", "direct 2", "direct 3", "nested 0",
"nested 2", "nested 3", "deep 0"],
+ ),
+ pytest.param(0, ["direct 0", "nested 0", "deep 0"]),
+ pytest.param(2, ["direct 2", "nested 2"]),
+ pytest.param(3, ["direct 3", "nested 3"]),
+ ],
+)
+def test_loop_family_legs_select_every_member_and_everything_nested_once(
+ family_rows, session, iteration, expected
+):
+ regions, rows = family_rows
+
+ query = select(TaskInstance.id).where(TaskInstance.task_id == "mapped")
+ legs = build_loop_family_legs(query, regions[1], iteration=iteration)
+
+ assert sorted(session.scalars(union_all(*legs))) == sorted(rows[name] for
name in expected)
+
+
def test_get_or_create_returns_the_region_that_filled_the_slot(regional_tis,
session):
run, _ = regional_tis
slot = {"dag_id": run.dag_id, "run_id": run.run_id, "node_id": "mapped",
"session": session}
@@ -653,6 +882,8 @@ def
test_get_or_create_finds_a_region_created_before_slot_keys_existed(regional_
dag_id=run.dag_id,
run_id=run.run_id,
node_id="legacy",
+ family_root_id=region_id,
+ path=b"",
resumes_from_index=0,
created_at=timezone.utcnow(),
)
diff --git a/airflow-core/tests/unit/utils/test_db_cleanup.py
b/airflow-core/tests/unit/utils/test_db_cleanup.py
index 57eeb9c76fd..38daeda56bf 100644
--- a/airflow-core/tests/unit/utils/test_db_cleanup.py
+++ b/airflow-core/tests/unit/utils/test_db_cleanup.py
@@ -1683,6 +1683,8 @@ class TestDBCleanup:
dag_id=COORDINATES["dag_id"],
run_id=COORDINATES["run_id"],
node_id="loop",
+ family_root_id=region_id,
+ path=b"",
created_at=NOW - timedelta(days=100),
)
)
