You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "jroachgolf84 (via GitHub)" <gi...@apache.org> on 2026/06/01 11:59:49 UTC
[PR] AIP-103: Passing `AssetStateAccessors` through to `BaseEventTrigger` [airflow]
jroachgolf84 opened a new pull request, #67839:
URL: https://github.com/apache/airflow/pull/67839
## Description
To ensure that classes inheriting from the `BaseEventTrigger` (typically used for Asset watching) are "Asset-aware", the work done in AIP-103 has been wired into the `BaseEventTrigger` class.
#65103 drafted an approach that "flipped" the definition of the `Asset` and `AssetWatcher`. However, for good reason, the approach was challenged. After conversation, the goal became to "pass" the `AssetStateAccessors` through to the `BaseEventTrigger` with some runtime magic.
With this model, defining an `Asset` and `AssetWatcher` remains the same as before.
```python
from airflow.sdk import Asset, AssetWatcher, DAG, task
from datetime import datetime
from triggers.event_triggers import GenericEventTrigger
generic_asset_watcher = AssetWatcher(
name="generic_asset_watcher",
trigger=GenericEventTrigger(
random_number=1,
waiter_delay=15
)
)
generic_asset = Asset(
name="generic_asset",
watchers=[generic_asset_watcher],
)
with DAG(
dag_id="aip_93_scoping",
start_date=datetime(2026, 1, 1),
schedule=[generic_asset]
) as dag:
@task
def downstream_task():
pass
downstream_task()
```
The two PR's below were closed in favor of this PR:
- #66595
## Testing
No unit tests have been written for this logic yet. However, testing has been performed E2E locally with `breeze`. The trigger authored below is what was used for testing. When this ran, the Asset `name` and `uri` were output in the Triggerer logs.
```python
class GenericEventTrigger(BaseEventTrigger):
def __init__(
self,
random_number,
waiter_delay,
asset_name,
**kwargs
):
super().__init__(**kwargs)
self.random_number = random_number
self.waiter_delay = waiter_delay
self.asset_name = asset_name
def serialize(self) -> tuple[str, dict[str, Any]]:
"""Serialize the Trigger, including the func, params, and waiter_delay."""
return (
self.__class__.__module__ + "." + self.__class__.__qualname__,
{
"random_number": self.random_number,
"waiter_delay": self.waiter_delay,
"asset_name": self.asset_name,
},
)
async def run(self) -> AsyncIterator[TriggerEvent]:
"""Logic that fires a TriggerEvent."""
logging.info(f"***** asset_states: {self.asset_states}")
while True:
result = random.randint(0, 5)
if result == self.random_number:
logging.info("yield'ing TriggerEvent")
yield TriggerEvent({"status": "success", "result": result})
break
logging.info(f"Sleeping for {self.waiter_delay} seconds")
await asyncio.sleep(self.waiter_delay)
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Passing `AssetStateAccessors` through to `BaseEventTrigger` [airflow]
Posted by "jroachgolf84 (via GitHub)" <gi...@apache.org>.
jroachgolf84 commented on PR #67839:
URL: https://github.com/apache/airflow/pull/67839#issuecomment-4592381809
cc: @cmarteepants
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Passing `AssetStateStoreAccessors` through to `BaseEventTrigger` [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh merged PR #67839:
URL: https://github.com/apache/airflow/pull/67839
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Passing `AssetStateAccessors` through to `BaseEventTrigger` [airflow]
Posted by "vincbeck (via GitHub)" <gi...@apache.org>.
vincbeck commented on code in PR #67839:
URL: https://github.com/apache/airflow/pull/67839#discussion_r3335356259
##########
airflow-core/src/airflow/jobs/triggerer_job_runner.py:
##########
@@ -1266,6 +1269,16 @@ async def create_triggers(self):
trigger_instance.triggerer_job_id = self.job_id
trigger_instance.timeout_after = workload.timeout_after
+ if isinstance(trigger_instance, BaseEventTrigger) and workload.watched_assets:
+ # Reconstruct AssetStateAccessors from watched_assets
+ from airflow.sdk.definitions.asset import Asset
+ from airflow.sdk.execution_time.context import AssetStateAccessors
+
+ # Potentially address Asset vs. AssetRef, AssetUriRef, etc.
+ trigger_instance.asset_state = AssetStateAccessors(
Review Comment:
nit: given https://github.com/apache/airflow/pull/67833 we might want to rename this variable
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Passing `AssetStateAccessors` through to `BaseEventTrigger` [airflow]
Posted by "jroachgolf84 (via GitHub)" <gi...@apache.org>.
jroachgolf84 commented on PR #67839:
URL: https://github.com/apache/airflow/pull/67839#issuecomment-4592705368
cc: @vincbeck and @amoghrajesh, I closed the previous PR (#66595) after renaming the branch. We'll use this one instead.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Passing `AssetStateAccessors` through to `BaseEventTrigger` [airflow]
Posted by "jroachgolf84 (via GitHub)" <gi...@apache.org>.
jroachgolf84 commented on PR #67839:
URL: https://github.com/apache/airflow/pull/67839#issuecomment-4622164180
Unless I get anymore feedback, I'm going to start to write tests for this.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Passing `AssetStateStoreAccessors` through to `BaseEventTrigger` [airflow]
Posted by "jroachgolf84 (via GitHub)" <gi...@apache.org>.
jroachgolf84 commented on PR #67839:
URL: https://github.com/apache/airflow/pull/67839#issuecomment-4752834440
@amoghrajesh - good to merge?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Passing `AssetStateStoreAccessors` through to `BaseEventTrigger` [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #67839:
URL: https://github.com/apache/airflow/pull/67839#discussion_r3439491208
##########
airflow-core/tests/unit/jobs/test_triggerer_job.py:
##########
@@ -561,6 +571,255 @@ def test_create_workload_uses_supervisor_id_without_job(jobless_supervisor, mock
assert factory.log_path == f"/logs/ti.trigger.{jobless_supervisor.id}.log"
+def test_create_workload_sets_watched_assets_for_asset_only_trigger(jobless_supervisor, mocker):
+ """_create_workload() should populate watched_assets when trigger.task_instance is None and assets exist."""
+ asset1 = mocker.Mock(spec=Asset)
+ asset1.name = "my_asset"
+ asset1.uri = "s3://bucket/key"
+
+ asset2 = mocker.Mock(spec=Asset)
+ asset2.name = "other_asset"
+ asset2.uri = "gs://bucket/path"
+
+ trigger = mocker.Mock(spec=BaseEventTrigger)
+ trigger.id = 42
+ trigger.classpath = "some.path.Trigger"
+ trigger.encrypted_kwargs = "encrypted"
+ trigger.task_instance = None # Not tied to a Task (similar to a BaseEventTrigger)
+ trigger.assets = [asset1, asset2]
+
+ workload = jobless_supervisor._create_workload(
+ trigger=trigger,
+ dag_bag=mocker.Mock(),
+ render_log_fname=mocker.Mock(),
+ session=mocker.Mock(),
+ )
+
+ assert workload is not None
+ assert workload.watched_assets == {"my_asset": "s3://bucket/key", "other_asset": "gs://bucket/path"}
+
+
+def test_create_workload_watched_assets_none_when_no_assets(jobless_supervisor, mocker):
+ """_create_workload() should set watched_assets=None when trigger.task_instance is None and assets is empty."""
+ trigger = mocker.Mock(spec=BaseEventTrigger)
+ trigger.id = 43
+ trigger.classpath = "some.path.Trigger"
+ trigger.encrypted_kwargs = "encrypted"
+ trigger.task_instance = None
+ trigger.assets = [] # No Assets are attached to the trigger
+
+ workload = jobless_supervisor._create_workload(
+ trigger=trigger,
+ dag_bag=mocker.Mock(),
+ render_log_fname=mocker.Mock(),
+ session=mocker.Mock(),
+ )
+
+ assert workload is not None
+ assert workload.watched_assets is None
+
+
+def test_run_trigger_workload_includes_watched_assets_field():
+ """RunTrigger workload should accept and store watched_assets."""
+ workload = RunTrigger(
+ id=1,
+ classpath="airflow.triggers.testing.SuccessTrigger",
+ encrypted_kwargs="fake",
+ watched_assets={"asset_a": "s3://a", "asset_b": "gs://b"},
+ )
+ assert workload.watched_assets == {"asset_a": "s3://a", "asset_b": "gs://b"}
+
+
+def test_run_trigger_workload_watched_assets_defaults_to_none():
+ """RunTrigger workload watched_assets should default to None."""
+ workload = RunTrigger(
+ id=1,
+ classpath="airflow.triggers.testing.SuccessTrigger",
+ encrypted_kwargs="fake",
+ )
+ assert workload.watched_assets is None
+
+
+@pytest.fixture
+def make_watcher_trigger():
+ """Factory fixture: call with a list to get a BaseEventTrigger subclass that appends each new instance."""
+
+ def factory(injected_instances):
+ class WatcherTrigger(BaseEventTrigger):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+ injected_instances.append(self)
+
+ def serialize(self):
+ return (f"{type(self).__module__}.{type(self).__qualname__}", {})
+
+ async def run(self):
+ yield TriggerEvent("done")
+
+ return WatcherTrigger
+
+ return factory
+
+
+@pytest.mark.asyncio
+@patch("airflow.jobs.triggerer_job_runner.TriggerRunner.get_trigger_by_classpath")
+async def test_create_triggers_injects_asset_state_store_for_base_event_trigger(
+ mock_get_classpath, session, make_watcher_trigger
+):
+ """asset_state_store is populated on BaseEventTrigger instances when watched_assets is set."""
+ injected_instances = []
+ mock_get_classpath.return_value = make_watcher_trigger(injected_instances)
+
+ runner = TriggerRunner()
+ runner.to_create.append(
+ workloads.RunTrigger.model_construct(
+ id=10,
+ ti=None,
+ classpath="fake.WatcherTrigger",
+ encrypted_kwargs="{}",
+ watched_assets={"my_asset": "s3://bucket/key"},
+ )
+ )
+
+ await runner.create_triggers()
+
+ # This is only testing that an exception was NOT thrown when creating the Trigger
+ assert 10 in runner.triggers
+
+ assert len(injected_instances) == 1
+ assert injected_instances[0].asset_state_store is not None
+ assert isinstance(injected_instances[0].asset_state_store, AssetStateStoreAccessors)
+
+ runner.triggers[10]["task"].cancel()
+ await runner.cleanup_finished_triggers()
+
+
+@pytest.mark.asyncio
+@patch("airflow.jobs.triggerer_job_runner.TriggerRunner.get_trigger_by_classpath")
+async def test_create_triggers_asset_state_store_none_when_no_watched_assets(
+ mock_get_classpath, session, make_watcher_trigger
+):
+ """asset_state_store stays None when watched_assets is not set on the workload."""
+ injected_instances = []
+ mock_get_classpath.return_value = make_watcher_trigger(injected_instances)
+
+ runner = TriggerRunner()
+ runner.to_create.append(
+ workloads.RunTrigger.model_construct(
+ id=11,
+ ti=None,
+ classpath="fake.WatcherTrigger",
+ encrypted_kwargs="{}",
+ watched_assets=None,
+ )
+ )
+
+ await runner.create_triggers()
+
+ assert len(injected_instances) == 1
+ assert injected_instances[0].asset_state_store is None
+
+ runner.triggers[11]["task"].cancel()
+ await runner.cleanup_finished_triggers()
+
+
+@pytest.mark.asyncio
+@patch("airflow.jobs.triggerer_job_runner.TriggerRunner.get_trigger_by_classpath")
+async def test_create_triggers_skips_asset_state_store_for_non_event_trigger(mock_get_classpath, session):
+ """asset_state_store injection is skipped for plain BaseTrigger (non-BaseEventTrigger) instances."""
+ mock_get_classpath.return_value = SuccessTrigger
+
+ runner = TriggerRunner()
+ runner.to_create.append(
+ workloads.RunTrigger.model_construct(
+ id=12, ti=None, classpath="airflow.triggers.testing.SuccessTrigger", encrypted_kwargs="{}"
+ )
+ )
+
+ await runner.create_triggers()
+
+ assert 12 in runner.triggers
+ assert not hasattr(runner.triggers[12]["task"], "asset_state_store")
Review Comment:
`runner.triggers[12]["task"]` is the asyncio `Task`, not the trigger instance, so this assertion is vacuously true: a `Task` never has an `asset_state_store` attribute regardless of the skip logic. To actually verify the non-event-trigger path is skipped, capture the trigger instance (like the `make_watcher_trigger` fixture does for the other tests) and assert `not hasattr(instance, "asset_state_store")` on it.
##########
airflow-core/src/airflow/jobs/triggerer_job_runner.py:
##########
@@ -1313,6 +1321,11 @@ async def create_triggers(self):
trigger_instance.triggerer_job_id = self.job_id
trigger_instance.timeout_after = workload.timeout_after
+ if isinstance(trigger_instance, BaseEventTrigger) and workload.watched_assets:
Review Comment:
Following up on the `_handle_request` thread above (marked resolved, but I don't think the runtime gap was actually closed out): the accessor injected here can't read or write at runtime. `AssetStateStoreAccessor.get()`/`.set()` send `GetAssetStateStoreByName`/`SetAssetStateStoreByName` over `SUPERVISOR_COMMS`, but the triggerer supervisor's `_handle_request` (line 543) has no branch for those message types and `ToTriggerSupervisor` (line 367) doesn't include them, so a real `self.asset_state_store[asset].get("key")` call lands on the `else: raise ValueError(f"Unknown message type ...")` fallthrough (line 626).
`test_create_triggers_asset_state_store_accessor_reads_and_writes` only mocks `SUPERVISOR_COMMS.send` and asserts the outgoing message, so it never exercises the supervisor side and stays green while the round trip is broken.
Was the get/set wiring meant to land in a follow-up? If so, worth calling out in the description that the accessor is injected but not yet functional. (Separately, these `send()` calls are synchronous inside `async def run()`, which the `init_comms` docstring at line 1229 flags as needing a `sync_to_async()` wrapper.)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Passing `AssetStateAccessors` through to `BaseEventTrigger` [airflow]
Posted by "jroachgolf84 (via GitHub)" <gi...@apache.org>.
jroachgolf84 commented on code in PR #67839:
URL: https://github.com/apache/airflow/pull/67839#discussion_r3341072074
##########
airflow-core/src/airflow/jobs/triggerer_job_runner.py:
##########
@@ -1266,6 +1269,16 @@ async def create_triggers(self):
trigger_instance.triggerer_job_id = self.job_id
trigger_instance.timeout_after = workload.timeout_after
+ if isinstance(trigger_instance, BaseEventTrigger) and workload.watched_assets:
+ # Reconstruct AssetStateAccessors from watched_assets
+ from airflow.sdk.definitions.asset import Asset
+ from airflow.sdk.execution_time.context import AssetStateAccessors
+
+ # Potentially address Asset vs. AssetRef, AssetUriRef, etc.
+ trigger_instance.asset_state = AssetStateAccessors(
Review Comment:
You bet, will do!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Passing `AssetStateStoreAccessors` through to `BaseEventTrigger` [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on PR #67839:
URL: https://github.com/apache/airflow/pull/67839#issuecomment-4764760128
Rebasing this one, to ensure CI stays green.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Passing `AssetStateStoreAccessors` through to `BaseEventTrigger` [airflow]
Posted by "jroachgolf84 (via GitHub)" <gi...@apache.org>.
jroachgolf84 commented on code in PR #67839:
URL: https://github.com/apache/airflow/pull/67839#discussion_r3343916715
##########
airflow-core/src/airflow/jobs/triggerer_job_runner.py:
##########
Review Comment:
I don't think so? What would that implementation look like?
##########
airflow-core/src/airflow/jobs/triggerer_job_runner.py:
##########
@@ -1266,6 +1269,16 @@ async def create_triggers(self):
trigger_instance.triggerer_job_id = self.job_id
trigger_instance.timeout_after = workload.timeout_after
+ if isinstance(trigger_instance, BaseEventTrigger) and workload.watched_assets:
+ # Reconstruct AssetStateAccessors from watched_assets
+ from airflow.sdk.definitions.asset import Asset
+ from airflow.sdk.execution_time.context import AssetStateAccessors
Review Comment:
Hmmm, not seeing that.
##########
airflow-core/src/airflow/jobs/triggerer_job_runner.py:
##########
@@ -1286,6 +1293,16 @@ async def create_triggers(self):
trigger_instance.triggerer_job_id = self.job_id
trigger_instance.timeout_after = workload.timeout_after
+ if isinstance(trigger_instance, BaseEventTrigger) and workload.watched_assets:
+ # Reconstruct AssetStateAccessors from watched_assets
+ from airflow.sdk.definitions.asset import Asset
+ from airflow.sdk.execution_time.context import AssetStoreAccessors
Review Comment:
Moved to top-level.
##########
airflow-core/src/airflow/jobs/triggerer_job_runner.py:
##########
@@ -741,11 +737,18 @@ def _create_workload(
render_log_fname: Callable[..., str],
session: Session,
) -> workloads.RunTrigger | None:
+ # Pass the "watched" Assets through for downstream use in BaseEventTrigger
if trigger.task_instance is None:
+ watched_assets: dict[str, str] | None = None
+
+ if trigger.asset_watchers:
+ watched_assets = {a.name: a.uri for a in trigger.assets}
Review Comment:
Resolved in next commit.
##########
airflow-core/tests/unit/jobs/test_triggerer_job.py:
##########
@@ -558,6 +558,248 @@ def test_create_workload_uses_supervisor_id_without_job(jobless_supervisor, mock
assert factory.log_path == f"/logs/ti.trigger.{jobless_supervisor.id}.log"
+def test_create_workload_sets_watched_assets_for_asset_only_trigger(jobless_supervisor, mocker):
Review Comment:
Implemented for tests in this PR.
##########
airflow-core/src/airflow/jobs/triggerer_job_runner.py:
##########
@@ -1286,6 +1293,16 @@ async def create_triggers(self):
trigger_instance.triggerer_job_id = self.job_id
trigger_instance.timeout_after = workload.timeout_after
+ if isinstance(trigger_instance, BaseEventTrigger) and workload.watched_assets:
+ # Reconstruct AssetStateAccessors from watched_assets
+ from airflow.sdk.definitions.asset import Asset
+ from airflow.sdk.execution_time.context import AssetStoreAccessors
+
+ # Potentially address Asset vs. AssetRef, AssetUriRef, etc.
Review Comment:
No, removing comment.
##########
airflow-core/tests/unit/triggers/test_base_trigger.py:
##########
@@ -232,3 +234,25 @@ async def stream():
payloads = [event.payload async for event in us.filter_shared_stream(stream())]
assert [p["region"] for p in payloads] == ["us", "us"]
+
+
+def test_base_event_trigger_asset_state_store_initialized_to_none():
+ """asset_state_store is None before it is set."""
+ trigger = _PlainEventTrigger()
+ assert trigger.asset_state_store is None
+
+
+def test_base_event_trigger_asset_state_store_can_be_set():
+ """asset_state_store can be set once the Trigger is initialized."""
+ trigger = _PlainEventTrigger()
+ mock_store = MagicMock()
+ trigger.asset_state_store = mock_store
+ assert trigger.asset_state_store is mock_store
+
+
+def test_base_event_trigger_asset_state_store_independent_across_instances():
+ """a.asset_state_store does not impact b.asset_state_store."""
+ a = _PlainEventTrigger(name="a")
+ b = _PlainEventTrigger(name="b")
+ a.asset_state_store = MagicMock()
+ assert b.asset_state_store is None
Review Comment:
Resolved in next commit.
##########
airflow-core/tests/unit/jobs/test_triggerer_job.py:
##########
@@ -558,6 +558,248 @@ def test_create_workload_uses_supervisor_id_without_job(jobless_supervisor, mock
assert factory.log_path == f"/logs/ti.trigger.{jobless_supervisor.id}.log"
+def test_create_workload_sets_watched_assets_for_asset_only_trigger(jobless_supervisor, mocker):
+ """_create_workload() should populate watched_assets when trigger.task_instance is None and assets exist."""
+ asset1 = mocker.Mock()
+ asset1.name = "my_asset"
+ asset1.uri = "s3://bucket/key"
+
+ asset2 = mocker.Mock()
+ asset2.name = "other_asset"
+ asset2.uri = "gs://bucket/path"
+
+ trigger = mocker.Mock()
+ trigger.id = 42
+ trigger.classpath = "some.path.Trigger"
+ trigger.encrypted_kwargs = "encrypted"
+ trigger.task_instance = None # Not tied to a Task (similar to a BaseEventTrigger)
+ trigger.assets = [asset1, asset2]
+
+ workload = jobless_supervisor._create_workload(
+ trigger=trigger,
+ dag_bag=mocker.Mock(),
+ render_log_fname=mocker.Mock(),
+ session=mocker.Mock(),
+ )
+
+ assert workload is not None
+ assert workload.watched_assets == {"my_asset": "s3://bucket/key", "other_asset": "gs://bucket/path"}
+
+
+def test_create_workload_watched_assets_none_when_no_assets(jobless_supervisor, mocker):
+ """_create_workload() should set watched_assets=None when trigger.task_instance is None and assets is empty."""
+ trigger = mocker.Mock()
+ trigger.id = 43
+ trigger.classpath = "some.path.Trigger"
+ trigger.encrypted_kwargs = "encrypted"
+ trigger.task_instance = None
+ trigger.assets = [] # No Assets are attached to the trigger
+
+ workload = jobless_supervisor._create_workload(
+ trigger=trigger,
+ dag_bag=mocker.Mock(),
+ render_log_fname=mocker.Mock(),
+ session=mocker.Mock(),
+ )
+
+ assert workload is not None
+ assert workload.watched_assets is None
+
+
+def test_run_trigger_workload_includes_watched_assets_field():
+ """RunTrigger workload should accept and store watched_assets."""
+ from airflow.executors.workloads.trigger import RunTrigger
+
+ workload = RunTrigger(
+ id=1,
+ classpath="airflow.triggers.testing.SuccessTrigger",
+ encrypted_kwargs="fake",
+ watched_assets={"asset_a": "s3://a", "asset_b": "gs://b"},
+ )
+ assert workload.watched_assets == {"asset_a": "s3://a", "asset_b": "gs://b"}
+
+
+def test_run_trigger_workload_watched_assets_defaults_to_none():
+ """RunTrigger workload watched_assets should default to None."""
+ from airflow.executors.workloads.trigger import RunTrigger
+
+ workload = RunTrigger(
+ id=1,
+ classpath="airflow.triggers.testing.SuccessTrigger",
+ encrypted_kwargs="fake",
+ )
+ assert workload.watched_assets is None
+
+
+@pytest.mark.asyncio
+@patch("airflow.jobs.triggerer_job_runner.Trigger._decrypt_kwargs", return_value={})
+@patch("airflow.jobs.triggerer_job_runner.TriggerRunner.get_trigger_by_classpath")
+async def test_create_triggers_injects_asset_store_for_base_event_trigger(
+ mock_get_classpath, mock_decrypt, session
+):
+ """asset_store is populated on BaseEventTrigger instances when watched_assets is set."""
+ from airflow.sdk.execution_time.context import AssetStoreAccessors
+ from airflow.triggers.base import BaseEventTrigger, TriggerEvent
+
+ injected_instances = []
+
+ class _WatcherTrigger(BaseEventTrigger):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+ injected_instances.append(self)
+
+ def serialize(self):
+ return (f"{type(self).__module__}.{type(self).__qualname__}", {})
+
+ async def run(self):
+ yield TriggerEvent("done")
+
+ mock_get_classpath.return_value = _WatcherTrigger
+
+ runner = TriggerRunner()
+ runner.to_create.append(
+ workloads.RunTrigger.model_construct(
+ id=10,
+ ti=None,
+ classpath="fake.WatcherTrigger",
+ encrypted_kwargs="fake",
+ watched_assets={"my_asset": "s3://bucket/key"},
+ )
+ )
+
+ await runner.create_triggers()
+
+ # This is only testing that an exception was NOT thrown when creating the Trigger
+ assert 10 in runner.triggers
+
+ assert len(injected_instances) == 1
+ assert injected_instances[0].asset_store is not None
+ assert isinstance(injected_instances[0].asset_store, AssetStoreAccessors)
+
+ runner.triggers[10]["task"].cancel()
+ await runner.cleanup_finished_triggers()
+
+
+@pytest.mark.asyncio
+@patch("airflow.jobs.triggerer_job_runner.Trigger._decrypt_kwargs", return_value={})
+@patch("airflow.jobs.triggerer_job_runner.TriggerRunner.get_trigger_by_classpath")
+async def test_create_triggers_asset_store_none_when_no_watched_assets(
+ mock_get_classpath, mock_decrypt, session
+):
+ """asset_store stays None when watched_assets is not set on the workload."""
+ from airflow.triggers.base import BaseEventTrigger, TriggerEvent
+
+ injected_instances = []
+
+ class _WatcherTrigger(BaseEventTrigger):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+ injected_instances.append(self)
+
+ def serialize(self):
+ return (f"{type(self).__module__}.{type(self).__qualname__}", {})
+
+ async def run(self):
+ yield TriggerEvent("done")
+
+ mock_get_classpath.return_value = _WatcherTrigger
+
+ runner = TriggerRunner()
+ runner.to_create.append(
+ workloads.RunTrigger.model_construct(
+ id=11,
+ ti=None,
+ classpath="fake.WatcherTrigger",
+ encrypted_kwargs="fake",
+ watched_assets=None,
+ )
+ )
+
+ await runner.create_triggers()
+
+ assert len(injected_instances) == 1
+ assert injected_instances[0].asset_store is None
+
+ runner.triggers[11]["task"].cancel()
+ await runner.cleanup_finished_triggers()
+
+
+@pytest.mark.asyncio
+@patch("airflow.jobs.triggerer_job_runner.Trigger._decrypt_kwargs", return_value={})
+@patch("airflow.jobs.triggerer_job_runner.TriggerRunner.get_trigger_by_classpath")
+async def test_create_triggers_skips_asset_store_for_non_event_trigger(
+ mock_get_classpath, mock_decrypt, session
+):
+ """asset_store injection is skipped for plain BaseTrigger (non-BaseEventTrigger) instances."""
+ from airflow.triggers.testing import SuccessTrigger
Review Comment:
Moved.
##########
airflow-core/src/airflow/jobs/triggerer_job_runner.py:
##########
@@ -1286,6 +1293,16 @@ async def create_triggers(self):
trigger_instance.triggerer_job_id = self.job_id
trigger_instance.timeout_after = workload.timeout_after
+ if isinstance(trigger_instance, BaseEventTrigger) and workload.watched_assets:
+ # Reconstruct AssetStateAccessors from watched_assets
Review Comment:
Removed.
##########
airflow-core/tests/unit/jobs/test_triggerer_job.py:
##########
@@ -558,6 +558,248 @@ def test_create_workload_uses_supervisor_id_without_job(jobless_supervisor, mock
assert factory.log_path == f"/logs/ti.trigger.{jobless_supervisor.id}.log"
+def test_create_workload_sets_watched_assets_for_asset_only_trigger(jobless_supervisor, mocker):
Review Comment:
Added test!
##########
airflow-core/tests/unit/jobs/test_triggerer_job.py:
##########
@@ -558,6 +558,248 @@ def test_create_workload_uses_supervisor_id_without_job(jobless_supervisor, mock
assert factory.log_path == f"/logs/ti.trigger.{jobless_supervisor.id}.log"
+def test_create_workload_sets_watched_assets_for_asset_only_trigger(jobless_supervisor, mocker):
+ """_create_workload() should populate watched_assets when trigger.task_instance is None and assets exist."""
+ asset1 = mocker.Mock()
+ asset1.name = "my_asset"
+ asset1.uri = "s3://bucket/key"
+
+ asset2 = mocker.Mock()
+ asset2.name = "other_asset"
+ asset2.uri = "gs://bucket/path"
+
+ trigger = mocker.Mock()
+ trigger.id = 42
+ trigger.classpath = "some.path.Trigger"
+ trigger.encrypted_kwargs = "encrypted"
+ trigger.task_instance = None # Not tied to a Task (similar to a BaseEventTrigger)
+ trigger.assets = [asset1, asset2]
+
+ workload = jobless_supervisor._create_workload(
+ trigger=trigger,
+ dag_bag=mocker.Mock(),
+ render_log_fname=mocker.Mock(),
+ session=mocker.Mock(),
+ )
+
+ assert workload is not None
+ assert workload.watched_assets == {"my_asset": "s3://bucket/key", "other_asset": "gs://bucket/path"}
+
+
+def test_create_workload_watched_assets_none_when_no_assets(jobless_supervisor, mocker):
+ """_create_workload() should set watched_assets=None when trigger.task_instance is None and assets is empty."""
+ trigger = mocker.Mock()
+ trigger.id = 43
+ trigger.classpath = "some.path.Trigger"
+ trigger.encrypted_kwargs = "encrypted"
+ trigger.task_instance = None
+ trigger.assets = [] # No Assets are attached to the trigger
+
+ workload = jobless_supervisor._create_workload(
+ trigger=trigger,
+ dag_bag=mocker.Mock(),
+ render_log_fname=mocker.Mock(),
+ session=mocker.Mock(),
+ )
+
+ assert workload is not None
+ assert workload.watched_assets is None
+
+
+def test_run_trigger_workload_includes_watched_assets_field():
+ """RunTrigger workload should accept and store watched_assets."""
+ from airflow.executors.workloads.trigger import RunTrigger
Review Comment:
Moved.
##########
airflow-core/tests/unit/jobs/test_triggerer_job.py:
##########
@@ -558,6 +558,248 @@ def test_create_workload_uses_supervisor_id_without_job(jobless_supervisor, mock
assert factory.log_path == f"/logs/ti.trigger.{jobless_supervisor.id}.log"
+def test_create_workload_sets_watched_assets_for_asset_only_trigger(jobless_supervisor, mocker):
+ """_create_workload() should populate watched_assets when trigger.task_instance is None and assets exist."""
+ asset1 = mocker.Mock()
+ asset1.name = "my_asset"
+ asset1.uri = "s3://bucket/key"
+
+ asset2 = mocker.Mock()
+ asset2.name = "other_asset"
+ asset2.uri = "gs://bucket/path"
+
+ trigger = mocker.Mock()
+ trigger.id = 42
+ trigger.classpath = "some.path.Trigger"
+ trigger.encrypted_kwargs = "encrypted"
+ trigger.task_instance = None # Not tied to a Task (similar to a BaseEventTrigger)
+ trigger.assets = [asset1, asset2]
+
+ workload = jobless_supervisor._create_workload(
+ trigger=trigger,
+ dag_bag=mocker.Mock(),
+ render_log_fname=mocker.Mock(),
+ session=mocker.Mock(),
+ )
+
+ assert workload is not None
+ assert workload.watched_assets == {"my_asset": "s3://bucket/key", "other_asset": "gs://bucket/path"}
+
+
+def test_create_workload_watched_assets_none_when_no_assets(jobless_supervisor, mocker):
+ """_create_workload() should set watched_assets=None when trigger.task_instance is None and assets is empty."""
+ trigger = mocker.Mock()
+ trigger.id = 43
+ trigger.classpath = "some.path.Trigger"
+ trigger.encrypted_kwargs = "encrypted"
+ trigger.task_instance = None
+ trigger.assets = [] # No Assets are attached to the trigger
+
+ workload = jobless_supervisor._create_workload(
+ trigger=trigger,
+ dag_bag=mocker.Mock(),
+ render_log_fname=mocker.Mock(),
+ session=mocker.Mock(),
+ )
+
+ assert workload is not None
+ assert workload.watched_assets is None
+
+
+def test_run_trigger_workload_includes_watched_assets_field():
+ """RunTrigger workload should accept and store watched_assets."""
+ from airflow.executors.workloads.trigger import RunTrigger
+
+ workload = RunTrigger(
+ id=1,
+ classpath="airflow.triggers.testing.SuccessTrigger",
+ encrypted_kwargs="fake",
+ watched_assets={"asset_a": "s3://a", "asset_b": "gs://b"},
+ )
+ assert workload.watched_assets == {"asset_a": "s3://a", "asset_b": "gs://b"}
+
+
+def test_run_trigger_workload_watched_assets_defaults_to_none():
+ """RunTrigger workload watched_assets should default to None."""
+ from airflow.executors.workloads.trigger import RunTrigger
+
+ workload = RunTrigger(
+ id=1,
+ classpath="airflow.triggers.testing.SuccessTrigger",
+ encrypted_kwargs="fake",
+ )
+ assert workload.watched_assets is None
+
+
+@pytest.mark.asyncio
+@patch("airflow.jobs.triggerer_job_runner.Trigger._decrypt_kwargs", return_value={})
+@patch("airflow.jobs.triggerer_job_runner.TriggerRunner.get_trigger_by_classpath")
+async def test_create_triggers_injects_asset_store_for_base_event_trigger(
+ mock_get_classpath, mock_decrypt, session
+):
+ """asset_store is populated on BaseEventTrigger instances when watched_assets is set."""
+ from airflow.sdk.execution_time.context import AssetStoreAccessors
+ from airflow.triggers.base import BaseEventTrigger, TriggerEvent
+
+ injected_instances = []
+
+ class _WatcherTrigger(BaseEventTrigger):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+ injected_instances.append(self)
+
+ def serialize(self):
+ return (f"{type(self).__module__}.{type(self).__qualname__}", {})
+
+ async def run(self):
+ yield TriggerEvent("done")
+
+ mock_get_classpath.return_value = _WatcherTrigger
+
+ runner = TriggerRunner()
+ runner.to_create.append(
+ workloads.RunTrigger.model_construct(
+ id=10,
+ ti=None,
+ classpath="fake.WatcherTrigger",
+ encrypted_kwargs="fake",
+ watched_assets={"my_asset": "s3://bucket/key"},
+ )
+ )
+
+ await runner.create_triggers()
+
+ # This is only testing that an exception was NOT thrown when creating the Trigger
+ assert 10 in runner.triggers
+
+ assert len(injected_instances) == 1
+ assert injected_instances[0].asset_store is not None
+ assert isinstance(injected_instances[0].asset_store, AssetStoreAccessors)
+
+ runner.triggers[10]["task"].cancel()
+ await runner.cleanup_finished_triggers()
+
+
+@pytest.mark.asyncio
+@patch("airflow.jobs.triggerer_job_runner.Trigger._decrypt_kwargs", return_value={})
+@patch("airflow.jobs.triggerer_job_runner.TriggerRunner.get_trigger_by_classpath")
+async def test_create_triggers_asset_store_none_when_no_watched_assets(
+ mock_get_classpath, mock_decrypt, session
+):
+ """asset_store stays None when watched_assets is not set on the workload."""
+ from airflow.triggers.base import BaseEventTrigger, TriggerEvent
Review Comment:
Moved.
##########
airflow-core/tests/unit/jobs/test_triggerer_job.py:
##########
@@ -558,6 +558,248 @@ def test_create_workload_uses_supervisor_id_without_job(jobless_supervisor, mock
assert factory.log_path == f"/logs/ti.trigger.{jobless_supervisor.id}.log"
+def test_create_workload_sets_watched_assets_for_asset_only_trigger(jobless_supervisor, mocker):
+ """_create_workload() should populate watched_assets when trigger.task_instance is None and assets exist."""
+ asset1 = mocker.Mock()
+ asset1.name = "my_asset"
+ asset1.uri = "s3://bucket/key"
+
+ asset2 = mocker.Mock()
+ asset2.name = "other_asset"
+ asset2.uri = "gs://bucket/path"
+
+ trigger = mocker.Mock()
+ trigger.id = 42
+ trigger.classpath = "some.path.Trigger"
+ trigger.encrypted_kwargs = "encrypted"
+ trigger.task_instance = None # Not tied to a Task (similar to a BaseEventTrigger)
+ trigger.assets = [asset1, asset2]
+
+ workload = jobless_supervisor._create_workload(
+ trigger=trigger,
+ dag_bag=mocker.Mock(),
+ render_log_fname=mocker.Mock(),
+ session=mocker.Mock(),
+ )
+
+ assert workload is not None
+ assert workload.watched_assets == {"my_asset": "s3://bucket/key", "other_asset": "gs://bucket/path"}
+
+
+def test_create_workload_watched_assets_none_when_no_assets(jobless_supervisor, mocker):
+ """_create_workload() should set watched_assets=None when trigger.task_instance is None and assets is empty."""
+ trigger = mocker.Mock()
+ trigger.id = 43
+ trigger.classpath = "some.path.Trigger"
+ trigger.encrypted_kwargs = "encrypted"
+ trigger.task_instance = None
+ trigger.assets = [] # No Assets are attached to the trigger
+
+ workload = jobless_supervisor._create_workload(
+ trigger=trigger,
+ dag_bag=mocker.Mock(),
+ render_log_fname=mocker.Mock(),
+ session=mocker.Mock(),
+ )
+
+ assert workload is not None
+ assert workload.watched_assets is None
+
+
+def test_run_trigger_workload_includes_watched_assets_field():
+ """RunTrigger workload should accept and store watched_assets."""
+ from airflow.executors.workloads.trigger import RunTrigger
+
+ workload = RunTrigger(
+ id=1,
+ classpath="airflow.triggers.testing.SuccessTrigger",
+ encrypted_kwargs="fake",
+ watched_assets={"asset_a": "s3://a", "asset_b": "gs://b"},
+ )
+ assert workload.watched_assets == {"asset_a": "s3://a", "asset_b": "gs://b"}
+
+
+def test_run_trigger_workload_watched_assets_defaults_to_none():
+ """RunTrigger workload watched_assets should default to None."""
+ from airflow.executors.workloads.trigger import RunTrigger
+
+ workload = RunTrigger(
+ id=1,
+ classpath="airflow.triggers.testing.SuccessTrigger",
+ encrypted_kwargs="fake",
+ )
+ assert workload.watched_assets is None
+
+
+@pytest.mark.asyncio
+@patch("airflow.jobs.triggerer_job_runner.Trigger._decrypt_kwargs", return_value={})
+@patch("airflow.jobs.triggerer_job_runner.TriggerRunner.get_trigger_by_classpath")
+async def test_create_triggers_injects_asset_store_for_base_event_trigger(
+ mock_get_classpath, mock_decrypt, session
+):
+ """asset_store is populated on BaseEventTrigger instances when watched_assets is set."""
+ from airflow.sdk.execution_time.context import AssetStoreAccessors
+ from airflow.triggers.base import BaseEventTrigger, TriggerEvent
+
+ injected_instances = []
+
+ class _WatcherTrigger(BaseEventTrigger):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+ injected_instances.append(self)
+
+ def serialize(self):
+ return (f"{type(self).__module__}.{type(self).__qualname__}", {})
+
+ async def run(self):
+ yield TriggerEvent("done")
+
+ mock_get_classpath.return_value = _WatcherTrigger
+
+ runner = TriggerRunner()
+ runner.to_create.append(
+ workloads.RunTrigger.model_construct(
+ id=10,
+ ti=None,
+ classpath="fake.WatcherTrigger",
+ encrypted_kwargs="fake",
+ watched_assets={"my_asset": "s3://bucket/key"},
+ )
+ )
+
+ await runner.create_triggers()
+
+ # This is only testing that an exception was NOT thrown when creating the Trigger
+ assert 10 in runner.triggers
+
+ assert len(injected_instances) == 1
+ assert injected_instances[0].asset_store is not None
+ assert isinstance(injected_instances[0].asset_store, AssetStoreAccessors)
+
+ runner.triggers[10]["task"].cancel()
+ await runner.cleanup_finished_triggers()
+
+
+@pytest.mark.asyncio
+@patch("airflow.jobs.triggerer_job_runner.Trigger._decrypt_kwargs", return_value={})
+@patch("airflow.jobs.triggerer_job_runner.TriggerRunner.get_trigger_by_classpath")
+async def test_create_triggers_asset_store_none_when_no_watched_assets(
+ mock_get_classpath, mock_decrypt, session
+):
+ """asset_store stays None when watched_assets is not set on the workload."""
+ from airflow.triggers.base import BaseEventTrigger, TriggerEvent
+
+ injected_instances = []
+
+ class _WatcherTrigger(BaseEventTrigger):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+ injected_instances.append(self)
+
+ def serialize(self):
+ return (f"{type(self).__module__}.{type(self).__qualname__}", {})
+
+ async def run(self):
+ yield TriggerEvent("done")
+
+ mock_get_classpath.return_value = _WatcherTrigger
+
+ runner = TriggerRunner()
+ runner.to_create.append(
+ workloads.RunTrigger.model_construct(
+ id=11,
+ ti=None,
+ classpath="fake.WatcherTrigger",
+ encrypted_kwargs="fake",
+ watched_assets=None,
+ )
+ )
+
+ await runner.create_triggers()
+
+ assert len(injected_instances) == 1
+ assert injected_instances[0].asset_store is None
+
+ runner.triggers[11]["task"].cancel()
+ await runner.cleanup_finished_triggers()
+
+
+@pytest.mark.asyncio
+@patch("airflow.jobs.triggerer_job_runner.Trigger._decrypt_kwargs", return_value={})
+@patch("airflow.jobs.triggerer_job_runner.TriggerRunner.get_trigger_by_classpath")
+async def test_create_triggers_skips_asset_store_for_non_event_trigger(
+ mock_get_classpath, mock_decrypt, session
+):
+ """asset_store injection is skipped for plain BaseTrigger (non-BaseEventTrigger) instances."""
+ from airflow.triggers.testing import SuccessTrigger
+
+ mock_get_classpath.return_value = SuccessTrigger
+
+ runner = TriggerRunner()
+ runner.to_create.append(
+ workloads.RunTrigger.model_construct(
+ id=12, ti=None, classpath="airflow.triggers.testing.SuccessTrigger", encrypted_kwargs="fake"
+ )
+ )
+
+ await runner.create_triggers()
+
+ assert 12 in runner.triggers
+ assert not hasattr(runner.triggers[12]["task"], "asset_store")
+
+ runner.triggers[12]["task"].cancel()
+ await runner.cleanup_finished_triggers()
+
+
+@pytest.mark.asyncio
+@patch("airflow.jobs.triggerer_job_runner.Trigger._decrypt_kwargs", return_value={})
+@patch("airflow.jobs.triggerer_job_runner.TriggerRunner.get_trigger_by_classpath")
+async def test_create_triggers_asset_store_contains_correct_assets(mock_get_classpath, mock_decrypt, session):
+ """AssetStoreAccessors built from watched_assets has entries for all provided name/URI pairs."""
+ from airflow.sdk.definitions.asset import Asset
+ from airflow.sdk.execution_time.context import AssetStoreAccessors
+ from airflow.triggers.base import BaseEventTrigger, TriggerEvent
Review Comment:
Moved.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Passing `AssetStateStoreAccessors` through to `BaseEventTrigger` [airflow]
Posted by "jroachgolf84 (via GitHub)" <gi...@apache.org>.
jroachgolf84 commented on PR #67839:
URL: https://github.com/apache/airflow/pull/67839#issuecomment-4695291313
Going to wait for @amoghrajesh to merge #68438 to make changes.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Passing `AssetStateAccessors` through to `BaseEventTrigger` [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #67839:
URL: https://github.com/apache/airflow/pull/67839#discussion_r3338985312
##########
airflow-core/src/airflow/jobs/triggerer_job_runner.py:
##########
Review Comment:
The `_handle_request` should be handling this in here?
##########
airflow-core/src/airflow/jobs/triggerer_job_runner.py:
##########
@@ -1266,6 +1269,16 @@ async def create_triggers(self):
trigger_instance.triggerer_job_id = self.job_id
trigger_instance.timeout_after = workload.timeout_after
+ if isinstance(trigger_instance, BaseEventTrigger) and workload.watched_assets:
+ # Reconstruct AssetStateAccessors from watched_assets
+ from airflow.sdk.definitions.asset import Asset
+ from airflow.sdk.execution_time.context import AssetStateAccessors
Review Comment:
Already a top level import
##########
airflow-core/src/airflow/jobs/triggerer_job_runner.py:
##########
@@ -741,11 +737,18 @@ def _create_workload(
render_log_fname: Callable[..., str],
session: Session,
) -> workloads.RunTrigger | None:
+ # Pass the "watched" Assets through for downstream use in BaseEventTrigger
if trigger.task_instance is None:
+ watched_assets: dict[str, str] | None = None
+
+ if trigger.asset_watchers:
+ watched_assets = {a.name: a.uri for a in trigger.assets}
Review Comment:
Either use trigger.assets or trigger.asset_watchers for if check and iteration
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Passing `AssetStateStoreAccessors` through to `BaseEventTrigger` [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #67839:
URL: https://github.com/apache/airflow/pull/67839#discussion_r3433402306
##########
airflow-core/src/airflow/jobs/triggerer_job_runner.py:
##########
@@ -751,11 +753,18 @@ def _create_workload(
render_log_fname: Callable[..., str],
session: Session,
) -> workloads.RunTrigger | None:
+ # Pass the "watched" Assets through for downstream use in BaseEventTrigger
Review Comment:
```suggestion
```
It's obvious here
##########
airflow-core/src/airflow/triggers/base.py:
##########
@@ -297,6 +298,12 @@ class BaseEventTrigger(BaseTrigger):
supports_triggerer_queue: bool = False
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+
+ # Injected by the triggerer before run() is called; mirrors how trigger_id is set
Review Comment:
```suggestion
# Injected by the triggerer before run() is called
```
##########
airflow-core/tests/unit/triggers/test_base_trigger.py:
##########
@@ -232,3 +234,25 @@ async def stream():
payloads = [event.payload async for event in us.filter_shared_stream(stream())]
assert [p["region"] for p in payloads] == ["us", "us"]
+
+
+def test_base_event_trigger_asset_state_store_initialized_to_none():
+ """asset_state_store is None before it is set."""
+ trigger = _PlainEventTrigger()
+ assert trigger.asset_state_store is None
+
+
+def test_base_event_trigger_asset_state_store_can_be_set():
+ """asset_state_store can be set once the Trigger is initialized."""
+ trigger = _PlainEventTrigger()
+ mock_store = MagicMock()
+ trigger.asset_state_store = mock_store
+ assert trigger.asset_state_store is mock_store
+
+
+def test_base_event_trigger_asset_state_store_independent_across_instances():
+ """a.asset_state_store does not impact b.asset_state_store."""
+ a = _PlainEventTrigger(name="a")
+ b = _PlainEventTrigger(name="b")
+ a.asset_state_store = MagicMock()
+ assert b.asset_state_store is None
Review Comment:
Use `spec` or `autospec` for mocks pls.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Passing `AssetStateStoreAccessors` through to `BaseEventTrigger` [airflow]
Posted by "jroachgolf84 (via GitHub)" <gi...@apache.org>.
jroachgolf84 commented on code in PR #67839:
URL: https://github.com/apache/airflow/pull/67839#discussion_r3439919099
##########
airflow-core/src/airflow/jobs/triggerer_job_runner.py:
##########
@@ -1313,6 +1321,11 @@ async def create_triggers(self):
trigger_instance.triggerer_job_id = self.job_id
trigger_instance.timeout_after = workload.timeout_after
+ if isinstance(trigger_instance, BaseEventTrigger) and workload.watched_assets:
Review Comment:
Just gave this a try, and as much as I wanted it not to be the case, you were right ;) . Here's the exception I was getting:
```
Traceback (most recent call last):
File "/opt/airflow/task-sdk/src/airflow/sdk/execution_time/supervisor.py", line 890, in handle_requests
msg = self.decoder.validate_python(self._deserialize_request(request.body))
File "/usr/python/lib/python3.10/site-packages/pydantic/type_adapter.py", line 441, in validate_python
return self.validator.validate_python(
pydantic_core._pydantic_core.ValidationError: 1 validation error for tagged-union[
TriggerStateChanges, GetConnection, DeleteVariable, GetVariable, GetVariableKeys,
PutVariable, DeleteXCom, GetXCom, SetXCom, GetTICount, GetTaskStates, GetDagRunState,
GetDRCount, GetPreviousTI, GetHITLDetailResponse, UpdateHITLDetail, MaskSecret
]
Input tag 'GetAssetStateStoreByName' found using 'type' does not match any of the expected tags:
'TriggerStateChanges', 'GetConnection', 'DeleteVariable', 'GetVariable', 'GetVariableKeys',
'PutVariable', 'DeleteXCom', 'GetXCom', 'SetXCom', 'GetTICount', 'GetTaskStates',
'GetDagRunState', 'GetDRCount', 'GetPreviousTI', 'GetHITLDetailResponse', 'UpdateHITLDetail',
'MaskSecret'
[type=union_tag_invalid, input_value={'name': 'generic_asset', ... 'type': 'GetAssetStateStoreByName'}, input_type=dict]
```
I'd like to get this out in a follow-up PR. I'll add that to the description here.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Passing `AssetStateStoreAccessors` through to `BaseEventTrigger` [airflow]
Posted by "jroachgolf84 (via GitHub)" <gi...@apache.org>.
jroachgolf84 commented on code in PR #67839:
URL: https://github.com/apache/airflow/pull/67839#discussion_r3439955611
##########
airflow-core/src/airflow/jobs/triggerer_job_runner.py:
##########
@@ -1313,6 +1321,11 @@ async def create_triggers(self):
trigger_instance.triggerer_job_id = self.job_id
trigger_instance.timeout_after = workload.timeout_after
+ if isinstance(trigger_instance, BaseEventTrigger) and workload.watched_assets:
Review Comment:
Note added, should be good to add this in a later PR.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Passing `AssetStateAccessors` through to `BaseEventTrigger` [airflow]
Posted by "jroachgolf84 (via GitHub)" <gi...@apache.org>.
jroachgolf84 commented on code in PR #67839:
URL: https://github.com/apache/airflow/pull/67839#discussion_r3341074213
##########
airflow-core/src/airflow/jobs/triggerer_job_runner.py:
##########
Review Comment:
What do you mean?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Passing `AssetStateAccessors` through to `BaseEventTrigger` [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #67839:
URL: https://github.com/apache/airflow/pull/67839#discussion_r3388013386
##########
airflow-core/src/airflow/jobs/triggerer_job_runner.py:
##########
@@ -1286,6 +1293,16 @@ async def create_triggers(self):
trigger_instance.triggerer_job_id = self.job_id
trigger_instance.timeout_after = workload.timeout_after
+ if isinstance(trigger_instance, BaseEventTrigger) and workload.watched_assets:
+ # Reconstruct AssetStateAccessors from watched_assets
+ from airflow.sdk.definitions.asset import Asset
+ from airflow.sdk.execution_time.context import AssetStoreAccessors
Review Comment:
Is there a need to have these imports inside here? Circular import or such
##########
airflow-core/src/airflow/jobs/triggerer_job_runner.py:
##########
@@ -1286,6 +1293,16 @@ async def create_triggers(self):
trigger_instance.triggerer_job_id = self.job_id
trigger_instance.timeout_after = workload.timeout_after
+ if isinstance(trigger_instance, BaseEventTrigger) and workload.watched_assets:
+ # Reconstruct AssetStateAccessors from watched_assets
+ from airflow.sdk.definitions.asset import Asset
+ from airflow.sdk.execution_time.context import AssetStoreAccessors
+
+ # Potentially address Asset vs. AssetRef, AssetUriRef, etc.
Review Comment:
This has to be done?
##########
airflow-core/tests/unit/jobs/test_triggerer_job.py:
##########
@@ -558,6 +558,248 @@ def test_create_workload_uses_supervisor_id_without_job(jobless_supervisor, mock
assert factory.log_path == f"/logs/ti.trigger.{jobless_supervisor.id}.log"
+def test_create_workload_sets_watched_assets_for_asset_only_trigger(jobless_supervisor, mocker):
Review Comment:
For all tests added, please use any Mocks with `spec`, `spec=Asset` or `create_autospec(Asset)` whatever deems right
##########
airflow-core/tests/unit/jobs/test_triggerer_job.py:
##########
@@ -558,6 +558,248 @@ def test_create_workload_uses_supervisor_id_without_job(jobless_supervisor, mock
assert factory.log_path == f"/logs/ti.trigger.{jobless_supervisor.id}.log"
+def test_create_workload_sets_watched_assets_for_asset_only_trigger(jobless_supervisor, mocker):
Review Comment:
Tests verify that `asset_store` is injected and is an `AssetStoreAccessors` instance, but none of them verify that it actually reads or writes, we should add those
##########
airflow-core/tests/unit/jobs/test_triggerer_job.py:
##########
@@ -558,6 +558,248 @@ def test_create_workload_uses_supervisor_id_without_job(jobless_supervisor, mock
assert factory.log_path == f"/logs/ti.trigger.{jobless_supervisor.id}.log"
+def test_create_workload_sets_watched_assets_for_asset_only_trigger(jobless_supervisor, mocker):
+ """_create_workload() should populate watched_assets when trigger.task_instance is None and assets exist."""
+ asset1 = mocker.Mock()
+ asset1.name = "my_asset"
+ asset1.uri = "s3://bucket/key"
+
+ asset2 = mocker.Mock()
+ asset2.name = "other_asset"
+ asset2.uri = "gs://bucket/path"
+
+ trigger = mocker.Mock()
+ trigger.id = 42
+ trigger.classpath = "some.path.Trigger"
+ trigger.encrypted_kwargs = "encrypted"
+ trigger.task_instance = None # Not tied to a Task (similar to a BaseEventTrigger)
+ trigger.assets = [asset1, asset2]
+
+ workload = jobless_supervisor._create_workload(
+ trigger=trigger,
+ dag_bag=mocker.Mock(),
+ render_log_fname=mocker.Mock(),
+ session=mocker.Mock(),
+ )
+
+ assert workload is not None
+ assert workload.watched_assets == {"my_asset": "s3://bucket/key", "other_asset": "gs://bucket/path"}
+
+
+def test_create_workload_watched_assets_none_when_no_assets(jobless_supervisor, mocker):
+ """_create_workload() should set watched_assets=None when trigger.task_instance is None and assets is empty."""
+ trigger = mocker.Mock()
+ trigger.id = 43
+ trigger.classpath = "some.path.Trigger"
+ trigger.encrypted_kwargs = "encrypted"
+ trigger.task_instance = None
+ trigger.assets = [] # No Assets are attached to the trigger
+
+ workload = jobless_supervisor._create_workload(
+ trigger=trigger,
+ dag_bag=mocker.Mock(),
+ render_log_fname=mocker.Mock(),
+ session=mocker.Mock(),
+ )
+
+ assert workload is not None
+ assert workload.watched_assets is None
+
+
+def test_run_trigger_workload_includes_watched_assets_field():
+ """RunTrigger workload should accept and store watched_assets."""
+ from airflow.executors.workloads.trigger import RunTrigger
+
+ workload = RunTrigger(
+ id=1,
+ classpath="airflow.triggers.testing.SuccessTrigger",
+ encrypted_kwargs="fake",
+ watched_assets={"asset_a": "s3://a", "asset_b": "gs://b"},
+ )
+ assert workload.watched_assets == {"asset_a": "s3://a", "asset_b": "gs://b"}
+
+
+def test_run_trigger_workload_watched_assets_defaults_to_none():
+ """RunTrigger workload watched_assets should default to None."""
+ from airflow.executors.workloads.trigger import RunTrigger
+
+ workload = RunTrigger(
+ id=1,
+ classpath="airflow.triggers.testing.SuccessTrigger",
+ encrypted_kwargs="fake",
+ )
+ assert workload.watched_assets is None
+
+
+@pytest.mark.asyncio
+@patch("airflow.jobs.triggerer_job_runner.Trigger._decrypt_kwargs", return_value={})
+@patch("airflow.jobs.triggerer_job_runner.TriggerRunner.get_trigger_by_classpath")
+async def test_create_triggers_injects_asset_store_for_base_event_trigger(
+ mock_get_classpath, mock_decrypt, session
+):
+ """asset_store is populated on BaseEventTrigger instances when watched_assets is set."""
+ from airflow.sdk.execution_time.context import AssetStoreAccessors
+ from airflow.triggers.base import BaseEventTrigger, TriggerEvent
+
+ injected_instances = []
+
+ class _WatcherTrigger(BaseEventTrigger):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+ injected_instances.append(self)
+
+ def serialize(self):
+ return (f"{type(self).__module__}.{type(self).__qualname__}", {})
+
+ async def run(self):
+ yield TriggerEvent("done")
+
+ mock_get_classpath.return_value = _WatcherTrigger
+
+ runner = TriggerRunner()
+ runner.to_create.append(
+ workloads.RunTrigger.model_construct(
+ id=10,
+ ti=None,
+ classpath="fake.WatcherTrigger",
+ encrypted_kwargs="fake",
+ watched_assets={"my_asset": "s3://bucket/key"},
+ )
+ )
+
+ await runner.create_triggers()
+
+ # This is only testing that an exception was NOT thrown when creating the Trigger
+ assert 10 in runner.triggers
+
+ assert len(injected_instances) == 1
+ assert injected_instances[0].asset_store is not None
+ assert isinstance(injected_instances[0].asset_store, AssetStoreAccessors)
+
+ runner.triggers[10]["task"].cancel()
+ await runner.cleanup_finished_triggers()
+
+
+@pytest.mark.asyncio
+@patch("airflow.jobs.triggerer_job_runner.Trigger._decrypt_kwargs", return_value={})
+@patch("airflow.jobs.triggerer_job_runner.TriggerRunner.get_trigger_by_classpath")
+async def test_create_triggers_asset_store_none_when_no_watched_assets(
+ mock_get_classpath, mock_decrypt, session
+):
+ """asset_store stays None when watched_assets is not set on the workload."""
+ from airflow.triggers.base import BaseEventTrigger, TriggerEvent
+
+ injected_instances = []
+
+ class _WatcherTrigger(BaseEventTrigger):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+ injected_instances.append(self)
+
+ def serialize(self):
+ return (f"{type(self).__module__}.{type(self).__qualname__}", {})
+
+ async def run(self):
+ yield TriggerEvent("done")
+
+ mock_get_classpath.return_value = _WatcherTrigger
+
+ runner = TriggerRunner()
+ runner.to_create.append(
+ workloads.RunTrigger.model_construct(
+ id=11,
+ ti=None,
+ classpath="fake.WatcherTrigger",
+ encrypted_kwargs="fake",
+ watched_assets=None,
+ )
+ )
+
+ await runner.create_triggers()
+
+ assert len(injected_instances) == 1
+ assert injected_instances[0].asset_store is None
+
+ runner.triggers[11]["task"].cancel()
+ await runner.cleanup_finished_triggers()
+
+
+@pytest.mark.asyncio
+@patch("airflow.jobs.triggerer_job_runner.Trigger._decrypt_kwargs", return_value={})
+@patch("airflow.jobs.triggerer_job_runner.TriggerRunner.get_trigger_by_classpath")
+async def test_create_triggers_skips_asset_store_for_non_event_trigger(
+ mock_get_classpath, mock_decrypt, session
+):
+ """asset_store injection is skipped for plain BaseTrigger (non-BaseEventTrigger) instances."""
+ from airflow.triggers.testing import SuccessTrigger
Review Comment:
Same here
##########
airflow-core/tests/unit/jobs/test_triggerer_job.py:
##########
@@ -558,6 +558,248 @@ def test_create_workload_uses_supervisor_id_without_job(jobless_supervisor, mock
assert factory.log_path == f"/logs/ti.trigger.{jobless_supervisor.id}.log"
+def test_create_workload_sets_watched_assets_for_asset_only_trigger(jobless_supervisor, mocker):
+ """_create_workload() should populate watched_assets when trigger.task_instance is None and assets exist."""
+ asset1 = mocker.Mock()
+ asset1.name = "my_asset"
+ asset1.uri = "s3://bucket/key"
+
+ asset2 = mocker.Mock()
+ asset2.name = "other_asset"
+ asset2.uri = "gs://bucket/path"
+
+ trigger = mocker.Mock()
+ trigger.id = 42
+ trigger.classpath = "some.path.Trigger"
+ trigger.encrypted_kwargs = "encrypted"
+ trigger.task_instance = None # Not tied to a Task (similar to a BaseEventTrigger)
+ trigger.assets = [asset1, asset2]
+
+ workload = jobless_supervisor._create_workload(
+ trigger=trigger,
+ dag_bag=mocker.Mock(),
+ render_log_fname=mocker.Mock(),
+ session=mocker.Mock(),
+ )
+
+ assert workload is not None
+ assert workload.watched_assets == {"my_asset": "s3://bucket/key", "other_asset": "gs://bucket/path"}
+
+
+def test_create_workload_watched_assets_none_when_no_assets(jobless_supervisor, mocker):
+ """_create_workload() should set watched_assets=None when trigger.task_instance is None and assets is empty."""
+ trigger = mocker.Mock()
+ trigger.id = 43
+ trigger.classpath = "some.path.Trigger"
+ trigger.encrypted_kwargs = "encrypted"
+ trigger.task_instance = None
+ trigger.assets = [] # No Assets are attached to the trigger
+
+ workload = jobless_supervisor._create_workload(
+ trigger=trigger,
+ dag_bag=mocker.Mock(),
+ render_log_fname=mocker.Mock(),
+ session=mocker.Mock(),
+ )
+
+ assert workload is not None
+ assert workload.watched_assets is None
+
+
+def test_run_trigger_workload_includes_watched_assets_field():
+ """RunTrigger workload should accept and store watched_assets."""
+ from airflow.executors.workloads.trigger import RunTrigger
+
+ workload = RunTrigger(
+ id=1,
+ classpath="airflow.triggers.testing.SuccessTrigger",
+ encrypted_kwargs="fake",
+ watched_assets={"asset_a": "s3://a", "asset_b": "gs://b"},
+ )
+ assert workload.watched_assets == {"asset_a": "s3://a", "asset_b": "gs://b"}
+
+
+def test_run_trigger_workload_watched_assets_defaults_to_none():
+ """RunTrigger workload watched_assets should default to None."""
+ from airflow.executors.workloads.trigger import RunTrigger
+
+ workload = RunTrigger(
+ id=1,
+ classpath="airflow.triggers.testing.SuccessTrigger",
+ encrypted_kwargs="fake",
+ )
+ assert workload.watched_assets is None
+
+
+@pytest.mark.asyncio
+@patch("airflow.jobs.triggerer_job_runner.Trigger._decrypt_kwargs", return_value={})
+@patch("airflow.jobs.triggerer_job_runner.TriggerRunner.get_trigger_by_classpath")
+async def test_create_triggers_injects_asset_store_for_base_event_trigger(
+ mock_get_classpath, mock_decrypt, session
+):
+ """asset_store is populated on BaseEventTrigger instances when watched_assets is set."""
+ from airflow.sdk.execution_time.context import AssetStoreAccessors
+ from airflow.triggers.base import BaseEventTrigger, TriggerEvent
+
+ injected_instances = []
+
+ class _WatcherTrigger(BaseEventTrigger):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+ injected_instances.append(self)
+
+ def serialize(self):
+ return (f"{type(self).__module__}.{type(self).__qualname__}", {})
+
+ async def run(self):
+ yield TriggerEvent("done")
+
+ mock_get_classpath.return_value = _WatcherTrigger
+
+ runner = TriggerRunner()
+ runner.to_create.append(
+ workloads.RunTrigger.model_construct(
+ id=10,
+ ti=None,
+ classpath="fake.WatcherTrigger",
+ encrypted_kwargs="fake",
+ watched_assets={"my_asset": "s3://bucket/key"},
+ )
+ )
+
+ await runner.create_triggers()
+
+ # This is only testing that an exception was NOT thrown when creating the Trigger
+ assert 10 in runner.triggers
+
+ assert len(injected_instances) == 1
+ assert injected_instances[0].asset_store is not None
+ assert isinstance(injected_instances[0].asset_store, AssetStoreAccessors)
+
+ runner.triggers[10]["task"].cancel()
+ await runner.cleanup_finished_triggers()
+
+
+@pytest.mark.asyncio
+@patch("airflow.jobs.triggerer_job_runner.Trigger._decrypt_kwargs", return_value={})
+@patch("airflow.jobs.triggerer_job_runner.TriggerRunner.get_trigger_by_classpath")
+async def test_create_triggers_asset_store_none_when_no_watched_assets(
+ mock_get_classpath, mock_decrypt, session
+):
+ """asset_store stays None when watched_assets is not set on the workload."""
+ from airflow.triggers.base import BaseEventTrigger, TriggerEvent
+
+ injected_instances = []
+
+ class _WatcherTrigger(BaseEventTrigger):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+ injected_instances.append(self)
+
+ def serialize(self):
+ return (f"{type(self).__module__}.{type(self).__qualname__}", {})
+
+ async def run(self):
+ yield TriggerEvent("done")
+
+ mock_get_classpath.return_value = _WatcherTrigger
+
+ runner = TriggerRunner()
+ runner.to_create.append(
+ workloads.RunTrigger.model_construct(
+ id=11,
+ ti=None,
+ classpath="fake.WatcherTrigger",
+ encrypted_kwargs="fake",
+ watched_assets=None,
+ )
+ )
+
+ await runner.create_triggers()
+
+ assert len(injected_instances) == 1
+ assert injected_instances[0].asset_store is None
+
+ runner.triggers[11]["task"].cancel()
+ await runner.cleanup_finished_triggers()
+
+
+@pytest.mark.asyncio
+@patch("airflow.jobs.triggerer_job_runner.Trigger._decrypt_kwargs", return_value={})
+@patch("airflow.jobs.triggerer_job_runner.TriggerRunner.get_trigger_by_classpath")
+async def test_create_triggers_skips_asset_store_for_non_event_trigger(
+ mock_get_classpath, mock_decrypt, session
+):
+ """asset_store injection is skipped for plain BaseTrigger (non-BaseEventTrigger) instances."""
+ from airflow.triggers.testing import SuccessTrigger
+
+ mock_get_classpath.return_value = SuccessTrigger
+
+ runner = TriggerRunner()
+ runner.to_create.append(
+ workloads.RunTrigger.model_construct(
+ id=12, ti=None, classpath="airflow.triggers.testing.SuccessTrigger", encrypted_kwargs="fake"
+ )
+ )
+
+ await runner.create_triggers()
+
+ assert 12 in runner.triggers
+ assert not hasattr(runner.triggers[12]["task"], "asset_store")
+
+ runner.triggers[12]["task"].cancel()
+ await runner.cleanup_finished_triggers()
+
+
+@pytest.mark.asyncio
+@patch("airflow.jobs.triggerer_job_runner.Trigger._decrypt_kwargs", return_value={})
+@patch("airflow.jobs.triggerer_job_runner.TriggerRunner.get_trigger_by_classpath")
+async def test_create_triggers_asset_store_contains_correct_assets(mock_get_classpath, mock_decrypt, session):
+ """AssetStoreAccessors built from watched_assets has entries for all provided name/URI pairs."""
+ from airflow.sdk.definitions.asset import Asset
+ from airflow.sdk.execution_time.context import AssetStoreAccessors
+ from airflow.triggers.base import BaseEventTrigger, TriggerEvent
Review Comment:
Same here
##########
airflow-core/src/airflow/jobs/triggerer_job_runner.py:
##########
@@ -1286,6 +1293,16 @@ async def create_triggers(self):
trigger_instance.triggerer_job_id = self.job_id
trigger_instance.timeout_after = workload.timeout_after
+ if isinstance(trigger_instance, BaseEventTrigger) and workload.watched_assets:
+ # Reconstruct AssetStateAccessors from watched_assets
Review Comment:
```suggestion
```
##########
airflow-core/tests/unit/jobs/test_triggerer_job.py:
##########
@@ -558,6 +558,248 @@ def test_create_workload_uses_supervisor_id_without_job(jobless_supervisor, mock
assert factory.log_path == f"/logs/ti.trigger.{jobless_supervisor.id}.log"
+def test_create_workload_sets_watched_assets_for_asset_only_trigger(jobless_supervisor, mocker):
+ """_create_workload() should populate watched_assets when trigger.task_instance is None and assets exist."""
+ asset1 = mocker.Mock()
+ asset1.name = "my_asset"
+ asset1.uri = "s3://bucket/key"
+
+ asset2 = mocker.Mock()
+ asset2.name = "other_asset"
+ asset2.uri = "gs://bucket/path"
+
+ trigger = mocker.Mock()
+ trigger.id = 42
+ trigger.classpath = "some.path.Trigger"
+ trigger.encrypted_kwargs = "encrypted"
+ trigger.task_instance = None # Not tied to a Task (similar to a BaseEventTrigger)
+ trigger.assets = [asset1, asset2]
+
+ workload = jobless_supervisor._create_workload(
+ trigger=trigger,
+ dag_bag=mocker.Mock(),
+ render_log_fname=mocker.Mock(),
+ session=mocker.Mock(),
+ )
+
+ assert workload is not None
+ assert workload.watched_assets == {"my_asset": "s3://bucket/key", "other_asset": "gs://bucket/path"}
+
+
+def test_create_workload_watched_assets_none_when_no_assets(jobless_supervisor, mocker):
+ """_create_workload() should set watched_assets=None when trigger.task_instance is None and assets is empty."""
+ trigger = mocker.Mock()
+ trigger.id = 43
+ trigger.classpath = "some.path.Trigger"
+ trigger.encrypted_kwargs = "encrypted"
+ trigger.task_instance = None
+ trigger.assets = [] # No Assets are attached to the trigger
+
+ workload = jobless_supervisor._create_workload(
+ trigger=trigger,
+ dag_bag=mocker.Mock(),
+ render_log_fname=mocker.Mock(),
+ session=mocker.Mock(),
+ )
+
+ assert workload is not None
+ assert workload.watched_assets is None
+
+
+def test_run_trigger_workload_includes_watched_assets_field():
+ """RunTrigger workload should accept and store watched_assets."""
+ from airflow.executors.workloads.trigger import RunTrigger
+
+ workload = RunTrigger(
+ id=1,
+ classpath="airflow.triggers.testing.SuccessTrigger",
+ encrypted_kwargs="fake",
+ watched_assets={"asset_a": "s3://a", "asset_b": "gs://b"},
+ )
+ assert workload.watched_assets == {"asset_a": "s3://a", "asset_b": "gs://b"}
+
+
+def test_run_trigger_workload_watched_assets_defaults_to_none():
+ """RunTrigger workload watched_assets should default to None."""
+ from airflow.executors.workloads.trigger import RunTrigger
+
+ workload = RunTrigger(
+ id=1,
+ classpath="airflow.triggers.testing.SuccessTrigger",
+ encrypted_kwargs="fake",
+ )
+ assert workload.watched_assets is None
+
+
+@pytest.mark.asyncio
+@patch("airflow.jobs.triggerer_job_runner.Trigger._decrypt_kwargs", return_value={})
+@patch("airflow.jobs.triggerer_job_runner.TriggerRunner.get_trigger_by_classpath")
+async def test_create_triggers_injects_asset_store_for_base_event_trigger(
+ mock_get_classpath, mock_decrypt, session
+):
+ """asset_store is populated on BaseEventTrigger instances when watched_assets is set."""
+ from airflow.sdk.execution_time.context import AssetStoreAccessors
+ from airflow.triggers.base import BaseEventTrigger, TriggerEvent
Review Comment:
Move import to top pls
##########
airflow-core/tests/unit/jobs/test_triggerer_job.py:
##########
@@ -558,6 +558,248 @@ def test_create_workload_uses_supervisor_id_without_job(jobless_supervisor, mock
assert factory.log_path == f"/logs/ti.trigger.{jobless_supervisor.id}.log"
+def test_create_workload_sets_watched_assets_for_asset_only_trigger(jobless_supervisor, mocker):
+ """_create_workload() should populate watched_assets when trigger.task_instance is None and assets exist."""
+ asset1 = mocker.Mock()
+ asset1.name = "my_asset"
+ asset1.uri = "s3://bucket/key"
+
+ asset2 = mocker.Mock()
+ asset2.name = "other_asset"
+ asset2.uri = "gs://bucket/path"
+
+ trigger = mocker.Mock()
+ trigger.id = 42
+ trigger.classpath = "some.path.Trigger"
+ trigger.encrypted_kwargs = "encrypted"
+ trigger.task_instance = None # Not tied to a Task (similar to a BaseEventTrigger)
+ trigger.assets = [asset1, asset2]
+
+ workload = jobless_supervisor._create_workload(
+ trigger=trigger,
+ dag_bag=mocker.Mock(),
+ render_log_fname=mocker.Mock(),
+ session=mocker.Mock(),
+ )
+
+ assert workload is not None
+ assert workload.watched_assets == {"my_asset": "s3://bucket/key", "other_asset": "gs://bucket/path"}
+
+
+def test_create_workload_watched_assets_none_when_no_assets(jobless_supervisor, mocker):
+ """_create_workload() should set watched_assets=None when trigger.task_instance is None and assets is empty."""
+ trigger = mocker.Mock()
+ trigger.id = 43
+ trigger.classpath = "some.path.Trigger"
+ trigger.encrypted_kwargs = "encrypted"
+ trigger.task_instance = None
+ trigger.assets = [] # No Assets are attached to the trigger
+
+ workload = jobless_supervisor._create_workload(
+ trigger=trigger,
+ dag_bag=mocker.Mock(),
+ render_log_fname=mocker.Mock(),
+ session=mocker.Mock(),
+ )
+
+ assert workload is not None
+ assert workload.watched_assets is None
+
+
+def test_run_trigger_workload_includes_watched_assets_field():
+ """RunTrigger workload should accept and store watched_assets."""
+ from airflow.executors.workloads.trigger import RunTrigger
Review Comment:
Top level import pls
##########
airflow-core/tests/unit/jobs/test_triggerer_job.py:
##########
@@ -558,6 +558,248 @@ def test_create_workload_uses_supervisor_id_without_job(jobless_supervisor, mock
assert factory.log_path == f"/logs/ti.trigger.{jobless_supervisor.id}.log"
+def test_create_workload_sets_watched_assets_for_asset_only_trigger(jobless_supervisor, mocker):
+ """_create_workload() should populate watched_assets when trigger.task_instance is None and assets exist."""
+ asset1 = mocker.Mock()
+ asset1.name = "my_asset"
+ asset1.uri = "s3://bucket/key"
+
+ asset2 = mocker.Mock()
+ asset2.name = "other_asset"
+ asset2.uri = "gs://bucket/path"
+
+ trigger = mocker.Mock()
+ trigger.id = 42
+ trigger.classpath = "some.path.Trigger"
+ trigger.encrypted_kwargs = "encrypted"
+ trigger.task_instance = None # Not tied to a Task (similar to a BaseEventTrigger)
+ trigger.assets = [asset1, asset2]
+
+ workload = jobless_supervisor._create_workload(
+ trigger=trigger,
+ dag_bag=mocker.Mock(),
+ render_log_fname=mocker.Mock(),
+ session=mocker.Mock(),
+ )
+
+ assert workload is not None
+ assert workload.watched_assets == {"my_asset": "s3://bucket/key", "other_asset": "gs://bucket/path"}
+
+
+def test_create_workload_watched_assets_none_when_no_assets(jobless_supervisor, mocker):
+ """_create_workload() should set watched_assets=None when trigger.task_instance is None and assets is empty."""
+ trigger = mocker.Mock()
+ trigger.id = 43
+ trigger.classpath = "some.path.Trigger"
+ trigger.encrypted_kwargs = "encrypted"
+ trigger.task_instance = None
+ trigger.assets = [] # No Assets are attached to the trigger
+
+ workload = jobless_supervisor._create_workload(
+ trigger=trigger,
+ dag_bag=mocker.Mock(),
+ render_log_fname=mocker.Mock(),
+ session=mocker.Mock(),
+ )
+
+ assert workload is not None
+ assert workload.watched_assets is None
+
+
+def test_run_trigger_workload_includes_watched_assets_field():
+ """RunTrigger workload should accept and store watched_assets."""
+ from airflow.executors.workloads.trigger import RunTrigger
+
+ workload = RunTrigger(
+ id=1,
+ classpath="airflow.triggers.testing.SuccessTrigger",
+ encrypted_kwargs="fake",
+ watched_assets={"asset_a": "s3://a", "asset_b": "gs://b"},
+ )
+ assert workload.watched_assets == {"asset_a": "s3://a", "asset_b": "gs://b"}
+
+
+def test_run_trigger_workload_watched_assets_defaults_to_none():
+ """RunTrigger workload watched_assets should default to None."""
+ from airflow.executors.workloads.trigger import RunTrigger
+
+ workload = RunTrigger(
+ id=1,
+ classpath="airflow.triggers.testing.SuccessTrigger",
+ encrypted_kwargs="fake",
+ )
+ assert workload.watched_assets is None
+
+
+@pytest.mark.asyncio
+@patch("airflow.jobs.triggerer_job_runner.Trigger._decrypt_kwargs", return_value={})
+@patch("airflow.jobs.triggerer_job_runner.TriggerRunner.get_trigger_by_classpath")
+async def test_create_triggers_injects_asset_store_for_base_event_trigger(
+ mock_get_classpath, mock_decrypt, session
+):
+ """asset_store is populated on BaseEventTrigger instances when watched_assets is set."""
+ from airflow.sdk.execution_time.context import AssetStoreAccessors
+ from airflow.triggers.base import BaseEventTrigger, TriggerEvent
+
+ injected_instances = []
+
+ class _WatcherTrigger(BaseEventTrigger):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+ injected_instances.append(self)
+
+ def serialize(self):
+ return (f"{type(self).__module__}.{type(self).__qualname__}", {})
+
+ async def run(self):
+ yield TriggerEvent("done")
+
+ mock_get_classpath.return_value = _WatcherTrigger
+
+ runner = TriggerRunner()
+ runner.to_create.append(
+ workloads.RunTrigger.model_construct(
+ id=10,
+ ti=None,
+ classpath="fake.WatcherTrigger",
+ encrypted_kwargs="fake",
+ watched_assets={"my_asset": "s3://bucket/key"},
+ )
+ )
+
+ await runner.create_triggers()
+
+ # This is only testing that an exception was NOT thrown when creating the Trigger
+ assert 10 in runner.triggers
+
+ assert len(injected_instances) == 1
+ assert injected_instances[0].asset_store is not None
+ assert isinstance(injected_instances[0].asset_store, AssetStoreAccessors)
+
+ runner.triggers[10]["task"].cancel()
+ await runner.cleanup_finished_triggers()
+
+
+@pytest.mark.asyncio
+@patch("airflow.jobs.triggerer_job_runner.Trigger._decrypt_kwargs", return_value={})
+@patch("airflow.jobs.triggerer_job_runner.TriggerRunner.get_trigger_by_classpath")
+async def test_create_triggers_asset_store_none_when_no_watched_assets(
+ mock_get_classpath, mock_decrypt, session
+):
+ """asset_store stays None when watched_assets is not set on the workload."""
+ from airflow.triggers.base import BaseEventTrigger, TriggerEvent
Review Comment:
Move imports to the top pls
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Passing `AssetStateStoreAccessors` through to `BaseEventTrigger` [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on PR #67839:
URL: https://github.com/apache/airflow/pull/67839#issuecomment-4738471869
Right, i am looking at this one now.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Passing `AssetStateStoreAccessors` through to `BaseEventTrigger` [airflow]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #67839:
URL: https://github.com/apache/airflow/pull/67839#issuecomment-4746032879
<!-- uv-lock-rebase-notice -->
`uv.lock` on `main` just moved via [#68710](https://github.com/apache/airflow/pull/68710) ("Add Ray constraints for Python 3.14 in Google provider"), commit [`9c49080`](https://github.com/apache/airflow/commit/9c4908019a3f8d386eaed56c0552adfde241e952) and this PR currently conflicts.
Quickest fix:
```bash
git fetch upstream main && git rebase upstream/main
rm uv.lock && uv lock
git add uv.lock && git rebase --continue
git push --force-with-lease
```
_Automated nudge — ignore if you're not ready to rebase. This comment is updated in place on future `uv.lock` bumps._
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Passing `AssetStateStoreAccessors` through to `BaseEventTrigger` [airflow]
Posted by "jroachgolf84 (via GitHub)" <gi...@apache.org>.
jroachgolf84 commented on code in PR #67839:
URL: https://github.com/apache/airflow/pull/67839#discussion_r3439952528
##########
airflow-core/tests/unit/jobs/test_triggerer_job.py:
##########
@@ -561,6 +571,255 @@ def test_create_workload_uses_supervisor_id_without_job(jobless_supervisor, mock
assert factory.log_path == f"/logs/ti.trigger.{jobless_supervisor.id}.log"
+def test_create_workload_sets_watched_assets_for_asset_only_trigger(jobless_supervisor, mocker):
+ """_create_workload() should populate watched_assets when trigger.task_instance is None and assets exist."""
+ asset1 = mocker.Mock(spec=Asset)
+ asset1.name = "my_asset"
+ asset1.uri = "s3://bucket/key"
+
+ asset2 = mocker.Mock(spec=Asset)
+ asset2.name = "other_asset"
+ asset2.uri = "gs://bucket/path"
+
+ trigger = mocker.Mock(spec=BaseEventTrigger)
+ trigger.id = 42
+ trigger.classpath = "some.path.Trigger"
+ trigger.encrypted_kwargs = "encrypted"
+ trigger.task_instance = None # Not tied to a Task (similar to a BaseEventTrigger)
+ trigger.assets = [asset1, asset2]
+
+ workload = jobless_supervisor._create_workload(
+ trigger=trigger,
+ dag_bag=mocker.Mock(),
+ render_log_fname=mocker.Mock(),
+ session=mocker.Mock(),
+ )
+
+ assert workload is not None
+ assert workload.watched_assets == {"my_asset": "s3://bucket/key", "other_asset": "gs://bucket/path"}
+
+
+def test_create_workload_watched_assets_none_when_no_assets(jobless_supervisor, mocker):
+ """_create_workload() should set watched_assets=None when trigger.task_instance is None and assets is empty."""
+ trigger = mocker.Mock(spec=BaseEventTrigger)
+ trigger.id = 43
+ trigger.classpath = "some.path.Trigger"
+ trigger.encrypted_kwargs = "encrypted"
+ trigger.task_instance = None
+ trigger.assets = [] # No Assets are attached to the trigger
+
+ workload = jobless_supervisor._create_workload(
+ trigger=trigger,
+ dag_bag=mocker.Mock(),
+ render_log_fname=mocker.Mock(),
+ session=mocker.Mock(),
+ )
+
+ assert workload is not None
+ assert workload.watched_assets is None
+
+
+def test_run_trigger_workload_includes_watched_assets_field():
+ """RunTrigger workload should accept and store watched_assets."""
+ workload = RunTrigger(
+ id=1,
+ classpath="airflow.triggers.testing.SuccessTrigger",
+ encrypted_kwargs="fake",
+ watched_assets={"asset_a": "s3://a", "asset_b": "gs://b"},
+ )
+ assert workload.watched_assets == {"asset_a": "s3://a", "asset_b": "gs://b"}
+
+
+def test_run_trigger_workload_watched_assets_defaults_to_none():
+ """RunTrigger workload watched_assets should default to None."""
+ workload = RunTrigger(
+ id=1,
+ classpath="airflow.triggers.testing.SuccessTrigger",
+ encrypted_kwargs="fake",
+ )
+ assert workload.watched_assets is None
+
+
+@pytest.fixture
+def make_watcher_trigger():
+ """Factory fixture: call with a list to get a BaseEventTrigger subclass that appends each new instance."""
+
+ def factory(injected_instances):
+ class WatcherTrigger(BaseEventTrigger):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+ injected_instances.append(self)
+
+ def serialize(self):
+ return (f"{type(self).__module__}.{type(self).__qualname__}", {})
+
+ async def run(self):
+ yield TriggerEvent("done")
+
+ return WatcherTrigger
+
+ return factory
+
+
+@pytest.mark.asyncio
+@patch("airflow.jobs.triggerer_job_runner.TriggerRunner.get_trigger_by_classpath")
+async def test_create_triggers_injects_asset_state_store_for_base_event_trigger(
+ mock_get_classpath, session, make_watcher_trigger
+):
+ """asset_state_store is populated on BaseEventTrigger instances when watched_assets is set."""
+ injected_instances = []
+ mock_get_classpath.return_value = make_watcher_trigger(injected_instances)
+
+ runner = TriggerRunner()
+ runner.to_create.append(
+ workloads.RunTrigger.model_construct(
+ id=10,
+ ti=None,
+ classpath="fake.WatcherTrigger",
+ encrypted_kwargs="{}",
+ watched_assets={"my_asset": "s3://bucket/key"},
+ )
+ )
+
+ await runner.create_triggers()
+
+ # This is only testing that an exception was NOT thrown when creating the Trigger
+ assert 10 in runner.triggers
+
+ assert len(injected_instances) == 1
+ assert injected_instances[0].asset_state_store is not None
+ assert isinstance(injected_instances[0].asset_state_store, AssetStateStoreAccessors)
+
+ runner.triggers[10]["task"].cancel()
+ await runner.cleanup_finished_triggers()
+
+
+@pytest.mark.asyncio
+@patch("airflow.jobs.triggerer_job_runner.TriggerRunner.get_trigger_by_classpath")
+async def test_create_triggers_asset_state_store_none_when_no_watched_assets(
+ mock_get_classpath, session, make_watcher_trigger
+):
+ """asset_state_store stays None when watched_assets is not set on the workload."""
+ injected_instances = []
+ mock_get_classpath.return_value = make_watcher_trigger(injected_instances)
+
+ runner = TriggerRunner()
+ runner.to_create.append(
+ workloads.RunTrigger.model_construct(
+ id=11,
+ ti=None,
+ classpath="fake.WatcherTrigger",
+ encrypted_kwargs="{}",
+ watched_assets=None,
+ )
+ )
+
+ await runner.create_triggers()
+
+ assert len(injected_instances) == 1
+ assert injected_instances[0].asset_state_store is None
+
+ runner.triggers[11]["task"].cancel()
+ await runner.cleanup_finished_triggers()
+
+
+@pytest.mark.asyncio
+@patch("airflow.jobs.triggerer_job_runner.TriggerRunner.get_trigger_by_classpath")
+async def test_create_triggers_skips_asset_state_store_for_non_event_trigger(mock_get_classpath, session):
+ """asset_state_store injection is skipped for plain BaseTrigger (non-BaseEventTrigger) instances."""
+ mock_get_classpath.return_value = SuccessTrigger
+
+ runner = TriggerRunner()
+ runner.to_create.append(
+ workloads.RunTrigger.model_construct(
+ id=12, ti=None, classpath="airflow.triggers.testing.SuccessTrigger", encrypted_kwargs="{}"
+ )
+ )
+
+ await runner.create_triggers()
+
+ assert 12 in runner.triggers
+ assert not hasattr(runner.triggers[12]["task"], "asset_state_store")
Review Comment:
Resolved in next commit.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org