Refactor: reduce some conditions in providers (#34440)

apache · Sep 19, 2023 · d20c32f · d20c32f
1 parent 26f6a51
commit d20c32f
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 26 deletions.
diff --git a/airflow/providers/google/cloud/triggers/dataproc.py b/airflow/providers/google/cloud/triggers/dataproc.py
@@ -98,11 +98,10 @@ async def run(self):
             )
             state = job.status.state
             self.log.info("Dataproc job: %s is in state: %s", self.job_id, state)
-            if state in (JobStatus.State.ERROR, JobStatus.State.DONE, JobStatus.State.CANCELLED):
-                if state in (JobStatus.State.DONE, JobStatus.State.CANCELLED):
-                    break
-                elif state == JobStatus.State.ERROR:
-                    raise AirflowException(f"Dataproc job execution failed {self.job_id}")
+            if state in (JobStatus.State.DONE, JobStatus.State.CANCELLED):
+                break
+            elif state == JobStatus.State.ERROR:
+                raise AirflowException(f"Dataproc job execution failed {self.job_id}")
             await asyncio.sleep(self.polling_interval_seconds)
         yield TriggerEvent({"job_id": self.job_id, "job_state": state})
 
@@ -316,21 +315,17 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
                 operation = await hook.get_operation(region=self.region, operation_name=self.name)
                 if operation.done:
                     if operation.error.message:
-                        yield TriggerEvent(
-                            {
-                                "operation_name": operation.name,
-                                "operation_done": operation.done,
-                                "status": "error",
-                                "message": operation.error.message,
-                            }
-                        )
-                        return
+                        status = "error"
+                        message = operation.error.message
+                    else:
+                        status = "success"
+                        message = "Operation is successfully ended."
                     yield TriggerEvent(
                         {
                             "operation_name": operation.name,
                             "operation_done": operation.done,
-                            "status": "success",
-                            "message": "Operation is successfully ended.",
+                            "status": status,
+                            "message": message,
                         }
                     )
                     return

diff --git a/airflow/providers/microsoft/azure/hooks/cosmos.py b/airflow/providers/microsoft/azure/hooks/cosmos.py
@@ -281,10 +281,7 @@ def upsert_document(self, document, database_name=None, collection_name=None, do
             raise AirflowBadRequest("You cannot insert a None document")
 
         # Add document id if isn't found
-        if "id" in document:
-            if document["id"] is None:
-                document["id"] = document_id
-        else:
+        if document.get("id") is None:
             document["id"] = document_id
 
         created_document = (

diff --git a/airflow/providers/singularity/operators/singularity.py b/airflow/providers/singularity/operators/singularity.py
@@ -155,9 +155,8 @@ def execute(self, context: Context) -> None:
         self.log.info("Stopping instance %s", self.instance)
         self.instance.stop()  # type: ignore[attr-defined]
 
-        if self.auto_remove is True:
-            if self.auto_remove and os.path.exists(self.image):
-                shutil.rmtree(self.image)
+        if self.auto_remove and os.path.exists(self.image):
+            shutil.rmtree(self.image)
 
         # If the container failed, raise the exception
         if result["return_code"] != 0:
@@ -179,6 +178,5 @@ def on_kill(self) -> None:
             self.instance.stop()
 
             # If an image exists, clean it up
-            if self.auto_remove is True:
-                if self.auto_remove and os.path.exists(self.image):
-                    shutil.rmtree(self.image)
+            if self.auto_remove and os.path.exists(self.image):
+                shutil.rmtree(self.image)