Repository: GoogleCloudPlatform/cloud-bigtable-client
Pull request #2070 (2019-02-06 16:33:44)

Using the client wrapper in CloudBigtableIO


Patch:
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/async/AsyncExecutor.java b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/async/AsyncExecutor.java
index f427804eb..729fec74d 100644
--- a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/async/AsyncExecutor.java
+++ b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/async/AsyncExecutor.java
@@ -16,25 +16,14 @@
 package com.google.cloud.bigtable.grpc.async;
 
 import java.io.IOException;
-import java.util.List;
 
-import com.google.bigtable.v2.CheckAndMutateRowRequest;
-import com.google.bigtable.v2.CheckAndMutateRowResponse;
-import com.google.bigtable.v2.MutateRowRequest;
-import com.google.bigtable.v2.MutateRowResponse;
-import com.google.bigtable.v2.MutateRowsRequest;
-import com.google.bigtable.v2.MutateRowsResponse;
 import com.google.bigtable.v2.ReadModifyWriteRowRequest;
 import com.google.bigtable.v2.ReadModifyWriteRowResponse;
-import com.google.bigtable.v2.ReadRowsRequest;
-import com.google.bigtable.v2.Row;
 import com.google.cloud.bigtable.config.Logger;
 import com.google.cloud.bigtable.grpc.BigtableDataClient;
-import com.google.cloud.bigtable.grpc.scanner.FlatRow;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
-import com.google.protobuf.MessageLite;
 
 /**
  * This class provides management of asynchronous Bigtable RPCs. It ensures that there aren't too
@@ -49,80 +38,6 @@ public class AsyncExecutor {
   /** Constant <code>LOG</code> */
   protected static final Logger LOG = new Logger(AsyncExecutor.class);
 
-  protected interface AsyncCall<RequestT, ResponseT> {
-    ListenableFuture<ResponseT> call(BigtableDataClient client, RequestT request);
-  }
-
-  /**
-   * Calls {@link BigtableDataClient#mutateRowAsync(MutateRowRequest)}.
-   */
-  protected static AsyncCall<MutateRowRequest, MutateRowResponse> MUTATE_ROW_ASYNC =
-      new AsyncCall<MutateRowRequest, MutateRowResponse>() {
-        @Override
-        public ListenableFuture<MutateRowResponse> call(BigtableDataClient client, MutateRowRequest request) {
-          return client.mutateRowAsync(request);
-        }
-      };
-
-  /**
-   * Calls {@link BigtableDataClient#mutateRowsAsync(MutateRowsRequest)}.
-   */
-  protected static AsyncCall<MutateRowsRequest, List<MutateRowsResponse>> MUTATE_ROWS_ASYNC =
-      new AsyncCall<MutateRowsRequest, List<MutateRowsResponse>>() {
-        @Override
-        public ListenableFuture<List<MutateRowsResponse>> call(BigtableDataClient client,
-            MutateRowsRequest request) {
-          return client.mutateRowsAsync(request);
-        }
-      };
-
-  /**
-   * Calls {@link BigtableDataClient#readModifyWriteRowAsync(ReadModifyWriteRowRequest)}.
-   */
-  protected static AsyncCall<ReadModifyWriteRowRequest, ReadModifyWriteRowResponse> READ_MODIFY_WRITE_ASYNC =
-      new AsyncCall<ReadModifyWriteRowRequest, ReadModifyWriteRowResponse>() {
-        @Override
-        public ListenableFuture<ReadModifyWriteRowResponse> call(BigtableDataClient client,
-            ReadModifyWriteRowRequest request) {
-          return client.readModifyWriteRowAsync(request);
-        }
-      };
-
-  /**
-   * Calls {@link BigtableDataClient#checkAndMutateRowAsync(CheckAndMutateRowRequest)}.
-   */
-  protected static AsyncCall<CheckAndMutateRowRequest, CheckAndMutateRowResponse> CHECK_AND_MUTATE_ASYNC =
-      new AsyncCall<CheckAndMutateRowRequest, CheckAndMutateRowResponse>() {
-        @Override
-        public ListenableFuture<CheckAndMutateRowResponse> call(BigtableDataClient client,
-            CheckAndMutateRowRequest request) {
-          return client.checkAndMutateRowAsync(request);
-        }
-      };
-
-  /**
-   * Calls {@link BigtableDataClient#readRowsAsync(ReadRowsRequest)}.
-   */
-  protected static AsyncCall<ReadRowsRequest, List<Row>> READ_ROWS_ASYNC =
-      new AsyncCall<ReadRowsRequest, List<Row>>() {
-        @Override
-        public ListenableFuture<List<Row>> call(BigtableDataClient client, ReadRowsRequest request) {
-          return client.readRowsAsync(request);
-        }
-      };
-
-
-  /**
-   * Calls {@link BigtableDataClient#readRowsAsync(ReadRowsRequest)}.
-   */
-  protected static AsyncCall<ReadRowsRequest, List<FlatRow>> READ_FLAT_ROWS_ASYNC =
-      new AsyncCall<ReadRowsRequest, List<FlatRow>>() {
-        @Override
-        public ListenableFuture<List<FlatRow>> call(BigtableDataClient client, ReadRowsRequest request) {
-          return client.readFlatRowsAsync(request);
-        }
-      };
-
   private final BigtableDataClient client;
   private final OperationAccountant operationsAccountant;
 
