Skip to content

Commit ffb1b8a

Browse files
amoghrajeshamoghrajesh
authored
AIP-103: Add Execution API endpoints for task and asset states (#66073)
1 parent f90f166 commit ffb1b8a

12 files changed

Lines changed: 857 additions & 3 deletions

File tree

‎airflow-core/src/airflow/api_fastapi/execution_api/datamodels/asset_state.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
18+
from __future__ import annotations
19+
20+
from airflow.api_fastapi.core_api.base import StrictBaseModel
21+
22+
23+
class AssetStateResponse(StrictBaseModel):
24+
"""Asset state value returned to a worker."""
25+
26+
value: str
27+
28+
29+
class AssetStatePutBody(StrictBaseModel):
30+
"""Request body for setting an asset state value."""
31+
32+
value: str

‎airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_state.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
18+
from __future__ import annotations
19+
20+
from airflow.api_fastapi.core_api.base import StrictBaseModel
21+
22+
23+
class TaskStateResponse(StrictBaseModel):
24+
"""Task state value returned to a worker."""
25+
26+
value: str
27+
28+
29+
class TaskStatePutBody(StrictBaseModel):
30+
"""Request body for setting a task state value."""
31+
32+
value: str

‎airflow-core/src/airflow/api_fastapi/execution_api/routes/__init__.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121

2222
from airflow.api_fastapi.execution_api.routes import (
2323
asset_events,
24+
asset_state,
2425
assets,
2526
connections,
2627
dag_runs,
@@ -29,6 +30,7 @@
2930
hitl,
3031
task_instances,
3132
task_reschedules,
33+
task_state,
3234
variables,
3335
xcoms,
3436
)
@@ -52,5 +54,7 @@
5254
authenticated_router.include_router(variables.router, prefix="/variables", tags=["Variables"])
5355
authenticated_router.include_router(xcoms.router, prefix="/xcoms", tags=["XComs"])
5456
authenticated_router.include_router(hitl.router, prefix="/hitlDetails", tags=["Human in the Loop"])
57+
authenticated_router.include_router(task_state.router, prefix="/state/ti", tags=["Task State"])
58+
authenticated_router.include_router(asset_state.router, prefix="/state/asset", tags=["Asset State"])
5559

5660
execution_api_router.include_router(authenticated_router)

‎airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_state.py‎

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
"""
18+
Execution API routes for asset state.
19+
20+
Asset state is keyed by asset *name* (not integer id) in the URL — asset names
21+
are unique, and callers (task SDK accessors) have the name from their Asset
22+
object without needing a DB lookup. The route resolves name → asset_id
23+
internally for the state backend scope.
24+
25+
Per-task asset registration checks are intentionally not implemented here
26+
(deferred to AIP-93 — see TODO comment below).
27+
"""
28+
29+
from __future__ import annotations
30+
31+
from typing import Annotated
32+
33+
from cadwyn import VersionedAPIRouter
34+
from fastapi import HTTPException, Query, status
35+
from sqlalchemy import select
36+
37+
from airflow._shared.state import AssetScope
38+
from airflow.api_fastapi.common.db.common import SessionDep
39+
from airflow.api_fastapi.execution_api.datamodels.asset_state import (
40+
AssetStatePutBody,
41+
AssetStateResponse,
42+
)
43+
from airflow.api_fastapi.execution_api.security import ExecutionAPIRoute
44+
from airflow.models.asset import AssetModel
45+
from airflow.state import get_state_backend
46+
47+
# TODO(AIP-103): enforce that the requesting task is registered with the asset
48+
# (via task_inlet_asset_reference or task_outlet_asset_reference) before
49+
# allowing reads/writes. Currently any task with a valid execution token can
50+
# access any asset's state — the same gap exists in /assets and /asset-events.
51+
# Proper fix is a unified asset-registration check across all asset routes,
52+
# not just here.
53+
router = VersionedAPIRouter(
54+
route_class=ExecutionAPIRoute,
55+
responses={
56+
status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
57+
status.HTTP_404_NOT_FOUND: {"description": "Not found"},
58+
},
59+
)
60+
61+
62+
def _resolve_asset_id(name: str, session: SessionDep) -> int:
63+
"""Resolve asset name → integer asset_id, 404 if not found."""
64+
asset_id = session.scalar(select(AssetModel.id).where(AssetModel.name == name, AssetModel.active.has()))
65+
if asset_id is None:
66+
raise HTTPException(
67+
status_code=status.HTTP_404_NOT_FOUND,
68+
detail={"reason": "not_found", "message": f"Asset {name!r} not found"},
69+
)
70+
return asset_id
71+
72+
73+
@router.get("/value")
74+
def get_asset_state(
75+
name: Annotated[str, Query(min_length=1)],
76+
key: Annotated[str, Query(min_length=1)],
77+
session: SessionDep,
78+
) -> AssetStateResponse:
79+
"""Get an asset state value."""
80+
asset_id = _resolve_asset_id(name, session)
81+
value = get_state_backend().get(AssetScope(asset_id=asset_id), key, session=session) # type: ignore[call-arg] # @provide_session adds session kwarg at runtime; BaseStateBackend signature omits it so mypy can't see it
82+
if value is None:
83+
raise HTTPException(
84+
status_code=status.HTTP_404_NOT_FOUND,
85+
detail={
86+
"reason": "not_found",
87+
"message": f"Asset state key {key!r} not found",
88+
},
89+
)
90+
return AssetStateResponse(value=value)
91+
92+
93+
@router.put("/value", status_code=status.HTTP_204_NO_CONTENT)
94+
def set_asset_state(
95+
name: Annotated[str, Query(min_length=1)],
96+
key: Annotated[str, Query(min_length=1)],
97+
body: AssetStatePutBody,
98+
session: SessionDep,
99+
) -> None:
100+
"""Set an asset state value."""
101+
asset_id = _resolve_asset_id(name, session)
102+
get_state_backend().set(AssetScope(asset_id=asset_id), key, body.value, session=session) # type: ignore[call-arg] # @provide_session adds session kwarg at runtime; BaseStateBackend signature omits it so mypy can't see it
103+
104+
105+
@router.delete("/value", status_code=status.HTTP_204_NO_CONTENT)
106+
def delete_asset_state(
107+
name: Annotated[str, Query(min_length=1)],
108+
key: Annotated[str, Query(min_length=1)],
109+
session: SessionDep,
110+
) -> None:
111+
"""Delete a single asset state key."""
112+
asset_id = _resolve_asset_id(name, session)
113+
get_state_backend().delete(AssetScope(asset_id=asset_id), key, session=session) # type: ignore[call-arg] # @provide_session adds session kwarg at runtime; BaseStateBackend signature omits it so mypy can't see it
114+
115+
116+
@router.delete("/clear", status_code=status.HTTP_204_NO_CONTENT)
117+
def clear_asset_state(
118+
name: Annotated[str, Query(min_length=1)],
119+
session: SessionDep,
120+
) -> None:
121+
"""Delete all state keys for an asset."""
122+
asset_id = _resolve_asset_id(name, session)
123+
get_state_backend().clear(AssetScope(asset_id=asset_id), session=session) # type: ignore[call-arg] # @provide_session adds session kwarg at runtime; BaseStateBackend signature omits it so mypy can't see it

‎airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state.py‎

Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
from __future__ import annotations
18+
19+
from typing import Annotated
20+
from uuid import UUID
21+
22+
from cadwyn import VersionedAPIRouter
23+
from fastapi import HTTPException, Path, Query, Security, status
24+
from sqlalchemy.orm import Session
25+
26+
from airflow._shared.state import TaskScope
27+
from airflow.api_fastapi.common.db.common import SessionDep
28+
from airflow.api_fastapi.execution_api.datamodels.task_state import (
29+
TaskStatePutBody,
30+
TaskStateResponse,
31+
)
32+
from airflow.api_fastapi.execution_api.security import ExecutionAPIRoute, require_auth
33+
from airflow.models.taskinstance import TaskInstance as TI
34+
from airflow.state import get_state_backend
35+
36+
router = VersionedAPIRouter(
37+
route_class=ExecutionAPIRoute,
38+
responses={
39+
status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
40+
status.HTTP_403_FORBIDDEN: {"description": "Access denied"},
41+
status.HTTP_404_NOT_FOUND: {"description": "Not found"},
42+
},
43+
dependencies=[Security(require_auth, scopes=["ti:self"])],
44+
)
45+
46+
47+
def _get_task_scope_for_ti(task_instance_id: UUID, session: Session) -> TaskScope:
48+
ti = session.get(TI, task_instance_id)
49+
if ti is None:
50+
raise HTTPException(
51+
status_code=status.HTTP_404_NOT_FOUND,
52+
detail={
53+
"reason": "not_found",
54+
"message": f"Task instance {task_instance_id} not found",
55+
},
56+
)
57+
return TaskScope(dag_id=ti.dag_id, run_id=ti.run_id, task_id=ti.task_id, map_index=ti.map_index)
58+
59+
60+
@router.get("/{task_instance_id}/{key}")
61+
def get_task_state(
62+
task_instance_id: UUID,
63+
key: Annotated[str, Path(min_length=1)],
64+
session: SessionDep,
65+
) -> TaskStateResponse:
66+
"""Get value for a task state."""
67+
scope = _get_task_scope_for_ti(task_instance_id, session)
68+
value = get_state_backend().get(scope, key, session=session) # type: ignore[call-arg] # @provide_session adds session kwarg at runtime; BaseStateBackend signature omits it so mypy can't see it
69+
if value is None:
70+
raise HTTPException(
71+
status_code=status.HTTP_404_NOT_FOUND,
72+
detail={
73+
"reason": "not_found",
74+
"message": f"Task state key {key!r} not found",
75+
},
76+
)
77+
return TaskStateResponse(value=value)
78+
79+
80+
@router.put("/{task_instance_id}/{key}", status_code=status.HTTP_204_NO_CONTENT)
81+
def set_task_state(
82+
task_instance_id: UUID,
83+
key: Annotated[str, Path(min_length=1)],
84+
body: TaskStatePutBody,
85+
session: SessionDep,
86+
) -> None:
87+
"""Set a task state key, creating or updating the row."""
88+
scope = _get_task_scope_for_ti(task_instance_id, session)
89+
get_state_backend().set(scope, key, body.value, session=session) # type: ignore[call-arg] # @provide_session adds session kwarg at runtime; BaseStateBackend signature omits it so mypy can't see it
90+
91+
92+
@router.delete("/{task_instance_id}/{key}", status_code=status.HTTP_204_NO_CONTENT)
93+
def delete_task_state(
94+
task_instance_id: UUID,
95+
key: Annotated[str, Path(min_length=1)],
96+
session: SessionDep,
97+
) -> None:
98+
"""Delete a single task state key."""
99+
scope = _get_task_scope_for_ti(task_instance_id, session)
100+
get_state_backend().delete(scope, key, session=session) # type: ignore[call-arg] # @provide_session adds session kwarg at runtime; BaseStateBackend signature omits it so mypy can't see it
101+
102+
103+
@router.delete("/{task_instance_id}", status_code=status.HTTP_204_NO_CONTENT)
104+
def clear_task_state(
105+
task_instance_id: UUID,
106+
session: SessionDep,
107+
all_map_indices: Annotated[bool, Query()] = False,
108+
) -> None:
109+
"""
110+
Delete all task state keys for this task instance.
111+
112+
By default, only keys for the requesting TI's exact ``map_index`` are
113+
cleared — same isolation as DELETE endpoint above.
114+
115+
Pass ``?all_map_indices=true`` to wipe state for every mapped sibling of
116+
the task in the same DAG run. This is intentionally fleet-wide: the
117+
``ti:self`` JWT authentication scope authenticates that the caller is
118+
a legitimate member of the mapped task group, and grants it authority
119+
to reset shared task state on behalf of the whole group.
120+
The SDK only forwards this flag when the user calls ``task_state.clear(all_map_indices=True)``
121+
explicitly, so the expanded scope is always an explicit opt-in by the task author.
122+
123+
For non-mapped tasks (``map_index=-1``), there is only ever one index, so
124+
``?all_map_indices=true`` is functionally identical to the default and is
125+
accepted without error.
126+
"""
127+
scope = _get_task_scope_for_ti(task_instance_id, session)
128+
get_state_backend().clear(scope, all_map_indices=all_map_indices, session=session) # type: ignore[call-arg] # @provide_session adds session kwarg at runtime; BaseStateBackend signature omits it so mypy can't see it

‎airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@
4040
MovePreviousRunEndpoint,
4141
RemoveUpstreamMapIndexesField,
4242
)
43-
from airflow.api_fastapi.execution_api.versions.v2026_04_17 import AddTeamNameField
43+
from airflow.api_fastapi.execution_api.versions.v2026_04_17 import AddStateEndpoints, AddTeamNameField
4444
from airflow.api_fastapi.execution_api.versions.v2026_06_16 import AddRetryPolicyFields
4545

