Implement JpaTransactionManager (#268)

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

@@ -0,0 +1,167 @@
+// 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.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 TransactionManager} for JPA compatible database. */
+@Singleton
+public class JpaTransactionManager implements TransactionManager {
+
+  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
+
+  // EntityManagerFactory is thread safe.
+  private final EntityManagerFactory emf;
+  private final Clock clock;
+  // TODO(shicong): Investigate alternatives for managing transaction information. ThreadLocal adds
+  //  an unnecessary restriction that each request has to be processed by one thread synchronously.
+  private final ThreadLocal<TransactionInfo> transactionInfo =
+      ThreadLocal.withInitial(TransactionInfo::new);
+
+  @Inject
+  JpaTransactionManager(@AppEngineEmf EntityManagerFactory emf, Clock clock) {
+    this.emf = emf;
+    this.clock = clock;
+  }
+
+  public EntityManager getEntityManager() {
+    if (transactionInfo.get().entityManager == null) {
+      throw new PersistenceException(
+          "No EntityManager has been initialized. getEntityManager() must be invoked in the scope"
+              + " of a transaction");
+    }
+    return transactionInfo.get().entityManager;
+  }
+
+  @Override
+  public boolean inTransaction() {
+    return transactionInfo.get().inTransaction;
+  }
+
+  @Override
+  public void assertInTransaction() {
+    if (!inTransaction()) {
+      throw new PersistenceException("Not in a transaction");
+    }
+  }
+
+  @Override
+  public <T> T transact(Work<T> work) {
+    // TODO(shicong): Investigate removing transactNew functionality after migration as it may
+    //  be same as this one.
+    if (inTransaction()) {
+      return work.run();
+    }
+    TransactionInfo txnInfo = transactionInfo.get();
+    txnInfo.entityManager = emf.createEntityManager();
+    EntityTransaction txn = txnInfo.entityManager.getTransaction();
+    try {
+      txn.begin();
+      txnInfo.inTransaction = true;
+      txnInfo.transactionTime = clock.nowUtc();
+      T result = work.run();
+      txn.commit();
+      return result;
+    } catch (Throwable transactionException) {
+      String rollbackMessage;
+      try {
+        txn.rollback();
+        rollbackMessage = "transaction rolled back";
+      } catch (Throwable rollbackException) {
+        logger.atSevere().withCause(rollbackException).log("Rollback failed, suppressing error");
+        rollbackMessage = "transaction rollback failed";
+      }
+      throw new PersistenceException(
+          "Error during transaction, " + rollbackMessage, transactionException);
+    } finally {
+      txnInfo.clear();
+    }
+  }
+
+  @Override
+  public void transact(Runnable work) {
+    transact(
+        () -> {
+          work.run();
+          return null;
+        });
+  }
+
+  @Override
+  public <T> T transactNew(Work<T> work) {
+    // TODO(shicong): Implements the functionality to start a new transaction.
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public void transactNew(Runnable work) {
+    // TODO(shicong): Implements the functionality to start a new transaction.
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public <T> T transactNewReadOnly(Work<T> work) {
+    // TODO(shicong): Implements read only transaction.
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public void transactNewReadOnly(Runnable work) {
+    // TODO(shicong): Implements read only transaction.
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public <T> T doTransactionless(Work<T> work) {
+    // TODO(shicong): Implements doTransactionless.
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public DateTime getTransactionTime() {
+    assertInTransaction();
+    TransactionInfo txnInfo = transactionInfo.get();
+    if (txnInfo.transactionTime == null) {
+      throw new PersistenceException("In a transaction but transactionTime is null");
+    }
+    return txnInfo.transactionTime;
+  }
+
+  private static class TransactionInfo {
+    EntityManager entityManager;
+    boolean inTransaction = false;
+    DateTime transactionTime;
+
+    private void clear() {
+      inTransaction = false;
+      transactionTime = null;
+      if (entityManager != null) {
+        // Close this EntityManager just let the connection pool be able to reuse it, it doesn't
+        // close the underlying database connection.
+        entityManager.close();
+        entityManager = null;
+      }
+    }
+  }
+}

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

@@ -14,24 +14,40 @@
 
 package google.registry.model.transaction;
 
+import com.google.common.annotations.VisibleForTesting;
 import google.registry.model.ofy.DatastoreTransactionManager;
+import google.registry.persistence.DaggerPersistenceComponent;
+import google.registry.persistence.PersistenceComponent;
 
 /** Factory class to create {@link TransactionManager} instance. */
+// TODO: Rename this to PersistenceFactory and move to persistence package.
 public class TransactionManagerFactory {
 
   private static final TransactionManager TM = createTransactionManager();
 
+  @VisibleForTesting static PersistenceComponent component = DaggerPersistenceComponent.create();
+
   private TransactionManagerFactory() {}
 
   private static TransactionManager createTransactionManager() {
-    // TODO: Conditionally returns the corresponding implementation once we have
-    //  CloudSqlTransactionManager
+    // TODO: Determine how to provision TransactionManager after the dual-write. During the
+    // dual-write transitional phase, we need the TransactionManager for both Datastore and Cloud
+    // SQL, and this method returns the one for Datastore.
     return new DatastoreTransactionManager(null);
   }
 
   /** 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 component.jpaTransactionManager();
+  }
 }

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

@@ -0,0 +1,38 @@
+// 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.persistence;
+
+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.JpaTransactionManager;
+import google.registry.util.UtilsModule;
+import javax.inject.Singleton;
+import javax.persistence.EntityManagerFactory;
+
+/** Dagger component to provide {@link EntityManagerFactory} instances. */
+@Singleton
+@Component(
+    modules = {
+      ConfigModule.class,
+      CredentialModule.class,
+      KmsModule.class,
+      PersistenceModule.class,
+      UtilsModule.class
+    })
+public interface PersistenceComponent {
+  JpaTransactionManager jpaTransactionManager();
+}

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

@@ -22,13 +22,25 @@ import static google.registry.config.RegistryConfig.getHibernateHikariMaximumPoo
 import static google.registry.config.RegistryConfig.getHibernateHikariMinimumIdle;
 import static google.registry.config.RegistryConfig.getHibernateLogSqlQueries;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Maps;
+import dagger.Module;
+import dagger.Provides;
+import google.registry.config.RegistryConfig.Config;
+import google.registry.keyring.kms.KmsKeyring;
+import java.lang.annotation.Documented;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.util.HashMap;
+import javax.inject.Qualifier;
 import javax.persistence.EntityManagerFactory;
 import javax.persistence.Persistence;
 import org.hibernate.cfg.Environment;
 
-/** Factory class to provide {@link EntityManagerFactory} instance. */
-public class EntityManagerFactoryProvider {
+/** Dagger module class for the persistence layer. */
+@Module
+public class PersistenceModule {
   // This name must be the same as the one defined in persistence.xml.
   public static final String PERSISTENCE_UNIT_NAME = "nomulus";
   public static final String HIKARI_CONNECTION_TIMEOUT = "hibernate.hikari.connectionTimeout";
@@ -36,7 +48,9 @@ public class EntityManagerFactoryProvider {
   public static final String HIKARI_MAXIMUM_POOL_SIZE = "hibernate.hikari.maximumPoolSize";
   public static final String HIKARI_IDLE_TIMEOUT = "hibernate.hikari.idleTimeout";
 
-  private static ImmutableMap<String, String> getDefaultProperties() {
+  @Provides
+  @DefaultHibernateConfigs
+  public static ImmutableMap<String, String> providesDefaultDatabaseConfigs() {
     ImmutableMap.Builder<String, String> properties = ImmutableMap.builder();
 
     properties.put(Environment.DRIVER, "org.postgresql.Driver");
@@ -59,18 +73,54 @@ public class EntityManagerFactoryProvider {
     return properties.build();
   }
 
+  @Provides
+  @AppEngineEmf
+  public static EntityManagerFactory providesAppEngineEntityManagerFactory(
+      @Config("cloudSqlJdbcUrl") String jdbcUrl,
+      @Config("cloudSqlUsername") String username,
+      @Config("cloudSqlInstanceConnectionName") String instanceConnectionName,
+      KmsKeyring kmsKeyring,
+      @DefaultHibernateConfigs ImmutableMap<String, String> defaultConfigs) {
+    String password = kmsKeyring.getCloudSqlPassword();
+
+    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;
+  }
+
   /** Constructs the {@link EntityManagerFactory} instance. */
-  public static EntityManagerFactory create(String jdbcUrl, String username, String password) {
-    ImmutableMap.Builder<String, String> properties = ImmutableMap.builder();
-    properties.putAll(getDefaultProperties());
+  @VisibleForTesting
+  public static EntityManagerFactory create(
+      String jdbcUrl, String username, String password, ImmutableMap<String, String> configs) {
+    HashMap<String, String> properties = Maps.newHashMap(configs);
     properties.put(Environment.URL, jdbcUrl);
     properties.put(Environment.USER, username);
     properties.put(Environment.PASS, password);
     EntityManagerFactory emf =
-        Persistence.createEntityManagerFactory(PERSISTENCE_UNIT_NAME, properties.build());
+        Persistence.createEntityManagerFactory(PERSISTENCE_UNIT_NAME, properties);
     checkState(
         emf != null,
         "Persistence.createEntityManagerFactory() returns a null EntityManagerFactory");
     return emf;
   }
+
+  /** Dagger qualifier for the {@link EntityManagerFactory} used for App Engine application. */
+  @Qualifier
+  @Documented
+  @Retention(RetentionPolicy.RUNTIME)
+  public @interface AppEngineEmf {}
+
+  /** 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 {}
 }