Add ability to specify retry for `Operation` and `polling.Future`. (#4922)

2 files changed, 19 insertions, 6 deletions

diff --git a/google/api_core/future/polling.py b/google/api_core/future/polling.py
index 9d46c99..b5aecde 100644
--- a/google/api_core/future/polling.py
+++ b/google/api_core/future/polling.py
@@ -28,6 +28,10 @@ class _OperationNotComplete(Exception):
     pass
 
 
+RETRY_PREDICATE = retry.if_exception_type(_OperationNotComplete)
+DEFAULT_RETRY = retry.Retry(predicate=RETRY_PREDICATE)
+
+
 class PollingFuture(base.Future):
     """A Future that needs to poll some service to check its status.
 
@@ -36,9 +40,16 @@ class PollingFuture(base.Future):
 
     .. note: Privacy here is intended to prevent the final class from
     overexposing, not to prevent subclasses from accessing methods.
+
+    Args:
+        retry (google.api_core.retry.Retry): The retry configuration used
+            when polling. This can be used to control how often :meth:`done`
+            is polled. Regardless of the retry's ``deadline``, it will be
+            overridden by the ``timeout`` argument to :meth:`result`.
     """
-    def __init__(self):
+    def __init__(self, retry=DEFAULT_RETRY):
         super(PollingFuture, self).__init__()
+        self._retry = retry
         self._result = None
         self._exception = None
         self._result_set = False
@@ -77,9 +88,7 @@ class PollingFuture(base.Future):
         if self._result_set:
             return
 
-        retry_ = retry.Retry(
-            predicate=retry.if_exception_type(_OperationNotComplete),
-            deadline=timeout)
+        retry_ = self._retry.with_deadline(timeout)
 
         try:
             retry_(self._done_or_raise)()
diff --git a/google/api_core/operation.py b/google/api_core/operation.py
index e18abee..51a7a96 100644
--- a/google/api_core/operation.py
+++ b/google/api_core/operation.py
@@ -61,12 +61,16 @@ class Operation(polling.PollingFuture):
             result.
         metadata_type (func:`type`): The protobuf type for the operation's
             metadata.
+        retry (google.api_core.retry.Retry): The retry configuration used
+            when polling. This can be used to control how often :meth:`done`
+            is polled. Regardless of the retry's ``deadline``, it will be
+            overridden by the ``timeout`` argument to :meth:`result`.
     """
 
     def __init__(
             self, operation, refresh, cancel,
-            result_type, metadata_type=None):
-        super(Operation, self).__init__()
+            result_type, metadata_type=None, retry=polling.DEFAULT_RETRY):
+        super(Operation, self).__init__(retry=retry)
         self._operation = operation
         self._refresh = refresh
         self._cancel = cancel