Use DB migration state to determine running async replay DS->SQL (#1175)

* Use DB migration state to determine running async replay DS->SQL

core/src/main/java/google/registry/backup/ReplayCommitLogsToSqlAction.java

@@ -27,7 +27,9 @@ import com.google.appengine.tools.cloudstorage.GcsFileMetadata;
 import com.google.appengine.tools.cloudstorage.GcsService;
 import com.google.common.collect.ImmutableList;
 import com.google.common.flogger.FluentLogger;
-import google.registry.config.RegistryConfig;
+import google.registry.model.common.DatabaseMigrationStateSchedule;
+import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState;
+import google.registry.model.common.DatabaseMigrationStateSchedule.ReplayDirection;
 import google.registry.model.server.Lock;
 import google.registry.model.translators.VKeyTranslatorFactory;
 import google.registry.persistence.VKey;
@@ -39,6 +41,7 @@ import google.registry.schema.replay.DatastoreOnlyEntity;
 import google.registry.schema.replay.NonReplicatedEntity;
 import google.registry.schema.replay.ReplaySpecializer;
 import google.registry.schema.replay.SqlReplayCheckpoint;
+import google.registry.util.Clock;
 import google.registry.util.RequestStatusChecker;
 import java.io.IOException;
 import java.io.InputStream;
@@ -69,15 +72,19 @@ public class ReplayCommitLogsToSqlAction implements Runnable {
   @Inject Response response;
   @Inject RequestStatusChecker requestStatusChecker;
   @Inject GcsDiffFileLister diffLister;
+  @Inject Clock clock;
 
   @Inject
   ReplayCommitLogsToSqlAction() {}
 
   @Override
   public void run() {
-    if (!RegistryConfig.getCloudSqlReplayCommitLogs()) {
-      String message = "ReplayCommitLogsToSqlAction was called but disabled in the config.";
-      logger.atWarning().log(message);
+    MigrationState state = DatabaseMigrationStateSchedule.getValueAtTime(clock.nowUtc());
+    if (!state.getReplayDirection().equals(ReplayDirection.DATASTORE_TO_SQL)) {
+      String message =
+          String.format(
+              "Skipping ReplayCommitLogsToSqlAction because we are in migration phase %s.", state);
+      logger.atInfo().log(message);
       // App Engine will retry on any non-2xx status code, which we don't want in this case.
       response.setStatus(SC_NO_CONTENT);
       response.setPayload(message);

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

@@ -1502,22 +1502,6 @@ public final class RegistryConfig {
     CONFIG_SETTINGS.get().cloudSql.replicateTransactions = replicateTransactions;
   }
 
-  /**
-   * Returns whether or not to replay commit logs to the SQL database after export to GCS.
-   *
-   * <p>If true, we will trigger the {@link google.registry.backup.ReplayCommitLogsToSqlAction}
-   * after the {@link google.registry.backup.ExportCommitLogDiffAction} to load the commit logs and
-   * replay them to SQL.
-   */
-  public static boolean getCloudSqlReplayCommitLogs() {
-    return CONFIG_SETTINGS.get().cloudSql.replayCommitLogs;
-  }
-
-  @VisibleForTesting
-  public static void overrideCloudSqlReplayCommitLogs(boolean replayCommitLogs) {
-    CONFIG_SETTINGS.get().cloudSql.replayCommitLogs = replayCommitLogs;
-  }
-
   /** 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

@@ -127,7 +127,6 @@ public class RegistryConfigSettings {
     public String username;
     public String instanceConnectionName;
     public boolean replicateTransactions;
-    public boolean replayCommitLogs;
   }
 
   /** Configuration for Apache Beam (Cloud Dataflow). */

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

@@ -230,8 +230,6 @@ cloudSql:
   # Set this to true to replicate cloud SQL transactions to datastore in the
   # background.
   replicateTransactions: false
-  # Set this to true to enable replay of commit logs to SQL
-  replayCommitLogs: false
 
 cloudDns:
   # Set both properties to null in Production.

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

@@ -50,20 +50,27 @@ public class DatabaseMigrationStateSchedule extends CrossTldSingleton
     DATASTORE
   }
 
+  public enum ReplayDirection {
+    NO_REPLAY,
+    DATASTORE_TO_SQL,
+    SQL_TO_DATASTORE
+  }
+
   /**
    * The current phase of the migration plus information about which database to use and whether or
    * not the phase is read-only.
    */
   public enum MigrationState {
-    DATASTORE_ONLY(PrimaryDatabase.DATASTORE, false),
-    DATASTORE_PRIMARY(PrimaryDatabase.DATASTORE, false),
-    DATASTORE_PRIMARY_READ_ONLY(PrimaryDatabase.DATASTORE, true),
-    SQL_PRIMARY_READ_ONLY(PrimaryDatabase.CLOUD_SQL, true),
-    SQL_PRIMARY(PrimaryDatabase.CLOUD_SQL, false),
-    SQL_ONLY(PrimaryDatabase.CLOUD_SQL, false);
+    DATASTORE_ONLY(PrimaryDatabase.DATASTORE, false, ReplayDirection.NO_REPLAY),
+    DATASTORE_PRIMARY(PrimaryDatabase.DATASTORE, false, ReplayDirection.DATASTORE_TO_SQL),
+    DATASTORE_PRIMARY_READ_ONLY(PrimaryDatabase.DATASTORE, true, ReplayDirection.DATASTORE_TO_SQL),
+    SQL_PRIMARY_READ_ONLY(PrimaryDatabase.CLOUD_SQL, true, ReplayDirection.SQL_TO_DATASTORE),
+    SQL_PRIMARY(PrimaryDatabase.CLOUD_SQL, false, ReplayDirection.SQL_TO_DATASTORE),
+    SQL_ONLY(PrimaryDatabase.CLOUD_SQL, false, ReplayDirection.NO_REPLAY);
 
     private final PrimaryDatabase primaryDatabase;
     private final boolean isReadOnly;
+    private final ReplayDirection replayDirection;
 
     public PrimaryDatabase getPrimaryDatabase() {
       return primaryDatabase;
@@ -73,9 +80,15 @@ public class DatabaseMigrationStateSchedule extends CrossTldSingleton
       return isReadOnly;
     }
 
-    MigrationState(PrimaryDatabase primaryDatabase, boolean isReadOnly) {
+    public ReplayDirection getReplayDirection() {
+      return replayDirection;
+    }
+
+    MigrationState(
+        PrimaryDatabase primaryDatabase, boolean isReadOnly, ReplayDirection replayDirection) {
       this.primaryDatabase = primaryDatabase;
       this.isReadOnly = isReadOnly;
+      this.replayDirection = replayDirection;
     }
   }
 
@@ -201,6 +214,11 @@ public class DatabaseMigrationStateSchedule extends CrossTldSingleton
     return CACHE.getUnchecked(DatabaseMigrationStateSchedule.class);
   }
 
+  /** Returns the database migration status at the given time. */
+  public static MigrationState getValueAtTime(DateTime dateTime) {
+    return get().getValueAtTime(dateTime);
+  }
+
   /** Loads the currently-set migration schedule from Datastore, or the default if none exists. */
   @VisibleForTesting
   static TimedTransitionProperty<MigrationState, MigrationStateTransition> getUncached() {

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

@@ -21,13 +21,16 @@ import static google.registry.backup.RestoreCommitLogsActionTest.GCS_BUCKET;
 import static google.registry.backup.RestoreCommitLogsActionTest.createCheckpoint;
 import static google.registry.backup.RestoreCommitLogsActionTest.saveDiffFile;
 import static google.registry.backup.RestoreCommitLogsActionTest.saveDiffFileNotToRestore;
+import static google.registry.model.common.DatabaseMigrationStateSchedule.DEFAULT_TRANSITION_MAP;
 import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
 import static google.registry.model.ofy.CommitLogBucket.getBucketKey;
 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.testing.DatabaseHelper.createTld;
 import static google.registry.testing.DatabaseHelper.newDomainBase;
 import static google.registry.testing.DatabaseHelper.persistActiveContact;
+import static google.registry.util.DateTimeUtils.START_OF_TIME;
 import static javax.servlet.http.HttpServletResponse.SC_NO_CONTENT;
 import static javax.servlet.http.HttpServletResponse.SC_OK;
 import static org.mockito.ArgumentMatchers.any;
@@ -40,9 +43,11 @@ import com.google.appengine.tools.cloudstorage.GcsServiceFactory;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.truth.Truth8;
 import com.googlecode.objectify.Key;
-import google.registry.config.RegistryConfig;
+import google.registry.model.common.DatabaseMigrationStateSchedule;
+import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState;
 import google.registry.model.contact.ContactResource;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.GracePeriod;
@@ -122,11 +127,20 @@ public class ReplayCommitLogsToSqlActionTest {
     action.gcsService = gcsService;
     action.response = response;
     action.requestStatusChecker = requestStatusChecker;
+    action.clock = fakeClock;
     action.diffLister = new GcsDiffFileLister();
     action.diffLister.gcsService = gcsService;
     action.diffLister.gcsBucket = GCS_BUCKET;
     action.diffLister.executor = newDirectExecutorService();
-    RegistryConfig.overrideCloudSqlReplayCommitLogs(true);
+    ofyTm()
+        .transact(
+            () ->
+                DatabaseMigrationStateSchedule.set(
+                    ImmutableSortedMap.of(
+                        START_OF_TIME,
+                        MigrationState.DATASTORE_ONLY,
+                        START_OF_TIME.plusMinutes(1),
+                        MigrationState.DATASTORE_PRIMARY)));
     TestObject.beforeSqlSaveCallCount = 0;
     TestObject.beforeSqlDeleteCallCount = 0;
   }
@@ -329,7 +343,7 @@ public class ReplayCommitLogsToSqlActionTest {
     ContactResource contactWithEdit =
         contact.asBuilder().setEmailAddress("replay@example.tld").build();
     CommitLogMutation contactMutation =
-        tm().transact(() -> CommitLogMutation.create(manifestKey, contactWithEdit));
+        ofyTm().transact(() -> CommitLogMutation.create(manifestKey, contactWithEdit));
 
     jpaTm().transact(() -> SqlReplayCheckpoint.set(now.minusMinutes(1).minusMillis(1)));
 
@@ -421,11 +435,13 @@ public class ReplayCommitLogsToSqlActionTest {
 
   @Test
   void testFailure_notEnabled() {
-    RegistryConfig.overrideCloudSqlReplayCommitLogs(false);
+    ofyTm().transact(() -> DatabaseMigrationStateSchedule.set(DEFAULT_TRANSITION_MAP.toValueMap()));
     action.run();
     assertThat(response.getStatus()).isEqualTo(SC_NO_CONTENT);
     assertThat(response.getPayload())
-        .isEqualTo("ReplayCommitLogsToSqlAction was called but disabled in the config.");
+        .isEqualTo(
+            "Skipping ReplayCommitLogsToSqlAction because we are in migration phase"
+                + " DATASTORE_ONLY.");
   }
 
   @Test