Skip to content

Commit

Permalink
Add 'retry' argument to '_AsyncJob.result'. (#6302)
Browse files Browse the repository at this point in the history
Pass it through to the '_begin' call.  Note that we need to modify
the 'api_core...PollingFuture' class before we can safely pass the
'retry' through to its 'result'.
  • Loading branch information
tseaver authored Dec 1, 2018
1 parent d0a3a7e commit be78a96
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 3 deletions.
7 changes: 5 additions & 2 deletions bigquery/google/cloud/bigquery/job.py
Original file line number Diff line number Diff line change
Expand Up @@ -678,14 +678,17 @@ def done(self, retry=DEFAULT_RETRY):
self.reload(retry=retry)
return self.state == _DONE_STATE

def result(self, timeout=None):
def result(self, timeout=None, retry=DEFAULT_RETRY):
"""Start the job and wait for it to complete and get the result.
:type timeout: float
:param timeout:
How long (in seconds) to wait for job to complete before raising
a :class:`concurrent.futures.TimeoutError`.
:type retry: :class:`google.api_core.retry.Retry`
:param retry: (Optional) How to retry the RPC.
:rtype: _AsyncJob
:returns: This instance.
Expand All @@ -695,7 +698,7 @@ def result(self, timeout=None):
not complete in the given timeout.
"""
if self.state is None:
self._begin()
self._begin(retry=retry)
# TODO: modify PollingFuture so it can pass a retry argument to done().
return super(_AsyncJob, self).result(timeout=timeout)

Expand Down
16 changes: 15 additions & 1 deletion bigquery/tests/unit/test_job.py
Original file line number Diff line number Diff line change
Expand Up @@ -789,13 +789,27 @@ def test_done_already(self):

@mock.patch("google.api_core.future.polling.PollingFuture.result")
def test_result_default_wo_state(self, result):
from google.cloud.bigquery.retry import DEFAULT_RETRY

client = _make_client(project=self.PROJECT)
job = self._make_one(self.JOB_ID, client)
begin = job._begin = mock.Mock()

self.assertIs(job.result(), result.return_value)

begin.assert_called_once()
begin.assert_called_once_with(retry=DEFAULT_RETRY)
result.assert_called_once_with(timeout=None)

@mock.patch('google.api_core.future.polling.PollingFuture.result')
def test_result_w_retry_wo_state(self, result):
client = _make_client(project=self.PROJECT)
job = self._make_one(self.JOB_ID, client)
begin = job._begin = mock.Mock()
retry = mock.Mock()

self.assertIs(job.result(retry=retry), result.return_value)

begin.assert_called_once_with(retry=retry)
result.assert_called_once_with(timeout=None)

@mock.patch("google.api_core.future.polling.PollingFuture.result")
Expand Down

0 comments on commit be78a96

Please sign in to comment.