Skip to content
This repository was archived by the owner on Sep 26, 2023. It is now read-only.

Commit 5081ec6

Browse files
arithmetic1728chanseokoh
arithmetic1728
and
chanseokoh
authored
feat: add api key support (#1436)
* feat: add api key support * Update gax/src/test/java/com/google/api/gax/rpc/ClientContextTest.java Co-authored-by: Chanseok Oh <chanseok@google.com> * Update gax/src/test/java/com/google/api/gax/rpc/ClientContextTest.java Co-authored-by: Chanseok Oh <chanseok@google.com> * update * update * Update gax/src/test/java/com/google/api/gax/rpc/ClientContextTest.java Co-authored-by: Chanseok Oh <chanseok@google.com> * Update gax/src/test/java/com/google/api/gax/rpc/ClientContextTest.java Co-authored-by: Chanseok Oh <chanseok@google.com> Co-authored-by: Chanseok Oh <chanseok@google.com>
1 parent f631a25 commit 5081ec6

File tree

4 files changed

+180
-11
lines changed

4 files changed

+180
-11
lines changed

‎gax/src/main/java/com/google/api/gax/rpc/ClientContext.java‎

Lines changed: 48 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -35,12 +35,14 @@
3535
import com.google.api.gax.core.BackgroundResource;
3636
import com.google.api.gax.core.ExecutorAsBackgroundResource;
3737
import com.google.api.gax.core.ExecutorProvider;
38+
import com.google.api.gax.rpc.internal.EnvironmentProvider;
3839
import com.google.api.gax.rpc.internal.QuotaProjectIdHidingCredentials;
3940
import com.google.api.gax.rpc.mtls.MtlsProvider;
4041
import com.google.api.gax.tracing.ApiTracerFactory;
4142
import com.google.api.gax.tracing.BaseApiTracerFactory;
4243
import com.google.auth.Credentials;
4344
import com.google.auto.value.AutoValue;
45+
import com.google.common.annotations.VisibleForTesting;
4446
import com.google.common.collect.ImmutableList;
4547
import com.google.common.collect.ImmutableMap;
4648
import com.google.common.collect.Sets;
@@ -65,6 +67,7 @@
6567
@AutoValue
6668
public abstract class ClientContext {
6769
private static final String QUOTA_PROJECT_ID_HEADER_KEY = "x-goog-user-project";
70+
private static final String API_KEY_HEADER_KEY = "x-goog-api-key";
6871

6972
/**
7073
* The objects that need to be closed in order to clean up the resources created in the process of
@@ -159,6 +162,32 @@ static String getEndpoint(
159162
return endpoint;
160163
}
161164

165+
/**
166+
* Retrieves the API key value and add it to the headers if API key exists. It first tries to
167+
* retrieve the value from the stub settings. If not found, it then tries the load the
168+
* GOOGLE_API_KEY environment variable. An IOException will be thrown if both GOOGLE_API_KEY and
169+
* GOOGLE_APPLICATION_CREDENTIALS environment variables are set.
170+
*/
171+
@VisibleForTesting
172+
static void addApiKeyToHeaders(
173+
StubSettings settings, EnvironmentProvider environmentProvider, Map<String, String> headers)
174+
throws IOException {
175+
if (settings.getApiKey() != null) {
176+
headers.put(API_KEY_HEADER_KEY, settings.getApiKey());
177+
return;
178+
}
179+
180+
String apiKey = environmentProvider.getenv("GOOGLE_API_KEY");
181+
String applicationCredentials = environmentProvider.getenv("GOOGLE_APPLICATION_CREDENTIALS");
182+
if (apiKey != null && applicationCredentials != null) {
183+
throw new IOException(
184+
"Environment variables GOOGLE_API_KEY and GOOGLE_APPLICATION_CREDENTIALS are mutually exclusive");
185+
}
186+
if (apiKey != null) {
187+
headers.put(API_KEY_HEADER_KEY, apiKey);
188+
}
189+
}
190+
162191
/**
163192
* Instantiates the executor, credentials, and transport context based on the given client
164193
* settings.
@@ -169,14 +198,21 @@ public static ClientContext create(StubSettings settings) throws IOException {
169198
ExecutorProvider backgroundExecutorProvider = settings.getBackgroundExecutorProvider();
170199
final ScheduledExecutorService backgroundExecutor = backgroundExecutorProvider.getExecutor();
171200

172-
Credentials credentials = settings.getCredentialsProvider().getCredentials();
201+
Credentials credentials = null;
202+
Map<String, String> headers = getHeadersFromSettingsAndEnvironment(settings, System::getenv);
173203

174-
if (settings.getQuotaProjectId() != null) {
175-
// If the quotaProjectId is set, wrap original credentials with correct quotaProjectId as
176-
// QuotaProjectIdHidingCredentials.
177-
// Ensure that a custom set quota project id takes priority over one detected by credentials.
178-
// Avoid the backend receiving possibly conflict values of quotaProjectId
179-
credentials = new QuotaProjectIdHidingCredentials(credentials);
204+
boolean hasApiKey = headers.containsKey(API_KEY_HEADER_KEY);
205+
if (!hasApiKey) {
206+
credentials = settings.getCredentialsProvider().getCredentials();
207+
208+
if (settings.getQuotaProjectId() != null) {
209+
// If the quotaProjectId is set, wrap original credentials with correct quotaProjectId as
210+
// QuotaProjectIdHidingCredentials.
211+
// Ensure that a custom set quota project id takes priority over one detected by
212+
// credentials.
213+
// Avoid the backend receiving possibly conflict values of quotaProjectId
214+
credentials = new QuotaProjectIdHidingCredentials(credentials);
215+
}
180216
}
181217

182218
TransportChannelProvider transportChannelProvider = settings.getTransportChannelProvider();
@@ -186,11 +222,11 @@ public static ClientContext create(StubSettings settings) throws IOException {
186222
if (transportChannelProvider.needsExecutor() && settings.getExecutorProvider() != null) {
187223
transportChannelProvider = transportChannelProvider.withExecutor(backgroundExecutor);
188224
}
189-
Map<String, String> headers = getHeadersFromSettings(settings);
225+
190226
if (transportChannelProvider.needsHeaders()) {
191227
transportChannelProvider = transportChannelProvider.withHeaders(headers);
192228
}
193-
if (transportChannelProvider.needsCredentials() && credentials != null) {
229+
if (!hasApiKey && transportChannelProvider.needsCredentials()) {
194230
transportChannelProvider = transportChannelProvider.withCredentials(credentials);
195231
}
196232
String endpoint =
@@ -260,7 +296,8 @@ public static ClientContext create(StubSettings settings) throws IOException {
260296
* Getting a header map from HeaderProvider and InternalHeaderProvider from settings with Quota
261297
* Project Id.
262298
*/
263-
private static Map<String, String> getHeadersFromSettings(StubSettings settings) {
299+
private static Map<String, String> getHeadersFromSettingsAndEnvironment(
300+
StubSettings settings, EnvironmentProvider environmentProvider) throws IOException {
264301
// Resolve conflicts when merging headers from multiple sources
265302
Map<String, String> userHeaders = settings.getHeaderProvider().getHeaders();
266303
Map<String, String> internalHeaders = settings.getInternalHeaderProvider().getHeaders();
@@ -286,6 +323,7 @@ private static Map<String, String> getHeadersFromSettings(StubSettings settings)
286323
effectiveHeaders.putAll(internalHeaders);
287324
effectiveHeaders.putAll(userHeaders);
288325
effectiveHeaders.putAll(conflictResolution);
326+
addApiKeyToHeaders(settings, environmentProvider, effectiveHeaders);
289327

290328
return ImmutableMap.copyOf(effectiveHeaders);
291329
}

‎gax/src/main/java/com/google/api/gax/rpc/StubSettings.java‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,7 @@ public abstract class StubSettings<SettingsT extends StubSettings<SettingsT>> {
7373
private final String endpoint;
7474
private final String mtlsEndpoint;
7575
private final String quotaProjectId;
76+
private final String apiKey;
7677
@Nullable private final WatchdogProvider streamWatchdogProvider;
7778
@Nonnull private final Duration streamWatchdogCheckInterval;
7879
@Nonnull private final ApiTracerFactory tracerFactory;
@@ -99,6 +100,7 @@ protected StubSettings(Builder builder) {
99100
this.mtlsEndpoint = builder.mtlsEndpoint;
100101
this.switchToMtlsEndpointAllowed = builder.switchToMtlsEndpointAllowed;
101102
this.quotaProjectId = builder.quotaProjectId;
103+
this.apiKey = builder.apiKey;
102104
this.streamWatchdogProvider = builder.streamWatchdogProvider;
103105
this.streamWatchdogCheckInterval = builder.streamWatchdogCheckInterval;
104106
this.tracerFactory = builder.tracerFactory;
@@ -154,6 +156,10 @@ public final String getQuotaProjectId() {
154156
return quotaProjectId;
155157
}
156158

159+
public final String getApiKey() {
160+
return apiKey;
161+
}
162+
157163
@BetaApi("The surface for streaming is not stable yet and may change in the future.")
158164
@Nullable
159165
public final WatchdogProvider getStreamWatchdogProvider() {
@@ -189,6 +195,7 @@ public String toString() {
189195
.add("mtlsEndpoint", mtlsEndpoint)
190196
.add("switchToMtlsEndpointAllowed", switchToMtlsEndpointAllowed)
191197
.add("quotaProjectId", quotaProjectId)
198+
.add("apiKey", apiKey)
192199
.add("streamWatchdogProvider", streamWatchdogProvider)
193200
.add("streamWatchdogCheckInterval", streamWatchdogCheckInterval)
194201
.add("tracerFactory", tracerFactory)
@@ -209,6 +216,7 @@ public abstract static class Builder<
209216
private String endpoint;
210217
private String mtlsEndpoint;
211218
private String quotaProjectId;
219+
private String apiKey;
212220
@Nullable private WatchdogProvider streamWatchdogProvider;
213221
@Nonnull private Duration streamWatchdogCheckInterval;
214222
@Nonnull private ApiTracerFactory tracerFactory;
@@ -234,6 +242,7 @@ protected Builder(StubSettings settings) {
234242
this.mtlsEndpoint = settings.mtlsEndpoint;
235243
this.switchToMtlsEndpointAllowed = settings.switchToMtlsEndpointAllowed;
236244
this.quotaProjectId = settings.quotaProjectId;
245+
this.apiKey = settings.apiKey;
237246
this.streamWatchdogProvider = settings.streamWatchdogProvider;
238247
this.streamWatchdogCheckInterval = settings.streamWatchdogCheckInterval;
239248
this.tracerFactory = settings.tracerFactory;
@@ -258,6 +267,7 @@ private static String getQuotaProjectIdFromClientContext(ClientContext clientCon
258267
}
259268

260269
protected Builder(ClientContext clientContext) {
270+
this.apiKey = null;
261271
if (clientContext == null) {
262272
this.backgroundExecutorProvider = InstantiatingExecutorProvider.newBuilder().build();
263273
this.transportChannelProvider = null;
@@ -432,6 +442,11 @@ public B setQuotaProjectId(String quotaProjectId) {
432442
return self();
433443
}
434444

445+
public B setApiKey(String apiKey) {
446+
this.apiKey = apiKey;
447+
return self();
448+
}
449+
435450
/**
436451
* Sets how often the {@link Watchdog} will check ongoing streaming RPCs. Defaults to 10 secs.
437452
* Use {@link Duration#ZERO} to disable.
@@ -513,6 +528,10 @@ public String getQuotaProjectId() {
513528
return quotaProjectId;
514529
}
515530

531+
public String getApiKey() {
532+
return apiKey;
533+
}
534+
516535
@BetaApi("The surface for streaming is not stable yet and may change in the future.")
517536
@Nonnull
518537
public Duration getStreamWatchdogCheckInterval() {
@@ -549,6 +568,7 @@ public String toString() {
549568
.add("mtlsEndpoint", mtlsEndpoint)
550569
.add("switchToMtlsEndpointAllowed", switchToMtlsEndpointAllowed)
551570
.add("quotaProjectId", quotaProjectId)
571+
.add("apiKey", apiKey)
552572
.add("streamWatchdogProvider", streamWatchdogProvider)
553573
.add("streamWatchdogCheckInterval", streamWatchdogCheckInterval)
554574
.add("tracerFactory", tracerFactory)

‎gax/src/main/java/com/google/api/gax/rpc/internal/EnvironmentProvider.java‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
/*
2+
* Copyright 2021 Google LLC
3+
*
4+
* Redistribution and use in source and binary forms, with or without
5+
* modification, are permitted provided that the following conditions are
6+
* met:
7+
*
8+
* * Redistributions of source code must retain the above copyright
9+
* notice, this list of conditions and the following disclaimer.
10+
* * Redistributions in binary form must reproduce the above
11+
* copyright notice, this list of conditions and the following disclaimer
12+
* in the documentation and/or other materials provided with the
13+
* distribution.
14+
* * Neither the name of Google LLC nor the names of its
15+
* contributors may be used to endorse or promote products derived from
16+
* this software without specific prior written permission.
17+
*
18+
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
19+
* "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
20+
* LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
21+
* A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
22+
* OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
23+
* SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
24+
* LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
25+
* DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
26+
* THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
27+
* (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
28+
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
29+
*/
30+
package com.google.api.gax.rpc.internal;
31+
32+
import com.google.api.core.InternalExtensionOnly;
33+
34+
/** Provides an interface to provide the environment variable values. */
35+
@InternalExtensionOnly
36+
public interface EnvironmentProvider {
37+
/** Returns the environment variable value. */
38+
String getenv(String name);
39+
}

‎gax/src/test/java/com/google/api/gax/rpc/ClientContextTest.java‎

Lines changed: 73 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import static com.google.common.truth.Truth.assertThat;
3333
import static org.junit.Assert.assertEquals;
3434
import static org.junit.Assert.assertFalse;
35+
import static org.junit.Assert.assertThrows;
3536
import static org.junit.Assert.assertTrue;
3637
import static org.junit.Assert.fail;
3738

@@ -41,6 +42,7 @@
4142
import com.google.api.gax.core.ExecutorProvider;
4243
import com.google.api.gax.core.FixedCredentialsProvider;
4344
import com.google.api.gax.core.FixedExecutorProvider;
45+
import com.google.api.gax.rpc.internal.EnvironmentProvider;
4446
import com.google.api.gax.rpc.mtls.MtlsProvider;
4547
import com.google.api.gax.rpc.mtls.MtlsProvider.MtlsEndpointUsagePolicy;
4648
import com.google.api.gax.rpc.testing.FakeChannel;
@@ -54,6 +56,7 @@
5456
import com.google.common.truth.Truth;
5557
import java.io.IOException;
5658
import java.util.Collections;
59+
import java.util.HashMap;
5760
import java.util.List;
5861
import java.util.Map;
5962
import java.util.concurrent.Executor;
@@ -176,7 +179,7 @@ public TransportChannelProvider withPoolSize(int size) {
176179

177180
@Override
178181
public TransportChannel getTransportChannel() throws IOException {
179-
if (needsCredentials()) {
182+
if (needsCredentials() && !headers.containsKey("x-goog-api-key")) {
180183
throw new IllegalStateException("Needs Credentials");
181184
}
182185
transport.setExecutor(executor);
@@ -769,4 +772,73 @@ public void testExecutorSettings() throws Exception {
769772
transportChannel = (FakeTransportChannel) context.getTransportChannel();
770773
assertThat(transportChannel.getExecutor()).isSameInstanceAs(executorProvider.getExecutor());
771774
}
775+
776+
@Test
777+
public void testAddApiKeyToHeadersFromStubSettings() throws IOException {
778+
StubSettings settings = new FakeStubSettings.Builder().setApiKey("stub-setting-key").build();
779+
EnvironmentProvider environmentProvider =
780+
name -> name.equals("GOOGLE_API_KEY") ? "env-key" : null;
781+
Map<String, String> headers = new HashMap<>();
782+
ClientContext.addApiKeyToHeaders(settings, environmentProvider, headers);
783+
assertThat(headers).containsEntry("x-goog-api-key", "stub-setting-key");
784+
}
785+
786+
@Test
787+
public void testAddApiKeyToHeadersFromEnvironmentProvider() throws IOException {
788+
StubSettings settings = new FakeStubSettings.Builder().build();
789+
EnvironmentProvider environmentProvider =
790+
name -> name.equals("GOOGLE_API_KEY") ? "env-key" : null;
791+
Map<String, String> headers = new HashMap<>();
792+
ClientContext.addApiKeyToHeaders(settings, environmentProvider, headers);
793+
assertThat(headers).containsEntry("x-goog-api-key", "env-key");
794+
}
795+
796+
@Test
797+
public void testAddApiKeyToHeadersNoApiKey() throws IOException {
798+
StubSettings settings = new FakeStubSettings.Builder().build();
799+
EnvironmentProvider environmentProvider = name -> null;
800+
Map<String, String> headers = new HashMap<>();
801+
ClientContext.addApiKeyToHeaders(settings, environmentProvider, headers);
802+
assertThat(headers).doesNotContainKey("x-goog-api-key");
803+
}
804+
805+
@Test
806+
public void testAddApiKeyToHeadersThrows() throws IOException {
807+
StubSettings settings = new FakeStubSettings.Builder().build();
808+
EnvironmentProvider environmentProvider =
809+
name -> name.equals("GOOGLE_API_KEY") ? "env-key" : "/path/to/adc/json";
810+
Map<String, String> headers = new HashMap<>();
811+
Exception ex =
812+
assertThrows(
813+
IOException.class,
814+
() -> ClientContext.addApiKeyToHeaders(settings, environmentProvider, headers));
815+
assertThat(ex)
816+
.hasMessageThat()
817+
.contains(
818+
"Environment variables GOOGLE_API_KEY and GOOGLE_APPLICATION_CREDENTIALS are mutually exclusive");
819+
}
820+
821+
@Test
822+
public void testApiKey() throws IOException {
823+
FakeStubSettings.Builder builder = new FakeStubSettings.Builder();
824+
825+
FakeTransportChannel transportChannel = FakeTransportChannel.create(new FakeChannel());
826+
FakeTransportProvider transportProvider =
827+
new FakeTransportProvider(transportChannel, null, true, null, null);
828+
builder.setTransportChannelProvider(transportProvider);
829+
830+
HeaderProvider headerProvider = Mockito.mock(HeaderProvider.class);
831+
Mockito.when(headerProvider.getHeaders()).thenReturn(ImmutableMap.of());
832+
builder.setHeaderProvider(headerProvider);
833+
834+
// Set API key.
835+
builder.setApiKey("key");
836+
837+
ClientContext context = ClientContext.create(builder.build());
838+
839+
// Check API key is in the transport channel's header.
840+
List<BackgroundResource> resources = context.getBackgroundResources();
841+
FakeTransportChannel fakeTransportChannel = (FakeTransportChannel) resources.get(0);
842+
assertThat(fakeTransportChannel.getHeaders()).containsEntry("x-goog-api-key", "key");
843+
}
772844
}

0 commit comments

Comments
 (0)