This is an automated email from the ASF dual-hosted git repository.
uranusjr pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 4495e92002c Decode state store vals through backend; treat null as
missing (#76009)
4495e92002c is described below
commit 4495e92002c2547f6bf10f494cb5a7b1895c0da6
Author: Tzu-ping Chung <[email protected]>
AuthorDate: Wed Nov 4 11:06:03 2026 +0800
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 +++++++++++
airflow-core/tests/unit/state/test_metastore.py | 7 ++++
shared/state/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
