This is an automated email from the ASF dual-hosted git repository.
jason810496 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new a327dd4a733 Return 410 for stale id for set rtif API (#68902)
a327dd4a733 is described below
commit a327dd4a73373f4f8927d22ed7512c43fa2bc23c
Author: PoAn Yang <[email protected]>
AuthorDate: Tue Jul 7 16:04:48 2026 +0900
Return 410 for stale id for set rtif API (#68902)
When a task whose operator sets overwrite_rtif_after_execution=True leaves
RUNNING via a retry or a clear, the server regenerates the task instance id
and archives the old one. finalize() still overwrites RTIF with the stale
id,
so the API server returns 404 and the worker logs a spurious
error traceback on top of the task's real outcome. The RTIF will be wrote
in next id.
Signed-off-by: PoAn Yang <[email protected]>
---
.../execution_api/routes/task_instances.py | 59 ++++++++++++----------
.../versions/head/test_task_instances.py | 32 +++++++++++-
.../src/airflow/sdk/execution_time/supervisor.py | 10 +++-
.../task_sdk/execution_time/test_supervisor.py | 39 ++++++++++++++
4 files changed, 110 insertions(+), 30 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 b813c4c8f7e..c1bac796023 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
@@ -22,7 +22,7 @@ import itertools
import json
from collections import defaultdict
from collections.abc import Iterator
-from typing import TYPE_CHECKING, Annotated, Any, cast
+from typing import TYPE_CHECKING, Annotated, Any, NoReturn, cast
from uuid import UUID
import attrs
@@ -860,6 +860,29 @@ def ti_skip_downstream(
log.info("Downstream tasks skipped", tasks_skipped=getattr(result,
"rowcount", 0))
+def _raise_ti_not_in_live_table(task_instance_id: UUID, session: SessionDep)
-> NoReturn:
+ """Raise 410 Gone if the missing TI id was archived to history, else 404
Not Found."""
+ if session.scalar(
+ select(func.count(TIH.task_instance_id)).where(TIH.task_instance_id ==
task_instance_id)
+ ):
+ log.error("TaskInstance not in live table but archived in history",
ti_id=str(task_instance_id))
+ raise HTTPException(
+ status_code=status.HTTP_410_GONE,
+ detail={
+ "reason": "not_found",
+ "message": "Task Instance not found, it may have been moved to
the Task Instance History table",
+ },
+ )
+ log.error("Task Instance not found", ti_id=str(task_instance_id))
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": "Task Instance not found",
+ },
+ )
+
+
@ti_id_router.put(
"/{task_instance_id}/heartbeat",
status_code=status.HTTP_204_NO_CONTENT,
@@ -920,29 +943,7 @@ def ti_heartbeat(
# Check if the TI exists in the Task Instance History table.
# If it does, it was likely cleared while running, so return 410 Gone
# instead of 404 Not Found to give the client a more specific signal.
- tih_exists = session.scalar(
-
select(func.count(TIH.task_instance_id)).where(TIH.task_instance_id ==
task_instance_id)
- )
- if tih_exists:
- log.error(
- "TaskInstance was previously cleared and archived in history,
heartbeat skipped",
- ti_id=str(task_instance_id),
- )
- raise HTTPException(
- status_code=status.HTTP_410_GONE,
- detail={
- "reason": "not_found",
- "message": "Task Instance not found, it may have been
moved to the Task Instance History table",
- },
- )
- log.error("Task Instance not found")
- raise HTTPException(
- status_code=status.HTTP_404_NOT_FOUND,
- detail={
- "reason": "not_found",
- "message": "Task Instance not found",
- },
- )
+ _raise_ti_not_in_live_table(task_instance_id, session)
if hostname != ti_payload.hostname or pid != ti_payload.pid:
log.warning(
@@ -989,6 +990,10 @@ def ti_heartbeat(
responses=create_openapi_http_exception_doc(
[
(status.HTTP_404_NOT_FOUND, "Task Instance not found"),
+ (
+ status.HTTP_410_GONE,
+ "Task Instance not found in the TI table but exists in the
Task Instance History table",
+ ),
(
HTTP_422_UNPROCESSABLE_CONTENT,
"Invalid payload for the setting rendered task instance
fields",
@@ -1007,10 +1012,8 @@ def ti_put_rtif(
task_instance = session.scalar(select(TI).where(TI.id == task_instance_id))
if not task_instance:
- log.error("Task Instance not found")
- raise HTTPException(
- status_code=status.HTTP_404_NOT_FOUND,
- )
+ # On retry/clear, the server regenerates the TI id. Return 410 for the
stale id.
+ _raise_ti_not_in_live_table(task_instance_id, session)
task_instance.update_rtif(put_rtif_payload, session=session)
log.debug("RenderedTaskInstanceFields updated successfully")
diff --git
a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
index da6ea74337d..542ce7eaaf1 100644
---
a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
+++
b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
@@ -2921,7 +2921,37 @@ class TestTIPutRTIF:
random_id = uuid6.uuid7()
response = client.put(f"/execution/task-instances/{random_id}/rtif",
json=payload)
assert response.status_code == 404
- assert response.json()["detail"] == "Not Found"
+ assert response.json()["detail"] == {
+ "reason": "not_found",
+ "message": "Task Instance not found",
+ }
+
+ def test_ti_put_rtif_archived_ti_returns_410(self, client, session,
create_task_instance):
+ ti = create_task_instance(
+ task_id="test_ti_put_rtif_archived",
+ state=State.RUNNING,
+ session=session,
+ )
+ session.commit()
+ old_ti_id = ti.id
+
+ # Archive the current try to TIH and assign a new UUID, mirroring
prepare_db_for_next_try().
+ ti.prepare_db_for_next_try(session)
+ session.commit()
+
+ assert session.get(TaskInstance, old_ti_id) is None
+ assert session.get(TaskInstanceHistory, old_ti_id) is not None
+
+ response = client.put(
+ f"/execution/task-instances/{old_ti_id}/rtif",
+ json={"field1": "rendered_value1"},
+ )
+
+ assert response.status_code == 410
+ assert response.json()["detail"] == {
+ "reason": "not_found",
+ "message": "Task Instance not found, it may have been moved to the
Task Instance History table",
+ }
class TestPreviousDagRun:
diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
index 5a5e3852fd7..a96523d5a6a 100644
--- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
+++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
@@ -1737,7 +1737,15 @@ class ActivitySubprocess(WatchedSubprocess):
elif isinstance(msg, PutVariable):
resp, dump_opts = handle_put_variable(self.client, msg)
elif isinstance(msg, SetRenderedFields):
- self.client.task_instances.set_rtif(self.id, msg.rendered_fields)
+ try:
+ self.client.task_instances.set_rtif(self.id,
msg.rendered_fields)
+ except ServerResponseError as e:
+ # On retry/clear the server replaces the TI id (archiving the
old one), so a late RTIF
+ # overwrite from finalize() lands on an id that no longer
exists. Supervisor kills such
+ # a worker when handling 410 heartbeat response. We only need
to skip this stale overwrite here.
+ if e.response.status_code != HTTPStatus.GONE:
+ raise
+ log.debug("Skipping RTIF overwrite; task instance archived on
retry/clear", ti_id=self.id)
elif isinstance(msg, SetRenderedMapIndex):
self.client.task_instances.set_rendered_map_index(self.id,
msg.rendered_map_index)
elif isinstance(msg, GetAssetByName):
diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
index 491a1b0e3fd..9bf0ef9d13b 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
@@ -3092,6 +3092,45 @@ class TestHandleRequest:
"detail": error.response.json(),
}
+ @pytest.mark.parametrize(
+ ("status_code", "expects_error"),
+ [
+ pytest.param(410, False, id="410_gone_is_swallowed"),
+ pytest.param(404, True, id="404_not_found_propagates"),
+ ],
+ )
+ def test_set_rendered_fields_swallows_410_but_propagates_404(
+ self, watched_subprocess, mocker, status_code, expects_error
+ ):
+ """A stale-id RTIF overwrite (410) is skipped silently; a bogus-id
(404) still propagates as an error."""
+ watched_subprocess, read_socket = watched_subprocess
+
+ error = ServerResponseError(
+ message="boom",
+ request=httpx.Request("PUT", "http://test";),
+ response=httpx.Response(status_code, json={"detail": "boom"}),
+ )
+ watched_subprocess.client.task_instances.set_rtif =
mocker.Mock(side_effect=error)
+
+ generator = watched_subprocess.handle_requests(log=mocker.Mock())
+ next(generator)
+
+ msg = SetRenderedFields(rendered_fields={"field1": "v1"})
+ req_frame = _RequestFrame(id=randint(1, 2**32 - 1),
body=msg.model_dump())
+ generator.send(req_frame)
+
+ read_socket.settimeout(0.1)
+ frame_len = int.from_bytes(read_socket.recv(4), "big")
+ frame =
msgspec.msgpack.Decoder(_ResponseFrame).decode(read_socket.recv(frame_len))
+
+ assert frame.id == req_frame.id
+ if expects_error:
+ assert frame.error is not None
+ assert frame.error["error"] == "API_SERVER_ERROR"
+ assert frame.error["detail"]["status_code"] == status_code
+ else:
+ assert frame.error is None
+
def test_handle_requests_network_exception_does_not_crash_loop(self,
watched_subprocess, mocker):
"""A transient network error must not crash the IPC generator.
