Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -130,6 +133,13 @@ public CompletableFuture<SessionReadModifyWriteRowResponse> readModifyWriteRow(
return f;
}

public CompletableFuture<SessionMutateRowsResponse> mutateRows(
SessionMutateRowsRequest req, Deadline deadline) {
UnaryResponseFuture<SessionMutateRowsResponse> f = new UnaryResponseFuture<>();
base.mutateRows(req, f, deadline);
return f;
}

@Override
public void close() {
this.base.close();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,7 @@ public static MaterializedViewAsync createAndStart(
/* mutateRowDescriptor= */ null, // Materialized views are read-only
/* checkAndMutateRowDescriptor= */ null,
/* readModifyWriteRowDescriptor= */ null,
/* mutateRowsDescriptor= */ null,
featureFlags,
clientInfo,
configManager,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -139,4 +142,11 @@ public CompletableFuture<SessionReadModifyWriteRowResponse> readModifyWriteRow(
base.readModifyWriteRow(req, f, deadline);
return f;
}

public CompletableFuture<SessionMutateRowsResponse> mutateRows(
SessionMutateRowsRequest req, Deadline deadline) {
UnaryResponseFuture<SessionMutateRowsResponse> f = new UnaryResponseFuture<>();
base.mutateRows(req, f, deadline);
return f;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -60,6 +66,8 @@ class TableBase implements AutoCloseable {
private final VRpcDescriptor<
?, SessionReadModifyWriteRowRequest, SessionReadModifyWriteRowResponse>
readModifyWriteRowDescriptor;
private final VRpcDescriptor<?, SessionMutateRowsRequest, SessionMutateRowsResponse>
mutateRowsDescriptor;

static <ReqT extends Message> TableBase createAndStart(
ReqT openReq,
Expand All @@ -70,6 +78,7 @@ static <ReqT extends Message> TableBase createAndStart(
checkAndMutateRowDescriptor,
VRpcDescriptor<?, SessionReadModifyWriteRowRequest, SessionReadModifyWriteRowResponse>
readModifyWriteRowDescriptor,
VRpcDescriptor<?, SessionMutateRowsRequest, SessionMutateRowsResponse> mutateRowsDescriptor,
FeatureFlags featureFlags,
ClientInfo clientInfo,
ClientConfigurationManager configManager,
Expand Down Expand Up @@ -108,6 +117,7 @@ static <ReqT extends Message> TableBase createAndStart(
mutateRowDescriptor,
checkAndMutateRowDescriptor,
readModifyWriteRowDescriptor,
mutateRowsDescriptor,
metrics,
timer,
userCallbackExecutor);
Expand All @@ -122,6 +132,7 @@ static <ReqT extends Message> TableBase createAndStart(
checkAndMutateRowDescriptor,
VRpcDescriptor<?, SessionReadModifyWriteRowRequest, SessionReadModifyWriteRowResponse>
readModifyWriteRowDescriptor,
VRpcDescriptor<?, SessionMutateRowsRequest, SessionMutateRowsResponse> mutateRowsDescriptor,
Metrics metrics,
BigtableTimer timer,
Executor userCallbackExecutor) {
Expand All @@ -130,6 +141,7 @@ static <ReqT extends Message> TableBase createAndStart(
this.mutateRowDescriptor = mutateRowDescriptor;
this.checkAndMutateRowDescriptor = checkAndMutateRowDescriptor;
this.readModifyWriteRowDescriptor = readModifyWriteRowDescriptor;
this.mutateRowsDescriptor = mutateRowsDescriptor;
this.metrics = metrics;
this.timer = timer;
this.userCallbackExecutor = userCallbackExecutor;
Expand Down Expand Up @@ -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<SessionMutateRowsResponse> listener,
Deadline deadline) {
RetryingVRpc<SessionMutateRowsRequest, SessionMutateRowsResponse> retry =
new RetryingVRpc<>(() -> sessionPool.newCall(mutateRowsDescriptor), timer);
ImmutableList<Mutation> 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);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -55,4 +56,10 @@ public <RowT> UnaryCallable<ReadModifyWriteRow, RowT> decorateReadModifyWriteRow
UnaryCallSettings<?, ?> settings) {
return classic;
}

@Override
public UnaryCallable<BulkMutation, Void> decorateMutateRows(
UnaryCallable<BulkMutation, Void> classic, UnaryCallSettings<?, ?> settings) {
return classic;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -42,4 +43,7 @@ <RowT> UnaryCallable<ReadModifyWriteRow, RowT> decorateReadModifyWriteRow(
UnaryCallable<ReadModifyWriteRow, RowT> classic,
RowAdapter<RowT> rowAdapter,
UnaryCallSettings<?, ?> settings);

UnaryCallable<BulkMutation, Void> decorateMutateRows(
UnaryCallable<BulkMutation, Void> classic, UnaryCallSettings<?, ?> settings);
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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);
}

/**
Expand Down Expand Up @@ -435,4 +439,16 @@ public <RowT> UnaryCallable<ReadModifyWriteRow, RowT> decorateReadModifyWriteRow
debugTagTracer,
SessionScope.READ_MODIFY_WRITE_ROW);
}

@Override
public UnaryCallable<BulkMutation, Void> decorateMutateRows(
UnaryCallable<BulkMutation, Void> classic, UnaryCallSettings<?, ?> settings) {
return new DivertingUnaryCallable<>(
configManager,
classic,
mutateRowsShim,
Util.extractTimeout(settings),
debugTagTracer,
SessionScope.MUTATE_ROWS);
}
}
Original file line number Diff line number Diff line change
@@ -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;
Comment on lines +23 to +24

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Import com.google.common.base.Preconditions to support defensive null-checking of parameters in the constructor.

Suggested change
import com.google.cloud.bigtable.data.v2.models.TargetId;
import io.grpc.Deadline;
import com.google.cloud.bigtable.data.v2.models.TargetId;
import com.google.common.base.Preconditions;
import io.grpc.Deadline;

import java.util.concurrent.CompletableFuture;

public class MutateRowsShim implements UnaryShim<BulkMutation, Void> {

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<Void> 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<Void> f = new CompletableFuture<>();
f.completeExceptionally(
new UnsupportedOperationException("Unsupported targetId type: " + targetId));
return f;
}
}
Loading
Loading