This is an automated email from the ASF dual-hosted git repository.
jason810496 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 77c0f6e35a8 Implement async state store methods for object storage
backend (#72131)
77c0f6e35a8 is described below
commit 77c0f6e35a88443b7bc2474b8b28a92f7d2f7163
Author: Jason(Zhe-You) Liu <[email protected]>
AuthorDate: Thu Sep 10 07:12:26 2026 +0800
Implement async state store methods for object storage backend (#72131)
* Implement async state store methods for object storage backend
* Document full async state store backend interface
---
.../task-and-asset-state-store.rst | 20 ++++++
.../providers/common/io/state_store/backend.py | 12 ++--
.../unit/common/io/state_store/test_backend.py | 84 ++++++++++++++++++++++
3 files changed, 112 insertions(+), 4 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 5ae6c0b1b33..04eb7174fed 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
@@ -138,6 +138,26 @@ 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)
+
+ async def aset(self, scope, key, value, *, expires_at=None,
session=None):
+ await asyncio.to_thread(self.set, scope, key, value,
expires_at=expires_at, session=session)
+
+ async def adelete(self, scope, key, *, session=None):
+ await asyncio.to_thread(self.delete, scope, key, session=session)
+
+ async def aclear(self, scope, *, all_map_indices=False, session=None):
+ await asyncio.to_thread(self.clear, scope,
all_map_indices=all_map_indices, session=session)
+
: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``:
diff --git
a/providers/common/io/src/airflow/providers/common/io/state_store/backend.py
b/providers/common/io/src/airflow/providers/common/io/state_store/backend.py
index dd9ad29cafe..a6b7a690510 100644
--- a/providers/common/io/src/airflow/providers/common/io/state_store/backend.py
+++ b/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 @@ class StateStoreObjectStorageBackend(BaseStoreBackend):
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 @@ class StateStoreObjectStorageBackend(BaseStoreBackend):
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)
diff --git
a/providers/common/io/tests/unit/common/io/state_store/test_backend.py
b/providers/common/io/tests/unit/common/io/state_store/test_backend.py
index 6e7c634b85f..f59c2e96441 100644
--- a/providers/common/io/tests/unit/common/io/state_store/test_backend.py
+++ b/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 @@ class TestStateStoreObjectStorageBackend:
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
