Skip to content

Commit de73e45

Browse files
plamutplamut
authored
feat(bigquery): add retry parameter to public methods where missing (#10026)
* Add retry to Client.get_service_account_email() * Add retry to job.cancel()
1 parent e41ed8a commit de73e45

File tree

4 files changed

+87
-11
lines changed

4 files changed

+87
-11
lines changed

‎bigquery/google/cloud/bigquery/client.py‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -208,7 +208,9 @@ def close(self):
208208
self._http._auth_request.session.close()
209209
self._http.close()
210210

211-
def get_service_account_email(self, project=None, timeout=None):
211+
def get_service_account_email(
212+
self, project=None, retry=DEFAULT_RETRY, timeout=None
213+
):
212214
"""Get the email address of the project's BigQuery service account
213215
214216
Note:
@@ -219,8 +221,10 @@ def get_service_account_email(self, project=None, timeout=None):
219221
project (str, optional):
220222
Project ID to use for retreiving service account email.
221223
Defaults to the client's project.
224+
retry (Optional[google.api_core.retry.Retry]): How to retry the RPC.
222225
timeout (Optional[float]):
223-
The number of seconds to wait for the API response.
226+
The number of seconds to wait for the underlying HTTP transport
227+
before using ``retry``.
224228
225229
Returns:
226230
str: service account email address
@@ -237,10 +241,7 @@ def get_service_account_email(self, project=None, timeout=None):
237241
project = self.project
238242
path = "/projects/%s/serviceAccount" % (project,)
239243

240-
# TODO: call thorugh self._call_api() and allow passing in a retry?
241-
api_response = self._connection.api_request(
242-
method="GET", path=path, timeout=timeout
243-
)
244+
api_response = self._call_api(retry, method="GET", path=path, timeout=timeout)
244245
return api_response["email"]
245246

246247
def list_projects(

‎bigquery/google/cloud/bigquery/job.py‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -710,7 +710,7 @@ def reload(self, client=None, retry=DEFAULT_RETRY, timeout=None):
710710
)
711711
self._set_properties(api_response)
712712

713-
def cancel(self, client=None, timeout=None):
713+
def cancel(self, client=None, retry=DEFAULT_RETRY, timeout=None):
714714
"""API call: cancel job via a POST request
715715
716716
See
@@ -720,8 +720,10 @@ def cancel(self, client=None, timeout=None):
720720
client (Optional[google.cloud.bigquery.client.Client]):
721721
the client to use. If not passed, falls back to the
722722
``client`` stored on the current dataset.
723+
retry (Optional[google.api_core.retry.Retry]): How to retry the RPC.
723724
timeout (Optional[float]):
724-
The number of seconds to wait for the API response.
725+
The number of seconds to wait for the underlying HTTP transport
726+
before using ``retry``
725727
726728
Returns:
727729
bool: Boolean indicating that the cancel request was sent.
@@ -732,10 +734,10 @@ def cancel(self, client=None, timeout=None):
732734
if self.location:
733735
extra_params["location"] = self.location
734736

735-
# TODO: call thorugh client._call_api() and allow passing in a retry?
736-
api_response = client._connection.api_request(
737+
api_response = client._call_api(
738+
retry,
737739
method="POST",
738-
path="%s/cancel" % (self.path,),
740+
path="{}/cancel".format(self.path),
739741
query_params=extra_params,
740742
timeout=timeout,
741743
)

‎bigquery/tests/unit/test_client.py‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -331,6 +331,42 @@ def test_get_service_account_email_w_alternate_project(self):
331331
conn.api_request.assert_called_once_with(method="GET", path=path, timeout=None)
332332
self.assertEqual(service_account_email, email)
333333

334+
def test_get_service_account_email_w_custom_retry(self):
335+
from google.cloud.bigquery.retry import DEFAULT_RETRY
336+
337+
api_path = "/projects/{}/serviceAccount".format(self.PROJECT)
338+
creds = _make_credentials()
339+
http = object()
340+
client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
341+
342+
resource = {
343+
"kind": "bigquery#getServiceAccountResponse",
344+
"email": "bq-123@bigquery-encryption.iam.gserviceaccount.com",
345+
}
346+
api_request_patcher = mock.patch.object(
347+
client._connection, "api_request", side_effect=[ValueError, resource],
348+
)
349+
350+
retry = DEFAULT_RETRY.with_deadline(1).with_predicate(
351+
lambda exc: isinstance(exc, ValueError)
352+
)
353+
354+
with api_request_patcher as fake_api_request:
355+
service_account_email = client.get_service_account_email(
356+
retry=retry, timeout=7.5
357+
)
358+
359+
self.assertEqual(
360+
service_account_email, "bq-123@bigquery-encryption.iam.gserviceaccount.com"
361+
)
362+
self.assertEqual(
363+
fake_api_request.call_args_list,
364+
[
365+
mock.call(method="GET", path=api_path, timeout=7.5),
366+
mock.call(method="GET", path=api_path, timeout=7.5), # was retried once
367+
],
368+
)
369+
334370
def test_list_projects_defaults(self):
335371
from google.cloud.bigquery.client import Project
336372

‎bigquery/tests/unit/test_job.py‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -831,6 +831,43 @@ def test_cancel_explicit(self):
831831
)
832832
self.assertEqual(job._properties, resource)
833833

834+
def test_cancel_w_custom_retry(self):
835+
from google.cloud.bigquery.retry import DEFAULT_RETRY
836+
837+
api_path = "/projects/{}/jobs/{}/cancel".format(self.PROJECT, self.JOB_ID)
838+
resource = {
839+
"jobReference": {
840+
"jobId": self.JOB_ID,
841+
"projectId": self.PROJECT,
842+
"location": None,
843+
},
844+
"configuration": {"test": True},
845+
}
846+
response = {"job": resource}
847+
job = self._set_properties_job()
848+
849+
api_request_patcher = mock.patch.object(
850+
job._client._connection, "api_request", side_effect=[ValueError, response],
851+
)
852+
retry = DEFAULT_RETRY.with_deadline(1).with_predicate(
853+
lambda exc: isinstance(exc, ValueError)
854+
)
855+
856+
with api_request_patcher as fake_api_request:
857+
result = job.cancel(retry=retry, timeout=7.5)
858+
859+
self.assertTrue(result)
860+
self.assertEqual(job._properties, resource)
861+
self.assertEqual(
862+
fake_api_request.call_args_list,
863+
[
864+
mock.call(method="POST", path=api_path, query_params={}, timeout=7.5),
865+
mock.call(
866+
method="POST", path=api_path, query_params={}, timeout=7.5,
867+
), # was retried once
868+
],
869+
)
870+
834871
def test__set_future_result_wo_done(self):
835872
client = _make_client(project=self.PROJECT)
836873
job = self._make_one(self.JOB_ID, client)

0 commit comments

Comments
 (0)