feat: add retry/timeout to manual surface

google/cloud/firestore_v1/batch.py

@@ -14,6 +14,7 @@
 
 """Helpers for batch requests to the Google Cloud Firestore API."""
 
+from google.api_core import retry as retries  # type: ignore
 
 from google.cloud.firestore_v1.base_batch import BaseWriteBatch
 
@@ -33,22 +34,36 @@ class WriteBatch(BaseWriteBatch):
     def __init__(self, client) -> None:
         super(WriteBatch, self).__init__(client=client)
 
-    def commit(self) -> list:
+    def commit(self, retry: retries.Retry = None, timeout: float = None) -> list:
         """Commit the changes accumulated in this batch.
 
+        Args
+            retry (google.api_core.retry.Retry): Designation of what errors, if any,
+                should be retried.
+            timeout (float): The timeout for this request.
+
         Returns:
             List[:class:`google.cloud.proto.firestore.v1.write.WriteResult`, ...]:
             The write results corresponding to the changes committed, returned
             in the same order as the changes were applied to this batch. A
             write result contains an ``update_time`` field.
         """
+        kwargs = {}
+
+        if retry is not None:
+            kwargs["retry"] = retry
+
+        if timeout is not None:
+            kwargs["timeout"] = timeout
+
         commit_response = self._client._firestore_api.commit(
             request={
                 "database": self._client._database_string,
                 "writes": self._write_pbs,
                 "transaction": None,
             },
             metadata=self._client._rpc_metadata,
+            **kwargs,
         )
 
         self._write_pbs = []

tests/unit/v1/test_batch.py

@@ -35,7 +35,7 @@ def test_constructor(self):
         self.assertIsNone(batch.write_results)
         self.assertIsNone(batch.commit_time)
 
-    def test_commit(self):
+    def _commit_helper(self, retry=None, timeout=None):
         from google.protobuf import timestamp_pb2
         from google.cloud.firestore_v1.types import firestore
         from google.cloud.firestore_v1.types import write
@@ -49,19 +49,27 @@ def test_commit(self):
         )
         firestore_api.commit.return_value = commit_response
 
+        kwargs = {}
+
+        if retry is not None:
+            kwargs["retry"] = retry
+
+        if timeout is not None:
+            kwargs["timeout"] = timeout
+
         # Attach the fake GAPIC to a real client.
         client = _make_client("grand")
         client._firestore_api_internal = firestore_api
 
         # Actually make a batch with some mutations and call commit().
         batch = self._make_one(client)
         document1 = client.document("a", "b")
-        batch.create(document1, {"ten": 10, "buck": u"ets"})
+        batch.create(document1, {"ten": 10, "buck": "ets"})
         document2 = client.document("c", "d", "e", "f")
         batch.delete(document2)
         write_pbs = batch._write_pbs[::]
 
-        write_results = batch.commit()
+        write_results = batch.commit(**kwargs)
         self.assertEqual(write_results, list(commit_response.write_results))
         self.assertEqual(batch.write_results, write_results)
         self.assertEqual(batch.commit_time.timestamp_pb(), timestamp)
@@ -76,8 +84,20 @@ def test_commit(self):
                 "transaction": None,
             },
             metadata=client._rpc_metadata,
+            **kwargs,
         )
 
+    def test_commit(self):
+        self._commit_helper()
+
+    def test_commit_w_retry_timeout(self):
+        from google.api_core.retry import Retry
+
+        retry = Retry(predicate=object())
+        timeout = 123.0
+
+        self._commit_helper(retry=retry, timeout=timeout)
+
     def test_as_context_mgr_wo_error(self):
         from google.protobuf import timestamp_pb2
         from google.cloud.firestore_v1.types import firestore
@@ -98,7 +118,7 @@ def test_as_context_mgr_wo_error(self):
 
         with batch as ctx_mgr:
             self.assertIs(ctx_mgr, batch)
-            ctx_mgr.create(document1, {"ten": 10, "buck": u"ets"})
+            ctx_mgr.create(document1, {"ten": 10, "buck": "ets"})
             ctx_mgr.delete(document2)
             write_pbs = batch._write_pbs[::]
 
@@ -127,7 +147,7 @@ def test_as_context_mgr_w_error(self):
 
         with self.assertRaises(RuntimeError):
             with batch as ctx_mgr:
-                ctx_mgr.create(document1, {"ten": 10, "buck": u"ets"})
+                ctx_mgr.create(document1, {"ten": 10, "buck": "ets"})
                 ctx_mgr.delete(document2)
                 raise RuntimeError("testing")