Specify explicit ofyTm usage in SetDatabaseTransitionScheduleCommand (#1081)

* Specify explicit ofyTm usage in SetDatabaseTransitionScheduleCommand

We cannot use the standard MutatingCommand because the DB schedule is
explicitly always stored in Datastore, and once we transition to
SQL-as-primary, MutatingCommand will stage the entity changes to SQL.

In addition, we remove the raw ofy() call from the test.

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

@@ -14,6 +14,8 @@
 
 package google.registry.tools;
 
+import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
+
 import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
 import com.google.common.collect.ImmutableSortedMap;
@@ -23,14 +25,13 @@ import google.registry.model.common.DatabaseTransitionSchedule.PrimaryDatabaseTr
 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 {
+public class SetDatabaseTransitionScheduleCommand extends ConfirmingCommand {
 
   @Parameter(
       names = "--transition_schedule",
@@ -48,16 +49,20 @@ public class SetDatabaseTransitionScheduleCommand extends MutatingCommand {
   private TransitionId transitionId;
 
   @Override
-  protected void init() {
-    Optional<DatabaseTransitionSchedule> currentSchedule =
-        DatabaseTransitionSchedule.get(transitionId);
+  protected String prompt() {
+    return String.format(
+        "Insert new schedule %s for transition ID %s?", transitionSchedule, transitionId);
+  }
 
+  @Override
+  protected String execute() {
     DatabaseTransitionSchedule newSchedule =
         DatabaseTransitionSchedule.create(
             transitionId,
             TimedTransitionProperty.fromValueMap(
                 transitionSchedule, PrimaryDatabaseTransition.class));
-
-    stageEntityChange(currentSchedule.orElse(null), newSchedule);
+    ofyTm().transact(() -> ofyTm().put(newSchedule));
+    return String.format(
+        "Inserted new schedule %s for transition ID %s.", transitionSchedule, transitionId);
   }
 }

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

@@ -16,7 +16,6 @@ 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;
@@ -24,12 +23,14 @@ 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;
@@ -39,11 +40,8 @@ import org.junit.jupiter.api.Test;
 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"));
   }
 
@@ -58,7 +56,13 @@ public class SetDatabaseTransitionScheduleCommandTest
 
   @Test
   void testSuccess_currentScheduleIsEmpty() throws Exception {
-    assertThat(ofy().load().key(key).now()).isNull();
+    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");
@@ -70,7 +74,10 @@ public class SetDatabaseTransitionScheduleCommandTest
                             .get()
                             .getPrimaryDatabase()))
         .isEqualTo(PrimaryDatabase.DATASTORE);
-    assertThat(command.prompt()).contains("Create DatabaseTransitionSchedule");
+    assertThat(command.prompt())
+        .isEqualTo(
+            "Insert new schedule {1970-01-01T00:00:00.000Z=DATASTORE} "
+                + "for transition ID SIGNED_MARK_REVOCATION_LIST?");
   }
 
   @Test
@@ -92,7 +99,8 @@ public class SetDatabaseTransitionScheduleCommandTest
         .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");
+        "--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(
@@ -117,6 +125,10 @@ public class SetDatabaseTransitionScheduleCommandTest
                             .get()
                             .getPrimaryDatabase()))
         .isEqualTo(PrimaryDatabase.DATASTORE);
-    assertThat(command.prompt()).contains("Update DatabaseTransitionSchedule");
+    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?");
   }
 }