feat: add api key support

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

@@ -35,7 +35,9 @@
 import com.google.api.gax.core.BackgroundResource;
 import com.google.api.gax.core.ExecutorAsBackgroundResource;
 import com.google.api.gax.core.ExecutorProvider;
+import com.google.api.gax.rpc.internal.EnvironmentProvider;
 import com.google.api.gax.rpc.internal.QuotaProjectIdHidingCredentials;
+import com.google.api.gax.rpc.internal.SystemEnvironmentProvider;
 import com.google.api.gax.rpc.mtls.MtlsProvider;
 import com.google.api.gax.tracing.ApiTracerFactory;
 import com.google.api.gax.tracing.BaseApiTracerFactory;
@@ -65,6 +67,7 @@
 @AutoValue
 public abstract class ClientContext {
   private static final String QUOTA_PROJECT_ID_HEADER_KEY = "x-goog-user-project";
+  private static final String API_KEY_HEADER_KEY = "x-goog-api-key";
 
   /**
    * The objects that need to be closed in order to clean up the resources created in the process of
@@ -159,6 +162,18 @@ static String getEndpoint(
     return endpoint;
   }
 
+  /**
+   * Returns the API key value. It first tries to retrieve the value from the stub settings. If not
+   * found, it then tries the load the GOOGLE_API_KEY environment variable. If API key is not found,
+   * null will be returned.
+   */
+  static String retrieveApiKey(StubSettings settings, EnvironmentProvider environmentProvider) {
+    if (settings.getApiKey() != null) {
+      return settings.getApiKey();
+    }
+    return environmentProvider.getenv("GOOGLE_API_KEY");
+  }
+
   /**
    * Instantiates the executor, credentials, and transport context based on the given client
    * settings.
@@ -169,14 +184,22 @@ public static ClientContext create(StubSettings settings) throws IOException {
     ExecutorProvider backgroundExecutorProvider = settings.getBackgroundExecutorProvider();
     final ScheduledExecutorService backgroundExecutor = backgroundExecutorProvider.getExecutor();
 
-    Credentials credentials = settings.getCredentialsProvider().getCredentials();
-
-    if (settings.getQuotaProjectId() != null) {
-      // If the quotaProjectId is set, wrap original credentials with correct quotaProjectId as
-      // QuotaProjectIdHidingCredentials.
-      // Ensure that a custom set quota project id takes priority over one detected by credentials.
-      // Avoid the backend receiving possibly conflict values of quotaProjectId
-      credentials = new QuotaProjectIdHidingCredentials(credentials);
+    Credentials credentials = null;
+    Map<String, String> headers = getHeadersFromSettings(settings);
+    boolean hasApiKey = headers.containsKey(API_KEY_HEADER_KEY);
+
+    // API key takes precedence, so we only load credentials if API key is not provided.
+    if (!hasApiKey) {
+      credentials = settings.getCredentialsProvider().getCredentials();
+
+      if (settings.getQuotaProjectId() != null) {
+        // If the quotaProjectId is set, wrap original credentials with correct quotaProjectId as
+        // QuotaProjectIdHidingCredentials.
+        // Ensure that a custom set quota project id takes priority over one detected by
+        // credentials.
+        // Avoid the backend receiving possibly conflict values of quotaProjectId
+        credentials = new QuotaProjectIdHidingCredentials(credentials);
+      }
     }
 
     TransportChannelProvider transportChannelProvider = settings.getTransportChannelProvider();
@@ -186,11 +209,11 @@ public static ClientContext create(StubSettings settings) throws IOException {
     if (transportChannelProvider.needsExecutor() && settings.getExecutorProvider() != null) {
       transportChannelProvider = transportChannelProvider.withExecutor(backgroundExecutor);
     }
-    Map<String, String> headers = getHeadersFromSettings(settings);
+
     if (transportChannelProvider.needsHeaders()) {
       transportChannelProvider = transportChannelProvider.withHeaders(headers);
     }
-    if (transportChannelProvider.needsCredentials() && credentials != null) {
+    if (!hasApiKey && transportChannelProvider.needsCredentials() && credentials != null) {
       transportChannelProvider = transportChannelProvider.withCredentials(credentials);
     }
     String endpoint =
@@ -258,7 +281,7 @@ public static ClientContext create(StubSettings settings) throws IOException {
 
   /**
    * Getting a header map from HeaderProvider and InternalHeaderProvider from settings with Quota
-   * Project Id.
+   * Project Id. API key will be added to the header if found.
    */
   private static Map<String, String> getHeadersFromSettings(StubSettings settings) {
     // Resolve conflicts when merging headers from multiple sources
@@ -287,6 +310,11 @@ private static Map<String, String> getHeadersFromSettings(StubSettings settings)
     effectiveHeaders.putAll(userHeaders);
     effectiveHeaders.putAll(conflictResolution);
 
+    String apiKey = retrieveApiKey(settings, new SystemEnvironmentProvider());
+    if (apiKey != null) {
+      effectiveHeaders.put(API_KEY_HEADER_KEY, apiKey);
+    }
+
     return ImmutableMap.copyOf(effectiveHeaders);
   }
 

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

@@ -73,6 +73,7 @@ public abstract class StubSettings<SettingsT extends StubSettings<SettingsT>> {
   private final String endpoint;
   private final String mtlsEndpoint;
   private final String quotaProjectId;
+  private final String apiKey;
   @Nullable private final WatchdogProvider streamWatchdogProvider;
   @Nonnull private final Duration streamWatchdogCheckInterval;
   @Nonnull private final ApiTracerFactory tracerFactory;
@@ -99,6 +100,7 @@ protected StubSettings(Builder builder) {
     this.mtlsEndpoint = builder.mtlsEndpoint;
     this.switchToMtlsEndpointAllowed = builder.switchToMtlsEndpointAllowed;
     this.quotaProjectId = builder.quotaProjectId;
+    this.apiKey = builder.apiKey;
     this.streamWatchdogProvider = builder.streamWatchdogProvider;
     this.streamWatchdogCheckInterval = builder.streamWatchdogCheckInterval;
     this.tracerFactory = builder.tracerFactory;
@@ -154,6 +156,10 @@ public final String getQuotaProjectId() {
     return quotaProjectId;
   }
 
+  public final String getApiKey() {
+    return apiKey;
+  }
+
   @BetaApi("The surface for streaming is not stable yet and may change in the future.")
   @Nullable
   public final WatchdogProvider getStreamWatchdogProvider() {
@@ -188,6 +194,7 @@ public String toString() {
         .add("mtlsEndpoint", mtlsEndpoint)
         .add("switchToMtlsEndpointAllowed", switchToMtlsEndpointAllowed)
         .add("quotaProjectId", quotaProjectId)
+        .add("apiKey", apiKey)
         .add("streamWatchdogProvider", streamWatchdogProvider)
         .add("streamWatchdogCheckInterval", streamWatchdogCheckInterval)
         .add("tracerFactory", tracerFactory)
@@ -208,6 +215,7 @@ public abstract static class Builder<
     private String endpoint;
     private String mtlsEndpoint;
     private String quotaProjectId;
+    private String apiKey;
     @Nullable private WatchdogProvider streamWatchdogProvider;
     @Nonnull private Duration streamWatchdogCheckInterval;
     @Nonnull private ApiTracerFactory tracerFactory;
@@ -233,6 +241,7 @@ protected Builder(StubSettings settings) {
       this.mtlsEndpoint = settings.mtlsEndpoint;
       this.switchToMtlsEndpointAllowed = settings.switchToMtlsEndpointAllowed;
       this.quotaProjectId = settings.quotaProjectId;
+      this.apiKey = settings.apiKey;
       this.streamWatchdogProvider = settings.streamWatchdogProvider;
       this.streamWatchdogCheckInterval = settings.streamWatchdogCheckInterval;
       this.tracerFactory = settings.tracerFactory;
@@ -257,6 +266,7 @@ private static String getQuotaProjectIdFromClientContext(ClientContext clientCon
     }
 
     protected Builder(ClientContext clientContext) {
+      this.apiKey = null;
       if (clientContext == null) {
         this.backgroundExecutorProvider = InstantiatingExecutorProvider.newBuilder().build();
         this.transportChannelProvider = null;
@@ -431,6 +441,11 @@ public B setQuotaProjectId(String quotaProjectId) {
       return self();
     }
 
+    public B setApiKey(String apiKey) {
+      this.apiKey = apiKey;
+      return self();
+    }
+
     /**
      * Sets how often the {@link Watchdog} will check ongoing streaming RPCs. Defaults to 10 secs.
      * Use {@link Duration#ZERO} to disable.
@@ -512,6 +527,10 @@ public String getQuotaProjectId() {
       return quotaProjectId;
     }
 
+    public String getApiKey() {
+      return apiKey;
+    }
+
     @BetaApi("The surface for streaming is not stable yet and may change in the future.")
     @Nonnull
     public Duration getStreamWatchdogCheckInterval() {
@@ -547,6 +566,7 @@ public String toString() {
           .add("mtlsEndpoint", mtlsEndpoint)
           .add("switchToMtlsEndpointAllowed", switchToMtlsEndpointAllowed)
           .add("quotaProjectId", quotaProjectId)
+          .add("apiKey", apiKey)
           .add("streamWatchdogProvider", streamWatchdogProvider)
           .add("streamWatchdogCheckInterval", streamWatchdogCheckInterval)
           .add("tracerFactory", tracerFactory)

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

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2021 Google LLC
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are
+ * met:
+ *
+ *     * Redistributions of source code must retain the above copyright
+ * notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above
+ * copyright notice, this list of conditions and the following disclaimer
+ * in the documentation and/or other materials provided with the
+ * distribution.
+ *     * Neither the name of Google LLC nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+ * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+ * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+ * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+ * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+package com.google.api.gax.rpc.internal;
+
+import com.google.api.core.InternalExtensionOnly;
+
+/** Provides an interface to provide the environment variable values. */
+@InternalExtensionOnly
+public interface EnvironmentProvider {
+  /** Returns the environment variable value. */
+  String getenv(String name);
+}

gax/src/main/java/com/google/api/gax/rpc/internal/SystemEnvironmentProvider.java

@@ -0,0 +1,41 @@
+/*
+ * Copyright 2021 Google LLC
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are
+ * met:
+ *
+ *     * Redistributions of source code must retain the above copyright
+ * notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above
+ * copyright notice, this list of conditions and the following disclaimer
+ * in the documentation and/or other materials provided with the
+ * distribution.
+ *     * Neither the name of Google LLC nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+ * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+ * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+ * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+ * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+package com.google.api.gax.rpc.internal;
+
+/**
+ * SystemEnvironmentProvider implements EnvironmentProvider interface and provides system
+ * environment variable values.
+ */
+public class SystemEnvironmentProvider implements EnvironmentProvider {
+  @Override
+  public String getenv(String name) {
+    return System.getenv(name);
+  }
+}

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

@@ -32,6 +32,7 @@
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
@@ -41,6 +42,7 @@
 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.internal.EnvironmentProvider;
 import com.google.api.gax.rpc.mtls.MtlsProvider;
 import com.google.api.gax.rpc.mtls.MtlsProvider.MtlsEndpointUsagePolicy;
 import com.google.api.gax.rpc.testing.FakeChannel;
@@ -190,7 +192,7 @@ public String getTransportName() {
 
     @Override
     public boolean needsCredentials() {
-      return credentials == null;
+      return credentials == null && !headers.containsKey("x-goog-api-key");
     }
 
     @Override
@@ -200,6 +202,19 @@ public TransportChannelProvider withCredentials(Credentials credentials) {
     }
   }
 
+  private static class FakeEnvironmentProvider implements EnvironmentProvider {
+    final String apiKey;
+
+    @Override
+    public String getenv(String name) {
+      return apiKey;
+    }
+
+    FakeEnvironmentProvider(String apiKey) {
+      this.apiKey = apiKey;
+    }
+  }
+
   @Test
   public void testNoAutoCloseContextNeedsNoExecutor() throws Exception {
     runTest(false, false, false, false);
@@ -769,4 +784,56 @@ public void testExecutorSettings() throws Exception {
     transportChannel = (FakeTransportChannel) context.getTransportChannel();
     assertThat(transportChannel.getExecutor()).isSameInstanceAs(executorProvider.getExecutor());
   }
+
+  @Test
+  public void testRetrieveApiKeyFromStubSettings() throws IOException {
+    StubSettings settings = new FakeStubSettings.Builder().setApiKey("stub-setting-key").build();
+    EnvironmentProvider environmentProvider = new FakeEnvironmentProvider("env-key");
+    assertEquals("stub-setting-key", ClientContext.retrieveApiKey(settings, environmentProvider));
+  }
+
+  @Test
+  public void testRetrieveApiKeyFromEnvironmentProvider() throws IOException {
+    StubSettings settings = new FakeStubSettings.Builder().build();
+    EnvironmentProvider environmentProvider = new FakeEnvironmentProvider("env-key");
+    assertEquals("env-key", ClientContext.retrieveApiKey(settings, environmentProvider));
+  }
+
+  @Test
+  public void testRetrieveApiKeyNoApiKey() throws IOException {
+    StubSettings settings = new FakeStubSettings.Builder().build();
+    EnvironmentProvider environmentProvider = new FakeEnvironmentProvider(null);
+    assertNull(ClientContext.retrieveApiKey(settings, environmentProvider));
+  }
+
+  @Test
+  public void testApiKey() throws IOException {
+    FakeStubSettings.Builder builder = new FakeStubSettings.Builder();
+
+    FakeTransportChannel transportChannel = FakeTransportChannel.create(new FakeChannel());
+    FakeTransportProvider transportProvider =
+        new FakeTransportProvider(transportChannel, null, true, null, null);
+    builder.setTransportChannelProvider(transportProvider);
+
+    HeaderProvider headerProvider = Mockito.mock(HeaderProvider.class);
+    Mockito.when(headerProvider.getHeaders()).thenReturn(ImmutableMap.<String, String>of());
+    builder.setHeaderProvider(headerProvider);
+
+    builder.setCredentialsProvider(
+        FixedCredentialsProvider.create(Mockito.mock(Credentials.class)));
+
+    // Set API key.
+    builder.setApiKey("key");
+
+    ClientContext context = ClientContext.create(builder.build());
+
+    // Since API key is provided, credentials shouldn't be loaded.
+    assertNull(context.getCredentials());
+
+    // Check API key is in the transport channel's header.
+    List<BackgroundResource> resources = context.getBackgroundResources();
+    FakeTransportChannel fakeTransportChannel = (FakeTransportChannel) resources.get(0);
+    assertTrue(fakeTransportChannel.getHeaders().containsKey("x-goog-api-key"));
+    assertEquals("key", fakeTransportChannel.getHeaders().get("x-goog-api-key"));
+  }
 }