fix: query() now warns when job_id is set and the default job_retry is ignored

Author: tswast

google/cloud/bigquery/_job_helpers.py

@@ -39,7 +39,9 @@
 import functools
 import os
 import uuid
+import textwrap
 from typing import Any, Dict, Optional, TYPE_CHECKING, Union
+import warnings
 
 import google.api_core.exceptions as core_exceptions
 from google.api_core import retry as retries
@@ -198,6 +200,44 @@ def _validate_job_config(request_body: Dict[str, Any], invalid_key: str):
         raise ValueError(f"got unexpected key {repr(invalid_key)} in job_config")
 
 
+def validate_job_retry(job_id: Optional[str], job_retry: Optional[retries.Retry]):
+    """Catch common mistakes, such as setting a job_id and job_retry at the same
+    time.
+    """
+    if job_id is not None and job_retry is not None:
+        # TODO(tswast): To avoid breaking changes but still allow a default
+        # query job retry, we currently only raise if they explicitly set a
+        # job_retry other than the default. In a future version, we may want to
+        # avoid this check for DEFAULT_JOB_RETRY and always raise.
+        if job_retry is not google.cloud.bigquery.retry.DEFAULT_JOB_RETRY:
+            raise TypeError(
+                textwrap.dedent(
+                    """
+                    `job_retry` was provided, but the returned job is
+                    not retryable, because a custom `job_id` was
+                    provided. To customize the job ID and allow for job
+                    retries, set job_id_prefix, instead.
+                    """
+                ).strip()
+            )
+        else:
+            warnings.warn(
+                textwrap.dedent(
+                    """
+                    job_retry must be explicitly set to None if job_id is set.
+                    BigQuery cannot retry a failed job by using the exact
+                    same ID. Setting job_id without explicitly disabling
+                    job_retry will raise an error in the future. To avoid this
+                    warning, either use job_id_prefix instead (preferred) or
+                    set job_retry=None.
+                    """
+                ).strip(),
+                category=FutureWarning,
+                # user code -> client.query / client.query_and_wait -> validate_job_retry
+                stacklevel=3,
+            )
+
+
 def _to_query_request(
     job_config: Optional[job.QueryJobConfig] = None,
     *,

google/cloud/bigquery/client.py

@@ -3388,7 +3388,7 @@ def query(
         project: Optional[str] = None,
         retry: retries.Retry = DEFAULT_RETRY,
         timeout: TimeoutType = DEFAULT_TIMEOUT,
-        job_retry: retries.Retry = DEFAULT_JOB_RETRY,
+        job_retry: Optional[retries.Retry] = DEFAULT_JOB_RETRY,
         api_method: Union[str, enums.QueryApiMethod] = enums.QueryApiMethod.INSERT,
     ) -> job.QueryJob:
         """Run a SQL query.
@@ -3455,18 +3455,9 @@ def query(
                 class, or if both ``job_id`` and non-``None`` non-default
                 ``job_retry`` are provided.
         """
-        job_id_given = job_id is not None
-        if (
-            job_id_given
-            and job_retry is not None
-            and job_retry is not DEFAULT_JOB_RETRY
-        ):
-            raise TypeError(
-                "`job_retry` was provided, but the returned job is"
-                " not retryable, because a custom `job_id` was"
-                " provided."
-            )
+        _job_helpers.validate_job_retry(job_id, job_retry)
 
+        job_id_given = job_id is not None
         if job_id_given and api_method == enums.QueryApiMethod.QUERY:
             raise TypeError(
                 "`job_id` was provided, but the 'QUERY' `api_method` was requested."
@@ -3524,7 +3515,7 @@ def query_and_wait(
         api_timeout: TimeoutType = DEFAULT_TIMEOUT,
         wait_timeout: Union[Optional[float], object] = POLLING_DEFAULT_VALUE,
         retry: retries.Retry = DEFAULT_RETRY,
-        job_retry: retries.Retry = DEFAULT_JOB_RETRY,
+        job_retry: Optional[retries.Retry] = DEFAULT_JOB_RETRY,
         page_size: Optional[int] = None,
         max_results: Optional[int] = None,
     ) -> RowIterator:

tests/unit/test_job_retry.py

@@ -511,26 +511,34 @@ def api_request(method, path, query_params=None, data=None, **kw):
 def test_raises_on_job_retry_on_query_with_non_retryable_jobs(client):
     with pytest.raises(
         TypeError,
-        match=re.escape(
+        match=(
             "`job_retry` was provided, but the returned job is"
             " not retryable, because a custom `job_id` was"
             " provided."
-        ),
+        ).replace(" ", r"\s"),
     ):
         client.query("select 42", job_id=42, job_retry=google.api_core.retry.Retry())
 
 
 def test_raises_on_job_retry_on_result_with_non_retryable_jobs(client):
     client._connection = make_connection({})
-    job = client.query("select 42", job_id=42)
+
+    with pytest.warns(
+        FutureWarning,
+        match=re.escape("job_retry must be explicitly set to None if job_id is set."),
+    ):
+        # Implicitly providing a job_retry is a warning and will be an error in the future.
+        job = client.query("select 42", job_id=42)
+
     with pytest.raises(
         TypeError,
-        match=re.escape(
+        match=(
             "`job_retry` was provided, but this job is"
             " not retryable, because a custom `job_id` was"
             " provided to the query that created this job."
-        ),
+        ).replace(" ", r"\s"),
     ):
+        # Expliclty providing a job_retry is an error.
         job.result(job_retry=google.api_core.retry.Retry())