This is an automated email from the ASF dual-hosted git repository.
vatsrahul1001 pushed a commit to branch v3-3-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-3-test by this push:
new 0d03ef4396d Add async asset store accessors for async tasks and
watcher triggers (#72127) (#72851)
0d03ef4396d is described below
commit 0d03ef4396d5d246db59d6e7b008f1d342fecaf7
Author: Jason(Zhe-You) Liu <[email protected]>
AuthorDate: Thu Sep 10 18:57:57 2026 +0800
Add async asset store accessors for async tasks and watcher triggers
(#72127) (#72851)
* Add async asset store accessors for async tasks and watcher triggers
* Clarify async asset state accessor docstrings
* Parametrize async asset state accessor tests
(cherry picked from commit 2aa581ea02ba9c39a6c85560b1308f4cc09d6c9d)
---
.../docs/core-concepts/asset-state-store.rst | 25 ++-
task-sdk/src/airflow/sdk/execution_time/context.py | 127 ++++++++----
.../tests/task_sdk/execution_time/test_context.py | 225 ++++++++++++++++++++-
3 files changed, 333 insertions(+), 44 deletions(-)
diff --git a/airflow-core/docs/core-concepts/asset-state-store.rst
b/airflow-core/docs/core-concepts/asset-state-store.rst
index 3aa7ff7719f..ddb3fe93fb4 100644
--- a/airflow-core/docs/core-concepts/asset-state-store.rst
+++ b/airflow-core/docs/core-concepts/asset-state-store.rst
@@ -19,6 +19,7 @@
.. spelling:word-list::
+ accessors
subscripted
subscripting
@@ -101,7 +102,7 @@ If the task has more than one concrete inlet or outlet,
calling the shorthand ra
API reference
-------------
-The following methods are available on both the per-asset accessor
(``context["asset_state_store"][my_asset]``) and the shorthand
(``context["asset_state_store"]``) when the task has exactly one inlet.
+The following methods are available on both the per-asset accessor
(``context["asset_state_store"][my_asset]``) and the shorthand
(``context["asset_state_store"]``) when the task has exactly one inlet. Each
has an ``a``-prefixed async counterpart — ``aget``, ``aset``, ``adelete``,
``aclear`` — for use inside ``async`` tasks and watcher triggers.
``get(key, default)``
~~~~~~~~~~~~~~~~~~~~~
@@ -150,6 +151,20 @@ Deletes *all* asset state store keys for the asset.
# Using context
context["asset_state_store"][my_asset].clear()
+``aget``, ``aset``, ``adelete``, ``aclear``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Async counterparts of ``get``, ``set``, ``delete``, and ``clear``. They take
the same arguments and behave identically to their synchronous siblings —
``aset`` has no ``retention`` parameter either — but ``await`` the round-trip
to the API server instead of blocking the event loop, so a coroutine can read
and advance asset state without stalling other concurrent work.
+
+.. code-block:: python
+
+ watermark = await context["asset_state_store"][my_asset].aget("watermark",
default="initial_watermark")
+ await context["asset_state_store"][my_asset].aset("watermark",
"2024-06-01T00:00:00Z")
+ await context["asset_state_store"][my_asset].adelete("watermark")
+ await context["asset_state_store"][my_asset].aclear()
+
+Calling the synchronous ``get``/``set``/``delete``/``clear`` from inside an
``async`` task blocks the event loop and defeats the concurrency the coroutine
was written for. Use the ``a``-prefixed methods there instead.
+
Using ``asset_state_store`` inside a Watcher Trigger
-----------------------------------------------------
@@ -184,19 +199,21 @@ Unlike task-based access (where the asset is identified
by an inlet or outlet de
async def run(self) -> AsyncIterator[TriggerEvent]:
while True:
- last_seen = self.asset_state_store.get("last_seen_id",
default=0)
+ last_seen = await self.asset_state_store.aget("last_seen_id",
default=0)
new_id = self._poll_for_new_record(
source=self.source,
last_seen=last_seen,
)
if new_id is not None:
- self.asset_state_store.set("last_seen_id", new_id)
+ await self.asset_state_store.aset("last_seen_id", new_id)
yield TriggerEvent({"status": "success", "record_id":
new_id})
return
await asyncio.sleep(self.waiter_delay)
+``run()`` is a coroutine, and every trigger on a triggerer shares one event
loop, so use the async accessors there. The synchronous methods also work, but
they hold the loop for the whole round-trip to the API server.
+
The corresponding :class:`~airflow.sdk.definitions.asset.AssetWatcher` wires
the trigger to the asset:
.. code-block:: python
@@ -217,7 +234,7 @@ The corresponding
:class:`~airflow.sdk.definitions.asset.AssetWatcher` wires the
...
-``self.asset_state_store`` behaves identically to the per-asset accessor
described in the task sections above: ``get``, ``set``, ``delete``, and
``clear`` are all available. Values written by the trigger are visible to any
task that declares ``my_asset`` as an inlet or outlet, and vice versa.
+``self.asset_state_store`` behaves identically to the per-asset accessor
described in the task sections above: ``get``, ``set``, ``delete``, ``clear``
and their async counterparts ``aget``, ``aset``, ``adelete``, ``aclear`` are
all available. Values written by the trigger are visible to any task that
declares ``my_asset`` as an inlet or outlet, and vice versa.
.. note::
diff --git a/task-sdk/src/airflow/sdk/execution_time/context.py
b/task-sdk/src/airflow/sdk/execution_time/context.py
index e9deba57d67..58c58e5945e 100644
--- a/task-sdk/src/airflow/sdk/execution_time/context.py
+++ b/task-sdk/src/airflow/sdk/execution_time/context.py
@@ -72,6 +72,7 @@ if TYPE_CHECKING:
OKResponse,
PrevSuccessfulDagRunResponse,
ReceiveMsgType,
+ ToSupervisor,
VariableResult,
)
from airflow.sdk.state import BaseStoreBackend
@@ -694,21 +695,31 @@ class AssetStateStoreAccessor:
def get(self, key: str, default: JsonValue = None) -> JsonValue:
"""Return the stored value, or ``default`` if the key does not
exist."""
- from airflow.sdk.execution_time.comms import (
- AssetStateStoreResult,
- ErrorResponse,
- GetAssetStateStoreByName,
- GetAssetStateStoreByUri,
- ToSupervisor,
- )
from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+ resp = SUPERVISOR_COMMS.send(self._build_get_message(key))
+ return self._extract_get_response(resp, key, default)
+
+ async def aget(self, key: str, default: JsonValue = None) -> JsonValue:
+ """Async version of `get` that awaits instead of blocking the event
loop."""
+ from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+
+ resp = await SUPERVISOR_COMMS.asend(self._build_get_message(key))
+ return self._extract_get_response(resp, key, default)
+
+ def _build_get_message(self, key: str) -> ToSupervisor:
+ from airflow.sdk.execution_time.comms import GetAssetStateStoreByName,
GetAssetStateStoreByUri
+
msg: ToSupervisor
if self._name:
msg = GetAssetStateStoreByName(name=self._name, key=key)
elif self._uri:
msg = GetAssetStateStoreByUri(uri=self._uri, key=key)
- resp = SUPERVISOR_COMMS.send(msg)
+ return msg
+
+ def _extract_get_response(self, resp: Any, key: str, default: JsonValue)
-> JsonValue:
+ from airflow.sdk.execution_time.comms import AssetStateStoreResult,
ErrorResponse
+
if isinstance(resp, ErrorResponse) and resp.error !=
ErrorType.ASSET_STORE_NOT_FOUND:
raise AirflowRuntimeError(resp)
if isinstance(resp, AssetStateStoreResult):
@@ -730,13 +741,19 @@ class AssetStateStoreAccessor:
def set(self, key: str, value: JsonValue) -> None:
"""Write or overwrite the value for the given key. ``value`` must not
be ``None``."""
- from airflow.sdk.execution_time.comms import (
- SetAssetStateStoreByName,
- SetAssetStateStoreByUri,
- ToSupervisor,
- )
from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+ SUPERVISOR_COMMS.send(self._build_set_message(key, value))
+
+ async def aset(self, key: str, value: JsonValue) -> None:
+ """Async version of `set` that awaits instead of blocking the event
loop."""
+ from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+
+ await SUPERVISOR_COMMS.asend(self._build_set_message(key, value))
+
+ def _build_set_message(self, key: str, value: JsonValue) -> ToSupervisor:
+ from airflow.sdk.execution_time.comms import SetAssetStateStoreByName,
SetAssetStateStoreByUri
+
if value is None:
raise ValueError("Cannot set value as None")
@@ -765,49 +782,66 @@ class AssetStateStoreAccessor:
msg = SetAssetStateStoreByName(name=self._name, key=key,
value=stored)
elif self._uri:
msg = SetAssetStateStoreByUri(uri=self._uri, key=key, value=stored)
- SUPERVISOR_COMMS.send(msg)
+ return msg
def delete(self, key: str) -> None:
"""Delete a single key. No-op if the key does not exist."""
- from airflow.sdk._shared.state import AssetScope
- from airflow.sdk.execution_time.comms import (
- DeleteAssetStateStoreByName,
- DeleteAssetStateStoreByUri,
- ToSupervisor,
- )
from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
- msg: ToSupervisor
- if self._name:
- msg = DeleteAssetStateStoreByName(name=self._name, key=key)
- elif self._uri:
- msg = DeleteAssetStateStoreByUri(uri=self._uri, key=key)
# DB ref first: if backend cleanup fails after this, the ref is gone
and
# deterministic keys are recoverable on next set().
- SUPERVISOR_COMMS.send(msg)
+ SUPERVISOR_COMMS.send(self._build_delete_message(key))
backend = _get_worker_state_store_backend()
if backend is not None:
backend.delete(AssetScope(name=self._name, uri=self._uri), key)
+ async def adelete(self, key: str) -> None:
+ """Async version of `delete` that awaits instead of blocking the event
loop."""
+ from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+
+ await SUPERVISOR_COMMS.asend(self._build_delete_message(key))
+ backend = _get_worker_state_store_backend()
+ if backend is not None:
+ await backend.adelete(AssetScope(name=self._name, uri=self._uri),
key)
+
+ def _build_delete_message(self, key: str) -> ToSupervisor:
+ from airflow.sdk.execution_time.comms import
DeleteAssetStateStoreByName, DeleteAssetStateStoreByUri
+
+ msg: ToSupervisor
+ if self._name:
+ msg = DeleteAssetStateStoreByName(name=self._name, key=key)
+ elif self._uri:
+ msg = DeleteAssetStateStoreByUri(uri=self._uri, key=key)
+ return msg
+
def clear(self) -> None:
"""Delete all state keys for this asset."""
- from airflow.sdk._shared.state import AssetScope
- from airflow.sdk.execution_time.comms import (
- ClearAssetStateStoreByName,
- ClearAssetStateStoreByUri,
- ToSupervisor,
- )
from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+ # DB ref first, same ordering rationale as delete().
+ SUPERVISOR_COMMS.send(self._build_clear_message())
+ backend = _get_worker_state_store_backend()
+ if backend is not None:
+ backend.clear(AssetScope(name=self._name, uri=self._uri))
+
+ async def aclear(self) -> None:
+ """Async version of `clear` that awaits instead of blocking the event
loop."""
+ from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+
+ await SUPERVISOR_COMMS.asend(self._build_clear_message())
+ backend = _get_worker_state_store_backend()
+ if backend is not None:
+ await backend.aclear(AssetScope(name=self._name, uri=self._uri))
+
+ def _build_clear_message(self) -> ToSupervisor:
+ from airflow.sdk.execution_time.comms import
ClearAssetStateStoreByName, ClearAssetStateStoreByUri
+
msg: ToSupervisor
if self._name:
msg = ClearAssetStateStoreByName(name=self._name)
elif self._uri:
msg = ClearAssetStateStoreByUri(uri=self._uri)
- SUPERVISOR_COMMS.send(msg)
- backend = _get_worker_state_store_backend()
- if backend is not None:
- backend.clear(AssetScope(name=self._name, uri=self._uri))
+ return msg
class AssetStateStoreAccessors:
@@ -818,7 +852,8 @@ class AssetStateStoreAccessors:
accessor as: ``context['asset_state_store'][MY_ASSET].get('watermark')``.
For tasks with exactly one concrete inlet or outlet, the accessor methods
(``get``,
- ``set``, ``delete``, ``clear``) can be called directly without
subscripting.
+ ``set``, ``delete``, ``clear``, and their async counterparts ``aget``,
``aset``,
+ ``adelete``, ``aclear``) can be called directly without subscripting.
"""
def __init__(self, inlets: list, outlets: list | None = None) -> None:
@@ -871,18 +906,34 @@ class AssetStateStoreAccessors:
"""Return the stored value for the single-inlet or single-outlet task,
or ``default`` if not found."""
return self._single_accessor().get(key, default)
+ async def aget(self, key: str, default: JsonValue = None) -> JsonValue:
+ """Async version of `get` that awaits instead of blocking the event
loop."""
+ return await self._single_accessor().aget(key, default)
+
def set(self, key: str, value: JsonValue) -> None:
"""Write or overwrite the value for the single-inlet task."""
self._single_accessor().set(key, value)
+ async def aset(self, key: str, value: JsonValue) -> None:
+ """Async version of `set` that awaits instead of blocking the event
loop."""
+ await self._single_accessor().aset(key, value)
+
def delete(self, key: str) -> None:
"""Delete a single key for the single-inlet task."""
self._single_accessor().delete(key)
+ async def adelete(self, key: str) -> None:
+ """Async version of `delete` that awaits instead of blocking the event
loop."""
+ await self._single_accessor().adelete(key)
+
def clear(self) -> None:
"""Delete all state keys for the single-inlet task."""
self._single_accessor().clear()
+ async def aclear(self) -> None:
+ """Async version of `clear` that awaits instead of blocking the event
loop."""
+ await self._single_accessor().aclear()
+
def __repr__(self) -> str:
parts = [f"name={k!r}" for k in self._by_name] + [f"uri={k!r}" for k
in self._by_uri]
return f"<AssetStateStoreAccessors [{', '.join(parts)}]>"
@@ -942,7 +993,6 @@ class _AssetRefResolutionMixin:
ErrorResponse,
GetAssetByName,
GetAssetByUri,
- ToSupervisor,
)
from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
@@ -1134,7 +1184,6 @@ class InletEventsAccessor(Sequence["AssetEventResult"]):
ErrorResponse,
GetAssetEventByAsset,
GetAssetEventByAssetAlias,
- ToSupervisor,
)
from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
diff --git a/task-sdk/tests/task_sdk/execution_time/test_context.py
b/task-sdk/tests/task_sdk/execution_time/test_context.py
index ba1ebb3694f..23e132432b2 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_context.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_context.py
@@ -27,7 +27,7 @@ import pytest
from pydantic import ValidationError
from airflow.sdk import BaseOperator, get_current_context, timezone
-from airflow.sdk._shared.state import TaskScope
+from airflow.sdk._shared.state import AssetScope, TaskScope
from airflow.sdk.api.datamodels._generated import (
AssetEventResponse,
AssetResponse,
@@ -1696,6 +1696,163 @@ class TestAssetStateStoreAccessor:
assert "max_value_storage_bytes" in
mock_log.warning.call_args[0][0]
mock_supervisor_comms.send.assert_called_once()
+ @pytest.mark.parametrize("lookup_by", ["name", "uri"])
+ @pytest.mark.asyncio
+ async def test_aget_returns_value(self, mock_supervisor_comms, lookup_by):
+ """aget awaits asend and returns the stored value, without touching
sync send."""
+ mock_supervisor_comms.asend.return_value =
AssetStateStoreResult(value="2026-04-30T00:00:00Z")
+ if lookup_by == "name":
+ accessor = AssetStateStoreAccessor(name=self.ASSET_NAME)
+ expected_message = GetAssetStateStoreByName(name=self.ASSET_NAME,
key="watermark")
+ else:
+ accessor = AssetStateStoreAccessor(uri=self.ASSET_URI)
+ expected_message = GetAssetStateStoreByUri(uri=self.ASSET_URI,
key="watermark")
+
+ result = await accessor.aget("watermark")
+
+ assert result == "2026-04-30T00:00:00Z"
+ mock_supervisor_comms.asend.assert_called_once_with(expected_message)
+ mock_supervisor_comms.send.assert_not_called()
+
+ @pytest.mark.asyncio
+ async def test_aget_returns_default_when_key_missing(self,
mock_supervisor_comms):
+ mock_supervisor_comms.asend.return_value = ErrorResponse(
+ error=ErrorType.ASSET_STORE_NOT_FOUND, detail={"key": "watermark"}
+ )
+
+ result = await AssetStateStoreAccessor(name=self.ASSET_NAME).aget(
+ "watermark", default="2026-01-01T00:00:00+00:00"
+ )
+
+ assert result == "2026-01-01T00:00:00+00:00"
+
+ @pytest.mark.asyncio
+ async def test_aget_raises_on_error(self, mock_supervisor_comms):
+ mock_supervisor_comms.asend.return_value = ErrorResponse(
+ error=ErrorType.GENERIC_ERROR, detail={"message": "server error"}
+ )
+
+ with pytest.raises(AirflowRuntimeError):
+ await
AssetStateStoreAccessor(name=self.ASSET_NAME).aget("some_key")
+
+ @pytest.mark.asyncio
+ async def test_aget_with_custom_backend_removes_decoration_marker(self,
mock_supervisor_comms):
+ """aget unwraps the external Store marker and resolves the ref via the
backend."""
+ mock_supervisor_comms.asend.return_value = AssetStateStoreResult(
+ value=_wrap_external_ref("s3://bucket/assets/orders/watermark")
+ )
+
+ backend = MagicMock(spec=BaseStoreBackend)
+ backend.deserialize_asset_state_store_from_ref.return_value =
"2026-05-01"
+
+ with patch(
+
"airflow.sdk.execution_time.context._get_worker_state_store_backend",
return_value=backend
+ ):
+ result = await
AssetStateStoreAccessor(name=self.ASSET_NAME).aget("watermark")
+
+ assert result == "2026-05-01"
+ backend.deserialize_asset_state_store_from_ref.assert_called_once_with(
+ "s3://bucket/assets/orders/watermark"
+ )
+
+ @pytest.mark.parametrize("lookup_by", ["name", "uri"])
+ @pytest.mark.asyncio
+ async def test_aset_operation(self, mock_supervisor_comms, lookup_by):
+ mock_supervisor_comms.asend.return_value = OKResponse(ok=True)
+ if lookup_by == "name":
+ accessor = AssetStateStoreAccessor(name=self.ASSET_NAME)
+ expected_message = SetAssetStateStoreByName(
+ name=self.ASSET_NAME, key="watermark",
value="2026-04-30T00:00:00Z"
+ )
+ else:
+ accessor = AssetStateStoreAccessor(uri=self.ASSET_URI)
+ expected_message = SetAssetStateStoreByUri(
+ uri=self.ASSET_URI, key="watermark",
value="2026-04-30T00:00:00Z"
+ )
+
+ await accessor.aset("watermark", "2026-04-30T00:00:00Z")
+
+ mock_supervisor_comms.asend.assert_called_once_with(expected_message)
+ mock_supervisor_comms.send.assert_not_called()
+
+ @pytest.mark.asyncio
+ async def test_aset_none_raises(self, mock_supervisor_comms):
+ with pytest.raises(ValueError, match="Cannot set value as None"):
+ await
AssetStateStoreAccessor(name=self.ASSET_NAME).aset("watermark", None)
+
+ mock_supervisor_comms.asend.assert_not_called()
+
+ @pytest.mark.asyncio
+ async def test_aset_with_custom_backend_decorates_value_with_marker(self,
mock_supervisor_comms):
+ """aset wraps the custom backend ref in the external Store marker
before sending."""
+ mock_supervisor_comms.asend.return_value = OKResponse(ok=True)
+
+ backend = MagicMock(spec=BaseStoreBackend)
+ backend.serialize_asset_state_store_to_ref.return_value =
"s3://bucket/assets/orders/watermark"
+
+ with patch(
+
"airflow.sdk.execution_time.context._get_worker_state_store_backend",
return_value=backend
+ ):
+ await
AssetStateStoreAccessor(name=self.ASSET_NAME).aset("watermark", "2026-05-01")
+
+ mock_supervisor_comms.asend.assert_called_once_with(
+ SetAssetStateStoreByName(
+ name=self.ASSET_NAME,
+ key="watermark",
+
value=_wrap_external_ref("s3://bucket/assets/orders/watermark"),
+ )
+ )
+
+ @pytest.mark.asyncio
+ async def test_adelete_awaits_asend(self, mock_supervisor_comms):
+ """adelete awaits asend without touching sync send."""
+ mock_supervisor_comms.asend.return_value = OKResponse(ok=True)
+
+ await
AssetStateStoreAccessor(name=self.ASSET_NAME).adelete("watermark")
+
+ mock_supervisor_comms.asend.assert_called_once_with(
+ DeleteAssetStateStoreByName(name=self.ASSET_NAME, key="watermark")
+ )
+ mock_supervisor_comms.send.assert_not_called()
+
+ @pytest.mark.asyncio
+ async def test_aclear_awaits_asend(self, mock_supervisor_comms):
+ """aclear awaits asend without touching sync send."""
+ mock_supervisor_comms.asend.return_value = OKResponse(ok=True)
+
+ await AssetStateStoreAccessor(uri=self.ASSET_URI).aclear()
+
+
mock_supervisor_comms.asend.assert_called_once_with(ClearAssetStateStoreByUri(uri=self.ASSET_URI))
+ mock_supervisor_comms.send.assert_not_called()
+
+ @pytest.mark.asyncio
+ async def test_adelete_purges_via_async_backend(self,
mock_supervisor_comms):
+ """adelete awaits the async backend instead of blocking on the sync
delete."""
+ mock_supervisor_comms.asend.return_value = OKResponse(ok=True)
+ backend = MagicMock(spec=BaseStoreBackend)
+
+ with patch(
+
"airflow.sdk.execution_time.context._get_worker_state_store_backend",
return_value=backend
+ ):
+ await
AssetStateStoreAccessor(name=self.ASSET_NAME).adelete("watermark")
+
+
backend.adelete.assert_awaited_once_with(AssetScope(name=self.ASSET_NAME,
uri=None), "watermark")
+ backend.delete.assert_not_called()
+
+ @pytest.mark.asyncio
+ async def test_aclear_purges_via_async_backend(self,
mock_supervisor_comms):
+ """aclear awaits the async backend instead of blocking on the sync
clear."""
+ mock_supervisor_comms.asend.return_value = OKResponse(ok=True)
+ backend = MagicMock(spec=BaseStoreBackend)
+
+ with patch(
+
"airflow.sdk.execution_time.context._get_worker_state_store_backend",
return_value=backend
+ ):
+ await AssetStateStoreAccessor(name=self.ASSET_NAME).aclear()
+
+
backend.aclear.assert_awaited_once_with(AssetScope(name=self.ASSET_NAME,
uri=None))
+ backend.clear.assert_not_called()
+
class TestAssetStateStoreAccessors:
ASSET_NAME = "my_asset"
@@ -1859,6 +2016,43 @@ class TestAssetStateStoreAccessors:
assert accessors._total == 0
mock_supervisor_comms.send.assert_not_called()
+ @pytest.mark.parametrize(
+ ("method_name", "expected_message"),
+ [
+ ("aget", GetAssetStateStoreByName(name=ASSET_NAME,
key="watermark")),
+ ("aset", SetAssetStateStoreByName(name=ASSET_NAME,
key="watermark", value="2026-05-01")),
+ ("adelete", DeleteAssetStateStoreByName(name=ASSET_NAME,
key="watermark")),
+ ("aclear", ClearAssetStateStoreByName(name=ASSET_NAME)),
+ ],
+ )
+ @pytest.mark.asyncio
+ async def test_single_inlet_async_shorthand(self, mock_supervisor_comms,
method_name, expected_message):
+ asset = Asset(name=self.ASSET_NAME, uri=f"s3://{self.ASSET_NAME}")
+ mock_supervisor_comms.asend.return_value = (
+ AssetStateStoreResult(value="v5") if method_name == "aget" else
OKResponse(ok=True)
+ )
+ accessors = AssetStateStoreAccessors([asset])
+
+ if method_name == "aget":
+ result = await accessors.aget("watermark")
+ assert result == "v5"
+ elif method_name == "aset":
+ await accessors.aset("watermark", "2026-05-01")
+ elif method_name == "adelete":
+ await accessors.adelete("watermark")
+ else:
+ await accessors.aclear()
+
+ mock_supervisor_comms.asend.assert_called_once_with(expected_message)
+
+ @pytest.mark.asyncio
+ async def test_double_reference_raises_for_async_accessor(self):
+ a1 = Asset(name="asset_one", uri="s3://one")
+ a2 = Asset(name="asset_two", uri="s3://two")
+
+ with pytest.raises(ValueError, match="2 concrete inlets and outlets"):
+ await AssetStateStoreAccessors([a1, a2]).aget("watermark")
+
class InMemoryStoreBackend(BaseStoreBackend):
"""Simple in-memory test backend."""
@@ -2043,3 +2237,32 @@ class TestAssetStateStoreAccessorWithCustomBackend:
assert "watermark" not in backend._actual_key_value_store
assert "file_count" not in backend._actual_key_value_store
mock_supervisor_comms.send.assert_any_call(ClearAssetStateStoreByName(name=self.ASSET_NAME))
+
+ @pytest.mark.asyncio
+ async def test_aset_sends_reference_not_value(self, mock_supervisor_comms,
backend):
+ """aset() stores actual value in backend and sends mem:// reference
via comms."""
+ mock_supervisor_comms.asend.return_value = OKResponse(ok=True)
+
+ await AssetStateStoreAccessor(name=self.ASSET_NAME).aset("watermark",
"2026-05-01")
+
+ expected_ref = f"mem://{self.ASSET_NAME}/watermark"
+ mock_supervisor_comms.asend.assert_called_once_with(
+ SetAssetStateStoreByName(
+ name=self.ASSET_NAME,
+ key="watermark",
+ value=_wrap_external_ref(expected_ref),
+ )
+ )
+ assert backend._actual_key_value_store["watermark"] == "2026-05-01"
+ assert backend.reference["watermark"] == expected_ref
+
+ @pytest.mark.asyncio
+ async def test_aget_resolves_reference_to_actual_value(self,
mock_supervisor_comms, backend):
+ """aget() fetches mem:// reference from DB, resolves it to actual
value via backend."""
+ ref = _wrap_external_ref(f"mem://{self.ASSET_NAME}/watermark")
+ backend._actual_key_value_store["watermark"] = "2026-05-01"
+ mock_supervisor_comms.asend.return_value =
AssetStateStoreResult(value=ref)
+
+ result = await
AssetStateStoreAccessor(name=self.ASSET_NAME).aget("watermark")
+
+ assert result == "2026-05-01"
