Bigtable 12. Implement SampleRowKeys

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

@@ -18,10 +18,13 @@
 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.Callables;
 import com.google.api.gax.rpc.ClientContext;
 import com.google.api.gax.rpc.ResponseObserver;
 import com.google.api.gax.rpc.ServerStreamingCallable;
 import com.google.api.gax.rpc.UnaryCallable;
+import com.google.bigtable.v2.SampleRowKeysRequest;
+import com.google.bigtable.v2.SampleRowKeysResponse;
 import com.google.cloud.bigtable.data.v2.internal.RequestContext;
 import com.google.cloud.bigtable.data.v2.models.DefaultRowAdapter;
 import com.google.cloud.bigtable.data.v2.models.KeyOffset;
@@ -64,6 +67,14 @@ public static EnhancedBigtableStub create(EnhancedBigtableStubSettings settings)
             .setEndpoint(settings.getEndpoint())
             .setCredentialsProvider(settings.getCredentialsProvider());
 
+    // SampleRowKeys retries are handled in the overlay: disable retries in the base layer (but make
+    // sure to preserve the exception callable settings.
+    baseSettingsBuilder
+        .sampleRowKeysSettings()
+        .setSimpleTimeoutNoRetries(
+            settings.sampleRowKeysSettings().getRetrySettings().getTotalTimeout())
+        .setRetryableCodes(settings.sampleRowKeysSettings().getRetryableCodes());
+
     BigtableStubSettings baseSettings = baseSettingsBuilder.build();
     ClientContext clientContext = ClientContext.create(baseSettings);
     GrpcBigtableStub stub = new GrpcBigtableStub(baseSettings, clientContext);
