Skip to content

Commit

Permalink
feat: Allow user-agents to be specified by both internal headers and …
Browse files Browse the repository at this point in the history
…user headers

This is primarily meant to address the googleapis/java-bigtable#404 (review). Where both the hbase adapter and the core client needs to set the UserAgent. To reconcile the conflict this PR takes grpc's approach and concatenates the values. The reason this needs to happen now is that for DirectPath we need to ship clients that send a UserAgent, however Bigtable needs to keep stats to differentiate client usage.
  • Loading branch information
igorbernstein2 committed Sep 23, 2020
1 parent 287cada commit 7e49fc0
Show file tree
Hide file tree
Showing 2 changed files with 107 additions and 18 deletions.
45 changes: 27 additions & 18 deletions gax/src/main/java/com/google/api/gax/rpc/ClientContext.java
Expand Up @@ -42,10 +42,13 @@
import com.google.auto.value.AutoValue;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Sets;
import java.io.IOException;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
Expand Down Expand Up @@ -222,27 +225,33 @@ public static ClientContext create(StubSettings settings) throws IOException {
* Project Id.
*/
private static Map<String, String> getHeadersFromSettings(StubSettings settings) {
ImmutableMap.Builder<String, String> headersBuilder = ImmutableMap.builder();
if (settings.getQuotaProjectId() != null) {
headersBuilder.put(QUOTA_PROJECT_ID_HEADER_KEY, settings.getQuotaProjectId());
for (Map.Entry<String, String> entry : settings.getHeaderProvider().getHeaders().entrySet()) {
if (entry.getKey().equals(QUOTA_PROJECT_ID_HEADER_KEY)) {
continue;
}
headersBuilder.put(entry);
// Resolve conflicts when merging headers from multiple sources
Map<String, String> userHeaders = settings.getHeaderProvider().getHeaders();
Map<String, String> internalHeaders = settings.getInternalHeaderProvider().getHeaders();
Map<String, String> conflictResolution = new HashMap<>();

Set<String> conflicts = Sets.intersection(userHeaders.keySet(), internalHeaders.keySet());
for (String key : conflicts) {
if ("user-agent".equals(key)) {
conflictResolution.put(key, userHeaders.get(key) + " " + internalHeaders.get(key));
continue;
}
for (Map.Entry<String, String> entry :
settings.getInternalHeaderProvider().getHeaders().entrySet()) {
if (entry.getKey().equals(QUOTA_PROJECT_ID_HEADER_KEY)) {
continue;
}
headersBuilder.put(entry);
// Backwards compat: quota project id can conflict if its overriden in settings
if (QUOTA_PROJECT_ID_HEADER_KEY.equals(key) && settings.getQuotaProjectId() != null) {
continue;
}
} else {
headersBuilder.putAll(settings.getHeaderProvider().getHeaders());
headersBuilder.putAll(settings.getInternalHeaderProvider().getHeaders());
throw new IllegalArgumentException("Header provider can't override the header: " + key);
}
if (settings.getQuotaProjectId() != null) {
conflictResolution.put(QUOTA_PROJECT_ID_HEADER_KEY, settings.getQuotaProjectId());
}
return headersBuilder.build();

Map<String, String> effectiveHeaders = new HashMap<>();
effectiveHeaders.putAll(internalHeaders);
effectiveHeaders.putAll(userHeaders);
effectiveHeaders.putAll(conflictResolution);

return ImmutableMap.copyOf(effectiveHeaders);
}

@AutoValue.Builder
Expand Down
80 changes: 80 additions & 0 deletions gax/src/test/java/com/google/api/gax/rpc/ClientContextTest.java
Expand Up @@ -36,6 +36,7 @@
import com.google.api.gax.core.CredentialsProvider;
import com.google.api.gax.core.ExecutorProvider;
import com.google.api.gax.core.FixedCredentialsProvider;
import com.google.api.gax.core.FixedExecutorProvider;
import com.google.api.gax.rpc.testing.FakeChannel;
import com.google.api.gax.rpc.testing.FakeClientSettings;
import com.google.api.gax.rpc.testing.FakeTransportChannel;
Expand Down Expand Up @@ -193,6 +194,14 @@ public TransportChannelProvider withCredentials(Credentials credentials) {
}
}

private class FakeChannelWithHeaders extends FakeChannel {
private final Map<String, String> headers;

FakeChannelWithHeaders(Map<String, String> headers) {
this.headers = headers;
}
}

@Test
public void testNoAutoCloseContextNeedsNoExecutor() throws Exception {
runTest(false, false, false, false);
Expand Down Expand Up @@ -526,4 +535,75 @@ public Credentials getCredentials() throws IOException {
ClientContext clientContext = ClientContext.create(builder.build());
assertThat(clientContext.getCredentials().getRequestMetadata(null)).isEqualTo(metaData);
}

@Test
public void testUserAgentInternalOnly() throws Exception {
TransportChannelProvider transportChannelProvider =
new FakeTransportProvider(
FakeTransportChannel.create(new FakeChannel()), null, true, null, null);

ClientSettings.Builder builder =
new FakeClientSettings.Builder()
.setExecutorProvider(
FixedExecutorProvider.create(Mockito.mock(ScheduledExecutorService.class)))
.setTransportChannelProvider(transportChannelProvider)
.setCredentialsProvider(
FixedCredentialsProvider.create(Mockito.mock(GoogleCredentials.class)));

builder.setInternalHeaderProvider(FixedHeaderProvider.create("user-agent", "internal-agent"));

ClientContext clientContext = ClientContext.create(builder.build());
FakeTransportChannel transportChannel =
(FakeTransportChannel) clientContext.getTransportChannel();

assertThat(transportChannel.getHeaders()).containsEntry("user-agent", "internal-agent");
}

@Test
public void testUserAgentExternalOnly() throws Exception {
TransportChannelProvider transportChannelProvider =
new FakeTransportProvider(
FakeTransportChannel.create(new FakeChannel()), null, true, null, null);

ClientSettings.Builder builder =
new FakeClientSettings.Builder()
.setExecutorProvider(
FixedExecutorProvider.create(Mockito.mock(ScheduledExecutorService.class)))
.setTransportChannelProvider(transportChannelProvider)
.setCredentialsProvider(
FixedCredentialsProvider.create(Mockito.mock(GoogleCredentials.class)));

builder.setHeaderProvider(FixedHeaderProvider.create("user-agent", "user-supplied-agent"));

ClientContext clientContext = ClientContext.create(builder.build());
FakeTransportChannel transportChannel =
(FakeTransportChannel) clientContext.getTransportChannel();

assertThat(transportChannel.getHeaders()).containsEntry("user-agent", "user-supplied-agent");
}

@Test
public void testUserAgentConcat() throws Exception {
TransportChannelProvider transportChannelProvider =
new FakeTransportProvider(
FakeTransportChannel.create(new FakeChannel()), null, true, null, null);

ClientSettings.Builder builder =
new FakeClientSettings.Builder()
.setExecutorProvider(
FixedExecutorProvider.create(Mockito.mock(ScheduledExecutorService.class)))
.setTransportChannelProvider(transportChannelProvider)
.setCredentialsProvider(
FixedCredentialsProvider.create(Mockito.mock(GoogleCredentials.class)));

builder.setHeaderProvider(FixedHeaderProvider.create("user-agent", "user-supplied-agent"));
builder.setInternalHeaderProvider(FixedHeaderProvider.create("user-agent", "internal-agent"));

ClientContext clientContext = ClientContext.create(builder.build());
FakeTransportChannel transportChannel =
(FakeTransportChannel) clientContext.getTransportChannel();

assertThat(transportChannel.getHeaders())
.containsEntry("user-agent", "user-supplied-agent internal-agent");
}
}

0 comments on commit 7e49fc0

Please sign in to comment.