feat: Add awaitOptimizeRestoredTable helper for Bigtable Admin

jinseopkim0 · jinseopkim0 · commit 55d7dec9f771 · 2026-02-13T16:51:52.000-05:00
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
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
@@ -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,38 @@ 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 restoreFuture The future returned by restoreTableAsync().
+  * @return An ApiFuture that tracks the optimization progress.
+  */
+  public ApiFuture<Empty> awaitOptimizeRestoredTable(
+      ApiFuture<RestoredTableResult> restoreFuture) {
+    // 1. Block and wait for the restore operation to complete
+    RestoredTableResult result;
+    try {
+      result = restoreFuture.get();
+    } catch (Exception e) {
+      throw new RuntimeException("Restore operation failed", e);
+    }
+
+    // 2. Extract the operation token from the result
+    // (RestoredTableResult already wraps the OptimizeRestoredTableOperationToken)
+    OptimizeRestoredTableOperationToken token = result.getOptimizeRestoredTableOperationToken();
+
+    if (token == null || Strings.isNullOrEmpty(token.getOperationName())) {
+      // 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(token.getOperationName());
+  }
+
   /**
    * 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
@@ -17,6 +17,9 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import com.google.bigtable.admin.v2.OptimizeRestoredTableMetadata;
+import com.google.cloud.bigtable.admin.v2.models.OptimizeRestoredTableOperationToken;
+import com.google.cloud.bigtable.admin.v2.models.RestoredTableResult;
 import com.google.api.core.ApiFuture;
 import com.google.api.core.ApiFutures;
 import com.google.api.gax.grpc.GrpcStatusCode;
@@ -285,6 +288,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 +1689,59 @@ 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";
+
+    // 1. Mock the Token
+    OptimizeRestoredTableOperationToken mockToken =
+        Mockito.mock(OptimizeRestoredTableOperationToken.class);
+    Mockito.when(mockToken.getOperationName()).thenReturn(optimizeToken);
+
+    // 2. Mock the Result (wrapping the token)
+    RestoredTableResult mockResult = Mockito.mock(RestoredTableResult.class);
+    Mockito.when(mockResult.getOptimizeRestoredTableOperationToken()).thenReturn(mockToken);
+
+    // 3. Mock the Input Future (returning the result)
+    ApiFuture<RestoredTableResult> mockRestoreFuture = Mockito.mock(ApiFuture.class);
+    Mockito.when(mockRestoreFuture.get()).thenReturn(mockResult);
+
+    // 4. Mock the Stub's behavior (resuming the Optimize Op)
+    OperationFuture<Empty, OptimizeRestoredTableMetadata> mockOptimizeOp =
+        Mockito.mock(OperationFuture.class);
+    Mockito.when(mockOptimizeRestoredTableCallable.resumeFutureCall(optimizeToken))
+        .thenReturn(mockOptimizeOp);
+
+    // Execute
+    ApiFuture<Empty> result = adminClient.awaitOptimizeRestoredTable(mockRestoreFuture);
+
+    // Verify
+    assertThat(result).isEqualTo(mockOptimizeOp);
+    Mockito.verify(mockOptimizeRestoredTableCallable).resumeFutureCall(optimizeToken);
+  }
+
+  @Test
+  public void testAwaitOptimizeRestoredTable_NoOp() throws Exception {
+    // Setup: Result with NO optimization token (null or empty)
+    RestoredTableResult mockResult = Mockito.mock(RestoredTableResult.class);
+    Mockito.when(mockResult.getOptimizeRestoredTableOperationToken()).thenReturn(null);
+
+    // Mock the Input Future
+    ApiFuture<RestoredTableResult> mockRestoreFuture = Mockito.mock(ApiFuture.class);
+    Mockito.when(mockRestoreFuture.get()).thenReturn(mockResult);
+
+    // Execute
+    ApiFuture<Empty> result = adminClient.awaitOptimizeRestoredTable(mockRestoreFuture);
+
+    // 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,
