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

Commit bff60a6

Browse files
jinseopkim0jinseopkim0
committed
PR feedback
1 parent d781aad commit bff60a6

3 files changed

Lines changed: 174 additions & 22 deletions

File tree

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

Lines changed: 168 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -15,18 +15,40 @@
1515
*/
1616
package com.google.cloud.bigtable.admin.v2;
1717

18+
import com.google.api.core.ApiFunction;
1819
import com.google.api.core.ApiFuture;
1920
import com.google.api.core.ApiFutures;
21+
import com.google.api.gax.grpc.GrpcCallSettings;
22+
import com.google.api.gax.grpc.GrpcCallableFactory;
23+
import com.google.api.gax.grpc.ProtoOperationTransformers.MetadataTransformer;
24+
import com.google.api.gax.grpc.ProtoOperationTransformers.ResponseTransformer;
25+
import com.google.api.gax.longrunning.OperationSnapshot;
26+
import com.google.api.gax.longrunning.OperationTimedPollAlgorithm;
27+
import com.google.api.gax.retrying.RetrySettings;
2028
import com.google.api.gax.rpc.ApiExceptions;
29+
import com.google.api.gax.rpc.ClientContext;
30+
import com.google.api.gax.rpc.OperationCallSettings;
31+
import com.google.api.gax.rpc.OperationCallable;
32+
import com.google.api.gax.rpc.UnaryCallSettings;
33+
import com.google.api.gax.rpc.UnaryCallable;
34+
import com.google.bigtable.admin.v2.OptimizeRestoredTableMetadata;
2135
import com.google.cloud.bigtable.admin.v2.models.ConsistencyRequest;
2236
import com.google.cloud.bigtable.admin.v2.models.OptimizeRestoredTableOperationToken;
2337
import com.google.cloud.bigtable.admin.v2.models.RestoredTableResult;
38+
import com.google.cloud.bigtable.admin.v2.stub.AwaitConsistencyCallable;
2439
import com.google.cloud.bigtable.admin.v2.stub.BigtableTableAdminStub;
25-
import com.google.cloud.bigtable.admin.v2.stub.EnhancedBigtableTableAdminStub;
40+
import com.google.cloud.bigtable.admin.v2.stub.BigtableTableAdminStubSettings;
2641
import com.google.common.base.Strings;
42+
import com.google.common.util.concurrent.MoreExecutors;
43+
import com.google.longrunning.Operation;
2744
import com.google.protobuf.Empty;
45+
import io.grpc.MethodDescriptor;
46+
import io.grpc.MethodDescriptor.Marshaller;
47+
import io.grpc.MethodDescriptor.MethodType;
2848
import java.io.IOException;
49+
import java.io.InputStream;
2950
import java.util.concurrent.ExecutionException;
51+
import org.threeten.bp.Duration;
3052

