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

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

@@ -2080,11 +2080,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

@@ -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
@@ -135,6 +140,9 @@ QueryRequest toPb() {
     if (reservation != null) {
       request.setReservation(reservation);
     }
+    if (jobTimeoutMs != null) {
+      request.setJobTimeoutMs(jobTimeoutMs);
+    }
     return request;
   }
 
@@ -156,6 +164,7 @@ public String toString() {
         .add("jobCreationMode", jobCreationMode)
         .add("formatOptions", formatOptions.getUseInt64Timestamp())
         .add("reservation", reservation)
+        .add("jobTimeoutMs", jobTimeoutMs)
         .toString();
   }
 
@@ -176,7 +185,8 @@ public int hashCode() {
         useLegacySql,
         jobCreationMode,
         formatOptions,
-        reservation);
+        reservation,
+        jobTimeoutMs);
   }
 
   @Override

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

@@ -237,6 +237,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))
@@ -534,7 +540,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;
@@ -2347,6 +2354,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 testFastQueryRequestCompletedWithLocation() throws InterruptedException, IOException {
     com.google.api.services.bigquery.model.QueryResponse queryResponsePb =
@@ -2975,7 +3025,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,21 +159,25 @@ 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());
 
   @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(false, REQUEST_INFO.isFastQuerySupported(jobIdNotSupported));
-    assertEquals(false, REQUEST_INFO_SUPPORTED.isFastQuerySupported(jobIdNotSupported));
+    assertFalse(REQUEST_INFO.isFastQuerySupported());
+    assertTrue(REQUEST_INFO_SUPPORTED.isFastQuerySupported());
+    assertTrue(REQUEST_INFO_WITH_TIMEOUT.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
@@ -185,6 +189,10 @@ 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);
   }
 
   @Test
@@ -228,5 +236,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());
   }
 }