You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by as...@apache.org on 2026/10/06 10:41:58 UTC
(airflow) 02/09: fixup! Keep retired task attempts and their data under the attempt UUID
This is an automated email from the ASF dual-hosted git repository.
ashb pushed a commit to branch store-historic-ti-ownership-data
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit c42c95c159492ed670ba1b9781ac12fb2ac848a2
Author: Ash Berlin-Taylor <as...@apache.org>
AuthorDate: Mon Oct 5 13:55:46 2026 +0100
fixup! Keep retired task attempts and their data under the attempt UUID
Keep older clients' responses, carry notes across attempts, and fix same-try history
---
.../execution_api/routes/task_instances.py | 4 +-
.../airflow/api_fastapi/execution_api/security.py | 13 +++---
.../execution_api/versions/v2026_10_30.py | 2 +-
.../0142_3_4_0_unify_task_attempt_ownership.py | 25 +++++++++++
airflow-core/src/airflow/models/taskinstance.py | 6 +++
.../core_api/routes/public/test_task_instances.py | 2 +-
.../api_fastapi/execution_api/test_security.py | 27 +++---------
.../versions/v2026_10_30/test_task_instances.py | 49 +++++++++++++++++++++
.../test_0142_unify_task_attempt_ownership.py | 50 +++++++++++++++++++---
.../tests/unit/models/test_taskinstance.py | 28 +++++++++++-
10 files changed, 164 insertions(+), 42 deletions(-)
diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
index cf7ecdcfadd..b9a36103e3b 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
@@ -1090,9 +1090,7 @@ async def ti_heartbeat(
if working_set is None:
# A retired attempt was likely cleared while running, so return 410 Gone
# instead of 404 Not Found to give the client a more specific signal.
- _raise_ti_not_in_live_table(
- task_instance_id, archived_in_history=IdentifyRetiredTaskStateUpdates.is_applied
- )
+ _raise_ti_not_in_live_table(task_instance_id, archived_in_history=True)
if hostname != ti_payload.hostname or pid != ti_payload.pid:
log.warning(
diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/security.py b/airflow-core/src/airflow/api_fastapi/execution_api/security.py
index 040467ce34f..1f10e1c04a5 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/security.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/security.py
@@ -233,17 +233,14 @@ async def require_auth(
# The versions package imports routes, which depend on this module.
from airflow.api_fastapi.execution_api.versions.v2026_10_30 import IdentifyRetiredTaskStateUpdates
- await _require_live_attempt(
- token,
- allow_callback="task_instance_id" not in request.path_params,
- identify_retired=IdentifyRetiredTaskStateUpdates.is_applied,
- )
- request.scope[_REQUEST_SCOPE_LIVE_ATTEMPT_KEY] = True
+ if IdentifyRetiredTaskStateUpdates.is_applied:
+ await _require_live_attempt(token, allow_callback="task_instance_id" not in request.path_params)
+ request.scope[_REQUEST_SCOPE_LIVE_ATTEMPT_KEY] = True
return token
-async def _require_live_attempt(token: TIToken, *, allow_callback: bool, identify_retired: bool) -> None:
+async def _require_live_attempt(token: TIToken, *, allow_callback: bool) -> None:
"""
Reject mutations from an attempt whose UUID is no longer in the working set.
@@ -264,7 +261,7 @@ async def _require_live_attempt(token: TIToken, *, allow_callback: bool, identif
and await session.scalar(select(Callback.id).where(Callback.id == token.id))
):
return
- archived = identify_retired and attempt is not None
+ archived = attempt is not None
raise HTTPException(
status_code=status.HTTP_410_GONE if archived else status.HTTP_404_NOT_FOUND,
detail={
diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py
index aef90d41506..8ca10d7c0bd 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py
@@ -46,7 +46,7 @@ class AddStoppedTaskReport(VersionChange):
class IdentifyRetiredTaskStateUpdates(VersionChangeWithSideEffects):
- """Return 410 for requests from archived attempts, preserving 404 for older clients and unknown attempts."""
+ """Reject every mutation from an archived attempt with 410; older clients keep each endpoint's own response."""
description = __doc__
instructions_to_migrate_to_previous_version = ()
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 689f2575081..0b39848899e 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
@@ -82,6 +82,7 @@ _COPY_COLUMNS = (
"retry_delay_override",
"retry_reason",
)
+_NOT_EXECUTING_STATES = ("success", "failed", "skipped", "upstream_failed", "removed", "restarting")
_HITL_COLUMNS = (
"options",
"subject",
@@ -223,6 +224,30 @@ def upgrade():
"legacy_task_data_owner",
not_valid=op.get_bind().dialect.name == "postgresql",
)
+ live = sa.table(
+ "task_instance",
+ *(sa.column(c) for c in _COORDINATES),
+ sa.column("try_number"),
+ sa.column("state"),
+ )
+ archived = sa.table(
+ "task_instance_history", *(sa.column(c) for c in _COORDINATES), sa.column("try_number")
+ )
+ same_coordinates = [live.c[c] == archived.c[c] for c in _COORDINATES]
+ # Clearing archived the try and left the replacement at it; only a row that is still executing that try is the same attempt.
+ op.execute(
+ live.update()
+ .where(
+ live.c.state.in_(_NOT_EXECUTING_STATES),
+ sa.exists().where(*same_coordinates, archived.c.try_number == live.c.try_number),
+ )
+ .values(
+ try_number=sa.select(sa.func.max(archived.c.try_number))
+ .where(*same_coordinates)
+ .scalar_subquery()
+ + 1
+ )
+ )
with op.batch_alter_table("task_instance_history") as batch:
batch.drop_constraint("task_instance_history_ti_fkey", type_="foreignkey")
with op.batch_alter_table("task_instance") as batch:
diff --git a/airflow-core/src/airflow/models/taskinstance.py b/airflow-core/src/airflow/models/taskinstance.py
index b3cc60ef1d3..e6a0055282f 100644
--- a/airflow-core/src/airflow/models/taskinstance.py
+++ b/airflow-core/src/airflow/models/taskinstance.py
@@ -1199,6 +1199,12 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
session.execute(insert(TaskInstance.__table__).values(values))
successor = session.scalars(select(TaskInstance).where(TaskInstance.id == values["id"])).one()
successor.task = self.task
+ if (
+ note := session.scalar(select(TaskInstanceNote).where(TaskInstanceNote.ti_id == self.id))
+ ) is not None:
+ carried = TaskInstanceNote(content=note.content, user_id=note.user_id)
+ carried.created_at, carried.updated_at = note.created_at, note.updated_at
+ successor.task_instance_note = carried
return successor
def complete_restart(self, *, session: Session) -> TaskInstance:
diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
index 3648fc28928..f1c4317b435 100644
--- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
+++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
@@ -4331,7 +4331,7 @@ class TestPostClearTaskInstances(TestTaskInstanceEndpoint):
ti_id = response_data["task_instances"][0]["id"]
assert ti_id != str(old_ti.id)
_check_task_instance_note(session, old_ti.id, {"content": "placeholder-note", "user_id": None})
- _check_task_instance_note(session, ti_id, None)
+ _check_task_instance_note(session, ti_id, {"content": "placeholder-note", "user_id": None})
@pytest.mark.db_test
def test_clear_dry_run_does_not_set_note(self, test_client, session):
diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/test_security.py b/airflow-core/tests/unit/api_fastapi/execution_api/test_security.py
index e80ccd18803..304d1e26f95 100644
--- a/airflow-core/tests/unit/api_fastapi/execution_api/test_security.py
+++ b/airflow-core/tests/unit/api_fastapi/execution_api/test_security.py
@@ -379,37 +379,20 @@ class TestAttemptLiveness:
("delete", "/execution/variables/key", None),
("post", "/execution/xcoms/{dag_id}/{run_id}/{task_id}/key", "late"),
]
- # Clients older than the Cadwyn change that identifies retired attempts get 404, not 410.
- OLD_CLIENT_ROUTES = [
- ("put", "/execution/task-instances/{ti_id}/heartbeat", {"hostname": "w", "pid": 1}),
- ("put", "/execution/task-instances/{ti_id}/rtif", {"field": "late"}),
- ("post", "/execution/xcoms/{dag_id}/{run_id}/{task_id}/key", "late"),
- ]
@pytest.mark.parametrize(
- ("method", "path", "body", "version", "expected_status"),
- [
- *(pytest.param(*route, None, 410, id=f"{route[0]}:{route[1]}") for route in RETIRED_ROUTES),
- *(
- pytest.param(*route, version, 404, id=f"{route[0]}:{route[1]}:{version}")
- for route in OLD_CLIENT_ROUTES
- for version in ("2025-04-11", "2026-06-30")
- ),
- ],
+ ("method", "path", "body"),
+ [pytest.param(*route, id=f"{route[0]}:{route[1]}") for route in RETIRED_ROUTES],
)
- def test_retired_attempt_rejected_before_mutation(
- self, client, caller, session, method, path, body, version, expected_status
- ):
+ def test_retired_attempt_rejected_before_mutation(self, client, caller, session, method, path, body):
ti, token = caller
path = path.format(ti_id=token.id, dag_id=ti.dag_id, run_id=ti.run_id, task_id=ti.task_id)
ti.prepare_db_for_next_try(session)
session.commit()
- response = client.request(
- method, path, json=body, headers={"Airflow-API-Version": version} if version else {}
- )
+ response = client.request(method, path, json=body)
- assert response.status_code == expected_status, response.text
+ assert response.status_code == 410, response.text
def test_callback_execution_token_can_mutate_variable(self, client, caller, session):
_, token = caller
diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py
index 7668e2528b6..eecffc755a4 100644
--- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py
+++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py
@@ -21,9 +21,12 @@ import pytest
from sqlalchemy import select
from airflow._shared.timezones import timezone
+from airflow.api_fastapi.execution_api.app import _jwt_generator
+from airflow.api_fastapi.execution_api.security import require_auth
from airflow.models.taskinstance import TaskInstance
from airflow.models.taskreschedule import TaskReschedule
from airflow.models.trigger import Trigger
+from airflow.models.xcom import XComModelV2
from airflow.sdk import task
from airflow.utils.state import State
@@ -268,3 +271,49 @@ class TestArgBindingsFieldBackwardCompat:
"from_default": True,
},
]
+
+
+@pytest.fixture
+def retired_attempt(client, exec_app, monkeypatch, create_task_instance, session):
+ """A running attempt, authenticated with a signed token, that has since been retried."""
+ ti = create_task_instance(state=State.RUNNING)
+ session.commit()
+ monkeypatch.delitem(exec_app.dependency_overrides, require_auth)
+ client.headers["Authorization"] = f"Bearer {_jwt_generator().generate({'sub': str(ti.id)})}"
+ successor = ti.prepare_db_for_next_try(session)
+ successor.state = State.UP_FOR_RETRY
+ session.commit()
+ return ti, successor
+
+
+@pytest.mark.parametrize(
+ ("version", "rtif_status", "heartbeat_status", "xcom_status", "xcom_kept"),
+ [
+ pytest.param("2025-04-11", 410, 410, 201, True, id="oldest"),
+ pytest.param("2026-06-30", 410, 410, 201, True, id="3.3"),
+ pytest.param("2026-10-30", 410, 410, 410, False, id="current"),
+ ],
+)
+def test_mutations_after_retry_response_by_version(
+ client, retired_attempt, session, version, rtif_status, heartbeat_status, xcom_status, xcom_kept
+):
+ ti, successor = retired_attempt
+ client.headers["Airflow-API-Version"] = version
+
+ rtif = client.put(f"/execution/task-instances/{ti.id}/rtif", json={"field": "late"})
+ heartbeat = client.put(
+ f"/execution/task-instances/{ti.id}/heartbeat", json={"hostname": "host", "pid": 1}
+ )
+ xcom = client.post(
+ f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/extra_link", json="https://example.com"
+ )
+
+ assert (rtif.status_code, heartbeat.status_code, xcom.status_code) == (
+ rtif_status,
+ heartbeat_status,
+ xcom_status,
+ )
+ session.expire_all()
+ stored = XComModelV2.get_for_attempt(ti.id, "extra_link", session=session)
+ assert (stored is not None) is xcom_kept
+ assert XComModelV2.get_for_attempt(successor.id, "extra_link", session=session) is None
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 d346d3c1d89..019376fb79a 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
@@ -19,7 +19,7 @@
from __future__ import annotations
from shutil import copyfile
-from uuid import uuid4
+from uuid import UUID, uuid4
import pytest
import sqlalchemy as sa
@@ -349,9 +349,8 @@ def test_upgrade_retains_history_and_legacy_owner(populated_predecessor):
assert foreign_key["referred_table"] == "legacy_task_data_owner"
-def test_upgrade_discards_only_history_conflicting_with_live_try(populated_predecessor):
- connection, config = populated_predecessor
- conflicting_id = uuid4()
+def archive_attempt_at_live_try(connection) -> UUID:
+ archived_id = uuid4()
connection.execute(
table(connection, "task_instance_history", "task_instance_id")
.insert()
@@ -360,7 +359,7 @@ def test_upgrade_discards_only_history_conflicting_with_live_try(populated_prede
task_id="task",
run_id="manual",
map_index=-1,
- task_instance_id=conflicting_id,
+ task_instance_id=archived_id,
try_number=2,
pool="default_pool",
pool_slots=1,
@@ -372,7 +371,7 @@ def test_upgrade_discards_only_history_conflicting_with_live_try(populated_prede
table(connection, "hitl_detail_history", "ti_history_id")
.insert()
.values(
- ti_history_id=conflicting_id,
+ ti_history_id=archived_id,
options=["yes"],
subject="conflicting review",
params={},
@@ -381,6 +380,12 @@ def test_upgrade_discards_only_history_conflicting_with_live_try(populated_prede
)
)
connection.commit()
+ return archived_id
+
+
+def test_upgrade_discards_only_history_conflicting_with_live_try(populated_predecessor):
+ connection, config = populated_predecessor
+ conflicting_id = archive_attempt_at_live_try(connection)
command.upgrade(config, REVISION)
@@ -401,6 +406,39 @@ def test_upgrade_discards_only_history_conflicting_with_live_try(populated_prede
assert conflicting_id not in rows
+@pytest.mark.parametrize(
+ "live_state", ["skipped", "upstream_failed", "removed", "success", "failed", "restarting"]
+)
+def test_upgrade_keeps_history_archived_at_the_try_of_a_live_row_that_never_ran_it(
+ populated_predecessor, live_state
+):
+ connection, config = populated_predecessor
+ archived_id = archive_attempt_at_live_try(connection)
+ live = table(connection, "task_instance", "id")
+ connection.execute(live.update().where(live.c.id == CURRENT_ID).values(state=live_state))
+ connection.commit()
+
+ command.upgrade(config, REVISION)
+
+ ti = table(connection, "task_instance", "id")
+ rows = {row.id: row for row in connection.execute(sa.select(ti))}
+ assert set(rows) == {CURRENT_ID, HISTORY_ID, archived_id}
+ assert (rows[CURRENT_ID].state, rows[CURRENT_ID].try_number, rows[CURRENT_ID].working_set) == (
+ live_state,
+ 3,
+ True,
+ )
+ assert (rows[archived_id].state, rows[archived_id].try_number, rows[archived_id].working_set) == (
+ "success",
+ 2,
+ None,
+ )
+ assert set(connection.scalars(sa.select(table(connection, "hitl_detail", "ti_id").c.ti_id))) == {
+ HISTORY_ID,
+ archived_id,
+ }
+
+
def test_upgrade_does_not_suppress_unrelated_history_identity_conflict(populated_predecessor):
connection, config = populated_predecessor
history = table(connection, "task_instance_history", "task_instance_id")
diff --git a/airflow-core/tests/unit/models/test_taskinstance.py b/airflow-core/tests/unit/models/test_taskinstance.py
index 5bb923e6048..31fbaa6aea2 100644
--- a/airflow-core/tests/unit/models/test_taskinstance.py
+++ b/airflow-core/tests/unit/models/test_taskinstance.py
@@ -2944,6 +2944,31 @@ class TestTaskInstance:
assert session.scalar(sa.select(sa.func.count()).select_from(TaskInstance)) == 3
assert successor.working_set is True
+ def test_retirement_carries_the_note_to_the_successor(self, ownership_session):
+ session = ownership_session
+ attempt = session.get(TaskInstance, CURRENT_ID)
+ attempt.note = "needs a look"
+ session.flush()
+
+ successor = attempt.prepare_db_for_next_try(session)
+ session.flush()
+ session.expire_all()
+
+ assert session.get(TaskInstance, successor.id).note == "needs a look"
+ assert session.get(TaskInstance, CURRENT_ID).note == "needs a look"
+
+ def test_completing_a_restart_carries_the_note_to_the_successor(self, ownership_session):
+ session = ownership_session
+ attempt = session.get(TaskInstance, CURRENT_ID)
+ attempt.state = TaskInstanceState.RESTARTING
+ attempt.note = "cleared while running"
+ session.flush()
+
+ successor = attempt.complete_restart(session=session)
+ session.expire_all()
+
+ assert session.get(TaskInstance, successor.id).note == "cleared while running"
+
@pytest.mark.parametrize(
("map_index", "deleted"),
[
@@ -3117,10 +3142,11 @@ class TestTaskInstance:
)
is not None
)
- for name in ("xcom_v1", "rtif_v1", "legacy_task_data_owner", "task_instance_note", "task_reschedule"):
+ for name in ("xcom_v1", "rtif_v1", "legacy_task_data_owner", "task_reschedule"):
assert session.scalar(sa.text(f"SELECT count(*) FROM {name}")) == int(
deleted_attempt == "current"
)
+ assert session.scalars(select(TaskInstanceNote.ti_id)).all() == [retained.id]
@pytest.mark.execution_timeout(10)
def test_coordinate_xcom_reads_choose_current_producer_and_exact_try_reads_history(