Skip to content

Commit

Permalink
Insert custom headers via HeaderInterceptor.
Browse files Browse the repository at this point in the history
  • Loading branch information
shinfan committed Mar 18, 2016
1 parent 9781444 commit 83402ca
Show file tree
Hide file tree
Showing 3 changed files with 135 additions and 0 deletions.
37 changes: 37 additions & 0 deletions src/main/java/com/google/api/gax/grpc/HeaderInterceptor.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
package com.google.api.gax.grpc;

import io.grpc.CallOptions;
import io.grpc.Channel;
import io.grpc.ClientCall;
import io.grpc.ClientInterceptor;
import io.grpc.ForwardingClientCall.SimpleForwardingClientCall;
import io.grpc.Metadata;
import io.grpc.MethodDescriptor;

/**
* An intercepter to handle custom header.
*/
public class HeaderInterceptor implements ClientInterceptor {
private static final Metadata.Key<String> HEADER_KEY =
Metadata.Key.of("x-google-apis-agent", Metadata.ASCII_STRING_MARSHALLER);
private final String header;

public HeaderInterceptor(String header) {
this.header = header;
}

@Override
public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
MethodDescriptor<ReqT, RespT> method,
CallOptions callOptions,
Channel next) {
ClientCall<ReqT, RespT> call = next.newCall(method, callOptions);
return new SimpleForwardingClientCall<ReqT, RespT>(call) {
@Override
public void start(ClientCall.Listener<RespT> responseListener, Metadata headers) {
headers.put(HEADER_KEY, header);
super.start(responseListener, headers);
}
};
}
}
29 changes: 29 additions & 0 deletions src/main/java/com/google/api/gax/grpc/ServiceApiSettings.java
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,8 @@
* and should not be used in production.
*/
public class ServiceApiSettings {
private String serviceGeneratorName;
private String serviceGeneratorVersion;
private ChannelProvider channelProvider;
private ExecutorProvider executorProvider;
private final ImmutableList<? extends ApiCallSettings> allMethods;
Expand All @@ -40,6 +42,12 @@ public class ServiceApiSettings {
*/
public static final int DEFAULT_EXECUTOR_THREADS = 4;

/**
* Default name and version of the service generator.
*/
private static final String DEFAULT_GENERATOR_NAME = "gapic";
private static final String DEFAULT_GENERATOR_VERSION = "0.0.0";

/**
* Constructs an instance of ServiceApiSettings.
*/
Expand Down Expand Up @@ -99,13 +107,26 @@ public ManagedChannel getChannel(Executor executor) throws IOException {

List<ClientInterceptor> interceptors = Lists.newArrayList();
interceptors.add(new ClientAuthInterceptor(settings.getCredentials(), executor));
interceptors.add(new HeaderInterceptor(serviceHeader()));

channel = NettyChannelBuilder.forAddress(settings.getServiceAddress(), settings.getPort())
.negotiationType(NegotiationType.TLS)
.intercept(interceptors)
.build();
return channel;
}

private String serviceHeader() {
// GAX version only works when the package is invoked as a jar.
String gaxVersion = ChannelProvider.class.getPackage().getImplementationVersion();
String javaVersion = Runtime.class.getPackage().getImplementationVersion();
String generatorName = serviceGeneratorVersion.isEmpty() ?
DEFAULT_GENERATOR_NAME : serviceGeneratorName;
String generatorVersion = serviceGeneratorVersion.isEmpty() ?
DEFAULT_GENERATOR_VERSION : serviceGeneratorVersion;
return String.format("gax-%s/java-%s/%s-%s",
gaxVersion, javaVersion, generatorName, generatorVersion);
}
};
return this;
}
Expand Down Expand Up @@ -190,4 +211,12 @@ public ServiceApiSettings setRetryParamsOnAllMethods(RetryParams retryParams) {
}
return this;
}

/**
* Sets the generator name and version for the GRPC custom header.
*/
public void setGeneratorHeader(String name, String version) {
this.serviceGeneratorName = name;
this.serviceGeneratorVersion = version;
}
}
69 changes: 69 additions & 0 deletions src/test/java/com/google/api/gax/grpc/HeaderInterceptorTest.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,69 @@
package com.google.api.gax.grpc;

import static org.junit.Assert.assertEquals;
import static org.mockito.Matchers.any;
import static org.mockito.Matchers.same;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
import org.mockito.ArgumentCaptor;
import org.mockito.Mock;
import org.mockito.Mockito;
import org.mockito.MockitoAnnotations;

import io.grpc.CallOptions;
import io.grpc.Channel;
import io.grpc.ClientCall;
import io.grpc.ClientInterceptors;
import io.grpc.Metadata;
import io.grpc.MethodDescriptor;

/**
* Tests for {@link HeaderInterceptor}.
*/
@RunWith(JUnit4.class)
public class HeaderInterceptorTest {

@Mock
private Channel channel;

@Mock
private ClientCall<String, Integer> call;

@Mock
private MethodDescriptor<String, Integer> method;

/**
* Sets up mocks.
*/
@Before public void setUp() {
MockitoAnnotations.initMocks(this);
when(channel.newCall(
Mockito.<MethodDescriptor<String, Integer>>any(), any(CallOptions.class)))
.thenReturn(call);
}

@Test
public void testInterceptor() {
final Metadata.Key<String> headerKey =
Metadata.Key.of("x-google-apis-agent", Metadata.ASCII_STRING_MARSHALLER);
String data = "abcd";
HeaderInterceptor interceptor = new HeaderInterceptor(data);
Channel intercepted = ClientInterceptors.intercept(channel, interceptor);
@SuppressWarnings("unchecked")
ClientCall.Listener<Integer> listener = mock(ClientCall.Listener.class);
ClientCall<String, Integer> interceptedCall = intercepted.newCall(method, CallOptions.DEFAULT);
// start() on the intercepted call will eventually reach the call created by the real channel
interceptedCall.start(listener, new Metadata());
// The headers passed to the real channel call will contain the information inserted by the
// interceptor.
ArgumentCaptor<Metadata> captor = ArgumentCaptor.forClass(Metadata.class);
verify(call).start(same(listener), captor.capture());
assertEquals(data, captor.getValue().get(headerKey));
}
}

0 comments on commit 83402ca

Please sign in to comment.