From 113b8f27419365c7277c6a300c5f07cea954cca1 Mon Sep 17 00:00:00 2001 From: Obada Alabbadi <76101898+obada-ab@users.noreply.github.com> Date: Tue, 24 Oct 2023 12:22:14 +0300 Subject: [PATCH] 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 --- .../com/google/cloud/bigquery/BigQuery.java | 10 ++++++ .../google/cloud/bigquery/BigQueryImpl.java | 10 ++++++ .../cloud/bigquery/BigQueryOptions.java | 11 ++++++ .../cloud/bigquery/QueryJobConfiguration.java | 36 ++++++++++++++++++- .../cloud/bigquery/QueryRequestInfo.java | 10 +++++- .../bigquery/QueryJobConfigurationTest.java | 11 ++++++ .../cloud/bigquery/QueryRequestInfoTest.java | 5 +++ .../cloud/bigquery/it/ITBigQueryTest.java | 26 ++++++++++++++ 8 files changed, 117 insertions(+), 2 deletions(-) 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; + } }