Add databaseTransitionSchedule entity and tool for updating (#926)

* Add databaseTransitionSchedule entitiy

* add UpdateDatabaseTransitionScheduleCommand

* small fixes

* change entity structure to no longer be singleton

* add get command

* fix getCache

* Change id to TransitionId enum

* more fixes

* Cleanup tests

* Add link to javadoc

* Add lastUpdateTime

* fix datatype of getCached

core/src/main/java/google/registry/model/EntityClasses.java

@@ -17,6 +17,7 @@ package google.registry.model;
 import com.google.common.collect.ImmutableSet;
 import google.registry.model.billing.BillingEvent;
 import google.registry.model.common.Cursor;
+import google.registry.model.common.DatabaseTransitionSchedule;
 import google.registry.model.common.EntityGroupRoot;
 import google.registry.model.common.GaeUserIdConverter;
 import google.registry.model.contact.ContactHistory;
@@ -74,6 +75,7 @@ public final class EntityClasses {
           ContactHistory.class,
           ContactResource.class,
           Cursor.class,
+          DatabaseTransitionSchedule.class,
           DomainBase.class,
           DomainHistory.class,
           EntityGroupRoot.class,

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

@@ -0,0 +1,159 @@
+// 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.base.Preconditions.checkNotNull;
+import static google.registry.config.RegistryConfig.getSingletonCacheRefreshDuration;
+import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
+import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+
+import com.google.common.cache.CacheBuilder;
+import com.google.common.cache.CacheLoader;
+import com.google.common.cache.LoadingCache;
+import com.google.common.collect.ImmutableSortedMap;
+import com.googlecode.objectify.Key;
+import com.googlecode.objectify.annotation.Embed;
+import com.googlecode.objectify.annotation.Entity;
+import com.googlecode.objectify.annotation.Id;
+import com.googlecode.objectify.annotation.Mapify;
+import com.googlecode.objectify.annotation.Parent;
+import google.registry.model.ImmutableObject;
+import google.registry.model.UpdateAutoTimestamp;
+import google.registry.model.common.TimedTransitionProperty.TimeMapper;
+import google.registry.model.common.TimedTransitionProperty.TimedTransition;
+import google.registry.model.registry.label.PremiumList;
+import google.registry.model.registry.label.ReservedList;
+import google.registry.model.smd.SignedMarkRevocationList;
+import google.registry.persistence.VKey;
+import google.registry.schema.replay.DatastoreOnlyEntity;
+import java.util.Optional;
+import javax.annotation.concurrent.Immutable;
+import org.joda.time.DateTime;
+
+@Entity
+@Immutable
+public class DatabaseTransitionSchedule extends ImmutableObject implements DatastoreOnlyEntity {
+
+  /**
+   * The name of the database to be treated as the primary database. The first entry in the schedule
+   * will always be Datastore.
+   */
+  public enum PrimaryDatabase {
+    CLOUD_SQL,
+    DATASTORE
+  }
+
+  /** The id of the transition schedule. */
+  public enum TransitionId {
+    /** The schedule for the migration of {@link PremiumList} and {@link ReservedList}. */
+    DOMAIN_LABEL_LISTS,
+    /** The schedule for the migration of the {@link SignedMarkRevocationList} entity. */
+    SIGNED_MARK_REVOCATION_LIST,
+  }
+
+  /**
+   * The transition to a specified primary database at a specific point in time, for use in a
+   * TimedTransitionProperty.
+   */
+  @Embed
+  public static class PrimaryDatabaseTransition extends TimedTransition<PrimaryDatabase> {
+    private PrimaryDatabase primaryDatabase;
+
+    @Override
+    protected PrimaryDatabase getValue() {
+      return primaryDatabase;
+    }
+
+    @Override
+    protected void setValue(PrimaryDatabase primaryDatabase) {
+      this.primaryDatabase = primaryDatabase;
+    }
+  }
+
+  @Parent Key<EntityGroupRoot> parent = getCrossTldKey();
+
+  @Id String transitionId;
+
+  /** An automatically managed timestamp of when this schedule was last written to Datastore. */
+  UpdateAutoTimestamp lastUpdateTime = UpdateAutoTimestamp.create(null);
+
+  /** A property that tracks the primary database for a dual-read/dual-write database migration. */
+  @Mapify(TimeMapper.class)
+  TimedTransitionProperty<PrimaryDatabase, PrimaryDatabaseTransition> databaseTransitions =
+      TimedTransitionProperty.forMapify(PrimaryDatabase.DATASTORE, PrimaryDatabaseTransition.class);
+
+  /** A cache that loads the {@link DatabaseTransitionSchedule} for a given id. */
+  private static final LoadingCache<String, Optional<DatabaseTransitionSchedule>> CACHE =
+      CacheBuilder.newBuilder()
+          .expireAfterWrite(
+              java.time.Duration.ofMillis(getSingletonCacheRefreshDuration().getMillis()))
+          .build(
+              new CacheLoader<String, Optional<DatabaseTransitionSchedule>>() {
+                @Override
+                public Optional<DatabaseTransitionSchedule> load(String transitionId) {
+                  return DatabaseTransitionSchedule.get(TransitionId.valueOf(transitionId));
+                }
+              });
+
+  public static DatabaseTransitionSchedule create(
+      TransitionId transitionId,
+      TimedTransitionProperty<PrimaryDatabase, PrimaryDatabaseTransition> databaseTransitions) {
+    checkNotNull(transitionId, "Id cannot be null");
+    checkNotNull(databaseTransitions, "databaseTransitions cannot be null");
+    databaseTransitions.checkValidity();
+    DatabaseTransitionSchedule instance = new DatabaseTransitionSchedule();
+    instance.transitionId = transitionId.name();
+    instance.databaseTransitions = databaseTransitions;
+    return instance;
+  }
+
+  /** Returns the database that is indicated as primary at the given time. */
+  public PrimaryDatabase getPrimaryDatabase() {
+    return databaseTransitions.getValueAtTime(tm().getTransactionTime());
+  }
+
+  /** Returns the database transitions as a map of start time to primary database. */
+  public ImmutableSortedMap<DateTime, PrimaryDatabase> getDatabaseTransitions() {
+    return databaseTransitions.toValueMap();
+  }
+
+  /**
+   * Returns the current cached schedule for the given id.
+   *
+   * <p>WARNING: The schedule returned by this method could be up to 10 minutes out of date.
+   */
+  public static Optional<DatabaseTransitionSchedule> getCached(TransitionId id) {
+    return CACHE.getUnchecked(id.name());
+  }
+
+  /** Returns the schedule for a given id. */
+  public static Optional<DatabaseTransitionSchedule> get(TransitionId transitionId) {
+    VKey<DatabaseTransitionSchedule> key =
+        VKey.create(
+            DatabaseTransitionSchedule.class,
+            transitionId,
+            Key.create(getCrossTldKey(), DatabaseTransitionSchedule.class, transitionId.name()));
+
+    return ofyTm().transact(() -> ofyTm().loadByKeyIfPresent(key));
+  }
+
+  @Override
+  public String toString() {
+    return String.format(
+        "%s(last updated at %s): %s",
+        transitionId, lastUpdateTime.getTimestamp(), databaseTransitions.toValueMap());
+  }
+}

core/src/main/java/google/registry/tools/GetDatabaseTransitionScheduleCommand.java

@@ -0,0 +1,44 @@
+// 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.util.PreconditionsUtils.checkArgumentPresent;
+
+import com.beust.jcommander.Parameter;
+import com.beust.jcommander.Parameters;
+import google.registry.model.common.DatabaseTransitionSchedule;
+import google.registry.model.common.DatabaseTransitionSchedule.TransitionId;
+import java.util.List;
+
+/** Command to show the {@link DatabaseTransitionSchedule} for a transition id. */
+@Parameters(separators = " =", commandDescription = "Show database transition schedule")
+final class GetDatabaseTransitionScheduleCommand implements CommandWithRemoteApi {
+
+  @Parameter(description = "Transition id(s) for the schedules to get", required = true)
+  private List<TransitionId> mainParameters;
+
+  @Override
+  public void run() {
+    for (TransitionId transitionId : mainParameters) {
+      DatabaseTransitionSchedule schedule =
+          checkArgumentPresent(
+              DatabaseTransitionSchedule.get(transitionId),
+              "A database transition schedule for %s does not exist",
+              transitionId);
+
+      System.out.println(schedule);
+    }
+  }
+}

core/src/main/java/google/registry/tools/RegistryTool.java

@@ -70,6 +70,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_transition_schedule", GetDatabaseTransitionScheduleCommand.class)
           .put("get_domain", GetDomainCommand.class)
           .put("get_history_entries", GetHistoryEntriesCommand.class)
           .put("get_host", GetHostCommand.class)
@@ -109,6 +110,7 @@ public final class RegistryTool {
           .put("resave_epp_resource", ResaveEppResourceCommand.class)
           .put("save_sql_credential", SaveSqlCredentialCommand.class)
           .put("send_escrow_report_to_icann", SendEscrowReportToIcannCommand.class)
+          .put("set_database_transition_schedule", SetDatabaseTransitionScheduleCommand.class)
           .put("set_num_instances", SetNumInstancesCommand.class)
           .put("set_sql_replay_checkpoint", SetSqlReplayCheckpointCommand.class)
           .put("setup_ote", SetupOteCommand.class)

core/src/main/java/google/registry/tools/SetDatabaseTransitionScheduleCommand.java

@@ -0,0 +1,62 @@
+// 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.Parameter;
+import com.beust.jcommander.Parameters;
+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.tools.params.TransitionListParameter.PrimaryDatabaseTransitions;
+import java.util.Optional;
+import org.joda.time.DateTime;
+
+/** Command to update {@link DatabaseTransitionSchedule}. */
+@Parameters(
+    separators = " =",
+    commandDescription = "Set the database transition schedule for transition id.")
+public class SetDatabaseTransitionScheduleCommand extends MutatingCommand {
+
+  @Parameter(
+      names = "--transition_schedule",
+      converter = PrimaryDatabaseTransitions.class,
+      validateWith = PrimaryDatabaseTransitions.class,
+      description =
+          "Comma-delimited list of database transitions, of the form"
+              + " <time>=<primary-database>[,<time>=<primary-database>]*")
+  ImmutableSortedMap<DateTime, PrimaryDatabase> transitionSchedule;
+
+  @Parameter(
+      names = "--transition_id",
+      description = "Transition id string for the schedule being updated")
+  private TransitionId transitionId;
+
+  @Override
+  protected void init() {
+    Optional<DatabaseTransitionSchedule> currentSchedule =
+        DatabaseTransitionSchedule.get(transitionId);
+
+    DatabaseTransitionSchedule newSchedule =
+        DatabaseTransitionSchedule.create(
+            transitionId,
+            TimedTransitionProperty.fromValueMap(
+                transitionSchedule, PrimaryDatabaseTransition.class));
+
+    stageEntityChange(currentSchedule.orElse(null), newSchedule);
+  }
+}

core/src/main/java/google/registry/tools/params/TransitionIdParameter.java

@@ -0,0 +1,23 @@
+// 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.params;
+
+import google.registry.model.common.DatabaseTransitionSchedule.TransitionId;
+
+/**
+ * {@link TransitionId} CLI parameter converter/validator. Required to support multi-value
+ * TransitionId parameters.
+ */
+public final class TransitionIdParameter extends EnumParameter<TransitionId> {}

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.DatabaseTransitionSchedule.PrimaryDatabase;
 import google.registry.model.domain.token.AllocationToken.TokenStatus;
 import google.registry.model.registry.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 primary database transitions. */
+  public static class PrimaryDatabaseTransitions extends TransitionListParameter<PrimaryDatabase> {
+    @Override
+    protected PrimaryDatabase parseValue(String value) {
+      return PrimaryDatabase.valueOf(value);
+    }
+  }
 }

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

@@ -0,0 +1,79 @@
+// 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);
+  }
+}

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

