Repository: GoogleCloudPlatform/cloud-bigtable-client
Pull request #2047 (2019-01-15 15:40:03)

Adding More methods in IBigtableDataClient
## What changes this PR contains

- Added below delegate methods in IBigtableDataClient & it's implementations
  - sampleRowKeys 
  - sampleRowKeysAsync
   - readRows
   - readRowsAsync
   - readFlatRowsList
  - readFlatRows
   - readFlatRowsAsync

 - Added FlatRowConverter#convertToModelRow to enable FlatRow to models.Row conversion.
This addition will enable us to replace grpc.BigtableDataClient reference from AsyncExecutor & other places as well.

rahulKQL: @sduskis @igorbernstein2 

My apologies for rebase history. Have rebased & updated DataClientWrapper to comply with GCJ-0.78-alpha version (basically KeyOffSet#getKey update).

Please have a look.


Patch:
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/core/IBigtableDataClient.java b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/core/IBigtableDataClient.java
index 2ac0c3c17..33b48f833 100644
--- a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/core/IBigtableDataClient.java
+++ b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/core/IBigtableDataClient.java
@@ -16,14 +16,20 @@
 package com.google.cloud.bigtable.core;
 
 import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation;
+import com.google.cloud.bigtable.data.v2.models.KeyOffset;
+import com.google.cloud.bigtable.data.v2.models.Query;
 import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow;
 import com.google.cloud.bigtable.data.v2.models.Row;
 import com.google.cloud.bigtable.data.v2.models.RowMutation;
+import com.google.cloud.bigtable.grpc.scanner.FlatRow;
+import com.google.cloud.bigtable.grpc.scanner.ResultScanner;
 import com.google.common.util.concurrent.ListenableFuture;
+import java.util.List;
 import java.util.concurrent.ExecutionException;
 
 /**
- * Interface to access Bigtable data service api.
+ * Interface to wrap {@link com.google.cloud.bigtable.grpc.BigtableDataClient} with
+ * Google-Cloud-java's models.
  */
 public interface IBigtableDataClient {
 
@@ -31,10 +37,8 @@ public interface IBigtableDataClient {
    * Mutate a row atomically.
    *
    * @param rowMutation a {@link RowMutation} model object.
-   * @throws ExecutionException if any.
-   * @throws InterruptedException if any.
    */
-  void mutateRow(RowMutation rowMutation) throws ExecutionException, InterruptedException;
+  void mutateRow(RowMutation rowMutation);
 
   /**
    * Mutate a row atomically.
@@ -42,21 +46,17 @@ public interface IBigtableDataClient {
    * @param rowMutation a {@link RowMutation} model object.
    * @return a {@link ListenableFuture} of type {@link Void} will be set when request is
    *     successful otherwise exception will be thrown.
-   * @throws InterruptedException if any.
    */
   //TODO(rahulkql): Once it is adapted to v2.models, change the return type to ApiFuture.
-  ListenableFuture<Void> mutateRowAsync(RowMutation rowMutation) throws InterruptedException;
+  ListenableFuture<Void> mutateRowAsync(RowMutation rowMutation);
 
   /**
    * Perform an atomic read-modify-write operation on a row.
    *
    * @param readModifyWriteRow a {@link ReadModifyWriteRow} model object.
    * @return {@link Row} a modified row.
-   * @throws ExecutionException if any.
-   * @throws InterruptedException if any.
    */
-  Row readModifyWriteRow(ReadModifyWriteRow readModifyWriteRow)
-      throws ExecutionException, InterruptedException;
+  Row readModifyWriteRow(ReadModifyWriteRow readModifyWriteRow);
 
   /**
    * Perform an atomic read-modify-write operation on a row.
@@ -64,10 +64,9 @@ public interface IBigtableDataClient {
    * @param readModifyWriteRow a {@link ReadModifyWriteRow} model object.
    * @return a {@link ListenableFuture} of type {@link Row} will be set when request is
    *     successful otherwise exception will be thrown.
-   * @throws InterruptedException if any.
    */
   //TODO(rahulkql): Once it is adapted to v2.models, change the return type to ApiFuture.
-  ListenableFuture<Row> readModifyWriteRowAsync(ReadModifyWriteRow readModifyWriteRow) throws InterruptedException;
+  ListenableFuture<Row> readModifyWriteRowAsync(ReadModifyWriteRow readModifyWriteRow);
 
   /**
    * Creates {@link IBulkMutation} batcher.
@@ -92,6 +91,66 @@ public interface IBigtableDataClient {
    * @throws ExecutionException if any.
    * @throws InterruptedException if any.
    */
-  Boolean checkAndMutateRow(ConditionalRowMutation conditionalRowMutation)
-      throws ExecutionException, InterruptedException;
+  Boolean checkAndMutateRow(ConditionalRowMutation conditionalRowMutation);
+
+  /**
+   * Sample row keys from a table.
+   *
+   * @param tableId a String object.
+   * @return an immutable {@link List} object.
+   */
+  List<KeyOffset> sampleRowKeys(String tableId);
+
+  /**
+   * Sample row keys from a table, returning a Future that will complete when the sampling has
+   * completed.
+   *
+   * @param tableId a String object.
+   * @return a {@link ListenableFuture} object.
+   */
+  //TODO(rahulkql): Once it is adapted to v2.models, change the return type to ApiFuture.
+  ListenableFuture< List<KeyOffset>> sampleRowKeysAsync(String tableId);
+
+  /**
+   * Perform a scan over {@link Row}s, in key order.
+   *
+   * @param request a {@link Query} object.
+   * @return a {@link Row} object.
+   */
+  ResultScanner<Row> readRows(Query request);
+
+  /**
+   * Read multiple {@link Row}s into an in-memory list, in key order.
+   *
+   * @return a {@link ListenableFuture} that will finish when
+   * all reads have completed.
+   * @param request a {@link Query} object.
+   */
+  ListenableFuture<List<Row>> readRowsAsync(Query request);
+
+  /**
+   * Returns a list of {@link FlatRow}s, in key order.
+   *
+   * @param request a {@link Query} object.
+   * @return a List with {@link FlatRow}s.
+   */
+  List<FlatRow> readFlatRowsList(Query request);
+
+  /**
+   * Perform a scan over {@link FlatRow}s, in key order.
+   *
+   * @param request a {@link Query} object.
+   * @return a {@link ResultScanner} object.
+   */
+  ResultScanner<FlatRow> readFlatRows(Query request);
+
+  /**
+   * Read multiple {@link FlatRow}s into an in-memory list, in key order.
+   *
+   * @return a {@link ListenableFuture} that will finish when
+   * all reads have completed.
+   * @param request a {@link Query} object.
+   */
+  //TODO(rahulkql): Once it is adapted to v2.models, change the return type to ApiFuture.
+  ListenableFuture<List<FlatRow>> readFlatRowsAsync(Query request);
 }
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/BigtableDataClientWrapper.java b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/BigtableDataClientWrapper.java
index bab3a03d3..6e96c49b1 100644
--- a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/BigtableDataClientWrapper.java
+++ b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/BigtableDataClientWrapper.java
@@ -23,26 +23,37 @@ import com.google.bigtable.v2.Family;
 import com.google.bigtable.v2.MutateRowRequest;
 import com.google.bigtable.v2.MutateRowResponse;
 import com.google.bigtable.v2.ReadModifyWriteRowResponse;
+import com.google.bigtable.v2.SampleRowKeysRequest;
+import com.google.bigtable.v2.SampleRowKeysResponse;
 import com.google.cloud.bigtable.config.BigtableOptions;
 import com.google.cloud.bigtable.core.IBigtableDataClient;
 import com.google.cloud.bigtable.core.IBulkMutation;
+import com.google.cloud.bigtable.data.v2.internal.NameUtil;
 import com.google.cloud.bigtable.data.v2.internal.RequestContext;
 import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation;
-import com.google.cloud.bigtable.data.v2.models.InstanceName;
+import com.google.cloud.bigtable.data.v2.models.KeyOffset;
+import com.google.cloud.bigtable.data.v2.models.Query;
 import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow;
 import com.google.cloud.bigtable.data.v2.models.Row;
 import com.google.cloud.bigtable.data.v2.models.RowCell;
 import com.google.cloud.bigtable.data.v2.models.RowMutation;
+import com.google.cloud.bigtable.grpc.scanner.FlatRow;
+import com.google.cloud.bigtable.grpc.scanner.FlatRowConverter;
+import com.google.cloud.bigtable.grpc.scanner.ResultScanner;
 import com.google.common.base.Function;
 import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.MoreExecutors;
+import java.io.IOException;
+import java.util.Collections;
+import java.util.List;
+import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 
 /**
- * This class implements the {@link IBigtableDataClient} interface which provides access to google cloud
- * java.
+ * This class implements the {@link IBigtableDataClient} interface and wraps
+ * {@link BigtableDataClient} with Google-cloud-java's models.
  */
 public class BigtableDataClientWrapper implements IBigtableDataClient {
 
@@ -52,11 +63,8 @@ public class BigtableDataClientWrapper implements IBigtableDataClient {
   public BigtableDataClientWrapper(BigtableDataClient bigtableDataClient,
       BigtableOptions options) {
     this.delegate = bigtableDataClient;
-    this.requestContext = RequestContext
-        .create(InstanceName.of(options.getProjectId(),
-            options.getInstanceId()),
-            options.getAppProfileId()
-        );
+    this.requestContext = RequestContext.create(options.getProjectId(), options.getInstanceId(),
+        options.getAppProfileId());
   }
 
   /** {@inheritDoc} */
@@ -131,6 +139,121 @@ public class BigtableDataClientWrapper implements IBigtableDataClient {
     return response.getPredicateMatched();
   }
 
+  /** {@inheritDoc} */
+  @Override
+  public List<KeyOffset> sampleRowKeys(String tableId) {
+    String fullTableName = NameUtil
+        .formatTableName(requestContext.getProjectId(), requestContext.getInstanceId(), tableId);
+    SampleRowKeysRequest requestProto = SampleRowKeysRequest.newBuilder()
+        .setTableName(fullTableName)
+        .build();
+    List<SampleRowKeysResponse> responseProto = delegate.sampleRowKeys(requestProto);
+    ImmutableList.Builder<KeyOffset> keyOffsetBuilder =
+        ImmutableList.builderWithExpectedSize(responseProto.size());
+    for(SampleRowKeysResponse rowKeys : responseProto){
+      keyOffsetBuilder.add(KeyOffset.create(rowKeys.getRowKey(), rowKeys.getOffsetBytes()));
+    }
+
+    return keyOffsetBuilder.build();
+  }
+
+  /** {@inheritDoc} */
+  @Override
+  public ListenableFuture<List<KeyOffset>> sampleRowKeysAsync(String tableId) {
+    String fullTableName = NameUtil
+        .formatTableName(requestContext.getProjectId(), requestContext.getInstanceId(), tableId);
+    SampleRowKeysRequest requestProto =
+        SampleRowKeysRequest.newBuilder().setTableName(fullTableName).build();
+    ListenableFuture<List<SampleRowKeysResponse>> responseProto =
+        delegate.sampleRowKeysAsync(requestProto);
+
+    return Futures.transform(responseProto, new Function<List<SampleRowKeysResponse>, List<KeyOffset>>() {
+      @Override
+      public List<KeyOffset> apply(@Nonnull List<SampleRowKeysResponse> rowKeysList) {
+        if(rowKeysList == null || rowKeysList.isEmpty()){
+          return Collections.EMPTY_LIST;
+        }
+        ImmutableList.Builder<KeyOffset> keyOffsetBuilder =
+            ImmutableList.builderWithExpectedSize(rowKeysList.size());
+        for(SampleRowKeysResponse rowKeys : rowKeysList){
+          keyOffsetBuilder.add(KeyOffset.create(rowKeys.getRowKey(), rowKeys.getOffsetBytes()));
+        }
+
+        return keyOffsetBuilder.build();
+      }
+    }, MoreExecutors.directExecutor());
+  }
+
+  /** {@inheritDoc} */
+  @Override
+  public ResultScanner<Row> readRows(Query request) {
+    final ResultScanner<FlatRow> delegate = readFlatRows(request);
+    return new ResultScanner<Row>() {
+
+      @Override
+      public Row next() throws IOException {
+        return FlatRowConverter.convertToModelRow(delegate.next());
+      }
+
+      @Override
+      public Row[] next(int count) throws IOException {
+        FlatRow[] flatRows = delegate.next(count);
+        Row[] rows = new Row[flatRows.length];
+        for (int i = 0; i < flatRows.length; i++) {
+          rows[i] = FlatRowConverter.convertToModelRow(flatRows[i]);
+        }
+        return rows;
+      }
+
+      @Override
+      public void close() throws IOException {
+        delegate.close();
+      }
+
+      @Override
+      public int available() {
+        return delegate.available();
+      }
+    };
+  }
+
+  /** {@inheritDoc} */
+  @Override
+  public ListenableFuture<List<Row>> readRowsAsync(Query request) {
+    ListenableFuture<List<FlatRow>> responseProto = readFlatRowsAsync(request);
+
+    return Futures.transform(responseProto, new Function<List<FlatRow>, List<Row>>() {
+      @Override
+      public List<Row> apply(List<FlatRow> flatRowList) {
+        ImmutableList.Builder<Row> rowBuilder =
+            ImmutableList.builderWithExpectedSize(flatRowList.size());
+        for(FlatRow flatRow : flatRowList){
+          rowBuilder.add(FlatRowConverter.convertToModelRow(flatRow));
+        }
+
+        return rowBuilder.build();
+      }
+    }, MoreExecutors.directExecutor());
+  }
+
+  /** {@inheritDoc} */
+  @Override
+  public List<FlatRow> readFlatRowsList(Query request) {
+    return delegate.readFlatRowsList(request.toProto(requestContext));
+  }
+
+  /** {@inheritDoc} */
+  @Override
+  public ResultScanner<FlatRow> readFlatRows(Query request) {
+    return delegate.readFlatRows(request.toProto(requestContext));
+  }
+
+  /** {@inheritDoc} */
+  @Override
+  public ListenableFuture<List<FlatRow>> readFlatRowsAsync(Query request) {
+    return delegate.readFlatRowsAsync(request.toProto(requestContext));
+  }
+
   /**
    * This method converts instances of {@link ReadModifyWriteRowResponse} to {@link Row}.
    *
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/scanner/FlatRowConverter.java b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/scanner/FlatRowConverter.java
index bd60041f4..e776b288f 100644
--- a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/scanner/FlatRowConverter.java
+++ b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/scanner/FlatRowConverter.java
@@ -19,6 +19,8 @@ import com.google.bigtable.v2.Cell;
 import com.google.bigtable.v2.Column;
 import com.google.bigtable.v2.Family;
 import com.google.bigtable.v2.Row;
+import com.google.cloud.bigtable.data.v2.models.RowCell;
+import com.google.common.collect.ImmutableList;
 import com.google.protobuf.ByteString;
 
 /**
@@ -95,4 +97,23 @@ public class FlatRowConverter {
     }
     return builder.build();
   }
-}
\ No newline at end of file
+
+  public static com.google.cloud.bigtable.data.v2.models.Row convertToModelRow(FlatRow row) {
+    if (row == null || row.getCells() == null) {
+      return null;
+    }
+    ImmutableList.Builder<RowCell> rowCellList =
+        ImmutableList.builderWithExpectedSize(row.getCells().size());
+    for (FlatRow.Cell cell : row.getCells()) {
+      rowCellList.add(toRowCell(cell));
+    }
+
+    return com.google.cloud.bigtable.data.v2.models.Row.create(row.getRowKey(),
+        rowCellList.build());
+  }
+
+  private static RowCell toRowCell(FlatRow.Cell cell) {
+    return RowCell.create(cell.getFamily(), cell.getQualifier(), cell.getTimestamp(),
+        cell.getLabels(), cell.getValue());
+  }
+}
diff --git a/bigtable-hbase-1.x-parent/bigtable-hbase-1.x-shaded/pom.xml b/bigtable-hbase-1.x-parent/bigtable-hbase-1.x-shaded/pom.xml
index 6aeb2a7da..6e25d8ac4 100644
--- a/bigtable-hbase-1.x-parent/bigtable-hbase-1.x-shaded/pom.xml
+++ b/bigtable-hbase-1.x-parent/bigtable-hbase-1.x-shaded/pom.xml
@@ -59,6 +59,11 @@ limitations under the License.
     <artifactId>metrics-core</artifactId>
     <version>${dropwizard.metrics.version}</version>
   </dependency>
+  <dependency>
+    <groupId>org.checkerframework</groupId>
+    <artifactId>checker-compat-qual</artifactId>
+    <version>${checkerframework.version}</version>
+  </dependency>
 
   <dependency>
     <groupId>io.opencensus</groupId>
@@ -166,6 +171,8 @@ limitations under the License.
                   <exclude>junit:junit</exclude>
                   <exclude>org.hamcrest:hamcrest-core</exclude>
                   <exclude>javax.inject:javax.inject</exclude>
+                  <exclude>javax.annotation:javax.annotation-api</exclude>
+                  <exclude>org.checkerframework:checker-compat-qual</exclude>
                 </excludes>
               </artifactSet>
               <relocations>
@@ -246,6 +253,14 @@ limitations under the License.
                   <pattern>org.apache.commons.codec</pattern>
                   <shadedPattern>com.google.bigtable.repackaged.org.apache.commons.codec</shadedPattern>
                 </relocation>
+                <relocation>
+                  <pattern>org.threeten</pattern>
+                  <shadedPattern>com.google.bigtable.repackaged.org.threeten</shadedPattern>
+                </relocation>
+                <relocation>
+                  <pattern>org.codehaus</pattern>
+                  <shadedPattern>com.google.bigtable.repackaged.org.codehaus</shadedPattern>
+                </relocation>
               </relocations>
             </configuration>
           </execution>
diff --git a/bigtable-hbase-2.x-parent/bigtable-hbase-2.x-shaded/pom.xml b/bigtable-hbase-2.x-parent/bigtable-hbase-2.x-shaded/pom.xml
index f11b1edaf..a1baf9813 100644
--- a/bigtable-hbase-2.x-parent/bigtable-hbase-2.x-shaded/pom.xml
+++ b/bigtable-hbase-2.x-parent/bigtable-hbase-2.x-shaded/pom.xml
@@ -63,6 +63,11 @@ limitations under the License.
     <artifactId>metrics-core</artifactId>
     <version>${dropwizard.metrics.version}</version>
   </dependency>
+    <dependency>
+        <groupId>org.checkerframework</groupId>
+        <artifactId>checker-compat-qual</artifactId>
+        <version>${checkerframework.version}</version>
+    </dependency>
 
   <dependency>
     <groupId>io.opencensus</groupId>
@@ -183,6 +188,8 @@ limitations under the License.
                   <exclude>javax.inject:javax.inject</exclude>
                   <exclude>org.apache.htrace:htrace-core4</exclude>
                   <exclude>org.apache.yetus:audience-annotations</exclude>
+                  <exclude>javax.annotation:javax.annotation-api</exclude>
+                  <exclude>org.checkerframework:checker-compat-qual</exclude>
                 </excludes>
               </artifactSet>
               <relocations>
@@ -264,6 +271,14 @@ limitations under the License.
                   <pattern>org.apache.commons.codec</pattern>
                   <shadedPattern>com.google.bigtable.repackaged.org.apache.commons.codec</shadedPattern>
                 </relocation>
+                <relocation>
+                  <pattern>org.threeten</pattern>
+                  <shadedPattern>com.google.bigtable.repackaged.org.threeten</shadedPattern>
+                </relocation>
+                <relocation>
+                  <pattern>org.codehaus</pattern>
+                   <shadedPattern>com.google.bigtable.repackaged.org.codehaus</shadedPattern>
+                </relocation>
               </relocations>
             </configuration>
           </execution>
diff --git a/pom.xml b/pom.xml
index df44c7e50..e25ec75cb 100644
--- a/pom.xml
+++ b/pom.xml
@@ -62,6 +62,7 @@ limitations under the License.
         <google-cloud.version>0.78.0-alpha</google-cloud.version>
         <opencensus.version>0.17.0</opencensus.version>
         <guava.version>26.0-android</guava.version>
+        <checkerframework.version>2.5.2</checkerframework.version>
 
         <!-- hbase dependency versions -->
         <hbase.version.1>1.4.5</hbase.version.1>


Test patch:
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/TestBigtableDataClientWrapper.java b/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/TestBigtableDataClientWrapper.java
index 9e5aea397..6c3aad5ea 100644
--- a/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/TestBigtableDataClientWrapper.java
+++ b/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/TestBigtableDataClientWrapper.java
@@ -31,18 +31,26 @@ import com.google.bigtable.v2.MutateRowResponse;
 import com.google.bigtable.v2.ReadModifyWriteRowRequest;
 import com.google.bigtable.v2.ReadModifyWriteRowResponse;
 import com.google.bigtable.v2.Row;
+import com.google.bigtable.v2.SampleRowKeysRequest;
+import com.google.bigtable.v2.SampleRowKeysResponse;
 import com.google.cloud.bigtable.config.BigtableOptions;
+import com.google.cloud.bigtable.data.v2.internal.NameUtil;
 import com.google.cloud.bigtable.data.v2.internal.RequestContext;
 import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation;
-import com.google.cloud.bigtable.data.v2.models.InstanceName;
+import com.google.cloud.bigtable.data.v2.models.KeyOffset;
 import com.google.cloud.bigtable.data.v2.models.Mutation;
+import com.google.cloud.bigtable.data.v2.models.Query;
 import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow;
 import com.google.cloud.bigtable.data.v2.models.RowCell;
 import com.google.cloud.bigtable.data.v2.models.RowMutation;
+import com.google.cloud.bigtable.grpc.scanner.FlatRow;
+import com.google.cloud.bigtable.grpc.scanner.FlatRowConverter;
+import com.google.cloud.bigtable.grpc.scanner.ResultScanner;
 import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.protobuf.ByteString;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import org.junit.Before;
@@ -51,6 +59,9 @@ import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.runners.MockitoJUnitRunner;
 
+/**
+ * Unit tests for the {@link BigtableDataClientWrapper}.
+ */
 @RunWith(MockitoJUnitRunner.class)
 public class TestBigtableDataClientWrapper {
 
@@ -67,7 +78,7 @@ public class TestBigtableDataClientWrapper {
   private static final ByteString ROW_KEY = ByteString.copyFromUtf8("test-key");
 
   private static final RequestContext REQUEST_CONTEXT =
-      RequestContext.create(InstanceName.of(PROJECT_ID, INSTANCE_ID), APP_PROFILE_ID);
+      RequestContext.create(PROJECT_ID, INSTANCE_ID, APP_PROFILE_ID);
 
   private BigtableOptions options =
       BigtableOptions.builder().setProjectId(PROJECT_ID).setInstanceId(INSTANCE_ID)
@@ -77,6 +88,9 @@ public class TestBigtableDataClientWrapper {
 
   private BigtableDataClientWrapper clientWrapper;
 
+  @Mock
+  private ResultScanner<FlatRow> mockFlatRow;
+
   @Before
   public void setUp() {
     clientWrapper = new BigtableDataClientWrapper(client, options);
@@ -279,4 +293,112 @@ public class TestBigtableDataClientWrapper {
     assertEquals(expectedModelRow, actualRow);
     verify(client).readModifyWriteRow(request);
   }
+
+  @Test
+  public void testSampleRowKeys(){
+    final ByteString ROW_KEY_1 = ByteString.copyFromUtf8("row-key-1");
+    final ByteString ROW_KEY_2 = ByteString.copyFromUtf8("row-key-2");
+    final ByteString ROW_KEY_3 = ByteString.copyFromUtf8("row-key-3");
+
+    String tableName = NameUtil.formatTableName(PROJECT_ID, INSTANCE_ID, TABLE_ID);
+    SampleRowKeysRequest requestProto =
+        SampleRowKeysRequest.newBuilder().setTableName(tableName).build();
+    ImmutableList.Builder sampleKeys = ImmutableList.builder();
+    sampleKeys.add(SampleRowKeysResponse.newBuilder()
+        .setRowKey(ROW_KEY_1).setOffsetBytes(11).build());
+    sampleKeys.add(SampleRowKeysResponse.newBuilder()
+        .setRowKey(ROW_KEY_2).setOffsetBytes(12).build());
+    sampleKeys.add(SampleRowKeysResponse.newBuilder()
+        .setRowKey(ROW_KEY_3).setOffsetBytes(13).build());
+    when(client.sampleRowKeys(requestProto)).thenReturn(sampleKeys.build());
+
+    List<KeyOffset> keyOffsetList = clientWrapper.sampleRowKeys(TABLE_ID);
+    assertEquals(keyOffsetList.get(0).getKey(), ROW_KEY_1);
+    assertEquals(keyOffsetList.get(1).getKey(), ROW_KEY_2);
+    assertEquals(keyOffsetList.get(2).getKey(), ROW_KEY_3);
+    verify(client).sampleRowKeys(requestProto);
+  }
+
+  @Test
+  public void testSampleRowKeysAsync() throws Exception{
+    final ByteString ROW_KEY_1 = ByteString.copyFromUtf8("row-key-1");
+    final ByteString ROW_KEY_2 = ByteString.copyFromUtf8("row-key-2");
+    final ByteString ROW_KEY_3 = ByteString.copyFromUtf8("row-key-3");
+
+    String tableName = NameUtil.formatTableName(PROJECT_ID, INSTANCE_ID, TABLE_ID);
+    SampleRowKeysRequest requestProto =
+        SampleRowKeysRequest.newBuilder().setTableName(tableName).build();
+    ImmutableList.Builder sampleKeys = ImmutableList.builder();
+    sampleKeys.add(SampleRowKeysResponse.newBuilder()
+        .setRowKey(ROW_KEY_1).setOffsetBytes(11).build());
+    sampleKeys.add(SampleRowKeysResponse.newBuilder()
+        .setRowKey(ROW_KEY_2).setOffsetBytes(12).build());
+    sampleKeys.add(SampleRowKeysResponse.newBuilder()
+        .setRowKey(ROW_KEY_3).setOffsetBytes(13).build());
+    List<SampleRowKeysResponse> responseProtos = sampleKeys.build();
+    when(client.sampleRowKeysAsync(requestProto)).thenReturn(Futures.immediateFuture(responseProtos));
+
+    List<KeyOffset> keyOffsetList = clientWrapper.sampleRowKeysAsync(TABLE_ID).get();
+    assertEquals(keyOffsetList.get(0).getKey(), ROW_KEY_1);
+    assertEquals(keyOffsetList.get(1).getKey(), ROW_KEY_2);
+    assertEquals(keyOffsetList.get(2).getKey(), ROW_KEY_3);
+    verify(client).sampleRowKeysAsync(requestProto);
+  }
+
+  @Test
+  public void testReadRows() throws Exception{
+    Query query = Query.create(TABLE_ID);
+    FlatRow row = FlatRow.newBuilder().withRowKey(ByteString.copyFromUtf8("key")).build();
+    when(client.readFlatRows(query.toProto(REQUEST_CONTEXT))).thenReturn(mockFlatRow);
+    when(mockFlatRow.next()).thenReturn(row);
+    ResultScanner<com.google.cloud.bigtable.data.v2.models.Row> actualResult =
+        clientWrapper.readRows(query);
+
+    assertEquals(FlatRowConverter.convertToModelRow(row), actualResult.next());
+    verify(mockFlatRow).next();
+    verify(client).readFlatRows(query.toProto(REQUEST_CONTEXT));
+  }
+
+  @Test
+  public void testReadRowsAsync() throws Exception{
+    Query query = Query.create(TABLE_ID);
+    FlatRow flatRow = FlatRow.newBuilder().withRowKey(ByteString.copyFromUtf8("key")).build();
+    List<FlatRow> listFlatRows = Arrays.asList(flatRow);
+    when(client.readFlatRowsAsync(query.toProto(REQUEST_CONTEXT)))
+        .thenReturn(Futures.immediateFuture(listFlatRows));
+    List<com.google.cloud.bigtable.data.v2.models.Row> actualResult =
+        clientWrapper.readRowsAsync(query).get();
+
+    assertEquals(Arrays.asList(FlatRowConverter.convertToModelRow(flatRow)), actualResult);
+    verify(client).readFlatRowsAsync(query.toProto(REQUEST_CONTEXT));
+  }
+
+  @Test
+  public void testReadFlatRowsList(){
+    Query query = Query.create(TABLE_ID);
+    FlatRow flatRow = FlatRow.newBuilder().withRowKey(ByteString.copyFromUtf8("key")).build();
+    List<FlatRow> listFlatRows = Arrays.asList(flatRow);
+    when(client.readFlatRowsList(query.toProto(REQUEST_CONTEXT))).thenReturn(listFlatRows);
+    clientWrapper.readFlatRowsList(query);
+    verify(client).readFlatRowsList(query.toProto(REQUEST_CONTEXT));
+  }
+
+  @Test
+  public void testReadFlatRows(){
+    Query query = Query.create(TABLE_ID).range("start", "end");
+    when(client.readFlatRows(query.toProto(REQUEST_CONTEXT))).thenReturn(mockFlatRow);
+    clientWrapper.readFlatRows(query);
+    verify(client).readFlatRows(query.toProto(REQUEST_CONTEXT));
+  }
+
+  @Test
+  public void testReadFlatRowsAsync(){
+    Query query = Query.create(TABLE_ID);
+    FlatRow flatRow = FlatRow.newBuilder().withRowKey(ByteString.copyFromUtf8("key")).build();
+    List<FlatRow> listFlatRows = Arrays.asList(flatRow);
+    when(client.readFlatRowsAsync(query.toProto(REQUEST_CONTEXT)))
+        .thenReturn(Futures.immediateFuture(listFlatRows));
+    clientWrapper.readFlatRowsAsync(query);
+    verify(client).readFlatRowsAsync(query.toProto(REQUEST_CONTEXT));
+  }
 }
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/scanner/TestFlatRowConverter.java b/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/scanner/TestFlatRowConverter.java
index 1db07a4e3..7c76fe982 100644
--- a/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/scanner/TestFlatRowConverter.java
+++ b/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/scanner/TestFlatRowConverter.java
@@ -15,6 +15,10 @@
  */
 package com.google.cloud.bigtable.grpc.scanner;
 
+import com.google.cloud.bigtable.data.v2.models.RowCell;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
 import org.junit.Assert;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -95,6 +99,40 @@ public class TestFlatRowConverter {
     testBothWays(simpleRow, expectedRow);
   }
 
+  @Test
+  public void testModelRowWithOneCell(){
+    FlatRow.Cell cell =
+        new FlatRow.Cell("family", toByteString("column"), 500, toByteString("value"), null);
+    FlatRow flatRow = FlatRow.newBuilder().withRowKey(toByteString("key")).addCell(cell).build();
+
+    com.google.cloud.bigtable.data.v2.models.Row expectedRow =
+        com.google.cloud.bigtable.data.v2.models.Row.create(flatRow.getRowKey(), Arrays.asList(
+            RowCell.create(cell.getFamily(), cell.getQualifier(), cell.getTimestamp(),
+                cell.getLabels(), cell.getValue())));
+
+    Assert.assertEquals(expectedRow, FlatRowConverter.convertToModelRow(flatRow));
+  }
+
+  @Test
+  public void testModelRowWithMultipleCell(){
+    FlatRow simpleRow = FlatRow.newBuilder()
+        .withRowKey(toByteString("key"))
+        .addCell("family1", toByteString("column"), 500, toByteString("value"), null)
+        .addCell("family1", toByteString("column2"), 500, toByteString("value"), null)
+        .addCell("family1", toByteString("column2"), 400, toByteString("value"), null)
+        .addCell("family2", toByteString("column"), 500, toByteString("value"), null)
+        .build();
+    List<RowCell> rowCells = new ArrayList<>();
+    for(FlatRow.Cell cell : simpleRow.getCells()){
+      rowCells.add(RowCell.create(cell.getFamily(), cell.getQualifier(), cell.getTimestamp(),
+          cell.getLabels(), cell.getValue()));
+    }
+    com.google.cloud.bigtable.data.v2.models.Row row =
+        com.google.cloud.bigtable.data.v2.models.Row.create(ByteString.copyFromUtf8("key"),
+            rowCells);
+    Assert.assertEquals(row, FlatRowConverter.convertToModelRow(simpleRow));
+  }
+
   private void testBothWays(FlatRow simpleRow, Row row) {
     Assert.assertEquals(row, FlatRowConverter.convert(simpleRow));
     Assert.assertEquals(simpleRow, FlatRowConverter.convert(row));