feat: Add header back to the client (#2016)

* feat: add back header setting * . * . * . * . * . * . * . * . * . * . * .

Author: yirutang

google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/ConnectionWorker.java

@@ -18,6 +18,7 @@
 import com.google.api.core.ApiFuture;
 import com.google.api.core.SettableApiFuture;
 import com.google.api.gax.batching.FlowController;
+import com.google.api.gax.rpc.FixedHeaderProvider;
 import com.google.auto.value.AutoValue;
 import com.google.cloud.bigquery.storage.v1.AppendRowsRequest.ProtoData;
 import com.google.cloud.bigquery.storage.v1.Exceptions.AppendSerializtionError;
@@ -77,6 +78,11 @@ class ConnectionWorker implements AutoCloseable {
  */
  private String streamName;
 
+ /*
+ * The location of this connection.
+ */
+ private String location = null;
+
  /*
  * The proto schema of rows to write. This schema can change during multiplexing.
  */
@@ -211,6 +217,7 @@ public static long getApiMaxRequestBytes() {
 
  public ConnectionWorker(
  String streamName,
+ String location,
  ProtoSchema writerSchema,
  long maxInflightRequests,
  long maxInflightBytes,
@@ -223,6 +230,9 @@ public ConnectionWorker(
  this.hasMessageInWaitingQueue = lock.newCondition();
  this.inflightReduced = lock.newCondition();
  this.streamName = streamName;
+ if (location != null && !location.isEmpty()) {
+ this.location = location;
+ }
  this.maxRetryDuration = maxRetryDuration;
  if (writerSchema == null) {
  throw new StatusRuntimeException(
@@ -236,6 +246,18 @@ public ConnectionWorker(
  this.waitingRequestQueue = new LinkedList<AppendRequestAndResponse>();
  this.inflightRequestQueue = new LinkedList<AppendRequestAndResponse>();
  // Always recreate a client for connection worker.
+ HashMap<String, String> newHeaders = new HashMap<>();
+ newHeaders.putAll(clientSettings.toBuilder().getHeaderProvider().getHeaders());
+ if (this.location == null) {
+ newHeaders.put("x-goog-request-params", "write_stream=" + this.streamName);
+ } else {
+ newHeaders.put("x-goog-request-params", "write_location=" + this.location);
+ }
+ BigQueryWriteSettings stubSettings =
+ clientSettings
+ .toBuilder()
+ .setHeaderProvider(FixedHeaderProvider.create(newHeaders))
+ .build();
  this.client = BigQueryWriteClient.create(clientSettings);
 
  this.appendThread =
@@ -297,6 +319,24 @@ public void run(Throwable finalStatus) {
 
  /** Schedules the writing of rows at given offset. */
  ApiFuture<AppendRowsResponse> append(StreamWriter streamWriter, ProtoRows rows, long offset) {
+ if (this.location != null && this.location != streamWriter.getLocation()) {
+ throw new StatusRuntimeException(
+ Status.fromCode(Code.INVALID_ARGUMENT)
+ .withDescription(
+ "StreamWriter with location "
+ + streamWriter.getLocation()
+ + " is scheduled to use a connection with location "
+ + this.location));
+ } else if (this.location == null && streamWriter.getStreamName() != this.streamName) {
+ // Location is null implies this is non-multiplexed connection.
+ throw new StatusRuntimeException(
+ Status.fromCode(Code.INVALID_ARGUMENT)
+ .withDescription(
+ "StreamWriter with stream name "
+ + streamWriter.getStreamName()
+ + " is scheduled to use a connection with stream name "
+ + this.streamName));
+ }
  Preconditions.checkNotNull(streamWriter);
  AppendRowsRequest.Builder requestBuilder = AppendRowsRequest.newBuilder();
  requestBuilder.setProtoRows(
@@ -322,6 +362,10 @@ Boolean isUserClosed() {
  }
  }
 
+ String getWriteLocation() {
+ return this.location;
+ }
+
  private ApiFuture<AppendRowsResponse> appendInternal(
  StreamWriter streamWriter, AppendRowsRequest message) {
  AppendRequestAndResponse requestWrapper = new AppendRequestAndResponse(message, streamWriter);

google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPool.java

@@ -288,7 +288,8 @@ private ConnectionWorker createOrReuseConnectionWorker(
  String streamReference = streamWriter.getStreamName();
  if (connectionWorkerPool.size() < currentMaxConnectionCount) {
  // Always create a new connection if we haven't reached current maximum.
- return createConnectionWorker(streamWriter.getStreamName(), streamWriter.getProtoSchema());
+ return createConnectionWorker(
+ streamWriter.getStreamName(), streamWriter.getLocation(), streamWriter.getProtoSchema());
  } else {
  ConnectionWorker existingBestConnection =
  pickBestLoadConnection(
@@ -304,7 +305,10 @@ private ConnectionWorker createOrReuseConnectionWorker(
  if (currentMaxConnectionCount > settings.maxConnectionsPerRegion()) {
  currentMaxConnectionCount = settings.maxConnectionsPerRegion();
  }
- return createConnectionWorker(streamWriter.getStreamName(), streamWriter.getProtoSchema());
+ return createConnectionWorker(
+ streamWriter.getStreamName(),
+ streamWriter.getLocation(),
+ streamWriter.getProtoSchema());
  } else {
  // Stick to the original connection if all the connections are overwhelmed.
  if (existingConnectionWorker != null) {
@@ -359,15 +363,16 @@ static ConnectionWorker pickBestLoadConnection(
  * a single stream reference. This is because createConnectionWorker(...) is called via
  * computeIfAbsent(...) which is at most once per key.
  */
- private ConnectionWorker createConnectionWorker(String streamName, ProtoSchema writeSchema)
- throws IOException {
+ private ConnectionWorker createConnectionWorker(
+ String streamName, String location, ProtoSchema writeSchema) throws IOException {
  if (enableTesting) {
  // Though atomic integer is super lightweight, add extra if check in case adding future logic.
  testValueCreateConnectionCount.getAndIncrement();
  }
  ConnectionWorker connectionWorker =
  new ConnectionWorker(
  streamName,
+ location,
  writeSchema,
  maxInflightRequests,
  maxInflightBytes,

google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/StreamWriter.java

@@ -208,6 +208,7 @@ private StreamWriter(Builder builder) throws IOException {
  SingleConnectionOrConnectionPool.ofSingleConnection(
  new ConnectionWorker(
  builder.streamName,
+ builder.location,
  builder.writerSchema,
  builder.maxInflightRequest,
  builder.maxInflightBytes,

google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPoolTest.java

@@ -430,6 +430,7 @@ private StreamWriter getTestStreamWriter(String streamName) throws IOException {
  return StreamWriter.newBuilder(streamName)
  .setWriterSchema(createProtoSchema())
  .setTraceId(TEST_TRACE_ID)
+ .setLocation("us")
  .setCredentialsProvider(NoCredentialsProvider.create())
  .setChannelProvider(serviceHelper.createChannelProvider())
  .build();

google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerTest.java

@@ -39,13 +39,15 @@
 import java.util.List;
 import java.util.UUID;
 import java.util.concurrent.ExecutionException;
+import java.util.logging.Logger;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
 @RunWith(JUnit4.class)
 public class ConnectionWorkerTest {
+ private static final Logger log = Logger.getLogger(StreamWriter.class.getName());
  private static final String TEST_STREAM_1 = "projects/p1/datasets/d1/tables/t1/streams/s1";
  private static final String TEST_STREAM_2 = "projects/p2/datasets/d2/tables/t2/streams/s2";
  private static final String TEST_TRACE_ID = "DATAFLOW:job_id";
@@ -84,10 +86,12 @@ public void testMultiplexedAppendSuccess() throws Exception {
  StreamWriter sw1 =
  StreamWriter.newBuilder(TEST_STREAM_1, client)
  .setWriterSchema(createProtoSchema("foo"))
+ .setLocation("us")
  .build();
  StreamWriter sw2 =
  StreamWriter.newBuilder(TEST_STREAM_2, client)
  .setWriterSchema(createProtoSchema("complicate"))
+ .setLocation("us")
  .build();
  // We do a pattern of:
  // send to stream1, string1
@@ -205,11 +209,20 @@ public void testAppendInSameStream_switchSchema() throws Exception {
  // send to stream1, schema1
  // ...
  StreamWriter sw1 =
- StreamWriter.newBuilder(TEST_STREAM_1, client).setWriterSchema(schema1).build();
+ StreamWriter.newBuilder(TEST_STREAM_1, client)
+ .setLocation("us")
+ .setWriterSchema(schema1)
+ .build();
  StreamWriter sw2 =
- StreamWriter.newBuilder(TEST_STREAM_1, client).setWriterSchema(schema2).build();
+ StreamWriter.newBuilder(TEST_STREAM_1, client)
+ .setLocation("us")
+ .setWriterSchema(schema2)
+ .build();
  StreamWriter sw3 =
- StreamWriter.newBuilder(TEST_STREAM_1, client).setWriterSchema(schema3).build();
+ StreamWriter.newBuilder(TEST_STREAM_1, client)
+ .setLocation("us")
+ .setWriterSchema(schema3)
+ .build();
  for (long i = 0; i < appendCount; i++) {
  switch ((int) i % 4) {
  case 0:
@@ -305,10 +318,14 @@ public void testAppendInSameStream_switchSchema() throws Exception {
  public void testAppendButInflightQueueFull() throws Exception {
  ProtoSchema schema1 = createProtoSchema("foo");
  StreamWriter sw1 =
- StreamWriter.newBuilder(TEST_STREAM_1, client).setWriterSchema(schema1).build();
+ StreamWriter.newBuilder(TEST_STREAM_1, client)
+ .setLocation("us")
+ .setWriterSchema(schema1)
+ .build();
  ConnectionWorker connectionWorker =
  new ConnectionWorker(
  TEST_STREAM_1,
+ "us",
  createProtoSchema("foo"),
  6,
  100000,
@@ -356,10 +373,14 @@ public void testAppendButInflightQueueFull() throws Exception {
  public void testThrowExceptionWhileWithinAppendLoop() throws Exception {
  ProtoSchema schema1 = createProtoSchema("foo");
  StreamWriter sw1 =
- StreamWriter.newBuilder(TEST_STREAM_1, client).setWriterSchema(schema1).build();
+ StreamWriter.newBuilder(TEST_STREAM_1, client)
+ .setLocation("us")
+ .setWriterSchema(schema1)
+ .build();
  ConnectionWorker connectionWorker =
  new ConnectionWorker(
  TEST_STREAM_1,
+ "us",
  createProtoSchema("foo"),
  100000,
  100000,
@@ -411,6 +432,69 @@ public void testThrowExceptionWhileWithinAppendLoop() throws Exception {
  assertThat(ex.getCause()).hasMessageThat().contains("Any exception can happen.");
  }
 
+ @Test
+ public void testLocationMismatch() throws Exception {
+ ProtoSchema schema1 = createProtoSchema("foo");
+ StreamWriter sw1 =
+ StreamWriter.newBuilder(TEST_STREAM_1, client)
+ .setWriterSchema(schema1)
+ .setLocation("eu")
+ .build();
+ ConnectionWorker connectionWorker =
+ new ConnectionWorker(
+ TEST_STREAM_1,
+ "us",
+ createProtoSchema("foo"),
+ 100000,
+ 100000,
+ Duration.ofSeconds(100),
+ FlowController.LimitExceededBehavior.Block,
+ TEST_TRACE_ID,
+ client.getSettings());
+ StatusRuntimeException ex =
+ assertThrows(
+ StatusRuntimeException.class,
+ () ->
+ sendTestMessage(
+ connectionWorker,
+ sw1,
+ createFooProtoRows(new String[] {String.valueOf(0)}),
+ 0));
+ assertEquals(
+ "INVALID_ARGUMENT: StreamWriter with location eu is scheduled to use a connection with location us",
+ ex.getMessage());
+ }
+
+ @Test
+ public void testStreamNameMismatch() throws Exception {
+ ProtoSchema schema1 = createProtoSchema("foo");
+ StreamWriter sw1 =
+ StreamWriter.newBuilder(TEST_STREAM_1, client).setWriterSchema(schema1).build();
+ ConnectionWorker connectionWorker =
+ new ConnectionWorker(
+ TEST_STREAM_2,
+ null,
+ createProtoSchema("foo"),
+ 100000,
+ 100000,
+ Duration.ofSeconds(100),
+ FlowController.LimitExceededBehavior.Block,
+ TEST_TRACE_ID,
+ client.getSettings());
+ StatusRuntimeException ex =
+ assertThrows(
+ StatusRuntimeException.class,
+ () ->
+ sendTestMessage(
+ connectionWorker,
+ sw1,
+ createFooProtoRows(new String[] {String.valueOf(0)}),
+ 0));
+ assertEquals(
+ "INVALID_ARGUMENT: StreamWriter with stream name projects/p1/datasets/d1/tables/t1/streams/s1 is scheduled to use a connection with stream name projects/p2/datasets/d2/tables/t2/streams/s2",
+ ex.getMessage());
+ }
+
  @Test
  public void testExponentialBackoff() throws Exception {
  assertThat(ConnectionWorker.calculateSleepTimeMilli(0)).isEqualTo(1);
@@ -440,6 +524,7 @@ private ConnectionWorker createConnectionWorker(
  throws IOException {
  return new ConnectionWorker(
  streamName,
+ "us",
  createProtoSchema("foo"),
  maxRequests,
  maxBytes,