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 80d6c5e1aadf2..8c99d50238671 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 65470ea6bc5e7..494cd09a9269c 100644 --- a/airflow-core/src/airflow/models/dynamic_region.py +++ b/airflow-core/src/airflow/models/dynamic_region.py @@ -28,7 +28,9 @@ import uuid6 from sqlalchemy import ( BINARY, + VARBINARY, CheckConstraint, + CompoundSelect, ForeignKeyConstraint, Index, Integer, @@ -44,7 +46,7 @@ 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 @@ 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 + + +@event.listens_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. +@event.listens_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 -@attrs.define(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 1bc12c932605e..0bb45fb62a83f 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 body(): 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 16fc5206149bc..27e489cc007cc 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 @@ 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 16ded6715be63..e28fbbc5af620 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 @@ 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 82f107abf1310..54c17bf9b2eda 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) + + +@pytest.mark.parametrize("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 + + +@pytest.mark.parametrize( + ("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) + + +@pytest.mark.parametrize( + ("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 + + +@pytest.fixture +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()} + + +@pytest.mark.parametrize( + ("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 57eeb9c76fd6c..38daeda56bfe0 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 @@ def test_dag_run_cleanup_archives_its_regions(self, ownership_session): 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), ) )