From 70bc4045e0786894914ce5317e828c803a1fae73 Mon Sep 17 00:00:00 2001 From: Tatiane Tosta <91583351+ttosta-google@users.noreply.github.com> Date: Fri, 20 Oct 2023 13:12:30 -0700 Subject: [PATCH] feat: add support for service account impersonation (#245) --- .envrc.example | 3 + .github/workflows/ci.yaml | 2 + README.md | 70 +++++++++++++++++++ alloydb-jdbc-connector/pom.xml | 11 ++- .../alloydb/AlloyDBAdminClientFactory.java | 5 +- .../cloud/alloydb/ConnectionConfig.java | 65 ++++++++++++++++- .../com/google/cloud/alloydb/Connector.java | 6 +- .../cloud/alloydb/ConnectorRegistry.java | 42 ++++++----- .../alloydb/CredentialsProviderFactory.java | 49 +++++++++++++ .../DefaultConnectionInfoRepository.java | 8 ++- .../google/cloud/alloydb/SocketFactory.java | 4 +- ...AccountImpersonationDataSourceFactory.java | 43 ++++++++++++ .../cloud/alloydb/ConnectionConfigTest.java | 27 ++++++- .../google/cloud/alloydb/ITConnectorTest.java | 18 ++--- ...ITDefaultConnectionInfoRepositoryTest.java | 6 +- .../ITServiceAccountImpersonationTest.java | 70 +++++++++++++++++++ 16 files changed, 383 insertions(+), 46 deletions(-) create mode 100644 alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/CredentialsProviderFactory.java create mode 100644 alloydb-jdbc-connector/src/test/java/com/google/cloud/alloydb/AlloyDbJdbcServiceAccountImpersonationDataSourceFactory.java create mode 100644 alloydb-jdbc-connector/src/test/java/com/google/cloud/alloydb/ITServiceAccountImpersonationTest.java diff --git a/.envrc.example b/.envrc.example index 3b7a4811..ce7de0ab 100644 --- a/.envrc.example +++ b/.envrc.example @@ -2,3 +2,6 @@ export ALLOYDB_DB="some-db" export ALLOYDB_USER="some-user" export ALLOYDB_PASS="some-password" export ALLOYDB_INSTANCE_NAME="projects//locations//clusters//instances/" +export ALLOYDB_INSTANCE_IP="some-IP-address" +export ALLOYDB_IAM_USER="some-user@my-project.iam" +export ALLOYDB_IMPERSONATED_USER="some-impersonated-IAM-user" diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 30cc83c4..5026b6d4 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -248,6 +248,7 @@ jobs: ALLOYDB_CLUSTER_PASS:${{ secrets.GOOGLE_CLOUD_PROJECT }}/ALLOYDB_CLUSTER_PASS ALLOYDB_IAM_USER:${{ secrets.GOOGLE_CLOUD_PROJECT }}/ALLOYDB_JAVA_IAM_USER ALLOYDB_INSTANCE_IP:${{ secrets.GOOGLE_CLOUD_PROJECT }}/ALLOYDB_INSTANCE_IP + ALLOYDB_IMPERSONATED_USER:${{ secrets.GOOGLE_CLOUD_PROJECT }}/IMPERSONATED_USER - name: Run tests env: @@ -257,6 +258,7 @@ jobs: ALLOYDB_PASS: '${{ steps.secrets.outputs.ALLOYDB_CLUSTER_PASS }}' ALLOYDB_INSTANCE_NAME: '${{ steps.secrets.outputs.ALLOYDB_INSTANCE_NAME }}' ALLOYDB_INSTANCE_IP: '${{ steps.secrets.outputs.ALLOYDB_INSTANCE_IP }}' + ALLOYDB_IMPERSONATED_USER: '${{ steps.secrets.outputs.ALLOYDB_IMPERSONATED_USER }}' JOB_TYPE: integration run: .kokoro/build.sh shell: bash diff --git a/README.md b/README.md index b3d39534..9dcc0244 100644 --- a/README.md +++ b/README.md @@ -157,6 +157,76 @@ performance from a connection pool. [e2e]: https://github.com/GoogleCloudPlatform/alloydb-java-connector/blob/main/alloydb-jdbc-connector/src/test/java/com/google/cloud/alloydb/ITSocketFactoryTest.java [pool-sizing]: https://github.com/brettwooldridge/HikariCP/wiki/About-Pool-Sizing +### Service Account Impersonation + +The Java Connector supports service account impersonation with the +`alloydbTargetPrincipal` JDBC connection property. When enabled, +all API requests are made impersonating the supplied service account. +The IAM principal must have the IAM role "Service Account Token Creator" +(i.e., `roles/iam.serviceAccounts.serviceAccountTokenCreator`) on the +service account provided in the `alloydbTargetPrincipal` property. + +#### Example + +``` java +import com.zaxxer.hikari.HikariConfig; +import com.zaxxer.hikari.HikariDataSource; + +public class ExampleApplication { + + private HikariDataSource dataSource; + + HikariDataSource getDataSource() { + HikariConfig config = new HikariConfig(); + + // There is no need to set a host on the JDBC URL + // since the Connector will resolve the correct IP address. + config.setJdbcUrl("jdbc:postgresql:///postgres"); + config.setUsername(System.getenv("ALLOYDB_USER")); + config.setPassword(System.getenv("ALLOYDB_PASS")); + + // Tell the driver to use the AlloyDB Java Connector's SocketFactory + // when connecting to an instance/ + config.addDataSourceProperty("socketFactory", + "com.google.cloud.alloydb.SocketFactory"); + // Tell the Java Connector which instance to connect to. + config.addDataSourceProperty("alloydbInstanceName", + System.getenv("ALLOYDB_INSTANCE_NAME")); + config.addDataSourceProperty("alloydbTargetPrincipal", + System.getenv("ALLOYDB_IMPERSONATED_USER")); + + this.dataSource = new HikariDataSource(config); + } + + // Use DataSource as usual ... + +} +``` + +#### Delegated Service Account Impersonation + +In addition, the `alloydbDelegates` property controls impersonation delegation. +The value is a comma-separated list of service accounts containing chained +list of delegates required to grant the final access_token. If set, +the sequence of identities must have "Service Account Token Creator" capability +granted to the preceding identity. For example, if set to +`"serviceAccountB,serviceAccountC"`, the IAM principal must have the +Token Creator role on serviceAccountB. Then serviceAccountB must have the +Token Creator role on serviceAccountC. Finally, serviceAccountC must have the +Token Creator role on `alloydbTargetPrincipal`. If unset, the IAM principal +must have the Token Creator role on `alloydbTargetPrincipal`. + +```java +config.addDataSourceProperty("alloydbTargetPrincipal", + "TARGET_SERVICE_ACCOUNT"); +config.addDataSourceProperty("alloydbDelegates", + "SERVICE_ACCOUNT_1,SERVICE_ACCOUNT_2"); +``` + +In this example, the IAM principal impersonates SERVICE_ACCOUNT_1 which +impersonates SERVICE_ACCOUNT_2 which then impersonates the +TARGET_SERVICE_ACCOUNT. + ## Support policy ### Major version lifecycle diff --git a/alloydb-jdbc-connector/pom.xml b/alloydb-jdbc-connector/pom.xml index 3f8c43ba..9e6347b4 100644 --- a/alloydb-jdbc-connector/pom.xml +++ b/alloydb-jdbc-connector/pom.xml @@ -117,6 +117,11 @@ provided + + com.google.auth + google-auth-library-oauth2-http + + org.slf4j @@ -160,11 +165,5 @@ test - - com.google.auth - google-auth-library-oauth2-http - test - - diff --git a/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/AlloyDBAdminClientFactory.java b/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/AlloyDBAdminClientFactory.java index 22605318..e8a2a492 100644 --- a/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/AlloyDBAdminClientFactory.java +++ b/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/AlloyDBAdminClientFactory.java @@ -15,6 +15,7 @@ */ package com.google.cloud.alloydb; +import com.google.api.gax.core.FixedCredentialsProvider; import com.google.api.gax.rpc.FixedHeaderProvider; import com.google.cloud.alloydb.v1beta.AlloyDBAdminClient; import com.google.cloud.alloydb.v1beta.AlloyDBAdminSettings; @@ -30,7 +31,8 @@ class AlloyDBAdminClientFactory { private static final String DEFAULT_ENDPOINT = "alloydb.googleapis.com:443"; - static AlloyDBAdminClient create() throws IOException { + static AlloyDBAdminClient create(FixedCredentialsProvider credentialsProvider) + throws IOException { AlloyDBAdminSettings.Builder settingsBuilder = AlloyDBAdminSettings.newBuilder(); Map headers = @@ -42,6 +44,7 @@ static AlloyDBAdminClient create() throws IOException { settingsBuilder .setEndpoint(DEFAULT_ENDPOINT) .setHeaderProvider(FixedHeaderProvider.create(headers)) + .setCredentialsProvider(credentialsProvider) .build(); return AlloyDBAdminClient.create(alloyDBAdminSettings); diff --git a/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/ConnectionConfig.java b/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/ConnectionConfig.java index bad3c890..834d8a25 100644 --- a/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/ConnectionConfig.java +++ b/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/ConnectionConfig.java @@ -17,39 +17,98 @@ package com.google.cloud.alloydb; import com.google.cloud.alloydb.v1beta.InstanceName; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; import java.util.Properties; public class ConnectionConfig { public static final String ALLOYDB_INSTANCE_NAME = "alloydbInstanceName"; + public static final String ALLOYDB_TARGET_PRINCIPAL = "alloydbTargetPrincipal"; + public static final String ALLOYDB_DELEGATES = "alloydbDelegates"; + public static final String ALLOYDB_NAMED_CONNECTION = "alloydbNamedConnection"; + public static final String DEFAULT_NAMED_CONNECTION = "default"; private final InstanceName instanceName; + private final String namedConnection; + private final String targetPrincipal; + private final List delegates; /** Create a new ConnectionConfig from the well known JDBC Connection properties. */ public static ConnectionConfig fromConnectionProperties(Properties props) { final String instanceNameStr = props.getProperty(ALLOYDB_INSTANCE_NAME, ""); final InstanceName instanceName = InstanceName.parse(instanceNameStr); + final String namedConnection = props.getProperty(ConnectionConfig.ALLOYDB_NAMED_CONNECTION); + final String targetPrincipal = props.getProperty(ConnectionConfig.ALLOYDB_TARGET_PRINCIPAL); + final String delegatesStr = props.getProperty(ConnectionConfig.ALLOYDB_DELEGATES); + final List delegates; + if (delegatesStr != null && !delegatesStr.isEmpty()) { + delegates = Arrays.asList(delegatesStr.split(",")); + } else { + delegates = Collections.emptyList(); + } - return new ConnectionConfig(instanceName); + return new ConnectionConfig(instanceName, namedConnection, targetPrincipal, delegates); } - private ConnectionConfig(InstanceName instanceName) { + private ConnectionConfig( + InstanceName instanceName, + String namedConnection, + String targetPrincipal, + List delegates) { this.instanceName = instanceName; + this.namedConnection = namedConnection; + this.targetPrincipal = targetPrincipal; + this.delegates = delegates; } public InstanceName getInstanceName() { return instanceName; } + public String getNamedConnection() { + if (namedConnection == null || namedConnection.isEmpty()) { + return DEFAULT_NAMED_CONNECTION; + } + return namedConnection; + } + + public String getTargetPrincipal() { + return targetPrincipal; + } + + public List getDelegates() { + return delegates; + } + /** The builder for the ConnectionConfig. */ public static class Builder { private InstanceName instanceName; + private String namedConnection; + private String targetPrincipal; + private List delegates; public Builder withInstanceName(InstanceName instanceName) { this.instanceName = instanceName; return this; } + public Builder withNamedConnection(String namedConnection) { + this.namedConnection = namedConnection; + return this; + } + + public Builder withTargetPrincipal(String targetPrincipal) { + this.targetPrincipal = targetPrincipal; + return this; + } + + public Builder withDelegates(List delegates) { + this.delegates = delegates; + return this; + } + public ConnectionConfig build() { - return new ConnectionConfig(instanceName); + return new ConnectionConfig(instanceName, namedConnection, targetPrincipal, delegates); } } } diff --git a/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/Connector.java b/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/Connector.java index d5ffaa79..e9216351 100644 --- a/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/Connector.java +++ b/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/Connector.java @@ -137,16 +137,16 @@ private static KeyManager[] initializeKeyManager( return keyManagerFactory.getKeyManagers(); } - Socket connect(ConnectionConfig config) throws IOException { + Socket connect(InstanceName instanceName) throws IOException { ConnectionInfoCache connectionInfoCache = instances.computeIfAbsent( - config.getInstanceName(), + instanceName, k -> { DefaultRateLimiter rateLimiter = new DefaultRateLimiter(RATE_LIMIT_PER_SEC); return connectionInfoCacheFactory.create( this.executor, this.connectionInfoRepo, - config.getInstanceName(), + instanceName, this.clientConnectorKeyPair, new RefreshCalculator(), rateLimiter); diff --git a/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/ConnectorRegistry.java b/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/ConnectorRegistry.java index e3078cd2..f9462746 100644 --- a/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/ConnectorRegistry.java +++ b/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/ConnectorRegistry.java @@ -15,6 +15,7 @@ */ package com.google.cloud.alloydb; +import com.google.api.gax.core.FixedCredentialsProvider; import com.google.cloud.alloydb.v1beta.AlloyDBAdminClient; import java.io.Closeable; import java.io.IOException; @@ -34,10 +35,7 @@ public enum ConnectorRegistry implements Closeable { private final ScheduledExecutorService executor; @SuppressWarnings("ImmutableEnumChecker") - private final AlloyDBAdminClient alloyDBAdminClient; - - @SuppressWarnings("ImmutableEnumChecker") - private final Connector connector; + private ConcurrentHashMap registeredConnectors; ConnectorRegistry() { // During refresh, each instance consumes 2 threads from the thread pool. By using 8 threads, @@ -45,27 +43,33 @@ public enum ConnectorRegistry implements Closeable { // configure 3 or fewer instances, requiring 6 threads during refresh. By setting // this to 8, it's enough threads for most users, plus a safety factor of 2. this.executor = Executors.newScheduledThreadPool(8); - try { - alloyDBAdminClient = AlloyDBAdminClient.create(); - } catch (IOException e) { - throw new RuntimeException(e); - } - this.connector = - new Connector( - executor, - new DefaultConnectionInfoRepository(executor, alloyDBAdminClient), - RsaKeyPairGenerator.generateKeyPair(), - new DefaultConnectionInfoCacheFactory(), - new ConcurrentHashMap<>()); + this.registeredConnectors = new ConcurrentHashMap<>(); } - Connector getConnector() { - return this.connector; + Connector getConnector(ConnectionConfig config) { + return registeredConnectors.computeIfAbsent( + config.getNamedConnection(), + k -> { + try { + FixedCredentialsProvider credentialsProvider = + CredentialsProviderFactory.create(config); + AlloyDBAdminClient alloyDBAdminClient = + AlloyDBAdminClientFactory.create(credentialsProvider); + + return new Connector( + executor, + new DefaultConnectionInfoRepository(executor, alloyDBAdminClient), + RsaKeyPairGenerator.generateKeyPair(), + new DefaultConnectionInfoCacheFactory(), + new ConcurrentHashMap<>()); + } catch (IOException e) { + throw new RuntimeException(e); + } + }); } @Override public void close() { this.executor.shutdown(); - this.alloyDBAdminClient.close(); } } diff --git a/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/CredentialsProviderFactory.java b/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/CredentialsProviderFactory.java new file mode 100644 index 00000000..952ad5e7 --- /dev/null +++ b/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/CredentialsProviderFactory.java @@ -0,0 +1,49 @@ +/* + * Copyright 2023 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.alloydb; + +import com.google.api.gax.core.FixedCredentialsProvider; +import com.google.auth.oauth2.GoogleCredentials; +import com.google.auth.oauth2.ImpersonatedCredentials; +import java.io.IOException; +import java.util.Arrays; + +class CredentialsProviderFactory { + + private static final String CLOUD_PLATFORM = "https://www.googleapis.com/auth/cloud-platform"; + private static final String ALLOYDB_LOGIN = "https://www.googleapis.com/auth/alloydb.login"; + + static FixedCredentialsProvider create(ConnectionConfig config) { + GoogleCredentials credentials; + try { + credentials = GoogleCredentials.getApplicationDefault(); + } catch (IOException e) { + throw new RuntimeException("failed to retrieve OAuth2 access token", e); + } + + if (config.getTargetPrincipal() != null && !config.getTargetPrincipal().isEmpty()) { + credentials = + ImpersonatedCredentials.newBuilder() + .setSourceCredentials(credentials) + .setTargetPrincipal(config.getTargetPrincipal()) + .setDelegates(config.getDelegates()) + .setScopes(Arrays.asList(ALLOYDB_LOGIN, CLOUD_PLATFORM)) + .build(); + } + + return FixedCredentialsProvider.create(credentials); + } +} diff --git a/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/DefaultConnectionInfoRepository.java b/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/DefaultConnectionInfoRepository.java index 87745ff0..767a61ee 100644 --- a/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/DefaultConnectionInfoRepository.java +++ b/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/DefaultConnectionInfoRepository.java @@ -25,6 +25,7 @@ import com.google.protobuf.ByteString; import com.google.protobuf.Duration; import java.io.ByteArrayInputStream; +import java.io.Closeable; import java.io.IOException; import java.io.StringWriter; import java.security.KeyPair; @@ -46,7 +47,7 @@ import org.bouncycastle.pkcs.jcajce.JcaPKCS10CertificationRequestBuilder; import org.bouncycastle.util.io.pem.PemObject; -class DefaultConnectionInfoRepository implements ConnectionInfoRepository { +class DefaultConnectionInfoRepository implements ConnectionInfoRepository, Closeable { private static final String CERTIFICATE_REQUEST = "CERTIFICATE REQUEST"; private static final String SHA_256_WITH_RSA = "SHA256WithRSA"; @@ -84,6 +85,11 @@ public ConnectionInfo getConnectionInfo(InstanceName instanceName, KeyPair keyPa info.getIpAddress(), info.getInstanceUid(), clientCertificate, certificateChain); } + @Override + public void close() { + this.alloyDBAdminClient.close(); + } + private com.google.cloud.alloydb.v1beta.ConnectionInfo getConnectionInfo( InstanceName instanceName) { return alloyDBAdminClient.getConnectionInfo(instanceName); diff --git a/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/SocketFactory.java b/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/SocketFactory.java index 57d1d3d9..22fcbe64 100644 --- a/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/SocketFactory.java +++ b/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/SocketFactory.java @@ -27,13 +27,13 @@ public class SocketFactory extends javax.net.SocketFactory { private final ConnectionConfig config; public SocketFactory(Properties properties) { - this.connector = ConnectorRegistry.INSTANCE.getConnector(); this.config = ConnectionConfig.fromConnectionProperties(properties); + this.connector = ConnectorRegistry.INSTANCE.getConnector(config); } @Override public Socket createSocket() throws IOException { - return this.connector.connect(this.config); + return this.connector.connect(this.config.getInstanceName()); } /* diff --git a/alloydb-jdbc-connector/src/test/java/com/google/cloud/alloydb/AlloyDbJdbcServiceAccountImpersonationDataSourceFactory.java b/alloydb-jdbc-connector/src/test/java/com/google/cloud/alloydb/AlloyDbJdbcServiceAccountImpersonationDataSourceFactory.java new file mode 100644 index 00000000..97e3176e --- /dev/null +++ b/alloydb-jdbc-connector/src/test/java/com/google/cloud/alloydb/AlloyDbJdbcServiceAccountImpersonationDataSourceFactory.java @@ -0,0 +1,43 @@ +/* + * Copyright 2023 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.alloydb; + +// [START alloydb_hikaricp_connect_connector_impersonated_user] +import com.zaxxer.hikari.HikariConfig; +import com.zaxxer.hikari.HikariDataSource; + +public class AlloyDbJdbcServiceAccountImpersonationDataSourceFactory { + + public static final String ALLOYDB_USER = System.getenv("ALLOYDB_USER"); + public static final String ALLOYDB_PASS = System.getenv("ALLOYDB_PASS"); + public static final String ALLOYDB_INSTANCE_NAME = System.getenv("ALLOYDB_INSTANCE_NAME"); + public static final String ALLOYDB_IMPERSONATED_USER = System.getenv("ALLOYDB_IMPERSONATED_USER"); + + static HikariDataSource createDataSource() { + HikariConfig config = new HikariConfig(); + + config.setJdbcUrl("jdbc:postgresql:///postgres"); + config.setUsername(ALLOYDB_USER); // e.g., "postgres" + config.setPassword(ALLOYDB_PASS); // e.g., "secret-password" + config.addDataSourceProperty("socketFactory", "com.google.cloud.alloydb.SocketFactory"); + // e.g., "projects/my-project/locations/us-central1/clusters/my-cluster/instances/my-instance" + config.addDataSourceProperty("alloydbInstanceName", ALLOYDB_INSTANCE_NAME); + config.addDataSourceProperty("alloydbTargetPrincipal", ALLOYDB_IMPERSONATED_USER); + + return new HikariDataSource(config); + } +} +// [END alloydb_hikaricp_connect_connector_impersonated_user] diff --git a/alloydb-jdbc-connector/src/test/java/com/google/cloud/alloydb/ConnectionConfigTest.java b/alloydb-jdbc-connector/src/test/java/com/google/cloud/alloydb/ConnectionConfigTest.java index 1f74cd0e..1ad6bad2 100644 --- a/alloydb-jdbc-connector/src/test/java/com/google/cloud/alloydb/ConnectionConfigTest.java +++ b/alloydb-jdbc-connector/src/test/java/com/google/cloud/alloydb/ConnectionConfigTest.java @@ -18,7 +18,10 @@ import static com.google.common.truth.Truth.assertThat; import com.google.cloud.alloydb.v1beta.InstanceName; +import java.util.Arrays; +import java.util.List; import java.util.Properties; +import java.util.stream.Collectors; import org.junit.Test; public class ConnectionConfigTest { @@ -28,21 +31,43 @@ public class ConnectionConfigTest { @Test public void testConfigFromProps() { + final String wantNamedConnection = "my-connection"; + final String wantTargetPrincipal = "test@example.com"; + final List wantDelegates = Arrays.asList("test1@example.com", "test2@example.com"); + final String delegates = wantDelegates.stream().collect(Collectors.joining(",")); Properties props = new Properties(); props.setProperty(ConnectionConfig.ALLOYDB_INSTANCE_NAME, INSTANCE_NAME); + props.setProperty(ConnectionConfig.ALLOYDB_NAMED_CONNECTION, wantNamedConnection); + props.setProperty(ConnectionConfig.ALLOYDB_TARGET_PRINCIPAL, wantTargetPrincipal); + props.setProperty(ConnectionConfig.ALLOYDB_DELEGATES, delegates); ConnectionConfig config = ConnectionConfig.fromConnectionProperties(props); assertThat(config.getInstanceName().toString()).isEqualTo(INSTANCE_NAME); + assertThat(config.getNamedConnection()).isEqualTo(wantNamedConnection); + assertThat(config.getTargetPrincipal()).isEqualTo(wantTargetPrincipal); + assertThat(config.getDelegates()).isEqualTo(wantDelegates); } @Test public void testConfigFromBuilder() { final InstanceName wantInstance = InstanceName.parse(INSTANCE_NAME); + final String wantNamedConnection = "my-connection"; + final String wantTargetPrincipal = "test@example.com"; + final List wantDelegates = Arrays.asList("test1@example.com", "test2@example.com"); - ConnectionConfig config = new ConnectionConfig.Builder().withInstanceName(wantInstance).build(); + ConnectionConfig config = + new ConnectionConfig.Builder() + .withInstanceName(wantInstance) + .withNamedConnection(wantNamedConnection) + .withTargetPrincipal(wantTargetPrincipal) + .withDelegates(wantDelegates) + .build(); assertThat(config.getInstanceName()).isEqualTo(wantInstance); + assertThat(config.getNamedConnection()).isEqualTo(wantNamedConnection); + assertThat(config.getTargetPrincipal()).isEqualTo(wantTargetPrincipal); + assertThat(config.getDelegates()).isEqualTo(wantDelegates); } } diff --git a/alloydb-jdbc-connector/src/test/java/com/google/cloud/alloydb/ITConnectorTest.java b/alloydb-jdbc-connector/src/test/java/com/google/cloud/alloydb/ITConnectorTest.java index 36884ec2..0add8453 100644 --- a/alloydb-jdbc-connector/src/test/java/com/google/cloud/alloydb/ITConnectorTest.java +++ b/alloydb-jdbc-connector/src/test/java/com/google/cloud/alloydb/ITConnectorTest.java @@ -17,6 +17,7 @@ import static com.google.common.truth.Truth.assertThat; +import com.google.api.gax.core.FixedCredentialsProvider; import com.google.cloud.alloydb.v1beta.AlloyDBAdminClient; import com.google.cloud.alloydb.v1beta.InstanceName; import com.google.common.base.Objects; @@ -44,7 +45,10 @@ public class ITConnectorTest { public void setUp() throws IOException { instanceName = System.getenv("ALLOYDB_INSTANCE_NAME"); // Create the client once and close it later. - alloydbAdminClient = AlloyDBAdminClient.create(); + ConnectionConfig config = + new ConnectionConfig.Builder().withInstanceName(InstanceName.parse(instanceName)).build(); + FixedCredentialsProvider credentialsProvider = CredentialsProviderFactory.create(config); + alloydbAdminClient = AlloyDBAdminClientFactory.create(credentialsProvider); } @After @@ -68,9 +72,7 @@ public void testConnect_createsSocketConnection() throws IOException { new DefaultConnectionInfoCacheFactory(), new ConcurrentHashMap<>()); - ConnectionConfig config = - new ConnectionConfig.Builder().withInstanceName(InstanceName.parse(instanceName)).build(); - socket = (SSLSocket) connector.connect(config); + socket = (SSLSocket) connector.connect(InstanceName.parse(instanceName)); assertThat(socket.getKeepAlive()).isTrue(); assertThat(socket.getTcpNoDelay()).isTrue(); @@ -104,18 +106,16 @@ public void testConnect_whenTlsHandshakeFails() SSLSocket socket = null; ScheduledThreadPoolExecutor executor = null; - try (AlloyDBAdminClient alloyDBAdminClient = AlloyDBAdminClientFactory.create()) { + try { executor = (ScheduledThreadPoolExecutor) Executors.newScheduledThreadPool(2); Connector connector = new Connector( executor, - new DefaultConnectionInfoRepository(executor, alloyDBAdminClient), + new DefaultConnectionInfoRepository(executor, alloydbAdminClient), clientConnectorKeyPair, connectionInfoCacheFactory, new ConcurrentHashMap<>()); - ConnectionConfig config = - new ConnectionConfig.Builder().withInstanceName(InstanceName.parse(instanceName)).build(); - socket = (SSLSocket) connector.connect(config); + socket = (SSLSocket) connector.connect(InstanceName.parse(instanceName)); } catch (ConnectException ignore) { // The socket connect will fail because it's trying to connect to localhost with TLS certs. // So ignore the exception here. diff --git a/alloydb-jdbc-connector/src/test/java/com/google/cloud/alloydb/ITDefaultConnectionInfoRepositoryTest.java b/alloydb-jdbc-connector/src/test/java/com/google/cloud/alloydb/ITDefaultConnectionInfoRepositoryTest.java index 425a9e02..d42582e1 100644 --- a/alloydb-jdbc-connector/src/test/java/com/google/cloud/alloydb/ITDefaultConnectionInfoRepositoryTest.java +++ b/alloydb-jdbc-connector/src/test/java/com/google/cloud/alloydb/ITDefaultConnectionInfoRepositoryTest.java @@ -19,6 +19,7 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertThrows; +import com.google.api.gax.core.FixedCredentialsProvider; import com.google.cloud.alloydb.v1beta.AlloyDBAdminClient; import com.google.cloud.alloydb.v1beta.InstanceName; import java.security.KeyPair; @@ -56,7 +57,10 @@ public void setUp() throws Exception { keyPair = generator.generateKeyPair(); executor = Executors.newSingleThreadExecutor(); - alloyDBAdminClient = AlloyDBAdminClientFactory.create(); + ConnectionConfig config = + new ConnectionConfig.Builder().withInstanceName(InstanceName.parse(instanceUri)).build(); + FixedCredentialsProvider credentialsProvider = CredentialsProviderFactory.create(config); + alloyDBAdminClient = AlloyDBAdminClientFactory.create(credentialsProvider); defaultConnectionInfoRepository = new DefaultConnectionInfoRepository(executor, alloyDBAdminClient); diff --git a/alloydb-jdbc-connector/src/test/java/com/google/cloud/alloydb/ITServiceAccountImpersonationTest.java b/alloydb-jdbc-connector/src/test/java/com/google/cloud/alloydb/ITServiceAccountImpersonationTest.java new file mode 100644 index 00000000..36e2fa81 --- /dev/null +++ b/alloydb-jdbc-connector/src/test/java/com/google/cloud/alloydb/ITServiceAccountImpersonationTest.java @@ -0,0 +1,70 @@ +/* + * Copyright 2023 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.alloydb; + +import static com.google.common.truth.Truth.assertThat; + +import com.google.common.collect.BoundType; +import com.google.common.collect.Range; +import com.zaxxer.hikari.HikariDataSource; +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Timestamp; +import java.time.Instant; +import java.time.temporal.ChronoUnit; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +public class ITServiceAccountImpersonationTest { + + private HikariDataSource dataSource; + + @Before + public void setUp() { + this.dataSource = AlloyDbJdbcServiceAccountImpersonationDataSourceFactory.createDataSource(); + } + + @After + public void tearDown() { + if (this.dataSource != null) { + dataSource.close(); + } + } + + @Test + public void testConnect() throws SQLException { + try (Connection connection = dataSource.getConnection()) { + try (PreparedStatement statement = connection.prepareStatement("SELECT NOW()")) { + ResultSet resultSet = statement.executeQuery(); + resultSet.next(); + Timestamp timestamp = resultSet.getTimestamp(1); + Instant databaseInstant = timestamp.toInstant(); + + Instant now = Instant.now(); + assertThat(databaseInstant) + .isIn( + Range.range( + now.minus(1, ChronoUnit.MINUTES), + BoundType.CLOSED, + now.plus(1, ChronoUnit.MINUTES), + BoundType.CLOSED)); + } + } + } +}