fix(bigquery): route JOB_CREATION_REQUIRED through fast query path

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;
@@ -2347,6 +2348,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

@@ -7335,10 +7335,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);
@@ -7402,9 +7401,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
@@ -7499,14 +7497,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.