Skip to content

Commit 284e17a

Browse files
Ilya GurovtswastHemangChothani
authored
feat: pass retry from Job.result() to Job.done() (#41)
* feat(bigquery): pass retry from Job.result() to Job.done(). * fix merge conflicts * drop the comment * use kwargs sentinel * check the mock retry * update dependencies * use kwargs pattern * feat: added unit test for retry * feat: added more exceptions Co-authored-by: Tim Swast <[email protected]> Co-authored-by: HemangChothani <[email protected]>
1 parent 8a8080b commit 284e17a

File tree

4 files changed

+69
-18
lines changed

4 files changed

+69
-18
lines changed

google/cloud/bigquery/job.py

+7-9
Original file line numberDiff line numberDiff line change
@@ -819,8 +819,9 @@ def result(self, retry=DEFAULT_RETRY, timeout=None):
819819
"""
820820
if self.state is None:
821821
self._begin(retry=retry, timeout=timeout)
822-
# TODO: modify PollingFuture so it can pass a retry argument to done().
823-
return super(_AsyncJob, self).result(timeout=timeout)
822+
823+
kwargs = {} if retry is DEFAULT_RETRY else {"retry": retry}
824+
return super(_AsyncJob, self).result(timeout=timeout, **kwargs)
824825

825826
def cancelled(self):
826827
"""Check if the job has been cancelled.
@@ -1845,7 +1846,7 @@ def destination(self):
18451846
"""
18461847
return TableReference.from_api_repr(
18471848
_helpers._get_sub_prop(
1848-
self._properties, ["configuration", "copy", "destinationTable"],
1849+
self._properties, ["configuration", "copy", "destinationTable"]
18491850
)
18501851
)
18511852

@@ -2043,10 +2044,7 @@ def __init__(self, job_id, source, destination_uris, client, job_config=None):
20432044
self._configuration = job_config
20442045

20452046
if source:
2046-
source_ref = {
2047-
"projectId": source.project,
2048-
"datasetId": source.dataset_id,
2049-
}
2047+
source_ref = {"projectId": source.project, "datasetId": source.dataset_id}
20502048

20512049
if isinstance(source, (Table, TableListItem, TableReference)):
20522050
source_ref["tableId"] = source.table_id
@@ -3138,10 +3136,10 @@ def done(self, retry=DEFAULT_RETRY, timeout=None, reload=True):
31383136

31393137
return self.state == _DONE_STATE
31403138

3141-
def _blocking_poll(self, timeout=None):
3139+
def _blocking_poll(self, timeout=None, **kwargs):
31423140
self._done_timeout = timeout
31433141
self._transport_timeout = timeout
3144-
super(QueryJob, self)._blocking_poll(timeout=timeout)
3142+
super(QueryJob, self)._blocking_poll(timeout=timeout, **kwargs)
31453143

31463144
@staticmethod
31473145
def _format_for_exception(query, job_id):

setup.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@
2929
# 'Development Status :: 5 - Production/Stable'
3030
release_status = "Development Status :: 5 - Production/Stable"
3131
dependencies = [
32-
"google-api-core[grpc] >= 1.22.2, < 2.0.0dev",
32+
"google-api-core[grpc] >= 1.23.0, < 2.0.0dev",
3333
"proto-plus >= 1.10.0",
3434
"google-cloud-core >= 1.4.1, < 2.0dev",
3535
"google-resumable-media >= 0.6.0, < 2.0dev",

testing/constraints-3.6.txt

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
google-api-core==1.22.2
1+
google-api-core==1.23.0
22
google-cloud-bigquery-storage==2.0.0
33
google-cloud-core==1.4.1
44
google-resumable-media==0.6.0

tests/unit/test_job.py

+60-7
Original file line numberDiff line numberDiff line change
@@ -864,7 +864,7 @@ def test_cancel_w_custom_retry(self):
864864
job = self._set_properties_job()
865865

866866
api_request_patcher = mock.patch.object(
867-
job._client._connection, "api_request", side_effect=[ValueError, response],
867+
job._client._connection, "api_request", side_effect=[ValueError, response]
868868
)
869869
retry = DEFAULT_RETRY.with_deadline(1).with_predicate(
870870
lambda exc: isinstance(exc, ValueError)
@@ -885,7 +885,7 @@ def test_cancel_w_custom_retry(self):
885885
[
886886
mock.call(method="POST", path=api_path, query_params={}, timeout=7.5),
887887
mock.call(
888-
method="POST", path=api_path, query_params={}, timeout=7.5,
888+
method="POST", path=api_path, query_params={}, timeout=7.5
889889
), # was retried once
890890
],
891891
)
@@ -1034,7 +1034,6 @@ def test_result_w_retry_wo_state(self):
10341034
custom_predicate = mock.Mock()
10351035
custom_predicate.return_value = True
10361036
custom_retry = google.api_core.retry.Retry(predicate=custom_predicate)
1037-
10381037
self.assertIs(job.result(retry=custom_retry), job)
10391038

10401039
begin_call = mock.call(
@@ -2757,7 +2756,7 @@ def test_cancel_w_bound_client(self):
27572756
final_attributes.assert_called_with({"path": PATH}, client, job)
27582757

27592758
conn.api_request.assert_called_once_with(
2760-
method="POST", path=PATH, query_params={}, timeout=None,
2759+
method="POST", path=PATH, query_params={}, timeout=None
27612760
)
27622761
self._verifyResourceProperties(job, RESOURCE)
27632762

@@ -2779,7 +2778,7 @@ def test_cancel_w_alternate_client(self):
27792778

27802779
conn1.api_request.assert_not_called()
27812780
conn2.api_request.assert_called_once_with(
2782-
method="POST", path=PATH, query_params={}, timeout=None,
2781+
method="POST", path=PATH, query_params={}, timeout=None
27832782
)
27842783
self._verifyResourceProperties(job, RESOURCE)
27852784

@@ -3205,7 +3204,7 @@ def test_exists_miss_w_bound_client(self):
32053204
final_attributes.assert_called_with({"path": PATH}, client, job)
32063205

32073206
conn.api_request.assert_called_once_with(
3208-
method="GET", path=PATH, query_params={"fields": "id"}, timeout=None,
3207+
method="GET", path=PATH, query_params={"fields": "id"}, timeout=None
32093208
)
32103209

32113210
def test_exists_hit_w_alternate_client(self):
@@ -3620,7 +3619,7 @@ def test_exists_miss_w_bound_client(self):
36203619
final_attributes.assert_called_with({"path": PATH}, client, job)
36213620

36223621
conn.api_request.assert_called_once_with(
3623-
method="GET", path=PATH, query_params={"fields": "id"}, timeout=None,
3622+
method="GET", path=PATH, query_params={"fields": "id"}, timeout=None
36243623
)
36253624

36263625
def test_exists_hit_w_alternate_client(self):
@@ -4812,6 +4811,60 @@ def test_result_with_max_results(self):
48124811
tabledata_list_request[1]["query_params"]["maxResults"], max_results
48134812
)
48144813

4814+
def test_result_w_retry(self):
4815+
from google.cloud.bigquery.table import RowIterator
4816+
4817+
query_resource = {
4818+
"jobComplete": False,
4819+
"jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID},
4820+
}
4821+
query_resource_done = {
4822+
"jobComplete": True,
4823+
"jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID},
4824+
"schema": {"fields": [{"name": "col1", "type": "STRING"}]},
4825+
"totalRows": "2",
4826+
}
4827+
job_resource = self._make_resource(started=True)
4828+
job_resource_done = self._make_resource(started=True, ended=True)
4829+
job_resource_done["configuration"]["query"]["destinationTable"] = {
4830+
"projectId": "dest-project",
4831+
"datasetId": "dest_dataset",
4832+
"tableId": "dest_table",
4833+
}
4834+
4835+
connection = _make_connection(
4836+
exceptions.NotFound("not normally retriable"),
4837+
query_resource,
4838+
exceptions.NotFound("not normally retriable"),
4839+
query_resource_done,
4840+
exceptions.NotFound("not normally retriable"),
4841+
job_resource_done,
4842+
)
4843+
client = _make_client(self.PROJECT, connection=connection)
4844+
job = self._get_target_class().from_api_repr(job_resource, client)
4845+
4846+
custom_predicate = mock.Mock()
4847+
custom_predicate.return_value = True
4848+
custom_retry = google.api_core.retry.Retry(predicate=custom_predicate)
4849+
4850+
self.assertIsInstance(job.result(retry=custom_retry), RowIterator)
4851+
query_results_call = mock.call(
4852+
method="GET",
4853+
path=f"/projects/{self.PROJECT}/queries/{self.JOB_ID}",
4854+
query_params={"maxResults": 0},
4855+
timeout=None,
4856+
)
4857+
reload_call = mock.call(
4858+
method="GET",
4859+
path=f"/projects/{self.PROJECT}/jobs/{self.JOB_ID}",
4860+
query_params={},
4861+
timeout=None,
4862+
)
4863+
4864+
connection.api_request.assert_has_calls(
4865+
[query_results_call, query_results_call, reload_call]
4866+
)
4867+
48154868
def test_result_w_empty_schema(self):
48164869
from google.cloud.bigquery.table import _EmptyRowIterator
48174870

0 commit comments

Comments
 (0)