feat(bigquery): job creation mode GA (#3804)

whuffman36 · web-flow · commit a21cde8994e9 · 2025-05-30T14:18:20.000-07:00
* feat(bigquery): job creation mode GA

* deprecate query preview

* Update samples

* rename snippet to job optional

* fix style

* rename sample test
diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java
@@ -40,7 +40,6 @@
 import com.google.cloud.Tuple;
 import com.google.cloud.bigquery.BigQueryRetryHelper.BigQueryRetryHelperException;
 import com.google.cloud.bigquery.InsertAllRequest.RowToInsert;
-import com.google.cloud.bigquery.QueryJobConfiguration.JobCreationMode;
 import com.google.cloud.bigquery.spi.v2.BigQueryRpc;
 import com.google.cloud.bigquery.spi.v2.HttpBigQueryRpc;
 import com.google.common.annotations.VisibleForTesting;
@@ -1407,12 +1406,10 @@ public TableResult query(QueryJobConfiguration configuration, JobOption... optio
       throws InterruptedException, JobException {
     Job.checkNotDryRun(configuration, "query");
 
-    if (getOptions().isQueryPreviewEnabled()) {
-      configuration =
-          configuration.toBuilder()
-              .setJobCreationMode(JobCreationMode.JOB_CREATION_OPTIONAL)
-              .build();
-    }
+    configuration =
+        configuration.toBuilder()
+            .setJobCreationMode(getOptions().getDefaultJobCreationMode())
+            .build();
 
     // If all parameters passed in configuration are supported by the query() method on the backend,
     // put on fast path
diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryOptions.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryOptions.java
@@ -20,10 +20,10 @@
 import com.google.cloud.ServiceOptions;
 import com.google.cloud.ServiceRpc;
 import com.google.cloud.TransportOptions;
+import com.google.cloud.bigquery.QueryJobConfiguration.JobCreationMode;
 import com.google.cloud.bigquery.spi.BigQueryRpcFactory;
 import com.google.cloud.bigquery.spi.v2.HttpBigQueryRpc;
 import com.google.cloud.http.HttpTransportOptions;
-import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableSet;
 import java.util.Set;
 
@@ -38,7 +38,7 @@ public class BigQueryOptions extends ServiceOptions<BigQuery, BigQueryOptions> {
   // set the option ThrowNotFound when you want to throw the exception when the value not found
   private boolean setThrowNotFound;
   private boolean useInt64Timestamps;
-  private String queryPreviewEnabled = System.getenv("QUERY_PREVIEW_ENABLED");
+  private JobCreationMode defaultJobCreationMode = JobCreationMode.JOB_CREATION_MODE_UNSPECIFIED;
 
   public static class DefaultBigQueryFactory implements BigQueryFactory {
 
@@ -139,8 +139,9 @@ public String getLocation() {
     return location;
   }
 
+  @Deprecated
   public boolean isQueryPreviewEnabled() {
-    return queryPreviewEnabled != null && queryPreviewEnabled.equalsIgnoreCase("TRUE");
+    return false;
   }
 
   public void setThrowNotFound(boolean setThrowNotFound) {
@@ -151,9 +152,11 @@ public void setUseInt64Timestamps(boolean useInt64Timestamps) {
     this.useInt64Timestamps = useInt64Timestamps;
   }
 
-  @VisibleForTesting
-  public void setQueryPreviewEnabled(String queryPreviewEnabled) {
-    this.queryPreviewEnabled = queryPreviewEnabled;
+  @Deprecated
+  public void setQueryPreviewEnabled(String queryPreviewEnabled) {}
+
+  public void setDefaultJobCreationMode(JobCreationMode jobCreationMode) {
+    this.defaultJobCreationMode = jobCreationMode;
   }
 
   public boolean getThrowNotFound() {
@@ -164,6 +167,10 @@ public boolean getUseInt64Timestamps() {
     return useInt64Timestamps;
   }
 
+  public JobCreationMode getDefaultJobCreationMode() {
+    return defaultJobCreationMode;
+  }
+
   @SuppressWarnings("unchecked")
   @Override
   public Builder toBuilder() {
diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryJobConfiguration.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryJobConfiguration.java
@@ -97,7 +97,7 @@ public enum Priority {
   }
 
   /** Job Creation Mode provides different options on job creation. */
-  enum JobCreationMode {
+  public enum JobCreationMode {
     /** Unspecified JobCreationMode, defaults to JOB_CREATION_REQUIRED. */
     JOB_CREATION_MODE_UNSPECIFIED,
     /** Default. Job creation is always required. */
@@ -683,7 +683,7 @@ public Builder setMaxResults(Long maxResults) {
      * Provides different options on job creation. If not specified the job creation mode is assumed
      * to be {@link JobCreationMode#JOB_CREATION_REQUIRED}.
      */
-    Builder setJobCreationMode(JobCreationMode jobCreationMode) {
+    public Builder setJobCreationMode(JobCreationMode jobCreationMode) {
       this.jobCreationMode = jobCreationMode;
       return this;
     }
@@ -959,7 +959,7 @@ public Long getMaxResults() {
   }
 
   /** Returns the job creation mode. */
-  JobCreationMode getJobCreationMode() {
+  public JobCreationMode getJobCreationMode() {
     return jobCreationMode;
   }
 
diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java
@@ -86,7 +86,8 @@ boolean isFastQuerySupported(JobId jobId) {
         && config.getTableDefinitions() == null
         && config.getTimePartitioning() == null
         && config.getUserDefinedFunctions() == null
-        && config.getWriteDisposition() == null;
+        && config.getWriteDisposition() == null
+        && config.getJobCreationMode() != JobCreationMode.JOB_CREATION_REQUIRED;
   }
 
   QueryRequest toPb() {
diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java
@@ -111,6 +111,7 @@
 import com.google.cloud.bigquery.PolicyTags;
 import com.google.cloud.bigquery.PrimaryKey;
 import com.google.cloud.bigquery.QueryJobConfiguration;
+import com.google.cloud.bigquery.QueryJobConfiguration.JobCreationMode;
 import com.google.cloud.bigquery.QueryJobConfiguration.Priority;
 import com.google.cloud.bigquery.QueryParameterValue;
 import com.google.cloud.bigquery.Range;
@@ -3483,6 +3484,7 @@ public void testExecuteSelectWithReadApi() throws SQLException {
     final int rowLimit = 5000;
     final String QUERY =
         "SELECT * FROM bigquery-public-data.new_york_taxi_trips.tlc_yellow_trips_2017 LIMIT %s";
+    bigquery.getOptions().setDefaultJobCreationMode(JobCreationMode.JOB_CREATION_REQUIRED);
     // Job timeout is somewhat arbitrary - just ensures that fast query is not used.
     // min result size and page row count ratio ensure that the ReadAPI is used.
     ConnectionSettings connectionSettingsReadAPIEnabledFastQueryDisabled =
@@ -7084,26 +7086,19 @@ public void testStatelessQueries() throws InterruptedException {
     RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create();
     BigQuery bigQuery = bigqueryHelper.getOptions().getService();
 
-    // Simulate setting the QUERY_PREVIEW_ENABLED environment variable.
-    bigQuery.getOptions().setQueryPreviewEnabled("TRUE");
+    // Stateless query should have no job id.
+    bigQuery.getOptions().setDefaultJobCreationMode(JobCreationMode.JOB_CREATION_OPTIONAL);
     TableResult tableResult = executeSimpleQuery(bigQuery);
     assertNotNull(tableResult.getQueryId());
     assertNull(tableResult.getJobId());
 
-    // The flag should be case-insensitive.
-    bigQuery.getOptions().setQueryPreviewEnabled("tRuE");
+    // Job creation takes over, no query id is created.
+    bigQuery.getOptions().setDefaultJobCreationMode(JobCreationMode.JOB_CREATION_REQUIRED);
     tableResult = executeSimpleQuery(bigQuery);
-    assertNotNull(tableResult.getQueryId());
-    assertNull(tableResult.getJobId());
-
-    // Any other values won't enable optional job creation mode.
-    bigQuery.getOptions().setQueryPreviewEnabled("test_value");
-    tableResult = executeSimpleQuery(bigQuery);
-    assertNotNull(tableResult.getQueryId());
+    assertNull(tableResult.getQueryId());
     assertNotNull(tableResult.getJobId());
 
-    // Reset the flag.
-    bigQuery.getOptions().setQueryPreviewEnabled(null);
+    bigQuery.getOptions().setDefaultJobCreationMode(JobCreationMode.JOB_CREATION_MODE_UNSPECIFIED);
     tableResult = executeSimpleQuery(bigQuery);
     assertNotNull(tableResult.getQueryId());
     assertNotNull(tableResult.getJobId());
@@ -7128,8 +7123,8 @@ public void testTableResultJobIdAndQueryId() throws InterruptedException {
     // Create local BigQuery for test scenario 1 to not contaminate global test parameters.
     RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create();
     BigQuery bigQuery = bigqueryHelper.getOptions().getService();
-    // Simulate setting the QUERY_PREVIEW_ENABLED environment variable.
-    bigQuery.getOptions().setQueryPreviewEnabled("TRUE");
+    // Allow queries to be stateless.
+    bigQuery.getOptions().setDefaultJobCreationMode(JobCreationMode.JOB_CREATION_OPTIONAL);
     String query = "SELECT 1 as one";
     QueryJobConfiguration configStateless = QueryJobConfiguration.newBuilder(query).build();
     TableResult result = bigQuery.query(configStateless);
@@ -7181,15 +7176,17 @@ public void testStatelessQueriesWithLocation() throws Exception {
               table.getTableId().getTable());
 
       // Test stateless query when BigQueryOption location matches dataset location.
-      bigQuery.getOptions().setQueryPreviewEnabled("TRUE");
+      bigQuery.getOptions().setDefaultJobCreationMode(JobCreationMode.JOB_CREATION_OPTIONAL);
       TableResult tb = bigQuery.query(QueryJobConfiguration.of(query));
       assertNull(tb.getJobId());
 
       // Test stateless query when BigQueryOption location does not match dataset location.
       try {
         BigQuery bigQueryWrongLocation =
             bigqueryHelper.getOptions().toBuilder().setLocation(wrongLocation).build().getService();
-        bigQueryWrongLocation.getOptions().setQueryPreviewEnabled("TRUE");
+        bigQueryWrongLocation
+            .getOptions()
+            .setDefaultJobCreationMode(JobCreationMode.JOB_CREATION_OPTIONAL);
         bigQueryWrongLocation.query(QueryJobConfiguration.of(query));
         fail("querying a table with wrong location shouldn't work");
       } catch (BigQueryException e) {
diff --git a/samples/snippets/src/main/java/com/example/bigquery/QueryJobOptional.java b/samples/snippets/src/main/java/com/example/bigquery/QueryJobOptional.java
@@ -16,34 +16,37 @@
 
 package com.example.bigquery;
 
-// [START bigquery_query_shortquery]
+// [START bigquery_query_job_optional]
 import com.google.cloud.bigquery.BigQuery;
 import com.google.cloud.bigquery.BigQueryException;
 import com.google.cloud.bigquery.BigQueryOptions;
 import com.google.cloud.bigquery.JobId;
 import com.google.cloud.bigquery.QueryJobConfiguration;
+import com.google.cloud.bigquery.QueryJobConfiguration.JobCreationMode;
 import com.google.cloud.bigquery.TableResult;
 
 // Sample demonstrating short mode query execution.
 //
-// While this feature is still in preview, it is controlled by
-// setting the environment variable QUERY_PREVIEW_ENABLED=TRUE
-// to request short mode execution.
-public class QueryShortMode {
+// This feature is controlled by setting the defaultJobCreationMode
+// field in the BigQueryOptions used for the client. JOB_CREATION_OPTIONAL
+// allows for the execution of queries without creating a job.
+public class QueryJobOptional {
 
   public static void main(String[] args) {
     String query =
         "SELECT name, gender, SUM(number) AS total FROM "
             + "bigquery-public-data.usa_names.usa_1910_2013 GROUP BY "
             + "name, gender ORDER BY total DESC LIMIT 10";
-    queryShortMode(query);
+    queryJobOptional(query);
   }
 
-  public static void queryShortMode(String query) {
+  public static void queryJobOptional(String query) {
     try {
       // Initialize client that will be used to send requests. This client only needs
       // to be created once, and can be reused for multiple requests.
-      BigQuery bigquery = BigQueryOptions.getDefaultInstance().getService();
+      BigQueryOptions options = BigQueryOptions.getDefaultInstance();
+      options.setDefaultJobCreationMode(JobCreationMode.JOB_CREATION_OPTIONAL);
+      BigQuery bigquery = options.getService();
 
       // Execute the query. The returned TableResult provides access information
       // about the query execution as well as query results.
@@ -72,4 +75,4 @@ public static void queryShortMode(String query) {
     }
   }
 }
-// [END bigquery_query_shortquery]
+// [END bigquery_query_job_optional]
diff --git a/samples/snippets/src/test/java/com/example/bigquery/QueryJobOptionalIT.java b/samples/snippets/src/test/java/com/example/bigquery/QueryJobOptionalIT.java
@@ -26,7 +26,7 @@
 import org.junit.Before;
 import org.junit.Test;
 
-public class QueryShortModeIT {
+public class QueryJobOptionalIT {
 
   private final Logger log = Logger.getLogger(this.getClass().getName());
   private ByteArrayOutputStream bout;
@@ -56,7 +56,7 @@ public void testQueryBatch() {
             + "bigquery-public-data.usa_names.usa_1910_2013 GROUP BY "
             + "name, gender ORDER BY total DESC LIMIT 10";
 
-    QueryShortMode.queryShortMode(query);
+    QueryJobOptional.queryJobOptional(query);
     assertThat(bout.toString()).contains("Query was run");
   }
 }
