googleapis · tswast · Apr 18, 2024 · Apr 12, 2024 · Apr 12, 2024 · Apr 13, 2024
@@ -17,6 +17,7 @@
 import concurrent.futures
 import copy
 import re
+import time
 import typing
 from typing import Any, Dict, Iterable, List, Optional, Union
 
@@ -1439,6 +1440,12 @@ def _done_or_raise(self, retry=DEFAULT_RETRY, timeout=None):
  # stored in _blocking_poll() in the process of polling for job completion.
  transport_timeout = timeout if timeout is not None else self._transport_timeout
 
+ # We could have been given a generic object() sentinel to represent a
+ # default timeout.
+ transport_timeout = (
+ transport_timeout if isinstance(timeout, (int, float)) else None
+ )
+
  try:
  self._reload_query_results(retry=retry, timeout=transport_timeout)
  except exceptions.GoogleAPIError as exc:
@@ -1528,6 +1535,10 @@ def result( # type: ignore # (incompatible with supertype)
  If Non-``None`` and non-default ``job_retry`` is
  provided and the job is not retryable.
  """
+ # Note: Since waiting for a query job to finish is more complex than
+ # refreshing the job state in a loop, we avoid calling the superclass
+ # in this method.
+
  if self.dry_run:
  return _EmptyRowIterator(
  project=self.project,
@@ -1548,46 +1559,90 @@ def result( # type: ignore # (incompatible with supertype)
  " provided to the query that created this job."
  )
 
- first = True
+ restart_query_job = False
 
- def do_get_result():
- nonlocal first
+ def is_job_done():
+ nonlocal restart_query_job
 
- if first:
- first = False
- else:
- # Note that we won't get here if retry_do_query is
- # None, because we won't use a retry.
+ if restart_query_job:
+ restart_query_job = False
 
  # The orinal job is failed. Create a new one.
+ #
+ # Note that we won't get here if retry_do_query is
+ # None, because we won't use a retry.
  job = retry_do_query()
 
- # If it's already failed, we might as well stop:
- if job.done() and job.exception() is not None:
- raise job.exception()
-
  # Become the new job:
  self.__dict__.clear()
  self.__dict__.update(job.__dict__)
 
- # This shouldn't be necessary, because once we have a good
- # job, it should stay good,and we shouldn't have to retry.
- # But let's be paranoid. :)
+ # It's possible the job fails again and we'll have to
+ # retry that too.
  self._retry_do_query = retry_do_query
  self._job_retry = job_retry
 
- super(QueryJob, self).result(retry=retry, timeout=timeout)
-
- # Since the job could already be "done" (e.g. got a finished job
- # via client.get_job), the superclass call to done() might not
- # set the self._query_results cache.
- if self._query_results is None or not self._query_results.complete:
- self._reload_query_results(retry=retry, timeout=timeout)
+ # Refresh the job status with jobs.get because some of the
+ # exceptions thrown by jobs.getQueryResults like timeout and
+ # rateLimitExceeded errors are ambiguous. We want to know if
+ # the query job failed and not just the call to
+ # jobs.getQueryResults.
+ if self.done(retry=retry, timeout=timeout):
+ # If it's already failed, we might as well stop.
+ if self.exception() is not None:
+ # Only try to restart the query job if the job failed for
+ # a retriable reason. For example, don't restart the query
+ # if the call to reload the job metadata within self.done()
+ # timed out.
+ restart_query_job = True
+ raise self.exception()
+ else:
+ # Optimization: avoid a call to jobs.getQueryResults if
+ # it's already been fetched, e.g. from jobs.query first
+ # page of results.
+ if (
+ self._query_results is None
+ or not self._query_results.complete
+ ):
+ self._reload_query_results(retry=retry, timeout=timeout)
+ return True
+
+ # Call jobs.getQueryResults with max results set to 0 just to
+ # wait for the query to finish. Unlike most methods,
+ # jobs.getQueryResults hangs as long as it can to ensure we
+ # know when the query has finished as soon as possible.
+ self._reload_query_results(retry=retry, timeout=timeout)
+
+ # Even if the query is finished now according to
+ # jobs.getQueryResults, we'll want to reload the job status if
+ # it's not already DONE.
+ return False
 
  if retry_do_query is not None and job_retry is not None:
- do_get_result = job_retry(do_get_result)
+ is_job_done = job_retry(is_job_done)
 
- do_get_result()
+ # timeout can be `None` or an object from our superclass
+ # indicating a default timeout.
+ remaining_timeout = timeout if isinstance(timeout, (float, int)) else None
+
+ if remaining_timeout is None:
+ # Since is_job_done() calls jobs.getQueryResults, which is a
+ # long-running API, don't delay the next request at all.
+ while not is_job_done():
+ pass
+ else:
+ # Use a monotonic clock since we don't actually care about
+ # daylight savings or similar, just the elapsed time.
+ previous_time = time.monotonic()
+
+ while not is_job_done():
+ current_time = time.monotonic()
+ elapsed_time = current_time - previous_time
+ remaining_timeout = remaining_timeout - elapsed_time
+ previous_time = current_time
+
+ if remaining_timeout < 0:
+ raise concurrent.futures.TimeoutError()
 
  except exceptions.GoogleAPICallError as exc:
  exc.message = _EXCEPTION_FOOTER_TEMPLATE.format(

@@ -970,7 +970,12 @@ def test_result(self):
  "rows": [{"f": [{"v": "abc"}]}],
  }
  conn = make_connection(
- query_resource, query_resource_done, job_resource_done, query_page_resource
+ job_resource,
+ query_resource,
+ job_resource,
+ query_resource_done,
+ job_resource_done,
+ query_page_resource,
  )
  client = _make_client(self.PROJECT, connection=conn)
  job = self._get_target_class().from_api_repr(job_resource, client)
@@ -1014,7 +1019,14 @@ def test_result(self):
  timeout=None,
  )
  conn.api_request.assert_has_calls(
- [query_results_call, query_results_call, reload_call, query_page_call]
+ [
+ reload_call,
+ query_results_call,
+ reload_call,
+ query_results_call,
+ reload_call,
+ query_page_call,
+ ]
  )
 
  def test_result_dry_run(self):
@@ -1254,9 +1266,13 @@ def test_result_w_retry(self):
  }
 
  connection = make_connection(
+ exceptions.NotFound("not normally retriable"),
+ job_resource,
  exceptions.NotFound("not normally retriable"),
  query_resource,
  exceptions.NotFound("not normally retriable"),
+ job_resource,
+ exceptions.NotFound("not normally retriable"),
  query_resource_done,
  exceptions.NotFound("not normally retriable"),
  job_resource_done,
@@ -1289,7 +1305,18 @@ def test_result_w_retry(self):
  )
 
  connection.api_request.assert_has_calls(
- [query_results_call, query_results_call, reload_call]
+ [
+ reload_call,
+ reload_call,
+ query_results_call,
+ query_results_call,
+ reload_call,
+ reload_call,
+ query_results_call,
+ query_results_call,
+ reload_call,
+ reload_call,
+ ]
  )
 
  def test_result_w_empty_schema(self):
@@ -1316,37 +1343,6 @@ def test_result_w_empty_schema(self):
  self.assertEqual(result.location, "asia-northeast1")
  self.assertEqual(result.query_id, "xyz-abc")
 
- def test_result_invokes_begins(self):
- begun_resource = self._make_resource()
- incomplete_resource = {
- "jobComplete": False,
- "jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID},
- "schema": {"fields": [{"name": "col1", "type": "STRING"}]},
- }
- query_resource = copy.deepcopy(incomplete_resource)
- query_resource["jobComplete"] = True
- done_resource = copy.deepcopy(begun_resource)
- done_resource["status"] = {"state": "DONE"}
- connection = make_connection(
- begun_resource,
- incomplete_resource,
- query_resource,
- done_resource,
- query_resource,
- )
- client = _make_client(project=self.PROJECT, connection=connection)
- job = self._make_one(self.JOB_ID, self.QUERY, client)
-
- job.result()
-
- self.assertEqual(len(connection.api_request.call_args_list), 4)
- begin_request = connection.api_request.call_args_list[0]
- query_request = connection.api_request.call_args_list[2]
- reload_request = connection.api_request.call_args_list[3]
- self.assertEqual(begin_request[1]["method"], "POST")
- self.assertEqual(query_request[1]["method"], "GET")
- self.assertEqual(reload_request[1]["method"], "GET")
-
  def test_result_w_timeout(self):
  import google.cloud.bigquery.client
 

@@ -896,18 +896,10 @@ def test_query_and_wait_caches_completed_query_results_more_pages():
  timeout=None,
  )
 
- # TODO(swast): Fetching job metadata isn't necessary in this case.
- jobs_get_path = "/projects/response-project/jobs/response-job-id"
- client._call_api.assert_any_call(
- None, # retry
- span_name="BigQuery.job.reload",
- span_attributes={"path": jobs_get_path},
- job_ref=mock.ANY,
- method="GET",
- path=jobs_get_path,
- query_params={"location": "response-location"},
- timeout=None,
- )
+ # Note: There is no get call to
+ # "/projects/response-project/jobs/response-job-id", because fetching job
+ # metadata isn't necessary in this case. The job already completed in
+ # jobs.query and we don't need the full job metadata in query_and_wait.
 
  # Fetch the remaining two pages.
  jobs_get_query_results_path = "/projects/response-project/queries/response-job-id"
@@ -944,6 +936,7 @@ def test_query_and_wait_incomplete_query():
  Client._list_rows_from_query_results, client
  )
  client._call_api.side_effect = (
+ # jobs.query
  {
  "jobReference": {
  "projectId": "response-project",
@@ -952,6 +945,16 @@ def test_query_and_wait_incomplete_query():
  },
  "jobComplete": False,
  },
+ # jobs.get
+ {
+ "jobReference": {
+ "projectId": "response-project",
+ "jobId": "response-job-id",
+ "location": "response-location",
+ },
+ "status": {"state": "RUNNING"},
+ },
+ # jobs.getQueryResults with max_results=0
  {
  "jobReference": {
  "projectId": "response-project",
@@ -968,13 +971,18 @@ def test_query_and_wait_incomplete_query():
  ],
  },
  },
+ # jobs.get
  {
  "jobReference": {
  "projectId": "response-project",
  "jobId": "response-job-id",
  "location": "response-location",
  },
+ "status": {"state": "DONE"},
  },
+ # jobs.getQueryResults
+ # Note: No more jobs.getQueryResults with max_results=0 because the
+ # previous call to jobs.getQueryResults returned with jobComplete=True.
  {
  "rows": [
  {"f": [{"v": "Whillma Phlyntstone"}, {"v": "27"}]},
@@ -987,6 +995,7 @@ def test_query_and_wait_incomplete_query():
  "totalRows": 2,
  "pageToken": "page-2",
  },
+ # jobs.getQueryResults
  {
  "rows": [
  {"f": [{"v": "Pearl Slaghoople"}, {"v": "53"}]},

@@ -187,8 +187,8 @@ def api_request(method, path, query_params=None, data=None, **kw):
  with pytest.raises(google.api_core.exceptions.RetryError):
  job.result()
 
- # We never got a successful job, so the job id never changed:
- assert job.job_id == orig_job_id
+ # We retried the job at least once, so we should have generated a new job ID.
+ assert job.job_id != orig_job_id
 
  # We failed because we couldn't succeed after 120 seconds.
  # But we can try again: