From 5cb2a0a430846a798c6dbe5a566834bd08275a8f Mon Sep 17 00:00:00 2001
From: Ben McIlwain <mcilwain@google.com>
Date: Tue, 12 Jul 2022 16:49:16 -0400
Subject: [PATCH] Re-add database migration state commands (#1702)

* Re-add database migration state commands

These were removed in PR #1661, but we do still need them for the time being
until we complete the ID migration as well.
---
 .../GetDatabaseMigrationStateCommand.java     |  36 ++++
 .../google/registry/tools/RegistryTool.java   |   2 +
 .../SetDatabaseMigrationStateCommand.java     |  72 +++++++
 .../tools/params/TransitionListParameter.java |   9 +
 .../DatabaseMigrationStateScheduleTest.java   | 187 ++++++++++++++++++
 .../registry/testing/DatabaseHelper.java      | 102 ++++++++++
 .../GetDatabaseMigrationStateCommandTest.java |  66 +++++++
 .../SetDatabaseMigrationStateCommandTest.java | 184 +++++++++++++++++
 8 files changed, 658 insertions(+)
 create mode 100644 core/src/main/java/google/registry/tools/GetDatabaseMigrationStateCommand.java
 create mode 100644 core/src/main/java/google/registry/tools/SetDatabaseMigrationStateCommand.java
 create mode 100644 core/src/test/java/google/registry/model/common/DatabaseMigrationStateScheduleTest.java
 create mode 100644 core/src/test/java/google/registry/tools/GetDatabaseMigrationStateCommandTest.java
 create mode 100644 core/src/test/java/google/registry/tools/SetDatabaseMigrationStateCommandTest.java

diff --git a/core/src/main/java/google/registry/tools/GetDatabaseMigrationStateCommand.java b/core/src/main/java/google/registry/tools/GetDatabaseMigrationStateCommand.java
new file mode 100644
index 000000000..8a4c680f2
--- /dev/null
+++ b/core/src/main/java/google/registry/tools/GetDatabaseMigrationStateCommand.java
@@ -0,0 +1,36 @@
+// Copyright 2021 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.tools;
+
+import com.beust.jcommander.Parameters;
+import google.registry.model.annotations.DeleteAfterMigration;
+import google.registry.model.common.DatabaseMigrationStateSchedule;
+import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState;
+import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationStateTransition;
+import google.registry.model.common.TimedTransitionProperty;
+
+/** A command to check the current Registry 3.0 migration state of the database. */
+@DeleteAfterMigration
+@Parameters(separators = " =", commandDescription = "Check current Registry 3.0 migration state")
+public class GetDatabaseMigrationStateCommand implements CommandWithRemoteApi {
+
+  @Override
+  public void run() throws Exception {
+    TimedTransitionProperty<MigrationState, MigrationStateTransition> migrationSchedule =
+        DatabaseMigrationStateSchedule.get();
+    System.out.println(
+        String.format("Current migration schedule: %s", migrationSchedule.toValueMap()));
+  }
+}
diff --git a/core/src/main/java/google/registry/tools/RegistryTool.java b/core/src/main/java/google/registry/tools/RegistryTool.java
index 3a2d4fdab..b3f1f5092 100644
--- a/core/src/main/java/google/registry/tools/RegistryTool.java
+++ b/core/src/main/java/google/registry/tools/RegistryTool.java
@@ -65,6 +65,7 @@ public final class RegistryTool {
           .put("get_allocation_token", GetAllocationTokenCommand.class)
           .put("get_claims_list", GetClaimsListCommand.class)
           .put("get_contact", GetContactCommand.class)
+          .put("get_database_migration_state", GetDatabaseMigrationStateCommand.class)
           .put("get_domain", GetDomainCommand.class)
           .put("get_history_entries", GetHistoryEntriesCommand.class)
           .put("get_host", GetHostCommand.class)
@@ -96,6 +97,7 @@ public final class RegistryTool {
           .put("resave_environment_entities", ResaveEnvironmentEntitiesCommand.class)
           .put("save_sql_credential", SaveSqlCredentialCommand.class)
           .put("send_escrow_report_to_icann", SendEscrowReportToIcannCommand.class)
+          .put("set_database_migration_state", SetDatabaseMigrationStateCommand.class)
           .put("set_num_instances", SetNumInstancesCommand.class)
           .put("setup_ote", SetupOteCommand.class)
           .put("uniform_rapid_suspension", UniformRapidSuspensionCommand.class)
diff --git a/core/src/main/java/google/registry/tools/SetDatabaseMigrationStateCommand.java b/core/src/main/java/google/registry/tools/SetDatabaseMigrationStateCommand.java
new file mode 100644
index 000000000..f880eced6
--- /dev/null
+++ b/core/src/main/java/google/registry/tools/SetDatabaseMigrationStateCommand.java
@@ -0,0 +1,72 @@
+// Copyright 2021 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.tools;
+
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+
+import com.beust.jcommander.Parameter;
+import com.beust.jcommander.Parameters;
+import com.google.common.collect.ImmutableSortedMap;
+import google.registry.model.annotations.DeleteAfterMigration;
+import google.registry.model.common.DatabaseMigrationStateSchedule;
+import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState;
+import google.registry.tools.params.TransitionListParameter.MigrationStateTransitions;
+import org.joda.time.DateTime;
+
+/** Command to set the Registry 3.0 database migration state schedule. */
+@DeleteAfterMigration
+@Parameters(
+    separators = " =",
+    commandDescription = "Set the current database migration state schedule.")
+public class SetDatabaseMigrationStateCommand extends ConfirmingCommand
+    implements CommandWithRemoteApi {
+
+  private static final String WARNING_MESSAGE =
+      "Attempting to change the schedule with an effect that would take place within the next 10 "
+          + "minutes. The cache expiration duration is 5 minutes so this MAY BE DANGEROUS.\n";
+
+  @Parameter(
+      names = "--migration_schedule",
+      converter = MigrationStateTransitions.class,
+      validateWith = MigrationStateTransitions.class,
+      required = true,
+      description =
+          "Comma-delimited list of database transitions, of the form"
+              + " <time>=<migration-state>[,<time>=<migration-state>]*")
+  ImmutableSortedMap<DateTime, MigrationState> transitionSchedule;
+
+  @Override
+  protected String prompt() {
+    return jpaTm()
+        .transact(
+            () -> {
+              StringBuilder result = new StringBuilder();
+              DateTime now = jpaTm().getTransactionTime();
+              DateTime nextTransition = transitionSchedule.ceilingKey(now);
+              if (nextTransition != null && nextTransition.isBefore(now.plusMinutes(10))) {
+                result.append(WARNING_MESSAGE);
+              }
+              return result
+                  .append(String.format("Set new migration state schedule %s?", transitionSchedule))
+                  .toString();
+            });
+  }
+
+  @Override
+  protected String execute() {
+    jpaTm().transact(() -> DatabaseMigrationStateSchedule.set(transitionSchedule));
+    return String.format("Successfully set new migration state schedule %s", transitionSchedule);
+  }
+}
diff --git a/core/src/main/java/google/registry/tools/params/TransitionListParameter.java b/core/src/main/java/google/registry/tools/params/TransitionListParameter.java
index 31f6ee30f..6cb27c261 100644
--- a/core/src/main/java/google/registry/tools/params/TransitionListParameter.java
+++ b/core/src/main/java/google/registry/tools/params/TransitionListParameter.java
@@ -19,6 +19,7 @@ import static com.google.common.base.Preconditions.checkArgument;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.collect.Ordering;
+import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState;
 import google.registry.model.domain.token.AllocationToken.TokenStatus;
 import google.registry.model.tld.Registry.TldState;
 import org.joda.money.Money;
@@ -72,4 +73,12 @@ public abstract class TransitionListParameter<V> extends KeyValueMapParameter<Da
       return TokenStatus.valueOf(value);
     }
   }
+
+  /** Converter-validator for states of the Registry 3.0 database migration. */
+  public static class MigrationStateTransitions extends TransitionListParameter<MigrationState> {
+    @Override
+    protected MigrationState parseValue(String value) {
+      return MigrationState.valueOf(value);
+    }
+  }
 }
