From 48903bb31962f94bbef3b330b9b094c2fccc034e Mon Sep 17 00:00:00 2001 From: Tres Seaver Date: Thu, 25 Oct 2018 12:11:23 -0400 Subject: [PATCH] Add 'retry' argument to '_AsyncJob.result'. 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'. --- bigquery/google/cloud/bigquery/job.py | 7 +++++-- bigquery/tests/unit/test_job.py | 16 +++++++++++++++- 2 files changed, 20 insertions(+), 3 deletions(-) diff --git a/bigquery/google/cloud/bigquery/job.py b/bigquery/google/cloud/bigquery/job.py index 5eec77de7e3a..8be147f9ab64 100644 --- a/bigquery/google/cloud/bigquery/job.py +++ b/bigquery/google/cloud/bigquery/job.py @@ -677,7 +677,7 @@ 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 @@ -685,6 +685,9 @@ def result(self, timeout=None): 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. @@ -694,7 +697,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) diff --git a/bigquery/tests/unit/test_job.py b/bigquery/tests/unit/test_job.py index edddebc05e31..48f5a7507c49 100644 --- a/bigquery/tests/unit/test_job.py +++ b/bigquery/tests/unit/test_job.py @@ -850,13 +850,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')