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

Toward #221

google/cloud/firestore_v1/document.py

@@ -137,7 +137,13 @@ def set(
         write_results = batch.commit(**kwargs)
         return _first_write_result(write_results)
 
-    def update(self, field_updates: dict, option: _helpers.WriteOption = None) -> Any:
+    def update(
+        self,
+        field_updates: dict,
+        option: _helpers.WriteOption = None,
+        retry: retries.Retry = None,
+        timeout: float = None,
+    ) -> Any:
         """Update an existing document in the Firestore database.
 
         By default, this method verifies that the document exists on the
@@ -271,6 +277,9 @@ def update(self, field_updates: dict, option: _helpers.WriteOption = None) -> An
             option (Optional[:class:`~google.cloud.firestore_v1.client.WriteOption`]):
                 A write option to make assertions / preconditions on the server
                 state of the document before applying changes.
+            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.types.WriteResult`:
@@ -282,7 +291,8 @@ def update(self, field_updates: dict, option: _helpers.WriteOption = None) -> An
         """
         batch = self._client.batch()
         batch.update(self, field_updates, option=option)
-        write_results = batch.commit()
+        kwargs = self._make_retry_timeout_kwargs(retry, timeout)
+        write_results = batch.commit(**kwargs)
         return _first_write_result(write_results)
 
     def delete(

tests/unit/v1/test_document.py

@@ -233,7 +233,7 @@ def _write_pb_for_update(document_path, update_values, field_paths):
             current_document=common.Precondition(exists=True),
         )
 
-    def _update_helper(self, **option_kwargs):
+    def _update_helper(self, retry=None, timeout=None, **option_kwargs):
         from google.cloud.firestore_v1.transforms import DELETE_FIELD
 
         # Create a minimal fake GAPIC with a dummy response.
@@ -250,12 +250,21 @@ def _update_helper(self, **option_kwargs):
         field_updates = collections.OrderedDict(
             (("hello", 1), ("then.do", False), ("goodbye", DELETE_FIELD))
         )
+
+        kwargs = {}
+
+        if retry is not None:
+            kwargs["retry"] = retry
+
+        if timeout is not None:
+            kwargs["timeout"] = timeout
+
         if option_kwargs:
             option = client.write_option(**option_kwargs)
-            write_result = document.update(field_updates, option=option)
+            write_result = document.update(field_updates, option=option, **kwargs)
         else:
             option = None
-            write_result = document.update(field_updates)
+            write_result = document.update(field_updates, **kwargs)
 
         # Verify the response and the mocks.
         self.assertIs(write_result, mock.sentinel.write_result)
@@ -276,6 +285,7 @@ def _update_helper(self, **option_kwargs):
                 "transaction": None,
             },
             metadata=client._rpc_metadata,
+            **kwargs,
         )
 
     def test_update_with_exists(self):
@@ -285,6 +295,13 @@ def test_update_with_exists(self):
     def test_update(self):
         self._update_helper()
 
+    def test_update_w_retry_timeout(self):
+        from google.api_core.retry import Retry
+
+        retry = Retry(predicate=object())
+        timeout = 123.0
+        self._update_helper(retry=retry, timeout=timeout)
+
     def test_update_with_precondition(self):
         from google.protobuf import timestamp_pb2