feat: add retry/timeout to 'collection.CollectionReference.add'

Toward #221

google/cloud/firestore_v1/collection.py

@@ -13,6 +13,9 @@
 # limitations under the License.
 
 """Classes for representing collections for the Google Cloud Firestore API."""
+
+from google.api_core import retry as retries  # type: ignore
+
 from google.cloud.firestore_v1.base_collection import (
     BaseCollectionReference,
     _auto_id,
@@ -64,7 +67,25 @@ def _query(self) -> query_mod.Query:
         """
         return query_mod.Query(self)
 
-    def add(self, document_data: dict, document_id: str = None) -> Tuple[Any, Any]:
+    @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 add(
+        self,
+        document_data: dict,
+        document_id: str = None,
+        retry: retries.Retry = None,
+        timeout: float = None,
+    ) -> Tuple[Any, Any]:
         """Create a document in the Firestore database with the provided data.
 
         Args:
@@ -75,6 +96,9 @@ def add(self, document_data: dict, document_id: str = None) -> Tuple[Any, Any]:
                 automatically assigned by the server (the assigned ID will be
                 a random 20 character string composed of digits,
                 uppercase and lowercase letters).
+            retry (google.api_core.retry.Retry): Designation of what errors, if any,
+                should be retried.
+            timeout (float): The timeout for this request.
 
         Returns:
             Tuple[:class:`google.protobuf.timestamp_pb2.Timestamp`, \
@@ -92,7 +116,8 @@ def add(self, document_data: dict, document_id: str = None) -> Tuple[Any, Any]:
             document_id = _auto_id()
 
         document_ref = self.document(document_id)
-        write_result = document_ref.create(document_data)
+        kwargs = self._make_retry_timeout_kwargs(retry, timeout)
+        write_result = document_ref.create(document_data, **kwargs)
         return write_result.update_time, document_ref
 
     def list_documents(self, page_size: int = None) -> Generator[Any, Any, None]:

tests/unit/v1/test_collection.py

@@ -138,7 +138,7 @@ def _write_pb_for_create(document_path, document_data):
             current_document=common.Precondition(exists=False),
         )
 
-    def test_add_explicit_id(self):
+    def _add_helper(self, retry=None, timeout=None):
         from google.cloud.firestore_v1.document import DocumentReference
 
         # Create a minimal fake GAPIC with a dummy response.
@@ -161,7 +161,18 @@ def test_add_explicit_id(self):
         collection = self._make_one("parent", client=client)
         document_data = {"zorp": 208.75, "i-did-not": b"know that"}
         doc_id = "child"
-        update_time, document_ref = collection.add(document_data, document_id=doc_id)
+
+        kwargs = {}
+
+        if retry is not None:
+            kwargs["retry"] = retry
+
+        if timeout is not None:
+            kwargs["timeout"] = timeout
+
+        update_time, document_ref = collection.add(
+            document_data, document_id=doc_id, **kwargs
+        )
 
         # Verify the response and the mocks.
         self.assertIs(update_time, mock.sentinel.update_time)
@@ -177,8 +188,19 @@ def test_add_explicit_id(self):
                 "transaction": None,
             },
             metadata=client._rpc_metadata,
+            **kwargs,
         )
 
+    def test_add_explicit_id(self):
+        self._add_helper()
+
+    def test_add_w_retry_timeout(self):
+        from google.api_core.retry import Retry
+
+        retry = Retry(predicate=object())
+        timeout = 123.0
+        self._add_helper(retry=retry, timeout=timeout)
+
     def _list_documents_helper(self, page_size=None):
         from google.api_core.page_iterator import Iterator
         from google.api_core.page_iterator import Page