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

Toward #3

google/cloud/datastore/transaction.py

@@ -240,16 +240,29 @@ def begin(self, retry=None, timeout=None):
             self._status = self._ABORTED
             raise
 
-    def rollback(self):
+    def rollback(self, retry=None, timeout=None):
         """Rolls back the current transaction.
 
         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:
             # No need to use the response it contains nothing.
-            self._client._datastore_api.rollback(self.project, self._id)
+            self._client._datastore_api.rollback(self.project, self._id, **kwargs)
         finally:
             super(Transaction, self).rollback()
             # Clear our own ID in case this gets accidentally reused.

tests/unit/test_transaction.py

@@ -148,10 +148,29 @@ def test_rollback(self):
         client = _Client(project, datastore_api=ds_api)
         xact = self._make_one(client)
         xact.begin()
+
         xact.rollback()
-        client._datastore_api.rollback.assert_called_once_with(project, id_)
+
         self.assertIsNone(xact.id)
-        ds_api.begin_transaction.assert_called_once_with(project)
+        ds_api.rollback.assert_called_once_with(project, id_)
+
+    def test_rollback_w_retry_w_timeout(self):
+        project = "PROJECT"
+        id_ = 239
+        retry = mock.Mock()
+        timeout = 100000
+
+        ds_api = _make_datastore_api(xact_id=id_)
+        client = _Client(project, datastore_api=ds_api)
+        xact = self._make_one(client)
+        xact.begin()
+
+        xact.rollback(retry=retry, timeout=timeout)
+
+        self.assertIsNone(xact.id)
+        ds_api.rollback.assert_called_once_with(
+            project, id_, retry=retry, timeout=timeout
+        )
 
     def test_commit_no_partial_keys(self):
         from google.cloud.datastore_v1.proto import datastore_pb2