Use the DatabaseMigrationSchedule to determine which TM to use (#1233)

* Use the DatabaseMigrationSchedule to determine which TM to use

We still allow the "manual" specification of a particular transaction
manager, most useful in @DualDatabaseTest classes. If that isn't
specified, we examine the migration schedule to see which to return.

Notes:
- This requires that any test that sets the migration schedule clean up
after itself so that it won't affect future test runs of other classes
(because the migration schedule cache is static)
- One alternative would, instead of having a "test override" for the
transaction manager, be to examine the registry environment and only
override the transaction manager in the UNIT_TEST environment. This
doesn't work because there are many instances in which tests simulate
non-test environment.

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

@@ -16,17 +16,22 @@ package google.registry.persistence.transaction;
 
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Preconditions.checkState;
+import static org.joda.time.DateTimeZone.UTC;
 
 import com.google.appengine.api.utils.SystemProperty;
 import com.google.appengine.api.utils.SystemProperty.Environment.Value;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Suppliers;
 import google.registry.config.RegistryEnvironment;
+import google.registry.model.common.DatabaseMigrationStateSchedule;
+import google.registry.model.common.DatabaseMigrationStateSchedule.PrimaryDatabase;
 import google.registry.model.ofy.DatastoreTransactionManager;
 import google.registry.persistence.DaggerPersistenceComponent;
 import google.registry.tools.RegistryToolEnvironment;
 import google.registry.util.NonFinalForTesting;
+import java.util.Optional;
 import java.util.function.Supplier;
+import org.joda.time.DateTime;
 
 /** Factory class to create {@link TransactionManager} instance. */
 // TODO: Rename this to PersistenceFactory and move to persistence package.
@@ -34,7 +39,8 @@ public class TransactionManagerFactory {
 
   private static final DatastoreTransactionManager ofyTm = createTransactionManager();
 
-  @NonFinalForTesting private static TransactionManager tm = ofyTm;
+  /** Optional override to manually set the transaction manager per-test. */
+  private static Optional<TransactionManager> tmForTest = Optional.empty();
 
   /** Supplier for jpaTm so that it is initialized only once, upon first usage. */
   @NonFinalForTesting
@@ -70,9 +76,19 @@ public class TransactionManagerFactory {
     return SystemProperty.environment.value() == Value.Production;
   }
 
-  /** Returns {@link TransactionManager} instance. */
+  /**
+   * Returns the {@link TransactionManager} instance.
+   *
+   * <p>Returns the {@link JpaTransactionManager} or {@link DatastoreTransactionManager} based on
+   * the migration schedule or the manually specified per-test transaction manager.
+   */
   public static TransactionManager tm() {
-    return tm;
+    if (tmForTest.isPresent()) {
+      return tmForTest.get();
+    }
+    PrimaryDatabase primaryDatabase =
+        DatabaseMigrationStateSchedule.getValueAtTime(DateTime.now(UTC)).getPrimaryDatabase();
+    return primaryDatabase.equals(PrimaryDatabase.DATASTORE) ? ofyTm() : jpaTm();
   }
 
   /**
@@ -111,9 +127,18 @@ public class TransactionManagerFactory {
     jpaTm = Suppliers.memoize(jpaTmSupplier::get);
   }
 
-  /** Sets the return of {@link #tm()} to the given instance of {@link TransactionManager}. */
+  /**
+   * Sets the return of {@link #tm()} to the given instance of {@link TransactionManager}.
+   *
+   * <p>Used when overriding the per-test transaction manager for dual-database tests.
+   */
   @VisibleForTesting
-  public static void setTm(TransactionManager newTm) {
+  public static void setTmForTest(TransactionManager newTm) {
-    tm = newTm;
+    tmForTest = Optional.of(newTm);
+  }
+
+  /** Resets the overridden transaction manager post-test. */
+  public static void removeTmOverrideForTest() {
+    tmForTest = Optional.empty();
   }
 }

core/src/test/java/google/registry/backup/ReplayCommitLogsToSqlActionTest.java

@@ -74,6 +74,7 @@ import google.registry.util.RequestStatusChecker;
 import java.io.IOException;
 import org.joda.time.DateTime;
 import org.joda.time.Duration;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -144,6 +145,11 @@ public class ReplayCommitLogsToSqlActionTest {
     TestObject.beforeSqlDeleteCallCount = 0;
   }
 
+  @AfterEach
+  void afterEach() {
+    ofyTm().transact(() -> DatabaseMigrationStateSchedule.set(DEFAULT_TRANSITION_MAP.toValueMap()));
+  }
+
   @Test
   void testReplay_multipleDiffFiles() throws Exception {
     jpaTm()

core/src/test/java/google/registry/beam/rde/RdePipelineTest.java

@@ -23,7 +23,7 @@ import static google.registry.model.common.Cursor.CursorType.RDE_STAGING;
 import static google.registry.model.rde.RdeMode.FULL;
 import static google.registry.model.rde.RdeMode.THIN;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
-import static google.registry.persistence.transaction.TransactionManagerFactory.setTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.setTmForTest;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.rde.RdeResourceType.CONTACT;
 import static google.registry.rde.RdeResourceType.DOMAIN;
@@ -122,7 +122,7 @@ public class RdePipelineTest {
   @BeforeEach
   void beforeEach() throws Exception {
     originalTm = tm();
-    setTm(jpaTm());
+    setTmForTest(jpaTm());
     loadInitialData();
 
     // Two real registrars have been created by loadInitialData(), named "New Registrar" and "The
@@ -169,7 +169,7 @@ public class RdePipelineTest {
 
   @AfterEach
   void afterEach() {
-    setTm(originalTm);
+    setTmForTest(originalTm);
   }
 
   @Test

core/src/test/java/google/registry/beam/spec11/Spec11PipelineTest.java

@@ -18,7 +18,7 @@ import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.ImmutableObjectSubject.immutableObjectCorrespondence;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
-import static google.registry.persistence.transaction.TransactionManagerFactory.setTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.setTmForTest;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.AppEngineExtension.makeRegistrar1;
 import static google.registry.testing.DatabaseHelper.createTld;
@@ -231,7 +231,7 @@ class Spec11PipelineTest {
 
   private void setupCloudSql() {
     TransactionManager originalTm = tm();
-    setTm(jpaTm());
+    setTmForTest(jpaTm());
     persistNewRegistrar("TheRegistrar");
     persistNewRegistrar("NewRegistrar");
     Registrar registrar1 =
@@ -271,7 +271,7 @@ class Spec11PipelineTest {
     persistResource(createDomain("no-email.com", "2A4BA9BBC-COM", registrar2, contact2));
     persistResource(
         createDomain("anti-anti-anti-virus.dev", "555666888-DEV", registrar3, contact3));
-    setTm(originalTm);
+    setTmForTest(originalTm);
   }
 
   private void verifySaveToGcs() throws Exception {

core/src/test/java/google/registry/model/common/DatabaseMigrationStateScheduleTest.java

@@ -15,6 +15,7 @@
 package google.registry.model.common;
 
 import static com.google.common.truth.Truth.assertThat;
+import static google.registry.model.common.DatabaseMigrationStateSchedule.DEFAULT_TRANSITION_MAP;
 import static google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState.DATASTORE_ONLY;
 import static google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState.DATASTORE_PRIMARY;
 import static google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState.DATASTORE_PRIMARY_READ_ONLY;
@@ -22,6 +23,7 @@ import static google.registry.model.common.DatabaseMigrationStateSchedule.Migrat
 import static google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState.SQL_PRIMARY;
 import static google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState.SQL_PRIMARY_READ_ONLY;
 import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
 import static org.junit.Assert.assertThrows;
 
@@ -30,6 +32,7 @@ import google.registry.model.EntityTestCase;
 import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState;
 import org.joda.time.DateTime;
 import org.joda.time.Duration;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
@@ -40,6 +43,11 @@ public class DatabaseMigrationStateScheduleTest extends EntityTestCase {
     fakeClock.setAutoIncrementByOneMilli();
   }
 
+  @AfterEach
+  void afterEach() {
+    ofyTm().transact(() -> DatabaseMigrationStateSchedule.set(DEFAULT_TRANSITION_MAP.toValueMap()));
+  }
+
   @Test
   void testEmpty_returnsDatastoreOnlyMap() {
     assertThat(DatabaseMigrationStateSchedule.getUncached())
@@ -134,6 +142,15 @@ public class DatabaseMigrationStateScheduleTest extends EntityTestCase {
         .isEqualTo("Must be called in a transaction");
   }
 
+  @Test
+  void testSuccess_factoryUsesSchedule() {
+    assertThat(tm().isOfy()).isTrue();
+    // set the schedule to have converted to SQL_PRIMARY in the past
+    fakeClock.setTo(START_OF_TIME.plusDays(1));
+    runValidTransition(DATASTORE_PRIMARY_READ_ONLY, SQL_PRIMARY);
+    assertThat(tm().isOfy()).isFalse();
+  }
+
   private void runValidTransition(MigrationState from, MigrationState to) {
     ImmutableSortedMap<DateTime, MigrationState> transitions =
         createMapEndingWithTransition(from, to);

core/src/test/java/google/registry/schema/replay/ReplicateToDatastoreActionTest.java

@@ -96,6 +96,13 @@ public class ReplicateToDatastoreActionTest {
   public void tearDown() {
     fakeClock.disableAutoIncrement();
     RegistryConfig.overrideCloudSqlReplicateTransactions(false);
+    ofyTm()
+        .transact(
+            () ->
+                ofyTm()
+                    .loadSingleton(DatabaseMigrationStateSchedule.class)
+                    .ifPresent(ofyTm()::delete));
+    DatabaseMigrationStateSchedule.CACHE.invalidateAll();
   }
 
   @RetryingTest(4)

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

@@ -154,15 +154,13 @@ class DualDatabaseTestInvocationContextProvider implements TestTemplateInvocatio
       context.getStore(NAMESPACE).put(ORIGINAL_TM_KEY, tm());
       DatabaseType databaseType =
           (DatabaseType) context.getStore(NAMESPACE).get(INJECTED_TM_SUPPLIER_KEY);
-      TransactionManagerFactory.setTm(databaseType.getTm());
+      TransactionManagerFactory.setTmForTest(databaseType.getTm());
     }
   }
 
   static void restoreTmAfterDualDatabaseTest(ExtensionContext context) {
     if (isDualDatabaseTest(context)) {
-      TransactionManager original =
+      TransactionManagerFactory.removeTmOverrideForTest();
-          (TransactionManager) context.getStore(NAMESPACE).get(ORIGINAL_TM_KEY);
-      TransactionManagerFactory.setTm(original);
     }
   }
 

core/src/test/java/google/registry/tools/GetDatabaseMigrationStateCommandTest.java

@@ -24,15 +24,15 @@ import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationStat
 import google.registry.testing.DualDatabaseTest;
 import google.registry.testing.TestOfyAndSql;
 import org.joda.time.DateTime;
-import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.AfterEach;
 
 /** Tests for {@link GetDatabaseMigrationStateCommand}. */
 @DualDatabaseTest
 public class GetDatabaseMigrationStateCommandTest
     extends CommandTestCase<GetDatabaseMigrationStateCommand> {
 
-  @BeforeEach
+  @AfterEach
-  void beforeEach() {
+  void afterEach() {
     ofyTm().transact(() -> DatabaseMigrationStateSchedule.set(DEFAULT_TRANSITION_MAP.toValueMap()));
   }
 

core/src/test/java/google/registry/tools/SetDatabaseMigrationStateCommandTest.java

@@ -28,6 +28,7 @@ import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationStat
 import google.registry.testing.DualDatabaseTest;
 import google.registry.testing.TestOfyAndSql;
 import org.joda.time.DateTime;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 
 /** Tests for {@link SetDatabaseMigrationStateCommand}. */
@@ -47,6 +48,11 @@ public class SetDatabaseMigrationStateCommandTest
     DatabaseMigrationStateSchedule.CACHE.invalidateAll();
   }
 
+  @AfterEach
+  void afterEach() {
+    ofyTm().transact(() -> DatabaseMigrationStateSchedule.set(DEFAULT_TRANSITION_MAP.toValueMap()));
+  }
+
   @TestOfyAndSql
   void testSuccess_setsBasicSchedule() throws Exception {
     assertThat(DatabaseMigrationStateSchedule.get()).isEqualTo(DEFAULT_TRANSITION_MAP);