feat: add support for preview features (#2923)

Enables preview query features, which currently only includes stateless queries (queries without jobId). These features won't always be enabled on the service side and there are additional checks and conditions. Fixes https://togithub.com/googleapis/java-bigquery/issues/2949

Author: obada-ab

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java

@@ -1596,6 +1596,16 @@ TableResult listTableData(
  * }
  * </pre>
  *
+ * This method supports query-related preview features via environmental variables (enabled by
+ * setting the {@code QUERY_PREVIEW_ENABLED} environment variable to "TRUE"). Specifically, this
+ * method supports:
+ *
+ * <ul>
+ * <li><b>Stateless queries</b>: query execution without corresponding job metadata
+ * </ul>
+ *
+ * The behaviour of these preview features is controlled by the bigquery service as well
+ *
  * @throws BigQueryException upon failure
  * @throws InterruptedException if the current thread gets interrupted while waiting for the query
  * to complete

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java

@@ -41,6 +41,7 @@
 import com.google.cloud.RetryHelper.RetryHelperException;
 import com.google.cloud.Tuple;
 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.common.annotations.VisibleForTesting;
 import com.google.common.base.Function;
@@ -1324,6 +1325,14 @@ 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();
+ }
+
  // If all parameters passed in configuration are supported by the query() method on the backend,
  // put on fast path
  QueryRequestInfo requestInfo = new QueryRequestInfo(configuration);
