Skip to content

Commit a403703

Browse files
mutianfmutianf
authored
feat(bigtable): route single-entry MutateRows through a point-write c… (#14028)
…allable Add MaybePointWriteCallable, mirroring MaybePointReadCallable: a BulkMutation with exactly one entry is converted to a RowMutation and dispatched through a point-write callable so it can benefit from the session-shim diversion, while multi-entry bulk mutations continue through the classic MutateRows path. The point-write callable falls back to the MutateRow RPC when the session diversion does not apply, but carries the caller's bulkMutateRowsSettings retry settings/codes so a single-entry bulk write retries the same way it would have as a MutateRows call.
1 parent 2a27c2c commit a403703

3 files changed

Lines changed: 243 additions & 1 deletion

File tree

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

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,7 @@
9494
import com.google.cloud.bigtable.data.v2.stub.metrics.StatsHeadersServerStreamingCallable;
9595
import com.google.cloud.bigtable.data.v2.stub.metrics.StatsHeadersUnaryCallable;
9696
import com.google.cloud.bigtable.data.v2.stub.mutaterows.BulkMutateRowsUserFacingCallable;
97+
import com.google.cloud.bigtable.data.v2.stub.mutaterows.MaybePointWriteCallable;
9798
import com.google.cloud.bigtable.data.v2.stub.mutaterows.MutateRowsAttemptResult;
9899
import com.google.cloud.bigtable.data.v2.stub.mutaterows.MutateRowsBatchingDescriptor;
99100
import com.google.cloud.bigtable.data.v2.stub.mutaterows.MutateRowsPartialErrorRetryAlgorithm;
@@ -200,8 +201,13 @@ public EnhancedBigtableStub(
200201
sampleRowKeysCallableWithRequest = createSampleRowKeysCallableWithRequest();
201202
mutateRowCallable = createMutateRowCallable();
202203
bulkMutateRowsCallable = createMutateRowsBaseCallable();
203-
externalBulkMutateRowsCallable =
204+
UnaryCallable<BulkMutation, Void> bulkMutateRowsVoidCallable =
204205
new MutateRowsErrorConverterUnaryCallable(bulkMutateRowsCallable);
206+
externalBulkMutateRowsCallable =
207+
new MaybePointWriteCallable(
208+
bulkMutateRowsVoidCallable,
209+
createPointWriteCallable(bulkMutateRowsVoidCallable),
210+
requestContext);
205211
checkAndMutateRowCallable = createCheckAndMutateRowCallable();
206212
readModifyWriteRowCallable = createReadModifyWriteRowCallable();
207213
generateInitialChangeStreamPartitionsCallable =
@@ -670,6 +676,40 @@ private UnaryCallable<RowMutation, Void> createMutateRowCallable() {
670676
.decorateMutateRow(classic, perOpSettings.mutateRowSettings);
671677
}
672678

679+
/**
680+
* Creates the point-write callable used by {@link MaybePointWriteCallable} to divert single-entry
681+
* {@link BulkMutation}s. This mirrors {@link #createPointReadCallable}: it exposes a single row
682+
* mutation through the session-shim diversion while preserving the bulk operation's retry
683+
* behavior.
684+
*
685+
* <p>Unlike point reads (where the single-row read is just {@code ReadRows} with a limit), {@code
686+
* MutateRow} and {@code MutateRows} are distinct RPCs. To preserve the existing wire behavior,
687+
* the fallback classic here delegates to the bulk {@code MutateRows} callable as a single-entry
688+
* batch rather than issuing a {@code MutateRow} RPC. So when the session shim does not divert
689+
* (e.g. a {@link com.google.cloud.bigtable.data.v2.internal.compat.DisabledShim}), a single-entry
690+
* bulk mutation still travels over {@code MutateRows} with the bulk operation's retry behavior;
691+
* only when the shim actively diverts does the mutation go to the session single-row write API.
692+
*/
693+
private UnaryCallable<RowMutation, Void> createPointWriteCallable(
694+
UnaryCallable<BulkMutation, Void> bulkMutateRowsVoidCallable) {
695+
UnaryCallSettings<RowMutation, Void> settings =
696+
perOpSettings.mutateRowSettings.toBuilder()
697+
.setRetrySettings(perOpSettings.bulkMutateRowsSettings.getRetrySettings())
698+
.setRetryableCodes(perOpSettings.bulkMutateRowsSettings.getRetryableCodes())
699+
.build();
700+
701+
UnaryCallable<RowMutation, Void> classic =
702+
new UnaryCallable<RowMutation, Void>() {
703+
@Override
704+
public ApiFuture<Void> futureCall(RowMutation request, ApiCallContext context) {
705+
return bulkMutateRowsVoidCallable.futureCall(
706+
BulkMutation.fromProto(request.toBulkProto(requestContext)), context);
707+
}
708+
};
709+
710+
return bigtableClientContext.getSessionShim().decorateMutateRow(classic, settings);
711+
}
712+
673713
/**
674714
* Creates a callable chain to handle MutatesRows RPCs. This is meant to be used for manual
675715
* batching. The chain will:

‎java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MaybePointWriteCallable.java‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.cloud.bigtable.data.v2.stub.mutaterows;
17+
18+
import com.google.api.core.ApiFuture;
19+
import com.google.api.core.InternalApi;
20+
import com.google.api.gax.rpc.ApiCallContext;
21+
import com.google.api.gax.rpc.UnaryCallable;
22+
import com.google.bigtable.v2.MutateRowRequest;
23+
import com.google.bigtable.v2.MutateRowsRequest;
24+
import com.google.cloud.bigtable.data.v2.internal.RequestContext;
25+
import com.google.cloud.bigtable.data.v2.models.BulkMutation;
26+
import com.google.cloud.bigtable.data.v2.models.RowMutation;
27+
28+
/**
29+
* Routes {@link BulkMutation}s that carry a single entry through a unary point-write callable,
30+
* letting them benefit from the same session-shim diversion as {@code MutateRow}. Bulk mutations
31+
* with more than one entry fall through to the classic {@code MutateRows} callable.
32+
*
33+
* <p>When the session diversion does not apply, the point-write callable falls back to the bulk
34+
* {@code MutateRows} RPC (as a single-entry batch), retaining the bulk operation's retry behavior,
35+
* so the single entry travels over the wire and retries exactly as it would have as part of a
36+
* {@code MutateRows} call.
37+
*/
38+
@InternalApi
39+
public class MaybePointWriteCallable extends UnaryCallable<BulkMutation, Void> {
40+
private final UnaryCallable<BulkMutation, Void> classic;
41+
private final UnaryCallable<RowMutation, Void> pointWriter;
42+
private final RequestContext requestContext;
43+
44+
public MaybePointWriteCallable(
45+
UnaryCallable<BulkMutation, Void> classic,
46+
UnaryCallable<RowMutation, Void> pointWriter,
47+
RequestContext requestContext) {
48+
this.classic = classic;
49+
this.pointWriter = pointWriter;
50+
this.requestContext = requestContext;
51+
}
52+
53+
@Override
54+
public ApiFuture<Void> futureCall(BulkMutation request, ApiCallContext context) {
55+
if (request.getEntryCount() != 1) {
56+
return classic.futureCall(request, context);
57+
}
58+
return pointWriter.futureCall(toRowMutation(request), context);
59+
}
60+
61+
private RowMutation toRowMutation(BulkMutation request) {
62+
MutateRowsRequest proto = request.toProto(requestContext);
63+
MutateRowsRequest.Entry entry = proto.getEntries(0);
64+
MutateRowRequest mutateRowRequest =
65+
MutateRowRequest.newBuilder()
66+
.setAppProfileId(proto.getAppProfileId())
67+
.setTableName(proto.getTableName())
68+
.setAuthorizedViewName(proto.getAuthorizedViewName())
69+
.setRowKey(entry.getRowKey())
70+
.addAllMutations(entry.getMutationsList())
71+
.build();
72+
return RowMutation.fromProto(mutateRowRequest);
73+
}
74+
}

‎java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MaybePointWriteCallableTest.java‎

Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.cloud.bigtable.data.v2.stub.mutaterows;
17+
18+
import static com.google.common.truth.Truth.assertThat;
19+
import static org.junit.jupiter.api.Assertions.assertThrows;
20+
21+
import com.google.api.core.ApiFuture;
22+
import com.google.api.core.SettableApiFuture;
23+
import com.google.api.gax.rpc.ApiCallContext;
24+
import com.google.api.gax.rpc.UnaryCallable;
25+
import com.google.cloud.bigtable.data.v2.internal.RequestContext;
26+
import com.google.cloud.bigtable.data.v2.models.BulkMutation;
27+
import com.google.cloud.bigtable.data.v2.models.Mutation;
28+
import com.google.cloud.bigtable.data.v2.models.RowMutation;
29+
import com.google.cloud.bigtable.data.v2.models.TableId;
30+
import java.util.concurrent.ExecutionException;
31+
import org.junit.jupiter.api.BeforeEach;
32+
import org.junit.jupiter.api.Test;
33+
34+
public class MaybePointWriteCallableTest {
35+
36+
private static final RequestContext REQUEST_CONTEXT =
37+
RequestContext.create("my-project", "my-instance", "my-profile");
38+
private static final TableId TABLE_ID = TableId.of("fake-table");
39+
40+
private FakeBulkCallable classic;
41+
private FakePointWriter pointWriter;
42+
private MaybePointWriteCallable callable;
43+
44+
@BeforeEach
45+
public void setUp() {
46+
classic = new FakeBulkCallable();
47+
pointWriter = new FakePointWriter();
48+
callable = new MaybePointWriteCallable(classic, pointWriter, REQUEST_CONTEXT);
49+
}
50+
51+
@Test
52+
public void singleEntry_routesToPointWriter() throws Exception {
53+
BulkMutation request =
54+
BulkMutation.create(TABLE_ID).add("row-key", Mutation.create().deleteRow());
55+
56+
ApiFuture<Void> future = callable.futureCall(request, null);
57+
pointWriter.response.set(null);
58+
59+
assertThat(future.get()).isNull();
60+
assertThat(classic.request).isNull();
61+
assertThat(pointWriter.request).isNotNull();
62+
// The single entry is converted back into a RowMutation targeting the same row.
63+
assertThat(pointWriter.request.getTargetId()).isEqualTo(TABLE_ID);
64+
}
65+
66+
@Test
67+
public void multipleEntries_fallsThroughToClassic() {
68+
BulkMutation request =
69+
BulkMutation.create(TABLE_ID)
70+
.add("row-a", Mutation.create().deleteRow())
71+
.add("row-b", Mutation.create().deleteRow());
72+
73+
callable.futureCall(request, null);
74+
75+
assertThat(pointWriter.request).isNull();
76+
assertThat(classic.request).isEqualTo(request);
77+
}
78+
79+
@Test
80+
public void pointWriterFails_propagates() {
81+
BulkMutation request =
82+
BulkMutation.create(TABLE_ID).add("row-key", Mutation.create().deleteRow());
83+
RuntimeException failure = new RuntimeException("point boom");
84+
85+
ApiFuture<Void> future = callable.futureCall(request, null);
86+
pointWriter.response.setException(failure);
87+
88+
ExecutionException thrown = assertThrows(ExecutionException.class, future::get);
89+
assertThat(thrown).hasCauseThat().isSameInstanceAs(failure);
90+
}
91+
92+
@Test
93+
public void classicFailure_propagates() {
94+
BulkMutation request =
95+
BulkMutation.create(TABLE_ID)
96+
.add("row-a", Mutation.create().deleteRow())
97+
.add("row-b", Mutation.create().deleteRow());
98+
RuntimeException failure = new RuntimeException("classic boom");
99+
classic.response.setException(failure);
100+
101+
ApiFuture<Void> future = callable.futureCall(request, null);
102+
103+
ExecutionException thrown = assertThrows(ExecutionException.class, future::get);
104+
assertThat(thrown).hasCauseThat().isSameInstanceAs(failure);
105+
}
106+
107+
private static class FakeBulkCallable extends UnaryCallable<BulkMutation, Void> {
108+
BulkMutation request;
109+
final SettableApiFuture<Void> response = SettableApiFuture.create();
110+
111+
@Override
112+
public ApiFuture<Void> futureCall(BulkMutation request, ApiCallContext context) {
113+
this.request = request;
114+
return response;
115+
}
116+
}
117+
118+
private static class FakePointWriter extends UnaryCallable<RowMutation, Void> {
119+
RowMutation request;
120+
final SettableApiFuture<Void> response = SettableApiFuture.create();
121+
122+
@Override
123+
public ApiFuture<Void> futureCall(RowMutation request, ApiCallContext context) {
124+
this.request = request;
125+
return response;
126+
}
127+
}
128+
}

0 commit comments

Comments
 (0)