This is an automated email from the ASF dual-hosted git repository.
vatsrahul1001 pushed a commit to branch v3-4-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-4-test by this push:
new 18a6d2daf0f Let asset watchers write asset state over the Execution
API (#76376) (#76469)
18a6d2daf0f is described below
commit 18a6d2daf0f14c485e30abf356717e623648cdb9
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"'
