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

Commit 555e0a2

Browse files
jinseopkim0jinseopkim0
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 555e0a2

2 files changed

Lines changed: 318 additions & 0 deletions

File tree

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

Lines changed: 177 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,177 @@
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+
com.google.cloud.bigtable.data.v2.internal.TableAdminRequestContext.create("", ""));
57+
// Pass the enhanced stub to the existing stub-based constructor
58+
return new BigtableTableAdminClientV2(stub);
59+
}
60+
61+
/** Constructs an instance of BigtableTableAdminClientV2 with the given stub. */
62+
public static final BigtableTableAdminClientV2 createClient(BigtableTableAdminStub stub) {
63+
return new BigtableTableAdminClientV2(stub);
64+
}
65+
66+
/**
67+
* Awaits the completion of the "Optimize Restored Table" operation.
68+
*
69+
* <p>This method blocks until the restore operation is complete, extracts the optimization token,
70+
* and returns an ApiFuture for the optimization phase.
71+
*
72+
* @param restoreFuture The future returned by restoreTableAsync().
73+
* @return An ApiFuture that tracks the optimization progress.
74+
*/
75+
public ApiFuture<Empty> awaitOptimizeRestoredTable(ApiFuture<RestoredTableResult> restoreFuture) {
76+
// 1. Block and wait for the restore operation to complete
77+
RestoredTableResult result;
78+
try {
79+
result = restoreFuture.get();
80+
} catch (Exception e) {
81+
throw new RuntimeException("Restore operation failed", e);
82+
}
83+
84+
// 2. Extract the operation token from the result
85+
// (RestoredTableResult already wraps the OptimizeRestoredTableOperationToken)
86+
OptimizeRestoredTableOperationToken token = result.getOptimizeRestoredTableOperationToken();
87+
88+
if (token == null || Strings.isNullOrEmpty(token.getOperationName())) {
89+
// If there is no optimization operation, return immediate success.
90+
return ApiFutures.immediateFuture(Empty.getDefaultInstance());
91+
}
92+
93+
// 3. Return the future for the optimization operation
94+
return ((EnhancedBigtableTableAdminStub) getStub()).awaitOptimizeRestoredTableCallable().resumeFutureCall(token.getOperationName());
95+
}
96+
97+
/**
98+
* Awaits a restored table is fully optimized.
99+
*
100+
* <p>Sample code
101+
*
102+
* <pre>{@code
103+
* RestoredTableResult result =
104+
* client.restoreTable(RestoreTableRequest.of(clusterId, backupId).setTableId(tableId));
105+
* client.awaitOptimizeRestoredTable(result.getOptimizeRestoredTableOperationToken());
106+
* }</pre>
107+
*/
108+
public void awaitOptimizeRestoredTable(OptimizeRestoredTableOperationToken token)
109+
throws ExecutionException, InterruptedException {
110+
awaitOptimizeRestoredTableAsync(token).get();
111+
}
112+
113+
/**
114+
* Awaits a restored table is fully optimized asynchronously.
115+
*
116+
* <p>Sample code
117+
*
118+
* <pre>{@code
119+
* RestoredTableResult result =
120+
* client.restoreTable(RestoreTableRequest.of(clusterId, backupId).setTableId(tableId));
121+
* ApiFuture<Void> future = client.awaitOptimizeRestoredTableAsync(
122+
* result.getOptimizeRestoredTableOperationToken());
123+
*
124+
* ApiFutures.addCallback(
125+
* future,
126+
* new ApiFutureCallback<Void>() {
127+
* public void onSuccess(Void unused) {
128+
* System.out.println("The optimization of the restored table is done.");
129+
* }
130+
*
131+
* public void onFailure(Throwable t) {
132+
* t.printStackTrace();
133+
* }
134+
* },
135+
* MoreExecutors.directExecutor()
136+
* );
137+
* }</pre>
138+
*/
139+
public ApiFuture<Void> awaitOptimizeRestoredTableAsync(
140+
OptimizeRestoredTableOperationToken token) {
141+
ApiFuture<Empty> emptyFuture =
142+
((EnhancedBigtableTableAdminStub) getStub()).awaitOptimizeRestoredTableCallable().resumeFutureCall(token.getOperationName());
143+
return ApiFutures.transform(
144+
emptyFuture,
145+
new com.google.api.core.ApiFunction<Empty, Void>() {
146+
@Override
147+
public Void apply(Empty input) {
148+
return null;
149+
}
150+
},
151+
com.google.common.util.concurrent.MoreExecutors.directExecutor());
152+
}
153+
154+
/**
155+
* Polls an existing consistency token until table replication is consistent across all clusters.
156+
* Useful for checking consistency of a token generated in a separate process. Blocks until
157+
* completion.
158+
*
159+
* @param tableId The table to check.
160+
* @param consistencyToken The token to poll.
161+
*/
162+
public void waitForConsistency(String tableId, String consistencyToken) {
163+
ApiExceptions.callAndTranslateApiException(waitForConsistencyAsync(tableId, consistencyToken));
164+
}
165+
166+
/**
167+
* Asynchronously polls the consistency token. Returns a future that completes when table
168+
* replication is consistent across all clusters.
169+
*
170+
* @param tableId The table to check.
171+
* @param consistencyToken The token to poll.
172+
*/
173+
public ApiFuture<Void> waitForConsistencyAsync(String tableId, String consistencyToken) {
174+
return ((EnhancedBigtableTableAdminStub) getStub()).awaitConsistencyCallable()
175+
.futureCall(ConsistencyRequest.forReplication(tableId, consistencyToken));
176+
}
177+
}

