Skip to content

Commit f3283e1

Browse files
authored
feat: add retry and timeout args to API methods (#67)
Closes #3
1 parent 8bb17ea commit f3283e1

File tree

8 files changed

+543
-80
lines changed

8 files changed

+543
-80
lines changed

google/cloud/datastore/batch.py

Lines changed: 23 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -236,7 +236,7 @@ def begin(self):
236236
raise ValueError("Batch already started previously.")
237237
self._status = self._IN_PROGRESS
238238

239-
def _commit(self):
239+
def _commit(self, retry, timeout):
240240
"""Commits the batch.
241241
242242
This is called by :meth:`commit`.
@@ -246,8 +246,16 @@ def _commit(self):
246246
else:
247247
mode = _datastore_pb2.CommitRequest.TRANSACTIONAL
248248

249+
kwargs = {}
250+
251+
if retry is not None:
252+
kwargs["retry"] = retry
253+
254+
if timeout is not None:
255+
kwargs["timeout"] = timeout
256+
249257
commit_response_pb = self._client._datastore_api.commit(
250-
self.project, mode, self._mutations, transaction=self._id
258+
self.project, mode, self._mutations, transaction=self._id, **kwargs
251259
)
252260
_, updated_keys = _parse_commit_response(commit_response_pb)
253261
# If the back-end returns without error, we are guaranteed that
@@ -257,21 +265,32 @@ def _commit(self):
257265
new_id = new_key_pb.path[-1].id
258266
entity.key = entity.key.completed_key(new_id)
259267

260-
def commit(self):
268+
def commit(self, retry=None, timeout=None):
261269
"""Commits the batch.
262270
263271
This is called automatically upon exiting a with statement,
264272
however it can be called explicitly if you don't want to use a
265273
context manager.
266274
275+
:type retry: :class:`google.api_core.retry.Retry`
276+
:param retry:
277+
A retry object used to retry requests. If ``None`` is specified,
278+
requests will be retried using a default configuration.
279+
280+
:type timeout: float
281+
:param timeout:
282+
Time, in seconds, to wait for the request to complete.
283+
Note that if ``retry`` is specified, the timeout applies
284+
to each individual attempt.
285+
267286
:raises: :class:`~exceptions.ValueError` if the batch is not
268287
in progress.
269288
"""
270289
if self._status != self._IN_PROGRESS:
271290
raise ValueError("Batch must be in progress to commit()")
272291

273292
try:
274-
self._commit()
293+
self._commit(retry=retry, timeout=timeout)
275294
finally:
276295
self._status = self._FINISHED
277296

0 commit comments

Comments
 (0)