Skip to content

Commit

Permalink
Browse files Browse the repository at this point in the history
feat: add retry and timeout args to API methods (#67)
Closes #3
  • Loading branch information
tseaver committed Aug 12, 2020
1 parent 8bb17ea commit f3283e1
Show file tree
Hide file tree
Showing 8 changed files with 543 additions and 80 deletions.
27 changes: 23 additions & 4 deletions google/cloud/datastore/batch.py
Expand Up @@ -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`.
Expand All @@ -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
Expand All @@ -257,21 +265,32 @@ 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.
"""
if self._status != self._IN_PROGRESS:
raise ValueError("Batch must be in progress to commit()")

try:
self._commit()
self._commit(retry=retry, timeout=timeout)
finally:
self._status = self._FINISHED

Expand Down

0 comments on commit f3283e1

Please sign in to comment.