sunank200 commented on code in PR #40084:
URL: https://github.com/apache/airflow/pull/40084#discussion_r1682878734
##########
airflow/models/taskinstance.py:
##########
@@ -715,6 +715,12 @@ 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 == TaskDeferred.TRIGGER_EXIT:
+ raise AirflowException(
+ "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:
I have removed `TaskDeferred.TRIGGER_EXIT` now
##########
airflow/sensors/date_time.py:
##########
Review Comment:
I think the same there is `DateTimeSensor` as well which works on
non-deferrable mode as well.
##########
airflow/dag_processing/manager.py:
##########
@@ -580,6 +580,7 @@ def _run_parsing_loop(self):
pass
elif isinstance(agent_signal, CallbackRequest):
self._add_callback_to_queue(agent_signal)
+ self.log.debug("_add_callback_to_queue; agent signal; %s",
agent_signal)
Review Comment:
removed it.
##########
airflow/triggers/base.py:
##########
@@ -137,3 +150,105 @@ def __eq__(self, other):
if isinstance(other, TriggerEvent):
return other.payload == self.payload
return False
+
+ @provide_session
+ def handle_submit(self, *, task_instance: TaskInstance, session: Session =
NEW_SESSION) -> None:
+ """
+ Handle the submit event for a given task instance.
+
+ This function sets the next method and next kwargs of the task
instance,
+ as well as its state to scheduled. It also adds the event's payload
+ into the kwargs for the task.
+
+ :param task_instance: The task instance to handle the submit event for.
+ :param session: The session to be used for the database callback sink.
+ """
+ # Get the next kwargs of the task instance, or an empty dictionary if
it doesn't exist
+ next_kwargs = task_instance.next_kwargs or {}
+
+ # Add the event's payload into the kwargs for the task
+ next_kwargs["event"] = self.payload
+
+ # Update the next kwargs of the task instance
+ task_instance.next_kwargs = next_kwargs
+
+ # Remove ourselves as its trigger
+ task_instance.trigger_id = None
+
+ # Set the state of the task instance to scheduled
+ task_instance.state = TaskInstanceState.SCHEDULED
+
+
+class BaseTaskEndEvent(TriggerEvent):
+ """Base event class to end the task without resuming on worker."""
+
+ task_instance_state: TaskInstanceState
+
+ def __init__(self, *, xcoms: dict[str, Any] | None = None, **kwargs) ->
None:
+ """
+ Initialize the class with the specified parameters.
+
+ :param xcoms: A dictionary of XComs or None.
+ :param kwargs: Additional keyword arguments.
+ """
+ if "payload" in kwargs:
+ raise ValueError("Param 'payload' not supported for this class.")
+ super().__init__(payload=self.task_instance_state)
+ self.xcoms = xcoms
+
+ @provide_session
+ def handle_submit(self, *, task_instance: TaskInstance, session: Session =
NEW_SESSION) -> None:
+ """
+ Submit event for the given task instance.
+
+ Marks the task with the state `task_instance_state` and optionally
pushes xcom if applicable.
+
+ :param task_instance: The task instance to be submitted.
+ :param session: The session to be used for the database callback sink.
+ """
+ # Mark the task with terminal state and prevent it from resuming on
worker
+ task_instance.trigger_id = None
+ task_instance.state = self.task_instance_state
+
+ self._submit_callback_if_necessary(task_instance=task_instance,
session=session)
+ self._push_xcoms_if_necessary(task_instance=task_instance)
+
+ def _submit_callback_if_necessary(self, *, task_instance: TaskInstance,
session) -> None:
+ """Submit a callback request if the task state is SUCCESS or FAILED."""
+ is_failure = self.task_instance_state == TaskInstanceState.FAILED
+ if self.task_instance_state in [TaskInstanceState.SUCCESS,
TaskInstanceState.FAILED]:
+ request = TaskCallbackRequest(
+ full_filepath=task_instance.dag_model.fileloc,
+ simple_task_instance=SimpleTaskInstance.from_ti(task_instance),
+ is_failure_callback=is_failure,
+ task_callback_type=self.task_instance_state,
+ )
+ log.info("Sending callback: %s", request)
+ try:
+ DatabaseCallbackSink().send(callback=request, session=session)
+ except Exception as e:
+ log.error("Failed to send callback: %s", e)
+
+ def _push_xcoms_if_necessary(self, *, task_instance: TaskInstance) -> None:
+ """Pushes XComs to the database if they are provided."""
+ if self.xcoms:
Review Comment:
When the task ends or exits from the trigger, the values from the output
needs to be pushed to xcom if applicable.
--
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]
