Load just the first yielded event in the triggerer

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

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:
         """

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

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):