Use a TimedTransitionProperty for the DB migration schedule (#1186)

This includes the following changes: - Convert the single-valued database migration state to a timed transition property, meaning that we can switch all instances over at the same time and schedule it in advance - Use a "cache" (technically an expiring memoized supplier) when retrieving the database migration state value - Delete the old DatabaseTransitionSchedule because it is no longer necessary. We took the idea from that and used it for the new DatabaseMigrationStateSchedule, though we cannot reuse the entity itself because the structure is fundamentally different. - Removed references to the DatabaseTransitionSchedule, mainly in the getter/setter commands+tests and a few odd references elsewhere.
2026-09-21 07:24:26 +00:00 · 2021-06-02 14:06:28 -04:00
parent 275f364dcb
commit 586189d7ee
23 changed files with 429 additions and 1000 deletions
@@ -0,0 +1,173 @@
+// 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_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.ofyTm;
+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 org.joda.time.DateTime;
+import org.joda.time.Duration;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+public class DatabaseMigrationStateScheduleTest extends EntityTestCase {
+
+  @BeforeEach
+  void beforeEach() {
+    fakeClock.setAutoIncrementByOneMilli();
+  }
+
+  @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_READ_ONLY);
+
+    runValidTransition(DATASTORE_PRIMARY_READ_ONLY, DATASTORE_ONLY);
+    runValidTransition(DATASTORE_PRIMARY_READ_ONLY, DATASTORE_PRIMARY);
+    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, 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_READ_ONLY)
+            .build();
+    assertThat(
+            assertThrows(
+                IllegalArgumentException.class,
+                () -> ofyTm().transact(() -> DatabaseMigrationStateSchedule.set(nowInvalidMap))))
+        .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() {
+    assertThat(
+            assertThrows(
+                IllegalStateException.class,
+                () ->
+                    DatabaseMigrationStateSchedule.set(
+                        DatabaseMigrationStateSchedule.DEFAULT_TRANSITION_MAP.toValueMap())))
+        .hasMessageThat()
+        .isEqualTo("Must be called in a transaction");
+  }
+
+  private void runValidTransition(MigrationState from, MigrationState to) {
+    ImmutableSortedMap<DateTime, MigrationState> transitions =
+        createMapEndingWithTransition(from, to);
+    ofyTm().transact(() -> DatabaseMigrationStateSchedule.set(transitions));
+    assertThat(DatabaseMigrationStateSchedule.getUncached().toValueMap())
+        .containsExactlyEntriesIn(transitions);
+  }
+
+  private void runInvalidTransition(MigrationState from, MigrationState to) {
+    ImmutableSortedMap<DateTime, MigrationState> transitions =
+        createMapEndingWithTransition(from, to);
+    assertThat(
+            assertThrows(
+                IllegalArgumentException.class,
+                () -> ofyTm().transact(() -> DatabaseMigrationStateSchedule.set(transitions))))
+        .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();
+  }
+}
@@ -1,111 +0,0 @@
-// 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.DatabaseMigrationStateWrapper.MigrationState.DATASTORE_ONLY;
-import static google.registry.model.common.DatabaseMigrationStateWrapper.MigrationState.DATASTORE_PRIMARY;
-import static google.registry.model.common.DatabaseMigrationStateWrapper.MigrationState.DATASTORE_PRIMARY_READ_ONLY;
-import static google.registry.model.common.DatabaseMigrationStateWrapper.MigrationState.SQL_ONLY;
-import static google.registry.model.common.DatabaseMigrationStateWrapper.MigrationState.SQL_PRIMARY;
-import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
-import static org.junit.jupiter.api.Assertions.assertThrows;
-
-import google.registry.model.common.DatabaseMigrationStateWrapper.MigrationState;
-import google.registry.testing.AppEngineExtension;
-import org.junit.jupiter.api.Test;
-import org.junit.jupiter.api.extension.RegisterExtension;
-
-public class DatabaseMigrationStateWrapperTest {
-
-  @RegisterExtension
-  public final AppEngineExtension appEngine =
-      AppEngineExtension.builder().withDatastoreAndCloudSql().build();
-
-  @Test
-  void testEmpty_returnsDatastore() {
-    assertThat(DatabaseMigrationStateWrapper.get()).isEqualTo(DATASTORE_ONLY);
-  }
-
-  @Test
-  void testEmpty_canChangeToDatastorePrimary() {
-    DatabaseMigrationStateWrapper.set(DATASTORE_PRIMARY);
-    assertThat(DatabaseMigrationStateWrapper.get()).isEqualTo(DATASTORE_PRIMARY);
-  }
-
-  @Test
-  void testValidTransitions() {
-    runValidTransition(DATASTORE_ONLY, DATASTORE_PRIMARY);
-
-    runValidTransition(DATASTORE_PRIMARY, DATASTORE_ONLY);
-    runValidTransition(DATASTORE_PRIMARY, DATASTORE_PRIMARY_READ_ONLY);
-
-    runValidTransition(DATASTORE_PRIMARY_READ_ONLY, DATASTORE_ONLY);
-    runValidTransition(DATASTORE_PRIMARY_READ_ONLY, DATASTORE_PRIMARY);
-    runValidTransition(DATASTORE_PRIMARY_READ_ONLY, SQL_PRIMARY);
-
-    runValidTransition(SQL_PRIMARY, DATASTORE_PRIMARY_READ_ONLY);
-    runValidTransition(SQL_PRIMARY, SQL_ONLY);
-
-    runValidTransition(SQL_ONLY, SQL_PRIMARY);
-  }
-
-  @Test
-  void testInvalidTransitions() {
-    runInvalidTransition(DATASTORE_ONLY, DATASTORE_ONLY);
-    runInvalidTransition(DATASTORE_ONLY, DATASTORE_PRIMARY_READ_ONLY);
-    runInvalidTransition(DATASTORE_ONLY, SQL_PRIMARY);
-    runInvalidTransition(DATASTORE_ONLY, SQL_ONLY);
-
-    runInvalidTransition(DATASTORE_PRIMARY, DATASTORE_PRIMARY);
-    runInvalidTransition(DATASTORE_PRIMARY, SQL_PRIMARY);
-    runInvalidTransition(DATASTORE_PRIMARY, SQL_ONLY);
-
-    runInvalidTransition(DATASTORE_PRIMARY_READ_ONLY, DATASTORE_PRIMARY_READ_ONLY);
-    runInvalidTransition(DATASTORE_PRIMARY_READ_ONLY, SQL_ONLY);
-
-    runInvalidTransition(SQL_PRIMARY, DATASTORE_ONLY);
-    runInvalidTransition(SQL_PRIMARY, DATASTORE_PRIMARY);
-    runInvalidTransition(SQL_PRIMARY, SQL_PRIMARY);
-
-    runInvalidTransition(SQL_ONLY, DATASTORE_ONLY);
-    runInvalidTransition(SQL_ONLY, DATASTORE_PRIMARY);
-    runInvalidTransition(SQL_ONLY, DATASTORE_PRIMARY_READ_ONLY);
-    runInvalidTransition(SQL_ONLY, SQL_ONLY);
-  }
-
-  private static void runValidTransition(MigrationState from, MigrationState to) {
-    setStateForced(from);
-    DatabaseMigrationStateWrapper.set(to);
-    assertThat(DatabaseMigrationStateWrapper.get()).isEqualTo(to);
-  }
-
-  private static void runInvalidTransition(MigrationState from, MigrationState to) {
-    setStateForced(from);
-    assertThat(
-            assertThrows(
-                IllegalArgumentException.class, () -> DatabaseMigrationStateWrapper.set(to)))
-        .hasMessageThat()
-        .isEqualTo(
-            String.format(
-                "Moving from migration state %s to %s is not a valid transition", from, to));
-  }
-
-  private static void setStateForced(MigrationState migrationState) {
-    DatabaseMigrationStateWrapper wrapper = new DatabaseMigrationStateWrapper(migrationState);
-    ofyTm().transact(() -> ofyTm().put(wrapper));
-    assertThat(DatabaseMigrationStateWrapper.get()).isEqualTo(migrationState);
-  }
-}
@@ -1,79 +0,0 @@
-// 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.persistence.transaction.TransactionManagerFactory.ofyTm;
-import static google.registry.util.DateTimeUtils.START_OF_TIME;
-import static org.junit.jupiter.api.Assertions.assertThrows;
-
-import com.google.common.collect.ImmutableSortedMap;
-import google.registry.model.EntityTestCase;
-import google.registry.model.common.DatabaseTransitionSchedule.PrimaryDatabase;
-import google.registry.model.common.DatabaseTransitionSchedule.PrimaryDatabaseTransition;
-import google.registry.model.common.DatabaseTransitionSchedule.TransitionId;
-import org.joda.time.Duration;
-import org.junit.jupiter.api.Test;
-
-/** Unit tests for {@link DatabaseTransitionSchedule}. */
-public class DatabaseTransitionScheduleTest extends EntityTestCase {
-
-  @Test
-  void testSuccess_persistence() {
-    TimedTransitionProperty<PrimaryDatabase, PrimaryDatabaseTransition> databaseTransitions =
-        TimedTransitionProperty.fromValueMap(
-            ImmutableSortedMap.of(START_OF_TIME, PrimaryDatabase.DATASTORE),
-            PrimaryDatabaseTransition.class);
-    DatabaseTransitionSchedule schedule =
-        DatabaseTransitionSchedule.create(
-            TransitionId.SIGNED_MARK_REVOCATION_LIST, databaseTransitions);
-    ofyTm().transactNew(() -> ofyTm().put(schedule));
-
-    assertThat(
-            DatabaseTransitionSchedule.get(TransitionId.SIGNED_MARK_REVOCATION_LIST)
-                .get()
-                .databaseTransitions)
-        .isEqualTo(databaseTransitions);
-  }
-
-  @Test
-  void testFailure_scheduleWithNoStartOfTime() {
-    assertThrows(
-        IllegalArgumentException.class,
-        () ->
-            DatabaseTransitionSchedule.create(
-                TransitionId.SIGNED_MARK_REVOCATION_LIST,
-                TimedTransitionProperty.fromValueMap(
-                    ImmutableSortedMap.of(fakeClock.nowUtc(), PrimaryDatabase.DATASTORE),
-                    PrimaryDatabaseTransition.class)));
-  }
-
-  @Test
-  void testSuccess_getPrimaryDatabase() {
-    DatabaseTransitionSchedule schedule =
-        DatabaseTransitionSchedule.create(
-            TransitionId.SIGNED_MARK_REVOCATION_LIST,
-            TimedTransitionProperty.fromValueMap(
-                ImmutableSortedMap.of(
-                    START_OF_TIME,
-                    PrimaryDatabase.DATASTORE,
-                    fakeClock.nowUtc().plusDays(1),
-                    PrimaryDatabase.CLOUD_SQL),
-                PrimaryDatabaseTransition.class));
-    assertThat(ofyTm().transact(schedule::getPrimaryDatabase)).isEqualTo(PrimaryDatabase.DATASTORE);
-    fakeClock.advanceBy(Duration.standardDays(5));
-    assertThat(ofyTm().transact(schedule::getPrimaryDatabase)).isEqualTo(PrimaryDatabase.CLOUD_SQL);
-  }
-}
@@ -17,12 +17,10 @@ package google.registry.testing;
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.io.Files.asCharSink;
 import static com.google.common.truth.Truth.assertWithMessage;
