From e2cd4f63ccc543e144f90e06eaadb2e96ce94943 Mon Sep 17 00:00:00 2001 From: Stephanie Wang Date: Fri, 4 Dec 2020 19:42:04 -0500 Subject: [PATCH] feat: set requestId for fast query path in QueryRequestInfo instead of QueryJobConfiguration (#987) As requested @epavan --- google-cloud-bigquery/pom.xml | 5 ++ .../cloud/bigquery/QueryJobConfiguration.java | 16 +--- .../cloud/bigquery/QueryRequestInfo.java | 5 +- .../cloud/bigquery/BigQueryImplTest.java | 85 +++++++++++-------- .../cloud/bigquery/QueryRequestInfoTest.java | 7 +- .../cloud/bigquery/it/ITBigQueryTest.java | 9 ++ pom.xml | 7 ++ 7 files changed, 78 insertions(+), 56 deletions(-) diff --git a/google-cloud-bigquery/pom.xml b/google-cloud-bigquery/pom.xml index 4b64bc7df..9128c564b 100644 --- a/google-cloud-bigquery/pom.xml +++ b/google-cloud-bigquery/pom.xml @@ -103,6 +103,11 @@ mockito-core test + + org.assertj + assertj-core + test + diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryJobConfiguration.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryJobConfiguration.java index 88ca5aa44..927ae846c 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryJobConfiguration.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryJobConfiguration.java @@ -35,7 +35,6 @@ import java.util.List; import java.util.Map; import java.util.Objects; -import java.util.UUID; /** * 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 { private final List connectionProperties; // maxResults is only used for fast query path private final Long maxResults; - private final String requestId; /** * Priority levels for a query. If not specified the priority is assumed to be {@link @@ -123,7 +121,6 @@ public static final class Builder private RangePartitioning rangePartitioning; private List connectionProperties; private Long maxResults; - private String requestId; private Builder() { super(Type.QUERY); @@ -157,7 +154,6 @@ private Builder(QueryJobConfiguration jobConfiguration) { this.rangePartitioning = jobConfiguration.rangePartitioning; this.connectionProperties = jobConfiguration.connectionProperties; this.maxResults = jobConfiguration.maxResults; - this.requestId = jobConfiguration.requestId; } private Builder(com.google.api.services.bigquery.model.JobConfiguration configurationPb) { @@ -625,11 +621,6 @@ public Builder setMaxResults(Long maxResults) { return this; } - Builder setRequestId(String requestId) { - this.requestId = requestId; - return this; - } - public QueryJobConfiguration build() { return new QueryJobConfiguration(this); } @@ -672,7 +663,6 @@ private QueryJobConfiguration(Builder builder) { this.rangePartitioning = builder.rangePartitioning; this.connectionProperties = builder.connectionProperties; this.maxResults = builder.maxResults; - this.requestId = builder.requestId; } /** @@ -875,10 +865,6 @@ public Long getMaxResults() { return maxResults; } - String getRequestId() { - return requestId; - } - @Override public Builder toBuilder() { return new Builder(this); @@ -1057,7 +1043,7 @@ com.google.api.services.bigquery.model.JobConfiguration toPb() { /** Creates a builder for a BigQuery Query Job given the query to be run. */ public static Builder newBuilder(String query) { checkArgument(!isNullOrEmpty(query), "Provided query is null or empty"); - return new Builder().setQuery(query).setRequestId(UUID.randomUUID().toString()); + return new Builder().setQuery(query); } /** diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java index 602404cc0..7492809d6 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java @@ -23,6 +23,7 @@ import com.google.common.collect.Lists; import java.util.List; import java.util.Map; +import java.util.UUID; final class QueryRequestInfo { @@ -35,9 +36,9 @@ final class QueryRequestInfo { private final Long maxResults; private final String query; private final List queryParameters; + private final String requestId; private final Boolean useQueryCache; private final Boolean useLegacySql; - private final String requestId; QueryRequestInfo(QueryJobConfiguration config) { this.config = config; @@ -49,9 +50,9 @@ final class QueryRequestInfo { this.maxResults = config.getMaxResults(); this.query = config.getQuery(); this.queryParameters = config.toPb().getQuery().getQueryParameters(); + this.requestId = UUID.randomUUID().toString(); this.useLegacySql = config.useLegacySql(); this.useQueryCache = config.useQueryCache(); - this.requestId = config.getRequestId(); } boolean isFastQuerySupported() { diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java index e1a48f620..7c970313f 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java @@ -1885,9 +1885,8 @@ public void testFastQueryRequestCompleted() throws InterruptedException { .setTotalBytesProcessed(42L) .setTotalRows(BigInteger.valueOf(1L)); - QueryRequestInfo requestInfo = new QueryRequestInfo(QUERY_JOB_CONFIGURATION_FOR_QUERY); - - when(bigqueryRpcMock.queryRpc(PROJECT, requestInfo.toPb())).thenReturn(queryResponsePb); + when(bigqueryRpcMock.queryRpc(eq(PROJECT), requestPbCapture.capture())) + .thenReturn(queryResponsePb); bigquery = options.getService(); TableResult result = bigquery.query(QUERY_JOB_CONFIGURATION_FOR_QUERY); @@ -1900,7 +1899,15 @@ public void testFastQueryRequestCompleted() throws InterruptedException { assertThat(row.get(0).getBooleanValue()).isFalse(); assertThat(row.get(1).getLongValue()).isEqualTo(1); } - verify(bigqueryRpcMock).queryRpc(PROJECT, requestInfo.toPb()); + + QueryRequest requestPb = requestPbCapture.getValue(); + assertEquals(QUERY_JOB_CONFIGURATION_FOR_QUERY.getQuery(), requestPb.getQuery()); + assertEquals( + QUERY_JOB_CONFIGURATION_FOR_QUERY.getDefaultDataset().getDataset(), + requestPb.getDefaultDataset().getDatasetId()); + assertEquals(QUERY_JOB_CONFIGURATION_FOR_QUERY.useQueryCache(), requestPb.getUseQueryCache()); + + verify(bigqueryRpcMock).queryRpc(eq(PROJECT), requestPbCapture.capture()); } @Test @@ -1937,14 +1944,22 @@ public void testFastQueryMultiplePages() throws InterruptedException { .setTotalBytesProcessed(42L) .setTotalRows(BigInteger.valueOf(1L)); - QueryRequestInfo requestInfo = new QueryRequestInfo(QUERY_JOB_CONFIGURATION_FOR_QUERY); - when(bigqueryRpcMock.queryRpc(PROJECT, requestInfo.toPb())).thenReturn(queryResponsePb); + when(bigqueryRpcMock.queryRpc(eq(PROJECT), requestPbCapture.capture())) + .thenReturn(queryResponsePb); bigquery = options.getService(); TableResult result = bigquery.query(QUERY_JOB_CONFIGURATION_FOR_QUERY); assertTrue(result.hasNextPage()); assertNotNull(result.getNextPageToken()); assertNotNull(result.getNextPage()); + + QueryRequest requestPb = requestPbCapture.getValue(); + assertEquals(QUERY_JOB_CONFIGURATION_FOR_QUERY.getQuery(), requestPb.getQuery()); + assertEquals( + QUERY_JOB_CONFIGURATION_FOR_QUERY.getDefaultDataset().getDataset(), + requestPb.getDefaultDataset().getDatasetId()); + assertEquals(QUERY_JOB_CONFIGURATION_FOR_QUERY.useQueryCache(), requestPb.getUseQueryCache()); + verify(bigqueryRpcMock).getJob(PROJECT, JOB, null, EMPTY_RPC_OPTIONS); verify(bigqueryRpcMock) .listTableData( @@ -1952,7 +1967,7 @@ public void testFastQueryMultiplePages() throws InterruptedException { DATASET, TABLE, BigQueryImpl.optionMap(BigQuery.TableDataListOption.pageToken(CURSOR))); - verify(bigqueryRpcMock).queryRpc(PROJECT, requestInfo.toPb()); + verify(bigqueryRpcMock).queryRpc(eq(PROJECT), requestPbCapture.capture()); } @Test @@ -1983,10 +1998,8 @@ public void testFastQuerySlowDdl() throws InterruptedException { .setTotalRows(BigInteger.valueOf(1L)) .setSchema(TABLE_SCHEMA.toPb()); - QueryRequestInfo requestInfo = new QueryRequestInfo(QUERY_JOB_CONFIGURATION_FOR_QUERY); - QueryRequest requestPb = requestInfo.toPb(); - - when(bigqueryRpcMock.queryRpc(PROJECT, requestPb)).thenReturn(queryResponsePb); + when(bigqueryRpcMock.queryRpc(eq(PROJECT), requestPbCapture.capture())) + .thenReturn(queryResponsePb); responseJob.getConfiguration().getQuery().setDestinationTable(TABLE_ID.toPb()); when(bigqueryRpcMock.getJob(PROJECT, JOB, null, EMPTY_RPC_OPTIONS)).thenReturn(responseJob); when(bigqueryRpcMock.getQueryResults( @@ -2004,7 +2017,14 @@ public void testFastQuerySlowDdl() throws InterruptedException { assertThat(row.get(1).getLongValue()).isEqualTo(1); } - verify(bigqueryRpcMock).queryRpc(PROJECT, requestInfo.toPb()); + QueryRequest requestPb = requestPbCapture.getValue(); + assertEquals(QUERY_JOB_CONFIGURATION_FOR_QUERY.getQuery(), requestPb.getQuery()); + assertEquals( + QUERY_JOB_CONFIGURATION_FOR_QUERY.getDefaultDataset().getDataset(), + requestPb.getDefaultDataset().getDatasetId()); + assertEquals(QUERY_JOB_CONFIGURATION_FOR_QUERY.useQueryCache(), requestPb.getUseQueryCache()); + + verify(bigqueryRpcMock).queryRpc(eq(PROJECT), requestPbCapture.capture()); verify(bigqueryRpcMock).getJob(PROJECT, JOB, null, EMPTY_RPC_OPTIONS); verify(bigqueryRpcMock) .getQueryResults( @@ -2287,9 +2307,6 @@ public void testFastQuerySQLShouldRetry() throws Exception { .setTotalRows(BigInteger.valueOf(1L)) .setSchema(TABLE_SCHEMA.toPb()); - QueryRequestInfo requestInfo = new QueryRequestInfo(QUERY_JOB_CONFIGURATION_FOR_QUERY); - QueryRequest requestPb = requestInfo.toPb(); - when(bigqueryRpcMock.queryRpc(eq(PROJECT), requestPbCapture.capture())) .thenThrow(new BigQueryException(500, "InternalError")) .thenThrow(new BigQueryException(502, "Bad Gateway")) @@ -2310,13 +2327,13 @@ public void testFastQuerySQLShouldRetry() throws Exception { List allRequests = requestPbCapture.getAllValues(); boolean idempotent = true; - String requestId = requestPb.getRequestId(); + String firstRequestId = allRequests.get(0).getRequestId(); for (QueryRequest request : allRequests) { - idempotent = request.getRequestId().equals(requestId); + idempotent = request.getRequestId().equals(firstRequestId); } assertTrue(idempotent); - verify(bigqueryRpcMock, times(5)).queryRpc(PROJECT, requestPb); + verify(bigqueryRpcMock, times(5)).queryRpc(eq(PROJECT), requestPbCapture.capture()); } @Test @@ -2331,9 +2348,6 @@ public void testFastQueryDMLShouldRetry() throws Exception { .setNumDmlAffectedRows(1L) .setSchema(TABLE_SCHEMA.toPb()); - QueryRequestInfo requestInfo = new QueryRequestInfo(QUERY_JOB_CONFIGURATION_FOR_DMLQUERY); - QueryRequest requestPb = requestInfo.toPb(); - when(bigqueryRpcMock.queryRpc(eq(PROJECT), requestPbCapture.capture())) .thenThrow(new BigQueryException(500, "InternalError")) .thenThrow(new BigQueryException(502, "Bad Gateway")) @@ -2354,13 +2368,13 @@ public void testFastQueryDMLShouldRetry() throws Exception { List allRequests = requestPbCapture.getAllValues(); boolean idempotent = true; - String requestId = requestPb.getRequestId(); + String firstRequestId = allRequests.get(0).getRequestId(); for (QueryRequest request : allRequests) { - idempotent = request.getRequestId().equals(requestId); + idempotent = request.getRequestId().equals(firstRequestId); } assertTrue(idempotent); - verify(bigqueryRpcMock, times(5)).queryRpc(PROJECT, requestPb); + verify(bigqueryRpcMock, times(5)).queryRpc(eq(PROJECT), requestPbCapture.capture()); } @Test @@ -2374,9 +2388,6 @@ public void testFastQueryDDLShouldRetry() throws Exception { .setTotalBytesProcessed(42L) .setSchema(TABLE_SCHEMA.toPb()); - QueryRequestInfo requestInfo = new QueryRequestInfo(QUERY_JOB_CONFIGURATION_FOR_DDLQUERY); - QueryRequest requestPb = requestInfo.toPb(); - when(bigqueryRpcMock.queryRpc(eq(PROJECT), requestPbCapture.capture())) .thenThrow(new BigQueryException(500, "InternalError")) .thenThrow(new BigQueryException(502, "Bad Gateway")) @@ -2397,13 +2408,13 @@ public void testFastQueryDDLShouldRetry() throws Exception { List allRequests = requestPbCapture.getAllValues(); boolean idempotent = true; - String requestId = requestPb.getRequestId(); + String firstRequestId = allRequests.get(0).getRequestId(); for (QueryRequest request : allRequests) { - idempotent = request.getRequestId().equals(requestId); + idempotent = request.getRequestId().equals(firstRequestId); } assertTrue(idempotent); - verify(bigqueryRpcMock, times(5)).queryRpc(PROJECT, requestPb); + verify(bigqueryRpcMock, times(5)).queryRpc(eq(PROJECT), requestPbCapture.capture()); } @Test @@ -2424,10 +2435,7 @@ public void testFastQueryBigQueryException() throws InterruptedException { .setPageToken(null) .setErrors(errorProtoList); - QueryRequestInfo requestInfo = new QueryRequestInfo(QUERY_JOB_CONFIGURATION_FOR_QUERY); - QueryRequest requestPb = requestInfo.toPb(); - - when(bigqueryRpcMock.queryRpc(PROJECT, requestPb)).thenReturn(responsePb); + when(bigqueryRpcMock.queryRpc(eq(PROJECT), requestPbCapture.capture())).thenReturn(responsePb); bigquery = options.getService(); try { @@ -2436,7 +2444,14 @@ public void testFastQueryBigQueryException() throws InterruptedException { } catch (BigQueryException ex) { assertEquals(Lists.transform(errorProtoList, BigQueryError.FROM_PB_FUNCTION), ex.getErrors()); } - verify(bigqueryRpcMock).queryRpc(PROJECT, requestPb); + + QueryRequest requestPb = requestPbCapture.getValue(); + assertEquals(QUERY_JOB_CONFIGURATION_FOR_QUERY.getQuery(), requestPb.getQuery()); + assertEquals( + QUERY_JOB_CONFIGURATION_FOR_QUERY.getDefaultDataset().getDataset(), + requestPb.getDefaultDataset().getDatasetId()); + assertEquals(QUERY_JOB_CONFIGURATION_FOR_QUERY.useQueryCache(), requestPb.getUseQueryCache()); + verify(bigqueryRpcMock).queryRpc(eq(PROJECT), requestPbCapture.capture()); } @Test diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryRequestInfoTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryRequestInfoTest.java index 48fb2dbc4..4adeeaf00 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryRequestInfoTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryRequestInfoTest.java @@ -16,6 +16,7 @@ package com.google.cloud.bigquery; +import static org.assertj.core.api.Assertions.*; import static org.junit.Assert.assertEquals; import com.google.api.services.bigquery.model.QueryRequest; @@ -31,7 +32,6 @@ public class QueryRequestInfoTest { - private static final String TEST_PROJECT_ID = "test-project-id"; private static final String QUERY = "BigQuery SQL"; private static final DatasetId DATASET_ID = DatasetId.of("dataset"); private static final TableId TABLE_ID = TableId.of("dataset", "table"); @@ -166,8 +166,7 @@ public void equalTo() { } private void compareQueryRequestInfo(QueryRequestInfo expected, QueryRequestInfo actual) { - assertEquals(expected, actual); - assertEquals(expected.hashCode(), actual.hashCode()); - assertEquals(expected.toString(), actual.toString()); + // requestId are expected to be different + assertThat(actual).isEqualToIgnoringGivenFields(expected, "requestId"); } } diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java index 07eeb9f3a..155de0366 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java @@ -1776,6 +1776,15 @@ public void testFastQueryMultipleRuns() throws InterruptedException { assertNull(result.getNextPageToken()); assertFalse(result.hasNextPage()); + // running the same QueryJobConfiguration with the same query again + TableResult result1Duplicate = bigquery.query(config); + assertEquals(QUERY_RESULT_SCHEMA, result1Duplicate.getSchema()); + assertEquals(2, result.getTotalRows()); + assertNull(result1Duplicate.getNextPage()); + assertNull(result1Duplicate.getNextPageToken()); + assertFalse(result1Duplicate.hasNextPage()); + + // running a new QueryJobConfiguration with the same query QueryJobConfiguration config2 = QueryJobConfiguration.newBuilder(query).setDefaultDataset(DatasetId.of(DATASET)).build(); TableResult result2 = bigquery.query(config2); diff --git a/pom.xml b/pom.xml index be7b1df65..d190e81c3 100644 --- a/pom.xml +++ b/pom.xml @@ -121,6 +121,13 @@ 1.113.4 test + + org.assertj + assertj-core + + 2.9.1 + test +