Skip to content

Commit

Permalink
feat: add support for preview features (#2923)
Browse files Browse the repository at this point in the history
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
  • Loading branch information
obada-ab committed Oct 24, 2023
1 parent f452cf4 commit 113b8f2
Show file tree
Hide file tree
Showing 8 changed files with 117 additions and 2 deletions.
Expand Up @@ -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
Expand Down
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
Expand Down
Expand Up @@ -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;

Expand All @@ -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 {

Expand Down Expand Up @@ -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;
}
Expand Down
Expand Up @@ -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
Expand All @@ -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> {

Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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);
}
Expand Down Expand Up @@ -699,6 +726,7 @@ private QueryJobConfiguration(Builder builder) {
this.rangePartitioning = builder.rangePartitioning;
this.connectionProperties = builder.connectionProperties;
this.maxResults = builder.maxResults;
this.jobCreationMode = builder.jobCreationMode;
}

/**
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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
Expand Down
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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) {
Expand Down Expand Up @@ -116,6 +119,9 @@ QueryRequest toPb() {
if (useQueryCache != null) {
request.setUseQueryCache(useQueryCache);
}
if (jobCreationMode != null) {
request.setJobCreationMode(jobCreationMode.toString());
}
return request;
}

Expand All @@ -134,6 +140,7 @@ public String toString() {
.add("createSession", createSession)
.add("useQueryCache", useQueryCache)
.add("useLegacySql", useLegacySql)
.add("jobCreationMode", jobCreationMode)
.toString();
}

Expand All @@ -151,7 +158,8 @@ public int hashCode() {
requestId,
createSession,
useQueryCache,
useLegacySql);
useLegacySql,
jobCreationMode);
}

@Override
Expand Down
Expand Up @@ -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;
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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() {
Expand Down Expand Up @@ -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);
Expand Down
Expand Up @@ -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;
Expand Down Expand Up @@ -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)
Expand All @@ -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 =
Expand Down Expand Up @@ -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"));
}
}
Expand Up @@ -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;
}
}

0 comments on commit 113b8f2

Please sign in to comment.