Fix bug in Dataflow hook when no jobs are returned (#18981)

apache · Oct 26, 2021 · 726a151 · 726a151
1 parent cdbb481
commit 726a151
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 3 deletions.
diff --git a/airflow/providers/google/cloud/hooks/dataflow.py b/airflow/providers/google/cloud/hooks/dataflow.py
@@ -356,18 +356,21 @@ def _fetch_all_jobs(self) -> List[dict]:
  .jobs()
  .list(projectId=self._project_number, location=self._job_location)
  )
- jobs: List[dict] = []
+ all_jobs: List[dict] = []
  while request is not None:
  response = request.execute(num_retries=self._num_retries)
- jobs.extend(response["jobs"])
+ jobs = response.get("jobs")
+ if jobs is None:
+ break
+ all_jobs.extend(jobs)
 
  request = (
  self._dataflow.projects()
  .locations()
  .jobs()
  .list_next(previous_request=request, previous_response=response)
  )
- return jobs
+ return all_jobs
 
  def _fetch_jobs_by_prefix_name(self, prefix_name: str) -> List[dict]:
  jobs = self._fetch_all_jobs()

diff --git a/tests/providers/google/cloud/hooks/test_dataflow.py b/tests/providers/google/cloud/hooks/test_dataflow.py
@@ -1697,6 +1697,27 @@ def test_fetch_list_job_messages_responses(self):
  )
  assert result == ["response_1"]
 
+ def test_fetch_all_jobs_when_no_jobs_returned(self):
+ # fmt: off
+ (
+ self.mock_dataflow
+ .projects.return_value
+ .locations.return_value
+ .jobs.return_value
+ .list.return_value
+ .execute.return_value
+ ) = {}
+ # fmt: on
+
+ jobs_controller = _DataflowJobsController(
+ dataflow=self.mock_dataflow,
+ project_number=TEST_PROJECT,
+ location=TEST_LOCATION,
+ job_id=TEST_JOB_ID,
+ )
+ result = jobs_controller._fetch_all_jobs()
+ assert result == []
+
  @mock.patch(DATAFLOW_STRING.format('_DataflowJobsController._fetch_list_job_messages_responses'))
  def test_fetch_job_messages_by_id(self, mock_fetch_responses):
  mock_fetch_responses.return_value = iter(