amoghrajesh commented on code in PR #75521:
URL: https://github.com/apache/airflow/pull/75521#discussion_r4376967979
##########
task-sdk/src/airflow/sdk/execution_time/context.py:
##########
@@ -1147,6 +1147,67 @@ def __repr__(self) -> str:
return f"<AssetStateStoreAccessors [{', '.join(parts)}]>"
+class TriggerAssetStateStoreAccessors(AssetStateStoreAccessors):
+ """
+ Asset state store accessors for a watcher trigger.
+
+ Watchers with the same trigger classpath and kwargs share one trigger, and
every event it
+ yields goes to all of its assets, so the trigger has one position to keep
no matter how many
+ assets it serves. Writes go to every asset; reads return the first asset
holding the key.
+ """
+
+ def _accessors(self) -> list[AssetStateStoreAccessor]:
+ # Sorted so reads pick the same asset on every triggerer, whatever
order the assets were loaded in.
+ return [self._by_name[n] for n in sorted(self._by_name)] + [
+ self._by_uri[u] for u in sorted(self._by_uri)
+ ]
+
+ def get(self, key: str, default: JsonValue = None) -> JsonValue:
+ """Return the value from the first asset holding ``key``, or
``default`` if none does."""
+ for accessor in self._accessors():
+ # set() rejects None, so None here means the asset does not hold
the key.
+ if (value := accessor.get(key)) is not None:
+ return value
+ return default
+
+ async def aget(self, key: str, default: JsonValue = None) -> JsonValue:
+ """Async version of `get` that awaits instead of blocking the event
loop."""
+ for accessor in self._accessors():
+ if (value := await accessor.aget(key)) is not None:
+ return value
+ return default
+
+ def set(self, key: str, value: JsonValue) -> None:
+ """Write the value for ``key`` to every asset the trigger watches."""
+ for accessor in self._accessors():
+ accessor.set(key, value)
+
+ async def aset(self, key: str, value: JsonValue) -> None:
+ """Async version of `set` that awaits instead of blocking the event
loop."""
+ for accessor in self._accessors():
+ await accessor.aset(key, value)
+
+ def delete(self, key: str) -> None:
+ """Delete ``key`` from every asset the trigger watches."""
+ for accessor in self._accessors():
+ accessor.delete(key)
+
+ async def adelete(self, key: str) -> None:
+ """Async version of `delete` that awaits instead of blocking the event
loop."""
+ for accessor in self._accessors():
+ await accessor.adelete(key)
+
+ def clear(self) -> None:
+ """Delete all state keys from every asset the trigger watches."""
+ for accessor in self._accessors():
+ accessor.clear()
+
+ async def aclear(self) -> None:
+ """Async version of `clear` that awaits instead of blocking the event
loop."""
+ for accessor in self._accessors():
+ await accessor.aclear()
Review Comment:
clear now wipes every key on every linked asset, including keys that tasks
wrote. Thats a big blast radius for one trigger call.
Suggest not overriding `clear`/`aclear` at all. Then the parent's
single-asset check still applies, and shared triggers use `delete(key)`
##########
airflow-core/tests/unit/jobs/test_triggerer_job.py:
##########
@@ -1237,6 +1237,43 @@ async def
test_create_triggers_asset_state_store_accessor_reads_and_writes(
await runner.cleanup_finished_triggers()
[email protected]
+@patch("airflow.jobs.triggerer_job_runner.TriggerRunner.get_trigger_by_classpath")
+async def test_create_triggers_asset_state_store_shared_by_several_assets(
+ mock_get_classpath, session, mock_supervisor_comms, make_watcher_trigger
+):
+ """A trigger shared by several assets can get and set state without
picking an asset."""
+ injected_instances = []
+ mock_get_classpath.return_value = make_watcher_trigger(injected_instances)
+
+ runner = TriggerRunner()
+ runner.to_create.append(
+ workloads.RunTrigger.model_construct(
+ id=15,
+ ti=None,
+ classpath="fake.WatcherTrigger",
+ encrypted_kwargs="{}",
+ watched_assets={"asset_a": "s3://bucket/a", "asset_b":
"s3://bucket/a"},
+ )
Review Comment:
Both assets have the same URI. Can we give them different ones so it looks
like a real setup?
##########
airflow-core/tests/unit/jobs/test_triggerer_job.py:
##########
@@ -1237,6 +1237,43 @@ async def
test_create_triggers_asset_state_store_accessor_reads_and_writes(
await runner.cleanup_finished_triggers()
[email protected]
+@patch("airflow.jobs.triggerer_job_runner.TriggerRunner.get_trigger_by_classpath")
+async def test_create_triggers_asset_state_store_shared_by_several_assets(
+ mock_get_classpath, session, mock_supervisor_comms, make_watcher_trigger
+):
+ """A trigger shared by several assets can get and set state without
picking an asset."""
+ injected_instances = []
+ mock_get_classpath.return_value = make_watcher_trigger(injected_instances)
+
+ runner = TriggerRunner()
+ runner.to_create.append(
+ workloads.RunTrigger.model_construct(
+ id=15,
+ ti=None,
+ classpath="fake.WatcherTrigger",
+ encrypted_kwargs="{}",
+ watched_assets={"asset_a": "s3://bucket/a", "asset_b":
"s3://bucket/a"},
+ )
+ )
+
+ await runner.create_triggers()
+ state_store = injected_instances[0].asset_state_store
+
+ mock_supervisor_comms.send.return_value =
AssetStateStoreResult(value="etag-1")
Review Comment:
The mock returns `etag-1` for every asset, so this passes whichever asset is
read. Fine as a sanity check, but can be better tested.
##########
airflow-core/docs/core-concepts/asset-state-store.rst:
##########
@@ -172,6 +172,8 @@ Using ``asset_state_store`` inside a Watcher Trigger
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.
+Watchers on several assets that use the same trigger class and arguments share
one trigger. The accessor then covers all of those assets: ``set``, ``delete``
and ``clear`` apply to every one of them, and ``get`` returns the value from
the first asset (by name) that holds the key.
Review Comment:
The intro above says "scoped to the asset" (one asset). Can we:
- change it to "the asset or assets this trigger watches"
- use f"last_seen_id:{self.source}" as the key in the example, so two
triggers on one asset don't share a key
- end this paragraph with a suggestion like "include what the trigger
watches in the key name"
And also optional: a two-asset snippet showing one shared trigger.
##########
task-sdk/src/airflow/sdk/execution_time/context.py:
##########
@@ -1147,6 +1147,67 @@ def __repr__(self) -> str:
return f"<AssetStateStoreAccessors [{', '.join(parts)}]>"
+class TriggerAssetStateStoreAccessors(AssetStateStoreAccessors):
+ """
+ Asset state store accessors for a watcher trigger.
+
+ Watchers with the same trigger classpath and kwargs share one trigger, and
every event it
+ yields goes to all of its assets, so the trigger has one position to keep
no matter how many
+ assets it serves. Writes go to every asset; reads return the first asset
holding the key.
+ """
+
+ def _accessors(self) -> list[AssetStateStoreAccessor]:
+ # Sorted so reads pick the same asset on every triggerer, whatever
order the assets were loaded in.
+ return [self._by_name[n] for n in sorted(self._by_name)] + [
+ self._by_uri[u] for u in sorted(self._by_uri)
+ ]
+
+ def get(self, key: str, default: JsonValue = None) -> JsonValue:
+ """Return the value from the first asset holding ``key``, or
``default`` if none does."""
+ for accessor in self._accessors():
Review Comment:
Umm, trying to settle for the first asset can return an old value. Example:
if assets a and b share a Kinesis trigger. `a` is removed, then added back. `a`
still has its old cursor and sorts first, so the trigger rewinds and replays
records.
I dont think its possible to tell that today, in which case making a note is
good enough in docs?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
