03. Surface: SampleRowKeys

Adds the surface to wrap SampleRowKeys RPC. The surface:
- hides the fact that the RPC is streaming
- wraps the request as a simple string table id
- wraps the response in KeyOffset

Author: igorbernstein2

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java

@@ -15,10 +15,14 @@
  */
 package com.google.cloud.bigtable.data.v2;
 
+import com.google.api.core.ApiFuture;
 import com.google.api.core.InternalApi;
+import com.google.api.gax.rpc.UnaryCallable;
 import com.google.bigtable.admin.v2.InstanceName;
 import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStub;
+import com.google.cloud.bigtable.data.v2.wrappers.KeyOffset;
 import java.io.IOException;
+import java.util.List;
 
 /**
  * Client for reading from and writing to existing Bigtable tables.
@@ -112,4 +116,46 @@ public static BigtableDataClient create(BigtableDataSettings settings) throws IO
   public void close() throws Exception {
     stub.close();
   }
+
+  /**
+   * Convenience method to asynchronously return a sample of row keys in the table. The returned row
+   * keys will delimit contiguous sections of the table of approximately equal size, which can be
+   * used to break up the data for distributed tasks like mapreduces. The returned callable object
+   * allows for customization of api invocation.
+   *
+   * <p>Sample code:
+   *
+   * <pre>{@code
+   * InstanceName instanceName = InstanceName.of("[PROJECT]", "[INSTANCE]");
+   * try (BigtableClient bigtableClient = BigtableClient.create(instanceName)) {
+   *   ApiFuture<List<KeyOffset>> keyOffsets = bigtableClient.sampleRowKeys("[TABLE]");
+   * }
+   * }</pre>
+   */
+  public ApiFuture<List<KeyOffset>> sampleRowKeys(String tableId) {
+    return sampleRowKeysCallable().futureCall(tableId);
+  }
+
+  /**
+   * Returns a sample of row keys in the table. The returned row keys will delimit contiguous
+   * sections of the table of approximately equal size, which can be used to break up the data for
+   * distributed tasks like mapreduces. The returned callable object allows for customization of api
+   * invocation.
+   *
+   * <p>Sample code:
+   *
+   * <pre>{@code
+   * InstanceName instanceName = InstanceName.of("[PROJECT]", "[INSTANCE]");
+   * try (BigtableClient bigtableClient = BigtableClient.create(instanceName)) {
+   *   // Synchronous invocation
+   *   List<KeyOffset> keyOffsets = bigtableClient.sampleRowKeysCallable().call("[TABLE]");
+   *
+   *   // Asynchronous invocation
+   *   ApiFuture<List<KeyOffset>> keyOffsets = bigtableClient.sampleRowKeysCallable().futureCall("[TABLE]");
+   * }
+   * }</pre>
+   */
+  public UnaryCallable<String, List<KeyOffset>> sampleRowKeysCallable() {
+    return stub.sampleRowKeysCallable();
+  }
 }

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataSettings.java

@@ -16,9 +16,12 @@
 package com.google.cloud.bigtable.data.v2;
 
 import com.google.api.gax.rpc.ClientSettings;
+import com.google.api.gax.rpc.UnaryCallSettings;
 import com.google.bigtable.admin.v2.InstanceName;
 import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStubSettings;
+import com.google.cloud.bigtable.data.v2.wrappers.KeyOffset;
 import java.io.IOException;
