feat: add retry/timeout to 'transaction.Transaction.get'

Toward #221

google/cloud/firestore_v1/transaction.py

@@ -169,19 +169,25 @@ def get_all(
         kwargs = self._make_retry_timeout_kwargs(retry, timeout)
         return self._client.get_all(references, transaction=self, **kwargs)
 
-    def get(self, ref_or_query) -> Any:
+    def get(
+        self, ref_or_query, retry: retries.Retry = None, timeout: float = None
+    ) -> Any:
         """
         Retrieve a document or a query result from the database.
         Args:
-            ref_or_query The document references or query object to return.
+            ref_or_query: The document references or query object to return.
+            retry (google.api_core.retry.Retry): Designation of what errors, if any,
+                should be retried.
+            timeout (float): The timeout for this request.
         Yields:
             .DocumentSnapshot: The next document snapshot that fulfills the
             query, or :data:`None` if the document does not exist.
         """
+        kwargs = self._make_retry_timeout_kwargs(retry, timeout)
         if isinstance(ref_or_query, DocumentReference):
-            return self._client.get_all([ref_or_query], transaction=self)
+            return self._client.get_all([ref_or_query], transaction=self, **kwargs)
         elif isinstance(ref_or_query, Query):
-            return ref_or_query.stream(transaction=self)
+            return ref_or_query.stream(transaction=self, **kwargs)
         else:
             raise ValueError(
                 'Value for argument "ref_or_query" must be a DocumentReference or a Query.'

tests/unit/v1/test_transaction.py

@@ -320,26 +320,66 @@ def test_get_all_w_retry_timeout(self):
         timeout = 123.0
         self._get_all_helper(retry=retry, timeout=timeout)
 
-    def test_get_document_ref(self):
+    def _get_document_ref_helper(self, retry=None, timeout=None):
         from google.cloud.firestore_v1.document import DocumentReference
 
         client = mock.Mock(spec=["get_all"])
         transaction = self._make_one(client)
         ref = DocumentReference("documents", "doc-id")
-        result = transaction.get(ref)
-        client.get_all.assert_called_once_with([ref], transaction=transaction)
+
+        kwargs = {}
+
+        if retry is not None:
+            kwargs["retry"] = retry
+
+        if timeout is not None:
+            kwargs["timeout"] = timeout
+
+        result = transaction.get(ref, **kwargs)
+
         self.assertIs(result, client.get_all.return_value)
+        client.get_all.assert_called_once_with([ref], transaction=transaction, **kwargs)
 
-    def test_get_w_query(self):
+    def test_get_document_ref(self):
+        self._get_document_ref_helper()
+
+    def test_get_document_ref_w_retry_timeout(self):
+        from google.api_core.retry import Retry
+
+        retry = Retry(predicate=object())
+        timeout = 123.0
+        self._get_document_ref_helper(retry=retry, timeout=timeout)
+
+    def _get_w_query_helper(self, retry=None, timeout=None):
         from google.cloud.firestore_v1.query import Query
 
         client = mock.Mock(spec=[])
         transaction = self._make_one(client)
         query = Query(parent=mock.Mock(spec=[]))
         query.stream = mock.MagicMock()
-        result = transaction.get(query)
-        query.stream.assert_called_once_with(transaction=transaction)
+
+        kwargs = {}
+
+        if retry is not None:
+            kwargs["retry"] = retry
+
+        if timeout is not None:
+            kwargs["timeout"] = timeout
+
+        result = transaction.get(query, **kwargs)
+
         self.assertIs(result, query.stream.return_value)
+        query.stream.assert_called_once_with(transaction=transaction, **kwargs)
+
+    def test_get_w_query(self):
+        self._get_w_query_helper()
+
+    def test_get_w_query_w_retry_timeout(self):
+        from google.api_core.retry import Retry
+
+        retry = Retry(predicate=object())
+        timeout = 123.0
+        self._get_w_query_helper(retry=retry, timeout=timeout)
 
     def test_get_failure(self):
         client = _make_client()