From 266329e89642bfc6be579e600d3f995f4416ae4e Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Thu, 24 Sep 2020 12:04:56 -0400 Subject: [PATCH] feat: Allow user-agents to be specified by both internal headers and user headers (#1190) * feat: Allow user-agents to be specified by both internal headers and user headers This is primarily meant to address the https://github.com/googleapis/java-bigtable/pull/404#pullrequestreview-480972135. 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. * remove stale code --- .../com/google/api/gax/rpc/ClientContext.java | 45 +++++++----- .../google/api/gax/rpc/ClientContextTest.java | 72 +++++++++++++++++++ 2 files changed, 99 insertions(+), 18 deletions(-) diff --git a/gax/src/main/java/com/google/api/gax/rpc/ClientContext.java b/gax/src/main/java/com/google/api/gax/rpc/ClientContext.java index 1dca243d9..ac2335235 100644 --- a/gax/src/main/java/com/google/api/gax/rpc/ClientContext.java +++ b/gax/src/main/java/com/google/api/gax/rpc/ClientContext.java @@ -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; @@ -222,27 +225,33 @@ public static ClientContext create(StubSettings settings) throws IOException { * Project Id. */ private static Map getHeadersFromSettings(StubSettings settings) { - ImmutableMap.Builder headersBuilder = ImmutableMap.builder(); - if (settings.getQuotaProjectId() != null) { - headersBuilder.put(QUOTA_PROJECT_ID_HEADER_KEY, settings.getQuotaProjectId()); - for (Map.Entry 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 userHeaders = settings.getHeaderProvider().getHeaders(); + Map internalHeaders = settings.getInternalHeaderProvider().getHeaders(); + Map conflictResolution = new HashMap<>(); + + Set 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 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 effectiveHeaders = new HashMap<>(); + effectiveHeaders.putAll(internalHeaders); + effectiveHeaders.putAll(userHeaders); + effectiveHeaders.putAll(conflictResolution); + + return ImmutableMap.copyOf(effectiveHeaders); } @AutoValue.Builder diff --git a/gax/src/test/java/com/google/api/gax/rpc/ClientContextTest.java b/gax/src/test/java/com/google/api/gax/rpc/ClientContextTest.java index 0f7305ca2..58a007644 100644 --- a/gax/src/test/java/com/google/api/gax/rpc/ClientContextTest.java +++ b/gax/src/test/java/com/google/api/gax/rpc/ClientContextTest.java @@ -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; @@ -526,4 +527,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"); + } }