feat: include User agent (googleapis#404)

* feat: populate UserAgent in addition to x-goog-api-client * add a comment
ad548 · Oct 6, 2020 · 0cd71b5 · 0cd71b5
1 parent 939fefa
commit 0cd71b5
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 3 deletions.
diff --git a/...le/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java b/...le/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java
@@ -20,9 +20,11 @@
 import com.google.api.gax.batching.BatchingSettings;
 import com.google.api.gax.batching.FlowControlSettings;
 import com.google.api.gax.batching.FlowController.LimitExceededBehavior;
+import com.google.api.gax.core.GaxProperties;
 import com.google.api.gax.core.GoogleCredentialsProvider;
 import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
 import com.google.api.gax.retrying.RetrySettings;
+import com.google.api.gax.rpc.FixedHeaderProvider;
 import com.google.api.gax.rpc.ServerStreamingCallSettings;
 import com.google.api.gax.rpc.StatusCode.Code;
 import com.google.api.gax.rpc.StubSettings;
@@ -38,8 +40,10 @@
 import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsBatchingDescriptor;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.logging.Logger;
 import javax.annotation.Nonnull;
@@ -530,8 +534,19 @@ private Builder() {
       setTransportChannelProvider(defaultTransportChannelProvider());
       setStreamWatchdogCheckInterval(baseDefaults.getStreamWatchdogCheckInterval());
       setStreamWatchdogProvider(baseDefaults.getStreamWatchdogProvider());
-      setInternalHeaderProvider(
-          BigtableStubSettings.defaultApiClientHeaderProviderBuilder().build());
+
+      // Inject the UserAgent in addition to api-client header
+      Map<String, String> headers =
+          ImmutableMap.<String, String>builder()
+              .putAll(
+                  BigtableStubSettings.defaultApiClientHeaderProviderBuilder().build().getHeaders())
+              // GrpcHeaderInterceptor treats the `user-agent` as a magic string
+              .put(
+                  "user-agent",
+                  "bigtable-java/"
+                      + GaxProperties.getLibraryVersion(EnhancedBigtableStubSettings.class))
+              .build();
+      setInternalHeaderProvider(FixedHeaderProvider.create(headers));
 
       // Per-method settings using baseSettings for defaults.
       readRowsSettings = ServerStreamingCallSettings.newBuilder();

diff --git a/...gtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java b/...gtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java
@@ -33,8 +33,13 @@
 import com.google.protobuf.ByteString;
 import com.google.protobuf.BytesValue;
 import com.google.protobuf.StringValue;
+import io.grpc.Metadata;
 import io.grpc.Server;
 import io.grpc.ServerBuilder;
+import io.grpc.ServerCall;
+import io.grpc.ServerCall.Listener;
+import io.grpc.ServerCallHandler;
+import io.grpc.ServerInterceptor;
 import io.grpc.stub.StreamObserver;
 import java.io.IOException;
 import java.net.ServerSocket;
@@ -56,6 +61,7 @@ public class EnhancedBigtableStubTest {
   private static final String APP_PROFILE_ID = "app-profile-id";
 
   private Server server;
+  private MetadataInterceptor metadataInterceptor;
   private FakeDataService fakeDataService;
   private EnhancedBigtableStubSettings defaultSettings;
   private EnhancedBigtableStub enhancedBigtableStub;
@@ -66,8 +72,13 @@ public void setUp() throws IOException, IllegalAccessException, InstantiationExc
     try (ServerSocket ss = new ServerSocket(0)) {
       port = ss.getLocalPort();
     }
+    metadataInterceptor = new MetadataInterceptor();
     fakeDataService = new FakeDataService();
-    server = ServerBuilder.forPort(port).addService(fakeDataService).build();
+    server =
+        ServerBuilder.forPort(port)
+            .intercept(metadataInterceptor)
+            .addService(fakeDataService)
+            .build();
     server.start();
 
     defaultSettings =
@@ -137,6 +148,33 @@ public void testChannelPrimerConfigured() throws IOException {
     }
   }
 
+  @Test
+  public void testUserAgent() throws InterruptedException {
+    ServerStreamingCallable<Query, Row> streamingCallable =
+        enhancedBigtableStub.createReadRowsCallable(new DefaultRowAdapter());
+
+    Query request = Query.create("table-id").rowKey("row-key");
+    streamingCallable.call(request).iterator().next();
+
+    assertThat(metadataInterceptor.headers).hasSize(1);
+    Metadata metadata = metadataInterceptor.headers.take();
+    assertThat(metadata.get(Metadata.Key.of("user-agent", Metadata.ASCII_STRING_MARSHALLER)))
+        .contains("bigtable-java/");
+  }
+
+  private static class MetadataInterceptor implements ServerInterceptor {
+    final BlockingQueue<Metadata> headers = Queues.newLinkedBlockingDeque();
+
+    @Override
+    public <ReqT, RespT> Listener<ReqT> interceptCall(
+        ServerCall<ReqT, RespT> serverCall,
+        Metadata metadata,
+        ServerCallHandler<ReqT, RespT> serverCallHandler) {
+      headers.add(metadata);
+      return serverCallHandler.startCall(serverCall, metadata);
+    }
+  }
+
   private static class FakeDataService extends BigtableGrpc.BigtableImplBase {
     final BlockingQueue<ReadRowsRequest> requests = Queues.newLinkedBlockingDeque();