diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/AuthorizedViewAsync.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/AuthorizedViewAsync.java index 0f2cf71a1a9d..0b3ef63e30b8 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/AuthorizedViewAsync.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/AuthorizedViewAsync.java @@ -22,6 +22,8 @@ import com.google.bigtable.v2.SessionCheckAndMutateRowResponse; import com.google.bigtable.v2.SessionMutateRowRequest; import com.google.bigtable.v2.SessionMutateRowResponse; +import com.google.bigtable.v2.SessionMutateRowsRequest; +import com.google.bigtable.v2.SessionMutateRowsResponse; import com.google.bigtable.v2.SessionReadModifyWriteRowRequest; import com.google.bigtable.v2.SessionReadModifyWriteRowResponse; import com.google.bigtable.v2.SessionReadRowRequest; @@ -80,6 +82,7 @@ static AuthorizedViewAsync createAndStart( VRpcDescriptor.MUTATE_ROW_AUTH_VIEW, VRpcDescriptor.CHECK_AND_MUTATE_ROW_AUTH_VIEW, VRpcDescriptor.READ_MODIFY_WRITE_ROW_AUTH_VIEW, + VRpcDescriptor.MUTATE_ROWS_AUTH_VIEW, featureFlags, clientInfo, configManager, @@ -130,6 +133,13 @@ public CompletableFuture readModifyWriteRow( return f; } + public CompletableFuture mutateRows( + SessionMutateRowsRequest req, Deadline deadline) { + UnaryResponseFuture f = new UnaryResponseFuture<>(); + base.mutateRows(req, f, deadline); + return f; + } + @Override public void close() { this.base.close(); diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/MaterializedViewAsync.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/MaterializedViewAsync.java index 856e424118f5..a8c7004fab89 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/MaterializedViewAsync.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/MaterializedViewAsync.java @@ -71,6 +71,7 @@ public static MaterializedViewAsync createAndStart( /* mutateRowDescriptor= */ null, // Materialized views are read-only /* checkAndMutateRowDescriptor= */ null, /* readModifyWriteRowDescriptor= */ null, + /* mutateRowsDescriptor= */ null, featureFlags, clientInfo, configManager, diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableAsync.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableAsync.java index 9d6a5b515ffd..086dd79bdf07 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableAsync.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableAsync.java @@ -23,6 +23,8 @@ import com.google.bigtable.v2.SessionCheckAndMutateRowResponse; import com.google.bigtable.v2.SessionMutateRowRequest; import com.google.bigtable.v2.SessionMutateRowResponse; +import com.google.bigtable.v2.SessionMutateRowsRequest; +import com.google.bigtable.v2.SessionMutateRowsResponse; import com.google.bigtable.v2.SessionReadModifyWriteRowRequest; import com.google.bigtable.v2.SessionReadModifyWriteRowResponse; import com.google.bigtable.v2.SessionReadRowRequest; @@ -78,6 +80,7 @@ public static TableAsync createAndStart( VRpcDescriptor.MUTATE_ROW, VRpcDescriptor.CHECK_AND_MUTATE_ROW, VRpcDescriptor.READ_MODIFY_WRITE_ROW, + VRpcDescriptor.MUTATE_ROWS, featureFlags, clientInfo, configManager, @@ -139,4 +142,11 @@ public CompletableFuture readModifyWriteRow( base.readModifyWriteRow(req, f, deadline); return f; } + + public CompletableFuture mutateRows( + SessionMutateRowsRequest req, Deadline deadline) { + UnaryResponseFuture f = new UnaryResponseFuture<>(); + base.mutateRows(req, f, deadline); + return f; + } } diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableBase.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableBase.java index f82861979ba6..f8d6ec880927 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableBase.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableBase.java @@ -15,12 +15,17 @@ */ package com.google.cloud.bigtable.data.v2.internal.api; +import static com.google.common.collect.ImmutableList.toImmutableList; + import com.google.bigtable.v2.CloseSessionRequest; import com.google.bigtable.v2.FeatureFlags; +import com.google.bigtable.v2.Mutation; import com.google.bigtable.v2.SessionCheckAndMutateRowRequest; import com.google.bigtable.v2.SessionCheckAndMutateRowResponse; import com.google.bigtable.v2.SessionMutateRowRequest; import com.google.bigtable.v2.SessionMutateRowResponse; +import com.google.bigtable.v2.SessionMutateRowsRequest; +import com.google.bigtable.v2.SessionMutateRowsResponse; import com.google.bigtable.v2.SessionReadModifyWriteRowRequest; import com.google.bigtable.v2.SessionReadModifyWriteRowResponse; import com.google.bigtable.v2.SessionReadRowRequest; @@ -40,6 +45,7 @@ import com.google.cloud.bigtable.data.v2.internal.session.VRpcDescriptor; import com.google.cloud.bigtable.data.v2.internal.util.ClientConfigurationManager; import com.google.common.annotations.VisibleForTesting; +import com.google.common.collect.ImmutableList; import com.google.protobuf.Message; import io.grpc.CallOptions; import io.grpc.Context; @@ -60,6 +66,8 @@ class TableBase implements AutoCloseable { private final VRpcDescriptor< ?, SessionReadModifyWriteRowRequest, SessionReadModifyWriteRowResponse> readModifyWriteRowDescriptor; + private final VRpcDescriptor + mutateRowsDescriptor; static TableBase createAndStart( ReqT openReq, @@ -70,6 +78,7 @@ static TableBase createAndStart( checkAndMutateRowDescriptor, VRpcDescriptor readModifyWriteRowDescriptor, + VRpcDescriptor mutateRowsDescriptor, FeatureFlags featureFlags, ClientInfo clientInfo, ClientConfigurationManager configManager, @@ -108,6 +117,7 @@ static TableBase createAndStart( mutateRowDescriptor, checkAndMutateRowDescriptor, readModifyWriteRowDescriptor, + mutateRowsDescriptor, metrics, timer, userCallbackExecutor); @@ -122,6 +132,7 @@ static TableBase createAndStart( checkAndMutateRowDescriptor, VRpcDescriptor readModifyWriteRowDescriptor, + VRpcDescriptor mutateRowsDescriptor, Metrics metrics, BigtableTimer timer, Executor userCallbackExecutor) { @@ -130,6 +141,7 @@ static TableBase createAndStart( this.mutateRowDescriptor = mutateRowDescriptor; this.checkAndMutateRowDescriptor = checkAndMutateRowDescriptor; this.readModifyWriteRowDescriptor = readModifyWriteRowDescriptor; + this.mutateRowsDescriptor = mutateRowsDescriptor; this.metrics = metrics; this.timer = timer; this.userCallbackExecutor = userCallbackExecutor; @@ -202,4 +214,23 @@ public void readModifyWriteRow( new VOperationImpl<>(retry, Context.current(), userCallbackExecutor, tracer, deadline, false) .start(req, listener); } + + public void mutateRows( + SessionMutateRowsRequest req, + VRpcListener listener, + Deadline deadline) { + RetryingVRpc retry = + new RetryingVRpc<>(() -> sessionPool.newCall(mutateRowsDescriptor), timer); + ImmutableList mutations = + req.getEntriesList().stream() + .flatMap(entry -> entry.getMutationsList().stream()) + .collect(toImmutableList()); + boolean idempotent = Util.isIdempotent(mutations); + VRpcTracer tracer = + metrics.newTableTracer(sessionPool.getInfo(), mutateRowsDescriptor, deadline); + + new VOperationImpl<>( + retry, Context.current(), userCallbackExecutor, tracer, deadline, idempotent) + .start(req, listener); + } } diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/DisabledShim.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/DisabledShim.java index b95064df164b..beabeb1d69b4 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/DisabledShim.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/DisabledShim.java @@ -17,6 +17,7 @@ import com.google.api.gax.rpc.UnaryCallSettings; import com.google.api.gax.rpc.UnaryCallable; +import com.google.cloud.bigtable.data.v2.models.BulkMutation; import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation; import com.google.cloud.bigtable.data.v2.models.Query; import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow; @@ -55,4 +56,10 @@ public UnaryCallable decorateReadModifyWriteRow UnaryCallSettings settings) { return classic; } + + @Override + public UnaryCallable decorateMutateRows( + UnaryCallable classic, UnaryCallSettings settings) { + return classic; + } } diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/Shim.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/Shim.java index 2fe6aade5cd0..2fbfde7a1e6e 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/Shim.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/Shim.java @@ -17,6 +17,7 @@ import com.google.api.gax.rpc.UnaryCallSettings; import com.google.api.gax.rpc.UnaryCallable; +import com.google.cloud.bigtable.data.v2.models.BulkMutation; import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation; import com.google.cloud.bigtable.data.v2.models.Query; import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow; @@ -42,4 +43,7 @@ UnaryCallable decorateReadModifyWriteRow( UnaryCallable classic, RowAdapter rowAdapter, UnaryCallSettings settings); + + UnaryCallable decorateMutateRows( + UnaryCallable classic, UnaryCallSettings settings); } diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java index 8db91d0a9308..5742f2906565 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java @@ -40,6 +40,7 @@ import com.google.cloud.bigtable.data.v2.internal.compat.ops.CheckAndMutateRowShim; import com.google.cloud.bigtable.data.v2.internal.compat.ops.DivertingUnaryCallable; import com.google.cloud.bigtable.data.v2.internal.compat.ops.MutateRowShim; +import com.google.cloud.bigtable.data.v2.internal.compat.ops.MutateRowsShim; import com.google.cloud.bigtable.data.v2.internal.compat.ops.ReadModifyWriteRowShim; import com.google.cloud.bigtable.data.v2.internal.compat.ops.ReadRowShim; import com.google.cloud.bigtable.data.v2.internal.compat.ops.RowBuilderShim; @@ -51,6 +52,7 @@ import com.google.cloud.bigtable.data.v2.internal.dp.DirectAccessInvestigator; import com.google.cloud.bigtable.data.v2.internal.util.ClientConfigurationManager; import com.google.cloud.bigtable.data.v2.internal.util.SessionScope; +import com.google.cloud.bigtable.data.v2.models.BulkMutation; import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation; import com.google.cloud.bigtable.data.v2.models.Query; import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow; @@ -110,6 +112,7 @@ public class ShimImpl implements Shim { private final MutateRowShim mutateRowShim; private final CheckAndMutateRowShim checkAndMutateRowShim; private final ReadModifyWriteRowShim readModifyWriteRowShim; + private final MutateRowsShim mutateRowsShim; public static Shim create( ClientInfo clientInfo, @@ -228,6 +231,7 @@ public ShimImpl( this.mutateRowShim = new MutateRowShim(sessionMap); this.checkAndMutateRowShim = new CheckAndMutateRowShim(sessionMap); this.readModifyWriteRowShim = new ReadModifyWriteRowShim(sessionMap); + this.mutateRowsShim = new MutateRowsShim(sessionMap); } /** @@ -435,4 +439,16 @@ public UnaryCallable decorateReadModifyWriteRow debugTagTracer, SessionScope.READ_MODIFY_WRITE_ROW); } + + @Override + public UnaryCallable decorateMutateRows( + UnaryCallable classic, UnaryCallSettings settings) { + return new DivertingUnaryCallable<>( + configManager, + classic, + mutateRowsShim, + Util.extractTimeout(settings), + debugTagTracer, + SessionScope.MUTATE_ROWS); + } } diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/MutateRowsShim.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/MutateRowsShim.java new file mode 100644 index 000000000000..ad843a02c584 --- /dev/null +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/MutateRowsShim.java @@ -0,0 +1,77 @@ +/* + * 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.internal.compat.ops; + +import com.google.bigtable.v2.SessionMutateRowsRequest; +import com.google.cloud.bigtable.data.v2.internal.session.SessionPool; +import com.google.cloud.bigtable.data.v2.models.AuthorizedViewId; +import com.google.cloud.bigtable.data.v2.models.BulkMutation; +import com.google.cloud.bigtable.data.v2.models.TableId; +import com.google.cloud.bigtable.data.v2.models.TargetId; +import io.grpc.Deadline; +import java.util.concurrent.CompletableFuture; + +public class MutateRowsShim implements UnaryShim { + + private final SessionPoolMap sessionMap; + + public MutateRowsShim(SessionPoolMap sessionMap) { + this.sessionMap = sessionMap; + } + + @Override + public void close() {} + + @Override + public boolean supports(BulkMutation request) { + TargetId targetId = request.getTargetId(); + SessionPool pool; + if (targetId instanceof TableId) { + pool = sessionMap.get((TableId) targetId, SessionKey.Permission.WRITE).getSessionPool(); + } else if (targetId instanceof AuthorizedViewId) { + pool = + sessionMap.get((AuthorizedViewId) targetId, SessionKey.Permission.WRITE).getSessionPool(); + } else { + return false; + } + return UnaryShim.shouldRouteToSession(pool); + } + + @Override + public CompletableFuture call(BulkMutation request, Deadline deadline) { + TargetId targetId = request.getTargetId(); + + SessionMutateRowsRequest innerReq = request.toSessionProto(); + + if (targetId instanceof TableId) { + return sessionMap.apply( + (TableId) targetId, + SessionKey.Permission.WRITE, + t -> t.mutateRows(innerReq, deadline).thenApply(r -> null)); + } + if (targetId instanceof AuthorizedViewId) { + return sessionMap.apply( + (AuthorizedViewId) targetId, + SessionKey.Permission.WRITE, + v -> v.mutateRows(innerReq, deadline).thenApply(r -> null)); + } + + CompletableFuture f = new CompletableFuture<>(); + f.completeExceptionally( + new UnsupportedOperationException("Unsupported targetId type: " + targetId)); + return f; + } +} diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/VRpcDescriptor.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/VRpcDescriptor.java index 575fb79fa26b..ad56311ed9a6 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/VRpcDescriptor.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/VRpcDescriptor.java @@ -23,6 +23,7 @@ import com.google.bigtable.v2.MaterializedViewRequest; import com.google.bigtable.v2.MaterializedViewResponse; import com.google.bigtable.v2.MutateRowRequest; +import com.google.bigtable.v2.MutateRowsRequest; import com.google.bigtable.v2.OpenAuthorizedViewRequest; import com.google.bigtable.v2.OpenMaterializedViewRequest; import com.google.bigtable.v2.OpenTableRequest; @@ -33,6 +34,8 @@ import com.google.bigtable.v2.SessionCheckAndMutateRowResponse; import com.google.bigtable.v2.SessionMutateRowRequest; import com.google.bigtable.v2.SessionMutateRowResponse; +import com.google.bigtable.v2.SessionMutateRowsRequest; +import com.google.bigtable.v2.SessionMutateRowsResponse; import com.google.bigtable.v2.SessionReadModifyWriteRowRequest; import com.google.bigtable.v2.SessionReadModifyWriteRowResponse; import com.google.bigtable.v2.SessionReadRowRequest; @@ -189,6 +192,30 @@ public final class VRpcDescriptor { .addAllRules(req.getRulesList()) .build()); + public static final VRpcDescriptor< + OpenTableRequest, SessionMutateRowsRequest, SessionMutateRowsResponse> + MUTATE_ROWS = + new VRpcDescriptor<>( + TABLE_SESSION, + MethodInfo.of("Bigtable.MutateRows", false), + createTableEncoder(TableRequest.Builder::setMutateRows), + createTableDecoder(TableResponse::getMutateRows), + (name, appProfileId, req) -> { + MutateRowsRequest.Builder builder = + MutateRowsRequest.newBuilder().setTableName(name).setAppProfileId(appProfileId); + for (SessionMutateRowsRequest.Entry entry : req.getEntriesList()) { + MutateRowsRequest.Entry.Builder entryBuilder = + MutateRowsRequest.Entry.newBuilder() + .setRowKey(entry.getKey()) + .addAllMutations(entry.getMutationsList()); + if (entry.hasIdempotency()) { + entryBuilder.setIdempotency(entry.getIdempotency()); + } + builder.addEntries(entryBuilder); + } + return builder.build(); + }); + public static final VRpcDescriptor< OpenAuthorizedViewRequest, SessionReadRowRequest, SessionReadRowResponse> READ_ROW_AUTH_VIEW = @@ -264,6 +291,32 @@ public final class VRpcDescriptor { .addAllRules(req.getRulesList()) .build()); + public static final VRpcDescriptor< + OpenAuthorizedViewRequest, SessionMutateRowsRequest, SessionMutateRowsResponse> + MUTATE_ROWS_AUTH_VIEW = + new VRpcDescriptor<>( + AUTHORIZED_VIEW_SESSION, + MethodInfo.of("Bigtable.MutateRows", false), + createAuthViewEncoder(AuthorizedViewRequest.Builder::setMutateRows), + createAuthViewDecoder(AuthorizedViewResponse::getMutateRows), + (name, appProfileId, req) -> { + MutateRowsRequest.Builder builder = + MutateRowsRequest.newBuilder() + .setAuthorizedViewName(name) + .setAppProfileId(appProfileId); + for (SessionMutateRowsRequest.Entry entry : req.getEntriesList()) { + MutateRowsRequest.Entry.Builder entryBuilder = + MutateRowsRequest.Entry.newBuilder() + .setRowKey(entry.getKey()) + .addAllMutations(entry.getMutationsList()); + if (entry.hasIdempotency()) { + entryBuilder.setIdempotency(entry.getIdempotency()); + } + builder.addEntries(entryBuilder); + } + return builder.build(); + }); + public static final VRpcDescriptor< OpenMaterializedViewRequest, SessionReadRowRequest, SessionReadRowResponse> READ_ROW_MAT_VIEW = diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/util/SessionScope.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/util/SessionScope.java index e89892a5eba8..b7e27928b2eb 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/util/SessionScope.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/util/SessionScope.java @@ -35,7 +35,8 @@ public enum SessionScope { CHECK_AND_MUTATE_ROW( "method:Bigtable.CheckAndMutateRow", /* fallbackToGlobalSessionLoad= */ false), READ_MODIFY_WRITE_ROW( - "method:Bigtable.ReadModifyWriteRow", /* fallbackToGlobalSessionLoad= */ false); + "method:Bigtable.ReadModifyWriteRow", /* fallbackToGlobalSessionLoad= */ false), + MUTATE_ROWS("method:Bigtable.MutateRows", /* fallbackToGlobalSessionLoad= */ false); private final String key; private final boolean fallbackToGlobalSessionLoad; diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/BulkMutation.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/BulkMutation.java index fb7b04269f07..767ffc770234 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/BulkMutation.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/BulkMutation.java @@ -20,6 +20,7 @@ import com.google.api.core.BetaApi; import com.google.api.core.InternalApi; import com.google.bigtable.v2.MutateRowsRequest; +import com.google.bigtable.v2.SessionMutateRowsRequest; import com.google.cloud.bigtable.data.v2.internal.NameUtil; import com.google.cloud.bigtable.data.v2.internal.RequestContext; import com.google.common.base.Preconditions; @@ -128,6 +129,27 @@ public int getEntryCount() { return builder.getEntriesCount(); } + @InternalApi + public TargetId getTargetId() { + return targetId; + } + + @InternalApi + public SessionMutateRowsRequest toSessionProto() { + SessionMutateRowsRequest.Builder sessionBuilder = SessionMutateRowsRequest.newBuilder(); + for (MutateRowsRequest.Entry entry : builder.getEntriesList()) { + SessionMutateRowsRequest.Entry.Builder entryBuilder = + SessionMutateRowsRequest.Entry.newBuilder() + .setKey(entry.getRowKey()) + .addAllMutations(entry.getMutationsList()); + if (entry.hasIdempotency()) { + entryBuilder.setIdempotency(entry.getIdempotency()); + } + sessionBuilder.addEntries(entryBuilder); + } + return sessionBuilder.build(); + } + @InternalApi public MutateRowsRequest toProto(RequestContext requestContext) { String resourceName = 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 cafc9e78e8b4..11f7cb18576d 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 @@ -203,11 +203,20 @@ public EnhancedBigtableStub( bulkMutateRowsCallable = createMutateRowsBaseCallable(); UnaryCallable bulkMutateRowsVoidCallable = new MutateRowsErrorConverterUnaryCallable(bulkMutateRowsCallable); - externalBulkMutateRowsCallable = + UnaryCallable maybePointWriteCallable = new MaybePointWriteCallable( bulkMutateRowsVoidCallable, createPointWriteCallable(bulkMutateRowsVoidCallable), requestContext); + externalBulkMutateRowsCallable = + bigtableClientContext + .getSessionShim() + .decorateMutateRows(maybePointWriteCallable, perOpSettings.bulkMutateRowsSettings) + .withDefaultCallContext( + bigtableClientContext + .getClientContext() + .getDefaultCallContext() + .withRetrySettings(perOpSettings.bulkMutateRowsSettings.getRetrySettings())); checkAndMutateRowCallable = createCheckAndMutateRowCallable(); readModifyWriteRowCallable = createReadModifyWriteRowCallable(); generateInitialChangeStreamPartitionsCallable = diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/api/TableBaseTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/api/TableBaseTest.java index 64b32d92c27d..1207a47da84e 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/api/TableBaseTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/api/TableBaseTest.java @@ -25,6 +25,7 @@ import com.google.bigtable.v2.SessionCheckAndMutateRowRequest; import com.google.bigtable.v2.SessionMutateRowRequest; import com.google.bigtable.v2.SessionMutateRowResponse; +import com.google.bigtable.v2.SessionMutateRowsRequest; import com.google.bigtable.v2.SessionReadModifyWriteRowRequest; import com.google.cloud.bigtable.data.v2.internal.csm.Metrics; import com.google.cloud.bigtable.data.v2.internal.csm.NoopMetrics; @@ -80,6 +81,7 @@ public void setup() { VRpcDescriptor.MUTATE_ROW, VRpcDescriptor.CHECK_AND_MUTATE_ROW, VRpcDescriptor.READ_MODIFY_WRITE_ROW, + VRpcDescriptor.MUTATE_ROWS, noopMetrics, mockTimer, MoreExecutors.directExecutor()); @@ -200,6 +202,38 @@ public void testReadModifyWriteRowNotIdempotent() { assertThat(fakeSessionPool.lastVRpc.ctx).isNotIdempotent(); } + @Test + public void testMutateRowsIdempotent() { + table.mutateRows( + SessionMutateRowsRequest.newBuilder() + .addEntries( + SessionMutateRowsRequest.Entry.newBuilder() + .addMutations( + Mutation.newBuilder() + .setSetCell(Mutation.SetCell.getDefaultInstance()) + .build())) + .build(), + new UnaryResponseFuture<>(), + deadline); + assertThat(fakeSessionPool.lastVRpc.ctx).isIdempotent(); + } + + @Test + public void testMutateRowsNotIdempotent() { + table.mutateRows( + SessionMutateRowsRequest.newBuilder() + .addEntries( + SessionMutateRowsRequest.Entry.newBuilder() + .addMutations( + Mutation.newBuilder() + .setAddToCell(Mutation.AddToCell.getDefaultInstance()) + .build())) + .build(), + new UnaryResponseFuture<>(), + deadline); + assertThat(fakeSessionPool.lastVRpc.ctx).isNotIdempotent(); + } + static class FakeSessionPool implements SessionPool { private FakeVRpc lastVRpc = null; diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/MutateRowsShimTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/MutateRowsShimTest.java new file mode 100644 index 000000000000..17ea5e6c7301 --- /dev/null +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/MutateRowsShimTest.java @@ -0,0 +1,149 @@ +/* + * 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.internal.compat.ops; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.google.bigtable.v2.SessionMutateRowsResponse; +import com.google.cloud.bigtable.data.v2.internal.api.AuthorizedViewAsync; +import com.google.cloud.bigtable.data.v2.internal.api.Client; +import com.google.cloud.bigtable.data.v2.internal.api.TableAsync; +import com.google.cloud.bigtable.data.v2.internal.session.SessionPool; +import com.google.cloud.bigtable.data.v2.models.AuthorizedViewId; +import com.google.cloud.bigtable.data.v2.models.BulkMutation; +import com.google.cloud.bigtable.data.v2.models.TableId; +import com.google.cloud.bigtable.data.v2.models.TargetId; +import io.grpc.Deadline; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +class MutateRowsShimTest { + + private Client client; + private TableAsync tableAsync; + private AuthorizedViewAsync authViewAsync; + private SessionPool tablePool; + private SessionPool authViewPool; + private MutateRowsShim shim; + + @BeforeEach + @SuppressWarnings("unchecked") + void setUp() { + client = mock(Client.class); + tableAsync = mock(TableAsync.class); + authViewAsync = mock(AuthorizedViewAsync.class); + tablePool = mock(SessionPool.class); + authViewPool = mock(SessionPool.class); + + when(client.openTableAsync(any(), any())).thenReturn(tableAsync); + when(client.openAuthorizedViewAsync(any(), any(), any())).thenReturn(authViewAsync); + doReturn(tablePool).when(tableAsync).getSessionPool(); + doReturn(authViewPool).when(authViewAsync).getSessionPool(); + + shim = new MutateRowsShim(new SessionPoolMap(client)); + } + + @Test + void supports_tableId_belowThreshold_returnsTrue() { + when(tablePool.getConsecutiveUnimplementedFailures()) + .thenReturn(UnaryShim.MAX_CONSECUTIVE_UNIMPLEMENTED_FAILURES - 1); + + BulkMutation request = BulkMutation.create(TableId.of("table1")); + assertThat(shim.supports(request)).isTrue(); + } + + @Test + void supports_tableId_atThresholdNoSession_returnsFalse() { + when(tablePool.getConsecutiveUnimplementedFailures()) + .thenReturn(UnaryShim.MAX_CONSECUTIVE_UNIMPLEMENTED_FAILURES); + when(tablePool.hasSession()).thenReturn(false); + + BulkMutation request = BulkMutation.create(TableId.of("table1")); + assertThat(shim.supports(request)).isFalse(); + } + + @Test + void supports_tableId_atThresholdWithSession_returnsTrue() { + when(tablePool.getConsecutiveUnimplementedFailures()) + .thenReturn(UnaryShim.MAX_CONSECUTIVE_UNIMPLEMENTED_FAILURES); + when(tablePool.hasSession()).thenReturn(true); + + BulkMutation request = BulkMutation.create(TableId.of("table1")); + assertThat(shim.supports(request)).isTrue(); + } + + @Test + void supports_authorizedViewId_belowThreshold_returnsTrue() { + when(authViewPool.getConsecutiveUnimplementedFailures()) + .thenReturn(UnaryShim.MAX_CONSECUTIVE_UNIMPLEMENTED_FAILURES - 1); + + BulkMutation request = BulkMutation.create(AuthorizedViewId.of("table1", "view1")); + assertThat(shim.supports(request)).isTrue(); + } + + @Test + void supports_unsupportedTargetId_returnsFalse() { + TargetId unsupported = mock(TargetId.class); + BulkMutation request = BulkMutation.create(unsupported); + assertThat(shim.supports(request)).isFalse(); + } + + @Test + void call_tableId_returnsNull() throws Exception { + SessionMutateRowsResponse response = SessionMutateRowsResponse.getDefaultInstance(); + when(tableAsync.mutateRows(any(), any())) + .thenReturn(CompletableFuture.completedFuture(response)); + + BulkMutation request = BulkMutation.create(TableId.of("table1")); + Deadline deadline = Deadline.after(1, TimeUnit.SECONDS); + + assertNull(shim.call(request, deadline).get()); + } + + @Test + void call_authorizedViewId_returnsNull() throws Exception { + SessionMutateRowsResponse response = SessionMutateRowsResponse.getDefaultInstance(); + when(authViewAsync.mutateRows(any(), any())) + .thenReturn(CompletableFuture.completedFuture(response)); + + BulkMutation request = BulkMutation.create(AuthorizedViewId.of("table1", "view1")); + Deadline deadline = Deadline.after(1, TimeUnit.SECONDS); + + assertNull(shim.call(request, deadline).get()); + } + + @Test + void call_unsupportedTargetId_futureFailsWithUnsupportedOperationException() { + TargetId unsupported = mock(TargetId.class); + BulkMutation request = BulkMutation.create(unsupported); + Deadline deadline = Deadline.after(1, TimeUnit.SECONDS); + + CompletableFuture future = shim.call(request, deadline); + + assertThat(future.isCompletedExceptionally()).isTrue(); + ExecutionException ex = assertThrows(ExecutionException.class, future::get); + assertThat(ex.getCause()).isInstanceOf(UnsupportedOperationException.class); + } +} diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/VRpcDescriptorTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/VRpcDescriptorTest.java index fa86d4dcd4b1..93bbfdb33e6b 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/VRpcDescriptorTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/VRpcDescriptorTest.java @@ -24,7 +24,9 @@ import com.google.bigtable.v2.CheckAndMutateRowRequest; import com.google.bigtable.v2.Column; import com.google.bigtable.v2.Family; +import com.google.bigtable.v2.Idempotency; import com.google.bigtable.v2.MutateRowRequest; +import com.google.bigtable.v2.MutateRowsRequest; import com.google.bigtable.v2.Mutation; import com.google.bigtable.v2.ReadModifyWriteRowRequest; import com.google.bigtable.v2.ReadModifyWriteRule; @@ -34,6 +36,8 @@ import com.google.bigtable.v2.SessionCheckAndMutateRowRequest; import com.google.bigtable.v2.SessionCheckAndMutateRowResponse; import com.google.bigtable.v2.SessionMutateRowRequest; +import com.google.bigtable.v2.SessionMutateRowsRequest; +import com.google.bigtable.v2.SessionMutateRowsResponse; import com.google.bigtable.v2.SessionReadModifyWriteRowRequest; import com.google.bigtable.v2.SessionReadModifyWriteRowResponse; import com.google.bigtable.v2.SessionReadRowRequest; @@ -42,6 +46,7 @@ import com.google.bigtable.v2.TableResponse; import com.google.protobuf.ByteString; import com.google.protobuf.InvalidProtocolBufferException; +import com.google.rpc.Status; import org.junit.jupiter.api.Test; class VRpcDescriptorTest { @@ -302,4 +307,104 @@ private static SessionReadModifyWriteRowRequest sampleReadModifyWriteRowRequest( .setIncrementAmount(1)) .build(); } + + @Test + void testMutateRows() throws InvalidProtocolBufferException { + assertThat(VRpcDescriptor.MUTATE_ROWS.getSessionDescriptor()) + .isEqualTo(VRpcDescriptor.TABLE_SESSION); + + TableResponse tableResp = + TableResponse.newBuilder() + .setMutateRows( + SessionMutateRowsResponse.newBuilder() + .addEntries( + SessionMutateRowsResponse.Entry.newBuilder() + .setIndex(0) + .setStatus(Status.newBuilder().setCode(0).build()))) + .build(); + assertThat(VRpcDescriptor.MUTATE_ROWS.decode(tableResp.toByteString())) + .isEqualTo(tableResp.getMutateRows()); + + SessionMutateRowsRequest req = sampleMutateRowsRequest(); + assertThat(TableRequest.parseFrom(VRpcDescriptor.MUTATE_ROWS.encode(req))) + .isEqualTo(TableRequest.newBuilder().setMutateRows(req).build()); + } + + @Test + void testMutateRowsAuthView() throws InvalidProtocolBufferException { + assertThat(VRpcDescriptor.MUTATE_ROWS_AUTH_VIEW.getSessionDescriptor()) + .isEqualTo(VRpcDescriptor.AUTHORIZED_VIEW_SESSION); + + AuthorizedViewResponse authViewResp = + AuthorizedViewResponse.newBuilder() + .setMutateRows( + SessionMutateRowsResponse.newBuilder() + .addEntries( + SessionMutateRowsResponse.Entry.newBuilder() + .setIndex(0) + .setStatus(Status.newBuilder().setCode(0).build()))) + .build(); + assertThat(VRpcDescriptor.MUTATE_ROWS_AUTH_VIEW.decode(authViewResp.toByteString())) + .isEqualTo(authViewResp.getMutateRows()); + + SessionMutateRowsRequest req = sampleMutateRowsRequest(); + assertThat(AuthorizedViewRequest.parseFrom(VRpcDescriptor.MUTATE_ROWS_AUTH_VIEW.encode(req))) + .isEqualTo(AuthorizedViewRequest.newBuilder().setMutateRows(req).build()); + } + + @Test + void testToLegacyProtoMutateRows() { + SessionMutateRowsRequest req = sampleMutateRowsRequest(); + MutateRowsRequest legacyReq = + (MutateRowsRequest) VRpcDescriptor.MUTATE_ROWS.toLegacyProto("table1", "app1", req); + assertThat(legacyReq.getTableName()).isEqualTo("table1"); + assertThat(legacyReq.getAuthorizedViewName()).isEmpty(); + assertThat(legacyReq.getAppProfileId()).isEqualTo("app1"); + assertThat(legacyReq.getEntriesCount()).isEqualTo(1); + assertThat(legacyReq.getEntries(0).getRowKey()).isEqualTo(req.getEntries(0).getKey()); + assertThat(legacyReq.getEntries(0).getMutationsList()) + .isEqualTo(req.getEntries(0).getMutationsList()); + } + + @Test + void testToLegacyProtoMutateRowsAuthView() { + SessionMutateRowsRequest req = sampleMutateRowsRequest(); + MutateRowsRequest legacyReq = + (MutateRowsRequest) + VRpcDescriptor.MUTATE_ROWS_AUTH_VIEW.toLegacyProto("view1", "app1", req); + assertThat(legacyReq.getAuthorizedViewName()).isEqualTo("view1"); + assertThat(legacyReq.getTableName()).isEmpty(); + assertThat(legacyReq.getAppProfileId()).isEqualTo("app1"); + assertThat(legacyReq.getEntriesCount()).isEqualTo(1); + assertThat(legacyReq.getEntries(0).getRowKey()).isEqualTo(req.getEntries(0).getKey()); + assertThat(legacyReq.getEntries(0).getMutationsList()) + .isEqualTo(req.getEntries(0).getMutationsList()); + } + + @Test + void testToLegacyProtoMutateRowsWithIdempotency() { + Idempotency idempotency = + Idempotency.newBuilder().setToken(ByteString.copyFromUtf8("token123")).build(); + SessionMutateRowsRequest req = + SessionMutateRowsRequest.newBuilder() + .addEntries( + SessionMutateRowsRequest.Entry.newBuilder() + .setKey(ByteString.copyFromUtf8("rowkey1")) + .addMutations(sampleSetCell()) + .setIdempotency(idempotency)) + .build(); + MutateRowsRequest legacyReq = + (MutateRowsRequest) VRpcDescriptor.MUTATE_ROWS.toLegacyProto("table1", "app1", req); + assertThat(legacyReq.getEntries(0).hasIdempotency()).isTrue(); + assertThat(legacyReq.getEntries(0).getIdempotency()).isEqualTo(idempotency); + } + + private static SessionMutateRowsRequest sampleMutateRowsRequest() { + return SessionMutateRowsRequest.newBuilder() + .addEntries( + SessionMutateRowsRequest.Entry.newBuilder() + .setKey(ByteString.copyFromUtf8("rowkey1")) + .addMutations(sampleSetCell())) + .build(); + } } diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/util/SessionScopeTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/util/SessionScopeTest.java index 190a03e59528..9bf93e3f4dab 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/util/SessionScopeTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/util/SessionScopeTest.java @@ -40,6 +40,7 @@ void noDiversionConfig_nonFallbackScopesIgnoreFlatSessionLoad() { assertThat(SessionScope.CHECK_AND_MUTATE_ROW.resolveSessionLoad(config)).isEqualTo(0f); assertThat(SessionScope.READ_MODIFY_WRITE_ROW.resolveSessionLoad(config)).isEqualTo(0f); + assertThat(SessionScope.MUTATE_ROWS.resolveSessionLoad(config)).isEqualTo(0f); } @Test @@ -94,7 +95,7 @@ void globalDiversionAppliesToEveryScope() { void unservedScope_doesNotCountAsDivertingTraffic() { SessionClientConfiguration config = SessionClientConfiguration.newBuilder() - .setSessionDiversionConfiguration(perScope("method:Bigtable.MutateRows", 1.0f)) + .setSessionDiversionConfiguration(perScope("method:Bigtable.ExecuteQuery", 1.0f)) .build(); for (SessionScope scope : SessionScope.values()) { diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/BulkMutationTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/BulkMutationTest.java index 498e3ee034c5..7c8591691790 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/BulkMutationTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/BulkMutationTest.java @@ -17,7 +17,9 @@ import static com.google.common.truth.Truth.assertThat; +import com.google.bigtable.v2.Idempotency; import com.google.bigtable.v2.MutateRowsRequest; +import com.google.bigtable.v2.SessionMutateRowsRequest; import com.google.cloud.bigtable.data.v2.internal.NameUtil; import com.google.cloud.bigtable.data.v2.internal.RequestContext; import com.google.protobuf.ByteString; @@ -314,4 +316,54 @@ public void testManyMutations() { } bulkMutation.add("key", mutation); } + + @Test + public void testGetTargetId() { + BulkMutation mutation = BulkMutation.create(TABLE_ID); + assertThat(mutation.getTargetId()).isEqualTo(TABLE_ID); + + AuthorizedViewId viewId = AuthorizedViewId.of(TABLE_ID, AUTHORIZED_VIEW_ID); + mutation = BulkMutation.create(viewId); + assertThat(mutation.getTargetId()).isEqualTo(viewId); + } + + @Test + public void testToSessionProto() { + BulkMutation m = + BulkMutation.create(TABLE_ID) + .add( + "key-a", + Mutation.create() + .setCell("fake-family1", "fake-qualifier1", 1_000, "fake-value1") + .setCell("fake-family2", "fake-qualifier2", 2_000, "fake-value2")) + .add( + ByteString.copyFromUtf8("key-b"), + Mutation.create().setCell("fake-family3", "fake-qualifier3", 3_000, "fake-value3")); + + SessionMutateRowsRequest sessionReq = m.toSessionProto(); + assertThat(sessionReq.getEntriesCount()).isEqualTo(2); + assertThat(sessionReq.getEntries(0).getKey()).isEqualTo(ByteString.copyFromUtf8("key-a")); + assertThat(sessionReq.getEntries(0).getMutationsCount()).isEqualTo(2); + assertThat(sessionReq.getEntries(1).getKey()).isEqualTo(ByteString.copyFromUtf8("key-b")); + assertThat(sessionReq.getEntries(1).getMutationsCount()).isEqualTo(1); + } + + @Test + public void testToSessionProtoWithIdempotency() { + Idempotency idempotency = + Idempotency.newBuilder().setToken(ByteString.copyFromUtf8("token123")).build(); + MutateRowsRequest proto = + MutateRowsRequest.newBuilder() + .setTableName("projects/p/instances/i/tables/t") + .addEntries( + MutateRowsRequest.Entry.newBuilder() + .setRowKey(ByteString.copyFromUtf8("key-a")) + .setIdempotency(idempotency)) + .build(); + BulkMutation m = BulkMutation.fromProto(proto); + SessionMutateRowsRequest sessionReq = m.toSessionProto(); + assertThat(sessionReq.getEntriesCount()).isEqualTo(1); + assertThat(sessionReq.getEntries(0).hasIdempotency()).isTrue(); + assertThat(sessionReq.getEntries(0).getIdempotency()).isEqualTo(idempotency); + } }