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 d69ec29e9256..0ff9eae331c3 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.SessionReadRowRequest; import com.google.bigtable.v2.SessionReadRowResponse; import com.google.cloud.bigtable.data.v2.internal.channels.ChannelPool; @@ -77,6 +79,7 @@ static AuthorizedViewAsync createAndStart( VRpcDescriptor.READ_ROW_AUTH_VIEW, VRpcDescriptor.MUTATE_ROW_AUTH_VIEW, VRpcDescriptor.CHECK_AND_MUTATE_ROW_AUTH_VIEW, + VRpcDescriptor.MUTATE_ROWS_AUTH_VIEW, featureFlags, clientInfo, configManager, @@ -120,6 +123,13 @@ public CompletableFuture checkAndMutateRow( 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 3a4537ab03a6..f4205911791e 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 @@ -68,7 +68,9 @@ public static MaterializedViewAsync createAndStart( openReq, VRpcDescriptor.MATERIALIZED_VIEW_SESSION, VRpcDescriptor.READ_ROW_MAT_VIEW, - // Materialized views are read-only, so mutateRow and checkAndMutateRow don't apply. + // Materialized views are read-only, so mutateRow, checkAndMutateRow, and mutateRows + // don't apply. + null, null, null, featureFlags, 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 adc4db3c793a..fe22b987d231 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.SessionReadRowRequest; import com.google.bigtable.v2.SessionReadRowResponse; import com.google.cloud.bigtable.data.v2.internal.channels.ChannelPool; @@ -75,6 +77,7 @@ public static TableAsync createAndStart( VRpcDescriptor.READ_ROW, VRpcDescriptor.MUTATE_ROW, VRpcDescriptor.CHECK_AND_MUTATE_ROW, + VRpcDescriptor.MUTATE_ROWS, featureFlags, clientInfo, configManager, @@ -132,4 +135,12 @@ public CompletableFuture checkAndMutateRow( base.checkAndMutateRow(req, f, deadline); return f; } + + // TODO: get deadline from compatibility layer + 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 c3696a51bee6..59423b3b0368 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 @@ -21,6 +21,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.SessionReadRowRequest; import com.google.bigtable.v2.SessionReadRowResponse; import com.google.cloud.bigtable.data.v2.internal.channels.ChannelPool; @@ -55,6 +57,8 @@ class TableBase implements AutoCloseable { mutateRowDescriptor; private final VRpcDescriptor checkAndMutateRowDescriptor; + private final VRpcDescriptor + mutateRowsDescriptor; static TableBase createAndStart( ReqT openReq, @@ -63,6 +67,7 @@ static TableBase createAndStart( VRpcDescriptor mutateRowDescriptor, VRpcDescriptor checkAndMutateRowDescriptor, + VRpcDescriptor mutateRowsDescriptor, FeatureFlags featureFlags, ClientInfo clientInfo, ClientConfigurationManager configManager, @@ -100,6 +105,7 @@ static TableBase createAndStart( readRowDescriptor, mutateRowDescriptor, checkAndMutateRowDescriptor, + mutateRowsDescriptor, metrics, timer, userCallbackExecutor); @@ -112,6 +118,7 @@ static TableBase createAndStart( VRpcDescriptor mutateRowDescriptor, VRpcDescriptor checkAndMutateRowDescriptor, + VRpcDescriptor mutateRowsDescriptor, Metrics metrics, BigtableTimer timer, Executor userCallbackExecutor) { @@ -119,6 +126,7 @@ static TableBase createAndStart( this.readRowDescriptor = readRowDescriptor; this.mutateRowDescriptor = mutateRowDescriptor; this.checkAndMutateRowDescriptor = checkAndMutateRowDescriptor; + this.mutateRowsDescriptor = mutateRowsDescriptor; this.metrics = metrics; this.timer = timer; this.userCallbackExecutor = userCallbackExecutor; @@ -174,4 +182,19 @@ public void checkAndMutateRow( 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); + boolean idempotent = Util.isIdempotent(req); + 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/api/Util.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Util.java index fba8bd0b8710..2cb7b09c436b 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Util.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Util.java @@ -18,6 +18,7 @@ import com.google.bigtable.v2.FeatureFlags; import com.google.bigtable.v2.Mutation; +import com.google.bigtable.v2.SessionMutateRowsRequest; import io.grpc.Metadata; import java.io.ByteArrayOutputStream; import java.io.IOException; @@ -76,6 +77,15 @@ public static boolean isIdempotent(List mutations) { return true; } + public static boolean isIdempotent(SessionMutateRowsRequest req) { + for (SessionMutateRowsRequest.Entry entry : req.getEntriesList()) { + if (!isIdempotent(entry.getMutationsList())) { + return false; + } + } + return true; + } + private static String urlEncode(String s) { try { return URLEncoder.encode(s, StandardCharsets.UTF_8.name()); 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 9b03ed9e0763..ae4f8e0abd18 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.RowAdapter; @@ -46,4 +47,10 @@ public UnaryCallable decorateCheckAndMutateRow( UnaryCallable classic, 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 e73168f452bb..28a0f29e3a38 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.RowAdapter; @@ -36,4 +37,7 @@ UnaryCallable decorateMutateRow( UnaryCallable decorateCheckAndMutateRow( UnaryCallable classic, 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 d5b96b291d60..a850ee3fd334 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 @@ -37,6 +37,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.ReadRowShim; import com.google.cloud.bigtable.data.v2.internal.compat.ops.ReadRowShimInner; import com.google.cloud.bigtable.data.v2.internal.csm.Metrics; @@ -45,6 +46,7 @@ import com.google.cloud.bigtable.data.v2.internal.csm.tracers.DirectPathCompatibleTracer; 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.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.RowAdapter; @@ -89,6 +91,7 @@ public class ShimImpl implements Shim { private final ReadRowShimInner readRowShimInner; private final MutateRowShim mutateRowShim; private final CheckAndMutateRowShim checkAndMutateRowShim; + private final MutateRowsShim mutateRowsShim; public static Shim create( ClientInfo clientInfo, @@ -205,6 +208,7 @@ public ShimImpl( this.readRowShimInner = new ReadRowShimInner(client); this.mutateRowShim = new MutateRowShim(client); this.checkAndMutateRowShim = new CheckAndMutateRowShim(client); + this.mutateRowsShim = new MutateRowsShim(client); } /** @@ -390,4 +394,11 @@ public UnaryCallable decorateCheckAndMutateRow( Util.extractTimeout(settings), debugTagTracer); } + + @Override + public UnaryCallable decorateMutateRows( + UnaryCallable classic, UnaryCallSettings settings) { + return new DivertingUnaryCallable<>( + configManager, classic, mutateRowsShim, Util.extractTimeout(settings), debugTagTracer); + } } 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..093ff4e294e0 --- /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,82 @@ +/* + * 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.OpenAuthorizedViewRequest; +import com.google.bigtable.v2.OpenTableRequest.Permission; +import com.google.bigtable.v2.SessionMutateRowsRequest; +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.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.io.IOException; +import java.util.concurrent.CompletableFuture; + +public class MutateRowsShim implements UnaryShim { + + private final SessionPoolMap tables; + private final SessionPoolMap authViews; + + public MutateRowsShim(Client client) { + tables = + new SessionPoolMap<>( + k -> client.openTableAsync(k.getTableId(), Permission.PERMISSION_WRITE)); + authViews = + new SessionPoolMap<>( + k -> + client.openAuthorizedViewAsync( + k.getTableId(), + k.getAuthorizedViewId(), + OpenAuthorizedViewRequest.Permission.PERMISSION_WRITE)); + } + + @Override + public void close() throws IOException { + tables.invalidateAll(); + authViews.invalidateAll(); + } + + @Override + public boolean supports(BulkMutation request) { + // TODO: enable when server side changes is rolled out + return false; + } + + @Override + public CompletableFuture call(BulkMutation request, Deadline deadline) { + TargetId targetId = request.getTargetId(); + + SessionMutateRowsRequest innerReq = request.toSessionProto(); + + if (targetId instanceof TableId) { + return tables.apply( + (TableId) targetId, t -> t.mutateRows(innerReq, deadline).thenApply(r -> null)); + } + if (targetId instanceof AuthorizedViewId) { + return authViews.apply( + (AuthorizedViewId) targetId, 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 8c668e6e9810..5bdfadb83b05 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; @@ -32,6 +33,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.SessionReadRowRequest; import com.google.bigtable.v2.SessionReadRowResponse; import com.google.bigtable.v2.SessionRequest; @@ -170,6 +173,30 @@ public final class VRpcDescriptor { return builder.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 = @@ -227,6 +254,32 @@ public final class VRpcDescriptor { return builder.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/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 b18f8572b91b..6984b05d5248 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,15 @@ 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); 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 41f526c8a0f2..6be7d76a83fb 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 @@ -24,6 +24,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.cloud.bigtable.data.v2.internal.csm.Metrics; import com.google.cloud.bigtable.data.v2.internal.csm.NoopMetrics; import com.google.cloud.bigtable.data.v2.internal.csm.attributes.ClientInfo; @@ -77,6 +78,7 @@ public void setup() { VRpcDescriptor.READ_ROW, VRpcDescriptor.MUTATE_ROW, VRpcDescriptor.CHECK_AND_MUTATE_ROW, + VRpcDescriptor.MUTATE_ROWS, noopMetrics, mockTimer, MoreExecutors.directExecutor()); @@ -185,6 +187,38 @@ public void testCheckAndMutateRowNotIdempotent() { 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/session/VRpcDescriptorTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/VRpcDescriptorTest.java index d3a3bb5c054f..225f21bb86bf 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.ReadRowsRequest; import com.google.bigtable.v2.Row; @@ -32,12 +34,15 @@ 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.SessionReadRowRequest; import com.google.bigtable.v2.SessionReadRowResponse; import com.google.bigtable.v2.TableRequest; 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 { @@ -216,4 +221,104 @@ private static Mutation sampleSetCell() { .setValue(ByteString.copyFromUtf8("v"))) .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/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); + } }