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

* feat(bigquery): job creation mode GA * deprecate query preview * Update samples * rename snippet to job optional * fix style * rename sample test

Author: whuffman36

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

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() {

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;
  }
 

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() {

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) {

samples/snippets/src/main/java/com/example/bigquery/QueryShortMode.java renamed to 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]

samples/snippets/src/test/java/com/example/bigquery/QueryShortModeIT.java renamed to 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");
  }
 }