You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "amoghrajesh (via GitHub)" <gi...@apache.org> on 2026/04/29 09:17:04 UTC
[PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66073:
URL: https://github.com/apache/airflow/pull/66073#discussion_r3173220458
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_state.py:
##########
@@ -0,0 +1,121 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Execution API routes for asset state.
+
+Asset state is keyed by asset *name* (not integer id) in the URL — asset names
+are unique, and callers (task SDK accessors) have the name from their Asset
+object without needing a DB lookup. The route resolves name → asset_id
+internally for the state backend scope.
+
+Per-task asset registration checks are intentionally not implemented here
+(deferred to AIP-93 — see TODO comment below).
+"""
+
+from __future__ import annotations
+
+from typing import Annotated
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, status
+from sqlalchemy import select
+
+from airflow._shared.state import AssetScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.asset_state import (
+ AssetStatePutBody,
+ AssetStateResponse,
+)
+from airflow.models.asset import AssetModel
+from airflow.state import get_state_backend
+
+# TODO(AIP-103): enforce that the requesting task is registered with the asset
+# (via task_inlet_asset_reference or task_outlet_asset_reference) before
+# allowing reads/writes. Currently any task with a valid execution token can
+# access any asset's state — the same gap exists in /assets and /asset-events.
+# Proper fix is a unified asset-registration check across all asset routes,
+# not just here.
+router = VersionedAPIRouter(
+ responses={
+ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
+ status.HTTP_404_NOT_FOUND: {"description": "Not found"},
+ },
+)
+
+
+def _resolve_asset_id(name: str, session: SessionDep) -> int:
+ """Resolve asset name → integer asset_id, 404 if not found."""
+ asset_id = session.scalar(select(AssetModel.id).where(AssetModel.name == name))
Review Comment:
Good catch, fixed it. Added `AssetModel.active.has()` to `_resolve_asset_id`, matching the pattern we already have. Updated the test fixture to also create an `asset_active` row so the tests reflect the actual constraint.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66073:
URL: https://github.com/apache/airflow/pull/66073#discussion_r3176669037
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state.py:
##########
@@ -0,0 +1,136 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Annotated
+from uuid import UUID
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, Query, Security, status
+from sqlalchemy.orm import Session
+
+from airflow._shared.state import TaskScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.task_state import (
+ TaskStatePutBody,
+ TaskStateResponse,
+)
+from airflow.api_fastapi.execution_api.security import require_auth
+from airflow.models.taskinstance import TaskInstance as TI
+from airflow.state import get_state_backend
+
+router = VersionedAPIRouter(
+ responses={
+ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
+ status.HTTP_403_FORBIDDEN: {"description": "Access denied"},
+ status.HTTP_404_NOT_FOUND: {"description": "Not found"},
+ },
+ dependencies=[Security(require_auth, scopes=["ti:self"])],
+)
+
+
+def _get_task_scope_for_ti(task_instance_id: UUID, session: Session) -> TaskScope:
+ ti = session.get(TI, task_instance_id)
+ if ti is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": f"Task instance {task_instance_id} not found",
+ },
+ )
+ return TaskScope(dag_id=ti.dag_id, run_id=ti.run_id, task_id=ti.task_id, map_index=ti.map_index)
+
+
+@router.get("/{task_instance_id}/{key}")
+def get_task_state(
+ task_instance_id: UUID,
+ key: Annotated[str, Path(min_length=1)],
+ session: SessionDep,
+) -> TaskStateResponse:
+ """Get value for a task state."""
+ scope = _get_task_scope_for_ti(task_instance_id, session)
+ value = get_state_backend().get(scope, key, session=session)
+ if value is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": f"Task state key {key!r} not found",
+ },
+ )
+ return TaskStateResponse(value=value)
+
+
+@router.put("/{task_instance_id}/{key}", status_code=status.HTTP_204_NO_CONTENT)
+def set_task_state(
+ task_instance_id: UUID,
+ key: Annotated[str, Path(min_length=1)],
+ body: TaskStatePutBody,
+ session: SessionDep,
+) -> None:
+ """Set a task state key, creating or updating the row."""
+ scope = _get_task_scope_for_ti(task_instance_id, session)
+ get_state_backend().set(scope, key, body.value, session=session)
+
+
+@router.delete("/{task_instance_id}/{key}", status_code=status.HTTP_204_NO_CONTENT)
+def delete_task_state(
+ task_instance_id: UUID,
+ key: Annotated[str, Path(min_length=1)],
+ session: SessionDep,
+) -> None:
+ """Delete a single task state key."""
+ scope = _get_task_scope_for_ti(task_instance_id, session)
+ get_state_backend().delete(scope, key, session=session)
+
+
+@router.delete("/{task_instance_id}", status_code=status.HTTP_204_NO_CONTENT)
+def clear_task_state(
+ task_instance_id: UUID,
+ session: SessionDep,
+ all_map_indices: Annotated[bool, Query()] = False,
+) -> None:
+ """
+ Delete all task state keys for this task instance.
+
+ By default, only keys for the requesting TI's exact ``map_index`` are
+ cleared — same isolation as DELETE endpoint above.
+
+ Pass ``?all_map_indices=true`` to wipe state for every mapped sibling of
+ the task in the same DAG run. This is intentionally fleet-wide: the
+ ``ti:self`` JWT authentication scope authenticates that the caller is
+ a legitimate member of the mapped task group, and grants it authority
+ to reset shared task state on behalf of the whole group.
+ The SDK only forwards this flag when the user calls ``task_state.clear(all_map_indices=True)``
+ explicitly, so the expanded scope is always an explicit opt-in by the task author.
+
+ For non-mapped tasks (``map_index=-1``), there is only ever one index, so
+ ``?all_map_indices=true`` is functionally identical to the default and is
+ accepted without error.
+ """
+ ti = session.get(TI, task_instance_id)
+ if ti is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": f"Task instance {task_instance_id} not found",
+ },
+ )
+ scope = TaskScope(dag_id=ti.dag_id, run_id=ti.run_id, task_id=ti.task_id, map_index=ti.map_index)
Review Comment:
Fixed — replaced the duplicated lookup with `_get_task_scope_for_ti`
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66073:
URL: https://github.com/apache/airflow/pull/66073#discussion_r3177951885
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state.py:
##########
@@ -0,0 +1,127 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Annotated
+from uuid import UUID
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, Query, Security, status
+from sqlalchemy.orm import Session
+
+from airflow._shared.state import TaskScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.task_state import (
+ TaskStatePutBody,
+ TaskStateResponse,
+)
+from airflow.api_fastapi.execution_api.security import require_auth
+from airflow.models.taskinstance import TaskInstance as TI
+from airflow.state import get_state_backend
+
+router = VersionedAPIRouter(
+ responses={
+ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
+ status.HTTP_403_FORBIDDEN: {"description": "Access denied"},
+ status.HTTP_404_NOT_FOUND: {"description": "Not found"},
+ },
+ dependencies=[Security(require_auth, scopes=["ti:self"])],
+)
+
+
+def _get_task_scope_for_ti(task_instance_id: UUID, session: Session) -> TaskScope:
+ ti = session.get(TI, task_instance_id)
+ if ti is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": f"Task instance {task_instance_id} not found",
+ },
+ )
+ return TaskScope(dag_id=ti.dag_id, run_id=ti.run_id, task_id=ti.task_id, map_index=ti.map_index)
+
+
+@router.get("/{task_instance_id}/{key}")
+def get_task_state(
+ task_instance_id: UUID,
+ key: Annotated[str, Path(min_length=1)],
+ session: SessionDep,
+) -> TaskStateResponse:
+ """Get value for a task state."""
+ scope = _get_task_scope_for_ti(task_instance_id, session)
+ value = get_state_backend().get(scope, key)
+ if value is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": f"Task state key {key!r} not found",
+ },
+ )
+ return TaskStateResponse(value=value)
+
+
+@router.put("/{task_instance_id}/{key}", status_code=status.HTTP_204_NO_CONTENT)
+def set_task_state(
+ task_instance_id: UUID,
+ key: Annotated[str, Path(min_length=1)],
+ body: TaskStatePutBody,
+ session: SessionDep,
+) -> None:
+ """Set a task state key, creating or updating the row."""
+ scope = _get_task_scope_for_ti(task_instance_id, session)
+ get_state_backend().set(scope, key, body.value)
Review Comment:
`get_state_backend()` returns `BaseStateBackend`, and the abstract interface intentionally does not declare a session param(didnt feel right to do it too) and the concrete `MetastoreStateBackend` gets it injected via `@provide_session` at runtime, which mypy can not see through.
I considered:
- Adding `**kwargs` to the base class — pollutes the interface and custom. backends will have to add it
- Settled on `# type: ignore[call-arg]` with a comment explaining why, which is the least invasive and keeps the design clean. Happy to hear if there's a better pattern for this, but for now proceeding with that in 8d0d6de99f
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66073:
URL: https://github.com/apache/airflow/pull/66073#discussion_r3173348087
##########
airflow-core/src/airflow/state/__init__.py:
##########
@@ -36,3 +36,14 @@ def resolve_state_backend() -> type[BaseStateBackend]:
f"Your custom state backend `{clazz.__name__}` is not a subclass of `BaseStateBackend`."
)
return clazz
+
+
+_backend_instance: BaseStateBackend | None = None
+
+
+def get_state_backend() -> BaseStateBackend:
+ """Return a cached instance of the configured state backend."""
+ global _backend_instance
+ if _backend_instance is None:
+ _backend_instance = resolve_state_backend()()
Review Comment:
Yeah I wanted to fix this later but didn't get to it, might be an issue for custom backends with file handling / network client calls.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66073:
URL: https://github.com/apache/airflow/pull/66073#discussion_r3173220458
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_state.py:
##########
@@ -0,0 +1,121 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Execution API routes for asset state.
+
+Asset state is keyed by asset *name* (not integer id) in the URL — asset names
+are unique, and callers (task SDK accessors) have the name from their Asset
+object without needing a DB lookup. The route resolves name → asset_id
+internally for the state backend scope.
+
+Per-task asset registration checks are intentionally not implemented here
+(deferred to AIP-93 — see TODO comment below).
+"""
+
+from __future__ import annotations
+
+from typing import Annotated
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, status
+from sqlalchemy import select
+
+from airflow._shared.state import AssetScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.asset_state import (
+ AssetStatePutBody,
+ AssetStateResponse,
+)
+from airflow.models.asset import AssetModel
+from airflow.state import get_state_backend
+
+# TODO(AIP-103): enforce that the requesting task is registered with the asset
+# (via task_inlet_asset_reference or task_outlet_asset_reference) before
+# allowing reads/writes. Currently any task with a valid execution token can
+# access any asset's state — the same gap exists in /assets and /asset-events.
+# Proper fix is a unified asset-registration check across all asset routes,
+# not just here.
+router = VersionedAPIRouter(
+ responses={
+ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
+ status.HTTP_404_NOT_FOUND: {"description": "Not found"},
+ },
+)
+
+
+def _resolve_asset_id(name: str, session: SessionDep) -> int:
+ """Resolve asset name → integer asset_id, 404 if not found."""
+ asset_id = session.scalar(select(AssetModel.id).where(AssetModel.name == name))
Review Comment:
Good catch, fixed it. Added `AssetModel.active.has()` to `_resolve_asset_id`, matching the pattern we already have. Updated the test fixture + added some test coverage for inactive assets too.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66073:
URL: https://github.com/apache/airflow/pull/66073#discussion_r3176674564
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state.py:
##########
@@ -0,0 +1,136 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Annotated
+from uuid import UUID
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, Query, Security, status
+from sqlalchemy.orm import Session
+
+from airflow._shared.state import TaskScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.task_state import (
+ TaskStatePutBody,
+ TaskStateResponse,
+)
+from airflow.api_fastapi.execution_api.security import require_auth
+from airflow.models.taskinstance import TaskInstance as TI
+from airflow.state import get_state_backend
+
+router = VersionedAPIRouter(
+ responses={
+ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
+ status.HTTP_403_FORBIDDEN: {"description": "Access denied"},
+ status.HTTP_404_NOT_FOUND: {"description": "Not found"},
+ },
+ dependencies=[Security(require_auth, scopes=["ti:self"])],
+)
+
+
+def _get_task_scope_for_ti(task_instance_id: UUID, session: Session) -> TaskScope:
+ ti = session.get(TI, task_instance_id)
+ if ti is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": f"Task instance {task_instance_id} not found",
+ },
+ )
+ return TaskScope(dag_id=ti.dag_id, run_id=ti.run_id, task_id=ti.task_id, map_index=ti.map_index)
+
+
+@router.get("/{task_instance_id}/{key}")
+def get_task_state(
+ task_instance_id: UUID,
+ key: Annotated[str, Path(min_length=1)],
+ session: SessionDep,
+) -> TaskStateResponse:
+ """Get value for a task state."""
+ scope = _get_task_scope_for_ti(task_instance_id, session)
+ value = get_state_backend().get(scope, key, session=session)
+ if value is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": f"Task state key {key!r} not found",
+ },
+ )
+ return TaskStateResponse(value=value)
+
+
+@router.put("/{task_instance_id}/{key}", status_code=status.HTTP_204_NO_CONTENT)
+def set_task_state(
+ task_instance_id: UUID,
+ key: Annotated[str, Path(min_length=1)],
+ body: TaskStatePutBody,
+ session: SessionDep,
+) -> None:
+ """Set a task state key, creating or updating the row."""
+ scope = _get_task_scope_for_ti(task_instance_id, session)
+ get_state_backend().set(scope, key, body.value, session=session)
+
+
+@router.delete("/{task_instance_id}/{key}", status_code=status.HTTP_204_NO_CONTENT)
+def delete_task_state(
+ task_instance_id: UUID,
+ key: Annotated[str, Path(min_length=1)],
+ session: SessionDep,
+) -> None:
+ """Delete a single task state key."""
+ scope = _get_task_scope_for_ti(task_instance_id, session)
+ get_state_backend().delete(scope, key, session=session)
+
+
+@router.delete("/{task_instance_id}", status_code=status.HTTP_204_NO_CONTENT)
+def clear_task_state(
+ task_instance_id: UUID,
+ session: SessionDep,
+ all_map_indices: Annotated[bool, Query()] = False,
+) -> None:
+ """
+ Delete all task state keys for this task instance.
+
+ By default, only keys for the requesting TI's exact ``map_index`` are
+ cleared — same isolation as DELETE endpoint above.
+
+ Pass ``?all_map_indices=true`` to wipe state for every mapped sibling of
+ the task in the same DAG run. This is intentionally fleet-wide: the
+ ``ti:self`` JWT authentication scope authenticates that the caller is
+ a legitimate member of the mapped task group, and grants it authority
+ to reset shared task state on behalf of the whole group.
+ The SDK only forwards this flag when the user calls ``task_state.clear(all_map_indices=True)``
+ explicitly, so the expanded scope is always an explicit opt-in by the task author.
+
+ For non-mapped tasks (``map_index=-1``), there is only ever one index, so
+ ``?all_map_indices=true`` is functionally identical to the default and is
+ accepted without error.
+ """
+ ti = session.get(TI, task_instance_id)
+ if ti is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": f"Task instance {task_instance_id} not found",
+ },
+ )
+ scope = TaskScope(dag_id=ti.dag_id, run_id=ti.run_id, task_id=ti.task_id, map_index=ti.map_index)
Review Comment:
Handled in 7bd1f4c64a
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state.py:
##########
@@ -0,0 +1,136 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Annotated
+from uuid import UUID
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, Query, Security, status
+from sqlalchemy.orm import Session
+
+from airflow._shared.state import TaskScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.task_state import (
+ TaskStatePutBody,
+ TaskStateResponse,
+)
+from airflow.api_fastapi.execution_api.security import require_auth
+from airflow.models.taskinstance import TaskInstance as TI
+from airflow.state import get_state_backend
+
+router = VersionedAPIRouter(
Review Comment:
Handled in 7bd1f4c64a
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66073:
URL: https://github.com/apache/airflow/pull/66073#discussion_r3173363504
##########
airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_state.py:
##########
@@ -0,0 +1,238 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+from uuid import uuid4
+
+import pytest
+from sqlalchemy import delete, select
+
+from airflow._shared.timezones import timezone
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import create_session
+
+if TYPE_CHECKING:
+ from fastapi.testclient import TestClient
+
+ from tests_common.pytest_plugin import CreateTaskInstance
+
+pytestmark = pytest.mark.db_test
+
+
+@pytest.fixture(autouse=True)
+def reset_state_tables():
+ with create_session() as session:
+ session.execute(delete(TaskStateModel))
+ session.execute(delete(DagRun))
+
+
+def _api_url(ti_id, key: str | None = None) -> str:
+ base = f"/execution/state/ti/{ti_id}"
+ return f"{base}/{key}" if key else base
+
+
+class TestGetTaskState:
+ def test_get_returns_value(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+ client.put(_api_url(ti.id, "job_id"), json={"value": "spark_001"})
+
+ response = client.get(_api_url(ti.id, "job_id"))
+
+ assert response.status_code == 200
+ assert response.json() == {"value": "spark_001"}
+
+ def test_get_missing_key_returns_404(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+
+ response = client.get(_api_url(ti.id, "never_set"))
+
+ assert response.status_code == 404
+ assert response.json()["detail"]["reason"] == "not_found"
+
+ def test_get_missing_ti_returns_404(self, client: TestClient):
+ response = client.get(_api_url(uuid4(), "any_key"))
+
+ assert response.status_code == 404
+ assert "Task instance" in response.json()["detail"]["message"]
+
+
+class TestPutTaskState:
+ def test_put_creates_row(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+
+ response = client.put(_api_url(ti.id, "job_id"), json={"value": "spark_001"})
+
+ assert response.status_code == 204
+ with create_session() as session:
+ row = session.scalar(
+ select(TaskStateModel).where(
+ TaskStateModel.dag_id == ti.dag_id,
+ TaskStateModel.run_id == ti.run_id,
+ TaskStateModel.task_id == ti.task_id,
+ TaskStateModel.key == "job_id",
+ )
+ )
+ assert row is not None
+ assert row.value == "spark_001"
+
+ def test_put_overwrites_existing(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+ client.put(_api_url(ti.id, "job_id"), json={"value": "spark_001"})
+
+ response = client.put(_api_url(ti.id, "job_id"), json={"value": "spark_002"})
+
+ assert response.status_code == 204
+ assert client.get(_api_url(ti.id, "job_id")).json() == {"value": "spark_002"}
+
+ def test_put_empty_body_returns_422(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+
+ response = client.put(_api_url(ti.id, "job_id"), json={})
+
+ assert response.status_code == 422
+
+ def test_put_extra_field_returns_422(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+
+ response = client.put(_api_url(ti.id, "job_id"), json={"value": "x", "extra": "y"})
+
+ assert response.status_code == 422
+
+ def test_put_null_value_returns_422(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+
+ response = client.put(_api_url(ti.id, "job_id"), json={"value": None})
+
+ assert response.status_code == 422
+
+ def test_put_missing_ti_returns_404(self, client: TestClient):
+ response = client.put(_api_url(uuid4(), "job_id"), json={"value": "x"})
+
+ assert response.status_code == 404
+
+
+class TestDeleteTaskState:
+ def test_delete_removes_key(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+ client.put(_api_url(ti.id, "job_id"), json={"value": "spark_001"})
+
+ response = client.delete(_api_url(ti.id, "job_id"))
+
+ assert response.status_code == 204
+ assert client.get(_api_url(ti.id, "job_id")).status_code == 404
+
+ def test_delete_missing_key_is_noop(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+
+ response = client.delete(_api_url(ti.id, "never_existed"))
+
+ assert response.status_code == 204
+
+ def test_delete_only_targets_one_key(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+ client.put(_api_url(ti.id, "job_id"), json={"value": "a"})
+ client.put(_api_url(ti.id, "checkpoint"), json={"value": "b"})
+
+ client.delete(_api_url(ti.id, "job_id"))
+
+ assert client.get(_api_url(ti.id, "job_id")).status_code == 404
+ assert client.get(_api_url(ti.id, "checkpoint")).json() == {"value": "b"}
+
+
+class TestClearTaskState:
+ def test_clear_removes_all_keys(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+ for k, v in [("job_id", "a"), ("checkpoint", "b"), ("retry_count", "c")]:
+ client.put(_api_url(ti.id, k), json={"value": v})
+
+ response = client.delete(_api_url(ti.id))
+
+ assert response.status_code == 204
+ with create_session() as session:
+ count = session.scalar(
+ select(TaskStateModel)
+ .where(TaskStateModel.dag_id == ti.dag_id, TaskStateModel.task_id == ti.task_id)
+ .with_only_columns(TaskStateModel.key)
+ )
+ assert count is None
Review Comment:
Good catch. Fixed to this with a more explicit check now
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on PR #66073:
URL: https://github.com/apache/airflow/pull/66073#issuecomment-4368264044
Merging this PR in right now to unblock further PRs, I have handled all the review comments on the PR, thanks for reviews @kaxil @jscheffl!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on code in PR #66073:
URL: https://github.com/apache/airflow/pull/66073#discussion_r3164140296
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_state.py:
##########
@@ -0,0 +1,95 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Annotated
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, status
+
+from airflow._shared.state import AssetScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.asset_state import (
+ AssetStatePutBody,
+ AssetStateResponse,
+)
+from airflow.models.asset import AssetModel
+from airflow.state import get_state_backend
+
+# TODO(AIP-103): enforce that the requesting task is registered with the asset
+# (via task_inlet_asset_reference or task_outlet_asset_reference) before
+# allowing reads/writes. Currently any task with a valid execution token can
+# access any asset's state — the same gap exists in /assets and /asset-events.
+# Proper fix is a unified asset-registration check across all asset routes,
+# not just here.
+router = VersionedAPIRouter(
+ responses={
+ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
+ status.HTTP_404_NOT_FOUND: {"description": "Not found"},
+ },
+)
+
+
+@router.get("/{asset_id}/{key}")
Review Comment:
Do we not need to add authentication to the routes?
Looked through the execution API and could not find it.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66073:
URL: https://github.com/apache/airflow/pull/66073#discussion_r3173355750
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state.py:
##########
@@ -0,0 +1,127 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Annotated
+from uuid import UUID
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, Query, Security, status
+from sqlalchemy.orm import Session
+
+from airflow._shared.state import TaskScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.task_state import (
+ TaskStatePutBody,
+ TaskStateResponse,
+)
+from airflow.api_fastapi.execution_api.security import require_auth
+from airflow.models.taskinstance import TaskInstance as TI
+from airflow.state import get_state_backend
+
+router = VersionedAPIRouter(
+ responses={
+ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
+ status.HTTP_403_FORBIDDEN: {"description": "Access denied"},
+ status.HTTP_404_NOT_FOUND: {"description": "Not found"},
+ },
+ dependencies=[Security(require_auth, scopes=["ti:self"])],
Review Comment:
The ti:self enforcement is tested centrally in throught another PR: https://github.com/apache/airflow/pull/66071 rather than per-route
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66073:
URL: https://github.com/apache/airflow/pull/66073#discussion_r3173257699
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state.py:
##########
@@ -0,0 +1,127 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Annotated
+from uuid import UUID
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, Query, Security, status
+from sqlalchemy.orm import Session
+
+from airflow._shared.state import TaskScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.task_state import (
+ TaskStatePutBody,
+ TaskStateResponse,
+)
+from airflow.api_fastapi.execution_api.security import require_auth
+from airflow.models.taskinstance import TaskInstance as TI
+from airflow.state import get_state_backend
+
+router = VersionedAPIRouter(
+ responses={
+ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
+ status.HTTP_403_FORBIDDEN: {"description": "Access denied"},
+ status.HTTP_404_NOT_FOUND: {"description": "Not found"},
+ },
+ dependencies=[Security(require_auth, scopes=["ti:self"])],
+)
+
+
+def _get_task_scope_for_ti(task_instance_id: UUID, session: Session) -> TaskScope:
+ ti = session.get(TI, task_instance_id)
+ if ti is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": f"Task instance {task_instance_id} not found",
+ },
+ )
+ return TaskScope(dag_id=ti.dag_id, run_id=ti.run_id, task_id=ti.task_id, map_index=ti.map_index)
+
+
+@router.get("/{task_instance_id}/{key}")
+def get_task_state(
+ task_instance_id: UUID,
+ key: Annotated[str, Path(min_length=1)],
+ session: SessionDep,
+) -> TaskStateResponse:
+ """Get value for a task state."""
+ scope = _get_task_scope_for_ti(task_instance_id, session)
+ value = get_state_backend().get(scope, key)
+ if value is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": f"Task state key {key!r} not found",
+ },
+ )
+ return TaskStateResponse(value=value)
+
+
+@router.put("/{task_instance_id}/{key}", status_code=status.HTTP_204_NO_CONTENT)
+def set_task_state(
+ task_instance_id: UUID,
+ key: Annotated[str, Path(min_length=1)],
+ body: TaskStatePutBody,
+ session: SessionDep,
+) -> None:
+ """Set a task state key, creating or updating the row."""
+ scope = _get_task_scope_for_ti(task_instance_id, session)
+ get_state_backend().set(scope, key, body.value)
Review Comment:
Hey, thanks for asking. It's intentional. `ti:self` confirms the caller is a valid member of the mapped task group and that is the right trust boundary for a fleet-wide clearing, since requiring each sibling's token would make the feature unusable. The SDK only sends the flag on explicit `clear(all_map_indices=True)` calls, so the wider scope is always author opt-in. I have updated the docstring to call this out.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66073:
URL: https://github.com/apache/airflow/pull/66073#discussion_r3176661277
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_state.py:
##########
@@ -0,0 +1,121 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Execution API routes for asset state.
+
+Asset state is keyed by asset *name* (not integer id) in the URL — asset names
+are unique, and callers (task SDK accessors) have the name from their Asset
+object without needing a DB lookup. The route resolves name → asset_id
+internally for the state backend scope.
+
+Per-task asset registration checks are intentionally not implemented here
+(deferred to AIP-93 — see TODO comment below).
+"""
+
+from __future__ import annotations
+
+from typing import Annotated
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, status
+from sqlalchemy import select
+
+from airflow._shared.state import AssetScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.asset_state import (
+ AssetStatePutBody,
+ AssetStateResponse,
+)
+from airflow.models.asset import AssetModel
+from airflow.state import get_state_backend
+
+# TODO(AIP-103): enforce that the requesting task is registered with the asset
+# (via task_inlet_asset_reference or task_outlet_asset_reference) before
+# allowing reads/writes. Currently any task with a valid execution token can
+# access any asset's state — the same gap exists in /assets and /asset-events.
+# Proper fix is a unified asset-registration check across all asset routes,
+# not just here.
+router = VersionedAPIRouter(
+ responses={
+ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
+ status.HTTP_404_NOT_FOUND: {"description": "Not found"},
+ },
+)
+
+
+def _resolve_asset_id(name: str, session: SessionDep) -> int:
+ """Resolve asset name → integer asset_id, 404 if not found."""
+ asset_id = session.scalar(select(AssetModel.id).where(AssetModel.name == name, AssetModel.active.has()))
+ if asset_id is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={"reason": "not_found", "message": f"Asset {name!r} not found"},
+ )
+ return asset_id
+
+
+@router.get("/{name}/{key}")
Review Comment:
Oh thats a really good call, I wasn't aware why Asset APIs did that.
I switched to query params trying to match the `by-name` pattern already used in assets API. I also added a test with `team/sales/orders` as the asset name to confirm it works.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66073:
URL: https://github.com/apache/airflow/pull/66073#discussion_r3176666737
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_state.py:
##########
@@ -0,0 +1,121 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Execution API routes for asset state.
+
+Asset state is keyed by asset *name* (not integer id) in the URL — asset names
+are unique, and callers (task SDK accessors) have the name from their Asset
+object without needing a DB lookup. The route resolves name → asset_id
+internally for the state backend scope.
+
+Per-task asset registration checks are intentionally not implemented here
+(deferred to AIP-93 — see TODO comment below).
+"""
+
+from __future__ import annotations
+
+from typing import Annotated
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, status
+from sqlalchemy import select
+
+from airflow._shared.state import AssetScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.asset_state import (
+ AssetStatePutBody,
+ AssetStateResponse,
+)
+from airflow.models.asset import AssetModel
+from airflow.state import get_state_backend
+
+# TODO(AIP-103): enforce that the requesting task is registered with the asset
Review Comment:
Nice, good call. I am planning to expand this and do it in the GC issue: https://github.com/apache/airflow/issues/65780 to also delete `asset_state` rows for assets with no asset_active entry. Already made a note of that
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66073:
URL: https://github.com/apache/airflow/pull/66073#discussion_r3173257699
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state.py:
##########
@@ -0,0 +1,127 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Annotated
+from uuid import UUID
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, Query, Security, status
+from sqlalchemy.orm import Session
+
+from airflow._shared.state import TaskScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.task_state import (
+ TaskStatePutBody,
+ TaskStateResponse,
+)
+from airflow.api_fastapi.execution_api.security import require_auth
+from airflow.models.taskinstance import TaskInstance as TI
+from airflow.state import get_state_backend
+
+router = VersionedAPIRouter(
+ responses={
+ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
+ status.HTTP_403_FORBIDDEN: {"description": "Access denied"},
+ status.HTTP_404_NOT_FOUND: {"description": "Not found"},
+ },
+ dependencies=[Security(require_auth, scopes=["ti:self"])],
+)
+
+
+def _get_task_scope_for_ti(task_instance_id: UUID, session: Session) -> TaskScope:
+ ti = session.get(TI, task_instance_id)
+ if ti is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": f"Task instance {task_instance_id} not found",
+ },
+ )
+ return TaskScope(dag_id=ti.dag_id, run_id=ti.run_id, task_id=ti.task_id, map_index=ti.map_index)
+
+
+@router.get("/{task_instance_id}/{key}")
+def get_task_state(
+ task_instance_id: UUID,
+ key: Annotated[str, Path(min_length=1)],
+ session: SessionDep,
+) -> TaskStateResponse:
+ """Get value for a task state."""
+ scope = _get_task_scope_for_ti(task_instance_id, session)
+ value = get_state_backend().get(scope, key)
+ if value is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": f"Task state key {key!r} not found",
+ },
+ )
+ return TaskStateResponse(value=value)
+
+
+@router.put("/{task_instance_id}/{key}", status_code=status.HTTP_204_NO_CONTENT)
+def set_task_state(
+ task_instance_id: UUID,
+ key: Annotated[str, Path(min_length=1)],
+ body: TaskStatePutBody,
+ session: SessionDep,
+) -> None:
+ """Set a task state key, creating or updating the row."""
+ scope = _get_task_scope_for_ti(task_instance_id, session)
+ get_state_backend().set(scope, key, body.value)
Review Comment:
Hey, thanks for asking. It's intentional. `ti:self` confirms the caller is a valid member of the mapped task group and that is the right trust boundary for a fleet-wide clearing, since requiring each sibling's token would make the feature unusable. The SDK only sends the flag on explicit `clear(all_map_indices=True)` calls, so the wider scope is always author opt-in. I have updated the docstring to call this out.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state.py:
##########
@@ -0,0 +1,127 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Annotated
+from uuid import UUID
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, Query, Security, status
+from sqlalchemy.orm import Session
+
+from airflow._shared.state import TaskScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.task_state import (
+ TaskStatePutBody,
+ TaskStateResponse,
+)
+from airflow.api_fastapi.execution_api.security import require_auth
+from airflow.models.taskinstance import TaskInstance as TI
+from airflow.state import get_state_backend
+
+router = VersionedAPIRouter(
+ responses={
+ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
+ status.HTTP_403_FORBIDDEN: {"description": "Access denied"},
+ status.HTTP_404_NOT_FOUND: {"description": "Not found"},
+ },
+ dependencies=[Security(require_auth, scopes=["ti:self"])],
+)
+
+
+def _get_task_scope_for_ti(task_instance_id: UUID, session: Session) -> TaskScope:
+ ti = session.get(TI, task_instance_id)
+ if ti is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": f"Task instance {task_instance_id} not found",
+ },
+ )
+ return TaskScope(dag_id=ti.dag_id, run_id=ti.run_id, task_id=ti.task_id, map_index=ti.map_index)
+
+
+@router.get("/{task_instance_id}/{key}")
+def get_task_state(
+ task_instance_id: UUID,
+ key: Annotated[str, Path(min_length=1)],
+ session: SessionDep,
+) -> TaskStateResponse:
+ """Get value for a task state."""
+ scope = _get_task_scope_for_ti(task_instance_id, session)
+ value = get_state_backend().get(scope, key)
+ if value is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": f"Task state key {key!r} not found",
+ },
+ )
+ return TaskStateResponse(value=value)
+
+
+@router.put("/{task_instance_id}/{key}", status_code=status.HTTP_204_NO_CONTENT)
+def set_task_state(
+ task_instance_id: UUID,
+ key: Annotated[str, Path(min_length=1)],
+ body: TaskStatePutBody,
+ session: SessionDep,
+) -> None:
+ """Set a task state key, creating or updating the row."""
+ scope = _get_task_scope_for_ti(task_instance_id, session)
+ get_state_backend().set(scope, key, body.value)
+
+
+@router.delete("/{task_instance_id}/{key}", status_code=status.HTTP_204_NO_CONTENT)
+def delete_task_state(
+ task_instance_id: UUID,
+ key: Annotated[str, Path(min_length=1)],
+ session: SessionDep,
+) -> None:
+ """Delete a single task state key."""
+ scope = _get_task_scope_for_ti(task_instance_id, session)
+ get_state_backend().delete(scope, key)
+
+
+@router.delete("/{task_instance_id}", status_code=status.HTTP_204_NO_CONTENT)
+def clear_task_state(
+ task_instance_id: UUID,
+ session: SessionDep,
+ all_map_indices: Annotated[bool, Query()] = False,
Review Comment:
Hey, thanks for asking. It's intentional. ti:self confirms the caller is a valid member of the mapped task group and that is the right trust boundary for a fleet-wide clearing, since requiring each sibling's token would make the feature unusable. The SDK only sends the flag on explicit clear(all_map_indices=True) calls, so the wider scope is always author opt-in. I have updated the docstring to call this out.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66073:
URL: https://github.com/apache/airflow/pull/66073#discussion_r3173382971
##########
airflow-core/src/airflow/state/__init__.py:
##########
@@ -36,3 +36,14 @@ def resolve_state_backend() -> type[BaseStateBackend]:
f"Your custom state backend `{clazz.__name__}` is not a subclass of `BaseStateBackend`."
)
return clazz
+
+
+_backend_instance: BaseStateBackend | None = None
+
+
+def get_state_backend() -> BaseStateBackend:
+ """Return a cached instance of the configured state backend."""
+ global _backend_instance
+ if _backend_instance is None:
+ _backend_instance = resolve_state_backend()()
Review Comment:
Handled in https://github.com/apache/airflow/commit/96e92bffdd17e48f86527833c96e172dadd41350
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state.py:
##########
@@ -0,0 +1,127 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Annotated
+from uuid import UUID
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, Query, Security, status
+from sqlalchemy.orm import Session
+
+from airflow._shared.state import TaskScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.task_state import (
+ TaskStatePutBody,
+ TaskStateResponse,
+)
+from airflow.api_fastapi.execution_api.security import require_auth
+from airflow.models.taskinstance import TaskInstance as TI
+from airflow.state import get_state_backend
+
+router = VersionedAPIRouter(
+ responses={
+ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
+ status.HTTP_403_FORBIDDEN: {"description": "Access denied"},
+ status.HTTP_404_NOT_FOUND: {"description": "Not found"},
+ },
+ dependencies=[Security(require_auth, scopes=["ti:self"])],
Review Comment:
Handled in https://github.com/apache/airflow/commit/96e92bffdd17e48f86527833c96e172dadd41350
##########
airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_state.py:
##########
@@ -0,0 +1,238 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+from uuid import uuid4
+
+import pytest
+from sqlalchemy import delete, select
+
+from airflow._shared.timezones import timezone
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import create_session
+
+if TYPE_CHECKING:
+ from fastapi.testclient import TestClient
+
+ from tests_common.pytest_plugin import CreateTaskInstance
+
+pytestmark = pytest.mark.db_test
+
+
+@pytest.fixture(autouse=True)
+def reset_state_tables():
+ with create_session() as session:
+ session.execute(delete(TaskStateModel))
+ session.execute(delete(DagRun))
+
+
+def _api_url(ti_id, key: str | None = None) -> str:
+ base = f"/execution/state/ti/{ti_id}"
+ return f"{base}/{key}" if key else base
+
+
+class TestGetTaskState:
+ def test_get_returns_value(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+ client.put(_api_url(ti.id, "job_id"), json={"value": "spark_001"})
+
+ response = client.get(_api_url(ti.id, "job_id"))
+
+ assert response.status_code == 200
+ assert response.json() == {"value": "spark_001"}
+
+ def test_get_missing_key_returns_404(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+
+ response = client.get(_api_url(ti.id, "never_set"))
+
+ assert response.status_code == 404
+ assert response.json()["detail"]["reason"] == "not_found"
+
+ def test_get_missing_ti_returns_404(self, client: TestClient):
+ response = client.get(_api_url(uuid4(), "any_key"))
+
+ assert response.status_code == 404
+ assert "Task instance" in response.json()["detail"]["message"]
+
+
+class TestPutTaskState:
+ def test_put_creates_row(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+
+ response = client.put(_api_url(ti.id, "job_id"), json={"value": "spark_001"})
+
+ assert response.status_code == 204
+ with create_session() as session:
+ row = session.scalar(
+ select(TaskStateModel).where(
+ TaskStateModel.dag_id == ti.dag_id,
+ TaskStateModel.run_id == ti.run_id,
+ TaskStateModel.task_id == ti.task_id,
+ TaskStateModel.key == "job_id",
+ )
+ )
+ assert row is not None
+ assert row.value == "spark_001"
+
+ def test_put_overwrites_existing(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+ client.put(_api_url(ti.id, "job_id"), json={"value": "spark_001"})
+
+ response = client.put(_api_url(ti.id, "job_id"), json={"value": "spark_002"})
+
+ assert response.status_code == 204
+ assert client.get(_api_url(ti.id, "job_id")).json() == {"value": "spark_002"}
+
+ def test_put_empty_body_returns_422(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+
+ response = client.put(_api_url(ti.id, "job_id"), json={})
+
+ assert response.status_code == 422
+
+ def test_put_extra_field_returns_422(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+
+ response = client.put(_api_url(ti.id, "job_id"), json={"value": "x", "extra": "y"})
+
+ assert response.status_code == 422
+
+ def test_put_null_value_returns_422(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+
+ response = client.put(_api_url(ti.id, "job_id"), json={"value": None})
+
+ assert response.status_code == 422
+
+ def test_put_missing_ti_returns_404(self, client: TestClient):
+ response = client.put(_api_url(uuid4(), "job_id"), json={"value": "x"})
+
+ assert response.status_code == 404
+
+
+class TestDeleteTaskState:
+ def test_delete_removes_key(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+ client.put(_api_url(ti.id, "job_id"), json={"value": "spark_001"})
+
+ response = client.delete(_api_url(ti.id, "job_id"))
+
+ assert response.status_code == 204
+ assert client.get(_api_url(ti.id, "job_id")).status_code == 404
+
+ def test_delete_missing_key_is_noop(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+
+ response = client.delete(_api_url(ti.id, "never_existed"))
+
+ assert response.status_code == 204
+
+ def test_delete_only_targets_one_key(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+ client.put(_api_url(ti.id, "job_id"), json={"value": "a"})
+ client.put(_api_url(ti.id, "checkpoint"), json={"value": "b"})
+
+ client.delete(_api_url(ti.id, "job_id"))
+
+ assert client.get(_api_url(ti.id, "job_id")).status_code == 404
+ assert client.get(_api_url(ti.id, "checkpoint")).json() == {"value": "b"}
+
+
+class TestClearTaskState:
+ def test_clear_removes_all_keys(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+ for k, v in [("job_id", "a"), ("checkpoint", "b"), ("retry_count", "c")]:
+ client.put(_api_url(ti.id, k), json={"value": v})
+
+ response = client.delete(_api_url(ti.id))
+
+ assert response.status_code == 204
+ with create_session() as session:
+ count = session.scalar(
+ select(TaskStateModel)
+ .where(TaskStateModel.dag_id == ti.dag_id, TaskStateModel.task_id == ti.task_id)
+ .with_only_columns(TaskStateModel.key)
+ )
+ assert count is None
Review Comment:
Handled in https://github.com/apache/airflow/commit/96e92bffdd17e48f86527833c96e172dadd41350
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_state.py:
##########
@@ -0,0 +1,121 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Execution API routes for asset state.
+
+Asset state is keyed by asset *name* (not integer id) in the URL — asset names
+are unique, and callers (task SDK accessors) have the name from their Asset
+object without needing a DB lookup. The route resolves name → asset_id
+internally for the state backend scope.
+
+Per-task asset registration checks are intentionally not implemented here
+(deferred to AIP-93 — see TODO comment below).
+"""
+
+from __future__ import annotations
+
+from typing import Annotated
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, status
+from sqlalchemy import select
+
+from airflow._shared.state import AssetScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.asset_state import (
+ AssetStatePutBody,
+ AssetStateResponse,
+)
+from airflow.models.asset import AssetModel
+from airflow.state import get_state_backend
+
+# TODO(AIP-103): enforce that the requesting task is registered with the asset
+# (via task_inlet_asset_reference or task_outlet_asset_reference) before
+# allowing reads/writes. Currently any task with a valid execution token can
+# access any asset's state — the same gap exists in /assets and /asset-events.
+# Proper fix is a unified asset-registration check across all asset routes,
+# not just here.
+router = VersionedAPIRouter(
+ responses={
+ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
+ status.HTTP_404_NOT_FOUND: {"description": "Not found"},
+ },
+)
+
+
+def _resolve_asset_id(name: str, session: SessionDep) -> int:
+ """Resolve asset name → integer asset_id, 404 if not found."""
+ asset_id = session.scalar(select(AssetModel.id).where(AssetModel.name == name))
+ if asset_id is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={"reason": "not_found", "message": f"Asset {name!r} not found"},
+ )
+ return asset_id
+
+
+@router.get("/{name}/{key}")
+def get_asset_state(
+ name: Annotated[str, Path(min_length=1)],
+ key: Annotated[str, Path(min_length=1)],
+ session: SessionDep,
+) -> AssetStateResponse:
+ """Get an asset state value."""
+ asset_id = _resolve_asset_id(name, session)
+ value = get_state_backend().get(AssetScope(asset_id=asset_id), key)
+ if value is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": f"Asset state key {key!r} not found",
+ },
+ )
+ return AssetStateResponse(value=value)
+
+
+@router.put("/{name}/{key}", status_code=status.HTTP_204_NO_CONTENT)
+def set_asset_state(
+ name: Annotated[str, Path(min_length=1)],
+ key: Annotated[str, Path(min_length=1)],
+ body: AssetStatePutBody,
+ session: SessionDep,
+) -> None:
+ """Set an asset state value."""
+ asset_id = _resolve_asset_id(name, session)
+ get_state_backend().set(AssetScope(asset_id=asset_id), key, body.value)
Review Comment:
Handled in https://github.com/apache/airflow/commit/96e92bffdd17e48f86527833c96e172dadd41350
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66073:
URL: https://github.com/apache/airflow/pull/66073#discussion_r3173366757
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_state.py:
##########
@@ -0,0 +1,121 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Execution API routes for asset state.
+
+Asset state is keyed by asset *name* (not integer id) in the URL — asset names
+are unique, and callers (task SDK accessors) have the name from their Asset
+object without needing a DB lookup. The route resolves name → asset_id
+internally for the state backend scope.
+
+Per-task asset registration checks are intentionally not implemented here
+(deferred to AIP-93 — see TODO comment below).
+"""
+
+from __future__ import annotations
+
+from typing import Annotated
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, status
+from sqlalchemy import select
+
+from airflow._shared.state import AssetScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.asset_state import (
+ AssetStatePutBody,
+ AssetStateResponse,
+)
+from airflow.models.asset import AssetModel
+from airflow.state import get_state_backend
+
+# TODO(AIP-103): enforce that the requesting task is registered with the asset
+# (via task_inlet_asset_reference or task_outlet_asset_reference) before
+# allowing reads/writes. Currently any task with a valid execution token can
+# access any asset's state — the same gap exists in /assets and /asset-events.
+# Proper fix is a unified asset-registration check across all asset routes,
+# not just here.
+router = VersionedAPIRouter(
+ responses={
+ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
+ status.HTTP_404_NOT_FOUND: {"description": "Not found"},
+ },
+)
+
+
+def _resolve_asset_id(name: str, session: SessionDep) -> int:
+ """Resolve asset name → integer asset_id, 404 if not found."""
+ asset_id = session.scalar(select(AssetModel.id).where(AssetModel.name == name))
+ if asset_id is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={"reason": "not_found", "message": f"Asset {name!r} not found"},
+ )
+ return asset_id
+
+
+@router.get("/{name}/{key}")
+def get_asset_state(
+ name: Annotated[str, Path(min_length=1)],
+ key: Annotated[str, Path(min_length=1)],
+ session: SessionDep,
+) -> AssetStateResponse:
+ """Get an asset state value."""
+ asset_id = _resolve_asset_id(name, session)
+ value = get_state_backend().get(AssetScope(asset_id=asset_id), key)
+ if value is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": f"Asset state key {key!r} not found",
+ },
+ )
+ return AssetStateResponse(value=value)
+
+
+@router.put("/{name}/{key}", status_code=status.HTTP_204_NO_CONTENT)
+def set_asset_state(
+ name: Annotated[str, Path(min_length=1)],
+ key: Annotated[str, Path(min_length=1)],
+ body: AssetStatePutBody,
+ session: SessionDep,
+) -> None:
+ """Set an asset state value."""
+ asset_id = _resolve_asset_id(name, session)
+ get_state_backend().set(AssetScope(asset_id=asset_id), key, body.value)
Review Comment:
This has been handled like aboev
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66073:
URL: https://github.com/apache/airflow/pull/66073#discussion_r3173214642
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_state.py:
##########
@@ -0,0 +1,95 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Annotated
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, status
+
+from airflow._shared.state import AssetScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.asset_state import (
+ AssetStatePutBody,
+ AssetStateResponse,
+)
+from airflow.models.asset import AssetModel
+from airflow.state import get_state_backend
+
+# TODO(AIP-103): enforce that the requesting task is registered with the asset
+# (via task_inlet_asset_reference or task_outlet_asset_reference) before
+# allowing reads/writes. Currently any task with a valid execution token can
+# access any asset's state — the same gap exists in /assets and /asset-events.
+# Proper fix is a unified asset-registration check across all asset routes,
+# not just here.
+router = VersionedAPIRouter(
+ responses={
+ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
+ status.HTTP_404_NOT_FOUND: {"description": "Not found"},
+ },
+)
+
+
+@router.get("/{asset_id}/{key}")
Review Comment:
@jscheffl both routers are included under `authenticated_router` on line 43 which has `Security(require_auth)` as a dependency.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66073:
URL: https://github.com/apache/airflow/pull/66073#discussion_r3173364689
##########
airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_state.py:
##########
@@ -0,0 +1,238 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+from uuid import uuid4
+
+import pytest
+from sqlalchemy import delete, select
+
+from airflow._shared.timezones import timezone
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import create_session
+
+if TYPE_CHECKING:
+ from fastapi.testclient import TestClient
+
+ from tests_common.pytest_plugin import CreateTaskInstance
+
+pytestmark = pytest.mark.db_test
+
+
+@pytest.fixture(autouse=True)
+def reset_state_tables():
+ with create_session() as session:
+ session.execute(delete(TaskStateModel))
+ session.execute(delete(DagRun))
+
+
+def _api_url(ti_id, key: str | None = None) -> str:
+ base = f"/execution/state/ti/{ti_id}"
+ return f"{base}/{key}" if key else base
+
+
+class TestGetTaskState:
+ def test_get_returns_value(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+ client.put(_api_url(ti.id, "job_id"), json={"value": "spark_001"})
+
+ response = client.get(_api_url(ti.id, "job_id"))
+
+ assert response.status_code == 200
+ assert response.json() == {"value": "spark_001"}
+
+ def test_get_missing_key_returns_404(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+
+ response = client.get(_api_url(ti.id, "never_set"))
+
+ assert response.status_code == 404
+ assert response.json()["detail"]["reason"] == "not_found"
+
+ def test_get_missing_ti_returns_404(self, client: TestClient):
+ response = client.get(_api_url(uuid4(), "any_key"))
+
+ assert response.status_code == 404
+ assert "Task instance" in response.json()["detail"]["message"]
+
+
+class TestPutTaskState:
+ def test_put_creates_row(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+
+ response = client.put(_api_url(ti.id, "job_id"), json={"value": "spark_001"})
+
+ assert response.status_code == 204
+ with create_session() as session:
+ row = session.scalar(
+ select(TaskStateModel).where(
+ TaskStateModel.dag_id == ti.dag_id,
+ TaskStateModel.run_id == ti.run_id,
+ TaskStateModel.task_id == ti.task_id,
+ TaskStateModel.key == "job_id",
+ )
+ )
+ assert row is not None
+ assert row.value == "spark_001"
+
+ def test_put_overwrites_existing(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+ client.put(_api_url(ti.id, "job_id"), json={"value": "spark_001"})
+
+ response = client.put(_api_url(ti.id, "job_id"), json={"value": "spark_002"})
+
+ assert response.status_code == 204
+ assert client.get(_api_url(ti.id, "job_id")).json() == {"value": "spark_002"}
+
+ def test_put_empty_body_returns_422(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+
+ response = client.put(_api_url(ti.id, "job_id"), json={})
+
+ assert response.status_code == 422
+
+ def test_put_extra_field_returns_422(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+
+ response = client.put(_api_url(ti.id, "job_id"), json={"value": "x", "extra": "y"})
+
+ assert response.status_code == 422
+
+ def test_put_null_value_returns_422(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+
+ response = client.put(_api_url(ti.id, "job_id"), json={"value": None})
+
+ assert response.status_code == 422
+
+ def test_put_missing_ti_returns_404(self, client: TestClient):
+ response = client.put(_api_url(uuid4(), "job_id"), json={"value": "x"})
+
+ assert response.status_code == 404
+
+
+class TestDeleteTaskState:
+ def test_delete_removes_key(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+ client.put(_api_url(ti.id, "job_id"), json={"value": "spark_001"})
+
+ response = client.delete(_api_url(ti.id, "job_id"))
+
+ assert response.status_code == 204
+ assert client.get(_api_url(ti.id, "job_id")).status_code == 404
+
+ def test_delete_missing_key_is_noop(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+
+ response = client.delete(_api_url(ti.id, "never_existed"))
+
+ assert response.status_code == 204
+
+ def test_delete_only_targets_one_key(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+ client.put(_api_url(ti.id, "job_id"), json={"value": "a"})
+ client.put(_api_url(ti.id, "checkpoint"), json={"value": "b"})
+
+ client.delete(_api_url(ti.id, "job_id"))
+
+ assert client.get(_api_url(ti.id, "job_id")).status_code == 404
+ assert client.get(_api_url(ti.id, "checkpoint")).json() == {"value": "b"}
+
+
+class TestClearTaskState:
+ def test_clear_removes_all_keys(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+ for k, v in [("job_id", "a"), ("checkpoint", "b"), ("retry_count", "c")]:
+ client.put(_api_url(ti.id, k), json={"value": v})
+
+ response = client.delete(_api_url(ti.id))
+
+ assert response.status_code == 204
+ with create_session() as session:
+ count = session.scalar(
+ select(TaskStateModel)
+ .where(TaskStateModel.dag_id == ti.dag_id, TaskStateModel.task_id == ti.task_id)
+ .with_only_columns(TaskStateModel.key)
+ )
+ assert count is None
Review Comment:
Went ahead with
```
remaining = session.scalars(
select(TaskStateModel.key).where(
TaskStateModel.dag_id == ti.dag_id,
TaskStateModel.task_id == ti.task_id,
)
).all()
assert remaining == []
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66073:
URL: https://github.com/apache/airflow/pull/66073#discussion_r3173351649
##########
airflow-core/src/airflow/state/__init__.py:
##########
@@ -36,3 +36,14 @@ def resolve_state_backend() -> type[BaseStateBackend]:
f"Your custom state backend `{clazz.__name__}` is not a subclass of `BaseStateBackend`."
)
return clazz
+
+
+_backend_instance: BaseStateBackend | None = None
+
+
+def get_state_backend() -> BaseStateBackend:
+ """Return a cached instance of the configured state backend."""
+ global _backend_instance
+ if _backend_instance is None:
+ _backend_instance = resolve_state_backend()()
Review Comment:
Fixed with locking
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66073:
URL: https://github.com/apache/airflow/pull/66073#discussion_r3176670526
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state.py:
##########
@@ -0,0 +1,136 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Annotated
+from uuid import UUID
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, Query, Security, status
+from sqlalchemy.orm import Session
+
+from airflow._shared.state import TaskScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.task_state import (
+ TaskStatePutBody,
+ TaskStateResponse,
+)
+from airflow.api_fastapi.execution_api.security import require_auth
+from airflow.models.taskinstance import TaskInstance as TI
+from airflow.state import get_state_backend
+
+router = VersionedAPIRouter(
Review Comment:
Good catch, yeah not an issue now but worth handling it, I added `route_class=ExecutionAPIRoute` to both routers, matching HITL and other routers.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66073:
URL: https://github.com/apache/airflow/pull/66073#discussion_r3173382126
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_state.py:
##########
@@ -0,0 +1,121 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Execution API routes for asset state.
+
+Asset state is keyed by asset *name* (not integer id) in the URL — asset names
+are unique, and callers (task SDK accessors) have the name from their Asset
+object without needing a DB lookup. The route resolves name → asset_id
+internally for the state backend scope.
+
+Per-task asset registration checks are intentionally not implemented here
+(deferred to AIP-93 — see TODO comment below).
+"""
+
+from __future__ import annotations
+
+from typing import Annotated
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, status
+from sqlalchemy import select
+
+from airflow._shared.state import AssetScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.asset_state import (
+ AssetStatePutBody,
+ AssetStateResponse,
+)
+from airflow.models.asset import AssetModel
+from airflow.state import get_state_backend
+
+# TODO(AIP-103): enforce that the requesting task is registered with the asset
+# (via task_inlet_asset_reference or task_outlet_asset_reference) before
+# allowing reads/writes. Currently any task with a valid execution token can
+# access any asset's state — the same gap exists in /assets and /asset-events.
+# Proper fix is a unified asset-registration check across all asset routes,
+# not just here.
+router = VersionedAPIRouter(
+ responses={
+ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
+ status.HTTP_404_NOT_FOUND: {"description": "Not found"},
+ },
+)
+
+
+def _resolve_asset_id(name: str, session: SessionDep) -> int:
+ """Resolve asset name → integer asset_id, 404 if not found."""
+ asset_id = session.scalar(select(AssetModel.id).where(AssetModel.name == name))
Review Comment:
Handled in https://github.com/apache/airflow/pull/66073/commits/96e92bffdd17e48f86527833c96e172dadd41350
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state.py:
##########
@@ -0,0 +1,127 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Annotated
+from uuid import UUID
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, Query, Security, status
+from sqlalchemy.orm import Session
+
+from airflow._shared.state import TaskScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.task_state import (
+ TaskStatePutBody,
+ TaskStateResponse,
+)
+from airflow.api_fastapi.execution_api.security import require_auth
+from airflow.models.taskinstance import TaskInstance as TI
+from airflow.state import get_state_backend
+
+router = VersionedAPIRouter(
+ responses={
+ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
+ status.HTTP_403_FORBIDDEN: {"description": "Access denied"},
+ status.HTTP_404_NOT_FOUND: {"description": "Not found"},
+ },
+ dependencies=[Security(require_auth, scopes=["ti:self"])],
+)
+
+
+def _get_task_scope_for_ti(task_instance_id: UUID, session: Session) -> TaskScope:
+ ti = session.get(TI, task_instance_id)
+ if ti is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": f"Task instance {task_instance_id} not found",
+ },
+ )
+ return TaskScope(dag_id=ti.dag_id, run_id=ti.run_id, task_id=ti.task_id, map_index=ti.map_index)
+
+
+@router.get("/{task_instance_id}/{key}")
+def get_task_state(
+ task_instance_id: UUID,
+ key: Annotated[str, Path(min_length=1)],
+ session: SessionDep,
+) -> TaskStateResponse:
+ """Get value for a task state."""
+ scope = _get_task_scope_for_ti(task_instance_id, session)
+ value = get_state_backend().get(scope, key)
+ if value is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": f"Task state key {key!r} not found",
+ },
+ )
+ return TaskStateResponse(value=value)
+
+
+@router.put("/{task_instance_id}/{key}", status_code=status.HTTP_204_NO_CONTENT)
+def set_task_state(
+ task_instance_id: UUID,
+ key: Annotated[str, Path(min_length=1)],
+ body: TaskStatePutBody,
+ session: SessionDep,
+) -> None:
+ """Set a task state key, creating or updating the row."""
+ scope = _get_task_scope_for_ti(task_instance_id, session)
+ get_state_backend().set(scope, key, body.value)
+
+
+@router.delete("/{task_instance_id}/{key}", status_code=status.HTTP_204_NO_CONTENT)
+def delete_task_state(
+ task_instance_id: UUID,
+ key: Annotated[str, Path(min_length=1)],
+ session: SessionDep,
+) -> None:
+ """Delete a single task state key."""
+ scope = _get_task_scope_for_ti(task_instance_id, session)
+ get_state_backend().delete(scope, key)
+
+
+@router.delete("/{task_instance_id}", status_code=status.HTTP_204_NO_CONTENT)
+def clear_task_state(
+ task_instance_id: UUID,
+ session: SessionDep,
+ all_map_indices: Annotated[bool, Query()] = False,
Review Comment:
Handled in https://github.com/apache/airflow/pull/66073/commits/96e92bffdd17e48f86527833c96e172dadd41350
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/__init__.py:
##########
@@ -52,5 +54,7 @@
authenticated_router.include_router(variables.router, prefix="/variables", tags=["Variables"])
authenticated_router.include_router(xcoms.router, prefix="/xcoms", tags=["XComs"])
authenticated_router.include_router(hitl.router, prefix="/hitlDetails", tags=["Human in the Loop"])
+authenticated_router.include_router(task_state.router, prefix="/state/ti", tags=["Task State"])
+authenticated_router.include_router(asset_state.router, prefix="/state/asset", tags=["Asset State"])
Review Comment:
Handled in https://github.com/apache/airflow/commit/96e92bffdd17e48f86527833c96e172dadd41350
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66073:
URL: https://github.com/apache/airflow/pull/66073#discussion_r3176674352
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state.py:
##########
@@ -0,0 +1,136 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Annotated
+from uuid import UUID
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, Query, Security, status
+from sqlalchemy.orm import Session
+
+from airflow._shared.state import TaskScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.task_state import (
+ TaskStatePutBody,
+ TaskStateResponse,
+)
+from airflow.api_fastapi.execution_api.security import require_auth
+from airflow.models.taskinstance import TaskInstance as TI
+from airflow.state import get_state_backend
+
+router = VersionedAPIRouter(
+ responses={
+ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
+ status.HTTP_403_FORBIDDEN: {"description": "Access denied"},
+ status.HTTP_404_NOT_FOUND: {"description": "Not found"},
+ },
+ dependencies=[Security(require_auth, scopes=["ti:self"])],
Review Comment:
Handled in 7bd1f4c64a
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_state.py:
##########
@@ -0,0 +1,121 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Execution API routes for asset state.
+
+Asset state is keyed by asset *name* (not integer id) in the URL — asset names
+are unique, and callers (task SDK accessors) have the name from their Asset
+object without needing a DB lookup. The route resolves name → asset_id
+internally for the state backend scope.
+
+Per-task asset registration checks are intentionally not implemented here
+(deferred to AIP-93 — see TODO comment below).
+"""
+
+from __future__ import annotations
+
+from typing import Annotated
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, status
+from sqlalchemy import select
+
+from airflow._shared.state import AssetScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.asset_state import (
+ AssetStatePutBody,
+ AssetStateResponse,
+)
+from airflow.models.asset import AssetModel
+from airflow.state import get_state_backend
+
+# TODO(AIP-103): enforce that the requesting task is registered with the asset
+# (via task_inlet_asset_reference or task_outlet_asset_reference) before
+# allowing reads/writes. Currently any task with a valid execution token can
+# access any asset's state — the same gap exists in /assets and /asset-events.
+# Proper fix is a unified asset-registration check across all asset routes,
+# not just here.
+router = VersionedAPIRouter(
+ responses={
+ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
+ status.HTTP_404_NOT_FOUND: {"description": "Not found"},
+ },
+)
+
+
+def _resolve_asset_id(name: str, session: SessionDep) -> int:
+ """Resolve asset name → integer asset_id, 404 if not found."""
+ asset_id = session.scalar(select(AssetModel.id).where(AssetModel.name == name, AssetModel.active.has()))
+ if asset_id is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={"reason": "not_found", "message": f"Asset {name!r} not found"},
+ )
+ return asset_id
+
+
+@router.get("/{name}/{key}")
Review Comment:
Handled in 7bd1f4c64a
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_state.py:
##########
@@ -0,0 +1,121 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Execution API routes for asset state.
+
+Asset state is keyed by asset *name* (not integer id) in the URL — asset names
+are unique, and callers (task SDK accessors) have the name from their Asset
+object without needing a DB lookup. The route resolves name → asset_id
+internally for the state backend scope.
+
+Per-task asset registration checks are intentionally not implemented here
+(deferred to AIP-93 — see TODO comment below).
+"""
+
+from __future__ import annotations
+
+from typing import Annotated
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, status
+from sqlalchemy import select
+
+from airflow._shared.state import AssetScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.asset_state import (
+ AssetStatePutBody,
+ AssetStateResponse,
+)
+from airflow.models.asset import AssetModel
+from airflow.state import get_state_backend
+
+# TODO(AIP-103): enforce that the requesting task is registered with the asset
Review Comment:
Handled in 7bd1f4c64a
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh closed pull request #66073: AIP-103: Add Execution API endpoints for task and asset states
URL: https://github.com/apache/airflow/pull/66073
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66073:
URL: https://github.com/apache/airflow/pull/66073#discussion_r3176656105
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state.py:
##########
@@ -0,0 +1,136 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Annotated
+from uuid import UUID
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, Query, Security, status
+from sqlalchemy.orm import Session
+
+from airflow._shared.state import TaskScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.task_state import (
+ TaskStatePutBody,
+ TaskStateResponse,
+)
+from airflow.api_fastapi.execution_api.security import require_auth
+from airflow.models.taskinstance import TaskInstance as TI
+from airflow.state import get_state_backend
+
+router = VersionedAPIRouter(
+ responses={
+ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
+ status.HTTP_403_FORBIDDEN: {"description": "Access denied"},
+ status.HTTP_404_NOT_FOUND: {"description": "Not found"},
+ },
+ dependencies=[Security(require_auth, scopes=["ti:self"])],
Review Comment:
Sure, added some tests in `TestTiSelfEnforcement` class with three tests that bypass the conftest mock by overriding `_jwt_bearer` directly instead of `require_auth`. This lets `require_auth` run properly and validates that a token bound to a different ti UUID gets 403 on GET, PUT, and clear. Covers the fleet wide all_map_indices path while I am at it
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh merged PR #66073:
URL: https://github.com/apache/airflow/pull/66073
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on PR #66073:
URL: https://github.com/apache/airflow/pull/66073#issuecomment-4368535451
No need to rebase
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #66073:
URL: https://github.com/apache/airflow/pull/66073#discussion_r3175794465
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state.py:
##########
@@ -0,0 +1,136 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Annotated
+from uuid import UUID
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, Query, Security, status
+from sqlalchemy.orm import Session
+
+from airflow._shared.state import TaskScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.task_state import (
+ TaskStatePutBody,
+ TaskStateResponse,
+)
+from airflow.api_fastapi.execution_api.security import require_auth
+from airflow.models.taskinstance import TaskInstance as TI
+from airflow.state import get_state_backend
+
+router = VersionedAPIRouter(
Review Comment:
This router and `asset_state.py:52` don't set `route_class=ExecutionAPIRoute`, but the other versioned routers in this module do (`hitl.py:36-37`, `task_instances.py:89-90`). The default at `security.py:170` makes `ti:self` work today, but as soon as someone wants to add a `Security(require_auth, scopes=["token:workload"])` route here, the workload-token branch won't be reachable because `allowed_token_types` falls back to `{"execution"}`. Worth aligning with the rest of the execution API now while these routers are new.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state.py:
##########
@@ -0,0 +1,136 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Annotated
+from uuid import UUID
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, Query, Security, status
+from sqlalchemy.orm import Session
+
+from airflow._shared.state import TaskScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.task_state import (
+ TaskStatePutBody,
+ TaskStateResponse,
+)
+from airflow.api_fastapi.execution_api.security import require_auth
+from airflow.models.taskinstance import TaskInstance as TI
+from airflow.state import get_state_backend
+
+router = VersionedAPIRouter(
+ responses={
+ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
+ status.HTTP_403_FORBIDDEN: {"description": "Access denied"},
+ status.HTTP_404_NOT_FOUND: {"description": "Not found"},
+ },
+ dependencies=[Security(require_auth, scopes=["ti:self"])],
+)
+
+
+def _get_task_scope_for_ti(task_instance_id: UUID, session: Session) -> TaskScope:
+ ti = session.get(TI, task_instance_id)
+ if ti is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": f"Task instance {task_instance_id} not found",
+ },
+ )
+ return TaskScope(dag_id=ti.dag_id, run_id=ti.run_id, task_id=ti.task_id, map_index=ti.map_index)
+
+
+@router.get("/{task_instance_id}/{key}")
+def get_task_state(
+ task_instance_id: UUID,
+ key: Annotated[str, Path(min_length=1)],
+ session: SessionDep,
+) -> TaskStateResponse:
+ """Get value for a task state."""
+ scope = _get_task_scope_for_ti(task_instance_id, session)
+ value = get_state_backend().get(scope, key, session=session)
+ if value is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": f"Task state key {key!r} not found",
+ },
+ )
+ return TaskStateResponse(value=value)
+
+
+@router.put("/{task_instance_id}/{key}", status_code=status.HTTP_204_NO_CONTENT)
+def set_task_state(
+ task_instance_id: UUID,
+ key: Annotated[str, Path(min_length=1)],
+ body: TaskStatePutBody,
+ session: SessionDep,
+) -> None:
+ """Set a task state key, creating or updating the row."""
+ scope = _get_task_scope_for_ti(task_instance_id, session)
+ get_state_backend().set(scope, key, body.value, session=session)
+
+
+@router.delete("/{task_instance_id}/{key}", status_code=status.HTTP_204_NO_CONTENT)
+def delete_task_state(
+ task_instance_id: UUID,
+ key: Annotated[str, Path(min_length=1)],
+ session: SessionDep,
+) -> None:
+ """Delete a single task state key."""
+ scope = _get_task_scope_for_ti(task_instance_id, session)
+ get_state_backend().delete(scope, key, session=session)
+
+
+@router.delete("/{task_instance_id}", status_code=status.HTTP_204_NO_CONTENT)
+def clear_task_state(
+ task_instance_id: UUID,
+ session: SessionDep,
+ all_map_indices: Annotated[bool, Query()] = False,
+) -> None:
+ """
+ Delete all task state keys for this task instance.
+
+ By default, only keys for the requesting TI's exact ``map_index`` are
+ cleared — same isolation as DELETE endpoint above.
+
+ Pass ``?all_map_indices=true`` to wipe state for every mapped sibling of
+ the task in the same DAG run. This is intentionally fleet-wide: the
+ ``ti:self`` JWT authentication scope authenticates that the caller is
+ a legitimate member of the mapped task group, and grants it authority
+ to reset shared task state on behalf of the whole group.
+ The SDK only forwards this flag when the user calls ``task_state.clear(all_map_indices=True)``
+ explicitly, so the expanded scope is always an explicit opt-in by the task author.
+
+ For non-mapped tasks (``map_index=-1``), there is only ever one index, so
+ ``?all_map_indices=true`` is functionally identical to the default and is
+ accepted without error.
+ """
+ ti = session.get(TI, task_instance_id)
+ if ti is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": f"Task instance {task_instance_id} not found",
+ },
+ )
+ scope = TaskScope(dag_id=ti.dag_id, run_id=ti.run_id, task_id=ti.task_id, map_index=ti.map_index)
Review Comment:
`clear_task_state` re-implements the TI lookup, 404 handling, and `TaskScope` build that already lives in `_get_task_scope_for_ti` (line 46). The two will drift: if anyone tweaks the 404 message, adds a permission check, or changes how the scope is derived in the helper, this branch silently misses it. Replacing lines 126-135 with `scope = _get_task_scope_for_ti(task_instance_id, session)` keeps the four route handlers in lockstep.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_state.py:
##########
@@ -0,0 +1,121 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Execution API routes for asset state.
+
+Asset state is keyed by asset *name* (not integer id) in the URL — asset names
+are unique, and callers (task SDK accessors) have the name from their Asset
+object without needing a DB lookup. The route resolves name → asset_id
+internally for the state backend scope.
+
+Per-task asset registration checks are intentionally not implemented here
+(deferred to AIP-93 — see TODO comment below).
+"""
+
+from __future__ import annotations
+
+from typing import Annotated
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, status
+from sqlalchemy import select
+
+from airflow._shared.state import AssetScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.asset_state import (
+ AssetStatePutBody,
+ AssetStateResponse,
+)
+from airflow.models.asset import AssetModel
+from airflow.state import get_state_backend
+
+# TODO(AIP-103): enforce that the requesting task is registered with the asset
+# (via task_inlet_asset_reference or task_outlet_asset_reference) before
+# allowing reads/writes. Currently any task with a valid execution token can
+# access any asset's state — the same gap exists in /assets and /asset-events.
+# Proper fix is a unified asset-registration check across all asset routes,
+# not just here.
+router = VersionedAPIRouter(
+ responses={
+ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
+ status.HTTP_404_NOT_FOUND: {"description": "Not found"},
+ },
+)
+
+
+def _resolve_asset_id(name: str, session: SessionDep) -> int:
+ """Resolve asset name → integer asset_id, 404 if not found."""
+ asset_id = session.scalar(select(AssetModel.id).where(AssetModel.name == name, AssetModel.active.has()))
+ if asset_id is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={"reason": "not_found", "message": f"Asset {name!r} not found"},
+ )
+ return asset_id
+
+
+@router.get("/{name}/{key}")
Review Comment:
Asset names accept `/` (only `self` and `context` are prohibited; see `_validate_asset_name` in `task-sdk/.../asset/__init__.py:206-210`). A legitimate asset named `team/sales` is unreachable at `/state/asset/team/sales/{key}` because FastAPI matches `name=team`, `key=sales`, and the trailing `/key` falls off the route entirely. Same for keys with `/`. Two existing options in this codebase: `assets.py` uses `/by-name?name=...` query parameters precisely to side-step this, or you can use a path converter that consumes the rest (`name: Annotated[str, Path(min_length=1, regex=...)]`). Worth doing here too, otherwise users with grouped/namespaced assets silently can't use this API at all.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_state.py:
##########
@@ -0,0 +1,121 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Execution API routes for asset state.
+
+Asset state is keyed by asset *name* (not integer id) in the URL — asset names
+are unique, and callers (task SDK accessors) have the name from their Asset
+object without needing a DB lookup. The route resolves name → asset_id
+internally for the state backend scope.
+
+Per-task asset registration checks are intentionally not implemented here
+(deferred to AIP-93 — see TODO comment below).
+"""
+
+from __future__ import annotations
+
+from typing import Annotated
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, status
+from sqlalchemy import select
+
+from airflow._shared.state import AssetScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.asset_state import (
+ AssetStatePutBody,
+ AssetStateResponse,
+)
+from airflow.models.asset import AssetModel
+from airflow.state import get_state_backend
+
+# TODO(AIP-103): enforce that the requesting task is registered with the asset
Review Comment:
Now that the resolver filters by `AssetModel.active.has()` (good change), state rows for assets that go inactive (DAG stops referencing them) become permanently unreachable through the API. `asset_state.asset_id` only cascades on `asset` row delete (`asset_state.py:53` `ondelete="CASCADE"`), but Airflow's GC removes the `asset_active` row, not the `asset` row. Net effect: orphan state rows accumulate forever for assets that ever went inactive. Worth deciding now whether to (a) cascade-delete state rows when `AssetActive` is removed, (b) add a periodic GC, or (c) document this and offer an admin-only "clear-all-inactive" endpoint. Otherwise this'll quietly grow forever in any long-lived deployment.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state.py:
##########
@@ -0,0 +1,136 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Annotated
+from uuid import UUID
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, Query, Security, status
+from sqlalchemy.orm import Session
+
+from airflow._shared.state import TaskScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.task_state import (
+ TaskStatePutBody,
+ TaskStateResponse,
+)
+from airflow.api_fastapi.execution_api.security import require_auth
+from airflow.models.taskinstance import TaskInstance as TI
+from airflow.state import get_state_backend
+
+router = VersionedAPIRouter(
+ responses={
+ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
+ status.HTTP_403_FORBIDDEN: {"description": "Access denied"},
+ status.HTTP_404_NOT_FOUND: {"description": "Not found"},
+ },
+ dependencies=[Security(require_auth, scopes=["ti:self"])],
Review Comment:
Following up on the prior comment on this line. The `ti:self` enforcement here is the primary protection on these routes, but the conftest at `airflow-core/tests/unit/api_fastapi/execution_api/conftest.py:56-60` overrides `require_auth` with a mock that always builds a `TIToken` whose `id` matches `path_params["task_instance_id"]`. With that override in place the cross-TI mismatch branch in `security.py:185-191` is structurally untestable here. Could you add at least one regression test that uses the real `require_auth` (a token bound to a different TI hits a 403 against this router)? The all-map-indices fleet-wide path is now also explicitly relying on this scope by docstring, which makes coverage even more important.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66073:
URL: https://github.com/apache/airflow/pull/66073#discussion_r3173342644
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state.py:
##########
@@ -0,0 +1,127 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Annotated
+from uuid import UUID
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, Query, Security, status
+from sqlalchemy.orm import Session
+
+from airflow._shared.state import TaskScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.task_state import (
+ TaskStatePutBody,
+ TaskStateResponse,
+)
+from airflow.api_fastapi.execution_api.security import require_auth
+from airflow.models.taskinstance import TaskInstance as TI
+from airflow.state import get_state_backend
+
+router = VersionedAPIRouter(
+ responses={
+ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
+ status.HTTP_403_FORBIDDEN: {"description": "Access denied"},
+ status.HTTP_404_NOT_FOUND: {"description": "Not found"},
+ },
+ dependencies=[Security(require_auth, scopes=["ti:self"])],
+)
+
+
+def _get_task_scope_for_ti(task_instance_id: UUID, session: Session) -> TaskScope:
+ ti = session.get(TI, task_instance_id)
+ if ti is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": f"Task instance {task_instance_id} not found",
+ },
+ )
+ return TaskScope(dag_id=ti.dag_id, run_id=ti.run_id, task_id=ti.task_id, map_index=ti.map_index)
+
+
+@router.get("/{task_instance_id}/{key}")
+def get_task_state(
+ task_instance_id: UUID,
+ key: Annotated[str, Path(min_length=1)],
+ session: SessionDep,
+) -> TaskStateResponse:
+ """Get value for a task state."""
+ scope = _get_task_scope_for_ti(task_instance_id, session)
+ value = get_state_backend().get(scope, key)
+ if value is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": f"Task state key {key!r} not found",
+ },
+ )
+ return TaskStateResponse(value=value)
+
+
+@router.put("/{task_instance_id}/{key}", status_code=status.HTTP_204_NO_CONTENT)
+def set_task_state(
+ task_instance_id: UUID,
+ key: Annotated[str, Path(min_length=1)],
+ body: TaskStatePutBody,
+ session: SessionDep,
+) -> None:
+ """Set a task state key, creating or updating the row."""
+ scope = _get_task_scope_for_ti(task_instance_id, session)
+ get_state_backend().set(scope, key, body.value)
Review Comment:
Fixed. All backend calls in both `task_state.py` and `asset_state.py` now pass `session=session` explicitly, so `@provide_session` reuses the injected session rather than opening a second one
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on PR #66073:
URL: https://github.com/apache/airflow/pull/66073#issuecomment-4366440136
Rerunning with full tests
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #66073:
URL: https://github.com/apache/airflow/pull/66073#issuecomment-4368268224
### Backport failed to create: v3-2-test. View the failure log <a href='https://github.com/apache/airflow/actions/runs/25301225672'> Run details </a>
Note: As of [Merging PRs targeted for Airflow 3.X](https://github.com/apache/airflow/blob/main/dev/README_AIRFLOW3_DEV.md#merging-prs-targeted-for-airflow-3x)
the committer who merges the PR is responsible for backporting the PRs that are bug fixes (generally speaking) to the maintenance branches.
In matter of doubt please ask in [#release-management](https://apache-airflow.slack.com/archives/C03G9H97MM2) Slack channel.
<table>
<tr>
<th>Status</th>
<th>Branch</th>
<th>Result</th>
</tr>
<tr>
<td>❌</td>
<td>v3-2-test</td>
<td><a href="https://github.com/apache/airflow/commit/ffb1b8a6d22b855da6f6a785c22125369732fd37"><img src='https://img.shields.io/badge/Commit-ffb1b8a-red' alt='Commit Link'></a></td>
</tr>
</table>
You can attempt to backport this manually by running:
```bash
cherry_picker ffb1b8a v3-2-test
```
This should apply the commit to the v3-2-test branch and leave the commit in conflict state marking
the files that need manual conflict resolution.
After you have resolved the conflicts, you can continue the backport process by running:
```bash
cherry_picker --continue
```
If you don't have cherry-picker installed, see the [installation guide](https://github.com/apache/airflow/blob/main/dev/README_AIRFLOW3_DEV.md#how-to-backport-pr-with-cherry-picker-cli).
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66073:
URL: https://github.com/apache/airflow/pull/66073#discussion_r3173283797
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/__init__.py:
##########
@@ -52,5 +54,7 @@
authenticated_router.include_router(variables.router, prefix="/variables", tags=["Variables"])
authenticated_router.include_router(xcoms.router, prefix="/xcoms", tags=["XComs"])
authenticated_router.include_router(hitl.router, prefix="/hitlDetails", tags=["Human in the Loop"])
+authenticated_router.include_router(task_state.router, prefix="/state/ti", tags=["Task State"])
+authenticated_router.include_router(asset_state.router, prefix="/state/asset", tags=["Asset State"])
Review Comment:
You are right on consistency grounds. As I understand it, technically new endpoints won't break old clients that don't know about them right? ie: old task-sdk workers have no code to call task_state or asset_state APIs.
But in any case, updating cadwyn migration rules as a record of what was introduced when is good practice and aligns with how we have done it before, so adding it.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Add Execution API endpoints for task and asset states [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #66073:
URL: https://github.com/apache/airflow/pull/66073#discussion_r3170621778
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_state.py:
##########
@@ -0,0 +1,121 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Execution API routes for asset state.
+
+Asset state is keyed by asset *name* (not integer id) in the URL — asset names
+are unique, and callers (task SDK accessors) have the name from their Asset
+object without needing a DB lookup. The route resolves name → asset_id
+internally for the state backend scope.
+
+Per-task asset registration checks are intentionally not implemented here
+(deferred to AIP-93 — see TODO comment below).
+"""
+
+from __future__ import annotations
+
+from typing import Annotated
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, status
+from sqlalchemy import select
+
+from airflow._shared.state import AssetScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.asset_state import (
+ AssetStatePutBody,
+ AssetStateResponse,
+)
+from airflow.models.asset import AssetModel
+from airflow.state import get_state_backend
+
+# TODO(AIP-103): enforce that the requesting task is registered with the asset
+# (via task_inlet_asset_reference or task_outlet_asset_reference) before
+# allowing reads/writes. Currently any task with a valid execution token can
+# access any asset's state — the same gap exists in /assets and /asset-events.
+# Proper fix is a unified asset-registration check across all asset routes,
+# not just here.
+router = VersionedAPIRouter(
+ responses={
+ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
+ status.HTTP_404_NOT_FOUND: {"description": "Not found"},
+ },
+)
+
+
+def _resolve_asset_id(name: str, session: SessionDep) -> int:
+ """Resolve asset name → integer asset_id, 404 if not found."""
+ asset_id = session.scalar(select(AssetModel.id).where(AssetModel.name == name))
Review Comment:
The other execution-API asset routes filter by `AssetModel.active.has()` when resolving by name (`assets.py:43,55`, `asset_events.py:89,91`). This one doesn't, so a row in `asset` that no longer has an `asset_active` partner (asset removed from its DAG but row still present) still resolves to an `asset_id` and lets workers read/write state against it. Worth matching the existing pattern.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state.py:
##########
@@ -0,0 +1,127 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Annotated
+from uuid import UUID
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, Query, Security, status
+from sqlalchemy.orm import Session
+
+from airflow._shared.state import TaskScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.task_state import (
+ TaskStatePutBody,
+ TaskStateResponse,
+)
+from airflow.api_fastapi.execution_api.security import require_auth
+from airflow.models.taskinstance import TaskInstance as TI
+from airflow.state import get_state_backend
+
+router = VersionedAPIRouter(
+ responses={
+ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
+ status.HTTP_403_FORBIDDEN: {"description": "Access denied"},
+ status.HTTP_404_NOT_FOUND: {"description": "Not found"},
+ },
+ dependencies=[Security(require_auth, scopes=["ti:self"])],
+)
+
+
+def _get_task_scope_for_ti(task_instance_id: UUID, session: Session) -> TaskScope:
+ ti = session.get(TI, task_instance_id)
+ if ti is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": f"Task instance {task_instance_id} not found",
+ },
+ )
+ return TaskScope(dag_id=ti.dag_id, run_id=ti.run_id, task_id=ti.task_id, map_index=ti.map_index)
+
+
+@router.get("/{task_instance_id}/{key}")
+def get_task_state(
+ task_instance_id: UUID,
+ key: Annotated[str, Path(min_length=1)],
+ session: SessionDep,
+) -> TaskStateResponse:
+ """Get value for a task state."""
+ scope = _get_task_scope_for_ti(task_instance_id, session)
+ value = get_state_backend().get(scope, key)
+ if value is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": f"Task state key {key!r} not found",
+ },
+ )
+ return TaskStateResponse(value=value)
+
+
+@router.put("/{task_instance_id}/{key}", status_code=status.HTTP_204_NO_CONTENT)
+def set_task_state(
+ task_instance_id: UUID,
+ key: Annotated[str, Path(min_length=1)],
+ body: TaskStatePutBody,
+ session: SessionDep,
+) -> None:
+ """Set a task state key, creating or updating the row."""
+ scope = _get_task_scope_for_ti(task_instance_id, session)
+ get_state_backend().set(scope, key, body.value)
+
+
+@router.delete("/{task_instance_id}/{key}", status_code=status.HTTP_204_NO_CONTENT)
+def delete_task_state(
+ task_instance_id: UUID,
+ key: Annotated[str, Path(min_length=1)],
+ session: SessionDep,
+) -> None:
+ """Delete a single task state key."""
+ scope = _get_task_scope_for_ti(task_instance_id, session)
+ get_state_backend().delete(scope, key)
+
+
+@router.delete("/{task_instance_id}", status_code=status.HTTP_204_NO_CONTENT)
+def clear_task_state(
+ task_instance_id: UUID,
+ session: SessionDep,
+ all_map_indices: Annotated[bool, Query()] = False,
Review Comment:
Router-level `Security(require_auth, scopes=["ti:self"])` confirms the JWT subject matches `path_params["task_instance_id"]`, but `?all_map_indices=true` then wipes state for every mapped sibling, which are different TIs with different tokens. A worker for `map_index=0` can clear state for `map_index=1..N`. Either the docstring should declare this as intentionally fleet-wide (and still safe because the caller passed an authenticated TI of the same task), or this branch needs a stronger scope (say `ti:fleet`) or a leader check (`map_index in (-1, 0)`).
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state.py:
##########
@@ -0,0 +1,127 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Annotated
+from uuid import UUID
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, Query, Security, status
+from sqlalchemy.orm import Session
+
+from airflow._shared.state import TaskScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.task_state import (
+ TaskStatePutBody,
+ TaskStateResponse,
+)
+from airflow.api_fastapi.execution_api.security import require_auth
+from airflow.models.taskinstance import TaskInstance as TI
+from airflow.state import get_state_backend
+
+router = VersionedAPIRouter(
+ responses={
+ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
+ status.HTTP_403_FORBIDDEN: {"description": "Access denied"},
+ status.HTTP_404_NOT_FOUND: {"description": "Not found"},
+ },
+ dependencies=[Security(require_auth, scopes=["ti:self"])],
+)
+
+
+def _get_task_scope_for_ti(task_instance_id: UUID, session: Session) -> TaskScope:
+ ti = session.get(TI, task_instance_id)
+ if ti is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": f"Task instance {task_instance_id} not found",
+ },
+ )
+ return TaskScope(dag_id=ti.dag_id, run_id=ti.run_id, task_id=ti.task_id, map_index=ti.map_index)
+
+
+@router.get("/{task_instance_id}/{key}")
+def get_task_state(
+ task_instance_id: UUID,
+ key: Annotated[str, Path(min_length=1)],
+ session: SessionDep,
+) -> TaskStateResponse:
+ """Get value for a task state."""
+ scope = _get_task_scope_for_ti(task_instance_id, session)
+ value = get_state_backend().get(scope, key)
+ if value is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": f"Task state key {key!r} not found",
+ },
+ )
+ return TaskStateResponse(value=value)
+
+
+@router.put("/{task_instance_id}/{key}", status_code=status.HTTP_204_NO_CONTENT)
+def set_task_state(
+ task_instance_id: UUID,
+ key: Annotated[str, Path(min_length=1)],
+ body: TaskStatePutBody,
+ session: SessionDep,
+) -> None:
+ """Set a task state key, creating or updating the row."""
+ scope = _get_task_scope_for_ti(task_instance_id, session)
+ get_state_backend().set(scope, key, body.value)
Review Comment:
`_get_task_scope_for_ti` uses the route's `SessionDep`, but `get_state_backend().set(...)` opens a fresh session via `@provide_session` on `MetastoreStateBackend.set` (`state/metastore.py:81-82`). Two consequences: (a) existence check and write run in different transactions, so the TI can be deleted or its `run_id` changed between them, leading to an orphan row or FK error, and (b) you pay an extra connection per request. Passing `session=session` into the backend call closes both. Same shape on `delete_task_state`, `clear_task_state`, and all four asset_state routes.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/__init__.py:
##########
@@ -52,5 +54,7 @@
authenticated_router.include_router(variables.router, prefix="/variables", tags=["Variables"])
authenticated_router.include_router(xcoms.router, prefix="/xcoms", tags=["XComs"])
authenticated_router.include_router(hitl.router, prefix="/hitlDetails", tags=["Human in the Loop"])
+authenticated_router.include_router(task_state.router, prefix="/state/ti", tags=["Task State"])
+authenticated_router.include_router(asset_state.router, prefix="/state/asset", tags=["Asset State"])
Review Comment:
8 new endpoints land here without a corresponding `endpoint(...).didnt_exist` entry in the unreleased version file (`versions/v2026_04_17.py`). A newer SDK calling these routes against an older API server won't fail cleanly through Cadwyn's version negotiation. See `contributing-docs/19_execution_api_versioning.rst`.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state.py:
##########
@@ -0,0 +1,127 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Annotated
+from uuid import UUID
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, Query, Security, status
+from sqlalchemy.orm import Session
+
+from airflow._shared.state import TaskScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.task_state import (
+ TaskStatePutBody,
+ TaskStateResponse,
+)
+from airflow.api_fastapi.execution_api.security import require_auth
+from airflow.models.taskinstance import TaskInstance as TI
+from airflow.state import get_state_backend
+
+router = VersionedAPIRouter(
+ responses={
+ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
+ status.HTTP_403_FORBIDDEN: {"description": "Access denied"},
+ status.HTTP_404_NOT_FOUND: {"description": "Not found"},
+ },
+ dependencies=[Security(require_auth, scopes=["ti:self"])],
Review Comment:
The new `ti:self` enforcement is the primary protection on these routes, but `airflow-core/tests/unit/api_fastapi/execution_api/conftest.py:56-60` overrides `require_auth` with a mock that always builds a `TIToken` whose `id` equals `path_params["task_instance_id"]`, so the mismatch branch in `security.py:185-191` is never exercised here. Worth a regression test that uses the real `require_auth` and asserts a token bound to a different TI gets a 403 against this router. Same gap means there's currently no asset_state coverage at all for the auth model that's planned to land via the AIP-93 follow-up.
##########
airflow-core/src/airflow/state/__init__.py:
##########
@@ -36,3 +36,14 @@ def resolve_state_backend() -> type[BaseStateBackend]:
f"Your custom state backend `{clazz.__name__}` is not a subclass of `BaseStateBackend`."
)
return clazz
+
+
+_backend_instance: BaseStateBackend | None = None
+
+
+def get_state_backend() -> BaseStateBackend:
+ """Return a cached instance of the configured state backend."""
+ global _backend_instance
+ if _backend_instance is None:
+ _backend_instance = resolve_state_backend()()
Review Comment:
Two threads racing through `if _backend_instance is None` can each call `resolve_state_backend()()`; whichever assigns last wins, the other gets GC'd. Mostly harmless for `MetastoreStateBackend`, but a custom backend with constructor side effects (file handle, network client, registered metrics, schema-creation hook) would double-init under load. A small `threading.Lock`, or wrapping a `functools.cache`'d helper, closes it.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_state.py:
##########
@@ -0,0 +1,121 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Execution API routes for asset state.
+
+Asset state is keyed by asset *name* (not integer id) in the URL — asset names
+are unique, and callers (task SDK accessors) have the name from their Asset
+object without needing a DB lookup. The route resolves name → asset_id
+internally for the state backend scope.
+
+Per-task asset registration checks are intentionally not implemented here
+(deferred to AIP-93 — see TODO comment below).
+"""
+
+from __future__ import annotations
+
+from typing import Annotated
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Path, status
+from sqlalchemy import select
+
+from airflow._shared.state import AssetScope
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.asset_state import (
+ AssetStatePutBody,
+ AssetStateResponse,
+)
+from airflow.models.asset import AssetModel
+from airflow.state import get_state_backend
+
+# TODO(AIP-103): enforce that the requesting task is registered with the asset
+# (via task_inlet_asset_reference or task_outlet_asset_reference) before
+# allowing reads/writes. Currently any task with a valid execution token can
+# access any asset's state — the same gap exists in /assets and /asset-events.
+# Proper fix is a unified asset-registration check across all asset routes,
+# not just here.
+router = VersionedAPIRouter(
+ responses={
+ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
+ status.HTTP_404_NOT_FOUND: {"description": "Not found"},
+ },
+)
+
+
+def _resolve_asset_id(name: str, session: SessionDep) -> int:
+ """Resolve asset name → integer asset_id, 404 if not found."""
+ asset_id = session.scalar(select(AssetModel.id).where(AssetModel.name == name))
+ if asset_id is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={"reason": "not_found", "message": f"Asset {name!r} not found"},
+ )
+ return asset_id
+
+
+@router.get("/{name}/{key}")
+def get_asset_state(
+ name: Annotated[str, Path(min_length=1)],
+ key: Annotated[str, Path(min_length=1)],
+ session: SessionDep,
+) -> AssetStateResponse:
+ """Get an asset state value."""
+ asset_id = _resolve_asset_id(name, session)
+ value = get_state_backend().get(AssetScope(asset_id=asset_id), key)
+ if value is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": f"Asset state key {key!r} not found",
+ },
+ )
+ return AssetStateResponse(value=value)
+
+
+@router.put("/{name}/{key}", status_code=status.HTTP_204_NO_CONTENT)
+def set_asset_state(
+ name: Annotated[str, Path(min_length=1)],
+ key: Annotated[str, Path(min_length=1)],
+ body: AssetStatePutBody,
+ session: SessionDep,
+) -> None:
+ """Set an asset state value."""
+ asset_id = _resolve_asset_id(name, session)
+ get_state_backend().set(AssetScope(asset_id=asset_id), key, body.value)
Review Comment:
Mirrors the task_state comment: `_resolve_asset_id` uses the route's `SessionDep`, then `get_state_backend().set(...)` opens a fresh session via `@provide_session`. The PUT/DELETE/clear handlers (lines 100, 110, 120) all share this shape. Passing `session=session` into the backend call keeps the lookup and the write in one transaction.
##########
airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_state.py:
##########
@@ -0,0 +1,238 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+from uuid import uuid4
+
+import pytest
+from sqlalchemy import delete, select
+
+from airflow._shared.timezones import timezone
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import create_session
+
+if TYPE_CHECKING:
+ from fastapi.testclient import TestClient
+
+ from tests_common.pytest_plugin import CreateTaskInstance
+
+pytestmark = pytest.mark.db_test
+
+
+@pytest.fixture(autouse=True)
+def reset_state_tables():
+ with create_session() as session:
+ session.execute(delete(TaskStateModel))
+ session.execute(delete(DagRun))
+
+
+def _api_url(ti_id, key: str | None = None) -> str:
+ base = f"/execution/state/ti/{ti_id}"
+ return f"{base}/{key}" if key else base
+
+
+class TestGetTaskState:
+ def test_get_returns_value(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+ client.put(_api_url(ti.id, "job_id"), json={"value": "spark_001"})
+
+ response = client.get(_api_url(ti.id, "job_id"))
+
+ assert response.status_code == 200
+ assert response.json() == {"value": "spark_001"}
+
+ def test_get_missing_key_returns_404(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+
+ response = client.get(_api_url(ti.id, "never_set"))
+
+ assert response.status_code == 404
+ assert response.json()["detail"]["reason"] == "not_found"
+
+ def test_get_missing_ti_returns_404(self, client: TestClient):
+ response = client.get(_api_url(uuid4(), "any_key"))
+
+ assert response.status_code == 404
+ assert "Task instance" in response.json()["detail"]["message"]
+
+
+class TestPutTaskState:
+ def test_put_creates_row(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+
+ response = client.put(_api_url(ti.id, "job_id"), json={"value": "spark_001"})
+
+ assert response.status_code == 204
+ with create_session() as session:
+ row = session.scalar(
+ select(TaskStateModel).where(
+ TaskStateModel.dag_id == ti.dag_id,
+ TaskStateModel.run_id == ti.run_id,
+ TaskStateModel.task_id == ti.task_id,
+ TaskStateModel.key == "job_id",
+ )
+ )
+ assert row is not None
+ assert row.value == "spark_001"
+
+ def test_put_overwrites_existing(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+ client.put(_api_url(ti.id, "job_id"), json={"value": "spark_001"})
+
+ response = client.put(_api_url(ti.id, "job_id"), json={"value": "spark_002"})
+
+ assert response.status_code == 204
+ assert client.get(_api_url(ti.id, "job_id")).json() == {"value": "spark_002"}
+
+ def test_put_empty_body_returns_422(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+
+ response = client.put(_api_url(ti.id, "job_id"), json={})
+
+ assert response.status_code == 422
+
+ def test_put_extra_field_returns_422(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+
+ response = client.put(_api_url(ti.id, "job_id"), json={"value": "x", "extra": "y"})
+
+ assert response.status_code == 422
+
+ def test_put_null_value_returns_422(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+
+ response = client.put(_api_url(ti.id, "job_id"), json={"value": None})
+
+ assert response.status_code == 422
+
+ def test_put_missing_ti_returns_404(self, client: TestClient):
+ response = client.put(_api_url(uuid4(), "job_id"), json={"value": "x"})
+
+ assert response.status_code == 404
+
+
+class TestDeleteTaskState:
+ def test_delete_removes_key(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+ client.put(_api_url(ti.id, "job_id"), json={"value": "spark_001"})
+
+ response = client.delete(_api_url(ti.id, "job_id"))
+
+ assert response.status_code == 204
+ assert client.get(_api_url(ti.id, "job_id")).status_code == 404
+
+ def test_delete_missing_key_is_noop(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+
+ response = client.delete(_api_url(ti.id, "never_existed"))
+
+ assert response.status_code == 204
+
+ def test_delete_only_targets_one_key(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+ client.put(_api_url(ti.id, "job_id"), json={"value": "a"})
+ client.put(_api_url(ti.id, "checkpoint"), json={"value": "b"})
+
+ client.delete(_api_url(ti.id, "job_id"))
+
+ assert client.get(_api_url(ti.id, "job_id")).status_code == 404
+ assert client.get(_api_url(ti.id, "checkpoint")).json() == {"value": "b"}
+
+
+class TestClearTaskState:
+ def test_clear_removes_all_keys(self, client: TestClient, create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+ for k, v in [("job_id", "a"), ("checkpoint", "b"), ("retry_count", "c")]:
+ client.put(_api_url(ti.id, k), json={"value": v})
+
+ response = client.delete(_api_url(ti.id))
+
+ assert response.status_code == 204
+ with create_session() as session:
+ count = session.scalar(
+ select(TaskStateModel)
+ .where(TaskStateModel.dag_id == ti.dag_id, TaskStateModel.task_id == ti.task_id)
+ .with_only_columns(TaskStateModel.key)
+ )
+ assert count is None
Review Comment:
`count` is `session.scalar(select(...).with_only_columns(TaskStateModel.key))`, so it returns the first matching `key` value (or `None` when empty), not a count. The test passes only because the table is empty after `clear`; it would also pass spuriously if a row's `key` were ever `None`, and it'd fail with a confusing message on any non-empty wrong result. `session.scalar(select(func.count()).select_from(TaskStateModel).where(...))` plus `assert count == 0` matches the name and the intent.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org