fix(bigquery): Update fast query path to allow jobTimeoutMs in the request (#13502)

Author: lqiu96

java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java

@@ -2169,11 +2169,15 @@ && getOptions().getOpenTelemetryTracer() != null) {
               .startSpan();
     }
     try (Scope queryScope = querySpan != null ? querySpan.makeCurrent() : null) {
-      // If all parameters passed in configuration are supported by the query() method on the
-      // backend, put on fast path
+      // The fast query path (jobs.query API) is preferred to reduce latency by avoiding
+      // the slow fallback path (jobs.insert API). We will opt to use it if the configuration
+      // and JobId allow (i.e. if all parameters passed in configuration are supported).
       QueryRequestInfo requestInfo =
           new QueryRequestInfo(configuration, getOptions().getDataFormatOptions());
-      if (requestInfo.isFastQuerySupported(jobId)) {
+      // Fast query path is not possible if job is specified in the JobID object.
+      // Respect Job field value in JobId specified by user.
+      // Specifying it will force the query to take the slower path.
+      if (requestInfo.isFastQuerySupported() && (jobId == null || jobId.getJob() == null)) {
         // Be careful when setting the projectID in JobId, if a projectID is specified in the JobId,
         // the job created by the query method will use that project. This may cause the query to
         // fail with "Access denied" if the project do not have enough permissions to run the job.

java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ConnectionImpl.java

@@ -245,7 +245,8 @@ private BigQueryResult getExecuteSelectResponse(
       labelMap = labels[0];
     }
     try {
-      // use jobs.query if possible
+      // The fast query path (jobs.query API) is preferred to reduce latency by avoiding
+      // the slow fallback path (jobs.insert API). We will opt to use it if possible.
       if (isFastQuerySupported()) {
         logger.log(Level.INFO, "\n Using Fast Query Path");
         final String projectId = bigQueryOptions.getProjectId();
@@ -810,7 +811,8 @@ void flagEndOfStream() { // package-private
             Level.WARNING,
             "\n"
                 + Thread.currentThread().getName()
-                + " Could not flag End of Stream, both the buffer types are null. This might happen when the connection is close without executing a query");
+                + " Could not flag End of Stream, both the buffer types are null. This might happen"
+                + " when the connection is close without executing a query");
       }
     } catch (InterruptedException e) {
       logger.log(
@@ -1260,7 +1262,6 @@ boolean isFastQuerySupported() {
         && connectionSettings.getCreateDisposition() == null
         && connectionSettings.getDestinationEncryptionConfiguration() == null
         && connectionSettings.getDestinationTable() == null
-        && connectionSettings.getJobTimeoutMs() == null
         && connectionSettings.getMaximumBillingTier() == null
         && connectionSettings.getPriority() == null
         && connectionSettings.getRangePartitioning() == null
@@ -1361,6 +1362,9 @@ QueryRequest createQueryRequest(
     content.setRequestId(requestId);
     // The new Connection interface only supports StandardSQL dialect
     content.setUseLegacySql(false);
+    if (connectionSettings.getJobTimeoutMs() != null) {
+      content.setJobTimeoutMs(connectionSettings.getJobTimeoutMs());
+    }
     return content;
   }
 

java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java

@@ -21,10 +21,10 @@
 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;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.UUID;
 
 final class QueryRequestInfo {
@@ -45,6 +45,7 @@ final class QueryRequestInfo {
   private final JobCreationMode jobCreationMode;
   private final DataFormatOptions formatOptions;
   private final String reservation;
+  private final Long jobTimeoutMs;
 
   QueryRequestInfo(
       QueryJobConfiguration config, com.google.cloud.bigquery.DataFormatOptions dataFormatOptions) {
@@ -64,22 +65,26 @@ final class QueryRequestInfo {
     this.jobCreationMode = config.getJobCreationMode();
     this.formatOptions = dataFormatOptions.toPb();
     this.reservation = config.getReservation();
+    this.jobTimeoutMs = config.getJobTimeoutMs();
   }
 
-  boolean isFastQuerySupported(JobId jobId) {
-    // Fast query path is not possible if job is specified in the JobID object
-    // Respect Job field value in JobId specified by user.
-    // Specifying it will force the query to take the slower path.
-    if (jobId != null) {
-      if (jobId.getJob() != null) {
-        return false;
-      }
-    }
+  /**
+   * Determines if the query can be executed via the "fast query" path (jobs.query API) instead of
+   * the "slow path" (jobs.insert API followed by jobs.getQueryResults).
+   *
+   * <p>The fast query path is preferred because it completes in a single RPC, significantly
+   * reducing end-to-end latency for small queries.
+   *
+   * <p>However, the jobs.query API does not support all configuration options available in
+   * jobs.insert (e.g., destination table, clustering, time partitioning). This method checks the
+   * QueryJobConfiguration for any unsupported options. If any are present, we must fall back to the
+   * jobs.insert path.
+   */
+  boolean isFastQuerySupported() {
     return config.getClustering() == null
         && config.getCreateDisposition() == null
         && config.getDestinationEncryptionConfiguration() == null
         && config.getDestinationTable() == null
-        && config.getJobTimeoutMs() == null
         && config.getMaximumBillingTier() == null
         && config.getPriority() == null
         && config.getRangePartitioning() == null
@@ -134,6 +139,9 @@ QueryRequest toPb() {
     if (reservation != null) {
       request.setReservation(reservation);
     }
+    if (jobTimeoutMs != null) {
+      request.setJobTimeoutMs(jobTimeoutMs);
+    }
     return request;
   }
 
@@ -155,12 +163,13 @@ public String toString() {
         .add("jobCreationMode", jobCreationMode)
         .add("formatOptions", formatOptions.getUseInt64Timestamp())
         .add("reservation", reservation)
+        .add("jobTimeoutMs", jobTimeoutMs)
         .toString();
   }
 
   @Override
   public int hashCode() {
-    return Objects.hashCode(
+    return Objects.hash(
         connectionProperties,
         defaultDataset,
         dryRun,
@@ -175,14 +184,34 @@ public int hashCode() {
         useLegacySql,
         jobCreationMode,
         formatOptions,
-        reservation);
+        reservation,
+        jobTimeoutMs);
   }
 
   @Override
   public boolean equals(Object obj) {
-    return obj == this
-        || obj != null
-            && obj.getClass().equals(QueryRequestInfo.class)
-            && java.util.Objects.equals(toPb(), ((QueryRequestInfo) obj).toPb());
+    if (obj == this) {
+      return true;
+    }
+    if (obj == null || !obj.getClass().equals(QueryRequestInfo.class)) {
+      return false;
+    }
+    QueryRequestInfo other = (QueryRequestInfo) obj;
+    return Objects.equals(connectionProperties, other.connectionProperties)
+        && Objects.equals(defaultDataset, other.defaultDataset)
+        && Objects.equals(dryRun, other.dryRun)
+        && Objects.equals(labels, other.labels)
+        && Objects.equals(maximumBytesBilled, other.maximumBytesBilled)
+        && Objects.equals(maxResults, other.maxResults)
+        && Objects.equals(query, other.query)
+        && Objects.equals(queryParameters, other.queryParameters)
+        && Objects.equals(requestId, other.requestId)
+        && Objects.equals(createSession, other.createSession)
+        && Objects.equals(useQueryCache, other.useQueryCache)
+        && Objects.equals(useLegacySql, other.useLegacySql)
+        && Objects.equals(jobCreationMode, other.jobCreationMode)
+        && Objects.equals(formatOptions, other.formatOptions)
+        && Objects.equals(reservation, other.reservation)
+        && Objects.equals(jobTimeoutMs, other.jobTimeoutMs);
   }
 }

java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java

@@ -240,6 +240,12 @@ public class BigQueryImplTest {
           .setDefaultDataset(DatasetId.of(PROJECT, DATASET))
           .setUseQueryCache(false)
           .build();
+  private static final QueryJobConfiguration QUERY_JOB_CONFIGURATION_WITH_TIMEOUT =
+      QueryJobConfiguration.newBuilder("SQL")
+          .setDefaultDataset(DatasetId.of(PROJECT, DATASET))
+          .setUseQueryCache(false)
+          .setJobTimeoutMs(1000L)
+          .build();
   private static final QueryJobConfiguration QUERY_JOB_CONFIGURATION_FOR_DMLQUERY =
       QueryJobConfiguration.newBuilder("DML")
           .setDefaultDataset(DatasetId.of(PROJECT, DATASET))
@@ -537,7 +543,8 @@ public class BigQueryImplTest {
   private HttpBigQueryRpc bigqueryRpcMock;
   private BigQuery bigquery;
   private static final String RATE_LIMIT_ERROR_MSG =
-      "Job exceeded rate limits: Your table exceeded quota for table update operations. For more information, see https://cloud.google.com/bigquery/docs/troubleshoot-quotas";
+      "Job exceeded rate limits: Your table exceeded quota for table update operations. For more"
+          + " information, see https://cloud.google.com/bigquery/docs/troubleshoot-quotas";
 
   @Captor private ArgumentCaptor<Map<BigQueryRpc.Option, Object>> capturedOptions;
   @Captor private ArgumentCaptor<com.google.api.services.bigquery.model.Job> jobCapture;
@@ -2394,6 +2401,49 @@ void testFastQueryRequestCompleted() throws InterruptedException, IOException {
         .queryRpcSkipExceptionTranslation(eq(PROJECT), requestPbCapture.capture());
   }
 
+  @Test
+  void testFastQueryRequestCompletedWithTimeout() throws InterruptedException, IOException {
+    com.google.api.services.bigquery.model.QueryResponse queryResponsePb =
+        new com.google.api.services.bigquery.model.QueryResponse()
+            .setCacheHit(false)
+            .setJobComplete(true)
+            .setKind("bigquery#queryResponse")
+            .setPageToken(null)
+            .setRows(ImmutableList.of(TABLE_ROW))
+            .setSchema(TABLE_SCHEMA.toPb())
+            .setTotalBytesProcessed(42L)
+            .setTotalRows(BigInteger.valueOf(1L));
+
+    when(bigqueryRpcMock.queryRpcSkipExceptionTranslation(eq(PROJECT), requestPbCapture.capture()))
+        .thenReturn(queryResponsePb);
+
+    bigquery = options.getService();
+    TableResult result = bigquery.query(QUERY_JOB_CONFIGURATION_WITH_TIMEOUT);
+    assertNull(result.getNextPage());
+    assertNull(result.getNextPageToken());
+    assertFalse(result.hasNextPage());
+    assertThat(result.getSchema()).isEqualTo(TABLE_SCHEMA);
+    assertThat(result.getTotalRows()).isEqualTo(1);
+    for (FieldValueList row : result.getValues()) {
+      assertThat(row.get(0).getBooleanValue()).isFalse();
+      assertThat(row.get(1).getLongValue()).isEqualTo(1);
+    }
+
+    QueryRequest requestPb = requestPbCapture.getValue();
+    assertEquals(QUERY_JOB_CONFIGURATION_WITH_TIMEOUT.getQuery(), requestPb.getQuery());
+    assertEquals(
+        QUERY_JOB_CONFIGURATION_WITH_TIMEOUT.getDefaultDataset().getDataset(),
+        requestPb.getDefaultDataset().getDatasetId());
+    assertEquals(
+        QUERY_JOB_CONFIGURATION_WITH_TIMEOUT.useQueryCache(), requestPb.getUseQueryCache());
+    assertEquals(
+        QUERY_JOB_CONFIGURATION_WITH_TIMEOUT.getJobTimeoutMs(), requestPb.getJobTimeoutMs());
+    assertNull(requestPb.getLocation());
+
+    verify(bigqueryRpcMock)
+        .queryRpcSkipExceptionTranslation(eq(PROJECT), requestPbCapture.capture());
+  }
+
   @Test
   void testQueryRequestRequiredJobCreationCompleted() throws InterruptedException, IOException {
     JobId queryJob = JobId.of(PROJECT, JOB);
@@ -3072,7 +3122,8 @@ void testFastQueryRateLimitIdempotency() throws Exception {
   @Test
   void testRateLimitRegEx() throws Exception {
     String msg2 =
-        "Job eceeded rate limits: Your table exceeded quota for table update operations. For more information, see https://cloud.google.com/bigquery/docs/troubleshoot-quotas";
+        "Job eceeded rate limits: Your table exceeded quota for table update operations. For more"
+            + " information, see https://cloud.google.com/bigquery/docs/troubleshoot-quotas";
     String msg3 = "exceeded rate exceeded quota for table update";
     String msg4 = "exceeded rate limits";
     assertTrue(

java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryRequestInfoTest.java

@@ -159,6 +159,12 @@ public class QueryRequestInfoTest {
   QueryRequestInfo REQUEST_INFO_SUPPORTED =
       new QueryRequestInfo(
           QUERY_JOB_CONFIGURATION_SUPPORTED, DataFormatOptions.newBuilder().build());
+  private static final QueryJobConfiguration QUERY_JOB_CONFIGURATION_WITH_TIMEOUT =
+      QUERY_JOB_CONFIGURATION_SUPPORTED.toBuilder().setJobTimeoutMs(TIMEOUT).build();
+  QueryRequestInfo REQUEST_INFO_WITH_TIMEOUT =
+      new QueryRequestInfo(
+          QUERY_JOB_CONFIGURATION_WITH_TIMEOUT, DataFormatOptions.newBuilder().build());
+
   private static final QueryJobConfiguration QUERY_JOB_CONFIGURATION_REQUIRED_SUPPORTED =
       QUERY_JOB_CONFIGURATION_SUPPORTED.toBuilder()
           .setJobCreationMode(JobCreationMode.JOB_CREATION_REQUIRED)
@@ -169,20 +175,18 @@ public class QueryRequestInfoTest {
 
   @Test
   public void testIsFastQuerySupported() {
-    JobId jobIdSupported = JobId.newBuilder().build();
-    JobId jobIdNotSupported = JobId.newBuilder().setJob("random-job-id").build();
-    assertEquals(false, REQUEST_INFO.isFastQuerySupported(jobIdSupported));
-    assertEquals(true, REQUEST_INFO_SUPPORTED.isFastQuerySupported(jobIdSupported));
-    assertEquals(true, REQUEST_INFO_REQUIRED_SUPPORTED.isFastQuerySupported(jobIdSupported));
-    assertEquals(false, REQUEST_INFO.isFastQuerySupported(jobIdNotSupported));
-    assertEquals(false, REQUEST_INFO_SUPPORTED.isFastQuerySupported(jobIdNotSupported));
-    assertEquals(false, REQUEST_INFO_REQUIRED_SUPPORTED.isFastQuerySupported(jobIdNotSupported));
+    assertFalse(REQUEST_INFO.isFastQuerySupported());
+    assertTrue(REQUEST_INFO_SUPPORTED.isFastQuerySupported());
+    assertTrue(REQUEST_INFO_WITH_TIMEOUT.isFastQuerySupported());
+    assertTrue(REQUEST_INFO_REQUIRED_SUPPORTED.isFastQuerySupported());
   }
 
   @Test
   public void testToPb() {
     QueryRequest requestPb = REQUEST_INFO.toPb();
     assertEquals(requestPb, REQUEST_INFO.toPb());
+    QueryRequest requestWithTimeoutPb = REQUEST_INFO_WITH_TIMEOUT.toPb();
+    assertEquals(TIMEOUT, requestWithTimeoutPb.getJobTimeoutMs());
   }
 
   @Test
@@ -194,6 +198,14 @@ public void equalTo() {
     compareQueryRequestInfo(
         new QueryRequestInfo(QUERY_JOB_CONFIGURATION, DataFormatOptions.newBuilder().build()),
         REQUEST_INFO);
+    compareQueryRequestInfo(
+        new QueryRequestInfo(
+            QUERY_JOB_CONFIGURATION_WITH_TIMEOUT, DataFormatOptions.newBuilder().build()),
+        REQUEST_INFO_WITH_TIMEOUT);
+    compareQueryRequestInfo(
+        new QueryRequestInfo(
+            QUERY_JOB_CONFIGURATION_REQUIRED_SUPPORTED, DataFormatOptions.newBuilder().build()),
+        REQUEST_INFO_REQUIRED_SUPPORTED);
   }
 
   @Test
@@ -237,5 +249,6 @@ private void compareQueryRequestInfo(QueryRequestInfo expected, QueryRequestInfo
     assertEquals(expectedQueryReq.get("jobCreationMode"), actualQueryReq.get("jobCreationMode"));
     assertEquals(expectedQueryReq.getFormatOptions(), actualQueryReq.getFormatOptions());
     assertEquals(expectedQueryReq.getReservation(), actualQueryReq.getReservation());
+    assertEquals(expectedQueryReq.getJobTimeoutMs(), actualQueryReq.getJobTimeoutMs());
   }
 }