feat: add routing header for multiplexed connection (#2035)

yirutang · web-flow · commit 1f2752f3988a · 2023-03-09T17:15:38.000-08:00
diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/ConnectionWorker.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/ConnectionWorker.java
@@ -51,6 +51,8 @@
 import java.util.concurrent.locks.ReentrantLock;
 import java.util.logging.Level;
 import java.util.logging.Logger;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 import javax.annotation.concurrent.GuardedBy;
 
 /**
@@ -219,11 +221,29 @@ class ConnectionWorker implements AutoCloseable {
   private RuntimeException testOnlyRunTimeExceptionInAppendLoop = null;
   private long testOnlyAppendLoopSleepTime = 0;
 
+  private static String projectMatching = "projects/[^/]+/";
+  private static Pattern streamPatternProject = Pattern.compile(projectMatching);
+
   /** The maximum size of one request. Defined by the API. */
   public static long getApiMaxRequestBytes() {
     return 10L * 1000L * 1000L; // 10 megabytes (https://en.wikipedia.org/wiki/Megabyte)
   }
 
+  static String extractProjectName(String streamName) {
+    Matcher streamMatcher = streamPatternProject.matcher(streamName);
+    if (streamMatcher.find()) {
+      return streamMatcher.group();
+    } else {
+      throw new IllegalStateException(
+          String.format("The passed in stream name does not match standard format %s", streamName));
+    }
+  }
+
+  static String getRoutingHeader(String streamName, String location) {
+    String project = extractProjectName(streamName);
+    return project + "locations/" + location;
+  }
+
   public ConnectionWorker(
       String streamName,
       String location,
@@ -259,6 +279,10 @@ public ConnectionWorker(
     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=" + getRoutingHeader(this.streamName, this.location));
     }
     BigQueryWriteSettings stubSettings =
         clientSettings
diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerTest.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerTest.java
@@ -710,4 +710,10 @@ public void testLongTimeIdleWontFail() throws Exception {
       assertEquals(i, futures.get(i).get().getAppendResult().getOffset().getValue());
     }
   }
+
+  @Test
+  public void testLocationName() throws Exception {
+    assertEquals(
+        "projects/p1/locations/us", ConnectionWorker.getRoutingHeader(TEST_STREAM_1, "us"));
+  }
 }
