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

Commit cc8299c

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 cc8299c

2 files changed

Lines changed: 312 additions & 0 deletions

File tree

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

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

‎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)