SameerMesiah97 commented on code in PR #74852:
URL: https://github.com/apache/airflow/pull/74852#discussion_r4405991801
##########
airflow-core/tests/unit/jobs/test_scheduler_job.py:
##########
@@ -13988,6 +13991,74 @@ def _produce_and_register_asset_event(
return apdr
[email protected]_serialized_dag
[email protected]("clear_asset_partition_rows")
[email protected]("temporal", [False, True])
[email protected]("event_kind", ["plain", "asset", "partitioned"])
+def test_watcher_event_creates_partitioned_consumer_run(dag_maker, session,
temporal, event_kind):
+ asset = Asset(name="watcher-partitioned-asset")
+ dag_id = "watcher-partitioned-consumer"
+ with dag_maker(
+ dag_id=dag_id,
+ schedule=PartitionedAssetTimetable(
+ assets=asset,
+ default_partition_mapper=StartOfDayMapper() if temporal else
IdentityMapper(),
+ ),
+ session=session,
+ ):
+ EmptyOperator(task_id="consume")
+ ordinary_dag_id = "watcher-ordinary-consumer"
+ with dag_maker(dag_id=ordinary_dag_id, schedule=asset, session=session):
+ EmptyOperator(task_id="consume")
+ session.commit()
+
+ trigger = Trigger(classpath="airflow.triggers.testing.SuccessTrigger",
kwargs={})
+ session.add(trigger)
+ session.flush()
+ asset_model = session.scalar(select(AssetModel).where(AssetModel.name ==
asset.name))
+ asset_model.add_trigger(trigger, "partition-watcher")
+ session.flush()
+ key = "2026-10-16T12:00:00"
+ payload = {"type": "AssetTriggerEvent", "partition_key": key}
+ event: TriggerEvent
+ if event_kind == "plain":
+ event = TriggerEvent(payload)
+ else:
+ event = AssetTriggerEvent(payload, partition_key=key if event_kind ==
"partitioned" else None)
+ adapter = TypeAdapter(ToTriggerSupervisor)
+ changes = messages.TriggerStateChanges(events=[(trigger.id, event, 7)])
+ decoded = adapter.validate_json(adapter.dump_json(changes))
+ Trigger.submit_event(trigger.id, decoded.events[0].event, session=session)
+ session.flush()
+
+ asset_event = session.scalar(select(AssetEvent).where(AssetEvent.asset_id
== asset_model.id))
+ partitioned = event_kind == "partitioned"
+ assert asset_event.partition_key == (key if partitioned else None)
+ assert asset_event.extra == {"from_trigger": True, "payload": payload}
+ ordinary_candidate = session.scalar(
+ select(AssetDagRunQueue).where(AssetDagRunQueue.target_dag_id ==
ordinary_dag_id)
+ )
+ assert (ordinary_candidate is not None) is not partitioned
+ runner = SchedulerJobRunner(
+ job=Job(job_type=SchedulerJobRunner.job_type),
executors=[MockExecutor(do_update=False)]
+ )
+ assert runner._create_dagruns_for_partitioned_asset_dags(session=session)
== (
+ {dag_id} if partitioned else set()
+ )
+ runs = session.scalars(select(DagRun).where(DagRun.dag_id == dag_id)).all()
+ assert len(runs) == int(partitioned)
+ if partitioned:
+ [run] = runs
+ assert run.partition_key == ("2026-10-16" if temporal else key)
+ assert run.logical_date is None
+ assert run.state == DagRunState.QUEUED
+ assert [event.id for event in run.consumed_asset_events] ==
[asset_event.id]
+ if temporal:
+ assert run.partition_date == pendulum.datetime(2026, 10, 16,
tz="UTC")
+ else:
+ assert run.partition_date is None
Review Comment:
This test could be a bit clearer:
- Parameterize the event and expected outcomes directly, instead of deriving
them from `event_kind`.
- Parameterize the mapper, expected partition key and date instead of using
`temporal`.
- Use explicit expectations instead of `is not partitioned` and
`int(partitioned)`.
- Remove the serialization round trip here since the dedicated test already
covers it.
- Add a comment explaining why the payload includes `type` and
`partition_key`.
##########
airflow-core/src/airflow/triggers/base.py:
##########
@@ -469,6 +470,29 @@ def __repr__(self) -> str:
return f"TriggerEvent<{self.payload!r}>"
+class AssetTriggerEvent(TriggerEvent):
+ """
+ An asset update emitted by a watcher, optionally identifying a native
partition.
+
+ The partition key applies to every asset associated with the trigger. The
payload
+ remains available in each asset event's extras and to any deferred task
resumed
+ by the same trigger.
Review Comment:
Could we add a test covering a deferred task and an asset watcher sharing
the same trigger? Submitting an `AssetTriggerEvent` should resume the task with
the original payload and persist the native partition key on the asset event.
This would cover the shared-trigger behavior documented here.
This is an optional improvement. Implement at your discreation.
##########
airflow-core/tests/unit/models/test_trigger.py:
##########
@@ -440,6 +441,26 @@ def test_submit_event(mock_callback_handle_event, session,
create_task_instance)
mock_callback_handle_event.assert_called_once_with(event, session)
[email protected]("partition_key", [None, "2026-10-16", " region-west
", "x" * 250])
+def test_submit_event_preserves_native_asset_partition(session, partition_key):
+ trigger = Trigger(classpath="airflow.triggers.testing.SuccessTrigger",
kwargs={})
+ session.add(trigger)
+ session.flush()
+ asset = AssetModel("partitioned")
+ asset.add_trigger(trigger, "partition_watcher")
Review Comment:
Could we extend this test to associate two assets with the same trigger and
assert that both asset events retain the native partition key and original
payload? The new class docstring explicitly promises that the key applies to
every asset associated with the trigger.
--
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]
