Skip to content

Commit cf15d45

Browse files
jinseopkim0cloud-java-bot
jinseopkim0
and
cloud-java-bot
authored
feat: Add awaitOptimizeRestoredTable helper for Bigtable Admin (#2781)
* feat: Add awaitOptimizeRestoredTable helper for Bigtable Admin Adds `awaitOptimizeRestoredTable` to simplify waiting for the secondary "Optimize" operation after a table restore. This method automatically extracts the operation token from the restore metadata and resumes the optimization LRO. This addresses the Long Running Sub-operations CUJ. Tracking Bug: b/475820271 * chore: generate libraries at Fri Feb 13 21:52:28 UTC 2026 --------- Co-authored-by: cloud-java-bot <[email protected]>
1 parent 2498287 commit cf15d45

2 files changed

Lines changed: 91 additions & 0 deletions

File tree

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

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@
7272
import com.google.cloud.bigtable.admin.v2.stub.EnhancedBigtableTableAdminStub;
7373
import com.google.cloud.bigtable.data.v2.internal.TableAdminRequestContext;
7474
import com.google.common.base.Preconditions;
75+
import com.google.common.base.Strings;
7576
import com.google.common.collect.ImmutableList;
7677
import com.google.common.collect.ImmutableMap;
7778
import com.google.common.collect.Lists;
@@ -1296,6 +1297,37 @@ public ApiFuture<RestoredTableResult> apply(com.google.bigtable.admin.v2.Table t
12961297
MoreExecutors.directExecutor());
12971298
}
12981299

