Bigtable: cleanup of futures + extras

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

@@ -15,6 +15,8 @@
  */
 package com.google.cloud.bigtable.admin.v2;
 
+import com.google.cloud.bigtable.admin.v2.BaseBigtableTableAdminClient.ListTablesPagedResponse;
+import com.google.common.util.concurrent.MoreExecutors;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
@@ -144,9 +146,7 @@ public void close() {
    * @see CreateTableRequest for createTable configurations
    */
   public Table createTable(CreateTableRequest request) {
-    com.google.bigtable.admin.v2.Table table =
-        this.stub.createTableCallable().call(request.toProto(instanceName));
-    return Table.fromProto(table);
+    return awaitFuture(createTableAsync(request));
   }
 
   /**
@@ -200,9 +200,7 @@ public ApiFuture<Table> createTableAsync(CreateTableRequest request) {
    * @see ModifyColumnFamiliesRequest for modifyFamily options
    */
   public Table modifyFamilies(ModifyColumnFamiliesRequest request) {
-    com.google.bigtable.admin.v2.Table table =
-        this.stub.modifyColumnFamiliesCallable().call(request.toProto(instanceName));
-    return Table.fromProto(table);
+    return awaitFuture(modifyFamiliesAsync(request));
   }
 
   /**
@@ -252,7 +250,7 @@ public ApiFuture<Table> modifyFamiliesAsync(ModifyColumnFamiliesRequest request)
    * }</pre>
    */
   public void deleteTable(String tableId) {
-    this.stub.deleteTableCallable().call(composeDeleteTableRequest(tableId));
+    awaitFuture(deleteTableAsync(tableId));
   }
 
   /**
@@ -283,9 +281,7 @@ public ApiFuture<Void> deleteTableAsync(String tableId) {
    * }</pre>
    */
   public Table getTable(String tableId) {
-    com.google.bigtable.admin.v2.Table table =
-        this.stub.getTableCallable().call(composeGetTableRequest(tableId));
-    return Table.fromProto(table);
+    return awaitFuture(getTableAsync(tableId));
   }
 
   /**
@@ -316,8 +312,7 @@ public ApiFuture<Table> getTableAsync(String tableId) {
    * }</pre>
    */
   public List<TableName> listTables() {
-    ListTablesResponse listResp = this.stub.listTablesCallable().call(composeListTableRequest());
-    return convertToTableNames(listResp);
+    return awaitFuture(listTablesAsync());
   }
 
   /**
@@ -332,17 +327,18 @@ public List<TableName> listTables() {
    * }</pre>
    */
   public ApiFuture<List<TableName>> listTablesAsync() {
-    ApiFuture<ListTablesResponse> listResp =
-        this.stub.listTablesCallable().futureCall(composeListTableRequest());
+    ApiFuture<ListTablesPagedResponse> listResp =
+        this.stub.listTablesPagedCallable().futureCall(composeListTableRequest());
 
     return ApiFutures.transform(
         listResp,
-        new ApiFunction<ListTablesResponse, List<TableName>>() {
+        new ApiFunction<ListTablesPagedResponse, List<TableName>>() {
           @Override
-          public List<TableName> apply(ListTablesResponse input) {
-            return convertToTableNames(input);
+          public List<TableName> apply(ListTablesPagedResponse input) {
+            return convertToTableNames(input.iterateAll());
           }
-        });
+        },
+        MoreExecutors.directExecutor());
   }
 
   /**
@@ -357,7 +353,7 @@ public List<TableName> apply(ListTablesResponse input) {
    * }</pre>
    */
   public void dropRowRange(String tableId, String rowKeyPrefix) {
-    dropRowRange(tableId, ByteString.copyFromUtf8(rowKeyPrefix));
+    awaitFuture(dropRowRangeAsync(tableId, rowKeyPrefix));
   }
 
   /**
@@ -387,7 +383,7 @@ public ApiFuture<Void> dropRowRangeAsync(String tableId, String rowKeyPrefix) {
    * }</pre>
    */
   public void dropRowRange(String tableId, ByteString rowKeyPrefix) {
-    this.stub.dropRowRangeCallable().call(composeDropRowRangeRequest(tableId, rowKeyPrefix, false));
+    awaitFuture(dropRowRangeAsync(tableId, rowKeyPrefix));
   }
 
   /**
@@ -420,7 +416,7 @@ public ApiFuture<Void> dropRowRangeAsync(String tableId, ByteString rowKeyPrefix
    * }</pre>
    */
   public void dropAllRows(String tableId) {
-    this.stub.dropRowRangeCallable().call(composeDropRowRangeRequest(tableId, null, true));
+    awaitFuture(dropAllRowsAsync(tableId));
   }
 
   /**
@@ -454,10 +450,7 @@ public ApiFuture<Void> dropAllRowsAsync(String tableId) {
    * }</pre>
    */
   public ConsistencyToken generateConsistencyToken(String tableId) {
-    return ConsistencyToken.fromProto(
-        this.stub
-            .generateConsistencyTokenCallable()
-            .call(composeGenerateConsistencyTokenRequest(tableId)));
+    return awaitFuture(generateConsistencyTokenAsync(tableId));
   }
 
   /**
@@ -485,7 +478,8 @@ public ApiFuture<ConsistencyToken> generateConsistencyTokenAsync(String tableId)
           public ConsistencyToken apply(GenerateConsistencyTokenResponse input) {
             return ConsistencyToken.fromProto(input);
           }
-        });
+        },
+        MoreExecutors.directExecutor());
   }
 
   /**
@@ -500,9 +494,7 @@ public ConsistencyToken apply(GenerateConsistencyTokenResponse input) {
    * }</pre>
    */
   public boolean isConsistent(String tableId, ConsistencyToken token) {
-    return stub.checkConsistencyCallable()
-        .call(token.toProto(getTableName(tableId)))
-        .getConsistent();
+    return awaitFuture(isConsistentAsync(tableId, token));
   }
 
   /**
@@ -527,7 +519,8 @@ public ApiFuture<Boolean> isConsistentAsync(String tableId, ConsistencyToken tok
           public Boolean apply(CheckConsistencyResponse input) {
             return input.getConsistent();
           }
-        });
+        },
+        MoreExecutors.directExecutor());
   }
 
   /**
@@ -591,10 +584,10 @@ GenerateConsistencyTokenRequest composeGenerateConsistencyTokenRequest(String ta
    * Helper method to convert ListTablesResponse to List<TableName>
    */
   @VisibleForTesting
-  static List<TableName> convertToTableNames(ListTablesResponse listTablesResponse) {
+  static List<TableName> convertToTableNames(Iterable<com.google.bigtable.admin.v2.Table> listTablesResponse) {
     List<TableName> tableNames = new ArrayList<>();
 
-    for (com.google.bigtable.admin.v2.Table table : listTablesResponse.getTablesList()) {
+    for (com.google.bigtable.admin.v2.Table table : listTablesResponse) {
       tableNames.add(TableName.parse(table.getName()));
     }
     return tableNames;
@@ -604,8 +597,7 @@ static List<TableName> convertToTableNames(ListTablesResponse listTablesResponse
   /**
    * Helper method to transform ApiFuture<com.google.bigtable.admin.v2.Table> to ApiFuture<Table>
    */
-  @VisibleForTesting
-  static ApiFuture<Table> transformToTableResponse(
+  private static ApiFuture<Table> transformToTableResponse(
       ApiFuture<com.google.bigtable.admin.v2.Table> future) {
     return ApiFutures.transform(
         future,
@@ -614,21 +606,31 @@ static ApiFuture<Table> transformToTableResponse(
           public Table apply(com.google.bigtable.admin.v2.Table table) {
             return Table.fromProto(table);
           }
-        });
+        },
+        MoreExecutors.directExecutor());
   }
 
   /**
    * Helper method to transform ApiFuture<Empty> to ApiFuture<Void>
    */
-  @VisibleForTesting
-  static ApiFuture<Void> transformToVoid(ApiFuture<Empty> future) {
+  private static ApiFuture<Void> transformToVoid(ApiFuture<Empty> future) {
     return ApiFutures.transform(
         future,
         new ApiFunction<Empty, Void>() {
           @Override
           public Void apply(Empty empty) {
             return null;
           }
-        });
+        },
+        MoreExecutors.directExecutor());
+  }
+
+  private <T> T awaitFuture(ApiFuture<T> future) {
+    try {
+      return future.get();
+    } catch(Throwable t) {
+      // TODO(igorbernstein2): figure out a better wrapper exception.
+      throw new RuntimeException(t);
+    }
   }
 }

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

@@ -17,6 +17,7 @@
 
 import com.google.bigtable.admin.v2.Table.ClusterState.ReplicationState;
 import com.google.common.base.MoreObjects;
+import com.google.common.base.Objects;
 
 // TODO(igorbernstein2): remove this class and promote Replication State to Table.
 /** Wrapper for {@link ClusterState} protocol buffer object */
@@ -47,6 +48,24 @@ public ReplicationState getReplicationState() {
     return replicationState;
   }
 
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) {
+      return true;
+    }
+    if (o == null || getClass() != o.getClass()) {
+      return false;
+    }
+    ClusterState that = (ClusterState) o;
+    return Objects.equal(id, that.id) &&
+        replicationState == that.replicationState;
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hashCode(id, replicationState);
+  }
+
   @Override
   public String toString() {
     return MoreObjects.toStringHelper(this)

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

@@ -22,6 +22,7 @@
 import com.google.bigtable.admin.v2.GcRule.RuleCase;
 import com.google.cloud.bigtable.admin.v2.models.GCRules.GCRule;
 import com.google.common.base.MoreObjects;
+import com.google.common.base.Objects;
 
 /** Wrapper for {@link ColumnFamily} protocol buffer object */
 public final class ColumnFamily {
@@ -63,6 +64,24 @@ public boolean hasGCRule() {
     return !RuleCase.RULE_NOT_SET.equals(rule.toProto().getRuleCase());
   }
 
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) {
+      return true;
+    }
+    if (o == null || getClass() != o.getClass()) {
+      return false;
+    }
+    ColumnFamily that = (ColumnFamily) o;
+    return Objects.equal(id, that.id) &&
+        Objects.equal(rule, that.rule);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hashCode(id, rule);
+  }
+
   @Override
   public String toString() {
     return MoreObjects.toStringHelper(this).add("id", id).add("GCRule", rule).toString();

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

@@ -20,6 +20,7 @@
 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;
 
 /**
  * Wrapper for {@link GenerateConsistencyTokenResponse#getConsistencyToken()}
@@ -53,6 +54,23 @@ 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/main/java/com/google/cloud/bigtable/admin/v2/models/CreateTableRequest.java

@@ -20,6 +20,7 @@
 import com.google.bigtable.admin.v2.InstanceName;
 import com.google.bigtable.admin.v2.Table;
 import com.google.cloud.bigtable.admin.v2.models.GCRules.GCRule;
+import com.google.common.base.Objects;
 import com.google.common.base.Preconditions;
 import com.google.protobuf.ByteString;
 
@@ -90,6 +91,24 @@ public CreateTableRequest addSplit(ByteString key) {
     return this;
   }
 
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) {
+      return true;
+    }
+    if (o == null || getClass() != o.getClass()) {
+      return false;
+    }
+    CreateTableRequest that = (CreateTableRequest) o;
+    return Objects.equal(createTableRequest, that.createTableRequest) &&
+        Objects.equal(tableRequest, that.tableRequest);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hashCode(createTableRequest, tableRequest);
+  }
+
   @InternalApi
   public com.google.bigtable.admin.v2.CreateTableRequest toProto(InstanceName instanceName) {
     Preconditions.checkNotNull(instanceName);