chaimt commented on a change in pull request #4633: AIRFLOW-3791: Dataflow -
Support check status if pipeline spans on multiple jobs
URL: https://github.com/apache/airflow/pull/4633#discussion_r303917061
##########
File path: airflow/contrib/hooks/gcp_dataflow_hook.py
##########
@@ -47,58 +47,131 @@ class DataflowJobStatus:
JOB_STATE_FAILED = "JOB_STATE_FAILED"
JOB_STATE_CANCELLED = "JOB_STATE_CANCELLED"
JOB_STATE_PENDING = "JOB_STATE_PENDING"
+ FAILED_END_STATES = {JOB_STATE_FAILED, JOB_STATE_CANCELLED}
+ SUCCEEDED_END_STATES = {JOB_STATE_DONE}
+ END_STATES = SUCCEEDED_END_STATES | FAILED_END_STATES
class _DataflowJob(LoggingMixin):
def __init__(self, dataflow, project_number, name, location, poll_sleep=10,
- job_id=None, num_retries=None):
+ job_id=None, num_retries=None, multiple_jobs=None):
self._dataflow = dataflow
self._project_number = project_number
self._job_name = name
self._job_location = location
+ self._multiple_jobs = multiple_jobs
self._job_id = job_id
self._num_retries = num_retries
- self._job = self._get_job()
+ if self._num_retries is None:
+ self._num_retries = 0
self._poll_sleep = poll_sleep
+ self._jobs = self._get_jobs()
+
+ def is_job_running(self):
+ """
+ Helper method to check if jos is still running in dataflow
+
+ :return: True if job is running.
+ :rtype: bool
+ """
+ for job in self._jobs:
+ if job['currentState'] not in DataflowJobStatus.END_STATES:
+ return True
+ return False
def _get_job_id_from_name(self):
+ """
+ Helper method to get list of jobs that start with job name
+
+ :return: list of jobs including id's
+ :rtype: list
+ """
jobs = self._dataflow.projects().locations().jobs().list(
projectId=self._project_number,
location=self._job_location
).execute(num_retries=self._num_retries)
- for job in jobs['jobs']:
- if job['name'].lower() == self._job_name.lower():
- self._job_id = job['id']
- return job
- return None
+ dataflow_jobs = []
+ if jobs:
+ for job in jobs['jobs']:
+ if job['name'].startswith(self._job_name.lower()):
+ dataflow_jobs.append(job)
+ if len(dataflow_jobs) == 1:
+ self._job_id = dataflow_jobs[0]['id']
+ return dataflow_jobs
+
+ def _get_jobs(self):
+ """
+ Helper method to get all jobs by name
- def _get_job(self):
- if self._job_id:
- job = self._dataflow.projects().locations().jobs().get(
+ :return: jobs
+ :rtype: list
+ """
+ if not self._multiple_jobs and self._job_id:
+ self._jobs = []
+ self._jobs.append(self._dataflow.projects().locations().jobs().get(
Review comment:
done
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