‎google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/admin/v2/BigtableTableAdminClientV2Test.java‎

Lines changed: 141 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,141 @@
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 static com.google.common.truth.Truth.assertThat;
19+
20+
import com.google.api.core.ApiFuture;
21+
import com.google.api.core.ApiFutures;
22+
import com.google.api.gax.longrunning.OperationFuture;
23+
import com.google.api.gax.rpc.OperationCallable;
24+
import com.google.api.gax.rpc.UnaryCallable;
25+
import com.google.bigtable.admin.v2.OptimizeRestoredTableMetadata;
26+
import com.google.cloud.bigtable.admin.v2.models.ConsistencyRequest;
27+
import com.google.cloud.bigtable.admin.v2.models.OptimizeRestoredTableOperationToken;
28+
import com.google.cloud.bigtable.admin.v2.models.RestoredTableResult;
29+
import com.google.cloud.bigtable.admin.v2.stub.EnhancedBigtableTableAdminStub;
30+
import com.google.protobuf.Empty;
31+
import java.util.concurrent.atomic.AtomicBoolean;
32+
import org.junit.Before;
33+
import org.junit.Rule;
34+
import org.junit.Test;
35+
import org.junit.runner.RunWith;
36+
import org.junit.runners.JUnit4;
37+
import org.mockito.Mock;
38+
import org.mockito.Mockito;
39+
import org.mockito.junit.MockitoJUnit;
40+
import org.mockito.junit.MockitoRule;
41+
import org.mockito.stubbing.Answer;
42+
43+
@RunWith(JUnit4.class)
44+
public class BigtableTableAdminClientV2Test {
45+
@Rule public final MockitoRule mockitoRule = MockitoJUnit.rule();
46+
47+
private static final String TABLE_ID = "my-table";
48+
49+
@Mock private EnhancedBigtableTableAdminStub mockStub;
50+
51+
@Mock
52+
private UnaryCallable<ConsistencyRequest, Void> mockAwaitConsistencyCallable;
53+
54+
@Mock
55+
private OperationCallable<Void, Empty, OptimizeRestoredTableMetadata>
56+
mockOptimizeRestoredTableCallable;
57+
58+
private BigtableTableAdminClientV2 client;
59+
60+
@Before
61+
public void setUp() {
62+
client = BigtableTableAdminClientV2.createClient(mockStub);
63+
}
64+
65+
@Test
66+
public void testWaitForConsistencyWithToken() {
67+
// Setup
68+
Mockito.when(mockStub.awaitConsistencyCallable()).thenReturn(mockAwaitConsistencyCallable);
69+
70+
String token = "my-token";
71+
ConsistencyRequest expectedRequest = ConsistencyRequest.forReplication(TABLE_ID, token);
72+
73+
final AtomicBoolean wasCalled = new AtomicBoolean(false);
74+
75+
Mockito.when(mockAwaitConsistencyCallable.futureCall(expectedRequest))
76+
.thenAnswer(
77+
(Answer<ApiFuture<Void>>)
78+
invocationOnMock -> {
79+
wasCalled.set(true);
80+
return ApiFutures.immediateFuture(null);
81+
});
82+
83+
// Execute
84+
client.waitForConsistency(TABLE_ID, token);
85+
86+
// Verify
87+
assertThat(wasCalled.get()).isTrue();
88+
}
89+
90+
@Test
91+
public void testAwaitOptimizeRestoredTable() throws Exception {
92+
// Setup
93+
Mockito.when(mockStub.awaitOptimizeRestoredTableCallable())
94+
.thenReturn(mockOptimizeRestoredTableCallable);
95+
96+
String optimizeToken = "my-optimization-token";
97+
98+
// 1. Mock the Token
99+
OptimizeRestoredTableOperationToken mockToken =
100+
Mockito.mock(OptimizeRestoredTableOperationToken.class);
101+
Mockito.when(mockToken.getOperationName()).thenReturn(optimizeToken);
102+
103+
// 2. Mock the Result (wrapping the token)
104+
RestoredTableResult mockResult = Mockito.mock(RestoredTableResult.class);
105+
Mockito.when(mockResult.getOptimizeRestoredTableOperationToken()).thenReturn(mockToken);
106+
107+
// 3. Mock the Input Future (returning the result)
108+
ApiFuture<RestoredTableResult> mockRestoreFuture = Mockito.mock(ApiFuture.class);
109+
Mockito.when(mockRestoreFuture.get()).thenReturn(mockResult);
110+
111+
// 4. Mock the Stub's behavior (resuming the Optimize Op)
112+
OperationFuture<Empty, OptimizeRestoredTableMetadata> mockOptimizeOp =
113+
Mockito.mock(OperationFuture.class);
114+
Mockito.when(mockOptimizeRestoredTableCallable.resumeFutureCall(optimizeToken))
115+
.thenReturn(mockOptimizeOp);
116+
117+
// Execute
118+
ApiFuture<Empty> result = client.awaitOptimizeRestoredTable(mockRestoreFuture);
119+
120+
// Verify
121+
assertThat(result).isEqualTo(mockOptimizeOp);
122+
Mockito.verify(mockOptimizeRestoredTableCallable).resumeFutureCall(optimizeToken);
123+
}
124+
125+
@Test
126+
public void testAwaitOptimizeRestoredTable_NoOp() throws Exception {
127+
// Setup: Result with NO optimization token (null or empty)
128+
RestoredTableResult mockResult = Mockito.mock(RestoredTableResult.class);
129+
Mockito.when(mockResult.getOptimizeRestoredTableOperationToken()).thenReturn(null);
130+
131+
// Mock the Input Future
132+
ApiFuture<RestoredTableResult> mockRestoreFuture = Mockito.mock(ApiFuture.class);
133+
Mockito.when(mockRestoreFuture.get()).thenReturn(mockResult);
134+
135+
// Execute
136+
ApiFuture<Empty> result = client.awaitOptimizeRestoredTable(mockRestoreFuture);
137+
138+
// Verify: Returns immediate success (Empty) without calling the stub
139+
assertThat(result.get()).isEqualTo(Empty.getDefaultInstance());
140+
}
141+
}

0 commit comments

Comments
 (0)