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),
+                }
+            )