This is an automated email from the ASF dual-hosted git repository.
ephraimbuddy 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 e3089bf977c Extract DB operations from TriggerRunnerSupervisor into
overridable methods (#65680)
e3089bf977c is described below
commit e3089bf977c1a936c1a54600dcecf1c7ea21d771
Author: Ephraim Anierobi <[email protected]>
AuthorDate: Mon Apr 27 20:56:52 2026 +0100
Extract DB operations from TriggerRunnerSupervisor into overridable methods
(#65680)
* Extract DB operations from TriggerRunnerSupervisor into overridable
methods
Add overridable hooks on TriggerRunnerSupervisor so the DB-backed
Trigger.* calls can be swapped for Execution API calls in a follow-up
for AIP-92 (isolate DAG processor, callback processor, and triggerer
from core services):
- on_trigger_event / on_trigger_failure wrap Trigger.submit_event /
Trigger.submit_failure.
- fetch_trigger_details / fetch_non_task_trigger_ids wrap the bulk
fetches used during update_triggers.
- build_trigger_workloads extracts the session-scoped workload
construction out of update_triggers, leaving update_triggers as a
pure diff/enqueue step. As a side effect, DBDagBag, the log
renderer, and the trigger fetches are now only created when there
are new triggers to process, which removes a small amount of idle
metadata-DB work on each loop iteration.
Also fix the return annotation on
Trigger.fetch_trigger_ids_with_non_task_associations from set[str]
to set[int] to match the actual PK type.
* Apply suggestion
Co-authored-by: Tzu-ping Chung <[email protected]>
---------
Co-authored-by: Tzu-ping Chung <[email protected]>
---
.../src/airflow/jobs/triggerer_job_runner.py | 78 +++++++++++++---------
airflow-core/src/airflow/models/trigger.py | 2 +-
airflow-core/tests/unit/jobs/test_triggerer_job.py | 57 ++++++++++++++++
3 files changed, 105 insertions(+), 32 deletions(-)
diff --git a/airflow-core/src/airflow/jobs/triggerer_job_runner.py
b/airflow-core/src/airflow/jobs/triggerer_job_runner.py
index 73c0ed906e8..2d0fc2ae224 100644
--- a/airflow-core/src/airflow/jobs/triggerer_job_runner.py
+++ b/airflow-core/src/airflow/jobs/triggerer_job_runner.py
@@ -616,8 +616,8 @@ class TriggerRunnerSupervisor(WatchedSubprocess):
def heartbeat_callback(self, session: Session | None = None) -> None:
Stats.incr("triggerer_heartbeat", 1, 1)
- def load_triggers(self):
- """Query the database for the triggers we're supposed to be running
and update the runner."""
+ def load_triggers(self) -> None:
+ """Assign triggers to this triggerer and update the runner with the
IDs it should run."""
Trigger.assign_unassigned(
self.job.id,
self.capacity,
@@ -633,12 +633,16 @@ class TriggerRunnerSupervisor(WatchedSubprocess):
# Get the event and its trigger ID
trigger_id, event = self.events.popleft()
# Tell the model to wake up its tasks
- Trigger.submit_event(trigger_id=trigger_id, event=event)
+ self.on_trigger_event(trigger_id=trigger_id, event=event)
# Emit stat event
Stats.incr("triggers.succeeded")
- def clean_unused(self):
- """Clean out unused or finished triggers."""
+ def on_trigger_event(self, trigger_id: int, event: TriggerEvent) -> None:
+ """Record that a trigger fired an event."""
+ Trigger.submit_event(trigger_id=trigger_id, event=event)
+
+ def clean_unused(self) -> None:
+ """Remove triggers that are no longer needed."""
Trigger.clean_unused()
def handle_failed_triggers(self):
@@ -649,11 +653,15 @@ class TriggerRunnerSupervisor(WatchedSubprocess):
"""
while self.failed_triggers:
# Tell the model to fail this trigger's deps
- trigger_id, saved_exc = self.failed_triggers.popleft()
- Trigger.submit_failure(trigger_id=trigger_id, exc=saved_exc)
+ trigger_id, exc = self.failed_triggers.popleft()
+ self.on_trigger_failure(trigger_id=trigger_id, exc=exc)
# Emit stat event
Stats.incr("triggers.failed")
+ def on_trigger_failure(self, trigger_id: int, exc: list[str] | None) ->
None:
+ """Record that a trigger failed."""
+ Trigger.submit_failure(trigger_id=trigger_id, exc=exc)
+
def emit_metrics(self):
DualStatsManager.gauge(
"triggers.running",
@@ -733,34 +741,22 @@ class TriggerRunnerSupervisor(WatchedSubprocess):
timeout_after=trigger.task_instance.trigger_timeout,
)
- def update_triggers(self, requested_trigger_ids: set[int]):
- """
- Request that we update what triggers we're running.
+ def fetch_trigger_details(self, trigger_ids: set[int], *, session:
Session) -> dict[int, Trigger]:
+ """Fetch trigger rows by ID."""
+ return Trigger.bulk_fetch(trigger_ids, session=session)
- Works out the differences - ones to add, and ones to remove - then
- adds them to the dequeues so the subprocess can actually mutate the
running
- trigger set.
- """
+ def fetch_non_task_trigger_ids(self, *, session: Session) -> set[int]:
+ """Fetch trigger IDs associated with non-task entities."""
+ return
Trigger.fetch_trigger_ids_with_non_task_associations(session=session)
+
+ def build_trigger_workloads(self, new_trigger_ids: set[int]) ->
list[workloads.RunTrigger]:
+ """Build workloads for new trigger IDs."""
dag_bag = DBDagBag()
render_log_fname = log_filename_template_renderer()
- known_trigger_ids = (
- self.running_triggers.union(x[0] for x in self.events)
- .union(self.cancelling_triggers)
- .union(trigger[0] for trigger in self.failed_triggers)
- .union(trigger.id for trigger in self.creating_triggers)
- )
- # Work out the two difference sets
- new_trigger_ids = requested_trigger_ids - known_trigger_ids
- cancel_trigger_ids = self.running_triggers - requested_trigger_ids
- # Bulk-fetch new trigger records
with create_session() as session:
- # Bulk-fetch new trigger records
- new_triggers = Trigger.bulk_fetch(new_trigger_ids, session=session)
- trigger_ids_with_non_task_associations =
Trigger.fetch_trigger_ids_with_non_task_associations(
- session=session
- )
+ new_triggers = self.fetch_trigger_details(new_trigger_ids,
session=session)
+ trigger_ids_with_non_task_associations =
self.fetch_non_task_trigger_ids(session=session)
to_create: list[workloads.RunTrigger] = []
- # Add in new triggers
for new_trigger_id in new_trigger_ids:
# Check it didn't vanish in the meantime
if new_trigger_id not in new_triggers:
@@ -794,7 +790,27 @@ class TriggerRunnerSupervisor(WatchedSubprocess):
):
to_create.append(workload)
- self.creating_triggers.extend(to_create)
+ return to_create
+
+ def update_triggers(self, requested_trigger_ids: set[int]):
+ """
+ Request that we update what triggers we're running.
+
+ Works out the differences - ones to add, and ones to remove - then
+ adds them to the dequeues so the subprocess can actually mutate the
running
+ trigger set.
+ """
+ known_trigger_ids = self.running_triggers.union(
+ (x[0] for x in self.events),
+ self.cancelling_triggers,
+ (trigger[0] for trigger in self.failed_triggers),
+ (trigger.id for trigger in self.creating_triggers),
+ )
+ # Work out the two difference sets
+ new_trigger_ids = requested_trigger_ids - known_trigger_ids
+ cancel_trigger_ids = self.running_triggers - requested_trigger_ids
+ if new_trigger_ids:
+
self.creating_triggers.extend(self.build_trigger_workloads(new_trigger_ids))
if cancel_trigger_ids:
# Enqueue orphaned triggers for cancellation
diff --git a/airflow-core/src/airflow/models/trigger.py
b/airflow-core/src/airflow/models/trigger.py
index 97079952885..589725ced2c 100644
--- a/airflow-core/src/airflow/models/trigger.py
+++ b/airflow-core/src/airflow/models/trigger.py
@@ -207,7 +207,7 @@ class Trigger(Base):
@classmethod
@provide_session
- def fetch_trigger_ids_with_non_task_associations(cls, session: Session =
NEW_SESSION) -> set[str]:
+ def fetch_trigger_ids_with_non_task_associations(cls, session: Session =
NEW_SESSION) -> set[int]:
"""Fetch all trigger IDs actively associated with non-task entities
like assets and callbacks."""
from airflow.models.callback import Callback # to avoid circular
import: Callback -> Trigger
diff --git a/airflow-core/tests/unit/jobs/test_triggerer_job.py
b/airflow-core/tests/unit/jobs/test_triggerer_job.py
index 4b047ab15e1..ef5b5a02c5c 100644
--- a/airflow-core/tests/unit/jobs/test_triggerer_job.py
+++ b/airflow-core/tests/unit/jobs/test_triggerer_job.py
@@ -1308,6 +1308,63 @@ def
test_update_triggers_prevents_duplicate_creation_queue_entries(session, supe
assert not any(trigger_id == trigger_orm.id for trigger_id, _ in
supervisor.failed_triggers)
+def test_update_triggers_delegates_workload_creation(supervisor_builder,
mocker):
+ supervisor = supervisor_builder()
+ supervisor.running_triggers = {1, 3}
+ workload = workloads.RunTrigger(id=2, classpath="some.trigger",
encrypted_kwargs="", ti=None)
+ build_trigger_workloads = mocker.patch.object(
+ TriggerRunnerSupervisor, "build_trigger_workloads", autospec=True,
return_value=[workload]
+ )
+
+ supervisor.update_triggers({2, 3})
+
+ build_trigger_workloads.assert_called_once_with(supervisor, {2})
+ assert list(supervisor.creating_triggers) == [workload]
+ assert supervisor.cancelling_triggers == {1}
+
+
+def test_update_triggers_uses_fetch_hooks(session, supervisor_builder, mocker):
+ trigger = TimeDeltaTrigger(datetime.timedelta(days=7))
+ _, _, trigger_orm, _ = create_trigger_in_db(session, trigger)
+
+ supervisor = supervisor_builder()
+ fetch_trigger_details_calls = []
+ fetch_non_task_trigger_ids_calls = []
+ original_fetch_trigger_details =
TriggerRunnerSupervisor.fetch_trigger_details
+ original_fetch_non_task_trigger_ids =
TriggerRunnerSupervisor.fetch_non_task_trigger_ids
+
+ def fetch_trigger_details(self, trigger_ids, *, session):
+ fetch_trigger_details_calls.append((trigger_ids, session))
+ return original_fetch_trigger_details(self, trigger_ids,
session=session)
+
+ def fetch_non_task_trigger_ids(self, *, session):
+ fetch_non_task_trigger_ids_calls.append(session)
+ return original_fetch_non_task_trigger_ids(self, session=session)
+
+ mocker.patch.object(
+ TriggerRunnerSupervisor, "fetch_trigger_details", autospec=True,
side_effect=fetch_trigger_details
+ )
+ mocker.patch.object(
+ TriggerRunnerSupervisor,
+ "fetch_non_task_trigger_ids",
+ autospec=True,
+ side_effect=fetch_non_task_trigger_ids,
+ )
+
+ supervisor.update_triggers({trigger_orm.id})
+
+ assert len(fetch_trigger_details_calls) == 1
+ assert len(fetch_non_task_trigger_ids_calls) == 1
+ assert fetch_trigger_details_calls == [({trigger_orm.id},
fetch_non_task_trigger_ids_calls[0])]
+ assert len(supervisor.creating_triggers) == 1
+ assert supervisor.creating_triggers[0].id == trigger_orm.id
+
+ supervisor.update_triggers({trigger_orm.id})
+
+ assert len(supervisor.creating_triggers) == 1
+ assert supervisor.creating_triggers[0].id == trigger_orm.id
+
+
def
test_update_triggers_prevents_duplicate_creation_queue_entries_with_multiple_triggers(
session, supervisor_builder, dag_maker
):
