From 0dac26e47b54ac279dbe0794dbb14d6aadd3b3fe Mon Sep 17 00:00:00 2001 From: Kenneth Jung Date: Mon, 25 Mar 2019 20:44:56 -0700 Subject: [PATCH 1/4] Pass through header provider in BQ storage client settings. This change modifies the enhanced stub for the BigQuery storage client to pass through the user-specified header provider to the storage settings base class at creationt ime. This addresses an issue where user-specified headers were being ignored when creating a client object. --- .../stub/EnhancedBigQueryStorageStub.java | 1 + .../v1beta1/stub/ResourceHeaderTest.java | 173 ++++++++++++++++++ 2 files changed, 174 insertions(+) create mode 100644 google-cloud-clients/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta1/stub/ResourceHeaderTest.java diff --git a/google-cloud-clients/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta1/stub/EnhancedBigQueryStorageStub.java b/google-cloud-clients/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta1/stub/EnhancedBigQueryStorageStub.java index d46d9c7bf6af..56ad2f96cd71 100644 --- a/google-cloud-clients/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta1/stub/EnhancedBigQueryStorageStub.java +++ b/google-cloud-clients/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta1/stub/EnhancedBigQueryStorageStub.java @@ -49,6 +49,7 @@ public static EnhancedBigQueryStorageStub create(EnhancedBigQueryStorageStubSett BigQueryStorageStubSettings.newBuilder() .setTransportChannelProvider(settings.getTransportChannelProvider()) .setEndpoint(settings.getEndpoint()) + .setHeaderProvider(settings.getHeaderProvider()) .setCredentialsProvider(settings.getCredentialsProvider()) .setStreamWatchdogCheckInterval(settings.getStreamWatchdogCheckInterval()) .setStreamWatchdogProvider(settings.getStreamWatchdogProvider()); diff --git a/google-cloud-clients/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta1/stub/ResourceHeaderTest.java b/google-cloud-clients/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta1/stub/ResourceHeaderTest.java new file mode 100644 index 000000000000..7a6de57bb081 --- /dev/null +++ b/google-cloud-clients/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta1/stub/ResourceHeaderTest.java @@ -0,0 +1,173 @@ +/* + * Copyright 2018 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.bigquery.storage.v1beta1.stub; + +import static com.google.common.truth.Truth.assertWithMessage; + +import com.google.api.gax.core.NoCredentialsProvider; +import com.google.api.gax.grpc.testing.InProcessServer; +import com.google.api.gax.grpc.testing.LocalChannelProvider; +import com.google.api.gax.rpc.FixedHeaderProvider; +import com.google.api.gax.rpc.UnimplementedException; +import com.google.cloud.bigquery.storage.v1beta1.BigQueryStorageClient; +import com.google.cloud.bigquery.storage.v1beta1.BigQueryStorageGrpc.BigQueryStorageImplBase; +import com.google.cloud.bigquery.storage.v1beta1.BigQueryStorageSettings; +import com.google.cloud.bigquery.storage.v1beta1.Storage.ReadRowsRequest; +import com.google.cloud.bigquery.storage.v1beta1.Storage.ReadSession; +import com.google.cloud.bigquery.storage.v1beta1.Storage.Stream; +import com.google.cloud.bigquery.storage.v1beta1.Storage.StreamPosition; +import com.google.cloud.bigquery.storage.v1beta1.TableReferenceProto.TableReference; +import java.util.regex.Pattern; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class ResourceHeaderTest { + + private static final TableReference TEST_TABLE_REFERENCE = + TableReference.newBuilder() + .setProjectId("project") + .setDatasetId("dataset") + .setTableId("table") + .build(); + + private static final ReadSession TEST_SESSION = + ReadSession.newBuilder().setName("sessionName").build(); + + private static final Stream TEST_STREAM = + Stream.newBuilder().setName("streamName").build(); + + private static final String NAME = "resource-header-test:123"; + + private static final String HEADER_NAME = "x-goog-request-params"; + + private static final Pattern DATASET_PATTERN = + Pattern.compile(".*" + "table_reference\\.dataset_id=dataset" + ".*"); + private static final Pattern ORIGINAL_STREAM_PATTERN = + Pattern.compile(".*" + "original_stream\\.name=streamName" + ".*"); + private static final Pattern PROJECT_PATTERN = + Pattern.compile(".*" + "table_reference\\.project_id=project" + ".*"); + private static final Pattern READ_POSITION_PATTERN = + Pattern.compile(".*" + "read_position\\.stream\\.name=streamName" + ".*"); + private static final Pattern SESSION_PATTERN = + Pattern.compile(".*" + "session\\.name=sessionName" + ".*"); + private static final Pattern STREAM_PATTERN = + Pattern.compile(".*" + "stream\\.name=streamName" + ".*"); + + private static final String TEST_HEADER_NAME = "simple-header-name"; + private static final String TEST_HEADER_VALUE = "simple-header-value"; + private static final Pattern TEST_PATTERN = Pattern.compile(".*" + TEST_HEADER_VALUE + ".*"); + + private InProcessServer server; + private LocalChannelProvider channelProvider; + private BigQueryStorageClient client; + + @Before + public void setUp() throws Exception { + server = new InProcessServer<>(new BigQueryStorageImplBase() {}, NAME); + server.start(); + channelProvider = LocalChannelProvider.create(NAME); + + BigQueryStorageSettings.Builder settingsBuilder = + BigQueryStorageSettings.newBuilder() + .setCredentialsProvider(NoCredentialsProvider.create()) + .setHeaderProvider(FixedHeaderProvider.create(TEST_HEADER_NAME, TEST_HEADER_VALUE)) + .setTransportChannelProvider(channelProvider); + + client = BigQueryStorageClient.create(settingsBuilder.build()); + } + + @After + public void tearDown() throws Exception { + client.close(); + server.stop(); + server.blockUntilShutdown(); + } + + @Test + public void createReadSessionTest() { + try { + client.createReadSession(TEST_TABLE_REFERENCE, "parents/project", 1); + } catch (UnimplementedException e) { + // Ignore the error: none of the methods are actually implemented. + } + verifyHeaderSent(PROJECT_PATTERN, DATASET_PATTERN); + } + + @Test + public void readRowsTest() { + try { + ReadRowsRequest request = + ReadRowsRequest.newBuilder() + .setReadPosition( + StreamPosition.newBuilder() + .setStream(TEST_STREAM) + .setOffset(125)) + .build(); + client.readRowsCallable().call(request); + } catch (UnimplementedException e) { + // Ignore the error: none of the methods are actually implemented. + } + + verifyHeaderSent(READ_POSITION_PATTERN); + } + + @Test + public void batchCreateReadStreamsForSessionTest() { + try { + client.batchCreateReadSessionStreams(TEST_SESSION, 1); + } catch (UnimplementedException e) { + // Ignore the error: none of the methods are actually implemented. + } + + verifyHeaderSent(SESSION_PATTERN); + } + + @Test + public void finalizeStreamTest() { + try { + client.finalizeStream(TEST_STREAM); + } catch (UnimplementedException e) { + // Ignore the error: none of the methods are actually implemented. + } + + verifyHeaderSent(STREAM_PATTERN); + } + + @Test + public void splitReadStreamTest() { + try { + client.splitReadStream(TEST_STREAM); + } catch (UnimplementedException e) { + // Ignore the error: none of the methods are actually implemented. + } + + verifyHeaderSent(ORIGINAL_STREAM_PATTERN); + } + + private void verifyHeaderSent(Pattern... patterns) { + for (Pattern pattern : patterns) { + boolean headerSent = channelProvider.isHeaderSent(HEADER_NAME, pattern); + assertWithMessage("Generated header was sent").that(headerSent).isTrue(); + } + boolean testHeaderSent = channelProvider.isHeaderSent(TEST_HEADER_NAME, TEST_PATTERN); + assertWithMessage("Provided header was sent").that(testHeaderSent).isTrue(); + } + +} From 8463f1814e48a013f5397cfdcf99db04c47e7365 Mon Sep 17 00:00:00 2001 From: Kenneth Jung Date: Mon, 25 Mar 2019 22:58:57 -0700 Subject: [PATCH 2/4] Apply formatting plugin --- .../storage/v1beta1/stub/ResourceHeaderTest.java | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/google-cloud-clients/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta1/stub/ResourceHeaderTest.java b/google-cloud-clients/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta1/stub/ResourceHeaderTest.java index 7a6de57bb081..c431fcccb0b3 100644 --- a/google-cloud-clients/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta1/stub/ResourceHeaderTest.java +++ b/google-cloud-clients/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta1/stub/ResourceHeaderTest.java @@ -50,8 +50,7 @@ public class ResourceHeaderTest { private static final ReadSession TEST_SESSION = ReadSession.newBuilder().setName("sessionName").build(); - private static final Stream TEST_STREAM = - Stream.newBuilder().setName("streamName").build(); + private static final Stream TEST_STREAM = Stream.newBuilder().setName("streamName").build(); private static final String NAME = "resource-header-test:123"; @@ -115,10 +114,7 @@ public void readRowsTest() { try { ReadRowsRequest request = ReadRowsRequest.newBuilder() - .setReadPosition( - StreamPosition.newBuilder() - .setStream(TEST_STREAM) - .setOffset(125)) + .setReadPosition(StreamPosition.newBuilder().setStream(TEST_STREAM).setOffset(125)) .build(); client.readRowsCallable().call(request); } catch (UnimplementedException e) { @@ -169,5 +165,4 @@ private void verifyHeaderSent(Pattern... patterns) { boolean testHeaderSent = channelProvider.isHeaderSent(TEST_HEADER_NAME, TEST_PATTERN); assertWithMessage("Provided header was sent").that(testHeaderSent).isTrue(); } - } From fa24a61e260740e536e1e572a0365290db200515 Mon Sep 17 00:00:00 2001 From: Kenneth Jung Date: Mon, 25 Mar 2019 23:50:26 -0700 Subject: [PATCH 3/4] Set up the in-proc server only once --- .../v1beta1/stub/ResourceHeaderTest.java | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/google-cloud-clients/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta1/stub/ResourceHeaderTest.java b/google-cloud-clients/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta1/stub/ResourceHeaderTest.java index c431fcccb0b3..7a9f9a4a44e9 100644 --- a/google-cloud-clients/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta1/stub/ResourceHeaderTest.java +++ b/google-cloud-clients/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta1/stub/ResourceHeaderTest.java @@ -32,7 +32,9 @@ import com.google.cloud.bigquery.storage.v1beta1.TableReferenceProto.TableReference; import java.util.regex.Pattern; import org.junit.After; +import org.junit.AfterClass; import org.junit.Before; +import org.junit.BeforeClass; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -73,28 +75,35 @@ public class ResourceHeaderTest { private static final String TEST_HEADER_VALUE = "simple-header-value"; private static final Pattern TEST_PATTERN = Pattern.compile(".*" + TEST_HEADER_VALUE + ".*"); - private InProcessServer server; + private static InProcessServer server; + private LocalChannelProvider channelProvider; private BigQueryStorageClient client; - @Before - public void setUp() throws Exception { + @BeforeClass + public static void setUpClass() throws Exception { server = new InProcessServer<>(new BigQueryStorageImplBase() {}, NAME); server.start(); - channelProvider = LocalChannelProvider.create(NAME); + } + @Before + public void setUp() throws Exception { + channelProvider = LocalChannelProvider.create(NAME); BigQueryStorageSettings.Builder settingsBuilder = BigQueryStorageSettings.newBuilder() .setCredentialsProvider(NoCredentialsProvider.create()) .setHeaderProvider(FixedHeaderProvider.create(TEST_HEADER_NAME, TEST_HEADER_VALUE)) .setTransportChannelProvider(channelProvider); - client = BigQueryStorageClient.create(settingsBuilder.build()); } @After public void tearDown() throws Exception { client.close(); + } + + @AfterClass + public static void tearDownClass() throws Exception { server.stop(); server.blockUntilShutdown(); } From c077c463c9905aaacb66f2252b390dc5ebb0b779 Mon Sep 17 00:00:00 2001 From: Kenneth Jung Date: Tue, 26 Mar 2019 10:05:34 -0700 Subject: [PATCH 4/4] It's 2019 now --- .../cloud/bigquery/storage/v1beta1/stub/ResourceHeaderTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/google-cloud-clients/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta1/stub/ResourceHeaderTest.java b/google-cloud-clients/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta1/stub/ResourceHeaderTest.java index 7a9f9a4a44e9..c2e2df1ebbf8 100644 --- a/google-cloud-clients/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta1/stub/ResourceHeaderTest.java +++ b/google-cloud-clients/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta1/stub/ResourceHeaderTest.java @@ -1,5 +1,5 @@ /* - * Copyright 2018 Google LLC + * Copyright 2019 Google LLC * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License.