Skip to content

Commit

Permalink
Rename SampleRowKeys to SampleRowKeysRequest
Browse files Browse the repository at this point in the history 
Change-Id: I8dda7ee1df31b184d04938cbc2f9f984d84138b4
  • Loading branch information
Lixia Chen committed Mar 27, 2024
1 parent 01de021 commit e5562af
Show file tree
Hide file tree
Showing 9 changed files with 162 additions and 134 deletions.
14 changes: 7 additions & 7 deletions ...le-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java
View file
Expand Up @@ -44,7 +44,7 @@
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.models.SampleRowKeys;
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.stub.EnhancedBigtableStub;
Expand Down Expand Up @@ -1432,7 +1432,7 @@ public ApiFuture<List<KeyOffset>> sampleRowKeysAsync(String tableId) {
* @see TableId
*/
public ApiFuture<List<KeyOffset>> sampleRowKeysAsync(TargetId targetId) {
return sampleRowKeysCallableWithRequest().futureCall(SampleRowKeys.create(targetId));
return sampleRowKeysCallableWithRequest().futureCall(SampleRowKeysRequest.create(targetId));
}

/**
Expand Down Expand Up @@ -1480,9 +1480,9 @@ public UnaryCallable<String, List<KeyOffset>> sampleRowKeysCallable() {
}

/**
* Returns a sample of row keys in the table. This callable allows takes a {@link SampleRowKeys}
* object rather than a String input, and thus can be used to sample against a specified {@link
* TargetId}.
* Returns a sample of row keys in the table. This callable allows takes a {@link
* SampleRowKeysRequest} object rather than a String input, and thus can be used to sample against
* a specified {@link TargetId}.
*
* <p>The returned row keys will delimit contiguous sections of the table of approximately equal
* size, which can be used to break up the data for distributed tasks like mapreduces. The
Expand All @@ -1492,7 +1492,7 @@ public UnaryCallable<String, List<KeyOffset>> sampleRowKeysCallable() {
*
* <pre>{@code
* try (BigtableDataClient bigtableDataClient = BigtableDataClient.create("[PROJECT]", "[INSTANCE]")) {
* SampleRowKeys sampleRowKeys = SampleRowKeys.create(TableId.of("[TABLE]"));
* SampleRowKeysRequest sampleRowKeys = SampleRowKeysRequest.create(TableId.of("[TABLE]"));
* // Synchronous invocation
* try {
* List<KeyOffset> keyOffsets = bigtableDataClient.sampleRowKeysCallableWithRequest().call(sampleRowKeys);
Expand Down Expand Up @@ -1523,7 +1523,7 @@ public UnaryCallable<String, List<KeyOffset>> sampleRowKeysCallable() {
* @see com.google.cloud.bigtable.data.v2.models.AuthorizedViewId
* @see TableId
*/
public UnaryCallable<SampleRowKeys, List<KeyOffset>> sampleRowKeysCallableWithRequest() {
public UnaryCallable<SampleRowKeysRequest, List<KeyOffset>> sampleRowKeysCallableWithRequest() {
return stub.sampleRowKeysCallableWithRequest();
}

Expand Down
29 changes: 15 additions & 14 deletions ...igtable/data/v2/models/SampleRowKeys.java → .../data/v2/models/SampleRowKeysRequest.java
View file
Expand Up @@ -17,7 +17,6 @@
package com.google.cloud.bigtable.data.v2.models;

import com.google.api.core.InternalApi;
import com.google.bigtable.v2.SampleRowKeysRequest;
import com.google.cloud.bigtable.data.v2.internal.NameUtil;
import com.google.cloud.bigtable.data.v2.internal.RequestContext;
import com.google.common.base.Objects;
Expand All @@ -26,22 +25,23 @@
import javax.annotation.Nonnull;

/** Wraps a {@link com.google.bigtable.v2.SampleRowKeysRequest}. */
public final class SampleRowKeys implements Serializable {
public final class SampleRowKeysRequest implements Serializable {
private final TargetId targetId;

private SampleRowKeys(TargetId targetId) {
private SampleRowKeysRequest(TargetId targetId) {
Preconditions.checkNotNull(targetId, "target id can't be null.");
this.targetId = targetId;
}

/** Creates a new instance of the sample row keys builder for the given target with targetId */
public static SampleRowKeys create(TargetId targetId) {
return new SampleRowKeys(targetId);
public static SampleRowKeysRequest create(TargetId targetId) {
return new SampleRowKeysRequest(targetId);
}

@InternalApi
public SampleRowKeysRequest toProto(RequestContext requestContext) {
SampleRowKeysRequest.Builder builder = SampleRowKeysRequest.newBuilder();
public com.google.bigtable.v2.SampleRowKeysRequest toProto(RequestContext requestContext) {
com.google.bigtable.v2.SampleRowKeysRequest.Builder builder =
com.google.bigtable.v2.SampleRowKeysRequest.newBuilder();
String resourceName =
targetId.toResourceName(requestContext.getProjectId(), requestContext.getInstanceId());
if (targetId.scopedForAuthorizedView()) {
Expand All @@ -53,19 +53,20 @@ public SampleRowKeysRequest toProto(RequestContext requestContext) {
}

/**
* Wraps the protobuf {@link SampleRowKeysRequest}.
* Wraps the protobuf {@link com.google.bigtable.v2.SampleRowKeysRequest}.
*
* <p>WARNING: Please note that the project id & instance id in the table/authorized view name
* will be overwritten by the configuration in the BigtableDataClient.
*/
public static SampleRowKeys fromProto(@Nonnull SampleRowKeysRequest request) {
public static SampleRowKeysRequest fromProto(
@Nonnull com.google.bigtable.v2.SampleRowKeysRequest request) {
String tableName = request.getTableName();
String authorizedViewName = request.getAuthorizedViewName();

SampleRowKeys sampleRowKeys =
SampleRowKeys.create(NameUtil.extractTargetId(tableName, authorizedViewName));
SampleRowKeysRequest sampleRowKeysRequest =
SampleRowKeysRequest.create(NameUtil.extractTargetId(tableName, authorizedViewName));

return sampleRowKeys;
return sampleRowKeysRequest;
}

@Override
Expand All @@ -76,7 +77,7 @@ public boolean equals(Object o) {
if (o == null || getClass() != o.getClass()) {
return false;
}
SampleRowKeys sampleRowKeys = (SampleRowKeys) o;
return Objects.equal(targetId, sampleRowKeys.targetId);
SampleRowKeysRequest sampleRowKeysRequest = (SampleRowKeysRequest) o;
return Objects.equal(targetId, sampleRowKeysRequest.targetId);
}
}
92 changes: 50 additions & 42 deletions ...d-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
View file
Expand Up @@ -66,7 +66,6 @@
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;
Expand All @@ -88,7 +87,7 @@
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.models.SampleRowKeys;
import com.google.cloud.bigtable.data.v2.models.SampleRowKeysRequest;
import com.google.cloud.bigtable.data.v2.models.TargetId;
import com.google.cloud.bigtable.data.v2.stub.changestream.ChangeStreamRecordMergingCallable;
import com.google.cloud.bigtable.data.v2.stub.changestream.GenerateInitialChangeStreamPartitionsUserCallable;
Expand Down Expand Up @@ -167,7 +166,8 @@ public class EnhancedBigtableStub implements AutoCloseable {
private final UnaryCallable<Query, Row> readRowCallable;
private final UnaryCallable<Query, List<Row>> bulkReadRowsCallable;
private final UnaryCallable<String, List<KeyOffset>> sampleRowKeysCallable;
private final UnaryCallable<SampleRowKeys, List<KeyOffset>> sampleRowKeysCallableWithRequest;
private final UnaryCallable<SampleRowKeysRequest, List<KeyOffset>>
sampleRowKeysCallableWithRequest;
private final UnaryCallable<RowMutation, Void> mutateRowCallable;
private final UnaryCallable<BulkMutation, Void> bulkMutateRowsCallable;
private final UnaryCallable<ConditionalRowMutation, Boolean> checkAndMutateRowCallable;
Expand Down Expand Up @@ -595,43 +595,49 @@ private <RowT> UnaryCallable<Query, List<RowT>> createBulkReadRowsCallable(
* Helper function that should only be used by createSampleRowKeysCallable() and
* createSampleRowKeysWithRequestCallable().
*/
private UnaryCallable<SampleRowKeysRequest, List<SampleRowKeysResponse>>
private UnaryCallable<com.google.bigtable.v2.SampleRowKeysRequest, List<SampleRowKeysResponse>>
createSampleRowKeysBaseCallable() {
ServerStreamingCallable<SampleRowKeysRequest, SampleRowKeysResponse> base =
GrpcRawCallableFactory.createServerStreamingCallable(
GrpcCallSettings.<SampleRowKeysRequest, SampleRowKeysResponse>newBuilder()
.setMethodDescriptor(BigtableGrpc.getSampleRowKeysMethod())
.setParamsExtractor(
new RequestParamsExtractor<SampleRowKeysRequest>() {
@Override
public Map<String, String> extract(
SampleRowKeysRequest sampleRowKeysRequest) {
String tableName = sampleRowKeysRequest.getTableName();
String authorizedViewName = sampleRowKeysRequest.getAuthorizedViewName();
if (tableName.isEmpty()) {
tableName =
NameUtil.extractTableNameFromAuthorizedViewName(authorizedViewName);
}
return ImmutableMap.of(
"table_name",
tableName,
"app_profile_id",
sampleRowKeysRequest.getAppProfileId());
}
})
.build(),
settings.sampleRowKeysSettings().getRetryableCodes());
ServerStreamingCallable<com.google.bigtable.v2.SampleRowKeysRequest, SampleRowKeysResponse>
base =
GrpcRawCallableFactory.createServerStreamingCallable(
GrpcCallSettings
.<com.google.bigtable.v2.SampleRowKeysRequest, SampleRowKeysResponse>
newBuilder()
.setMethodDescriptor(BigtableGrpc.getSampleRowKeysMethod())
.setParamsExtractor(
new RequestParamsExtractor<com.google.bigtable.v2.SampleRowKeysRequest>() {
@Override
public Map<String, String> extract(
com.google.bigtable.v2.SampleRowKeysRequest sampleRowKeysRequest) {
String tableName = sampleRowKeysRequest.getTableName();
String authorizedViewName =
sampleRowKeysRequest.getAuthorizedViewName();
if (tableName.isEmpty()) {
tableName =
NameUtil.extractTableNameFromAuthorizedViewName(
authorizedViewName);
}
return ImmutableMap.of(
"table_name",
tableName,
"app_profile_id",
sampleRowKeysRequest.getAppProfileId());
}
})
.build(),
settings.sampleRowKeysSettings().getRetryableCodes());

UnaryCallable<SampleRowKeysRequest, List<SampleRowKeysResponse>> spoolable = base.all();
UnaryCallable<com.google.bigtable.v2.SampleRowKeysRequest, List<SampleRowKeysResponse>>
spoolable = base.all();

UnaryCallable<SampleRowKeysRequest, List<SampleRowKeysResponse>> withStatsHeaders =
new StatsHeadersUnaryCallable<>(spoolable);
UnaryCallable<com.google.bigtable.v2.SampleRowKeysRequest, List<SampleRowKeysResponse>>
withStatsHeaders = new StatsHeadersUnaryCallable<>(spoolable);

UnaryCallable<SampleRowKeysRequest, List<SampleRowKeysResponse>> withBigtableTracer =
new BigtableTracerUnaryCallable<>(withStatsHeaders);
UnaryCallable<com.google.bigtable.v2.SampleRowKeysRequest, List<SampleRowKeysResponse>>
withBigtableTracer = new BigtableTracerUnaryCallable<>(withStatsHeaders);

UnaryCallable<SampleRowKeysRequest, List<SampleRowKeysResponse>> retryable =
withRetries(withBigtableTracer, settings.sampleRowKeysSettings());
UnaryCallable<com.google.bigtable.v2.SampleRowKeysRequest, List<SampleRowKeysResponse>>
retryable = withRetries(withBigtableTracer, settings.sampleRowKeysSettings());

return retryable;
}
Expand All @@ -651,8 +657,8 @@ public Map<String, String> extract(
private UnaryCallable<String, List<KeyOffset>> createSampleRowKeysCallable() {
String methodName = "SampleRowKeys";

UnaryCallable<SampleRowKeysRequest, List<SampleRowKeysResponse>> baseCallable =
createSampleRowKeysBaseCallable();
UnaryCallable<com.google.bigtable.v2.SampleRowKeysRequest, List<SampleRowKeysResponse>>
baseCallable = createSampleRowKeysBaseCallable();
return createUserFacingUnaryCallable(
methodName, new SampleRowKeysCallable(baseCallable, requestContext));
}
Expand All @@ -661,19 +667,21 @@ private UnaryCallable<String, List<KeyOffset>> createSampleRowKeysCallable() {
* Creates a callable chain to handle SampleRowKeys RPcs. The chain will:
*
* <ul>
* <li>Convert a {@link SampleRowKeys} to a {@link com.google.bigtable.v2.SampleRowKeysRequest}.
* <li>Convert a {@link SampleRowKeysRequest} to a {@link
* com.google.bigtable.v2.SampleRowKeysRequest}.
* <li>Dispatch the request to the GAPIC's {@link BigtableStub#sampleRowKeysCallable()}.
* <li>Spool responses into a list.
* <li>Retry on failure.
* <li>Convert the responses into {@link KeyOffset}s.
* <li>Add tracing & metrics.
* </ul>
*/
private UnaryCallable<SampleRowKeys, List<KeyOffset>> createSampleRowKeysCallableWithRequest() {
private UnaryCallable<SampleRowKeysRequest, List<KeyOffset>>
createSampleRowKeysCallableWithRequest() {
String methodName = "SampleRowKeys";

UnaryCallable<SampleRowKeysRequest, List<SampleRowKeysResponse>> baseCallable =
createSampleRowKeysBaseCallable();
UnaryCallable<com.google.bigtable.v2.SampleRowKeysRequest, List<SampleRowKeysResponse>>
baseCallable = createSampleRowKeysBaseCallable();
return createUserFacingUnaryCallable(
methodName, new SampleRowKeysCallableWithRequest(baseCallable, requestContext));
}
Expand Down Expand Up @@ -1281,7 +1289,7 @@ public UnaryCallable<String, List<KeyOffset>> sampleRowKeysCallable() {
return sampleRowKeysCallable;
}

public UnaryCallable<SampleRowKeys, List<KeyOffset>> sampleRowKeysCallableWithRequest() {
public UnaryCallable<SampleRowKeysRequest, List<KeyOffset>> sampleRowKeysCallableWithRequest() {
return sampleRowKeysCallableWithRequest;
}

Expand Down
15 changes: 9 additions & 6 deletions ...rc/main/java/com/google/cloud/bigtable/data/v2/stub/SampleRowKeysCallableWithRequest.java
View file
Expand Up @@ -21,30 +21,33 @@
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.RequestContext;
import com.google.cloud.bigtable.data.v2.models.KeyOffset;
import com.google.cloud.bigtable.data.v2.models.SampleRowKeys;
import com.google.cloud.bigtable.data.v2.models.SampleRowKeysRequest;
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 SampleRowKeysCallableWithRequest extends UnaryCallable<SampleRowKeys, List<KeyOffset>> {
class SampleRowKeysCallableWithRequest
extends UnaryCallable<SampleRowKeysRequest, List<KeyOffset>> {
private final RequestContext requestContext;
private final UnaryCallable<SampleRowKeysRequest, List<SampleRowKeysResponse>> inner;
private final UnaryCallable<
com.google.bigtable.v2.SampleRowKeysRequest, List<SampleRowKeysResponse>>
inner;

SampleRowKeysCallableWithRequest(
UnaryCallable<SampleRowKeysRequest, List<SampleRowKeysResponse>> inner,
UnaryCallable<com.google.bigtable.v2.SampleRowKeysRequest, List<SampleRowKeysResponse>> inner,
RequestContext requestContext) {

this.requestContext = requestContext;
this.inner = inner;
}

@Override
public ApiFuture<List<KeyOffset>> futureCall(SampleRowKeys request, ApiCallContext context) {
public ApiFuture<List<KeyOffset>> futureCall(
SampleRowKeysRequest request, ApiCallContext context) {
ApiFuture<List<SampleRowKeysResponse>> rawResponse =
inner.futureCall(request.toProto(requestContext), context);

Expand Down
19 changes: 11 additions & 8 deletions ...oud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientTests.java
View file
Expand Up @@ -39,7 +39,7 @@
import com.google.cloud.bigtable.data.v2.models.RowCell;
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.SampleRowKeys;
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.stub.EnhancedBigtableStub;
Expand Down Expand Up @@ -79,7 +79,10 @@ public class BigtableDataClientTests {

@Mock private UnaryCallable<Query, Row> mockReadRowCallable;
@Mock private UnaryCallable<String, List<KeyOffset>> mockSampleRowKeysCallable;
@Mock private UnaryCallable<SampleRowKeys, List<KeyOffset>> mockSampleRowKeysCallableWithRequest;

@Mock
private UnaryCallable<SampleRowKeysRequest, List<KeyOffset>> mockSampleRowKeysCallableWithRequest;

@Mock private UnaryCallable<RowMutation, Void> mockMutateRowCallable;
@Mock private UnaryCallable<ConditionalRowMutation, Boolean> mockCheckAndMutateRowCallable;
@Mock private UnaryCallable<ReadModifyWriteRow, Row> mockReadModifyWriteRowCallable;
Expand Down Expand Up @@ -637,7 +640,7 @@ public void proxySampleRowKeysTest() {

bigtableDataClient.sampleRowKeysAsync("fake-table");
Mockito.verify(mockSampleRowKeysCallableWithRequest)
.futureCall(SampleRowKeys.create(TableId.of("fake-table")));
.futureCall(SampleRowKeysRequest.create(TableId.of("fake-table")));
}

@Test
Expand All @@ -649,7 +652,7 @@ public void proxySampleRowKeysOnAuthorizedViewTest() {
AuthorizedViewId.of("fake-table", "fake-authorized-view"));
Mockito.verify(mockSampleRowKeysCallableWithRequest)
.futureCall(
SampleRowKeys.create(AuthorizedViewId.of("fake-table", "fake-authorized-view")));
SampleRowKeysRequest.create(AuthorizedViewId.of("fake-table", "fake-authorized-view")));
}

@Test
Expand All @@ -659,11 +662,11 @@ public void sampleRowKeysTest() {

Mockito.when(
mockSampleRowKeysCallableWithRequest.futureCall(
ArgumentMatchers.any(SampleRowKeys.class)))
ArgumentMatchers.any(SampleRowKeysRequest.class)))
.thenReturn(ApiFutures.immediateFuture(Collections.<KeyOffset>emptyList()));
bigtableDataClient.sampleRowKeys("fake-table");
Mockito.verify(mockSampleRowKeysCallableWithRequest)
.futureCall(SampleRowKeys.create(TableId.of("fake-table")));
.futureCall(SampleRowKeysRequest.create(TableId.of("fake-table")));
}

@Test
Expand All @@ -673,12 +676,12 @@ public void sampleRowKeysOnAuthorizedViewTest() {

Mockito.when(
mockSampleRowKeysCallableWithRequest.futureCall(
ArgumentMatchers.any(SampleRowKeys.class)))
ArgumentMatchers.any(SampleRowKeysRequest.class)))
.thenReturn(ApiFutures.immediateFuture(Collections.<KeyOffset>emptyList()));
bigtableDataClient.sampleRowKeys(AuthorizedViewId.of("fake-table", "fake-authorized-view"));
Mockito.verify(mockSampleRowKeysCallableWithRequest)
.futureCall(
SampleRowKeys.create(AuthorizedViewId.of("fake-table", "fake-authorized-view")));
SampleRowKeysRequest.create(AuthorizedViewId.of("fake-table", "fake-authorized-view")));
}

@Test
Expand Down

0 comments on commit e5562af

Please sign in to comment.