This is an automated email from the ASF dual-hosted git repository.
guan404ming 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 e4ca33ee9b9 Add async task store accessors for async tasks (#68232)
e4ca33ee9b9 is described below
commit e4ca33ee9b9382d9ab9542aec27e03944684902d
Author: Guan-Ming Chiu <[email protected]>
AuthorDate: Thu Jul 9 14:45:17 2026 +0800
Add async task store accessors for async tasks (#68232)
* Add async task store accessors for async tasks
* Add adelete and aclear
Signed-off-by: Guan-Ming (Wesley) Chiu
<[email protected]>
* Handle review comments
Signed-off-by: Guan-Ming (Wesley) Chiu
<[email protected]>
---------
Signed-off-by: Guan-Ming (Wesley) Chiu
<[email protected]>
---
.../docs/core-concepts/resumable-tasks.rst | 5 +-
.../docs/core-concepts/task-state-store.rst | 41 ++++-
task-sdk/src/airflow/sdk/execution_time/context.py | 125 ++++++++------
.../tests/task_sdk/execution_time/test_context.py | 190 +++++++++++++++++++++
4 files changed, 303 insertions(+), 58 deletions(-)
diff --git a/airflow-core/docs/core-concepts/resumable-tasks.rst
b/airflow-core/docs/core-concepts/resumable-tasks.rst
index ecb2bcaf4c0..703bf445362 100644
--- a/airflow-core/docs/core-concepts/resumable-tasks.rst
+++ b/airflow-core/docs/core-concepts/resumable-tasks.rst
@@ -143,8 +143,9 @@ operations (HTTP requests, database queries, file reads)
within a single task ex
blocking the event loop while waiting for each one.
The worker slot is held for the full duration of the task. Async tasks are not
designed for
-long external waits or crash recovery; they are designed for high-throughput
I/O work that
-completes within a single execution.
+long external waits; they are designed for high-throughput I/O work that
completes within a
+single execution. To survive a worker crash, an async task can checkpoint its
progress to
+:doc:`task state store <task-state-store>` with ``aget``/``aset`` and resume
from the checkpoint on retry.
For guidance on when to use async tasks versus deferrable operators, see
:doc:`task-sdk:deferred-vs-async-operators`.
diff --git a/airflow-core/docs/core-concepts/task-state-store.rst
b/airflow-core/docs/core-concepts/task-state-store.rst
index 7702fb90f68..cd5ffbc8b5d 100644
--- a/airflow-core/docs/core-concepts/task-state-store.rst
+++ b/airflow-core/docs/core-concepts/task-state-store.rst
@@ -22,6 +22,8 @@
intra
Intra
checkpointing
+ Checkpointing
+ accessors
Task State Store
================
@@ -30,7 +32,7 @@ Task State Store
Task store is a persistent key/value store scoped to a single task instance
(``dag_id`` + ``run_id`` + ``task_id`` + ``map_index``). It survives worker
crashes and task retries within the same Dag run, making it suitable for
storing external job IDs, intra-task checkpoints, and progress metadata.
-Data persisted via task state store is accessed through the task context via
``context["task_state_store"]`` and exposes four methods: ``get``, ``set``,
``delete``, and ``clear``.
+Data persisted via task state store is accessed through the task context via
``context["task_state_store"]`` and exposes the synchronous methods ``get``,
``set``, ``delete``, and ``clear``, plus the async counterparts ``aget``,
``aset``, ``adelete``, and ``aclear`` for use inside ``async`` tasks.
Accessing task state store
@@ -133,6 +135,19 @@ Deletes *all* task state store keys for this task instance.
task_state_store.clear()
+``aget``, ``aset``, ``adelete``, ``aclear``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Async counterparts of ``get``, ``set``, ``delete``, and ``clear`` for use
inside ``async`` tasks. They take the same arguments and behave identically to
their synchronous siblings, but ``await`` the round-trip to the API server
instead of blocking the event loop, so a coroutine can checkpoint its progress
without stalling other concurrent work.
+
+.. code-block:: python
+
+ value = await task_state_store.aget("job_id", default="123456789")
+ await task_state_store.aset("job_id", job_id, retention=NEVER_EXPIRE)
+ await task_state_store.adelete("job_id")
+ await task_state_store.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.
Some Example Use Cases
----------------------
@@ -197,6 +212,30 @@ For tasks that process paginated or batched data, store
the last-completed offse
On a retry, the task reads ``last_page`` and skips pages that were already
processed.
+Checkpointing from an async task
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+An ``async`` task gains concurrency by awaiting I/O instead of blocking on it.
To keep that benefit while checkpointing, use the async accessors
``aget``/``aset`` so the task store round-trip does not stall the event loop.
This is the same paginated-ingest pattern as above, written for a coroutine
that paginates through an external API.
+
+.. code-block:: python
+
+ from airflow.sdk import DAG, task
+
+ with DAG("async_paginated_ingest", schedule="@daily"):
+
+ @task
+ async def ingest_pages(**context):
+ task_state_store = context["task_state_store"]
+ last_page = await task_state_store.aget("last_page")
+
+ start_page = last_page + 1 if last_page is not None else 1
+
+ for page in range(start_page, total_pages + 1):
+ await fetch_and_load(page)
+ await task_state_store.aset("last_page", page)
+
+After a crash the task resumes from the stored ``last_page`` instead of
restarting at page 1, without ever blocking the event loop.
+
Progress metadata
~~~~~~~~~~~~~~~~~
diff --git a/task-sdk/src/airflow/sdk/execution_time/context.py
b/task-sdk/src/airflow/sdk/execution_time/context.py
index 7b89164f895..2828525a7bb 100644
--- a/task-sdk/src/airflow/sdk/execution_time/context.py
+++ b/task-sdk/src/airflow/sdk/execution_time/context.py
@@ -52,6 +52,37 @@ from airflow.sdk.exceptions import (
AirflowSecretsBackendAccessDenied,
ErrorType,
)
+from airflow.sdk.execution_time.comms import (
+ AssetsByAliasResult,
+ AssetStateStoreResult,
+ ClearAssetStateStoreByName,
+ ClearAssetStateStoreByUri,
+ ClearTaskStateStore,
+ DeleteAssetStateStoreByName,
+ DeleteAssetStateStoreByUri,
+ DeleteTaskStateStore,
+ DeleteVariable,
+ ErrorResponse,
+ GetAssetByName,
+ GetAssetByUri,
+ GetAssetEventByAsset,
+ GetAssetEventByAssetAlias,
+ GetAssetsByAlias,
+ GetAssetStateStoreByName,
+ GetAssetStateStoreByUri,
+ GetPrevSuccessfulDagRun,
+ GetTaskStateStore,
+ GetVariableKeys,
+ PrevSuccessfulDagRunResponse,
+ PrevSuccessfulDagRunResult,
+ PutVariable,
+ SetAssetStateStoreByName,
+ SetAssetStateStoreByUri,
+ SetTaskStateStore,
+ TaskStateStoreResult,
+ ToSupervisor,
+ VariableKeysResult,
+)
from airflow.sdk.log import amask_secret, mask_secret
if TYPE_CHECKING:
@@ -70,7 +101,6 @@ if TYPE_CHECKING:
AssetResult,
ConnectionResult,
OKResponse,
- PrevSuccessfulDagRunResponse,
ReceiveMsgType,
VariableResult,
)
@@ -140,7 +170,6 @@ T = TypeVar("T")
def _process_connection_result_conn(conn_result: ReceiveMsgType | None) ->
Connection:
from airflow.sdk.definitions.connection import Connection
- from airflow.sdk.execution_time.comms import ErrorResponse
if isinstance(conn_result, ErrorResponse):
raise AirflowRuntimeError(conn_result)
@@ -333,7 +362,6 @@ def _get_variable(key: str, deserialize_json: bool) -> Any:
# If no backend found the variable, raise a not found error (mirrors
_get_connection)
from airflow.sdk.exceptions import AirflowRuntimeError, ErrorType
- from airflow.sdk.execution_time.comms import ErrorResponse
raise AirflowRuntimeError(
ErrorResponse(error=ErrorType.VARIABLE_NOT_FOUND, detail={"message":
f"Variable {key} not found"})
@@ -345,11 +373,6 @@ _VARIABLE_KEYS_PAGE_SIZE = 1000
def _get_variable_keys(prefix: str | None = None) -> list[str]:
from airflow.sdk.exceptions import AirflowRuntimeError
- from airflow.sdk.execution_time.comms import (
- ErrorResponse,
- GetVariableKeys,
- VariableKeysResult,
- )
from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
all_keys: list[str] = []
@@ -378,7 +401,6 @@ def _set_variable(key: str, value: Any, description: str |
None = None, serializ
import json
from airflow.sdk.execution_time.cache import SecretCache
- from airflow.sdk.execution_time.comms import PutVariable
from airflow.sdk.execution_time.secrets.execution_api import
ExecutionAPISecretsBackend
from airflow.sdk.execution_time.supervisor import
ensure_secrets_backend_loaded
from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
@@ -425,7 +447,6 @@ def _delete_variable(key: str) -> None:
# will make that module depend on Task SDK, which is not ideal because
we intend to
# keep Task SDK as a separate package than execution time mods.
from airflow.sdk.execution_time.cache import SecretCache
- from airflow.sdk.execution_time.comms import DeleteVariable
from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
msg = SUPERVISOR_COMMS.send(DeleteVariable(key=key))
@@ -544,11 +565,19 @@ class TaskStateStoreAccessor:
``datetime`` is not JSON-serializable; store it as
``value.isoformat()`` and
parse it back with ``datetime.fromisoformat(result)``.
"""
- from airflow.sdk.execution_time.comms import ErrorResponse,
GetTaskStateStore, TaskStateStoreResult
from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
resp = SUPERVISOR_COMMS.send(GetTaskStateStore(ti_id=self._ti_id,
key=key))
+ return self._extract_get_response(resp, key, default)
+
+ async def aget(self, key: str, default: JsonValue = None) -> JsonValue:
+ """Async version of :meth:`get` that awaits instead of blocking the
event loop."""
+ from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+
+ resp = await
SUPERVISOR_COMMS.asend(GetTaskStateStore(ti_id=self._ti_id, key=key))
+ return self._extract_get_response(resp, key, default)
+ def _extract_get_response(self, resp: Any, key: str, default: JsonValue)
-> JsonValue:
if isinstance(resp, ErrorResponse) and resp.error !=
ErrorType.TASK_STORE_NOT_FOUND:
raise AirflowRuntimeError(resp)
if isinstance(resp, TaskStateStoreResult):
@@ -579,9 +608,19 @@ class TaskStateStoreAccessor:
- ``NEVER_EXPIRE`` — key never expires, regardless of the global
config and is skipped by garbage collection.
- ``None`` (default) — use the global ``[state_store]
default_retention_days`` config.
"""
- from airflow.sdk.execution_time.comms import SetTaskStateStore
from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+ SUPERVISOR_COMMS.send(self._build_set_message(key, value, retention))
+
+ async def aset(self, key: str, value: JsonValue, *, retention: timedelta |
None = None) -> None:
+ """Async version of :meth:`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,
retention))
+
+ def _build_set_message(
+ self, key: str, value: JsonValue, retention: timedelta | None
+ ) -> SetTaskStateStore:
if value is None:
raise ValueError("Cannot set value as None")
@@ -622,11 +661,10 @@ class TaskStateStoreAccessor:
limit,
)
- 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.execution_time.comms import DeleteTaskStateStore
from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
# cleanup the DB ref first, if backend cleanup fails after this, the
ref is gone and
@@ -636,9 +674,17 @@ class TaskStateStoreAccessor:
if backend is not None:
backend.delete(self._scope, key)
+ async def adelete(self, key: str) -> None:
+ """Async version of :meth:`delete` that awaits instead of blocking the
event loop."""
+ from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+
+ await SUPERVISOR_COMMS.asend(DeleteTaskStateStore(ti_id=self._ti_id,
key=key))
+ backend = _get_worker_state_store_backend()
+ if backend is not None:
+ await backend.adelete(self._scope, key)
+
def clear(self) -> None:
"""Delete all keys for this task instance."""
- from airflow.sdk.execution_time.comms import ClearTaskStateStore
from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
# cleanup the DB ref first, if backend cleanup fails after this, the
ref is gone and
@@ -648,6 +694,15 @@ class TaskStateStoreAccessor:
if backend is not None:
backend.clear(self._scope)
+ async def aclear(self) -> None:
+ """Async version of :meth:`clear` that awaits instead of blocking the
event loop."""
+ from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+
+ await SUPERVISOR_COMMS.asend(ClearTaskStateStore(ti_id=self._ti_id))
+ backend = _get_worker_state_store_backend()
+ if backend is not None:
+ await backend.aclear(self._scope)
+
def _clear_backend_only(self) -> None:
"""
Clear external storage via the worker backend without sending a comms
message.
@@ -689,13 +744,6 @@ 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
msg: ToSupervisor
@@ -725,11 +773,6 @@ 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
if value is None:
@@ -765,11 +808,6 @@ class AssetStateStoreAccessor:
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
@@ -787,11 +825,6 @@ class AssetStateStoreAccessor:
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
msg: ToSupervisor
@@ -826,7 +859,6 @@ class AssetStateStoreAccessors:
elif isinstance(inlet, AssetUriRef):
self._by_uri[inlet.uri] =
AssetStateStoreAccessor(uri=inlet.uri)
elif isinstance(inlet, AssetAlias):
- from airflow.sdk.execution_time.comms import
AssetsByAliasResult, GetAssetsByAlias
from airflow.sdk.execution_time.task_runner import
SUPERVISOR_COMMS
resp =
SUPERVISOR_COMMS.send(GetAssetsByAlias(alias_name=inlet.name))
@@ -933,12 +965,6 @@ class _AssetRefResolutionMixin:
@staticmethod
def _get_asset_from_db(name: str | None = None, uri: str | None = None) ->
Asset:
from airflow.sdk.definitions.asset import Asset
- from airflow.sdk.execution_time.comms import (
- ErrorResponse,
- GetAssetByName,
- GetAssetByUri,
- ToSupervisor,
- )
from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
msg: ToSupervisor
@@ -1132,12 +1158,6 @@ class InletEventsAccessor(Sequence["AssetEventResult"]):
@functools.cached_property
def _asset_events(self) -> list[AssetEventResult]:
- from airflow.sdk.execution_time.comms import (
- ErrorResponse,
- GetAssetEventByAsset,
- GetAssetEventByAssetAlias,
- ToSupervisor,
- )
from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
query_dict: dict[str, Any] = {
@@ -1300,11 +1320,6 @@ class TriggeringAssetEventsAccessor(
@cache # Prevent multiple API access.
def get_previous_dagrun_success(ti_id: UUID) -> PrevSuccessfulDagRunResponse:
from airflow.sdk.execution_time import task_runner
- from airflow.sdk.execution_time.comms import (
- GetPrevSuccessfulDagRun,
- PrevSuccessfulDagRunResponse,
- PrevSuccessfulDagRunResult,
- )
msg =
task_runner.SUPERVISOR_COMMS.send(GetPrevSuccessfulDagRun(ti_id=ti_id))
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 4bb88b713eb..8d6e52c4b7f 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_context.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_context.py
@@ -1405,6 +1405,161 @@ class TestTaskStateStoreAccessor:
assert result == {"rows": 123}
backend.deserialize_task_state_store_from_ref.assert_called_once_with("s3://bucket/ti_123/job_id")
+ @pytest.mark.asyncio
+ async def test_aget_returns_value(self, mock_supervisor_comms):
+ """aget awaits asend and returns the stored value, without touching
sync send."""
+ mock_supervisor_comms.asend.return_value =
TaskStateStoreResult(value="app_001")
+
+ result = await TaskStateStoreAccessor(ti_id=self.TI_ID,
scope=self.SCOPE).aget("job_id")
+
+ assert result == "app_001"
+
mock_supervisor_comms.asend.assert_called_once_with(GetTaskStateStore(ti_id=self.TI_ID,
key="job_id"))
+ 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.TASK_STORE_NOT_FOUND, detail={"key": "job_id"}
+ )
+
+ result = await TaskStateStoreAccessor(ti_id=self.TI_ID,
scope=self.SCOPE).aget(
+ "job_id", default="default-id"
+ )
+
+ assert result == "default-id"
+
+ @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 TaskStateStoreAccessor(ti_id=self.TI_ID,
scope=self.SCOPE).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 = TaskStateStoreResult(
+ value=_wrap_external_ref("s3://bucket/ti_123/job_id")
+ )
+
+ backend = MagicMock(spec=BaseStoreBackend)
+ backend.deserialize_task_state_store_from_ref.return_value = {"rows":
123}
+
+ with patch(
+
"airflow.sdk.execution_time.context._get_worker_state_store_backend",
return_value=backend
+ ):
+ result = await TaskStateStoreAccessor(ti_id=self.TI_ID,
scope=self.SCOPE).aget("job_id")
+
+ assert result == {"rows": 123}
+
backend.deserialize_task_state_store_from_ref.assert_called_once_with("s3://bucket/ti_123/job_id")
+
+ @pytest.mark.asyncio
+ async def test_aset_with_global_retention(self, mock_supervisor_comms,
time_machine):
+ """aset awaits asend with the message built from the global retention
config."""
+ mock_supervisor_comms.asend.return_value = OKResponse(ok=True)
+ now = datetime(2026, 5, 14, 12, 0, 0, tzinfo=dt_timezone.utc)
+ time_machine.move_to(now, tick=False)
+
+ with conf_vars({("state_store", "default_retention_days"): "30"}):
+ await TaskStateStoreAccessor(ti_id=self.TI_ID,
scope=self.SCOPE).aset("job_id", "app_001")
+
+ mock_supervisor_comms.asend.assert_called_once_with(
+ SetTaskStateStore(
+ ti_id=self.TI_ID,
+ key="job_id",
+ value="app_001",
+ expires_at=datetime(2026, 6, 13, 12, 0, 0,
tzinfo=dt_timezone.utc),
+ )
+ )
+ 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 TaskStateStoreAccessor(ti_id=self.TI_ID,
scope=self.SCOPE).aset("job_id", 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_task_state_store_to_ref.return_value =
"s3://bucket/ti_123/job_id"
+
+ with (
+ patch(
+
"airflow.sdk.execution_time.context._get_worker_state_store_backend",
+ return_value=backend,
+ ),
+ conf_vars({("state_store", "default_retention_days"): "0"}),
+ ):
+ await TaskStateStoreAccessor(ti_id=self.TI_ID,
scope=self.SCOPE).aset("job_id", "spark_001")
+
+ mock_supervisor_comms.asend.assert_called_once_with(
+ SetTaskStateStore(
+ ti_id=self.TI_ID,
+ key="job_id",
+ value=_wrap_external_ref("s3://bucket/ti_123/job_id"),
+ expires_at=None,
+ )
+ )
+
+ @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 TaskStateStoreAccessor(ti_id=self.TI_ID,
scope=self.SCOPE).adelete("job_id")
+
+ mock_supervisor_comms.asend.assert_called_once_with(
+ DeleteTaskStateStore(ti_id=self.TI_ID, key="job_id")
+ )
+ 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 TaskStateStoreAccessor(ti_id=self.TI_ID,
scope=self.SCOPE).aclear()
+
+
mock_supervisor_comms.asend.assert_called_once_with(ClearTaskStateStore(ti_id=self.TI_ID))
+ 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 TaskStateStoreAccessor(ti_id=self.TI_ID,
scope=self.SCOPE).adelete("job_id")
+
+ backend.adelete.assert_awaited_once_with(self.SCOPE, "job_id")
+ 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 TaskStateStoreAccessor(ti_id=self.TI_ID,
scope=self.SCOPE).aclear()
+
+ backend.aclear.assert_awaited_once_with(self.SCOPE)
+ backend.clear.assert_not_called()
+
class TestAssetStateStoreAccessor:
ASSET_NAME = "debug_watcher_asset"
@@ -1850,6 +2005,41 @@ class TestTaskStateStoreAccessorWithCustomBackend:
assert "checkpoint" not in backend._actual_key_value_store
mock_supervisor_comms.send.assert_any_call(ClearTaskStateStore(ti_id=self.TI_ID))
+ @pytest.mark.asyncio
+ async def test_aset_returns_reference_to_storage(self,
mock_supervisor_comms, backend, time_machine):
+ """aset() stores actual value in backend and sends mem:// reference
via comms."""
+ mock_supervisor_comms.asend.return_value = OKResponse(ok=True)
+ expected_ref =
f"mem://{self.SCOPE.dag_id}/{self.SCOPE.run_id}/{self.SCOPE.task_id}/{self.SCOPE.map_index}/job_id"
+
+ frozen_dt = datetime(2026, 1, 1, 12, 0, 0, tzinfo=dt_timezone.utc)
+ time_machine.move_to(frozen_dt, tick=False)
+
+ await TaskStateStoreAccessor(ti_id=self.TI_ID,
scope=self.SCOPE).aset("job_id", "app_001")
+
+ mock_supervisor_comms.asend.assert_called_once_with(
+ SetTaskStateStore(
+ ti_id=self.TI_ID,
+ key="job_id",
+ value=_wrap_external_ref(expected_ref),
+ expires_at=frozen_dt + timedelta(days=30),
+ )
+ )
+ assert backend._actual_key_value_store["job_id"] == "app_001"
+ assert backend.reference["job_id"] == 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.SCOPE.dag_id}/{self.SCOPE.run_id}/{self.SCOPE.task_id}/{self.SCOPE.map_index}/job_id"
+ )
+ backend._actual_key_value_store["job_id"] = "app_001"
+ mock_supervisor_comms.asend.return_value =
TaskStateStoreResult(value=ref)
+
+ result = await TaskStateStoreAccessor(ti_id=self.TI_ID,
scope=self.SCOPE).aget("job_id")
+
+ assert result == "app_001"
+
class TestAssetStateStoreAccessorWithCustomBackend:
ASSET_NAME = "my_asset"
