amoghrajesh opened a new pull request, #66073:
URL: https://github.com/apache/airflow/pull/66073
<!-- 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:
http://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.
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".
-->
- [x] Yes - claude sonnet 4.6
<!--
Generated-by: [Tool Name] following [the
guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
-->
Only the last commit is relevant, this has been built on top of
https://github.com/apache/airflow/pull/65759
closes: https://github.com/apache/airflow/issues/66069
This PR is part of AIP-103 (Task State Management) and adds the execution
API endpoints for task and asset state.
## Summary
- New endpoints under `/execution/state/`:
- `GET/PUT/DELETE /state/ti/{task_instance_id}/{key}` and
`DELETE /state/ti/{task_instance_id}` to clear all keys for a task
- `GET/PUT/DELETE /state/asset/{asset_id}/{key}` and
`DELETE /state/asset/{asset_id}` to clear all keys for an asset
- Routes call the configured `[state_store] backend` (defaults to
`MetastoreStateBackend` from PR 1). Custom worker-side backends bypass the API
entirely (although I think I should hardcode `MetastoreStateBackend` in API
server since we do not want state backend credentials there?)
- Task state is JWT-scoped via `Security(require_auth, scopes=["ti:self"])`
on the router — a task can only access its own state.
- Single backend instance shared across both routers via a cached
`get_state_backend()` in `airflow.state`.
## A few design choices worth flagging
Why routes call `BaseStateBackend` instead of querying the tables directly.
- Why we don't pass the route's `session` into the backend?
- First attempt threaded `session=session` from `SessionDep` into
`backend.set(...)`. mypy rightly complained — `BaseStateBackend.set` doesn't
declare a `session` kwarg, only `MetastoreStateBackend.set` does (via
`@provide_session`). I went back and forth on adding `*, session: Any = None`
to the abstract methods, but it leaks a
SQLAlchemy concept into an interface that S3 / Redis / GCS backends have no
use for.
So:
- Backend manages its own session via `@provide_session` →
`create_session()`
- Routes keep `SessionDep` for the non-backend lookups (TaskInstance /
AssetModel existence) — those are reads, no commit needed.
Cost is one extra session per state op, which should be negligible. Worth it
to keep the abstraction clean.
- **No Cadwyn migration.** These are net-new endpoints. Old clients that
don't call them are unaffected.
- **Mapped task semantics:** PUT, GET, and single-key DELETE stay scoped to
one mapped instance — the `task_instance_id` resolves to a unique `(dag_id,
run_id, task_id, map_index)` so reads and writes never cross siblings.
DELETE-all is the exception: it wipes state across every `map_index` of the
task. Clear-all is a "task is done with its state" operation, not a
per-instance reset; per-instance reset is already covered by single-key DELETE.
## What's deferred / todo so far
- **Per-task asset registration check.** AIP-103 defined that the JWT
security model "should be updated to only allow tasks to modify those Assets
they are already registered with". Same gap exists in `/assets` and
`/asset-events` today — the proper fix is a unified check across all asset
routes, which I will probably try before merge here. TODO in `asset_state.py`.
## Testing Manually
### Setup
1. Breeze running: breeze start-airflow --backend postgres
--load-example-dags
2. A DAG triggered so at least one TaskInstance exists
3. Use this script and run in breeze container and use that token:
```python
def main() -> None:
from airflow.configuration import conf
secret = conf.get("api_auth", "jwt_secret", fallback=None)
pkey_path = conf.get("api_auth", "jwt_private_key_path", fallback=None)
print(
f"[diag] jwt_secret first 8 chars: {(secret or '')[:8]!r} (empty
means using private key file or auto-gen)"
)
print(f"[diag] jwt_private_key_path: {pkey_path!r}")
print(f"[diag] jwt_audience: {conf.get('execution_api', 'jwt_audience',
fallback=None)!r}")
with create_session() as session:
ti = session.scalar(select(TaskInstance).limit(1))
if ti is None:
print(
"No task instances found. Trigger a DAG first (e.g. `airflow
dags trigger example_bash_operator`)."
)
return
gen = _jwt_generator()
token = gen.generate(extras={"sub": str(ti.id), "scope":
"execution"})
print(f"ti_id: {ti.id}")
print(f"dag_id: {ti.dag_id} run_id: {ti.run_id} task_id:
{ti.task_id} map_index: {ti.map_index}")
print(f"token: {token}")
if __name__ == "__main__":
main()
```
Example:
```shell script
[Breeze:3.10.20] root@f43294940632:/opt/airflow$ python dev/gen_ti_token.py
2026-04-29T08:26:29.764581Z [info ] setup plugin
alembic.autogenerate.schemas [alembic.runtime.plugins] loc=plugins.py:37
2026-04-29T08:26:29.765573Z [info ] setup plugin
alembic.autogenerate.tables [alembic.runtime.plugins] loc=plugins.py:37
2026-04-29T08:26:29.766389Z [info ] setup plugin
alembic.autogenerate.types [alembic.runtime.plugins] loc=plugins.py:37
2026-04-29T08:26:29.766492Z [info ] setup plugin
alembic.autogenerate.constraints [alembic.runtime.plugins] loc=plugins.py:37
2026-04-29T08:26:29.766558Z [info ] setup plugin
alembic.autogenerate.defaults [alembic.runtime.plugins] loc=plugins.py:37
2026-04-29T08:26:29.766622Z [info ] setup plugin
alembic.autogenerate.comments [alembic.runtime.plugins] loc=plugins.py:37
[diag] jwt_secret first 8 chars: 'tIyIg3gY' (empty means using private key
file or auto-gen)
[diag] jwt_private_key_path: None
[diag] jwt_audience: 'urn:airflow.apache.org:task'
2026-04-29T08:26:30.272426Z [warning ] The HMAC key is 24 bytes long, which
is below the minimum recommended length of 64 bytes for SHA512. See RFC 7518
Section 3.2. [py.warnings] category=InsecureKeyLengthWarning
filename=/usr/python/lib/python3.10/site-packages/jwt/api_jwt.py lineno=147
ti_id: 019dd858-888f-7160-b16f-033390d4386a
dag_id: my_dag run_id: manual__2026-04-29T08:26:14.011123+00:00 task_id:
t1 map_index: -1
token:
eyJhbGciOiJIUzUxMiIsImtpZCI6Im5vdC11c2VkIiwidHlwIjoiSldUIn0.eyJzdWIiOiIwMTlkZDg1OC04ODhmLTcxNjAtYjE2Zi0wMzMzOTBkNDM4NmEiLCJzY29wZSI6ImV4ZWN1dGlvbiIsImp0aSI6Ijc2NmM2MDgyZDEzZTQ1MzU4Y2U5YjNkMmZmMjZhOWRiIiwiaXNzIjoiYWlyZmxvdyIsImF1ZCI6InVybjphaXJmbG93LmFwYWNoZS5vcmc6dGFzayIsIm5iZiI6MTc3NzQ1MTE5MCwiZXhwIjoxNzc3NDUxNzkwLCJpYXQiOjE3Nzc0NTExOTB9.Y-nAYstD5zcBiX7JcXqCml1-sIFlmjqjkAYSNbOOQNaRmdMaYPmxcTIAX1fKr0a_fgRyla3sGOMhjdACBD_sEQ
```
4. I created a postman collections with AIP 103 related API endpoints, but I
will showcase running them below.
### Testing
#### Regular Tasks
1. Put on task_state
<img width="985" height="691" alt="image"
src="https://github.com/user-attachments/assets/b04213a2-73c4-4d40-98b7-a74836dc5181";
/>
Validation:
<img width="1728" height="728" alt="image"
src="https://github.com/user-attachments/assets/061c0faf-288b-4de5-93e7-fe4b28835365";
/>
2. Get task state
<img width="1728" height="728" alt="image"
src="https://github.com/user-attachments/assets/23d66346-ec09-49b1-af0c-c3fb4b052122";
/>
3. Delete task state
<img width="1728" height="728" alt="image"
src="https://github.com/user-attachments/assets/427f4a9f-d1ad-4fb3-93ff-bd6034600622";
/>
Validation:
<img width="1728" height="728" alt="image"
src="https://github.com/user-attachments/assets/c632b2f3-202d-4716-b8b3-eff7ca4451a2";
/>
4. Write and then task state for same task
<img width="1728" height="728" alt="image"
src="https://github.com/user-attachments/assets/01dc76f4-53bf-493f-9ed8-e0e2ce81c4de";
/>
<img width="1728" height="728" alt="image"
src="https://github.com/user-attachments/assets/9d70cc0b-1d72-455c-9332-b808d1a4a142";
/>
<img width="1728" height="728" alt="image"
src="https://github.com/user-attachments/assets/89a5ccc4-d8df-41e2-b21b-799d35d22314";
/>
5. Demonstrating clear_all
- First adding one more key
<img width="1728" height="728" alt="image"
src="https://github.com/user-attachments/assets/6924fc67-7e26-499d-9628-e2beca345734";
/>
<img width="1728" height="728" alt="image"
src="https://github.com/user-attachments/assets/9b52e2c3-ba33-4deb-9980-bd55ede56a06";
/>
- Sending API call
<img width="1728" height="728" alt="image"
src="https://github.com/user-attachments/assets/adc41320-0ab0-47dd-8c81-e19bab6cdc8a";
/>
All gone:
<img width="1728" height="728" alt="image"
src="https://github.com/user-attachments/assets/ff0b2b70-7363-495d-8746-abdfdcb91ba9";
/>
#### Mapped Tasks
- Everything related to PUT, GET, DELETE remains same because ti_id is
unique for mapped tasks and it will populate those values right but DELETE-all
is intentionally different. A worker calling `DELETE /state/ti/{ti_id}` wipes
state across every `map_index` of the task, not just its own. The reasoning:
clear-all is a "task is done with its state" operation, not a per-instance
reset. Per-instance reset is already covered by DELETE. Showcasing that below.
Pushed in data for map_index: 0,1,2
<img width="1206" height="728" alt="image"
src="https://github.com/user-attachments/assets/760b06b7-4e40-43a6-9989-53a63c927738";
/>
1. Now trying DELETE ALL with ti_id: `00000000-0000-0000-0000-000000000001`
<img width="994" height="728" alt="image"
src="https://github.com/user-attachments/assets/66e925e7-1a7b-413e-bfdf-97d88bd09f28";
/>
2. Now trying DELETE ALL with ti id as one of the mapped tasks:
`019dd865-5ef5-73ad-8226-8c14ecf3c0ce`
<img width="994" height="728" alt="image"
src="https://github.com/user-attachments/assets/b6582682-a374-435f-8544-089754136e7e";
/>
All values erased:
<img width="994" height="728" alt="image"
src="https://github.com/user-attachments/assets/1bb7ee1c-d479-45ae-ab90-f4209ee08f91";
/>
---
* 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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
