Merge branch 'main' into PR #2614 to update

Author: lqiu96

.github/CODEOWNERS

@@ -4,11 +4,11 @@
 # For syntax help see:
 # https://help.github.com/en/github/creating-cloning-and-archiving-repositories/about-code-owners#codeowners-syntax
 
-# The @googleapis/api-bigtable @googleapis/api-bigtable-partners is the default owner for changes in this repo
-*                       @googleapis/cloud-sdk-java-team @googleapis/api-bigtable @googleapis/api-bigtable-partners
+# The @googleapis/api-bigtable is the default owner for changes in this repo
+*                       @googleapis/cloud-sdk-java-team @googleapis/api-bigtable
 
 # for handwritten libraries, keep codeowner_team in .repo-metadata.json as owner
-**/*.java               @googleapis/api-bigtable @googleapis/api-bigtable-partners
+**/*.java               @googleapis/api-bigtable @googleapis/cloud-sdk-java-team
 
 
 # The java-samples-reviewers team is the default owner for samples changes
@@ -18,5 +18,5 @@ samples/**/*.java       @googleapis/java-samples-reviewers
 samples/snippets/generated/       @googleapis/cloud-sdk-java-team
 
 # Admin Module (Cloud Java Team ownership)
-**/com/google/cloud/bigtable/admin/** @googleapis/api-bigtable @googleapis/api-bigtable-partners @googleapis/cloud-sdk-java-team
-**/com/google/bigtable/admin/** @googleapis/api-bigtable @googleapis/api-bigtable-partners @googleapis/cloud-sdk-java-team
+**/com/google/cloud/bigtable/admin/** @googleapis/api-bigtable @googleapis/cloud-sdk-java-team
+**/com/google/bigtable/admin/** @googleapis/api-bigtable @googleapis/cloud-sdk-java-team

.repo-metadata.json

@@ -13,7 +13,7 @@
   "api_id": "bigtable.googleapis.com",
   "library_type": "GAPIC_COMBO",
   "requires_billing": true,
-  "codeowner_team": "@googleapis/api-bigtable @googleapis/api-bigtable-partners",
+  "codeowner_team": "@googleapis/api-bigtable",
   "excluded_poms": "google-cloud-bigtable-bom",
   "issue_tracker": "https://issuetracker.google.com/savedsearches/559777",
   "extra_versioned_modules": "google-cloud-bigtable-emulator,google-cloud-bigtable-emulator-core",

generation_config.yaml

@@ -27,7 +27,7 @@ libraries:
   issue_tracker: https://issuetracker.google.com/savedsearches/559777
   release_level: stable
   distribution_name: com.google.cloud:google-cloud-bigtable
-  codeowner_team: '@googleapis/api-bigtable @googleapis/api-bigtable-partners'
+  codeowner_team: '@googleapis/api-bigtable'
   api_id: bigtable.googleapis.com
   library_type: GAPIC_COMBO
   extra_versioned_modules: google-cloud-bigtable-emulator,google-cloud-bigtable-emulator-core

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

@@ -1460,6 +1460,30 @@ public ApiFuture<Void> awaitReplicationAsync(final String tableId) {
     return stub.awaitReplicationCallable().futureCall(tableName);
   }
 
+  /**
+   * Polls an existing consistency token until table replication is consistent across all clusters.
+   * Useful for checking consistency of a token generated in a separate process. Blocks until
+   * completion.
+   *
+   * @param tableId The table to check.
+   * @param consistencyToken The token to poll.
+   */
+  public void waitForConsistency(String tableId, String consistencyToken) {
+    ApiExceptions.callAndTranslateApiException(waitForConsistencyAsync(tableId, consistencyToken));
+  }
+
+  /**
+   * Asynchronously polls the consistency token. Returns a future that completes when table
+   * replication is consistent across all clusters.
+   *
+   * @param tableId The table to check.
+   * @param consistencyToken The token to poll.
+   */
+  public ApiFuture<Void> waitForConsistencyAsync(String tableId, String consistencyToken) {
+    return stub.awaitConsistencyCallable()
+        .futureCall(ConsistencyRequest.forReplication(tableId, consistencyToken));
+  }
+
   /**
    * Creates a new authorized view with the specified configuration.
    *

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

@@ -23,7 +23,9 @@
 import com.google.bigtable.admin.v2.StandardReadRemoteWrites;
 import com.google.bigtable.admin.v2.TableName;
 import com.google.cloud.bigtable.data.v2.internal.TableAdminRequestContext;
+import com.google.common.base.Preconditions;
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 
 @AutoValue
 public abstract class ConsistencyRequest {
@@ -33,14 +35,36 @@ public abstract class ConsistencyRequest {
   @Nonnull
   protected abstract CheckConsistencyRequest.ModeCase getMode();
 
+  /**
+   * Internal accessor for the consistency token. Must be public to be accessible from the stub
+   * package.
+   */
+  @InternalApi
+  @Nullable
+  public abstract String getConsistencyToken();
+
   public static ConsistencyRequest forReplication(String tableId) {
     return new AutoValue_ConsistencyRequest(
-        tableId, CheckConsistencyRequest.ModeCase.STANDARD_READ_REMOTE_WRITES);
+        tableId, CheckConsistencyRequest.ModeCase.STANDARD_READ_REMOTE_WRITES, null);
+  }
+
+  /**
+   * Creates a request to check consistency using an existing token.
+   *
+   * @param tableId The table ID.
+   * @param consistencyToken The token to check. Must not be null.
+   * @throws NullPointerException if consistencyToken is null.
+   */
+  public static ConsistencyRequest forReplication(String tableId, String consistencyToken) {
+    Preconditions.checkNotNull(consistencyToken, "consistencyToken must not be null");
+
+    return new AutoValue_ConsistencyRequest(
+        tableId, CheckConsistencyRequest.ModeCase.STANDARD_READ_REMOTE_WRITES, consistencyToken);
   }
 
   public static ConsistencyRequest forDataBoost(String tableId) {
     return new AutoValue_ConsistencyRequest(
-        tableId, CheckConsistencyRequest.ModeCase.DATA_BOOST_READ_LOCAL_WRITES);
+        tableId, CheckConsistencyRequest.ModeCase.DATA_BOOST_READ_LOCAL_WRITES, null);
   }
 
   @InternalApi

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

