Move the try outside the loop when this is possible in Google provider (

#33976) * Move the try outside the loop when this is possible in Google provider --------- Co-authored-by: Tzu-ping Chung <[email protected]>
apache · Sep 3, 2023 · 4f20b0f · 4f20b0f
1 parent 8918b43
commit 4f20b0f
Show file tree

Hide file tree

Showing 10 changed files with 103 additions and 118 deletions.
diff --git a/airflow/providers/google/cloud/triggers/bigquery.py b/airflow/providers/google/cloud/triggers/bigquery.py
@@ -75,8 +75,8 @@ async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
  """Gets current job execution status and yields a TriggerEvent."""
  """Gets current job execution status and yields a TriggerEvent."""
  hook = self._get_async_hook()
- while True:
- try:
+ try:
+ while True:
  job_status = await hook.get_job_status(job_id=self.job_id, project_id=self.project_id)
  if job_status == "success":
  yield TriggerEvent(
@@ -95,10 +95,9 @@ async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
  "Bigquery job status is %s. Sleeping for %s seconds.", job_status, self.poll_interval
  )
  await asyncio.sleep(self.poll_interval)
- except Exception as e:
- self.log.exception("Exception occurred while checking for query completion")
- yield TriggerEvent({"status": "error", "message": str(e)})
- return
+ except Exception as e:
+ self.log.exception("Exception occurred while checking for query completion")
+ yield TriggerEvent({"status": "error", "message": str(e)})
 
  def _get_async_hook(self) -> BigQueryAsyncHook:
  return BigQueryAsyncHook(gcp_conn_id=self.conn_id)
@@ -124,8 +123,8 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
  async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
  """Gets current job execution status and yields a TriggerEvent."""
  hook = self._get_async_hook()
- while True:
- try:
+ try:
+ while True:
  # Poll for job execution status
  job_status = await hook.get_job_status(job_id=self.job_id, project_id=self.project_id)
  if job_status == "success":
@@ -160,10 +159,9 @@ async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
  "Bigquery job status is %s. Sleeping for %s seconds.", job_status, self.poll_interval
  )
  await asyncio.sleep(self.poll_interval)
- except Exception as e:
- self.log.exception("Exception occurred while checking for query completion")
- yield TriggerEvent({"status": "error", "message": str(e)})
- return
+ except Exception as e:
+ self.log.exception("Exception occurred while checking for query completion")
+ yield TriggerEvent({"status": "error", "message": str(e)})
 
 
 class BigQueryGetDataTrigger(BigQueryInsertJobTrigger):
@@ -196,8 +194,8 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
  async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
  """Gets current job execution status and yields a TriggerEvent with response data."""
  hook = self._get_async_hook()
- while True:
- try:
+ try:
+ while True:
  # Poll for job execution status
  job_status = await hook.get_job_status(job_id=self.job_id, project_id=self.project_id)
  if job_status == "success":
@@ -220,10 +218,9 @@ async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
  "Bigquery job status is %s. Sleeping for %s seconds.", job_status, self.poll_interval
  )
  await asyncio.sleep(self.poll_interval)
- except Exception as e:
- self.log.exception("Exception occurred while checking for query completion")
- yield TriggerEvent({"status": "error", "message": str(e)})
- return
+ except Exception as e:
+ self.log.exception("Exception occurred while checking for query completion")
+ yield TriggerEvent({"status": "error", "message": str(e)})
 
 
 class BigQueryIntervalCheckTrigger(BigQueryInsertJobTrigger):
@@ -302,8 +299,8 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
  async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
  """Gets current job execution status and yields a TriggerEvent."""
  hook = self._get_async_hook()
- while True:
- try:
+ try:
+ while True:
  first_job_response_from_hook = await hook.get_job_status(
  job_id=self.first_job_id, project_id=self.project_id
  )
@@ -365,10 +362,9 @@ async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
  )
  return
 
- except Exception as e:
- self.log.exception("Exception occurred while checking for query completion")
- yield TriggerEvent({"status": "error", "message": str(e)})
- return
+ except Exception as e:
+ self.log.exception("Exception occurred while checking for query completion")
+ yield TriggerEvent({"status": "error", "message": str(e)})
 
 
 class BigQueryValueCheckTrigger(BigQueryInsertJobTrigger):
@@ -430,8 +426,8 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
  async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
  """Gets current job execution status and yields a TriggerEvent."""
  hook = self._get_async_hook()
- while True:
- try:
+ try:
+ while True:
  # Poll for job execution status
  response_from_hook = await hook.get_job_status(job_id=self.job_id, project_id=self.project_id)
  if response_from_hook == "success":
@@ -448,10 +444,9 @@ async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
  else:
  yield TriggerEvent({"status": "error", "message": response_from_hook, "records": None})
  return
- except Exception as e:
- self.log.exception("Exception occurred while checking for query completion")
- yield TriggerEvent({"status": "error", "message": str(e)})
- return
+ except Exception as e:
+ self.log.exception("Exception occurred while checking for query completion")
+ yield TriggerEvent({"status": "error", "message": str(e)})
 
 
 class BigQueryTableExistenceTrigger(BaseTrigger):