3153
/**
3254
* Modern Cloud Bigtable Table Admin Client.
@@ -36,25 +58,135 @@
3658
* generator cannot handle natively (e.g., chained Long Running Operations, Consistency Polling).
3759
*/
3860
public class BigtableTableAdminClientV2 extends BaseBigtableTableAdminClient {
61+
private final AwaitConsistencyCallable awaitConsistencyCallable;
62+
private final OperationCallable<Void, Empty, OptimizeRestoredTableMetadata>
63+
optimizeRestoredTableOperationBaseCallable;
3964

4065
protected BigtableTableAdminClientV2(BaseBigtableTableAdminSettings settings) throws IOException {
4166
super(settings);
67+
this.awaitConsistencyCallable =
68+
createAwaitConsistencyCallable((BigtableTableAdminStubSettings) settings.getStubSettings());
69+
this.optimizeRestoredTableOperationBaseCallable =
70+
createOptimizeRestoredTableOperationBaseCallable(
71+
(BigtableTableAdminStubSettings) settings.getStubSettings());
4272
}
4373

4474
protected BigtableTableAdminClientV2(BigtableTableAdminStub stub) {
4575
super(stub);
76+
this.awaitConsistencyCallable = null;
77+
this.optimizeRestoredTableOperationBaseCallable = null;
78+
}
79+
80+
private AwaitConsistencyCallable createAwaitConsistencyCallable(
81+
BigtableTableAdminStubSettings settings) throws IOException {
82+
ClientContext clientContext = ClientContext.create(settings);
83+
// TODO(igorbernstein2): expose polling settings
84+
RetrySettings pollingSettings =
85+
RetrySettings.newBuilder()
86+
.setTotalTimeout(
87+
settings.checkConsistencySettings().getRetrySettings().getTotalTimeout())
88+
.setInitialRetryDelay(Duration.ofSeconds(10))
89+
.setRetryDelayMultiplier(1.0)
90+
.setMaxRetryDelay(Duration.ofSeconds(10))
91+
.setInitialRpcTimeout(Duration.ZERO)
92+
.setMaxRpcTimeout(Duration.ZERO)
93+
.setRpcTimeoutMultiplier(1.0)
94+
.build();
95+
96+
return AwaitConsistencyCallable.create(
97+
getStub().generateConsistencyTokenCallable(),
98+
getStub().checkConsistencyCallable(),
99+
clientContext,
100+
pollingSettings);
101+
}
102+
103+
private OperationCallable<Void, Empty, OptimizeRestoredTableMetadata>
104+
createOptimizeRestoredTableOperationBaseCallable(BigtableTableAdminStubSettings settings)
105+
throws IOException {
106+
ClientContext clientContext = ClientContext.create(settings);
107+
108+
GrpcCallSettings<Void, Operation> unusedInitialCallSettings =
109+
GrpcCallSettings.create(
110+
MethodDescriptor.<Void, Operation>newBuilder()
111+
.setType(MethodType.UNARY)
112+
.setFullMethodName(
113+
"google.bigtable.admin.v2.BigtableTableAdmin/OptimizeRestoredTable")
114+
.setRequestMarshaller(
115+
new Marshaller<Void>() {
116+
@Override
117+
public InputStream stream(Void value) {
118+
throw new UnsupportedOperationException("not used");
119+
}
120+
121+
@Override
122+
public Void parse(InputStream stream) {
123+
throw new UnsupportedOperationException("not used");
124+
}
125+
})
126+
.setResponseMarshaller(
127+
new Marshaller<Operation>() {
128+
@Override
129+
public InputStream stream(Operation value) {
130+
throw new UnsupportedOperationException("not used");
131+
}
132+
133+
@Override
134+
public Operation parse(InputStream stream) {
135+
throw new UnsupportedOperationException("not used");
136+
}
137+
})
138+
.build());
139+
140+
final MetadataTransformer<OptimizeRestoredTableMetadata> protoMetadataTransformer =
141+
MetadataTransformer.create(OptimizeRestoredTableMetadata.class);
142+
143+
final ResponseTransformer<com.google.protobuf.Empty> protoResponseTransformer =
144+
ResponseTransformer.create(com.google.protobuf.Empty.class);
145+
146+
OperationCallSettings<Void, Empty, OptimizeRestoredTableMetadata> operationCallSettings =
147+
OperationCallSettings.<Void, Empty, OptimizeRestoredTableMetadata>newBuilder()
148+
.setInitialCallSettings(
149+
UnaryCallSettings.<Void, OperationSnapshot>newUnaryCallSettingsBuilder()
150+
.setSimpleTimeoutNoRetries(Duration.ZERO)
151+
.build())
152+
.setMetadataTransformer(
153+
new ApiFunction<OperationSnapshot, OptimizeRestoredTableMetadata>() {
154+
@Override
155+
public OptimizeRestoredTableMetadata apply(OperationSnapshot input) {
156+
return protoMetadataTransformer.apply(input);
157+
}
158+
})
159+
.setResponseTransformer(
160+
new ApiFunction<OperationSnapshot, Empty>() {
161+
@Override
162+
public Empty apply(OperationSnapshot input) {
163+
return protoResponseTransformer.apply(input);
164+
}
165+
})
166+
.setPollingAlgorithm(
167+
OperationTimedPollAlgorithm.create(
168+
RetrySettings.newBuilder()
169+
.setInitialRetryDelay(Duration.ofMillis(500L))
170+
.setRetryDelayMultiplier(1.5)
171+
.setMaxRetryDelay(Duration.ofMillis(5000L))
172+
.setInitialRpcTimeout(Duration.ZERO)
173+
.setRpcTimeoutMultiplier(1.0)
174+
.setMaxRpcTimeout(Duration.ZERO)
175+
.setTotalTimeout(Duration.ofMillis(600000L))
176+
.build()))
177+
.build();
178+
179+
return GrpcCallableFactory.createOperationCallable(
180+
unusedInitialCallSettings,
181+
operationCallSettings,
182+
clientContext,
183+
getStub().getOperationsStub());
46184
}
47185

48186
/** Constructs an instance of BigtableTableAdminClientV2 with the given settings. */
49187
public static final BigtableTableAdminClientV2 createClient(
50188
BaseBigtableTableAdminSettings settings) throws IOException {
51-
// Explicitly create the enhanced stub
52-
EnhancedBigtableTableAdminStub stub =
53-
EnhancedBigtableTableAdminStub.createEnhanced(
54-
(com.google.cloud.bigtable.admin.v2.stub.BigtableTableAdminStubSettings)
55-
settings.getStubSettings());
56-
// Pass the enhanced stub to the existing stub-based constructor
57-
return new BigtableTableAdminClientV2(stub);
189+
return new BigtableTableAdminClientV2(settings);
58190
}
59191

60192
/** Constructs an instance of BigtableTableAdminClientV2 with the given stub. */
@@ -90,9 +222,7 @@ public ApiFuture<Empty> awaitOptimizeRestoredTable(ApiFuture<RestoredTableResult
90222
}
91223

92224
// 3. Return the future for the optimization operation
93-
return ((EnhancedBigtableTableAdminStub) getStub())
94-
.awaitOptimizeRestoredTableCallable()
95-
.resumeFutureCall(token.getOperationName());
225+
return getOptimizeRestoredTableCallable().resumeFutureCall(token.getOperationName());
96226
}
97227