diff --git a/core/src/test/java/google/registry/model/common/DatabaseMigrationStateScheduleTest.java b/core/src/test/java/google/registry/model/common/DatabaseMigrationStateScheduleTest.java
new file mode 100644
index 000000000..e2f43bf73
--- /dev/null
+++ b/core/src/test/java/google/registry/model/common/DatabaseMigrationStateScheduleTest.java
@@ -0,0 +1,187 @@
+// Copyright 2021 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.common;
+
+import static com.google.common.truth.Truth.assertThat;
+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_NO_ASYNC;
+import static google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState.DATASTORE_PRIMARY_READ_ONLY;
+import static google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState.SQL_ONLY;
+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.jpaTm;
+import static google.registry.util.DateTimeUtils.START_OF_TIME;
+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;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+/** Tests for {@link DatabaseMigrationStateSchedule}. */
+public class DatabaseMigrationStateScheduleTest extends EntityTestCase {
+
+  @BeforeEach
+  void beforeEach() {
+    fakeClock.setAutoIncrementByOneMilli();
+  }
+
+  @AfterEach
+  void afterEach() {
+    DatabaseHelper.removeDatabaseMigrationSchedule();
+  }
+
+  @Test
+  void testEmpty_returnsDatastoreOnlyMap() {
+    assertThat(DatabaseMigrationStateSchedule.getUncached())
+        .isEqualTo(DatabaseMigrationStateSchedule.DEFAULT_TRANSITION_MAP);
+  }
+
+  @Test
+  void testValidTransitions() {
+    // First, verify that no-ops are safe
+    for (MigrationState migrationState : MigrationState.values()) {
+      runValidTransition(migrationState, migrationState);
+    }
+
+    // Next, the transitions that will actually cause a change
+    runValidTransition(DATASTORE_ONLY, DATASTORE_PRIMARY);
+
+    runValidTransition(DATASTORE_PRIMARY, DATASTORE_ONLY);
+    runValidTransition(DATASTORE_PRIMARY, DATASTORE_PRIMARY_NO_ASYNC);
+    runValidTransition(DATASTORE_PRIMARY_NO_ASYNC, DATASTORE_PRIMARY_READ_ONLY);
+
+    runValidTransition(DATASTORE_PRIMARY_READ_ONLY, DATASTORE_ONLY);
+    runValidTransition(DATASTORE_PRIMARY_READ_ONLY, DATASTORE_PRIMARY);
+    runValidTransition(DATASTORE_PRIMARY_READ_ONLY, DATASTORE_PRIMARY_NO_ASYNC);
+    runValidTransition(DATASTORE_PRIMARY_READ_ONLY, SQL_PRIMARY_READ_ONLY);
+    runValidTransition(DATASTORE_PRIMARY_READ_ONLY, SQL_PRIMARY);
+
+    runValidTransition(SQL_PRIMARY_READ_ONLY, DATASTORE_PRIMARY_READ_ONLY);
+    runValidTransition(SQL_PRIMARY_READ_ONLY, SQL_PRIMARY);
+
+    runValidTransition(SQL_PRIMARY, SQL_PRIMARY_READ_ONLY);
+    runValidTransition(SQL_PRIMARY, SQL_ONLY);
+
+    runValidTransition(SQL_ONLY, SQL_PRIMARY);
+  }
+
+  @Test
+  void testInvalidTransitions() {
+    runInvalidTransition(DATASTORE_ONLY, DATASTORE_PRIMARY_READ_ONLY);
+    runInvalidTransition(DATASTORE_ONLY, SQL_PRIMARY_READ_ONLY);
+    runInvalidTransition(DATASTORE_ONLY, SQL_PRIMARY);
+    runInvalidTransition(DATASTORE_ONLY, SQL_ONLY);
+
+    runInvalidTransition(DATASTORE_PRIMARY, DATASTORE_PRIMARY_READ_ONLY);
+    runInvalidTransition(DATASTORE_PRIMARY, SQL_PRIMARY_READ_ONLY);
+    runInvalidTransition(DATASTORE_PRIMARY, SQL_PRIMARY);
+    runInvalidTransition(DATASTORE_PRIMARY, SQL_ONLY);
+
+    runInvalidTransition(DATASTORE_PRIMARY_READ_ONLY, SQL_ONLY);
+
+    runInvalidTransition(SQL_PRIMARY_READ_ONLY, DATASTORE_ONLY);
+    runInvalidTransition(SQL_PRIMARY_READ_ONLY, DATASTORE_PRIMARY);
+    runInvalidTransition(SQL_PRIMARY_READ_ONLY, SQL_ONLY);
+
+    runInvalidTransition(SQL_PRIMARY, DATASTORE_ONLY);
+    runInvalidTransition(SQL_PRIMARY, DATASTORE_PRIMARY);
+    runInvalidTransition(SQL_PRIMARY, DATASTORE_PRIMARY_READ_ONLY);
+
+    runInvalidTransition(SQL_ONLY, DATASTORE_ONLY);
+    runInvalidTransition(SQL_ONLY, DATASTORE_PRIMARY);
+    runInvalidTransition(SQL_ONLY, DATASTORE_PRIMARY_READ_ONLY);
+  }
+
+  @Test
+  void testFailure_newMapImpliesInvalidChangeNow() {
+    DateTime startTime = fakeClock.nowUtc();
+    fakeClock.advanceBy(Duration.standardHours(6));
+
+    // The new map is valid by itself, but not with the current state of DATASTORE_ONLY because the
+    // new map implies that the current state is DATASTORE_PRIMARY_READ_ONLY
+    ImmutableSortedMap<DateTime, MigrationState> nowInvalidMap =
+        ImmutableSortedMap.<DateTime, MigrationState>naturalOrder()
+            .put(START_OF_TIME, DATASTORE_ONLY)
+            .put(startTime.plusHours(1), DATASTORE_PRIMARY)
+            .put(startTime.plusHours(2), DATASTORE_PRIMARY_NO_ASYNC)
+            .put(startTime.plusHours(3), DATASTORE_PRIMARY_READ_ONLY)
+            .build();
+    IllegalArgumentException thrown =
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> jpaTm().transact(() -> DatabaseMigrationStateSchedule.set(nowInvalidMap)));
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo(
+            "Cannot transition from current state-as-of-now DATASTORE_ONLY "
+                + "to new state-as-of-now DATASTORE_PRIMARY_READ_ONLY");
+  }
+
+  @Test
+  void testFailure_notInTransaction() {
+    IllegalStateException thrown =
+        assertThrows(
+            IllegalStateException.class,
+            () ->
+                DatabaseMigrationStateSchedule.set(
+                    DatabaseMigrationStateSchedule.DEFAULT_TRANSITION_MAP.toValueMap()));
+    assertThat(thrown).hasMessageThat().isEqualTo("Not in a transaction");
+  }
+
+  private void runValidTransition(MigrationState from, MigrationState to) {
+    ImmutableSortedMap<DateTime, MigrationState> transitions =
+        createMapEndingWithTransition(from, to);
+    jpaTm().transact(() -> DatabaseMigrationStateSchedule.set(transitions));
+    assertThat(DatabaseMigrationStateSchedule.getUncached().toValueMap())
+        .containsExactlyEntriesIn(transitions);
+  }
+
+  private void runInvalidTransition(MigrationState from, MigrationState to) {
+    ImmutableSortedMap<DateTime, MigrationState> transitions =
+        createMapEndingWithTransition(from, to);
+    IllegalArgumentException thrown =
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> jpaTm().transact(() -> DatabaseMigrationStateSchedule.set(transitions)));
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo(
+            String.format("validStateTransitions map cannot transition from %s to %s.", from, to));
+  }
+
+  // Create a transition map that is valid up to the "from" transition, then add the "to" transition
+  private ImmutableSortedMap<DateTime, MigrationState> createMapEndingWithTransition(
+      MigrationState from, MigrationState to) {
+    ImmutableSortedMap.Builder<DateTime, MigrationState> builder =
+        ImmutableSortedMap.naturalOrder();
+    builder.put(START_OF_TIME, DATASTORE_ONLY);
+    MigrationState[] allMigrationStates = MigrationState.values();
+    for (int i = 0; i < allMigrationStates.length; i++) {
+      builder.put(fakeClock.nowUtc().plusMinutes(i), allMigrationStates[i]);
+      if (allMigrationStates[i].equals(from)) {
+        break;
+      }
+    }
+    builder.put(fakeClock.nowUtc().plusDays(1), to);
+    return builder.build();
+  }
+}
diff --git a/core/src/test/java/google/registry/testing/DatabaseHelper.java b/core/src/test/java/google/registry/testing/DatabaseHelper.java
index 491e52415..fcece0542 100644
--- a/core/src/test/java/google/registry/testing/DatabaseHelper.java
+++ b/core/src/test/java/google/registry/testing/DatabaseHelper.java
@@ -71,6 +71,8 @@ import google.registry.model.billing.BillingEvent.Flag;
 import google.registry.model.billing.BillingEvent.Reason;
 import google.registry.model.billing.BillingEvent.Recurring;
 import google.registry.model.billing.BillingEvent.RenewalPriceBehavior;
