From 005a5c0ed097bd0e05b585b19bac2e757021f7ca Mon Sep 17 00:00:00 2001 From: Teng Zhong Date: Fri, 15 Jul 2022 17:15:06 -0400 Subject: [PATCH 1/6] feat: Add ListChangeStreamPartitions callable --- .../bigtable/data/v2/BigtableDataClient.java | 136 ++++++++++++++++++ .../data/v2/stub/EnhancedBigtableStub.java | 77 ++++++++++ .../v2/stub/EnhancedBigtableStubSettings.java | 36 +++++ ...istChangeStreamPartitionsUserCallable.java | 92 ++++++++++++ .../data/v2/BigtableDataClientTests.java | 34 +++++ .../EnhancedBigtableStubSettingsTest.java | 1 + ...hangeStreamPartitionsUserCallableTest.java | 83 +++++++++++ 7 files changed, 459 insertions(+) create mode 100644 google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/changestream/ListChangeStreamPartitionsUserCallable.java create mode 100644 google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/changestream/ListChangeStreamPartitionsUserCallableTest.java diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java index ce9a57fa7ec6..424d6fe50f24 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java @@ -35,6 +35,7 @@ import com.google.cloud.bigtable.data.v2.models.Filters.Filter; import com.google.cloud.bigtable.data.v2.models.KeyOffset; import com.google.cloud.bigtable.data.v2.models.Query; +import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange; import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow; import com.google.cloud.bigtable.data.v2.models.Row; import com.google.cloud.bigtable.data.v2.models.RowAdapter; @@ -1489,6 +1490,141 @@ public UnaryCallable readModifyWriteRowCallable() { return stub.readModifyWriteRowCallable(); } + /** + * Convenience method for synchronously streaming the partitions of a table. The returned + * ServerStream instance is not threadsafe, it can only be used from single thread. + * + *

Sample code: + * + * 

{@code
+   * try (BigtableDataClient bigtableDataClient = BigtableDataClient.create("[PROJECT]", "[INSTANCE]")) {
+   *   String tableId = "[TABLE]";
+   *
+   *   try {
+   *     ServerStream stream = bigtableDataClient.listChangeStreamPartitions(tableId);
+   *     int count = 0;
+   *
+   *     // Iterator style
+   *     for (ByteStringRange partition : stream) {
+   *       if (++count > 10) {
+   *         stream.cancel();
+   *         break;
+   *       }
+   *       // Do something with partition
+   *     }
+   *   } catch (NotFoundException e) {
+   *     System.out.println("Tried to read a non-existent table");
+   *   } catch (RuntimeException e) {
+   *     e.printStackTrace();
+   *   }
+   * }
+   * }
+ * + * @see ServerStreamingCallable For call styles. + */ + public ServerStream listChangeStreamPartitions(String tableId) { + return listChangeStreamPartitionsCallable().call(tableId); + } + + /** + * Convenience method for asynchronously streaming the partitions of a table. + * + *

Sample code: + * + * 

{@code
+   * try (BigtableDataClient bigtableDataClient = BigtableDataClient.create("[PROJECT]", "[INSTANCE]")) {
+   *   String tableId = "[TABLE]";
+   *
+   *   bigtableDataClient.listChangeStreamPartitionsAsync(tableId, new ResponseObserver() {
+   *     StreamController controller;
+   *     int count = 0;
+   *
+   *     public void onStart(StreamController controller) {
+   *       this.controller = controller;
+   *     }
+   *     public void onResponse(ByteStringRange partition) {
+   *       if (++count > 10) {
+   *         controller.cancel();
+   *         return;
+   *       }
+   *       // Do something with partition
+   *     }
+   *     public void onError(Throwable t) {
+   *       if (t instanceof NotFoundException) {
+   *         System.out.println("Tried to read a non-existent table");
+   *       } else {
+   *         t.printStackTrace();
+   *       }
+   *     }
+   *     public void onComplete() {
+   *       // Handle stream completion
+   *     }
+   *   });
+   * }
+   * }
+ */ + public void listChangeStreamPartitionsAsync( + String tableId, ResponseObserver observer) { + listChangeStreamPartitionsCallable().call(tableId, observer); + } + + /** + * Streams back the results of the query. The returned callable object allows for customization of + * api invocation. + * + *

Sample code: + * + * 

{@code
+   * try (BigtableDataClient bigtableDataClient = BigtableDataClient.create("[PROJECT]", "[INSTANCE]")) {
+   *   String tableId = "[TABLE]";
+   *
+   *   // Iterator style
+   *   try {
+   *     for(ByteStringRange partition : bigtableDataClient.listChangeStreamPartitionsCallable().call(tableId)) {
+   *       // Do something with partition
+   *     }
+   *   } catch (NotFoundException e) {
+   *     System.out.println("Tried to read a non-existent table");
+   *   } catch (RuntimeException e) {
+   *     e.printStackTrace();
+   *   }
+   *
+   *   // Sync style
+   *   try {
+   *     List partitions = bigtableDataClient.listChangeStreamPartitionsCallable().all().call(tableId);
+   *   } catch (NotFoundException e) {
+   *     System.out.println("Tried to read a non-existent table");
+   *   } catch (RuntimeException e) {
+   *     e.printStackTrace();
+   *   }
+   *
+   *   // Point look up
+   *   ApiFuture partitionFuture =
+   *     bigtableDataClient.listChangeStreamPartitionsCallable().first().futureCall(tableId);
+   *
+   *   ApiFutures.addCallback(partitionFuture, new ApiFutureCallback() {
+   *     public void onFailure(Throwable t) {
+   *       if (t instanceof NotFoundException) {
+   *         System.out.println("Tried to read a non-existent table");
+   *       } else {
+   *         t.printStackTrace();
+   *       }
+   *     }
+   *     public void onSuccess(ByteStringRange result) {
+   *       System.out.println("Got partition: " + result);
+   *     }
+   *   }, MoreExecutors.directExecutor());
+   *
+   *   // etc
+   * }
+   * }
+ * + * @see ServerStreamingCallable For call styles. + */ + public ServerStreamingCallable listChangeStreamPartitionsCallable() { + return stub.listChangeStreamPartitionsCallable(); + } + /** Close the clients and releases all associated resources. */ @Override public void close() { 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 ec237aabf73a..4e21c7980cf8 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 @@ -47,6 +47,8 @@ import com.google.bigtable.v2.BigtableGrpc; import com.google.bigtable.v2.CheckAndMutateRowRequest; import com.google.bigtable.v2.CheckAndMutateRowResponse; +import com.google.bigtable.v2.ListChangeStreamPartitionsRequest; +import com.google.bigtable.v2.ListChangeStreamPartitionsResponse; import com.google.bigtable.v2.MutateRowRequest; import com.google.bigtable.v2.MutateRowResponse; import com.google.bigtable.v2.MutateRowsRequest; @@ -65,11 +67,13 @@ import com.google.cloud.bigtable.data.v2.models.DefaultRowAdapter; import com.google.cloud.bigtable.data.v2.models.KeyOffset; import com.google.cloud.bigtable.data.v2.models.Query; +import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange; import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow; import com.google.cloud.bigtable.data.v2.models.Row; import com.google.cloud.bigtable.data.v2.models.RowAdapter; import com.google.cloud.bigtable.data.v2.models.RowMutation; import com.google.cloud.bigtable.data.v2.models.RowMutationEntry; +import com.google.cloud.bigtable.data.v2.stub.changestream.ListChangeStreamPartitionsUserCallable; import com.google.cloud.bigtable.data.v2.stub.metrics.BigtableTracerStreamingCallable; import com.google.cloud.bigtable.data.v2.stub.metrics.BigtableTracerUnaryCallable; import com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsTracerFactory; @@ -142,6 +146,8 @@ public class EnhancedBigtableStub implements AutoCloseable { private final UnaryCallable checkAndMutateRowCallable; private final UnaryCallable readModifyWriteRowCallable; + private final ServerStreamingCallable listChangeStreamPartitionsCallable; + public static EnhancedBigtableStub create(EnhancedBigtableStubSettings settings) throws IOException { settings = finalizeSettings(settings, Tags.getTagger(), Stats.getStatsRecorder()); @@ -284,6 +290,7 @@ public EnhancedBigtableStub(EnhancedBigtableStubSettings settings, ClientContext bulkMutateRowsCallable = createBulkMutateRowsCallable(); checkAndMutateRowCallable = createCheckAndMutateRowCallable(); readModifyWriteRowCallable = createReadModifyWriteRowCallable(); + listChangeStreamPartitionsCallable = createListChangeStreamPartitionsCallable(); } // @@ -798,6 +805,71 @@ public Map extract(ReadModifyWriteRowRequest request) { methodName, new ReadModifyWriteRowCallable(retrying, requestContext)); } + /** + * Creates a callable chain to handle streaming ListChangeStreamPartitions RPCs. The chain will: + * + *
    + *
  • Convert a String format tableId into a {@link + * com.google.bigtable.v2.ReadChangeStreamRequest} and dispatch the RPC. + *
  • Upon receiving the response stream, it will convert the {@link + * com.google.bigtable.v2.ListChangeStreamPartitionsResponse}s into {@link ByteStringRange}. + *
+ */ + private ServerStreamingCallable + createListChangeStreamPartitionsCallable() { + ServerStreamingCallable + base = + GrpcRawCallableFactory.createServerStreamingCallable( + GrpcCallSettings + . + newBuilder() + .setMethodDescriptor(BigtableGrpc.getListChangeStreamPartitionsMethod()) + .setParamsExtractor( + new RequestParamsExtractor() { + @Override + public Map extract( + ListChangeStreamPartitionsRequest listChangeStreamPartitionsRequest) { + return ImmutableMap.of( + "table_name", + listChangeStreamPartitionsRequest.getTableName(), + "app_profile_id", + listChangeStreamPartitionsRequest.getAppProfileId()); + } + }) + .build(), + settings.listChangeStreamPartitionsSettings().getRetryableCodes()); + + ServerStreamingCallable userCallable = + new ListChangeStreamPartitionsUserCallable(base, requestContext); + + ServerStreamingCallable withStatsHeaders = + new StatsHeadersServerStreamingCallable<>(userCallable); + + // Copy settings for the middle String -> ByteStringRange callable (as opposed to the inner + // ListChangeStreamPartitionsRequest -> ListChangeStreamPartitionsResponse callable). + ServerStreamingCallSettings innerSettings = + ServerStreamingCallSettings.newBuilder() + .setRetryableCodes(settings.listChangeStreamPartitionsSettings().getRetryableCodes()) + .setRetrySettings(settings.listChangeStreamPartitionsSettings().getRetrySettings()) + .setIdleTimeout(settings.listChangeStreamPartitionsSettings().getIdleTimeout()) + .build(); + + ServerStreamingCallable watched = + Callables.watched(withStatsHeaders, innerSettings, clientContext); + + ServerStreamingCallable withBigtableTracer = + new BigtableTracerStreamingCallable<>(watched); + + ServerStreamingCallable retrying = + Callables.retrying(withBigtableTracer, innerSettings, clientContext); + + SpanName span = getSpanName("ListChangeStreamPartitions"); + ServerStreamingCallable traced = + new TracedServerStreamingCallable<>(retrying, clientContext.getTracerFactory(), span); + + return traced.withDefaultCallContext(clientContext.getDefaultCallContext()); + } + /** * Wraps a callable chain in a user presentable callable that will inject the default call context * and trace the call. @@ -854,6 +926,11 @@ public UnaryCallable checkAndMutateRowCallable( public UnaryCallable readModifyWriteRowCallable() { return readModifyWriteRowCallable; } + + /** Returns a streaming list change stream partitions callable */ + public ServerStreamingCallable listChangeStreamPartitionsCallable() { + return listChangeStreamPartitionsCallable; + } // private SpanName getSpanName(String methodName) { 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 395ba52b08d2..e46061dde04c 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 @@ -37,6 +37,7 @@ import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation; import com.google.cloud.bigtable.data.v2.models.KeyOffset; import com.google.cloud.bigtable.data.v2.models.Query; +import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange; import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow; import com.google.cloud.bigtable.data.v2.models.Row; import com.google.cloud.bigtable.data.v2.models.RowMutation; @@ -137,6 +138,22 @@ public class EnhancedBigtableStubSettings extends StubSettings LIST_CHANGE_STREAM_PARTITIONS_RETRY_CODES = + ImmutableSet.builder().addAll(IDEMPOTENT_RETRY_CODES).add(Code.ABORTED).build(); + + private static final RetrySettings LIST_CHANGE_STREAM_PARTITIONS_RETRY_SETTINGS = + RetrySettings.newBuilder() + .setInitialRetryDelay(Duration.ofMillis(10)) + .setRetryDelayMultiplier(2.0) + .setMaxRetryDelay(Duration.ofMinutes(1)) + .setMaxAttempts(10) + .setJittered(true) + .setInitialRpcTimeout(Duration.ofMinutes(5)) + .setRpcTimeoutMultiplier(2.0) + .setMaxRpcTimeout(Duration.ofMinutes(5)) + .setTotalTimeout(Duration.ofHours(12)) + .build(); + /** * Scopes that are equivalent to JWT's audience. * @@ -174,6 +191,9 @@ public class EnhancedBigtableStubSettings extends StubSettings checkAndMutateRowSettings; private final UnaryCallSettings readModifyWriteRowSettings; + private final ServerStreamingCallSettings + listChangeStreamPartitionsSettings; + private EnhancedBigtableStubSettings(Builder builder) { super(builder); @@ -208,6 +228,7 @@ private EnhancedBigtableStubSettings(Builder builder) { bulkReadRowsSettings = builder.bulkReadRowsSettings.build(); checkAndMutateRowSettings = builder.checkAndMutateRowSettings.build(); readModifyWriteRowSettings = builder.readModifyWriteRowSettings.build(); + listChangeStreamPartitionsSettings = builder.listChangeStreamPartitionsSettings.build(); } /** Create a new builder. */ @@ -491,6 +512,10 @@ public UnaryCallSettings readModifyWriteRowSettings() { return readModifyWriteRowSettings; } + public ServerStreamingCallSettings listChangeStreamPartitionsSettings() { + return listChangeStreamPartitionsSettings; + } + /** Returns a builder containing all the values of this settings class. */ public Builder toBuilder() { return new Builder(this); @@ -516,6 +541,9 @@ public static class Builder extends StubSettings.Builder readModifyWriteRowSettings; + private final ServerStreamingCallSettings.Builder + listChangeStreamPartitionsSettings; + /** * Initializes a new Builder with sane defaults for all settings. * @@ -626,6 +654,12 @@ private Builder() { readModifyWriteRowSettings = UnaryCallSettings.newUnaryCallSettingsBuilder(); copyRetrySettings(baseDefaults.readModifyWriteRowSettings(), readModifyWriteRowSettings); + + listChangeStreamPartitionsSettings = ServerStreamingCallSettings.newBuilder(); + listChangeStreamPartitionsSettings + .setRetryableCodes(LIST_CHANGE_STREAM_PARTITIONS_RETRY_CODES) + .setRetrySettings(LIST_CHANGE_STREAM_PARTITIONS_RETRY_SETTINGS) + .setIdleTimeout(Duration.ofMinutes(5)); } private Builder(EnhancedBigtableStubSettings settings) { @@ -646,6 +680,7 @@ private Builder(EnhancedBigtableStubSettings settings) { bulkReadRowsSettings = settings.bulkReadRowsSettings.toBuilder(); checkAndMutateRowSettings = settings.checkAndMutateRowSettings.toBuilder(); readModifyWriteRowSettings = settings.readModifyWriteRowSettings.toBuilder(); + listChangeStreamPartitionsSettings = settings.listChangeStreamPartitionsSettings.toBuilder(); } // @@ -857,6 +892,7 @@ public String toString() { .add("bulkReadRowsSettings", bulkReadRowsSettings) .add("checkAndMutateRowSettings", checkAndMutateRowSettings) .add("readModifyWriteRowSettings", readModifyWriteRowSettings) + .add("listChangeStreamPartitionsSettings", listChangeStreamPartitionsSettings) .add("parent", super.toString()) .toString(); } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/changestream/ListChangeStreamPartitionsUserCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/changestream/ListChangeStreamPartitionsUserCallable.java new file mode 100644 index 000000000000..b9c51b5635cc --- /dev/null +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/changestream/ListChangeStreamPartitionsUserCallable.java @@ -0,0 +1,92 @@ +/* + * 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.bigtable.data.v2.stub.changestream; + +import com.google.api.gax.rpc.ApiCallContext; +import com.google.api.gax.rpc.ResponseObserver; +import com.google.api.gax.rpc.ServerStreamingCallable; +import com.google.api.gax.rpc.StreamController; +import com.google.bigtable.v2.*; +import com.google.cloud.bigtable.data.v2.internal.NameUtil; +import com.google.cloud.bigtable.data.v2.internal.RequestContext; +import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange; + +/** Simple wrapper for ListChangeStreamPartitions to wrap the request and response protobufs. */ +public class ListChangeStreamPartitionsUserCallable + extends ServerStreamingCallable { + private final RequestContext requestContext; + private final ServerStreamingCallable< + ListChangeStreamPartitionsRequest, ListChangeStreamPartitionsResponse> + inner; + + public ListChangeStreamPartitionsUserCallable( + ServerStreamingCallable + inner, + RequestContext requestContext) { + this.requestContext = requestContext; + this.inner = inner; + } + + @Override + public void call( + String tableId, ResponseObserver responseObserver, ApiCallContext context) { + String tableName = + NameUtil.formatTableName( + requestContext.getProjectId(), requestContext.getInstanceId(), tableId); + ListChangeStreamPartitionsRequest request = + ListChangeStreamPartitionsRequest.newBuilder() + .setTableName(tableName) + .setAppProfileId(requestContext.getAppProfileId()) + .build(); + + inner.call(request, new ConvertPartitionToRangeObserver(responseObserver), context); + } + + private class ConvertPartitionToRangeObserver + implements ResponseObserver { + + private final ResponseObserver outerObserver; + + ConvertPartitionToRangeObserver(ResponseObserver observer) { + this.outerObserver = observer; + } + + @Override + public void onStart(final StreamController controller) { + outerObserver.onStart(controller); + } + + @Override + public void onResponse(ListChangeStreamPartitionsResponse response) { + ByteStringRange range = + ByteStringRange.unbounded() + .of( + response.getPartition().getRowRange().getStartKeyClosed(), + response.getPartition().getRowRange().getEndKeyOpen()); + outerObserver.onResponse(range); + } + + @Override + public void onError(Throwable t) { + outerObserver.onError(t); + } + + @Override + public void onComplete() { + outerObserver.onComplete(); + } + } +} diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientTests.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientTests.java index 34c9a29d71b7..603e7e8b6714 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientTests.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientTests.java @@ -30,6 +30,7 @@ import com.google.cloud.bigtable.data.v2.models.KeyOffset; import com.google.cloud.bigtable.data.v2.models.Mutation; import com.google.cloud.bigtable.data.v2.models.Query; +import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange; import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow; import com.google.cloud.bigtable.data.v2.models.Row; import com.google.cloud.bigtable.data.v2.models.RowCell; @@ -79,6 +80,9 @@ public class BigtableDataClientTests { @Mock private Batcher mockBulkMutationBatcher; @Mock private Batcher mockBulkReadRowsBatcher; + @Mock(answer = Answers.RETURNS_DEEP_STUBS) + private ServerStreamingCallable mockListChangeStreamPartitionsCallable; + private BigtableDataClient bigtableDataClient; @Before @@ -153,6 +157,14 @@ public void proxyReadRowsCallableTest() { assertThat(bigtableDataClient.readRowsCallable()).isSameInstanceAs(mockReadRowsCallable); } + @Test + public void proxyListChangeStreamPartitionsCallableTest() { + Mockito.when(mockStub.listChangeStreamPartitionsCallable()) + .thenReturn(mockListChangeStreamPartitionsCallable); + assertThat(bigtableDataClient.listChangeStreamPartitionsCallable()) + .isSameInstanceAs(mockListChangeStreamPartitionsCallable); + } + @Test public void proxyReadRowAsyncTest() { Mockito.when(mockStub.readRowCallable()).thenReturn(mockReadRowCallable); @@ -300,6 +312,28 @@ public void proxyReadRowsAsyncTest() { Mockito.verify(mockReadRowsCallable).call(query, mockObserver); } + @Test + public void proxyListChangeStreamPartitionsSyncTest() { + Mockito.when(mockStub.listChangeStreamPartitionsCallable()) + .thenReturn(mockListChangeStreamPartitionsCallable); + + bigtableDataClient.listChangeStreamPartitions("fake-table"); + + Mockito.verify(mockListChangeStreamPartitionsCallable).call("fake-table"); + } + + @Test + public void proxyListChangeStreamPartitionsAsyncTest() { + Mockito.when(mockStub.listChangeStreamPartitionsCallable()) + .thenReturn(mockListChangeStreamPartitionsCallable); + + @SuppressWarnings("unchecked") + ResponseObserver mockObserver = Mockito.mock(ResponseObserver.class); + bigtableDataClient.listChangeStreamPartitionsAsync("fake-table", mockObserver); + + Mockito.verify(mockListChangeStreamPartitionsCallable).call("fake-table", mockObserver); + } + @Test public void proxySampleRowKeysCallableTest() { Mockito.when(mockStub.sampleRowKeysCallable()).thenReturn(mockSampleRowKeysCallable); 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 c4e5ea2e4076..32ab93d1f2dd 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 @@ -701,6 +701,7 @@ public void isRefreshingChannelFalseValueTest() { "bulkReadRowsSettings", "checkAndMutateRowSettings", "readModifyWriteRowSettings", + "listChangeStreamPartitionsSettings", }; @Test diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/changestream/ListChangeStreamPartitionsUserCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/changestream/ListChangeStreamPartitionsUserCallableTest.java new file mode 100644 index 000000000000..34ec1e98aa3f --- /dev/null +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/changestream/ListChangeStreamPartitionsUserCallableTest.java @@ -0,0 +1,83 @@ +/* + * 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.bigtable.data.v2.stub.changestream; + +import static com.google.common.truth.Truth.assertThat; + +import com.google.bigtable.v2.ListChangeStreamPartitionsRequest; +import com.google.bigtable.v2.ListChangeStreamPartitionsResponse; +import com.google.bigtable.v2.RowRange; +import com.google.bigtable.v2.StreamPartition; +import com.google.cloud.bigtable.data.v2.internal.NameUtil; +import com.google.cloud.bigtable.data.v2.internal.RequestContext; +import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange; +import com.google.cloud.bigtable.gaxx.testing.FakeStreamingApi; +import com.google.common.collect.Lists; +import com.google.common.truth.Truth; +import com.google.protobuf.ByteString; +import java.util.List; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class ListChangeStreamPartitionsUserCallableTest { + private final RequestContext requestContext = + RequestContext.create("my-project", "my-instance", "my-profile"); + + @Test + public void requestIsCorrect() { + FakeStreamingApi.ServerStreamingStashCallable< + ListChangeStreamPartitionsRequest, ListChangeStreamPartitionsResponse> + inner = new FakeStreamingApi.ServerStreamingStashCallable<>(Lists.newArrayList()); + ListChangeStreamPartitionsUserCallable listChangeStreamPartitionsUserCallable = + new ListChangeStreamPartitionsUserCallable(inner, requestContext); + + listChangeStreamPartitionsUserCallable.all().call("my-table"); + assertThat(inner.getActualRequest()) + .isEqualTo( + ListChangeStreamPartitionsRequest.newBuilder() + .setTableName( + NameUtil.formatTableName( + requestContext.getProjectId(), requestContext.getInstanceId(), "my-table")) + .setAppProfileId(requestContext.getAppProfileId()) + .build()); + } + + @Test + public void responseIsConverted() { + FakeStreamingApi.ServerStreamingStashCallable< + ListChangeStreamPartitionsRequest, ListChangeStreamPartitionsResponse> + inner = + new FakeStreamingApi.ServerStreamingStashCallable<>( + Lists.newArrayList( + ListChangeStreamPartitionsResponse.newBuilder() + .setPartition( + StreamPartition.newBuilder() + .setRowRange( + RowRange.newBuilder() + .setStartKeyClosed(ByteString.copyFromUtf8("apple")) + .setEndKeyOpen(ByteString.copyFromUtf8("banana")) + .build()) + .build()) + .build())); + ListChangeStreamPartitionsUserCallable listChangeStreamPartitionsUserCallable = + new ListChangeStreamPartitionsUserCallable(inner, requestContext); + + List results = listChangeStreamPartitionsUserCallable.all().call("my-table"); + Truth.assertThat(results).containsExactly(ByteStringRange.create("apple", "banana")); + } +} From 603bec93e4b1936eee8c2640c1398e3d83c8c0e0 Mon Sep 17 00:00:00 2001 From: Teng Zhong Date: Tue, 19 Jul 2022 14:40:03 -0400 Subject: [PATCH 2/6] feat: Change return type of ListChangeStreamPartitions to RowRange --- .../bigtable/data/v2/BigtableDataClient.java | 27 ++++++------ .../data/v2/stub/EnhancedBigtableStub.java | 44 ++++++------------- .../v2/stub/EnhancedBigtableStubSettings.java | 15 +++---- ...istChangeStreamPartitionsUserCallable.java | 21 +++++---- .../data/v2/BigtableDataClientTests.java | 6 +-- ...hangeStreamPartitionsUserCallableTest.java | 10 +++-- 6 files changed, 54 insertions(+), 69 deletions(-) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java index 424d6fe50f24..0dc9d40afb5d 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java @@ -29,13 +29,13 @@ import com.google.api.gax.rpc.ServerStream; import com.google.api.gax.rpc.ServerStreamingCallable; import com.google.api.gax.rpc.UnaryCallable; +import com.google.bigtable.v2.RowRange; import com.google.cloud.bigtable.data.v2.models.BulkMutation; import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation; import com.google.cloud.bigtable.data.v2.models.Filters; import com.google.cloud.bigtable.data.v2.models.Filters.Filter; import com.google.cloud.bigtable.data.v2.models.KeyOffset; import com.google.cloud.bigtable.data.v2.models.Query; -import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange; import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow; import com.google.cloud.bigtable.data.v2.models.Row; import com.google.cloud.bigtable.data.v2.models.RowAdapter; @@ -1501,11 +1501,11 @@ public UnaryCallable readModifyWriteRowCallable() { * String tableId = "[TABLE]"; * * try { - * ServerStream stream = bigtableDataClient.listChangeStreamPartitions(tableId); + * ServerStream stream = bigtableDataClient.listChangeStreamPartitions(tableId); * int count = 0; * * // Iterator style - * for (ByteStringRange partition : stream) { + * for (RowRange partition : stream) { * if (++count > 10) { * stream.cancel(); * break; @@ -1522,7 +1522,7 @@ public UnaryCallable readModifyWriteRowCallable() { * * @see ServerStreamingCallable For call styles. */ - public ServerStream listChangeStreamPartitions(String tableId) { + public ServerStream listChangeStreamPartitions(String tableId) { return listChangeStreamPartitionsCallable().call(tableId); } @@ -1535,14 +1535,14 @@ public ServerStream listChangeStreamPartitions(String tableId) * try (BigtableDataClient bigtableDataClient = BigtableDataClient.create("[PROJECT]", "[INSTANCE]")) { * String tableId = "[TABLE]"; * - * bigtableDataClient.listChangeStreamPartitionsAsync(tableId, new ResponseObserver() { + * bigtableDataClient.listChangeStreamPartitionsAsync(tableId, new ResponseObserver() { * StreamController controller; * int count = 0; * * public void onStart(StreamController controller) { * this.controller = controller; * } - * public void onResponse(ByteStringRange partition) { + * public void onResponse(RowRange partition) { * if (++count > 10) { * controller.cancel(); * return; @@ -1563,8 +1563,7 @@ public ServerStream listChangeStreamPartitions(String tableId) * } * } */ - public void listChangeStreamPartitionsAsync( - String tableId, ResponseObserver observer) { + public void listChangeStreamPartitionsAsync(String tableId, ResponseObserver observer) { listChangeStreamPartitionsCallable().call(tableId, observer); } @@ -1580,7 +1579,7 @@ public void listChangeStreamPartitionsAsync( * * // Iterator style * try { - * for(ByteStringRange partition : bigtableDataClient.listChangeStreamPartitionsCallable().call(tableId)) { + * for(RowRange partition : bigtableDataClient.listChangeStreamPartitionsCallable().call(tableId)) { * // Do something with partition * } * } catch (NotFoundException e) { @@ -1591,7 +1590,7 @@ public void listChangeStreamPartitionsAsync( * * // Sync style * try { - * List partitions = bigtableDataClient.listChangeStreamPartitionsCallable().all().call(tableId); + * List partitions = bigtableDataClient.listChangeStreamPartitionsCallable().all().call(tableId); * } catch (NotFoundException e) { * System.out.println("Tried to read a non-existent table"); * } catch (RuntimeException e) { @@ -1599,10 +1598,10 @@ public void listChangeStreamPartitionsAsync( * } * * // Point look up - * ApiFuture partitionFuture = + * ApiFuture partitionFuture = * bigtableDataClient.listChangeStreamPartitionsCallable().first().futureCall(tableId); * - * ApiFutures.addCallback(partitionFuture, new ApiFutureCallback() { + * ApiFutures.addCallback(partitionFuture, new ApiFutureCallback() { * public void onFailure(Throwable t) { * if (t instanceof NotFoundException) { * System.out.println("Tried to read a non-existent table"); @@ -1610,7 +1609,7 @@ public void listChangeStreamPartitionsAsync( * t.printStackTrace(); * } * } - * public void onSuccess(ByteStringRange result) { + * public void onSuccess(RowRange result) { * System.out.println("Got partition: " + result); * } * }, MoreExecutors.directExecutor()); @@ -1621,7 +1620,7 @@ public void listChangeStreamPartitionsAsync( * * @see ServerStreamingCallable For call styles. */ - public ServerStreamingCallable listChangeStreamPartitionsCallable() { + public ServerStreamingCallable listChangeStreamPartitionsCallable() { return stub.listChangeStreamPartitionsCallable(); } 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 4e21c7980cf8..12469daef168 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 @@ -44,21 +44,7 @@ import com.google.api.gax.tracing.TracedUnaryCallable; import com.google.auth.Credentials; import com.google.auth.oauth2.ServiceAccountJwtAccessCredentials; -import com.google.bigtable.v2.BigtableGrpc; -import com.google.bigtable.v2.CheckAndMutateRowRequest; -import com.google.bigtable.v2.CheckAndMutateRowResponse; -import com.google.bigtable.v2.ListChangeStreamPartitionsRequest; -import com.google.bigtable.v2.ListChangeStreamPartitionsResponse; -import com.google.bigtable.v2.MutateRowRequest; -import com.google.bigtable.v2.MutateRowResponse; -import com.google.bigtable.v2.MutateRowsRequest; -import com.google.bigtable.v2.MutateRowsResponse; -import com.google.bigtable.v2.ReadModifyWriteRowRequest; -import com.google.bigtable.v2.ReadModifyWriteRowResponse; -import com.google.bigtable.v2.ReadRowsRequest; -import com.google.bigtable.v2.ReadRowsResponse; -import com.google.bigtable.v2.SampleRowKeysRequest; -import com.google.bigtable.v2.SampleRowKeysResponse; +import com.google.bigtable.v2.*; import com.google.cloud.bigtable.Version; import com.google.cloud.bigtable.data.v2.internal.JwtCredentialsWithAudience; import com.google.cloud.bigtable.data.v2.internal.RequestContext; @@ -67,7 +53,6 @@ import com.google.cloud.bigtable.data.v2.models.DefaultRowAdapter; import com.google.cloud.bigtable.data.v2.models.KeyOffset; import com.google.cloud.bigtable.data.v2.models.Query; -import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange; import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow; import com.google.cloud.bigtable.data.v2.models.Row; import com.google.cloud.bigtable.data.v2.models.RowAdapter; @@ -146,7 +131,7 @@ public class EnhancedBigtableStub implements AutoCloseable { private final UnaryCallable checkAndMutateRowCallable; private final UnaryCallable readModifyWriteRowCallable; - private final ServerStreamingCallable listChangeStreamPartitionsCallable; + private final ServerStreamingCallable listChangeStreamPartitionsCallable; public static EnhancedBigtableStub create(EnhancedBigtableStubSettings settings) throws IOException { @@ -810,13 +795,12 @@ public Map extract(ReadModifyWriteRowRequest request) { * *
    *
  • Convert a String format tableId into a {@link - * com.google.bigtable.v2.ReadChangeStreamRequest} and dispatch the RPC. + * com.google.bigtable.v2.ListChangeStreamPartitionsRequest} and dispatch the RPC. *
  • Upon receiving the response stream, it will convert the {@link - * com.google.bigtable.v2.ListChangeStreamPartitionsResponse}s into {@link ByteStringRange}. + * com.google.bigtable.v2.ListChangeStreamPartitionsResponse}s into {@link RowRange}. *
*/ - private ServerStreamingCallable - createListChangeStreamPartitionsCallable() { + private ServerStreamingCallable createListChangeStreamPartitionsCallable() { ServerStreamingCallable base = GrpcRawCallableFactory.createServerStreamingCallable( @@ -839,32 +823,32 @@ public Map extract( .build(), settings.listChangeStreamPartitionsSettings().getRetryableCodes()); - ServerStreamingCallable userCallable = + ServerStreamingCallable userCallable = new ListChangeStreamPartitionsUserCallable(base, requestContext); - ServerStreamingCallable withStatsHeaders = + ServerStreamingCallable withStatsHeaders = new StatsHeadersServerStreamingCallable<>(userCallable); // Copy settings for the middle String -> ByteStringRange callable (as opposed to the inner // ListChangeStreamPartitionsRequest -> ListChangeStreamPartitionsResponse callable). - ServerStreamingCallSettings innerSettings = - ServerStreamingCallSettings.newBuilder() + ServerStreamingCallSettings innerSettings = + ServerStreamingCallSettings.newBuilder() .setRetryableCodes(settings.listChangeStreamPartitionsSettings().getRetryableCodes()) .setRetrySettings(settings.listChangeStreamPartitionsSettings().getRetrySettings()) .setIdleTimeout(settings.listChangeStreamPartitionsSettings().getIdleTimeout()) .build(); - ServerStreamingCallable watched = + ServerStreamingCallable watched = Callables.watched(withStatsHeaders, innerSettings, clientContext); - ServerStreamingCallable withBigtableTracer = + ServerStreamingCallable withBigtableTracer = new BigtableTracerStreamingCallable<>(watched); - ServerStreamingCallable retrying = + ServerStreamingCallable retrying = Callables.retrying(withBigtableTracer, innerSettings, clientContext); SpanName span = getSpanName("ListChangeStreamPartitions"); - ServerStreamingCallable traced = + ServerStreamingCallable traced = new TracedServerStreamingCallable<>(retrying, clientContext.getTracerFactory(), span); return traced.withDefaultCallContext(clientContext.getDefaultCallContext()); @@ -928,7 +912,7 @@ public UnaryCallable readModifyWriteRowCallable() { } /** Returns a streaming list change stream partitions callable */ - public ServerStreamingCallable listChangeStreamPartitionsCallable() { + public ServerStreamingCallable listChangeStreamPartitionsCallable() { return listChangeStreamPartitionsCallable; } // 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 e46061dde04c..83f0445bc559 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 @@ -33,11 +33,11 @@ import com.google.api.gax.rpc.TransportChannelProvider; import com.google.api.gax.rpc.UnaryCallSettings; import com.google.auth.Credentials; +import com.google.bigtable.v2.RowRange; import com.google.cloud.bigtable.Version; import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation; import com.google.cloud.bigtable.data.v2.models.KeyOffset; import com.google.cloud.bigtable.data.v2.models.Query; -import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange; import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow; import com.google.cloud.bigtable.data.v2.models.Row; import com.google.cloud.bigtable.data.v2.models.RowMutation; @@ -148,10 +148,10 @@ public class EnhancedBigtableStubSettings extends StubSettings checkAndMutateRowSettings; private final UnaryCallSettings readModifyWriteRowSettings; - private final ServerStreamingCallSettings - listChangeStreamPartitionsSettings; + private final ServerStreamingCallSettings listChangeStreamPartitionsSettings; private EnhancedBigtableStubSettings(Builder builder) { super(builder); @@ -512,7 +511,7 @@ public UnaryCallSettings readModifyWriteRowSettings() { return readModifyWriteRowSettings; } - public ServerStreamingCallSettings listChangeStreamPartitionsSettings() { + public ServerStreamingCallSettings listChangeStreamPartitionsSettings() { return listChangeStreamPartitionsSettings; } @@ -541,7 +540,7 @@ public static class Builder extends StubSettings.Builder readModifyWriteRowSettings; - private final ServerStreamingCallSettings.Builder + private final ServerStreamingCallSettings.Builder listChangeStreamPartitionsSettings; /** diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/changestream/ListChangeStreamPartitionsUserCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/changestream/ListChangeStreamPartitionsUserCallable.java index b9c51b5635cc..559d1e8ecfae 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/changestream/ListChangeStreamPartitionsUserCallable.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/changestream/ListChangeStreamPartitionsUserCallable.java @@ -22,11 +22,10 @@ import com.google.bigtable.v2.*; import com.google.cloud.bigtable.data.v2.internal.NameUtil; import com.google.cloud.bigtable.data.v2.internal.RequestContext; -import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange; /** Simple wrapper for ListChangeStreamPartitions to wrap the request and response protobufs. */ public class ListChangeStreamPartitionsUserCallable - extends ServerStreamingCallable { + extends ServerStreamingCallable { private final RequestContext requestContext; private final ServerStreamingCallable< ListChangeStreamPartitionsRequest, ListChangeStreamPartitionsResponse> @@ -42,7 +41,7 @@ public ListChangeStreamPartitionsUserCallable( @Override public void call( - String tableId, ResponseObserver responseObserver, ApiCallContext context) { + String tableId, ResponseObserver responseObserver, ApiCallContext context) { String tableName = NameUtil.formatTableName( requestContext.getProjectId(), requestContext.getInstanceId(), tableId); @@ -58,9 +57,9 @@ public void call( private class ConvertPartitionToRangeObserver implements ResponseObserver { - private final ResponseObserver outerObserver; + private final ResponseObserver outerObserver; - ConvertPartitionToRangeObserver(ResponseObserver observer) { + ConvertPartitionToRangeObserver(ResponseObserver observer) { this.outerObserver = observer; } @@ -71,12 +70,12 @@ public void onStart(final StreamController controller) { @Override public void onResponse(ListChangeStreamPartitionsResponse response) { - ByteStringRange range = - ByteStringRange.unbounded() - .of( - response.getPartition().getRowRange().getStartKeyClosed(), - response.getPartition().getRowRange().getEndKeyOpen()); - outerObserver.onResponse(range); + RowRange rowRange = + RowRange.newBuilder() + .setStartKeyClosed(response.getPartition().getRowRange().getStartKeyClosed()) + .setEndKeyOpen(response.getPartition().getRowRange().getEndKeyOpen()) + .build(); + outerObserver.onResponse(rowRange); } @Override diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientTests.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientTests.java index 603e7e8b6714..fcbcc15e308e 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientTests.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientTests.java @@ -24,13 +24,13 @@ import com.google.api.gax.rpc.ResponseObserver; import com.google.api.gax.rpc.ServerStreamingCallable; import com.google.api.gax.rpc.UnaryCallable; +import com.google.bigtable.v2.RowRange; import com.google.cloud.bigtable.data.v2.models.BulkMutation; import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation; import com.google.cloud.bigtable.data.v2.models.Filters.Filter; import com.google.cloud.bigtable.data.v2.models.KeyOffset; import com.google.cloud.bigtable.data.v2.models.Mutation; import com.google.cloud.bigtable.data.v2.models.Query; -import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange; import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow; import com.google.cloud.bigtable.data.v2.models.Row; import com.google.cloud.bigtable.data.v2.models.RowCell; @@ -81,7 +81,7 @@ public class BigtableDataClientTests { @Mock private Batcher mockBulkReadRowsBatcher; @Mock(answer = Answers.RETURNS_DEEP_STUBS) - private ServerStreamingCallable mockListChangeStreamPartitionsCallable; + private ServerStreamingCallable mockListChangeStreamPartitionsCallable; private BigtableDataClient bigtableDataClient; @@ -328,7 +328,7 @@ public void proxyListChangeStreamPartitionsAsyncTest() { .thenReturn(mockListChangeStreamPartitionsCallable); @SuppressWarnings("unchecked") - ResponseObserver mockObserver = Mockito.mock(ResponseObserver.class); + ResponseObserver mockObserver = Mockito.mock(ResponseObserver.class); bigtableDataClient.listChangeStreamPartitionsAsync("fake-table", mockObserver); Mockito.verify(mockListChangeStreamPartitionsCallable).call("fake-table", mockObserver); diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/changestream/ListChangeStreamPartitionsUserCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/changestream/ListChangeStreamPartitionsUserCallableTest.java index 34ec1e98aa3f..03db35f8d6b0 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/changestream/ListChangeStreamPartitionsUserCallableTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/changestream/ListChangeStreamPartitionsUserCallableTest.java @@ -23,7 +23,6 @@ import com.google.bigtable.v2.StreamPartition; import com.google.cloud.bigtable.data.v2.internal.NameUtil; import com.google.cloud.bigtable.data.v2.internal.RequestContext; -import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange; import com.google.cloud.bigtable.gaxx.testing.FakeStreamingApi; import com.google.common.collect.Lists; import com.google.common.truth.Truth; @@ -77,7 +76,12 @@ public void responseIsConverted() { ListChangeStreamPartitionsUserCallable listChangeStreamPartitionsUserCallable = new ListChangeStreamPartitionsUserCallable(inner, requestContext); - List results = listChangeStreamPartitionsUserCallable.all().call("my-table"); - Truth.assertThat(results).containsExactly(ByteStringRange.create("apple", "banana")); + List results = listChangeStreamPartitionsUserCallable.all().call("my-table"); + Truth.assertThat(results) + .containsExactly( + RowRange.newBuilder() + .setStartKeyClosed(ByteString.copyFromUtf8("apple")) + .setEndKeyOpen(ByteString.copyFromUtf8("banana")) + .build()); } } From 941bedf4bb9c297ddc4ae3e0a94fa6d2229bcdc3 Mon Sep 17 00:00:00 2001 From: Teng Zhong Date: Tue, 19 Jul 2022 14:49:57 -0400 Subject: [PATCH 3/6] feat: Fix format for ListChangeStreamPartitions --- .../data/v2/stub/EnhancedBigtableStub.java | 17 ++++++++++++++++- .../ListChangeStreamPartitionsUserCallable.java | 4 +++- 2 files changed, 19 insertions(+), 2 deletions(-) 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 12469daef168..bc6ff0a05cb7 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 @@ -44,7 +44,22 @@ import com.google.api.gax.tracing.TracedUnaryCallable; import com.google.auth.Credentials; import com.google.auth.oauth2.ServiceAccountJwtAccessCredentials; -import com.google.bigtable.v2.*; +import com.google.bigtable.v2.BigtableGrpc; +import com.google.bigtable.v2.CheckAndMutateRowRequest; +import com.google.bigtable.v2.CheckAndMutateRowResponse; +import com.google.bigtable.v2.ListChangeStreamPartitionsRequest; +import com.google.bigtable.v2.ListChangeStreamPartitionsResponse; +import com.google.bigtable.v2.MutateRowRequest; +import com.google.bigtable.v2.MutateRowResponse; +import com.google.bigtable.v2.MutateRowsRequest; +import com.google.bigtable.v2.MutateRowsResponse; +import com.google.bigtable.v2.ReadModifyWriteRowRequest; +import com.google.bigtable.v2.ReadModifyWriteRowResponse; +import com.google.bigtable.v2.ReadRowsRequest; +import com.google.bigtable.v2.ReadRowsResponse; +import com.google.bigtable.v2.RowRange; +import com.google.bigtable.v2.SampleRowKeysRequest; +import com.google.bigtable.v2.SampleRowKeysResponse; import com.google.cloud.bigtable.Version; import com.google.cloud.bigtable.data.v2.internal.JwtCredentialsWithAudience; import com.google.cloud.bigtable.data.v2.internal.RequestContext; diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/changestream/ListChangeStreamPartitionsUserCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/changestream/ListChangeStreamPartitionsUserCallable.java index 559d1e8ecfae..1d3393bb2b9f 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/changestream/ListChangeStreamPartitionsUserCallable.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/changestream/ListChangeStreamPartitionsUserCallable.java @@ -19,7 +19,9 @@ import com.google.api.gax.rpc.ResponseObserver; import com.google.api.gax.rpc.ServerStreamingCallable; import com.google.api.gax.rpc.StreamController; -import com.google.bigtable.v2.*; +import com.google.bigtable.v2.ListChangeStreamPartitionsRequest; +import com.google.bigtable.v2.ListChangeStreamPartitionsResponse; +import com.google.bigtable.v2.RowRange; import com.google.cloud.bigtable.data.v2.internal.NameUtil; import com.google.cloud.bigtable.data.v2.internal.RequestContext; From e00577dfd584d60531129caac85893cc144a2bc5 Mon Sep 17 00:00:00 2001 From: Teng Zhong Date: Tue, 19 Jul 2022 21:01:14 -0400 Subject: [PATCH 4/6] fix: Address comments for ListChangeStreamPartitionsCallable --- .../bigtable/data/v2/BigtableDataClient.java | 3 + ...va => ConvertStreamExceptionCallable.java} | 27 +++---- .../data/v2/stub/EnhancedBigtableStub.java | 15 ++-- .../ConvertStreamExceptionCallableTest.java | 76 +++++++++++++++++++ 4 files changed, 103 insertions(+), 18 deletions(-) rename google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/{readrows/ReadRowsConvertExceptionCallable.java => ConvertStreamExceptionCallable.java} (69%) create mode 100644 google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/ConvertStreamExceptionCallableTest.java diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java index 0dc9d40afb5d..2acb53d4dcf2 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java @@ -1522,6 +1522,7 @@ public UnaryCallable readModifyWriteRowCallable() { * * @see ServerStreamingCallable For call styles. */ + @InternalApi public ServerStream listChangeStreamPartitions(String tableId) { return listChangeStreamPartitionsCallable().call(tableId); } @@ -1563,6 +1564,7 @@ public ServerStream listChangeStreamPartitions(String tableId) { * } * } */ + @InternalApi public void listChangeStreamPartitionsAsync(String tableId, ResponseObserver observer) { listChangeStreamPartitionsCallable().call(tableId, observer); } @@ -1620,6 +1622,7 @@ public void listChangeStreamPartitionsAsync(String tableId, ResponseObserver listChangeStreamPartitionsCallable() { return stub.listChangeStreamPartitionsCallable(); } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsConvertExceptionCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/ConvertStreamExceptionCallable.java similarity index 69% rename from google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsConvertExceptionCallable.java rename to google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/ConvertStreamExceptionCallable.java index 0c58f6644172..55a0d390fbb8 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsConvertExceptionCallable.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/ConvertStreamExceptionCallable.java @@ -13,7 +13,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package com.google.cloud.bigtable.data.v2.stub.readrows; +package com.google.cloud.bigtable.data.v2.stub; import com.google.api.core.InternalApi; import com.google.api.gax.rpc.ApiCallContext; @@ -27,29 +27,30 @@ * This callable converts the "Received rst stream" exception into a retryable {@link ApiException}. */ @InternalApi -public final class ReadRowsConvertExceptionCallable - extends ServerStreamingCallable { +public final class ConvertStreamExceptionCallable + extends ServerStreamingCallable { - private final ServerStreamingCallable innerCallable; + private final ServerStreamingCallable innerCallable; - public ReadRowsConvertExceptionCallable( - ServerStreamingCallable innerCallable) { + public ConvertStreamExceptionCallable( + ServerStreamingCallable innerCallable) { this.innerCallable = innerCallable; } @Override public void call( - ReadRowsRequest request, ResponseObserver responseObserver, ApiCallContext context) { - ReadRowsConvertExceptionResponseObserver observer = - new ReadRowsConvertExceptionResponseObserver<>(responseObserver); + RequestT request, ResponseObserver responseObserver, ApiCallContext context) { + ConvertStreamExceptionResponseObserver observer = + new ConvertStreamExceptionResponseObserver<>(responseObserver); innerCallable.call(request, observer, context); } - private class ReadRowsConvertExceptionResponseObserver implements ResponseObserver { + private class ConvertStreamExceptionResponseObserver + implements ResponseObserver { - private final ResponseObserver outerObserver; + private final ResponseObserver outerObserver; - ReadRowsConvertExceptionResponseObserver(ResponseObserver outerObserver) { + ConvertStreamExceptionResponseObserver(ResponseObserver outerObserver) { this.outerObserver = outerObserver; } @@ -59,7 +60,7 @@ public void onStart(StreamController controller) { } @Override - public void onResponse(RowT response) { + public void onResponse(ResponseT response) { outerObserver.onResponse(response); } 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 bc6ff0a05cb7..7d2cd85b6510 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 @@ -88,7 +88,6 @@ import com.google.cloud.bigtable.data.v2.stub.mutaterows.MutateRowsRetryingCallable; import com.google.cloud.bigtable.data.v2.stub.readrows.FilterMarkerRowsCallable; import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsBatchingDescriptor; -import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsConvertExceptionCallable; import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsResumptionStrategy; import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsRetryCompletedCallable; import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsUserCallable; @@ -417,7 +416,7 @@ public Map extract(ReadRowsRequest readRowsRequest) { // should be treated similar to UNAVAILABLE. However, this exception has an INTERNAL error code // which by default is not retryable. Convert the exception so it can be retried in the client. ServerStreamingCallable convertException = - new ReadRowsConvertExceptionCallable<>(withStatsHeaders); + new ConvertStreamExceptionCallable<>(withStatsHeaders); ServerStreamingCallable merging = new RowMergingCallable<>(convertException, rowAdapter); @@ -844,8 +843,14 @@ public Map extract( ServerStreamingCallable withStatsHeaders = new StatsHeadersServerStreamingCallable<>(userCallable); - // Copy settings for the middle String -> ByteStringRange callable (as opposed to the inner - // ListChangeStreamPartitionsRequest -> ListChangeStreamPartitionsResponse callable). + // Sometimes ListChangeStreamPartitions connections are disconnected via an RST frame. This + // error is transient and should be treated similar to UNAVAILABLE. However, this exception + // has an INTERNAL error code which by default is not retryable. Convert the exception so it + // can be retried in the client. + ServerStreamingCallable convertException = + new ConvertStreamExceptionCallable<>(withStatsHeaders); + + // Copy idle timeout settings for watchdog. ServerStreamingCallSettings innerSettings = ServerStreamingCallSettings.newBuilder() .setRetryableCodes(settings.listChangeStreamPartitionsSettings().getRetryableCodes()) @@ -854,7 +859,7 @@ public Map extract( .build(); ServerStreamingCallable watched = - Callables.watched(withStatsHeaders, innerSettings, clientContext); + Callables.watched(convertException, innerSettings, clientContext); ServerStreamingCallable withBigtableTracer = new BigtableTracerStreamingCallable<>(watched); diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/ConvertStreamExceptionCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/ConvertStreamExceptionCallableTest.java new file mode 100644 index 000000000000..8f08e15b235b --- /dev/null +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/ConvertStreamExceptionCallableTest.java @@ -0,0 +1,76 @@ +/* + * 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.bigtable.data.v2.stub; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +import com.google.api.gax.grpc.GrpcStatusCode; +import com.google.api.gax.rpc.ApiCallContext; +import com.google.api.gax.rpc.ApiException; +import com.google.api.gax.rpc.InternalException; +import com.google.api.gax.rpc.ResponseObserver; +import com.google.api.gax.rpc.ServerStreamingCallable; +import io.grpc.Status; +import io.grpc.StatusRuntimeException; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class ConvertStreamExceptionCallableTest { + + @Test + public void rstStreamExceptionConvertedToRetryableTest() { + ApiException originalException = + new InternalException( + new StatusRuntimeException( + Status.INTERNAL.withDescription( + "INTERNAL: HTTP/2 error code: INTERNAL_ERROR\nReceived Rst Stream")), + GrpcStatusCode.of(Status.Code.INTERNAL), + false); + assertFalse(originalException.isRetryable()); + SettableExceptionCallable settableExceptionCallable = + new SettableExceptionCallable<>(originalException); + ConvertStreamExceptionCallable convertStreamExceptionCallable = + new ConvertStreamExceptionCallable<>(settableExceptionCallable); + + Throwable actualError = null; + try { + convertStreamExceptionCallable.all().call("fake-request"); + } catch (Throwable t) { + actualError = t; + } + assert actualError instanceof InternalException; + InternalException actualException = (InternalException) actualError; + assertTrue(actualException.isRetryable()); + } + + private static final class SettableExceptionCallable + extends ServerStreamingCallable { + private final Throwable throwable; + + public SettableExceptionCallable(Throwable throwable) { + this.throwable = throwable; + } + + @Override + public void call( + RequestT request, ResponseObserver responseObserver, ApiCallContext context) { + responseObserver.onError(throwable); + } + } +} From 4a12339c4cb862808ab229b9a750cdbb24785cac Mon Sep 17 00:00:00 2001 From: Teng Zhong Date: Wed, 20 Jul 2022 12:10:46 -0400 Subject: [PATCH 5/6] feat: Add comments for IntervalApi for ListChangeStreamPartitions --- .../google/cloud/bigtable/data/v2/BigtableDataClient.java | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java index 2acb53d4dcf2..38bc4dc811a0 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java @@ -1522,7 +1522,7 @@ public UnaryCallable readModifyWriteRowCallable() { * * @see ServerStreamingCallable For call styles. */ - @InternalApi + @InternalApi("Used in Changestream beam pipeline.") public ServerStream listChangeStreamPartitions(String tableId) { return listChangeStreamPartitionsCallable().call(tableId); } @@ -1564,7 +1564,7 @@ public ServerStream listChangeStreamPartitions(String tableId) { * } * } */ - @InternalApi + @InternalApi("Used in Changestream beam pipeline.") public void listChangeStreamPartitionsAsync(String tableId, ResponseObserver observer) { listChangeStreamPartitionsCallable().call(tableId, observer); } @@ -1622,7 +1622,7 @@ public void listChangeStreamPartitionsAsync(String tableId, ResponseObserver listChangeStreamPartitionsCallable() { return stub.listChangeStreamPartitionsCallable(); } From 54d4770df66e8977aba9286a215993740dd1f907 Mon Sep 17 00:00:00 2001 From: Teng Zhong Date: Wed, 20 Jul 2022 14:32:34 -0400 Subject: [PATCH 6/6] feat: Ignore renaming of ReadRowsConvertExceptionCallable --- google-cloud-bigtable/clirr-ignored-differences.xml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/google-cloud-bigtable/clirr-ignored-differences.xml b/google-cloud-bigtable/clirr-ignored-differences.xml index 588327d0de77..3fa8f3ee1e58 100644 --- a/google-cloud-bigtable/clirr-ignored-differences.xml +++ b/google-cloud-bigtable/clirr-ignored-differences.xml @@ -39,6 +39,11 @@ 8001 com/google/cloud/bigtable/data/v2/stub/metrics/CompositeTracerFactory + + + 8001 + com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsConvertExceptionCallable + 8001