diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java index 705b3027ed..15b16e32f0 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java @@ -185,11 +185,13 @@ public static EnhancedBigtableStubSettings finalizeSettings( // workaround JWT audience issues patchCredentials(builder); - // patch cookies interceptor - InstantiatingGrpcChannelProvider.Builder transportProvider = null; - if (builder.getTransportChannelProvider() instanceof InstantiatingGrpcChannelProvider) { - transportProvider = - ((InstantiatingGrpcChannelProvider) builder.getTransportChannelProvider()).toBuilder(); + InstantiatingGrpcChannelProvider.Builder transportProvider = + builder.getTransportChannelProvider() instanceof InstantiatingGrpcChannelProvider + ? ((InstantiatingGrpcChannelProvider) builder.getTransportChannelProvider()).toBuilder() + : null; + + if (builder.getEnableRoutingCookie() && transportProvider != null) { + // patch cookies interceptor transportProvider.setInterceptorProvider(() -> ImmutableList.of(new CookiesInterceptor())); } @@ -371,9 +373,12 @@ public ServerStreamingCallable createReadRowsCallable( new TracedServerStreamingCallable<>( readRowsUserCallable, clientContext.getTracerFactory(), span); - // CookieHolder needs to be injected to the CallOptions outside of retries, otherwise retry - // attempts won't see a CookieHolder. - ServerStreamingCallable withCookie = new CookiesServerStreamingCallable<>(traced); + ServerStreamingCallable withCookie = traced; + if (settings.getEnableRoutingCookie()) { + // CookieHolder needs to be injected to the CallOptions outside of retries, otherwise retry + // attempts won't see a CookieHolder. + withCookie = new CookiesServerStreamingCallable<>(traced); + } return withCookie.withDefaultCallContext(clientContext.getDefaultCallContext()); } @@ -411,7 +416,10 @@ public UnaryCallable createReadRowCallable(RowAdapter new TracedUnaryCallable<>( firstRow, clientContext.getTracerFactory(), getSpanName("ReadRow")); - UnaryCallable withCookie = new CookiesUnaryCallable<>(traced); + UnaryCallable withCookie = traced; + if (settings.getEnableRoutingCookie()) { + withCookie = new CookiesUnaryCallable<>(traced); + } return withCookie.withDefaultCallContext(clientContext.getDefaultCallContext()); } @@ -654,7 +662,10 @@ private UnaryCallable createBulkMutateRowsCallable() { new TracedUnaryCallable<>( tracedBatcherUnaryCallable, clientContext.getTracerFactory(), spanName); - UnaryCallable withCookie = new CookiesUnaryCallable<>(traced); + UnaryCallable withCookie = traced; + if (settings.getEnableRoutingCookie()) { + withCookie = new CookiesUnaryCallable<>(traced); + } return withCookie.withDefaultCallContext(clientContext.getDefaultCallContext()); } @@ -938,8 +949,10 @@ public Map extract( ServerStreamingCallable traced = new TracedServerStreamingCallable<>(retrying, clientContext.getTracerFactory(), span); - ServerStreamingCallable withCookie = - new CookiesServerStreamingCallable<>(traced); + ServerStreamingCallable withCookie = traced; + if (settings.getEnableRoutingCookie()) { + withCookie = new CookiesServerStreamingCallable<>(traced); + } return withCookie.withDefaultCallContext(clientContext.getDefaultCallContext()); } @@ -1021,8 +1034,10 @@ public Map extract( new TracedServerStreamingCallable<>( readChangeStreamUserCallable, clientContext.getTracerFactory(), span); - ServerStreamingCallable withCookie = - new CookiesServerStreamingCallable<>(traced); + ServerStreamingCallable withCookie = traced; + if (settings.getEnableRoutingCookie()) { + withCookie = new CookiesServerStreamingCallable<>(traced); + } return withCookie.withDefaultCallContext(clientContext.getDefaultCallContext()); } @@ -1037,9 +1052,12 @@ private UnaryCallable createUserFacin UnaryCallable traced = new TracedUnaryCallable<>(inner, clientContext.getTracerFactory(), getSpanName(methodName)); - // CookieHolder needs to be injected to the CallOptions outside of retries, otherwise retry - // attempts won't see a CookieHolder. - UnaryCallable withCookie = new CookiesUnaryCallable<>(traced); + UnaryCallable withCookie = traced; + if (settings.getEnableRoutingCookie()) { + // CookieHolder needs to be injected to the CallOptions outside of retries, otherwise retry + // attempts won't see a CookieHolder. + withCookie = new CookiesUnaryCallable<>(traced); + } return withCookie.withDefaultCallContext(clientContext.getDefaultCallContext()); } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java index cffd9c85df..d099311ffd 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java @@ -211,6 +211,7 @@ public class EnhancedBigtableStubSettings extends StubSettings primedTableIds; private final Map jwtAudienceMapping; + private final boolean enableRoutingCookie; private final ServerStreamingCallSettings readRowsSettings; private final UnaryCallSettings readRowSettings; @@ -252,6 +253,7 @@ private EnhancedBigtableStubSettings(Builder builder) { isRefreshingChannel = builder.isRefreshingChannel; primedTableIds = builder.primedTableIds; jwtAudienceMapping = builder.jwtAudienceMapping; + enableRoutingCookie = builder.enableRoutingCookie; // Per method settings. readRowsSettings = builder.readRowsSettings.build(); @@ -313,6 +315,14 @@ public Map getJwtAudienceMapping() { return jwtAudienceMapping; } + /** + * Gets if routing cookie is enabled. If true, client will retry a request with extra metadata + * server sent back. + */ + public boolean getEnableRoutingCookie() { + return enableRoutingCookie; + } + /** Returns a builder for the default ChannelProvider for this service. */ public static InstantiatingGrpcChannelProvider.Builder defaultGrpcTransportProviderBuilder() { return BigtableStubSettings.defaultGrpcTransportProviderBuilder() @@ -595,6 +605,7 @@ public static class Builder extends StubSettings.Builder primedTableIds; private Map jwtAudienceMapping; + private boolean enableRoutingCookie; private final ServerStreamingCallSettings.Builder readRowsSettings; private final UnaryCallSettings.Builder readRowSettings; @@ -627,6 +638,7 @@ private Builder() { primedTableIds = ImmutableList.of(); jwtAudienceMapping = DEFAULT_JWT_AUDIENCE_MAPPING; setCredentialsProvider(defaultCredentialsProviderBuilder().build()); + this.enableRoutingCookie = true; // Defaults provider BigtableStubSettings.Builder baseDefaults = BigtableStubSettings.newBuilder(); @@ -745,6 +757,7 @@ private Builder(EnhancedBigtableStubSettings settings) { isRefreshingChannel = settings.isRefreshingChannel; primedTableIds = settings.primedTableIds; jwtAudienceMapping = settings.jwtAudienceMapping; + enableRoutingCookie = settings.enableRoutingCookie; // Per method settings. readRowsSettings = settings.readRowsSettings.toBuilder(); @@ -893,6 +906,23 @@ public Map getJwtAudienceMapping() { return jwtAudienceMapping; } + /** + * Sets if routing cookie is enabled. If true, client will retry a request with extra metadata + * server sent back. + */ + public Builder setEnableRoutingCookie(boolean enableRoutingCookie) { + this.enableRoutingCookie = enableRoutingCookie; + return this; + } + + /** + * Gets if routing cookie is enabled. If true, client will retry a request with extra metadata + * server sent back. + */ + public boolean getEnableRoutingCookie() { + return enableRoutingCookie; + } + /** Returns the builder for the settings used for calls to readRows. */ public ServerStreamingCallSettings.Builder readRowsSettings() { return readRowsSettings; @@ -1019,6 +1049,7 @@ public String toString() { .add("isRefreshingChannel", isRefreshingChannel) .add("primedTableIds", primedTableIds) .add("jwtAudienceMapping", jwtAudienceMapping) + .add("enableRoutingCookie", enableRoutingCookie) .add("readRowsSettings", readRowsSettings) .add("readRowSettings", readRowSettings) .add("sampleRowKeysSettings", sampleRowKeysSettings) diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/CookiesHolderTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/CookiesHolderTest.java index 5dac053523..a52bd84592 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/CookiesHolderTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/CookiesHolderTest.java @@ -58,6 +58,7 @@ import io.grpc.Status; import io.grpc.StatusRuntimeException; import io.grpc.stub.StreamObserver; +import java.io.IOException; import java.util.ArrayList; import java.util.HashSet; import java.util.List; @@ -83,6 +84,7 @@ public class CookiesHolderTest { private Server server; private final FakeService fakeService = new FakeService(); + private BigtableDataSettings.Builder settings; private BigtableDataClient client; private final List serverMetadata = new ArrayList<>(); @@ -138,6 +140,8 @@ public ServerCall.Listener interceptCall( .build()) .setRetryableCodes(StatusCode.Code.UNAVAILABLE); + this.settings = settings; + client = BigtableDataClient.create(settings.build()); } @@ -379,7 +383,7 @@ public void sendHeaders(Metadata headers) { } @Test - public void testAllMethodsAreCalled() throws InterruptedException { + public void testAllMethodsAreCalled() { // This test ensures that all methods respect the retry cookie except for the ones that are // explicitly added to the methods list. It requires that any newly method is exercised in this // test. This is enforced by introspecting grpc method descriptors. @@ -422,6 +426,53 @@ public void testAllMethodsAreCalled() throws InterruptedException { assertThat(methods).containsExactlyElementsIn(expected); } + @Test + public void testDisableRoutingCookie() throws IOException { + // This test disables routing cookie in the client settings and ensures that none of the routing + // cookie + // is added. + settings.stubSettings().setEnableRoutingCookie(false); + try (BigtableDataClient client = BigtableDataClient.create(settings.build())) { + client.readRows(Query.create("fake-table")).iterator().hasNext(); + assertThat(fakeService.count.get()).isEqualTo(2); + fakeService.count.set(0); + + client.mutateRow(RowMutation.create("fake-table", "key").setCell("cf", "q", "v")); + assertThat(fakeService.count.get()).isEqualTo(2); + fakeService.count.set(0); + + client.bulkMutateRows( + BulkMutation.create("fake-table") + .add(RowMutationEntry.create("key").setCell("cf", "q", "v"))); + assertThat(fakeService.count.get()).isEqualTo(2); + fakeService.count.set(0); + + client.sampleRowKeys("fake-table"); + assertThat(fakeService.count.get()).isEqualTo(2); + fakeService.count.set(0); + + client.checkAndMutateRow( + ConditionalRowMutation.create("fake-table", "key") + .then(Mutation.create().setCell("cf", "q", "v"))); + assertThat(fakeService.count.get()).isEqualTo(2); + fakeService.count.set(0); + + client.readModifyWriteRow( + ReadModifyWriteRow.create("fake-table", "key").append("cf", "q", "v")); + assertThat(fakeService.count.get()).isEqualTo(2); + fakeService.count.set(0); + + client.generateInitialChangeStreamPartitions("fake-table").iterator().hasNext(); + assertThat(fakeService.count.get()).isEqualTo(2); + fakeService.count.set(0); + + client.readChangeStream(ReadChangeStreamQuery.create("fake-table")).iterator().hasNext(); + assertThat(fakeService.count.get()).isEqualTo(2); + + assertThat(methods).isEmpty(); + } + } + static class FakeService extends BigtableGrpc.BigtableImplBase { private boolean returnCookie = true; diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettingsTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettingsTest.java index fbd6442e0c..4dd763cc49 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettingsTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettingsTest.java @@ -77,6 +77,7 @@ public void settingsAreNotLostTest() { CredentialsProvider credentialsProvider = Mockito.mock(CredentialsProvider.class); WatchdogProvider watchdogProvider = Mockito.mock(WatchdogProvider.class); Duration watchdogInterval = Duration.ofSeconds(12); + boolean enableRoutingCookie = false; EnhancedBigtableStubSettings.Builder builder = EnhancedBigtableStubSettings.newBuilder() @@ -87,7 +88,8 @@ public void settingsAreNotLostTest() { .setEndpoint(endpoint) .setCredentialsProvider(credentialsProvider) .setStreamWatchdogProvider(watchdogProvider) - .setStreamWatchdogCheckInterval(watchdogInterval); + .setStreamWatchdogCheckInterval(watchdogInterval) + .setEnableRoutingCookie(enableRoutingCookie); verifyBuilder( builder, @@ -98,7 +100,8 @@ public void settingsAreNotLostTest() { endpoint, credentialsProvider, watchdogProvider, - watchdogInterval); + watchdogInterval, + enableRoutingCookie); verifySettings( builder.build(), projectId, @@ -108,7 +111,8 @@ public void settingsAreNotLostTest() { endpoint, credentialsProvider, watchdogProvider, - watchdogInterval); + watchdogInterval, + enableRoutingCookie); verifyBuilder( builder.build().toBuilder(), projectId, @@ -118,7 +122,8 @@ public void settingsAreNotLostTest() { endpoint, credentialsProvider, watchdogProvider, - watchdogInterval); + watchdogInterval, + enableRoutingCookie); } private void verifyBuilder( @@ -130,7 +135,8 @@ private void verifyBuilder( String endpoint, CredentialsProvider credentialsProvider, WatchdogProvider watchdogProvider, - Duration watchdogInterval) { + Duration watchdogInterval, + boolean enableRoutingCookie) { assertThat(builder.getProjectId()).isEqualTo(projectId); assertThat(builder.getInstanceId()).isEqualTo(instanceId); assertThat(builder.getAppProfileId()).isEqualTo(appProfileId); @@ -139,6 +145,7 @@ private void verifyBuilder( assertThat(builder.getCredentialsProvider()).isEqualTo(credentialsProvider); assertThat(builder.getStreamWatchdogProvider()).isSameInstanceAs(watchdogProvider); assertThat(builder.getStreamWatchdogCheckInterval()).isEqualTo(watchdogInterval); + assertThat(builder.getEnableRoutingCookie()).isEqualTo(enableRoutingCookie); } private void verifySettings( @@ -150,7 +157,8 @@ private void verifySettings( String endpoint, CredentialsProvider credentialsProvider, WatchdogProvider watchdogProvider, - Duration watchdogInterval) { + Duration watchdogInterval, + boolean enableRoutingCookie) { assertThat(settings.getProjectId()).isEqualTo(projectId); assertThat(settings.getInstanceId()).isEqualTo(instanceId); assertThat(settings.getAppProfileId()).isEqualTo(appProfileId); @@ -159,6 +167,7 @@ private void verifySettings( assertThat(settings.getCredentialsProvider()).isEqualTo(credentialsProvider); assertThat(settings.getStreamWatchdogProvider()).isSameInstanceAs(watchdogProvider); assertThat(settings.getStreamWatchdogCheckInterval()).isEqualTo(watchdogInterval); + assertThat(settings.getEnableRoutingCookie()).isEqualTo(enableRoutingCookie); } @Test @@ -781,6 +790,33 @@ public void isRefreshingChannelFalseValueTest() { assertThat(builder.build().toBuilder().isRefreshingChannel()).isFalse(); } + @Test + public void routingCookieIsEnabled() { + String dummyProjectId = "my-project"; + String dummyInstanceId = "my-instance"; + EnhancedBigtableStubSettings.Builder builder = + EnhancedBigtableStubSettings.newBuilder() + .setProjectId(dummyProjectId) + .setInstanceId(dummyInstanceId); + assertThat(builder.getEnableRoutingCookie()).isTrue(); + assertThat(builder.build().getEnableRoutingCookie()).isTrue(); + assertThat(builder.build().toBuilder().getEnableRoutingCookie()).isTrue(); + } + + @Test + public void routingCookieFalseValueSet() { + String dummyProjectId = "my-project"; + String dummyInstanceId = "my-instance"; + EnhancedBigtableStubSettings.Builder builder = + EnhancedBigtableStubSettings.newBuilder() + .setProjectId(dummyProjectId) + .setInstanceId(dummyInstanceId) + .setEnableRoutingCookie(false); + assertThat(builder.getEnableRoutingCookie()).isFalse(); + assertThat(builder.build().getEnableRoutingCookie()).isFalse(); + assertThat(builder.build().toBuilder().getEnableRoutingCookie()).isFalse(); + } + static final String[] SETTINGS_LIST = { "projectId", "instanceId", @@ -788,6 +824,7 @@ public void isRefreshingChannelFalseValueTest() { "isRefreshingChannel", "primedTableIds", "jwtAudienceMapping", + "enableRoutingCookie", "readRowsSettings", "readRowSettings", "sampleRowKeysSettings",