Skip to content

Commit 86ef016

Browse files
authored
Make timeout Optional for wait_for_operation (#20981)
1 parent 2af0f70 commit 86ef016

File tree

2 files changed

+2
-4
lines changed

2 files changed

+2
-4
lines changed

β€Žairflow/providers/google/cloud/hooks/dataproc.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -277,7 +277,7 @@ def get_batch_client(
277277
credentials=self._get_credentials(), client_info=self.client_info, client_options=client_options
278278
)
279279

280-
def wait_for_operation(self, timeout: float, operation: Operation):
280+
def wait_for_operation(self, operation: Operation, timeout: Optional[float] = None):
281281
"""Waits for long-lasting operation to complete."""
282282
try:
283283
return operation.result(timeout=timeout)

β€Žairflow/providers/google/cloud/operators/dataproc.py

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2101,9 +2101,7 @@ def execute(self, context: 'Context'):
21012101
timeout=self.timeout,
21022102
metadata=self.metadata,
21032103
)
2104-
if self.timeout is None:
2105-
raise AirflowException('Timeout should be set here')
2106-
result = hook.wait_for_operation(self.timeout, self.operation)
2104+
result = hook.wait_for_operation(timeout=self.timeout, operation=self.operation)
21072105
self.log.info("Batch %s created", self.batch_id)
21082106
except AlreadyExists:
21092107
self.log.info("Batch with given id already exists")

0 commit comments

Comments
 (0)