@@ -152,48 +67,6 @@ public class AsyncExecutor {
     this.operationsAccountant = operationAccountant;
   }
 
-  /**
-   * Performs a {@link com.google.cloud.bigtable.grpc.BigtableDataClient#mutateRowAsync(MutateRowRequest)} on the
-   * {@link com.google.bigtable.v2.MutateRowRequest}. This method may block if
-   * {@link OperationAccountant#registerOperation(ListenableFuture)} blocks.
-   *
-   * @param request The {@link com.google.bigtable.v2.MutateRowRequest} to send.
-   * @return a {@link com.google.common.util.concurrent.ListenableFuture} which can be listened to for completion events.
-   * @throws java.lang.InterruptedException if any.
-   */
-  public ListenableFuture<MutateRowResponse> mutateRowAsync(MutateRowRequest request)
-      throws InterruptedException {
-    return call(MUTATE_ROW_ASYNC, request);
-  }
-
-  /**
-   * Performs a {@link com.google.cloud.bigtable.grpc.BigtableDataClient#mutateRowsAsync(MutateRowsRequest)} on the
-   * {@link com.google.bigtable.v2.MutateRowsRequest}. This method may block if
-   * {@link OperationAccountant#registerOperation(ListenableFuture)} blocks.
-   *
-   * @param request The {@link com.google.bigtable.v2.MutateRowRequest} to send.
-   * @return a {@link com.google.common.util.concurrent.ListenableFuture} which can be listened to for completion events.
-   * @throws java.lang.InterruptedException if any.
-   */
-  public ListenableFuture<List<MutateRowsResponse>> mutateRowsAsync(MutateRowsRequest request)
-      throws InterruptedException {
-    return call(MUTATE_ROWS_ASYNC, request);
-  }
-
-  /**
-   * Performs a {@link com.google.cloud.bigtable.grpc.BigtableDataClient#checkAndMutateRowAsync(CheckAndMutateRowRequest)} on the
-   * {@link com.google.bigtable.v2.CheckAndMutateRowRequest}. This method may block if
-   * {@link OperationAccountant#registerOperation(ListenableFuture)} blocks.
-   *
-   * @param request The {@link com.google.bigtable.v2.CheckAndMutateRowRequest} to send.
-   * @return a {@link com.google.common.util.concurrent.ListenableFuture} which can be listened to for completion events.
-   * @throws java.lang.InterruptedException if any.
-   */
-  public ListenableFuture<CheckAndMutateRowResponse> checkAndMutateRowAsync(
-      CheckAndMutateRowRequest request) throws InterruptedException {
-    return call(CHECK_AND_MUTATE_ASYNC, request);
-  }
-
   /**
    * Performs a {@link com.google.cloud.bigtable.grpc.BigtableDataClient#readModifyWriteRow(ReadModifyWriteRowRequest)} on the
    * {@link com.google.bigtable.v2.ReadModifyWriteRowRequest}. This method may block if
@@ -204,45 +77,12 @@ public class AsyncExecutor {
    * @throws java.lang.InterruptedException if any.
    */
   public ListenableFuture<ReadModifyWriteRowResponse>
