feat: job creation mode GA (#2190)

shollyman · web-flow · commit 64cd39fb395c · 2025-05-27T10:53:49.000-07:00
* feat: job creation mode GA

This PR makes the underlying functionality related to how queries can
optionally avoid job creation a GA feature.  It does the following:

* no longer uses the preview QUERY_PREVIEW_ENABLED environment variable
  to control job creation
* adds a new argument to Client instantiation to control job creation
  mode
* adds a property/setter to Client to control job creation mode

This PR also updates/renames the sample demonstrating how to leverage
job creation mode with Client.query_and_wait.
diff --git 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
@@ -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
@@ -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_job_optional.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_job_optional.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
@@ -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
@@ -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
 
