From 7c96d31af41730c1103a9bff9fa578fb493e32ba Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Mon, 10 Aug 2026 17:34:38 +0000 Subject: [PATCH 1/2] feat(bigtable): route single-entry MutateRows through a point-write callable 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. --- .../data/v2/stub/EnhancedBigtableStub.java | 37 ++++- .../mutaterows/MaybePointWriteCallable.java | 73 ++++++++++ .../MaybePointWriteCallableTest.java | 128 ++++++++++++++++++ 3 files changed, 237 insertions(+), 1 deletion(-) create mode 100644 java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MaybePointWriteCallable.java create mode 100644 java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MaybePointWriteCallableTest.java diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java index 1f1a45c46b05..fc62046bbf24 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java @@ -94,6 +94,7 @@ import com.google.cloud.bigtable.data.v2.stub.metrics.StatsHeadersServerStreamingCallable; import com.google.cloud.bigtable.data.v2.stub.metrics.StatsHeadersUnaryCallable; import com.google.cloud.bigtable.data.v2.stub.mutaterows.BulkMutateRowsUserFacingCallable; +import com.google.cloud.bigtable.data.v2.stub.mutaterows.MaybePointWriteCallable; import com.google.cloud.bigtable.data.v2.stub.mutaterows.MutateRowsAttemptResult; import com.google.cloud.bigtable.data.v2.stub.mutaterows.MutateRowsBatchingDescriptor; import com.google.cloud.bigtable.data.v2.stub.mutaterows.MutateRowsPartialErrorRetryAlgorithm; @@ -200,8 +201,11 @@ public EnhancedBigtableStub( sampleRowKeysCallableWithRequest = createSampleRowKeysCallableWithRequest(); mutateRowCallable = createMutateRowCallable(); bulkMutateRowsCallable = createMutateRowsBaseCallable(); - externalBulkMutateRowsCallable = + UnaryCallable bulkMutateRowsVoidCallable = new MutateRowsErrorConverterUnaryCallable(bulkMutateRowsCallable); + externalBulkMutateRowsCallable = + new MaybePointWriteCallable( + bulkMutateRowsVoidCallable, createPointWriteCallable(), requestContext); checkAndMutateRowCallable = createCheckAndMutateRowCallable(); readModifyWriteRowCallable = createReadModifyWriteRowCallable(); generateInitialChangeStreamPartitionsCallable = @@ -670,6 +674,37 @@ private UnaryCallable createMutateRowCallable() { .decorateMutateRow(classic, perOpSettings.mutateRowSettings); } + /** + * Creates the point-write callable used by {@link MaybePointWriteCallable} to divert single-entry + * {@link BulkMutation}s. This mirrors {@link #createPointReadCallable}: it exposes a single row + * mutation through the session-shim diversion while preserving the bulk operation's retry + * behavior. + * + *

Unlike the plain {@link #createMutateRowCallable()}, this uses the {@code MutateRow} RPC + * (the natural single-row write) but carries the caller's {@code bulkMutateRowsSettings} retry + * settings/codes, so a single-entry bulk write retries the same way it would have as a bulk + * operation. + */ + private UnaryCallable createPointWriteCallable() { + UnaryCallSettings settings = + perOpSettings.mutateRowSettings.toBuilder() + .setRetrySettings(perOpSettings.bulkMutateRowsSettings.getRetrySettings()) + .setRetryableCodes(perOpSettings.bulkMutateRowsSettings.getRetryableCodes()) + .build(); + + UnaryCallable classic = + createUnaryCallable( + BigtableGrpc.getMutateRowMethod(), + req -> + composeRequestParams( + req.getAppProfileId(), req.getTableName(), req.getAuthorizedViewName()), + settings, + req -> req.toProto(requestContext), + resp -> null); + + return bigtableClientContext.getSessionShim().decorateMutateRow(classic, settings); + } + /** * Creates a callable chain to handle MutatesRows RPCs. This is meant to be used for manual * batching. The chain will: diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MaybePointWriteCallable.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MaybePointWriteCallable.java new file mode 100644 index 000000000000..5e0ba964fad6 --- /dev/null +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MaybePointWriteCallable.java @@ -0,0 +1,73 @@ +/* + * Copyright 2026 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.mutaterows; + +import com.google.api.core.ApiFuture; +import com.google.api.core.InternalApi; +import com.google.api.gax.rpc.ApiCallContext; +import com.google.api.gax.rpc.UnaryCallable; +import com.google.bigtable.v2.MutateRowRequest; +import com.google.bigtable.v2.MutateRowsRequest; +import com.google.cloud.bigtable.data.v2.internal.RequestContext; +import com.google.cloud.bigtable.data.v2.models.BulkMutation; +import com.google.cloud.bigtable.data.v2.models.RowMutation; + +/** + * Routes {@link BulkMutation}s that carry a single entry through a unary point-write callable, + * letting them benefit from the same session-shim diversion as {@code MutateRow}. Bulk mutations + * with more than one entry fall through to the classic {@code MutateRows} callable. + * + *

The point-write callable itself falls back to the {@code MutateRow} RPC when the session + * diversion does not apply, while retaining the bulk operation's retry behavior, so the single + * entry retries the same way it would have as part of a {@code MutateRows} call. + */ +@InternalApi +public class MaybePointWriteCallable extends UnaryCallable { + private final UnaryCallable classic; + private final UnaryCallable pointWriter; + private final RequestContext requestContext; + + public MaybePointWriteCallable( + UnaryCallable classic, + UnaryCallable pointWriter, + RequestContext requestContext) { + this.classic = classic; + this.pointWriter = pointWriter; + this.requestContext = requestContext; + } + + @Override + public ApiFuture futureCall(BulkMutation request, ApiCallContext context) { + if (request.getEntryCount() != 1) { + return classic.futureCall(request, context); + } + return pointWriter.futureCall(toRowMutation(request), context); + } + + private RowMutation toRowMutation(BulkMutation request) { + MutateRowsRequest proto = request.toProto(requestContext); + MutateRowsRequest.Entry entry = proto.getEntries(0); + MutateRowRequest mutateRowRequest = + MutateRowRequest.newBuilder() + .setAppProfileId(proto.getAppProfileId()) + .setTableName(proto.getTableName()) + .setAuthorizedViewName(proto.getAuthorizedViewName()) + .setRowKey(entry.getRowKey()) + .addAllMutations(entry.getMutationsList()) + .build(); + return RowMutation.fromProto(mutateRowRequest); + } +} diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MaybePointWriteCallableTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MaybePointWriteCallableTest.java new file mode 100644 index 000000000000..1115fa1eb269 --- /dev/null +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MaybePointWriteCallableTest.java @@ -0,0 +1,128 @@ +/* + * Copyright 2026 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.mutaterows; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import com.google.api.core.ApiFuture; +import com.google.api.core.SettableApiFuture; +import com.google.api.gax.rpc.ApiCallContext; +import com.google.api.gax.rpc.UnaryCallable; +import com.google.cloud.bigtable.data.v2.internal.RequestContext; +import com.google.cloud.bigtable.data.v2.models.BulkMutation; +import com.google.cloud.bigtable.data.v2.models.Mutation; +import com.google.cloud.bigtable.data.v2.models.RowMutation; +import com.google.cloud.bigtable.data.v2.models.TableId; +import java.util.concurrent.ExecutionException; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +public class MaybePointWriteCallableTest { + + private static final RequestContext REQUEST_CONTEXT = + RequestContext.create("my-project", "my-instance", "my-profile"); + private static final TableId TABLE_ID = TableId.of("fake-table"); + + private FakeBulkCallable classic; + private FakePointWriter pointWriter; + private MaybePointWriteCallable callable; + + @BeforeEach + public void setUp() { + classic = new FakeBulkCallable(); + pointWriter = new FakePointWriter(); + callable = new MaybePointWriteCallable(classic, pointWriter, REQUEST_CONTEXT); + } + + @Test + public void singleEntry_routesToPointWriter() throws Exception { + BulkMutation request = + BulkMutation.create(TABLE_ID).add("row-key", Mutation.create().deleteRow()); + + ApiFuture future = callable.futureCall(request, null); + pointWriter.response.set(null); + + assertThat(future.get()).isNull(); + assertThat(classic.request).isNull(); + assertThat(pointWriter.request).isNotNull(); + // The single entry is converted back into a RowMutation targeting the same row. + assertThat(pointWriter.request.getTargetId()).isEqualTo(TABLE_ID); + } + + @Test + public void multipleEntries_fallsThroughToClassic() { + BulkMutation request = + BulkMutation.create(TABLE_ID) + .add("row-a", Mutation.create().deleteRow()) + .add("row-b", Mutation.create().deleteRow()); + + callable.futureCall(request, null); + + assertThat(pointWriter.request).isNull(); + assertThat(classic.request).isEqualTo(request); + } + + @Test + public void pointWriterFails_propagates() { + BulkMutation request = + BulkMutation.create(TABLE_ID).add("row-key", Mutation.create().deleteRow()); + RuntimeException failure = new RuntimeException("point boom"); + + ApiFuture future = callable.futureCall(request, null); + pointWriter.response.setException(failure); + + ExecutionException thrown = assertThrows(ExecutionException.class, future::get); + assertThat(thrown).hasCauseThat().isSameInstanceAs(failure); + } + + @Test + public void classicFailure_propagates() { + BulkMutation request = + BulkMutation.create(TABLE_ID) + .add("row-a", Mutation.create().deleteRow()) + .add("row-b", Mutation.create().deleteRow()); + RuntimeException failure = new RuntimeException("classic boom"); + classic.response.setException(failure); + + ApiFuture future = callable.futureCall(request, null); + + ExecutionException thrown = assertThrows(ExecutionException.class, future::get); + assertThat(thrown).hasCauseThat().isSameInstanceAs(failure); + } + + private static class FakeBulkCallable extends UnaryCallable { + BulkMutation request; + final SettableApiFuture response = SettableApiFuture.create(); + + @Override + public ApiFuture futureCall(BulkMutation request, ApiCallContext context) { + this.request = request; + return response; + } + } + + private static class FakePointWriter extends UnaryCallable { + RowMutation request; + final SettableApiFuture response = SettableApiFuture.create(); + + @Override + public ApiFuture futureCall(RowMutation request, ApiCallContext context) { + this.request = request; + return response; + } + } +} From bed383a1b56b703e58db685228fd5426924a76c7 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Mon, 10 Aug 2026 19:45:28 +0000 Subject: [PATCH 2/2] fix(bigtable): keep single-entry bulk mutation fallback on MutateRows The point-write callable used by MaybePointWriteCallable previously fell back to the MutateRow RPC when the session shim did not divert traffic. Unlike point reads (where the single-row read is just ReadRows with a limit), MutateRow and MutateRows are distinct RPCs, so this changed the wire behavior for single-entry bulk mutations and broke tests that only implement MutateRows. Build the fallback classic so it delegates to the bulk MutateRows callable as a single-entry batch, then let the shim decorate it. When the shim diverts, the mutation goes to the session single-row write API; otherwise it stays on MutateRows with the bulk operation's retry behavior. This removes the need for the nullable/reference-equality check on the shim's return value. --- .../data/v2/stub/EnhancedBigtableStub.java | 33 +++++++++++-------- .../mutaterows/MaybePointWriteCallable.java | 7 ++-- 2 files changed, 23 insertions(+), 17 deletions(-) diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java index fc62046bbf24..90c49b61518f 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java @@ -205,7 +205,9 @@ public EnhancedBigtableStub( new MutateRowsErrorConverterUnaryCallable(bulkMutateRowsCallable); externalBulkMutateRowsCallable = new MaybePointWriteCallable( - bulkMutateRowsVoidCallable, createPointWriteCallable(), requestContext); + bulkMutateRowsVoidCallable, + createPointWriteCallable(bulkMutateRowsVoidCallable), + requestContext); checkAndMutateRowCallable = createCheckAndMutateRowCallable(); readModifyWriteRowCallable = createReadModifyWriteRowCallable(); generateInitialChangeStreamPartitionsCallable = @@ -680,12 +682,16 @@ private UnaryCallable createMutateRowCallable() { * mutation through the session-shim diversion while preserving the bulk operation's retry * behavior. * - *

Unlike the plain {@link #createMutateRowCallable()}, this uses the {@code MutateRow} RPC - * (the natural single-row write) but carries the caller's {@code bulkMutateRowsSettings} retry - * settings/codes, so a single-entry bulk write retries the same way it would have as a bulk - * operation. + *

Unlike point reads (where the single-row read is just {@code ReadRows} with a limit), {@code + * MutateRow} and {@code MutateRows} are distinct RPCs. To preserve the existing wire behavior, + * the fallback classic here delegates to the bulk {@code MutateRows} callable as a single-entry + * batch rather than issuing a {@code MutateRow} RPC. So when the session shim does not divert + * (e.g. a {@link com.google.cloud.bigtable.data.v2.internal.compat.DisabledShim}), a single-entry + * bulk mutation still travels over {@code MutateRows} with the bulk operation's retry behavior; + * only when the shim actively diverts does the mutation go to the session single-row write API. */ - private UnaryCallable createPointWriteCallable() { + private UnaryCallable createPointWriteCallable( + UnaryCallable bulkMutateRowsVoidCallable) { UnaryCallSettings settings = perOpSettings.mutateRowSettings.toBuilder() .setRetrySettings(perOpSettings.bulkMutateRowsSettings.getRetrySettings()) @@ -693,14 +699,13 @@ private UnaryCallable createPointWriteCallable() { .build(); UnaryCallable classic = - createUnaryCallable( - BigtableGrpc.getMutateRowMethod(), - req -> - composeRequestParams( - req.getAppProfileId(), req.getTableName(), req.getAuthorizedViewName()), - settings, - req -> req.toProto(requestContext), - resp -> null); + new UnaryCallable() { + @Override + public ApiFuture futureCall(RowMutation request, ApiCallContext context) { + return bulkMutateRowsVoidCallable.futureCall( + BulkMutation.fromProto(request.toBulkProto(requestContext)), context); + } + }; return bigtableClientContext.getSessionShim().decorateMutateRow(classic, settings); } diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MaybePointWriteCallable.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MaybePointWriteCallable.java index 5e0ba964fad6..efbbb076657e 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MaybePointWriteCallable.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MaybePointWriteCallable.java @@ -30,9 +30,10 @@ * letting them benefit from the same session-shim diversion as {@code MutateRow}. Bulk mutations * with more than one entry fall through to the classic {@code MutateRows} callable. * - *

The point-write callable itself falls back to the {@code MutateRow} RPC when the session - * diversion does not apply, while retaining the bulk operation's retry behavior, so the single - * entry retries the same way it would have as part of a {@code MutateRows} call. + *

When the session diversion does not apply, the point-write callable falls back to the bulk + * {@code MutateRows} RPC (as a single-entry batch), retaining the bulk operation's retry behavior, + * so the single entry travels over the wire and retries exactly as it would have as part of a + * {@code MutateRows} call. */ @InternalApi public class MaybePointWriteCallable extends UnaryCallable {