feat: add default timeout for Client.get_job() (#1935)

* feat: add default timeout for Client.get_job() * change timeout type detection * lint * fix unit test and coverage * add type hint * fix type hint * change import style and add comments * remove sentinel value in client * type hint * typo * add sentinel for query_and_wait() * add unit tests * fix unit test * Update google/cloud/bigquery/job/query.py Co-authored-by: Tim Sweña (Swast) <[email protected]> * Update google/cloud/bigquery/job/query.py Co-authored-by: Tim Sweña (Swast) <[email protected]> * address comments * typo * type hint * typos --------- Co-authored-by: Tim Sweña (Swast) <[email protected]>
googleapis · May 31, 2024 · 9fbad76 · sasha-gitg · Jul 10, 2024 · 9fbad76
1 parent 94d61b0
commit 9fbad76
Show file tree

Hide file tree

Showing 14 changed files with 421 additions and 209 deletions.
diff --git a/google/cloud/bigquery/_helpers.py b/google/cloud/bigquery/_helpers.py
@@ -33,6 +33,8 @@
 from google.auth import credentials as ga_credentials  # type: ignore
 from google.api_core import client_options as client_options_lib
 
+TimeoutType = Union[float, None]
+
 _RFC3339_MICROS_NO_ZULU = "%Y-%m-%dT%H:%M:%S.%f"
 _TIMEONLY_WO_MICROS = "%H:%M:%S"
 _TIMEONLY_W_MICROS = "%H:%M:%S.%f"

diff --git a/google/cloud/bigquery/_job_helpers.py b/google/cloud/bigquery/_job_helpers.py
@@ -39,14 +39,15 @@
 import functools
 import os
 import uuid
-from typing import Any, Dict, TYPE_CHECKING, Optional
+from typing import Any, Dict, Optional, TYPE_CHECKING, Union
 
 import google.api_core.exceptions as core_exceptions
 from google.api_core import retry as retries
 
 from google.cloud.bigquery import job
 import google.cloud.bigquery.query
 from google.cloud.bigquery import table
+from google.cloud.bigquery.retry import POLLING_DEFAULT_VALUE
 
 # Avoid circular imports
 if TYPE_CHECKING:  # pragma: NO COVER
@@ -328,7 +329,7 @@ def query_and_wait(
     location: Optional[str],
     project: str,
     api_timeout: Optional[float] = None,
-    wait_timeout: Optional[float] = None,
+    wait_timeout: Optional[Union[float, object]] = POLLING_DEFAULT_VALUE,
     retry: Optional[retries.Retry],
     job_retry: Optional[retries.Retry],
     page_size: Optional[int] = None,
@@ -364,10 +365,12 @@ def query_and_wait(
         api_timeout (Optional[float]):
             The number of seconds to wait for the underlying HTTP transport
             before using ``retry``.
-        wait_timeout (Optional[float]):
+        wait_timeout (Optional[Union[float, object]]):
             The number of seconds to wait for the query to finish. If the
             query doesn't finish before this timeout, the client attempts
-            to cancel the query.
+            to cancel the query. If unset, the underlying Client.get_job() API
+            call has timeout, but we still wait indefinitely for the job to
+            finish.
         retry (Optional[google.api_core.retry.Retry]):
             How to retry the RPC.  This only applies to making RPC
             calls.  It isn't used to retry failed jobs.  This has
@@ -545,7 +548,7 @@ def _supported_by_jobs_query(request_body: Dict[str, Any]) -> bool:
 def _wait_or_cancel(
     job: job.QueryJob,
     api_timeout: Optional[float],
-    wait_timeout: Optional[float],
+    wait_timeout: Optional[Union[object, float]],
     retry: Optional[retries.Retry],
     page_size: Optional[int],
     max_results: Optional[int],

diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py
@@ -82,6 +82,7 @@
 from google.cloud.bigquery._helpers import _DEFAULT_UNIVERSE
 from google.cloud.bigquery._helpers import _validate_universe
 from google.cloud.bigquery._helpers import _get_client_universe
+from google.cloud.bigquery._helpers import TimeoutType
 from google.cloud.bigquery._job_helpers import make_job_id as _make_job_id
 from google.cloud.bigquery.dataset import Dataset
 from google.cloud.bigquery.dataset import DatasetListItem
@@ -107,6 +108,7 @@
     DEFAULT_JOB_RETRY,
     DEFAULT_RETRY,
     DEFAULT_TIMEOUT,
+    DEFAULT_GET_JOB_TIMEOUT,
 )
 from google.cloud.bigquery.routine import Routine
 from google.cloud.bigquery.routine import RoutineReference
@@ -123,7 +125,6 @@
     _versions_helpers.PANDAS_VERSIONS.try_import()
 )  # mypy check fails because pandas import is outside module, there are type: ignore comments related to this
 
-TimeoutType = Union[float, None]
 ResumableTimeoutType = Union[
     None, float, Tuple[float, float]
 ]  # for resumable media methods
@@ -2139,7 +2140,7 @@ def get_job(
         project: Optional[str] = None,
         location: Optional[str] = None,
         retry: retries.Retry = DEFAULT_RETRY,
-        timeout: TimeoutType = DEFAULT_TIMEOUT,
+        timeout: TimeoutType = DEFAULT_GET_JOB_TIMEOUT,
     ) -> Union[job.LoadJob, job.CopyJob, job.ExtractJob, job.QueryJob, job.UnknownJob]:
         """Fetch a job for the project associated with this client.
 

diff --git a/google/cloud/bigquery/job/base.py b/google/cloud/bigquery/job/base.py
@@ -26,8 +26,11 @@
 import google.api_core.future.polling
 
 from google.cloud.bigquery import _helpers
-from google.cloud.bigquery.retry import DEFAULT_RETRY
 from google.cloud.bigquery._helpers import _int_or_none
+from google.cloud.bigquery.retry import (
+    DEFAULT_GET_JOB_TIMEOUT,
+    DEFAULT_RETRY,
+)
 
 
 _DONE_STATE = "DONE"
@@ -801,7 +804,7 @@ def reload(
         self,
         client=None,
         retry: "retries.Retry" = DEFAULT_RETRY,
-        timeout: Optional[float] = None,
+        timeout: Optional[float] = DEFAULT_GET_JOB_TIMEOUT,
     ):
         """API call:  refresh job properties via a GET request.
 
@@ -820,22 +823,14 @@ def reload(
         """
         client = self._require_client(client)
 
-        extra_params = {}
-        if self.location:
-            extra_params["location"] = self.location
-        span_attributes = {"path": self.path}
-
-        api_response = client._call_api(
-            retry,
-            span_name="BigQuery.job.reload",
-            span_attributes=span_attributes,
-            job_ref=self,
-            method="GET",
-            path=self.path,
-            query_params=extra_params,
+        got_job = client.get_job(
+            self,
+            project=self.project,
+            location=self.location,
+            retry=retry,
             timeout=timeout,
         )
-        self._set_properties(api_response)
+        self._set_properties(got_job._properties)
 
     def cancel(
         self,
@@ -913,7 +908,7 @@ def _set_future_result(self):
     def done(
         self,
         retry: "retries.Retry" = DEFAULT_RETRY,
-        timeout: Optional[float] = None,
+        timeout: Optional[float] = DEFAULT_GET_JOB_TIMEOUT,
         reload: bool = True,
     ) -> bool:
         """Checks if the job is complete.

diff --git a/google/cloud/bigquery/job/query.py b/google/cloud/bigquery/job/query.py
@@ -40,7 +40,11 @@
     StructQueryParameter,
     UDFResource,
 )
-from google.cloud.bigquery.retry import DEFAULT_RETRY, DEFAULT_JOB_RETRY
+from google.cloud.bigquery.retry import (
+    DEFAULT_RETRY,
+    DEFAULT_JOB_RETRY,
+    POLLING_DEFAULT_VALUE,
+)
 from google.cloud.bigquery.routine import RoutineReference
 from google.cloud.bigquery.schema import SchemaField
 from google.cloud.bigquery.table import _EmptyRowIterator
@@ -1437,7 +1441,7 @@ def result(  # type: ignore  # (incompatible with supertype)
         page_size: Optional[int] = None,
         max_results: Optional[int] = None,
         retry: Optional[retries.Retry] = DEFAULT_RETRY,
-        timeout: Optional[float] = None,
+        timeout: Optional[Union[float, object]] = POLLING_DEFAULT_VALUE,
         start_index: Optional[int] = None,
         job_retry: Optional[retries.Retry] = DEFAULT_JOB_RETRY,
     ) -> Union["RowIterator", _EmptyRowIterator]:
@@ -1457,11 +1461,14 @@ def result(  # type: ignore  # (incompatible with supertype)
                 is ``DONE``, retrying is aborted early even if the
                 results are not available, as this will not change
                 anymore.
-            timeout (Optional[float]):
+            timeout (Optional[Union[float, \
+                google.api_core.future.polling.PollingFuture._DEFAULT_VALUE, \
+            ]]):
                 The number of seconds to wait for the underlying HTTP transport
-                before using ``retry``.
-                If multiple requests are made under the hood, ``timeout``
-                applies to each individual request.
+                before using ``retry``. If ``None``, wait indefinitely
+                unless an error is returned. If unset, only the
+                underlying API calls have their default timeouts, but we still
+                wait indefinitely for the job to finish.
             start_index (Optional[int]):
                 The zero-based index of the starting row to read.
             job_retry (Optional[google.api_core.retry.Retry]):
@@ -1507,6 +1514,13 @@ def result(  # type: ignore  # (incompatible with supertype)
                 # Intentionally omit job_id and query_id since this doesn't
                 # actually correspond to a finished query job.
             )
+
+        # When timeout has default sentinel value ``object()``, do not pass
+        # anything to invoke default timeouts in subsequent calls.
+        kwargs: Dict[str, Union[_helpers.TimeoutType, object]] = {}
+        if type(timeout) is not object:
+            kwargs["timeout"] = timeout
+
         try:
             retry_do_query = getattr(self, "_retry_do_query", None)
             if retry_do_query is not None:
@@ -1548,7 +1562,7 @@ def is_job_done():
                 # 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 self.done(retry=retry, **kwargs):
                     # If it's already failed, we might as well stop.
                     job_failed_exception = self.exception()
                     if job_failed_exception is not None:
@@ -1585,14 +1599,14 @@ def is_job_done():
                         # response from the REST API. This ensures we aren't
                         # making any extra API calls if the previous loop
                         # iteration fetched the finished job.
-                        self._reload_query_results(retry=retry, timeout=timeout)
+                        self._reload_query_results(retry=retry, **kwargs)
                         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)
+                self._reload_query_results(retry=retry, **kwargs)
 
                 # Even if the query is finished now according to
                 # jobs.getQueryResults, we'll want to reload the job status if
@@ -1682,10 +1696,10 @@ def is_job_done():
             max_results=max_results,
             start_index=start_index,
             retry=retry,
-            timeout=timeout,
             query_id=self.query_id,
             first_page_response=first_page_response,
             num_dml_affected_rows=self._query_results.num_dml_affected_rows,
+            **kwargs,
         )
         rows._preserve_order = _contains_order_by(self.query)
         return rows

diff --git a/google/cloud/bigquery/retry.py b/google/cloud/bigquery/retry.py
@@ -14,6 +14,7 @@
 
 from google.api_core import exceptions
 from google.api_core import retry
+import google.api_core.future.polling
 from google.auth import exceptions as auth_exceptions  # type: ignore
 import requests.exceptions
 
@@ -140,3 +141,13 @@ def _job_should_retry(exc):
 """
 The default job retry object.
 """
+
+DEFAULT_GET_JOB_TIMEOUT = 128
+"""
+Default timeout for Client.get_job().
+"""
+
+POLLING_DEFAULT_VALUE = google.api_core.future.polling.PollingFuture._DEFAULT_VALUE
+"""
+Default value defined in google.api_core.future.polling.PollingFuture.
+"""