Add support for nomulus tool to connect to Cloud SQL

core/src/main/java/google/registry/model/transaction/JpaTransactionManagerImpl.java

@@ -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 @@ 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;
   }

core/src/main/java/google/registry/model/transaction/TransactionManagerFactory.java

@@ -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;
   }
 }

core/src/main/java/google/registry/persistence/PersistenceComponent.java

@@ -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();
 }

core/src/main/java/google/registry/persistence/PersistenceModule.java

@@ -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;
@@ -75,24 +78,55 @@ public static ImmutableMap<String, String> providesDefaultDatabaseConfigs() {
   }
 
   @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. */
@@ -104,29 +138,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 {}
 }

core/src/main/java/google/registry/tools/CommandWithCloudSql.java

@@ -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 {}

core/src/main/java/google/registry/tools/RegistryCli.java

@@ -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();
   }