Skip to content

Commit 0cd71b5

Browse files
feat: include User agent (#404)
* feat: populate UserAgent in addition to x-goog-api-client * add a comment
1 parent 939fefa commit 0cd71b5

File tree

2 files changed

+56
-3
lines changed

2 files changed

+56
-3
lines changed

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java

+17-2
Original file line numberDiff line numberDiff line change
@@ -20,9 +20,11 @@
2020
import com.google.api.gax.batching.BatchingSettings;
2121
import com.google.api.gax.batching.FlowControlSettings;
2222
import com.google.api.gax.batching.FlowController.LimitExceededBehavior;
23+
import com.google.api.gax.core.GaxProperties;
2324
import com.google.api.gax.core.GoogleCredentialsProvider;
2425
import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
2526
import com.google.api.gax.retrying.RetrySettings;
27+
import com.google.api.gax.rpc.FixedHeaderProvider;
2628
import com.google.api.gax.rpc.ServerStreamingCallSettings;
2729
import com.google.api.gax.rpc.StatusCode.Code;
2830
import com.google.api.gax.rpc.StubSettings;
@@ -38,8 +40,10 @@
3840
import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsBatchingDescriptor;
3941
import com.google.common.base.Preconditions;
4042
import com.google.common.collect.ImmutableList;
43+
import com.google.common.collect.ImmutableMap;
4144
import com.google.common.collect.ImmutableSet;
4245
import java.util.List;
46+
import java.util.Map;
4347
import java.util.Set;
4448
import java.util.logging.Logger;
4549
import javax.annotation.Nonnull;
@@ -530,8 +534,19 @@ private Builder() {
530534
setTransportChannelProvider(defaultTransportChannelProvider());
531535
setStreamWatchdogCheckInterval(baseDefaults.getStreamWatchdogCheckInterval());
532536
setStreamWatchdogProvider(baseDefaults.getStreamWatchdogProvider());
533-
setInternalHeaderProvider(
534-
BigtableStubSettings.defaultApiClientHeaderProviderBuilder().build());
537+
538+
// Inject the UserAgent in addition to api-client header
539+
Map<String, String> headers =
540+
ImmutableMap.<String, String>builder()
541+
.putAll(
542+
BigtableStubSettings.defaultApiClientHeaderProviderBuilder().build().getHeaders())
543+
// GrpcHeaderInterceptor treats the `user-agent` as a magic string
544+
.put(
545+
"user-agent",
546+
"bigtable-java/"
547+
+ GaxProperties.getLibraryVersion(EnhancedBigtableStubSettings.class))
548+
.build();
549+
setInternalHeaderProvider(FixedHeaderProvider.create(headers));
535550

536551
// Per-method settings using baseSettings for defaults.
537552
readRowsSettings = ServerStreamingCallSettings.newBuilder();

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java

+39-1
Original file line numberDiff line numberDiff line change
@@ -33,8 +33,13 @@
3333
import com.google.protobuf.ByteString;
3434
import com.google.protobuf.BytesValue;
3535
import com.google.protobuf.StringValue;
36+
import io.grpc.Metadata;
3637
import io.grpc.Server;
3738
import io.grpc.ServerBuilder;
39+
import io.grpc.ServerCall;
40+
import io.grpc.ServerCall.Listener;
41+
import io.grpc.ServerCallHandler;
42+
import io.grpc.ServerInterceptor;
3843
import io.grpc.stub.StreamObserver;
3944
import java.io.IOException;
4045
import java.net.ServerSocket;
@@ -56,6 +61,7 @@ public class EnhancedBigtableStubTest {
5661
private static final String APP_PROFILE_ID = "app-profile-id";
5762

5863
private Server server;
64+
private MetadataInterceptor metadataInterceptor;
5965
private FakeDataService fakeDataService;
6066
private EnhancedBigtableStubSettings defaultSettings;
6167
private EnhancedBigtableStub enhancedBigtableStub;
@@ -66,8 +72,13 @@ public void setUp() throws IOException, IllegalAccessException, InstantiationExc
6672
try (ServerSocket ss = new ServerSocket(0)) {
6773
port = ss.getLocalPort();
6874
}
75+
metadataInterceptor = new MetadataInterceptor();
6976
fakeDataService = new FakeDataService();
70-
server = ServerBuilder.forPort(port).addService(fakeDataService).build();
77+
server =
78+
ServerBuilder.forPort(port)
79+
.intercept(metadataInterceptor)
80+
.addService(fakeDataService)
81+
.build();
7182
server.start();
7283

7384
defaultSettings =
@@ -137,6 +148,33 @@ public void testChannelPrimerConfigured() throws IOException {
137148
}
138149
}
139150

151+
@Test
152+
public void testUserAgent() throws InterruptedException {
153+
ServerStreamingCallable<Query, Row> streamingCallable =
154+
enhancedBigtableStub.createReadRowsCallable(new DefaultRowAdapter());
155+
156+
Query request = Query.create("table-id").rowKey("row-key");
157+
streamingCallable.call(request).iterator().next();
158+
159+
assertThat(metadataInterceptor.headers).hasSize(1);
160+
Metadata metadata = metadataInterceptor.headers.take();
161+
assertThat(metadata.get(Metadata.Key.of("user-agent", Metadata.ASCII_STRING_MARSHALLER)))
162+
.contains("bigtable-java/");
163+
}
164+
165+
private static class MetadataInterceptor implements ServerInterceptor {
166+
final BlockingQueue<Metadata> headers = Queues.newLinkedBlockingDeque();
167+
168+
@Override
169+
public <ReqT, RespT> Listener<ReqT> interceptCall(
170+
ServerCall<ReqT, RespT> serverCall,
171+
Metadata metadata,
172+
ServerCallHandler<ReqT, RespT> serverCallHandler) {
173+
headers.add(metadata);
174+
return serverCallHandler.startCall(serverCall, metadata);
175+
}
176+
}
177+
140178
private static class FakeDataService extends BigtableGrpc.BigtableImplBase {
141179
final BlockingQueue<ReadRowsRequest> requests = Queues.newLinkedBlockingDeque();
142180

0 commit comments

Comments
 (0)