From 85301932d11c7021b03380870e7c6c66d856e817 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 27 Aug 2026 05:32:15 +0000 Subject: [PATCH 1/3] Add async asset store accessors for async tasks and watcher triggers --- .../docs/core-concepts/asset-state-store.rst | 25 +- .../src/airflow/sdk/execution_time/context.py | 92 +++++-- .../task_sdk/execution_time/test_context.py | 246 +++++++++++++++++- 3 files changed, 343 insertions(+), 20 deletions(-) diff --git a/airflow-core/docs/core-concepts/asset-state-store.rst b/airflow-core/docs/core-concepts/asset-state-store.rst index 3aa7ff7719f8a..ddb3fe93fb46c 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 7ddc6edf1a3aa..4b00d6380ae68 100644 --- a/task-sdk/src/airflow/sdk/execution_time/context.py +++ b/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 :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(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 :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)) + + 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 :meth:`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 :meth:`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 :meth:`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 :meth:`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 :meth:`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 :meth:`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"" 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 31009408ce4df..c2af0fc70c665 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 @@ 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, @@ -1896,6 +1896,170 @@ def test_set_warns_when_value_exceeds_limit(self, mock_supervisor_comms): assert "max_value_storage_bytes" in mock_log.warning.call_args[0][0] mock_supervisor_comms.send.assert_called_once() + @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 = AssetStateStoreResult(value="2026-04-30T00:00:00Z") + + result = await AssetStateStoreAccessor(name=self.ASSET_NAME).aget("watermark") + + assert result == "2026-04-30T00:00:00Z" + mock_supervisor_comms.asend.assert_called_once_with( + GetAssetStateStoreByName(name=self.ASSET_NAME, key="watermark") + ) + mock_supervisor_comms.send.assert_not_called() + + @pytest.mark.asyncio + async def test_aget_by_uri(self, mock_supervisor_comms): + mock_supervisor_comms.asend.return_value = AssetStateStoreResult(value="2026-04-30T00:00:00Z") + + result = await AssetStateStoreAccessor(uri=self.ASSET_URI).aget("watermark") + + assert result == "2026-04-30T00:00:00Z" + mock_supervisor_comms.asend.assert_called_once_with( + GetAssetStateStoreByUri(uri=self.ASSET_URI, key="watermark") + ) + + @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.asyncio + async def test_aset_operation(self, mock_supervisor_comms): + mock_supervisor_comms.asend.return_value = OKResponse(ok=True) + + await AssetStateStoreAccessor(name=self.ASSET_NAME).aset("watermark", "2026-04-30T00:00:00Z") + + mock_supervisor_comms.asend.assert_called_once_with( + SetAssetStateStoreByName(name=self.ASSET_NAME, key="watermark", value="2026-04-30T00:00:00Z") + ) + mock_supervisor_comms.send.assert_not_called() + + @pytest.mark.asyncio + async def test_aset_by_uri(self, mock_supervisor_comms): + mock_supervisor_comms.asend.return_value = OKResponse(ok=True) + + await AssetStateStoreAccessor(uri=self.ASSET_URI).aset("watermark", "2026-04-30T00:00:00Z") + + mock_supervisor_comms.asend.assert_called_once_with( + SetAssetStateStoreByUri(uri=self.ASSET_URI, key="watermark", value="2026-04-30T00:00:00Z") + ) + + @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" @@ -2059,6 +2223,57 @@ def test_outlet_alias_is_ignored(self, mock_supervisor_comms): assert accessors._total == 0 mock_supervisor_comms.send.assert_not_called() + @pytest.mark.asyncio + async def test_aget_single_inlet_simplified(self, mock_supervisor_comms): + asset = Asset(name=self.ASSET_NAME, uri=f"s3://{self.ASSET_NAME}") + mock_supervisor_comms.asend.return_value = AssetStateStoreResult(value="v5") + + result = await AssetStateStoreAccessors([asset]).aget("watermark") + + assert result == "v5" + mock_supervisor_comms.asend.assert_called_once_with( + GetAssetStateStoreByName(name=self.ASSET_NAME, key="watermark") + ) + + @pytest.mark.asyncio + async def test_aset_single_inlet_simplified(self, mock_supervisor_comms): + asset = Asset(name=self.ASSET_NAME, uri=f"s3://{self.ASSET_NAME}") + mock_supervisor_comms.asend.return_value = OKResponse(ok=True) + + await AssetStateStoreAccessors([asset]).aset("watermark", "2026-05-01") + + mock_supervisor_comms.asend.assert_called_once_with( + SetAssetStateStoreByName(name=self.ASSET_NAME, key="watermark", value="2026-05-01") + ) + + @pytest.mark.asyncio + async def test_adelete_single_inlet_simplified(self, mock_supervisor_comms): + asset = Asset(name=self.ASSET_NAME, uri=f"s3://{self.ASSET_NAME}") + mock_supervisor_comms.asend.return_value = OKResponse(ok=True) + + await AssetStateStoreAccessors([asset]).adelete("watermark") + + mock_supervisor_comms.asend.assert_called_once_with( + DeleteAssetStateStoreByName(name=self.ASSET_NAME, key="watermark") + ) + + @pytest.mark.asyncio + async def test_aclear_single_inlet_simplified(self, mock_supervisor_comms): + asset = Asset(name=self.ASSET_NAME, uri=f"s3://{self.ASSET_NAME}") + mock_supervisor_comms.asend.return_value = OKResponse(ok=True) + + await AssetStateStoreAccessors([asset]).aclear() + + mock_supervisor_comms.asend.assert_called_once_with(ClearAssetStateStoreByName(name=self.ASSET_NAME)) + + @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.""" @@ -2278,3 +2493,32 @@ def test_clear_purges_all_from_backend_and_clears_db(self, mock_supervisor_comms 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" From 613737849e29b2e04c27c6c55d73494c50524cc0 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 9 Sep 2026 10:53:18 +0000 Subject: [PATCH 2/3] Clarify async asset state accessor docstrings --- .../src/airflow/sdk/execution_time/context.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/task-sdk/src/airflow/sdk/execution_time/context.py b/task-sdk/src/airflow/sdk/execution_time/context.py index 4b00d6380ae68..f23179a50d660 100644 --- a/task-sdk/src/airflow/sdk/execution_time/context.py +++ b/task-sdk/src/airflow/sdk/execution_time/context.py @@ -755,7 +755,7 @@ def get(self, key: str, default: JsonValue = None) -> JsonValue: 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.""" + """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)) @@ -796,7 +796,7 @@ def set(self, key: str, value: JsonValue) -> None: SUPERVISOR_COMMS.send(self._build_set_message(key, value)) async def aset(self, key: str, value: JsonValue) -> None: - """Async version of :meth:`set` that awaits instead of blocking the event loop.""" + """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)) @@ -844,7 +844,7 @@ def delete(self, key: str) -> None: backend.delete(AssetScope(name=self._name, uri=self._uri), key) async def adelete(self, key: str) -> None: - """Async version of :meth:`delete` that awaits instead of blocking the event loop.""" + """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)) @@ -871,7 +871,7 @@ def clear(self) -> None: backend.clear(AssetScope(name=self._name, uri=self._uri)) async def aclear(self) -> None: - """Async version of :meth:`clear` that awaits instead of blocking the event loop.""" + """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()) @@ -950,7 +950,7 @@ def get(self, key: str, default: JsonValue = None) -> JsonValue: return self._single_accessor().get(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.""" + """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: @@ -958,7 +958,7 @@ def set(self, key: str, value: JsonValue) -> None: self._single_accessor().set(key, value) async def aset(self, key: str, value: JsonValue) -> None: - """Async version of :meth:`set` that awaits instead of blocking the event loop.""" + """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: @@ -966,7 +966,7 @@ def delete(self, key: str) -> None: self._single_accessor().delete(key) async def adelete(self, key: str) -> None: - """Async version of :meth:`delete` that awaits instead of blocking the event loop.""" + """Async version of `delete` that awaits instead of blocking the event loop.""" await self._single_accessor().adelete(key) def clear(self) -> None: @@ -974,7 +974,7 @@ def clear(self) -> None: self._single_accessor().clear() async def aclear(self) -> None: - """Async version of :meth:`clear` that awaits instead of blocking the event loop.""" + """Async version of `clear` that awaits instead of blocking the event loop.""" await self._single_accessor().aclear() def __repr__(self) -> str: From 5e6a3e7bed77d3f2fcc5e1f27010a86ce2e16090 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 9 Sep 2026 10:55:28 +0000 Subject: [PATCH 3/3] Parametrize async asset state accessor tests --- .../task_sdk/execution_time/test_context.py | 119 ++++++++---------- 1 file changed, 49 insertions(+), 70 deletions(-) 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 c2af0fc70c665..f28d8e9c3e55e 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_context.py +++ b/task-sdk/tests/task_sdk/execution_time/test_context.py @@ -1896,30 +1896,24 @@ def test_set_warns_when_value_exceeds_limit(self, mock_supervisor_comms): 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): + 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 AssetStateStoreAccessor(name=self.ASSET_NAME).aget("watermark") + result = await accessor.aget("watermark") assert result == "2026-04-30T00:00:00Z" - mock_supervisor_comms.asend.assert_called_once_with( - GetAssetStateStoreByName(name=self.ASSET_NAME, key="watermark") - ) + mock_supervisor_comms.asend.assert_called_once_with(expected_message) mock_supervisor_comms.send.assert_not_called() - @pytest.mark.asyncio - async def test_aget_by_uri(self, mock_supervisor_comms): - mock_supervisor_comms.asend.return_value = AssetStateStoreResult(value="2026-04-30T00:00:00Z") - - result = await AssetStateStoreAccessor(uri=self.ASSET_URI).aget("watermark") - - assert result == "2026-04-30T00:00:00Z" - mock_supervisor_comms.asend.assert_called_once_with( - GetAssetStateStoreByUri(uri=self.ASSET_URI, key="watermark") - ) - @pytest.mark.asyncio async def test_aget_returns_default_when_key_missing(self, mock_supervisor_comms): mock_supervisor_comms.asend.return_value = ErrorResponse( @@ -1961,27 +1955,26 @@ async def test_aget_with_custom_backend_removes_decoration_marker(self, mock_sup "s3://bucket/assets/orders/watermark" ) + @pytest.mark.parametrize("lookup_by", ["name", "uri"]) @pytest.mark.asyncio - async def test_aset_operation(self, mock_supervisor_comms): + 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 AssetStateStoreAccessor(name=self.ASSET_NAME).aset("watermark", "2026-04-30T00:00:00Z") + await accessor.aset("watermark", "2026-04-30T00:00:00Z") - mock_supervisor_comms.asend.assert_called_once_with( - SetAssetStateStoreByName(name=self.ASSET_NAME, key="watermark", value="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_by_uri(self, mock_supervisor_comms): - mock_supervisor_comms.asend.return_value = OKResponse(ok=True) - - await AssetStateStoreAccessor(uri=self.ASSET_URI).aset("watermark", "2026-04-30T00:00:00Z") - - mock_supervisor_comms.asend.assert_called_once_with( - SetAssetStateStoreByUri(uri=self.ASSET_URI, key="watermark", value="2026-04-30T00:00:00Z") - ) - @pytest.mark.asyncio async def test_aset_none_raises(self, mock_supervisor_comms): with pytest.raises(ValueError, match="Cannot set value as None"): @@ -2223,48 +2216,34 @@ def test_outlet_alias_is_ignored(self, mock_supervisor_comms): 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_aget_single_inlet_simplified(self, mock_supervisor_comms): - asset = Asset(name=self.ASSET_NAME, uri=f"s3://{self.ASSET_NAME}") - mock_supervisor_comms.asend.return_value = AssetStateStoreResult(value="v5") - - result = await AssetStateStoreAccessors([asset]).aget("watermark") - - assert result == "v5" - mock_supervisor_comms.asend.assert_called_once_with( - GetAssetStateStoreByName(name=self.ASSET_NAME, key="watermark") - ) - - @pytest.mark.asyncio - async def test_aset_single_inlet_simplified(self, mock_supervisor_comms): - asset = Asset(name=self.ASSET_NAME, uri=f"s3://{self.ASSET_NAME}") - mock_supervisor_comms.asend.return_value = OKResponse(ok=True) - - await AssetStateStoreAccessors([asset]).aset("watermark", "2026-05-01") - - mock_supervisor_comms.asend.assert_called_once_with( - SetAssetStateStoreByName(name=self.ASSET_NAME, key="watermark", value="2026-05-01") - ) - - @pytest.mark.asyncio - async def test_adelete_single_inlet_simplified(self, mock_supervisor_comms): - asset = Asset(name=self.ASSET_NAME, uri=f"s3://{self.ASSET_NAME}") - mock_supervisor_comms.asend.return_value = OKResponse(ok=True) - - await AssetStateStoreAccessors([asset]).adelete("watermark") - - mock_supervisor_comms.asend.assert_called_once_with( - DeleteAssetStateStoreByName(name=self.ASSET_NAME, key="watermark") - ) - - @pytest.mark.asyncio - async def test_aclear_single_inlet_simplified(self, mock_supervisor_comms): + 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 = OKResponse(ok=True) - - await AssetStateStoreAccessors([asset]).aclear() + 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(ClearAssetStateStoreByName(name=self.ASSET_NAME)) + mock_supervisor_comms.asend.assert_called_once_with(expected_message) @pytest.mark.asyncio async def test_double_reference_raises_for_async_accessor(self):