From edaaa45213b20a6684e0bc67395675b1f60b9b3f Mon Sep 17 00:00:00 2001
From: Tres Seaver <tseaver@palladion.com>
Date: Tue, 11 Aug 2020 15:47:20 -0400
Subject: [PATCH] feat: add 'retry'/'timeout' args to 'Batch.commit'

Toward #3
---
 google/cloud/datastore/batch.py | 27 ++++++++++++++++++----
 tests/unit/test_batch.py        | 40 +++++++++++++++++++++++++++++++++
 2 files changed, 63 insertions(+), 4 deletions(-)

diff --git a/google/cloud/datastore/batch.py b/google/cloud/datastore/batch.py
index dc8463f3..294c1b45 100644
--- a/google/cloud/datastore/batch.py
+++ b/google/cloud/datastore/batch.py
@@ -236,7 +236,7 @@ def begin(self):
             raise ValueError("Batch already started previously.")
         self._status = self._IN_PROGRESS
 
-    def _commit(self):
+    def _commit(self, retry, timeout):
         """Commits the batch.
 
         This is called by :meth:`commit`.
@@ -246,8 +246,16 @@ def _commit(self):
         else:
             mode = _datastore_pb2.CommitRequest.TRANSACTIONAL
 
+        kwargs = {}
+
+        if retry is not None:
+            kwargs["retry"] = retry
+
+        if timeout is not None:
+            kwargs["timeout"] = timeout
+
         commit_response_pb = self._client._datastore_api.commit(
-            self.project, mode, self._mutations, transaction=self._id
+            self.project, mode, self._mutations, transaction=self._id, **kwargs
         )
         _, updated_keys = _parse_commit_response(commit_response_pb)
         # If the back-end returns without error, we are guaranteed that
@@ -257,13 +265,24 @@ def _commit(self):
             new_id = new_key_pb.path[-1].id
             entity.key = entity.key.completed_key(new_id)
 
-    def commit(self):
+    def commit(self, retry=None, timeout=None):
         """Commits the batch.
 
         This is called automatically upon exiting a with statement,
         however it can be called explicitly if you don't want to use a
         context manager.
 
+        :type retry: :class:`google.api_core.retry.Retry`
+        :param retry:
+            A retry object used to retry requests. If ``None`` is specified,
+            requests will be retried using a default configuration.
+
+        :type timeout: float
+        :param timeout:
+            Time, in seconds, to wait for the request to complete.
+            Note that if ``retry`` is specified, the timeout applies
+            to each individual attempt.
+
         :raises: :class:`~exceptions.ValueError` if the batch is not
                  in progress.
         """
@@ -271,7 +290,7 @@ def commit(self):
             raise ValueError("Batch must be in progress to commit()")
 
         try:
-            self._commit()
+            self._commit(retry=retry, timeout=timeout)
         finally:
             self._status = self._FINISHED
 
diff --git a/tests/unit/test_batch.py b/tests/unit/test_batch.py
index 8516e78c..7ad2aeab 100644
--- a/tests/unit/test_batch.py
+++ b/tests/unit/test_batch.py
@@ -229,6 +229,46 @@ def test_commit(self):
         mode = datastore_pb2.CommitRequest.NON_TRANSACTIONAL
         commit_method.assert_called_with(project, mode, [], transaction=None)
 
+    def test_commit_w_timeout(self):
+        from google.cloud.datastore_v1.proto import datastore_pb2
+
+        project = "PROJECT"
+        client = _Client(project)
+        batch = self._make_one(client)
+        timeout = 100000
+
+        self.assertEqual(batch._status, batch._INITIAL)
+        batch.begin()
+        self.assertEqual(batch._status, batch._IN_PROGRESS)
+        batch.commit(timeout=timeout)
+        self.assertEqual(batch._status, batch._FINISHED)
+
+        commit_method = client._datastore_api.commit
+        mode = datastore_pb2.CommitRequest.NON_TRANSACTIONAL
+        commit_method.assert_called_with(
+            project, mode, [], transaction=None, timeout=timeout
+        )
+
+    def test_commit_w_retry(self):
+        from google.cloud.datastore_v1.proto import datastore_pb2
+
+        project = "PROJECT"
+        client = _Client(project)
+        batch = self._make_one(client)
+        retry = mock.Mock()
+
+        self.assertEqual(batch._status, batch._INITIAL)
+        batch.begin()
+        self.assertEqual(batch._status, batch._IN_PROGRESS)
+        batch.commit(retry=retry)
+        self.assertEqual(batch._status, batch._FINISHED)
+
+        commit_method = client._datastore_api.commit
+        mode = datastore_pb2.CommitRequest.NON_TRANSACTIONAL
+        commit_method.assert_called_with(
+            project, mode, [], transaction=None, retry=retry
+        )
+
     def test_commit_wrong_status(self):
         project = "PROJECT"
         client = _Client(project)