-      readModifyWriteRowAsync(ReadModifyWriteRowRequest request) throws InterruptedException {
-    return call(READ_MODIFY_WRITE_ASYNC, request);
-  }
-
-  /**
-   * Performs a {@link com.google.cloud.bigtable.grpc.BigtableDataClient#readRowsAsync(ReadRowsRequest)} on the
-   * {@link com.google.bigtable.v2.ReadRowsRequest}. This method may block if
-   * {@link OperationAccountant#registerOperation(ListenableFuture)} blocks.
-   *
-   * @param request The {@link com.google.bigtable.v2.ReadRowsRequest} to send.
-   * @return a {@link com.google.common.util.concurrent.ListenableFuture} which can be listened to for completion events.
-   * @throws java.lang.InterruptedException if any.
-   */
-  public ListenableFuture<List<Row>> readRowsAsync(ReadRowsRequest request)
-      throws InterruptedException {
-    return call(READ_ROWS_ASYNC, request);
-  }
-
-  /**
-   * Performs a {@link com.google.cloud.bigtable.grpc.BigtableDataClient#readRowsAsync(ReadRowsRequest)} on the
-   * {@link com.google.bigtable.v2.ReadRowsRequest}. This method may block if
-   * {@link OperationAccountant#registerOperation(ListenableFuture)} blocks.
-   *
-   * @param request The {@link com.google.bigtable.v2.ReadRowsRequest} to send.
-   * @return a {@link com.google.common.util.concurrent.ListenableFuture} which can be listened to for completion events.
-   * @throws java.lang.InterruptedException if any.
-   */
-  public ListenableFuture<List<FlatRow>> readFlatRowsAsync(ReadRowsRequest request)
-      throws InterruptedException {
-    return call(READ_FLAT_ROWS_ASYNC, request);
-  }
-
-  private <RequestT extends MessageLite, ResponseT> ListenableFuture<ResponseT> call(
-      AsyncCall<RequestT, ResponseT> rpc, RequestT request) {
+      readModifyWriteRowAsync(ReadModifyWriteRowRequest request) {
     // Wait until both the memory and rpc count maximum requirements are achieved before getting a
     // unique id used to track this request.
-    ListenableFuture<ResponseT> future;
+    ListenableFuture<ReadModifyWriteRowResponse> future;
     try {
-      future = rpc.call(client, request);
+      future = client.readModifyWriteRowAsync(request);
     } catch (Throwable e) {
       future = Futures.immediateFailedFuture(e);
     }
@@ -278,23 +118,4 @@ public class AsyncExecutor {
     return operationsAccountant.hasInflightOperations();
   }
 
-  /**
-   * <p>
-   * Getter for the field <code>client</code>.
-   * </p>
-   * @return a {@link com.google.cloud.bigtable.grpc.BigtableDataClient} object.
-   */
-  public BigtableDataClient getClient() {
-    return client;
-  }
-
-  /**
-   * <p>
-   * Getter for the field <code>operationsAccountant</code>.
-   * </p>
-   * @return a {@link com.google.cloud.bigtable.grpc.async.OperationAccountant} object.
-   */
-  public OperationAccountant getOperationAccountant() {
-    return operationsAccountant;
-  }
 }
diff --git a/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/BatchExecutor.java b/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/BatchExecutor.java
index 83f5ba0f9..beaeb867d 100644
--- a/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/BatchExecutor.java
+++ b/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/BatchExecutor.java
@@ -220,9 +220,6 @@ public class BatchExecutor {
       } else if (row instanceof RowMutations) {
         return bulkOperation.bulkMutation.add(requestAdapter.adaptEntry((RowMutations) row));
       }
-    } catch (InterruptedException e) {
-      Thread.currentThread().interrupt();
-      return Futures.immediateFailedFuture(new IOException("Could not process the batch due to interrupt", e));
     } catch (Throwable e) {
       return Futures.immediateFailedFuture(new IOException("Could not process the batch", e));
     }
diff --git a/bigtable-dataflow-parent/bigtable-hbase-beam/src/main/java/com/google/cloud/bigtable/batch/common/CloudBigtableService.java b/bigtable-dataflow-parent/bigtable-hbase-beam/src/main/java/com/google/cloud/bigtable/batch/common/CloudBigtableService.java
index 1f2490901..20da82636 100755
--- a/bigtable-dataflow-parent/bigtable-hbase-beam/src/main/java/com/google/cloud/bigtable/batch/common/CloudBigtableService.java
+++ b/bigtable-dataflow-parent/bigtable-hbase-beam/src/main/java/com/google/cloud/bigtable/batch/common/CloudBigtableService.java
@@ -15,7 +15,7 @@
  */
 package com.google.cloud.bigtable.batch.common;
 
-import com.google.bigtable.repackaged.com.google.bigtable.v2.SampleRowKeysResponse;
+import com.google.bigtable.repackaged.com.google.cloud.bigtable.data.v2.models.KeyOffset;
 import com.google.cloud.bigtable.beam.CloudBigtableIO;
 import com.google.cloud.bigtable.beam.CloudBigtableTableConfiguration;
 
@@ -27,6 +27,5 @@ import java.util.List;
  * Cloud Bigtable
  */
 public interface CloudBigtableService {
-  List<SampleRowKeysResponse> getSampleRowKeys(CloudBigtableTableConfiguration config)
-      throws IOException;
+  List<KeyOffset> getSampleRowKeys(CloudBigtableTableConfiguration config) throws IOException;
 }
diff --git a/bigtable-dataflow-parent/bigtable-hbase-beam/src/main/java/com/google/cloud/bigtable/batch/common/CloudBigtableServiceImpl.java b/bigtable-dataflow-parent/bigtable-hbase-beam/src/main/java/com/google/cloud/bigtable/batch/common/CloudBigtableServiceImpl.java
index f07457eed..712876315 100755
--- a/bigtable-dataflow-parent/bigtable-hbase-beam/src/main/java/com/google/cloud/bigtable/batch/common/CloudBigtableServiceImpl.java
+++ b/bigtable-dataflow-parent/bigtable-hbase-beam/src/main/java/com/google/cloud/bigtable/batch/common/CloudBigtableServiceImpl.java
@@ -15,11 +15,8 @@
  */
 package com.google.cloud.bigtable.batch.common;
 
