diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java index ac6262e69..80fd6618d 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java @@ -1596,6 +1596,16 @@ TableResult listTableData( * } * * + * 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: + * + * + * + * 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 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 index ef7e8cb8b..0d5842724 100644 --- 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 @@ -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); 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 index 2e22ba922..e53439f02 100644 --- 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 @@ -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 { 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; } 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 index cc726bdd1..0ad85137b 100644 --- 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 @@ -73,6 +73,7 @@ public final class QueryJobConfiguration extends JobConfiguration { private final List 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 { @@ -125,6 +141,7 @@ public static final class Builder private RangePartitioning rangePartitioning; private List 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 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 index 00a898363..00a11f723 100644 --- 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 @@ -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 diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryJobConfigurationTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryJobConfigurationTest.java index 9a20219d6..f71e152e6 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryJobConfigurationTest.java +++ b/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 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.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); diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryRequestInfoTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryRequestInfoTest.java index 456475597..0d9464c76 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryRequestInfoTest.java +++ b/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 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")); } } 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 index 909500be8..8cada3e08 100644 --- 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 @@ -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; + } }