kaxil commented on code in PR #74407:
URL: https://github.com/apache/airflow/pull/74407#discussion_r4211328572
##########
airflow-core/src/airflow/migrations/versions/0142_3_4_0_unify_task_attempt_ownership.py:
##########
@@ -370,24 +370,244 @@ 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),
Review Comment:
I know we've discussed this, and per #73953 these tests stop paying off once
the migration is released. 0142 isn't released yet, though, and it's still
changing: this PR rewrote the downgrade, and
https://github.com/apache/airflow/pull/74407#discussion_r4211011934 comment
would change it again. A one-off downgrade run against a DB with this shape
would cover it just as well if you'd rather not add the test. I mostly want it
run once with real data.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