+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;
@@ -124,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 {
@@ -1435,5 +1438,104 @@ public class DatabaseHelper {
     return entity;
   }
 
+  /**
+   * Sets a DATASTORE_PRIMARY_NO_ASYNC 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
+   * DATASTORE_PRIMARY_READ_ONLY 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 setMigrationScheduleToDatastorePrimaryNoAsync(FakeClock fakeClock) {
+    DateTime now = fakeClock.nowUtc();
+    jpaTm()
+        .transact(
+            () ->
+                DatabaseMigrationStateSchedule.set(
+                    ImmutableSortedMap.of(
+                        START_OF_TIME,
+                        MigrationState.DATASTORE_ONLY,
+                        now.plusMillis(1),
+                        MigrationState.DATASTORE_PRIMARY,
+                        now.plusMillis(2),
+                        MigrationState.DATASTORE_PRIMARY_NO_ASYNC)));
+    fakeClock.advanceBy(Duration.standardSeconds(1));
+  }
+
+  /**
+   * Sets a DATASTORE_PRIMARY_READ_ONLY 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
+   * DATASTORE_PRIMARY_READ_ONLY 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 setMigrationScheduleToDatastorePrimaryReadOnly(FakeClock fakeClock) {
+    DateTime now = fakeClock.nowUtc();
+    jpaTm()
+        .transact(
+            () ->
+                DatabaseMigrationStateSchedule.set(
+                    ImmutableSortedMap.of(
+                        START_OF_TIME,
+                        MigrationState.DATASTORE_ONLY,
+                        now.plusMillis(1),
+                        MigrationState.DATASTORE_PRIMARY,
+                        now.plusMillis(2),
+                        MigrationState.DATASTORE_PRIMARY_NO_ASYNC,
+                        now.plusMillis(3),
+                        MigrationState.DATASTORE_PRIMARY_READ_ONLY)));
+    fakeClock.advanceBy(Duration.standardSeconds(1));
+  }
+
+  /**
+   * 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();
+    jpaTm()
+        .transact(
+            () ->
+                DatabaseMigrationStateSchedule.set(
+                    ImmutableSortedMap.of(
+                        START_OF_TIME,
+                        MigrationState.DATASTORE_ONLY,
+                        now.plusMillis(1),
+                        MigrationState.DATASTORE_PRIMARY,
+                        now.plusMillis(2),
+                        MigrationState.DATASTORE_PRIMARY_NO_ASYNC,
+                        now.plusMillis(3),
+                        MigrationState.DATASTORE_PRIMARY_READ_ONLY,
+                        now.plusMillis(4),
+                        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
+    jpaTm()
+        .transact(
+            () ->
+                jpaTm()
+                    .put(
+                        new DatabaseMigrationStateSchedule(
+                            DatabaseMigrationStateSchedule.DEFAULT_TRANSITION_MAP)));
+    DatabaseMigrationStateSchedule.CACHE.invalidateAll();
+  }
+
   private DatabaseHelper() {}
 }
diff --git a/core/src/test/java/google/registry/tools/GetDatabaseMigrationStateCommandTest.java b/core/src/test/java/google/registry/tools/GetDatabaseMigrationStateCommandTest.java
new file mode 100644
index 000000000..33658237c
--- /dev/null
+++ b/core/src/test/java/google/registry/tools/GetDatabaseMigrationStateCommandTest.java
@@ -0,0 +1,66 @@
+// Copyright 2021 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.tools;
+
+import static google.registry.model.common.DatabaseMigrationStateSchedule.DEFAULT_TRANSITION_MAP;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+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 org.joda.time.DateTime;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+
+/** Tests for {@link GetDatabaseMigrationStateCommand}. */
+public class GetDatabaseMigrationStateCommandTest
+    extends CommandTestCase<GetDatabaseMigrationStateCommand> {
+
+  @AfterEach
+  void afterEach() {
+    DatabaseHelper.removeDatabaseMigrationSchedule();
+  }
+
+  @Test
+  void testInitial_returnsDatastoreOnly() throws Exception {
+    runCommand();
+    assertStdoutIs(
+        String.format("Current migration schedule: %s\n", DEFAULT_TRANSITION_MAP.toValueMap()));
+  }
+
+  @Test
+  void testFullSchedule() throws Exception {
+    DateTime now = fakeClock.nowUtc();
+    ImmutableSortedMap<DateTime, MigrationState> transitions =
+        ImmutableSortedMap.of(
+            START_OF_TIME,
+            MigrationState.DATASTORE_ONLY,
+            now.plusHours(1),
+            MigrationState.DATASTORE_PRIMARY,
+            now.plusHours(2),
+            MigrationState.DATASTORE_PRIMARY_NO_ASYNC,
+            now.plusHours(3),
+            MigrationState.DATASTORE_PRIMARY_READ_ONLY,
+            now.plusHours(4),
+            MigrationState.SQL_PRIMARY,
+            now.plusHours(5),
+            MigrationState.SQL_ONLY);
+    jpaTm().transact(() -> DatabaseMigrationStateSchedule.set(transitions));
+    runCommand();
+    assertStdoutIs(String.format("Current migration schedule: %s\n", transitions));
+  }
+}
diff --git a/core/src/test/java/google/registry/tools/SetDatabaseMigrationStateCommandTest.java b/core/src/test/java/google/registry/tools/SetDatabaseMigrationStateCommandTest.java
new file mode 100644
index 000000000..48f56b829
--- /dev/null
+++ b/core/src/test/java/google/registry/tools/SetDatabaseMigrationStateCommandTest.java
@@ -0,0 +1,184 @@
+// Copyright 2021 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.tools;
+
+import static com.google.common.truth.Truth.assertThat;
+import static com.google.common.truth.Truth8.assertThat;
+import static google.registry.model.common.DatabaseMigrationStateSchedule.DEFAULT_TRANSITION_MAP;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.util.DateTimeUtils.START_OF_TIME;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+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 org.joda.time.DateTime;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+
+/** Tests for {@link SetDatabaseMigrationStateCommand}. */
+public class SetDatabaseMigrationStateCommandTest
+    extends CommandTestCase<SetDatabaseMigrationStateCommand> {
+
+  @AfterEach
+  void afterEach() {
+    DatabaseHelper.removeDatabaseMigrationSchedule();
+  }
+
+  @Test
+  void testSuccess_setsBasicSchedule() throws Exception {
+    assertThat(DatabaseMigrationStateSchedule.get()).isEqualTo(DEFAULT_TRANSITION_MAP);
+    assertThat(jpaTm().transact(() -> jpaTm().loadSingleton(DatabaseMigrationStateSchedule.class)))
+        .isEmpty();
+    runCommandForced("--migration_schedule=1970-01-01T00:00:00.000Z=DATASTORE_ONLY");
+    // use a raw ofy call to check what's in the DB
+    jpaTm()
+        .transact(
+            () ->
+                assertThat(
+                        jpaTm()
+                            .loadSingleton(DatabaseMigrationStateSchedule.class)
+                            .get()
+                            .migrationTransitions)
+                    .isEqualTo(DEFAULT_TRANSITION_MAP));
+    assertThat(DatabaseMigrationStateSchedule.get()).isEqualTo(DEFAULT_TRANSITION_MAP);
+  }
+
+  @Test
+  void testSuccess_fullSchedule() throws Exception {
+    DateTime now = fakeClock.nowUtc();
+    DateTime datastorePrimary = now.plusHours(1);
+    DateTime datastorePrimaryNoAsync = now.plusHours(2);
+    DateTime datastorePrimaryReadOnly = now.plusHours(3);
+    DateTime sqlPrimary = now.plusHours(4);
+    DateTime sqlOnly = now.plusHours(5);
+    runCommandForced(
+        String.format(
+            "--migration_schedule=%s=DATASTORE_ONLY,%s=DATASTORE_PRIMARY,"
+                + "%s=DATASTORE_PRIMARY_NO_ASYNC,%s=DATASTORE_PRIMARY_READ_ONLY,"
+                + "%s=SQL_PRIMARY,%s=SQL_ONLY",
+            START_OF_TIME,
+            datastorePrimary,
+            datastorePrimaryNoAsync,
+            datastorePrimaryReadOnly,
+            sqlPrimary,
+            sqlOnly));
+    assertThat(DatabaseMigrationStateSchedule.get().toValueMap())
+        .containsExactlyEntriesIn(
+            ImmutableSortedMap.of(
+                START_OF_TIME,
+                MigrationState.DATASTORE_ONLY,
+                datastorePrimary,
+                MigrationState.DATASTORE_PRIMARY,
+                datastorePrimaryNoAsync,
+                MigrationState.DATASTORE_PRIMARY_NO_ASYNC,
+                datastorePrimaryReadOnly,
+                MigrationState.DATASTORE_PRIMARY_READ_ONLY,
+                sqlPrimary,
+                MigrationState.SQL_PRIMARY,
+                sqlOnly,
+                MigrationState.SQL_ONLY));
+  }
+
+  @Test
+  void testSuccess_warnsOnChangeSoon() throws Exception {
+    DateTime now = fakeClock.nowUtc();
+    runCommandForced(
+        String.format(
+            "--migration_schedule=%s=DATASTORE_ONLY,%s=DATASTORE_PRIMARY",
+            START_OF_TIME, now.plusMinutes(1)));
+    assertThat(DatabaseMigrationStateSchedule.get().toValueMap())
+        .containsExactlyEntriesIn(
+            ImmutableSortedMap.of(
+                START_OF_TIME,
+                MigrationState.DATASTORE_ONLY,
+                now.plusMinutes(1),
+                MigrationState.DATASTORE_PRIMARY));
+    assertInStdout("MAY BE DANGEROUS");
+  }
+
+  @Test
+  void testSuccess_goesBackward() throws Exception {
+    DateTime now = fakeClock.nowUtc();
+    runCommandForced(
+        String.format(
+            "--migration_schedule=%s=DATASTORE_ONLY,%s=DATASTORE_PRIMARY,"
+                + "%s=DATASTORE_PRIMARY_NO_ASYNC,%s=DATASTORE_PRIMARY_READ_ONLY,"
+                + "%s=DATASTORE_PRIMARY",
+            START_OF_TIME, now.plusHours(1), now.plusHours(2), now.plusHours(3), now.plusHours(4)));
+    assertThat(DatabaseMigrationStateSchedule.get().toValueMap())
+        .containsExactlyEntriesIn(
+            ImmutableSortedMap.of(
+                START_OF_TIME,
+                MigrationState.DATASTORE_ONLY,
+                now.plusHours(1),
+                MigrationState.DATASTORE_PRIMARY,
+                now.plusHours(2),
+                MigrationState.DATASTORE_PRIMARY_NO_ASYNC,
+                now.plusHours(3),
+                MigrationState.DATASTORE_PRIMARY_READ_ONLY,
+                now.plusHours(4),
+                MigrationState.DATASTORE_PRIMARY));
+  }
+
+  @Test
+  void testFailure_invalidTransition() {
+    IllegalArgumentException thrown =
+        assertThrows(
+            IllegalArgumentException.class,
+            () ->
+                runCommandForced(
+                    String.format(
+                        "--migration_schedule=%s=DATASTORE_ONLY,%s=DATASTORE_PRIMARY_READ_ONLY",
+                        START_OF_TIME, START_OF_TIME.plusHours(1))));
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo(
+            "validStateTransitions map cannot transition from DATASTORE_ONLY "
+                + "to DATASTORE_PRIMARY_READ_ONLY.");
+  }
+
+  @Test
+  void testFailure_invalidTransitionFromOldToNew() {
+    // The map we pass in is valid by itself, but we can't go from DATASTORE_ONLY now to
+    // DATASTORE_PRIMARY_READ_ONLY now
+    DateTime now = fakeClock.nowUtc();
+    IllegalArgumentException thrown =
+        assertThrows(
+            IllegalArgumentException.class,
+            () ->
+                runCommandForced(
+                    String.format(
+                        "--migration_schedule=%s=DATASTORE_ONLY,%s=DATASTORE_PRIMARY,"
+                            + "%s=DATASTORE_PRIMARY_NO_ASYNC,%s=DATASTORE_PRIMARY_READ_ONLY",
+                        START_OF_TIME, now.minusHours(3), now.minusHours(2), now.minusHours(1))));
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo(
+            "Cannot transition from current state-as-of-now DATASTORE_ONLY "
+                + "to new state-as-of-now DATASTORE_PRIMARY_READ_ONLY");
+  }
+
+  @Test
+  void testFailure_invalidParams() {
+    assertThrows(ParameterException.class, this::runCommandForced);
+    assertThrows(ParameterException.class, () -> runCommandForced("--migration_schedule=FOOBAR"));
+    assertThrows(
+        ParameterException.class,
+        () -> runCommandForced("--migration_schedule=1970-01-01T00:00:00.000Z=FOOBAR"));
+  }
+}