-import com.google.bigtable.repackaged.com.google.bigtable.v2.SampleRowKeysRequest;
-import com.google.bigtable.repackaged.com.google.bigtable.v2.SampleRowKeysResponse;
-import com.google.bigtable.repackaged.com.google.cloud.bigtable.config.BigtableOptions;
+import com.google.bigtable.repackaged.com.google.cloud.bigtable.data.v2.models.KeyOffset;
 import com.google.bigtable.repackaged.com.google.cloud.bigtable.grpc.BigtableSession;
-import com.google.bigtable.repackaged.com.google.cloud.bigtable.grpc.BigtableTableName;
 import com.google.cloud.bigtable.beam.CloudBigtableTableConfiguration;
 
 import java.io.IOException;
@@ -28,15 +25,10 @@ import java.util.List;
 public class CloudBigtableServiceImpl implements CloudBigtableService {
 
   @Override
-  public List<SampleRowKeysResponse> getSampleRowKeys(CloudBigtableTableConfiguration config)
+  public List<KeyOffset> getSampleRowKeys(CloudBigtableTableConfiguration config)
       throws IOException {
-    BigtableOptions bigtableOptions = config.toBigtableOptions();
-    try (BigtableSession session = new BigtableSession(bigtableOptions)) {
-      BigtableTableName tableName =
-          bigtableOptions.getInstanceName().toTableName(config.getTableId());
-      SampleRowKeysRequest request =
-          SampleRowKeysRequest.newBuilder().setTableName(tableName.toString()).build();
-      return session.getDataClient().sampleRowKeys(request);
+    try (BigtableSession session = new BigtableSession(config.toBigtableOptions())) {
+      return session.getClientWrapper().sampleRowKeys(config.getTableId());
     }
   }
 }
diff --git a/bigtable-dataflow-parent/bigtable-hbase-beam/src/main/java/com/google/cloud/bigtable/beam/CloudBigtableIO.java b/bigtable-dataflow-parent/bigtable-hbase-beam/src/main/java/com/google/cloud/bigtable/beam/CloudBigtableIO.java
index 651c0c54f..74c314957 100755
--- a/bigtable-dataflow-parent/bigtable-hbase-beam/src/main/java/com/google/cloud/bigtable/beam/CloudBigtableIO.java
+++ b/bigtable-dataflow-parent/bigtable-hbase-beam/src/main/java/com/google/cloud/bigtable/beam/CloudBigtableIO.java
@@ -26,6 +26,8 @@ import java.util.NoSuchElementException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
 
+import com.google.bigtable.repackaged.com.google.cloud.bigtable.data.v2.models.KeyOffset;
+import com.google.bigtable.repackaged.com.google.cloud.bigtable.data.v2.models.Query;
 import org.apache.beam.sdk.annotations.Experimental;
 import org.apache.beam.sdk.coders.CannotProvideCoderException;
 import org.apache.beam.sdk.coders.Coder;
@@ -66,7 +68,6 @@ import org.slf4j.LoggerFactory;
 
 import com.google.bigtable.repackaged.com.google.bigtable.v2.SampleRowKeysResponse;
 import com.google.bigtable.repackaged.com.google.cloud.bigtable.config.BulkOptions;
-import com.google.bigtable.repackaged.com.google.cloud.bigtable.grpc.BigtableDataClient;
 import com.google.bigtable.repackaged.com.google.cloud.bigtable.grpc.BigtableInstanceName;
 import com.google.bigtable.repackaged.com.google.cloud.bigtable.grpc.BigtableSession;
 import com.google.bigtable.repackaged.com.google.cloud.bigtable.grpc.BigtableSessionSharedThreadPools;
@@ -163,7 +164,7 @@ public class CloudBigtableIO {
      */
     private final CloudBigtableScanConfiguration configuration;
 
-    private transient List<SampleRowKeysResponse> sampleRowKeys;
+    private transient List<KeyOffset> sampleRowKeys;
 
     AbstractSource(CloudBigtableScanConfiguration configuration) {
       this.configuration = configuration;
@@ -185,8 +186,8 @@ public class CloudBigtableIO {
       List<SourceWithKeys> splits = new ArrayList<>();
       byte[] startKey = HConstants.EMPTY_START_ROW;
       long lastOffset = 0;
-      for (SampleRowKeysResponse response : getSampleRowKeys()) {
-        byte[] endKey = response.getRowKey().toByteArray();
+      for (KeyOffset response : getSampleRowKeys()) {
+        byte[] endKey = response.getKey().toByteArray();
         // Avoid empty regions.
         if (Bytes.equals(startKey, endKey) && startKey.length > 0) {
           continue;
@@ -274,11 +275,11 @@ public class CloudBigtableIO {
 
     /**
      * Performs a call to get sample row keys from
-     * {@link BigtableDataClient#sampleRowKeys(com.google.bigtable.repackaged.com.google.bigtable.v2.SampleRowKeysRequest)}
+     * {@link CloudBigtableServiceImpl#getSampleRowKeys(CloudBigtableTableConfiguration)}
      * if they are not yet cached. The sample row keys give information about tablet key boundaries
      * and estimated sizes.
      */
-    public synchronized List<SampleRowKeysResponse> getSampleRowKeys() throws IOException {
+    public synchronized List<KeyOffset> getSampleRowKeys() throws IOException {
       if (sampleRowKeys == null) {
         sampleRowKeys = new CloudBigtableServiceImpl().getSampleRowKeys(getConfiguration());
       }
@@ -286,7 +287,7 @@ public class CloudBigtableIO {
     }
 
     @VisibleForTesting
-    void setSampleRowKeys(List<SampleRowKeysResponse> sampleRowKeys) {
+    void setSampleRowKeys(List<KeyOffset> sampleRowKeys) {
       this.sampleRowKeys = sampleRowKeys;
     }
 
@@ -317,8 +318,8 @@ public class CloudBigtableIO {
 
       byte[] startKey = HConstants.EMPTY_START_ROW;
       long lastOffset = 0;
-      for (SampleRowKeysResponse response : getSampleRowKeys()) {
-        byte[] currentEndKey = response.getRowKey().toByteArray();
+      for (KeyOffset response : getSampleRowKeys()) {
+        byte[] currentEndKey = response.getKey().toByteArray();
         // Avoid empty regions.
         if (Bytes.equals(startKey, currentEndKey) && startKey.length != 0) {
           continue;
@@ -418,7 +419,7 @@ public class CloudBigtableIO {
     /**
      * Splits the table based on keys that belong to tablets, known as "regions" in the HBase API.
      * The current implementation uses the HBase {@link RegionLocator} interface, which calls
-     * {@link BigtableDataClient#sampleRowKeys(com.google.bigtable.repackaged.com.google.bigtable.v2.SampleRowKeysRequest)}
+     * {@link  CloudBigtableServiceImpl#getSampleRowKeys(CloudBigtableTableConfiguration)}
      * under the covers. A {@link SourceWithKeys} may correspond to a single region or a portion of
      * a region.
      * <p>
@@ -451,7 +452,7 @@ public class CloudBigtableIO {
 
     /**
      * Gets an estimated size based on data returned from
-     * {@link BigtableDataClient#sampleRowKeys(com.google.bigtable.repackaged.com.google.bigtable.v2.SampleRowKeysRequest)}.
+     * {@link CloudBigtableServiceImpl#getSampleRowKeys(CloudBigtableTableConfiguration)}.
      * The estimate will be high if a {@link Scan} is set on the
      * {@link CloudBigtableScanConfiguration}; in such cases, the estimate will not take the Scan
      * into account, and will return a larger estimate than what the {@link CloudBigtableIO.Reader}
@@ -468,8 +469,8 @@ public class CloudBigtableIO {
 
       byte[] startKey = HConstants.EMPTY_START_ROW;
       long lastOffset = 0;
-      for (SampleRowKeysResponse response : getSampleRowKeys()) {
-        byte[] currentEndKey = response.getRowKey().toByteArray();
+      for (KeyOffset response : getSampleRowKeys()) {
+        byte[] currentEndKey = response.getKey().toByteArray();
         // Avoid empty regions.
         if (Bytes.equals(startKey, currentEndKey) && startKey.length != 0) {
           continue;
@@ -620,7 +621,7 @@ public class CloudBigtableIO {
 
       // This will use cached data channels under the covers.
       session = new BigtableSession(BigtableOptionsFactory.fromConfiguration(config));
-      scanner = session.getDataClient().readFlatRows(source.getConfiguration().getRequest());
+      scanner = session.getClientWrapper().readFlatRows(Query.fromProto(source.getConfiguration().getRequest()));
     }
 
     /**


Test patch:
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/async/TestAsyncExecutor.java b/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/async/TestAsyncExecutor.java
index e2e2c5683..8cec21fde 100644
--- a/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/async/TestAsyncExecutor.java
+++ b/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/async/TestAsyncExecutor.java
@@ -62,26 +62,9 @@ public class TestAsyncExecutor {
     Assert.assertFalse(underTest.hasInflightRequests());
   }
 
-  @Test
-  public void testMutation() throws InterruptedException {
-    when(client.mutateRowAsync(any(MutateRowRequest.class))).thenReturn(future);
-    underTest.mutateRowAsync(MutateRowRequest.getDefaultInstance());
-    Assert.assertTrue(underTest.hasInflightRequests());
-    future.set("");
-    Assert.assertFalse(underTest.hasInflightRequests());
-  }
-
-  @Test
-  public void testCheckAndMutate() throws InterruptedException {
-    when(client.checkAndMutateRowAsync(any(CheckAndMutateRowRequest.class))).thenReturn(future);
-    underTest.checkAndMutateRowAsync(CheckAndMutateRowRequest.getDefaultInstance());
-    Assert.assertTrue(underTest.hasInflightRequests());
-    future.set("");
-    Assert.assertFalse(underTest.hasInflightRequests());
-  }
 
   @Test
-  public void testReadWriteModify() throws InterruptedException {
+  public void testReadWriteModify()  {
     when(client.readModifyWriteRowAsync(any(ReadModifyWriteRowRequest.class))).thenReturn(future);
     underTest.readModifyWriteRowAsync(ReadModifyWriteRowRequest.getDefaultInstance());
     Assert.assertTrue(underTest.hasInflightRequests());
@@ -90,30 +73,9 @@ public class TestAsyncExecutor {
   }
 
   @Test
-  public void testReadRowsAsync() throws InterruptedException {
-    when(client.readRowsAsync(any(ReadRowsRequest.class))).thenReturn(future);
-    underTest.readRowsAsync(ReadRowsRequest.getDefaultInstance());
-    Assert.assertTrue(underTest.hasInflightRequests());
-    future.set("");
-    Assert.assertFalse(underTest.hasInflightRequests());
-  }
-
-  @Test
-  public void testReadFlatRowsAsync() throws InterruptedException {
-    when(client.readFlatRowsAsync(any(ReadRowsRequest.class))).thenReturn(future);
-    underTest.readFlatRowsAsync(ReadRowsRequest.getDefaultInstance());
-    Assert.assertTrue(underTest.hasInflightRequests());
-    future.set("");
-    Assert.assertFalse(underTest.hasInflightRequests());
-  }
-
-  @Test
-  public void testInvalidMutation() throws Exception {
-    try {
-      when(client.mutateRowAsync(any(MutateRowRequest.class))).thenThrow(new RuntimeException());
-      underTest.mutateRowAsync(MutateRowRequest.getDefaultInstance());
-    } catch(Exception ignored) {
-    }
+  public void testInvalidMutation() {
+    when(client.readModifyWriteRowAsync(any(ReadModifyWriteRowRequest.class))).thenThrow(new RuntimeException());
+    underTest.readModifyWriteRowAsync(ReadModifyWriteRowRequest.getDefaultInstance());
     future.set(ReadRowsResponse.getDefaultInstance());
     Assert.assertFalse(underTest.hasInflightRequests());
   }
diff --git a/bigtable-client-core-parent/bigtable-hbase/src/test/java/com/google/cloud/bigtable/hbase/TestBatchExecutor.java b/bigtable-client-core-parent/bigtable-hbase/src/test/java/com/google/cloud/bigtable/hbase/TestBatchExecutor.java
index 3e7b71b6b..9164a0d2d 100644
--- a/bigtable-client-core-parent/bigtable-hbase/src/test/java/com/google/cloud/bigtable/hbase/TestBatchExecutor.java
+++ b/bigtable-client-core-parent/bigtable-hbase/src/test/java/com/google/cloud/bigtable/hbase/TestBatchExecutor.java
@@ -24,7 +24,6 @@ import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
-import com.google.bigtable.v2.MutateRowRequest;
 import com.google.bigtable.v2.MutateRowsRequest;
 import com.google.bigtable.v2.ReadModifyWriteRowRequest;
 import com.google.bigtable.v2.ReadRowsRequest;
@@ -70,7 +69,6 @@ import org.mockito.MockitoAnnotations;
 import org.mockito.invocation.InvocationOnMock;
 import org.mockito.stubbing.Answer;
 
-import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
@@ -134,8 +132,9 @@ public class TestBatchExecutor {
   private HBaseRequestAdapter requestAdapter;
 
   private BigtableOptions options;
+
   @Before
-  public void setup() throws InterruptedException {
+  public void setup() {
     options = BigtableOptions.builder()
         .setProjectId("projectId")
         .setInstanceId("instanceId")
@@ -196,28 +195,25 @@ public class TestBatchExecutor {
 
   @Test
   public void testShutdownService() throws Exception {
-    when(mockAsyncExecutor.mutateRowAsync(any(MutateRowRequest.class)))
-        .thenThrow(new IllegalStateException("closed"));
+    when(mockFuture.get()).thenThrow(new IllegalStateException("closed"));
     try {
       batch(Arrays.asList(randomPut()));
     } catch (RetriesExhaustedWithDetailsException e) {
       Assert.assertEquals(1, e.getCauses().size());
-      Assert.assertEquals(IOException.class, e.getCause(0).getClass());
+      Assert.assertEquals(IllegalStateException.class, e.getCause(0).getClass());
     }
   }
 
   @Test
   public void testAsyncException() throws Exception {
     String message = "Something bad happened";
-    when(mockAsyncExecutor.mutateRowAsync(any(MutateRowRequest.class)))
-        .thenThrow(new RuntimeException(message));
+    when(mockFuture.get()).thenThrow(new RuntimeException(message));
     try {
       batch(Arrays.asList(randomPut()));
     } catch (RetriesExhaustedWithDetailsException e) {
       Assert.assertEquals(1, e.getCauses().size());
-      Assert.assertEquals(IOException.class, e.getCause(0).getClass());
-      Assert.assertEquals(RuntimeException.class, e.getCause(0).getCause().getClass());
-      Assert.assertEquals(message, e.getCause(0).getCause().getMessage());
+      Assert.assertEquals(RuntimeException.class, e.getCause(0).getClass());
+      Assert.assertEquals(message, e.getCause(0).getMessage());
     }
   }
 
diff --git a/bigtable-client-core-parent/bigtable-hbase/src/test/java/com/google/cloud/bigtable/hbase/TestBigtableBufferedMutator.java b/bigtable-client-core-parent/bigtable-hbase/src/test/java/com/google/cloud/bigtable/hbase/TestBigtableBufferedMutator.java
index 7ebc3fa3c..a9f6eb444 100644
--- a/bigtable-client-core-parent/bigtable-hbase/src/test/java/com/google/cloud/bigtable/hbase/TestBigtableBufferedMutator.java
+++ b/bigtable-client-core-parent/bigtable-hbase/src/test/java/com/google/cloud/bigtable/hbase/TestBigtableBufferedMutator.java
@@ -42,7 +42,6 @@ import org.junit.runners.JUnit4;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
-import com.google.bigtable.v2.MutateRowRequest;
 import com.google.bigtable.v2.MutateRowResponse;
 import com.google.bigtable.v2.MutateRowsRequest;
 import com.google.bigtable.v2.ReadModifyWriteRowRequest;
@@ -112,25 +111,23 @@ public class TestBigtableBufferedMutator {
   }
 
   @Test
-  public void testPut() throws IOException, InterruptedException {
+  public void testPut() throws IOException{
     when(mockBulkMutation.add(any(MutateRowsRequest.Entry.class))).thenReturn(future);
     BigtableBufferedMutator underTest = createMutator(new Configuration(false));
     underTest.mutate(SIMPLE_PUT);
-    verify(mockAsyncExecutor, times(0)).mutateRowAsync(any(MutateRowRequest.class));
     verify(mockBulkMutation, times(1)).add(any(MutateRowsRequest.Entry.class));
   }
 
   @Test
-  public void testDelete() throws IOException, InterruptedException {
+  public void testDelete() throws IOException {
     when(mockBulkMutation.add(any(MutateRowsRequest.Entry.class))).thenReturn(future);
     BigtableBufferedMutator underTest = createMutator(new Configuration(false));
     underTest.mutate(new Delete(EMPTY_BYTES));
-    verify(mockAsyncExecutor, times(0)).mutateRowAsync(any(MutateRowRequest.class));
     verify(mockBulkMutation, times(1)).add(any(MutateRowsRequest.Entry.class));
   }
 
   @Test
-  public void testIncrement() throws IOException, InterruptedException {
+  public void testIncrement() throws IOException {
     when(mockAsyncExecutor.readModifyWriteRowAsync(any(ReadModifyWriteRowRequest.class)))
         .thenReturn(future);
     BigtableBufferedMutator underTest = createMutator(new Configuration(false));
@@ -140,7 +137,7 @@ public class TestBigtableBufferedMutator {
   }
 
   @Test
-  public void testAppend() throws IOException, InterruptedException {
+  public void testAppend() throws IOException {
     when(mockAsyncExecutor.readModifyWriteRowAsync(any(ReadModifyWriteRowRequest.class)))
         .thenReturn(future);
     BigtableBufferedMutator underTest = createMutator(new Configuration(false));
@@ -168,7 +165,6 @@ public class TestBigtableBufferedMutator {
     when(mockBulkMutation.add(any(MutateRowsRequest.Entry.class))).thenReturn(future);
     final BigtableBufferedMutator underTest = createMutator(config);
     underTest.mutate(SIMPLE_PUT);
-    verify(mockAsyncExecutor, times(0)).mutateRowAsync(any(MutateRowRequest.class));
     verify(mockBulkMutation, times(1)).add(any(MutateRowsRequest.Entry.class));
     underTest.flush();
     verify(mockBulkMutation, times(1)).flush();
diff --git a/bigtable-dataflow-parent/bigtable-hbase-beam/src/test/java/com/google/cloud/bigtable/beam/CloudBigtableIOIntegrationTest.java b/bigtable-dataflow-parent/bigtable-hbase-beam/src/test/java/com/google/cloud/bigtable/beam/CloudBigtableIOIntegrationTest.java
index 315218788..621e9604a 100755
--- a/bigtable-dataflow-parent/bigtable-hbase-beam/src/test/java/com/google/cloud/bigtable/beam/CloudBigtableIOIntegrationTest.java
+++ b/bigtable-dataflow-parent/bigtable-hbase-beam/src/test/java/com/google/cloud/bigtable/beam/CloudBigtableIOIntegrationTest.java
@@ -17,6 +17,7 @@
 
 import com.google.bigtable.repackaged.com.google.bigtable.v2.SampleRowKeysResponse;
 import com.google.bigtable.repackaged.com.google.cloud.bigtable.config.Logger;
+import com.google.bigtable.repackaged.com.google.cloud.bigtable.data.v2.models.KeyOffset;
 import com.google.cloud.bigtable.hbase.BigtableConfiguration;
 import org.apache.beam.sdk.io.BoundedSource;
 import org.apache.beam.sdk.io.BoundedSource.BoundedReader;
@@ -252,10 +253,10 @@ public class CloudBigtableIOIntegrationTest {
       try {
         CloudBigtableIO.Source source =
             (CloudBigtableIO.Source) CloudBigtableIO.read(createScanConfig(tableName));
-        List<SampleRowKeysResponse> sampleRowKeys = source.getSampleRowKeys();
+        List<KeyOffset> sampleRowKeys = source.getSampleRowKeys();
         LOG.info("Creating BoundedSource in testEstimatedAndSplitForSmallTable()");
         long estimatedSizeBytes = source.getEstimatedSizeBytes(null);
-        SampleRowKeysResponse lastSample = sampleRowKeys.get(sampleRowKeys.size() - 1);
+        KeyOffset lastSample = sampleRowKeys.get(sampleRowKeys.size() - 1);
         Assert.assertEquals(lastSample.getOffsetBytes(), estimatedSizeBytes);
 
         LOG.info("Creating Bundles in testEstimatedAndSplitForSmallTable()");
@@ -294,10 +295,10 @@ public class CloudBigtableIOIntegrationTest {
         LOG.info("Getting Source in testEstimatedAndSplitForLargeTable()");
         CloudBigtableIO.Source source =
             (CloudBigtableIO.Source) CloudBigtableIO.read(createScanConfig(tableName));
-        List<SampleRowKeysResponse> sampleRowKeys = source.getSampleRowKeys();
+        List<KeyOffset> sampleRowKeys = source.getSampleRowKeys();
         LOG.info("Getting estimated size in testEstimatedAndSplitForLargeTable()");
         long estimatedSizeBytes = source.getEstimatedSizeBytes(null);
-        SampleRowKeysResponse lastSample = sampleRowKeys.get(sampleRowKeys.size() - 1);
+        KeyOffset lastSample = sampleRowKeys.get(sampleRowKeys.size() - 1);
         Assert.assertEquals(lastSample.getOffsetBytes(), estimatedSizeBytes);
 
         LOG.info("Getting Bundles in testEstimatedAndSplitForLargeTable()");
diff --git a/bigtable-dataflow-parent/bigtable-hbase-beam/src/test/java/com/google/cloud/bigtable/beam/CloudBigtableIOTest.java b/bigtable-dataflow-parent/bigtable-hbase-beam/src/test/java/com/google/cloud/bigtable/beam/CloudBigtableIOTest.java
index 27589c057..0ec463651 100755
--- a/bigtable-dataflow-parent/bigtable-hbase-beam/src/test/java/com/google/cloud/bigtable/beam/CloudBigtableIOTest.java
+++ b/bigtable-dataflow-parent/bigtable-hbase-beam/src/test/java/com/google/cloud/bigtable/beam/CloudBigtableIOTest.java
@@ -25,6 +25,7 @@ import java.util.Comparator;
 import java.util.Iterator;
 import java.util.List;
 
+import com.google.bigtable.repackaged.com.google.cloud.bigtable.data.v2.models.KeyOffset;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.coders.CannotProvideCoderException;
 import org.apache.beam.sdk.coders.Coder;
@@ -109,7 +110,7 @@ public class CloudBigtableIOTest {
 
   @Test
   public void testSampleRowKeys() throws Exception {
-    List<SampleRowKeysResponse> sampleRowKeys = new ArrayList<>();
+    List<KeyOffset> sampleRowKeys = new ArrayList<>();
     int count = (int) (AbstractSource.COUNT_MAX_SPLIT_COUNT * 3 - 5);
     byte[][] keys = Bytes.split("A".getBytes(), "Z".getBytes(), count-2);
     long tabletSize = 2L * 1024L * 1024L * 1024L;
@@ -117,10 +118,7 @@ public class CloudBigtableIOTest {
     for (byte[] currentKey : keys) {
       boundary += tabletSize;
       try {
-        sampleRowKeys.add(SampleRowKeysResponse.newBuilder()
-          .setRowKey(ByteString.copyFrom(currentKey))
-          .setOffsetBytes(boundary)
-          .build());
+        sampleRowKeys.add(KeyOffset.create(ByteString.copyFrom(currentKey), boundary));
       } catch (NoClassDefFoundError e) {
         // This could cause some problems for javadoc or cobertura because of the shading magic we
         // do.