Decode state store vals through backend; treat null as missing (#76009)

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

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)

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)

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")
 

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"})

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
     ):

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.

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