feat(bigtable): route single-entry MutateRows through a point-write c…

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,13 @@ public EnhancedBigtableStub(
     sampleRowKeysCallableWithRequest = createSampleRowKeysCallableWithRequest();
     mutateRowCallable = createMutateRowCallable();
     bulkMutateRowsCallable = createMutateRowsBaseCallable();
-    externalBulkMutateRowsCallable =
+    UnaryCallable<BulkMutation, Void> bulkMutateRowsVoidCallable =
         new MutateRowsErrorConverterUnaryCallable(bulkMutateRowsCallable);
+    externalBulkMutateRowsCallable =
+        new MaybePointWriteCallable(
+            bulkMutateRowsVoidCallable,
+            createPointWriteCallable(bulkMutateRowsVoidCallable),
+            requestContext);
     checkAndMutateRowCallable = createCheckAndMutateRowCallable();
     readModifyWriteRowCallable = createReadModifyWriteRowCallable();
     generateInitialChangeStreamPartitionsCallable =
@@ -670,6 +676,40 @@ private UnaryCallable<RowMutation, Void> 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.
+   *
+   * <p>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<RowMutation, Void> createPointWriteCallable(
+      UnaryCallable<BulkMutation, Void> bulkMutateRowsVoidCallable) {
+    UnaryCallSettings<RowMutation, Void> settings =
+        perOpSettings.mutateRowSettings.toBuilder()
+            .setRetrySettings(perOpSettings.bulkMutateRowsSettings.getRetrySettings())
+            .setRetryableCodes(perOpSettings.bulkMutateRowsSettings.getRetryableCodes())
+            .build();
+
+    UnaryCallable<RowMutation, Void> classic =
+        new UnaryCallable<RowMutation, Void>() {
+          @Override
+          public ApiFuture<Void> futureCall(RowMutation request, ApiCallContext context) {
+            return bulkMutateRowsVoidCallable.futureCall(
+                BulkMutation.fromProto(request.toBulkProto(requestContext)), context);
+          }
+        };
+
+    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:

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

@@ -0,0 +1,74 @@
+/*
+ * 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.
+ *
+ * <p>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<BulkMutation, Void> {
+  private final UnaryCallable<BulkMutation, Void> classic;
+  private final UnaryCallable<RowMutation, Void> pointWriter;
+  private final RequestContext requestContext;
+
+  public MaybePointWriteCallable(
+      UnaryCallable<BulkMutation, Void> classic,
+      UnaryCallable<RowMutation, Void> pointWriter,
+      RequestContext requestContext) {
+    this.classic = classic;
+    this.pointWriter = pointWriter;
+    this.requestContext = requestContext;
+  }
+
+  @Override
+  public ApiFuture<Void> 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);
+  }
+}

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<Void> 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<Void> 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<Void> future = callable.futureCall(request, null);
+
+    ExecutionException thrown = assertThrows(ExecutionException.class, future::get);
+    assertThat(thrown).hasCauseThat().isSameInstanceAs(failure);
+  }
+
+  private static class FakeBulkCallable extends UnaryCallable<BulkMutation, Void> {
+    BulkMutation request;
+    final SettableApiFuture<Void> response = SettableApiFuture.create();
+
+    @Override
+    public ApiFuture<Void> futureCall(BulkMutation request, ApiCallContext context) {
+      this.request = request;
+      return response;
+    }
+  }
+
+  private static class FakePointWriter extends UnaryCallable<RowMutation, Void> {
+    RowMutation request;
+    final SettableApiFuture<Void> response = SettableApiFuture.create();
+
+    @Override
+    public ApiFuture<Void> futureCall(RowMutation request, ApiCallContext context) {
+      this.request = request;
+      return response;
+    }
+  }
+}