diff --git a/airflow/jobs/triggerer_job_runner.py b/airflow/jobs/triggerer_job_runner.py index b25ce6a7d3790..0638c6d5d8967 100644 --- a/airflow/jobs/triggerer_job_runner.py +++ b/airflow/jobs/triggerer_job_runner.py @@ -17,6 +17,7 @@ from __future__ import annotations import asyncio +import inspect import logging import os import signal @@ -28,7 +29,7 @@ from contextlib import suppress from copy import copy from queue import SimpleQueue -from typing import TYPE_CHECKING, TypeVar +from typing import TYPE_CHECKING, AsyncIterator, TypeVar from sqlalchemy import func, select @@ -600,10 +601,17 @@ async def run_trigger(self, trigger_id, trigger): self.log.info("trigger %s starting", name) try: self.set_individual_trigger_logging(trigger) - async for event in trigger.run(): - self.log.info("Trigger %s fired: %s", self.triggers[trigger_id]["name"], event) + result: TriggerEvent | AsyncIterator[TriggerEvent] = trigger.run() + if inspect.isasyncgen(result): + async for event in result: + self.log.info("Trigger %s fired: %s", self.triggers[trigger_id]["name"], event) + self.triggers[trigger_id]["events"] += 1 + self.events.append((trigger_id, event)) + break # should we break here? + else: + self.log.info("Trigger %s fired: %s", self.triggers[trigger_id]["name"], result) self.triggers[trigger_id]["events"] += 1 - self.events.append((trigger_id, event)) + self.events.append((trigger_id, await result)) except asyncio.CancelledError: if timeout := trigger.task_instance.trigger_timeout: timeout = timeout.replace(tzinfo=timezone.utc) if not timeout.tzinfo else timeout diff --git a/airflow/triggers/base.py b/airflow/triggers/base.py index 0d239af0cafd4..87a1fc5eb8a70 100644 --- a/airflow/triggers/base.py +++ b/airflow/triggers/base.py @@ -56,7 +56,7 @@ def serialize(self) -> tuple[str, dict[str, Any]]: raise NotImplementedError("Triggers must implement serialize()") @abc.abstractmethod - async def run(self) -> AsyncIterator[TriggerEvent]: + async def run(self) -> TriggerEvent | AsyncIterator[TriggerEvent]: """ Run the trigger in an asynchronous context. @@ -73,7 +73,6 @@ async def run(self) -> AsyncIterator[TriggerEvent]: and then rely on cleanup() being called when they are no longer needed. """ raise NotImplementedError("Triggers must implement run()") - yield # To convince Mypy this is an async iterator. async def cleanup(self) -> None: """ diff --git a/airflow/triggers/external_task.py b/airflow/triggers/external_task.py index 98305ea4f1800..52ccb30326633 100644 --- a/airflow/triggers/external_task.py +++ b/airflow/triggers/external_task.py @@ -186,7 +186,7 @@ def serialize(self) -> tuple[str, dict[str, typing.Any]]: }, ) - async def run(self) -> typing.AsyncIterator[TriggerEvent]: + async def run(self) -> TriggerEvent: """ Check periodically in the database to see if the dag exists and is in the running state. @@ -203,16 +203,14 @@ async def run(self) -> typing.AsyncIterator[TriggerEvent]: self.log.info("Waiting for DAG to start execution...") await asyncio.sleep(self.poll_interval) else: - yield TriggerEvent({"status": "timeout"}) - return + return TriggerEvent({"status": "timeout"}) # mypy confuses typing here if await self.count_tasks() == len(self.execution_dates): # type: ignore[call-arg] - yield TriggerEvent({"status": "success"}) - return + return TriggerEvent({"status": "success"}) self.log.info("Task is still running, sleeping for %s seconds...", self.poll_interval) await asyncio.sleep(self.poll_interval) except Exception: - yield TriggerEvent({"status": "failed"}) + return TriggerEvent({"status": "failed"}) @sync_to_async @provide_session @@ -282,13 +280,13 @@ def serialize(self) -> tuple[str, dict[str, typing.Any]]: }, ) - async def run(self) -> typing.AsyncIterator[TriggerEvent]: + async def run(self) -> TriggerEvent: """Check periodically if the dag run exists, and has hit one of the states yet, or not.""" while True: # mypy confuses typing here num_dags = await self.count_dags() # type: ignore[call-arg] if num_dags == len(self.execution_dates): - yield TriggerEvent(self.serialize()) + return TriggerEvent(self.serialize()) await asyncio.sleep(self.poll_interval) @sync_to_async diff --git a/airflow/triggers/temporal.py b/airflow/triggers/temporal.py index 18bdd80bff385..e7af9b705507c 100644 --- a/airflow/triggers/temporal.py +++ b/airflow/triggers/temporal.py @@ -47,7 +47,7 @@ def __init__(self, moment: datetime.datetime): def serialize(self) -> tuple[str, dict[str, Any]]: return ("airflow.triggers.temporal.DateTimeTrigger", {"moment": self.moment}) - async def run(self): + async def run(self) -> TriggerEvent: """ Loop until the relevant time is met. @@ -69,8 +69,8 @@ async def run(self): self.log.info("sleeping 1 second...") await asyncio.sleep(1) # Send our single event and then we're done - self.log.info("yielding event with payload %r", self.moment) - yield TriggerEvent(self.moment) + self.log.info("returning event with payload %r", self.moment) + return TriggerEvent(self.moment) class TimeDeltaTrigger(DateTimeTrigger):