Skip to content
This repository was archived by the owner on May 8, 2026. It is now read-only.

Commit e0b0d8a

Browse files
committed
feat: add BigtableTableAdminClientV2 to support Selective GAPIC
This commit introduces `BigtableTableAdminClientV2`, a new client class that extends the auto-generated `BaseBigtableTableAdminClient`. It relocates the manual wrappers for CUJs from the legacy `BigtableTableAdminClient`. b/502616786
1 parent 5c4a2c7 commit e0b0d8a

9 files changed

Lines changed: 364 additions & 87 deletions

File tree

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/BigtableTableAdminClient.java

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -184,10 +184,8 @@ public static BigtableTableAdminClient create(
184184
/** Constructs an instance of BigtableTableAdminClient with the given settings. */
185185
public static BigtableTableAdminClient create(@Nonnull BigtableTableAdminSettings settings)
186186
throws IOException {
187-
TableAdminRequestContext requestContext =
188-
TableAdminRequestContext.create(settings.getProjectId(), settings.getInstanceId());
189187
EnhancedBigtableTableAdminStub stub =
190-
EnhancedBigtableTableAdminStub.createEnhanced(settings.getStubSettings(), requestContext);
188+
EnhancedBigtableTableAdminStub.createEnhanced(settings.getStubSettings());
191189
return create(settings.getProjectId(), settings.getInstanceId(), stub);
192190
}
193191

@@ -1697,7 +1695,9 @@ public void waitForConsistency(String tableId, String consistencyToken) {
16971695
*/
16981696
public ApiFuture<Void> waitForConsistencyAsync(String tableId, String consistencyToken) {
16991697
return stub.awaitConsistencyCallable()
1700-
.futureCall(ConsistencyRequest.forReplication(tableId, consistencyToken));
1698+
.futureCall(
1699+
ConsistencyRequest.forReplication(
1700+
NameUtil.formatTableName(projectId, instanceId, tableId), consistencyToken));
17011701
}
17021702

17031703
/**
Lines changed: 176 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,176 @@
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.admin.v2;
17+
18+
import com.google.api.core.ApiFuture;
19+
import com.google.api.core.ApiFutures;
20+
import com.google.api.gax.rpc.ApiExceptions;
21+
import com.google.cloud.bigtable.admin.v2.models.ConsistencyRequest;
22+
import com.google.cloud.bigtable.admin.v2.models.OptimizeRestoredTableOperationToken;
23+
import com.google.cloud.bigtable.admin.v2.models.RestoredTableResult;
24+
import com.google.cloud.bigtable.admin.v2.stub.BigtableTableAdminStub;
25+
import com.google.cloud.bigtable.admin.v2.stub.EnhancedBigtableTableAdminStub;
26+
import com.google.common.base.Strings;
27+
import com.google.protobuf.Empty;
28+
import java.io.IOException;
29+
import java.util.concurrent.ExecutionException;
30+
import javax.annotation.Nonnull;
31+
32+
/**
33+
* Modern Cloud Bigtable Table Admin Client.
34+
*
35+
* <p>This client extends the auto-generated {@link BaseBigtableTableAdminClient} to provide manual
36+
* overrides and additional convenience methods for Critical User Journeys (CUJs) that the GAPIC
37+
* generator cannot handle natively (e.g., chained Long Running Operations, Consistency Polling).
38+
*/
39+
public class BigtableTableAdminClientV2 extends BaseBigtableTableAdminClient {
40+
41+
protected BigtableTableAdminClientV2(BaseBigtableTableAdminSettings settings) throws IOException {
42+
super(settings);
43+
}
44+
45+
protected BigtableTableAdminClientV2(BigtableTableAdminStub stub) {
46+
super(stub);
47+
}
48+
49+
/** Constructs an instance of BigtableTableAdminClientV2 with the given settings. */
50+
public static final BigtableTableAdminClientV2 createClient(BaseBigtableTableAdminSettings settings)
51+
throws IOException {
52+
// Explicitly create the enhanced stub
53+
EnhancedBigtableTableAdminStub stub =
54+
EnhancedBigtableTableAdminStub.createEnhanced(
55+
(com.google.cloud.bigtable.admin.v2.stub.BigtableTableAdminStubSettings) settings.getStubSettings());
56+
// Pass the enhanced stub to the existing stub-based constructor
57+
return new BigtableTableAdminClientV2(stub);
58+
}
59+
60+
/** Constructs an instance of BigtableTableAdminClientV2 with the given stub. */
61+
public static final BigtableTableAdminClientV2 createClient(BigtableTableAdminStub stub) {
62+
return new BigtableTableAdminClientV2(stub);
63+
}
64+
65+
/**
66+
* Awaits the completion of the "Optimize Restored Table" operation.
67+
*
68+
* <p>This method blocks until the restore operation is complete, extracts the optimization token,
69+
* and returns an ApiFuture for the optimization phase.
70+
*
71+
* @param restoreFuture The future returned by restoreTableAsync().
72+
* @return An ApiFuture that tracks the optimization progress.
73+
*/
74+
public ApiFuture<Empty> awaitOptimizeRestoredTable(ApiFuture<RestoredTableResult> restoreFuture) {
75+
// 1. Block and wait for the restore operation to complete
76+
RestoredTableResult result;
77+
try {
78+
result = restoreFuture.get();
79+
} catch (Exception e) {
80+
throw new RuntimeException("Restore operation failed", e);
81+
}
82+
83+
// 2. Extract the operation token from the result
84+
// (RestoredTableResult already wraps the OptimizeRestoredTableOperationToken)
85+
OptimizeRestoredTableOperationToken token = result.getOptimizeRestoredTableOperationToken();
86+
87+
if (token == null || Strings.isNullOrEmpty(token.getOperationName())) {
88+
// If there is no optimization operation, return immediate success.
89+
return ApiFutures.immediateFuture(Empty.getDefaultInstance());
90+
}
91+
92+
// 3. Return the future for the optimization operation
93+
return ((EnhancedBigtableTableAdminStub) getStub()).awaitOptimizeRestoredTableCallable().resumeFutureCall(token.getOperationName());
94+
}
95+
96+
/**
97+
* Awaits a restored table is fully optimized.
98+
*
99+
* <p>Sample code
100+
*
101+
* <pre>{@code
102+
* RestoredTableResult result =
103+
* client.restoreTable(RestoreTableRequest.of(clusterId, backupId).setTableId(tableId));
104+
* client.awaitOptimizeRestoredTable(result.getOptimizeRestoredTableOperationToken());
105+
* }</pre>
106+
*/
107+
public void awaitOptimizeRestoredTable(OptimizeRestoredTableOperationToken token)
108+
throws ExecutionException, InterruptedException {
109+
awaitOptimizeRestoredTableAsync(token).get();
110+
}
111+
112+
/**
113+
* Awaits a restored table is fully optimized asynchronously.
114+
*
115+
* <p>Sample code
116+
*
117+
* <pre>{@code
118+
* RestoredTableResult result =
119+
* client.restoreTable(RestoreTableRequest.of(clusterId, backupId).setTableId(tableId));
120+
* ApiFuture<Void> future = client.awaitOptimizeRestoredTableAsync(
121+
* result.getOptimizeRestoredTableOperationToken());
122+
*
123+
* ApiFutures.addCallback(
124+
* future,
125+
* new ApiFutureCallback<Void>() {
126+
* public void onSuccess(Void unused) {
127+
* System.out.println("The optimization of the restored table is done.");
128+
* }
129+
*
130+
* public void onFailure(Throwable t) {
131+
* t.printStackTrace();
132+
* }
133+
* },
134+
* MoreExecutors.directExecutor()
135+
* );
136+
* }</pre>
137+
*/
138+
public ApiFuture<Void> awaitOptimizeRestoredTableAsync(
139+
OptimizeRestoredTableOperationToken token) {
140+
ApiFuture<Empty> emptyFuture =
141+
((EnhancedBigtableTableAdminStub) getStub()).awaitOptimizeRestoredTableCallable().resumeFutureCall(token.getOperationName());
142+
return ApiFutures.transform(
143+
emptyFuture,
144+
new com.google.api.core.ApiFunction<Empty, Void>() {
145+
@Override
146+
public Void apply(Empty input) {
147+
return null;
148+
}
149+
},
150+
com.google.common.util.concurrent.MoreExecutors.directExecutor());
151+
}
152+
153+
/**
154+
* Polls an existing consistency token until table replication is consistent across all clusters.
155+
* Useful for checking consistency of a token generated in a separate process. Blocks until
156+
* completion.
157+
*
158+
* @param tableName The fully qualified table name to check.
159+
* @param consistencyToken The token to poll.
160+
*/
161+
public void waitForConsistency(String tableName, String consistencyToken) {
162+
ApiExceptions.callAndTranslateApiException(waitForConsistencyAsync(tableName, consistencyToken));
163+
}
164+
165+
/**
166+
* Asynchronously polls the consistency token. Returns a future that completes when table
167+
* replication is consistent across all clusters.
168+
*
169+
* @param tableName The fully qualified table name to check.
170+
* @param consistencyToken The token to poll.
171+
*/
172+
public ApiFuture<Void> waitForConsistencyAsync(String tableName, String consistencyToken) {
173+
return ((EnhancedBigtableTableAdminStub) getStub()).awaitConsistencyCallable()
174+
.futureCall(ConsistencyRequest.forReplication(tableName, consistencyToken));
175+
}
176+
}

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/models/ConsistencyRequest.java

Lines changed: 12 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@
3030
@AutoValue
3131
public abstract class ConsistencyRequest {
3232
@Nonnull
33-
protected abstract String getTableId();
33+
protected abstract String getTableName();
3434

3535
@Nonnull
3636
protected abstract CheckConsistencyRequest.ModeCase getMode();
@@ -43,53 +43,46 @@ public abstract class ConsistencyRequest {
4343
@Nullable
4444
public abstract String getConsistencyToken();
4545

46-
public static ConsistencyRequest forReplication(String tableId) {
46+
public static ConsistencyRequest forReplication(String tableName) {
4747
return new AutoValue_ConsistencyRequest(
48-
tableId, CheckConsistencyRequest.ModeCase.STANDARD_READ_REMOTE_WRITES, null);
48+
tableName, CheckConsistencyRequest.ModeCase.STANDARD_READ_REMOTE_WRITES, null);
4949
}
5050

5151
/**
5252
* Creates a request to check consistency using an existing token.
5353
*
54-
* @param tableId The table ID.
54+
* @param tableName The table name.
5555
* @param consistencyToken The token to check. Must not be null.
5656
* @throws NullPointerException if consistencyToken is null.
5757
*/
58-
public static ConsistencyRequest forReplication(String tableId, String consistencyToken) {
58+
public static ConsistencyRequest forReplication(String tableName, String consistencyToken) {
5959
Preconditions.checkNotNull(consistencyToken, "consistencyToken must not be null");
6060

6161
return new AutoValue_ConsistencyRequest(
62-
tableId, CheckConsistencyRequest.ModeCase.STANDARD_READ_REMOTE_WRITES, consistencyToken);
62+
tableName, CheckConsistencyRequest.ModeCase.STANDARD_READ_REMOTE_WRITES, consistencyToken);
6363
}
6464

65-
public static ConsistencyRequest forDataBoost(String tableId) {
65+
public static ConsistencyRequest forDataBoost(String tableName) {
6666
return new AutoValue_ConsistencyRequest(
67-
tableId, CheckConsistencyRequest.ModeCase.DATA_BOOST_READ_LOCAL_WRITES, null);
67+
tableName, CheckConsistencyRequest.ModeCase.DATA_BOOST_READ_LOCAL_WRITES, null);
6868
}
6969

7070
@InternalApi
71-
public CheckConsistencyRequest toCheckConsistencyProto(
72-
TableAdminRequestContext requestContext, String token) {
71+
public CheckConsistencyRequest toCheckConsistencyProto(String token) {
7372
CheckConsistencyRequest.Builder builder = CheckConsistencyRequest.newBuilder();
74-
TableName tableName =
75-
TableName.of(requestContext.getProjectId(), requestContext.getInstanceId(), getTableId());
7673

7774
if (getMode().equals(CheckConsistencyRequest.ModeCase.STANDARD_READ_REMOTE_WRITES)) {
7875
builder.setStandardReadRemoteWrites(StandardReadRemoteWrites.newBuilder().build());
7976
} else {
8077
builder.setDataBoostReadLocalWrites(DataBoostReadLocalWrites.newBuilder().build());
8178
}
8279

83-
return builder.setName(tableName.toString()).setConsistencyToken(token).build();
80+
return builder.setName(getTableName()).setConsistencyToken(token).build();
8481
}
8582

8683
@InternalApi
87-
public GenerateConsistencyTokenRequest toGenerateTokenProto(
88-
TableAdminRequestContext requestContext) {
84+
public GenerateConsistencyTokenRequest toGenerateTokenProto() {
8985
GenerateConsistencyTokenRequest.Builder builder = GenerateConsistencyTokenRequest.newBuilder();
90-
TableName tableName =
91-
TableName.of(requestContext.getProjectId(), requestContext.getInstanceId(), getTableId());
92-
93-
return builder.setName(tableName.toString()).build();
86+
return builder.setName(getTableName()).build();
9487
}
9588
}

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/stub/AwaitConsistencyCallable.java

Lines changed: 6 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -56,15 +56,12 @@ class AwaitConsistencyCallable extends UnaryCallable<ConsistencyRequest, Void> {
5656
private final UnaryCallable<CheckConsistencyRequest, CheckConsistencyResponse> checkCallable;
5757
private final RetryingExecutor<CheckConsistencyResponse> executor;
5858

59-
private final TableAdminRequestContext requestContext;
60-
6159
static AwaitConsistencyCallable create(
6260
UnaryCallable<GenerateConsistencyTokenRequest, GenerateConsistencyTokenResponse>
6361
generateCallable,
6462
UnaryCallable<CheckConsistencyRequest, CheckConsistencyResponse> checkCallable,
6563
ClientContext clientContext,
66-
RetrySettings pollingSettings,
67-
TableAdminRequestContext requestContext) {
64+
RetrySettings pollingSettings) {
6865

6966
RetryAlgorithm<CheckConsistencyResponse> retryAlgorithm =
7067
new RetryAlgorithm<>(
@@ -75,20 +72,18 @@ static AwaitConsistencyCallable create(
7572
new ScheduledRetryingExecutor<>(retryAlgorithm, clientContext.getExecutor());
7673

7774
return new AwaitConsistencyCallable(
78-
generateCallable, checkCallable, retryingExecutor, requestContext);
75+
generateCallable, checkCallable, retryingExecutor);
7976
}
8077

8178
@VisibleForTesting
8279
AwaitConsistencyCallable(
8380
UnaryCallable<GenerateConsistencyTokenRequest, GenerateConsistencyTokenResponse>
8481
generateCallable,
8582
UnaryCallable<CheckConsistencyRequest, CheckConsistencyResponse> checkCallable,
86-
RetryingExecutor<CheckConsistencyResponse> executor,
87-
TableAdminRequestContext requestContext) {
83+
RetryingExecutor<CheckConsistencyResponse> executor) {
8884
this.generateCallable = generateCallable;
8985
this.checkCallable = checkCallable;
9086
this.executor = executor;
91-
this.requestContext = requestContext;
9287
}
9388

9489
@Override
@@ -98,22 +93,20 @@ public ApiFuture<Void> futureCall(
9893
// If the token is already provided, skip generation and poll directly.
9994
if (consistencyRequest.getConsistencyToken() != null) {
10095
CheckConsistencyRequest request =
101-
consistencyRequest.toCheckConsistencyProto(
102-
requestContext, consistencyRequest.getConsistencyToken());
96+
consistencyRequest.toCheckConsistencyProto(consistencyRequest.getConsistencyToken());
10397
return pollToken(request, apiCallContext);
10498
}
10599

106100
ApiFuture<GenerateConsistencyTokenResponse> tokenFuture =
107-
generateToken(consistencyRequest.toGenerateTokenProto(requestContext), apiCallContext);
101+
generateToken(consistencyRequest.toGenerateTokenProto(), apiCallContext);
108102

109103
return ApiFutures.transformAsync(
110104
tokenFuture,
111105
new ApiAsyncFunction<GenerateConsistencyTokenResponse, Void>() {
112106
@Override
113107
public ApiFuture<Void> apply(GenerateConsistencyTokenResponse input) {
114108
CheckConsistencyRequest request =
115-
consistencyRequest.toCheckConsistencyProto(
116-
requestContext, input.getConsistencyToken());
109+
consistencyRequest.toCheckConsistencyProto(input.getConsistencyToken());
117110
return pollToken(request, apiCallContext);
118111
}
119112
},

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/stub/AwaitReplicationCallable.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ static AwaitReplicationCallable create(AwaitConsistencyCallable awaitConsistency
4242

4343
@Override
4444
public ApiFuture<Void> futureCall(final TableName tableName, final ApiCallContext context) {
45-
ConsistencyRequest consistencyRequest = ConsistencyRequest.forReplication(tableName.getTable());
45+
ConsistencyRequest consistencyRequest = ConsistencyRequest.forReplication(tableName.toString());
4646

4747
return awaitConsistencyCallable.futureCall(consistencyRequest, context);
4848
}

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/stub/EnhancedBigtableTableAdminStub.java

Lines changed: 4 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -54,31 +54,27 @@ public class EnhancedBigtableTableAdminStub extends GrpcBigtableTableAdminStub {
5454
private final BigtableTableAdminStubSettings settings;
5555
private final ClientContext clientContext;
5656

57-
private final TableAdminRequestContext requestContext;
58-
5957
@Deprecated private final AwaitReplicationCallable awaitReplicationCallable;
6058

6159
private final AwaitConsistencyCallable awaitConsistencyCallable;
6260
private final OperationCallable<Void, Empty, OptimizeRestoredTableMetadata>
6361
optimizeRestoredTableOperationBaseCallable;
6462

6563
public static EnhancedBigtableTableAdminStub createEnhanced(
66-
BigtableTableAdminStubSettings settings, TableAdminRequestContext requestContext)
64+
BigtableTableAdminStubSettings settings)
6765
throws IOException {
6866
return new EnhancedBigtableTableAdminStub(
69-
settings, ClientContext.create(settings), requestContext);
67+
settings, ClientContext.create(settings));
7068
}
7169

7270
private EnhancedBigtableTableAdminStub(
7371
BigtableTableAdminStubSettings settings,
74-
ClientContext clientContext,
75-
TableAdminRequestContext requestContext)
72+
ClientContext clientContext)
7673
throws IOException {
7774
super(settings, clientContext);
7875

7976
this.settings = settings;
8077
this.clientContext = clientContext;
81-
this.requestContext = requestContext;
8278
this.awaitConsistencyCallable = createAwaitConsistencyCallable();
8379
this.awaitReplicationCallable = createAwaitReplicationCallable();
8480
this.optimizeRestoredTableOperationBaseCallable =
@@ -113,8 +109,7 @@ private AwaitConsistencyCallable createAwaitConsistencyCallable() {
113109
generateConsistencyTokenCallable(),
114110
checkConsistencyCallable(),
115111
clientContext,
116-
pollingSettings,
117-
requestContext);
112+
pollingSettings);
118113
}
119114

120115
// Plug into gax operation infrastructure

0 commit comments

Comments
 (0)