1616
1717import abc
1818import concurrent .futures
19- import functools
20- import operator
21-
22- import six
23- import tenacity
2419
20+ from google .api .core import exceptions
21+ from google .api .core import retry
2522from google .api .core .future import _helpers
2623from google .api .core .future import base
2724
2825
26+ class _OperationNotComplete (Exception ):
27+ """Private exception used for polling via retry."""
28+ pass
29+
30+
2931class PollingFuture (base .Future ):
3032 """A Future that needs to poll some service to check its status.
3133
@@ -55,6 +57,11 @@ def done(self):
5557 # pylint: disable=redundant-returns-doc, missing-raises-doc
5658 raise NotImplementedError ()
5759
60+ def _done_or_raise (self ):
61+ """Check if the future is done and raise if it's not."""
62+ if not self .done ():
63+ raise _OperationNotComplete ()
64+
5865 def running (self ):
5966 """True if the operation is currently running."""
6067 return not self .done ()
@@ -69,29 +76,16 @@ def _blocking_poll(self, timeout=None):
6976 if self ._result_set :
7077 return
7178
72- retry_on = tenacity .retry_if_result (
73- functools .partial (operator .is_not , True ))
74- # Use exponential backoff with jitter.
75- wait_on = (
76- tenacity .wait_exponential (multiplier = 1 , max = 10 ) +
77- tenacity .wait_random (0 , 1 ))
78-
79- if timeout is None :
80- retry = tenacity .retry (retry = retry_on , wait = wait_on )
81- else :
82- retry = tenacity .retry (
83- retry = retry_on ,
84- wait = wait_on ,
85- stop = tenacity .stop_after_delay (timeout ))
79+ retry_ = retry .Retry (
80+ predicate = retry .if_exception_type (_OperationNotComplete ),
81+ deadline = timeout )
8682
8783 try :
88- retry (self .done )()
89- except tenacity .RetryError as exc :
90- six .raise_from (
91- concurrent .futures .TimeoutError (
92- 'Operation did not complete within the designated '
93- 'timeout.' ),
94- exc )
84+ retry_ (self ._done_or_raise )()
85+ except exceptions .RetryError :
86+ raise concurrent .futures .TimeoutError (
87+ 'Operation did not complete within the designated '
88+ 'timeout.' )
9589
9690 def result (self , timeout = None ):
9791 """Get the result of the operation, blocking if necessary.
0 commit comments