Skip to content
Permalink
Browse files
feat: include User agent (#404)
* feat: populate UserAgent in addition to x-goog-api-client

* add a comment
  • Loading branch information
igorbernstein2 committed Oct 6, 2020
1 parent 939fefa commit 0cd71b59305cd7a336c223faff68402a8bee4639
@@ -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();
@@ -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();

0 comments on commit 0cd71b5

Please sign in to comment.