Skip to content
Merged
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 41 additions & 1 deletion google/cloud/bigquery/_job_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
*,
Expand Down Expand Up @@ -308,7 +348,7 @@ def query_jobs_query(
project: str,
retry: retries.Retry,
timeout: Optional[float],
job_retry: retries.Retry,
job_retry: Optional[retries.Retry],
) -> job.QueryJob:
"""Initiate a query using jobs.query with jobCreationMode=JOB_CREATION_REQUIRED.

Expand Down
15 changes: 3 additions & 12 deletions google/cloud/bigquery/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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."
Expand Down
18 changes: 13 additions & 5 deletions tests/unit/test_job_retry.py
Original file line number Diff line number Diff line change
Expand Up @@ -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())


Expand Down