You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "ashb (via GitHub)" <gi...@apache.org> on 2026/10/07 14:18:17 UTC
[PR] Allow downgrading past the TI ownership migration when data is present. [airflow]
ashb opened a new pull request, #74407:
URL: https://github.com/apache/airflow/pull/74407
The downgrade refused whenever XCom_v2 or rendered field data had rows, or
any task retried since the upgrade, so a rollback after any real use of 3.4
was impossible.
This fixes it, as long as the data can be downgraded non-destructively
"as if 3.3 wrote it". This means you might loose data that 3.3 wouldn't
have kept for archived tries.
Document the cost and caveats in the 3.4.0 release notes.
<!-- SPDX-License-Identifier: Apache-2.0
https://www.apache.org/licenses/LICENSE-2.0 -->
<!--
Thank you for contributing!
Please provide above a brief description of the changes made in this pull request.
Write a good git commit message following this guide: https://chris.beams.io/posts/git-commit/
Please make sure that your code changes are covered with tests.
And in case of new features or big changes remember to adjust the documentation.
For user-facing UI changes, please attach before/after screenshots (or a short
screen recording) so reviewers can assess the visual impact.
Feel free to ping (in general) for the review if you do not see reaction for a few days
(72 Hours is the minimum reaction time you can expect from volunteers) - we sometimes miss notifications.
In case of an existing issue, reference it using one of the following:
* closes: #ISSUE
* related: #ISSUE
-->
---
##### Was generative AI tooling used to co-author this PR?
<!--
If generative AI tooling has been used in the process of authoring this PR, please
change below checkbox to `[X]` followed by the name of the tool, uncomment the "Generated-by".
-->
- [ ] Yes (please specify the tool below)
<!--
Generated-by: [Tool Name] following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
-->
---
* Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
* For fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
* When adding dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
* For significant user-facing changes create newsfragment: `{pr_number}.significant.rst`, in [airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments). You can add this file in a follow-up commit after the PR is created so you know the PR number.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Allow downgrading past the TI ownership migration when data is present. [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #74407:
URL: https://github.com/apache/airflow/pull/74407#discussion_r4281530683
##########
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),
+ ti.c.id.in_(sa.select(xcom_v2.c.task_instance_id))
+ | ti.c.id.in_(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():
+ _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)
Review Comment:
Try number is part of 0138 downgrade not this one.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Allow downgrading past the TI ownership migration when data is present. [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #74407:
URL: https://github.com/apache/airflow/pull/74407#discussion_r4217447996
##########
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:
Yes, I agree, this is exactly where unit tests are useful still.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Allow downgrading past the TI ownership migration when data is present. [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb merged PR #74407:
URL: https://github.com/apache/airflow/pull/74407
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Allow downgrading past the TI ownership migration when data is present. [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
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 the L595 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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Allow downgrading past the TI ownership migration when data is present. [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #74407:
URL: https://github.com/apache/airflow/pull/74407#discussion_r4208048147
##########
RELEASE_NOTES.rst:
##########
@@ -24,6 +24,21 @@
.. towncrier release notes start
+Airflow 3.4.0 (unreleased)
Review Comment:
@vatsrahul1001 Of especial note for you here:
When we use towncrier to generate the release notes it will create a new 3.4.0 header. Please sort out the duplicate headers and put this note "somewhere appropriate near the top"
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Allow downgrading past the TI ownership migration when data is present. [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #74407:
URL: https://github.com/apache/airflow/pull/74407#discussion_r4269807555
##########
airflow-core/src/airflow/models/taskinstance.py:
##########
@@ -2581,6 +2581,20 @@ def expand_mapped_task(self, *, session: Session) -> tuple[Sequence[TaskInstance
else:
dr = unmapped_ti.dag_run
state = unmapped_ti.state
+ if unmapped_ti.try_number > 0:
+ # Archived attempts only exist once the placeholder has been retried. They would have
Review Comment:
These aren't always placeholder attempts. `_check_for_removed_or_restored_tasks` keeps a sole -1 TI for expansion when a task becomes mapped, so an unmapped task that ran and was then removed and restored as mapped (restore goes through `prepare_db_for_next_try`), or was `up_for_retry` when the DAG changed, reaches here with archived attempts that really ran. This deletes them along with their XCom and rendered fields, where 3.3 moved those history rows to index 0 through the `ON UPDATE CASCADE`. I haven't run either path. It's probably rare enough to accept, but then the comment should say so rather than calling them placeholders.
##########
RELEASE_NOTES.rst:
##########
@@ -24,6 +24,25 @@
.. towncrier release notes start
+Airflow 3.4.0 (unreleased)
+--------------------------
+
+.. 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
Review Comment:
This still reads as if the cost scales with rows written after the upgrade. On Postgres the downgrade re-adds `xcom_task_instance_fkey` and `rtif_ti_fkey` as validated constraints through `_redirect_legacy`, which scans every row of the legacy tables, pre-upgrade ones included, and SQLite's batch path copies them. On the 300m row table mentioned here that's the slow part of the rollback. Could the note say the cost on Postgres and SQLite depends on the total size of `xcom` and `rendered_task_instance_fields`?
##########
airflow-core/tests/unit/models/test_mappedoperator.py:
##########
@@ -228,6 +228,66 @@ def test_expand_mapped_task_failed_state_in_db(dag_maker, session):
queue_workload.assert_called_once()
+@pytest.mark.parametrize("zero_index_exists", [False, True])
+def test_expand_mapped_task_deletes_archived_attempts_of_the_placeholder(
+ 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)
+ survivors = {zero_index.id, other_task_ti.id, other_run_ti.id}
Review Comment:
Nothing here binds the `dag_id` filter on the new delete. `other_task_ti` and `other_run_ti` cover `task_id` and `run_id`, but no other DAG with the same task_id and run_id is seeded. Run IDs like `scheduled__...` collide across DAGs all the time, so losing that filter would delete another DAG's attempts. Could a second DAG with the same mapped task_id and run_id go into `survivors`?
##########
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),
+ ti.c.id.in_(sa.select(xcom_v2.c.task_instance_id))
+ | ti.c.id.in_(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():
+ _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)
Review Comment:
This keeps the current row's `try_number`, but 3.4 has already bumped it for a pending retry while 3.3.2 bumps it again in `schedule_tis` (`TI.try_number + 1`). With `retries=2` and try 1 failed on 3.4, the TI comes back as try 2 in `up_for_retry`, 3.3 runs it as try 3, and `3 <= 2` fails the task one retry early. Same for a TI that was cleared but not scheduled yet. Could the downgrade decrement `try_number` on current rows in `up_for_retry` or NULL state that have an archived attempt at `try_number - 1`, plus a test for each?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Allow downgrading past the TI ownership migration when data is present. [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #74407:
URL: https://github.com/apache/airflow/pull/74407#discussion_r4211011947
##########
RELEASE_NOTES.rst:
##########
@@ -24,6 +24,21 @@
.. towncrier release notes start
+Airflow 3.4.0 (unreleased)
+--------------------------
+
+.. note::
+ Downgrading from this release to 3.4.0 can be slow, depending on the number of
Review Comment:
"from this release to 3.4.0" should say 3.3.x, I think.
The cost also isn't only proportional to rows written after the upgrade. On Postgres `_redirect_legacy` re-adds the `xcom` and `rendered_task_instance_fields` FKs as validated constraints, which scans the whole legacy table, and that can't be `NOT VALID` because a re-upgrade's `_check_source` requires a validated FK.
The PR description says the caveats are documented here, but the note doesn't mention that XCom and rendered fields of every attempt except the latest are dropped, or that the downgrade can still refuse (orphaned archived attempts, owners whose coordinates moved, offline `--sql`). Could those go in too, so nobody finds out mid-rollback?
##########
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),
+ ti.c.id.in_(sa.select(xcom_v2.c.task_instance_id))
+ | ti.c.id.in_(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):
Review Comment:
This refusal is reachable from a normal workflow after the upgrade, not only from hand-deleted rows. If a mapped task's `map_index=-1` placeholder is `upstream_failed` and someone clears the failed upstream plus downstream, `clear_task_instances` calls `prepare_db_for_next_try` on it, which archives the -1 row and inserts a new current -1 row. Once the upstream succeeds, `expand_mapped_task` rewrites that current row to `map_index=0` (or deletes it when index 0 already exists), and nothing moves the archived -1 row along with it. The downgrade then counts it as an orphan and tells the user to delete it.
On 3.3 the `ON UPDATE CASCADE` / `ON DELETE CASCADE` on `task_instance_history_ti_fkey` carried that history row to index 0 or dropped it. Since 3.4 hasn't shipped, would it be simpler to have `expand_mapped_task` move (or delete) the archived -1 attempts together with the placeholder, so this guard stays a true "can't happen"? Otherwise the downgrade needs to fold these itself.
##########
airflow-core/tests/unit/migrations/test_0142_unify_task_attempt_ownership.py:
##########
@@ -630,43 +628,442 @@ def test_upgrade_preserves_uuid_children_and_legacy_cascades(populated_predecess
assert connection.scalars(sa.select(data.c.task_instance_id)).all() == []
-@pytest.mark.parametrize("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")
+
+
+@pytest.mark.parametrize("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"):
Review Comment:
The removed refusal tests also asserted the schema was untouched afterwards (`xcom_v1` still present). Without that, moving a guard below `_create_history_tables()` would still pass here, and on MySQL that leaves `task_instance_history` behind so the next attempt fails with "table already exists". Could this one and the refusals at L789 and L803 keep the `get_table_names()` check?
##########
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:
As far as I can see no test binds this `working_set` filter. The only archived attempt with `xcom_v2`/`rtif_v2` rows in the tests is `HISTORY_ID`, which sits at a coordinate that already has an owner, so `~exists(owner)` skips it either way. The shape the filter protects is probably the most common one: a task first run after the upgrade that retried, with both tries writing XCom. Without the filter both attempts get an owner row for the same coordinates and `legacy_task_data_owner_pkey` aborts the downgrade. Could we add a test for that (archived try 1 and current try 2 at a post-upgrade coordinate, both writing `return_value` and rendered fields, only try 2 survives)?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Allow downgrading past the TI ownership migration when data is present. [airflow]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #74407:
URL: https://github.com/apache/airflow/pull/74407#issuecomment-6191122462
### Backport failed to create: v3-4-test. View the failure log <a href='https://github.com/apache/airflow/actions/runs/38728441080'> Run details </a>
Note: As of [Merging PRs targeted for Airflow 3.X](https://github.com/apache/airflow/blob/main/dev/README_AIRFLOW3_DEV.md#merging-prs-targeted-for-airflow-3x)
the committer who merges the PR is responsible for backporting the PRs that are bug fixes (generally speaking) to the maintenance branches.
In matter of doubt please ask in [#release-management](https://apache-airflow.slack.com/archives/C03G9H97MM2) Slack channel.
<table>
<tr>
<th>Status</th>
<th>Branch</th>
<th>Result</th>
</tr>
<tr>
<td>❌</td>
<td>v3-4-test</td>
<td><a href="https://github.com/apache/airflow/commit/a87c6ddb52a07795f1fd9799fc4396ee1418e4a9"><img src='https://img.shields.io/badge/Commit-a87c6dd-red' alt='Commit Link'></a></td>
</tr>
</table>
You can attempt to backport this manually by running:
```bash
cherry_picker a87c6dd v3-4-test
```
This should apply the commit to the v3-4-test branch and leave the commit in conflict state marking
the files that need manual conflict resolution.
After you have resolved the conflicts, you can continue the backport process by running:
```bash
cherry_picker --continue
```
If you don't have cherry-picker installed, see the [installation guide](https://github.com/apache/airflow/blob/main/dev/README.md#how-to-backport-pr-with-cherry-picker-cli).
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Allow downgrading past the TI ownership migration when data is present. [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #74407:
URL: https://github.com/apache/airflow/pull/74407#discussion_r4281823410
##########
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),
+ ti.c.id.in_(sa.select(xcom_v2.c.task_instance_id))
+ | ti.c.id.in_(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():
+ _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)
Review Comment:
Yup, my bad
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Allow downgrading past the TI ownership migration when data is present. [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on PR #74407:
URL: https://github.com/apache/airflow/pull/74407#issuecomment-6193614088
Pushed to 2af7b35dd33 on v3-4-test.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Allow downgrading past the TI ownership migration when data is present. [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org