feat: allow using existing OAuth token for JDBC connection (#37)

* feat: allow using existing OAuth token for JDBC connection

Allow the user to specify an existing OAuth token to use for a JDBC connection,
instead of requiring the user to specify a credentials file or using the default
credentials of the environment.

Fixes #29

* tests: fix test cases

* fix: remove unused method

* fix: use default credentials when no file found

* fix: fix faulty merge

Author: olavloite

src/main/java/com/google/cloud/spanner/jdbc/ConnectionOptions.java

@@ -17,6 +17,7 @@
 package com.google.cloud.spanner.jdbc;
 
 import com.google.auth.Credentials;
+import com.google.auth.oauth2.AccessToken;
 import com.google.auth.oauth2.GoogleCredentials;
 import com.google.auth.oauth2.ServiceAccountCredentials;
 import com.google.cloud.NoCredentials;
@@ -140,6 +141,7 @@ public String[] getValidValues() {
   static final boolean DEFAULT_READONLY = false;
   static final boolean DEFAULT_RETRY_ABORTS_INTERNALLY = true;
   private static final String DEFAULT_CREDENTIALS = null;
+  private static final String DEFAULT_OAUTH_TOKEN = null;
   private static final String DEFAULT_NUM_CHANNELS = null;
   private static final String DEFAULT_USER_AGENT = null;
 
@@ -156,6 +158,10 @@ public String[] getValidValues() {
   public static final String RETRY_ABORTS_INTERNALLY_PROPERTY_NAME = "retryAbortsInternally";
   /** Name of the 'credentials' connection property. */
   public static final String CREDENTIALS_PROPERTY_NAME = "credentials";
+  /**
+   * OAuth token to use for authentication. Cannot be used in combination with a credentials file.
+   */
+  public static final String OAUTH_TOKEN_PROPERTY_NAME = "oauthToken";
   /** Name of the 'numChannels' connection property. */
   public static final String NUM_CHANNELS_PROPERTY_NAME = "numChannels";
   /** Custom user agent string is only for other Google libraries. */
@@ -173,6 +179,7 @@ public String[] getValidValues() {
                   ConnectionProperty.createBooleanProperty(
                       RETRY_ABORTS_INTERNALLY_PROPERTY_NAME, "", DEFAULT_RETRY_ABORTS_INTERNALLY),
                   ConnectionProperty.createStringProperty(CREDENTIALS_PROPERTY_NAME, ""),
+                  ConnectionProperty.createStringProperty(OAUTH_TOKEN_PROPERTY_NAME, ""),
                   ConnectionProperty.createStringProperty(NUM_CHANNELS_PROPERTY_NAME, ""),
                   ConnectionProperty.createBooleanProperty(
                       USE_PLAIN_TEXT_PROPERTY_NAME, "", DEFAULT_USE_PLAIN_TEXT),
@@ -223,6 +230,7 @@ public static void closeSpanner() {
   public static class Builder {
     private String uri;
     private String credentialsUrl;
+    private String oauthToken;
     private Credentials credentials;
     private List<StatementExecutionInterceptor> statementExecutionInterceptors =
         Collections.emptyList();
@@ -308,6 +316,22 @@ public Builder setCredentialsUrl(String credentialsUrl) {
       return this;
     }
 
+    /**
+     * Sets the OAuth token to use with this connection. The token must be a valid token with access
+     * to the resources (project/instance/database) that the connection will be accessing. This
+     * authentication method cannot be used in combination with a credentials file. If both an OAuth
+     * token and a credentials file is specified, the {@link #build()} method will throw an
+     * exception.
+     *
+     * @param oauthToken A valid OAuth token for the Google Cloud project that is used by this
+     *     connection.
+     * @return this builder
+     */
+    public Builder setOAuthToken(String oauthToken) {
+      this.oauthToken = oauthToken;
+      return this;
+    }
+
     @VisibleForTesting
     Builder setStatementExecutionInterceptors(List<StatementExecutionInterceptor> interceptors) {
       this.statementExecutionInterceptors = interceptors;
@@ -339,6 +363,7 @@ public static Builder newBuilder() {
 
   private final String uri;
   private final String credentialsUrl;
+  private final String oauthToken;
 
   private final boolean usePlainText;
   private final String host;
@@ -363,6 +388,13 @@ private ConnectionOptions(Builder builder) {
     this.uri = builder.uri;
     this.credentialsUrl =
         builder.credentialsUrl != null ? builder.credentialsUrl : parseCredentials(builder.uri);
+    this.oauthToken =
+        builder.oauthToken != null ? builder.oauthToken : parseOAuthToken(builder.uri);
+    // Check that not both credentials and an OAuth token have been specified.
+    Preconditions.checkArgument(
+        (builder.credentials == null && this.credentialsUrl == null) || this.oauthToken == null,
+        "Cannot specify both credentials and an OAuth token.");
+
     this.usePlainText = parseUsePlainText(this.uri);
     this.userAgent = parseUserAgent(this.uri);
 
@@ -376,8 +408,13 @@ 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 && this.credentialsUrl == null && this.usePlainText) {
+    if (builder.credentials == null
+        && this.credentialsUrl == null
+        && this.oauthToken == null
+        && this.usePlainText) {
       this.credentials = NoCredentials.getInstance();
+    } else if (this.oauthToken != null) {
+      this.credentials = new GoogleCredentials(new AccessToken(oauthToken, null));
     } else {
       this.credentials =
           builder.credentials == null
@@ -446,6 +483,12 @@ static String parseCredentials(String uri) {
     return value != null ? value : DEFAULT_CREDENTIALS;
   }
 
+  @VisibleForTesting
+  static String parseOAuthToken(String uri) {
+    String value = parseUriProperty(uri, OAUTH_TOKEN_PROPERTY_NAME);
+    return value != null ? value : DEFAULT_OAUTH_TOKEN;
+  }
+
   @VisibleForTesting
   static String parseNumChannels(String uri) {
     String value = parseUriProperty(uri, NUM_CHANNELS_PROPERTY_NAME);

src/main/java/com/google/cloud/spanner/jdbc/JdbcDriver.java

@@ -74,6 +74,12 @@
  *   <li>credentials (String): URL for the credentials file to use for the connection. If you do not
  *       specify any credentials at all, the default credentials of the environment as returned by
  *       {@link GoogleCredentials#getApplicationDefault()} will be used.
+ *   <li>oauthtoken (String): A valid OAuth2 token to use for the JDBC connection. The token must
+ *       have been obtained with one or both of the scopes
+ *       'https://www.googleapis.com/auth/spanner.admin' and/or
+ *       'https://www.googleapis.com/auth/spanner.data'. If you specify both a credentials file and
+ *       an OAuth token, the JDBC driver will throw an exception when you try to obtain a
+ *       connection.
  *   <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 false.
  *   <li>retryAbortsInternally (boolean): Sets the initial retryAbortsInternally mode for the

src/test/java/com/google/cloud/spanner/jdbc/ConnectionOptionsTest.java

@@ -19,6 +19,7 @@
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.fail;
 
+import com.google.auth.oauth2.GoogleCredentials;
 import com.google.auth.oauth2.ServiceAccountCredentials;
 import com.google.cloud.spanner.SpannerOptions;
 import java.util.Arrays;
@@ -323,4 +324,66 @@ public void testParsePropertiesSpecifiedMultipleTimes() {
             "cloudspanner:/projects/test-project-123/instances/test-instance/databases/test-database"
                 + ";autocommit=false;readonly=false;autocommit=true");
   }
+
+  @Test
+  public void testParseOAuthToken() {
+    assertThat(
+            ConnectionOptions.parseUriProperty(
+                "cloudspanner:/projects/test-project-123/instances/test-instance/databases/test-database"
+                    + "?oauthtoken=RsT5OjbzRn430zqMLgV3Ia",
+                "OAuthToken"))
+        .isEqualTo("RsT5OjbzRn430zqMLgV3Ia");
+    // Try to use both credentials and an OAuth token. That should fail.
+    ConnectionOptions.Builder builder =
+        ConnectionOptions.newBuilder()
+            .setUri(
+                "cloudspanner:/projects/test-project-123/instances/test-instance/databases/test-database"
+                    + "?OAuthToken=RsT5OjbzRn430zqMLgV3Ia;credentials=/path/to/credentials.json");
+    try {
+      builder.build();
+      fail("missing expected exception");
+    } catch (IllegalArgumentException e) {
+      assertThat(e.getMessage()).contains("Cannot specify both credentials and an OAuth token");
+    }
+
+    // Now try to use only an OAuth token.
+    builder =
+        ConnectionOptions.newBuilder()
+            .setUri(
+                "cloudspanner:/projects/test-project-123/instances/test-instance/databases/test-database"
+                    + "?OAuthToken=RsT5OjbzRn430zqMLgV3Ia");
+    ConnectionOptions options = builder.build();
+    assertThat(options.getCredentials()).isInstanceOf(GoogleCredentials.class);
+    GoogleCredentials credentials = (GoogleCredentials) options.getCredentials();
+    assertThat(credentials.getAccessToken().getTokenValue()).isEqualTo("RsT5OjbzRn430zqMLgV3Ia");
+  }
+
+  @Test
+  public void testSetOAuthToken() {
+    ConnectionOptions options =
+        ConnectionOptions.newBuilder()
+            .setUri(
+                "cloudspanner:/projects/test-project-123/instances/test-instance/databases/test-database")
+            .setOAuthToken("RsT5OjbzRn430zqMLgV3Ia")
+            .build();
+    assertThat(options.getCredentials()).isInstanceOf(GoogleCredentials.class);
+    GoogleCredentials credentials = (GoogleCredentials) options.getCredentials();
+    assertThat(credentials.getAccessToken()).isNotNull();
+    assertThat(credentials.getAccessToken().getTokenValue()).isEqualTo("RsT5OjbzRn430zqMLgV3Ia");
+  }
+
+  @Test
+  public void testSetOAuthTokenAndCredentials() {
+    try {
+      ConnectionOptions.newBuilder()
+          .setUri(
+              "cloudspanner:/projects/test-project-123/instances/test-instance/databases/test-database")
+          .setOAuthToken("RsT5OjbzRn430zqMLgV3Ia")
+          .setCredentialsUrl(FILE_TEST_PATH)
+          .build();
+      fail("missing expected exception");
+    } catch (IllegalArgumentException e) {
+      assertThat(e.getMessage()).contains("Cannot specify both credentials and an OAuth token");
+    }
+  }
 }

src/test/java/com/google/cloud/spanner/jdbc/JdbcDriverTest.java

@@ -17,6 +17,7 @@
 package com.google.cloud.spanner.jdbc;
 
 import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.fail;
 
 import com.google.cloud.spanner.MockSpannerServiceImpl;
 import io.grpc.Server;
@@ -87,4 +88,17 @@ public void testInvalidConnect() throws SQLException {
       assertThat(connection.isClosed()).isFalse();
     }
   }
+
+  @Test
+  public void testConnectWithCredentialsAndOAuthToken() throws SQLException {
+    try (Connection connection =
+        DriverManager.getConnection(
+            String.format(
+                "jdbc:cloudspanner://localhost:%d/projects/test-project/instances/static-test-instance/databases/test-database;usePlainText=true;credentials=%s;OAuthToken=%s",
+                server.getPort(), TEST_KEY_PATH, "some-token"))) {
+      fail("missing expected exception");
+    } catch (SQLException e) {
+      assertThat(e.getMessage()).contains("Cannot specify both credentials and an OAuth token");
+    }
+  }
 }

src/test/java/com/google/cloud/spanner/jdbc/it/ITJdbcConnectTest.java

@@ -16,14 +16,16 @@
 
 package com.google.cloud.spanner.jdbc.it;
 
-import static org.hamcrest.CoreMatchers.equalTo;
-import static org.hamcrest.CoreMatchers.is;
-import static org.hamcrest.MatcherAssert.assertThat;
+import static com.google.common.truth.Truth.assertThat;
 
+import com.google.auth.oauth2.AccessToken;
+import com.google.auth.oauth2.GoogleCredentials;
 import com.google.cloud.spanner.IntegrationTest;
+import com.google.cloud.spanner.SpannerOptions;
 import com.google.cloud.spanner.jdbc.CloudSpannerJdbcConnection;
 import com.google.cloud.spanner.jdbc.ITAbstractJdbcTest;
 import com.google.cloud.spanner.jdbc.JdbcDataSource;
+import java.io.FileInputStream;
 import java.sql.Connection;
 import java.sql.DriverManager;
 import java.sql.ResultSet;
@@ -57,30 +59,30 @@ private String createBaseUrl() {
   }
 
   private void testDefaultConnection(Connection connection) throws SQLException {
-    assertThat(connection.isWrapperFor(CloudSpannerJdbcConnection.class), is(true));
+    assertThat(connection.isWrapperFor(CloudSpannerJdbcConnection.class)).isTrue();
     CloudSpannerJdbcConnection cs = connection.unwrap(CloudSpannerJdbcConnection.class);
-    assertThat(cs.getAutoCommit(), is(true));
-    assertThat(cs.isReadOnly(), is(false));
+    assertThat(cs.getAutoCommit()).isTrue();
+    assertThat(cs.isReadOnly()).isFalse();
     try (ResultSet rs = connection.createStatement().executeQuery("SELECT 1")) {
-      assertThat(rs.next(), is(true));
-      assertThat(rs.getInt(1), is(equalTo(1)));
+      assertThat(rs.next()).isTrue();
+      assertThat(rs.getInt(1)).isEqualTo(1);
     }
     cs.setAutoCommit(false);
-    assertThat(cs.isRetryAbortsInternally(), is(true));
+    assertThat(cs.isRetryAbortsInternally()).isTrue();
   }
 
   private void testNonDefaultConnection(Connection connection) throws SQLException {
-    assertThat(connection.isWrapperFor(CloudSpannerJdbcConnection.class), is(true));
+    assertThat(connection.isWrapperFor(CloudSpannerJdbcConnection.class)).isTrue();
     CloudSpannerJdbcConnection cs = connection.unwrap(CloudSpannerJdbcConnection.class);
-    assertThat(cs.getAutoCommit(), is(false));
-    assertThat(cs.isReadOnly(), is(true));
+    assertThat(cs.getAutoCommit()).isFalse();
+    assertThat(cs.isReadOnly()).isTrue();
     try (ResultSet rs = connection.createStatement().executeQuery("SELECT 1")) {
-      assertThat(rs.next(), is(true));
-      assertThat(rs.getInt(1), is(equalTo(1)));
+      assertThat(rs.next()).isTrue();
+      assertThat(rs.getInt(1)).isEqualTo(1);
     }
     connection.commit();
     cs.setReadOnly(false);
-    assertThat(cs.isRetryAbortsInternally(), is(false));
+    assertThat(cs.isRetryAbortsInternally()).isFalse();
   }
 
   @Test
@@ -194,4 +196,21 @@ public void testConnectWithDataSourceWithConflictingValues() throws SQLException
       testNonDefaultConnection(connection);
     }
   }
+
+  @Test
+  public void testConnectWithOAuthToken() throws Exception {
+    GoogleCredentials credentials;
+    if (hasValidKeyFile()) {
+      credentials = GoogleCredentials.fromStream(new FileInputStream(getKeyFile()));
+    } else {
+      credentials = GoogleCredentials.getApplicationDefault();
+    }
+    credentials = credentials.createScoped(SpannerOptions.getDefaultInstance().getScopes());
+    AccessToken token = credentials.refreshAccessToken();
+    String urlWithOAuth = createBaseUrl() + "?OAuthToken=" + token.getTokenValue();
+    try (Connection connectionWithOAuth = DriverManager.getConnection(urlWithOAuth)) {
+      // Try to do a query using the connection created with an OAuth token.
+      testDefaultConnection(connectionWithOAuth);
+    }
+  }
 }