Add async asset store accessors for async tasks and watcher triggers

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

task-sdk/src/airflow/sdk/execution_time/context.py

@@ -751,12 +751,25 @@ 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.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:
         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:
         if isinstance(resp, ErrorResponse) and resp.error != ErrorType.ASSET_STORE_NOT_FOUND:
             raise AirflowRuntimeError(resp)
         if isinstance(resp, AssetStateStoreResult):
@@ -780,6 +793,15 @@ 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.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:
         if value is None:
             raise ValueError("Cannot set value as None")
 
@@ -808,39 +830,62 @@ def set(self, key: str, value: JsonValue) -> None:
             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.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:
+        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.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:
         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:
@@ -851,7 +896,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:
@@ -903,18 +949,34 @@ def get(self, key: str, default: JsonValue = None) -> JsonValue:
         """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)}]>"