Implement async state store methods for object storage backend

airflow-core/docs/administration-and-deployment/task-and-asset-state-store.rst

@@ -138,6 +138,17 @@ Each method receives a ``scope`` argument that is either a :class:`~airflow.sdk.
             elif isinstance(scope, AssetScope):
                 return self._asset_store.get(scope, key)
 
+If the storage client is synchronous, implement the async methods by offloading the sync work to a worker thread rather than calling it inline, so callers on an event loop (``async`` tasks and watcher triggers) are not blocked:
+
+.. code-block:: python
+
+    import asyncio
+
+
+    class MyBackend(BaseStoreBackend):
+        async def aget(self, scope, key, *, session=None):
+            return await asyncio.to_thread(self.get, scope, key)
+
 :class:`~airflow.sdk.state.AssetScope` has three optional fields: ``asset_id`` (integer, server-side only), ``name``, and ``uri``. At least one must be set. Server-side operations (REST API calls) provide ``asset_id``. Worker-side operations provide ``name`` or ``uri`` (workers do not have access to the integer ``asset_id``).
 
 Configure the class via ``[state_store] backend``:

providers/common/io/src/airflow/providers/common/io/state_store/backend.py

@@ -16,6 +16,7 @@
 # under the License.
 from __future__ import annotations
 
+import asyncio
 import json
 from functools import cache
 from typing import TYPE_CHECKING
@@ -187,8 +188,11 @@ def clear(
             case _:
                 raise TypeError(f"Unknown scope type: {type(scope)}")
 
+    # fsspec is synchronous even for filesystems whose transport is async underneath, so the
+    # a-prefixed methods offload to a worker thread to keep the caller's event loop free.
+    # ``session`` is unused throughout: this backend never touches the metastore.
     async def aget(self, scope: StoreScope, key: str, *, session: AsyncSession | None = None) -> str | None:
-        raise NotImplementedError
+        return await asyncio.to_thread(self.get, scope, key)
 
     async def aset(
         self,
@@ -199,15 +203,15 @@ async def aset(
         expires_at: datetime | None = None,
         session: AsyncSession | None = None,
     ) -> None:
-        raise NotImplementedError
+        await asyncio.to_thread(self.set, scope, key, value, expires_at=expires_at)
 
     async def adelete(self, scope: StoreScope, key: str, *, session: AsyncSession | None = None) -> None:
-        raise NotImplementedError
+        await asyncio.to_thread(self.delete, scope, key)
 
     async def aclear(
         self, scope: StoreScope, *, all_map_indices: bool = False, session: AsyncSession | None = None
     ) -> None:
-        raise NotImplementedError
+        await asyncio.to_thread(self.clear, scope, all_map_indices=all_map_indices)
 
     def serialize_task_state_store_to_ref(self, *, value: JsonValue, key: str, scope: TaskScope) -> str:
         serialized = json.dumps(value)

providers/common/io/tests/unit/common/io/state_store/test_backend.py

@@ -16,6 +16,9 @@
 # under the License.
 from __future__ import annotations
 
+import threading
+from unittest import mock
+
 import pytest
 
 from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS
@@ -272,3 +275,84 @@ def test_negative_threshold_raises(self, base_path):
             backend._get_threshold.cache_clear()
             with pytest.raises(ValueError, match="must be non-negative"):
                 backend._get_threshold()
+
+    @pytest.mark.asyncio
+    async def test_aset_and_aget_task(self, store, task_scope):
+        await store.aset(task_scope, "k", "hello")
+        assert await store.aget(task_scope, "k") == "hello"
+        # the async round-trip is visible to the sync API and vice versa
+        assert store.get(task_scope, "k") == "hello"
+
+    @pytest.mark.asyncio
+    async def test_aget_missing_returns_none(self, store, task_scope):
+        assert await store.aget(task_scope, "missing") is None
+
+    @pytest.mark.asyncio
+    async def test_adelete_task(self, store, task_scope):
+        store.set(task_scope, "k", "v")
+        await store.adelete(task_scope, "k")
+        assert store.get(task_scope, "k") is None
+
+    @pytest.mark.asyncio
+    async def test_adelete_missing_is_noop(self, store, task_scope):
+        await store.adelete(task_scope, "does_not_exist")
+
+    @pytest.mark.asyncio
+    async def test_aclear_task_single_map_index(self, store, task_scope):
+        store.set(task_scope, "k1", "v1")
+        store.set(task_scope, "k2", "v2")
+        await store.aclear(task_scope)
+        assert store.get(task_scope, "k1") is None
+        assert store.get(task_scope, "k2") is None
+
+    @pytest.mark.asyncio
+    async def test_aclear_task_all_map_indices(self, store):
+        scope0 = TaskScope(dag_id="d", run_id="r", task_id="t", map_index=0)
+        scope1 = TaskScope(dag_id="d", run_id="r", task_id="t", map_index=1)
+        store.set(scope0, "k", "v0")
+        store.set(scope1, "k", "v1")
+        await store.aclear(scope0, all_map_indices=True)
+        assert store.get(scope0, "k") is None
+        assert store.get(scope1, "k") is None
+
+    @pytest.mark.asyncio
+    async def test_adelete_asset(self, store, asset_scope):
+        store.set(asset_scope, "watermark", "2026-05-01")
+        await store.adelete(asset_scope, "watermark")
+        assert store.get(asset_scope, "watermark") is None
+
+    @pytest.mark.asyncio
+    async def test_aclear_asset(self, store, asset_scope):
+        store.set(asset_scope, "k1", "v1")
+        store.set(asset_scope, "k2", "v2")
+        await store.aclear(asset_scope)
+        assert store.get(asset_scope, "k1") is None
+        assert store.get(asset_scope, "k2") is None
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        ("async_call", "sync_method"),
+        [
+            (lambda store, scope: store.aget(scope, "k"), "get"),
+            (lambda store, scope: store.aset(scope, "k", "v"), "set"),
+            (lambda store, scope: store.adelete(scope, "k"), "delete"),
+            (lambda store, scope: store.aclear(scope), "clear"),
+        ],
+        ids=["aget", "aset", "adelete", "aclear"],
+    )
+    async def test_async_methods_run_blocking_work_off_the_loop_thread(
+        self, store, task_scope, async_call, sync_method
+    ):
+        loop_thread = threading.get_ident()
+        call_threads = []
+        original = getattr(store, sync_method)
+
+        def record_thread(*args, **kwargs):
+            call_threads.append(threading.get_ident())
+            return original(*args, **kwargs)
+
+        with mock.patch.object(store, sync_method, record_thread):
+            await async_call(store, task_scope)
+
+        assert call_threads
+        assert loop_thread not in call_threads