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..b5ee9d90ea 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; @@ -1296,6 +1297,37 @@ 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 restoreFuture The future returned by restoreTableAsync(). + * @return An ApiFuture that tracks the optimization progress. + */ + public ApiFuture awaitOptimizeRestoredTable(ApiFuture 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 index e89bd8fbb5..c1d5da6592 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 @@ -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 + 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 mockRestoreFuture = Mockito.mock(ApiFuture.class); + Mockito.when(mockRestoreFuture.get()).thenReturn(mockResult); + + // 4. Mock the Stub's behavior (resuming the Optimize Op) + OperationFuture mockOptimizeOp = + Mockito.mock(OperationFuture.class); + Mockito.when(mockOptimizeRestoredTableCallable.resumeFutureCall(optimizeToken)) + .thenReturn(mockOptimizeOp); + + // Execute + ApiFuture 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 mockRestoreFuture = Mockito.mock(ApiFuture.class); + Mockito.when(mockRestoreFuture.get()).thenReturn(mockResult); + + // Execute + ApiFuture result = adminClient.awaitOptimizeRestoredTable(mockRestoreFuture); + + // 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, diff --git a/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/CbtTestProxy.java b/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/CbtTestProxy.java index f977a13e58..d2147e6167 100644 --- a/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/CbtTestProxy.java +++ b/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 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 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(); }