Bigtable: clean up consistency token

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

@@ -465,7 +465,7 @@ public ConsistencyToken generateConsistencyToken(String tableId) {
    * }
    * }</pre>
    */
-  public ApiFuture<ConsistencyToken> generateConsistencyTokenAsync(String tableId) {
+  public ApiFuture<ConsistencyToken> generateConsistencyTokenAsync(final String tableId) {
     ApiFuture<GenerateConsistencyTokenResponse> tokenResp =
         this.stub
             .generateConsistencyTokenCallable()
@@ -475,8 +475,9 @@ public ApiFuture<ConsistencyToken> generateConsistencyTokenAsync(String tableId)
         tokenResp,
         new ApiFunction<GenerateConsistencyTokenResponse, ConsistencyToken>() {
           @Override
-          public ConsistencyToken apply(GenerateConsistencyTokenResponse input) {
-            return ConsistencyToken.fromProto(input);
+          public ConsistencyToken apply(GenerateConsistencyTokenResponse proto) {
+            TableName tableName = TableName.of(instanceName.getProject(), instanceName.getInstance(), tableId);
+            return ConsistencyToken.of(tableName, proto.getConsistencyToken());
           }
         },
         MoreExecutors.directExecutor());
@@ -489,28 +490,25 @@ public ConsistencyToken apply(GenerateConsistencyTokenResponse input) {
    *
    * <pre>{@code
    * try(BigtableTableAdminClient client =  BigtableTableAdminClient.create(InstanceName.of("[PROJECT]", "[INSTANCE]"))) {
-   *   boolean consistent = client.isConsistent("tableId", token);
-   * }
-   * }</pre>
-   */
-  public boolean isConsistent(String tableId, ConsistencyToken token) {
-    return awaitFuture(isConsistentAsync(tableId, token));
-  }
-
-  /**
-   * Checks replication consistency for the specified token consistency token asynchronously
+   *   // Perform some mutations.
    *
-   * <p>Sample code:
+   *   ConsistencyToken token = client.generateConsistencyToken("table-id");
+   *   while(!client.isConsistent(token)) {
+   *     Thread.sleep(100);
+   *   }
    *
-   * <pre>{@code
-   * try(BigtableTableAdminClient client =  BigtableTableAdminClient.create(InstanceName.of("[PROJECT]", "[INSTANCE]"))) {
-   *   boolean consistent = client.isConsistentAsync("tableId", token);
+   *   // Now all clusters are consistent
    * }
    * }</pre>
    */
-  public ApiFuture<Boolean> isConsistentAsync(String tableId, ConsistencyToken token) {
-    ApiFuture<CheckConsistencyResponse> checkConsResp =
-        stub.checkConsistencyCallable().futureCall(token.toProto(getTableName(tableId)));
+  public boolean isConsistent(ConsistencyToken token) {
+    return awaitFuture(isConsistentAsync(token));
+  }
+
+  @VisibleForTesting
+  ApiFuture<Boolean> isConsistentAsync(ConsistencyToken token) {
+    ApiFuture<CheckConsistencyResponse> checkConsResp = stub.checkConsistencyCallable()
+        .futureCall(token.toProto(instanceName));
 
     return ApiFutures.transform(
         checkConsResp,
@@ -523,6 +521,8 @@ public Boolean apply(CheckConsistencyResponse input) {
         MoreExecutors.directExecutor());
   }
 
+  // TODO(igorbernstein2): add awaitConsist() & awaitConsistAsync() that generate & poll a token
+
   /**
    * Helper method to construct the table name in format:
    * projects/{project}/instances/{instance}/tables/{tableId}

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

@@ -16,63 +16,39 @@
 package com.google.cloud.bigtable.admin.v2.models;
 
 import com.google.api.core.InternalApi;
+import com.google.api.core.InternalExtensionOnly;
+import com.google.auto.value.AutoValue;
 import com.google.bigtable.admin.v2.CheckConsistencyRequest;
 import com.google.bigtable.admin.v2.GenerateConsistencyTokenResponse;
-import com.google.common.annotations.VisibleForTesting;
-import com.google.common.base.MoreObjects;
-import com.google.common.base.Objects;
+import com.google.bigtable.admin.v2.InstanceName;
+import com.google.bigtable.admin.v2.TableName;
+import com.google.common.base.Preconditions;
 
 /**
  * Wrapper for {@link GenerateConsistencyTokenResponse#getConsistencyToken()}
  *
  * <p>Cannot be created. They are obtained by invoking {@link
  * com.google.cloud.bigtable.admin.v2.BigtableTableAdminClient#generateConsistencyToken(String)}
  */
-public final class ConsistencyToken {
-  private final String token;
-
-  @InternalApi
-  public static ConsistencyToken fromProto(GenerateConsistencyTokenResponse proto) {
-    return new ConsistencyToken(proto.getConsistencyToken());
+@InternalExtensionOnly
+@AutoValue
+public abstract class ConsistencyToken {
+  public static ConsistencyToken of(TableName tableName, String token) {
+    return new AutoValue_ConsistencyToken(tableName, token);
   }
 
-  private ConsistencyToken(String token) {
-    this.token = token;
-  }
+  abstract TableName getTableName();
+  abstract String getToken();
 
-  // TODO(igorbernstein): tableName should be part of the token and be parameterized
   @InternalApi
-  public CheckConsistencyRequest toProto(String tableName) {
+  public CheckConsistencyRequest toProto(InstanceName instanceName) {
+    Preconditions.checkArgument(
+        instanceName.equals(InstanceName.of(getTableName().getProject(), getTableName().getInstance())),
+        "Consistency tokens are only valid within a single instance.");
+
     return CheckConsistencyRequest.newBuilder()
-        .setName(tableName)
-        .setConsistencyToken(token)
+        .setName(getTableName().toString())
+        .setConsistencyToken(getToken())
         .build();
   }
-
-  @VisibleForTesting
-  String getToken() {
-    return token;
-  }
-
-  @Override
-  public boolean equals(Object o) {
-    if (this == o) {
-      return true;
-    }
-    if (o == null || getClass() != o.getClass()) {
-      return false;
-    }
-    ConsistencyToken that = (ConsistencyToken) o;
-    return Objects.equal(token, that.token);
-  }
-
-  @Override
-  public int hashCode() {
-    return Objects.hashCode(token);
-  }
-
-  @Override
-  public String toString() {
-    return MoreObjects.toStringHelper(this).add("token", token).toString();
-  }
 }

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

@@ -496,7 +496,7 @@ public void testGenerateConsistencyToken() {
     ConsistencyToken actualResult = adminClient.generateConsistencyToken(TABLE_NAME.getTable());
 
     // Verify
-    assertThat(actualResult).isEqualTo(ConsistencyToken.fromProto(expectedResponse));
+    assertThat(actualResult).isEqualTo(ConsistencyToken.of(TABLE_NAME, "fakeToken"));
   }
 
   @Test
@@ -519,7 +519,7 @@ public void testGenerateConsistencyTokenAsync() throws Exception {
         .generateConsistencyTokenAsync(TABLE_NAME.getTable());
 
     // Verify
-    assertThat(actualResult.get()).isEqualTo(ConsistencyToken.fromProto(expectedResponse));
+    assertThat(actualResult.get()).isEqualTo(ConsistencyToken.of(TABLE_NAME, "fakeToken"));
   }
 
   @Test
@@ -538,13 +538,9 @@ public void testCheckConsistencyToken() {
         .thenReturn(ApiFutures.immediateFuture(expectedResponse));
 
     // Execute
-    ConsistencyToken actualToken = ConsistencyToken.fromProto(
-        GenerateConsistencyTokenResponse.newBuilder()
-            .setConsistencyToken("fakeToken")
-            .build()
-    );
+    ConsistencyToken actualToken = ConsistencyToken.of(TABLE_NAME, "fakeToken");
 
-    boolean actualResult = adminClient.isConsistent(TABLE_NAME.getTable(), actualToken);
+    boolean actualResult = adminClient.isConsistent(actualToken);
 
     // Verify
     assertThat(actualResult).isTrue();

google-cloud-clients/google-cloud-bigtable-admin/src/test/java/com/google/cloud/bigtable/admin/v2/it/BigtableTableAdminClientIT.java

@@ -235,7 +235,7 @@ public void checkConsistency() {
       tableAdmin.createTable(CreateTableRequest.of(tableId));
       ConsistencyToken consistencyToken = tableAdmin.generateConsistencyToken(tableId);
       assertNotNull(consistencyToken);
-      boolean consistent = tableAdmin.isConsistent(tableId, consistencyToken);
+      boolean consistent = tableAdmin.isConsistent(consistencyToken);
       assertTrue(consistent);
     } finally {
       tableAdmin.deleteTable(tableId);

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

@@ -15,23 +15,47 @@
  */
 package com.google.cloud.bigtable.admin.v2.models;
 
-import static org.junit.Assert.assertEquals;
+import static com.google.common.truth.Truth.assertThat;
 
-import com.google.bigtable.admin.v2.GenerateConsistencyTokenResponse;
+import com.google.bigtable.admin.v2.CheckConsistencyRequest;
+import com.google.bigtable.admin.v2.InstanceName;
+import com.google.bigtable.admin.v2.TableName;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
 @RunWith(JUnit4.class)
 public class ConsistencyTokenTest {
+  private static final InstanceName INSTANCE_NAME = InstanceName.of("my-project", "my-instance");
+  private static final TableName TABLE_NAME = TableName.of(INSTANCE_NAME.getProject(), INSTANCE_NAME.getInstance(), "my-table");
+  private static final String TOKEN_VALUE = "87282hgwd8yg";
+
+  @Test
+  public void testToProto() {
+    ConsistencyToken token = ConsistencyToken.of(TABLE_NAME, TOKEN_VALUE);
+
+    assertThat(token.toProto(INSTANCE_NAME)).isEqualTo(
+        CheckConsistencyRequest.newBuilder()
+          .setName(TABLE_NAME.toString())
+          .setConsistencyToken(TOKEN_VALUE)
+          .build()
+    );
+  }
+
   @Test
-  public void fromJsonTest() {
-    ConsistencyToken tokenResponse =
-        ConsistencyToken.fromProto(
-            GenerateConsistencyTokenResponse.newBuilder()
-                .setConsistencyToken("87282hgwd8yg")
-                .build());
-
-    assertEquals("87282hgwd8yg", tokenResponse.getToken());
+  public void testInstanceMismatch() {
+    ConsistencyToken token = ConsistencyToken.of(TABLE_NAME, TOKEN_VALUE);
+
+    InstanceName otherInstanceName = InstanceName.of("my-project", "other-instance");
+
+    Exception actualError = null;
+
+    try {
+      token.toProto(otherInstanceName);
+    } catch (Exception e) {
+      actualError = e;
+    }
+
+    assertThat(actualError).isInstanceOf(IllegalArgumentException.class);
   }
 }