Fix issues related to Cloud SQL connection (#321)

* Reenable JpaTransactionManager for Alpha and Crash

* Make UploadClaimsListCommand implement CommandWithCloudSql

* Fix wrong call to get password

* Use Cloud SQL Socket library to provision TransactionManager

* Change to use dataSource configs

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

@@ -14,8 +14,15 @@
 
 package google.registry.model.transaction;
 
+import static google.registry.config.RegistryEnvironment.ALPHA;
+import static google.registry.config.RegistryEnvironment.CRASH;
+
+import com.google.appengine.api.utils.SystemProperty;
+import com.google.appengine.api.utils.SystemProperty.Environment.Value;
 import com.google.common.annotations.VisibleForTesting;
+import google.registry.config.RegistryEnvironment;
 import google.registry.model.ofy.DatastoreTransactionManager;
+import google.registry.persistence.DaggerPersistenceComponent;
 
 /** Factory class to create {@link TransactionManager} instance. */
 // TODO: Rename this to PersistenceFactory and move to persistence package.
@@ -27,20 +34,12 @@ public class TransactionManagerFactory {
   private TransactionManagerFactory() {}
 
   private static JpaTransactionManager createJpaTransactionManager() {
-    // TODO(shicong): There is currently no environment where we want to create a real JPA
+    if (shouldEnableJpaTm() && isInAppEngine()) {
-    // transaction manager here.  The unit tests that require one are all set up using
+      return DaggerPersistenceComponent.create().appEngineJpaTransactionManager();
-    // JpaTransactionManagerRule which launches its own PostgreSQL instance.  When we actually have
+    } else {
-    // PostgreSQL tables in production, ensure that all of the test environments are set up
-    // correctly and restore the code that creates a JpaTransactionManager when
-    // RegistryEnvironment.get() != UNITTEST.
-    //
-    // We removed the original code because it didn't work in sandbox (due to the absence of the
-    // persistence.xml file, which has since been added), and then (after adding this) didn't work
-    // in crash because the postgresql password hadn't been set up.  Prior to restoring, we'll need
-    // to do setup in all environments, and we probably only want to do this once we're actually
-    // using Cloud SQL for one of the new tables.
       return DummyJpaTransactionManager.create();
     }
+  }
 
   private static TransactionManager createTransactionManager() {
     // TODO: Determine how to provision TransactionManager after the dual-write. During the
@@ -54,8 +53,27 @@ public class TransactionManagerFactory {
    * {@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
+    if (shouldEnableJpaTm()) {
-    // jpaTm = DaggerPersistenceComponent.create().nomulusToolJpaTransactionManager();
+      jpaTm = DaggerPersistenceComponent.create().nomulusToolJpaTransactionManager();
+    }
+  }
+
+  // TODO(shicong): Enable JpaTm for all environments and remove this function
+  private static boolean shouldEnableJpaTm() {
+    return RegistryEnvironment.get() == ALPHA || RegistryEnvironment.get() == CRASH;
+  }
+
+  /**
+   * This function uses App Engine API to determine if the current runtime environment is App
+   * Engine.
+   *
+   * @see <a
+   *     href="https://cloud.google.com/appengine/docs/standard/java/javadoc/com/google/appengine/api/utils/SystemProperty">App
+   *     Engine API public doc</a>
+   */
+  private static boolean isInAppEngine() {
+    // SystemProperty.environment.value() returns null if the current runtime is local JVM
+    return SystemProperty.environment.value() == Value.Production;
   }
 
   /** Returns {@link TransactionManager} instance. */

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

@@ -51,6 +51,10 @@ public class PersistenceModule {
   public static final String HIKARI_MAXIMUM_POOL_SIZE = "hibernate.hikari.maximumPoolSize";
   public static final String HIKARI_IDLE_TIMEOUT = "hibernate.hikari.idleTimeout";
 
+  public static final String HIKARI_DS_SOCKET_FACTORY = "hibernate.hikari.dataSource.socketFactory";
+  public static final String HIKARI_DS_CLOUD_SQL_INSTANCE =
+      "hibernate.hikari.dataSource.cloudSqlInstance";
+
   @Provides
   @DefaultHibernateConfigs
   public static ImmutableMap<String, String> providesDefaultDatabaseConfigs() {
@@ -79,27 +83,31 @@ public class PersistenceModule {
     return properties.build();
   }
 
+  @Provides
+  @Singleton
+  @PartialCloudSqlConfigs
+  public static ImmutableMap<String, String> providesPartialCloudSqlConfigs(
+      @Config("cloudSqlJdbcUrl") String jdbcUrl,
+      @Config("cloudSqlInstanceConnectionName") String instanceConnectionName,
+      @DefaultHibernateConfigs ImmutableMap<String, String> defaultConfigs) {
+    HashMap<String, String> overrides = Maps.newHashMap(defaultConfigs);
+    overrides.put(Environment.URL, jdbcUrl);
+    overrides.put(HIKARI_DS_SOCKET_FACTORY, "com.google.cloud.sql.postgres.SocketFactory");
+    overrides.put(HIKARI_DS_CLOUD_SQL_INSTANCE, instanceConnectionName);
+    return ImmutableMap.copyOf(overrides);
+  }
+
   @Provides
   @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,
+      @PartialCloudSqlConfigs ImmutableMap<String, String> cloudSqlConfigs,
       Clock clock) {
-    HashMap<String, String> overrides = Maps.newHashMap(defaultConfigs);
+    HashMap<String, String> overrides = Maps.newHashMap(cloudSqlConfigs);
-
-    // 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);
-
-    overrides.put(Environment.URL, jdbcUrl);
     overrides.put(Environment.USER, username);
     overrides.put(Environment.PASS, kmsKeyring.getCloudSqlPassword());
-
     return new JpaTransactionManagerImpl(create(overrides), clock);
   }
 
@@ -107,27 +115,13 @@ public class PersistenceModule {
   @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,
+      @PartialCloudSqlConfigs ImmutableMap<String, String> cloudSqlConfigs,
       Clock clock) {
-
+    HashMap<String, String> overrides = Maps.newHashMap(cloudSqlConfigs);
-    // Cloud SQL JDBC Socket Factory library requires the jdbc url to include all connection
+    overrides.put(Environment.USER, username);
-    // information, otherwise the connection initialization will fail. See here for more details:
+    overrides.put(Environment.PASS, kmsKeyring.getToolsCloudSqlPassword());
-    // 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);
   }
 
@@ -166,6 +160,11 @@ public class PersistenceModule {
   @Documented
   @interface NomulusToolJpaTm {}
 
+  /** Dagger qualifier for the partial Cloud SQL configs. */
+  @Qualifier
+  @Documented
+  @interface PartialCloudSqlConfigs {}
+
   /** 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

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

@@ -32,7 +32,7 @@ import java.util.List;
 
 /** A command to upload a {@link ClaimsListShard}. */
 @Parameters(separators = " =", commandDescription = "Manually upload a new claims list file")
-final class UploadClaimsListCommand extends ConfirmingCommand implements CommandWithRemoteApi {
+final class UploadClaimsListCommand extends ConfirmingCommand implements CommandWithCloudSql {
 
   @Parameter(description = "Claims list filename")
   private List<String> mainParameters = new ArrayList<>();