Add an object to store database migration stages (#1170)

* Add an object to store database migration stages

go/registry-3.0-stage-management for more details

This basically boils down to storing an enum in the database so that we
can tell what stage of the migration we're in.

We use a cross-TLD parent so that we can have strong transactional
consistency on retrieval.

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.DatabaseMigrationStateWrapper;
 import google.registry.model.common.DatabaseTransitionSchedule;
 import google.registry.model.common.EntityGroupRoot;
 import google.registry.model.common.GaeUserIdConverter;
@@ -75,6 +76,7 @@ public final class EntityClasses {
           ContactHistory.class,
           ContactResource.class,
           Cursor.class,
+          DatabaseMigrationStateWrapper.class,
           DatabaseTransitionSchedule.class,
           DomainBase.class,
           DomainHistory.class,

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

@@ -0,0 +1,122 @@
+// 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.ofyTm;
+
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import com.googlecode.objectify.annotation.Entity;
+import google.registry.model.annotations.InCrossTld;
+import google.registry.model.common.DatabaseTransitionSchedule.PrimaryDatabase;
+import google.registry.schema.replay.DatastoreOnlyEntity;
+
+/**
+ * A wrapper object representing the current stage of the Registry 3.0 Cloud SQL migration.
+ *
+ * <p>The entity is stored in Datastore throughout the entire migration so as to have a single point
+ * of access (avoiding a two-phase commit problem).
+ */
+@Entity
+@InCrossTld
+public class DatabaseMigrationStateWrapper extends CrossTldSingleton
+    implements DatastoreOnlyEntity {
+
+  /**
+   * 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_ONLY(PrimaryDatabase.DATASTORE, false),
+    DATASTORE_PRIMARY(PrimaryDatabase.DATASTORE, false),
+    DATASTORE_PRIMARY_READ_ONLY(PrimaryDatabase.DATASTORE, true),
+    SQL_PRIMARY(PrimaryDatabase.CLOUD_SQL, false),
+    SQL_ONLY(PrimaryDatabase.CLOUD_SQL, false);
+
+    private final PrimaryDatabase primaryDatabase;
+    private final boolean readOnly;
+
+    public PrimaryDatabase getPrimaryDatabase() {
+      return primaryDatabase;
+    }
+
+    public boolean isReadOnly() {
+      return readOnly;
+    }
+
+    MigrationState(PrimaryDatabase primaryDatabase, boolean readOnly) {
+      this.primaryDatabase = primaryDatabase;
+      this.readOnly = readOnly;
+    }
+  }
+
+  private MigrationState migrationState;
+
+  // Required for Objectify initialization
+  private DatabaseMigrationStateWrapper() {}
+
+  DatabaseMigrationStateWrapper(MigrationState migrationState) {
+    this.migrationState = migrationState;
+  }
+
+  // The valid state transitions. Basically, at state N, state N+1 is valid as well as all previous
+  // states, with one type of exception: when in either of the SQL states, we can only move back
+  // one step so that we can make sure that any modifications have been replayed back to Datastore.
+  private static final ImmutableMap<MigrationState, ImmutableSet<MigrationState>>
+      VALID_STATE_TRANSITIONS =
+          ImmutableMap.of(
+              MigrationState.DATASTORE_ONLY,
+              ImmutableSet.of(MigrationState.DATASTORE_PRIMARY),
+              MigrationState.DATASTORE_PRIMARY,
+              ImmutableSet.of(
+                  MigrationState.DATASTORE_ONLY, MigrationState.DATASTORE_PRIMARY_READ_ONLY),
+              MigrationState.DATASTORE_PRIMARY_READ_ONLY,
+              ImmutableSet.of(
+                  MigrationState.DATASTORE_ONLY,
+                  MigrationState.DATASTORE_PRIMARY,
+                  MigrationState.SQL_PRIMARY),
+              MigrationState.SQL_PRIMARY,
+              ImmutableSet.of(MigrationState.DATASTORE_PRIMARY_READ_ONLY, MigrationState.SQL_ONLY),
+              MigrationState.SQL_ONLY,
+              ImmutableSet.of(MigrationState.SQL_PRIMARY));
+
+  private static boolean isValidStateTransition(MigrationState from, MigrationState to) {
+    return VALID_STATE_TRANSITIONS.get(from).contains(to);
+  }
+
+  /** Sets and persists to Datastore the current database migration state. */
+  public static void set(MigrationState newState) {
+    MigrationState currentState = get();
+    checkArgument(
+        isValidStateTransition(currentState, newState),
+        "Moving from migration state %s to %s is not a valid transition",
+        currentState,
+        newState);
+    DatabaseMigrationStateWrapper wrapper = new DatabaseMigrationStateWrapper(newState);
+    ofyTm().transact(() -> ofyTm().put(wrapper));
+  }
+
+  /** Retrieves the current state of the migration (or DATASTORE_ONLY if it hasn't started). */
+  public static MigrationState get() {
+    return ofyTm()
+        .transact(
+            () ->
+                ofyTm()
+                    .loadSingleton(DatabaseMigrationStateWrapper.class)
+                    .map(s -> s.migrationState)
+                    .orElse(MigrationState.DATASTORE_ONLY));
+  }
+}

core/src/test/java/google/registry/model/common/DatabaseMigrationStateWrapperTest.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.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);
+  }
+}

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

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

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

@@ -1,6 +1,7 @@
 ClaimsList
 ClaimsListSingleton
 Cursor
+DatabaseMigrationStateWrapper
 DatabaseTransitionSchedule
 KmsSecret
 KmsSecretRevision

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

@@ -78,6 +78,18 @@ class google.registry.model.common.Cursor {
   google.registry.model.UpdateAutoTimestamp lastUpdateTime;
   org.joda.time.DateTime cursorTime;
 }
+class google.registry.model.common.DatabaseMigrationStateWrapper {
+  @Id long id;
+  @Parent com.googlecode.objectify.Key<google.registry.model.common.EntityGroupRoot> parent;
+  google.registry.model.common.DatabaseMigrationStateWrapper$MigrationState migrationState;
+}
+enum google.registry.model.common.DatabaseMigrationStateWrapper$MigrationState {
+  DATASTORE_ONLY;
+  DATASTORE_PRIMARY;
+  DATASTORE_PRIMARY_READ_ONLY;
+  SQL_ONLY;
+  SQL_PRIMARY;
+}
 class google.registry.model.common.DatabaseTransitionSchedule {
   @Id java.lang.String transitionId;
   @Parent com.googlecode.objectify.Key<google.registry.model.common.EntityGroupRoot> parent;