diff --git a/google-cloud-clients/google-cloud-bigtable-admin/src/main/java/com/google/cloud/bigtable/admin/v2/BigtableTableAdminClient.java b/google-cloud-clients/google-cloud-bigtable-admin/src/main/java/com/google/cloud/bigtable/admin/v2/BigtableTableAdminClient.java index a0d2f911aa7e..d5abcb4972e2 100644 --- a/google-cloud-clients/google-cloud-bigtable-admin/src/main/java/com/google/cloud/bigtable/admin/v2/BigtableTableAdminClient.java +++ b/google-cloud-clients/google-cloud-bigtable-admin/src/main/java/com/google/cloud/bigtable/admin/v2/BigtableTableAdminClient.java @@ -454,10 +454,7 @@ public ApiFuture dropAllRowsAsync(String tableId) { * } */ public ConsistencyToken generateConsistencyToken(String tableId) { - return ConsistencyToken.fromProto( - this.stub - .generateConsistencyTokenCallable() - .call(composeGenerateConsistencyTokenRequest(tableId))); + return awaitFuture(generateConsistencyTokenAsync(tableId)); } /** @@ -472,7 +469,7 @@ public ConsistencyToken generateConsistencyToken(String tableId) { * } * } */ - public ApiFuture generateConsistencyTokenAsync(String tableId) { + public ApiFuture generateConsistencyTokenAsync(final String tableId) { ApiFuture tokenResp = this.stub .generateConsistencyTokenCallable() @@ -482,8 +479,9 @@ public ApiFuture generateConsistencyTokenAsync(String tableId) tokenResp, new ApiFunction() { @Override - public ConsistencyToken apply(GenerateConsistencyTokenResponse input) { - return ConsistencyToken.fromProto(input); + public ConsistencyToken apply(GenerateConsistencyTokenResponse proto) { + TableName tableName = TableName.of(instanceName.getProject(), instanceName.getProject(), tableId); + return ConsistencyToken.of(tableName, proto.getConsistencyToken()); } }); } @@ -495,30 +493,25 @@ public ConsistencyToken apply(GenerateConsistencyTokenResponse input) { * * 
{@code
    * try(BigtableTableAdminClient client =  BigtableTableAdminClient.create(InstanceName.of("[PROJECT]", "[INSTANCE]"))) {
-   *   boolean consistent = client.isConsistent("tableId", token);
-   * }
-   * }
- */ - public boolean isConsistent(String tableId, ConsistencyToken token) { - return stub.checkConsistencyCallable() - .call(token.toProto(getTableName(tableId))) - .getConsistent(); - } - - /** - * Checks replication consistency for the specified token consistency token asynchronously + * // Perform some mutations. * - *

Sample code: + * ConsistencyToken token = client.generateConsistencyToken("table-id"); + * while(!client.isConsistent(token)) { + * Thread.sleep(100); + * } * - * 

{@code
-   * try(BigtableTableAdminClient client =  BigtableTableAdminClient.create(InstanceName.of("[PROJECT]", "[INSTANCE]"))) {
-   *   boolean consistent = client.isConsistentAsync("tableId", token);
+   *   // Now all clusters are consistent
    * }
    * }
*/ - public ApiFuture isConsistentAsync(String tableId, ConsistencyToken token) { - ApiFuture checkConsResp = - stub.checkConsistencyCallable().futureCall(token.toProto(getTableName(tableId))); + public boolean isConsistent(ConsistencyToken token) { + return awaitFuture(isConsistentAsync(token)); + } + + @VisibleForTesting + ApiFuture isConsistentAsync(ConsistencyToken token) { + ApiFuture checkConsResp = stub.checkConsistencyCallable() + .futureCall(token.toProto(instanceName)); return ApiFutures.transform( checkConsResp, @@ -530,6 +523,8 @@ public Boolean apply(CheckConsistencyResponse input) { }); } + // 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} @@ -631,4 +626,13 @@ public Void apply(Empty empty) { } }); } + + private T awaitFuture(ApiFuture future) { + try { + return future.get(); + } catch(Throwable t) { + // TODO(igorbernstein2): figure out a better wrapper exception. + throw new RuntimeException(t); + } + } } diff --git a/google-cloud-clients/google-cloud-bigtable-admin/src/main/java/com/google/cloud/bigtable/admin/v2/models/ConsistencyToken.java b/google-cloud-clients/google-cloud-bigtable-admin/src/main/java/com/google/cloud/bigtable/admin/v2/models/ConsistencyToken.java index 9b3d32be5a59..74b42c30db00 100644 --- a/google-cloud-clients/google-cloud-bigtable-admin/src/main/java/com/google/cloud/bigtable/admin/v2/models/ConsistencyToken.java +++ b/google-cloud-clients/google-cloud-bigtable-admin/src/main/java/com/google/cloud/bigtable/admin/v2/models/ConsistencyToken.java @@ -18,8 +18,11 @@ import com.google.api.core.InternalApi; import com.google.bigtable.admin.v2.CheckConsistencyRequest; import com.google.bigtable.admin.v2.GenerateConsistencyTokenResponse; +import com.google.bigtable.admin.v2.InstanceName; +import com.google.bigtable.admin.v2.TableName; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.MoreObjects; +import com.google.common.base.Preconditions; /** * Wrapper for {@link GenerateConsistencyTokenResponse#getConsistencyToken()} @@ -28,22 +31,26 @@ * com.google.cloud.bigtable.admin.v2.BigtableTableAdminClient#generateConsistencyToken(String)} */ public final class ConsistencyToken { + private final TableName tableName; private final String token; - @InternalApi - public static ConsistencyToken fromProto(GenerateConsistencyTokenResponse proto) { - return new ConsistencyToken(proto.getConsistencyToken()); + public static ConsistencyToken of(TableName tableName, String token) { + return new ConsistencyToken(tableName, token); } - private ConsistencyToken(String token) { + public ConsistencyToken(TableName tableName, String token) { + this.tableName = tableName; this.token = token; } - // 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(tableName.getProject(), tableName.getInstance())), + "Consistency tokens are only valid within a single instance."); + return CheckConsistencyRequest.newBuilder() - .setName(tableName) + .setName(tableName.toString()) .setConsistencyToken(token) .build(); } diff --git a/google-cloud-clients/google-cloud-bigtable-admin/src/test/java/com/google/cloud/bigtable/admin/v2/BigtableTableAdminClientTest.java b/google-cloud-clients/google-cloud-bigtable-admin/src/test/java/com/google/cloud/bigtable/admin/v2/BigtableTableAdminClientTest.java index 38d078b115d9..a39282501589 100644 --- a/google-cloud-clients/google-cloud-bigtable-admin/src/test/java/com/google/cloud/bigtable/admin/v2/BigtableTableAdminClientTest.java +++ b/google-cloud-clients/google-cloud-bigtable-admin/src/test/java/com/google/cloud/bigtable/admin/v2/BigtableTableAdminClientTest.java @@ -257,7 +257,7 @@ public void generateAndCheckConsistency() { Mockito.when(mockCheckConsistencyCallable.call(any(CheckConsistencyRequest.class))) .thenReturn(consistencyResp); - adminClient.isConsistent("tableId", consistencyToken); + adminClient.isConsistent(consistencyToken); Mockito.verify(mockCheckConsistencyCallable).call(requestCaptor.capture()); } @@ -282,7 +282,7 @@ public void generateAndCheckConsistencyAsync() throws Exception { Mockito.when(mockCheckConsistencyCallable.futureCall(any(CheckConsistencyRequest.class))) .thenReturn(consistencyResp); - adminClient.isConsistentAsync("tableId", consistencyTokenFuture.get()); + adminClient.isConsistentAsync(consistencyTokenFuture.get()); Mockito.verify(mockCheckConsistencyCallable).futureCall(requestCaptor.capture()); } diff --git a/google-cloud-clients/google-cloud-bigtable-admin/src/test/java/com/google/cloud/bigtable/admin/v2/it/BigtableTableAdminClientIT.java b/google-cloud-clients/google-cloud-bigtable-admin/src/test/java/com/google/cloud/bigtable/admin/v2/it/BigtableTableAdminClientIT.java index 19d28b892978..c5c578aece41 100644 --- a/google-cloud-clients/google-cloud-bigtable-admin/src/test/java/com/google/cloud/bigtable/admin/v2/it/BigtableTableAdminClientIT.java +++ b/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); diff --git a/google-cloud-clients/google-cloud-bigtable-admin/src/test/java/com/google/cloud/bigtable/admin/v2/models/ConsistencyTokenTest.java b/google-cloud-clients/google-cloud-bigtable-admin/src/test/java/com/google/cloud/bigtable/admin/v2/models/ConsistencyTokenTest.java index 508bf546c63b..b426c822300f 100644 --- a/google-cloud-clients/google-cloud-bigtable-admin/src/test/java/com/google/cloud/bigtable/admin/v2/models/ConsistencyTokenTest.java +++ b/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); } }