@@ -1416,6 +1425,7 @@ public com.google.api.services.bigquery.model.QueryResponse call() {
  public TableResult query(QueryJobConfiguration configuration, JobId jobId, JobOption... options)
  throws InterruptedException, JobException {
  Job.checkNotDryRun(configuration, "query");
+
  // If all parameters passed in configuration are supported by the query() method on the backend,
  // put on fast path
  QueryRequestInfo requestInfo = new QueryRequestInfo(configuration);

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryOptions.java

@@ -24,6 +24,7 @@
 import com.google.cloud.bigquery.spi.v2.BigQueryRpc;
 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;
 
@@ -37,6 +38,7 @@ public class BigQueryOptions extends ServiceOptions<BigQuery, BigQueryOptions> {
  private final String location;
  // set the option ThrowNotFound when you want to throw the exception when the value not found
  private boolean setThrowNotFound;
+ private String queryPreviewEnabled = System.getenv("QUERY_PREVIEW_ENABLED");
 
  public static class DefaultBigQueryFactory implements BigQueryFactory {
 
@@ -130,10 +132,19 @@ public String getLocation() {
  return location;
  }
 
+ public boolean isQueryPreviewEnabled() {
+ return queryPreviewEnabled != null && queryPreviewEnabled.equalsIgnoreCase("TRUE");
+ }
+
  public void setThrowNotFound(boolean setThrowNotFound) {
  this.setThrowNotFound = setThrowNotFound;
  }
 
+ @VisibleForTesting
+ public void setQueryPreviewEnabled(String queryPreviewEnabled) {
+ this.queryPreviewEnabled = queryPreviewEnabled;
+ }
+
  public boolean getThrowNotFound() {
  return setThrowNotFound;
  }

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryJobConfiguration.java

@@ -73,6 +73,7 @@ public final class QueryJobConfiguration extends JobConfiguration {
  private final List<ConnectionProperty> connectionProperties;
  // maxResults is only used for fast query path
  private final Long maxResults;
+ private final JobCreationMode jobCreationMode;
 
  /**
  * Priority levels for a query. If not specified the priority is assumed to be {@link
@@ -94,6 +95,21 @@ public enum Priority {
  BATCH
  }
 
+ /** Job Creation Mode provides different options on job creation. */
+ enum JobCreationMode {
+ /** Unspecified JobCreationMode, defaults to JOB_CREATION_REQUIRED. */
+ JOB_CREATION_MODE_UNSPECIFIED,
+ /** Default. Job creation is always required. */
+ JOB_CREATION_REQUIRED,
+ /**
+ * 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. If Job creation is required,
+ * JOB_CREATION_REQUIRED mode should be used, which is the default.
+ */
+ JOB_CREATION_OPTIONAL,
+ }
+
  public static final class Builder
  extends JobConfiguration.Builder<QueryJobConfiguration, Builder> {
 
@@ -125,6 +141,7 @@ public static final class Builder
  private RangePartitioning rangePartitioning;
  private List<ConnectionProperty> connectionProperties;
  private Long maxResults;
+ private JobCreationMode jobCreationMode;
 
  private Builder() {
  super(Type.QUERY);
@@ -160,6 +177,7 @@ private Builder(QueryJobConfiguration jobConfiguration) {
  this.rangePartitioning = jobConfiguration.rangePartitioning;
  this.connectionProperties = jobConfiguration.connectionProperties;
  this.maxResults = jobConfiguration.maxResults;
+ this.jobCreationMode = jobConfiguration.jobCreationMode;
  }
 
  private Builder(com.google.api.services.bigquery.model.JobConfiguration configurationPb) {
@@ -655,6 +673,15 @@ public Builder setMaxResults(Long maxResults) {
  return this;
  }
 
+ /**
+ * 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) {
+ this.jobCreationMode = jobCreationMode;
+ return this;
+ }
+
  public QueryJobConfiguration build() {
  return new QueryJobConfiguration(this);
  }
@@ -699,6 +726,7 @@ private QueryJobConfiguration(Builder builder) {
  this.rangePartitioning = builder.rangePartitioning;
  this.connectionProperties = builder.connectionProperties;
  this.maxResults = builder.maxResults;
+ this.jobCreationMode = builder.jobCreationMode;
  }
 
  /**
@@ -910,6 +938,11 @@ public Long getMaxResults() {
  return maxResults;
  }
 
+ /** Returns the job creation mode. */
+ JobCreationMode getJobCreationMode() {
+ return jobCreationMode;
+ }
+
  @Override
  public Builder toBuilder() {
  return new Builder(this);
@@ -944,7 +977,8 @@ ToStringHelper toStringHelper() {
  .add("jobTimeoutMs", jobTimeoutMs)
  .add("labels", labels)
  .add("rangePartitioning", rangePartitioning)
- .add("connectionProperties", connectionProperties);
+ .add("connectionProperties", connectionProperties)
+ .add("jobCreationMode", jobCreationMode);
  }
 
  @Override

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java

@@ -18,6 +18,7 @@
 
 import com.google.api.services.bigquery.model.QueryParameter;
 import com.google.api.services.bigquery.model.QueryRequest;
+import com.google.cloud.bigquery.QueryJobConfiguration.JobCreationMode;
 import com.google.common.base.MoreObjects;
 import com.google.common.base.Objects;
 import com.google.common.collect.Lists;
@@ -40,6 +41,7 @@ final class QueryRequestInfo {
  private final Boolean createSession;
  private final Boolean useQueryCache;
  private final Boolean useLegacySql;
+ private final JobCreationMode jobCreationMode;
 
  QueryRequestInfo(QueryJobConfiguration config) {
  this.config = config;
@@ -55,6 +57,7 @@ final class QueryRequestInfo {
  this.createSession = config.createSession();
  this.useLegacySql = config.useLegacySql();
  this.useQueryCache = config.useQueryCache();
+ this.jobCreationMode = config.getJobCreationMode();
  }
 
  boolean isFastQuerySupported(JobId jobId) {
@@ -116,6 +119,9 @@ QueryRequest toPb() {
  if (useQueryCache != null) {
  request.setUseQueryCache(useQueryCache);
  }
+ if (jobCreationMode != null) {
+ request.setJobCreationMode(jobCreationMode.toString());
+ }
  return request;
  }
 
@@ -134,6 +140,7 @@ public String toString() {
  .add("createSession", createSession)
  .add("useQueryCache", useQueryCache)
  .add("useLegacySql", useLegacySql)
+ .add("jobCreationMode", jobCreationMode)
  .toString();
  }
 
@@ -151,7 +158,8 @@ public int hashCode() {
  requestId,
  createSession,
  useQueryCache,
- useLegacySql);
+ useLegacySql,
+ jobCreationMode);
  }
 
  @Override

google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryJobConfigurationTest.java

@@ -23,6 +23,7 @@
 import com.google.cloud.bigquery.JobInfo.CreateDisposition;
 import com.google.cloud.bigquery.JobInfo.SchemaUpdateOption;
 import com.google.cloud.bigquery.JobInfo.WriteDisposition;
+import com.google.cloud.bigquery.QueryJobConfiguration.JobCreationMode;
 import com.google.cloud.bigquery.QueryJobConfiguration.Priority;
 import com.google.cloud.bigquery.TimePartitioning.Type;
 import com.google.common.collect.ImmutableList;
@@ -110,6 +111,7 @@ public class QueryJobConfigurationTest {
  private static final Map<String, QueryParameterValue> NAME_PARAMETER =
  ImmutableMap.of("string", STRING_PARAMETER, "timestamp", TIMESTAMP_PARAMETER);
  private static final String PARAMETER_MODE = "POSITIONAL";
+ private static final JobCreationMode JOB_CREATION_MODE = JobCreationMode.JOB_CREATION_OPTIONAL;
  private static final QueryJobConfiguration QUERY_JOB_CONFIGURATION =
  QueryJobConfiguration.newBuilder(QUERY)
  .setUseQueryCache(USE_QUERY_CACHE)
@@ -150,6 +152,8 @@ public class QueryJobConfigurationTest {
  .setPositionalParameters(ImmutableList.<QueryParameterValue>of())
  .setNamedParameters(NAME_PARAMETER)
  .build();
+ private static final QueryJobConfiguration QUERY_JOB_CONFIGURATION_SET_JOB_CREATION_MODE =
+ QUERY_JOB_CONFIGURATION.toBuilder().setJobCreationMode(JOB_CREATION_MODE).build();
 
  @Test
  public void testToBuilder() {
@@ -230,6 +234,13 @@ public void testNamedParameter() {
  QUERY_JOB_CONFIGURATION_SET_NAME_PARAMETER.toBuilder().build());
  }
 
+ @Test
+ public void testJobCreationMode() {
+ compareQueryJobConfiguration(
+ QUERY_JOB_CONFIGURATION_SET_JOB_CREATION_MODE,
+ QUERY_JOB_CONFIGURATION_SET_JOB_CREATION_MODE.toBuilder().build());
+ }
+
  private void compareQueryJobConfiguration(
  QueryJobConfiguration expected, QueryJobConfiguration value) {
  assertEquals(expected, value);

google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryRequestInfoTest.java

@@ -23,6 +23,7 @@
 import com.google.cloud.bigquery.JobInfo.CreateDisposition;
 import com.google.cloud.bigquery.JobInfo.SchemaUpdateOption;
 import com.google.cloud.bigquery.JobInfo.WriteDisposition;
+import com.google.cloud.bigquery.QueryJobConfiguration.JobCreationMode;
 import com.google.cloud.bigquery.QueryJobConfiguration.Priority;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -105,6 +106,8 @@ public class QueryRequestInfoTest {
  ImmutableList.of(STRING_PARAMETER, TIMESTAMP_PARAMETER);
  private static final Map<String, QueryParameterValue> NAME_PARAMETER =
  ImmutableMap.of("string", STRING_PARAMETER, "timestamp", TIMESTAMP_PARAMETER);
+ private static final JobCreationMode jobCreationModeRequired =
+ JobCreationMode.JOB_CREATION_REQUIRED;
  private static final QueryJobConfiguration QUERY_JOB_CONFIGURATION =
  QueryJobConfiguration.newBuilder(QUERY)
  .setUseQueryCache(USE_QUERY_CACHE)
@@ -131,6 +134,7 @@ public class QueryRequestInfoTest {
  .setConnectionProperties(CONNECTION_PROPERTIES)
  .setPositionalParameters(POSITIONAL_PARAMETER)
  .setMaxResults(100L)
+ .setJobCreationMode(jobCreationModeRequired)
  .build();
  QueryRequestInfo REQUEST_INFO = new QueryRequestInfo(QUERY_JOB_CONFIGURATION);
  private static final QueryJobConfiguration QUERY_JOB_CONFIGURATION_SUPPORTED =
@@ -194,5 +198,6 @@ private void compareQueryRequestInfo(QueryRequestInfo expected, QueryRequestInfo
  assertEquals(expectedQueryReq.getCreateSession(), actualQueryReq.getCreateSession());
  assertEquals(expectedQueryReq.getUseQueryCache(), actualQueryReq.getUseQueryCache());
  assertEquals(expectedQueryReq.getUseLegacySql(), actualQueryReq.getUseLegacySql());
+ assertEquals(expectedQueryReq.get("jobCreationMode"), actualQueryReq.get("jobCreationMode"));
  }
 }

google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java

@@ -6188,4 +6188,30 @@ public void testAlreadyExistJobExceptionHandling() throws InterruptedException {
  }
  }
  }
+
+ @Test
+ public void testStatelessQueries() throws InterruptedException {
+ // simulate setting the QUERY_PREVIEW_ENABLED environment variable
+ bigquery.getOptions().setQueryPreviewEnabled("TRUE");
+ assertNull(executeSimpleQuery().getJobId());
+
+ // the flag should be case-insensitive
+ bigquery.getOptions().setQueryPreviewEnabled("tRuE");
+ assertNull(executeSimpleQuery().getJobId());
+
+ // any other values won't enable optional job creation mode
+ bigquery.getOptions().setQueryPreviewEnabled("test_value");
+ assertNotNull(executeSimpleQuery().getJobId());
+
+ // reset the flag
+ bigquery.getOptions().setQueryPreviewEnabled(null);
+ assertNotNull(executeSimpleQuery().getJobId());
+ }
+
+ private TableResult executeSimpleQuery() throws InterruptedException {
+ String query = "SELECT 1 as one";
+ QueryJobConfiguration config = QueryJobConfiguration.newBuilder(query).build();
+ TableResult result = bigquery.query(config);
+ return result;
+ }
 }