-import static google.registry.model.ofy.ObjectifyService.auditedOfy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatabaseHelper.persistSimpleResources;
 import static google.registry.testing.DualDatabaseTestInvocationContextProvider.injectTmForDualDatabaseTest;
 import static google.registry.testing.DualDatabaseTestInvocationContextProvider.restoreTmAfterDualDatabaseTest;
-import static google.registry.util.DateTimeUtils.START_OF_TIME;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 import static google.registry.util.ResourceUtils.readResourceUtf8;
 import static java.nio.charset.StandardCharsets.UTF_8;
@@ -42,16 +40,10 @@ import com.google.common.base.Joiner;
 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.collect.Sets;
 import com.google.common.io.Files;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.ObjectifyFilter;
-import google.registry.model.common.DatabaseTransitionSchedule;
-import google.registry.model.common.DatabaseTransitionSchedule.PrimaryDatabase;
-import google.registry.model.common.DatabaseTransitionSchedule.PrimaryDatabaseTransition;
-import google.registry.model.common.DatabaseTransitionSchedule.TransitionId;
-import google.registry.model.common.TimedTransitionProperty;
 import google.registry.model.ofy.ObjectifyService;
 import google.registry.model.registrar.Registrar;
 import google.registry.model.registrar.Registrar.State;
