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);
+    }
+  }
 }