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")
