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 <[email protected]> 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, }, ] + + [email protected] +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 + + [email protected]( + ("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 [email protected]( + "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(
