From 820c75108194299e1481aa52776fc7b3b5366bc8 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Mon, 9 Feb 2026 17:08:49 -0500 Subject: [PATCH 1/2] 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 --- .../admin/v2/BigtableTableAdminClient.java | 39 ++++++++++++ .../v2/BigtableTableAdminClientTests.java | 62 +++++++++++++++++++ 2 files changed, 101 insertions(+) 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..999175ce69 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 @@ -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 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..9afc032534 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 @@ -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 + 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 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, @@ -1706,3 +1767,4 @@ private String getResourceFilePath(String filePath) throws URISyntaxException { return Paths.get(protoSchema.toURI()).toAbsolutePath().toString(); } } + From 0c17745363ecfed592386e3101def00e6d019edd Mon Sep 17 00:00:00 2001 From: cloud-java-bot Date: Mon, 9 Feb 2026 22:12:35 +0000 Subject: [PATCH 2/2] chore: generate libraries at Mon Feb 9 22:10:07 UTC 2026 --- .../cloud/bigtable/admin/v2/BigtableTableAdminClient.java | 4 +--- .../bigtable/admin/v2/BigtableTableAdminClientTests.java | 4 +--- 2 files changed, 2 insertions(+), 6 deletions(-) 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 999175ce69..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; @@ -88,9 +89,6 @@ 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 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 9afc032534..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 @@ -17,8 +17,6 @@ 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; @@ -47,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; @@ -1767,4 +1766,3 @@ private String getResourceFilePath(String filePath) throws URISyntaxException { return Paths.get(protoSchema.toURI()).toAbsolutePath().toString(); } } -