fix: request id for fast query path (#781)

Praful Makani · web-flow · commit 15dc9f0a95eb · 2020-10-06T11:45:28.000-04:00
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
@@ -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<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
@@ -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);
@@ -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
@@ -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;
@@ -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;
@@ -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
@@ -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 =
