diff --git a/google/cloud/bigquery/_job_helpers.py b/google/cloud/bigquery/_job_helpers.py index 4a884ada5..888dc1e73 100644 --- a/google/cloud/bigquery/_job_helpers.py +++ b/google/cloud/bigquery/_job_helpers.py @@ -37,7 +37,6 @@ import copy import functools -import os import uuid import textwrap from typing import Any, Dict, Optional, TYPE_CHECKING, Union @@ -400,12 +399,6 @@ def query_and_wait( ) -> table.RowIterator: """Run the query, wait for it to finish, and return the results. - While ``jobCreationMode=JOB_CREATION_OPTIONAL`` is in preview in the - ``jobs.query`` REST API, use the default ``jobCreationMode`` unless - the environment variable ``QUERY_PREVIEW_ENABLED=true``. After - ``jobCreationMode`` is GA, this method will always use - ``jobCreationMode=JOB_CREATION_OPTIONAL``. See: - https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query Args: client: @@ -500,9 +493,8 @@ def query_and_wait( request_body["maxResults"] = min(page_size, max_results) elif page_size is not None or max_results is not None: request_body["maxResults"] = page_size or max_results - - if os.getenv("QUERY_PREVIEW_ENABLED", "").casefold() == "true": - request_body["jobCreationMode"] = "JOB_CREATION_OPTIONAL" + if client.default_job_creation_mode: + request_body["jobCreationMode"] = client.default_job_creation_mode def do_query(): request_body["requestId"] = make_job_id() diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py index 067b389a5..c6873545b 100644 --- a/google/cloud/bigquery/client.py +++ b/google/cloud/bigquery/client.py @@ -221,6 +221,10 @@ class Client(ClientWithProject): client_options (Optional[Union[google.api_core.client_options.ClientOptions, Dict]]): Client options used to set user options on the client. API Endpoint should be set through client_options. + default_job_creation_mode (Optional[str]): + Sets the default job creation mode used by query methods such as + query_and_wait(). For lightweight queries, JOB_CREATION_OPTIONAL is + generally recommended. Raises: google.auth.exceptions.DefaultCredentialsError: @@ -243,6 +247,7 @@ def __init__( client_options: Optional[ Union[google.api_core.client_options.ClientOptions, Dict[str, Any]] ] = None, + default_job_creation_mode: Optional[str] = None, ) -> None: if client_options is None: client_options = {} @@ -277,6 +282,7 @@ def __init__( self._connection = Connection(self, **kw_args) self._location = location self._default_load_job_config = copy.deepcopy(default_load_job_config) + self.default_job_creation_mode = default_job_creation_mode # Use property setter so validation can run. self.default_query_job_config = default_query_job_config @@ -286,6 +292,15 @@ def location(self): """Default location for jobs / datasets / tables.""" return self._location + @property + def default_job_creation_mode(self): + """Default job creation mode used for query execution.""" + return self._default_job_creation_mode + + @default_job_creation_mode.setter + def default_job_creation_mode(self, value: Optional[str]): + self._default_job_creation_mode = value + @property def default_query_job_config(self) -> Optional[QueryJobConfig]: """Default ``QueryJobConfig`` or ``None``. @@ -3532,13 +3547,6 @@ def query_and_wait( ) -> RowIterator: """Run the query, wait for it to finish, and return the results. - While ``jobCreationMode=JOB_CREATION_OPTIONAL`` is in preview in the - ``jobs.query`` REST API, use the default ``jobCreationMode`` unless - the environment variable ``QUERY_PREVIEW_ENABLED=true``. After - ``jobCreationMode`` is GA, this method will always use - ``jobCreationMode=JOB_CREATION_OPTIONAL``. See: - https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query - Args: query (str): SQL query to be executed. Defaults to the standard SQL diff --git a/google/cloud/bigquery/enums.py b/google/cloud/bigquery/enums.py index 203ea3c7b..4cb7a056d 100644 --- a/google/cloud/bigquery/enums.py +++ b/google/cloud/bigquery/enums.py @@ -407,3 +407,22 @@ class BigLakeTableFormat(object): ICEBERG = "ICEBERG" """Apache Iceberg format.""" + + +class JobCreationMode(object): + """Documented values for Job Creation Mode.""" + + JOB_CREATION_MODE_UNSPECIFIED = "JOB_CREATION_MODE_UNSPECIFIED" + """Job creation mode is unspecified.""" + + JOB_CREATION_REQUIRED = "JOB_CREATION_REQUIRED" + """Job creation is always required.""" + + JOB_CREATION_OPTIONAL = "JOB_CREATION_OPTIONAL" + """Job creation is optional. + + Returning immediate results is prioritized. + BigQuery will automatically determine if a Job needs to be created. + The conditions under which BigQuery can decide to not create a Job are + subject to change. + """ diff --git a/samples/client_query_shortmode.py b/samples/client_query_job_optional.py similarity index 69% rename from samples/client_query_shortmode.py rename to samples/client_query_job_optional.py index 50446dc48..6321aea35 100644 --- a/samples/client_query_shortmode.py +++ b/samples/client_query_job_optional.py @@ -13,16 +13,18 @@ # limitations under the License. -def client_query_shortmode() -> None: - # [START bigquery_query_shortquery] - # This example demonstrates issuing a query that may be run in short query mode. - # - # To enable the short query mode preview feature, the QUERY_PREVIEW_ENABLED - # environmental variable should be set to `TRUE`. +def client_query_job_optional() -> None: + # [START bigquery_query_job_optional] + # This example demonstrates executing a query without requiring an associated + # job. from google.cloud import bigquery + from google.cloud.bigquery.enums import JobCreationMode - # Construct a BigQuery client object. - client = bigquery.Client() + # Construct a BigQuery client object, specifying that the library should + # avoid creating jobs when possible. + client = bigquery.Client( + default_job_creation_mode=JobCreationMode.JOB_CREATION_OPTIONAL + ) query = """ SELECT @@ -44,10 +46,12 @@ def client_query_shortmode() -> None: if rows.job_id is not None: print("Query was run with job state. Job ID: {}".format(rows.job_id)) else: - print("Query was run in short mode. Query ID: {}".format(rows.query_id)) + print( + "Query was run without creating a job. Query ID: {}".format(rows.query_id) + ) print("The query data:") for row in rows: # Row values can be accessed by field name or index. print("name={}, gender={}, total={}".format(row[0], row[1], row["total"])) - # [END bigquery_query_shortquery] + # [END bigquery_query_job_optional] diff --git a/samples/tests/test_client_query_shortmode.py b/samples/tests/test_client_query_job_optional.py similarity index 85% rename from samples/tests/test_client_query_shortmode.py rename to samples/tests/test_client_query_job_optional.py index 41132f24c..0e0b2cf19 100644 --- a/samples/tests/test_client_query_shortmode.py +++ b/samples/tests/test_client_query_job_optional.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC +# Copyright 2025 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -14,13 +14,13 @@ import typing -from .. import client_query_shortmode +from .. import client_query_job_optional if typing.TYPE_CHECKING: import pytest def test_client_query_shortmode(capsys: "pytest.CaptureFixture[str]") -> None: - client_query_shortmode.client_query_shortmode() + client_query_job_optional.client_query_job_optional() out, err = capsys.readouterr() assert "Query was run" in out diff --git a/tests/unit/test__job_helpers.py b/tests/unit/test__job_helpers.py index 4fa093c69..417f911b8 100644 --- a/tests/unit/test__job_helpers.py +++ b/tests/unit/test__job_helpers.py @@ -554,13 +554,9 @@ def test_query_and_wait_retries_job_times_out(): ) -def test_query_and_wait_sets_job_creation_mode(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv( - "QUERY_PREVIEW_ENABLED", - # The comparison should be case insensitive. - "TrUe", - ) +def test_query_and_wait_sets_job_creation_mode(): client = mock.create_autospec(Client) + client.default_job_creation_mode = "JOB_CREATION_OPTIONAL" client._call_api.return_value = { "jobReference": { "projectId": "response-project", @@ -642,6 +638,7 @@ def test_query_and_wait_sets_location(): "useInt64Timestamp": True, }, "requestId": mock.ANY, + "jobCreationMode": mock.ANY, }, timeout=None, ) @@ -658,6 +655,7 @@ def test_query_and_wait_sets_location(): ) def test_query_and_wait_sets_max_results(max_results, page_size, expected): client = mock.create_autospec(Client) + client.default_job_creation_mode = None client._call_api.return_value = { "jobReference": { "projectId": "response-project", @@ -703,6 +701,7 @@ def test_query_and_wait_sets_max_results(max_results, page_size, expected): def test_query_and_wait_caches_completed_query_results_one_page(): client = mock.create_autospec(Client) + client.default_job_creation_mode = None client._call_api.return_value = { "jobReference": { "projectId": "response-project", @@ -768,6 +767,7 @@ def test_query_and_wait_caches_completed_query_results_one_page(): def test_query_and_wait_caches_completed_query_results_one_page_no_rows(): client = mock.create_autospec(Client) + client.default_job_creation_mode = None client._call_api.return_value = { "jobReference": { "projectId": "response-project", diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 468068321..8ce8d2cbd 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -213,6 +213,17 @@ def test_ctor_w_client_options_universe(self): ) self.assertEqual(client._connection.API_BASE_URL, "https://bigquery.foo.com") + def test_ctor_w_job_creation_mode(self): + creds = _make_credentials() + http = object() + client = self._make_one( + project=self.PROJECT, + credentials=creds, + _http=http, + default_job_creation_mode="foo", + ) + self.assertEqual(client.default_job_creation_mode, "foo") + def test_ctor_w_location(self): from google.cloud.bigquery._http import Connection