@@ -0,0 +1,127 @@
+// 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.FakeClock;
+import google.registry.testing.InjectExtension;
+import org.joda.time.DateTime;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+/** Unit tests for {@link GetDatabaseTransitionScheduleCommand} */
+public class GetDatabaseTransitionScheduleCommandTest
+    extends CommandTestCase<GetDatabaseTransitionScheduleCommand> {
+
+  @RegisterExtension public final InjectExtension inject = new InjectExtension();
+
+  @BeforeEach
+  void beforeEach() {
+    inject.setStaticField(
+        Ofy.class, "clock", new FakeClock(DateTime.parse("1984-12-21T06:07:08.789Z")));
+  }
+
+  @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 {
+    fakeClock.setTo(DateTime.parse("2020-10-01T00:00:00Z"));
+    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));
+    TimedTransitionProperty<PrimaryDatabase, PrimaryDatabaseTransition> databaseTransitions2 =
+        TimedTransitionProperty.fromValueMap(
+            ImmutableSortedMap.of(
+                START_OF_TIME,
+                PrimaryDatabase.DATASTORE,
+                fakeClock.nowUtc(),
+                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.789Z):"
+            + " {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");
+  }
+}

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

@@ -0,0 +1,111 @@
+// 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.model.ofy.ObjectifyService.ofy;
+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 com.google.common.collect.ImmutableSortedMap;
+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 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> {
+
+  Key<DatabaseTransitionSchedule> key;
+
+  @BeforeEach
+  void setup() {
+    key = Key.create(getCrossTldKey(), DatabaseTransitionSchedule.class, "test");
+    fakeClock.setTo(DateTime.parse("2020-12-01T00:00:00Z"));
+  }
+
+  @Test
+  void testSuccess_currentScheduleIsEmpty() throws Exception {
+    assertThat(ofy().load().key(key).now()).isNull();
+    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()).contains("Create DatabaseTransitionSchedule");
+  }
+
+  @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()).contains("Update DatabaseTransitionSchedule");
+  }
+}

core/src/test/resources/google/registry/export/backup_kinds.txt

@@ -2,6 +2,7 @@ AllocationToken
 Cancellation
 ContactResource
 Cursor
+DatabaseTransitionSchedule
 DomainBase
 EntityGroupRoot
 EppResourceIndex

core/src/test/resources/google/registry/model/schema.txt

@@ -78,6 +78,20 @@ class google.registry.model.common.Cursor {
   google.registry.model.UpdateAutoTimestamp lastUpdateTime;
   org.joda.time.DateTime cursorTime;
 }
+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;
+  org.joda.time.DateTime transitionTime;
+}
 class google.registry.model.common.EntityGroupRoot {
   @Id java.lang.String id;
   google.registry.model.UpdateAutoTimestamp updateTimestamp;