+import java.util.List;
 import javax.annotation.Nonnull;
 
 /**
@@ -64,6 +67,11 @@ public String getAppProfileId() {
     return getTypedStubSettings().getAppProfileId();
   }
 
+  /** Returns the object with the settings used for calls to sampleRowKeys. */
+  public UnaryCallSettings<String, List<KeyOffset>> sampleRowKeysSettings() {
+    return getTypedStubSettings().sampleRowKeysSettings();
+  }
+
   @SuppressWarnings("unchecked")
   EnhancedBigtableStubSettings getTypedStubSettings() {
     return (EnhancedBigtableStubSettings) getStubSettings();
@@ -124,6 +132,11 @@ public String getAppProfileId() {
       return getTypedStubSettings().getAppProfileId();
     }
 
+    /** Returns the builder for the settings used for calls to SampleRowKeysSettings. */
+    public UnaryCallSettings.Builder<String, List<KeyOffset>> sampleRowKeysSettings() {
+      return getTypedStubSettings().sampleRowKeysSettings();
+    }
+
     @SuppressWarnings("unchecked")
     private EnhancedBigtableStubSettings.Builder getTypedStubSettings() {
       return (EnhancedBigtableStubSettings.Builder) getStubSettings();

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java

@@ -15,10 +15,15 @@
  */
 package com.google.cloud.bigtable.data.v2.stub;
 
+import com.google.api.core.ApiFuture;
 import com.google.api.core.InternalApi;
+import com.google.api.gax.rpc.ApiCallContext;
 import com.google.api.gax.rpc.ClientContext;
+import com.google.api.gax.rpc.UnaryCallable;
 import com.google.cloud.bigtable.data.v2.internal.RequestContext;
+import com.google.cloud.bigtable.data.v2.wrappers.KeyOffset;
 import java.io.IOException;
+import java.util.List;
 
 /**
  * The core client that converts method calls to RPCs.
@@ -39,6 +44,8 @@ public class EnhancedBigtableStub implements AutoCloseable {
   private final ClientContext clientContext;
   private final RequestContext requestContext;
 
+  private final UnaryCallable<String, List<KeyOffset>> sampleRowKeysCallable;
+
   public static EnhancedBigtableStub create(EnhancedBigtableStubSettings settings)
       throws IOException {
     // Configure the base settings
@@ -63,6 +70,8 @@ public static EnhancedBigtableStub create(EnhancedBigtableStubSettings settings)
     this.stub = stub;
     this.requestContext =
         RequestContext.create(settings.getInstanceName(), settings.getAppProfileId());
+
+    sampleRowKeysCallable = createSampleRowKeysCallable();
   }
 
   @Override
@@ -71,8 +80,31 @@ public void close() throws Exception {
   }
 
   // <editor-fold desc="Callable creators">
+
+  /**
+   * Creates a callable chain to handle SampleRowKeys RPcs. The chain will:
+   *
+   * <ul>
+   *   <li>Convert a table id to a {@link com.google.bigtable.v2.SampleRowKeysRequest}.
+   *   <li>Dispatch the request to the GAPIC's {@link BigtableStub#sampleRowKeysCallable()}.
+   *   <li>Spool responses into a list.
+   *   <li>Retry on failure.
+   *   <li>Convert the responses into {@link KeyOffset}s.
+   * </ul>
+   */
+  private UnaryCallable<String, List<KeyOffset>> createSampleRowKeysCallable() {
+    return new UnaryCallable<String, List<KeyOffset>>() {
+      @Override
+      public ApiFuture<List<KeyOffset>> futureCall(String request, ApiCallContext context) {
+        throw new UnsupportedOperationException("todo");
+      }
+    };
+  }
   // </editor-fold>
 
   // <editor-fold desc="Callable accessors">
+  public UnaryCallable<String, List<KeyOffset>> sampleRowKeysCallable() {
+    return sampleRowKeysCallable;
+  }
   // </editor-fold>
 }

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java

@@ -17,11 +17,16 @@
 
 import com.google.api.core.InternalApi;
 import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
+import com.google.api.gax.retrying.RetrySettings;
+import com.google.api.gax.rpc.StatusCode.Code;
 import com.google.api.gax.rpc.StubSettings;
 import com.google.api.gax.rpc.UnaryCallSettings;
 import com.google.bigtable.admin.v2.InstanceName;
+import com.google.cloud.bigtable.data.v2.wrappers.KeyOffset;
 import com.google.common.base.Preconditions;
+import java.util.List;
 import javax.annotation.Nonnull;
+import org.threeten.bp.Duration;
 
 /**
  * Settings class to configure an instance of {@link EnhancedBigtableStub}.
@@ -58,12 +63,15 @@ public class EnhancedBigtableStubSettings extends StubSettings<EnhancedBigtableS
   private final InstanceName instanceName;
   private final String appProfileId;
 
+  private final UnaryCallSettings<String, List<KeyOffset>> sampleRowKeysSettings;
+
   private EnhancedBigtableStubSettings(Builder builder) {
     super(builder);
     instanceName = builder.instanceName;
     appProfileId = builder.appProfileId;
 
     // Per method settings.
+    sampleRowKeysSettings = builder.sampleRowKeysSettings.build();
   }
 
   public static Builder newBuilder() {
@@ -80,6 +88,11 @@ public String getAppProfileId() {
     return appProfileId;
   }
 
+  /** Returns the object with the settings used for calls to SampleRowKeys. */
+  public UnaryCallSettings<String, List<KeyOffset>> sampleRowKeysSettings() {
+    return sampleRowKeysSettings;
+  }
+
   /** Returns a builder containing all the values of this settings class. */
   public Builder toBuilder() {
     return new Builder(this);
@@ -90,6 +103,8 @@ public static class Builder extends StubSettings.Builder<EnhancedBigtableStubSet
     private InstanceName instanceName;
     private String appProfileId;
 
+    private final UnaryCallSettings.Builder<String, List<KeyOffset>> sampleRowKeysSettings;
+
     /**
      * Initializes a new Builder with sane defaults for all settings.
      *
@@ -114,6 +129,22 @@ private Builder() {
               .build());
 
       // Per-method settings using baseSettings for defaults.
+
+      sampleRowKeysSettings = UnaryCallSettings.newUnaryCallSettingsBuilder();
+      /* TODO: copy retryCodes & retrySettings from baseSettings.sampleRowKeysSettings once it exists in GAPIC */
+      sampleRowKeysSettings
+          .setRetryableCodes(Code.DEADLINE_EXCEEDED, Code.UNAVAILABLE, Code.ABORTED)
+          .setRetrySettings(
+              RetrySettings.newBuilder()
+                  .setMaxAttempts(10)
+                  .setTotalTimeout(Duration.ofMinutes(1))
+                  .setInitialRetryDelay(Duration.ofMillis(100))
+                  .setRetryDelayMultiplier(1.3)
+                  .setMaxRetryDelay(Duration.ofMinutes(1))
+                  .setInitialRpcTimeout(Duration.ofSeconds(20))
+                  .setRpcTimeoutMultiplier(1)
+                  .setMaxRpcTimeout(Duration.ofSeconds(20))
+                  .build());
     }
 
     private Builder(EnhancedBigtableStubSettings settings) {
@@ -122,6 +153,7 @@ private Builder(EnhancedBigtableStubSettings settings) {
       appProfileId = settings.appProfileId;
 
       // Per method settings.
+      sampleRowKeysSettings = settings.sampleRowKeysSettings.toBuilder();
     }
 
     // <editor-fold desc="Private Helpers">
@@ -169,6 +201,11 @@ public String getAppProfileId() {
       return appProfileId;
     }
 
+    /** Returns the builder for the settings used for calls to SampleRowKeysSettings. */
+    public UnaryCallSettings.Builder<String, List<KeyOffset>> sampleRowKeysSettings() {
+      return sampleRowKeysSettings;
+    }
+
     @SuppressWarnings("unchecked")
     public EnhancedBigtableStubSettings build() {
       Preconditions.checkState(instanceName != null, "InstanceName must be set");

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/wrappers/KeyOffset.java

@@ -0,0 +1,47 @@
+/*
+ * Copyright 2018 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.cloud.bigtable.data.v2.wrappers;
+
+import com.google.api.core.InternalApi;
+import com.google.auto.value.AutoValue;
+import com.google.protobuf.ByteString;
+
+/** Represents the offset of a row key in a table. */
+@AutoValue
+public abstract class KeyOffset {
+  @InternalApi
+  public static KeyOffset create(ByteString key, long offsetBytes) {
+    return new AutoValue_KeyOffset(key, offsetBytes);
+  }
+
+  /**
+   * Sorted streamed sequence of sample row keys in the table. The table might have contents before
+   * the first row key in the list and after the last one, but a key containing the empty string
+   * indicates "end of table" and will be the last response given, if present.
+   *
+   * <p>Note that row keys in this list may not have ever been written to or read from, and users
+   * should therefore not make any assumptions about the row key structure that are specific to
+   * their use case.
+   */
+  public abstract ByteString key();
+
+  /**
+   * Approximate total storage space used by all rows in the table which precede {@link #key()}.
+   * Buffering the contents of all rows between two subsequent samples would require space roughly
+   * equal to the difference in their {@link #offsetBytes()} fields.
+   */
+  public abstract long offsetBytes();
+}

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientTest.java

@@ -15,7 +15,12 @@
  */
 package com.google.cloud.bigtable.data.v2;
 
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.api.gax.rpc.UnaryCallable;
 import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStub;
+import com.google.cloud.bigtable.data.v2.wrappers.KeyOffset;
+import java.util.List;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -26,11 +31,13 @@
 @RunWith(MockitoJUnitRunner.class)
 public class BigtableDataClientTest {
   @Mock private EnhancedBigtableStub mockStub;
+  @Mock private UnaryCallable<String, List<KeyOffset>> mockSampleRowKeysCallable;
 
   private BigtableDataClient bigtableDataClient;
 
   @Before
   public void setUp() {
+    Mockito.when(mockStub.sampleRowKeysCallable()).thenReturn(mockSampleRowKeysCallable);
     bigtableDataClient = new BigtableDataClient(mockStub);
   }
 
@@ -39,4 +46,15 @@ public void proxyCloseTest() throws Exception {
     bigtableDataClient.close();
     Mockito.verify(mockStub).close();
   }
+
+  @Test
+  public void proxySampleRowKeysCallableTest() {
+    assertThat(bigtableDataClient.sampleRowKeysCallable()).isSameAs(mockSampleRowKeysCallable);
+  }
+
+  @Test
+  public void proxySampleRowKeysTest() {
+    bigtableDataClient.sampleRowKeys("fake-table");
+    Mockito.verify(mockSampleRowKeysCallable).futureCall("fake-table");
+  }
 }