This is an automated email from the ASF dual-hosted git repository. vatsrahul1001 pushed a commit to branch changes-3.4.0rc3 in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 232c15e9d0da3cbbb079fa0fb1a5d07301364814 Author: Rahul Vats <[email protected]> AuthorDate: Wed Nov 4 22:44:25 2026 +0530 Let asset watchers write asset state over the Execution API (#76376) (#76469) (cherry picked from commit 4ef770b91d699800d5b8af5fce72e541f659c67d) Co-authored-by: Jed Cunningham <[email protected]> Co-authored-by: Amogh Desai <[email protected]> Co-authored-by: Tzu-ping Chung <[email protected]> --- .../api_fastapi/execution_api/in_process.py | 3 +- .../execution_api/routes/asset_state_store.py | 18 +++-- .../airflow/api_fastapi/execution_api/security.py | 39 +++++++++-- .../execution_api/versions/v2026_10_30.py | 11 ++- .../api_fastapi/execution_api/test_security.py | 6 +- .../versions/head/test_asset_state_store.py | 80 ++++++++++++++++++---- 6 files changed, 128 insertions(+), 29 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/in_process.py b/airflow-core/src/airflow/api_fastapi/execution_api/in_process.py index 3807e3f7cc6..eb428196d43 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/in_process.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/in_process.py @@ -61,6 +61,7 @@ def _configure_in_process_auth(app: FastAPI, *, triggerer_id: int | None = None) from airflow.api_fastapi.execution_api.security import ( _IN_PROCESS_NON_TI_CALLER, _IN_PROCESS_TRIGGERER_ID, + ASSET_WATCHER_ID, _jwt_bearer, ) @@ -70,7 +71,7 @@ def _configure_in_process_auth(app: FastAPI, *, triggerer_id: int | None = None) ti_id = UUID( request.path_params.get("task_instance_id") or request.headers.get("X-Airflow-In-Process-Attempt-Id") - or "00000000-0000-0000-0000-000000000000" + or str(ASSET_WATCHER_ID) ) # Watchers and other trusted in-process callers may have no task identity. request.scope[_IN_PROCESS_NON_TI_CALLER] = ( 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 f5d387b495a..ef167c63d70 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 @@ -43,15 +43,18 @@ from airflow.api_fastapi.execution_api.datamodels.asset_state_store import ( AssetStateStoreResponse, ) from airflow.api_fastapi.execution_api.datamodels.token import TIToken -from airflow.api_fastapi.execution_api.security import CurrentTIToken, ExecutionAPIRoute +from airflow.api_fastapi.execution_api.security import ( + ASSET_WATCHER_ID, + CurrentTIToken, + ExecutionAPIRoute, + allow_asset_watcher, +) from airflow.models.asset import AssetModel from airflow.models.task_coordinates import public_map_index_expression from airflow.models.taskinstance import TaskInstance from airflow.state import get_state_backend from airflow.state.metastore import MetastoreBackend -NULL_UUID = UUID(int=0) - class _TIWriterFields(NamedTuple): dag_id: str @@ -147,8 +150,7 @@ def _put_asset_state_store( ) -> None: backend = get_state_backend() if isinstance(backend, MetastoreBackend): - if token.id == NULL_UUID: - # Since the asset state store routes do not have `task_instance_id` in their path params, the default kicks in which is"00000000-0000-0000-0000-000000000000" + if token.id == ASSET_WATCHER_ID: backend.set_asset_state_store( scope, key, @@ -179,6 +181,7 @@ def _put_asset_state_store( @router.put("/by-name/value", status_code=status.HTTP_204_NO_CONTENT) +@allow_asset_watcher def set_asset_state_store_by_name( name: Annotated[str, Query(min_length=1)], key: Annotated[str, Query(min_length=1)], @@ -193,6 +196,7 @@ def set_asset_state_store_by_name( @router.delete("/by-name/value", status_code=status.HTTP_204_NO_CONTENT) +@allow_asset_watcher def delete_asset_state_store_by_name( name: Annotated[str, Query(min_length=1)], key: Annotated[str, Query(min_length=1)], @@ -204,6 +208,7 @@ def delete_asset_state_store_by_name( @router.delete("/by-name/clear", status_code=status.HTTP_204_NO_CONTENT) +@allow_asset_watcher def clear_asset_state_store_by_name( name: Annotated[str, Query(min_length=1)], session: SessionDep, @@ -231,6 +236,7 @@ def get_asset_state_store_by_uri( @router.put("/by-uri/value", status_code=status.HTTP_204_NO_CONTENT) +@allow_asset_watcher def set_asset_state_store_by_uri( uri: Annotated[str, Query(min_length=1)], key: Annotated[str, Query(min_length=1)], @@ -245,6 +251,7 @@ def set_asset_state_store_by_uri( @router.delete("/by-uri/value", status_code=status.HTTP_204_NO_CONTENT) +@allow_asset_watcher def delete_asset_state_store_by_uri( uri: Annotated[str, Query(min_length=1)], key: Annotated[str, Query(min_length=1)], @@ -256,6 +263,7 @@ def delete_asset_state_store_by_uri( @router.delete("/by-uri/clear", status_code=status.HTTP_204_NO_CONTENT) +@allow_asset_watcher def clear_asset_state_store_by_uri( uri: Annotated[str, Query(min_length=1)], session: SessionDep, diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/security.py b/airflow-core/src/airflow/api_fastapi/execution_api/security.py index 09bb22f2910..023aedca6ec 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/security.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/security.py @@ -73,6 +73,7 @@ Why ``ExecutionAPIRoute`` is needed: from collections.abc import Callable from typing import Any, ParamSpec, TypeVar, get_args +from uuid import UUID import structlog import svcs @@ -99,10 +100,17 @@ _REQUEST_SCOPE_LIVE_ATTEMPT_KEY = "live_attempt_checked" _IN_PROCESS_NON_TI_CALLER = "airflow_in_process_non_ti_caller" _IN_PROCESS_TRIGGERER_ID = "airflow_in_process_triggerer_id" _SKIP_AUTO_TI_ATTEMPT_LIVE = "skip_auto_ti_attempt_live" -_ALLOW_NON_TI_WRITER = "allow_non_ti_writer" +_LIVENESS_EXEMPTION = "liveness_exemption" _P = ParamSpec("_P") _R = TypeVar("_R") +# A triggerer that reaches the Execution API over HTTP authenticates watchers +# with an execution token setting 'sub' to this value. Routes marked with +# allow_asset_watcher exempt it from the attempt-liveness check, so such a +# token can write and clear the state of any asset. The in-process Execution +# API also falls back to this ID for every caller that has no task identity. +ASSET_WATCHER_ID = UUID(int=0) + def skip_auto_ti_attempt_live(endpoint: Callable[_P, _R]) -> Callable[_P, _R]: """Mark a route that checks attempt liveness itself, so ``require_auth`` skips its automatic check.""" @@ -110,10 +118,28 @@ def skip_auto_ti_attempt_live(endpoint: Callable[_P, _R]) -> Callable[_P, _R]: return endpoint +def _mark_liveness_exemption( + endpoint: Callable[_P, _R], accepts: Callable[[TIToken], bool] +) -> Callable[_P, _R]: + """Let tokens accepted by ``accepts`` skip ``require_auth``'s automatic liveness check on this route.""" + setattr(endpoint, _LIVENESS_EXEMPTION, accepts) + return endpoint + + def allow_non_ti_writer(endpoint: Callable[_P, _R]) -> Callable[_P, _R]: """Mark a route whose data no task attempt owns, so ``require_auth`` skips its automatic check.""" - setattr(endpoint, _ALLOW_NON_TI_WRITER, True) - return endpoint + return _mark_liveness_exemption(endpoint, lambda _: True) + + +def allow_asset_watcher(endpoint: Callable[_P, _R]) -> Callable[_P, _R]: + """ + Mark a route that accepts mutations from an asset watcher token in place of a live attempt. + + Unlike ``allow_non_ti_writer``, other tokens without a live attempt are still rejected. The asset + state store's delete and clear handlers never read the token, so they would let any such token + wipe an asset's state. + """ + return _mark_liveness_exemption(endpoint, lambda token: token.id == ASSET_WATCHER_ID) class JWTBearer(HTTPBearer): @@ -232,6 +258,7 @@ async def require_auth( detail="Token subject does not match callback ID", ) + liveness_exemption = getattr(route, _LIVENESS_EXEMPTION, None) if ( request.method not in {"GET", "HEAD", "OPTIONS"} and token_scope != "callback" @@ -239,7 +266,7 @@ async def require_auth( and not request.scope.get(_IN_PROCESS_NON_TI_CALLER) and not request.scope.get(_REQUEST_SCOPE_LIVE_ATTEMPT_KEY) and not getattr(route, _SKIP_AUTO_TI_ATTEMPT_LIVE, False) - and not getattr(route, _ALLOW_NON_TI_WRITER, False) + and not (liveness_exemption and liveness_exemption(token)) ): # The versions package imports routes, which depend on this module. from airflow.api_fastapi.execution_api.versions.v2026_10_30 import IdentifyArchivedTaskStateUpdates @@ -312,12 +339,12 @@ class ExecutionAPIRoute(APIRoute): allowed_token_types: frozenset[str] skip_auto_ti_attempt_live: bool - allow_non_ti_writer: bool + liveness_exemption: Callable[[TIToken], bool] | None def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) self.skip_auto_ti_attempt_live = bool(getattr(self.endpoint, _SKIP_AUTO_TI_ATTEMPT_LIVE, False)) - self.allow_non_ti_writer = bool(getattr(self.endpoint, _ALLOW_NON_TI_WRITER, False)) + self.liveness_exemption = getattr(self.endpoint, _LIVENESS_EXEMPTION, None) all_scopes: set[str] = set() for dep in self.dependencies: diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py index 7c6942aa8ef..65a0f0ef25b 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py @@ -17,6 +17,8 @@ from __future__ import annotations +import inspect + from cadwyn import ( ResponseInfo, VersionChange, @@ -72,9 +74,14 @@ class AddStoppedTaskReport(VersionChange): class IdentifyArchivedTaskStateUpdates(VersionChangeWithSideEffects): - """Reject mutations from an archived attempt with 410, except Variable writes; older clients keep each endpoint's own response.""" + """ + Reject mutations from an archived attempt with 410. - description = __doc__ + Variable writes and asset-watcher writes to asset state are exempt. Older + clients keep each endpoint's own response. + """ + + description = inspect.cleandoc(__doc__) instructions_to_migrate_to_previous_version = () diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/test_security.py b/airflow-core/tests/unit/api_fastapi/execution_api/test_security.py index 7d24a1cd446..39177c83d6f 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/test_security.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/test_security.py @@ -36,6 +36,7 @@ from airflow.api_fastapi.execution_api import security from airflow.api_fastapi.execution_api.app import InProcessExecutionAPI from airflow.api_fastapi.execution_api.datamodels.token import TIClaims, TIToken, TokenScope from airflow.api_fastapi.execution_api.security import ( + ASSET_WATCHER_ID, ExecutionAPIRoute, _jwt_bearer, get_team_name_dep, @@ -358,9 +359,10 @@ class TestAttemptLiveness: assert client.delete(path).status_code == 410 assert client.get(path).status_code == 410 - def test_unknown_execution_identity_cannot_write_xcom(self, client, caller, session): + @pytest.mark.parametrize("identity", ["unknown", "asset-watcher"]) + def test_non_ti_identity_cannot_write_xcom(self, client, caller, session, identity): ti, token = caller - token.id = uuid4() + token.id = ASSET_WATCHER_ID if identity == "asset-watcher" else uuid4() path = f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/result" assert client.post(path, json="unowned").status_code == 404 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 97ed0fe1562..70c1220a3b5 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 @@ -17,14 +17,15 @@ from __future__ import annotations import json -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any +from uuid import uuid4 import pytest from fastapi import Request from sqlalchemy import delete, select from airflow.api_fastapi.execution_api.datamodels.token import TIClaims, TIToken -from airflow.api_fastapi.execution_api.security import require_auth +from airflow.api_fastapi.execution_api.security import ASSET_WATCHER_ID, _jwt_bearer, require_auth from airflow.models.asset import AssetActive, AssetModel from airflow.models.asset_state_store import AssetStateStoreModel from airflow.models.dynamic_region import DynamicRegion @@ -35,6 +36,9 @@ from airflow.utils.session import create_session from tests_common.test_utils.mock_operators import MockOperator if TYPE_CHECKING: + from collections.abc import Callable + from uuid import UUID + from fastapi import FastAPI from fastapi.testclient import TestClient from sqlalchemy.orm import Session @@ -59,6 +63,17 @@ _BY_URI_VALUE = "/execution/store/asset/by-uri/value" _BY_URI_CLEAR = "/execution/store/asset/by-uri/clear" +def _authenticate_as(exec_app: FastAPI, monkeypatch: pytest.MonkeyPatch, token_id: UUID) -> None: + """Authenticate as ``token_id`` while keeping the real ``require_auth`` checks.""" + + async def _token() -> TIToken: + return TIToken(id=token_id, claims=TIClaims(scope="execution")) + + jwt_bearer: Callable[..., Any] = _jwt_bearer + monkeypatch.delitem(exec_app.dependency_overrides, require_auth) + monkeypatch.setitem(exec_app.dependency_overrides, jwt_bearer, _token) + + @pytest.fixture(autouse=True) def reset_state_tables(): with create_session() as session: @@ -434,19 +449,12 @@ class TestInactiveAssetRejected: assert response.status_code == 404 -class TestPutAssetStateWatcherWrite: - """PUT routes record WATCHER kind when no TI backs the token (WATCHER/triggerer path).""" +class TestAssetStateWatcherWrite: + """An asset watcher token, which has no task instance behind it, can mutate asset state.""" @pytest.fixture(autouse=True) - def setup_null_uuid_auth(self, exec_app: FastAPI): - from uuid import UUID - - async def _watcher_auth(request: Request) -> TIToken: - return TIToken( - id=UUID("00000000-0000-0000-0000-000000000000"), claims=TIClaims(scope="execution") - ) - - exec_app.dependency_overrides[require_auth] = _watcher_auth + def setup_watcher_auth(self, exec_app: FastAPI, monkeypatch): + _authenticate_as(exec_app, monkeypatch, ASSET_WATCHER_ID) def test_put_by_name_creates_row(self, client: TestClient, asset: AssetModel, session): response = client.put( @@ -504,3 +512,49 @@ class TestPutAssetStateWatcherWrite: assert row.last_updated_by_run_id is None assert row.last_updated_by_task_id is None assert row.last_updated_by_map_index is None + + def test_delete_and_clear(self, client: TestClient, asset: AssetModel): + for k in ("watermark", "last_id"): + _create_asset_store_row(asset.id, k, "x") + + assert client.delete(_BY_URI_VALUE, params={"uri": asset.uri, "key": "watermark"}).status_code == 204 + assert client.get(_BY_NAME_VALUE, params={"name": asset.name, "key": "watermark"}).status_code == 404 + assert client.delete(_BY_NAME_CLEAR, params={"name": asset.name}).status_code == 204 + assert client.get(_BY_NAME_VALUE, params={"name": asset.name, "key": "last_id"}).status_code == 404 + + +class TestAssetStateUnknownIdentityRejected: + """A token that is neither a live attempt nor the asset watcher cannot mutate asset state.""" + + @pytest.fixture(autouse=True) + def setup_unknown_auth(self, exec_app: FastAPI, monkeypatch): + _authenticate_as(exec_app, monkeypatch, uuid4()) + + @pytest.mark.parametrize( + ("method", "path", "by", "body"), + [ + pytest.param("put", _BY_NAME_VALUE, "name", {"value": "stale"}, id="put-by-name"), + pytest.param("delete", _BY_NAME_VALUE, "name", None, id="delete-by-name"), + pytest.param("delete", _BY_NAME_CLEAR, "name", None, id="clear-by-name"), + pytest.param("put", _BY_URI_VALUE, "uri", {"value": "stale"}, id="put-by-uri"), + pytest.param("delete", _BY_URI_VALUE, "uri", None, id="delete-by-uri"), + pytest.param("delete", _BY_URI_CLEAR, "uri", None, id="clear-by-uri"), + ], + ) + def test_mutation_rejected(self, client: TestClient, asset: AssetModel, method, path, by, body): + _create_asset_store_row(asset.id, "watermark", "original") + params = {by: getattr(asset, by)} + if path in (_BY_NAME_VALUE, _BY_URI_VALUE): + params["key"] = "watermark" + + response = client.request(method, path, params=params, json=body) + + assert response.status_code == 404, response.text + assert response.json()["detail"]["message"] == "Task Instance not found" + with create_session() as session: + value = session.scalar( + select(AssetStateStoreModel.value).where( + AssetStateStoreModel.asset_id == asset.id, AssetStateStoreModel.key == "watermark" + ) + ) + assert value == '"original"'
