Skip to content

Commit

Permalink
feat: support CredentialsProvider in Connection API (#1869)
Browse files Browse the repository at this point in the history 
* 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
  • Loading branch information
@olavloite
olavloite authored May 5, 2022
1 parent f74c77d commit f1d2d3e
Show file tree
Hide file tree
Showing 6 changed files with 332 additions and 88 deletions.
5 changes: 5 additions & 0 deletions google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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())
Expand Down
97 changes: 80 additions & 17 deletions ...le-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionOptions.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;

/**
Expand Down Expand Up @@ -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.
*/
Expand Down Expand Up @@ -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."),
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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;

Expand Down Expand Up @@ -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();
Expand All @@ -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) {
Expand Down Expand Up @@ -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;
}
Expand Down Expand Up @@ -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;
Expand Down
28 changes: 18 additions & 10 deletions google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SpannerPool.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;

/**
Expand Down Expand Up @@ -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) {
Expand All @@ -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);
Expand Down
2 changes: 1 addition & 1 deletion ...oud-spanner/src/test/java/com/google/cloud/spanner/connection/AbstractMockServerTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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())
Expand Down
Loading

0 comments on commit f1d2d3e

Please sign in to comment.