feat: Support building a BigQueryWriteClient within the StreamWriterV2 (#876)

yayi-google · web-flow · commit 237c82711649 · 2021-02-23T21:04:12.000-05:00
* Support building a BigQueryWriteClient within the StreamWriterV2.

* feat: Support building a BigQueryWriteClient within the StreamWriterV2
diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta2/StreamWriterV2.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta2/StreamWriterV2.java
@@ -17,13 +17,16 @@
 
 import com.google.api.core.ApiFuture;
 import com.google.api.core.SettableApiFuture;
+import com.google.api.gax.core.CredentialsProvider;
+import com.google.api.gax.rpc.TransportChannelProvider;
 import com.google.cloud.bigquery.storage.v1beta2.StreamConnection.DoneCallback;
 import com.google.cloud.bigquery.storage.v1beta2.StreamConnection.RequestCallback;
 import com.google.common.base.Preconditions;
 import com.google.common.util.concurrent.Uninterruptibles;
 import io.grpc.Status;
 import io.grpc.Status.Code;
 import io.grpc.StatusRuntimeException;
+import java.io.IOException;
 import java.util.Deque;
 import java.util.LinkedList;
 import java.util.concurrent.TimeUnit;
@@ -36,8 +39,6 @@
 /**
  * A BigQuery Stream Writer that can be used to write data into BigQuery Table.
  *
- * <p>TODO: Add credential support.
- *
  * <p>TODO: Attach schema.
  *
  * <p>TODO: Attach traceId.
@@ -104,6 +105,16 @@ public class StreamWriterV2 implements AutoCloseable {
   @GuardedBy("lock")
   private final Deque<AppendRequestAndResponse> inflightRequestQueue;
 
+  /*
+   * A client used to interact with BigQuery.
+   */
+  private BigQueryWriteClient client;
+
+  /*
+   * If true, the client above is created by this writer and should be closed.
+   */
+  private boolean ownsBigQueryWriteClient = false;
+
   /*
    * Wraps the underlying bi-directional stream connection with server.
    */
@@ -119,7 +130,7 @@ public static long getApiMaxRequestBytes() {
     return 8L * 1000L * 1000L; // 8 megabytes (https://en.wikipedia.org/wiki/Megabyte)
   }
 
-  private StreamWriterV2(Builder builder) {
+  private StreamWriterV2(Builder builder) throws IOException {
     this.lock = new ReentrantLock();
     this.hasMessageInWaitingQueue = lock.newCondition();
     this.inflightReduced = lock.newCondition();
@@ -128,9 +139,22 @@ private StreamWriterV2(Builder builder) {
     this.maxInflightBytes = builder.maxInflightBytes;
     this.waitingRequestQueue = new LinkedList<AppendRequestAndResponse>();
     this.inflightRequestQueue = new LinkedList<AppendRequestAndResponse>();
+    if (builder.client == null) {
+      BigQueryWriteSettings stubSettings =
+          BigQueryWriteSettings.newBuilder()
+              .setCredentialsProvider(builder.credentialsProvider)
+              .setTransportChannelProvider(builder.channelProvider)
+              .setEndpoint(builder.endpoint)
+              .build();
+      this.client = BigQueryWriteClient.create(stubSettings);
+      this.ownsBigQueryWriteClient = true;
+    } else {
+      this.client = builder.client;
+      this.ownsBigQueryWriteClient = false;
+    }
     this.streamConnection =
         new StreamConnection(
-            builder.client,
+            this.client,
             new RequestCallback() {
               @Override
               public void run(AppendRowsResponse response) {
@@ -261,6 +285,9 @@ public void close() {
       log.warning(
           "Append handler join is interrupted. Stream: " + streamName + " Error: " + e.toString());
     }
+    if (this.ownsBigQueryWriteClient) {
+      this.client.close();
+    }
   }
 
   /*
@@ -405,6 +432,11 @@ public static StreamWriterV2.Builder newBuilder(String streamName, BigQueryWrite
     return new StreamWriterV2.Builder(streamName, client);
   }
 
+  /** Constructs a new {@link StreamWriterV2.Builder} using the given stream. */
+  public static StreamWriterV2.Builder newBuilder(String streamName) {
+    return new StreamWriterV2.Builder(streamName);
+  }
+
   /** A builder of {@link StreamWriterV2}s. */
   public static final class Builder {
 
@@ -420,6 +452,19 @@ public static final class Builder {
 
     private long maxInflightBytes = DEFAULT_MAX_INFLIGHT_BYTES;
 
+    private String endpoint = BigQueryWriteSettings.getDefaultEndpoint();
+
+    private TransportChannelProvider channelProvider =
+        BigQueryWriteSettings.defaultGrpcTransportProviderBuilder().setChannelsPerCpu(1).build();
+
+    private CredentialsProvider credentialsProvider =
+        BigQueryWriteSettings.defaultCredentialsProviderBuilder().build();
+
+    private Builder(String streamName) {
+      this.streamName = Preconditions.checkNotNull(streamName);
+      this.client = null;
+    }
+
     private Builder(String streamName, BigQueryWriteClient client) {
       this.streamName = Preconditions.checkNotNull(streamName);
       this.client = Preconditions.checkNotNull(client);
@@ -435,8 +480,34 @@ public Builder setMaxInflightBytes(long value) {
       return this;
     }
 
+    /** Gives the ability to override the gRPC endpoint. */
+    public Builder setEndpoint(String endpoint) {
+      this.endpoint = Preconditions.checkNotNull(endpoint, "Endpoint is null.");
+      return this;
+    }
+
+    /**
+     * {@code ChannelProvider} to use to create Channels, which must point at Cloud BigQuery Storage
+     * API endpoint.
+     *
+     * <p>For performance, this client benefits from having multiple underlying connections. See
+     * {@link com.google.api.gax.grpc.InstantiatingGrpcChannelProvider.Builder#setPoolSize(int)}.
+     */
+    public Builder setChannelProvider(TransportChannelProvider channelProvider) {
+      this.channelProvider =
+          Preconditions.checkNotNull(channelProvider, "ChannelProvider is null.");
+      return this;
+    }
+
+    /** {@code CredentialsProvider} to use to create Credentials to authenticate calls. */
+    public Builder setCredentialsProvider(CredentialsProvider credentialsProvider) {
+      this.credentialsProvider =
+          Preconditions.checkNotNull(credentialsProvider, "CredentialsProvider is null.");
+      return this;
+    }
+
     /** Builds the {@code StreamWriterV2}. */
-    public StreamWriterV2 build() {
+    public StreamWriterV2 build() throws IOException {
       return new StreamWriterV2(this);
     }
   }
diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/StreamWriterV2Test.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/StreamWriterV2Test.java
@@ -31,6 +31,7 @@
 import com.google.protobuf.Int64Value;
 import io.grpc.Status;
 import io.grpc.StatusRuntimeException;
+import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
@@ -81,7 +82,7 @@ public void tearDown() throws Exception {
     serviceHelper.stop();
   }
 
-  private StreamWriterV2 getTestStreamWriterV2() {
+  private StreamWriterV2 getTestStreamWriterV2() throws IOException {
     return StreamWriterV2.newBuilder(TEST_STREAM, client).build();
   }
 
@@ -158,6 +159,20 @@ public void run() {
     appendThread.interrupt();
   }
 
+  @Test
+  public void testBuildBigQueryWriteClientInWriter() throws Exception {
+    StreamWriterV2 writer =
+        StreamWriterV2.newBuilder(TEST_STREAM)
+            .setCredentialsProvider(NoCredentialsProvider.create())
+            .setChannelProvider(serviceHelper.createChannelProvider())
+            .build();
+
+    testBigQueryWrite.addResponse(createAppendResponse(0));
+    ApiFuture<AppendRowsResponse> appendFuture1 = sendTestMessage(writer, new String[] {"A"});
+    assertEquals(0, appendFuture1.get().getAppendResult().getOffset().getValue());
+    writer.close();
+  }
+
   @Test
   public void testAppendSuccess() throws Exception {
     StreamWriterV2 writer = getTestStreamWriterV2();
@@ -371,7 +386,7 @@ public void testAppendsWithTinyMaxInflightBytes() throws Exception {
   }
 
   @Test
-  public void testMessageTooLarge() {
+  public void testMessageTooLarge() throws Exception {
     StreamWriterV2 writer = getTestStreamWriterV2();
 
     String oversized = Strings.repeat("a", (int) (StreamWriterV2.getApiMaxRequestBytes() + 1));
