Bigtable: add a helper method to easily connect to an emulator (#4506)

igorbernstein2 · web-flow · commit e6cf8d32987d · 2019-02-20T10:25:23.000-05:00
* Bigtable: add a helper method to easily connect to an emulator

* fix bad merge

* format

* use emulator helper in tests
diff --git a/google-cloud-clients/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/BigtableTableAdminSettings.java b/google-cloud-clients/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/BigtableTableAdminSettings.java
@@ -15,9 +15,13 @@
  */
 package com.google.cloud.bigtable.admin.v2;
 
+import com.google.api.core.ApiFunction;
+import com.google.api.gax.core.NoCredentialsProvider;
+import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
 import com.google.cloud.bigtable.admin.v2.stub.BigtableTableAdminStubSettings;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Verify;
+import io.grpc.ManagedChannelBuilder;
 import java.io.IOException;
 import javax.annotation.Nullable;
 
@@ -79,6 +83,29 @@ public static Builder newBuilder() {
     return new Builder();
   }
 
+  /** Create a new builder preconfigured to connect to the Bigtable emulator. */
+  public static Builder newBuilderForEmulator(int port) {
+    Builder builder = newBuilder().setProjectId("fake-project").setInstanceId("fake-instance");
+
+    builder
+        .stubSettings()
+        .setCredentialsProvider(NoCredentialsProvider.create())
+        .setEndpoint("localhost:" + port)
+        .setTransportChannelProvider(
+            InstantiatingGrpcChannelProvider.newBuilder()
+                .setPoolSize(1)
+                .setChannelConfigurator(
+                    new ApiFunction<ManagedChannelBuilder, ManagedChannelBuilder>() {
+                      @Override
+                      public ManagedChannelBuilder apply(ManagedChannelBuilder input) {
+                        return input.usePlaintext();
+                      }
+                    })
+                .build());
+
+    return builder;
+  }
+
   /** Builder for BigtableTableAdminSettings. */
   public static final class Builder {
     @Nullable private String projectId;
diff --git a/google-cloud-clients/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataSettings.java b/google-cloud-clients/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataSettings.java
@@ -15,6 +15,9 @@
  */
 package com.google.cloud.bigtable.data.v2;
 
+import com.google.api.core.ApiFunction;
+import com.google.api.gax.core.NoCredentialsProvider;
+import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
 import com.google.api.gax.rpc.BatchingCallSettings;
 import com.google.api.gax.rpc.ClientSettings;
 import com.google.api.gax.rpc.ServerStreamingCallSettings;
@@ -26,6 +29,7 @@
 import com.google.cloud.bigtable.data.v2.models.Row;
 import com.google.cloud.bigtable.data.v2.models.RowMutation;
 import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStubSettings;
+import io.grpc.ManagedChannelBuilder;
 import java.io.IOException;
 import java.util.List;
 import javax.annotation.Nonnull;
@@ -67,6 +71,32 @@ public static Builder newBuilder() {
     return new Builder();
   }
 
+  /** Create a new builder preconfigured to connect to the Bigtable emulator. */
+  public static Builder newBuilderForEmulator(int port) {
+    Builder builder = newBuilder();
+
+    builder
+        .getTypedStubSettings()
+        .setProjectId("fake-project")
+        .setInstanceId("fake-instance")
+        .setCredentialsProvider(NoCredentialsProvider.create())
+        .setEndpoint("localhost:" + port)
+        .setTransportChannelProvider(
+            InstantiatingGrpcChannelProvider.newBuilder()
+                .setMaxInboundMessageSize(256 * 1024 * 1024)
+                .setPoolSize(1)
+                .setChannelConfigurator(
+                    new ApiFunction<ManagedChannelBuilder, ManagedChannelBuilder>() {
+                      @Override
+                      public ManagedChannelBuilder apply(ManagedChannelBuilder input) {
+                        return input.usePlaintext();
+                      }
+                    })
+                .build());
+
+    return builder;
+  }
+
   /** Returns the target project id. */
   public String getProjectId() {
     return getTypedStubSettings().getProjectId();
diff --git a/google-cloud-clients/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/it/env/Emulator.java b/google-cloud-clients/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/it/env/Emulator.java
@@ -15,17 +15,11 @@
  */
 package com.google.cloud.bigtable.data.v2.it.env;
 
-import com.google.api.gax.core.NoCredentialsProvider;
-import com.google.api.gax.grpc.GrpcTransportChannel;
-import com.google.api.gax.rpc.ClientSettings;
-import com.google.api.gax.rpc.FixedTransportChannelProvider;
-import com.google.bigtable.admin.v2.BigtableTableAdminGrpc;
-import com.google.bigtable.admin.v2.BigtableTableAdminGrpc.BigtableTableAdminBlockingStub;
+import com.google.cloud.bigtable.admin.v2.BigtableTableAdminClient;
+import com.google.cloud.bigtable.admin.v2.BigtableTableAdminSettings;
 import com.google.cloud.bigtable.data.v2.BigtableDataClient;
 import com.google.cloud.bigtable.data.v2.BigtableDataSettings;
 import com.google.common.io.CharStreams;
-import io.grpc.ManagedChannel;
-import io.grpc.ManagedChannelBuilder;
 import java.io.BufferedReader;
 import java.io.IOException;
 import java.io.InputStream;
@@ -52,8 +46,7 @@ class Emulator {
   private final Path executable;
   private Process process;
   private boolean isStopped = true;
-  private ManagedChannel channel;
-  private BigtableTableAdminBlockingStub tableAdminClient;
+  private BigtableTableAdminClient tableAdminClient;
   private BigtableDataClient dataClient;
 
   private static final String PROJECT_ID = "fake-project";
@@ -94,16 +87,18 @@ void start() throws Exception {
 
     waitForPort(availablePort);
 
-    channel = createChannel(availablePort);
-
-    tableAdminClient = BigtableTableAdminGrpc.newBlockingStub(channel);
+    tableAdminClient =
+        BigtableTableAdminClient.create(
+            BigtableTableAdminSettings.newBuilderForEmulator(availablePort)
+                .setProjectId(PROJECT_ID)
+                .setInstanceId(INSTANCE_ID)
+                .build());
 
     dataClient =
         BigtableDataClient.create(
-            configureClient(
-                    BigtableDataSettings.newBuilder()
-                        .setProjectId(PROJECT_ID)
-                        .setInstanceId(INSTANCE_ID))
+            BigtableDataSettings.newBuilderForEmulator(availablePort)
+                .setProjectId(PROJECT_ID)
+                .setInstanceId(INSTANCE_ID)
                 .build());
 
     Runtime.getRuntime()
@@ -122,8 +117,7 @@ public void run() {
   void stop() throws Exception {
     try {
       dataClient.close();
-      channel.shutdownNow();
-      channel.awaitTermination(1, TimeUnit.MINUTES);
+      tableAdminClient.close();
     } finally {
       isStopped = true;
       process.destroy();
@@ -134,7 +128,7 @@ BigtableDataClient getDataClient() {
     return dataClient;
   }
 
-  BigtableTableAdminBlockingStub getTableAdminClient() {
+  BigtableTableAdminClient getTableAdminClient() {
     return tableAdminClient;
   }
 
@@ -171,22 +165,6 @@ private void waitForPort(int port) throws InterruptedException, TimeoutException
     throw new TimeoutException("Timed out waiting for server to start");
   }
 
-  private ManagedChannel createChannel(int port) {
-    return ManagedChannelBuilder.forAddress("localhost", port)
-        .usePlaintext()
-        .maxInboundMessageSize(256 * 1024 * 1024)
-        .build();
-  }
-
-  private <T extends ClientSettings.Builder<?, ?>> T configureClient(T settings) {
-    settings
-        .setCredentialsProvider(new NoCredentialsProvider())
-        .setTransportChannelProvider(
-            FixedTransportChannelProvider.create(GrpcTransportChannel.create(channel)));
-
-    return settings;
-  }
-
   private static void pipeStreamToLog(final InputStream stream, final Level level) {
     final BufferedReader reader = new BufferedReader(new InputStreamReader(stream));
 
diff --git a/google-cloud-clients/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/it/env/EmulatorEnv.java b/google-cloud-clients/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/it/env/EmulatorEnv.java
@@ -15,11 +15,8 @@
  */
 package com.google.cloud.bigtable.data.v2.it.env;
 
-import com.google.bigtable.admin.v2.ColumnFamily;
-import com.google.bigtable.admin.v2.CreateTableRequest;
-import com.google.bigtable.admin.v2.Table;
+import com.google.cloud.bigtable.admin.v2.models.CreateTableRequest;
 import com.google.cloud.bigtable.data.v2.BigtableDataClient;
-import com.google.cloud.bigtable.data.v2.internal.NameUtil;
 
 public class EmulatorEnv implements TestEnv {
   private static final String PROJECT_ID = "fake-project";
@@ -36,14 +33,7 @@ public void start() throws Exception {
 
     emulator
         .getTableAdminClient()
-        .createTable(
-            CreateTableRequest.newBuilder()
-                .setParent(NameUtil.formatInstanceName(PROJECT_ID, INSTANCE_ID))
-                .setTableId(TABLE_ID)
-                .setTable(
-                    Table.newBuilder()
-                        .putColumnFamilies(FAMILY_ID, ColumnFamily.getDefaultInstance()))
-                .build());
+        .createTable(CreateTableRequest.of(TABLE_ID).addFamily(FAMILY_ID));
   }
 
   @Override
diff --git a/google-cloud-testing/google-cloud-bigtable-emulator/README.md b/google-cloud-testing/google-cloud-bigtable-emulator/README.md
@@ -110,27 +110,10 @@ public class ExampleTest {
   @Before
   public void setUp() throws IOException {
     // Initialize the clients to connect to the emulator
-    BigtableTableAdminSettings.Builder tableAdminSettings = BigtableTableAdminSettings.newBuilder()
-        .setProjectId("fake-project")
-        .setInstanceId("fake-instance");
-    tableAdminSettings.stubSettings()
-        .setCredentialsProvider(NoCredentialsProvider.create())
-        .setTransportChannelProvider(
-            FixedTransportChannelProvider.create(
-                GrpcTransportChannel.create(bigtableEmulator.getAdminChannel())
-            )
-        );
+    BigtableTableAdminSettings.Builder tableAdminSettings = BigtableTableAdminSettings.newBuilderForEmulator(bigtableEmulator.getPort());
     tableAdminClient = BigtableTableAdminClient.create(tableAdminSettings.build());
 
-    BigtableDataSettings.Builder dataSettings = BigtableDataSettings.newBuilder()
-        .setProjectId("fake-project")
-        .setInstanceId("fake-instance")
-        .setCredentialsProvider(NoCredentialsProvider.create())
-        .setTransportChannelProvider(
-            FixedTransportChannelProvider.create(
-                GrpcTransportChannel.create(bigtableEmulator.getDataChannel())
-            )
-        );
+    BigtableDataSettings.Builder dataSettings = BigtableDataSettings.newBuilderForEmulator(bigtableEmulator.getPort());
     dataClient = BigtableDataClient.create(dataSettings.build());
 
     // Create a test table that can be used in tests