@@ -501,8 +496,8 @@ def _get_async_hook(self) -> BigQueryTableAsyncHook:
 
  async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
  """Will run until the table exists in the Google Big Query."""
- while True:
- try:
+ try:
+ while True:
  hook = self._get_async_hook()
  response = await self._table_exists(
  hook=hook, dataset=self.dataset_id, table_id=self.table_id, project_id=self.project_id
@@ -511,10 +506,9 @@ async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
  yield TriggerEvent({"status": "success", "message": "success"})
  return
  await asyncio.sleep(self.poll_interval)
- except Exception as e:
- self.log.exception("Exception occurred while checking for Table existence")
- yield TriggerEvent({"status": "error", "message": str(e)})
- return
+ except Exception as e:
+ self.log.exception("Exception occurred while checking for Table existence")
+ yield TriggerEvent({"status": "error", "message": str(e)})
 
  async def _table_exists(
  self, hook: BigQueryTableAsyncHook, dataset: str, table_id: str, project_id: str

diff --git a/airflow/providers/google/cloud/triggers/bigquery_dts.py b/airflow/providers/google/cloud/triggers/bigquery_dts.py
@@ -83,8 +83,8 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
  async def run(self) -> AsyncIterator[TriggerEvent]:
  """If the Transfer Run is in a terminal state, then yield TriggerEvent object."""
  hook = self._get_async_hook()
- while True:
- try:
+ try:
+ while True:
  transfer_run: TransferRun = await hook.get_transfer_run(
  project_id=self.project_id,
  config_id=self.config_id,
@@ -129,14 +129,13 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
  self.log.info("Job is still working...")
  self.log.info("Waiting for %s seconds", self.poll_interval)
  await asyncio.sleep(self.poll_interval)
- except Exception as e:
- yield TriggerEvent(
- {
- "status": "failed",
- "message": f"Trigger failed with exception: {e}",
- }
- )
- return
+ except Exception as e:
+ yield TriggerEvent(
+ {
+ "status": "failed",
+ "message": f"Trigger failed with exception: {e}",
+ }
+ )
 
  def _get_async_hook(self) -> AsyncBiqQueryDataTransferServiceHook:
  return AsyncBiqQueryDataTransferServiceHook(

diff --git a/airflow/providers/google/cloud/triggers/cloud_batch.py b/airflow/providers/google/cloud/triggers/cloud_batch.py
@@ -92,9 +92,8 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
  """
  timeout = self.timeout
  hook = self._get_async_hook()
- while timeout is None or timeout > 0:
-
- try:
+ try:
+ while timeout is None or timeout > 0:
  job: Job = await hook.get_batch_job(job_name=self.job_name)
 
  status: JobStatus.State = job.status.state
@@ -134,10 +133,10 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
  if timeout is None or timeout > 0:
  await asyncio.sleep(self.polling_period_seconds)
 
-  except Exception as e:
-  self.log.exception("Exception occurred while checking for job completion.")
-  yield TriggerEvent({"status": "error", "message": str(e)})
-  return
+ except Exception as e:
+ self.log.exception("Exception occurred while checking for job completion.")
+ yield TriggerEvent({"status": "error", "message": str(e)})
+ return
 
  self.log.exception(f"Job with name [{self.job_name}] timed out")
  yield TriggerEvent(
@@ -147,7 +146,6 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
  "message": f"Batch job with name {self.job_name} timed out",
  }
  )
- return
 
  def _get_async_hook(self) -> CloudBatchAsyncHook:
  return CloudBatchAsyncHook(

diff --git a/airflow/providers/google/cloud/triggers/cloud_build.py b/airflow/providers/google/cloud/triggers/cloud_build.py
@@ -78,8 +78,8 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
  async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
  """Gets current build execution status and yields a TriggerEvent."""
  hook = self._get_async_hook()
- while True:
- try:
+ try:
+ while True:
  # Poll for job execution status
  cloud_build_instance = await hook.get_cloud_build(
  id_=self.id_,
@@ -119,10 +119,9 @@ async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
  )
  return
 
- except Exception as e:
- self.log.exception("Exception occurred while checking for Cloud Build completion")
- yield TriggerEvent({"status": "error", "message": str(e)})
- return
+ except Exception as e:
+ self.log.exception("Exception occurred while checking for Cloud Build completion")
+ yield TriggerEvent({"status": "error", "message": str(e)})
 
  def _get_async_hook(self) -> CloudBuildAsyncHook:
  return CloudBuildAsyncHook(gcp_conn_id=self.gcp_conn_id)
diff --git a/airflow/providers/google/cloud/triggers/cloud_sql.py b/airflow/providers/google/cloud/triggers/cloud_sql.py
@@ -64,8 +64,8 @@ def serialize(self):
  )
 
  async def run(self):
- while True:
- try:
+ try:
+ while True:
  operation = await self.hook.get_operation(
  project_id=self.project_id, operation_name=self.operation_name
  )
@@ -93,11 +93,11 @@ async def run(self):
  self.poke_interval,
  )
  await asyncio.sleep(self.poke_interval)
-  except Exception as e:
-  self.log.exception("Exception occurred while checking operation status.")
-  yield TriggerEvent(
-  {
-  "status": "failed",
-  "message": str(e),
-  }
-  )
+ except Exception as e:
+ self.log.exception("Exception occurred while checking operation status.")
+ yield TriggerEvent(
+ {
+ "status": "failed",
+ "message": str(e),
+ }
+ )
diff --git a/airflow/providers/google/cloud/triggers/dataflow.py b/airflow/providers/google/cloud/triggers/dataflow.py
@@ -92,8 +92,8 @@ async def run(self):
  amount of time stored in self.poll_sleep variable.
  """
  hook = self._get_async_hook()
- while True:
- try:
+ try:
+ while True:
  status = await hook.get_job_status(
  project_id=self.project_id,
  job_id=self.job_id,
@@ -129,10 +129,9 @@ async def run(self):
  self.log.info("Current job status is: %s", status)
  self.log.info("Sleeping for %s seconds.", self.poll_sleep)
  await asyncio.sleep(self.poll_sleep)
- except Exception as e:
- self.log.exception("Exception occurred while checking for job completion.")
- yield TriggerEvent({"status": "error", "message": str(e)})
- return
+ except Exception as e:
+ self.log.exception("Exception occurred while checking for job completion.")
+ yield TriggerEvent({"status": "error", "message": str(e)})
 
  def _get_async_hook(self) -> AsyncDataflowHook:
  return AsyncDataflowHook(

diff --git a/airflow/providers/google/cloud/triggers/datafusion.py b/airflow/providers/google/cloud/triggers/datafusion.py
@@ -83,8 +83,8 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
  async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
  """Gets current pipeline status and yields a TriggerEvent."""
  hook = self._get_async_hook()
- while True:
- try:
+ try:
+ while True:
  # Poll for job execution status
  response_from_hook = await hook.get_pipeline_status(
  success_states=self.success_states,
@@ -109,10 +109,9 @@ async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
  else:
  yield TriggerEvent({"status": "error", "message": response_from_hook})
  return
- except Exception as e:
- self.log.exception("Exception occurred while checking for pipeline state")
- yield TriggerEvent({"status": "error", "message": str(e)})
- return
+ except Exception as e:
+ self.log.exception("Exception occurred while checking for pipeline state")
+ yield TriggerEvent({"status": "error", "message": str(e)})
 
  def _get_async_hook(self) -> DataFusionAsyncHook:
  return DataFusionAsyncHook(

diff --git a/airflow/providers/google/cloud/triggers/dataproc.py b/airflow/providers/google/cloud/triggers/dataproc.py
@@ -263,8 +263,8 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
 
  async def run(self) -> AsyncIterator[TriggerEvent]:
  """Wait until cluster is deleted completely."""
- while self.end_time > time.time():
- try:
+ try:
+ while self.end_time > time.time():
  cluster = await self.get_async_hook().get_cluster(
  region=self.region, # type: ignore[arg-type]
  cluster_name=self.cluster_name,
@@ -277,13 +277,12 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
  self.polling_interval_seconds,
  )
  await asyncio.sleep(self.polling_interval_seconds)
- except NotFound:
- yield TriggerEvent({"status": "success", "message": ""})
- return
- except Exception as e:
- yield TriggerEvent({"status": "error", "message": str(e)})
- return
- yield TriggerEvent({"status": "error", "message": "Timeout"})
+ except NotFound:
+ yield TriggerEvent({"status": "success", "message": ""})
+ except Exception as e:
+ yield TriggerEvent({"status": "error", "message": str(e)})
+ else:
+ yield TriggerEvent({"status": "error", "message": "Timeout"})
 
 
 class DataprocWorkflowTrigger(DataprocBaseTrigger):
@@ -312,8 +311,8 @@ def serialize(self):
 
  async def run(self) -> AsyncIterator[TriggerEvent]:
  hook = self.get_async_hook()
- while True:
- try:
+ try:
+ while True:
  operation = await hook.get_operation(region=self.region, operation_name=self.name)
  if operation.done:
  if operation.error.message:
@@ -338,12 +337,11 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
  else:
  self.log.info("Sleeping for %s seconds.", self.polling_interval_seconds)
  await asyncio.sleep(self.polling_interval_seconds)
- except Exception as e:
- self.log.exception("Exception occurred while checking operation status.")
- yield TriggerEvent(
- {
- "status": "failed",
- "message": str(e),
- }
- )
- return
+ except Exception as e:
+ self.log.exception("Exception occurred while checking operation status.")
+ yield TriggerEvent(
+ {
+ "status": "failed",
+ "message": str(e),
+ }
+ )