Repository: GoogleCloudPlatform/cloud-bigtable-client
Pull request #2064 (2019-02-06 01:57:10)

BulkRead uses the new client wrapper


Patch:
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/BigtableSession.java b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/BigtableSession.java
index a0471711d..226418c23 100644
--- a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/BigtableSession.java
+++ b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/BigtableSession.java
@@ -361,7 +361,7 @@ public class BigtableSession implements Closeable {
    * @return a {@link com.google.cloud.bigtable.grpc.async.BulkRead} object.
    */
   public BulkRead createBulkRead(BigtableTableName tableName) {
-    return new BulkRead(dataClient, tableName, options.getBulkOptions().getBulkMaxRowKeyCount(),
+    return new BulkRead(clientWrapper, tableName, options.getBulkOptions().getBulkMaxRowKeyCount(),
         BigtableSessionSharedThreadPools.getInstance().getBatchThreadPool()
     );
   }
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/async/BulkRead.java b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/async/BulkRead.java
index 0dc6d1c2a..def25d851 100644
--- a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/async/BulkRead.java
+++ b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/async/BulkRead.java
@@ -24,12 +24,14 @@ import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
 
+import com.google.cloud.bigtable.data.v2.models.Filters;
+import com.google.cloud.bigtable.data.v2.models.Query;
+import com.google.cloud.bigtable.grpc.BigtableDataClientWrapper;
 import com.google.common.base.Preconditions;
 import com.google.bigtable.v2.ReadRowsRequest;
 import com.google.bigtable.v2.RowFilter;
 import com.google.bigtable.v2.RowSet;
 import com.google.cloud.bigtable.config.Logger;
-import com.google.cloud.bigtable.grpc.BigtableDataClient;
 import com.google.cloud.bigtable.grpc.BigtableTableName;
 import com.google.cloud.bigtable.grpc.scanner.FlatRow;
 import com.google.cloud.bigtable.grpc.scanner.ResultScanner;
@@ -65,7 +67,7 @@ public class BulkRead {
         }
       };
 
-  private final BigtableDataClient client;
+  private final BigtableDataClientWrapper client;
   private final int batchSizes;
   private final ExecutorService threadPool;
   private final String tableName;
@@ -74,12 +76,12 @@ public class BulkRead {
 
   /**
    * Constructor for BulkRead.
-   * @param client a {@link BigtableDataClient} object.
+   * @param client a {@link BigtableDataClientWrapper} object.
    * @param tableName a {@link BigtableTableName} object.
    * @param batchSizes The number of keys to lookup per RPC.
    * @param threadPool the {@link ExecutorService} to execute the batched reads on
    */
-  public BulkRead(BigtableDataClient client, BigtableTableName tableName, int batchSizes,
+  public BulkRead(BigtableDataClientWrapper client, BigtableTableName tableName, int batchSizes,
       ExecutorService threadPool) {
     this.client = client;
     this.tableName = tableName.toString();
@@ -176,12 +178,15 @@ public class BulkRead {
     @Override
     public void run() {
       try {
-        ResultScanner<FlatRow> scanner = client.readFlatRows(ReadRowsRequest.newBuilder()
+        Query query = Query.fromProto(ReadRowsRequest.newBuilder()
             .setTableName(tableName)
             .setFilter(filter)
+
+            // TODO: Query.addRowKeys(List<ByteString>) would be super helpful, and allow
+            // this class to construct a Query without passing in a proto
             .setRows(RowSet.newBuilder().addAllRowKeys(futures.keys()).build())
-            .build()
-        );
+            .build());
+        ResultScanner<FlatRow> scanner = client.readFlatRows(query);
         while (true) {
           FlatRow row = scanner.next();
           if (row == null) {
diff --git a/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/adapters/DeleteAdapter.java b/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/adapters/DeleteAdapter.java
index bd33f1e11..0f1559675 100644
--- a/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/adapters/DeleteAdapter.java
+++ b/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/adapters/DeleteAdapter.java
@@ -122,7 +122,7 @@ public class DeleteAdapter extends MutationAdapter<Delete> {
 
   @Override
   /** {@inheritDoc} */
-  protected void adaptMutations(Delete operation, com.google.cloud.bigtable.data.v2.models.MutationApi<?> mutation) {
+  public void adapt(Delete operation, com.google.cloud.bigtable.data.v2.models.MutationApi<?> mutation) {
     if (operation.getFamilyCellMap().isEmpty()) {
       throwIfUnsupportedDeleteRow(operation);
 
diff --git a/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/adapters/HBaseMutationAdapter.java b/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/adapters/HBaseMutationAdapter.java
index d1245bb62..6b8325c25 100644
--- a/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/adapters/HBaseMutationAdapter.java
+++ b/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/adapters/HBaseMutationAdapter.java
@@ -75,13 +75,13 @@ public class HBaseMutationAdapter extends MutationAdapter<Mutation> {
 
   @Override
   /** {@inheritDoc} */
-  protected void adaptMutations(Mutation mutation, com.google.cloud.bigtable.data.v2.models.MutationApi<?> mutationModel) {
+  public void adapt(Mutation mutation, com.google.cloud.bigtable.data.v2.models.MutationApi<?> mutationModel) {
     MutationAdapter<Mutation> adapter = adapterMap.get(mutation.getClass());
     if (adapter == null) {
       throw new UnsupportedOperationException(
           String.format(
               "Cannot adapt mutation of type %s.", mutation.getClass().getCanonicalName()));
     }
-    adapter.adaptMutations(mutation, mutationModel);
+    adapter.adapt(mutation, mutationModel);
   }
 }
diff --git a/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/adapters/MutationAdapter.java b/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/adapters/MutationAdapter.java
index 0665d3f5d..eb78f7d07 100644
--- a/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/adapters/MutationAdapter.java
+++ b/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/adapters/MutationAdapter.java
@@ -38,16 +38,6 @@ public abstract class MutationAdapter<T extends Row>
     return ByteStringer.extract(bs);
   }
 
-  /** {@inheritDoc} */
-  @Override
-  public final void adapt(T operation, com.google.cloud.bigtable.data.v2.models.MutationApi<?> mutation) {
-    adaptMutations(operation, mutation);
-  }
-
-  public final void toEntry(T operation, com.google.cloud.bigtable.data.v2.models.MutationApi<?> mutation) {
-    adaptMutations(operation, mutation);
-  }
-
   /**
    * Converts an HBase {@link Mutation} which represents a set of changes to a single row from an
    * HBase perspective to a Google Cloud Java {@link com.google.cloud.bigtable.data.v2.models.Mutation}
@@ -59,7 +49,8 @@ public abstract class MutationAdapter<T extends Row>
    * @param operation The HBase {@link Mutation} to convert
    * @param mutation The model {@link com.google.cloud.bigtable.data.v2.models.MutationApi}
    */
-  protected abstract void adaptMutations(
+  @Override
+  public abstract void adapt(
       T operation,
       com.google.cloud.bigtable.data.v2.models.MutationApi<?> mutation
   );
diff --git a/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/adapters/PutAdapter.java b/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/adapters/PutAdapter.java
index 7ad346d1d..6a29c7dee 100644
--- a/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/adapters/PutAdapter.java
+++ b/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/adapters/PutAdapter.java
@@ -74,7 +74,8 @@ public class PutAdapter extends MutationAdapter<Put> {
   }
 
   @Override
-  protected void adaptMutations(Put operation, com.google.cloud.bigtable.data.v2.models.MutationApi<?> mutation) {
+  /** {@inheritDoc} */
+  public void adapt(Put operation, com.google.cloud.bigtable.data.v2.models.MutationApi<?> mutation) {
     if (operation.isEmpty()) {
       throw new IllegalArgumentException("No columns to insert");
     }
diff --git a/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/adapters/RowMutationsAdapter.java b/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/adapters/RowMutationsAdapter.java
index 9bcb37d60..ae6d39654 100644
--- a/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/adapters/RowMutationsAdapter.java
+++ b/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/adapters/RowMutationsAdapter.java
@@ -39,12 +39,13 @@ public class RowMutationsAdapter extends MutationAdapter<RowMutations>{
   }
 
   @Override
-  protected void adaptMutations(
+  /** {@inheritDoc} */
+  public void adapt(
       RowMutations operation,
       com.google.cloud.bigtable.data.v2.models.MutationApi<?> mutationApi
   ) {
     for (Mutation mutation : operation.getMutations()) {
-      mutationAdapter.adaptMutations(mutation, mutationApi);
+      mutationAdapter.adapt(mutation, mutationApi);
     }
   }
 }
diff --git a/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/adapters/UnsupportedMutationAdapter.java b/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/adapters/UnsupportedMutationAdapter.java
index c8e90e318..ee6dff19c 100644
--- a/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/adapters/UnsupportedMutationAdapter.java
+++ b/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/adapters/UnsupportedMutationAdapter.java
@@ -43,7 +43,7 @@ public class UnsupportedMutationAdapter<T extends Mutation> extends MutationAdap
    * Adapt a single HBase Operation to a single Bigtable generated message.
    */
   @Override
-  protected void adaptMutations(T operation, com.google.cloud.bigtable.data.v2.models.MutationApi<?> mutation) {
+  public void adapt(T operation, com.google.cloud.bigtable.data.v2.models.MutationApi<?> mutation) {
     throw new UnsupportedOperationException(
       String.format("The %s operation is unsupported.", operationDescription));
 }


Test patch:
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/async/TestBulkRead.java b/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/async/TestBulkRead.java
index 070ac13b5..3167a756f 100644
--- a/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/async/TestBulkRead.java
+++ b/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/async/TestBulkRead.java
@@ -30,6 +30,9 @@ import java.util.Set;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.TimeUnit;
 
+import com.google.cloud.bigtable.data.v2.internal.RequestContext;
+import com.google.cloud.bigtable.data.v2.models.Query;
+import com.google.cloud.bigtable.grpc.BigtableDataClientWrapper;
 import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
@@ -44,7 +47,6 @@ import org.mockito.stubbing.Answer;
 import com.google.bigtable.v2.MutateRowsRequest;
 import com.google.bigtable.v2.ReadRowsRequest;
 import com.google.bigtable.v2.RowSet;
-import com.google.cloud.bigtable.grpc.BigtableDataClient;
 import com.google.cloud.bigtable.grpc.BigtableTableName;
 import com.google.cloud.bigtable.grpc.scanner.FlatRow;
 import com.google.cloud.bigtable.grpc.scanner.ResultScanner;
@@ -59,9 +61,10 @@ public class TestBulkRead {
 
   private static final BigtableTableName TABLE_NAME =
       new BigtableTableName("projects/SomeProject/instances/SomeInstance/tables/SomeTable");
+  private static final RequestContext CONTEXT = RequestContext.create("SomeProject", "SomeInstance", "");
 
   @Mock
-  BigtableDataClient mockClient;
+  BigtableDataClientWrapper mockClient;
 
   @Mock
   ExecutorService mockThreadPool;
@@ -93,12 +96,12 @@ public class TestBulkRead {
    */
   @Test
   public void testBatchOfOne() throws Exception {
-    when(mockClient.readFlatRows(any(ReadRowsRequest.class))).thenReturn(mockScanner);
+    when(mockClient.readFlatRows(any(Query.class))).thenReturn(mockScanner);
     FlatRow row = createRow(ByteString.copyFromUtf8("Key"));
     when(mockScanner.next()).thenReturn(row).thenReturn(null);
     ListenableFuture<FlatRow> future = underTest.add(createRequest(row.getRowKey()));
     underTest.flush();
-    verify(mockClient, times(1)).readFlatRows(any(ReadRowsRequest.class));
+    verify(mockClient, times(1)).readFlatRows(any(Query.class));
     Assert.assertEquals(row, future.get(10, TimeUnit.MILLISECONDS));
   }
 
@@ -107,14 +110,14 @@ public class TestBulkRead {
    */
   @Test
   public void testDuplicateKey() throws Exception {
-    when(mockClient.readFlatRows(any(ReadRowsRequest.class))).thenReturn(mockScanner);
+    when(mockClient.readFlatRows(any(Query.class))).thenReturn(mockScanner);
     FlatRow row = createRow(ByteString.copyFromUtf8("Key"));
     when(mockScanner.next()).thenReturn(row).thenReturn(null);
     ReadRowsRequest request = createRequest(row.getRowKey());
     ListenableFuture<FlatRow> future1 = underTest.add(request);
     ListenableFuture<FlatRow> future2 = underTest.add(request);
     underTest.flush();
-    verify(mockClient, times(1)).readFlatRows(any(ReadRowsRequest.class));
+    verify(mockClient, times(1)).readFlatRows(any(Query.class));
     Assert.assertEquals(row, future1.get(10, TimeUnit.MILLISECONDS));
     Assert.assertEquals(row, future2.get(10, TimeUnit.MILLISECONDS));
   }
@@ -129,8 +132,8 @@ public class TestBulkRead {
         addRows(rowKeys, new Answer<ResultScanner<FlatRow>>() {
           @Override
           public ResultScanner<FlatRow> answer(InvocationOnMock invocation) throws Throwable {
-            ReadRowsRequest request = invocation.getArgumentAt(0, ReadRowsRequest.class);
-            List<ByteString> list = new ArrayList<>(request.getRows().getRowKeysList());
+            Query request = invocation.getArgumentAt(0, Query.class);
+            List<ByteString> list = new ArrayList<>(request.toProto(CONTEXT).getRows().getRowKeysList());
             Collections.shuffle(list);
             return createMockScanner(list.iterator());
           }
@@ -153,8 +156,8 @@ public class TestBulkRead {
     List<ListenableFuture<FlatRow>> futures = addRows(rowKeys, new Answer<ResultScanner<FlatRow>>() {
       @Override
       public ResultScanner<FlatRow> answer(InvocationOnMock invocation) throws Throwable {
-        ReadRowsRequest request = invocation.getArgumentAt(0, ReadRowsRequest.class);
-        ArrayList<ByteString> rowKeysList = new ArrayList<>(request.getRows().getRowKeysList());
+        Query request = invocation.getArgumentAt(0, Query.class);
+        List<ByteString> rowKeysList = new ArrayList<>(request.toProto(CONTEXT).getRows().getRowKeysList());
         missing.add(rowKeysList.remove((int) Math.random() * rowKeysList.size()));
         return createMockScanner(rowKeysList.iterator());
       }
@@ -183,7 +186,7 @@ public class TestBulkRead {
    */
   private List<ListenableFuture<FlatRow>> addRows(
       List<ByteString> rowKeys, Answer<ResultScanner<FlatRow>> scannerGenerator) {
-    when(mockClient.readFlatRows(any(ReadRowsRequest.class)))
+    when(mockClient.readFlatRows(any(Query.class)))
         .thenAnswer(scannerGenerator);
 
     List<ListenableFuture<FlatRow>> futures = new ArrayList<>();
@@ -192,7 +195,7 @@ public class TestBulkRead {
     }
     underTest.flush();
     verify(mockClient, times(rowKeys.size() / underTest.getBatchSizes()))
-        .readFlatRows(any(ReadRowsRequest.class));
+        .readFlatRows(any(Query.class));
     return futures;
   }
 
diff --git a/bigtable-client-core-parent/bigtable-hbase/src/test/java/com/google/cloud/bigtable/hbase/adapters/TestHBaseMutationAdapter.java b/bigtable-client-core-parent/bigtable-hbase/src/test/java/com/google/cloud/bigtable/hbase/adapters/TestHBaseMutationAdapter.java
index 40a42a146..57db4d8e7 100644
--- a/bigtable-client-core-parent/bigtable-hbase/src/test/java/com/google/cloud/bigtable/hbase/adapters/TestHBaseMutationAdapter.java
+++ b/bigtable-client-core-parent/bigtable-hbase/src/test/java/com/google/cloud/bigtable/hbase/adapters/TestHBaseMutationAdapter.java
@@ -81,37 +81,37 @@ public class TestHBaseMutationAdapter {
   @Test
   public void testPutIsAdapted() {
     Put put = new Put(dataHelper.randomData("rk1"));
-    adapter.adaptMutations(put, mutation);
+    adapter.adapt(put, mutation);
 
     Mockito.verify(putAdapter, Mockito.times(1))
-        .adaptMutations(Mockito.any(Put.class), Mockito.eq(mutation));
+        .adapt(Mockito.any(Put.class), Mockito.eq(mutation));
   }
 
   @Test
   public void testDeleteIsAdapted() {
     Delete delete = new Delete(dataHelper.randomData("rk1"));
-    adapter.adaptMutations(delete, mutation);
+    adapter.adapt(delete, mutation);
 
     Mockito.verify(deleteAdapter, Mockito.times(1))
-        .adaptMutations(Mockito.any(Delete.class), Mockito.eq(mutation));
+        .adapt(Mockito.any(Delete.class), Mockito.eq(mutation));
   }
 
   @Test
   public void testAppendIsAdapted() {
     Append append = new Append(dataHelper.randomData("rk1"));
-    adapter.adaptMutations(append, mutation);
+    adapter.adapt(append, mutation);
 
     Mockito.verify(appendAdapter, Mockito.times(1))
-        .adaptMutations(Mockito.any(Append.class), Mockito.eq(mutation));
+        .adapt(Mockito.any(Append.class), Mockito.eq(mutation));
   }
 
   @Test
   public void testIncrementIsAdapted() {
     Increment increment = new Increment(dataHelper.randomData("rk1"));
-    adapter.adaptMutations(increment, mutation);
+    adapter.adapt(increment, mutation);
 
     Mockito.verify(incrementAdapter, Mockito.times(1))
-        .adaptMutations(Mockito.any(Increment.class), Mockito.eq(mutation));
+        .adapt(Mockito.any(Increment.class), Mockito.eq(mutation));
   }
 
   @Test
@@ -122,6 +122,6 @@ public class TestHBaseMutationAdapter {
     expectedException.expectMessage("Cannot adapt mutation of type");
     expectedException.expectMessage("UnknownMutation");
 
-    adapter.adaptMutations(unknownMutation, mutation);
+    adapter.adapt(unknownMutation, mutation);
   }
 }
diff --git a/bigtable-client-core-parent/bigtable-hbase/src/test/java/com/google/cloud/bigtable/hbase/adapters/TestRowMutationsAdapter.java b/bigtable-client-core-parent/bigtable-hbase/src/test/java/com/google/cloud/bigtable/hbase/adapters/TestRowMutationsAdapter.java
index 7ed3ce296..57a1b9cc0 100644
--- a/bigtable-client-core-parent/bigtable-hbase/src/test/java/com/google/cloud/bigtable/hbase/adapters/TestRowMutationsAdapter.java
+++ b/bigtable-client-core-parent/bigtable-hbase/src/test/java/com/google/cloud/bigtable/hbase/adapters/TestRowMutationsAdapter.java
@@ -104,7 +104,7 @@ public class TestRowMutationsAdapter {
     MutateRowRequest result = toMutateRowRequest(rowKey, mutation);
     Assert.assertArrayEquals(rowKey, result.getRowKey().toByteArray());
 
-    Mockito.verify(mutationAdapter, times(2)).adaptMutations(any(
+    Mockito.verify(mutationAdapter, times(2)).adapt(any(
         org.apache.hadoop.hbase.client.Mutation.class),Matchers.eq(mutation));
   }