Visitar URL original
fix(bigquery): route JOB_CREATION_REQUIRED through fast query path (#… · googleapis/google-cloud-java@aa8e704 · GitHub
Skip to content

Commit aa8e704

Browse files
jinseopkim0lqiu96
jinseopkim0
authored and
lqiu96
committed
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
1 parent 9ada1dc commit aa8e704

4 files changed

Lines changed: 66 additions & 9 deletions

File tree

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -87,8 +87,7 @@ boolean isFastQuerySupported(JobId jobId) {
8787
&& config.getTableDefinitions() == null
8888
&& config.getTimePartitioning() == null
8989
&& config.getUserDefinedFunctions() == null
90-
&& config.getWriteDisposition() == null
91-
&& config.getJobCreationMode() != JobCreationMode.JOB_CREATION_REQUIRED;
90+
&& config.getWriteDisposition() == null;
9291
}
9392

9493
QueryRequest toPb() {

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

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import static org.mockito.ArgumentMatchers.nullable;
3333
import static org.mockito.Mockito.doReturn;
3434
import static org.mockito.Mockito.mock;
35+
import static org.mockito.Mockito.never;
3536
import static org.mockito.Mockito.times;
3637
import static org.mockito.Mockito.verify;
3738
import static org.mockito.Mockito.when;
@@ -2393,6 +2394,56 @@ void testFastQueryRequestCompleted() throws InterruptedException, IOException {
23932394
.queryRpcSkipExceptionTranslation(eq(PROJECT), requestPbCapture.capture());
23942395
}
23952396

2397+
@Test
2398+
void testQueryRequestRequiredJobCreationCompleted() throws InterruptedException, IOException {
2399+
JobId queryJob = JobId.of(PROJECT, JOB);
2400+
com.google.api.services.bigquery.model.QueryResponse queryResponsePb =
2401+
new com.google.api.services.bigquery.model.QueryResponse()
2402+
.setCacheHit(false)
2403+
.setJobComplete(true)
2404+
.setKind("bigquery#queryResponse")
2405+
.setPageToken(null)
2406+
.setRows(ImmutableList.of(TABLE_ROW))
2407+
.setSchema(TABLE_SCHEMA.toPb())
2408+
.setTotalBytesProcessed(42L)
2409+
.setTotalRows(BigInteger.valueOf(1L))
2410+
.setJobReference(queryJob.toPb());
2411+
2412+
QueryJobConfiguration config =
2413+
QUERY_JOB_CONFIGURATION_FOR_QUERY.toBuilder()
2414+
.setJobCreationMode(QueryJobConfiguration.JobCreationMode.JOB_CREATION_REQUIRED)
2415+
.build();
2416+
2417+
when(bigqueryRpcMock.queryRpcSkipExceptionTranslation(eq(PROJECT), requestPbCapture.capture()))
2418+
.thenReturn(queryResponsePb);
2419+
2420+
bigquery = options.getService();
2421+
TableResult result = bigquery.query(config);
2422+
assertNull(result.getNextPage());
2423+
assertNull(result.getNextPageToken());
2424+
assertFalse(result.hasNextPage());
2425+
assertThat(result.getSchema()).isEqualTo(TABLE_SCHEMA);
2426+
assertThat(result.getTotalRows()).isEqualTo(1);
2427+
assertThat(result.getJobId()).isEqualTo(queryJob);
2428+
for (FieldValueList row : result.getValues()) {
2429+
assertThat(row.get(0).getBooleanValue()).isFalse();
2430+
assertThat(row.get(1).getLongValue()).isEqualTo(1);
2431+
}
2432+
2433+
QueryRequest requestPb = requestPbCapture.getValue();
2434+
assertEquals(config.getQuery(), requestPb.getQuery());
2435+
assertEquals(
2436+
config.getDefaultDataset().getDataset(), requestPb.getDefaultDataset().getDatasetId());
2437+
assertEquals(config.useQueryCache(), requestPb.getUseQueryCache());
2438+
assertNull(requestPb.getLocation());
2439+
2440+
verify(bigqueryRpcMock)
2441+
.queryRpcSkipExceptionTranslation(eq(PROJECT), requestPbCapture.capture());
2442+
verify(bigqueryRpcMock, never())
2443+
.createSkipExceptionTranslation(
2444+
any(com.google.api.services.bigquery.model.Job.class), any());
2445+
}
2446+
23962447
@Test
23972448
void testFastQueryRequestCompletedWithLocation() throws InterruptedException, IOException {
23982449
com.google.api.services.bigquery.model.QueryResponse queryResponsePb =

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -159,15 +159,24 @@ public class QueryRequestInfoTest {
159159
QueryRequestInfo REQUEST_INFO_SUPPORTED =
160160
new QueryRequestInfo(
161161
QUERY_JOB_CONFIGURATION_SUPPORTED, DataFormatOptions.newBuilder().build());
162+
private static final QueryJobConfiguration QUERY_JOB_CONFIGURATION_REQUIRED_SUPPORTED =
163+
QUERY_JOB_CONFIGURATION_SUPPORTED.toBuilder()
164+
.setJobCreationMode(JobCreationMode.JOB_CREATION_REQUIRED)
165+
.build();
166+
QueryRequestInfo REQUEST_INFO_REQUIRED_SUPPORTED =
167+
new QueryRequestInfo(
168+
QUERY_JOB_CONFIGURATION_REQUIRED_SUPPORTED, DataFormatOptions.newBuilder().build());
162169

163170
@Test
164171
public void testIsFastQuerySupported() {
165172
JobId jobIdSupported = JobId.newBuilder().build();
166173
JobId jobIdNotSupported = JobId.newBuilder().setJob("random-job-id").build();
167174
assertEquals(false, REQUEST_INFO.isFastQuerySupported(jobIdSupported));
168175
assertEquals(true, REQUEST_INFO_SUPPORTED.isFastQuerySupported(jobIdSupported));
176+
assertEquals(true, REQUEST_INFO_REQUIRED_SUPPORTED.isFastQuerySupported(jobIdSupported));
169177
assertEquals(false, REQUEST_INFO.isFastQuerySupported(jobIdNotSupported));
170178
assertEquals(false, REQUEST_INFO_SUPPORTED.isFastQuerySupported(jobIdNotSupported));
179+
assertEquals(false, REQUEST_INFO_REQUIRED_SUPPORTED.isFastQuerySupported(jobIdNotSupported));
171180
}
172181

173182
@Test

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

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -7344,10 +7344,9 @@ void testStatelessQueries() throws InterruptedException {
73447344
(tableResult.getJobId() != null) ^ (tableResult.getQueryId() != null),
73457345
"Exactly one of jobId or queryId should be non-null");
73467346

7347-
// Job creation takes over, no query id is created.
73487347
bigQuery.getOptions().setDefaultJobCreationMode(JobCreationMode.JOB_CREATION_REQUIRED);
73497348
tableResult = executeSimpleQuery(bigQuery);
7350-
assertNull(tableResult.getQueryId());
7349+
assertNotNull(tableResult.getQueryId());
73517350
assertNotNull(tableResult.getJobId());
73527351

73537352
bigQuery.getOptions().setDefaultJobCreationMode(JobCreationMode.JOB_CREATION_MODE_UNSPECIFIED);
@@ -7411,9 +7410,8 @@ void testTableResultJobIdAndQueryId() throws InterruptedException {
74117410
.setJobCreationMode(JobCreationMode.JOB_CREATION_REQUIRED)
74127411
.build();
74137412
result = bigQuery.query(configWithJob);
7414-
result = job.getQueryResults();
74157413
assertNotNull(result.getJobId());
7416-
assertNull(result.getQueryId());
7414+
assertNotNull(result.getQueryId());
74177415
}
74187416

74197417
@Test
@@ -7508,14 +7506,14 @@ void testQueryWithTimeout() throws InterruptedException {
75087506
// Allow 2 seconds of timeout value to account for random delays
75097507
assertTrue(millis < 1_000_000 * 2);
75107508

7511-
// Stateful query returns Job
7512-
// Test scenario 3 to ensure job is created if JobCreationMode is set.
7509+
// Test scenario 3 to ensure TableResult is returned with JobId if JobCreationMode is REQUIRED
75137510
config =
75147511
QueryJobConfiguration.newBuilder(query)
75157512
.setJobCreationMode(JobCreationMode.JOB_CREATION_REQUIRED)
75167513
.build();
75177514
result = bigQuery.queryWithTimeout(config, null, null);
7518-
assertTrue(result instanceof Job);
7515+
assertTrue(result instanceof TableResult);
7516+
assertNotNull(((TableResult) result).getJobId());
75197517

75207518
// Stateful query returns Job
75217519
// Test scenario 4 to ensure job is created if Query is long running.

0 commit comments

Comments
 (0)