Add support for nomulus tool to connect to Cloud SQL (#303)

This commit is contained in:
Shicong Huang 2019-10-09 10:30:35 -04:00 committed by GitHub
parent 022f397cd9
commit a694e247cd
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 106 additions and 29 deletions

View file

@ -15,10 +15,7 @@
package google.registry.model.transaction;
import com.google.common.flogger.FluentLogger;
import google.registry.persistence.PersistenceModule.AppEngineEmf;
import google.registry.util.Clock;
import javax.inject.Inject;
import javax.inject.Singleton;
import javax.persistence.EntityManager;
import javax.persistence.EntityManagerFactory;
import javax.persistence.EntityTransaction;
@ -26,7 +23,6 @@ import javax.persistence.PersistenceException;
import org.joda.time.DateTime;
/** Implementation of {@link JpaTransactionManager} for JPA compatible database. */
@Singleton
public class JpaTransactionManagerImpl implements JpaTransactionManager {
private static final FluentLogger logger = FluentLogger.forEnclosingClass();
@ -39,8 +35,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
private final ThreadLocal<TransactionInfo> transactionInfo =
ThreadLocal.withInitial(TransactionInfo::new);
@Inject
JpaTransactionManagerImpl(@AppEngineEmf EntityManagerFactory emf, Clock clock) {
public JpaTransactionManagerImpl(EntityManagerFactory emf, Clock clock) {
this.emf = emf;
this.clock = clock;
}

View file

@ -49,6 +49,15 @@ public class TransactionManagerFactory {
return new DatastoreTransactionManager(null);
}
/**
* Sets jpaTm to the implementation for Nomulus tool. Note that this method should be only used by
* {@link google.registry.tools.RegistryCli} to initialize jpaTm.
*/
public static void initForTool() {
// TODO(shicong): Uncomment the line below when we set up Cloud SQL instance in all environments
// jpaTm = DaggerPersistenceComponent.create().nomulusToolJpaTransactionManager();
}
/** Returns {@link TransactionManager} instance. */
public static TransactionManager tm() {
return TM;
@ -56,11 +65,6 @@ public class TransactionManagerFactory {
/** Returns {@link JpaTransactionManager} instance. */
public static JpaTransactionManager jpaTm() {
// TODO: Returns corresponding TransactionManager based on the runtime environment.
// We have 3 kinds of runtime environment:
// 1. App Engine
// 2. Local JVM used by nomulus tool
// 3. Unit test
return jpaTm;
}
}

View file

@ -18,7 +18,9 @@ import dagger.Component;
import google.registry.config.CredentialModule;
import google.registry.config.RegistryConfig.ConfigModule;
import google.registry.keyring.kms.KmsModule;
import google.registry.model.transaction.JpaTransactionManagerImpl;
import google.registry.model.transaction.JpaTransactionManager;
import google.registry.persistence.PersistenceModule.AppEngineJpaTm;
import google.registry.persistence.PersistenceModule.NomulusToolJpaTm;
import google.registry.util.UtilsModule;
import javax.inject.Singleton;
import javax.persistence.EntityManagerFactory;
@ -34,5 +36,10 @@ import javax.persistence.EntityManagerFactory;
UtilsModule.class
})
public interface PersistenceComponent {
JpaTransactionManagerImpl jpaTransactionManager();
@AppEngineJpaTm
JpaTransactionManager appEngineJpaTransactionManager();
@NomulusToolJpaTm
JpaTransactionManager nomulusToolJpaTransactionManager();
}

View file

@ -29,11 +29,14 @@ import dagger.Module;
import dagger.Provides;
import google.registry.config.RegistryConfig.Config;
import google.registry.keyring.kms.KmsKeyring;
import google.registry.model.transaction.JpaTransactionManager;
import google.registry.model.transaction.JpaTransactionManagerImpl;
import google.registry.util.Clock;
import java.lang.annotation.Documented;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.util.HashMap;
import java.util.Map;
import javax.inject.Qualifier;
import javax.inject.Singleton;
import javax.persistence.EntityManagerFactory;
import javax.persistence.Persistence;
import org.hibernate.cfg.Environment;
@ -77,24 +80,55 @@ public class PersistenceModule {
}
@Provides
@AppEngineEmf
public static EntityManagerFactory providesAppEngineEntityManagerFactory(
@Singleton
@AppEngineJpaTm
public static JpaTransactionManager providesAppEngineJpaTm(
@Config("cloudSqlJdbcUrl") String jdbcUrl,
@Config("cloudSqlUsername") String username,
@Config("cloudSqlInstanceConnectionName") String instanceConnectionName,
KmsKeyring kmsKeyring,
@DefaultHibernateConfigs ImmutableMap<String, String> defaultConfigs) {
String password = kmsKeyring.getCloudSqlPassword();
@DefaultHibernateConfigs ImmutableMap<String, String> defaultConfigs,
Clock clock) {
HashMap<String, String> overrides = Maps.newHashMap(defaultConfigs);
// For Java users, the Cloud SQL JDBC Socket Factory can provide authenticated connections.
// See https://github.com/GoogleCloudPlatform/cloud-sql-jdbc-socket-factory for details.
overrides.put("socketFactory", "com.google.cloud.sql.postgres.SocketFactory");
overrides.put("cloudSqlInstance", instanceConnectionName);
EntityManagerFactory emf = create(jdbcUrl, username, password, ImmutableMap.copyOf(overrides));
Runtime.getRuntime().addShutdownHook(new Thread(emf::close));
return emf;
overrides.put(Environment.URL, jdbcUrl);
overrides.put(Environment.USER, username);
overrides.put(Environment.PASS, kmsKeyring.getCloudSqlPassword());
return new JpaTransactionManagerImpl(create(overrides), clock);
}
@Provides
@Singleton
@NomulusToolJpaTm
public static JpaTransactionManager providesNomulusToolJpaTm(
@Config("toolsCloudSqlJdbcUrl") String jdbcUrl,
@Config("toolsCloudSqlUsername") String username,
@Config("cloudSqlInstanceConnectionName") String instanceConnectionName,
KmsKeyring kmsKeyring,
@DefaultHibernateConfigs ImmutableMap<String, String> defaultConfigs,
Clock clock) {
// Cloud SQL JDBC Socket Factory library requires the jdbc url to include all connection
// information, otherwise the connection initialization will fail. See here for more details:
// https://github.com/GoogleCloudPlatform/cloud-sql-jdbc-socket-factory
String fullJdbcUrl =
new StringBuilder(jdbcUrl)
.append("?cloudSqlInstance=" + instanceConnectionName)
.append("&socketFactory=com.google.cloud.sql.postgres.SocketFactory")
.append("&user=" + username)
.append("&password=" + kmsKeyring.getCloudSqlPassword())
.toString();
HashMap<String, String> overrides = Maps.newHashMap(defaultConfigs);
overrides.put(Environment.URL, fullJdbcUrl);
return new JpaTransactionManagerImpl(create(overrides), clock);
}
/** Constructs the {@link EntityManagerFactory} instance. */
@ -106,29 +140,36 @@ public class PersistenceModule {
properties.put(Environment.USER, username);
properties.put(Environment.PASS, password);
return create(ImmutableMap.copyOf(properties));
}
private static EntityManagerFactory create(Map<String, String> properties) {
// If there are no annotated classes, we can create the EntityManagerFactory from the generic
// method. Otherwise we have to use a more tailored approach. Note that this adds to the set
// of annotated classes defined in the configuration, it does not override them.
EntityManagerFactory emf =
Persistence.createEntityManagerFactory(PERSISTENCE_UNIT_NAME, properties);
Persistence.createEntityManagerFactory(
PERSISTENCE_UNIT_NAME, ImmutableMap.copyOf(properties));
checkState(
emf != null,
"Persistence.createEntityManagerFactory() returns a null EntityManagerFactory");
return emf;
}
/** Dagger qualifier for the {@link EntityManagerFactory} used for App Engine application. */
/** Dagger qualifier for {@link JpaTransactionManager} used for App Engine application. */
@Qualifier
@Documented
@Retention(RetentionPolicy.RUNTIME)
public @interface AppEngineEmf {}
@interface AppEngineJpaTm {}
/** Dagger qualifier for {@link JpaTransactionManager} used for Nomulus tool. */
@Qualifier
@Documented
@interface NomulusToolJpaTm {}
/** Dagger qualifier for the default Hibernate configurations. */
// TODO(shicong): Change annotations in this class to none public or put them in a top level
// package
@Qualifier
@Documented
@Retention(RetentionPolicy.RUNTIME)
public @interface DefaultHibernateConfigs {}
}

View file

@ -0,0 +1,23 @@
// Copyright 2019 The Nomulus Authors. All Rights Reserved.
//
// 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
//
// http://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 google.registry.tools;
/**
* Marker interface for commands that use Cloud Sql.
*
* <p>Just implementing this is sufficient to use Cloud Sql; {@link RegistryTool} will install it as
* needed.
*/
interface CommandWithCloudSql extends CommandWithRemoteApi {}

View file

@ -31,6 +31,7 @@ import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Iterables;
import google.registry.config.RegistryConfig;
import google.registry.model.ofy.ObjectifyService;
import google.registry.model.transaction.TransactionManagerFactory;
import google.registry.tools.AuthModule.LoginRequiredException;
import google.registry.tools.params.ParameterFactory;
import java.io.ByteArrayInputStream;
@ -210,6 +211,8 @@ final class RegistryCli implements AutoCloseable, CommandRunner {
}
// CommandWithRemoteApis need to have the remote api installed to work.
// CommandWithCloudSql extends CommandWithRemoteApi so the command will also get the remote
// api installed. This is because the DB password is stored in Datastore.
if (command instanceof CommandWithRemoteApi) {
if (installer == null) {
installer = new RemoteApiInstaller();
@ -233,6 +236,10 @@ final class RegistryCli implements AutoCloseable, CommandRunner {
ofy().clearSessionCache();
}
if (command instanceof CommandWithCloudSql) {
TransactionManagerFactory.initForTool();
}
command.run();
}