SameerMesiah97 commented on code in PR #65743:
URL: https://github.com/apache/airflow/pull/65743#discussion_r3136636407
##########
providers/google/src/airflow/providers/google/cloud/triggers/dataproc.py:
##########
@@ -123,111 +123,31 @@ def serialize(self):
},
)
- if not AIRFLOW_V_3_0_PLUS:
Review Comment:
I don't understand why you removed this. This will break cancellation
semantics for this trigger for the Airflow 3.x versions before 3.30 (which
appears to be the version where `on_kill()` is being introduced in
`BaseTrigger`).
##########
providers/google/src/airflow/providers/google/cloud/triggers/dataproc.py:
##########
@@ -123,111 +123,31 @@ def serialize(self):
},
)
- if not AIRFLOW_V_3_0_PLUS:
-
- @provide_session
- def get_task_instance(self, session: Session) -> TaskInstance:
- """
- Get the task instance for the current task.
-
- :param session: Sqlalchemy session
- """
- task_instance = session.scalar(
- select(TaskInstance).where(
- TaskInstance.dag_id == self.task_instance.dag_id,
- TaskInstance.task_id == self.task_instance.task_id,
- TaskInstance.run_id == self.task_instance.run_id,
- TaskInstance.map_index == self.task_instance.map_index,
- )
- )
- if task_instance is None:
- raise AirflowException(
- "TaskInstance with dag_id: %s,task_id: %s, run_id: %s and
map_index: %s is not found",
- self.task_instance.dag_id,
- self.task_instance.task_id,
- self.task_instance.run_id,
- self.task_instance.map_index,
- )
- return task_instance
-
- async def get_task_state(self):
- from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
-
- task_states_response = await
sync_to_async(RuntimeTaskInstance.get_task_states)(
- dag_id=self.task_instance.dag_id,
- task_ids=[self.task_instance.task_id],
- run_ids=[self.task_instance.run_id],
- map_index=self.task_instance.map_index,
+ async def on_kill(self) -> None:
+ """Cancel the Dataproc job when the trigger is cancelled by a user
action."""
+ if not (self.job_id and self.cancel_on_kill):
+ return
+ self.log.info("Cancelling the job: %s", self.job_id)
+ await sync_to_async(self.get_sync_hook().cancel_job)(
+ job_id=self.job_id, project_id=self.project_id, region=self.region
)
- try:
- task_state =
task_states_response[self.task_instance.run_id][self.task_instance.task_id]
- except Exception:
- raise AirflowException(
- "TaskInstance with dag_id: %s, task_id: %s, run_id: %s and
map_index: %s is not found",
- self.task_instance.dag_id,
- self.task_instance.task_id,
- self.task_instance.run_id,
- self.task_instance.map_index,
- )
- return task_state
-
- async def safe_to_cancel(self) -> bool:
- """
- Whether it is safe to cancel the external job which is being executed
by this trigger.
-
- This is to avoid the case that `asyncio.CancelledError` is called
because the trigger itself is stopped.
- Because in those cases, we should NOT cancel the external job.
- """
- if AIRFLOW_V_3_0_PLUS:
- task_state = await self.get_task_state()
- else:
- # Database query is needed to get the latest state of the task
instance.
- task_instance = self.get_task_instance() # type: ignore[call-arg]
- task_state = task_instance.state
- return task_state != TaskInstanceState.DEFERRED
Review Comment:
You’re removing quite a bit of existing code here, which introduces a
backwards compatibility issue for pre-3.30 Airflow. In those versions
`on_kill()` isn’t invoked, so cancellation would no longer go through. I’d have
another look at PR #65672; there `on_kill()` is added without removing the
existing `CancelledError` path.
I think it would be better to just add `on_kill()` (with tests) and leave
the current logic in `run() `as-is. That way older versions continue to use the
existing code path, and newer versions can make use of `on_kill()` when it’s
available.
--
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]
