From 3fa44fd605824b89898708c1ec65d2d69475cae9 Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Mon, 4 Nov 2024 15:04:19 -0500 Subject: [PATCH] chore: simplify sample rows key callable chain (#2396) Previously there were 3 chain creators: 1. createSampleRowKeysBaseCallable 2. createSampleRowKeysWithRequestCallable 3. createSampleRowKeysCallable The primary reason for this is that SampleRowKeysWithRequest was introduced after createSampleRowKeysCallable because it supports authorized views. This pr simplifies the logic by moving everything into createSampleRowKeysWithRequestCallable and makes createSampleRowKeysCallable be a tiny shim to convert a String tableId into a SampleRowKeysRequest --- .../data/v2/stub/EnhancedBigtableStub.java | 89 ++++++------- .../data/v2/stub/SampleRowKeysCallable.java | 79 ----------- .../v2/stub/SampleRowKeysCallableTest.java | 125 ------------------ 3 files changed, 40 insertions(+), 253 deletions(-) delete mode 100644 google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/SampleRowKeysCallable.java delete mode 100644 google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/SampleRowKeysCallableTest.java 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 91c63c2b85..3a9344268f 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 @@ -21,6 +21,7 @@ import static com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsConstants.INSTANCE_ID_KEY; import com.google.api.core.ApiFunction; +import com.google.api.core.ApiFuture; import com.google.api.core.BetaApi; import com.google.api.core.InternalApi; import com.google.api.gax.batching.Batcher; @@ -39,6 +40,7 @@ import com.google.api.gax.retrying.RetryAlgorithm; import com.google.api.gax.retrying.RetryingExecutorWithContext; import com.google.api.gax.retrying.ScheduledRetryingExecutor; +import com.google.api.gax.rpc.ApiCallContext; import com.google.api.gax.rpc.Callables; import com.google.api.gax.rpc.ClientContext; import com.google.api.gax.rpc.RequestParamsExtractor; @@ -98,6 +100,7 @@ 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.models.SampleRowKeysRequest; +import com.google.cloud.bigtable.data.v2.models.TableId; import com.google.cloud.bigtable.data.v2.models.TargetId; import com.google.cloud.bigtable.data.v2.models.sql.Statement; import com.google.cloud.bigtable.data.v2.stub.changestream.ChangeStreamRecordMergingCallable; @@ -194,7 +197,7 @@ public class EnhancedBigtableStub implements AutoCloseable { private final ServerStreamingCallable readRowsCallable; private final UnaryCallable readRowCallable; private final UnaryCallable> bulkReadRowsCallable; - private final UnaryCallable> sampleRowKeysCallable; + @Deprecated private final UnaryCallable> sampleRowKeysCallable; private final UnaryCallable> sampleRowKeysCallableWithRequest; private final UnaryCallable mutateRowCallable; @@ -698,11 +701,40 @@ private UnaryCallable> createBulkReadRowsCallable( } /** - * Helper function that should only be used by createSampleRowKeysCallable() and - * createSampleRowKeysWithRequestCallable(). + * Simple wrapper around {@link #createSampleRowKeysCallableWithRequest()} to provide backwards + * compatibility + * + * @deprecated + */ + @Deprecated + private UnaryCallable> createSampleRowKeysCallable() { + UnaryCallable> baseCallable = + createSampleRowKeysCallableWithRequest(); + return new UnaryCallable>() { + @Override + public ApiFuture> futureCall(String s, ApiCallContext apiCallContext) { + return baseCallable.futureCall(SampleRowKeysRequest.create(TableId.of(s)), apiCallContext); + } + }; + } + + /** + * Creates a callable chain to handle SampleRowKeys RPcs. The chain will: + * + *
    + *
  • Convert a {@link SampleRowKeysRequest} to a {@link + * com.google.bigtable.v2.SampleRowKeysRequest}. + *
  • Dispatch the request to the GAPIC's {@link BigtableStub#sampleRowKeysCallable()}. + *
  • Spool responses into a list. + *
  • Retry on failure. + *
  • Convert the responses into {@link KeyOffset}s. + *
  • Add tracing & metrics. + *
*/ - private UnaryCallable> - createSampleRowKeysBaseCallable() { + private UnaryCallable> + createSampleRowKeysCallableWithRequest() { + String methodName = "SampleRowKeys"; + ServerStreamingCallable base = GrpcRawCallableFactory.createServerStreamingCallable( @@ -745,51 +777,8 @@ public Map extract( UnaryCallable> retryable = withRetries(withBigtableTracer, settings.sampleRowKeysSettings()); - return retryable; - } - - /** - * Creates a callable chain to handle SampleRowKeys RPcs. The chain will: - * - *
    - *
  • Convert a table id to a {@link com.google.bigtable.v2.SampleRowKeysRequest}. - *
  • Dispatch the request to the GAPIC's {@link BigtableStub#sampleRowKeysCallable()}. - *
  • Spool responses into a list. - *
  • Retry on failure. - *
  • Convert the responses into {@link KeyOffset}s. - *
  • Add tracing & metrics. - *
- */ - private UnaryCallable> createSampleRowKeysCallable() { - String methodName = "SampleRowKeys"; - - UnaryCallable> - baseCallable = createSampleRowKeysBaseCallable(); - return createUserFacingUnaryCallable( - methodName, new SampleRowKeysCallable(baseCallable, requestContext)); - } - - /** - * Creates a callable chain to handle SampleRowKeys RPcs. The chain will: - * - *
    - *
  • Convert a {@link SampleRowKeysRequest} to a {@link - * com.google.bigtable.v2.SampleRowKeysRequest}. - *
  • Dispatch the request to the GAPIC's {@link BigtableStub#sampleRowKeysCallable()}. - *
  • Spool responses into a list. - *
  • Retry on failure. - *
  • Convert the responses into {@link KeyOffset}s. - *
  • Add tracing & metrics. - *
- */ - private UnaryCallable> - createSampleRowKeysCallableWithRequest() { - String methodName = "SampleRowKeys"; - - UnaryCallable> - baseCallable = createSampleRowKeysBaseCallable(); return createUserFacingUnaryCallable( - methodName, new SampleRowKeysCallableWithRequest(baseCallable, requestContext)); + methodName, new SampleRowKeysCallableWithRequest(retryable, requestContext)); } /** @@ -1470,6 +1459,8 @@ public UnaryCallable readRowCallable() { return readRowCallable; } + /** Deprecated, please use {@link #sampleRowKeysCallableWithRequest} */ + @Deprecated public UnaryCallable> sampleRowKeysCallable() { return sampleRowKeysCallable; } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/SampleRowKeysCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/SampleRowKeysCallable.java deleted file mode 100644 index 7658e41492..0000000000 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/SampleRowKeysCallable.java +++ /dev/null @@ -1,79 +0,0 @@ -/* - * 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 com.google.api.core.ApiFunction; -import com.google.api.core.ApiFuture; -import com.google.api.core.ApiFutures; -import com.google.api.gax.rpc.ApiCallContext; -import com.google.api.gax.rpc.UnaryCallable; -import com.google.bigtable.v2.SampleRowKeysRequest; -import com.google.bigtable.v2.SampleRowKeysResponse; -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.KeyOffset; -import com.google.common.collect.ImmutableList; -import com.google.common.util.concurrent.MoreExecutors; -import java.util.List; - -/** Simple wrapper for SampleRowKeys to wrap the request and response protobufs. */ -class SampleRowKeysCallable extends UnaryCallable> { - private final RequestContext requestContext; - private final UnaryCallable> inner; - - SampleRowKeysCallable( - UnaryCallable> inner, - RequestContext requestContext) { - - this.requestContext = requestContext; - this.inner = inner; - } - - @Override - public ApiFuture> futureCall(String tableId, ApiCallContext context) { - String tableName = - NameUtil.formatTableName( - requestContext.getProjectId(), requestContext.getInstanceId(), tableId); - - SampleRowKeysRequest request = - SampleRowKeysRequest.newBuilder() - .setTableName(tableName) - .setAppProfileId(requestContext.getAppProfileId()) - .build(); - - ApiFuture> rawResponse = inner.futureCall(request, context); - - return ApiFutures.transform( - rawResponse, - new ApiFunction, List>() { - @Override - public List apply(List rawResponse) { - return convert(rawResponse); - } - }, - MoreExecutors.directExecutor()); - } - - private static List convert(List rawResponse) { - ImmutableList.Builder results = ImmutableList.builder(); - - for (SampleRowKeysResponse element : rawResponse) { - results.add(KeyOffset.create(element.getRowKey(), element.getOffsetBytes())); - } - - return results.build(); - } -} diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/SampleRowKeysCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/SampleRowKeysCallableTest.java deleted file mode 100644 index 40a30d3263..0000000000 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/SampleRowKeysCallableTest.java +++ /dev/null @@ -1,125 +0,0 @@ -/* - * 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 com.google.common.truth.Truth.assertThat; - -import com.google.api.core.ApiFuture; -import com.google.api.core.SettableApiFuture; -import com.google.api.gax.grpc.GrpcStatusCode; -import com.google.api.gax.rpc.ApiCallContext; -import com.google.api.gax.rpc.NotFoundException; -import com.google.api.gax.rpc.UnaryCallable; -import com.google.bigtable.v2.SampleRowKeysRequest; -import com.google.bigtable.v2.SampleRowKeysResponse; -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.KeyOffset; -import com.google.common.collect.ImmutableList; -import com.google.protobuf.ByteString; -import io.grpc.Status.Code; -import java.util.List; -import java.util.concurrent.ExecutionException; -import java.util.concurrent.TimeUnit; -import org.junit.Before; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; - -@RunWith(JUnit4.class) -public class SampleRowKeysCallableTest { - - private final RequestContext requestContext = - RequestContext.create("my-project", "my-instance", "my-profile"); - private FakeCallable inner; - private SampleRowKeysCallable callable; - - @Before - public void setUp() { - inner = new FakeCallable(); - callable = new SampleRowKeysCallable(inner, requestContext); - } - - @Test - public void requestIsCorrect() { - callable.futureCall("my-table"); - - assertThat(inner.request) - .isEqualTo( - SampleRowKeysRequest.newBuilder() - .setTableName( - NameUtil.formatTableName( - requestContext.getProjectId(), requestContext.getInstanceId(), "my-table")) - .setAppProfileId(requestContext.getAppProfileId()) - .build()); - } - - @Test - public void responseCorrectlyTransformed() throws Exception { - ApiFuture> result = callable.futureCall("my-table"); - - inner.response.set( - ImmutableList.of( - SampleRowKeysResponse.newBuilder() - .setRowKey(ByteString.copyFromUtf8("key1")) - .setOffsetBytes(100) - .build(), - SampleRowKeysResponse.newBuilder() - .setRowKey(ByteString.copyFromUtf8("")) - .setOffsetBytes(1000) - .build())); - - assertThat(result.get(1, TimeUnit.SECONDS)) - .isEqualTo( - ImmutableList.of( - KeyOffset.create(ByteString.copyFromUtf8("key1"), 100), - KeyOffset.create(ByteString.EMPTY, 1000))); - } - - @Test - public void errorIsPropagated() throws Exception { - ApiFuture> result = callable.futureCall("my-table"); - - Throwable expectedError = - new NotFoundException("fake error", null, GrpcStatusCode.of(Code.NOT_FOUND), false); - inner.response.setException(expectedError); - - Throwable actualError = null; - try { - result.get(1, TimeUnit.SECONDS); - } catch (ExecutionException e) { - actualError = e.getCause(); - } - - assertThat(actualError).isEqualTo(expectedError); - } - - static class FakeCallable - extends UnaryCallable> { - SampleRowKeysRequest request; - ApiCallContext callContext; - SettableApiFuture> response = SettableApiFuture.create(); - - @Override - public ApiFuture> futureCall( - SampleRowKeysRequest request, ApiCallContext context) { - this.request = request; - this.callContext = context; - - return response; - } - } -}