Implement JpaTransactionManager (#268)

core/build.gradle

@@ -56,7 +56,7 @@ def dockerIncompatibleTestPatterns = [
     // every file is read/write-able. There is no way to exclude specific test
     // methods, so we exclude the whole test class.
     "google/registry/tools/params/PathParameterTest.*",
-    "google/registry/persistence/EntityManagerFactoryProviderTest.*",
+    "google/registry/persistence/PersistenceModuleTest.*",
 ]
 
 // Tests that conflict with members of both the main test suite and the

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
+    // TODO: Determine how to provision TransactionManager after the dual-write. During the
-    //  CloudSqlTransactionManager
+    // 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. */
+/** Dagger module class for the persistence layer. */
-public class EntityManagerFactoryProvider {
+@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) {
+  @VisibleForTesting
-    ImmutableMap.Builder<String, String> properties = ImmutableMap.builder();
+  public static EntityManagerFactory create(
-    properties.putAll(getDefaultProperties());
+      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 {}
 }

core/src/test/java/google/registry/model/transaction/JpaTransactionManagerTest.java

@@ -0,0 +1,211 @@
+// 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 static com.google.common.truth.Truth.assertThat;
+import static google.registry.testing.JUnitBackports.assertThrows;
+import static google.registry.testing.TestDataHelper.fileClassPath;
+import static org.joda.time.DateTimeZone.UTC;
+
+import google.registry.persistence.PersistenceModule;
+import google.registry.testing.FakeClock;
+import google.registry.util.Clock;
+import java.math.BigInteger;
+import javax.persistence.EntityManager;
+import javax.persistence.EntityManagerFactory;
+import javax.persistence.PersistenceException;
+import org.joda.time.DateTime;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.testcontainers.containers.JdbcDatabaseContainer;
+import org.testcontainers.containers.PostgreSQLContainer;
+
+/** Unit tests for {@link JpaTransactionManager}. */
+@RunWith(JUnit4.class)
+public class JpaTransactionManagerTest {
+  @Rule
+  public JdbcDatabaseContainer database =
+      new PostgreSQLContainer().withInitScript(fileClassPath(getClass(), "test_schema.sql"));
+
+  private DateTime now = DateTime.now(UTC);
+  private Clock clock = new FakeClock(now);
+  private EntityManagerFactory emf;
+  private JpaTransactionManager txnManager;
+
+  @Before
+  public void init() {
+    emf =
+        PersistenceModule.create(
+            database.getJdbcUrl(),
+            database.getUsername(),
+            database.getPassword(),
+            PersistenceModule.providesDefaultDatabaseConfigs());
+    txnManager = new JpaTransactionManager(emf, clock);
+  }
+
+  @After
+  public void clear() {
+    if (emf != null) {
+      emf.close();
+    }
+  }
+
+  @Test
+  public void inTransaction_returnsCorrespondingResult() {
+    assertThat(txnManager.inTransaction()).isFalse();
+    txnManager.transact(() -> assertThat(txnManager.inTransaction()).isTrue());
+    assertThat(txnManager.inTransaction()).isFalse();
+  }
+
+  @Test
+  public void assertInTransaction_throwsExceptionWhenNotInTransaction() {
+    assertThrows(PersistenceException.class, () -> txnManager.assertInTransaction());
+    txnManager.transact(() -> txnManager.assertInTransaction());
+    assertThrows(PersistenceException.class, () -> txnManager.assertInTransaction());
+  }
+
+  @Test
+  public void getTransactionTime_throwsExceptionWhenNotInTransaction() {
+    assertThrows(PersistenceException.class, () -> txnManager.getTransactionTime());
+    txnManager.transact(() -> assertThat(txnManager.getTransactionTime()).isEqualTo(now));
+    assertThrows(PersistenceException.class, () -> txnManager.getTransactionTime());
+  }
+
+  @Test
+  public void transact_succeeds() {
+    assertPersonEmpty();
+    assertCompanyEmpty();
+    txnManager.transact(
+        () -> {
+          insertPerson(10);
+          insertCompany("Foo");
+          insertCompany("Bar");
+        });
+    assertPersonCount(1);
+    assertPersonExist(10);
+    assertCompanyCount(2);
+    assertCompanyExist("Foo");
+    assertCompanyExist("Bar");
+  }
+
+  @Test
+  public void transact_hasNoEffectWithPartialSuccess() {
+    assertPersonEmpty();
+    assertCompanyEmpty();
+    assertThrows(
+        RuntimeException.class,
+        () ->
+            txnManager.transact(
+                () -> {
+                  insertPerson(10);
+                  insertCompany("Foo");
+                  throw new RuntimeException();
+                }));
+    assertPersonEmpty();
+    assertCompanyEmpty();
+  }
+
+  @Test
+  public void transact_reusesExistingTransaction() {
+    assertPersonEmpty();
+    assertCompanyEmpty();
+    txnManager.transact(
+        () ->
+            txnManager.transact(
+                () -> {
+                  insertPerson(10);
+                  insertCompany("Foo");
+                  insertCompany("Bar");
+                }));
+    assertPersonCount(1);
+    assertPersonExist(10);
+    assertCompanyCount(2);
+    assertCompanyExist("Foo");
+    assertCompanyExist("Bar");
+  }
+
+  private void insertPerson(int age) {
+    txnManager
+        .getEntityManager()
+        .createNativeQuery(String.format("INSERT INTO Person (age) VALUES (%d)", age))
+        .executeUpdate();
+  }
+
+  private void insertCompany(String name) {
+    txnManager
+        .getEntityManager()
+        .createNativeQuery(String.format("INSERT INTO Company (name) VALUES ('%s')", name))
+        .executeUpdate();
+  }
+
+  private void assertPersonExist(int age) {
+    txnManager.transact(
+        () -> {
+          EntityManager em = txnManager.getEntityManager();
+          Integer maybeAge =
+              (Integer)
+                  em.createNativeQuery(String.format("SELECT age FROM Person WHERE age = %d", age))
+                      .getSingleResult();
+          assertThat(maybeAge).isEqualTo(age);
+        });
+  }
+
+  private void assertCompanyExist(String name) {
+    txnManager.transact(
+        () -> {
+          String maybeName =
+              (String)
+                  txnManager
+                      .getEntityManager()
+                      .createNativeQuery(
+                          String.format("SELECT name FROM Company WHERE name = '%s'", name))
+                      .getSingleResult();
+          assertThat(maybeName).isEqualTo(name);
+        });
+  }
+
+  private void assertPersonCount(int count) {
+    assertThat(countTable("Person")).isEqualTo(count);
+  }
+
+  private void assertCompanyCount(int count) {
+    assertThat(countTable("Company")).isEqualTo(count);
+  }
+
+  private void assertPersonEmpty() {
+    assertPersonCount(0);
+  }
+
+  private void assertCompanyEmpty() {
+    assertCompanyCount(0);
+  }
+
+  private int countTable(String tableName) {
+    return txnManager.transact(
+        () -> {
+          BigInteger colCount =
+              (BigInteger)
+                  txnManager
+                      .getEntityManager()
+                      .createNativeQuery(String.format("SELECT COUNT(*) FROM %s", tableName))
+                      .getSingleResult();
+          return colCount.intValue();
+        });
+  }
+}

core/src/test/java/google/registry/persistence/PersistenceModuleTest.java

@@ -26,9 +26,9 @@ import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 import org.testcontainers.containers.PostgreSQLContainer;
 
-/** Unit tests for {@link EntityManagerFactoryProvider}. */
+/** Unit tests for {@link PersistenceModule}. */
 @RunWith(JUnit4.class)
-public class EntityManagerFactoryProviderTest {
+public class PersistenceModuleTest {
   @Rule public PostgreSQLContainer database = new PostgreSQLContainer();
 
   private EntityManagerFactory emf;
@@ -36,8 +36,11 @@ public class EntityManagerFactoryProviderTest {
   @Before
   public void init() {
     emf =
-        EntityManagerFactoryProvider.create(
+        PersistenceModule.create(
-            database.getJdbcUrl(), database.getUsername(), database.getPassword());
+            database.getJdbcUrl(),
+            database.getUsername(),
+            database.getPassword(),
+            PersistenceModule.providesDefaultDatabaseConfigs());
   }
 
   @After

core/src/test/java/google/registry/testing/TestDataHelper.java

@@ -114,6 +114,23 @@ public final class TestDataHelper {
     return String.format("src/test/resources/%s/%s", packagePath, filename);
   }
 
+  /**
+   * Constructs the relative classpath for the given {@code filename} under the {@code context}'s
+   * package.
+   *
+   * <p>For example, if the {@code context} is a class with name com.google.registry.FileClasspath,
+   * and the given {@code filename} is testdata.txt, then the return value would be
+   * com/google/registry/testdata.txt.
+   *
+   * <p>This function is useful when you just need a relative path starting from the Java root
+   * package to the given {@code filename}. The other utility functions in this class either return
+   * an absolute path or a relative path but starting from src/ directory.
+   */
+  public static String fileClassPath(Class<?> context, String filename) {
+    String packagePath = context.getPackage().getName().replace('.', '/');
+    return String.format("%s/%s", packagePath, filename);
+  }
+
   /** Returns a recursive iterable of all files in the given directory. */
   public static Iterable<Path> listFiles(Class<?> context, String directory) throws Exception {
     URI dir = Resources.getResource(context, directory).toURI();

core/src/test/resources/google/registry/model/transaction/test_schema.sql

@@ -0,0 +1,21 @@
+-- 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.
+
+CREATE TABLE Person (
+  age INT NOT NULL
+);
+
+CREATE TABLE Company (
+  name TEXT NOT NULL
+);