feat(bigquery): add retry parameter to public methods where missing (#…

…10026)

* Add retry to Client.get_service_account_email()

* Add retry to job.cancel()

bigquery/google/cloud/bigquery/client.py

@@ -208,7 +208,9 @@ def close(self):
         self._http._auth_request.session.close()
         self._http.close()
 
-    def get_service_account_email(self, project=None, timeout=None):
+    def get_service_account_email(
+        self, project=None, retry=DEFAULT_RETRY, timeout=None
+    ):
         """Get the email address of the project's BigQuery service account
 
         Note:
@@ -219,8 +221,10 @@ def get_service_account_email(self, project=None, timeout=None):
             project (str, optional):
                 Project ID to use for retreiving service account email.
                 Defaults to the client's project.
+            retry (Optional[google.api_core.retry.Retry]): How to retry the RPC.
             timeout (Optional[float]):
-                The number of seconds to wait for the API response.
+                The number of seconds to wait for the underlying HTTP transport
+                before using ``retry``.
 
         Returns:
             str: service account email address
@@ -237,10 +241,7 @@ def get_service_account_email(self, project=None, timeout=None):
             project = self.project
         path = "/projects/%s/serviceAccount" % (project,)
 
-        # TODO: call thorugh self._call_api() and allow passing in a retry?
-        api_response = self._connection.api_request(
-            method="GET", path=path, timeout=timeout
-        )
+        api_response = self._call_api(retry, method="GET", path=path, timeout=timeout)
         return api_response["email"]
 
     def list_projects(

bigquery/google/cloud/bigquery/job.py

@@ -710,7 +710,7 @@ def reload(self, client=None, retry=DEFAULT_RETRY, timeout=None):
         )
         self._set_properties(api_response)
 
-    def cancel(self, client=None, timeout=None):
+    def cancel(self, client=None, retry=DEFAULT_RETRY, timeout=None):
         """API call:  cancel job via a POST request
 
         See
@@ -720,8 +720,10 @@ def cancel(self, client=None, timeout=None):
             client (Optional[google.cloud.bigquery.client.Client]):
                 the client to use.  If not passed, falls back to the
                 ``client`` stored on the current dataset.
+            retry (Optional[google.api_core.retry.Retry]): How to retry the RPC.
             timeout (Optional[float]):
-                The number of seconds to wait for the API response.
+                The number of seconds to wait for the underlying HTTP transport
+                before using ``retry``
 
         Returns:
             bool: Boolean indicating that the cancel request was sent.
@@ -732,10 +734,10 @@ def cancel(self, client=None, timeout=None):
         if self.location:
             extra_params["location"] = self.location
 
-        # TODO: call thorugh client._call_api() and allow passing in a retry?
-        api_response = client._connection.api_request(
+        api_response = client._call_api(
+            retry,
             method="POST",
-            path="%s/cancel" % (self.path,),
+            path="{}/cancel".format(self.path),
             query_params=extra_params,
             timeout=timeout,
         )

bigquery/tests/unit/test_client.py

@@ -331,6 +331,42 @@ def test_get_service_account_email_w_alternate_project(self):
         conn.api_request.assert_called_once_with(method="GET", path=path, timeout=None)
         self.assertEqual(service_account_email, email)
 
+    def test_get_service_account_email_w_custom_retry(self):
+        from google.cloud.bigquery.retry import DEFAULT_RETRY
+
+        api_path = "/projects/{}/serviceAccount".format(self.PROJECT)
+        creds = _make_credentials()
+        http = object()
+        client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
+
+        resource = {
+            "kind": "bigquery#getServiceAccountResponse",
+            "email": "bq-123@bigquery-encryption.iam.gserviceaccount.com",
+        }
+        api_request_patcher = mock.patch.object(
+            client._connection, "api_request", side_effect=[ValueError, resource],
+        )
+
+        retry = DEFAULT_RETRY.with_deadline(1).with_predicate(
+            lambda exc: isinstance(exc, ValueError)
+        )
+
+        with api_request_patcher as fake_api_request:
+            service_account_email = client.get_service_account_email(
+                retry=retry, timeout=7.5
+            )
+
+        self.assertEqual(
+            service_account_email, "bq-123@bigquery-encryption.iam.gserviceaccount.com"
+        )
+        self.assertEqual(
+            fake_api_request.call_args_list,
+            [
+                mock.call(method="GET", path=api_path, timeout=7.5),
+                mock.call(method="GET", path=api_path, timeout=7.5),  # was retried once
+            ],
+        )
+
     def test_list_projects_defaults(self):
         from google.cloud.bigquery.client import Project
 

bigquery/tests/unit/test_job.py

@@ -831,6 +831,43 @@ def test_cancel_explicit(self):
         )
         self.assertEqual(job._properties, resource)
 
+    def test_cancel_w_custom_retry(self):
+        from google.cloud.bigquery.retry import DEFAULT_RETRY
+
+        api_path = "/projects/{}/jobs/{}/cancel".format(self.PROJECT, self.JOB_ID)
+        resource = {
+            "jobReference": {
+                "jobId": self.JOB_ID,
+                "projectId": self.PROJECT,
+                "location": None,
+            },
+            "configuration": {"test": True},
+        }
+        response = {"job": resource}
+        job = self._set_properties_job()
+
+        api_request_patcher = mock.patch.object(
+            job._client._connection, "api_request", side_effect=[ValueError, response],
+        )
+        retry = DEFAULT_RETRY.with_deadline(1).with_predicate(
+            lambda exc: isinstance(exc, ValueError)
+        )
+
+        with api_request_patcher as fake_api_request:
+            result = job.cancel(retry=retry, timeout=7.5)
+
+        self.assertTrue(result)
+        self.assertEqual(job._properties, resource)
+        self.assertEqual(
+            fake_api_request.call_args_list,
+            [
+                mock.call(method="POST", path=api_path, query_params={}, timeout=7.5),
+                mock.call(
+                    method="POST", path=api_path, query_params={}, timeout=7.5,
+                ),  # was retried once
+            ],
+        )
+
     def test__set_future_result_wo_done(self):
         client = _make_client(project=self.PROJECT)
         job = self._make_one(self.JOB_ID, client)