Lee-W commented on code in PR #40084:
URL: https://github.com/apache/airflow/pull/40084#discussion_r1684527807
##########
airflow/models/taskinstance.py:
##########
@@ -715,6 +716,13 @@ def _execute_task(task_instance: TaskInstance |
TaskInstancePydantic, context: C
if not task_instance.next_kwargs:
task_instance.next_kwargs = {}
task_instance.next_kwargs[f"{task_to_execute.__class__.__name__}__sentinel"] =
_sentinel
+ elif task_instance.next_method == "__trigger_exit__":
Review Comment:
Where does this "__trigger_exit__" come from? I thought it was a constant.
Even if we move it out of the `TaskDeferred`, I think we probably still should
make it a constant for op author to use. Also this behavior seems to be missed
in the doc
##########
airflow/sensors/time_delta.py:
##########
@@ -59,28 +59,34 @@ class TimeDeltaSensorAsync(TimeDeltaSensor):
Will defers itself to avoid taking up a worker slot while it is waiting.
:param delta: time length to wait after the data interval before
succeeding.
+ :param end_from_trigger: End the task directly from the triggerer without
going into the worker.
.. seealso::
For more information on how to use this sensor, take a look at the
guide:
:ref:`howto/operator:TimeDeltaSensorAsync`
"""
+ def __init__(self, *, end_from_trigger: bool = False, delta, **kwargs):
+ super().__init__(delta=delta, **kwargs)
+ self.end_from_trigger = end_from_trigger
+ self.delta = delta
+
def execute(self, context: Context) -> bool | NoReturn:
target_dttm = context["data_interval_end"]
target_dttm += self.delta
if timezone.utcnow() > target_dttm:
# If the target datetime is in the past, return immediately
return True
try:
- trigger = DateTimeTrigger(moment=target_dttm)
+ trigger = DateTimeTrigger(moment=target_dttm,
end_from_trigger=self.end_from_trigger)
except (TypeError, ValueError) as e:
if self.soft_fail:
raise AirflowSkipException("Skipping due to soft_fail is set
to True.") from e
raise
self.defer(trigger=trigger, method_name="execute_complete")
- def execute_complete(self, context, event=None) -> None:
- """Execute for when the trigger fires - return immediately."""
+ def execute_complete(self, context, event=None):
Review Comment:
```suggestion
def execute_complete(self, context, event=None):
def execute_complete(self, context: Context, event: dict[str, Any] |
None = None) -> None:
```
##########
airflow/models/taskinstance.py:
##########
@@ -715,6 +716,13 @@ def _execute_task(task_instance: TaskInstance |
TaskInstancePydantic, context: C
if not task_instance.next_kwargs:
task_instance.next_kwargs = {}
task_instance.next_kwargs[f"{task_to_execute.__class__.__name__}__sentinel"] =
_sentinel
+ elif task_instance.next_method == "__trigger_exit__":
+ log.error("Task is resuming from deferral without next_method
specified.")
+ raise TaskDeferralError(
+ "Task is resuming from deferral without next_method specified.
"
+ "You must either set `method_name` when deferring, or use a
trigger "
+ "that is designed to exit the task."
+ )
Review Comment:
```suggestion
error_msg = "Task is resuming from deferral without next_method
specified."
log.error(error_msg)
raise TaskDeferralError(
f"{error_msg} "
"You must either set `method_name` when deferring, or use a
trigger "
"that is designed to exit the task."
)
```
nitpick: not sure whether it's better, but I would probably try this
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]