@@ -109,12 +120,16 @@ public void call(
    * </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");
-      }
-    };
+    UnaryCallable<SampleRowKeysRequest, List<SampleRowKeysResponse>> spoolable =
+        stub.sampleRowKeysCallable().all();
+
+    UnaryCallable<SampleRowKeysRequest, List<SampleRowKeysResponse>> retryable =
+        Callables.retrying(spoolable, settings.sampleRowKeysSettings(), clientContext);
+
+    UnaryCallable<SampleRowKeysRequest, List<SampleRowKeysResponse>> withContext =
+        retryable.withDefaultCallContext(clientContext.getDefaultCallContext());
+
+    return new SampleRowKeysCallable(withContext, requestContext);
   }
 
   private UnaryCallable<RowMutation, Void> createMutateRowCallable() {

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

@@ -0,0 +1,79 @@
+/*
+ * 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.stub;
+
+import com.google.api.core.ApiFunction;
+import com.google.api.core.ApiFuture;
+import com.google.api.core.ApiFutures;
+import com.google.api.gax.rpc.ApiCallContext;
+import com.google.api.gax.rpc.UnaryCallable;
+import com.google.bigtable.v2.SampleRowKeysRequest;
+import com.google.bigtable.v2.SampleRowKeysResponse;
+import com.google.bigtable.v2.TableName;
+import com.google.cloud.bigtable.data.v2.internal.RequestContext;
+import com.google.cloud.bigtable.data.v2.models.KeyOffset;
+import com.google.common.collect.ImmutableList;
+import java.util.List;
+
+/** Simple wrapper for SampleRowKeys to wrap the request and response protobufs. */
+class SampleRowKeysCallable extends UnaryCallable<String, List<KeyOffset>> {
+  private final RequestContext requestContext;
+  private final UnaryCallable<SampleRowKeysRequest, List<SampleRowKeysResponse>> inner;
+
+  SampleRowKeysCallable(
+      UnaryCallable<SampleRowKeysRequest, List<SampleRowKeysResponse>> inner,
+      RequestContext requestContext) {
+
+    this.requestContext = requestContext;
+    this.inner = inner;
+  }
+
+  @Override
+  public ApiFuture<List<KeyOffset>> futureCall(String tableId, ApiCallContext context) {
+    TableName tableName =
+        TableName.of(
+            requestContext.getInstanceName().getProject(),
+            requestContext.getInstanceName().getInstance(),
+            tableId);
+
+    SampleRowKeysRequest request =
+        SampleRowKeysRequest.newBuilder()
+            .setTableName(tableName.toString())
+            .setAppProfileId(requestContext.getAppProfileId())
+            .build();
+
+    ApiFuture<List<SampleRowKeysResponse>> rawResponse = inner.futureCall(request, context);
+
+    return ApiFutures.transform(
+        rawResponse,
+        new ApiFunction<List<SampleRowKeysResponse>, List<KeyOffset>>() {
+          @Override
+          public List<KeyOffset> apply(List<SampleRowKeysResponse> rawResponse) {
+            return convert(rawResponse);
+          }
+        });
+  }
+
+  private static List<KeyOffset> convert(List<SampleRowKeysResponse> rawResponse) {
+    ImmutableList.Builder<KeyOffset> results = ImmutableList.builder();
+
+    for (SampleRowKeysResponse element : rawResponse) {
+      results.add(KeyOffset.create(element.getRowKey(), element.getOffsetBytes()));
+    }
+
+    return results.build();
+  }
+}

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

@@ -0,0 +1,122 @@
+/*
+ * 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.stub;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.api.core.ApiFuture;
+import com.google.api.core.SettableApiFuture;
+import com.google.api.gax.grpc.GrpcStatusCode;
+import com.google.api.gax.rpc.ApiCallContext;
+import com.google.api.gax.rpc.NotFoundException;
+import com.google.api.gax.rpc.UnaryCallable;
+import com.google.bigtable.admin.v2.InstanceName;
+import com.google.bigtable.v2.SampleRowKeysRequest;
+import com.google.bigtable.v2.SampleRowKeysResponse;
+import com.google.cloud.bigtable.data.v2.internal.RequestContext;
+import com.google.cloud.bigtable.data.v2.models.KeyOffset;
+import com.google.common.collect.ImmutableList;
+import com.google.protobuf.ByteString;
+import io.grpc.Status.Code;
+import java.util.List;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class SampleRowKeysCallableTest {
+  private final RequestContext requestContext =
+      RequestContext.create(InstanceName.of("my-project", "my-instance"), "my-app-profile");
+  private FakeCallable inner;
+  private SampleRowKeysCallable callable;
+
+  @Before
+  public void setUp() {
+    inner = new FakeCallable();
+    callable = new SampleRowKeysCallable(inner, requestContext);
+  }
+
+  @Test
+  public void requestIsCorrect() {
+    callable.futureCall("my-table");
+
+    assertThat(inner.request)
+        .isEqualTo(
+            SampleRowKeysRequest.newBuilder()
+                .setTableName(requestContext.getInstanceName() + "/tables/my-table")
+                .setAppProfileId(requestContext.getAppProfileId())
+                .build());
+  }
+
+  @Test
+  public void responseCorrectlyTransformed() throws Exception {
+    ApiFuture<List<KeyOffset>> result = callable.futureCall("my-table");
+
+    inner.response.set(
+        ImmutableList.of(
+            SampleRowKeysResponse.newBuilder()
+                .setRowKey(ByteString.copyFromUtf8("key1"))
+                .setOffsetBytes(100)
+                .build(),
+            SampleRowKeysResponse.newBuilder()
+                .setRowKey(ByteString.copyFromUtf8(""))
+                .setOffsetBytes(1000)
+                .build()));
+
+    assertThat(result.get(1, TimeUnit.SECONDS))
+        .isEqualTo(
+            ImmutableList.of(
+                KeyOffset.create(ByteString.copyFromUtf8("key1"), 100),
+                KeyOffset.create(ByteString.EMPTY, 1000)));
+  }
+
+  @Test
+  public void errorIsPropagated() throws Exception {
+    ApiFuture<List<KeyOffset>> result = callable.futureCall("my-table");
+
+    Throwable expectedError =
+        new NotFoundException("fake error", null, GrpcStatusCode.of(Code.NOT_FOUND), false);
+    inner.response.setException(expectedError);
+
+    Throwable actualError = null;
+    try {
+      result.get(1, TimeUnit.SECONDS);
+    } catch (ExecutionException e) {
+      actualError = e.getCause();
+    }
+
+    assertThat(actualError).isEqualTo(expectedError);
+  }
+
+  static class FakeCallable
+      extends UnaryCallable<SampleRowKeysRequest, List<SampleRowKeysResponse>> {
+    SampleRowKeysRequest request;
+    ApiCallContext callContext;
+    SettableApiFuture<List<SampleRowKeysResponse>> response = SettableApiFuture.create();
+
+    @Override
+    public ApiFuture<List<SampleRowKeysResponse>> futureCall(
+        SampleRowKeysRequest request, ApiCallContext context) {
+      this.request = request;
+      this.callContext = context;
+
+      return response;
+    }
+  }
+}