feat: add 'retry'/'timeout' args to 'Transaction.commit'

Toward #3

google/cloud/datastore/transaction.py

@@ -268,7 +268,7 @@ def rollback(self, retry=None, timeout=None):
             # Clear our own ID in case this gets accidentally reused.
             self._id = None
 
-    def commit(self):
+    def commit(self, retry=None, timeout=None):
         """Commits the transaction.
 
         This is called automatically upon exiting a with statement,
@@ -278,9 +278,22 @@ def commit(self):
         This method has necessary side-effects:
 
         - Sets the current transaction's ID to None.
+
+        :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.
         """
+        kwargs = _make_retry_timeout_kwargs(retry, timeout)
+
         try:
-            super(Transaction, self).commit()
+            super(Transaction, self).commit(**kwargs)
         finally:
             # Clear our own ID in case this gets accidentally reused.
             self._id = None

tests/unit/test_transaction.py

@@ -177,42 +177,48 @@ def test_commit_no_partial_keys(self):
 
         project = "PROJECT"
         id_ = 1002930
+        mode = datastore_pb2.CommitRequest.TRANSACTIONAL
+
         ds_api = _make_datastore_api(xact_id=id_)
         client = _Client(project, datastore_api=ds_api)
         xact = self._make_one(client)
         xact.begin()
         xact.commit()
 
-        mode = datastore_pb2.CommitRequest.TRANSACTIONAL
-        client._datastore_api.commit.assert_called_once_with(
-            project, mode, [], transaction=id_
-        )
+        ds_api.commit.assert_called_once_with(project, mode, [], transaction=id_)
         self.assertIsNone(xact.id)
-        ds_api.begin_transaction.assert_called_once_with(project)
 
-    def test_commit_w_partial_keys(self):
+    def test_commit_w_partial_keys_w_retry_w_timeout(self):
         from google.cloud.datastore_v1.proto import datastore_pb2
 
         project = "PROJECT"
         kind = "KIND"
         id1 = 123
+        mode = datastore_pb2.CommitRequest.TRANSACTIONAL
         key = _make_key(kind, id1, project)
         id2 = 234
+        retry = mock.Mock()
+        timeout = 100000
+
         ds_api = _make_datastore_api(key, xact_id=id2)
         client = _Client(project, datastore_api=ds_api)
         xact = self._make_one(client)
         xact.begin()
         entity = _Entity()
+
         xact.put(entity)
-        xact.commit()
+        xact.commit(retry=retry, timeout=timeout)
 
-        mode = datastore_pb2.CommitRequest.TRANSACTIONAL
         ds_api.commit.assert_called_once_with(
-            project, mode, xact.mutations, transaction=id2
+            project,
+            mode,
+            xact.mutations,
+            transaction=id2,
+            retry=retry,
+            timeout=timeout,
         )
         self.assertIsNone(xact.id)
         self.assertEqual(entity.key.path, [{"kind": kind, "id": id1}])
-        ds_api.begin_transaction.assert_called_once_with(project)
 
     def test_context_manager_no_raise(self):
         from google.cloud.datastore_v1.proto import datastore_pb2
@@ -274,6 +280,7 @@ def test_put_read_only(self):
         entity = _Entity()
         xact = self._make_one(client, read_only=True)
         xact.begin()
+
         with self.assertRaises(RuntimeError):
             xact.put(entity)