refactor: Updated IterableOperator with DeferredIterable

apache · dabla · Sep 29, 2024 · Sep 29, 2024 · Sep 29, 2024 · Sep 30, 2024
commit 97c64467e56f573c1823b2484a73743e83c2881d
@@ -18,13 +18,14 @@
 from __future__ import annotations
 
 import asyncio
-from collections.abc import Generator, Iterator, Sequence
+from collections.abc import Generator, Iterable, Iterator, Sequence
 from contextlib import contextmanager, suppress
 from typing import TYPE_CHECKING, Any
 
 from airflow.exceptions import AirflowException
 from airflow.serialization import serde
 from airflow.utils.log.logging_mixin import LoggingMixin
+from airflow.utils.mixins import ResolveMixin
 from airflow.utils.module_loading import import_string
 from airflow.utils.xcom import XCOM_RETURN_KEY
 
@@ -117,7 +118,7 @@ def deserialize(cls, data: dict, version: int):
         return XComIterable(**data)
 
 
-class DeferredIterable(Iterator, LoggingMixin):
+class DeferredIterable(Iterator, ResolveMixin, LoggingMixin):
     """An iterable that lazily fetches XCom values one by one instead of loading all at once."""
 
     def __init__(
@@ -136,7 +137,12 @@ def __init__(
         self.context = context
         self.index = 0
 
-    def resolve(self, context: Context) -> DeferredIterable:
+    def iter_references(self) -> Iterable[tuple[Operator, str]]:
+        yield self.operator, XCOM_RETURN_KEY
+
+    def resolve(
+        self, context: Context, *, include_xcom: bool = True
+    ) -> DeferredIterable:
         return DeferredIterable(
             results=self.results,
             trigger=self.trigger,
@@ -159,10 +165,25 @@ def __next__(self):
 
         self.log.info("No more results. Running trigger: %s", self.trigger)
 
+        if not self.context:
+            raise AirflowException("Context is required to run the trigger.")
+
+        results = self._execute_trigger()
+
+        if isinstance(results, (list, set)):
+            self.results.extend(results)
+        else:
+            self.results.append(results)
+
+        self.index += 1
+        return self.results[-1]
+
+    def _execute_trigger(self):
         try:
             with event_loop() as loop:
                 self.log.info("Running trigger: %s", self.trigger)
                 event = loop.run_until_complete(run_trigger(self.trigger))
+                self.operator.render_template_fields(context=self.context)
                 next_method = getattr(self.operator, self.next_method)
                 self.log.info("Triggering next method: %s", self.next_method)
                 results = next_method(self.context, event.payload)
@@ -172,15 +193,13 @@ def __next__(self):
 
         if isinstance(results, DeferredIterable):
             self.trigger = results.trigger
-            self.results.extend(results.results)
-        else:
-            self.trigger = None
-            self.results.extend(results)
+            return results.results
 
-        self.index += 1
-        return self.results[-1]
+        self.trigger = None
+        return results
 
     def __len__(self):
+        # TODO: maybe we should raise an exception here as you can't know the total length of an iterable in advance
         return len(self.results)
 
     def __getitem__(self, index: int):
@@ -201,14 +220,15 @@ def serialize(self):
         }
 
     @classmethod
-    def get_operator_from_dag(cls, dag_fileloc: str, dag_id: str, task_id: str) -> Operator:
+    def get_operator_from_dag(
+        cls, dag_fileloc: str, dag_id: str, task_id: str
+    ) -> Operator:
         """Loads a DAG using DagBag and gets the operator by task_id."""
 
         from airflow.models import DagBag
 
         dag_bag = DagBag(dag_folder=None)  # Avoid loading all DAGs
-        processed_dags = dag_bag.process_file(dag_fileloc)
-        cls.logger().info("processed_dags: %s", processed_dags)
+        dag_bag.process_file(dag_fileloc)
         cls.logger().info("dag_bag: %s", dag_bag)
         cls.logger().info("dags: %s", dag_bag.dags)
         return dag_bag.dags[dag_id].get_task(task_id)
@@ -219,9 +239,14 @@ def deserialize(cls, data: dict, version: int):
 
         trigger_class = import_string(data["trigger"][0])
         trigger = trigger_class(**data["trigger"][1])
-        operator = cls.get_operator_from_dag(data["dag_fileloc"], data["dag_id"], data["task_id"])
+        operator = cls.get_operator_from_dag(
+            data["dag_fileloc"], data["dag_id"], data["task_id"]
+        )
         return DeferredIterable(
-            results=data["results"], trigger=trigger, operator=operator, next_method=data["next_method"]
+            results=data["results"],
+            trigger=trigger,
+            operator=operator,
+            next_method=data["next_method"],
         )
 
 

@@ -77,6 +77,7 @@ def __init__(
         self.__context = dict(context.items())
         self._task_instance = task_instance
         self._is_async_mode: bool = False  # Flag to track sync/async mode
+        self._result: Any | None = None
 
     @property
     def task_instance(self) -> TaskInstance:
@@ -99,8 +100,16 @@ def mode(self) -> str:
         return "async" if self._is_async_mode else "sync"
 
     @abstractmethod
+    def execute(self, *args, **kwargs):
+        raise NotImplementedError
+
     def run(self, *args, **kwargs):
-        raise NotImplementedError()
+        self._result = self.execute(*args, **kwargs)
+        return self._result
+
+    async def run_deferred(self, *args, **kwargs):
+        self._result = await self.execute(*args, **kwargs)
+        return self._result
 
     def __enter__(self):
         if self.log.isEnabledFor(logging.INFO):
@@ -145,7 +154,9 @@ def __exit__(self, exc_type, exc_value, traceback):
 
             self.task_instance.set_state(TaskInstanceState.FAILED)
             raise exc_value
-        self.operator.post_execute(context=self.context)
+        if self.operator.do_xcom_push:
+            self.task_instance.xcom_push(key=XCOM_RETURN_KEY, value=self._result)
+        self.operator.post_execute(context=self.context, result=self._result)
         self.task_instance.set_state(TaskInstanceState.SUCCESS)
         if self.log.isEnabledFor(logging.INFO):
             self.log.info(
@@ -170,7 +181,7 @@ class OperatorExecutor(TaskExecutor):
     :meta private:
     """
 
-    def run(self, *args, **kwargs):
+    def execute(self, *args, **kwargs):
         outlet_events = context_get_outlet_events(self.context)
         # TODO: change back to operator.execute once ExecutorSafeguard is fixed
         if hasattr(self.operator.execute, "__wrapped__"):
@@ -197,7 +208,7 @@ class TriggerExecutor(TaskExecutor):
     :meta private:
     """
 
-    async def run(self, task_deferred: TaskDeferred):
+    async def execute(self, task_deferred: TaskDeferred):
         event = await run_trigger(task_deferred.trigger)
 
         self.log.debug("event: %s", event)
@@ -217,7 +228,7 @@ async def run(self, task_deferred: TaskDeferred):
                         logger=self.log,
                     ).run(self.context, event.payload)
                 except TaskDeferred as task_deferred:
-                    return await self.run(task_deferred=task_deferred)
+                    return await self.execute(task_deferred=task_deferred)
 
 
 class IterableOperator(BaseOperator):
@@ -321,9 +332,18 @@ def _resolve_expand_input(self, context: Context, session: Session):
         self.log.debug("resolved_input: %s", resolved_input)
 
         if isinstance(resolved_input, _MapResult):
-            self._mapped_kwargs = map(lambda value: self._resolve(value=value, context=context, session=session), resolved_input)
+            self._mapped_kwargs = map(
+                lambda value: self._resolve(
+                    value=value, context=context, session=session
+                ),
+                resolved_input,
+            )
         else:
-            self._mapped_kwargs = iter(self._lazy_mapped_kwargs(input=resolved_input, context=context, session=session))
+            self._mapped_kwargs = iter(
+                self._lazy_mapped_kwargs(
+                    input=resolved_input, context=context, session=session
+                )
+            )
 
         self.log.debug("mapped_kwargs: %s", self._mapped_kwargs)
 
@@ -345,7 +365,9 @@ def _run_tasks(
         prev_futures_count = 0
         futures: dict[ApplyResult, TaskInstance] = {}
         failed_tasks: list[TaskInstance] = []
-        chunked_tasks: Iterator[Iterable[TaskInstance]] = ichunked(tasks, (self.max_active_tis_per_dag * 2))
+        chunked_tasks: Iterator[Iterable[TaskInstance]] = ichunked(
+            tasks, (self.max_active_tis_per_dag * 2)
+        )
 
         with ThreadPool(processes=self.max_active_tis_per_dag) as pool:
             for task in next(chunked_tasks, []):
@@ -385,7 +407,9 @@ def _run_tasks(
                         if task.next_try_number > self.retries:
                             exception = AirflowTaskTimeout(e)
                         else:
-                            reschedule_date = min(reschedule_date, task.next_retry_datetime())
+                            reschedule_date = min(
+                                reschedule_date, task.next_retry_datetime()
+                            )
                             failed_tasks.append(task)
                     except AirflowRescheduleTaskInstanceException as e:
                         reschedule_date = min(reschedule_date, e.reschedule_date)
@@ -405,15 +429,20 @@ def _run_tasks(
 
                     with event_loop() as loop:
                         for result in loop.run_until_complete(
-                            gather(*[loop.create_task(task) for task in deferred_tasks], return_exceptions=True)
+                            gather(
+                                *[loop.create_task(task) for task in deferred_tasks],
+                                return_exceptions=True,
+                            )
                         ):
                             self.log.debug("result: %s", result)
 
                             if isinstance(result, Exception):
                                 if isinstance(
                                     result, AirflowRescheduleTaskInstanceException
                                 ):
-                                    reschedule_date = min(reschedule_date, result.reschedule_date)
+                                    reschedule_date = min(
+                                        reschedule_date, result.reschedule_date
+                                    )
                                     failed_tasks.append(result.task)
                                 else:
                                     exception = result
@@ -449,7 +478,9 @@ def _run_tasks(
     @classmethod
     def _xcom_pull(cls, task_instance: TaskInstance):
         with suppress(JSONDecodeError):
-            return task_instance.xcom_pull(task_ids=task_instance.task_id, dag_id=task_instance.dag_id)
+            return task_instance.xcom_pull(
+                task_ids=task_instance.task_id, dag_id=task_instance.dag_id
+            )
         return None
 
     def _run_operator(self, context: Context, task_instance: TaskInstance):
@@ -462,11 +493,11 @@ def _run_operator(self, context: Context, task_instance: TaskInstance):
                 with OperatorExecutor(
                     context=context, task_instance=task_instance
                 ) as executor:
-                    result = executor.run()
-                    if self.do_xcom_push:
-                        task_instance.xcom_push(key=XCOM_RETURN_KEY, value=result)
+                    return executor.run()
             else:
-                self.log.info("Task instance %s already completed.", task_instance.task_id)
+                self.log.info(
+                    "Task instance %s already completed.", task_instance.task_id
+                )
             return result
         except TaskDeferred as task_deferred:
             return task_deferred
@@ -478,12 +509,11 @@ async def _run_deferrable(
             async with TriggerExecutor(
                 context=context, task_instance=task_instance
             ) as executor:
-                result = await executor.run(task_deferred)
-                if self.do_xcom_push:
-                    task_instance.xcom_push(key=XCOM_RETURN_KEY, value=result)
-                return result
+                return await executor.run_deferred(task_deferred)
 
-    def _create_task(self, run_id: str, index: int, mapped_kwargs: dict) -> TaskInstance:
+    def _create_task(
+        self, run_id: str, index: int, mapped_kwargs: dict
+    ) -> TaskInstance:
         operator = self._unmap_operator(index, mapped_kwargs)
         return TaskInstance(
             task=operator,
@@ -495,7 +525,9 @@ def execute(self, context: Context):
             context=context,
             tasks=iter(
                 map(
-                    lambda mapped_kwargs: self._create_task(context["ti"].run_id, mapped_kwargs[0], mapped_kwargs[1]),
+                    lambda mapped_kwargs: self._create_task(
+                        context["ti"].run_id, mapped_kwargs[0], mapped_kwargs[1]
+                    ),
                     enumerate(self._mapped_kwargs),
                 )
             ),

@@ -33,8 +33,6 @@
 from airflow.utils.trigger_rule import TriggerRule
 from airflow.utils.xcom import XCOM_RETURN_KEY
 
-from airflow.models.iterable import DeferredIterable
-
 if TYPE_CHECKING:
     from airflow.sdk.bases.operator import BaseOperator
     from airflow.sdk.definitions.edges import EdgeModifier
@@ -358,6 +356,8 @@ def resolve(self, context: Mapping[str, Any]) -> Any:
                 default=NOTSET,
             )
         if not isinstance(result, ArgNotSet):
+            if isinstance(result, ResolveMixin):
+                return result.resolve(context)
             return result
         if self.key == XCOM_RETURN_KEY:
             return None
@@ -450,7 +450,7 @@ def resolve(self, context: Mapping[str, Any]) -> Any:
         value = self.arg.resolve(context)
         if not isinstance(value, (Sequence, Iterable, dict)):
             raise ValueError(f"XCom map expects sequence or dict, not {type(value).__name__}")
-        if isinstance(value, DeferredIterable):
+        if isinstance(value, ResolveMixin):
             value = value.resolve(context)
         return _MapResult(value, self.callables)