File tree

3 files changed

+104
-7
lines changed

3 files changed

+104
-7
lines changed
Original file line numberDiff line numberDiff line change
@@ -989,25 +989,43 @@ def done(self, retry=DEFAULT_RETRY, timeout=None, reload=True):
989989
unfinished jobs before checking. Default ``True``.
990990
991991
Returns:
992-
bool: True if the job is complete, False otherwise.
992+
bool: ``True`` if the job is complete or if fetching its status resulted in
993+
an error, ``False`` otherwise.
993994
"""
994995
# Do not refresh if the state is already done, as the job will not
995996
# change once complete.
996997
is_done = self.state == _DONE_STATE
997998
if not reload or is_done:
998999
return is_done
9991000

1000-
self._reload_query_results(retry=retry, timeout=timeout)
1001-
10021001
# If an explicit timeout is not given, fall back to the transport timeout
10031002
# stored in _blocking_poll() in the process of polling for job completion.
10041003
transport_timeout = timeout if timeout is not None else self._transport_timeout
10051004

1005+
try:
1006+
self._reload_query_results(retry=retry, timeout=transport_timeout)
1007+
except exceptions.GoogleAPIError as exc:
1008+
# Reloading also updates error details on self, thus no need for an
1009+
# explicit self.set_exception() call if reloading succeeds.
1010+
try:
1011+
self.reload(retry=retry, timeout=transport_timeout)
1012+
except exceptions.GoogleAPIError:
1013+
# Use the query results reload exception, as it generally contains
1014+
# much more useful error information.
1015+
self.set_exception(exc)
1016+
return True
1017+
else:
1018+
return self.state == _DONE_STATE
1019+
10061020
# Only reload the job once we know the query is complete.
10071021
# This will ensure that fields such as the destination table are
10081022
# correctly populated.
10091023
if self._query_results.complete:
1010-
self.reload(retry=retry, timeout=transport_timeout)
1024+
try:
1025+
self.reload(retry=retry, timeout=transport_timeout)
1026+
except exceptions.GoogleAPIError as exc:
1027+
self.set_exception(exc)
1028+
return True
10111029

10121030
return self.state == _DONE_STATE
10131031

Original file line numberDiff line numberDiff line change
@@ -967,7 +967,7 @@ def test_result_w_retry_wo_state(self):
967967
custom_predicate = mock.Mock()
968968
custom_predicate.return_value = True
969969
custom_retry = google.api_core.retry.Retry(
970-
predicate=custom_predicate, initial=0.001, maximum=0.001, deadline=0.001,
970+
predicate=custom_predicate, initial=0.001, maximum=0.001, deadline=0.1,
971971
)
972972
self.assertIs(job.result(retry=custom_retry), job)
973973

Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
import copy
1717
import http
1818
import textwrap
19+
import types
1920

2021
import freezegun
2122
from google.api_core import exceptions
@@ -308,7 +309,7 @@ def test_cancelled(self):
308309

309310
self.assertTrue(job.cancelled())
310311

311-
def test_done(self):
312+
def test_done_job_complete(self):
312313
client = _make_client(project=self.PROJECT)
313314
resource = self._make_resource(ended=True)
314315
job = self._get_target_class().from_api_repr(resource, client)
@@ -356,6 +357,84 @@ def test_done_w_timeout_and_longer_internal_api_timeout(self):
356357
call_args = fake_reload.call_args
357358
self.assertAlmostEqual(call_args.kwargs.get("timeout"), expected_timeout)
358359

360+
def test_done_w_query_results_error_reload_ok_job_finished(self):
361+
client = _make_client(project=self.PROJECT)
362+
bad_request_error = exceptions.BadRequest("Error in query")
363+
client._get_query_results = mock.Mock(side_effect=bad_request_error)
364+
365+
resource = self._make_resource(ended=False)
366+
job = self._get_target_class().from_api_repr(resource, client)
367+
job._exception = None
368+
369+
def fake_reload(self, *args, **kwargs):
370+
self._properties["status"]["state"] = "DONE"
371+
self.set_exception(copy.copy(bad_request_error))
372+
373+
fake_reload_method = types.MethodType(fake_reload, job)
374+
375+
with mock..object(job, "reload", new=fake_reload_method):
376+
is_done = job.done()
377+
378+
assert is_done
379+
assert isinstance(job._exception, exceptions.BadRequest)
380+
381+
def test_done_w_query_results_error_reload_ok_job_still_running(self):
382+
client = _make_client(project=self.PROJECT)
383+
retry_error = exceptions.RetryError("Too many retries", cause=TimeoutError)
384+
client._get_query_results = mock.Mock(side_effect=retry_error)
385+
386+
resource = self._make_resource(ended=False)
387+
job = self._get_target_class().from_api_repr(resource, client)
388+
job._exception = None
389+
390+
def fake_reload(self, *args, **kwargs):
391+
self._properties["status"]["state"] = "RUNNING"
392+
393+
fake_reload_method = types.MethodType(fake_reload, job)
394+
395+
with mock..object(job, "reload", new=fake_reload_method):
396+
is_done = job.done()
397+
398+
assert not is_done
399+
assert job._exception is None
400+
401+
def test_done_w_query_results_error_reload_error(self):
402+
client = _make_client(project=self.PROJECT)
403+
bad_request_error = exceptions.BadRequest("Error in query")
404+
client._get_query_results = mock.Mock(side_effect=bad_request_error)
405+
406+
resource = self._make_resource(ended=False)
407+
job = self._get_target_class().from_api_repr(resource, client)
408+
reload_error = exceptions.DataLoss("Oops, sorry!")
409+
job.reload = mock.Mock(side_effect=reload_error)
410+
job._exception = None
411+
412+
is_done = job.done()
413+
414+
assert is_done
415+
assert job._exception is bad_request_error
416+
417+
def test_done_w_job_query_results_ok_reload_error(self):
418+
client = _make_client(project=self.PROJECT)
419+
query_results = google.cloud.bigquery.query._QueryResults(
420+
properties={
421+
"jobComplete": True,
422+
"jobReference": {"projectId": self.PROJECT, "jobId": "12345"},
423+
}
424+
)
425+
client._get_query_results = mock.Mock(return_value=query_results)
426+
427+
resource = self._make_resource(ended=False)
428+
job = self._get_target_class().from_api_repr(resource, client)
429+
retry_error = exceptions.RetryError("Too many retries", cause=TimeoutError)
430+
job.reload = mock.Mock(side_effect=retry_error)
431+
job._exception = None
432+
433+
is_done = job.done()
434+
435+
assert is_done
436+
assert job._exception is retry_error
437+
359438
def test_query_plan(self):
360439
from google.cloud._helpers import _RFC3339_MICROS
361440
from google.cloud.bigquery.job import QueryPlanEntry
@@ -973,7 +1052,7 @@ def test_result_w_retry(self):
9731052
initial=0.001,
9741053
maximum=0.001,
9751054
multiplier=1.0,
976-
deadline=0.001,
1055+
deadline=0.1,
9771056
predicate=custom_predicate,
9781057
)
9791058

0 commit comments

Comments
 (0)