Skip to content

Commit

Permalink
fix: request id for fast query path (#781)
Browse files Browse the repository at this point in the history
  • Loading branch information
Praful Makani authored Oct 6, 2020
1 parent 1778e74 commit 15dc9f0
Show file tree
Hide file tree
Showing 3 changed files with 43 additions and 6 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -71,6 +72,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 String requestId;

/**
* Priority levels for a query. If not specified the priority is assumed to be {@link
Expand Down Expand Up @@ -121,6 +123,7 @@ public static final class Builder
private RangePartitioning rangePartitioning;
private List<ConnectionProperty> connectionProperties;
private Long maxResults;
private String requestId;

private Builder() {
super(Type.QUERY);
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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);
}
Expand Down Expand Up @@ -663,6 +672,7 @@ private QueryJobConfiguration(Builder builder) {
this.rangePartitioning = builder.rangePartitioning;
this.connectionProperties = builder.connectionProperties;
this.maxResults = builder.maxResults;
this.requestId = builder.requestId;
}

/**
Expand Down Expand Up @@ -865,6 +875,10 @@ public Long getMaxResults() {
return maxResults;
}

String getRequestId() {
return requestId;
}

@Override
public Builder toBuilder() {
return new Builder(this);
Expand Down Expand Up @@ -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());
}

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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<ConnectionProperty> connectionProperties;
private final DatasetId defaultDataset;
Expand All @@ -39,6 +37,7 @@ final class QueryRequestInfo {
private final List<QueryParameter> queryParameters;
private final Boolean useQueryCache;
private final Boolean useLegacySql;
private final String requestId;

QueryRequestInfo(QueryJobConfiguration config) {
this.config = config;
Expand All @@ -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() {
Expand Down Expand Up @@ -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);
}
Expand All @@ -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)
Expand All @@ -133,7 +133,7 @@ public int hashCode() {
maxResults,
query,
queryParameters,
REQUEST_ID,
requestId,
useQueryCache,
useLegacySql);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 =
Expand Down

0 comments on commit 15dc9f0

Please sign in to comment.