From 4495e92002c2547f6bf10f494cb5a7b1895c0da6 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Wed, 4 Nov 2026 11:06:03 +0800 Subject: [PATCH] Decode state store vals through backend; treat null as missing (#76009) --- .../task-and-asset-state-store.rst | 2 + .../execution_api/routes/asset_state_store.py | 10 ++--- .../execution_api/routes/task_state_store.py | 5 +-- .../versions/head/test_asset_state_store.py | 19 +++++++- .../versions/head/test_task_state_store.py | 22 +++++++++ .../tests/unit/state/test_metastore.py | 7 +++ .../src/airflow_shared/state/__init__.py | 30 +++++++++++-- shared/state/tests/state/test_state.py | 45 +++++++++++++++++++ 8 files changed, 127 insertions(+), 13 deletions(-) diff --git a/airflow-core/docs/administration-and-deployment/task-and-asset-state-store.rst b/airflow-core/docs/administration-and-deployment/task-and-asset-state-store.rst index 9b13da2abc1..0d88a3d28b7 100644 --- a/airflow-core/docs/administration-and-deployment/task-and-asset-state-store.rst +++ b/airflow-core/docs/administration-and-deployment/task-and-asset-state-store.rst @@ -128,6 +128,8 @@ Custom backends A custom backend must subclass :class:`~airflow.sdk.state.BaseStoreBackend` and implement its abstract methods: ``get``, ``set``, ``delete``, and ``clear`` for synchronous callers and the ``aget``, ``aset``, ``adelete``, and ``aclear`` async equivalents. Refer to :class:`~airflow.sdk.state.BaseStoreBackend` for the full API. +Values are passed to :meth:`~airflow.sdk.state.BaseStoreBackend.set` and returned from :meth:`~airflow.sdk.state.BaseStoreBackend.get` as JSON encoded strings, and are never the JSON ``null``. The Execution API reads values through :meth:`~airflow.sdk.state.BaseStoreBackend.get_value` (:meth:`~airflow.sdk.state.BaseStoreBackend.aget_value` is the async equivalent); a backend that stores values natively can override them to skip decoding on reads. + Each method receives a ``scope`` argument that is either a :class:`~airflow.sdk.state.TaskScope` or an :class:`~airflow.sdk.state.AssetScope`. The union of the two is exported as ``airflow.sdk.state.StoreScope`` for type annotations. Use ``isinstance`` to dispatch: .. code-block:: python 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..466b5f22b74 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 @@ -129,13 +129,12 @@ def get_asset_state_store_by_name( ) -> AssetStateStoreResponse: """Get an asset state store value by asset name.""" asset_id = _resolve_asset_id_by_name(name, session) - value = get_state_backend().get(AssetScope(asset_id=asset_id), key, session=session) - if value is None: + if (value := get_state_backend().get_value(AssetScope(asset_id=asset_id), key, session=session)) is None: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail={"reason": "not_found", "message": f"Asset state store key {key!r} not found"}, ) - return AssetStateStoreResponse(value=json.loads(value)) + return AssetStateStoreResponse(value=value) def _put_asset_state_store( @@ -221,13 +220,12 @@ def get_asset_state_store_by_uri( ) -> AssetStateStoreResponse: """Get an asset state store value by asset URI.""" asset_id = _resolve_asset_id_by_uri(uri, session) - value = get_state_backend().get(AssetScope(asset_id=asset_id), key, session=session) - if value is None: + if (value := get_state_backend().get_value(AssetScope(asset_id=asset_id), key, session=session)) is None: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail={"reason": "not_found", "message": f"Asset state store key {key!r} not found"}, ) - return AssetStateStoreResponse(value=json.loads(value)) + return AssetStateStoreResponse(value=value) @router.put("/by-uri/value", status_code=status.HTTP_204_NO_CONTENT) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state_store.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state_store.py index 94b4ec01391..4289af2be30 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state_store.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state_store.py @@ -72,8 +72,7 @@ def get_task_state_store( ) -> TaskStateStoreResponse: """Get value for a task state store key.""" scope = _get_task_scope_for_ti(task_instance_id, session) - value = get_state_backend().get(scope, key, session=session) - if value is None: + if (value := get_state_backend().get_value(scope, key, session=session)) is None: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail={ @@ -81,7 +80,7 @@ def get_task_state_store( "message": f"Task state store key {key!r} not found", }, ) - return TaskStateStoreResponse(value=json.loads(value)) + return TaskStateStoreResponse(value=value) @router.put("/{task_instance_id}/{key:path}", status_code=status.HTTP_204_NO_CONTENT) 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..d24c6e50d61 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 @@ -37,6 +37,7 @@ from tests_common.test_utils.mock_operators import MockOperator if TYPE_CHECKING: from fastapi import FastAPI from fastapi.testclient import TestClient + from pydantic import JsonValue from sqlalchemy.orm import Session from airflow.models.taskinstance import TaskInstance @@ -48,7 +49,7 @@ pytestmark = pytest.mark.db_test _BY_NAME_VALUE = "/execution/store/asset/by-name/value" -def _create_asset_store_row(asset_id: int, key: str, value: str) -> None: +def _create_asset_store_row(asset_id: int, key: str, value: JsonValue) -> None: """Insert an AssetStateStoreModel row directly in the database.""" with create_session() as s: s.add(AssetStateStoreModel(asset_id=asset_id, key=key, value=json.dumps(value))) @@ -86,6 +87,14 @@ def inactive_asset(session: Session) -> AssetModel: class TestGetAssetStateByName: + def test_get_stored_null_returns_404(self, client: TestClient, asset: AssetModel): + _create_asset_store_row(asset.id, "watermark", None) + + response = client.get(_BY_NAME_VALUE, params={"name": asset.name, "key": "watermark"}) + + assert response.status_code == 404 + assert response.json()["detail"]["reason"] == "not_found" + def test_get_returns_value(self, client: TestClient, asset: AssetModel): _create_asset_store_row(asset.id, "watermark", "2026-04-29") @@ -320,6 +329,14 @@ class TestClearAssetStateByName: class TestGetAssetStateByUri: + def test_get_stored_null_returns_404(self, client: TestClient, asset: AssetModel): + _create_asset_store_row(asset.id, "watermark", None) + + response = client.get(_BY_URI_VALUE, params={"uri": asset.uri, "key": "watermark"}) + + assert response.status_code == 404 + assert response.json()["detail"]["reason"] == "not_found" + def test_get_returns_value(self, client: TestClient, asset: AssetModel): _create_asset_store_row(asset.id, "watermark", "2026-04-29") diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_state_store.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_state_store.py index 042ae4ee653..f621c2b78fe 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_state_store.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_state_store.py @@ -79,6 +79,28 @@ def test_task_uuid_isolates_state_in_different_regions(client, create_task_insta class TestGetTaskState: + def test_get_stored_null_returns_404(self, client: TestClient, create_task_instance: CreateTaskInstance): + ti = create_task_instance() + with create_session() as session: + session.add( + TaskStateStoreModel( + dag_run_id=ti.dag_run.id, + dag_id=ti.dag_id, + run_id=ti.run_id, + task_id=ti.task_id, + region_id=ti.region_id, + region_index=ti.region_index, + key="job_id", + value="null", + updated_at=timezone.utcnow(), + ) + ) + + response = client.get(_api_url(ti.id, "job_id")) + + assert response.status_code == 404 + assert response.json()["detail"]["reason"] == "not_found" + 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"}) diff --git a/airflow-core/tests/unit/state/test_metastore.py b/airflow-core/tests/unit/state/test_metastore.py index d353949374f..de9335dc83b 100644 --- a/airflow-core/tests/unit/state/test_metastore.py +++ b/airflow-core/tests/unit/state/test_metastore.py @@ -400,6 +400,13 @@ class TestMetastoreBackendTaskScope: class TestMetastoreBackendAssetScope: + def test_get_value_without_session(self, backend: MetastoreBackend, asset_committed: AssetModel): + scope = AssetScope(asset_id=asset_committed.id) + backend.set(scope, "watermark", '"2026-04-29"') + + assert backend.get_value(scope, "watermark") == "2026-04-29" + assert backend.get_value(scope, "missing") is None + def test_get_returns_none_for_missing_key( self, session: Session, backend: MetastoreBackend, asset: AssetModel ): diff --git a/shared/state/src/airflow_shared/state/__init__.py b/shared/state/src/airflow_shared/state/__init__.py index 5c0ed914b59..555d109f1e8 100644 --- a/shared/state/src/airflow_shared/state/__init__.py +++ b/shared/state/src/airflow_shared/state/__init__.py @@ -167,8 +167,9 @@ class BaseStoreBackend(ABC): """ Return the stored JSON encoded value string, or None if the key does not exist. - Must handle both ``TaskScope`` and ``AssetScope``. The execution API calls - ``json.loads`` on the returned string from here, so it must be a valid JSON document. + Must handle both ``TaskScope`` and ``AssetScope``. :meth:`get_value` decodes the + returned string with ``json.loads``, so it must be a valid JSON document. It is + never the JSON ``null``, since :meth:`set` never receives one. """ @abstractmethod @@ -186,7 +187,8 @@ class BaseStoreBackend(ABC): Must handle both ``TaskScope`` and ``AssetScope``. ``value`` is always a JSON encoded string (the execution API calls ``json.dumps`` before passing it - here); store it verbatim so ``get`` can return it unchanged. + here); store it verbatim so ``get`` can return it unchanged. It is never the + JSON ``null``; callers reject a null value before it reaches the backend. ``expires_at`` is an absolute UTC datetime after which the row may be deleted. Pass ``None`` (default) for a key that should never expire — stored as ``NULL``, @@ -269,6 +271,28 @@ class BaseStoreBackend(ABC): If ``None``, implementations manage their own async session internally. """ + def get_value(self, scope: StoreScope, key: str, *, session: Session | None = None) -> JsonValue: + """ + Return the decoded value for the given key, or None if the key does not exist. + + The default implementation decodes what :meth:`get` returns. A stored JSON ``null`` + breaks the :meth:`set` contract, so it is treated as a missing key. A backend that + stores values natively, rather than as JSON strings, can override this to skip + decoding on reads. + """ + kwargs = {"session": session} if session is not None else {} + if (value := self.get(scope, key, **kwargs)) is None: + return None + return json.loads(value) + + async def aget_value( + self, scope: StoreScope, key: str, *, session: AsyncSession | None = None + ) -> JsonValue: + """Async variant of :meth:`get_value`.""" + if (value := await self.aget(scope, key, session=session)) is None: + return None + return json.loads(value) + def cleanup(self) -> None: """ Remove expired and orphaned state records. diff --git a/shared/state/tests/state/test_state.py b/shared/state/tests/state/test_state.py index 578ae302212..f5782988eda 100644 --- a/shared/state/tests/state/test_state.py +++ b/shared/state/tests/state/test_state.py @@ -208,3 +208,48 @@ class TestBaseStoreBackend: b.deserialize_asset_state_store_from_ref("s3://bucket/assets/my_asset/watermark") == "resolved:s3://bucket/assets/my_asset/watermark" ) + + +class _StoredValueBackend(BaseStoreBackend): + """Backend whose ``get`` and ``aget`` return one fixed stored string.""" + + def __init__(self, stored: str | None): + self.stored = stored + + def get(self, scope, key, *, session=None): + return self.stored + + def set(self, scope, key, value, *, expires_at=None, session=None): ... + def delete(self, scope, key, *, session=None): ... + def clear(self, scope, *, all_map_indices=False, session=None): ... + + async def aget(self, scope, key, *, session=None): + return self.stored + + async def aset(self, scope, key, value, *, expires_at=None, session=None): ... + async def adelete(self, scope, key, *, session=None): ... + async def aclear(self, scope, *, all_map_indices=False, session=None): ... + + +_GET_VALUE_CASES = [ + pytest.param(None, None, id="missing"), + pytest.param("null", None, id="stored-null"), + pytest.param('"x"', "x", id="string"), + pytest.param("0", 0, id="falsy-number"), + pytest.param('{"a": [1]}', {"a": [1]}, id="object"), +] + + +class TestBaseStoreBackendGetValue: + @pytest.mark.parametrize(("stored", "expected"), _GET_VALUE_CASES) + def test_get_value(self, stored, expected): + backend = _StoredValueBackend(stored) + + assert backend.get_value(TaskScope("dag", "run", "task"), "key") == expected + + @pytest.mark.asyncio + @pytest.mark.parametrize(("stored", "expected"), _GET_VALUE_CASES) + async def test_aget_value(self, stored, expected): + backend = _StoredValueBackend(stored) + + assert await backend.aget_value(AssetScope(name="asset"), "key") == expected