diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/BigtableTableAdminClient.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/BigtableTableAdminClient.java index 0e5a4c9433..a5f125e324 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/BigtableTableAdminClient.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/BigtableTableAdminClient.java @@ -72,6 +72,7 @@ import com.google.cloud.bigtable.admin.v2.stub.EnhancedBigtableTableAdminStub; import com.google.cloud.bigtable.data.v2.internal.TableAdminRequestContext; import com.google.common.base.Preconditions; +import com.google.common.base.Strings; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.Lists; @@ -1296,6 +1297,42 @@ public ApiFuture apply(com.google.bigtable.admin.v2.Table t MoreExecutors.directExecutor()); } + /** + * Awaits the completion of the "Optimize Restored Table" operation. + * + *

This method blocks until the restore operation is complete, extracts the optimization token, + * and returns an ApiFuture for the optimization phase. + * + * @param restoreOp The OperationFuture returned by restoreTableAsync(). + * @return An ApiFuture that tracks the optimization progress. + */ + public ApiFuture awaitOptimizeRestoredTable( + OperationFuture restoreOp) { + // 1. Block and wait for the restore operation to complete + // (We accept blocking here per agreement that admin operations can be synchronous) + try { + restoreOp.get(); + } catch (Exception e) { + throw new RuntimeException("Restore operation failed", e); + } + + // 2. Extract the operation name from the metadata + String optimizeOpName; + try { + optimizeOpName = restoreOp.getMetadata().get().getOptimizeTableOperationName(); + } catch (Exception e) { + throw new RuntimeException("Failed to extract optimization token from restore metadata", e); + } + + if (Strings.isNullOrEmpty(optimizeOpName)) { + // If there is no optimization operation, return immediate success. + return ApiFutures.immediateFuture(Empty.getDefaultInstance()); + } + + // 3. Return the future for the optimization operation + return stub.awaitOptimizeRestoredTableCallable().resumeFutureCall(optimizeOpName); + } + /** * Awaits a restored table is fully optimized. * diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/admin/v2/BigtableTableAdminClientTests.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/admin/v2/BigtableTableAdminClientTests.java index e89bd8fbb5..5586dc6337 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/admin/v2/BigtableTableAdminClientTests.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/admin/v2/BigtableTableAdminClientTests.java @@ -45,6 +45,7 @@ import com.google.bigtable.admin.v2.ListBackupsRequest; import com.google.bigtable.admin.v2.ListTablesRequest; import com.google.bigtable.admin.v2.ModifyColumnFamiliesRequest.Modification; +import com.google.bigtable.admin.v2.OptimizeRestoredTableMetadata; import com.google.bigtable.admin.v2.RestoreSourceType; import com.google.bigtable.admin.v2.RestoreTableMetadata; import com.google.bigtable.admin.v2.SchemaBundleName; @@ -285,6 +286,10 @@ public class BigtableTableAdminClientTests { com.google.iam.v1.TestIamPermissionsRequest, com.google.iam.v1.TestIamPermissionsResponse> mockTestIamPermissionsCallable; + @Mock + private OperationCallable + mockOptimizeRestoredTableCallable; + @Before public void setUp() { adminClient = BigtableTableAdminClient.create(PROJECT_ID, INSTANCE_ID, mockStub); @@ -1682,6 +1687,61 @@ public void testWaitForConsistencyWithToken() { assertThat(wasCalled.get()).isTrue(); } + @Test + public void testAwaitOptimizeRestoredTable() throws Exception { + // Setup + Mockito.when(mockStub.awaitOptimizeRestoredTableCallable()) + .thenReturn(mockOptimizeRestoredTableCallable); + + String optimizeToken = "my-optimization-token"; + RestoreTableMetadata restoreMetadata = + RestoreTableMetadata.newBuilder().setOptimizeTableOperationName(optimizeToken).build(); + + // Mock the input OperationFuture (Restore Op) + OperationFuture mockRestoreOp = + Mockito.mock(OperationFuture.class); + Mockito.when(mockRestoreOp.get()) + .thenReturn( + com.google.bigtable.admin.v2.Table.getDefaultInstance()); // Return proto, not wrapper + Mockito.when(mockRestoreOp.getMetadata()) + .thenReturn(ApiFutures.immediateFuture(restoreMetadata)); + + // Mock the Stub's behavior (Optimize Op) + OperationFuture mockOptimizeOp = + Mockito.mock(OperationFuture.class); + Mockito.when(mockOptimizeRestoredTableCallable.resumeFutureCall(optimizeToken)) + .thenReturn(mockOptimizeOp); + + // Execute + ApiFuture result = adminClient.awaitOptimizeRestoredTable(mockRestoreOp); + + // Verify + assertThat(result).isEqualTo(mockOptimizeOp); + Mockito.verify(mockOptimizeRestoredTableCallable).resumeFutureCall(optimizeToken); + } + + @Test + public void testAwaitOptimizeRestoredTable_NoOp() throws Exception { + // Setup: Metadata with NO optimization token + RestoreTableMetadata restoreMetadata = + RestoreTableMetadata.newBuilder().setOptimizeTableOperationName("").build(); + + OperationFuture mockRestoreOp = + Mockito.mock(OperationFuture.class); + Mockito.when(mockRestoreOp.get()) + .thenReturn( + com.google.bigtable.admin.v2.Table.getDefaultInstance()); // Return proto, not wrapper + Mockito.when(mockRestoreOp.getMetadata()) + .thenReturn(ApiFutures.immediateFuture(restoreMetadata)); + + // Execute + ApiFuture result = adminClient.awaitOptimizeRestoredTable(mockRestoreOp); + + // Verify: Returns immediate success (Empty) without calling the stub + assertThat(result.get()).isEqualTo(Empty.getDefaultInstance()); + Mockito.verifyNoInteractions(mockStub); + } + private void mockOperationResult( OperationCallable callable, ReqT request,