feat: support CredentialsProvider in Connection API (#1869)

olavloite · web-flow · commit f1d2d3ef1dbd · 2022-05-05T10:01:50.000+05:30
* feat: support CredentialsProvider in Connection API

Adds suppport for setting a CredentialsProvider instead of a
credentialsUrl in a connection string. The CredentialsProvider reference
must be a class name to a public class with a public no-arg constructor.
This option is available in the Connection API, which means that any
client that uses that API can directly benefit from it (this effectively
means the JDBC driver).

Fixes b/231174409
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java
@@ -259,6 +259,11 @@ static UnitOfWorkType of(TransactionMode transactionMode) {
     setDefaultTransactionOptions();
   }
 
+  @VisibleForTesting
+  Spanner getSpanner() {
+    return this.spanner;
+  }
+
   private DdlClient createDdlClient() {
     return DdlClient.newBuilder()
         .setDatabaseAdminClient(spanner.getDatabaseAdminClient())
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionOptions.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionOptions.java
@@ -17,6 +17,7 @@
 package com.google.cloud.spanner.connection;
 
 import com.google.api.core.InternalApi;
+import com.google.api.gax.core.CredentialsProvider;
 import com.google.api.gax.rpc.TransportChannelProvider;
 import com.google.auth.Credentials;
 import com.google.auth.oauth2.AccessToken;
@@ -36,15 +37,20 @@
 import com.google.common.base.Preconditions;
 import com.google.common.collect.Sets;
 import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions;
+import java.io.IOException;
+import java.lang.reflect.Constructor;
+import java.lang.reflect.InvocationTargetException;
 import java.net.URL;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Objects;
 import java.util.Set;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
+import java.util.stream.Stream;
 import javax.annotation.Nullable;
 
 /**
@@ -182,6 +188,8 @@ public String[] getValidValues() {
   public static final String CREDENTIALS_PROPERTY_NAME = "credentials";
   /** Name of the 'encodedCredentials' connection property. */
   public static final String ENCODED_CREDENTIALS_PROPERTY_NAME = "encodedCredentials";
+  /** Name of the 'credentialsProvider' connection property. */
+  public static final String CREDENTIALS_PROVIDER_PROPERTY_NAME = "credentialsProvider";
   /**
    * OAuth token to use for authentication. Cannot be used in combination with a credentials file.
    */
@@ -231,6 +239,9 @@ public String[] getValidValues() {
                   ConnectionProperty.createStringProperty(
                       ENCODED_CREDENTIALS_PROPERTY_NAME,
                       "Base64-encoded credentials to use for this connection. If neither this property or a credentials location are set, the connection will use the default Google Cloud credentials for the runtime environment."),
+                  ConnectionProperty.createStringProperty(
+                      CREDENTIALS_PROVIDER_PROPERTY_NAME,
+                      "The class name of the com.google.api.gax.core.CredentialsProvider implementation that should be used to obtain credentials for connections."),
                   ConnectionProperty.createStringProperty(
                       OAUTH_TOKEN_PROPERTY_NAME,
                       "A valid pre-existing OAuth token to use for authentication for this connection. Setting this property will take precedence over any value set for a credentials file."),
@@ -386,6 +397,12 @@ private boolean isValidUri(String uri) {
      *   <li>encodedCredentials (String): A Base64 encoded string containing the Google credentials
      *       to use. You should only set either this property or the `credentials` (file location)
      *       property, but not both at the same time.
+     *   <li>credentialsProvider (String): Class name of the {@link
+     *       com.google.api.gax.core.CredentialsProvider} that should be used to get credentials for
+     *       a connection that is created by this {@link ConnectionOptions}. The credentials will be
+     *       retrieved from the {@link com.google.api.gax.core.CredentialsProvider} when a new
+     *       connection is created. A connection will use the credentials that were obtained at
+     *       creation during its lifetime.
      *   <li>autocommit (boolean): Sets the initial autocommit mode for the connection. Default is
      *       true.
      *   <li>readonly (boolean): Sets the initial readonly mode for the connection. Default is
@@ -501,6 +518,7 @@ public static Builder newBuilder() {
   private final String warnings;
   private final String credentialsUrl;
   private final String encodedCredentials;
+  private final CredentialsProvider credentialsProvider;
   private final String oauthToken;
   private final Credentials fixedCredentials;
 
@@ -537,22 +555,22 @@ private ConnectionOptions(Builder builder) {
     this.credentialsUrl =
         builder.credentialsUrl != null ? builder.credentialsUrl : parseCredentials(builder.uri);
     this.encodedCredentials = parseEncodedCredentials(builder.uri);
-    // Check that not both a credentials location and encoded credentials have been specified in the
-    // connection URI.
-    Preconditions.checkArgument(
-        this.credentialsUrl == null || this.encodedCredentials == null,
-        "Cannot specify both a credentials URL and encoded credentials. Only set one of the properties.");
-
+    this.credentialsProvider = parseCredentialsProvider(builder.uri);
     this.oauthToken =
         builder.oauthToken != null ? builder.oauthToken : parseOAuthToken(builder.uri);
-    this.fixedCredentials = builder.credentials;
-    // Check that not both credentials and an OAuth token have been specified.
+    // Check that at most one of credentials location, encoded credentials, credentials provider and
+    // OUAuth token has been specified in the connection URI.
     Preconditions.checkArgument(
-        (builder.credentials == null
-                && this.credentialsUrl == null
-                && this.encodedCredentials == null)
-            || this.oauthToken == null,
-        "Cannot specify both credentials and an OAuth token.");
+        Stream.of(
+                    this.credentialsUrl,
+                    this.encodedCredentials,
+                    this.credentialsProvider,
+                    this.oauthToken)
+                .filter(Objects::nonNull)
+                .count()
+            <= 1,
+        "Specify only one of credentialsUrl, encodedCredentials, credentialsProvider and OAuth token");
+    this.fixedCredentials = builder.credentials;
 
     this.userAgent = parseUserAgent(this.uri);
     QueryOptions.Builder queryOptionsBuilder = QueryOptions.newBuilder();
@@ -570,14 +588,24 @@ private ConnectionOptions(Builder builder) {
     // Using credentials on a plain text connection is not allowed, so if the user has not specified
     // any credentials and is using a plain text connection, we should not try to get the
     // credentials from the environment, but default to NoCredentials.
-    if (builder.credentials == null
+    if (this.fixedCredentials == null
         && this.credentialsUrl == null
         && this.encodedCredentials == null
+        && this.credentialsProvider == null
         && this.oauthToken == null
         && this.usePlainText) {
       this.credentials = NoCredentials.getInstance();
     } else if (this.oauthToken != null) {
       this.credentials = new GoogleCredentials(new AccessToken(oauthToken, null));
+    } else if (this.credentialsProvider != null) {
+      try {
+        this.credentials = this.credentialsProvider.getCredentials();
+      } catch (IOException exception) {
+        throw SpannerExceptionFactory.newSpannerException(
+            ErrorCode.INVALID_ARGUMENT,
+            "Failed to get credentials from CredentialsProvider: " + exception.getMessage(),
+            exception);
+      }
     } else if (this.fixedCredentials != null) {
       this.credentials = fixedCredentials;
     } else if (this.encodedCredentials != null) {
@@ -691,18 +719,49 @@ static boolean parseRetryAbortsInternally(String uri) {
   }
 
   @VisibleForTesting
-  static String parseCredentials(String uri) {
+  static @Nullable String parseCredentials(String uri) {
     String value = parseUriProperty(uri, CREDENTIALS_PROPERTY_NAME);
     return value != null ? value : DEFAULT_CREDENTIALS;
   }
 
   @VisibleForTesting
-  static String parseEncodedCredentials(String uri) {
+  static @Nullable String parseEncodedCredentials(String uri) {
     return parseUriProperty(uri, ENCODED_CREDENTIALS_PROPERTY_NAME);
   }
 
   @VisibleForTesting
-  static String parseOAuthToken(String uri) {
+  static @Nullable CredentialsProvider parseCredentialsProvider(String uri) {
+    String name = parseUriProperty(uri, CREDENTIALS_PROVIDER_PROPERTY_NAME);
+    if (name != null) {
+      try {
+        Class<? extends CredentialsProvider> clazz =
+            (Class<? extends CredentialsProvider>) Class.forName(name);
+        Constructor<? extends CredentialsProvider> constructor = clazz.getDeclaredConstructor();
+        return constructor.newInstance();
+      } catch (ClassNotFoundException classNotFoundException) {
+        throw SpannerExceptionFactory.newSpannerException(
+            ErrorCode.INVALID_ARGUMENT,
+            "Unknown or invalid CredentialsProvider class name: " + name,
+            classNotFoundException);
+      } catch (NoSuchMethodException noSuchMethodException) {
+        throw SpannerExceptionFactory.newSpannerException(
+            ErrorCode.INVALID_ARGUMENT,
+            "Credentials provider " + name + " does not have a public no-arg constructor.",
+            noSuchMethodException);
+      } catch (InvocationTargetException
+          | InstantiationException
+          | IllegalAccessException exception) {
+        throw SpannerExceptionFactory.newSpannerException(
+            ErrorCode.INVALID_ARGUMENT,
+            "Failed to create an instance of " + name + ": " + exception.getMessage(),
+            exception);
+      }
+    }
+    return null;
+  }
+
+  @VisibleForTesting
+  static @Nullable String parseOAuthToken(String uri) {
     String value = parseUriProperty(uri, OAUTH_TOKEN_PROPERTY_NAME);
     return value != null ? value : DEFAULT_OAUTH_TOKEN;
   }
@@ -849,6 +908,10 @@ Credentials getFixedCredentials() {
     return this.fixedCredentials;
   }
 
+  CredentialsProvider getCredentialsProvider() {
+    return this.credentialsProvider;
+  }
+
   /** The {@link SessionPoolOptions} of this {@link ConnectionOptions}. */
   public SessionPoolOptions getSessionPoolOptions() {
     return sessionPoolOptions;
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SpannerPool.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SpannerPool.java
@@ -27,12 +27,10 @@
 import com.google.common.base.Function;
 import com.google.common.base.MoreObjects;
 import com.google.common.base.Preconditions;
-import com.google.common.base.Predicates;
 import com.google.common.base.Ticker;
-import com.google.common.collect.Iterables;
 import io.grpc.ManagedChannelBuilder;
+import java.io.IOException;
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -43,6 +41,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.logging.Level;
 import java.util.logging.Logger;
+import java.util.stream.Stream;
 import javax.annotation.concurrent.GuardedBy;
 
 /**
@@ -120,15 +119,17 @@ static class CredentialsKey {
     static final Object DEFAULT_CREDENTIALS_KEY = new Object();
     final Object key;
 
-    static CredentialsKey create(ConnectionOptions options) {
+    static CredentialsKey create(ConnectionOptions options) throws IOException {
       return new CredentialsKey(
-          Iterables.find(
-              Arrays.asList(
+          Stream.of(
                   options.getOAuthToken(),
+                  options.getCredentialsProvider() == null ? null : options.getCredentials(),
                   options.getFixedCredentials(),
                   options.getCredentialsUrl(),
-                  DEFAULT_CREDENTIALS_KEY),
-              Predicates.notNull()));
+                  DEFAULT_CREDENTIALS_KEY)
+              .filter(Objects::nonNull)
+              .findFirst()
+              .get());
     }
 
     private CredentialsKey(Object key) {
@@ -155,10 +156,17 @@ static class SpannerPoolKey {
 
     @VisibleForTesting
     static SpannerPoolKey of(ConnectionOptions options) {
-      return new SpannerPoolKey(options);
+      try {
+        return new SpannerPoolKey(options);
+      } catch (IOException ioException) {
+        throw SpannerExceptionFactory.newSpannerException(
+            ErrorCode.INVALID_ARGUMENT,
+            "Failed to get credentials: " + ioException.getMessage(),
+            ioException);
+      }
     }
 
-    private SpannerPoolKey(ConnectionOptions options) {
+    private SpannerPoolKey(ConnectionOptions options) throws IOException {
       this.host = options.getHost();
       this.projectId = options.getProjectId();
       this.credentialsKey = CredentialsKey.create(options);
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AbstractMockServerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AbstractMockServerTest.java
@@ -173,7 +173,7 @@ public static void stopServer() {
     try {
       SpannerPool.INSTANCE.checkAndCloseSpanners(
           CheckAndCloseSpannersMode.ERROR,
-          new ForceCloseSpannerFunction(100L, TimeUnit.MILLISECONDS));
+          new ForceCloseSpannerFunction(500L, TimeUnit.MILLISECONDS));
     } finally {
       Logger.getLogger(AbstractFuture.class.getName()).setUseParentHandlers(futureParentHandlers);
       Logger.getLogger(LogExceptionRunnable.class.getName())
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionOptionsTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionOptionsTest.java
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/CredentialsProviderTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/CredentialsProviderTest.java