@@ -93,6 +93,15 @@ static AwaitConsistencyCallable create(
   @Override
   public ApiFuture<Void> futureCall(
       final ConsistencyRequest consistencyRequest, final ApiCallContext apiCallContext) {
+
+    // If the token is already provided, skip generation and poll directly.
+    if (consistencyRequest.getConsistencyToken() != null) {
+      CheckConsistencyRequest request =
+          consistencyRequest.toCheckConsistencyProto(
+              requestContext, consistencyRequest.getConsistencyToken());
+      return pollToken(request, apiCallContext);
+    }
+
     ApiFuture<GenerateConsistencyTokenResponse> tokenFuture =
         generateToken(consistencyRequest.toGenerateTokenProto(requestContext), apiCallContext);
 

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/admin/v2/BigtableTableAdminClientTests.java

@@ -1657,6 +1657,31 @@ public void testTestIamPermissions() {
     assertThat(actualResult).containsExactly("bigtable.backups.get");
   }
 
+  @Test
+  public void testWaitForConsistencyWithToken() {
+    // Setup
+    Mockito.when(mockStub.awaitConsistencyCallable()).thenReturn(mockAwaitConsistencyCallable);
+
+    String token = "my-token";
+    ConsistencyRequest expectedRequest = ConsistencyRequest.forReplication(TABLE_ID, token);
+
+    final AtomicBoolean wasCalled = new AtomicBoolean(false);
+
+    Mockito.when(mockAwaitConsistencyCallable.futureCall(expectedRequest))
+        .thenAnswer(
+            (Answer<ApiFuture<Void>>)
+                invocationOnMock -> {
+                  wasCalled.set(true);
+                  return ApiFutures.immediateFuture(null);
+                });
+
+    // Execute
+    adminClient.waitForConsistency(TABLE_ID, token);
+
+    // Verify
+    assertThat(wasCalled.get()).isTrue();
+  }
+
   private <ReqT, RespT, MetaT> void mockOperationResult(
       OperationCallable<ReqT, RespT, MetaT> callable,
       ReqT request,

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/admin/v2/models/ConsistencyRequestTest.java

@@ -79,4 +79,22 @@ public void testToGenerateTokenProto() {
     assertThat(generateRequest.getName())
         .isEqualTo(NameUtil.formatTableName(PROJECT_ID, INSTANCE_ID, TABLE_ID));
   }
+
+  @Test
+  public void testToCheckConsistencyProtoWithToken() {
+    ConsistencyRequest consistencyRequest =
+        ConsistencyRequest.forReplication(TABLE_ID, CONSISTENCY_TOKEN);
+
+    TableAdminRequestContext requestContext =
+        TableAdminRequestContext.create(PROJECT_ID, INSTANCE_ID);
+
+    CheckConsistencyRequest checkConsistencyRequest =
+        consistencyRequest.toCheckConsistencyProto(requestContext, CONSISTENCY_TOKEN);
+
+    assertThat(checkConsistencyRequest.getName())
+        .isEqualTo(NameUtil.formatTableName(PROJECT_ID, INSTANCE_ID, TABLE_ID));
+    assertThat(checkConsistencyRequest.getConsistencyToken()).isEqualTo(CONSISTENCY_TOKEN);
+    assertThat(checkConsistencyRequest.getModeCase())
+        .isEqualTo(CheckConsistencyRequest.ModeCase.STANDARD_READ_REMOTE_WRITES);
+  }
 }

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/admin/v2/stub/AwaitConsistencyCallableTest.java

@@ -42,6 +42,7 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
+import org.mockito.ArgumentMatchers;
 import org.mockito.Mock;
 import org.mockito.Mockito;
 import org.mockito.junit.MockitoJUnit;
@@ -325,4 +326,37 @@ public void testAwaitReplicationCallablePolling() throws Exception {
 
     consistentFuture.get(1, TimeUnit.SECONDS);
   }
+
+  @Test
+  public void testWithProvidedToken() throws Exception {
+    // 1. Setup: Request with a pre-existing token
+    String existingToken = "existing-token";
+    ConsistencyRequest consistencyRequest =
+        ConsistencyRequest.forReplication(TABLE_ID, existingToken);
+
+    // 2. Setup: Mock the check operation to succeed immediately
+    CheckConsistencyRequest expectedCheckRequest =
+        CheckConsistencyRequest.newBuilder()
+            .setName(TABLE_NAME.toString())
+            .setConsistencyToken(existingToken)
+            .setStandardReadRemoteWrites(StandardReadRemoteWrites.newBuilder().build())
+            .build();
+    CheckConsistencyResponse expectedResponse =
+        CheckConsistencyResponse.newBuilder().setConsistent(true).build();
+
+    Mockito.when(mockCheckConsistencyCallable.futureCall(expectedCheckRequest, CALL_CONTEXT))
+        .thenReturn(ApiFutures.immediateFuture(expectedResponse));
+
+    // 3. Execute
+    ApiFuture<Void> future = awaitConsistencyCallable.futureCall(consistencyRequest, CALL_CONTEXT);
+    future.get(1, TimeUnit.SECONDS);
+
+    // 4. Verify: Generate was NEVER called, Check WAS called
+    Mockito.verify(mockGenerateConsistencyTokenCallable, Mockito.never())
+        .futureCall(
+            ArgumentMatchers.any(GenerateConsistencyTokenRequest.class),
+            ArgumentMatchers.any(ApiCallContext.class));
+    Mockito.verify(mockCheckConsistencyCallable, Mockito.times(1))
+        .futureCall(expectedCheckRequest, CALL_CONTEXT);
+  }
 }