ColtenOuO opened a new pull request, #70983:
URL: https://github.com/apache/airflow/pull/70983
## Sumarry
The task state store is a per-task-instance key/value store that survives
worker
crashes and retries within a Dag run — the place operators keep external job
IDs,
checkpoints and progress metadata. Its REST API is addressed by four values
that
together identify one task instance: `dag_id` + `dag_run_id` + `task_id` +
`map_index`.
Only two of the six endpoints ever checked that the task instance those four
values
name actually exists.
`_require_ti()` was written as a plain helper that each handler had to
remember to
call, and it was called from `PUT` and `PATCH` only. The read and delete
handlers
went straight to the query, and because those queries filter on the four
columns
directly — they never join `task_instance` — a request naming a task
instance that
doesn't exist simply matches zero rows. At the SQL level that is
indistinguishable
from "this task instance exists but has no state stored", so the API
answered with
the latter.
All six endpoints declare `404` in `create_openapi_http_exception_doc(...)`,
so for
three of them the spec described a status the handler could never return.
## What that looked like
A typo in `task_id` was silently successful:
```console
$ curl .../taskInstances/proces/state-store # note the missing 's'
200 {"task_state_store": [], "total_entries": 0}
```
The caller cannot tell "this task has no state" from "this task does not
exist".
Worse, a delete reported success while touching nothing:
```console
$ curl -X DELETE .../taskInstances/proces/state-store
204 No Content
```
That is the damaging one. Clearing state store is typically a *pre-rerun
cleanup*
step — clearing a task instance in the UI deliberately does **not** wipe its
state
store, so a stale checkpoint survives into the next attempt. An operator
wiping a
bad external job ID before re-running gets `204`, believes the state is gone,
re-runs, and the task resumes from exactly the state they thought they had
removed.
Nothing in the response hints that the `task_id` was wrong.
The same resource also answered the same question three different ways
depending on
the verb:
```console
$ curl -X PUT .../taskInstances/proces/state-store/job_id -d
'{"value":"x"}'
404 {"detail": "Task instance not found for dag_id='etl', ...,
task_id='proces', map_index=-1"}
$ curl .../taskInstances/proces/state-store/job_id
404 {"detail": "Task state store key 'job_id' not found"} # right
status, wrong reason
$ curl -X DELETE .../taskInstances/proces/state-store/job_id
204 No Content
```
The `GET` case is the most misleading: it happens to return `404`, but
blames the
key, sending the reader off to check their key name when the real problem is
the
task ID.
## Before / after
| Endpoint | Before | After |
| --- | --- | --- |
| `GET .../state-store` | `200` + empty collection | `404` |
| `GET .../state-store/{key}` | `404`, message blames the key | `404`,
message names the task instance |
| `PUT .../state-store/{key}` | `404` | unchanged |
| `PATCH .../state-store/{key}` | `404` | unchanged |
| `DELETE .../state-store/{key}` | `204` | `404` |
| `DELETE .../state-store` | `204` | `404` |
A missing *key* on an existing task instance still reports the key,
unchanged:
`Task state store key 'job_id' not found`.
## The fix
`asset_state_store.py` — the sibling module from the same feature — already
solves
this correctly: `_get_asset_or_404` is wired up as `AssetIdDep =
Annotated[int,
Depends(...)]` and **every** one of its endpoints takes `asset_id:
AssetIdDep`, so
the check rides on the parameter and cannot be forgotten when an endpoint is
added.
This applies the same pattern on the task side. `_require_ti` becomes a
dependency
returning the `TaskScope` the handlers already needed, which also absorbs
the old
`_get_scope()` helper — each handler drops from four path/query parameters
plus two
helper calls to a single `scope`. The diff removes more lines from the
handlers than
it adds.
### Why `clear` needs its own dependency
`DELETE .../state-store?all_map_indices=true` means "clear this task at
every map
index", and documents that `map_index` is ignored. `TaskScopeDep` validates
one
*exact* map index, defaulting to `-1` — and an expanded mapped task has no
`map_index=-1` instance: `TaskMap.expand_mapped_task` either converts that
row into
`map_index=0` or deletes it. Validating with `TaskScopeDep` would therefore
reject a
perfectly valid request:
```console
# 'process' expanded to map_index 0, 1, 2 — there is no map_index=-1 row
$ curl -X DELETE .../taskInstances/process/state-store?all_map_indices=true
404 Task instance not found ... map_index=-1 # wrong
```
A dependency only receives the parameters it declares, so `TaskScopeDep`
cannot see
`all_map_indices` and can only ever take the strict path. `ClearTaskScopeDep`
declares the flag, and passes `map_index=None` when it is set, which drops
the
`map_index` predicate and matches any instance of that task in the run. The
strict
path is untouched: `?map_index=7` on a task expanded to 0..2 still returns
`404`.
## Generated files
`v2-rest-api-generated.yaml` and `ui/openapi-gen/` are regenerated.
## Tests
New `TestUnknownTaskInstance` covers all seven request shapes against both a
bad
`task_id` and a bad `run_id`, a non-existent `map_index`, and a regression
test that
a genuinely missing key still reports the key rather than the task instance.
`test_all_map_indices_clears_across_mapped_instances` inserted state rows at
map
indices that have no task instance and then read them back through `GET`. It
now
asserts on the rows directly — the test is about `clear`, not about `get`,
and the
read endpoints legitimately reject a map index with no task instance.
76 → 97 tests in this file; `test_asset_state_store.py` passes unchanged.
---
##### Was generative AI tooling used to co-author this PR?
- [X] Yes — Claude Code (Opus 5)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
