Repository: apache/airflow
File: providers/google/src/airflow/providers/google/cloud/triggers/dataproc.py (Python)

Code before:
                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

    async def run(self):
        hook = self.get_async_hook()
        # Trigger client cache with sync call get_credentials(), evaluated once.
        await hook.get_job_client(region=self.region)

        try:
            while True:
                job = await hook.get_job(
                    project_id=self.project_id, region=self.region, job_id=self.job_id
                )
                state = job.status.state
                self.log.info("Dataproc job: %s is in state: %s", self.job_id, state)
                if state in (JobStatus.State.DONE, JobStatus.State.CANCELLED, JobStatus.State.ERROR):
                    break
                await asyncio.sleep(self.polling_interval_seconds)
            yield TriggerEvent(
                {"job_id": self.job_id, "job_state": JobStatus.State(state).name, "job": Job.to_dict(job)}
            )
        except asyncio.CancelledError:
            self.log.info("Task got cancelled.")
            try:
                if self.job_id and self.cancel_on_kill and await self.safe_to_cancel():
                    self.log.info(
                        "Cancelling the job as it is safe to do so. Note that the airflow TaskInstance is not"
                        " in deferred state."
                    )
                    self.log.info("Cancelling the job: %s", self.job_id)
                    # The synchronous hook is utilized to delete the cluster when a task is cancelled. This
                    # is because the asynchronous hook deletion is not awaited when the trigger task is

Review comment (bug):
Since the client is now initialized once and reused for the trigger lifetime, do we rely on google-auth’s automatic credential refresh? Just checking that long-running triggers won’t fail due to token expiry.

Code after:
                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

    async def run(self):
        hook = self.get_async_hook()
        # Trigger client cache with sync call get_credentials(), evaluated once.
        await hook.get_job_client(region=self.region)

        try:
            while True:
                job = await hook.get_job(project_id=self.project_id, region=self.region, job_id=self.job_id)
                state = job.status.state
                self.log.info("Dataproc job: %s is in state: %s", self.job_id, state)
                if state in (JobStatus.State.DONE, JobStatus.State.CANCELLED, JobStatus.State.ERROR):
                    break
                await asyncio.sleep(self.polling_interval_seconds)
            yield TriggerEvent(
                {"job_id": self.job_id, "job_state": JobStatus.State(state).name, "job": Job.to_dict(job)}
            )
        except asyncio.CancelledError:
            self.log.info("Task got cancelled.")
            try:
                if self.job_id and self.cancel_on_kill and await self.safe_to_cancel():
                    self.log.info(
                        "Cancelling the job as it is safe to do so. Note that the airflow TaskInstance is not"
                        " in deferred state."
                    )
                    self.log.info("Cancelling the job: %s", self.job_id)
                    # The synchronous hook is utilized to delete the cluster when a task is cancelled. This
                    # is because the asynchronous hook deletion is not awaited when the trigger task is
                    # cancelled. The call for deleting the cluster or job through the sync hook is not a
                    # blocking call, which means it does not wait until the cluster or job is deleted.