Make timeout Optional for wait_for_operation (#20981)

apache · Jan 23, 2022 · 86ef016 · 86ef016
1 parent 2af0f70
commit 86ef016
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/airflow/providers/google/cloud/hooks/dataproc.py b/airflow/providers/google/cloud/hooks/dataproc.py
@@ -277,7 +277,7 @@ def get_batch_client(
             credentials=self._get_credentials(), client_info=self.client_info, client_options=client_options
         )
 
-    def wait_for_operation(self, timeout: float, operation: Operation):
+    def wait_for_operation(self, operation: Operation, timeout: Optional[float] = None):
         """Waits for long-lasting operation to complete."""
         try:
             return operation.result(timeout=timeout)

diff --git a/airflow/providers/google/cloud/operators/dataproc.py b/airflow/providers/google/cloud/operators/dataproc.py
@@ -2101,9 +2101,7 @@ def execute(self, context: 'Context'):
                 timeout=self.timeout,
                 metadata=self.metadata,
             )
-            if self.timeout is None:
-                raise AirflowException('Timeout should be set here')
-            result = hook.wait_for_operation(self.timeout, self.operation)
+            result = hook.wait_for_operation(timeout=self.timeout, operation=self.operation)
             self.log.info("Batch %s created", self.batch_id)
         except AlreadyExists:
             self.log.info("Batch with given id already exists")