feat: add retry/timeout to 'document.Document.get'

Toward #221

google/cloud/firestore_v1/document.py

@@ -14,6 +14,8 @@
 
 """Classes for representing documents for the Google Cloud Firestore API."""
 
+from google.api_core import retry as retries  # type: ignore
+
 from google.cloud.firestore_v1.base_document import (
     BaseDocumentReference,
     DocumentSnapshot,
@@ -55,6 +57,18 @@ class DocumentReference(BaseDocumentReference):
     def __init__(self, *path, **kwargs) -> None:
         super(DocumentReference, self).__init__(*path, **kwargs)
 
+    @staticmethod
+    def _make_retry_timeout_kwargs(retry, timeout):
+        kwargs = {}
+
+        if retry is not None:
+            kwargs["retry"] = retry
+
+        if timeout is not None:
+            kwargs["timeout"] = timeout
+
+        return kwargs
+
     def create(self, document_data) -> Any:
         """Create the current document in the Firestore database.
 
@@ -283,7 +297,11 @@ def delete(self, option: _helpers.WriteOption = None) -> Any:
         return commit_response.commit_time
 
     def get(
-        self, field_paths: Iterable[str] = None, transaction=None
+        self,
+        field_paths: Iterable[str] = None,
+        transaction=None,
+        retry: retries.Retry = None,
+        timeout: float = None,
     ) -> DocumentSnapshot:
         """Retrieve a snapshot of the current document.
 
@@ -302,6 +320,9 @@ def get(
             transaction (Optional[:class:`~google.cloud.firestore_v1.transaction.Transaction`]):
                 An existing transaction that this reference
                 will be retrieved in.
+            retry (google.api_core.retry.Retry): Designation of what errors, if any,
+                should be retried.
+            timeout (float): The timeout for this request.
 
         Returns:
             :class:`~google.cloud.firestore_v1.base_document.DocumentSnapshot`:
@@ -318,6 +339,7 @@ def get(
             mask = common.DocumentMask(field_paths=sorted(field_paths))
         else:
             mask = None
+        kwargs = self._make_retry_timeout_kwargs(retry, timeout)
 
         firestore_api = self._client._firestore_api
         try:
@@ -328,6 +350,7 @@ def get(
                     "transaction": _helpers.get_transaction_id(transaction),
                 },
                 metadata=self._client._rpc_metadata,
+                **kwargs,
             )
         except exceptions.NotFound:
             data = None

tests/unit/v1/test_document.py

@@ -313,7 +313,14 @@ def test_delete_with_option(self):
         timestamp_pb = timestamp_pb2.Timestamp(seconds=1058655101, nanos=100022244)
         self._delete_helper(last_update_time=timestamp_pb)
 
-    def _get_helper(self, field_paths=None, use_transaction=False, not_found=False):
+    def _get_helper(
+        self,
+        field_paths=None,
+        use_transaction=False,
+        not_found=False,
+        retry=None,
+        timeout=None,
+    ):
         from google.api_core.exceptions import NotFound
         from google.cloud.firestore_v1.types import common
         from google.cloud.firestore_v1.types import document
@@ -344,7 +351,17 @@ def _get_helper(self, field_paths=None, use_transaction=False, not_found=False):
         else:
             transaction = None
 
-        snapshot = document.get(field_paths=field_paths, transaction=transaction)
+        kwargs = {}
+
+        if retry is not None:
+            kwargs["retry"] = retry
+
+        if timeout is not None:
+            kwargs["timeout"] = timeout
+
+        snapshot = document.get(
+            field_paths=field_paths, transaction=transaction, **kwargs
+        )
 
         self.assertIs(snapshot.reference, document)
         if not_found:
@@ -378,6 +395,7 @@ def _get_helper(self, field_paths=None, use_transaction=False, not_found=False):
                 "transaction": expected_transaction_id,
             },
             metadata=client._rpc_metadata,
+            **kwargs
         )
 
     def test_get_not_found(self):
@@ -386,6 +404,13 @@ def test_get_not_found(self):
     def test_get_default(self):
         self._get_helper()
 
+    def test_get_w_retry_timeout(self):
+        from google.api_core.retry import Retry
+
+        retry = Retry(predicate=object())
+        timeout = 123.0
+        self._get_helper(retry=retry, timeout=timeout)
+
     def test_get_w_string_field_path(self):
         with self.assertRaises(ValueError):
             self._get_helper(field_paths="foo")