From e1c594b17c5f0f772d4f703cd9e8f89a6afa23ca Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Thu, 11 Jun 2026 13:45:59 -0400 Subject: [PATCH] fix(bigquery): route JOB_CREATION_REQUIRED through fast query path --- .../cloud/bigquery/QueryRequestInfo.java | 3 +- .../cloud/bigquery/BigQueryImplTest.java | 51 +++++++++++++++++++ .../cloud/bigquery/QueryRequestInfoTest.java | 9 ++++ .../cloud/bigquery/it/ITBigQueryTest.java | 12 ++--- 4 files changed, 66 insertions(+), 9 deletions(-) diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java index c7033817c367..4e203a26af58 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java +++ b/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() { diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java index 20a6ef679e89..00ec9f466890 100644 --- a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java +++ b/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 = diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryRequestInfoTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryRequestInfoTest.java index be1f0e1982f9..496676e5b788 100644 --- a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryRequestInfoTest.java +++ b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryRequestInfoTest.java @@ -159,6 +159,13 @@ 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() { @@ -166,8 +173,10 @@ public void testIsFastQuerySupported() { 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 diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java index dd7319af2e4c..16a1b7dcf10f 100644 --- a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java +++ b/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.