feat: Add awaitOptimizeRestoredTable helper for Bigtable Admin

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<RestoredTableResult> apply(com.google.bigtable.admin.v2.Table t
         MoreExecutors.directExecutor());
   }
 
+  /**
+   * Awaits the completion of the "Optimize Restored Table" operation.
+   *
+   * <p>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<Empty> awaitOptimizeRestoredTable(
+      OperationFuture<com.google.bigtable.admin.v2.Table, RestoreTableMetadata> 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.
    *

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<Void, Empty, OptimizeRestoredTableMetadata>
+      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<com.google.bigtable.admin.v2.Table, RestoreTableMetadata> 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<Empty, OptimizeRestoredTableMetadata> mockOptimizeOp =
+        Mockito.mock(OperationFuture.class);
+    Mockito.when(mockOptimizeRestoredTableCallable.resumeFutureCall(optimizeToken))
+        .thenReturn(mockOptimizeOp);
+
+    // Execute
+    ApiFuture<Empty> 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<com.google.bigtable.admin.v2.Table, RestoreTableMetadata> 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<Empty> result = adminClient.awaitOptimizeRestoredTable(mockRestoreOp);
+
+    // Verify: Returns immediate success (Empty) without calling the stub
+    assertThat(result.get()).isEqualTo(Empty.getDefaultInstance());
+    Mockito.verifyNoInteractions(mockStub);
+  }
+
   private <ReqT, RespT, MetaT> void mockOperationResult(
       OperationCallable<ReqT, RespT, MetaT> callable,
       ReqT request,