Skip to content
Permalink
Browse files

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

  • Loading branch information...
hstonec committed Oct 9, 2019
1 parent 022f397 commit a694e247cd2c6b0d4e7731a00029f92c1f0b963d
@@ -15,18 +15,14 @@
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;
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 @@
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;
}
@@ -49,18 +49,22 @@ private static TransactionManager createTransactionManager() {
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;
}

/** 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;
}
}
@@ -18,7 +18,9 @@
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 @@
UtilsModule.class
})
public interface PersistenceComponent {
JpaTransactionManagerImpl jpaTransactionManager();

@AppEngineJpaTm
JpaTransactionManager appEngineJpaTransactionManager();

@NomulusToolJpaTm
JpaTransactionManager nomulusToolJpaTransactionManager();
}
@@ -29,11 +29,14 @@
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 @@
}

@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 static EntityManagerFactory create(
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
@interface AppEngineJpaTm {}

/** Dagger qualifier for {@link JpaTransactionManager} used for Nomulus tool. */
@Qualifier
@Documented
@Retention(RetentionPolicy.RUNTIME)
public @interface AppEngineEmf {}
@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 {}
}
@@ -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 {}
@@ -31,6 +31,7 @@
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 @@ private void runCommand(Command command) throws Exception {
}

// 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 @@ private void runCommand(Command command) throws Exception {
ofy().clearSessionCache();
}

if (command instanceof CommandWithCloudSql) {
TransactionManagerFactory.initForTool();
}

command.run();
}

0 comments on commit a694e24

Please sign in to comment.
You can’t perform that action at this time.