4646
bundle = VersionBundle(
@@ -49,6 +49,7 @@
4949
Version(
5050
"2026-04-17",
5151
AddTeamNameField,
52+
AddStateEndpoints,
5253
),
5354
Version(
5455
"2026-04-06",

‎airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_04_17.py‎

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717

1818
from __future__ import annotations
1919

20-
from cadwyn import ResponseInfo, VersionChange, convert_response_to_previous_version_for, schema
20+
from cadwyn import ResponseInfo, VersionChange, convert_response_to_previous_version_for, endpoint, schema
2121

2222
from airflow.api_fastapi.execution_api.datamodels.taskinstance import DagRun, TIRunContext
2323

@@ -34,3 +34,20 @@ def remove_team_name_field(response: ResponseInfo) -> None: # type: ignore[misc
3434
"""Remove the ``team_name`` field from dag_run for older API versions."""
3535
if "dag_run" in response.body and isinstance(response.body["dag_run"], dict):
3636
response.body["dag_run"].pop("team_name", None)
37+
38+
39+
class AddStateEndpoints(VersionChange):
40+
"""Add task state and asset state CRUD endpoints."""
41+
42+
description = __doc__
43+
44+
instructions_to_migrate_to_previous_version = (
45+
endpoint("/state/ti/{task_instance_id}/{key}", ["GET"]).didnt_exist,
46+
endpoint("/state/ti/{task_instance_id}/{key}", ["PUT"]).didnt_exist,
47+
endpoint("/state/ti/{task_instance_id}/{key}", ["DELETE"]).didnt_exist,
48+
endpoint("/state/ti/{task_instance_id}", ["DELETE"]).didnt_exist,
49+
endpoint("/state/asset/value", ["GET"]).didnt_exist,
50+
endpoint("/state/asset/value", ["PUT"]).didnt_exist,
51+
endpoint("/state/asset/value", ["DELETE"]).didnt_exist,
52+
endpoint("/state/asset/clear", ["DELETE"]).didnt_exist,
53+
)

‎airflow-core/src/airflow/migrations/versions/0113_3_3_0_add_retry_policy_fields_to_ti.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@
3030
is a metadata-only operation (no table rewrite).
3131
3232
Revision ID: b8f3e4a1d2c9
33-
Revises: 9fabad868fdb
33+
Revises: fde9ed84d07b
3434
Create Date: 2026-04-16 12:00:00.000000
3535
3636
"""

0 commit comments

Comments
 (0)