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,37 @@ 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.
    *

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;
@@ -76,6 +77,7 @@
 import com.google.cloud.bigtable.admin.v2.models.CreateTableRequest;
 import com.google.cloud.bigtable.admin.v2.models.EncryptionInfo;
 import com.google.cloud.bigtable.admin.v2.models.ModifyColumnFamiliesRequest;
+import com.google.cloud.bigtable.admin.v2.models.OptimizeRestoredTableOperationToken;
 import com.google.cloud.bigtable.admin.v2.models.RestoreTableRequest;
 import com.google.cloud.bigtable.admin.v2.models.RestoredTableResult;
 import com.google.cloud.bigtable.admin.v2.models.SchemaBundle;
@@ -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,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,

test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/CbtTestProxy.java

@@ -289,8 +289,7 @@ public void mutateRow(
       // This response is empty.
       client.dataClient().mutateRow(mutation);
     } catch (ApiException e) {
-      responseObserver.onNext(
-          MutateRowResult.newBuilder().setStatus(convertStatus(e)).build());
+      responseObserver.onNext(MutateRowResult.newBuilder().setStatus(convertStatus(e)).build());
       responseObserver.onCompleted();
       return;
     } catch (StatusRuntimeException e) {
@@ -342,8 +341,7 @@ public void bulkMutateRows(
       responseObserver.onCompleted();
       return;
     } catch (ApiException e) {
-      responseObserver.onNext(
-          MutateRowsResult.newBuilder().setStatus(convertStatus(e)).build());
+      responseObserver.onNext(MutateRowsResult.newBuilder().setStatus(convertStatus(e)).build());
       responseObserver.onCompleted();
       return;
     } catch (StatusRuntimeException e) {
@@ -395,8 +393,7 @@ public void readRow(ReadRowRequest request, StreamObserver<RowResult> responseOb
         logger.info(String.format("readRow() did not find row: %s", request.getRowKey()));
       }
     } catch (ApiException e) {
-      responseObserver.onNext(
-          RowResult.newBuilder().setStatus(convertStatus(e)).build());
+      responseObserver.onNext(RowResult.newBuilder().setStatus(convertStatus(e)).build());
       responseObserver.onCompleted();
       return;
     } catch (StatusRuntimeException e) {
@@ -440,8 +437,7 @@ public void readRows(ReadRowsRequest request, StreamObserver<RowsResult> respons
           // Note that the default instance == OK
           resultBuilder.setStatus(com.google.rpc.Status.getDefaultInstance()).build());
     } catch (ApiException e) {
-      responseObserver.onNext(
-          RowsResult.newBuilder().setStatus(convertStatus(e)).build());
+      responseObserver.onNext(RowsResult.newBuilder().setStatus(convertStatus(e)).build());
       responseObserver.onCompleted();
       return;
     } catch (StatusRuntimeException e) {
@@ -558,8 +554,7 @@ public void sampleRowKeys(
     try {
       keyOffsets = client.dataClient().sampleRowKeys(tableId);
     } catch (ApiException e) {
-      responseObserver.onNext(
-          SampleRowKeysResult.newBuilder().setStatus(convertStatus(e)).build());
+      responseObserver.onNext(SampleRowKeysResult.newBuilder().setStatus(convertStatus(e)).build());
       responseObserver.onCompleted();
       return;
     } catch (StatusRuntimeException e) {
@@ -643,8 +638,7 @@ public void readModifyWriteRow(
                 "readModifyWriteRow() did not find row: %s", request.getRequest().getRowKey()));
       }
     } catch (ApiException e) {
-      responseObserver.onNext(
-          RowResult.newBuilder().setStatus(convertStatus(e)).build());
+      responseObserver.onNext(RowResult.newBuilder().setStatus(convertStatus(e)).build());
       responseObserver.onCompleted();
       return;
     } catch (StatusRuntimeException e) {
@@ -700,8 +694,7 @@ public void executeQuery(
       responseObserver.onError(e);
       return;
     } catch (ApiException e) {
-      responseObserver.onNext(
-          ExecuteQueryResult.newBuilder().setStatus(convertStatus(e)).build());
+      responseObserver.onNext(ExecuteQueryResult.newBuilder().setStatus(convertStatus(e)).build());
       responseObserver.onCompleted();
       return;
     } catch (StatusRuntimeException e) {
@@ -803,7 +796,9 @@ private static com.google.rpc.Status convertStatus(ApiException e) {
       return status;
     }
 
-    return com.google.rpc.Status.newBuilder().setCode(e.getStatusCode().getCode().ordinal()).setMessage(e.getMessage())
+    return com.google.rpc.Status.newBuilder()
+        .setCode(e.getStatusCode().getCode().ordinal())
+        .setMessage(e.getMessage())
         .build();
   }