updating the header name

googleapis · Oct 19, 2023 · 678bc35 · 678bc35
1 parent 9333ba0
commit 678bc35
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 26 deletions.
diff --git a/...loud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/CookieInterceptor.java b/...loud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/CookieInterceptor.java
@@ -18,7 +18,6 @@
 import static com.google.cloud.bigtable.data.v2.stub.CookiesHolder.COOKIES_HOLDER_KEY;
 import static com.google.cloud.bigtable.data.v2.stub.CookiesHolder.ROUTING_COOKIE_METADATA_KEY;
 
-import com.google.protobuf.ByteString;
 import io.grpc.CallOptions;
 import io.grpc.Channel;
 import io.grpc.ClientCall;
@@ -44,7 +43,7 @@ public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
       public void start(Listener<RespT> responseListener, Metadata headers) {
         CookiesHolder cookie = callOptions.getOption(COOKIES_HOLDER_KEY);
         if (cookie != null && cookie.getRoutingCookie() != null) {
-          headers.put(ROUTING_COOKIE_METADATA_KEY, cookie.getRoutingCookie().toByteArray());
+          headers.put(ROUTING_COOKIE_METADATA_KEY, cookie.getRoutingCookie());
         }
         super.start(new UpdateCookieListener<>(responseListener, callOptions), headers);
       }
@@ -63,10 +62,12 @@ static class UpdateCookieListener<RespT>
 
     @Override
     public void onClose(Status status, Metadata trailers) {
-      if (trailers != null && trailers.containsKey(ROUTING_COOKIE_METADATA_KEY)) {
-        byte[] cookie = trailers.get(ROUTING_COOKIE_METADATA_KEY);
-        CookiesHolder cookieHolder = callOptions.getOption(COOKIES_HOLDER_KEY);
-        cookieHolder.setRoutingCookie(ByteString.copyFrom(cookie));
+      CookiesHolder cookiesHolder = callOptions.getOption(COOKIES_HOLDER_KEY);
+      if (cookiesHolder != null
+          && trailers != null
+          && trailers.containsKey(ROUTING_COOKIE_METADATA_KEY)) {
+        String cookie = trailers.get(ROUTING_COOKIE_METADATA_KEY);
+        cookiesHolder.setRoutingCookie(cookie);
       }
       super.onClose(status, trailers);
     }

diff --git a/...le-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/CookiesHolder.java b/...le-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/CookiesHolder.java
@@ -15,7 +15,6 @@
  */
 package com.google.cloud.bigtable.data.v2.stub;
 
-import com.google.protobuf.ByteString;
 import io.grpc.CallOptions;
 import io.grpc.Metadata;
 import javax.annotation.Nullable;
@@ -28,16 +27,16 @@ class CookiesHolder {
 
   static final String ROUTING_COOKIE_KEY = "x-goog-cbt-cookie-routing";
 
-  static final Metadata.Key<byte[]> ROUTING_COOKIE_METADATA_KEY =
-      Metadata.Key.of(ROUTING_COOKIE_KEY + "-bin", Metadata.BINARY_BYTE_MARSHALLER);
+  static final Metadata.Key<String> ROUTING_COOKIE_METADATA_KEY =
+      Metadata.Key.of(ROUTING_COOKIE_KEY, Metadata.ASCII_STRING_MARSHALLER);
 
-  @Nullable private ByteString routingCookie;
+  @Nullable private String routingCookie;
 
-  void setRoutingCookie(@Nullable ByteString routingCookie) {
+  void setRoutingCookie(@Nullable String routingCookie) {
     this.routingCookie = routingCookie;
   }
 
-  ByteString getRoutingCookie() {
+  String getRoutingCookie() {
     return this.routingCookie;
   }
 }
diff --git a/...cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/CookieHolderTest.java b/...cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/CookieHolderTest.java
@@ -29,7 +29,6 @@
 import com.google.cloud.bigtable.data.v2.FakeServiceBuilder;
 import com.google.cloud.bigtable.data.v2.models.Query;
 import com.google.common.collect.ImmutableList;
-import com.google.protobuf.ByteString;
 import io.grpc.Metadata;
 import io.grpc.Server;
 import io.grpc.ServerCall;
@@ -38,7 +37,6 @@
 import io.grpc.Status;
 import io.grpc.StatusRuntimeException;
 import io.grpc.stub.StreamObserver;
-import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -52,10 +50,9 @@
 public class CookieHolderTest {
   private Server server;
   private FakeService fakeService = new FakeService();
-
   private BigtableDataClient client;
-
   private List<Metadata> serverMetadata = new ArrayList<>();
+  private String testCookie = "test-routing-cookie";
 
   @Before
   public void setup() throws Exception {
@@ -97,9 +94,9 @@ public void testReadRowsRetryCookieIsForwarded() {
 
     assertThat(fakeService.count.get()).isGreaterThan(1);
     assertThat(serverMetadata.size()).isEqualTo(fakeService.count.get());
-    byte[] bytes = serverMetadata.get(1).get(ROUTING_COOKIE_METADATA_KEY);
+    String bytes = serverMetadata.get(1).get(ROUTING_COOKIE_METADATA_KEY);
     assertThat(bytes).isNotNull();
-    assertThat(new String(bytes, StandardCharsets.UTF_8)).isEqualTo("test-routing-cookie");
+    assertThat(bytes).isEqualTo(testCookie);
 
     serverMetadata.clear();
   }
@@ -111,9 +108,10 @@ public void testSampleRowKeysRetryCookieIsForwarded() {
 
     assertThat(fakeService.count.get()).isGreaterThan(1);
     assertThat(serverMetadata.size()).isEqualTo(fakeService.count.get());
-    byte[] bytes = serverMetadata.get(1).get(ROUTING_COOKIE_METADATA_KEY);
+    String bytes = serverMetadata.get(1).get(ROUTING_COOKIE_METADATA_KEY);
+
     assertThat(bytes).isNotNull();
-    assertThat(new String(bytes, StandardCharsets.UTF_8)).isEqualTo("test-routing-cookie");
+    assertThat(bytes).isEqualTo(testCookie);
 
     serverMetadata.clear();
   }
@@ -133,9 +131,7 @@ public void readRows(
         ReadRowsRequest request, StreamObserver<ReadRowsResponse> responseObserver) {
       if (count.getAndIncrement() < 1) {
         Metadata trailers = new Metadata();
-        trailers.put(
-            ROUTING_COOKIE_METADATA_KEY,
-            ByteString.copyFromUtf8("test-routing-cookie").toByteArray());
+        trailers.put(ROUTING_COOKIE_METADATA_KEY, testCookie);
         StatusRuntimeException exception = new StatusRuntimeException(Status.UNAVAILABLE, trailers);
         responseObserver.onError(exception);
         return;
@@ -149,9 +145,7 @@ public void sampleRowKeys(
         SampleRowKeysRequest request, StreamObserver<SampleRowKeysResponse> responseObserver) {
       if (count.getAndIncrement() < 1) {
         Metadata trailers = new Metadata();
-        trailers.put(
-            ROUTING_COOKIE_METADATA_KEY,
-            ByteString.copyFromUtf8("test-routing-cookie").toByteArray());
+        trailers.put(ROUTING_COOKIE_METADATA_KEY, testCookie);
         StatusRuntimeException exception = new StatusRuntimeException(Status.UNAVAILABLE, trailers);
         responseObserver.onError(exception);
         return;