diff --git a/core/src/main/java/google/registry/model/ofy/DatastoreTransactionManager.java b/core/src/main/java/google/registry/model/ofy/DatastoreTransactionManager.java
index ded206aa1..8fac0fcda 100644
--- a/core/src/main/java/google/registry/model/ofy/DatastoreTransactionManager.java
+++ b/core/src/main/java/google/registry/model/ofy/DatastoreTransactionManager.java
@@ -14,7 +14,6 @@
package google.registry.model.ofy;
-
import static google.registry.model.ofy.ObjectifyService.ofy;
import google.registry.model.transaction.TransactionManager;
diff --git a/core/src/main/java/google/registry/model/transaction/DummyJpaTransactionManager.java b/core/src/main/java/google/registry/model/transaction/DummyJpaTransactionManager.java
new file mode 100644
index 000000000..6e0ba81ab
--- /dev/null
+++ b/core/src/main/java/google/registry/model/transaction/DummyJpaTransactionManager.java
@@ -0,0 +1,44 @@
+// 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 java.lang.reflect.Proxy;
+
+/**
+ * A dummy implementation for {@link JpaTransactionManager} which throws exception when any of its
+ * method is invoked.
+ *
+ *
This is used to initialize the {@link TransactionManagerFactory#jpaTm} when running unit
+ * tests, because obviously we cannot connect to the actual Cloud SQL backend in a unit test.
+ *
+ *
If a unit test needs to access the Cloud SQL database, it must add JpaTransactionManagerRule
+ * as a JUnit rule in the test class.
+ */
+public class DummyJpaTransactionManager {
+
+ /** Constructs a dummy {@link JpaTransactionManager} instance. */
+ public static JpaTransactionManager create() {
+ return (JpaTransactionManager)
+ Proxy.newProxyInstance(
+ JpaTransactionManager.class.getClassLoader(),
+ new Class[] {JpaTransactionManager.class},
+ (proxy, method, args) -> {
+ throw new UnsupportedOperationException(
+ "JpaTransactionManager was not initialized as the runtime is detected as"
+ + " Unittest. Add JpaTransactionManagerRule in the unit test for"
+ + " initialization.");
+ });
+ }
+}
diff --git a/core/src/main/java/google/registry/model/transaction/JpaTransactionManager.java b/core/src/main/java/google/registry/model/transaction/JpaTransactionManager.java
index 3f5d50111..3a2290521 100644
--- a/core/src/main/java/google/registry/model/transaction/JpaTransactionManager.java
+++ b/core/src/main/java/google/registry/model/transaction/JpaTransactionManager.java
@@ -14,154 +14,10 @@
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 =
- 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 transact(Work 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 transactNew(Work 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 transactNewReadOnly(Work 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 doTransactionless(Work 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;
- }
- }
- }
+/** Sub-interface of {@link TransactionManager} which defines JPA related methods. */
+public interface JpaTransactionManager extends TransactionManager {
+ /** Returns the {@link EntityManager} for the current request. */
+ EntityManager getEntityManager();
}
diff --git a/core/src/main/java/google/registry/model/transaction/JpaTransactionManagerImpl.java b/core/src/main/java/google/registry/model/transaction/JpaTransactionManagerImpl.java
new file mode 100644
index 000000000..f58ca6482
--- /dev/null
+++ b/core/src/main/java/google/registry/model/transaction/JpaTransactionManagerImpl.java
@@ -0,0 +1,168 @@
+// 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 JpaTransactionManager} for JPA compatible database. */
+@Singleton
+public class JpaTransactionManagerImpl implements JpaTransactionManager {
+
+ 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 =
+ ThreadLocal.withInitial(TransactionInfo::new);
+
+ @Inject
+ JpaTransactionManagerImpl(@AppEngineEmf EntityManagerFactory emf, Clock clock) {
+ this.emf = emf;
+ this.clock = clock;
+ }
+
+ @Override
+ 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 transact(Work 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 transactNew(Work 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 transactNewReadOnly(Work 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 doTransactionless(Work 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;
+ }
+ }
+ }
+}
diff --git a/core/src/main/java/google/registry/model/transaction/TransactionManagerFactory.java b/core/src/main/java/google/registry/model/transaction/TransactionManagerFactory.java
index 5deb851fe..77ddcf9d8 100644
--- a/core/src/main/java/google/registry/model/transaction/TransactionManagerFactory.java
+++ b/core/src/main/java/google/registry/model/transaction/TransactionManagerFactory.java
@@ -14,21 +14,29 @@
package google.registry.model.transaction;
+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.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();
+ @VisibleForTesting static JpaTransactionManager jpaTm = createJpaTransactionManager();
private TransactionManagerFactory() {}
+ private static JpaTransactionManager createJpaTransactionManager() {
+ if (SystemProperty.environment.value() == Value.Production) {
+ return DaggerPersistenceComponent.create().jpaTransactionManager();
+ } else {
+ return DummyJpaTransactionManager.create();
+ }
+ }
+
private static TransactionManager createTransactionManager() {
// 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
@@ -48,6 +56,6 @@ public class TransactionManagerFactory {
// 1. App Engine
// 2. Local JVM used by nomulus tool
// 3. Unit test
- return component.jpaTransactionManager();
+ return jpaTm;
}
}
diff --git a/core/src/main/java/google/registry/persistence/PersistenceComponent.java b/core/src/main/java/google/registry/persistence/PersistenceComponent.java
index 5cb880651..fb5c1f0e3 100644
--- a/core/src/main/java/google/registry/persistence/PersistenceComponent.java
+++ b/core/src/main/java/google/registry/persistence/PersistenceComponent.java
@@ -18,7 +18,7 @@ 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.model.transaction.JpaTransactionManagerImpl;
import google.registry.util.UtilsModule;
import javax.inject.Singleton;
import javax.persistence.EntityManagerFactory;
@@ -34,5 +34,5 @@ import javax.persistence.EntityManagerFactory;
UtilsModule.class
})
public interface PersistenceComponent {
- JpaTransactionManager jpaTransactionManager();
+ JpaTransactionManagerImpl jpaTransactionManager();
}
diff --git a/core/src/test/java/google/registry/model/transaction/DummyJpaTransactionManagerTest.java b/core/src/test/java/google/registry/model/transaction/DummyJpaTransactionManagerTest.java
new file mode 100644
index 000000000..173aefc17
--- /dev/null
+++ b/core/src/test/java/google/registry/model/transaction/DummyJpaTransactionManagerTest.java
@@ -0,0 +1,33 @@
+// 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 google.registry.model.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.testing.JUnitBackports.assertThrows;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** JUnit test for {@link DummyJpaTransactionManager} */
+@RunWith(JUnit4.class)
+public class DummyJpaTransactionManagerTest {
+
+ @Test
+ public void throwsExceptionWhenAnyMethodIsInvoked() {
+ assertThrows(UnsupportedOperationException.class, () -> jpaTm().transact(() -> null));
+ assertThrows(UnsupportedOperationException.class, () -> jpaTm().getTransactionTime());
+ }
+}
diff --git a/core/src/test/java/google/registry/model/transaction/JpaTransactionManagerImplTest.java b/core/src/test/java/google/registry/model/transaction/JpaTransactionManagerImplTest.java
new file mode 100644
index 000000000..4e60b9d69
--- /dev/null
+++ b/core/src/test/java/google/registry/model/transaction/JpaTransactionManagerImplTest.java
@@ -0,0 +1,192 @@
+// 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.model.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.testing.JUnitBackports.assertThrows;
+import static google.registry.testing.TestDataHelper.fileClassPath;
+
+import google.registry.testing.FakeClock;
+import java.math.BigInteger;
+import javax.persistence.EntityManager;
+import javax.persistence.PersistenceException;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Unit tests for {@link JpaTransactionManagerImpl}. */
+@RunWith(JUnit4.class)
+public class JpaTransactionManagerImplTest {
+ @Rule
+ public final JpaTransactionManagerRule jpaTmRule =
+ new JpaTransactionManagerRule.Builder()
+ .withInitScript(fileClassPath(getClass(), "test_schema.sql"))
+ .build();
+
+ @Test
+ public void inTransaction_returnsCorrespondingResult() {
+ assertThat(jpaTm().inTransaction()).isFalse();
+ jpaTm().transact(() -> assertThat(jpaTm().inTransaction()).isTrue());
+ assertThat(jpaTm().inTransaction()).isFalse();
+ }
+
+ @Test
+ public void assertInTransaction_throwsExceptionWhenNotInTransaction() {
+ assertThrows(PersistenceException.class, () -> jpaTm().assertInTransaction());
+ jpaTm().transact(() -> jpaTm().assertInTransaction());
+ assertThrows(PersistenceException.class, () -> jpaTm().assertInTransaction());
+ }
+
+ @Test
+ public void getTransactionTime_throwsExceptionWhenNotInTransaction() {
+ FakeClock txnClock = jpaTmRule.getTxnClock();
+ txnClock.advanceOneMilli();
+ assertThrows(PersistenceException.class, () -> jpaTm().getTransactionTime());
+ jpaTm().transact(() -> assertThat(jpaTm().getTransactionTime()).isEqualTo(txnClock.nowUtc()));
+ assertThrows(PersistenceException.class, () -> jpaTm().getTransactionTime());
+ }
+
+ @Test
+ public void transact_succeeds() {
+ assertPersonEmpty();
+ assertCompanyEmpty();
+ jpaTm()
+ .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,
+ () ->
+ jpaTm()
+ .transact(
+ () -> {
+ insertPerson(10);
+ insertCompany("Foo");
+ throw new RuntimeException();
+ }));
+ assertPersonEmpty();
+ assertCompanyEmpty();
+ }
+
+ @Test
+ public void transact_reusesExistingTransaction() {
+ assertPersonEmpty();
+ assertCompanyEmpty();
+ jpaTm()
+ .transact(
+ () ->
+ jpaTm()
+ .transact(
+ () -> {
+ insertPerson(10);
+ insertCompany("Foo");
+ insertCompany("Bar");
+ }));
+ assertPersonCount(1);
+ assertPersonExist(10);
+ assertCompanyCount(2);
+ assertCompanyExist("Foo");
+ assertCompanyExist("Bar");
+ }
+
+ private void insertPerson(int age) {
+ jpaTm()
+ .getEntityManager()
+ .createNativeQuery(String.format("INSERT INTO Person (age) VALUES (%d)", age))
+ .executeUpdate();
+ }
+
+ private void insertCompany(String name) {
+ jpaTm()
+ .getEntityManager()
+ .createNativeQuery(String.format("INSERT INTO Company (name) VALUES ('%s')", name))
+ .executeUpdate();
+ }
+
+ private void assertPersonExist(int age) {
+ jpaTm()
+ .transact(
+ () -> {
+ EntityManager em = jpaTm().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) {
+ jpaTm()
+ .transact(
+ () -> {
+ String maybeName =
+ (String)
+ jpaTm()
+ .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 jpaTm()
+ .transact(
+ () -> {
+ BigInteger colCount =
+ (BigInteger)
+ jpaTm()
+ .getEntityManager()
+ .createNativeQuery(String.format("SELECT COUNT(*) FROM %s", tableName))
+ .getSingleResult();
+ return colCount.intValue();
+ });
+ }
+}
diff --git a/core/src/test/java/google/registry/model/transaction/JpaTransactionManagerRule.java b/core/src/test/java/google/registry/model/transaction/JpaTransactionManagerRule.java
new file mode 100644
index 000000000..a0ba698ff
--- /dev/null
+++ b/core/src/test/java/google/registry/model/transaction/JpaTransactionManagerRule.java
@@ -0,0 +1,109 @@
+// 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 org.joda.time.DateTimeZone.UTC;
+
+import google.registry.persistence.PersistenceModule;
+import google.registry.testing.FakeClock;
+import javax.persistence.EntityManagerFactory;
+import org.joda.time.DateTime;
+import org.junit.rules.ExternalResource;
+import org.junit.rules.RuleChain;
+import org.junit.runner.Description;
+import org.junit.runners.model.Statement;
+import org.testcontainers.containers.JdbcDatabaseContainer;
+import org.testcontainers.containers.PostgreSQLContainer;
+
+/**
+ * JUnit Rule to provision {@link JpaTransactionManagerImpl} backed by {@link PostgreSQLContainer}.
+ *
+ * This rule also replaces the {@link JpaTransactionManagerImpl} provided by {@link
+ * TransactionManagerFactory} with the {@link JpaTransactionManagerImpl} generated by the rule
+ * itself, so that all SQL queries will be sent to the database instance created by {@link
+ * PostgreSQLContainer} to achieve test purpose.
+ */
+public class JpaTransactionManagerRule extends ExternalResource {
+ private static final String SCHEMA_GOLDEN_SQL = "sql/schema/nomulus.golden.sql";
+
+ private final DateTime now = DateTime.now(UTC);
+ private final FakeClock clock = new FakeClock(now);
+ private final String initScript;
+ private JdbcDatabaseContainer database;
+ private EntityManagerFactory emf;
+ private JpaTransactionManager cachedTm;
+
+ private JpaTransactionManagerRule(String initScript) {
+ this.initScript = initScript;
+ }
+
+ /** Wraps {@link JpaTransactionManagerRule} in a {@link PostgreSQLContainer}. */
+ @Override
+ public Statement apply(Statement base, Description description) {
+ database = new PostgreSQLContainer().withInitScript(initScript);
+ return RuleChain.outerRule(database)
+ .around(JpaTransactionManagerRule.super::apply)
+ .apply(base, description);
+ }
+
+ @Override
+ public void before() {
+ emf =
+ PersistenceModule.create(
+ database.getJdbcUrl(),
+ database.getUsername(),
+ database.getPassword(),
+ PersistenceModule.providesDefaultDatabaseConfigs());
+ JpaTransactionManagerImpl txnManager = new JpaTransactionManagerImpl(emf, clock);
+ cachedTm = TransactionManagerFactory.jpaTm;
+ TransactionManagerFactory.jpaTm = txnManager;
+ }
+
+ @Override
+ public void after() {
+ TransactionManagerFactory.jpaTm = cachedTm;
+ if (emf != null) {
+ emf.close();
+ }
+ cachedTm = null;
+ }
+
+ /** Returns the {@link FakeClock} used by the underlying {@link JpaTransactionManagerImpl}. */
+ public FakeClock getTxnClock() {
+ return clock;
+ }
+
+ /** Builder for {@link JpaTransactionManagerRule}. */
+ public static class Builder {
+ private String initScript;
+
+ /**
+ * Sets the SQL script to be used to initialize the database. If not set,
+ * sql/schema/nomulus.golden.sql will be used.
+ */
+ public Builder withInitScript(String initScript) {
+ this.initScript = initScript;
+ return this;
+ }
+
+ /** Builds a {@link JpaTransactionManagerRule} instance. */
+ public JpaTransactionManagerRule build() {
+ if (initScript == null) {
+ initScript = SCHEMA_GOLDEN_SQL;
+ }
+ return new JpaTransactionManagerRule(initScript);
+ }
+ }
+}
diff --git a/core/src/test/java/google/registry/model/transaction/JpaTransactionManagerRuleTest.java b/core/src/test/java/google/registry/model/transaction/JpaTransactionManagerRuleTest.java
new file mode 100644
index 000000000..59d36b6b7
--- /dev/null
+++ b/core/src/test/java/google/registry/model/transaction/JpaTransactionManagerRuleTest.java
@@ -0,0 +1,59 @@
+// 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.model.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.testing.JUnitBackports.assertThrows;
+
+import java.util.List;
+import javax.persistence.PersistenceException;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** JUnit test for {@link JpaTransactionManagerRule} */
+@RunWith(JUnit4.class)
+public class JpaTransactionManagerRuleTest {
+
+ @Rule
+ public final JpaTransactionManagerRule jpaTmRule =
+ new JpaTransactionManagerRule.Builder().build();
+
+ @Test
+ public void verifiesRuleWorks() {
+ assertThrows(
+ PersistenceException.class,
+ () ->
+ jpaTm()
+ .transact(
+ () ->
+ jpaTm()
+ .getEntityManager()
+ .createNativeQuery("SELECT * FROM NoneExistentTable")
+ .getResultList()));
+ jpaTm()
+ .transact(
+ () -> {
+ List results =
+ jpaTm()
+ .getEntityManager()
+ .createNativeQuery("SELECT * FROM \"ClaimsList\"")
+ .getResultList();
+ assertThat(results).isEmpty();
+ });
+ }
+}
diff --git a/core/src/test/java/google/registry/model/transaction/JpaTransactionManagerTest.java b/core/src/test/java/google/registry/model/transaction/JpaTransactionManagerTest.java
deleted file mode 100644
index 449e17d9d..000000000
--- a/core/src/test/java/google/registry/model/transaction/JpaTransactionManagerTest.java
+++ /dev/null
@@ -1,211 +0,0 @@
-// 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();
- });
- }
-}