From 29fa162353ef2fddb33ba237a55884ae919696b9 Mon Sep 17 00:00:00 2001 From: Vincent Hsiao <124506982+fat-catTW@users.noreply.github.com> Date: Sat, 11 Jul 2026 16:28:25 +0000 Subject: [PATCH] Restrict asset state store access to registered tasks --- .../execution_api/routes/asset_state_store.py | 78 +++++++++--- .../versions/head/test_asset_state_store.py | 115 +++++++++++++++++- 2 files changed, 175 insertions(+), 18 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_state_store.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_state_store.py index 4424e4ad3f180..91416838701c3 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_state_store.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_state_store.py @@ -22,8 +22,9 @@ whichever identifier their inlet type carries: ``Asset``/``AssetNameRef`` use the name routes, ``AssetUriRef`` uses the URI routes. -Per-task asset registration checks are intentionally not implemented here -(deferred to AIP-93 — see TODO comment below). +Task tokens must be registered with an asset as an inlet or outlet before +accessing that asset's state store. Watcher tokens are not backed by task +instances and are allowed through this route separately. """ from __future__ import annotations @@ -34,7 +35,7 @@ from cadwyn import VersionedAPIRouter from fastapi import HTTPException, Query, status -from sqlalchemy import select +from sqlalchemy import exists, or_, select from airflow._shared.state import AssetScope, AssetStateStoreWriterKind from airflow.api_fastapi.common.db.common import SessionDep @@ -44,7 +45,7 @@ ) from airflow.api_fastapi.execution_api.datamodels.token import TIToken from airflow.api_fastapi.execution_api.security import CurrentTIToken, ExecutionAPIRoute -from airflow.models.asset import AssetModel +from airflow.models.asset import AssetModel, TaskInletAssetReference, TaskOutletAssetReference from airflow.models.taskinstance import TaskInstance from airflow.state import get_state_backend from airflow.state.metastore import MetastoreBackend @@ -71,12 +72,6 @@ def _fetch_ti_writer_fields(token: TIToken, session: SessionDep) -> _TIWriterFie return row.dag_id, row.run_id, row.task_id, row.map_index -# 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 store — 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( route_class=ExecutionAPIRoute, responses={ @@ -106,14 +101,51 @@ def _resolve_asset_id_by_uri(uri: str, session: SessionDep) -> int: return asset_id +def _authorize_asset_state_store_access( + token: TIToken, asset_id: int, session: SessionDep +) -> _TIWriterFields | None: + if token.id == NULL_UUID: + return None + + ti_fields = _fetch_ti_writer_fields(token, session) + dag_id, _, task_id, _ = ti_fields + is_registered = session.scalar( + select( + or_( + exists().where( + TaskInletAssetReference.asset_id == asset_id, + TaskInletAssetReference.dag_id == dag_id, + TaskInletAssetReference.task_id == task_id, + ), + exists().where( + TaskOutletAssetReference.asset_id == asset_id, + TaskOutletAssetReference.dag_id == dag_id, + TaskOutletAssetReference.task_id == task_id, + ), + ) + ) + ) + if not is_registered: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "reason": "forbidden", + "message": "Task is not registered with this asset as an inlet or outlet", + }, + ) + return ti_fields + + @router.get("/by-name/value") def get_asset_state_store_by_name( name: Annotated[str, Query(min_length=1)], key: Annotated[str, Query(min_length=1)], session: SessionDep, + token: TIToken = CurrentTIToken, ) -> AssetStateStoreResponse: """Get an asset state store value by asset name.""" asset_id = _resolve_asset_id_by_name(name, session) + _authorize_asset_state_store_access(token, asset_id, session) value = get_state_backend().get(AssetScope(asset_id=asset_id), key, session=session) if value is None: raise HTTPException( @@ -129,6 +161,7 @@ def _put_asset_state_store( body: AssetStateStorePutBody, token: TIToken, session: SessionDep, + ti_fields: _TIWriterFields | None, ) -> None: backend = get_state_backend() if isinstance(backend, MetastoreBackend): @@ -142,7 +175,8 @@ def _put_asset_state_store( session=session, ) else: - ti_fields = _fetch_ti_writer_fields(token, session) + if ti_fields is None: + ti_fields = _fetch_ti_writer_fields(token, session) dag_id, run_id, task_id, map_index = ti_fields backend.set_asset_state_store( @@ -169,9 +203,9 @@ def set_asset_state_store_by_name( token: TIToken = CurrentTIToken, ) -> None: """Set an asset state store value by asset name.""" - _put_asset_state_store( - AssetScope(asset_id=_resolve_asset_id_by_name(name, session)), key, body, token, session - ) + asset_id = _resolve_asset_id_by_name(name, session) + ti_fields = _authorize_asset_state_store_access(token, asset_id, session) + _put_asset_state_store(AssetScope(asset_id=asset_id), key, body, token, session, ti_fields) @router.delete("/by-name/value", status_code=status.HTTP_204_NO_CONTENT) @@ -179,9 +213,11 @@ def delete_asset_state_store_by_name( name: Annotated[str, Query(min_length=1)], key: Annotated[str, Query(min_length=1)], session: SessionDep, + token: TIToken = CurrentTIToken, ) -> None: """Delete a single asset state store key by asset name.""" asset_id = _resolve_asset_id_by_name(name, session) + _authorize_asset_state_store_access(token, asset_id, session) get_state_backend().delete(AssetScope(asset_id=asset_id), key, session=session) @@ -189,9 +225,11 @@ def delete_asset_state_store_by_name( def clear_asset_state_store_by_name( name: Annotated[str, Query(min_length=1)], session: SessionDep, + token: TIToken = CurrentTIToken, ) -> None: """Delete all state store keys for an asset by asset name.""" asset_id = _resolve_asset_id_by_name(name, session) + _authorize_asset_state_store_access(token, asset_id, session) get_state_backend().clear(AssetScope(asset_id=asset_id), session=session) @@ -200,9 +238,11 @@ def get_asset_state_store_by_uri( uri: Annotated[str, Query(min_length=1)], key: Annotated[str, Query(min_length=1)], session: SessionDep, + token: TIToken = CurrentTIToken, ) -> AssetStateStoreResponse: """Get an asset state store value by asset URI.""" asset_id = _resolve_asset_id_by_uri(uri, session) + _authorize_asset_state_store_access(token, asset_id, session) value = get_state_backend().get(AssetScope(asset_id=asset_id), key, session=session) if value is None: raise HTTPException( @@ -221,9 +261,9 @@ def set_asset_state_store_by_uri( token: TIToken = CurrentTIToken, ) -> None: """Set an asset state store value by asset URI.""" - _put_asset_state_store( - AssetScope(asset_id=_resolve_asset_id_by_uri(uri, session)), key, body, token, session - ) + asset_id = _resolve_asset_id_by_uri(uri, session) + ti_fields = _authorize_asset_state_store_access(token, asset_id, session) + _put_asset_state_store(AssetScope(asset_id=asset_id), key, body, token, session, ti_fields) @router.delete("/by-uri/value", status_code=status.HTTP_204_NO_CONTENT) @@ -231,9 +271,11 @@ def delete_asset_state_store_by_uri( uri: Annotated[str, Query(min_length=1)], key: Annotated[str, Query(min_length=1)], session: SessionDep, + token: TIToken = CurrentTIToken, ) -> None: """Delete a single asset state store key by asset URI.""" asset_id = _resolve_asset_id_by_uri(uri, session) + _authorize_asset_state_store_access(token, asset_id, session) get_state_backend().delete(AssetScope(asset_id=asset_id), key, session=session) @@ -241,7 +283,9 @@ def delete_asset_state_store_by_uri( def clear_asset_state_store_by_uri( uri: Annotated[str, Query(min_length=1)], session: SessionDep, + token: TIToken = CurrentTIToken, ) -> None: """Delete all state store keys for an asset by asset URI.""" asset_id = _resolve_asset_id_by_uri(uri, session) + _authorize_asset_state_store_access(token, asset_id, session) get_state_backend().clear(AssetScope(asset_id=asset_id), session=session) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_asset_state_store.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_asset_state_store.py index 86a64b34a9aa2..620111ae7ceb6 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_asset_state_store.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_asset_state_store.py @@ -25,7 +25,7 @@ from airflow.api_fastapi.execution_api.datamodels.token import TIClaims, TIToken from airflow.api_fastapi.execution_api.security import require_auth -from airflow.models.asset import AssetActive, AssetModel +from airflow.models.asset import AssetActive, AssetModel, TaskInletAssetReference, TaskOutletAssetReference from airflow.models.asset_state_store import AssetStateStoreModel from airflow.utils.session import create_session @@ -49,6 +49,25 @@ def _create_asset_store_row(asset_id: int, key: str, value: str) -> None: s.add(AssetStateStoreModel(asset_id=asset_id, key=key, value=json.dumps(value))) +def _create_task_asset_reference( + session: Session, + task_instance: TaskInstance, + asset: AssetModel, + reference_model=TaskOutletAssetReference, +) -> None: + session.add( + reference_model(asset_id=asset.id, dag_id=task_instance.dag_id, task_id=task_instance.task_id) + ) + session.commit() + + +def _request_asset_state_store(client: TestClient, method: str, path: str, *, params, json_body=None): + kwargs = {"params": params} + if json_body is not None: + kwargs["json"] = json_body + return getattr(client, method)(path, **kwargs) + + _BY_NAME_CLEAR = "/execution/store/asset/by-name/clear" _BY_URI_VALUE = "/execution/store/asset/by-uri/value" _BY_URI_CLEAR = "/execution/store/asset/by-uri/clear" @@ -58,6 +77,8 @@ def _create_asset_store_row(asset_id: int, key: str, value: str) -> None: def reset_state_tables(): with create_session() as session: session.execute(delete(AssetStateStoreModel)) + session.execute(delete(TaskInletAssetReference)) + session.execute(delete(TaskOutletAssetReference)) session.execute(delete(AssetModel)) @@ -129,6 +150,10 @@ async def _auth(request: Request) -> TIToken: exec_app.dependency_overrides[require_auth] = _auth + @pytest.fixture(autouse=True) + def setup_asset_reference(self, setup_ti_auth, asset: AssetModel, session: Session): + _create_task_asset_reference(session, self._ti, asset) + def test_put_creates_row(self, client: TestClient, asset: AssetModel, session: Session): response = client.put( _BY_NAME_VALUE, params={"name": asset.name, "key": "watermark"}, json={"value": "2026-04-29"} @@ -302,6 +327,10 @@ async def _auth(request: Request) -> TIToken: exec_app.dependency_overrides[require_auth] = _auth + @pytest.fixture(autouse=True) + def setup_asset_reference(self, setup_ti_auth, asset: AssetModel, session: Session): + _create_task_asset_reference(session, self._ti, asset) + def test_put_creates_row(self, client: TestClient, asset: AssetModel, session: Session): response = client.put( _BY_URI_VALUE, params={"uri": asset.uri, "key": "watermark"}, json={"value": "2026-04-29"} @@ -370,6 +399,90 @@ def test_clear_removes_all_keys(self, client: TestClient, asset: AssetModel): assert row is None +class TestTaskAssetRegistration: + @pytest.fixture(autouse=True) + def setup_ti_auth( + self, + exec_app: FastAPI, + create_task_instance: CreateTaskInstance, + ): + self._ti: TaskInstance = create_task_instance() + + async def _auth(request: Request) -> TIToken: + return TIToken(id=self._ti.id, claims=TIClaims(scope="execution")) + + exec_app.dependency_overrides[require_auth] = _auth + + @pytest.mark.parametrize( + ("method", "path", "params", "json_body"), + [ + ("get", _BY_NAME_VALUE, {"name": "test_asset", "key": "watermark"}, None), + ("put", _BY_NAME_VALUE, {"name": "test_asset", "key": "watermark"}, {"value": "x"}), + ("delete", _BY_NAME_VALUE, {"name": "test_asset", "key": "watermark"}, None), + ("delete", _BY_NAME_CLEAR, {"name": "test_asset"}, None), + ("get", _BY_URI_VALUE, {"uri": "s3://bucket/test", "key": "watermark"}, None), + ("put", _BY_URI_VALUE, {"uri": "s3://bucket/test", "key": "watermark"}, {"value": "x"}), + ("delete", _BY_URI_VALUE, {"uri": "s3://bucket/test", "key": "watermark"}, None), + ("delete", _BY_URI_CLEAR, {"uri": "s3://bucket/test"}, None), + ], + ) + def test_unregistered_task_is_forbidden( + self, + client: TestClient, + asset: AssetModel, + method: str, + path: str, + params: dict, + json_body: dict | None, + ): + response = _request_asset_state_store(client, method, path, params=params, json_body=json_body) + + assert response.status_code == 403 + assert response.json()["detail"]["reason"] == "forbidden" + + @pytest.mark.parametrize( + ("method", "path", "params", "json_body", "expected_status"), + [ + ("get", _BY_NAME_VALUE, {"name": "test_asset", "key": "watermark"}, None, 200), + ("put", _BY_NAME_VALUE, {"name": "test_asset", "key": "watermark"}, {"value": "x"}, 204), + ("delete", _BY_NAME_VALUE, {"name": "test_asset", "key": "watermark"}, None, 204), + ("delete", _BY_NAME_CLEAR, {"name": "test_asset"}, None, 204), + ("get", _BY_URI_VALUE, {"uri": "s3://bucket/test", "key": "watermark"}, None, 200), + ("put", _BY_URI_VALUE, {"uri": "s3://bucket/test", "key": "watermark"}, {"value": "x"}, 204), + ("delete", _BY_URI_VALUE, {"uri": "s3://bucket/test", "key": "watermark"}, None, 204), + ("delete", _BY_URI_CLEAR, {"uri": "s3://bucket/test"}, None, 204), + ], + ) + def test_outlet_registered_task_can_access_asset_state_store( + self, + client: TestClient, + asset: AssetModel, + session: Session, + method: str, + path: str, + params: dict, + json_body: dict | None, + expected_status: int, + ): + _create_task_asset_reference(session, self._ti, asset) + _create_asset_store_row(asset.id, "watermark", "2026-04-29") + + response = _request_asset_state_store(client, method, path, params=params, json_body=json_body) + + assert response.status_code == expected_status + + def test_inlet_registered_task_can_access_asset_state_store( + self, client: TestClient, asset: AssetModel, session: Session + ): + _create_task_asset_reference(session, self._ti, asset, TaskInletAssetReference) + _create_asset_store_row(asset.id, "watermark", "2026-04-29") + + response = client.get(_BY_NAME_VALUE, params={"name": asset.name, "key": "watermark"}) + + assert response.status_code == 200 + assert response.json() == {"value": "2026-04-29"} + + class TestInactiveAssetRejected: """An asset row without a corresponding asset_active entry is treated as not found."""