This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch v3-3-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-3-test by this push:
new 650efd91af4 [v3-3-test] AIP-103: Allow for
`get`/`set`/`delete`/`clear` to run for `AssetStateStoreAccessor` on Triggerer
(#68900) (#68966)
650efd91af4 is described below
commit 650efd91af4957ccb9ba87927feb839f4bb9e7e6
Author: github-actions[bot]
<41898282+github-actions[bot]@users.noreply.github.com>
AuthorDate: Fri Jun 26 14:55:18 2026 -0400
[v3-3-test] AIP-103: Allow for `get`/`set`/`delete`/`clear` to run for
`AssetStateStoreAccessor` on Triggerer (#68900) (#68966)
(cherry picked from commit feea09d6107bc877c675e00c91b3d42d234502d5)
Co-authored-by: Jake McGrath
<[email protected]>
Co-authored-by: Amogh Desai <[email protected]>
---
.../docs/core-concepts/asset-state-store.rst | 74 +++++++++
.../execution_api/routes/asset_state_store.py | 38 +++--
.../src/airflow/jobs/triggerer_job_runner.py | 48 ++++++
.../versions/head/test_asset_state_store.py | 72 +++++++++
airflow-core/tests/unit/jobs/test_triggerer_job.py | 117 +++++++++++++--
.../airflow/sdk/execution_time/request_handlers.py | 94 ++++++++++++
.../execution_time/test_request_handlers.py | 167 +++++++++++++++++++++
7 files changed, 587 insertions(+), 23 deletions(-)
diff --git a/airflow-core/docs/core-concepts/asset-state-store.rst
b/airflow-core/docs/core-concepts/asset-state-store.rst
index 1f019d89d22..3aa7ff7719f 100644
--- a/airflow-core/docs/core-concepts/asset-state-store.rst
+++ b/airflow-core/docs/core-concepts/asset-state-store.rst
@@ -150,6 +150,80 @@ Deletes *all* asset state store keys for the asset.
# Using context
context["asset_state_store"][my_asset].clear()
+Using ``asset_state_store`` inside a Watcher Trigger
+-----------------------------------------------------
+
+:class:`~airflow.triggers.base.BaseEventTrigger` subclasses (watcher triggers)
can read and write asset state store directly from within ``run()``. The
triggerer injects ``self.asset_state_store`` before ``run()`` is called, scoped
to the asset the trigger is watching. It is not available during ``__init__``
or ``serialize()``, only access it from within ``run()``.
+
+Unlike task-based access (where the asset is identified by an inlet or outlet
declaration), the accessor in a watcher trigger is automatically bound to the
watched asset, so no subscripting is needed.
+
+.. code-block:: python
+
+ import asyncio
+ from collections.abc import AsyncIterator
+ from typing import Any
+
+ from airflow.triggers.base import BaseEventTrigger, TriggerEvent
+
+
+ class PollEventsTrigger(BaseEventTrigger):
+ def __init__(self, source: str, waiter_delay: int, **kwargs):
+ super().__init__(**kwargs)
+ self.source = source
+ self.waiter_delay = waiter_delay
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ return (
+ f"{self.__class__.__module__}.{self.__class__.__qualname__}",
+ {"source": self.source, "waiter_delay": self.waiter_delay},
+ )
+
+ def _poll_for_new_record(self, source: str, last_seen: str) -> str |
None:
+ ... # Add logic for polling a certain source
+ return None
+
+ async def run(self) -> AsyncIterator[TriggerEvent]:
+ while True:
+ last_seen = self.asset_state_store.get("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)
+ yield TriggerEvent({"status": "success", "record_id":
new_id})
+ return
+
+ await asyncio.sleep(self.waiter_delay)
+
+The corresponding :class:`~airflow.sdk.definitions.asset.AssetWatcher` wires
the trigger to the asset:
+
+.. code-block:: python
+
+ from airflow.sdk import Asset, AssetWatcher
+
+ from my_dag.triggers import PollEventsTrigger
+
+ my_asset = Asset(
+ name="orders_api",
+ watchers=[
+ AssetWatcher(
+ name="orders_api_watcher",
+ trigger=PollEventsTrigger(source="orders", waiter_delay=30),
+ )
+ ],
+ )
+
+ ...
+
+``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.
+
+.. note::
+
+ ``self.asset_state_store`` is only available inside
:class:`~airflow.triggers.base.BaseEventTrigger` subclasses. Plain
:class:`~airflow.triggers.base.BaseTrigger` subclasses (used for task deferral)
do not have access to asset state store.
+
+
Some Example Use cases
----------------------
diff --git
a/airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_state_store.py
b/airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_state_store.py
index f602962528e..4424e4ad3f1 100644
---
a/airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_state_store.py
+++
b/airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_state_store.py
@@ -30,6 +30,7 @@ from __future__ import annotations
import json
from typing import Annotated
+from uuid import UUID
from cadwyn import VersionedAPIRouter
from fastapi import HTTPException, Query, status
@@ -49,6 +50,7 @@ from airflow.state import get_state_backend
from airflow.state.metastore import MetastoreBackend
_TIWriterFields = tuple[str, str, str, int]
+NULL_UUID = UUID(int=0)
def _fetch_ti_writer_fields(token: TIToken, session: SessionDep) ->
_TIWriterFields:
@@ -130,18 +132,30 @@ def _put_asset_state_store(
) -> None:
backend = get_state_backend()
if isinstance(backend, MetastoreBackend):
- dag_id, run_id, task_id, map_index = _fetch_ti_writer_fields(token,
session)
- backend.set_asset_state_store(
- scope,
- key,
- json.dumps(body.value),
- kind=AssetStateStoreWriterKind.TASK,
- dag_id=dag_id,
- run_id=run_id,
- task_id=task_id,
- map_index=map_index,
- session=session,
- )
+ if token.id == NULL_UUID:
+ # Since the asset state store routes do not have
`task_instance_id` in their path params, the default kicks in which
is"00000000-0000-0000-0000-000000000000"
+ backend.set_asset_state_store(
+ scope,
+ key,
+ json.dumps(body.value),
+ kind=AssetStateStoreWriterKind.WATCHER,
+ session=session,
+ )
+ else:
+ ti_fields = _fetch_ti_writer_fields(token, session)
+ dag_id, run_id, task_id, map_index = ti_fields
+
+ backend.set_asset_state_store(
+ scope,
+ key,
+ json.dumps(body.value),
+ kind=AssetStateStoreWriterKind.TASK,
+ dag_id=dag_id,
+ run_id=run_id,
+ task_id=task_id,
+ map_index=map_index,
+ session=session,
+ )
else:
backend.set(scope, key, json.dumps(body.value), session=session)
diff --git a/airflow-core/src/airflow/jobs/triggerer_job_runner.py
b/airflow-core/src/airflow/jobs/triggerer_job_runner.py
index b29978e52a9..eccd9d2fc04 100644
--- a/airflow-core/src/airflow/jobs/triggerer_job_runner.py
+++ b/airflow-core/src/airflow/jobs/triggerer_job_runner.py
@@ -60,13 +60,20 @@ from airflow.observability.metrics import stats_utils
from airflow.sdk.api.datamodels._generated import HITLDetailResponse
from airflow.sdk.definitions.asset import Asset
from airflow.sdk.execution_time.comms import (
+ AssetStateStoreResult,
+ ClearAssetStateStoreByName,
+ ClearAssetStateStoreByUri,
CommsDecoder,
ConnectionResult,
DagRunStateResult,
+ DeleteAssetStateStoreByName,
+ DeleteAssetStateStoreByUri,
DeleteVariable,
DeleteXCom,
DRCount,
ErrorResponse,
+ GetAssetStateStoreByName,
+ GetAssetStateStoreByUri,
GetConnection,
GetDagRunState,
GetDRCount,
@@ -80,6 +87,8 @@ from airflow.sdk.execution_time.comms import (
MaskSecret,
OKResponse,
PutVariable,
+ SetAssetStateStoreByName,
+ SetAssetStateStoreByUri,
SetXCom,
TaskStatesResult,
TICount,
@@ -92,8 +101,14 @@ from airflow.sdk.execution_time.comms import (
)
from airflow.sdk.execution_time.context import AssetStateStoreAccessors
from airflow.sdk.execution_time.request_handlers import (
+ handle_clear_asset_state_store_by_name,
+ handle_clear_asset_state_store_by_uri,
+ handle_delete_asset_state_store_by_name,
+ handle_delete_asset_state_store_by_uri,
handle_delete_variable,
handle_delete_xcom,
+ handle_get_asset_state_store_by_name,
+ handle_get_asset_state_store_by_uri,
handle_get_connection,
handle_get_dag_run_state,
handle_get_dr_count,
@@ -105,6 +120,8 @@ from airflow.sdk.execution_time.request_handlers import (
handle_get_xcom,
handle_mask_secret,
handle_put_variable,
+ handle_set_asset_state_store_by_name,
+ handle_set_asset_state_store_by_uri,
handle_set_xcom,
)
from airflow.sdk.execution_time.supervisor import WatchedSubprocess,
make_buffered_socket_reader
@@ -353,6 +370,7 @@ ToTriggerRunner = Annotated[
| DRCount
| TICount
| TaskStatesResult
+ | AssetStateStoreResult
| HITLDetailResponseResult
| ErrorResponse
| OKResponse,
@@ -376,6 +394,14 @@ ToTriggerSupervisor = Annotated[
| SetXCom
| GetTICount
| GetTaskStates
+ | ClearAssetStateStoreByName
+ | ClearAssetStateStoreByUri
+ | DeleteAssetStateStoreByName
+ | DeleteAssetStateStoreByUri
+ | GetAssetStateStoreByName
+ | GetAssetStateStoreByUri
+ | SetAssetStateStoreByName
+ | SetAssetStateStoreByUri
| GetDagRunState
| GetDRCount
| GetPreviousTI
@@ -622,6 +648,28 @@ class TriggerRunnerSupervisor(WatchedSubprocess):
resp =
HITLDetailResponseResult.from_api_response(response=api_resp)
elif isinstance(msg, MaskSecret):
handle_mask_secret(msg)
+ elif isinstance(msg, ClearAssetStateStoreByName):
+ handle_clear_asset_state_store_by_name(self.client, msg)
+ resp = OKResponse(ok=True)
+ elif isinstance(msg, ClearAssetStateStoreByUri):
+ handle_clear_asset_state_store_by_uri(self.client, msg)
+ resp = OKResponse(ok=True)
+ elif isinstance(msg, DeleteAssetStateStoreByName):
+ handle_delete_asset_state_store_by_name(self.client, msg)
+ resp = OKResponse(ok=True)
+ elif isinstance(msg, DeleteAssetStateStoreByUri):
+ handle_delete_asset_state_store_by_uri(self.client, msg)
+ resp = OKResponse(ok=True)
+ elif isinstance(msg, GetAssetStateStoreByName):
+ resp, dump_opts =
handle_get_asset_state_store_by_name(self.client, msg)
+ elif isinstance(msg, GetAssetStateStoreByUri):
+ resp, dump_opts = handle_get_asset_state_store_by_uri(self.client,
msg)
+ elif isinstance(msg, SetAssetStateStoreByName):
+ handle_set_asset_state_store_by_name(self.client, msg)
+ resp = OKResponse(ok=True)
+ elif isinstance(msg, SetAssetStateStoreByUri):
+ handle_set_asset_state_store_by_uri(self.client, msg)
+ resp = OKResponse(ok=True)
else:
raise ValueError(f"Unknown message type {type(msg)}")
diff --git
a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_asset_state_store.py
b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_asset_state_store.py
index 57b46233653..86a64b34a9a 100644
---
a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_asset_state_store.py
+++
b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_asset_state_store.py
@@ -380,3 +380,75 @@ class TestInactiveAssetRejected:
def test_get_inactive_asset_by_uri_returns_404(self, client: TestClient,
inactive_asset: AssetModel):
response = client.get(_BY_URI_VALUE, params={"uri":
inactive_asset.uri, "key": "watermark"})
assert response.status_code == 404
+
+
+class TestPutAssetStateWatcherWrite:
+ """PUT routes record WATCHER kind when no TI backs the token
(WATCHER/triggerer path)."""
+
+ @pytest.fixture(autouse=True)
+ def setup_null_uuid_auth(self, exec_app: FastAPI):
+ from uuid import UUID
+
+ async def _watcher_auth(request: Request) -> TIToken:
+ return TIToken(
+ id=UUID("00000000-0000-0000-0000-000000000000"),
claims=TIClaims(scope="execution")
+ )
+
+ exec_app.dependency_overrides[require_auth] = _watcher_auth
+
+ def test_put_by_name_creates_row(self, client: TestClient, asset:
AssetModel, session):
+ response = client.put(
+ _BY_NAME_VALUE, params={"name": asset.name, "key": "watermark"},
json={"value": "2026-01-01"}
+ )
+
+ assert response.status_code == 204
+
+ row = session.scalar(
+ select(AssetStateStoreModel).where(
+ AssetStateStoreModel.asset_id == asset.id,
+ AssetStateStoreModel.key == "watermark",
+ )
+ )
+
+ assert row is not None
+ assert row.value == '"2026-01-01"'
+
+ def test_put_by_name_records_watcher_kind(self, client: TestClient, asset:
AssetModel, session):
+ client.put(
+ _BY_NAME_VALUE, params={"name": asset.name, "key": "watermark"},
json={"value": "2026-01-01"}
+ )
+
+ row = session.scalar(
+ select(AssetStateStoreModel).where(
+ AssetStateStoreModel.asset_id == asset.id,
+ AssetStateStoreModel.key == "watermark",
+ )
+ )
+
+ assert row is not None
+ assert row.last_updated_by_kind == "watcher"
+ assert row.last_updated_by_dag_id is None
+ assert row.last_updated_by_run_id is None
+ assert row.last_updated_by_task_id is None
+ assert row.last_updated_by_map_index is None
+
+ def test_put_by_uri_records_watcher_kind(self, client: TestClient, asset:
AssetModel, session):
+ response = client.put(
+ _BY_URI_VALUE, params={"uri": asset.uri, "key": "watermark"},
json={"value": "2026-06-01"}
+ )
+
+ assert response.status_code == 204
+
+ row = session.scalar(
+ select(AssetStateStoreModel).where(
+ AssetStateStoreModel.asset_id == asset.id,
+ AssetStateStoreModel.key == "watermark",
+ )
+ )
+
+ assert row is not None
+ assert row.last_updated_by_kind == "watcher"
+ assert row.last_updated_by_dag_id is None
+ assert row.last_updated_by_run_id is None
+ assert row.last_updated_by_task_id is None
+ assert row.last_updated_by_map_index is None
diff --git a/airflow-core/tests/unit/jobs/test_triggerer_job.py
b/airflow-core/tests/unit/jobs/test_triggerer_job.py
index 0bfdbc171c2..9ad31b4e5c2 100644
--- a/airflow-core/tests/unit/jobs/test_triggerer_job.py
+++ b/airflow-core/tests/unit/jobs/test_triggerer_job.py
@@ -76,10 +76,21 @@ from airflow.providers.standard.operators.python import
PythonOperator
from airflow.providers.standard.triggers.file import FileDeleteTrigger
from airflow.providers.standard.triggers.temporal import DateTimeTrigger,
TimeDeltaTrigger
from airflow.sdk import DAG, Asset, BaseHook, BaseOperator
+from airflow.sdk.api.client import Client
+from airflow.sdk.api.datamodels._generated import AssetStateStoreResponse
+from airflow.sdk.exceptions import ErrorType
from airflow.sdk.execution_time.comms import (
AssetStateStoreResult,
+ ClearAssetStateStoreByName,
+ ClearAssetStateStoreByUri,
+ DeleteAssetStateStoreByName,
+ DeleteAssetStateStoreByUri,
+ ErrorResponse,
GetAssetStateStoreByName,
+ GetAssetStateStoreByUri,
+ OKResponse,
SetAssetStateStoreByName,
+ SetAssetStateStoreByUri,
ToSupervisor,
ToTask,
_RequestFrame,
@@ -834,6 +845,99 @@ async def
test_create_triggers_asset_state_store_accessor_reads_and_writes(
await runner.cleanup_finished_triggers()
+class TestTriggerSupervisorAssetStateStore:
+ """Supervisor side of the asset-state-store round trip.
+
+ The accessor-only test above asserts the messages a watcher trigger emits,
but never
+ exercises ``TriggerRunnerSupervisor._handle_request``. These drive the
supervisor
+ directly so a broken (or missing) branch can't pass silently.
+ """
+
+ @pytest.fixture
+ def supervisor(self, jobless_supervisor, mocker):
+ jobless_supervisor.client = mocker.MagicMock(spec=Client)
+ mocker.patch.object(TriggerRunnerSupervisor, "send_msg")
+ return jobless_supervisor
+
+ @staticmethod
+ def _handle(supervisor, msg):
+ supervisor._handle_request(msg,
log=MagicMock(spec=FilteringBoundLogger), req_id=7)
+
+ def test_get_by_name_wraps_response_and_replies(self, supervisor):
+ supervisor.client.asset_state_store.get.return_value =
AssetStateStoreResponse(value="2026-01-01")
+
+ self._handle(supervisor, GetAssetStateStoreByName(name="asset_a",
key="watermark"))
+
+
supervisor.client.asset_state_store.get.assert_called_once_with(key="watermark",
name="asset_a")
+ supervisor.send_msg.assert_called_once_with(
+ AssetStateStoreResult(value="2026-01-01"), request_id=7, error=None
+ )
+
+ def test_get_by_uri_wraps_response_and_replies(self, supervisor):
+ supervisor.client.asset_state_store.get.return_value =
AssetStateStoreResponse(value="2026-01-01")
+
+ self._handle(supervisor, GetAssetStateStoreByUri(uri="s3://bucket/a",
key="watermark"))
+
+
supervisor.client.asset_state_store.get.assert_called_once_with(key="watermark",
uri="s3://bucket/a")
+ supervisor.send_msg.assert_called_once_with(
+ AssetStateStoreResult(value="2026-01-01"), request_id=7, error=None
+ )
+
+ def test_get_passes_through_not_found_error(self, supervisor):
+ err = ErrorResponse(error=ErrorType.ASSET_STORE_NOT_FOUND,
detail={"key": "watermark"})
+ supervisor.client.asset_state_store.get.return_value = err
+
+ self._handle(supervisor, GetAssetStateStoreByName(name="asset_a",
key="watermark"))
+
+ supervisor.send_msg.assert_called_once_with(err, request_id=7,
error=None)
+
+ def test_set_by_name(self, supervisor):
+ self._handle(
+ supervisor, SetAssetStateStoreByName(name="asset_a",
key="watermark", value="2026-01-01")
+ )
+
+ supervisor.client.asset_state_store.set.assert_called_once_with(
+ key="watermark", value="2026-01-01", name="asset_a"
+ )
+ supervisor.send_msg.assert_called_once_with(OKResponse(ok=True),
request_id=7, error=None)
+
+ def test_set_by_uri(self, supervisor):
+ self._handle(
+ supervisor, SetAssetStateStoreByUri(uri="s3://bucket/a",
key="watermark", value="2026-01-01")
+ )
+
+ supervisor.client.asset_state_store.set.assert_called_once_with(
+ key="watermark", value="2026-01-01", uri="s3://bucket/a"
+ )
+ supervisor.send_msg.assert_called_once_with(OKResponse(ok=True),
request_id=7, error=None)
+
+ def test_delete_by_name(self, supervisor):
+ self._handle(supervisor, DeleteAssetStateStoreByName(name="asset_a",
key="watermark"))
+
+
supervisor.client.asset_state_store.delete.assert_called_once_with(key="watermark",
name="asset_a")
+ supervisor.send_msg.assert_called_once_with(OKResponse(ok=True),
request_id=7, error=None)
+
+ def test_delete_by_uri(self, supervisor):
+ self._handle(supervisor,
DeleteAssetStateStoreByUri(uri="s3://bucket/a", key="watermark"))
+
+ supervisor.client.asset_state_store.delete.assert_called_once_with(
+ key="watermark", uri="s3://bucket/a"
+ )
+ supervisor.send_msg.assert_called_once_with(OKResponse(ok=True),
request_id=7, error=None)
+
+ def test_clear_by_name(self, supervisor):
+ self._handle(supervisor, ClearAssetStateStoreByName(name="asset_a"))
+
+
supervisor.client.asset_state_store.clear.assert_called_once_with(name="asset_a")
+ supervisor.send_msg.assert_called_once_with(OKResponse(ok=True),
request_id=7, error=None)
+
+ def test_clear_by_uri(self, supervisor):
+ self._handle(supervisor,
ClearAssetStateStoreByUri(uri="s3://bucket/a"))
+
+
supervisor.client.asset_state_store.clear.assert_called_once_with(uri="s3://bucket/a")
+ supervisor.send_msg.assert_called_once_with(OKResponse(ok=True),
request_id=7, error=None)
+
+
def test_trigger_lifecycle(spy_agency: SpyAgency, session, testing_dag_bundle):
"""
Checks that the triggerer will correctly see a new Trigger in the database
@@ -2368,19 +2472,11 @@ class TestTriggererMessageTypes:
"CreateHITLDetailPayload",
"SetRenderedMapIndex",
"GetDag",
- # AIP-103 task/asset store — triggerer has no task execution
context.
+ # AIP-103 task store — triggerer has no task execution context.
"GetTaskStateStore",
"SetTaskStateStore",
"DeleteTaskStateStore",
"ClearTaskStateStore",
- "GetAssetStateStoreByName",
- "GetAssetStateStoreByUri",
- "SetAssetStateStoreByName",
- "SetAssetStateStoreByUri",
- "DeleteAssetStateStoreByName",
- "DeleteAssetStateStoreByUri",
- "ClearAssetStateStoreByName",
- "ClearAssetStateStoreByUri",
}
in_task_but_not_in_trigger_runner = {
@@ -2402,9 +2498,8 @@ class TestTriggererMessageTypes:
"PreviousTIResult",
"HITLDetailRequestResult",
"DagResult",
- # AIP-103 task/asset store results — worker-only responses to the
above messages.
+ # AIP-103 task store result — worker-only response to the above
messages.
"TaskStateStoreResult",
- "AssetStateStoreResult",
}
supervisor_diff = (
diff --git a/task-sdk/src/airflow/sdk/execution_time/request_handlers.py
b/task-sdk/src/airflow/sdk/execution_time/request_handlers.py
index 959be43fe93..a31596b3d2b 100644
--- a/task-sdk/src/airflow/sdk/execution_time/request_handlers.py
+++ b/task-sdk/src/airflow/sdk/execution_time/request_handlers.py
@@ -31,6 +31,7 @@ from typing import TYPE_CHECKING
from uuid import UUID
from airflow.sdk.api.datamodels._generated import (
+ AssetStateStoreResponse,
ConnectionResponse,
DagRunStateResponse,
TaskStatesResponse,
@@ -40,10 +41,17 @@ from airflow.sdk.api.datamodels._generated import (
XComSequenceSliceResponse,
)
from airflow.sdk.execution_time.comms import (
+ AssetStateStoreResult,
+ ClearAssetStateStoreByName,
+ ClearAssetStateStoreByUri,
ConnectionResult,
DagRunStateResult,
+ DeleteAssetStateStoreByName,
+ DeleteAssetStateStoreByUri,
DeleteVariable,
DeleteXCom,
+ GetAssetStateStoreByName,
+ GetAssetStateStoreByUri,
GetConnection,
GetDagRunState,
GetDRCount,
@@ -60,6 +68,8 @@ from airflow.sdk.execution_time.comms import (
MaskSecret,
PrevSuccessfulDagRunResult,
PutVariable,
+ SetAssetStateStoreByName,
+ SetAssetStateStoreByUri,
SetXCom,
TaskStatesResult,
VariableKeysResult,
@@ -272,3 +282,87 @@ def handle_get_xcom(client: Client, msg: GetXCom) ->
tuple[BaseModel | None, dic
xcom_result = XComResult.from_xcom_response(xcom)
return xcom_result, {"exclude_unset": True}
return xcom, {}
+
+
+def handle_get_asset_state_store_by_name(
+ client: Client, msg: GetAssetStateStoreByName
+) -> tuple[BaseModel | None, dict[str, bool]]:
+ asset_state = client.asset_state_store.get(
+ key=msg.key,
+ name=msg.name,
+ )
+
+ if isinstance(asset_state, AssetStateStoreResponse):
+ return
AssetStateStoreResult.from_asset_state_store_response(asset_state), {}
+
+ return asset_state, {}
+
+
+def handle_get_asset_state_store_by_uri(
+ client: Client, msg: GetAssetStateStoreByUri
+) -> tuple[BaseModel | None, dict[str, bool]]:
+ asset_state = client.asset_state_store.get(
+ key=msg.key,
+ uri=msg.uri,
+ )
+
+ if isinstance(asset_state, AssetStateStoreResponse):
+ return
AssetStateStoreResult.from_asset_state_store_response(asset_state), {}
+
+ return asset_state, {}
+
+
+def handle_set_asset_state_store_by_name(
+ client: Client, msg: SetAssetStateStoreByName
+) -> tuple[BaseModel | None, dict[str, bool]]:
+ client.asset_state_store.set(
+ key=msg.key,
+ value=msg.value,
+ name=msg.name,
+ )
+ return None, {}
+
+
+def handle_set_asset_state_store_by_uri(
+ client: Client, msg: SetAssetStateStoreByUri
+) -> tuple[BaseModel | None, dict[str, bool]]:
+ client.asset_state_store.set(
+ key=msg.key,
+ value=msg.value,
+ uri=msg.uri,
+ )
+ return None, {}
+
+
+def handle_delete_asset_state_store_by_name(
+ client: Client, msg: DeleteAssetStateStoreByName
+) -> tuple[BaseModel | None, dict[str, bool]]:
+ client.asset_state_store.delete(
+ key=msg.key,
+ name=msg.name,
+ )
+ return None, {}
+
+
+def handle_delete_asset_state_store_by_uri(
+ client: Client, msg: DeleteAssetStateStoreByUri
+) -> tuple[BaseModel | None, dict[str, bool]]:
+ client.asset_state_store.delete(
+ key=msg.key,
+ uri=msg.uri,
+ )
+ return None, {}
+
+
+def handle_clear_asset_state_store_by_name(
+ client: Client, msg: ClearAssetStateStoreByName
+) -> tuple[BaseModel | None, dict[str, bool]]:
+ client.asset_state_store.clear(name=msg.name)
+ return None, {}
+
+
+def handle_clear_asset_state_store_by_uri(
+ client: Client, msg: ClearAssetStateStoreByUri
+) -> tuple[BaseModel | None, dict[str, bool]]:
+ client.asset_state_store.clear(uri=msg.uri)
+ return None, {}
diff --git a/task-sdk/tests/task_sdk/execution_time/test_request_handlers.py
b/task-sdk/tests/task_sdk/execution_time/test_request_handlers.py
new file mode 100644
index 00000000000..c3999f9baca
--- /dev/null
+++ b/task-sdk/tests/task_sdk/execution_time/test_request_handlers.py
@@ -0,0 +1,167 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import pytest
+
+from airflow.sdk.api import client as sdk_client
+from airflow.sdk.api.datamodels._generated import AssetStateStoreResponse
+from airflow.sdk.exceptions import ErrorType
+from airflow.sdk.execution_time.comms import (
+ AssetStateStoreResult,
+ ClearAssetStateStoreByName,
+ ClearAssetStateStoreByUri,
+ DeleteAssetStateStoreByName,
+ DeleteAssetStateStoreByUri,
+ ErrorResponse,
+ GetAssetStateStoreByName,
+ GetAssetStateStoreByUri,
+ SetAssetStateStoreByName,
+ SetAssetStateStoreByUri,
+)
+from airflow.sdk.execution_time.request_handlers import (
+ handle_clear_asset_state_store_by_name,
+ handle_clear_asset_state_store_by_uri,
+ handle_delete_asset_state_store_by_name,
+ handle_delete_asset_state_store_by_uri,
+ handle_get_asset_state_store_by_name,
+ handle_get_asset_state_store_by_uri,
+ handle_set_asset_state_store_by_name,
+ handle_set_asset_state_store_by_uri,
+)
+
+
[email protected]
+def client():
+ return MagicMock(spec=sdk_client.Client)
+
+
+def test_get_asset_state_store_by_name_wraps_response_as_result(client):
+ client.asset_state_store.get.return_value =
AssetStateStoreResponse(value="2026-01-01")
+
+ result, dump_opts = handle_get_asset_state_store_by_name(
+ client, GetAssetStateStoreByName(name="asset_a", key="watermark")
+ )
+
+ client.asset_state_store.get.assert_called_once_with(key="watermark",
name="asset_a")
+ assert result == AssetStateStoreResult(value="2026-01-01")
+ assert dump_opts == {}
+
+
+def test_get_asset_state_store_by_name_passes_through_error_response(client):
+ err = ErrorResponse(error=ErrorType.ASSET_STORE_NOT_FOUND, detail={"key":
"watermark"})
+ client.asset_state_store.get.return_value = err
+
+ result, dump_opts = handle_get_asset_state_store_by_name(
+ client, GetAssetStateStoreByName(name="asset_a", key="watermark")
+ )
+
+ assert result is err
+ assert dump_opts == {}
+
+
+def test_get_asset_state_store_by_uri_wraps_response_as_result(client):
+ client.asset_state_store.get.return_value =
AssetStateStoreResponse(value="2026-01-01")
+
+ result, dump_opts = handle_get_asset_state_store_by_uri(
+ client, GetAssetStateStoreByUri(uri="s3://bucket/a", key="watermark")
+ )
+
+ client.asset_state_store.get.assert_called_once_with(key="watermark",
uri="s3://bucket/a")
+ assert result == AssetStateStoreResult(value="2026-01-01")
+ assert dump_opts == {}
+
+
+def test_get_asset_state_store_by_uri_passes_through_error_response(client):
+ err = ErrorResponse(error=ErrorType.ASSET_STORE_NOT_FOUND, detail={"key":
"watermark"})
+ client.asset_state_store.get.return_value = err
+
+ result, dump_opts = handle_get_asset_state_store_by_uri(
+ client, GetAssetStateStoreByUri(uri="s3://bucket/a", key="watermark")
+ )
+
+ assert result is err
+ assert dump_opts == {}
+
+
[email protected](
+ ("handler", "msg", "call_kwargs", "method"),
+ [
+ (
+ handle_set_asset_state_store_by_name,
+ SetAssetStateStoreByName,
+ {
+ "name": "asset_a",
+ "key": "watermark",
+ "value": "2026-01-01",
+ },
+ "set",
+ ),
+ (
+ handle_set_asset_state_store_by_uri,
+ SetAssetStateStoreByUri,
+ {
+ "uri": "s3://bucket/a",
+ "key": "watermark",
+ "value": "2026-01-01",
+ },
+ "set",
+ ),
+ (
+ handle_delete_asset_state_store_by_name,
+ DeleteAssetStateStoreByName,
+ {
+ "name": "asset_a",
+ "key": "watermark",
+ },
+ "delete",
+ ),
+ (
+ handle_delete_asset_state_store_by_uri,
+ DeleteAssetStateStoreByUri,
+ {
+ "uri": "s3://bucket/a",
+ "key": "watermark",
+ },
+ "delete",
+ ),
+ (
+ handle_clear_asset_state_store_by_name,
+ ClearAssetStateStoreByName,
+ {
+ "name": "asset_a",
+ },
+ "clear",
+ ),
+ (
+ handle_clear_asset_state_store_by_uri,
+ ClearAssetStateStoreByUri,
+ {
+ "uri": "s3://bucket/a",
+ },
+ "clear",
+ ),
+ ],
+)
+def test_asset_store_delegates_to_client(client, handler, msg, call_kwargs,
method):
+ result, dump_opts = handler(client, msg(**call_kwargs))
+
+ getattr(client.asset_state_store,
method).assert_called_once_with(**call_kwargs)
+ assert result is None
+ assert dump_opts == {}
