fix(trigger): add next_kwargs to StartTriggerArgs

apache · Jun 21, 2024 · b222b9a · b222b9a
1 parent 06eabdd
commit b222b9a
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py
@@ -1606,15 +1606,15 @@ def _defer_task(
 
     if exception is not None:
         trigger_row = Trigger.from_object(exception.trigger)
-        trigger_kwargs = exception.kwargs
         next_method = exception.method_name
+        next_kwargs = exception.kwargs
         timeout = exception.timeout
     elif ti.task is not None and ti.task.start_trigger_args is not None:
         trigger_row = Trigger(
             classpath=ti.task.start_trigger_args.trigger_cls,
             kwargs=ti.task.start_trigger_args.trigger_kwargs or {},
         )
-        trigger_kwargs = ti.task.start_trigger_args.trigger_kwargs
+        next_kwargs = ti.task.start_trigger_args.next_kwargs
         next_method = ti.task.start_trigger_args.next_method
         timeout = ti.task.start_trigger_args.timeout
     else:
@@ -1635,7 +1635,7 @@ def _defer_task(
     ti.state = TaskInstanceState.DEFERRED
     ti.trigger_id = trigger_row.id
     ti.next_method = next_method
-    ti.next_kwargs = trigger_kwargs or {}
+    ti.next_kwargs = next_kwargs or {}
 
     # Calculate timeout too if it was passed
     if timeout is not None:

diff --git a/airflow/triggers/base.py b/airflow/triggers/base.py
@@ -31,13 +31,15 @@ class StartTriggerArgs:
     trigger_cls: str
     next_method: str
     trigger_kwargs: dict[str, Any] | None = None
+    next_kwargs: dict[str, Any] | None = None
     timeout: timedelta | None = None
 
     def serialize(self):
         return {
             "trigger_cls": self.trigger_cls,
             "trigger_kwargs": self.trigger_kwargs,
             "next_method": self.next_method,
+            "next_kwargs": self.next_kwargs,
             "timeout": self.timeout,
         }