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

Toward #3

google/cloud/datastore/transaction.py

@@ -18,6 +18,19 @@
 from google.cloud.datastore_v1.types import TransactionOptions
 
 
+def _make_retry_timeout_kwargs(retry, timeout):
+    """Helper: make optional retry / timeout kwargs dict."""
+    kwargs = {}
+
+    if retry is not None:
+        kwargs["retry"] = retry
+
+    if timeout is not None:
+        kwargs["timeout"] = timeout
+
+    return kwargs
+
+
 class Transaction(Batch):
     """An abstraction representing datastore Transactions.
 
@@ -193,19 +206,35 @@ def current(self):
         if isinstance(top, Transaction):
             return top
 
-    def begin(self):
+    def begin(self, retry=None, timeout=None):
         """Begins a transaction.
 
         This method is called automatically when entering 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 transaction has
                  already begun.
         """
         super(Transaction, self).begin()
+
+        kwargs = _make_retry_timeout_kwargs(retry, timeout)
+
         try:
-            response_pb = self._client._datastore_api.begin_transaction(self.project)
+            response_pb = self._client._datastore_api.begin_transaction(
+                self.project, **kwargs
+            )
             self._id = response_pb.transaction
         except:  # noqa: E722 do not use bare except, specify exception instead
             self._status = self._ABORTED

tests/unit/test_transaction.py

@@ -94,6 +94,23 @@ def test_begin(self):
         self.assertEqual(xact.id, id_)
         ds_api.begin_transaction.assert_called_once_with(project)
 
+    def test_begin_w_retry_w_timeout(self):
+        project = "PROJECT"
+        id_ = 889
+        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(retry=retry, timeout=timeout)
+
+        self.assertEqual(xact.id, id_)
+        ds_api.begin_transaction.assert_called_once_with(
+            project, retry=retry, timeout=timeout
+        )
+
     def test_begin_tombstoned(self):
         project = "PROJECT"
         id_ = 1094