Skip to content

Commit e2cd4f6

Browse files
stephaniewang526stephaniewang526
authored
feat: set requestId for fast query path in QueryRequestInfo instead of QueryJobConfiguration (googleapis#987)
As requested @epavan
1 parent a692fe1 commit e2cd4f6

File tree

7 files changed

+78
-56
lines changed

7 files changed

+78
-56
lines changed

‎google-cloud-bigquery/pom.xml

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,11 @@
103103
<artifactId>mockito-core</artifactId>
104104
<scope>test</scope>
105105
</dependency>
106+
<dependency>
107+
<groupId>org.assertj</groupId>
108+
<artifactId>assertj-core</artifactId>
109+
<scope>test</scope>
110+
</dependency>
106111
</dependencies>
107112

108113
<build>

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

Lines changed: 1 addition & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,6 @@
3535
import java.util.List;
3636
import java.util.Map;
3737
import java.util.Objects;
38-
import java.util.UUID;
3938

4039
/**
4140
* Google BigQuery Query Job configuration. A Query Job runs a query against BigQuery data. Query
@@ -72,7 +71,6 @@ public final class QueryJobConfiguration extends JobConfiguration {
7271
private final List<ConnectionProperty> connectionProperties;
7372
// maxResults is only used for fast query path
7473
private final Long maxResults;
75-
private final String requestId;
7674

7775
/**
7876
* Priority levels for a query. If not specified the priority is assumed to be {@link
@@ -123,7 +121,6 @@ public static final class Builder
123121
private RangePartitioning rangePartitioning;
124122
private List<ConnectionProperty> connectionProperties;
125123
private Long maxResults;
126-
private String requestId;
127124

128125
private Builder() {
129126
super(Type.QUERY);
@@ -157,7 +154,6 @@ private Builder(QueryJobConfiguration jobConfiguration) {
157154
this.rangePartitioning = jobConfiguration.rangePartitioning;
158155
this.connectionProperties = jobConfiguration.connectionProperties;
159156
this.maxResults = jobConfiguration.maxResults;
160-
this.requestId = jobConfiguration.requestId;
161157
}
162158

163159
private Builder(com.google.api.services.bigquery.model.JobConfiguration configurationPb) {
@@ -625,11 +621,6 @@ public Builder setMaxResults(Long maxResults) {
625621
return this;
626622
}
627623

628-
Builder setRequestId(String requestId) {
629-
this.requestId = requestId;
630-
return this;
631-
}
632-
633624
public QueryJobConfiguration build() {
634625
return new QueryJobConfiguration(this);
635626
}
@@ -672,7 +663,6 @@ private QueryJobConfiguration(Builder builder) {
672663
this.rangePartitioning = builder.rangePartitioning;
673664
this.connectionProperties = builder.connectionProperties;
674665
this.maxResults = builder.maxResults;
675-
this.requestId = builder.requestId;
676666
}
677667

678668
/**
@@ -875,10 +865,6 @@ public Long getMaxResults() {
875865
return maxResults;
876866
}
877867

878-
String getRequestId() {
879-
return requestId;
880-
}
881-
882868
@Override
883869
public Builder toBuilder() {
884870
return new Builder(this);
@@ -1057,7 +1043,7 @@ com.google.api.services.bigquery.model.JobConfiguration toPb() {
10571043
/** Creates a builder for a BigQuery Query Job given the query to be run. */
10581044
public static Builder newBuilder(String query) {
10591045
checkArgument(!isNullOrEmpty(query), "Provided query is null or empty");
1060-
return new Builder().setQuery(query).setRequestId(UUID.randomUUID().toString());
1046+
return new Builder().setQuery(query);
10611047
}
10621048

10631049
/**

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import com.google.common.collect.Lists;
2424
import java.util.List;
2525
import java.util.Map;
26+
import java.util.UUID;
2627

2728
final class QueryRequestInfo {
2829

@@ -35,9 +36,9 @@ final class QueryRequestInfo {
3536
private final Long maxResults;
3637
private final String query;
3738
private final List<QueryParameter> queryParameters;
39+
private final String requestId;
3840
private final Boolean useQueryCache;
3941
private final Boolean useLegacySql;
40-
private final String requestId;
4142

4243
QueryRequestInfo(QueryJobConfiguration config) {
4344
this.config = config;
@@ -49,9 +50,9 @@ final class QueryRequestInfo {
4950
this.maxResults = config.getMaxResults();
5051
this.query = config.getQuery();
5152
this.queryParameters = config.toPb().getQuery().getQueryParameters();
53+
this.requestId = UUID.randomUUID().toString();
5254
this.useLegacySql = config.useLegacySql();
5355
this.useQueryCache = config.useQueryCache();
54-
this.requestId = config.getRequestId();
5556
}
5657

5758
boolean isFastQuerySupported() {

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

Lines changed: 50 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -1885,9 +1885,8 @@ public void testFastQueryRequestCompleted() throws InterruptedException {
18851885
.setTotalBytesProcessed(42L)
18861886
.setTotalRows(BigInteger.valueOf(1L));
18871887

1888-
QueryRequestInfo requestInfo = new QueryRequestInfo(QUERY_JOB_CONFIGURATION_FOR_QUERY);
1889-
1890-
when(bigqueryRpcMock.queryRpc(PROJECT, requestInfo.toPb())).thenReturn(queryResponsePb);
1888+
when(bigqueryRpcMock.queryRpc(eq(PROJECT), requestPbCapture.capture()))
1889+
.thenReturn(queryResponsePb);
18911890

18921891
bigquery = options.getService();
18931892
TableResult result = bigquery.query(QUERY_JOB_CONFIGURATION_FOR_QUERY);
@@ -1900,7 +1899,15 @@ public void testFastQueryRequestCompleted() throws InterruptedException {
19001899
assertThat(row.get(0).getBooleanValue()).isFalse();
19011900
assertThat(row.get(1).getLongValue()).isEqualTo(1);
19021901
}
1903-
verify(bigqueryRpcMock).queryRpc(PROJECT, requestInfo.toPb());
1902+
1903+
QueryRequest requestPb = requestPbCapture.getValue();
1904+
assertEquals(QUERY_JOB_CONFIGURATION_FOR_QUERY.getQuery(), requestPb.getQuery());
1905+
assertEquals(
1906+
QUERY_JOB_CONFIGURATION_FOR_QUERY.getDefaultDataset().getDataset(),
1907+
requestPb.getDefaultDataset().getDatasetId());
1908+
assertEquals(QUERY_JOB_CONFIGURATION_FOR_QUERY.useQueryCache(), requestPb.getUseQueryCache());
1909+
1910+
verify(bigqueryRpcMock).queryRpc(eq(PROJECT), requestPbCapture.capture());
19041911
}
19051912

19061913
@Test
@@ -1937,22 +1944,30 @@ public void testFastQueryMultiplePages() throws InterruptedException {
19371944
.setTotalBytesProcessed(42L)
19381945
.setTotalRows(BigInteger.valueOf(1L));
19391946

1940-
QueryRequestInfo requestInfo = new QueryRequestInfo(QUERY_JOB_CONFIGURATION_FOR_QUERY);
1941-
when(bigqueryRpcMock.queryRpc(PROJECT, requestInfo.toPb())).thenReturn(queryResponsePb);
1947+
when(bigqueryRpcMock.queryRpc(eq(PROJECT), requestPbCapture.capture()))
1948+
.thenReturn(queryResponsePb);
19421949

19431950
bigquery = options.getService();
19441951
TableResult result = bigquery.query(QUERY_JOB_CONFIGURATION_FOR_QUERY);
19451952
assertTrue(result.hasNextPage());
19461953
assertNotNull(result.getNextPageToken());
19471954
assertNotNull(result.getNextPage());
1955+
1956+
QueryRequest requestPb = requestPbCapture.getValue();
1957+
assertEquals(QUERY_JOB_CONFIGURATION_FOR_QUERY.getQuery(), requestPb.getQuery());
1958+
assertEquals(
1959+
QUERY_JOB_CONFIGURATION_FOR_QUERY.getDefaultDataset().getDataset(),
1960+
requestPb.getDefaultDataset().getDatasetId());
1961+
assertEquals(QUERY_JOB_CONFIGURATION_FOR_QUERY.useQueryCache(), requestPb.getUseQueryCache());
1962+
19481963
verify(bigqueryRpcMock).getJob(PROJECT, JOB, null, EMPTY_RPC_OPTIONS);
19491964
verify(bigqueryRpcMock)
19501965
.listTableData(
19511966
PROJECT,
19521967
DATASET,
19531968
TABLE,
19541969
BigQueryImpl.optionMap(BigQuery.TableDataListOption.pageToken(CURSOR)));
1955-
verify(bigqueryRpcMock).queryRpc(PROJECT, requestInfo.toPb());
1970+
verify(bigqueryRpcMock).queryRpc(eq(PROJECT), requestPbCapture.capture());
19561971
}
19571972

19581973
@Test
@@ -1983,10 +1998,8 @@ public void testFastQuerySlowDdl() throws InterruptedException {
19831998
.setTotalRows(BigInteger.valueOf(1L))
19841999
.setSchema(TABLE_SCHEMA.toPb());
19852000

1986-
QueryRequestInfo requestInfo = new QueryRequestInfo(QUERY_JOB_CONFIGURATION_FOR_QUERY);
1987-
QueryRequest requestPb = requestInfo.toPb();
1988-
1989-
when(bigqueryRpcMock.queryRpc(PROJECT, requestPb)).thenReturn(queryResponsePb);
2001+
when(bigqueryRpcMock.queryRpc(eq(PROJECT), requestPbCapture.capture()))
2002+
.thenReturn(queryResponsePb);
19902003
responseJob.getConfiguration().getQuery().setDestinationTable(TABLE_ID.toPb());
19912004
when(bigqueryRpcMock.getJob(PROJECT, JOB, null, EMPTY_RPC_OPTIONS)).thenReturn(responseJob);
19922005
when(bigqueryRpcMock.getQueryResults(
@@ -2004,7 +2017,14 @@ public void testFastQuerySlowDdl() throws InterruptedException {
20042017
assertThat(row.get(1).getLongValue()).isEqualTo(1);
20052018
}
20062019

2007-
verify(bigqueryRpcMock).queryRpc(PROJECT, requestInfo.toPb());
2020+
QueryRequest requestPb = requestPbCapture.getValue();
2021+
assertEquals(QUERY_JOB_CONFIGURATION_FOR_QUERY.getQuery(), requestPb.getQuery());
2022+
assertEquals(
2023+
QUERY_JOB_CONFIGURATION_FOR_QUERY.getDefaultDataset().getDataset(),
2024+
requestPb.getDefaultDataset().getDatasetId());
2025+
assertEquals(QUERY_JOB_CONFIGURATION_FOR_QUERY.useQueryCache(), requestPb.getUseQueryCache());
2026+
2027+
verify(bigqueryRpcMock).queryRpc(eq(PROJECT), requestPbCapture.capture());
20082028
verify(bigqueryRpcMock).getJob(PROJECT, JOB, null, EMPTY_RPC_OPTIONS);
20092029
verify(bigqueryRpcMock)
20102030
.getQueryResults(
@@ -2287,9 +2307,6 @@ public void testFastQuerySQLShouldRetry() throws Exception {
22872307
.setTotalRows(BigInteger.valueOf(1L))
22882308
.setSchema(TABLE_SCHEMA.toPb());
22892309

2290-
QueryRequestInfo requestInfo = new QueryRequestInfo(QUERY_JOB_CONFIGURATION_FOR_QUERY);
2291-
QueryRequest requestPb = requestInfo.toPb();
2292-
22932310
when(bigqueryRpcMock.queryRpc(eq(PROJECT), requestPbCapture.capture()))
22942311
.thenThrow(new BigQueryException(500, "InternalError"))
22952312
.thenThrow(new BigQueryException(502, "Bad Gateway"))
@@ -2310,13 +2327,13 @@ public void testFastQuerySQLShouldRetry() throws Exception {
23102327

23112328
List<QueryRequest> allRequests = requestPbCapture.getAllValues();
23122329
boolean idempotent = true;
2313-
String requestId = requestPb.getRequestId();
2330+
String firstRequestId = allRequests.get(0).getRequestId();
23142331
for (QueryRequest request : allRequests) {
2315-
idempotent = request.getRequestId().equals(requestId);
2332+
idempotent = request.getRequestId().equals(firstRequestId);
23162333
}
23172334
assertTrue(idempotent);
23182335

2319-
verify(bigqueryRpcMock, times(5)).queryRpc(PROJECT, requestPb);
2336+
verify(bigqueryRpcMock, times(5)).queryRpc(eq(PROJECT), requestPbCapture.capture());
23202337
}
23212338

23222339
@Test
@@ -2331,9 +2348,6 @@ public void testFastQueryDMLShouldRetry() throws Exception {
23312348
.setNumDmlAffectedRows(1L)
23322349
.setSchema(TABLE_SCHEMA.toPb());
23332350

2334-
QueryRequestInfo requestInfo = new QueryRequestInfo(QUERY_JOB_CONFIGURATION_FOR_DMLQUERY);
2335-
QueryRequest requestPb = requestInfo.toPb();
2336-
23372351
when(bigqueryRpcMock.queryRpc(eq(PROJECT), requestPbCapture.capture()))
23382352
.thenThrow(new BigQueryException(500, "InternalError"))
23392353
.thenThrow(new BigQueryException(502, "Bad Gateway"))
@@ -2354,13 +2368,13 @@ public void testFastQueryDMLShouldRetry() throws Exception {
23542368

23552369
List<QueryRequest> allRequests = requestPbCapture.getAllValues();
23562370
boolean idempotent = true;
2357-
String requestId = requestPb.getRequestId();
2371+
String firstRequestId = allRequests.get(0).getRequestId();
23582372
for (QueryRequest request : allRequests) {
2359-
idempotent = request.getRequestId().equals(requestId);
2373+
idempotent = request.getRequestId().equals(firstRequestId);
23602374
}
23612375
assertTrue(idempotent);
23622376

2363-
verify(bigqueryRpcMock, times(5)).queryRpc(PROJECT, requestPb);
2377+
verify(bigqueryRpcMock, times(5)).queryRpc(eq(PROJECT), requestPbCapture.capture());
23642378
}
23652379

23662380
@Test
@@ -2374,9 +2388,6 @@ public void testFastQueryDDLShouldRetry() throws Exception {
23742388
.setTotalBytesProcessed(42L)
23752389
.setSchema(TABLE_SCHEMA.toPb());
23762390

2377-
QueryRequestInfo requestInfo = new QueryRequestInfo(QUERY_JOB_CONFIGURATION_FOR_DDLQUERY);
2378-
QueryRequest requestPb = requestInfo.toPb();
2379-
23802391
when(bigqueryRpcMock.queryRpc(eq(PROJECT), requestPbCapture.capture()))
23812392
.thenThrow(new BigQueryException(500, "InternalError"))
23822393
.thenThrow(new BigQueryException(502, "Bad Gateway"))
@@ -2397,13 +2408,13 @@ public void testFastQueryDDLShouldRetry() throws Exception {
23972408

23982409
List<QueryRequest> allRequests = requestPbCapture.getAllValues();
23992410
boolean idempotent = true;
2400-
String requestId = requestPb.getRequestId();
2411+
String firstRequestId = allRequests.get(0).getRequestId();
24012412
for (QueryRequest request : allRequests) {
2402-
idempotent = request.getRequestId().equals(requestId);
2413+
idempotent = request.getRequestId().equals(firstRequestId);
24032414
}
24042415
assertTrue(idempotent);
24052416

2406-
verify(bigqueryRpcMock, times(5)).queryRpc(PROJECT, requestPb);
2417+
verify(bigqueryRpcMock, times(5)).queryRpc(eq(PROJECT), requestPbCapture.capture());
24072418
}
24082419

24092420
@Test
@@ -2424,10 +2435,7 @@ public void testFastQueryBigQueryException() throws InterruptedException {
24242435
.setPageToken(null)
24252436
.setErrors(errorProtoList);
24262437

2427-
QueryRequestInfo requestInfo = new QueryRequestInfo(QUERY_JOB_CONFIGURATION_FOR_QUERY);
2428-
QueryRequest requestPb = requestInfo.toPb();
2429-
2430-
when(bigqueryRpcMock.queryRpc(PROJECT, requestPb)).thenReturn(responsePb);
2438+
when(bigqueryRpcMock.queryRpc(eq(PROJECT), requestPbCapture.capture())).thenReturn(responsePb);
24312439

24322440
bigquery = options.getService();
24332441
try {
@@ -2436,7 +2444,14 @@ public void testFastQueryBigQueryException() throws InterruptedException {
24362444
} catch (BigQueryException ex) {
24372445
assertEquals(Lists.transform(errorProtoList, BigQueryError.FROM_PB_FUNCTION), ex.getErrors());
24382446
}
2439-
verify(bigqueryRpcMock).queryRpc(PROJECT, requestPb);
2447+
2448+
QueryRequest requestPb = requestPbCapture.getValue();
2449+
assertEquals(QUERY_JOB_CONFIGURATION_FOR_QUERY.getQuery(), requestPb.getQuery());
2450+
assertEquals(
2451+
QUERY_JOB_CONFIGURATION_FOR_QUERY.getDefaultDataset().getDataset(),
2452+
requestPb.getDefaultDataset().getDatasetId());
2453+
assertEquals(QUERY_JOB_CONFIGURATION_FOR_QUERY.useQueryCache(), requestPb.getUseQueryCache());
2454+
verify(bigqueryRpcMock).queryRpc(eq(PROJECT), requestPbCapture.capture());
24402455
}
24412456

24422457
@Test

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

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
package com.google.cloud.bigquery;
1818

19+
import static org.assertj.core.api.Assertions.*;
1920
import static org.junit.Assert.assertEquals;
2021

2122
import com.google.api.services.bigquery.model.QueryRequest;
@@ -31,7 +32,6 @@
3132

3233
public class QueryRequestInfoTest {
3334

34-
private static final String TEST_PROJECT_ID = "test-project-id";
3535
private static final String QUERY = "BigQuery SQL";
3636
private static final DatasetId DATASET_ID = DatasetId.of("dataset");
3737
private static final TableId TABLE_ID = TableId.of("dataset", "table");
@@ -166,8 +166,7 @@ public void equalTo() {
166166
}
167167

168168
private void compareQueryRequestInfo(QueryRequestInfo expected, QueryRequestInfo actual) {
169-
assertEquals(expected, actual);
170-
assertEquals(expected.hashCode(), actual.hashCode());
171-
assertEquals(expected.toString(), actual.toString());
169+
// requestId are expected to be different
170+
assertThat(actual).isEqualToIgnoringGivenFields(expected, "requestId");
172171
}
173172
}

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1776,6 +1776,15 @@ public void testFastQueryMultipleRuns() throws InterruptedException {
17761776
assertNull(result.getNextPageToken());
17771777
assertFalse(result.hasNextPage());
17781778

1779+
// running the same QueryJobConfiguration with the same query again
1780+
TableResult result1Duplicate = bigquery.query(config);
1781+
assertEquals(QUERY_RESULT_SCHEMA, result1Duplicate.getSchema());
1782+
assertEquals(2, result.getTotalRows());
1783+
assertNull(result1Duplicate.getNextPage());
1784+
assertNull(result1Duplicate.getNextPageToken());
1785+
assertFalse(result1Duplicate.hasNextPage());
1786+
1787+
// running a new QueryJobConfiguration with the same query
17791788
QueryJobConfiguration config2 =
17801789
QueryJobConfiguration.newBuilder(query).setDefaultDataset(DatasetId.of(DATASET)).build();
17811790
TableResult result2 = bigquery.query(config2);

‎pom.xml

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,13 @@
121121
<version>1.113.4</version>
122122
<scope>test</scope>
123123
</dependency>
124+
<dependency>
125+
<groupId>org.assertj</groupId>
126+
<artifactId>assertj-core</artifactId>
127+
<!-- use 2.9.1 for Java 7 projects -->
128+
<version>2.9.1</version>
129+
<scope>test</scope>
130+
</dependency>
124131
</dependencies>
125132
</dependencyManagement>
126133

0 commit comments

Comments
 (0)