@@ -402,18 +394,8 @@ public final class AppEngineExtension implements BeforeEachCallback, AfterEachCa
      if (withDatastore && !withoutCannedData) {
        loadInitialData();
      }
-    } else {
-      // If we're using SQL, set replayed entities to use SQL
-      DatabaseTransitionSchedule schedule =
-          DatabaseTransitionSchedule.create(
-              TransitionId.REPLAYED_ENTITIES,
-              TimedTransitionProperty.fromValueMap(
-                  ImmutableSortedMap.of(START_OF_TIME, PrimaryDatabase.CLOUD_SQL),
-                  PrimaryDatabaseTransition.class));
-      tm().transactNew(() -> auditedOfy().saveWithoutBackup().entity(schedule).now());
-      if (withCloudSql && !withJpaUnitTest && !withoutCannedData) {
-        loadInitialData();
-      }
+    } else if (withCloudSql && !withJpaUnitTest && !withoutCannedData) {
+      loadInitialData();
    }
  }

@@ -1,132 +0,0 @@
-// 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 google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
-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.DatabaseTransitionSchedule;
-import google.registry.model.common.DatabaseTransitionSchedule.PrimaryDatabase;
-import google.registry.model.common.DatabaseTransitionSchedule.PrimaryDatabaseTransition;
-import google.registry.model.common.DatabaseTransitionSchedule.TransitionId;
-import google.registry.model.common.TimedTransitionProperty;
-import google.registry.model.ofy.Ofy;
-import google.registry.testing.InjectExtension;
-import google.registry.testing.SetClockExtension;
-import org.joda.time.DateTime;
-import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Order;
-import org.junit.jupiter.api.Test;
-import org.junit.jupiter.api.extension.RegisterExtension;
-
-/** Unit tests for {@link GetDatabaseTransitionScheduleCommand} */
-public class GetDatabaseTransitionScheduleCommandTest
-    extends CommandTestCase<GetDatabaseTransitionScheduleCommand> {
-
-  @Order(value = Order.DEFAULT - 3)
-  @RegisterExtension
-  final SetClockExtension setClockExtension =
-      new SetClockExtension(fakeClock, "1984-12-21T06:07:08.789Z");
-
-  @RegisterExtension public final InjectExtension inject = new InjectExtension();
-
-  @BeforeEach
-  void beforeEach() {
-    inject.setStaticField(Ofy.class, "clock", fakeClock);
-  }
-
-  @Test
-  void testSuccess() throws Exception {
-    TimedTransitionProperty<PrimaryDatabase, PrimaryDatabaseTransition> databaseTransitions =
-        TimedTransitionProperty.fromValueMap(
-            ImmutableSortedMap.of(START_OF_TIME, PrimaryDatabase.DATASTORE),
-            PrimaryDatabaseTransition.class);
-    DatabaseTransitionSchedule schedule =
-        DatabaseTransitionSchedule.create(
-            TransitionId.SIGNED_MARK_REVOCATION_LIST, databaseTransitions);
-    ofyTm().transactNew(() -> ofyTm().put(schedule));
-    runCommand("SIGNED_MARK_REVOCATION_LIST");
-    assertStdoutIs(
-        "SIGNED_MARK_REVOCATION_LIST(last updated at 1984-12-21T06:07:08.789Z):"
-            + " {1970-01-01T00:00:00.000Z=DATASTORE}\n");
-  }
-
-  @Test
-  void testSuccess_multipleArguments() throws Exception {
-    TimedTransitionProperty<PrimaryDatabase, PrimaryDatabaseTransition> databaseTransitions =
-        TimedTransitionProperty.fromValueMap(
-            ImmutableSortedMap.of(START_OF_TIME, PrimaryDatabase.DATASTORE),
-            PrimaryDatabaseTransition.class);
-    DatabaseTransitionSchedule schedule =
-        DatabaseTransitionSchedule.create(TransitionId.DOMAIN_LABEL_LISTS, databaseTransitions);
-    ofyTm().transactNew(() -> ofyTm().put(schedule));
-    fakeClock.advanceOneMilli(); // Now 1984-12-21T06:07:08.790Z
-    TimedTransitionProperty<PrimaryDatabase, PrimaryDatabaseTransition> databaseTransitions2 =
-        TimedTransitionProperty.fromValueMap(
-            ImmutableSortedMap.of(
-                START_OF_TIME,
-                PrimaryDatabase.DATASTORE,
-                DateTime.parse("2020-10-01T00:00:00Z"),
-                PrimaryDatabase.CLOUD_SQL),
-            PrimaryDatabaseTransition.class);
-    DatabaseTransitionSchedule schedule2 =
-        DatabaseTransitionSchedule.create(
-            TransitionId.SIGNED_MARK_REVOCATION_LIST, databaseTransitions2);
-    ofyTm().transactNew(() -> ofyTm().put(schedule2));
-    runCommand("DOMAIN_LABEL_LISTS", "SIGNED_MARK_REVOCATION_LIST");
-    assertStdoutIs(
-        "DOMAIN_LABEL_LISTS(last updated at 1984-12-21T06:07:08.789Z):"
-            + " {1970-01-01T00:00:00.000Z=DATASTORE}\n"
-            + "SIGNED_MARK_REVOCATION_LIST(last updated at 1984-12-21T06:07:08.790Z):"
-            + " {1970-01-01T00:00:00.000Z=DATASTORE, 2020-10-01T00:00:00.000Z=CLOUD_SQL}\n");
-  }
-
-  @Test
-  void testFailure_scheduleDoesNotExist() {
-    IllegalArgumentException thrown =
-        assertThrows(
-            IllegalArgumentException.class, () -> runCommand("SIGNED_MARK_REVOCATION_LIST"));
-    assertThat(thrown)
-        .hasMessageThat()
-        .contains("A database transition schedule for SIGNED_MARK_REVOCATION_LIST does not exist");
-  }
-
-  @Test
-  void testFailure_noIdGiven() {
-    assertThrows(ParameterException.class, this::runCommand);
-  }
-
-  @Test
-  void testFailure_oneScheduleDoesNotExist() {
-    TimedTransitionProperty<PrimaryDatabase, PrimaryDatabaseTransition> databaseTransitions =
-        TimedTransitionProperty.fromValueMap(
-            ImmutableSortedMap.of(START_OF_TIME, PrimaryDatabase.DATASTORE),
-            PrimaryDatabaseTransition.class);
-    DatabaseTransitionSchedule schedule =
-        DatabaseTransitionSchedule.create(TransitionId.DOMAIN_LABEL_LISTS, databaseTransitions);
-    ofyTm().transactNew(() -> ofyTm().put(schedule));
-    IllegalArgumentException thrown =
-        assertThrows(
-            IllegalArgumentException.class,
-            () -> runCommand("DOMAIN_LABEL_LISTS", "SIGNED_MARK_REVOCATION_LIST"));
-    assertThat(thrown)
-        .hasMessageThat()
-        .contains("A database transition schedule for SIGNED_MARK_REVOCATION_LIST does not exist");
-  }
-}
@@ -1,134 +0,0 @@
-// 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 google.registry.model.common.EntityGroupRoot.getCrossTldKey;
-import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
-import static google.registry.testing.DatabaseHelper.persistResource;
-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 com.google.common.truth.Truth8;
-import com.googlecode.objectify.Key;
-import google.registry.model.common.DatabaseTransitionSchedule;
-import google.registry.model.common.DatabaseTransitionSchedule.PrimaryDatabase;
-import google.registry.model.common.DatabaseTransitionSchedule.PrimaryDatabaseTransition;
-import google.registry.model.common.DatabaseTransitionSchedule.TransitionId;
-import google.registry.model.common.TimedTransitionProperty;
-import google.registry.persistence.VKey;
-import org.joda.time.DateTime;
-import org.joda.time.Duration;
-import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
-
-/** Unit tests for {@link SetDatabaseTransitionScheduleCommand}. */
-public class SetDatabaseTransitionScheduleCommandTest
-    extends CommandTestCase<SetDatabaseTransitionScheduleCommand> {
-
-  @BeforeEach
-  void setup() {
-    fakeClock.setTo(DateTime.parse("2020-12-01T00:00:00Z"));
-  }
-
-  @Test
-  void testFailure_noTransitionId() throws Exception {
-    ParameterException thrown =
-        assertThrows(
-            ParameterException.class,
-            () -> runCommandForced("--transition_schedule=2021-04-14T00:00:00.000Z=DATASTORE"));
-    assertThat(thrown).hasMessageThat().contains("--transition_id");
-  }
-
-  @Test
-  void testSuccess_currentScheduleIsEmpty() throws Exception {
-    Truth8.assertThat(
-            ofyTm()
-                .loadByKeyIfPresent(
-                    VKey.createOfy(
-                        DatabaseTransitionSchedule.class,
-                        Key.create(getCrossTldKey(), DatabaseTransitionSchedule.class, "test"))))
-        .isEmpty();
-    runCommandForced(
-        "--transition_id=SIGNED_MARK_REVOCATION_LIST",
-        "--transition_schedule=1970-01-01T00:00:00.000Z=DATASTORE");
-    assertThat(
-            ofyTm()
-                .transact(
-                    () ->
-                        DatabaseTransitionSchedule.get(TransitionId.SIGNED_MARK_REVOCATION_LIST)
-                            .get()
-                            .getPrimaryDatabase()))
-        .isEqualTo(PrimaryDatabase.DATASTORE);
-    assertThat(command.prompt())
-        .isEqualTo(
-            "Insert new schedule {1970-01-01T00:00:00.000Z=DATASTORE} "
-                + "for transition ID SIGNED_MARK_REVOCATION_LIST?");
-  }
-
-  @Test
-  void testSuccess() throws Exception {
-    ImmutableSortedMap<DateTime, PrimaryDatabase> transitionMap =
-        ImmutableSortedMap.of(
-            START_OF_TIME,
-            PrimaryDatabase.DATASTORE,
-            fakeClock.nowUtc().minusDays(1),
-            PrimaryDatabase.CLOUD_SQL);
-    persistResource(
-        DatabaseTransitionSchedule.create(
-            TransitionId.SIGNED_MARK_REVOCATION_LIST,
-            TimedTransitionProperty.fromValueMap(transitionMap, PrimaryDatabaseTransition.class)));
-    assertThat(
-            DatabaseTransitionSchedule.get(TransitionId.SIGNED_MARK_REVOCATION_LIST)
-                .get()
-                .getDatabaseTransitions())
-        .isEqualTo(transitionMap);
-    runCommandForced(
-        "--transition_id=SIGNED_MARK_REVOCATION_LIST",
-        "--transition_schedule=1970-01-01T00:00:00.000Z=DATASTORE,"
-            + "2020-11-30T00:00:00.000Z=CLOUD_SQL,2020-12-06T00:00:00.000Z=DATASTORE");
-    ImmutableSortedMap<DateTime, PrimaryDatabase> retrievedTransitionMap =
-        ofyTm()
-            .transact(
-                () ->
-                    DatabaseTransitionSchedule.get(TransitionId.SIGNED_MARK_REVOCATION_LIST)
-                        .get()
-                        .getDatabaseTransitions());
-    assertThat(retrievedTransitionMap)
-        .containsExactly(
-            START_OF_TIME,
-            PrimaryDatabase.DATASTORE,
-            fakeClock.nowUtc().minusDays(1),
-            PrimaryDatabase.CLOUD_SQL,
-            fakeClock.nowUtc().plusDays(5),
-            PrimaryDatabase.DATASTORE);
-    fakeClock.advanceBy(Duration.standardDays(5));
-    assertThat(
-            ofyTm()
-                .transact(
-                    () ->
-                        DatabaseTransitionSchedule.get(TransitionId.SIGNED_MARK_REVOCATION_LIST)
-                            .get()
-                            .getPrimaryDatabase()))
-        .isEqualTo(PrimaryDatabase.DATASTORE);
-    assertThat(command.prompt())
-        .isEqualTo(
-            "Insert new schedule {1970-01-01T00:00:00.000Z=DATASTORE, "
-                + "2020-11-30T00:00:00.000Z=CLOUD_SQL, 2020-12-06T00:00:00.000Z=DATASTORE} "
-                + "for transition ID SIGNED_MARK_REVOCATION_LIST?");
-  }
-}
@@ -2,8 +2,7 @@ AllocationToken
 Cancellation
 ContactResource
 Cursor
