def poke(self, context: dict) -> bool: hook = DataprocHook(gcp_conn_id=self.gcp_conn_id) job = hook.get_job(job_id=self.dataproc_job_id, location=self.location, project_id=self.project_id) state = job.status.state if state == JobStatus.ERROR: raise AirflowException('Job failed:\n{}'.format(job)) elif state in {JobStatus.CANCELLED, JobStatus.CANCEL_PENDING, JobStatus.CANCEL_STARTED}: raise AirflowException('Job was cancelled:\n{}'.format(job)) elif JobStatus.DONE == state: self.log.debug("Job %s completed successfully.", self.dataproc_job_id) return True elif JobStatus.ATTEMPT_FAILURE == state: self.log.debug("Job %s attempt has failed.", self.dataproc_job_id) self.log.info("Waiting for job %s to complete.", self.dataproc_job_id) return False
def poke(self, context: "Context") -> bool: hook = DataprocHook(gcp_conn_id=self.gcp_conn_id) if self.wait_timeout: try: job = hook.get_job(job_id=self.dataproc_job_id, region=self.region, project_id=self.project_id) except ServerError as err: self.log.info(f"DURATION RUN: {self._duration()}") if self._duration() > self.wait_timeout: raise AirflowException( f"Timeout: dataproc job {self.dataproc_job_id} " f"is not ready after {self.wait_timeout}s") self.log.info( "Retrying. Dataproc API returned server error when waiting for job: %s", err) return False else: job = hook.get_job(job_id=self.dataproc_job_id, region=self.region, project_id=self.project_id) state = job.status.state if state == JobStatus.State.ERROR: raise AirflowException(f'Job failed:\n{job}') elif state in { JobStatus.State.CANCELLED, JobStatus.State.CANCEL_PENDING, JobStatus.State.CANCEL_STARTED, }: raise AirflowException(f'Job was cancelled:\n{job}') elif JobStatus.State.DONE == state: self.log.debug("Job %s completed successfully.", self.dataproc_job_id) return True elif JobStatus.State.ATTEMPT_FAILURE == state: self.log.debug("Job %s attempt has failed.", self.dataproc_job_id) self.log.info("Waiting for job %s to complete.", self.dataproc_job_id) return False
def setUp(self): with mock.patch(BASE_STRING.format("GoogleBaseHook.__init__"), new=mock_init): self.hook = DataprocHook(gcp_conn_id="test")