Use the DB migration schedule for SQL->DS replay (#1242)

This is instead of the current configuration parameter.

In addition, this adds some helpers to DatabaseHelper to make the
transitions easier, since we more frequently need to alter + reset the
schedule.

core/src/main/java/google/registry/config/RegistryConfig.java

@@ -1494,21 +1494,6 @@ public final class RegistryConfig {
     return CONFIG_SETTINGS.get().hibernate.hikariIdleTimeout;
   }
 
-  /**
-   * Returns whether to replicate cloud SQL transactions to datastore.
-   *
-   * <p>If true, all cloud SQL transactions will be persisted as TransactionEntity objects in the
-   * Transaction table and replayed against datastore in a cron job.
-   */
-  public static boolean getCloudSqlReplicateTransactions() {
-    return CONFIG_SETTINGS.get().cloudSql.replicateTransactions;
-  }
-
-  @VisibleForTesting
-  public static void overrideCloudSqlReplicateTransactions(boolean replicateTransactions) {
-    CONFIG_SETTINGS.get().cloudSql.replicateTransactions = replicateTransactions;
-  }
-
   /** Returns the roid suffix to be used for the roids of all contacts and hosts. */
   public static String getContactAndHostRoidSuffix() {
     return CONFIG_SETTINGS.get().registryPolicy.contactAndHostRoidSuffix;

core/src/main/java/google/registry/config/RegistryConfigSettings.java

@@ -126,7 +126,6 @@ public class RegistryConfigSettings {
     // TODO(05012021): remove username field after it is removed from all yaml files.
     public String username;
     public String instanceConnectionName;
-    public boolean replicateTransactions;
   }
 
   /** Configuration for Apache Beam (Cloud Dataflow). */

core/src/main/java/google/registry/config/files/default-config.yaml

@@ -227,9 +227,6 @@ cloudSql:
   jdbcUrl: jdbc:postgresql://localhost
   # This name is used by Cloud SQL when connecting to the database.
   instanceConnectionName: project-id:region:instance-id
-  # Set this to true to replicate cloud SQL transactions to datastore in the
-  # background.
-  replicateTransactions: false
 
 cloudDns:
   # Set both properties to null in Production.

core/src/main/java/google/registry/model/common/DatabaseMigrationStateSchedule.java

@@ -225,7 +225,7 @@ public class DatabaseMigrationStateSchedule extends CrossTldSingleton
   @VisibleForTesting
   static TimedTransitionProperty<MigrationState, MigrationStateTransition> getUncached() {
     return ofyTm()
-        .transact(
+        .transactNew(
             () ->
                 ofyTm()
                     .loadSingleton(DatabaseMigrationStateSchedule.class)

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

@@ -30,8 +30,9 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Streams;
 import com.google.common.flogger.FluentLogger;
-import google.registry.config.RegistryConfig;
 import google.registry.model.ImmutableObject;
+import google.registry.model.common.DatabaseMigrationStateSchedule;
+import google.registry.model.common.DatabaseMigrationStateSchedule.ReplayDirection;
 import google.registry.model.index.EppResourceIndex;
 import google.registry.model.index.ForeignKeyIndex.ForeignKeyContactIndex;
 import google.registry.model.index.ForeignKeyIndex.ForeignKeyDomainIndex;
@@ -103,6 +104,10 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   // synchronously.
   private final ThreadLocal<TransactionInfo> transactionInfo =
       ThreadLocal.withInitial(TransactionInfo::new);
+  // If this value is present, use it to determine whether or not to replay SQL transactions to
+  // Datastore, rather than using the schedule stored in Datastore.
+  private static ThreadLocal<Optional<Boolean>> replaySqlToDatastoreOverrideForTest =
+      ThreadLocal.withInitial(Optional::empty);
 
   public JpaTransactionManagerImpl(EntityManagerFactory emf, Clock clock) {
     this.emf = emf;
@@ -736,6 +741,16 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     return entity;
   }
 
+  /** Sets the override to always/never replay SQL transactions to Datastore. */
+  public static void setReplaySqlToDatastoreOverrideForTest(boolean replaySqlToDs) {
+    replaySqlToDatastoreOverrideForTest.set(Optional.of(replaySqlToDs));
+  }
+
+  /** Removes the replay-SQL-to-Datastore override; the migration schedule will then be used. */
+  public static void removeReplaySqlToDsOverrideForTest() {
+    replaySqlToDatastoreOverrideForTest.set(Optional.empty());
+  }
+
   private static class TransactionInfo {
     EntityManager entityManager;
     boolean inTransaction = false;
@@ -755,7 +770,14 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
       checkArgumentNotNull(clock);
       inTransaction = true;
       transactionTime = clock.nowUtc();
-      if (withBackup && RegistryConfig.getCloudSqlReplicateTransactions()) {
+      if (withBackup
+          && replaySqlToDatastoreOverrideForTest
+              .get()
+              .orElseGet(
+                  () ->
+                      DatabaseMigrationStateSchedule.getValueAtTime(transactionTime)
+                          .getReplayDirection()
+                          .equals(ReplayDirection.SQL_TO_DATASTORE))) {
         contentsBuilder = new Transaction.Builder();
       }
     }

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

@@ -67,6 +67,7 @@ import google.registry.persistence.transaction.TransactionManagerFactory;
 import google.registry.schema.replay.SqlReplayCheckpoint;
 import google.registry.schema.tld.PremiumEntry;
 import google.registry.testing.AppEngineExtension;
+import google.registry.testing.DatabaseHelper;
 import google.registry.testing.FakeClock;
 import google.registry.testing.FakeResponse;
 import google.registry.testing.TestObject;
@@ -149,7 +150,7 @@ public class ReplayCommitLogsToSqlActionTest {
 
   @AfterEach
   void afterEach() {
-    ofyTm().transact(() -> DatabaseMigrationStateSchedule.set(DEFAULT_TRANSITION_MAP.toValueMap()));
+    DatabaseHelper.removeDatabaseMigrationSchedule();
   }
 
   @Test

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

@@ -15,7 +15,6 @@
 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;
@@ -30,6 +29,7 @@ import static org.junit.Assert.assertThrows;
 import com.google.common.collect.ImmutableSortedMap;
 import google.registry.model.EntityTestCase;
 import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState;
+import google.registry.testing.DatabaseHelper;
 import org.joda.time.DateTime;
 import org.joda.time.Duration;
 import org.junit.jupiter.api.AfterEach;
@@ -45,7 +45,7 @@ public class DatabaseMigrationStateScheduleTest extends EntityTestCase {
 
   @AfterEach
   void afterEach() {
-    ofyTm().transact(() -> DatabaseMigrationStateSchedule.set(DEFAULT_TRANSITION_MAP.toValueMap()));
+    DatabaseHelper.removeDatabaseMigrationSchedule();
   }
 
   @Test

core/src/test/java/google/registry/persistence/transaction/JpaTransactionManagerExtension.java

@@ -202,11 +202,13 @@ abstract class JpaTransactionManagerExtension implements BeforeEachCallback, Aft
     JpaTransactionManagerImpl txnManager = new JpaTransactionManagerImpl(emf, clock);
     cachedTm = TransactionManagerFactory.jpaTm();
     TransactionManagerFactory.setJpaTm(Suppliers.ofInstance(txnManager));
+    JpaTransactionManagerImpl.setReplaySqlToDatastoreOverrideForTest(false);
   }
 
   @Override
   public void afterEach(ExtensionContext context) {
     TransactionManagerFactory.setJpaTm(Suppliers.ofInstance(cachedTm));
+    JpaTransactionManagerImpl.removeReplaySqlToDsOverrideForTest();
     cachedTm = null;
   }
 

core/src/test/java/google/registry/persistence/transaction/TransactionTest.java

@@ -22,36 +22,53 @@ import static org.junit.jupiter.api.Assertions.assertThrows;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Entity;
 import com.googlecode.objectify.annotation.Id;
-import google.registry.config.RegistryConfig;
 import google.registry.model.ImmutableObject;
+import google.registry.model.ofy.Ofy;
 import google.registry.persistence.VKey;
 import google.registry.testing.AppEngineExtension;
+import google.registry.testing.DatabaseHelper;
+import google.registry.testing.FakeClock;
+import google.registry.testing.InjectExtension;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.ObjectOutputStream;
 import java.io.StreamCorruptedException;
+import org.joda.time.DateTime;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
 class TransactionTest {
 
+  private final FakeClock fakeClock = new FakeClock(DateTime.parse("2000-01-01TZ"));
+
   @RegisterExtension
   final AppEngineExtension appEngine =
       AppEngineExtension.builder()
           .withDatastoreAndCloudSql()
+          .withClock(fakeClock)
           .withOfyTestEntities(TestEntity.class)
           .withJpaUnitTestEntities(TestEntity.class)
           .build();
 
+  @RegisterExtension public final InjectExtension inject = new InjectExtension();
+
   private TestEntity fooEntity, barEntity;
 
   @BeforeEach
   void beforeEach() {
+    JpaTransactionManagerImpl.removeReplaySqlToDsOverrideForTest();
+    inject.setStaticField(Ofy.class, "clock", fakeClock);
     fooEntity = new TestEntity("foo");
     barEntity = new TestEntity("bar");
   }
 
+  @AfterEach
+  void afterEach() {
+    DatabaseHelper.removeDatabaseMigrationSchedule();
+  }
+
   @Test
   void testTransactionReplay() {
     Transaction txn = new Transaction.Builder().addUpdate(fooEntity).addUpdate(barEntity).build();
@@ -102,8 +119,7 @@ class TransactionTest {
 
   @Test
   void testTransactionSerialization() throws IOException {
-    RegistryConfig.overrideCloudSqlReplicateTransactions(true);
+    DatabaseHelper.setMigrationScheduleToSqlPrimary(fakeClock);
-    try {
     jpaTm()
         .transact(
             () -> {
@@ -125,9 +141,6 @@ class TransactionTest {
       assertThat(
               jpaTm().transact(() -> jpaTm().exists(VKey.createSql(TransactionEntity.class, 2L))))
           .isFalse();
-    } finally {
-      RegistryConfig.overrideCloudSqlReplicateTransactions(false);
-    }
   }
 
   @Test

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

@@ -26,15 +26,16 @@ import com.google.common.testing.TestLogHandler;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Entity;
 import com.googlecode.objectify.annotation.Id;
-import google.registry.config.RegistryConfig;
 import google.registry.model.ImmutableObject;
 import google.registry.model.common.DatabaseMigrationStateSchedule;
 import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState;
 import google.registry.model.ofy.CommitLogBucket;
 import google.registry.model.ofy.Ofy;
 import google.registry.persistence.VKey;
+import google.registry.persistence.transaction.JpaTransactionManagerImpl;
 import google.registry.persistence.transaction.TransactionEntity;
 import google.registry.testing.AppEngineExtension;
+import google.registry.testing.DatabaseHelper;
 import google.registry.testing.FakeClock;
 import google.registry.testing.InjectExtension;
 import java.util.List;
@@ -68,41 +69,21 @@ public class ReplicateToDatastoreActionTest {
 
   @BeforeEach
   public void setUp() {
+    JpaTransactionManagerImpl.removeReplaySqlToDsOverrideForTest();
     injectExtension.setStaticField(Ofy.class, "clock", fakeClock);
     // Use a single bucket to expose timestamp inversion problems.
     injectExtension.setStaticField(
         CommitLogBucket.class, "bucketIdSupplier", Suppliers.ofInstance(1));
     fakeClock.setAutoIncrementByOneMilli();
-    RegistryConfig.overrideCloudSqlReplicateTransactions(true);
+    DatabaseHelper.setMigrationScheduleToSqlPrimary(fakeClock);
     Logger.getLogger(ReplicateToDatastoreAction.class.getCanonicalName()).addHandler(logHandler);
-    DateTime now = fakeClock.nowUtc();
-    ofyTm()
-        .transact(
-            () ->
-                DatabaseMigrationStateSchedule.set(
-                    ImmutableSortedMap.of(
-                        START_OF_TIME,
-                        MigrationState.DATASTORE_ONLY,
-                        now,
-                        MigrationState.DATASTORE_PRIMARY,
-                        now.plusHours(1),
-                        MigrationState.DATASTORE_PRIMARY_READ_ONLY,
-                        now.plusHours(2),
-                        MigrationState.SQL_PRIMARY)));
     fakeClock.advanceBy(Duration.standardDays(1));
   }
 
   @AfterEach
   public void tearDown() {
+    DatabaseHelper.removeDatabaseMigrationSchedule();
     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/DatabaseHelper.java

@@ -35,6 +35,7 @@ import static google.registry.model.ResourceTransferUtils.createTransferResponse
 import static google.registry.model.ofy.ObjectifyService.auditedOfy;
 import static google.registry.model.registry.Registry.TldState.GENERAL_AVAILABILITY;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.persistence.transaction.TransactionManagerUtil.ofyTmOrDoNothing;
 import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
@@ -71,6 +72,8 @@ import google.registry.model.EppResourceUtils;
 import google.registry.model.billing.BillingEvent;
 import google.registry.model.billing.BillingEvent.Flag;
 import google.registry.model.billing.BillingEvent.Reason;
+import google.registry.model.common.DatabaseMigrationStateSchedule;
+import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState;
 import google.registry.model.contact.ContactAuthInfo;
 import google.registry.model.contact.ContactHistory;
 import google.registry.model.contact.ContactResource;
@@ -123,6 +126,7 @@ import org.joda.money.CurrencyUnit;
 import org.joda.money.Money;
 import org.joda.time.DateTime;
 import org.joda.time.DateTimeZone;
+import org.joda.time.Duration;
 
 /** Static utils for setting up test resources. */
 public class DatabaseHelper {
@@ -1360,5 +1364,45 @@ public class DatabaseHelper {
     return entity;
   }
 
+  /**
+   * Sets a SQL_PRIMARY state on the {@link DatabaseMigrationStateSchedule}.
+   *
+   * <p>In order to allow for tests to manipulate the clock how they need, we start the transitions
+   * one millisecond after the clock's current time (in case the clock's current value is
+   * START_OF_TIME). We then advance the clock one second so that we're in the SQL_PRIMARY phase.
+   *
+   * <p>We must use the current time, otherwise the setting of the migration state will fail due to
+   * an invalid transition.
+   */
+  public static void setMigrationScheduleToSqlPrimary(FakeClock fakeClock) {
+    DateTime now = fakeClock.nowUtc();
+    ofyTm()
+        .transact(
+            () ->
+                DatabaseMigrationStateSchedule.set(
+                    ImmutableSortedMap.of(
+                        START_OF_TIME,
+                        MigrationState.DATASTORE_ONLY,
+                        now.plusMillis(1),
+                        MigrationState.DATASTORE_PRIMARY,
+                        now.plusMillis(2),
+                        MigrationState.DATASTORE_PRIMARY_READ_ONLY,
+                        now.plusMillis(3),
+                        MigrationState.SQL_PRIMARY)));
+    fakeClock.advanceBy(Duration.standardSeconds(1));
+  }
+
+  /** Removes the database migration schedule, in essence transitioning to DATASTORE_ONLY. */
+  public static void removeDatabaseMigrationSchedule() {
+    // use the raw calls because going SQL_PRIMARY -> DATASTORE_ONLY is not valid
+    ofyTm()
+        .transact(
+            () ->
+                ofyTm()
+                    .loadSingleton(DatabaseMigrationStateSchedule.class)
+                    .ifPresent(ofyTm()::delete));
+    DatabaseMigrationStateSchedule.CACHE.invalidateAll();
+  }
+
   private DatabaseHelper() {}
 }

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

@@ -22,12 +22,12 @@ import com.google.common.base.Suppliers;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
-import google.registry.config.RegistryConfig;
 import google.registry.model.ImmutableObject;
 import google.registry.model.ofy.CommitLogBucket;
 import google.registry.model.ofy.ReplayQueue;
 import google.registry.model.ofy.TransactionInfo;
 import google.registry.persistence.VKey;
+import google.registry.persistence.transaction.JpaTransactionManagerImpl;
 import google.registry.persistence.transaction.TransactionEntity;
 import google.registry.schema.replay.DatastoreEntity;
 import google.registry.schema.replay.ReplicateToDatastoreAction;
@@ -92,7 +92,7 @@ public class ReplayExtension implements BeforeEachCallback, AfterEachCallback {
     // transaction manager gets injected.
     inOfyContext = DualDatabaseTestInvocationContextProvider.inOfyContext(context);
     if (sqlToDsReplicator != null && !inOfyContext) {
-      RegistryConfig.overrideCloudSqlReplicateTransactions(true);
+      JpaTransactionManagerImpl.setReplaySqlToDatastoreOverrideForTest(true);
     }
 
     context.getStore(ExtensionContext.Namespace.GLOBAL).put(ReplayExtension.class, this);
@@ -105,7 +105,7 @@ public class ReplayExtension implements BeforeEachCallback, AfterEachCallback {
     replay();
     injectExtension.afterEach(context);
     if (sqlToDsReplicator != null) {
-      RegistryConfig.overrideCloudSqlReplicateTransactions(false);
+      JpaTransactionManagerImpl.setReplaySqlToDatastoreOverrideForTest(false);
     }
   }
 

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

@@ -21,6 +21,7 @@ import static google.registry.util.DateTimeUtils.START_OF_TIME;
 import com.google.common.collect.ImmutableSortedMap;
 import google.registry.model.common.DatabaseMigrationStateSchedule;
 import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState;
+import google.registry.testing.DatabaseHelper;
 import google.registry.testing.DualDatabaseTest;
 import google.registry.testing.TestOfyAndSql;
 import org.joda.time.DateTime;
@@ -33,7 +34,7 @@ public class GetDatabaseMigrationStateCommandTest
 
   @AfterEach
   void afterEach() {
-    ofyTm().transact(() -> DatabaseMigrationStateSchedule.set(DEFAULT_TRANSITION_MAP.toValueMap()));
+    DatabaseHelper.removeDatabaseMigrationSchedule();
   }
 
   @TestOfyAndSql

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

@@ -25,32 +25,20 @@ import com.beust.jcommander.ParameterException;
 import com.google.common.collect.ImmutableSortedMap;
 import google.registry.model.common.DatabaseMigrationStateSchedule;
 import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState;
+import google.registry.testing.DatabaseHelper;
 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}. */
 @DualDatabaseTest
 public class SetDatabaseMigrationStateCommandTest
     extends CommandTestCase<SetDatabaseMigrationStateCommand> {
 
-  @BeforeEach
-  void beforeEach() {
-    // clear out any static state that may have been persisted
-    ofyTm()
-        .transact(
-            () ->
-                ofyTm()
-                    .loadSingleton(DatabaseMigrationStateSchedule.class)
-                    .ifPresent(ofyTm()::delete));
-    DatabaseMigrationStateSchedule.CACHE.invalidateAll();
-  }
-
   @AfterEach
   void afterEach() {
-    ofyTm().transact(() -> DatabaseMigrationStateSchedule.set(DEFAULT_TRANSITION_MAP.toValueMap()));
+    DatabaseHelper.removeDatabaseMigrationSchedule();
   }
 
   @TestOfyAndSql