This is an automated email from the ASF dual-hosted git repository.
ashb pushed a commit to branch v3-4-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-4-test by this push:
new 2af7b35dd33 [v3-4-test] Allow downgrading past the TI ownership
migration when data is present (#74407)
2af7b35dd33 is described below
commit 2af7b35dd33d4b6d79afd5262e807358cad30421
Author: Ash Berlin-Taylor <[email protected]>
AuthorDate: Thu Oct 15 11:43:14 2026 +0100
[v3-4-test] Allow downgrading past the TI ownership migration when data is
present (#74407)
The 0142 downgrade refused whenever xcom_v2 or rtif_v2 held rows or any
task had retried since the upgrade, so a rollback after any real use of
3.4 was impossible. It now folds that data back into the legacy tables as
if 3.3 had written it. That is only lossless for the latest attempt: the
legacy tables are keyed by task coordinates and can hold values for one
attempt, so earlier attempts' XCom and rendered fields are dropped while
the attempts themselves move back to task_instance_history. The release
note records that cost and the runtime, which grows with the size of xcom
and rendered_task_instance_fields.
Data 3.3 cannot represent (several attempts sharing coordinates, archived
attempts with no current attempt, owners that moved while legacy rows
exist) is refused by read-only checks before the first write, because
MySQL cannot roll back DDL. Past that point a failure on MySQL leaves the
schema partly downgraded and needs a restore from backup.
The downgrade leaves try_number alone. Migration 0138's downgrade already
subtracts one from pending retries and runs after this one on any
downgrade to 3.3, so rewinding here as well made a retried task reuse an
archived try number. Only a test that downgrades through 0138 can see it.
Expansion now deletes the archived attempts of an unmapped task instance.
Once it is promoted or removed they have no current attempt, and a
cleared mapped task would make every later downgrade refuse. The delete is
scoped to the placeholder's own region so expanding one loop pass cannot
remove a sibling pass's archived attempts. It discards real earlier runs
of a task that later became mapped.
(cherry picked from commit a87c6ddb52a07795f1fd9799fc4396ee1418e4a9)
---
RELEASE_NOTES.rst | 22 +
.../0142_3_4_0_unify_task_attempt_ownership.py | 263 +++++++++-
airflow-core/src/airflow/models/taskinstance.py | 16 +
.../test_0142_unify_task_attempt_ownership.py | 530 +++++++++++++++++++--
.../tests/unit/models/test_mappedoperator.py | 126 +++++
.../tests/unit/models/test_taskinstance.py | 2 +-
reproducible_build.yaml | 4 +-
7 files changed, 905 insertions(+), 58 deletions(-)
diff --git a/RELEASE_NOTES.rst b/RELEASE_NOTES.rst
index 88e5ddff469..262bd752479 100644
--- a/RELEASE_NOTES.rst
+++ b/RELEASE_NOTES.rst
@@ -42,6 +42,28 @@ under a per-attempt UUID so each iteration's data is
preserved and can be cleare
For detailed usage instructions, see :doc:`/authoring-and-scheduling/loops`.
+.. note::
+ Downgrading to any version earlier than 3.4.0 can be slow, depending on the
number of
+ XCom rows that have been written after upgrading. Upgrading was tested on a
+ 300m row XCom table, so please don't worry about upgrade time (we made sure
+ that the XCom *data* didn't get touched or read during migration to keep it
+ fast).
+
+ This is because this release contains a migration that starts storing new
XCom
+ rows in an xcom_v2 table, and in order to downgrade we have to move those
over
+ to the old table, and there is no way of doing this without copying the
+ written xcom values.
+
+ The downgrade refuses to continue when it finds data it cannot map back
+ safely. Each migration commits separately, so the later 3.4.0 migrations
+ have already been reversed at that point: fix the data and re-run the
+ downgrade, or run ``airflow db migrate`` to return to 3.4.
+
+ On MySQL a failure after those checks leaves the schema partly downgraded,
+ as MySQL cannot roll back DDL. Taking a backup before attempting the
+ downgrade on MySQL is advised, so it can be restored in that case.
+
+
Dag Importer for non-Python Dags (#72369, #73728, #73853, #74020, #74084,
#74333, #74586)
"""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
diff --git
a/airflow-core/src/airflow/migrations/versions/0142_3_4_0_unify_task_attempt_ownership.py
b/airflow-core/src/airflow/migrations/versions/0142_3_4_0_unify_task_attempt_ownership.py
index c4a6e9a0401..41a1e0e055a 100644
---
a/airflow-core/src/airflow/migrations/versions/0142_3_4_0_unify_task_attempt_ownership.py
+++
b/airflow-core/src/airflow/migrations/versions/0142_3_4_0_unify_task_attempt_ownership.py
@@ -350,24 +350,246 @@ def upgrade():
)
+def _build_task_instance_table():
+ return sa.table(
+ "task_instance",
+ sa.column("id"),
+ *(sa.column(c) for c in _COPY_COLUMNS),
+ sa.column("working_set", sa.Boolean()),
+ sa.column("dag_version_id"),
+ )
+
+
+def _build_legacy_owner_table():
+ return sa.table(
+ "legacy_task_data_owner", *(sa.column(c) for c in _COORDINATES),
sa.column("task_instance_id")
+ )
+
+
+def _select_latest_attempts(ti):
+ newer = ti.alias("newer")
+ # An anti-join, because grouping by max(try_number) aggregates the whole
table however few rows need copying.
+ return (
+ sa.select(ti.c.id, *(ti.c[c] for c in _COORDINATES))
+ .where(
+ ~sa.exists().where(
+ *(newer.c[c] == ti.c[c] for c in _COORDINATES),
newer.c.try_number > ti.c.try_number
+ )
+ )
+ .subquery()
+ )
+
+
+def _has_duplicate_attempts(bind, ti) -> bool:
+ coordinates = [ti.c[c] for c in _COORDINATES]
+ for where, group_by in (
+ (ti.c.working_set.is_not(None), coordinates),
+ (sa.true(), [*coordinates, ti.c.try_number]),
+ ):
+ duplicated =
sa.select(1).select_from(ti).where(where).group_by(*group_by).having(sa.func.count()
> 1)
+ if bind.scalar(duplicated.limit(1)):
+ return True
+ return False
+
+
+def _select_archived_ids(ti):
+ return sa.select(ti.c.id).where(ti.c.working_set.is_(None))
+
+
+def _count_orphaned_archived_attempts(bind, ti) -> int:
+ current = ti.alias("current")
+ orphaned = (
+ sa.select(sa.func.count())
+ .select_from(ti)
+ .where(
+ ti.c.working_set.is_(None),
+ ~sa.exists().where(
+ *(current.c[c] == ti.c[c] for c in _COORDINATES),
current.c.working_set.is_not(None)
+ ),
+ )
+ )
+ return bind.scalar(orphaned)
+
+
+def _has_moved_owners_with_legacy_rows(bind, ti) -> bool:
+ owner = _build_legacy_owner_table()
+ moved = sa.or_(*(owner.c[c] != ti.c[c] for c in _COORDINATES))
+ for name in ("xcom_v1", "rtif_v1"):
+ legacy = sa.table(name, *(sa.column(c) for c in _COORDINATES))
+ found = (
+ sa.select(1)
+ .select_from(owner.join(ti, ti.c.id == owner.c.task_instance_id))
+ .where(moved, sa.exists().where(*(legacy.c[c] == owner.c[c] for c
in _COORDINATES)))
+ )
+ if bind.scalar(found.limit(1)):
+ return True
+ return False
+
+
+def _delete_moved_owners(ti):
+ owner = _build_legacy_owner_table()
+ op.execute(
+ owner.delete().where(
+ sa.exists().where(
+ ti.c.id == owner.c.task_instance_id,
+ sa.or_(*(ti.c[c] != owner.c[c] for c in _COORDINATES)),
+ )
+ )
+ )
+
+
+def _delete_legacy_rows_of_archived_attempts(ti):
+ owner = _build_legacy_owner_table()
+ archived_ids = _select_archived_ids(ti)
+ # Legacy rows owned by an archived attempt are hidden from its successor,
so keeping them would expose them again.
+ for name in ("xcom_v1", "rtif_v1"):
+ legacy = sa.table(name, *(sa.column(c) for c in _COORDINATES))
+ op.execute(
+ legacy.delete().where(
+ sa.exists().where(
+ *(owner.c[c] == legacy.c[c] for c in _COORDINATES),
+ owner.c.task_instance_id.in_(archived_ids),
+ )
+ )
+ )
+
+
+def _move_owners_to_current_attempts(ti):
+ owner = _build_legacy_owner_table()
+ current = ti.alias("current")
+ archived_ids = _select_archived_ids(ti)
+ xcom_v2 = sa.table("xcom_v2", sa.column("task_instance_id"))
+ rtif_v2 = sa.table("rtif_v2", sa.column("task_instance_id"))
+ # Deleting archived attempts must not cascade into legacy rows through an
owner that points at one.
+ op.execute(
+ owner.update()
+ .where(owner.c.task_instance_id.in_(archived_ids))
+ .values(
+ task_instance_id=sa.select(current.c.id)
+ .where(*(current.c[c] == owner.c[c] for c in _COORDINATES),
current.c.working_set.is_not(None))
+ .scalar_subquery()
+ )
+ )
+ # Legacy tables reference owners by coordinates, so attempts created after
the upgrade need an owner row first.
+ op.execute(
+ owner.insert().from_select(
+ [*_COORDINATES, "task_instance_id"],
+ sa.select(*(ti.c[c] for c in _COORDINATES), ti.c.id).where(
+ ti.c.working_set.is_not(None),
+ # An OR of two IN subqueries cannot become a semi-join, so
PostgreSQL rescans both once they outgrow work_mem.
+ ti.c.id.in_(
+
sa.select(xcom_v2.c.task_instance_id).union(sa.select(rtif_v2.c.task_instance_id))
+ ),
+ ~sa.exists().where(*(owner.c[c] == ti.c[c] for c in
_COORDINATES)),
+ ),
+ )
+ )
+
+
+def _delete_replaced_legacy_rows(v1, v2, latest, *matching):
+ op.execute(
+ v1.delete().where(
+ sa.exists().where(
+ v2.c.task_instance_id == latest.c.id,
+ *(latest.c[c] == v1.c[c] for c in _COORDINATES),
+ *matching,
+ )
+ )
+ )
+
+
+def _copy_xcom_to_legacy(latest):
+ columns = ("key", "value", "timestamp", "dag_result", "mapped_length")
+ v2 = sa.table("xcom_v2", sa.column("task_instance_id"), *(sa.column(c) for
c in columns))
+ v1 = sa.table("xcom_v1", *(sa.column(c) for c in (*_COORDINATES, *columns,
"dag_run_id")))
+ dag_run = sa.table("dag_run", sa.column("id"), sa.column("dag_id"),
sa.column("run_id"))
+ _delete_replaced_legacy_rows(v1, v2, latest, v2.c.key == v1.c.key)
+ source = v2.join(latest, v2.c.task_instance_id == latest.c.id).join(
+ dag_run, sa.and_(dag_run.c.dag_id == latest.c.dag_id, dag_run.c.run_id
== latest.c.run_id)
+ )
+ op.execute(
+ v1.insert().from_select(
+ [*_COORDINATES, *columns, "dag_run_id"],
+ sa.select(
+ *(latest.c[c] for c in _COORDINATES), *(v2.c[c] for c in
columns), dag_run.c.id
+ ).select_from(source),
+ )
+ )
+
+
+def _copy_rendered_fields_to_legacy(latest):
+ columns = ("rendered_fields", "k8s_pod_yaml")
+ v2 = sa.table("rtif_v2", sa.column("task_instance_id"), *(sa.column(c) for
c in columns))
+ v1 = sa.table("rtif_v1", *(sa.column(c) for c in (*_COORDINATES,
*columns)))
+ _delete_replaced_legacy_rows(v1, v2, latest)
+ op.execute(
+ v1.insert().from_select(
+ [*_COORDINATES, *columns],
+ sa.select(*(latest.c[c] for c in _COORDINATES), *(v2.c[c] for c in
columns)).select_from(
+ v2.join(latest, v2.c.task_instance_id == latest.c.id)
+ ),
+ )
+ )
+
+
+def _move_archived_attempts_to_history(ti):
+ history = sa.table(
+ "task_instance_history",
+ sa.column("task_instance_id"),
+ *(sa.column(c) for c in _COPY_COLUMNS),
+ sa.column("dag_version_id"),
+ )
+ archived = ti.c.working_set.is_(None)
+ archived_ids = _select_archived_ids(ti)
+ op.execute(
+ history.insert().from_select(
+ ["task_instance_id", *_COPY_COLUMNS, "dag_version_id"],
+ sa.select(ti.c.id, *(ti.c[c] for c in _COPY_COLUMNS),
ti.c.dag_version_id).where(archived),
+ )
+ )
+ hitl = sa.table("hitl_detail", sa.column("ti_id"), *(sa.column(c) for c in
_HITL_COLUMNS))
+ hitl_history = sa.table(
+ "hitl_detail_history", sa.column("ti_history_id"), *(sa.column(c) for
c in _HITL_COLUMNS)
+ )
+ op.execute(
+ hitl_history.insert().from_select(
+ ["ti_history_id", *_HITL_COLUMNS],
+ sa.select(hitl.c.ti_id, *(hitl.c[c] for c in _HITL_COLUMNS))
+ .join(ti, ti.c.id == hitl.c.ti_id)
+ .where(archived),
+ )
+ )
+ # SQLite runs this with foreign keys off, so cascades cannot be relied on.
+ for child in ("hitl_detail", "task_instance_note", "task_reschedule"):
+ child_table = sa.table(child, sa.column("ti_id"))
+
op.execute(child_table.delete().where(child_table.c.ti_id.in_(archived_ids)))
+ op.execute(ti.delete().where(archived))
+
+
def downgrade():
- """Refuse downgrade when the predecessor cannot represent retained
ownership."""
+ """Fold the latest attempt's data back into the legacy tables and restore
archived attempts to history."""
if op.get_context().as_sql:
raise RuntimeError("Offline downgrade cannot verify retained attempt
ownership")
bind = op.get_bind()
- if bind.scalar(sa.text("SELECT 1 FROM task_instance WHERE working_set IS
NULL LIMIT 1")):
- raise RuntimeError("Cannot downgrade attempt ownership with historical
attempts")
- for name in ("xcom_v2", "rtif_v2"):
- if bind.scalar(sa.text(f"SELECT 1 FROM {name} LIMIT 1")):
- raise RuntimeError(f"Cannot downgrade attempt ownership with data
in {name}")
- if bind.scalar(
- sa.text(
- "SELECT 1 FROM legacy_task_data_owner o JOIN task_instance t ON
t.id=o.task_instance_id "
- "WHERE o.dag_id<>t.dag_id OR o.task_id<>t.task_id OR
o.run_id<>t.run_id OR o.map_index<>t.map_index LIMIT 1"
+ ti = _build_task_instance_table()
+ if _has_duplicate_attempts(bind, ti):
+ raise RuntimeError("Cannot downgrade attempt ownership with multiple
attempts sharing coordinates")
+ if orphans := _count_orphaned_archived_attempts(bind, ti):
+ raise RuntimeError(
+ f"Cannot downgrade attempt ownership: {orphans} archived attempts
(task_instance rows with "
+ "working_set IS NULL) have no current attempt at their
coordinates; delete them first"
)
- ):
+ if _has_moved_owners_with_legacy_rows(bind, ti):
raise RuntimeError("Cannot downgrade attempt ownership after legacy
owners changed coordinates")
with sqlite_rebuilds(op):
+ _create_history_tables()
+ _delete_moved_owners(ti)
+ _delete_legacy_rows_of_archived_attempts(ti)
+ _move_owners_to_current_attempts(ti)
+ latest = _select_latest_attempts(ti)
+ _copy_xcom_to_legacy(latest)
+ _copy_rendered_fields_to_legacy(latest)
+ _move_archived_attempts_to_history(ti)
op.drop_table("rtif_v2")
op.drop_table("xcom_v2")
op.drop_index("ti_current_state", table_name="task_instance")
@@ -378,16 +600,24 @@ def downgrade():
batch.create_unique_constraint("task_instance_composite_key",
list(_COORDINATES))
for column in ("working_set", "archived_reason"):
batch.drop_column(column)
+ with op.batch_alter_table("task_instance_history") as batch:
+ batch.create_foreign_key(
+ "task_instance_history_ti_fkey",
+ "task_instance",
+ list(_COORDINATES),
+ list(_COORDINATES),
+ ondelete="CASCADE",
+ onupdate="CASCADE",
+ )
for old_name, name, constraint, onupdate in _LEGACY:
_redirect_legacy(name, constraint, "task_instance",
onupdate=onupdate)
op.rename_table(name, old_name)
op.drop_table("legacy_task_data_owner")
with op.batch_alter_table("log") as batch:
batch.drop_column("task_instance_id")
- _restore_history_tables()
-def _restore_history_tables():
+def _create_history_tables():
op.create_table(
"task_instance_history",
sa.Column("task_instance_id", sa.Uuid(), nullable=False),
@@ -429,13 +659,6 @@ def _restore_history_tables():
sa.Column("retry_reason", sa.String(500), nullable=True),
sa.PrimaryKeyConstraint("task_instance_id",
name="task_instance_history_pkey"),
sa.UniqueConstraint(*_COORDINATES, "try_number",
name="task_instance_history_dtrt_uq"),
- sa.ForeignKeyConstraint(
- list(_COORDINATES),
- [f"task_instance.{c}" for c in _COORDINATES],
- name="task_instance_history_ti_fkey",
- ondelete="CASCADE",
- onupdate="CASCADE",
- ),
)
op.create_index("idx_tih_dag_run", "task_instance_history", ["dag_id",
"run_id"])
op.create_table(
diff --git a/airflow-core/src/airflow/models/taskinstance.py
b/airflow-core/src/airflow/models/taskinstance.py
index fe68cfe0e48..7881b8a1b52 100644
--- a/airflow-core/src/airflow/models/taskinstance.py
+++ b/airflow-core/src/airflow/models/taskinstance.py
@@ -3175,6 +3175,7 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
total_length = None
region_id = self.region_id
+ unmapped_region_id = region_id
state: str | None = None
unmapped_ti: TaskInstance | None = session.scalars(
select(TaskInstance).where(
@@ -3230,6 +3231,21 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
else:
dr = unmapped_ti.dag_run
state = unmapped_ti.state
+ if unmapped_ti.try_number > 0:
+ # Once the unmapped coordinate is gone its archived
attempts are unreachable, and the
+ # 0142 downgrade rejects them; their XCom and rendered
fields are deleted with them.
+ session.execute(
+ delete(TaskInstance)
+ .where(
+ TaskInstance.dag_id == task.dag_id,
+ TaskInstance.task_id == task.task_id,
+ TaskInstance.run_id == run_id,
+ TaskInstance.region_id == unmapped_region_id,
+ TaskInstance.map_index == -1,
+ TaskInstance.working_set.is_(None),
+ )
+ .execution_options(synchronize_session=False,
include_all_attempts=True)
+ )
zero_index_ti_exists = exists_query(
TaskInstance.working_set.is_(True),
TaskInstance.dag_id == task.dag_id,
diff --git
a/airflow-core/tests/unit/migrations/test_0142_unify_task_attempt_ownership.py
b/airflow-core/tests/unit/migrations/test_0142_unify_task_attempt_ownership.py
index 4fa7aab1ce2..f264e4f4393 100644
---
a/airflow-core/tests/unit/migrations/test_0142_unify_task_attempt_ownership.py
+++
b/airflow-core/tests/unit/migrations/test_0142_unify_task_attempt_ownership.py
@@ -24,12 +24,15 @@ from uuid import UUID, uuid4
import pytest
import sqlalchemy as sa
from alembic import command
+from alembic.migration import MigrationContext
+from alembic.operations import Operations
from sqlalchemy.exc import IntegrityError, OperationalError
from airflow import settings
from airflow.utils.db import _get_alembic_config
from tests_common.test_utils.attempt_ownership import (
+ COORDINATES,
CURRENT_ID,
DANGLING_VERSION,
HISTORY_ID,
@@ -41,6 +44,11 @@ pytestmark = [pytest.mark.db_test,
pytest.mark.execution_timeout(60)]
PREDECESSOR = "90e4d18ccadf"
REVISION = "e7c2a91bd540"
+PREVIOUS_RELEASE = "c9f4b3e7a218"
+DAG_RUN_ID = 41
+SECOND_DAG_RUN_ID = 42
+THIRD_DAG_RUN_ID = 43
+FOURTH_DAG_RUN_ID = 44
@pytest.fixture(scope="module")
@@ -133,7 +141,7 @@ def populated_predecessor(predecessor):
dag_run = table(connection, "dag_run")
connection.execute(
dag_run.insert().values(
- id=41,
+ id=DAG_RUN_ID,
dag_id="ownership",
run_id="manual",
run_type="manual",
@@ -176,7 +184,7 @@ def populated_predecessor(predecessor):
.insert()
.values(
**coordinates,
- dag_run_id=41,
+ dag_run_id=DAG_RUN_ID,
key="return_value",
value={"legacy": True},
timestamp=NOW,
@@ -497,15 +505,6 @@ def
test_upgrade_normalizes_only_null_historical_retry_budget(
assert target_column["nullable"] is False
-def
test_downgrade_refuses_retained_history_before_changing_schema(populated_predecessor):
- connection, config = populated_predecessor
- command.upgrade(config, REVISION)
- with pytest.raises(RuntimeError, match="histor"):
- command.downgrade(config, PREDECESSOR)
- assert "xcom_v1" in sa.inspect(connection).get_table_names()
- assert "xcom" not in sa.inspect(connection).get_table_names()
-
-
def test_empty_downgrade_restores_predecessor_constraints(predecessor):
connection, config = predecessor
@@ -630,43 +629,504 @@ def
test_upgrade_preserves_uuid_children_and_legacy_cascades(populated_predecess
assert connection.scalars(sa.select(data.c.task_instance_id)).all() ==
[]
[email protected]("unsafe_data", ["xcom_v2", "rtif_v2", "moved_owner"])
-def test_downgrade_rejects_unrepresentable_ownership(populated_predecessor,
unsafe_data):
+COORDINATE_KEY = ("dag_id", "task_id", "run_id", "map_index")
+NEIGHBOURS = {
+ "task_id": {"task_id": "neighbour"},
+ "run_id": {"run_id": "second"},
+ "map_index": {"map_index": 0},
+ "dag_id": {"dag_id": "other"},
+}
+
+
+def insert_attempt(connection, **values) -> UUID:
+ attempt_id = uuid4()
+ row = (
+ COORDINATES | {"try_number": 1, "state": "success", "pool":
"default_pool", "pool_slots": 1} | values
+ )
+ connection.execute(table(connection, "task_instance",
"id").insert().values(id=attempt_id, **row))
+ return attempt_id
+
+
+def insert_dag_run(connection, dag_run_id, dag_id, run_id):
+ connection.execute(
+ table(connection, "dag_run")
+ .insert()
+ .values(
+ id=dag_run_id,
+ dag_id=dag_id,
+ run_id=run_id,
+ run_type="manual",
+ run_after=NOW,
+ state="running",
+ start_date=NOW,
+ )
+ )
+
+
+def insert_neighbour_dag_runs(connection) -> dict[tuple[str, str], int]:
+ insert_dag_run(connection, SECOND_DAG_RUN_ID, "ownership", "second")
+ insert_dag_run(connection, THIRD_DAG_RUN_ID, "other", "manual")
+ return {
+ ("ownership", "manual"): DAG_RUN_ID,
+ ("ownership", "second"): SECOND_DAG_RUN_ID,
+ ("other", "manual"): THIRD_DAG_RUN_ID,
+ }
+
+
+def insert_xcom_v2(connection, attempt_id, key, value, **values):
+ connection.execute(
+ table(connection, "xcom_v2", "id", "task_instance_id")
+ .insert()
+ .values(id=uuid4(), task_instance_id=attempt_id, key=key, value=value,
timestamp=NOW, **values)
+ )
+
+
+def insert_rtif_v2(connection, attempt_id, rendered_fields, **values):
+ connection.execute(
+ table(connection, "rtif_v2", "id", "task_instance_id")
+ .insert()
+ .values(id=uuid4(), task_instance_id=attempt_id,
rendered_fields=rendered_fields, **values)
+ )
+
+
+def retry_current_attempt(connection, **values) -> UUID:
+ ti = table(connection, "task_instance", "id")
+ connection.execute(
+ ti.update().where(ti.c.id == CURRENT_ID).values(working_set=None,
archived_reason="retry")
+ )
+ return insert_attempt(connection, try_number=3, state="running", **values)
+
+
+def read_legacy_xcom(connection):
+ rows = connection.execute(table(connection, "xcom").select())
+ return {(row.dag_id, row.run_id, row.task_id, row.map_index, row.key): row
for row in rows}
+
+
+def read_legacy_rendered_fields(connection):
+ rows = connection.execute(table(connection,
"rendered_task_instance_fields").select())
+ return {(row.dag_id, row.run_id, row.task_id, row.map_index): row for row
in rows}
+
+
+def drop_unique_constraint(connection, name):
+ sqlite = connection.dialect.name == "sqlite"
+ connection.commit()
+ if sqlite:
+ connection.exec_driver_sql("PRAGMA foreign_keys=OFF")
+ with (
+ Operations.context(MigrationContext.configure(connection)) as ops,
+ ops.batch_alter_table("task_instance") as batch,
+ ):
+ batch.drop_constraint(name, type_="unique")
+ connection.commit()
+ if sqlite:
+ connection.exec_driver_sql("PRAGMA foreign_keys=ON")
+
+
+def assert_upgraded_schema_untouched(connection):
+ tables = sa.inspect(connection).get_table_names()
+ assert {"xcom_v1", "legacy_task_data_owner"} <= set(tables)
+ assert "task_instance_history" not in tables
+ assert "hitl_detail_history" not in tables
+
+
[email protected]("legacy_store", ["xcom",
"rendered_task_instance_fields"])
+def
test_downgrade_rejects_owner_that_changed_coordinates(populated_predecessor,
legacy_store):
connection, config = populated_predecessor
- connection.execute(table(connection, "hitl_detail_history").delete())
- connection.execute(table(connection, "task_instance_history").delete())
+ other_store = {"xcom", "rendered_task_instance_fields"} - {legacy_store}
+ for name in ("hitl_detail_history", "task_instance_history", *other_store):
+ connection.execute(table(connection, name).delete())
connection.commit()
command.upgrade(config, REVISION)
- if unsafe_data == "moved_owner":
- ti = table(connection, "task_instance", "id")
- connection.execute(ti.update().where(ti.c.id ==
CURRENT_ID).values(map_index=0))
- elif unsafe_data == "xcom_v2":
+ ti = table(connection, "task_instance", "id")
+ connection.execute(ti.update().where(ti.c.id ==
CURRENT_ID).values(map_index=0))
+ connection.commit()
+ with pytest.raises(RuntimeError, match="legacy owners changed
coordinates"):
+ command.downgrade(config, PREDECESSOR)
+ assert_upgraded_schema_untouched(connection)
+
+
+def
test_downgrade_deletes_stale_owner_of_an_attempt_retried_after_promotion(populated_predecessor):
+ connection, config = populated_predecessor
+ for name in ("xcom", "rendered_task_instance_fields",
"hitl_detail_history", "task_instance_history"):
+ connection.execute(table(connection, name).delete())
+ connection.commit()
+ command.upgrade(config, REVISION)
+ ti = table(connection, "task_instance", "id")
+ connection.execute(ti.update().where(ti.c.id ==
CURRENT_ID).values(map_index=0))
+ successor_id = retry_current_attempt(connection, map_index=0)
+ insert_xcom_v2(connection, successor_id, "k", "v")
+ connection.commit()
+
+ command.downgrade(config, PREDECESSOR)
+
+ assert {key: row.value for key, row in
read_legacy_xcom(connection).items()} == {
+ ("ownership", "manual", "task", 0, "k"): "v"
+ }
+
+
+def
test_downgrade_ignores_owner_that_changed_coordinates_without_legacy_rows(populated_predecessor):
+ connection, config = populated_predecessor
+ for name in ("xcom", "rendered_task_instance_fields",
"hitl_detail_history", "task_instance_history"):
+ connection.execute(table(connection, name).delete())
+ connection.commit()
+ command.upgrade(config, REVISION)
+ ti = table(connection, "task_instance", "id")
+ connection.execute(ti.update().where(ti.c.id ==
CURRENT_ID).values(map_index=0))
+ connection.commit()
+
+ command.downgrade(config, PREDECESSOR)
+
+ live = table(connection, "task_instance", "id")
+ assert [(row.id, row.map_index) for row in
connection.execute(live.select())] == [(CURRENT_ID, 0)]
+
+
[email protected](
+ ("dropped_constraint", "duplicate"),
+ [
+ pytest.param("task_instance_current_key", {"try_number": 5},
id="second-current-attempt"),
+ pytest.param(
+ "task_instance_try_key", {"try_number": 1, "working_set": None},
id="repeated-try-number"
+ ),
+ ],
+)
+def test_downgrade_rejects_attempts_sharing_coordinates(populated_predecessor,
dropped_constraint, duplicate):
+ connection, config = populated_predecessor
+ command.upgrade(config, REVISION)
+ drop_unique_constraint(connection, dropped_constraint)
+ insert_attempt(connection, **duplicate)
+ connection.commit()
+ with pytest.raises(RuntimeError, match="multiple attempts sharing
coordinates"):
+ command.downgrade(config, PREDECESSOR)
+ assert_upgraded_schema_untouched(connection)
+
+
+def
test_downgrade_rejects_archived_attempts_without_a_current_attempt(populated_predecessor):
+ connection, config = populated_predecessor
+ insert_neighbour_dag_runs(connection)
+ connection.commit()
+ command.upgrade(config, REVISION)
+ for overrides in NEIGHBOURS.values():
+ insert_attempt(connection, **(COORDINATES | overrides))
+ ti = table(connection, "task_instance", "id")
+ connection.execute(ti.delete().where(ti.c.id == CURRENT_ID))
+ connection.commit()
+ with pytest.raises(RuntimeError, match="1 archived attempts"):
+ command.downgrade(config, PREDECESSOR)
+ assert_upgraded_schema_untouched(connection)
+
+
+def
test_downgrade_keeps_only_the_latest_attempt_of_a_task_first_run_after_upgrade(populated_predecessor):
+ connection, config = populated_predecessor
+ command.upgrade(config, REVISION)
+ coordinates = COORDINATES | {"task_id": "neighbour"}
+ first_id = insert_attempt(connection, **coordinates, working_set=None,
archived_reason="retry")
+ second_id = insert_attempt(connection, **coordinates, try_number=2)
+ for attempt_id, value in ((first_id, "first"), (second_id, "second")):
+ insert_xcom_v2(connection, attempt_id, "return_value", value)
+ insert_rtif_v2(connection, attempt_id, {"f": value})
+ connection.commit()
+
+ command.downgrade(config, PREDECESSOR)
+
+ assert {key: row.value for key, row in
read_legacy_xcom(connection).items() if key[2] == "neighbour"} == {
+ ("ownership", "manual", "neighbour", -1, "return_value"): "second"
+ }
+ assert {
+ key: row.rendered_fields
+ for key, row in read_legacy_rendered_fields(connection).items()
+ if key[2] == "neighbour"
+ } == {("ownership", "manual", "neighbour", -1): {"f": "second"}}
+
+
+def test_downgrade_replaces_legacy_xcom_key_by_key(populated_predecessor):
+ connection, config = populated_predecessor
+ dag_runs = insert_neighbour_dag_runs(connection)
+ legacy_xcom = table(connection, "xcom")
+ legacy_coordinates = [
+ COORDINATES | {"task_id": "other"},
+ *(
+ COORDINATES | {"task_id": "mapped", "run_id": run_id, "map_index":
map_index}
+ for run_id in ("manual", "second")
+ for map_index in (0, 1)
+ ),
+ ]
+ for coordinates in legacy_coordinates:
+ insert_attempt(connection, **coordinates)
connection.execute(
- table(connection, unsafe_data, "task_instance_id", "id")
- .insert()
- .values(
- id=uuid4(),
- task_instance_id=CURRENT_ID,
- key="new",
- value=1,
+ legacy_xcom.insert().values(
+ **coordinates,
+ dag_run_id=dag_runs[(coordinates["dag_id"],
coordinates["run_id"])],
+ key="legacy",
+ value="legacy",
timestamp=NOW,
)
)
- else:
+ connection.commit()
+ command.upgrade(config, REVISION)
+
+ ti = table(connection, "task_instance", "id")
+ mapped_attempts = {
+ (row.run_id, row.map_index): row.id
+ for row in connection.execute(ti.select().where(ti.c.task_id ==
"mapped"))
+ }
+ for attempt_id, key, value in (
+ (HISTORY_ID, "from_archived_attempt", "archived"),
+ (CURRENT_ID, "return_value", "latest"),
+ (CURRENT_ID, "added", "added"),
+ (mapped_attempts[("manual", 0)], "extra", "extra"),
+ (mapped_attempts[("second", 1)], "legacy", "replaced"),
+ ):
+ insert_xcom_v2(connection, attempt_id, key, value)
+ connection.commit()
+
+ command.downgrade(config, PREDECESSOR)
+
+ xcom = read_legacy_xcom(connection)
+ assert {key: (row.value, row.dag_run_id) for key, row in xcom.items()} == {
+ ("ownership", "manual", "task", -1, "return_value"): ("latest",
DAG_RUN_ID),
+ ("ownership", "manual", "task", -1, "added"): ("added", DAG_RUN_ID),
+ ("ownership", "manual", "other", -1, "legacy"): ("legacy", DAG_RUN_ID),
+ ("ownership", "manual", "mapped", 0, "legacy"): ("legacy", DAG_RUN_ID),
+ ("ownership", "manual", "mapped", 0, "extra"): ("extra", DAG_RUN_ID),
+ ("ownership", "manual", "mapped", 1, "legacy"): ("legacy", DAG_RUN_ID),
+ ("ownership", "second", "mapped", 0, "legacy"): ("legacy",
SECOND_DAG_RUN_ID),
+ ("ownership", "second", "mapped", 1, "legacy"): ("replaced",
SECOND_DAG_RUN_ID),
+ }
+
+
+def test_downgrade_carries_payload_columns_back(populated_predecessor):
+ connection, config = populated_predecessor
+ command.upgrade(config, REVISION)
+ created_id = insert_attempt(connection, task_id="created_after_upgrade")
+ insert_xcom_v2(connection, CURRENT_ID, "return_value", "latest",
dag_result=True, mapped_length=3)
+ insert_xcom_v2(connection, created_id, "return_value", "created")
+ insert_rtif_v2(connection, HISTORY_ID, {"f": "archived"})
+ insert_rtif_v2(connection, CURRENT_ID, {"f": "latest"},
k8s_pod_yaml={"kind": "Pod"})
+ insert_rtif_v2(connection, created_id, {"f": "created"})
+ connection.commit()
+
+ command.downgrade(config, PREDECESSOR)
+
+ xcom = read_legacy_xcom(connection)
+ assert {key: (row.dag_result, row.mapped_length) for key, row in
xcom.items()} == {
+ ("ownership", "manual", "task", -1, "return_value"): (True, 3),
+ ("ownership", "manual", "created_after_upgrade", -1, "return_value"):
(None, None),
+ }
+ rendered_fields = read_legacy_rendered_fields(connection)
+ assert {key: (row.rendered_fields, row.k8s_pod_yaml) for key, row in
rendered_fields.items()} == {
+ ("ownership", "manual", "task", -1): ({"f": "latest"}, {"kind":
"Pod"}),
+ ("ownership", "manual", "created_after_upgrade", -1): ({"f":
"created"}, None),
+ }
+
+
+def
test_downgrade_copies_attempts_that_wrote_only_one_store(populated_predecessor):
+ connection, config = populated_predecessor
+ command.upgrade(config, REVISION)
+ insert_xcom_v2(connection, insert_attempt(connection,
task_id="only_xcom"), "k", "v")
+ insert_rtif_v2(connection, insert_attempt(connection,
task_id="only_rtif"), {"f": "v"})
+ connection.commit()
+
+ command.downgrade(config, PREDECESSOR)
+
+ assert {key[2:] for key in read_legacy_xcom(connection) if key[2] !=
"task"} == {("only_xcom", -1, "k")}
+ assert {key[2:] for key in read_legacy_rendered_fields(connection) if
key[2] != "task"} == {
+ ("only_rtif", -1)
+ }
+
+
+def
test_downgrade_copies_attempts_created_next_to_an_owned_coordinate(populated_predecessor):
+ connection, config = populated_predecessor
+ insert_neighbour_dag_runs(connection)
+ connection.commit()
+ command.upgrade(config, REVISION)
+ for overrides in NEIGHBOURS.values():
+ insert_xcom_v2(connection, insert_attempt(connection, **(COORDINATES |
overrides)), "k", "v")
+ connection.commit()
+
+ command.downgrade(config, PREDECESSOR)
+
+ legacy_xcom = table(connection, "xcom")
+ copied = connection.execute(legacy_xcom.select().where(legacy_xcom.c.key
== "k")).all()
+ assert sorted(tuple(getattr(row, name) for name in COORDINATE_KEY) for row
in copied) == sorted(
+ tuple((COORDINATES | overrides)[name] for name in COORDINATE_KEY) for
overrides in NEIGHBOURS.values()
+ )
+
+
+def
test_downgrade_keeps_legacy_data_of_neighbouring_coordinates(populated_predecessor):
+ connection, config = populated_predecessor
+ dag_runs = insert_neighbour_dag_runs(connection)
+ for overrides in NEIGHBOURS.values():
+ coordinates = COORDINATES | overrides
+ insert_attempt(connection, **coordinates)
connection.execute(
- table(connection, unsafe_data, "task_instance_id", "id")
+ table(connection, "xcom")
.insert()
.values(
- id=uuid4(),
- task_instance_id=CURRENT_ID,
- rendered_fields={},
+ **coordinates,
+ dag_run_id=dag_runs[(coordinates["dag_id"],
coordinates["run_id"])],
+ key="legacy",
+ value="legacy",
+ timestamp=NOW,
)
)
+ connection.execute(
+ table(connection, "rendered_task_instance_fields")
+ .insert()
+ .values(**coordinates, rendered_fields={"f": "legacy"})
+ )
connection.commit()
- with pytest.raises(RuntimeError, match="Cannot downgrade"):
- command.downgrade(config, PREDECESSOR)
- assert "xcom_v1" in sa.inspect(connection).get_table_names()
- assert "rtif_v2" in sa.inspect(connection).get_table_names()
+ command.upgrade(config, REVISION)
+ retry_current_attempt(connection)
+ connection.commit()
+
+ command.downgrade(config, PREDECESSOR)
+
+ expected = sorted(tuple((COORDINATES | o)[name] for name in
COORDINATE_KEY) for o in NEIGHBOURS.values())
+ xcom = connection.execute(table(connection, "xcom").select()).all()
+ assert sorted(tuple(getattr(row, name) for name in COORDINATE_KEY) for row
in xcom) == expected
+ rendered_fields = connection.execute(table(connection,
"rendered_task_instance_fields").select()).all()
+ assert sorted(tuple(getattr(row, name) for name in COORDINATE_KEY) for row
in rendered_fields) == expected
+
+
+def test_downgrade_keeps_dags_sharing_a_run_id_apart(populated_predecessor):
+ connection, config = populated_predecessor
+ insert_neighbour_dag_runs(connection)
+ insert_dag_run(connection, FOURTH_DAG_RUN_ID, "third", "manual")
+ insert_attempt(connection, dag_id="third")
+ connection.execute(
+ table(connection, "xcom")
+ .insert()
+ .values(
+ **(COORDINATES | {"dag_id": "third"}),
+ dag_run_id=FOURTH_DAG_RUN_ID,
+ key="legacy",
+ value="third",
+ timestamp=NOW,
+ )
+ )
+ connection.commit()
+ command.upgrade(config, REVISION)
+ other_id = insert_attempt(connection, dag_id="other", try_number=2)
+ insert_xcom_v2(connection, CURRENT_ID, "k", "ownership")
+ insert_xcom_v2(connection, other_id, "k", "other")
+ connection.commit()
+
+ command.downgrade(config, PREDECESSOR)
+
+ xcom = read_legacy_xcom(connection)
+ assert {(key[0], key[4]): (row.value, row.dag_run_id) for key, row in
xcom.items()} == {
+ ("other", "k"): ("other", THIRD_DAG_RUN_ID),
+ ("ownership", "k"): ("ownership", DAG_RUN_ID),
+ ("ownership", "return_value"): ({"legacy": True}, DAG_RUN_ID),
+ ("third", "legacy"): ("third", FOURTH_DAG_RUN_ID),
+ }
+
+
+def
test_downgrade_drops_legacy_data_of_a_retried_attempt_whose_successor_wrote_nothing(
+ populated_predecessor,
+):
+ connection, config = populated_predecessor
+ command.upgrade(config, REVISION)
+ retry_current_attempt(connection)
+ connection.commit()
+
+ command.downgrade(config, PREDECESSOR)
+
+ assert read_legacy_xcom(connection) == {}
+ assert read_legacy_rendered_fields(connection) == {}
+
+
+def
test_downgrade_replaces_legacy_data_of_a_retried_attempt_with_its_successors(populated_predecessor):
+ connection, config = populated_predecessor
+ command.upgrade(config, REVISION)
+ successor_id = retry_current_attempt(connection)
+ insert_xcom_v2(connection, successor_id, "from_successor", "new")
+ insert_rtif_v2(connection, successor_id, {"f": "new"})
+ connection.commit()
+
+ command.downgrade(config, PREDECESSOR)
+
+ assert {key: row.value for key, row in
read_legacy_xcom(connection).items()} == {
+ ("ownership", "manual", "task", -1, "from_successor"): "new"
+ }
+ assert {key: row.rendered_fields for key, row in
read_legacy_rendered_fields(connection).items()} == {
+ ("ownership", "manual", "task", -1): {"f": "new"}
+ }
+
+
+def
test_downgrade_restores_archived_attempts_to_history(populated_predecessor):
+ connection, config = populated_predecessor
+ history = table(connection, "task_instance_history", "task_instance_id",
"dag_version_id")
+ hitl_history = table(connection, "hitl_detail_history", "ti_history_id")
+ source_history = connection.execute(history.select()).one()._mapping
+ source_hitl = connection.execute(hitl_history.select()).one()._mapping
+ connection.commit()
+ command.upgrade(config, REVISION)
+ connection.execute(
+ table(connection, "task_instance_note", "ti_id")
+ .insert()
+ .values(ti_id=HISTORY_ID, content="archived note", created_at=NOW,
updated_at=NOW)
+ )
+ connection.execute(
+ table(connection, "task_reschedule", "ti_id")
+ .insert()
+ .values(ti_id=HISTORY_ID, start_date=NOW, end_date=NOW, duration=0,
reschedule_date=NOW)
+ )
+ connection.execute(
+ table(connection, "hitl_detail", "ti_id")
+ .insert()
+ .values(ti_id=CURRENT_ID, options=["a"], subject="live", params={},
params_input={}, created_at=NOW)
+ )
+ connection.commit()
+
+ command.downgrade(config, PREDECESSOR)
+
+ restored = connection.execute(history.select()).one()._mapping
+ for name, value in source_history.items():
+ if name not in {"trigger_id", "dag_version_id", "max_tries"}:
+ assert restored[name] == value
+ assert connection.execute(hitl_history.select()).one()._mapping ==
source_hitl
+ live = table(connection, "task_instance", "id")
+ assert connection.scalars(sa.select(live.c.id)).all() == [CURRENT_ID]
+ for name in ("task_instance_note", "task_reschedule", "hitl_detail"):
+ child = table(connection, name, "ti_id")
+ assert connection.scalars(sa.select(child.c.ti_id)).all() ==
[CURRENT_ID]
+
+ connection.commit()
+ command.upgrade(config, REVISION)
+ live = table(connection, "task_instance", "id")
+ assert
connection.scalars(sa.select(live.c.id).where(live.c.working_set.is_(None))).all()
== [HISTORY_ID]
+ hitl = table(connection, "hitl_detail", "ti_id")
+ assert sorted(connection.scalars(sa.select(hitl.c.ti_id)).all()) ==
sorted([CURRENT_ID, HISTORY_ID])
+
+
[email protected](
+ ("state", "expected_try_number"),
+ [
+ pytest.param("up_for_retry", 2, id="retried"),
+ pytest.param(None, 2, id="cleared"),
+ pytest.param("scheduled", 3, id="scheduled"),
+ pytest.param("running", 3, id="running"),
+ ],
+)
+def
test_downgrade_to_the_previous_release_holds_a_pending_attempt_at_its_last_archived_try(
+ populated_predecessor, state, expected_try_number
+):
+ connection, config = populated_predecessor
+ command.upgrade(config, REVISION)
+ successor_id = retry_current_attempt(connection)
+ ti = table(connection, "task_instance", "id")
+ connection.execute(ti.update().where(ti.c.id ==
successor_id).values(state=state))
+ connection.commit()
+
+ command.downgrade(config, PREVIOUS_RELEASE)
+
+ ti = table(connection, "task_instance", "id")
+ history = table(connection, "task_instance_history", "task_instance_id")
+ assert connection.scalars(sa.select(ti.c.try_number).where(ti.c.id ==
successor_id)).one() == (
+ expected_try_number
+ )
+ assert sorted(connection.scalars(sa.select(history.c.try_number))) == [1,
2]
@pytest.mark.parametrize("source_problem", ["unvalidated", "reanchored"])
diff --git a/airflow-core/tests/unit/models/test_mappedoperator.py
b/airflow-core/tests/unit/models/test_mappedoperator.py
index 0f280daed02..1d908059894 100644
--- a/airflow-core/tests/unit/models/test_mappedoperator.py
+++ b/airflow-core/tests/unit/models/test_mappedoperator.py
@@ -331,6 +331,75 @@ def test_expand_mapped_task_failed_state_in_db(dag_maker,
session):
queue_workload.assert_called_once()
[email protected]("zero_index_exists", [False, True])
+def
test_expand_mapped_task_deletes_archived_attempts_of_the_unmapped_task_instance(
+ dag_maker, session, zero_index_exists
+):
+ with dag_maker(session=session, serialized=True) as dag:
+ task1 = BaseOperator(task_id="op1")
+ other = BaseOperator(task_id="other")
+ mapped =
MockOperator.partial(task_id="task_2").expand(arg2=task1.output)
+ dr = dag_maker.create_dagrun()
+ other_dr = dag_maker.create_dagrun(run_id="other_run",
logical_date=DEFAULT_DATE + timedelta(days=1))
+ mapped_deser = dag.task_dict[mapped.task_id]
+ dr.get_task_instance(mapped.task_id,
session=session).prepare_db_for_next_try(
+ session
+ ).prepare_db_for_next_try(session)
+
+ zero_index = TaskInstance(
+ mapped_deser,
+ run_id=dr.run_id,
+ map_index=0,
+ state=TaskInstanceState.SUCCESS,
+ dag_version_id=DagVersion.get_latest_version(dr.dag_id).id,
+ )
+ session.add(zero_index)
+ session.flush()
+ if zero_index_exists:
+ zero_index.prepare_db_for_next_try(session)
+ else:
+ zero_index.archive(reason="retry", session=session)
+ other_task_ti = dr.get_task_instance(other.task_id, session=session)
+ other_task_ti.prepare_db_for_next_try(session)
+ other_run_ti = other_dr.get_task_instance(mapped.task_id, session=session)
+ other_run_ti.prepare_db_for_next_try(session)
+ with dag_maker(dag_id="other_dag", session=session, serialized=True):
+ BaseOperator(task_id=mapped.task_id)
+ other_dag_ti = dag_maker.create_dagrun(run_id=dr.run_id).get_task_instance(
+ mapped.task_id, session=session
+ )
+ other_dag_ti.prepare_db_for_next_try(session)
+ survivors = {zero_index.id, other_task_ti.id, other_run_ti.id,
other_dag_ti.id}
+ push_mapped_length(dr.get_task_instance(task1.task_id, session=session),
[1, 2], session=session)
+ session.flush()
+
+ expand_mapped_task_instances(mapped_deser, dr.run_id, session=session)
+
+ def get_ids(*conditions):
+ return set(
+ session.scalars(
+
select(TaskInstance.id).where(*conditions).execution_options(include_all_attempts=True)
+ )
+ )
+
+ placeholders = (
+ TaskInstance.dag_id == dr.dag_id,
+ TaskInstance.run_id == dr.run_id,
+ TaskInstance.task_id == mapped.task_id,
+ TaskInstance.map_index == -1,
+ )
+ assert get_ids(*placeholders) == set()
+ assert get_ids(TaskInstance.id.in_(survivors)) == survivors
+ current_indexes = session.scalars(
+ select(TaskInstance.map_index).where(
+ TaskInstance.dag_id == dr.dag_id,
+ TaskInstance.run_id == dr.run_id,
+ TaskInstance.task_id == mapped.task_id,
+ )
+ )
+ assert set(current_indexes) == {0, 1}
+
+
def test_stable_mapped_indexes_do_not_query_historical_max_try(dag_maker,
session):
with dag_maker(session=session, serialized=True) as dag:
upstream = BaseOperator(task_id="upstream")
@@ -2349,3 +2418,60 @@ def
test_expansion_nested_under_a_loop_pass_is_a_child_region_that_revises_in_pl
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]
+
+
[email protected]("sibling_end_state", [TaskInstanceState.SKIPPED,
TaskInstanceState.UPSTREAM_FAILED])
+def test_expansion_keeps_the_archived_attempts_of_a_sibling_region_placeholder(
+ dag_maker, session, sibling_end_state
+):
+ 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.get_or_create(
+ dag_id=dr.dag_id, run_id=dr.run_id, node_id="loop", session=session
+ )
+ session.add(loop_region)
+ session.flush()
+ placeholders: 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,
+ )
+ placeholders.append(ti)
+ yield ti
+
+ for loop_pass in (0, 1):
+ session.add_all(
+ list(
+ dr._create_tasks(
+ [serialized], creator, session=session,
parent_region=(loop_region.id, loop_pass)
+ )
+ )
+ )
+ session.flush()
+ first, second = placeholders
+ first.prepare_db_for_next_try(session).state = sibling_end_state
+ expanding = second.prepare_db_for_next_try(session)
+ session.flush()
+
+ expanding.expand_mapped_task(session=session)
+
+ archived_ids = session.scalars(
+ select(TaskInstance.id)
+ .where(
+ TaskInstance.working_set.is_(None),
+ TaskInstance.task_id == mapped.task_id,
+ TaskInstance.region_id == first.region_id,
+ )
+ .execution_options(include_all_attempts=True)
+ ).all()
+ assert archived_ids == [first.id]
diff --git a/airflow-core/tests/unit/models/test_taskinstance.py
b/airflow-core/tests/unit/models/test_taskinstance.py
index c11c011ac2c..4a7721b3a41 100644
--- a/airflow-core/tests/unit/models/test_taskinstance.py
+++ b/airflow-core/tests/unit/models/test_taskinstance.py
@@ -3992,7 +3992,7 @@ 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")
- with assert_queries_count(11, margin=2):
+ with assert_queries_count(10):
mapped_tis, max_map_index = expand_mapped_task_instances(
show_task, dag_run.run_id, session=session
)
diff --git a/reproducible_build.yaml b/reproducible_build.yaml
index a0bd2678b32..1f2faf497cb 100644
--- a/reproducible_build.yaml
+++ b/reproducible_build.yaml
@@ -1,2 +1,2 @@
-release-notes-hash: 1f594372ec3c483e2b104aa4d46361e7
-source-date-epoch: 1791890181
+release-notes-hash: 4a6502f86dfc61b17c0cea5d58457d54
+source-date-epoch: 1792070450
