Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: request id for fast query path #781

Merged
merged 1 commit into from Oct 6, 2020
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
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
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
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