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 <swast@google.com>
Co-authored-by: HemangChothani <hemang.chothani@qlogic.io>

google/cloud/bigquery/job.py

@@ -819,8 +819,9 @@ def result(self, retry=DEFAULT_RETRY, timeout=None):
         """
         if self.state is None:
             self._begin(retry=retry, timeout=timeout)
-        # TODO: modify PollingFuture so it can pass a retry argument to done().
-        return super(_AsyncJob, self).result(timeout=timeout)
+
+        kwargs = {} if retry is DEFAULT_RETRY else {"retry": retry}
+        return super(_AsyncJob, self).result(timeout=timeout, **kwargs)
 
     def cancelled(self):
         """Check if the job has been cancelled.
@@ -1845,7 +1846,7 @@ def destination(self):
         """
         return TableReference.from_api_repr(
             _helpers._get_sub_prop(
-                self._properties, ["configuration", "copy", "destinationTable"],
+                self._properties, ["configuration", "copy", "destinationTable"]
             )
         )
 
@@ -2043,10 +2044,7 @@ def __init__(self, job_id, source, destination_uris, client, job_config=None):
         self._configuration = job_config
 
         if source:
-            source_ref = {
-                "projectId": source.project,
-                "datasetId": source.dataset_id,
-            }
+            source_ref = {"projectId": source.project, "datasetId": source.dataset_id}
 
             if isinstance(source, (Table, TableListItem, TableReference)):
                 source_ref["tableId"] = source.table_id
@@ -3138,10 +3136,10 @@ def done(self, retry=DEFAULT_RETRY, timeout=None, reload=True):
 
         return self.state == _DONE_STATE
 
-    def _blocking_poll(self, timeout=None):
+    def _blocking_poll(self, timeout=None, **kwargs):
         self._done_timeout = timeout
         self._transport_timeout = timeout
-        super(QueryJob, self)._blocking_poll(timeout=timeout)
+        super(QueryJob, self)._blocking_poll(timeout=timeout, **kwargs)
 
     @staticmethod
     def _format_for_exception(query, job_id):

setup.py

@@ -29,7 +29,7 @@
 # 'Development Status :: 5 - Production/Stable'
 release_status = "Development Status :: 5 - Production/Stable"
 dependencies = [
-    "google-api-core[grpc] >= 1.22.2, < 2.0.0dev",
+    "google-api-core[grpc] >= 1.23.0, < 2.0.0dev",
     "proto-plus >= 1.10.0",
     "google-cloud-core >= 1.4.1, < 2.0dev",
     "google-resumable-media >= 0.6.0, < 2.0dev",

testing/constraints-3.6.txt

@@ -1,4 +1,4 @@
-google-api-core==1.22.2
+google-api-core==1.23.0
 google-cloud-bigquery-storage==2.0.0
 google-cloud-core==1.4.1
 google-resumable-media==0.6.0

tests/unit/test_job.py

@@ -864,7 +864,7 @@ def test_cancel_w_custom_retry(self):
         job = self._set_properties_job()
 
         api_request_patcher = mock.patch.object(
-            job._client._connection, "api_request", side_effect=[ValueError, response],
+            job._client._connection, "api_request", side_effect=[ValueError, response]
         )
         retry = DEFAULT_RETRY.with_deadline(1).with_predicate(
             lambda exc: isinstance(exc, ValueError)
@@ -885,7 +885,7 @@ def test_cancel_w_custom_retry(self):
             [
                 mock.call(method="POST", path=api_path, query_params={}, timeout=7.5),
                 mock.call(
-                    method="POST", path=api_path, query_params={}, timeout=7.5,
+                    method="POST", path=api_path, query_params={}, timeout=7.5
                 ),  # was retried once
             ],
         )
@@ -1034,7 +1034,6 @@ def test_result_w_retry_wo_state(self):
         custom_predicate = mock.Mock()
         custom_predicate.return_value = True
         custom_retry = google.api_core.retry.Retry(predicate=custom_predicate)
-
         self.assertIs(job.result(retry=custom_retry), job)
 
         begin_call = mock.call(
@@ -2757,7 +2756,7 @@ def test_cancel_w_bound_client(self):
         final_attributes.assert_called_with({"path": PATH}, client, job)
 
         conn.api_request.assert_called_once_with(
-            method="POST", path=PATH, query_params={}, timeout=None,
+            method="POST", path=PATH, query_params={}, timeout=None
         )
         self._verifyResourceProperties(job, RESOURCE)
 
@@ -2779,7 +2778,7 @@ def test_cancel_w_alternate_client(self):
 
         conn1.api_request.assert_not_called()
         conn2.api_request.assert_called_once_with(
-            method="POST", path=PATH, query_params={}, timeout=None,
+            method="POST", path=PATH, query_params={}, timeout=None
         )
         self._verifyResourceProperties(job, RESOURCE)
 
@@ -3205,7 +3204,7 @@ def test_exists_miss_w_bound_client(self):
         final_attributes.assert_called_with({"path": PATH}, client, job)
 
         conn.api_request.assert_called_once_with(
-            method="GET", path=PATH, query_params={"fields": "id"}, timeout=None,
+            method="GET", path=PATH, query_params={"fields": "id"}, timeout=None
         )
 
     def test_exists_hit_w_alternate_client(self):
@@ -3620,7 +3619,7 @@ def test_exists_miss_w_bound_client(self):
         final_attributes.assert_called_with({"path": PATH}, client, job)
 
         conn.api_request.assert_called_once_with(
-            method="GET", path=PATH, query_params={"fields": "id"}, timeout=None,
+            method="GET", path=PATH, query_params={"fields": "id"}, timeout=None
         )
 
     def test_exists_hit_w_alternate_client(self):
@@ -4812,6 +4811,60 @@ def test_result_with_max_results(self):
             tabledata_list_request[1]["query_params"]["maxResults"], max_results
         )
 
+    def test_result_w_retry(self):
+        from google.cloud.bigquery.table import RowIterator
+
+        query_resource = {
+            "jobComplete": False,
+            "jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID},
+        }
+        query_resource_done = {
+            "jobComplete": True,
+            "jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID},
+            "schema": {"fields": [{"name": "col1", "type": "STRING"}]},
+            "totalRows": "2",
+        }
+        job_resource = self._make_resource(started=True)
+        job_resource_done = self._make_resource(started=True, ended=True)
+        job_resource_done["configuration"]["query"]["destinationTable"] = {
+            "projectId": "dest-project",
+            "datasetId": "dest_dataset",
+            "tableId": "dest_table",
+        }
+
+        connection = _make_connection(
+            exceptions.NotFound("not normally retriable"),
+            query_resource,
+            exceptions.NotFound("not normally retriable"),
+            query_resource_done,
+            exceptions.NotFound("not normally retriable"),
+            job_resource_done,
+        )
+        client = _make_client(self.PROJECT, connection=connection)
+        job = self._get_target_class().from_api_repr(job_resource, client)
+
+        custom_predicate = mock.Mock()
+        custom_predicate.return_value = True
+        custom_retry = google.api_core.retry.Retry(predicate=custom_predicate)
+
+        self.assertIsInstance(job.result(retry=custom_retry), RowIterator)
+        query_results_call = mock.call(
+            method="GET",
+            path=f"/projects/{self.PROJECT}/queries/{self.JOB_ID}",
+            query_params={"maxResults": 0},
+            timeout=None,
+        )
+        reload_call = mock.call(
+            method="GET",
+            path=f"/projects/{self.PROJECT}/jobs/{self.JOB_ID}",
+            query_params={},
+            timeout=None,
+        )
+
+        connection.api_request.assert_has_calls(
+            [query_results_call, query_results_call, reload_call]
+        )
+
     def test_result_w_empty_schema(self):
         from google.cloud.bigquery.table import _EmptyRowIterator