Revert "Remove DatabaseMigrationStateSchedule (#1689)" (#1698)

This reverts commit 18d51738ce.

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

@@ -0,0 +1,272 @@
+// 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.checkArgument;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.util.DateTimeUtils.START_OF_TIME;
+
+import com.github.benmanes.caffeine.cache.LoadingCache;
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.ImmutableMultimap;
+import com.google.common.collect.ImmutableSortedMap;
+import com.google.common.flogger.FluentLogger;
+import google.registry.config.RegistryEnvironment;
+import google.registry.model.CacheUtils;
+import google.registry.model.annotations.DeleteAfterMigration;
+import google.registry.model.common.TimedTransitionProperty.TimedTransition;
+import java.time.Duration;
+import java.util.Arrays;
+import javax.persistence.Entity;
+import javax.persistence.PersistenceException;
+import org.joda.time.DateTime;
+
+/**
+ * A wrapper object representing the stage-to-time mapping of the Registry 3.0 Cloud SQL migration.
+ *
+ * <p>The entity is stored in SQL throughout the entire migration so as to have a single point of
+ * access.
+ */
+@DeleteAfterMigration
+@Entity
+public class DatabaseMigrationStateSchedule extends CrossTldSingleton {
+
+  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
+
+  public enum PrimaryDatabase {
+    CLOUD_SQL,
+    DATASTORE
+  }
+
+  public enum ReplayDirection {
+    NO_REPLAY,
+    DATASTORE_TO_SQL,
+    SQL_TO_DATASTORE
+  }
+
+  /**
+   * The current phase of the migration plus information about which database to use and whether or
+   * not the phase is read-only.
+   */
+  public enum MigrationState {
+    /** Datastore is the only DB being used. */
+    DATASTORE_ONLY(PrimaryDatabase.DATASTORE, false, ReplayDirection.NO_REPLAY),
+
+    /** Datastore is the primary DB, with changes replicated to Cloud SQL. */
+    DATASTORE_PRIMARY(PrimaryDatabase.DATASTORE, false, ReplayDirection.DATASTORE_TO_SQL),
+
+    /** Datastore is the primary DB, with replication, and async actions are disallowed. */
+    DATASTORE_PRIMARY_NO_ASYNC(PrimaryDatabase.DATASTORE, false, ReplayDirection.DATASTORE_TO_SQL),
+
+    /** Datastore is the primary DB, with replication, and all mutating actions are disallowed. */
+    DATASTORE_PRIMARY_READ_ONLY(PrimaryDatabase.DATASTORE, true, ReplayDirection.DATASTORE_TO_SQL),
+
+    /**
+     * Cloud SQL is the primary DB, with replication back to Datastore, and all mutating actions are
+     * disallowed.
+     */
+    SQL_PRIMARY_READ_ONLY(PrimaryDatabase.CLOUD_SQL, true, ReplayDirection.SQL_TO_DATASTORE),
+
+    /** Cloud SQL is the primary DB, with changes replicated to Datastore. */
+    SQL_PRIMARY(PrimaryDatabase.CLOUD_SQL, false, ReplayDirection.SQL_TO_DATASTORE),
+
+    /** Cloud SQL is the only DB being used. */
+    SQL_ONLY(PrimaryDatabase.CLOUD_SQL, false, ReplayDirection.NO_REPLAY);
+
+    private final PrimaryDatabase primaryDatabase;
+    private final boolean isReadOnly;
+    private final ReplayDirection replayDirection;
+
+    public PrimaryDatabase getPrimaryDatabase() {
+      return primaryDatabase;
+    }
+
+    public boolean isReadOnly() {
+      return isReadOnly;
+    }
+
+    public ReplayDirection getReplayDirection() {
+      return replayDirection;
+    }
+
+    MigrationState(
+        PrimaryDatabase primaryDatabase, boolean isReadOnly, ReplayDirection replayDirection) {
+      this.primaryDatabase = primaryDatabase;
+      this.isReadOnly = isReadOnly;
+      this.replayDirection = replayDirection;
+    }
+  }
+
+  public static class MigrationStateTransition extends TimedTransition<MigrationState> {
+    private MigrationState migrationState;
+
+    @Override
+    public MigrationState getValue() {
+      return migrationState;
+    }
+
+    @Override
+    protected void setValue(MigrationState migrationState) {
+      this.migrationState = migrationState;
+    }
+  }
+
+  /**
+   * Cache of the current migration schedule. The key is meaningless; this is essentially a memoized
+   * Supplier that can be reset for testing purposes and after writes.
+   */
+  @VisibleForTesting
+  public static final LoadingCache<
+          Class<DatabaseMigrationStateSchedule>,
+          TimedTransitionProperty<MigrationState, MigrationStateTransition>>
+      // Each instance should cache the migration schedule for five minutes before reloading
+      CACHE =
+          CacheUtils.newCacheBuilder(Duration.ofMinutes(5))
+              .build(singletonClazz -> DatabaseMigrationStateSchedule.getUncached());
+
+  // Restrictions on the state transitions, e.g. no going from DATASTORE_ONLY to SQL_ONLY
+  private static final ImmutableMultimap<MigrationState, MigrationState> VALID_STATE_TRANSITIONS =
+      createValidStateTransitions();
+
+  /**
+   * The valid state transitions. Generally, one can advance the state one step or move backward any
+   * number of steps, as long as the step we're moving back to has the same primary database as the
+   * one we're in. Otherwise, we must move to the corresponding READ_ONLY stage first.
+   */
+  private static ImmutableMultimap<MigrationState, MigrationState> createValidStateTransitions() {
+    ImmutableMultimap.Builder<MigrationState, MigrationState> builder =
+        new ImmutableMultimap.Builder<MigrationState, MigrationState>()
+            .put(MigrationState.DATASTORE_ONLY, MigrationState.DATASTORE_PRIMARY)
+            .putAll(
+                MigrationState.DATASTORE_PRIMARY,
+                MigrationState.DATASTORE_ONLY,
+                MigrationState.DATASTORE_PRIMARY_NO_ASYNC)
+            .putAll(
+                MigrationState.DATASTORE_PRIMARY_NO_ASYNC,
+                MigrationState.DATASTORE_ONLY,
+                MigrationState.DATASTORE_PRIMARY,
+                MigrationState.DATASTORE_PRIMARY_READ_ONLY)
+            .putAll(
+                MigrationState.DATASTORE_PRIMARY_READ_ONLY,
+                MigrationState.DATASTORE_ONLY,
+                MigrationState.DATASTORE_PRIMARY,
+                MigrationState.DATASTORE_PRIMARY_NO_ASYNC,
+                MigrationState.SQL_PRIMARY_READ_ONLY,
+                MigrationState.SQL_PRIMARY)
+            .putAll(
+                MigrationState.SQL_PRIMARY_READ_ONLY,
+                MigrationState.DATASTORE_PRIMARY_READ_ONLY,
+                MigrationState.SQL_PRIMARY)
+            .putAll(
+                MigrationState.SQL_PRIMARY,
+                MigrationState.SQL_PRIMARY_READ_ONLY,
+                MigrationState.SQL_ONLY)
+            .putAll(
+                MigrationState.SQL_ONLY,
+                MigrationState.SQL_PRIMARY_READ_ONLY,
+                MigrationState.SQL_PRIMARY);
+
+    // In addition, we can always transition from a state to itself (useful when updating the map).
+    Arrays.stream(MigrationState.values()).forEach(state -> builder.put(state, state));
+    return builder.build();
+  }
+
+  // Default map to return if we have never saved any -- only use Datastore.
+  @VisibleForTesting
+  public static final TimedTransitionProperty<MigrationState, MigrationStateTransition>
+      DEFAULT_TRANSITION_MAP =
+          TimedTransitionProperty.fromValueMap(
+              ImmutableSortedMap.of(START_OF_TIME, MigrationState.DATASTORE_ONLY),
+              MigrationStateTransition.class);
+
+  @VisibleForTesting
+  public TimedTransitionProperty<MigrationState, MigrationStateTransition> migrationTransitions =
+      TimedTransitionProperty.forMapify(
+          MigrationState.DATASTORE_ONLY, MigrationStateTransition.class);
+
+  // Required for Objectify initialization
+  private DatabaseMigrationStateSchedule() {}
+
+  @VisibleForTesting
+  public DatabaseMigrationStateSchedule(
+      TimedTransitionProperty<MigrationState, MigrationStateTransition> migrationTransitions) {
+    this.migrationTransitions = migrationTransitions;
+  }
+
+  /** Sets and persists to SQL the provided migration transition schedule. */
+  public static void set(ImmutableSortedMap<DateTime, MigrationState> migrationTransitionMap) {
+    jpaTm().assertInTransaction();
+    TimedTransitionProperty<MigrationState, MigrationStateTransition> transitions =
+        TimedTransitionProperty.make(
+            migrationTransitionMap,
+            MigrationStateTransition.class,
+            VALID_STATE_TRANSITIONS,
+            "validStateTransitions",
+            MigrationState.DATASTORE_ONLY,
+            "migrationTransitionMap must start with DATASTORE_ONLY");
+    validateTransitionAtCurrentTime(transitions);
+    jpaTm().putWithoutBackup(new DatabaseMigrationStateSchedule(transitions));
+    CACHE.invalidateAll();
+  }
+
+  /** Loads the currently-set migration schedule from the cache, or the default if none exists. */
+  public static TimedTransitionProperty<MigrationState, MigrationStateTransition> get() {
+    return CACHE.get(DatabaseMigrationStateSchedule.class);
+  }
+
+  /** Returns the database migration status at the given time. */
+  public static MigrationState getValueAtTime(DateTime dateTime) {
+    return get().getValueAtTime(dateTime);
+  }
+
+  /** Loads the currently-set migration schedule from SQL, or the default if none exists. */
+  @VisibleForTesting
+  static TimedTransitionProperty<MigrationState, MigrationStateTransition> getUncached() {
+    return jpaTm()
+        .transactWithoutBackup(
+            () -> {
+              try {
+                return jpaTm()
+                    .loadSingleton(DatabaseMigrationStateSchedule.class)
+                    .map(s -> s.migrationTransitions)
+                    .orElse(DEFAULT_TRANSITION_MAP);
+              } catch (PersistenceException e) {
+                if (!RegistryEnvironment.get().equals(RegistryEnvironment.UNITTEST)) {
+                  throw e;
+                }
+                logger.atWarning().withCause(e).log(
+                    "Error when retrieving migration schedule; this should only happen in tests.");
+                return DEFAULT_TRANSITION_MAP;
+              }
+            });
+  }
+
+  /**
+   * A provided map of transitions may be valid by itself (i.e. it shifts states properly, doesn't
+   * skip states, and doesn't backtrack incorrectly) while still being invalid. In addition to the
+   * transitions in the map being valid, the single transition from the current map at the current
+   * time to the new map at the current time must also be valid.
+   */
+  private static void validateTransitionAtCurrentTime(
+      TimedTransitionProperty<MigrationState, MigrationStateTransition> newTransitions) {
+    MigrationState currentValue = getUncached().getValueAtTime(jpaTm().getTransactionTime());
+    MigrationState nextCurrentValue = newTransitions.getValueAtTime(jpaTm().getTransactionTime());
+    checkArgument(
+        VALID_STATE_TRANSITIONS.get(currentValue).contains(nextCurrentValue),
+        "Cannot transition from current state-as-of-now %s to new state-as-of-now %s",
+        currentValue,
+        nextCurrentValue);
+  }
+}

core/src/main/java/google/registry/persistence/converter/DatabaseMigrationScheduleTransitionConverter.java

@@ -0,0 +1,48 @@
+// 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.persistence.converter;
+
+import com.google.common.collect.Maps;
+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 java.util.Map;
+import javax.persistence.Converter;
+import org.joda.time.DateTime;
+
+/** JPA converter for {@link DatabaseMigrationStateSchedule} transitions. */
+@DeleteAfterMigration
+@Converter(autoApply = true)
+public class DatabaseMigrationScheduleTransitionConverter
+    extends TimedTransitionPropertyConverterBase<MigrationState, MigrationStateTransition> {
+
+  @Override
+  Map.Entry<String, String> convertToDatabaseMapEntry(
+      Map.Entry<DateTime, MigrationStateTransition> entry) {
+    return Maps.immutableEntry(entry.getKey().toString(), entry.getValue().getValue().name());
+  }
+
+  @Override
+  Map.Entry<DateTime, MigrationState> convertToEntityMapEntry(Map.Entry<String, String> entry) {
+    return Maps.immutableEntry(
+        DateTime.parse(entry.getKey()), MigrationState.valueOf(entry.getValue()));
+  }
+
+  @Override
+  Class<MigrationStateTransition> getTimedTransitionSubclass() {
+    return MigrationStateTransition.class;
+  }
+}

core/src/main/resources/META-INF/persistence.xml

@@ -42,6 +42,7 @@
     <class>google.registry.model.billing.BillingEvent$OneTime</class>
     <class>google.registry.model.billing.BillingEvent$Recurring</class>
     <class>google.registry.model.common.Cursor</class>
+    <class>google.registry.model.common.DatabaseMigrationStateSchedule</class>
     <class>google.registry.model.contact.ContactHistory</class>
     <class>google.registry.model.contact.ContactResource</class>
     <class>google.registry.model.domain.DomainBase</class>
@@ -84,6 +85,7 @@
     <class>google.registry.persistence.converter.CidrAddressBlockListConverter</class>
     <class>google.registry.persistence.converter.CurrencyToBillingConverter</class>
     <class>google.registry.persistence.converter.CurrencyUnitConverter</class>
+    <class>google.registry.persistence.converter.DatabaseMigrationScheduleTransitionConverter</class>
     <class>google.registry.persistence.converter.DateTimeConverter</class>
     <class>google.registry.persistence.converter.DurationConverter</class>
     <class>google.registry.persistence.converter.InetAddressSetConverter</class>

core/src/test/java/google/registry/persistence/converter/DatabaseMigrationScheduleTransitionConverterTest.java

@@ -0,0 +1,91 @@
+// 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.persistence.converter;
+
+import static com.google.common.truth.Truth.assertThat;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.testing.DatabaseHelper.insertInDb;
+import static google.registry.util.DateTimeUtils.START_OF_TIME;
+
+import com.google.common.collect.ImmutableSortedMap;
+import google.registry.model.ImmutableObject;
+import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState;
+import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationStateTransition;
+import google.registry.model.common.TimedTransitionProperty;
+import google.registry.persistence.transaction.JpaTestExtensions;
+import google.registry.persistence.transaction.JpaTestExtensions.JpaUnitTestExtension;
+import javax.persistence.Entity;
+import javax.persistence.Id;
+import org.joda.time.DateTime;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+/** Unit tests for {@link DatabaseMigrationScheduleTransitionConverter}. */
+public class DatabaseMigrationScheduleTransitionConverterTest {
+
+  @RegisterExtension
+  public final JpaUnitTestExtension jpa =
+      new JpaTestExtensions.Builder()
+          .withEntityClass(DatabaseMigrationScheduleTransitionConverterTestEntity.class)
+          .buildUnitTestExtension();
+
+  private static final ImmutableSortedMap<DateTime, MigrationState> values =
+      ImmutableSortedMap.of(
+          START_OF_TIME,
+          MigrationState.DATASTORE_ONLY,
+          DateTime.parse("2001-01-01T00:00:00.0Z"),
+          MigrationState.DATASTORE_PRIMARY,
+          DateTime.parse("2002-01-01T01:00:00.0Z"),
+          MigrationState.DATASTORE_PRIMARY_NO_ASYNC,
+          DateTime.parse("2002-01-01T02:00:00.0Z"),
+          MigrationState.DATASTORE_PRIMARY_READ_ONLY,
+          DateTime.parse("2002-01-02T00:00:00.0Z"),
+          MigrationState.SQL_PRIMARY,
+          DateTime.parse("2002-01-03T00:00:00.0Z"),
+          MigrationState.SQL_ONLY);
+
+  @Test
+  void roundTripConversion_returnsSameTimedTransitionProperty() {
+    TimedTransitionProperty<MigrationState, MigrationStateTransition> timedTransitionProperty =
+        TimedTransitionProperty.fromValueMap(values, MigrationStateTransition.class);
+    DatabaseMigrationScheduleTransitionConverterTestEntity testEntity =
+        new DatabaseMigrationScheduleTransitionConverterTestEntity(timedTransitionProperty);
+    insertInDb(testEntity);
+    DatabaseMigrationScheduleTransitionConverterTestEntity persisted =
+        jpaTm()
+            .transact(
+                () ->
+                    jpaTm()
+                        .getEntityManager()
+                        .find(DatabaseMigrationScheduleTransitionConverterTestEntity.class, "id"));
+    assertThat(persisted.timedTransitionProperty).containsExactlyEntriesIn(timedTransitionProperty);
+  }
+
+  @Entity
+  private static class DatabaseMigrationScheduleTransitionConverterTestEntity
+      extends ImmutableObject {
+
+    @Id String name = "id";
+
+    TimedTransitionProperty<MigrationState, MigrationStateTransition> timedTransitionProperty;
+
+    private DatabaseMigrationScheduleTransitionConverterTestEntity() {}
+
+    private DatabaseMigrationScheduleTransitionConverterTestEntity(
+        TimedTransitionProperty<MigrationState, MigrationStateTransition> timedTransitionProperty) {
+      this.timedTransitionProperty = timedTransitionProperty;
+    }
+  }
+}

core/src/test/java/google/registry/persistence/transaction/JpaEntityCoverageExtension.java

@@ -41,6 +41,11 @@ public class JpaEntityCoverageExtension implements BeforeEachCallback, AfterEach
 
   private static final ImmutableSet<String> IGNORE_ENTITIES =
       ImmutableSet.of(
+          // DatabaseMigrationStateSchedule is persisted in tests, however any test that sets it
+          // needs to remove it in order to avoid affecting any other tests running in the same JVM.
+          // TODO(gbrodman): remove this when we implement proper read-only modes for the
+          // transaction managers.
+          "DatabaseMigrationStateSchedule",
           // TransactionEntity is trivial; its persistence is tested in TransactionTest.
           "TransactionEntity");
 

db/src/main/resources/sql/flyway.txt

@@ -116,4 +116,3 @@ V115__add_renewal_columns_to_billing_recurrence.sql
 V116__add_renewal_column_to_allocation_token.sql
 V117__add_billing_recurrence_last_expansion_column.sql
 V118__drop_billing_identifier_column_from_registrar.sql
-V119__drop_database_migration_schedule.sql

db/src/main/resources/sql/flyway/V119__drop_database_migration_schedule.sql

@@ -1,15 +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.
-
-DROP TABLE IF EXISTS "DatabaseMigrationStateSchedule" CASCADE;

db/src/main/resources/sql/schema/db-schema.sql.generated

@@ -242,6 +242,12 @@
         primary key (scope, type)
     );
 
