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 927ae846c..88ca5aa44 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 @@ -35,6 +35,7 @@ import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.UUID; /** * Google BigQuery Query Job configuration. A Query Job runs a query against BigQuery data. Query @@ -71,6 +72,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 String requestId; /** * Priority levels for a query. If not specified the priority is assumed to be {@link @@ -121,6 +123,7 @@ public static final class Builder private RangePartitioning rangePartitioning; private List connectionProperties; private Long maxResults; + private String requestId; private Builder() { super(Type.QUERY); @@ -154,6 +157,7 @@ private Builder(QueryJobConfiguration jobConfiguration) { this.rangePartitioning = jobConfiguration.rangePartitioning; this.connectionProperties = jobConfiguration.connectionProperties; this.maxResults = jobConfiguration.maxResults; + this.requestId = jobConfiguration.requestId; } private Builder(com.google.api.services.bigquery.model.JobConfiguration configurationPb) { @@ -621,6 +625,11 @@ public Builder setMaxResults(Long maxResults) { return this; } + Builder setRequestId(String requestId) { + this.requestId = requestId; + return this; + } + public QueryJobConfiguration build() { return new QueryJobConfiguration(this); } @@ -663,6 +672,7 @@ private QueryJobConfiguration(Builder builder) { this.rangePartitioning = builder.rangePartitioning; this.connectionProperties = builder.connectionProperties; this.maxResults = builder.maxResults; + this.requestId = builder.requestId; } /** @@ -865,6 +875,10 @@ public Long getMaxResults() { return maxResults; } + String getRequestId() { + return requestId; + } + @Override public Builder toBuilder() { return new Builder(this); @@ -1043,7 +1057,7 @@ com.google.api.services.bigquery.model.JobConfiguration toPb() { /** Creates a builder for a BigQuery Query Job given the query to be run. */ public static Builder newBuilder(String query) { checkArgument(!isNullOrEmpty(query), "Provided query is null or empty"); - return new Builder().setQuery(query); + return new Builder().setQuery(query).setRequestId(UUID.randomUUID().toString()); } /** 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 2882e788b..602404cc0 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 @@ -23,11 +23,9 @@ import com.google.common.collect.Lists; import java.util.List; import java.util.Map; -import java.util.UUID; final class QueryRequestInfo { - private static final String REQUEST_ID = UUID.randomUUID().toString(); private QueryJobConfiguration config; private final List connectionProperties; private final DatasetId defaultDataset; @@ -39,6 +37,7 @@ final class QueryRequestInfo { private final List queryParameters; private final Boolean useQueryCache; private final Boolean useLegacySql; + private final String requestId; QueryRequestInfo(QueryJobConfiguration config) { this.config = config; @@ -52,6 +51,7 @@ final class QueryRequestInfo { this.queryParameters = config.toPb().getQuery().getQueryParameters(); this.useLegacySql = config.useLegacySql(); this.useQueryCache = config.useQueryCache(); + this.requestId = config.getRequestId(); } boolean isFastQuerySupported() { @@ -92,7 +92,7 @@ QueryRequest toPb() { request.setMaxResults(maxResults); } request.setQuery(query); - request.setRequestId(REQUEST_ID); + request.setRequestId(requestId); if (queryParameters != null) { request.setQueryParameters(queryParameters); } @@ -115,7 +115,7 @@ public String toString() { .add("maximumBytesBilled", maximumBytesBilled) .add("maxResults", maxResults) .add("query", query) - .add("requestId", REQUEST_ID) + .add("requestId", requestId) .add("queryParameters", queryParameters) .add("useQueryCache", useQueryCache) .add("useLegacySql", useLegacySql) @@ -133,7 +133,7 @@ public int hashCode() { maxResults, query, queryParameters, - REQUEST_ID, + requestId, useQueryCache, useLegacySql); } 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 d7497b918..14b4f7a88 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 @@ -1560,6 +1560,29 @@ public void testQueryCaseInsensitiveSchemaFieldByGetName() throws InterruptedExc assertEquals(2, rowCount); } + @Test + public void testFastQueryMultipleRuns() throws InterruptedException { + String query = + "SELECT TimestampField, StringField, BooleanField FROM " + TABLE_ID_FASTQUERY.getTable(); + QueryJobConfiguration config = + QueryJobConfiguration.newBuilder(query).setDefaultDataset(DatasetId.of(DATASET)).build(); + TableResult result = bigquery.query(config); + assertEquals(QUERY_RESULT_SCHEMA, result.getSchema()); + assertEquals(2, result.getTotalRows()); + assertNull(result.getNextPage()); + assertNull(result.getNextPageToken()); + assertFalse(result.hasNextPage()); + + QueryJobConfiguration config2 = + QueryJobConfiguration.newBuilder(query).setDefaultDataset(DatasetId.of(DATASET)).build(); + TableResult result2 = bigquery.query(config2); + assertEquals(QUERY_RESULT_SCHEMA, result2.getSchema()); + assertEquals(2, result2.getTotalRows()); + assertNull(result2.getNextPage()); + assertNull(result2.getNextPageToken()); + assertFalse(result2.hasNextPage()); + } + @Test public void testFastSQLQuery() throws InterruptedException { String query =