Skip to content

Commit e2111ee

Browse files
alexoneillalexoneill
committed
feat(bigtable): support SessionMutateRows on the session client
1 parent ecd27c5 commit e2111ee

17 files changed

Lines changed: 585 additions & 3 deletions

File tree

‎java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/AuthorizedViewAsync.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@
2222
import com.google.bigtable.v2.SessionCheckAndMutateRowResponse;
2323
import com.google.bigtable.v2.SessionMutateRowRequest;
2424
import com.google.bigtable.v2.SessionMutateRowResponse;
25+
import com.google.bigtable.v2.SessionMutateRowsRequest;
26+
import com.google.bigtable.v2.SessionMutateRowsResponse;
2527
import com.google.bigtable.v2.SessionReadModifyWriteRowRequest;
2628
import com.google.bigtable.v2.SessionReadModifyWriteRowResponse;
2729
import com.google.bigtable.v2.SessionReadRowRequest;
@@ -80,6 +82,7 @@ static AuthorizedViewAsync createAndStart(
8082
VRpcDescriptor.MUTATE_ROW_AUTH_VIEW,
8183
VRpcDescriptor.CHECK_AND_MUTATE_ROW_AUTH_VIEW,
8284
VRpcDescriptor.READ_MODIFY_WRITE_ROW_AUTH_VIEW,
85+
VRpcDescriptor.MUTATE_ROWS_AUTH_VIEW,
8386
featureFlags,
8487
clientInfo,
8588
configManager,
@@ -130,6 +133,13 @@ public CompletableFuture<SessionReadModifyWriteRowResponse> readModifyWriteRow(
130133
return f;
131134
}
132135

136+
public CompletableFuture<SessionMutateRowsResponse> mutateRows(
137+
SessionMutateRowsRequest req, Deadline deadline) {
138+
UnaryResponseFuture<SessionMutateRowsResponse> f = new UnaryResponseFuture<>();
139+
base.mutateRows(req, f, deadline);
140+
return f;
141+
}
142+
133143
@Override
134144
public void close() {
135145
this.base.close();

‎java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/MaterializedViewAsync.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,7 @@ public static MaterializedViewAsync createAndStart(
7171
/* mutateRowDescriptor= */ null, // Materialized views are read-only
7272
/* checkAndMutateRowDescriptor= */ null,
7373
/* readModifyWriteRowDescriptor= */ null,
74+
/* mutateRowsDescriptor= */ null,
7475
featureFlags,
7576
clientInfo,
7677
configManager,

‎java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableAsync.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,8 @@
2323
import com.google.bigtable.v2.SessionCheckAndMutateRowResponse;
2424
import com.google.bigtable.v2.SessionMutateRowRequest;
2525
import com.google.bigtable.v2.SessionMutateRowResponse;
26+
import com.google.bigtable.v2.SessionMutateRowsRequest;
27+
import com.google.bigtable.v2.SessionMutateRowsResponse;
2628
import com.google.bigtable.v2.SessionReadModifyWriteRowRequest;
2729
import com.google.bigtable.v2.SessionReadModifyWriteRowResponse;
2830
import com.google.bigtable.v2.SessionReadRowRequest;
@@ -78,6 +80,7 @@ public static TableAsync createAndStart(
7880
VRpcDescriptor.MUTATE_ROW,
7981
VRpcDescriptor.CHECK_AND_MUTATE_ROW,
8082
VRpcDescriptor.READ_MODIFY_WRITE_ROW,
83+
VRpcDescriptor.MUTATE_ROWS,
8184
featureFlags,
8285
clientInfo,
8386
configManager,
@@ -139,4 +142,11 @@ public CompletableFuture<SessionReadModifyWriteRowResponse> readModifyWriteRow(
139142
base.readModifyWriteRow(req, f, deadline);
140143
return f;
141144
}
145+
146+
public CompletableFuture<SessionMutateRowsResponse> mutateRows(
147+
SessionMutateRowsRequest req, Deadline deadline) {
148+
UnaryResponseFuture<SessionMutateRowsResponse> f = new UnaryResponseFuture<>();
149+
base.mutateRows(req, f, deadline);
150+
return f;
151+
}
142152
}

‎java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableBase.java‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,12 +15,17 @@
1515
*/
1616
package com.google.cloud.bigtable.data.v2.internal.api;
1717

18+
import static com.google.common.collect.ImmutableList.toImmutableList;
19+
1820
import com.google.bigtable.v2.CloseSessionRequest;
1921
import com.google.bigtable.v2.FeatureFlags;
22+
import com.google.bigtable.v2.Mutation;
2023
import com.google.bigtable.v2.SessionCheckAndMutateRowRequest;
2124
import com.google.bigtable.v2.SessionCheckAndMutateRowResponse;
2225
import com.google.bigtable.v2.SessionMutateRowRequest;
2326
import com.google.bigtable.v2.SessionMutateRowResponse;
27+
import com.google.bigtable.v2.SessionMutateRowsRequest;
28+
import com.google.bigtable.v2.SessionMutateRowsResponse;
2429
import com.google.bigtable.v2.SessionReadModifyWriteRowRequest;
2530
import com.google.bigtable.v2.SessionReadModifyWriteRowResponse;
2631
import com.google.bigtable.v2.SessionReadRowRequest;
@@ -40,6 +45,7 @@
4045
import com.google.cloud.bigtable.data.v2.internal.session.VRpcDescriptor;
4146
import com.google.cloud.bigtable.data.v2.internal.util.ClientConfigurationManager;
4247
import com.google.common.annotations.VisibleForTesting;
48+
import com.google.common.collect.ImmutableList;
4349
import com.google.protobuf.Message;
4450
import io.grpc.CallOptions;
4551
import io.grpc.Context;
@@ -60,6 +66,8 @@ class TableBase implements AutoCloseable {
6066
private final VRpcDescriptor<
6167
?, SessionReadModifyWriteRowRequest, SessionReadModifyWriteRowResponse>
6268
readModifyWriteRowDescriptor;
69+
private final VRpcDescriptor<?, SessionMutateRowsRequest, SessionMutateRowsResponse>
70+
mutateRowsDescriptor;
6371

6472
static <ReqT extends Message> TableBase createAndStart(
6573
ReqT openReq,
@@ -70,6 +78,7 @@ static <ReqT extends Message> TableBase createAndStart(
7078
checkAndMutateRowDescriptor,
7179
VRpcDescriptor<?, SessionReadModifyWriteRowRequest, SessionReadModifyWriteRowResponse>
7280
readModifyWriteRowDescriptor,
81+
VRpcDescriptor<?, SessionMutateRowsRequest, SessionMutateRowsResponse> mutateRowsDescriptor,
7382
FeatureFlags featureFlags,
7483
ClientInfo clientInfo,
7584
ClientConfigurationManager configManager,
@@ -108,6 +117,7 @@ static <ReqT extends Message> TableBase createAndStart(
108117
mutateRowDescriptor,
109118
checkAndMutateRowDescriptor,
110119
readModifyWriteRowDescriptor,
120+
mutateRowsDescriptor,
111121
metrics,
112122
timer,
113123
userCallbackExecutor);
@@ -122,6 +132,7 @@ static <ReqT extends Message> TableBase createAndStart(
122132
checkAndMutateRowDescriptor,
123133
VRpcDescriptor<?, SessionReadModifyWriteRowRequest, SessionReadModifyWriteRowResponse>
124134
readModifyWriteRowDescriptor,
135+
VRpcDescriptor<?, SessionMutateRowsRequest, SessionMutateRowsResponse> mutateRowsDescriptor,
125136
Metrics metrics,
126137
BigtableTimer timer,
127138
Executor userCallbackExecutor) {
@@ -130,6 +141,7 @@ static <ReqT extends Message> TableBase createAndStart(
130141
this.mutateRowDescriptor = mutateRowDescriptor;
131142
this.checkAndMutateRowDescriptor = checkAndMutateRowDescriptor;
132143
this.readModifyWriteRowDescriptor = readModifyWriteRowDescriptor;
144+
this.mutateRowsDescriptor = mutateRowsDescriptor;
133145
this.metrics = metrics;
134146
this.timer = timer;
135147
this.userCallbackExecutor = userCallbackExecutor;
@@ -202,4 +214,23 @@ public void readModifyWriteRow(
202214
new VOperationImpl<>(retry, Context.current(), userCallbackExecutor, tracer, deadline, false)
203215
.start(req, listener);
204216
}
217+
218+
public void mutateRows(
219+
SessionMutateRowsRequest req,
220+
VRpcListener<SessionMutateRowsResponse> listener,
221+
Deadline deadline) {
222+
RetryingVRpc<SessionMutateRowsRequest, SessionMutateRowsResponse> retry =
223+
new RetryingVRpc<>(() -> sessionPool.newCall(mutateRowsDescriptor), timer);
224+
ImmutableList<Mutation> mutations =
225+
req.getEntriesList().stream()
226+
.flatMap(entry -> entry.getMutationsList().stream())
227+
.collect(toImmutableList());
228+
boolean idempotent = Util.isIdempotent(mutations);
229+
VRpcTracer tracer =
230+
metrics.newTableTracer(sessionPool.getInfo(), mutateRowsDescriptor, deadline);
231+
232+
new VOperationImpl<>(
233+
retry, Context.current(), userCallbackExecutor, tracer, deadline, idempotent)
234+
.start(req, listener);
235+
}
205236
}

‎java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/DisabledShim.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
import com.google.api.gax.rpc.UnaryCallSettings;
1919
import com.google.api.gax.rpc.UnaryCallable;
20+
import com.google.cloud.bigtable.data.v2.models.BulkMutation;
2021
import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation;
2122
import com.google.cloud.bigtable.data.v2.models.Query;
2223
import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow;
@@ -55,4 +56,10 @@ public <RowT> UnaryCallable<ReadModifyWriteRow, RowT> decorateReadModifyWriteRow
5556
UnaryCallSettings<?, ?> settings) {
5657
return classic;
5758
}
59+
60+
@Override
61+
public UnaryCallable<BulkMutation, Void> decorateMutateRows(
62+
UnaryCallable<BulkMutation, Void> classic, UnaryCallSettings<?, ?> settings) {
63+
return classic;
64+
}
5865
}

‎java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/Shim.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
import com.google.api.gax.rpc.UnaryCallSettings;
1919
import com.google.api.gax.rpc.UnaryCallable;
20+
import com.google.cloud.bigtable.data.v2.models.BulkMutation;
2021
import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation;
2122
import com.google.cloud.bigtable.data.v2.models.Query;
2223
import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow;
@@ -42,4 +43,7 @@ <RowT> UnaryCallable<ReadModifyWriteRow, RowT> decorateReadModifyWriteRow(
4243
UnaryCallable<ReadModifyWriteRow, RowT> classic,
4344
RowAdapter<RowT> rowAdapter,
4445
UnaryCallSettings<?, ?> settings);
46+
47+
UnaryCallable<BulkMutation, Void> decorateMutateRows(
48+
UnaryCallable<BulkMutation, Void> classic, UnaryCallSettings<?, ?> settings);
4549
}

‎java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@
4040
import com.google.cloud.bigtable.data.v2.internal.compat.ops.CheckAndMutateRowShim;
4141
import com.google.cloud.bigtable.data.v2.internal.compat.ops.DivertingUnaryCallable;
4242
import com.google.cloud.bigtable.data.v2.internal.compat.ops.MutateRowShim;
43+
import com.google.cloud.bigtable.data.v2.internal.compat.ops.MutateRowsShim;
4344
import com.google.cloud.bigtable.data.v2.internal.compat.ops.ReadModifyWriteRowShim;
4445
import com.google.cloud.bigtable.data.v2.internal.compat.ops.ReadRowShim;
4546
import com.google.cloud.bigtable.data.v2.internal.compat.ops.RowBuilderShim;
@@ -51,6 +52,7 @@
5152
import com.google.cloud.bigtable.data.v2.internal.dp.DirectAccessInvestigator;
5253
import com.google.cloud.bigtable.data.v2.internal.util.ClientConfigurationManager;
5354
import com.google.cloud.bigtable.data.v2.internal.util.SessionScope;
55+
import com.google.cloud.bigtable.data.v2.models.BulkMutation;
5456
import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation;
5557
import com.google.cloud.bigtable.data.v2.models.Query;
5658
import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow;
@@ -110,6 +112,7 @@ public class ShimImpl implements Shim {
110112
private final MutateRowShim mutateRowShim;
111113
private final CheckAndMutateRowShim checkAndMutateRowShim;
112114
private final ReadModifyWriteRowShim readModifyWriteRowShim;
115+
private final MutateRowsShim mutateRowsShim;
113116

114117
public static Shim create(
115118
ClientInfo clientInfo,
@@ -228,6 +231,7 @@ public ShimImpl(
228231
this.mutateRowShim = new MutateRowShim(sessionMap);
229232
this.checkAndMutateRowShim = new CheckAndMutateRowShim(sessionMap);
230233
this.readModifyWriteRowShim = new ReadModifyWriteRowShim(sessionMap);
234+
this.mutateRowsShim = new MutateRowsShim(sessionMap);
231235
}
232236

233237
/**
@@ -435,4 +439,16 @@ public <RowT> UnaryCallable<ReadModifyWriteRow, RowT> decorateReadModifyWriteRow
435439
debugTagTracer,
436440
SessionScope.READ_MODIFY_WRITE_ROW);
437441
}
442+
443+
@Override
444+
public UnaryCallable<BulkMutation, Void> decorateMutateRows(
445+
UnaryCallable<BulkMutation, Void> classic, UnaryCallSettings<?, ?> settings) {
446+
return new DivertingUnaryCallable<>(
447+
configManager,
448+
classic,
449+
mutateRowsShim,
450+
Util.extractTimeout(settings),
451+
debugTagTracer,
452+
SessionScope.MUTATE_ROWS);
453+
}
438454
}

‎java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/MutateRowsShim.java‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.cloud.bigtable.data.v2.internal.compat.ops;
17+
18+
import com.google.bigtable.v2.SessionMutateRowsRequest;
19+
import com.google.cloud.bigtable.data.v2.internal.session.SessionPool;
20+
import com.google.cloud.bigtable.data.v2.models.AuthorizedViewId;
21+
import com.google.cloud.bigtable.data.v2.models.BulkMutation;
22+
import com.google.cloud.bigtable.data.v2.models.TableId;
23+
import com.google.cloud.bigtable.data.v2.models.TargetId;
24+
import io.grpc.Deadline;
25+
import java.util.concurrent.CompletableFuture;
26+
27+
public class MutateRowsShim implements UnaryShim<BulkMutation, Void> {
28+
29+
private final SessionPoolMap sessionMap;
30+
31+
public MutateRowsShim(SessionPoolMap sessionMap) {
32+
this.sessionMap = sessionMap;
33+
}
34+
35+
@Override
36+
public void close() {}
37+
38+
@Override
39+
public boolean supports(BulkMutation request) {
40+
TargetId targetId = request.getTargetId();
41+
SessionPool<?> pool;
42+
if (targetId instanceof TableId) {
43+
pool = sessionMap.get((TableId) targetId, SessionKey.Permission.WRITE).getSessionPool();
44+
} else if (targetId instanceof AuthorizedViewId) {
45+
pool =
46+
sessionMap.get((AuthorizedViewId) targetId, SessionKey.Permission.WRITE).getSessionPool();
47+
} else {
48+
return false;
49+
}
50+
return UnaryShim.shouldRouteToSession(pool);
51+
}
52+
53+
@Override
54+
public CompletableFuture<Void> call(BulkMutation request, Deadline deadline) {
55+
TargetId targetId = request.getTargetId();
56+
57+
SessionMutateRowsRequest innerReq = request.toSessionProto();
58+
59+
if (targetId instanceof TableId) {
60+
return sessionMap.apply(
61+
(TableId) targetId,
62+
SessionKey.Permission.WRITE,
63+
t -> t.mutateRows(innerReq, deadline).thenApply(r -> null));
64+
}
65+
if (targetId instanceof AuthorizedViewId) {
66+
return sessionMap.apply(
67+
(AuthorizedViewId) targetId,
68+
SessionKey.Permission.WRITE,
69+
v -> v.mutateRows(innerReq, deadline).thenApply(r -> null));
70+
}
71+
72+
CompletableFuture<Void> f = new CompletableFuture<>();
73+
f.completeExceptionally(
74+
new UnsupportedOperationException("Unsupported targetId type: " + targetId));
75+
return f;
76+
}
77+
}

0 commit comments

Comments
 (0)