+    create table "DatabaseMigrationStateSchedule" (
+       id int8 not null,
+        migration_transitions hstore,
+        primary key (id)
+    );
+
     create table "DelegationSignerData" (
        algorithm int4 not null,
         digest bytea not null,

db/src/main/resources/sql/schema/nomulus.golden.sql

@@ -320,6 +320,16 @@ CREATE TABLE public."Cursor" (
 );
 
 
+--
+-- Name: DatabaseMigrationStateSchedule; Type: TABLE; Schema: public; Owner: -
+--
+
+CREATE TABLE public."DatabaseMigrationStateSchedule" (
+    id bigint NOT NULL,
+    migration_transitions public.hstore
+);
+
+
 --
 -- Name: DelegationSignerData; Type: TABLE; Schema: public; Owner: -
 --
@@ -1215,6 +1225,14 @@ ALTER TABLE ONLY public."Cursor"
     ADD CONSTRAINT "Cursor_pkey" PRIMARY KEY (scope, type);
 
 
+--
+-- Name: DatabaseMigrationStateSchedule DatabaseMigrationStateSchedule_pkey; Type: CONSTRAINT; Schema: public; Owner: -
+--
+
+ALTER TABLE ONLY public."DatabaseMigrationStateSchedule"
+    ADD CONSTRAINT "DatabaseMigrationStateSchedule_pkey" PRIMARY KEY (id);
+
+
 --
 -- Name: DelegationSignerData DelegationSignerData_pkey; Type: CONSTRAINT; Schema: public; Owner: -
 --
@@ -1462,6 +1480,13 @@ ALTER TABLE ONLY public."DomainHistoryHost"
 CREATE INDEX allocation_token_domain_name_idx ON public."AllocationToken" USING btree (domain_name);
 
 
+--
+-- Name: database_migration_state_schedule_singleton; Type: INDEX; Schema: public; Owner: -
+--
+
+CREATE UNIQUE INDEX database_migration_state_schedule_singleton ON public."DatabaseMigrationStateSchedule" USING btree ((true));
+
+
 --
 -- Name: domain_dns_refresh_request_time_idx; Type: INDEX; Schema: public; Owner: -
 --