-DatabaseMigrationStateWrapper
-DatabaseTransitionSchedule
+DatabaseMigrationStateSchedule
 DomainBase
 EntityGroupRoot
 EppResourceIndex
@@ -1,8 +1,7 @@
 ClaimsList
 ClaimsListSingleton
 Cursor
-DatabaseMigrationStateWrapper
-DatabaseTransitionSchedule
+DatabaseMigrationStateSchedule
 KmsSecret
 KmsSecretRevision
 PremiumList
@@ -78,30 +78,21 @@ class google.registry.model.common.Cursor {
  google.registry.model.UpdateAutoTimestamp lastUpdateTime;
  org.joda.time.DateTime cursorTime;
 }
-class google.registry.model.common.DatabaseMigrationStateWrapper {
+class google.registry.model.common.DatabaseMigrationStateSchedule {
  @Id long id;
  @Parent com.googlecode.objectify.Key<google.registry.model.common.EntityGroupRoot> parent;
-  google.registry.model.common.DatabaseMigrationStateWrapper$MigrationState migrationState;
+  google.registry.model.common.TimedTransitionProperty<google.registry.model.common.DatabaseMigrationStateSchedule$MigrationState, google.registry.model.common.DatabaseMigrationStateSchedule$MigrationStateTransition> migrationTransitions;
 }
-enum google.registry.model.common.DatabaseMigrationStateWrapper$MigrationState {
+enum google.registry.model.common.DatabaseMigrationStateSchedule$MigrationState {
  DATASTORE_ONLY;
  DATASTORE_PRIMARY;
  DATASTORE_PRIMARY_READ_ONLY;
  SQL_ONLY;
  SQL_PRIMARY;
+  SQL_PRIMARY_READ_ONLY;
 }
-class google.registry.model.common.DatabaseTransitionSchedule {
-  @Id java.lang.String transitionId;
-  @Parent com.googlecode.objectify.Key<google.registry.model.common.EntityGroupRoot> parent;
-  google.registry.model.UpdateAutoTimestamp lastUpdateTime;
-  google.registry.model.common.TimedTransitionProperty<google.registry.model.common.DatabaseTransitionSchedule$PrimaryDatabase, google.registry.model.common.DatabaseTransitionSchedule$PrimaryDatabaseTransition> databaseTransitions;
-}
-enum google.registry.model.common.DatabaseTransitionSchedule$PrimaryDatabase {
-  CLOUD_SQL;
-  DATASTORE;
-}
-class google.registry.model.common.DatabaseTransitionSchedule$PrimaryDatabaseTransition {
-  google.registry.model.common.DatabaseTransitionSchedule$PrimaryDatabase primaryDatabase;
+class google.registry.model.common.DatabaseMigrationStateSchedule$MigrationStateTransition {
+  google.registry.model.common.DatabaseMigrationStateSchedule$MigrationState migrationState;
  org.joda.time.DateTime transitionTime;
 }
 class google.registry.model.common.EntityGroupRoot {