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

Author: jinseopkim0

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

@@ -88,6 +88,9 @@
 import java.util.Map;
 import java.util.concurrent.ExecutionException;
 import javax.annotation.Nonnull;
+import com.google.bigtable.admin.v2.OptimizeRestoredTableMetadata;
+import com.google.bigtable.admin.v2.RestoreTableMetadata;
+import com.google.common.base.Strings;
 
 /**
  * Client for creating, configuring, and deleting Cloud Bigtable tables
@@ -1296,6 +1299,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

@@ -17,6 +17,8 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import com.google.bigtable.admin.v2.OptimizeRestoredTableMetadata;
+import com.google.bigtable.admin.v2.RestoreTableMetadata;
 import com.google.api.core.ApiFuture;
 import com.google.api.core.ApiFutures;
 import com.google.api.gax.grpc.GrpcStatusCode;
@@ -285,6 +287,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 +1688,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,
@@ -1706,3 +1767,4 @@ private String getResourceFilePath(String filePath) throws URISyntaxException {
     return Paths.get(protoSchema.toURI()).toAbsolutePath().toString();
   }
 }
+