Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: include User agent #404

Merged
merged 2 commits into from Oct 6, 2020
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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();
Expand Down
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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 =
Expand Down Expand Up @@ -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();

Expand Down