This is an automated email from the ASF dual-hosted git repository.
hussein-awala pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 1b792598b6d Support native asset partitions in watcher events (#74852)
1b792598b6d is described below
commit 1b792598b6d0baaf70985a062bddf3f8e3b058f5
Author: antonlin1 <[email protected]>
AuthorDate: Wed Nov 4 15:46:59 2026 -0500
Support native asset partitions in watcher events (#74852)
* Support native asset partitions in watcher events
Partition-aware Dags cannot consume native partition scopes from watcher
events when the key is retained only as payload metadata.
* Fix the partitioned asset event documentation heading
* Trim watcher partition publication documentation
* Scope watcher partition metadata to asset trigger events
Watchers need to identify native partitions without adding asset-specific
fields to task resumption and task-end events.
* Clarify why watcher events explicitly set their type
* Keep watcher partition guidance alongside trigger authoring
* Verify watcher persistence before shared-stream acknowledgement
* Keep partitioned asset event guidance concise
* Strengthen partitioned watcher event regression tests
* Test shared task and asset watcher events explicitly
* Fix partitioned watcher event documentation heading
---
.../docs/authoring-and-scheduling/assets.rst | 2 +
.../authoring-and-scheduling/event-scheduling.rst | 13 ++
airflow-core/src/airflow/models/trigger.py | 3 +-
airflow-core/src/airflow/triggers/base.py | 31 +++-
.../versions/head/test_asset_events.py | 49 +++++-
airflow-core/tests/unit/jobs/test_scheduler_job.py | 107 ++++++++++++
airflow-core/tests/unit/jobs/test_triggerer_job.py | 182 +++++++++++++++------
airflow-core/tests/unit/models/test_trigger.py | 53 ++++++
.../tests/unit/triggers/test_base_trigger.py | 15 +-
9 files changed, 398 insertions(+), 57 deletions(-)
diff --git a/airflow-core/docs/authoring-and-scheduling/assets.rst
b/airflow-core/docs/authoring-and-scheduling/assets.rst
index 4082e41ad70..f2a5d147bdf 100644
--- a/airflow-core/docs/authoring-and-scheduling/assets.rst
+++ b/airflow-core/docs/authoring-and-scheduling/assets.rst
@@ -538,6 +538,8 @@ depend on whether the producer and consumer have a team
association:
When Multi-Team mode is disabled, ``access_control`` is ignored and all asset
events are delivered to all
consuming Dags, preserving backward compatibility.
+.. _asset-partitions:
+
Asset partitions
----------------
diff --git a/airflow-core/docs/authoring-and-scheduling/event-scheduling.rst
b/airflow-core/docs/authoring-and-scheduling/event-scheduling.rst
index ba9967ad1e4..4e3672c0e8d 100644
--- a/airflow-core/docs/authoring-and-scheduling/event-scheduling.rst
+++ b/airflow-core/docs/authoring-and-scheduling/event-scheduling.rst
@@ -64,6 +64,19 @@ event-driven scheduling, then a new trigger must be created.
This new trigger must inherit ``BaseEventTrigger`` and ensure it properly
works with event-driven scheduling.
It might inherit from the existing trigger as well if both triggers share some
common code.
+Publishing partitioned asset events
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+To publish an update for a specific asset partition, yield
``AssetTriggerEvent`` with a ``partition_key``::
+
+ from airflow.triggers.base import AssetTriggerEvent
+
+ yield AssetTriggerEvent({"changed": True}, partition_key="2026-10-16")
+
+Airflow uses this key for partition-aware scheduling.
+
+See :ref:`asset-partitions` for configuring downstream consumers.
+
Sharing one poll across sibling triggers
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
diff --git a/airflow-core/src/airflow/models/trigger.py
b/airflow-core/src/airflow/models/trigger.py
index 43ad7b815b3..de6467a5851 100644
--- a/airflow-core/src/airflow/models/trigger.py
+++ b/airflow-core/src/airflow/models/trigger.py
@@ -50,7 +50,7 @@ from airflow.models.asset import AssetWatcherModel
from airflow.models.base import Base, StringID
from airflow.models.taskinstance import TaskInstance
from airflow.serialization.enums import stringify_encoding_keys
-from airflow.triggers.base import BaseTaskEndEvent
+from airflow.triggers.base import AssetTriggerEvent, BaseTaskEndEvent
from airflow.utils.session import NEW_SESSION, provide_session
from airflow.utils.sqlalchemy import UtcDateTime, with_row_locks
from airflow.utils.state import TaskInstanceState
@@ -399,6 +399,7 @@ class Trigger(Base):
AssetManager.register_asset_change(
asset=asset.to_serialized(),
extra={"from_trigger": True, "payload": event.payload},
+ partition_key=event.partition_key if isinstance(event,
AssetTriggerEvent) else None,
session=session,
)
if trigger.callback:
diff --git a/airflow-core/src/airflow/triggers/base.py
b/airflow-core/src/airflow/triggers/base.py
index 01b0e04522f..9e7615c2b98 100644
--- a/airflow-core/src/airflow/triggers/base.py
+++ b/airflow-core/src/airflow/triggers/base.py
@@ -21,13 +21,14 @@ import json
from collections.abc import AsyncIterator, Hashable
from dataclasses import dataclass
from datetime import timedelta
-from typing import TYPE_CHECKING, Annotated, Any
+from typing import TYPE_CHECKING, Annotated, Any, Literal
import structlog
from pydantic import (
BaseModel,
Discriminator,
JsonValue,
+ StringConstraints,
Tag,
model_serializer,
)
@@ -469,6 +470,29 @@ class TriggerEvent(BaseModel):
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.
+ """
+
+ type: Literal["AssetTriggerEvent"] = "AssetTriggerEvent"
+ partition_key: Annotated[str, StringConstraints(pattern=r"\S",
max_length=250)] | None = None
+
+ def __init__(self, payload, *, partition_key: str | None = None, **kwargs):
+ # Explicitly set type so serialization includes it. The triggerer
supervisor
+ # needs it to reconstruct AssetTriggerEvent and preserve the partition
key.
+ super().__init__(
+ payload=payload,
+ type=kwargs.pop("type", "AssetTriggerEvent"),
+ partition_key=partition_key,
+ **kwargs,
+ )
+
+
class BaseTaskEndEvent(TriggerEvent):
"""
Base event class to end the task without resuming on worker.
@@ -521,13 +545,18 @@ class TaskSkippedEvent(BaseTaskEndEvent):
def trigger_event_discriminator(v):
if isinstance(v, dict):
+ if v.get("type") == "AssetTriggerEvent":
+ return "_asset_"
return v.get("task_instance_state", "_event_")
+ if isinstance(v, AssetTriggerEvent):
+ return "_asset_"
if isinstance(v, TriggerEvent):
return getattr(v, "task_instance_state", "_event_")
DiscrimatedTriggerEvent = Annotated[
Annotated[TriggerEvent, Tag("_event_")]
+ | Annotated[AssetTriggerEvent, Tag("_asset_")]
| Annotated[TaskSuccessEvent, Tag(TaskInstanceState.SUCCESS)]
| Annotated[TaskFailedEvent, Tag(TaskInstanceState.FAILED)]
| Annotated[TaskSkippedEvent, Tag(TaskInstanceState.SKIPPED)],
diff --git
a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_asset_events.py
b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_asset_events.py
index 8cf59a440c5..76dd42b8866 100644
---
a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_asset_events.py
+++
b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_asset_events.py
@@ -19,11 +19,16 @@ from __future__ import annotations
from datetime import datetime
+import httpx2
import pytest
+from sqlalchemy import delete
from airflow._shared.timezones import timezone
-from airflow.models.asset import AssetActive, AssetAliasModel, AssetEvent,
AssetModel
+from airflow.models.asset import AssetActive, AssetAliasModel, AssetEvent,
AssetModel, AssetWatcherModel
from airflow.models.dagrun import DagRun
+from airflow.models.trigger import Trigger
+from airflow.sdk.api.client import Client as ExecutionAPIClient
+from airflow.triggers.base import AssetTriggerEvent
from airflow.utils.state import DagRunState
from airflow.utils.types import DagRunType
@@ -96,6 +101,48 @@ def test_asset_alias(session, test_asset_events,
test_asset):
class TestGetAssetEventByAsset:
+ @pytest.mark.parametrize("partition_key", [None, "2026-10-16"])
+ def test_get_watcher_event_through_task_sdk(self, client, session,
test_asset, partition_key, request):
+ trigger = Trigger(classpath="airflow.triggers.testing.SuccessTrigger",
kwargs={})
+ session.add(trigger)
+ session.flush()
+
+ def cleanup():
+ session.execute(delete(AssetEvent).where(AssetEvent.asset_id ==
test_asset.id))
+
session.execute(delete(AssetWatcherModel).where(AssetWatcherModel.trigger_id ==
trigger.id))
+ session.execute(delete(Trigger).where(Trigger.id == trigger.id))
+ session.commit()
+
+ request.addfinalizer(cleanup)
+ test_asset.add_trigger(trigger, "partition_watcher")
+ session.flush()
+ payload = {"partition_key": "payload-only", "changed": True}
+ event = AssetTriggerEvent(payload, partition_key=partition_key)
+ Trigger.submit_event(trigger.id, event, session=session)
+ session.commit()
+ session.expire_all()
+
+ def read_event(request: httpx2.Request) -> httpx2.Response:
+ response = client.get(
+ request.url.path,
+ params=list(request.url.params.multi_items()),
+ headers=dict(request.headers),
+ )
+ return httpx2.Response(response.status_code,
headers=response.headers, content=response.content)
+
+ with ExecutionAPIClient(
+ base_url="http://testserver/execution/";,
+ token="fake",
+ transport=httpx2.MockTransport(read_event),
+ ) as sdk_client:
+ result = sdk_client.asset_events.get(name=test_asset.name)
+
+ assert len(result.asset_events) == 1
+ returned_event = result.asset_events[0]
+ assert returned_event.partition_key == partition_key
+ assert returned_event.extra == {"from_trigger": True, "payload":
payload}
+ assert isinstance(returned_event.timestamp, datetime)
+
@pytest.mark.usefixtures("test_asset")
def test_get_by_asset_with_created_dagrun_without_start_date(self, client,
session, test_asset_events):
created_dagrun = DagRun(
diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py
b/airflow-core/tests/unit/jobs/test_scheduler_job.py
index 0b6bd0fee54..c35b62287d6 100644
--- a/airflow-core/tests/unit/jobs/test_scheduler_job.py
+++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py
@@ -162,6 +162,7 @@ from airflow.timetables.simple import (
PartitionedAssetTimetable as CorePartitionedAssetTimetable,
PartitionedAtRuntime,
)
+from airflow.triggers.base import AssetTriggerEvent, TriggerEvent
from airflow.utils.session import NEW_SESSION, create_session, provide_session
from airflow.utils.sqlalchemy import CommitProhibitorGuard, with_row_locks
from airflow.utils.state import CallbackState, DagRunState,
DagSchedulingState, State, TaskInstanceState
@@ -14345,6 +14346,112 @@ def _produce_and_register_asset_event(
return apdr
[email protected]_serialized_dag
[email protected]("clear_asset_partition_rows")
[email protected](
+ ("mapper", "expected_partition_key", "expected_partition_date"),
+ [
+ pytest.param(IdentityMapper(), "2026-10-16T12:00:00", None,
id="identity"),
+ pytest.param(StartOfDayMapper(), "2026-10-16", pendulum.datetime(2026,
10, 16, tz="UTC"), id="daily"),
+ ],
+)
+# Payload fields resemble native event fields to prove metadata cannot select
the event type or partition.
[email protected](
+ (
+ "event",
+ "expected_native_key",
+ "expected_ordinary_candidate",
+ "expected_created_dag_ids",
+ "expected_run_count",
+ ),
+ [
+ pytest.param(
+ TriggerEvent({"type": "AssetTriggerEvent", "partition_key":
"payload-only"}),
+ None,
+ True,
+ set(),
+ 0,
+ id="ordinary",
+ ),
+ pytest.param(
+ AssetTriggerEvent({"type": "AssetTriggerEvent", "partition_key":
"payload-only"}),
+ None,
+ True,
+ set(),
+ 0,
+ id="asset-without-partition",
+ ),
+ pytest.param(
+ AssetTriggerEvent(
+ {"type": "AssetTriggerEvent", "partition_key": "payload-only"},
+ partition_key="2026-10-16T12:00:00",
+ ),
+ "2026-10-16T12:00:00",
+ False,
+ {"watcher-partitioned-consumer"},
+ 1,
+ id="partitioned",
+ ),
+ ],
+)
+def test_watcher_event_creates_partitioned_consumer_run(
+ dag_maker,
+ session,
+ mapper,
+ expected_partition_key,
+ expected_partition_date,
+ event,
+ expected_native_key,
+ expected_ordinary_candidate,
+ expected_created_dag_ids,
+ expected_run_count,
+):
+ 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=mapper,
+ ),
+ 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()
+ Trigger.submit_event(trigger.id, event, session=session)
+ session.flush()
+
+ asset_event = session.scalar(select(AssetEvent).where(AssetEvent.asset_id
== asset_model.id))
+ assert asset_event.partition_key == expected_native_key
+ assert asset_event.extra == {"from_trigger": True, "payload":
event.payload}
+ ordinary_candidate = session.scalar(
+ select(AssetDagRunQueue).where(AssetDagRunQueue.target_dag_id ==
ordinary_dag_id)
+ )
+ assert (ordinary_candidate is not None) is expected_ordinary_candidate
+ runner = SchedulerJobRunner(
+ job=Job(job_type=SchedulerJobRunner.job_type),
executors=[MockExecutor(do_update=False)]
+ )
+ assert runner._create_dagruns_for_partitioned_asset_dags(session=session)
== expected_created_dag_ids
+ runs = session.scalars(select(DagRun).where(DagRun.dag_id == dag_id)).all()
+ assert len(runs) == expected_run_count
+ for run in runs:
+ assert run.partition_key == expected_partition_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]
+ assert run.partition_date == expected_partition_date
+
+
@pytest.mark.need_serialized_dag
@pytest.mark.usefixtures("clear_asset_partition_rows")
@pytest.mark.parametrize("batch_asset_events", [True, False])
diff --git a/airflow-core/tests/unit/jobs/test_triggerer_job.py
b/airflow-core/tests/unit/jobs/test_triggerer_job.py
index a614b30cc78..e38fbe0343b 100644
--- a/airflow-core/tests/unit/jobs/test_triggerer_job.py
+++ b/airflow-core/tests/unit/jobs/test_triggerer_job.py
@@ -76,7 +76,7 @@ from airflow.jobs.triggerer_job_runner import (
messages,
)
from airflow.models import Connection, DagModel, DagRun, Trigger, Variable
-from airflow.models.asset import AssetModel
+from airflow.models.asset import AssetEvent, AssetModel
from airflow.models.dag_version import DagVersion
from airflow.models.dagbag import DBDagBag
from airflow.models.dagbundle import DagBundleModel
@@ -112,12 +112,21 @@ from airflow.sdk.execution_time.comms import (
SetAssetStateStoreByUri,
ToSupervisor,
ToTask,
+ _new_encoder,
_RequestFrame,
_ResponseFrame,
)
from airflow.sdk.execution_time.context import AssetStateStoreAccessors
from airflow.serialization.serialized_objects import LazyDeserializedDAG
-from airflow.triggers.base import BaseEventTrigger, BaseTrigger, TriggerEvent
+from airflow.triggers.base import (
+ AssetTriggerEvent,
+ BaseEventTrigger,
+ BaseTrigger,
+ TaskFailedEvent,
+ TaskSkippedEvent,
+ TaskSuccessEvent,
+ TriggerEvent,
+)
from airflow.triggers.shared_stream import SharedStreamProducer
from airflow.triggers.testing import FailureTrigger, SuccessTrigger
from airflow.utils.state import State, TaskInstanceState
@@ -125,6 +134,7 @@ from airflow.utils.types import DagRunType
from tests_common.test_utils.config import conf_vars
from tests_common.test_utils.db import (
+ clear_db_assets,
clear_db_connections,
clear_db_dag_bundles,
clear_db_dags,
@@ -1837,11 +1847,25 @@ class TestTriggerRunner:
# Group is torn down on unsubscribe.
assert trigger_runner._shared_streams._groups == {}
- def test_shared_stream_ack_mode_integration(self) -> None:
- """A BaseEventTrigger whose producer factory is overridden advances
only after the event's derived trigger event is confirmed persisted."""
- advanced: list[Any] = []
- # Container lets _drive() write trigger_runner back for post-run
assertions.
- state: dict[str, Any] = {}
+ @pytest.mark.parametrize("partition_key", [None, "2026-10-16"])
+ @pytest.mark.parametrize("commit_fails", [False, True])
+ def test_shared_stream_ack_mode_integration(
+ self, session, request, jobless_supervisor, partition_key: str | None,
commit_fails: bool
+ ) -> None:
+ """Broker advancement requires successful commit of the event's native
partition key."""
+ trigger_model =
Trigger(classpath="airflow.triggers.testing.SuccessTrigger", kwargs={})
+ session.add(trigger_model)
+ session.flush()
+ asset = AssetModel("shared-stream-event")
+ asset.add_trigger(trigger_model, "shared-stream-watcher")
+ session.add(asset)
+ session.commit()
+ trigger_id, asset_id = trigger_model.id, asset.id
+ request.addfinalizer(clear_db_assets)
+
+ advanced: list[str] = []
+ advance_called = asyncio.Event()
+ event_queued = asyncio.Event()
class _AckSharedProducer(SharedStreamProducer):
async def open_stream(self):
@@ -1851,7 +1875,7 @@ class TestTriggerRunner:
async def advance(self, batch):
advanced.extend(item.broker_payload for item in batch)
- state["advance_called"].set()
+ advance_called.set()
class _AckSharedTrigger(BaseEventTrigger):
def __init__(self, queue_url: str, region: str | None = None):
@@ -1875,62 +1899,83 @@ class TestTriggerRunner:
async def filter_shared_stream(self, shared_stream):
async for raw in shared_stream:
if self.region is None or raw["region"] == self.region:
- yield TriggerEvent(raw)
+ if partition_key is None:
+ yield TriggerEvent(raw)
+ else:
+ yield AssetTriggerEvent(raw,
partition_key=partition_key)
+ event_queued.set()
async def run(self): # pragma: no cover - replaced by
filter_shared_stream
yield TriggerEvent({})
async def _drive():
- advance_called = asyncio.Event()
- state["advance_called"] = advance_called
-
trigger_runner = TriggerRunner()
trigger_runner.triggers = {
- 1: {"task": MagicMock(spec=asyncio.Task), "is_watcher": True,
"name": "us", "events": 0}
+ trigger_id: {
+ "task": MagicMock(spec=asyncio.Task),
+ "is_watcher": True,
+ "name": "us",
+ "events": 0,
+ }
}
trigger = _AckSharedTrigger(queue_url="https://ack-q";, region="us")
trigger.task_instance = MagicMock()
trigger.task_instance.map_index = -1
- state["trigger_runner"] = trigger_runner
-
- run_task = asyncio.create_task(trigger_runner.run_trigger(1,
trigger))
- # Wait until the event lands on the outbound queue with its
- # persist-confirmation seq.
- for _ in range(100):
- if trigger_runner.events:
- break
- await asyncio.sleep(0.01)
- assert trigger_runner.events, "the trigger event must reach the
outbound queue"
- _trigger_id, _event, seq = trigger_runner.events[0]
- assert seq is not None, "a shared ack trigger's event must carry a
seq"
-
- # The subscriber moved past the event, but the advance is gated
- # on persistence: nothing may advance before the confirmation
arrives.
- assert advanced == [], "advance must not run before the persist
confirmation"
- # Simulate the supervisor confirming the persist on the next sync.
- trigger_runner._shared_streams.confirm_persisted([seq])
- await asyncio.wait_for(advance_called.wait(), timeout=2.0)
-
- run_task.cancel()
- with contextlib.suppress(asyncio.CancelledError):
- await run_task
+ async def exchange(msg: messages.TriggerStateChanges) ->
messages.TriggerStateSync:
+ wire = _new_encoder().encode(msg.model_dump())
+ decoded =
TypeAdapter(ToTriggerSupervisor).validate_python(msgspec.msgpack.decode(wire))
+ with mock.patch.object(TriggerRunnerSupervisor, "send_msg",
autospec=True) as send_msg:
+ jobless_supervisor._handle_request(decoded,
log=jobless_supervisor.process_log, req_id=0)
+ response = send_msg.call_args.args[1]
+ wire = _new_encoder().encode(response.model_dump())
+ decoded_response =
TypeAdapter(ToTriggerRunner).validate_python(msgspec.msgpack.decode(wire))
+ assert isinstance(decoded_response, messages.TriggerStateSync)
+ return decoded_response
+
+ run_task =
asyncio.create_task(trigger_runner.run_trigger(trigger_id, trigger))
+ try:
+ await asyncio.wait_for(event_queued.wait(), timeout=2.0)
+ assert trigger_runner.events[0].persist_seq is not None
+ with mock.patch.object(trigger_runner, "asend", autospec=True,
side_effect=exchange):
+ await
trigger_runner.sync_state_to_supervisor(finished_ids=[])
+ assert advanced == []
+ if commit_fails:
+ with mock.patch(
+ "sqlalchemy.orm.Session.commit",
+ autospec=True,
+ side_effect=RuntimeError("commit failed"),
+ ):
+ with pytest.raises(RuntimeError, match="commit
failed"):
+ jobless_supervisor.handle_events()
+ else:
+ jobless_supervisor.handle_events()
+
+ asset_event =
session.scalar(select(AssetEvent).where(AssetEvent.asset_id == asset_id))
+ if commit_fails:
+ assert asset_event is None
+ assert list(jobless_supervisor.persisted_event_seqs)
== []
+ else:
+ assert asset_event is not None
+ assert asset_event.partition_key == partition_key
+ assert asset_event.extra == {"from_trigger": True,
"payload": {"region": "us"}}
+
+ assert advanced == []
+ await
trigger_runner.sync_state_to_supervisor(finished_ids=[])
+ if commit_fails:
+ assert advanced == []
+ assert not advance_called.is_set()
+ else:
+ await asyncio.wait_for(advance_called.wait(),
timeout=2.0)
+ finally:
+ run_task.cancel()
+ with contextlib.suppress(asyncio.CancelledError):
+ await run_task
+
+ assert trigger_runner._shared_streams._groups == {}
asyncio.run(_drive())
- trigger_runner = state["trigger_runner"]
-
- events = list(trigger_runner.events)
- assert len(events) == 1
- trigger_id, event, _seq = events[0]
- assert trigger_id == 1
- assert event.payload == {"region": "us"}
-
- # The producer's advance must have been called with the broker_payload.
- assert advanced == ["receipt-handle-1"], (
- "producer.advance must be called once the event's persistence is
confirmed"
- )
- # Group is torn down on unsubscribe.
- assert trigger_runner._shared_streams._groups == {}
+ assert advanced == ([] if commit_fails else ["receipt-handle-1"])
def test_run_trigger_on_kill_timeout_does_not_block_cleanup(self, session)
-> None:
"""A hanging on_kill() is interrupted after the timeout and cleanup
still runs."""
@@ -3719,10 +3764,41 @@ def test_trigger_state_messages_round_trip_with_seqs():
assert decoded_sync.to_cancel == {5}
-def test_handle_events_records_persist_confirmations(jobless_supervisor):
[email protected]("persist_seq", [None, 7])
[email protected]("wire_format", ["json", "msgpack"])
[email protected](
+ "event",
+ [
+ AssetTriggerEvent({"changed": True}, partition_key="2026-10-16"),
+ AssetTriggerEvent(True),
+ TriggerEvent({"type": "AssetTriggerEvent", "partition_key":
"payload-only"}),
+ TaskSuccessEvent(),
+ TaskFailedEvent(),
+ TaskSkippedEvent(),
+ ],
+)
+def test_asset_trigger_event_state_message_round_trip(persist_seq,
wire_format, event):
+ adapter = TypeAdapter(ToTriggerSupervisor)
+ changes = messages.TriggerStateChanges(type="TriggerStateChanges",
events=[(1, event, persist_seq)])
+
+ if wire_format == "json":
+ decoded = adapter.validate_json(adapter.dump_json(changes))
+ else:
+ decoded =
adapter.validate_python(msgspec.msgpack.decode(_new_encoder().encode(changes)))
+
+ assert isinstance(decoded, messages.TriggerStateChanges)
+ assert decoded.events is not None
+ assert decoded.events[0].event == event
+ assert type(decoded.events[0].event) is type(event)
+ assert decoded.events[0].event.payload == event.payload
+ assert decoded.events[0].persist_seq == persist_seq
+
+
[email protected]("event_class", [TriggerEvent, AssetTriggerEvent])
+def test_handle_events_records_persist_confirmations(jobless_supervisor,
event_class):
"""handle_events confirms the seq of a persisted event; events without a
seq are not confirmed."""
- event_with_seq = TriggerEvent(True)
- event_without_seq = TriggerEvent("x")
+ event_with_seq = event_class(True)
+ event_without_seq = event_class("x")
jobless_supervisor.events.append(TriggerEventEntry(1, event_with_seq, 7))
jobless_supervisor.events.append(TriggerEventEntry(2, event_without_seq,
None))
diff --git a/airflow-core/tests/unit/models/test_trigger.py
b/airflow-core/tests/unit/models/test_trigger.py
index de89888b9d8..5963f18fa62 100644
--- a/airflow-core/tests/unit/models/test_trigger.py
+++ b/airflow-core/tests/unit/models/test_trigger.py
@@ -45,6 +45,7 @@ from airflow.sdk.definitions.callback import AsyncCallback
from airflow.serialization.encoders import encode_trigger
from airflow.serialization.serialized_objects import BaseSerialization
from airflow.triggers.base import (
+ AssetTriggerEvent,
BaseEventTrigger,
BaseTrigger,
TaskFailedEvent,
@@ -455,6 +456,32 @@ def test_submit_event(mock_callback_handle_event, session,
create_task_instance)
mock_callback_handle_event.assert_called_once_with(event, session)
+def test_submit_asset_event_to_shared_task_and_asset_watcher(session,
create_task_instance):
+ trigger = Trigger(classpath="airflow.triggers.testing.SuccessTrigger",
kwargs={})
+ session.add(trigger)
+ session.flush()
+ task_instance = create_task_instance(
+ session=session, logical_date=timezone.datetime(2026, 10, 16),
state=State.DEFERRED
+ )
+ task_instance.trigger_id = trigger.id
+ task_instance.next_kwargs = {"resume_arg": "kept"}
+ asset = AssetModel("shared-trigger-asset")
+ asset.add_trigger(trigger, "shared-trigger-watcher")
+ session.add(asset)
+ session.commit()
+ payload = {"partition_key": "payload-only", "changed": True}
+
+ Trigger.submit_event(trigger.id, AssetTriggerEvent(payload,
partition_key="2026-10-16"), session=session)
+ session.flush()
+ session.refresh(task_instance)
+
+ assert task_instance.state == State.SCHEDULED
+ assert task_instance.next_kwargs == {"event": payload, "resume_arg":
"kept"}
+ asset_event = session.scalar(select(AssetEvent).where(AssetEvent.asset_id
== asset.id))
+ assert asset_event.partition_key == "2026-10-16"
+ assert asset_event.extra == {"from_trigger": True, "payload": payload}
+
+
@pytest.mark.parametrize("owner", ["task", "asset_watcher", "callback",
"none"])
def test_submit_event_records_the_event_time_for_every_kind(session,
create_task_instance, owner):
t0 = timezone.datetime(2026, 10, 21, 12, 0, 0)
@@ -482,6 +509,32 @@ def
test_submit_event_records_the_event_time_for_every_kind(session, create_task
assert session.get(Trigger, trigger.id).last_event_at == t0
[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()
+ assets = [AssetModel("partitioned-first"),
AssetModel("partitioned-second")]
+ for asset in assets:
+ asset.add_trigger(trigger, f"partition_watcher_{asset.name}")
+ session.add_all(assets)
+ session.flush()
+ payload = {"partition_key": "payload-only", "changed": True}
+
+ Trigger.submit_event(trigger.id, AssetTriggerEvent(payload,
partition_key=partition_key), session=session)
+ session.flush()
+ session.expire_all()
+
+ asset_events = session.scalars(
+ select(AssetEvent).where(AssetEvent.asset_id.in_([a.id for a in
assets]))
+ ).all()
+ assert len(asset_events) == 2
+ assert {event.asset_id for event in asset_events} == {asset.id for asset
in assets}
+ for asset_event in asset_events:
+ assert asset_event.partition_key == partition_key
+ assert asset_event.extra == {"from_trigger": True, "payload": payload}
+
+
@pytest.mark.parametrize(("asset_count", "expected_query_count"), [(1, 6), (5,
6)])
@patch("airflow.models.trigger.AssetManager.register_asset_change")
def test_submit_event_no_n_plus_one_for_assets(_, session, asset_count,
expected_query_count):
diff --git a/airflow-core/tests/unit/triggers/test_base_trigger.py
b/airflow-core/tests/unit/triggers/test_base_trigger.py
index f65d56517b9..c30a21408d9 100644
--- a/airflow-core/tests/unit/triggers/test_base_trigger.py
+++ b/airflow-core/tests/unit/triggers/test_base_trigger.py
@@ -20,10 +20,23 @@ from __future__ import annotations
from unittest.mock import create_autospec
import pytest
+from pydantic import ValidationError
from airflow.sdk.bases.operator import BaseOperator
from airflow.sdk.execution_time.context import AssetStateStoreAccessors
-from airflow.triggers.base import BaseEventTrigger, BaseTrigger,
StartTriggerArgs, TriggerEvent
+from airflow.triggers.base import (
+ AssetTriggerEvent,
+ BaseEventTrigger,
+ BaseTrigger,
+ StartTriggerArgs,
+ TriggerEvent,
+)
+
+
[email protected]("partition_key", ["", " \t\n", "x" * 251, 42])
+def test_asset_trigger_event_rejects_invalid_partition_key(partition_key):
+ with pytest.raises(ValidationError, match="partition_key"):
+ AssetTriggerEvent({"changed": True}, partition_key=partition_key)
class DummyOperator(BaseOperator):
