fix(bigquery): route JOB_CREATION_REQUIRED through fast query path (#13437)

Routes queries under `JobCreationMode.JOB_CREATION_REQUIRED` to the fast
query path (`jobs.query` API / 1 RPC) to avoid the slow fallback path
(`jobs.insert` API / 2 RPCs).


https://docs.cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#QueryResponse
says:
> object
([JobReference](https://docs.cloud.google.com/bigquery/docs/reference/rest/v2/JobReference))
> 
> Reference to the Job that was created to run the query. This field
will be present even if the original request timed out, in which case
jobs.getQueryResults can be used to read the results once the query has
completed. Since this API only returns the first page of results,
subsequent pages can be fetched via the same mechanism
(jobs.getQueryResults).
> 
> If jobCreationMode was set to JOB_CREATION_OPTIONAL and the query
completes without creating a job, this field will be empty.

Thus, routing `JOB_CREATION_REQUIRED` through the fast path is
preferred.

b/522363981

Author: jinseopkim0

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

@@ -87,8 +87,7 @@ boolean isFastQuerySupported(JobId jobId) {
         && config.getTableDefinitions() == null
         && config.getTimePartitioning() == null
         && config.getUserDefinedFunctions() == null
-        && config.getWriteDisposition() == null
-        && config.getJobCreationMode() != JobCreationMode.JOB_CREATION_REQUIRED;
+        && config.getWriteDisposition() == null;
   }
 
   QueryRequest toPb() {

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

@@ -32,6 +32,7 @@
 import static org.mockito.ArgumentMatchers.nullable;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -2393,6 +2394,56 @@ void testFastQueryRequestCompleted() throws InterruptedException, IOException {
         .queryRpcSkipExceptionTranslation(eq(PROJECT), requestPbCapture.capture());
   }
 
+  @Test
+  void testQueryRequestRequiredJobCreationCompleted() throws InterruptedException, IOException {
+    JobId queryJob = JobId.of(PROJECT, JOB);
+    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))
+            .setJobReference(queryJob.toPb());
+
+    QueryJobConfiguration config =
+        QUERY_JOB_CONFIGURATION_FOR_QUERY.toBuilder()
+            .setJobCreationMode(QueryJobConfiguration.JobCreationMode.JOB_CREATION_REQUIRED)
+            .build();
+
+    when(bigqueryRpcMock.queryRpcSkipExceptionTranslation(eq(PROJECT), requestPbCapture.capture()))
+        .thenReturn(queryResponsePb);
+
+    bigquery = options.getService();
+    TableResult result = bigquery.query(config);
+    assertNull(result.getNextPage());
+    assertNull(result.getNextPageToken());
+    assertFalse(result.hasNextPage());
+    assertThat(result.getSchema()).isEqualTo(TABLE_SCHEMA);
+    assertThat(result.getTotalRows()).isEqualTo(1);
+    assertThat(result.getJobId()).isEqualTo(queryJob);
+    for (FieldValueList row : result.getValues()) {
+      assertThat(row.get(0).getBooleanValue()).isFalse();
+      assertThat(row.get(1).getLongValue()).isEqualTo(1);
+    }
+
+    QueryRequest requestPb = requestPbCapture.getValue();
+    assertEquals(config.getQuery(), requestPb.getQuery());
+    assertEquals(
+        config.getDefaultDataset().getDataset(), requestPb.getDefaultDataset().getDatasetId());
+    assertEquals(config.useQueryCache(), requestPb.getUseQueryCache());
+    assertNull(requestPb.getLocation());
+
+    verify(bigqueryRpcMock)
+        .queryRpcSkipExceptionTranslation(eq(PROJECT), requestPbCapture.capture());
+    verify(bigqueryRpcMock, never())
+        .createSkipExceptionTranslation(
+            any(com.google.api.services.bigquery.model.Job.class), any());
+  }
+
   @Test
   void testFastQueryRequestCompletedWithLocation() throws InterruptedException, IOException {
     com.google.api.services.bigquery.model.QueryResponse queryResponsePb =

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

@@ -159,15 +159,24 @@ public class QueryRequestInfoTest {
   QueryRequestInfo REQUEST_INFO_SUPPORTED =
       new QueryRequestInfo(
           QUERY_JOB_CONFIGURATION_SUPPORTED, DataFormatOptions.newBuilder().build());
+  private static final QueryJobConfiguration QUERY_JOB_CONFIGURATION_REQUIRED_SUPPORTED =
+      QUERY_JOB_CONFIGURATION_SUPPORTED.toBuilder()
+          .setJobCreationMode(JobCreationMode.JOB_CREATION_REQUIRED)
+          .build();
+  QueryRequestInfo REQUEST_INFO_REQUIRED_SUPPORTED =
+      new QueryRequestInfo(
+          QUERY_JOB_CONFIGURATION_REQUIRED_SUPPORTED, 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(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));
   }
 
   @Test

java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java

@@ -7344,10 +7344,9 @@ void testStatelessQueries() throws InterruptedException {
         (tableResult.getJobId() != null) ^ (tableResult.getQueryId() != null),
         "Exactly one of jobId or queryId should be non-null");
 
-    // Job creation takes over, no query id is created.
     bigQuery.getOptions().setDefaultJobCreationMode(JobCreationMode.JOB_CREATION_REQUIRED);
     tableResult = executeSimpleQuery(bigQuery);
-    assertNull(tableResult.getQueryId());
+    assertNotNull(tableResult.getQueryId());
     assertNotNull(tableResult.getJobId());
 
     bigQuery.getOptions().setDefaultJobCreationMode(JobCreationMode.JOB_CREATION_MODE_UNSPECIFIED);
@@ -7411,9 +7410,8 @@ void testTableResultJobIdAndQueryId() throws InterruptedException {
             .setJobCreationMode(JobCreationMode.JOB_CREATION_REQUIRED)
             .build();
     result = bigQuery.query(configWithJob);
-    result = job.getQueryResults();
     assertNotNull(result.getJobId());
-    assertNull(result.getQueryId());
+    assertNotNull(result.getQueryId());
   }
 
   @Test
@@ -7508,14 +7506,14 @@ void testQueryWithTimeout() throws InterruptedException {
     // Allow 2 seconds of timeout value to account for random delays
     assertTrue(millis < 1_000_000 * 2);
 
-    // Stateful query returns Job
-    // Test scenario 3 to ensure job is created if JobCreationMode is set.
+    // Test scenario 3 to ensure TableResult is returned with JobId if JobCreationMode is REQUIRED
     config =
         QueryJobConfiguration.newBuilder(query)
             .setJobCreationMode(JobCreationMode.JOB_CREATION_REQUIRED)
             .build();
     result = bigQuery.queryWithTimeout(config, null, null);
-    assertTrue(result instanceof Job);
+    assertTrue(result instanceof TableResult);
+    assertNotNull(((TableResult) result).getJobId());
 
     // Stateful query returns Job
     // Test scenario 4 to ensure job is created if Query is long running.