1300+
/**
1301+
* Awaits the completion of the "Optimize Restored Table" operation.
1302+
*
1303+
* <p>This method blocks until the restore operation is complete, extracts the optimization token,
1304+
* and returns an ApiFuture for the optimization phase.
1305+
*
1306+
* @param restoreFuture The future returned by restoreTableAsync().
1307+
* @return An ApiFuture that tracks the optimization progress.
1308+
*/
1309+
public ApiFuture<Empty> awaitOptimizeRestoredTable(ApiFuture<RestoredTableResult> restoreFuture) {
1310+
// 1. Block and wait for the restore operation to complete
1311+
RestoredTableResult result;
1312+
try {
1313+
result = restoreFuture.get();
1314+
} catch (Exception e) {
1315+
throw new RuntimeException("Restore operation failed", e);
1316+
}
1317+
1318+
// 2. Extract the operation token from the result
1319+
// (RestoredTableResult already wraps the OptimizeRestoredTableOperationToken)
1320+
OptimizeRestoredTableOperationToken token = result.getOptimizeRestoredTableOperationToken();
1321+
1322+
if (token == null || Strings.isNullOrEmpty(token.getOperationName())) {
1323+
// If there is no optimization operation, return immediate success.
1324+
return ApiFutures.immediateFuture(Empty.getDefaultInstance());
1325+
}
1326+
1327+
// 3. Return the future for the optimization operation
1328+
return stub.awaitOptimizeRestoredTableCallable().resumeFutureCall(token.getOperationName());
1329+
}
1330+
12991331
/**
13001332
* Awaits a restored table is fully optimized.
13011333
*

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

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
import com.google.bigtable.admin.v2.ListBackupsRequest;
4646
import com.google.bigtable.admin.v2.ListTablesRequest;
4747
import com.google.bigtable.admin.v2.ModifyColumnFamiliesRequest.Modification;
48+
import com.google.bigtable.admin.v2.OptimizeRestoredTableMetadata;
4849
import com.google.bigtable.admin.v2.RestoreSourceType;
4950
import com.google.bigtable.admin.v2.RestoreTableMetadata;
5051
import com.google.bigtable.admin.v2.SchemaBundleName;
@@ -76,6 +77,7 @@
7677
import com.google.cloud.bigtable.admin.v2.models.CreateTableRequest;
7778
import com.google.cloud.bigtable.admin.v2.models.EncryptionInfo;
7879
import com.google.cloud.bigtable.admin.v2.models.ModifyColumnFamiliesRequest;
80+
import com.google.cloud.bigtable.admin.v2.models.OptimizeRestoredTableOperationToken;
7981
import com.google.cloud.bigtable.admin.v2.models.RestoreTableRequest;
8082
import com.google.cloud.bigtable.admin.v2.models.RestoredTableResult;
8183
import com.google.cloud.bigtable.admin.v2.models.SchemaBundle;
@@ -285,6 +287,10 @@ public class BigtableTableAdminClientTests {
285287
com.google.iam.v1.TestIamPermissionsRequest, com.google.iam.v1.TestIamPermissionsResponse>
286288
mockTestIamPermissionsCallable;
287289

290+
@Mock
291+
private OperationCallable<Void, Empty, OptimizeRestoredTableMetadata>
292+
mockOptimizeRestoredTableCallable;
293+
288294
@Before
289295
public void setUp() {
290296
adminClient = BigtableTableAdminClient.create(PROJECT_ID, INSTANCE_ID, mockStub);
@@ -1682,6 +1688,59 @@ public void testWaitForConsistencyWithToken() {
16821688
assertThat(wasCalled.get()).isTrue();
16831689
}
16841690

1691+
@Test
1692+
public void testAwaitOptimizeRestoredTable() throws Exception {
1693+
// Setup
1694+
Mockito.when(mockStub.awaitOptimizeRestoredTableCallable())
1695+
.thenReturn(mockOptimizeRestoredTableCallable);
1696+
1697+
String optimizeToken = "my-optimization-token";
1698+
1699+
// 1. Mock the Token
1700+
OptimizeRestoredTableOperationToken mockToken =
1701+
Mockito.mock(OptimizeRestoredTableOperationToken.class);
1702+
Mockito.when(mockToken.getOperationName()).thenReturn(optimizeToken);
1703+
1704+
// 2. Mock the Result (wrapping the token)
1705+
RestoredTableResult mockResult = Mockito.mock(RestoredTableResult.class);
1706+
Mockito.when(mockResult.getOptimizeRestoredTableOperationToken()).thenReturn(mockToken);
1707+
1708+
// 3. Mock the Input Future (returning the result)
1709+
ApiFuture<RestoredTableResult> mockRestoreFuture = Mockito.mock(ApiFuture.class);
1710+
Mockito.when(mockRestoreFuture.get()).thenReturn(mockResult);
1711+
1712+
// 4. Mock the Stub's behavior (resuming the Optimize Op)
1713+
OperationFuture<Empty, OptimizeRestoredTableMetadata> mockOptimizeOp =
1714+
Mockito.mock(OperationFuture.class);
1715+
Mockito.when(mockOptimizeRestoredTableCallable.resumeFutureCall(optimizeToken))
1716+
.thenReturn(mockOptimizeOp);
1717+
1718+
// Execute
1719+
ApiFuture<Empty> result = adminClient.awaitOptimizeRestoredTable(mockRestoreFuture);
1720+
1721+
// Verify
1722+
assertThat(result).isEqualTo(mockOptimizeOp);
1723+
Mockito.verify(mockOptimizeRestoredTableCallable).resumeFutureCall(optimizeToken);
1724+
}
1725+
1726+
@Test
1727+
public void testAwaitOptimizeRestoredTable_NoOp() throws Exception {
1728+
// Setup: Result with NO optimization token (null or empty)
1729+
RestoredTableResult mockResult = Mockito.mock(RestoredTableResult.class);
1730+
Mockito.when(mockResult.getOptimizeRestoredTableOperationToken()).thenReturn(null);
1731+
1732+
// Mock the Input Future
1733+
ApiFuture<RestoredTableResult> mockRestoreFuture = Mockito.mock(ApiFuture.class);
1734+
Mockito.when(mockRestoreFuture.get()).thenReturn(mockResult);
1735+
1736+
// Execute
1737+
ApiFuture<Empty> result = adminClient.awaitOptimizeRestoredTable(mockRestoreFuture);
1738+
1739+
// Verify: Returns immediate success (Empty) without calling the stub
1740+
assertThat(result.get()).isEqualTo(Empty.getDefaultInstance());
1741+
Mockito.verifyNoInteractions(mockStub);
1742+
}
1743+
16851744
private <ReqT, RespT, MetaT> void mockOperationResult(
16861745
OperationCallable<ReqT, RespT, MetaT> callable,
16871746
ReqT request,

0 commit comments

Comments
 (0)