Skip to content

Commit b40828c

Browse files
chore: refactor unary callable chains to avoid code duplication (#2394)
This is primarily a code clean up change that extracts all of the unary chain building logic into a helper method and removes a significant amount of code duplication.
1 parent 3fa44fd commit b40828c

File tree

8 files changed

+234
-681
lines changed

8 files changed

+234
-681
lines changed

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/CheckAndMutateRowCallable.java

Lines changed: 0 additions & 56 deletions
This file was deleted.

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java

Lines changed: 80 additions & 112 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222

2323
import com.google.api.core.ApiFunction;
2424
import com.google.api.core.ApiFuture;
25+
import com.google.api.core.ApiFutures;
2526
import com.google.api.core.BetaApi;
2627
import com.google.api.core.InternalApi;
2728
import com.google.api.gax.batching.Batcher;
@@ -57,22 +58,17 @@
5758
import com.google.auth.Credentials;
5859
import com.google.auth.oauth2.ServiceAccountJwtAccessCredentials;
5960
import com.google.bigtable.v2.BigtableGrpc;
60-
import com.google.bigtable.v2.CheckAndMutateRowRequest;
6161
import com.google.bigtable.v2.CheckAndMutateRowResponse;
6262
import com.google.bigtable.v2.ExecuteQueryRequest;
6363
import com.google.bigtable.v2.ExecuteQueryResponse;
6464
import com.google.bigtable.v2.GenerateInitialChangeStreamPartitionsRequest;
6565
import com.google.bigtable.v2.GenerateInitialChangeStreamPartitionsResponse;
66-
import com.google.bigtable.v2.MutateRowRequest;
67-
import com.google.bigtable.v2.MutateRowResponse;
6866
import com.google.bigtable.v2.MutateRowsRequest;
6967
import com.google.bigtable.v2.MutateRowsResponse;
7068
import com.google.bigtable.v2.PingAndWarmRequest;
7169
import com.google.bigtable.v2.PingAndWarmResponse;
7270
import com.google.bigtable.v2.ReadChangeStreamRequest;
7371
import com.google.bigtable.v2.ReadChangeStreamResponse;
74-
import com.google.bigtable.v2.ReadModifyWriteRowRequest;
75-
import com.google.bigtable.v2.ReadModifyWriteRowResponse;
7672
import com.google.bigtable.v2.ReadRowsRequest;
7773
import com.google.bigtable.v2.ReadRowsResponse;
7874
import com.google.bigtable.v2.RowRange;
@@ -144,8 +140,10 @@
144140
import com.google.common.base.Preconditions;
145141
import com.google.common.collect.ImmutableList;
146142
import com.google.common.collect.ImmutableMap;
143+
import com.google.common.util.concurrent.MoreExecutors;
147144
import com.google.protobuf.ByteString;
148145
import io.grpc.ManagedChannelBuilder;
146+
import io.grpc.MethodDescriptor;
149147
import io.opencensus.stats.Stats;
150148
import io.opencensus.stats.StatsRecorder;
151149
import io.opencensus.tags.TagKey;
@@ -162,6 +160,7 @@
162160
import java.util.Map;
163161
import java.util.Set;
164162
import java.util.concurrent.TimeUnit;
163+
import java.util.function.Function;
165164
import java.util.logging.Level;
166165
import java.util.logging.Logger;
167166
import javax.annotation.Nonnull;
@@ -790,42 +789,14 @@ public Map<String, String> extract(
790789
* </ul>
791790
*/
792791
private UnaryCallable<RowMutation, Void> createMutateRowCallable() {
793-
String methodName = "MutateRow";
794-
UnaryCallable<MutateRowRequest, MutateRowResponse> base =
795-
GrpcRawCallableFactory.createUnaryCallable(
796-
GrpcCallSettings.<MutateRowRequest, MutateRowResponse>newBuilder()
797-
.setMethodDescriptor(BigtableGrpc.getMutateRowMethod())
798-
.setParamsExtractor(
799-
new RequestParamsExtractor<MutateRowRequest>() {
800-
@Override
801-
public Map<String, String> extract(MutateRowRequest mutateRowRequest) {
802-
String tableName = mutateRowRequest.getTableName();
803-
String authorizedViewName = mutateRowRequest.getAuthorizedViewName();
804-
if (tableName.isEmpty()) {
805-
tableName =
806-
NameUtil.extractTableNameFromAuthorizedViewName(authorizedViewName);
807-
}
808-
return ImmutableMap.of(
809-
"table_name",
810-
tableName,
811-
"app_profile_id",
812-
mutateRowRequest.getAppProfileId());
813-
}
814-
})
815-
.build(),
816-
settings.mutateRowSettings().getRetryableCodes());
817-
818-
UnaryCallable<MutateRowRequest, MutateRowResponse> withStatsHeaders =
819-
new StatsHeadersUnaryCallable<>(base);
820-
821-
UnaryCallable<MutateRowRequest, MutateRowResponse> withBigtableTracer =
822-
new BigtableTracerUnaryCallable<>(withStatsHeaders);
823-
824-
UnaryCallable<MutateRowRequest, MutateRowResponse> retrying =
825-
withRetries(withBigtableTracer, settings.mutateRowSettings());
826-
827-
return createUserFacingUnaryCallable(
828-
methodName, new MutateRowCallable(retrying, requestContext));
792+
return createUnaryCallable(
793+
BigtableGrpc.getMutateRowMethod(),
794+
req ->
795+
composeRequestParams(
796+
req.getAppProfileId(), req.getTableName(), req.getAuthorizedViewName()),
797+
settings.mutateRowSettings(),
798+
req -> req.toProto(requestContext),
799+
resp -> null);
829800
}
830801

831802
/**
@@ -1045,44 +1016,14 @@ public Batcher<ByteString, Row> newBulkReadRowsBatcher(
10451016
* </ul>
10461017
*/
10471018
private UnaryCallable<ConditionalRowMutation, Boolean> createCheckAndMutateRowCallable() {
1048-
String methodName = "CheckAndMutateRow";
1049-
UnaryCallable<CheckAndMutateRowRequest, CheckAndMutateRowResponse> base =
1050-
GrpcRawCallableFactory.createUnaryCallable(
1051-
GrpcCallSettings.<CheckAndMutateRowRequest, CheckAndMutateRowResponse>newBuilder()
1052-
.setMethodDescriptor(BigtableGrpc.getCheckAndMutateRowMethod())
1053-
.setParamsExtractor(
1054-
new RequestParamsExtractor<CheckAndMutateRowRequest>() {
1055-
@Override
1056-
public Map<String, String> extract(
1057-
CheckAndMutateRowRequest checkAndMutateRowRequest) {
1058-
String tableName = checkAndMutateRowRequest.getTableName();
1059-
String authorizedViewName =
1060-
checkAndMutateRowRequest.getAuthorizedViewName();
1061-
if (tableName.isEmpty()) {
1062-
tableName =
1063-
NameUtil.extractTableNameFromAuthorizedViewName(authorizedViewName);
1064-
}
1065-
return ImmutableMap.of(
1066-
"table_name",
1067-
tableName,
1068-
"app_profile_id",
1069-
checkAndMutateRowRequest.getAppProfileId());
1070-
}
1071-
})
1072-
.build(),
1073-
settings.checkAndMutateRowSettings().getRetryableCodes());
1074-
1075-
UnaryCallable<CheckAndMutateRowRequest, CheckAndMutateRowResponse> withStatsHeaders =
1076-
new StatsHeadersUnaryCallable<>(base);
1077-
1078-
UnaryCallable<CheckAndMutateRowRequest, CheckAndMutateRowResponse> withBigtableTracer =
1079-
new BigtableTracerUnaryCallable<>(withStatsHeaders);
1080-
1081-
UnaryCallable<CheckAndMutateRowRequest, CheckAndMutateRowResponse> retrying =
1082-
withRetries(withBigtableTracer, settings.checkAndMutateRowSettings());
1083-
1084-
return createUserFacingUnaryCallable(
1085-
methodName, new CheckAndMutateRowCallable(retrying, requestContext));
1019+
return createUnaryCallable(
1020+
BigtableGrpc.getCheckAndMutateRowMethod(),
1021+
req ->
1022+
composeRequestParams(
1023+
req.getAppProfileId(), req.getTableName(), req.getAuthorizedViewName()),
1024+
settings.checkAndMutateRowSettings(),
1025+
req -> req.toProto(requestContext),
1026+
CheckAndMutateRowResponse::getPredicateMatched);
10861027
}
10871028

10881029
/**
@@ -1096,39 +1037,16 @@ public Map<String, String> extract(
10961037
* </ul>
10971038
*/
10981039
private UnaryCallable<ReadModifyWriteRow, Row> createReadModifyWriteRowCallable() {
1099-
UnaryCallable<ReadModifyWriteRowRequest, ReadModifyWriteRowResponse> base =
1100-
GrpcRawCallableFactory.createUnaryCallable(
1101-
GrpcCallSettings.<ReadModifyWriteRowRequest, ReadModifyWriteRowResponse>newBuilder()
1102-
.setMethodDescriptor(BigtableGrpc.getReadModifyWriteRowMethod())
1103-
.setParamsExtractor(
1104-
new RequestParamsExtractor<ReadModifyWriteRowRequest>() {
1105-
@Override
1106-
public Map<String, String> extract(ReadModifyWriteRowRequest request) {
1107-
String tableName = request.getTableName();
1108-
String authorizedViewName = request.getAuthorizedViewName();
1109-
if (tableName.isEmpty()) {
1110-
tableName =
1111-
NameUtil.extractTableNameFromAuthorizedViewName(authorizedViewName);
1112-
}
1113-
return ImmutableMap.of(
1114-
"table_name", tableName, "app_profile_id", request.getAppProfileId());
1115-
}
1116-
})
1117-
.build(),
1118-
settings.readModifyWriteRowSettings().getRetryableCodes());
1119-
1120-
UnaryCallable<ReadModifyWriteRowRequest, ReadModifyWriteRowResponse> withStatsHeaders =
1121-
new StatsHeadersUnaryCallable<>(base);
1122-
1123-
String methodName = "ReadModifyWriteRow";
1124-
UnaryCallable<ReadModifyWriteRowRequest, ReadModifyWriteRowResponse> withBigtableTracer =
1125-
new BigtableTracerUnaryCallable<>(withStatsHeaders);
1126-
1127-
UnaryCallable<ReadModifyWriteRowRequest, ReadModifyWriteRowResponse> retrying =
1128-
withRetries(withBigtableTracer, settings.readModifyWriteRowSettings());
1129-
1130-
return createUserFacingUnaryCallable(
1131-
methodName, new ReadModifyWriteRowCallable(retrying, requestContext));
1040+
DefaultRowAdapter rowAdapter = new DefaultRowAdapter();
1041+
1042+
return createUnaryCallable(
1043+
BigtableGrpc.getReadModifyWriteRowMethod(),
1044+
req ->
1045+
composeRequestParams(
1046+
req.getAppProfileId(), req.getTableName(), req.getAuthorizedViewName()),
1047+
settings.readModifyWriteRowSettings(),
1048+
req -> req.toProto(requestContext),
1049+
resp -> rowAdapter.createRowFromProto(resp.getRow()));
11321050
}
11331051

11341052
/**
@@ -1393,6 +1311,56 @@ private <RequestT, ResponseT> UnaryCallable<RequestT, ResponseT> createUserFacin
13931311
return traced.withDefaultCallContext(clientContext.getDefaultCallContext());
13941312
}
13951313

1314+
private Map<String, String> composeRequestParams(
1315+
String appProfileId, String tableName, String authorizedViewName) {
1316+
if (tableName.isEmpty()) {
1317+
tableName = NameUtil.extractTableNameFromAuthorizedViewName(authorizedViewName);
1318+
}
1319+
return ImmutableMap.of("table_name", tableName, "app_profile_id", appProfileId);
1320+
}
1321+
1322+
private <BaseReqT, BaseRespT, ReqT, RespT> UnaryCallable<ReqT, RespT> createUnaryCallable(
1323+
MethodDescriptor<BaseReqT, BaseRespT> methodDescriptor,
1324+
RequestParamsExtractor<BaseReqT> headerParamsFn,
1325+
UnaryCallSettings<ReqT, RespT> callSettings,
1326+
Function<ReqT, BaseReqT> requestTransformer,
1327+
Function<BaseRespT, RespT> responseTranformer) {
1328+
1329+
UnaryCallable<BaseReqT, BaseRespT> base =
1330+
GrpcRawCallableFactory.createUnaryCallable(
1331+
GrpcCallSettings.<BaseReqT, BaseRespT>newBuilder()
1332+
.setMethodDescriptor(methodDescriptor)
1333+
.setParamsExtractor(headerParamsFn)
1334+
.build(),
1335+
callSettings.getRetryableCodes());
1336+
1337+
UnaryCallable<BaseReqT, BaseRespT> withStatsHeaders = new StatsHeadersUnaryCallable<>(base);
1338+
1339+
UnaryCallable<BaseReqT, BaseRespT> withBigtableTracer =
1340+
new BigtableTracerUnaryCallable<>(withStatsHeaders);
1341+
1342+
UnaryCallable<BaseReqT, BaseRespT> retrying = withRetries(withBigtableTracer, callSettings);
1343+
1344+
UnaryCallable<ReqT, RespT> transformed =
1345+
new UnaryCallable<ReqT, RespT>() {
1346+
@Override
1347+
public ApiFuture<RespT> futureCall(ReqT reqT, ApiCallContext apiCallContext) {
1348+
ApiFuture<BaseRespT> f =
1349+
retrying.futureCall(requestTransformer.apply(reqT), apiCallContext);
1350+
return ApiFutures.transform(
1351+
f, responseTranformer::apply, MoreExecutors.directExecutor());
1352+
}
1353+
};
1354+
1355+
UnaryCallable<ReqT, RespT> traced =
1356+
new TracedUnaryCallable<>(
1357+
transformed,
1358+
clientContext.getTracerFactory(),
1359+
getSpanName(methodDescriptor.getBareMethodName()));
1360+
1361+
return traced.withDefaultCallContext(clientContext.getDefaultCallContext());
1362+
}
1363+
13961364
private UnaryCallable<PingAndWarmRequest, PingAndWarmResponse> createPingAndWarmCallable() {
13971365
UnaryCallable<PingAndWarmRequest, PingAndWarmResponse> pingAndWarm =
13981366
GrpcRawCallableFactory.createUnaryCallable(

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/MutateRowCallable.java

Lines changed: 0 additions & 56 deletions
This file was deleted.

0 commit comments

Comments
 (0)