98228
/**
@@ -140,9 +270,7 @@ public void awaitOptimizeRestoredTable(OptimizeRestoredTableOperationToken token
140270
public ApiFuture<Void> awaitOptimizeRestoredTableAsync(
141271
OptimizeRestoredTableOperationToken token) {
142272
ApiFuture<Empty> emptyFuture =
143-
((EnhancedBigtableTableAdminStub) getStub())
144-
.awaitOptimizeRestoredTableCallable()
145-
.resumeFutureCall(token.getOperationName());
273+
getOptimizeRestoredTableCallable().resumeFutureCall(token.getOperationName());
146274
return ApiFutures.transform(
147275
emptyFuture,
148276
new com.google.api.core.ApiFunction<Empty, Void>() {
@@ -175,8 +303,32 @@ public void waitForConsistency(String tableName, String consistencyToken) {
175303
* @param consistencyToken The token to poll.
176304
*/
177305
public ApiFuture<Void> waitForConsistencyAsync(String tableName, String consistencyToken) {
178-
return ((EnhancedBigtableTableAdminStub) getStub())
179-
.awaitConsistencyCallable()
306+
return getAwaitConsistencyCallable()
180307
.futureCall(ConsistencyRequest.forReplicationFromTableName(tableName, consistencyToken));
181308
}
309+
310+
private UnaryCallable<ConsistencyRequest, Void> getAwaitConsistencyCallable() {
311+
if (awaitConsistencyCallable != null) {
312+
return awaitConsistencyCallable;
313+
}
314+
// Fallback for tests or stub-based initialization
315+
if (getStub() instanceof com.google.cloud.bigtable.admin.v2.stub.EnhancedBigtableTableAdminStub) {
316+
return ((com.google.cloud.bigtable.admin.v2.stub.EnhancedBigtableTableAdminStub) getStub())
317+
.awaitConsistencyCallable();
318+
}
319+
throw new IllegalStateException("AwaitConsistencyCallable not initialized.");
320+
}
321+
322+
private OperationCallable<Void, Empty, OptimizeRestoredTableMetadata>
323+
getOptimizeRestoredTableCallable() {
324+
if (optimizeRestoredTableOperationBaseCallable != null) {
325+
return optimizeRestoredTableOperationBaseCallable;
326+
}
327+
// Fallback for tests or stub-based initialization
328+
if (getStub() instanceof com.google.cloud.bigtable.admin.v2.stub.EnhancedBigtableTableAdminStub) {
329+
return ((com.google.cloud.bigtable.admin.v2.stub.EnhancedBigtableTableAdminStub) getStub())
330+
.awaitOptimizeRestoredTableCallable();
331+
}
332+
throw new IllegalStateException("OptimizeRestoredTableCallable not initialized.");
333+
}
182334
}

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -51,15 +51,15 @@
5151
* <p>This callable wraps GenerateConsistencyToken and CheckConsistency RPCs. It will generate a
5252
* token then poll until isConsistent is true.
5353
*/
54-
class AwaitConsistencyCallable extends UnaryCallable<ConsistencyRequest, Void> {
54+
public class AwaitConsistencyCallable extends UnaryCallable<ConsistencyRequest, Void> {
5555
private final UnaryCallable<GenerateConsistencyTokenRequest, GenerateConsistencyTokenResponse>
5656
generateCallable;
5757
private final UnaryCallable<CheckConsistencyRequest, CheckConsistencyResponse> checkCallable;
5858
private final RetryingExecutor<CheckConsistencyResponse> executor;
5959

6060
@Nullable private final TableAdminRequestContext requestContext;
6161

62-
static AwaitConsistencyCallable create(
62+
public static AwaitConsistencyCallable create(
6363
UnaryCallable<GenerateConsistencyTokenRequest, GenerateConsistencyTokenResponse>
6464
generateCallable,
6565
UnaryCallable<CheckConsistencyRequest, CheckConsistencyResponse> checkCallable,
@@ -79,7 +79,7 @@ static AwaitConsistencyCallable create(
7979
generateCallable, checkCallable, retryingExecutor, requestContext);
8080
}
8181

82-
static AwaitConsistencyCallable create(
82+
public static AwaitConsistencyCallable create(
8383
UnaryCallable<GenerateConsistencyTokenRequest, GenerateConsistencyTokenResponse>
8484
generateCallable,
8585
UnaryCallable<CheckConsistencyRequest, CheckConsistencyResponse> checkCallable,

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -148,9 +148,9 @@ public void testCreateClientWithSettings() throws Exception {
148148
.build();
149149
try (BigtableTableAdminClientV2 settingsClient =
150150
BigtableTableAdminClientV2.createClient(settings)) {
151-
// Verify that the underlying stub was correctly instantiated as the Enhanced stub
152-
// so that downcasts in the CUJs do not throw a ClassCastException.
153-
assertThat(settingsClient.getStub()).isInstanceOf(EnhancedBigtableTableAdminStub.class);
151+
// Verify that the underlying stub is NOT an Enhanced stub by default
152+
// but the client has successfully initialized its own callables.
153+
assertThat(settingsClient.getStub()).isNotInstanceOf(EnhancedBigtableTableAdminStub.class);
154154
}
155155
}
156156
}

0 commit comments

Comments
 (0)