Fix access to a nullable field in HistoryEntry (#1193)

* Fix access to a nullable field in HistoryEntry

core/src/main/java/google/registry/backup/ReplayCommitLogsToSqlAction.java

@@ -27,7 +27,9 @@ import com.google.appengine.tools.cloudstorage.GcsFileMetadata;
 import com.google.appengine.tools.cloudstorage.GcsService;
 import com.google.common.collect.ImmutableList;
 import com.google.common.flogger.FluentLogger;
-import google.registry.config.RegistryConfig;
+import google.registry.model.common.DatabaseMigrationStateSchedule;
+import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState;
+import google.registry.model.common.DatabaseMigrationStateSchedule.ReplayDirection;
 import google.registry.model.server.Lock;
 import google.registry.model.translators.VKeyTranslatorFactory;
 import google.registry.persistence.VKey;
@@ -39,6 +41,7 @@ import google.registry.schema.replay.DatastoreOnlyEntity;
 import google.registry.schema.replay.NonReplicatedEntity;
 import google.registry.schema.replay.ReplaySpecializer;
 import google.registry.schema.replay.SqlReplayCheckpoint;
+import google.registry.util.Clock;
 import google.registry.util.RequestStatusChecker;
 import java.io.IOException;
 import java.io.InputStream;
@@ -69,15 +72,19 @@ public class ReplayCommitLogsToSqlAction implements Runnable {
   @Inject Response response;
   @Inject RequestStatusChecker requestStatusChecker;
   @Inject GcsDiffFileLister diffLister;
+  @Inject Clock clock;
 
   @Inject
   ReplayCommitLogsToSqlAction() {}
 
   @Override
   public void run() {
-    if (!RegistryConfig.getCloudSqlReplayCommitLogs()) {
-      String message = "ReplayCommitLogsToSqlAction was called but disabled in the config.";
-      logger.atWarning().log(message);
+    MigrationState state = DatabaseMigrationStateSchedule.getValueAtTime(clock.nowUtc());
+    if (!state.getReplayDirection().equals(ReplayDirection.DATASTORE_TO_SQL)) {
+      String message =
+          String.format(
+              "Skipping ReplayCommitLogsToSqlAction because we are in migration phase %s.", state);
+      logger.atInfo().log(message);
       // App Engine will retry on any non-2xx status code, which we don't want in this case.
       response.setStatus(SC_NO_CONTENT);
       response.setPayload(message);

core/src/main/java/google/registry/batch/DeleteExpiredDomainsAction.java

@@ -17,8 +17,8 @@ package google.registry.batch;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.net.MediaType.PLAIN_TEXT_UTF_8;
 import static google.registry.flows.FlowUtils.marshalWithLenientRetry;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 import static google.registry.util.DateTimeUtils.END_OF_TIME;
 import static google.registry.util.ResourceUtils.readResourceUtf8;
 import static java.nio.charset.StandardCharsets.UTF_8;
@@ -36,6 +36,7 @@ import google.registry.flows.StatelessRequestSessionMetadata;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.eppcommon.ProtocolDefinition;
 import google.registry.model.eppoutput.EppOutput;
+import google.registry.persistence.transaction.QueryComposer.Comparator;
 import google.registry.request.Action;
 import google.registry.request.Response;
 import google.registry.request.auth.Auth;
@@ -128,12 +129,15 @@ public class DeleteExpiredDomainsAction implements Runnable {
     logger.atInfo().log(
         "Deleting non-renewing domains with autorenew end times up through %s.", runTime);
 
-    // Note: This query is (and must be) non-transactional, and thus, is only eventually consistent.
+    // Note: in Datastore, this query is (and must be) non-transactional, and thus, is only
+    // eventually consistent.
     ImmutableList<DomainBase> domainsToDelete =
-        ofy().load().type(DomainBase.class).filter("autorenewEndTime <=", runTime).list().stream()
-            // Datastore can't do two inequalities in one query, so the second happens in-memory.
-            .filter(d -> d.getDeletionTime().isEqual(END_OF_TIME))
-            .collect(toImmutableList());
+        transactIfJpaTm(
+            () ->
+                tm().createQueryComposer(DomainBase.class)
+                    .where("autorenewEndTime", Comparator.LTE, runTime)
+                    .where("deletionTime", Comparator.EQ, END_OF_TIME)
+                    .list());
     if (domainsToDelete.isEmpty()) {
       logger.atInfo().log("Found 0 domains to delete.");
       response.setPayload("Found 0 domains to delete.");

core/src/main/java/google/registry/batch/ExpandRecurringBillingEventsAction.java

@@ -21,9 +21,12 @@ import static google.registry.mapreduce.MapreduceRunner.PARAM_DRY_RUN;
 import static google.registry.mapreduce.inputs.EppResourceInputs.createChildEntityInput;
 import static google.registry.model.common.Cursor.CursorType.RECURRING_BILLING;
 import static google.registry.model.domain.Period.Unit.YEARS;
-import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.model.ofy.ObjectifyService.auditedOfy;
 import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_AUTORENEW;
+import static google.registry.persistence.transaction.QueryComposer.Comparator.EQ;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 import static google.registry.pricing.PricingEngineProxy.getDomainRenewCost;
 import static google.registry.util.CollectionUtils.union;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
@@ -38,10 +41,8 @@ import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Range;
 import com.google.common.collect.Streams;
 import com.google.common.flogger.FluentLogger;
-import com.googlecode.objectify.Key;
 import google.registry.mapreduce.MapreduceRunner;
 import google.registry.mapreduce.inputs.NullInput;
-import google.registry.model.EppResource;
 import google.registry.model.ImmutableObject;
 import google.registry.model.billing.BillingEvent;
 import google.registry.model.billing.BillingEvent.Flag;
@@ -54,7 +55,7 @@ import google.registry.model.domain.Period;
 import google.registry.model.registry.Registry;
 import google.registry.model.reporting.DomainTransactionRecord;
 import google.registry.model.reporting.DomainTransactionRecord.TransactionReportField;
-import google.registry.model.reporting.HistoryEntry;
+import google.registry.persistence.VKey;
 import google.registry.request.Action;
 import google.registry.request.Parameter;
 import google.registry.request.Response;
@@ -92,31 +93,87 @@ public class ExpandRecurringBillingEventsAction implements Runnable {
 
   @Override
   public void run() {
-    Cursor cursor =
-        tm().loadByKeyIfPresent(Cursor.createGlobalVKey(RECURRING_BILLING)).orElse(null);
     DateTime executeTime = clock.nowUtc();
-    DateTime persistedCursorTime = (cursor == null ? START_OF_TIME : cursor.getCursorTime());
+    DateTime persistedCursorTime =
+        transactIfJpaTm(
+            () ->
+                tm().loadByKeyIfPresent(Cursor.createGlobalVKey(RECURRING_BILLING))
+                    .orElse(Cursor.createGlobal(RECURRING_BILLING, START_OF_TIME))
+                    .getCursorTime());
     DateTime cursorTime = cursorTimeParam.orElse(persistedCursorTime);
     checkArgument(
-        cursorTime.isBefore(executeTime),
-        "Cursor time must be earlier than execution time.");
+        cursorTime.isBefore(executeTime), "Cursor time must be earlier than execution time.");
     logger.atInfo().log(
         "Running Recurring billing event expansion for billing time range [%s, %s).",
         cursorTime, executeTime);
-    mrRunner
-        .setJobName("Expand Recurring billing events into synthetic OneTime events.")
-        .setModuleName("backend")
-        .runMapreduce(
-            new ExpandRecurringBillingEventsMapper(isDryRun, cursorTime, clock.nowUtc()),
-            new ExpandRecurringBillingEventsReducer(isDryRun, persistedCursorTime),
-            // Add an extra shard that maps over a null recurring event (see the mapper for why).
-            ImmutableList.of(
-                new NullInput<>(),
-                createChildEntityInput(
-                    ImmutableSet.of(DomainBase.class), ImmutableSet.of(Recurring.class))))
-        .sendLinkToMapreduceConsole(response);
-  }
+    if (tm().isOfy()) {
+      mrRunner
+          .setJobName("Expand Recurring billing events into synthetic OneTime events.")
+          .setModuleName("backend")
+          .runMapreduce(
+              new ExpandRecurringBillingEventsMapper(isDryRun, cursorTime, clock.nowUtc()),
+              new ExpandRecurringBillingEventsReducer(isDryRun, persistedCursorTime),
+              // Add an extra shard that maps over a null recurring event (see the mapper for why).
+              ImmutableList.of(
+                  new NullInput<>(),
+                  createChildEntityInput(
+                      ImmutableSet.of(DomainBase.class), ImmutableSet.of(Recurring.class))))
+          .sendLinkToMapreduceConsole(response);
+    } else {
+      int numBillingEventsSaved =
+          jpaTm()
+              .transact(
+                  () ->
+                      jpaTm()
+                          .query(
+                              "FROM BillingRecurrence "
+                                  + "WHERE event_time <= :executeTime "
+                                  + "AND event_time < recurrence_end_time",
+                              Recurring.class)
+                          .setParameter("executeTime", executeTime.toDate())
+                          // Need to get a list from the transaction and then convert it to a stream
+                          // for further processing. If we get a stream directly, each elements gets
+                          // processed downstream eagerly but Hibernate returns a
+                          // ScrollableResultsIterator that cannot be advanced outside the
+                          // transaction, resulting in an exception.
+                          .getResultList())
+              .stream()
+              .map(
+                  recurring ->
+                      jpaTm()
+                          .transact(
+                              () ->
+                                  expandBillingEvent(recurring, executeTime, cursorTime, isDryRun)))
+              .reduce(0, Integer::sum);
 
+      if (!isDryRun) {
+        logger.atInfo().log("Saved OneTime billing events", numBillingEventsSaved);
+      } else {
+        logger.atInfo().log("Generated OneTime billing events (dry run)", numBillingEventsSaved);
+      }
+      logger.atInfo().log(
+          "Recurring event expansion %s complete for billing event range [%s, %s).",
+          isDryRun ? "(dry run) " : "", cursorTime, executeTime);
+      tm().transact(
+              () -> {
+                // Check for the unlikely scenario where the cursor has been altered during the
+                // expansion.
+                DateTime currentCursorTime =
+                    tm().loadByKeyIfPresent(Cursor.createGlobalVKey(RECURRING_BILLING))
+                        .orElse(Cursor.createGlobal(RECURRING_BILLING, START_OF_TIME))
+                        .getCursorTime();
+                if (!currentCursorTime.equals(persistedCursorTime)) {
+                  throw new IllegalStateException(
+                      String.format(
+                          "Current cursor position %s does not match persisted cursor position %s.",
+                          currentCursorTime, persistedCursorTime));
+                }
+                if (!isDryRun) {
+                  tm().put(Cursor.createGlobal(RECURRING_BILLING, executeTime));
+                }
+              });
+    }
+  }
   /** Mapper to expand {@link Recurring} billing events into synthetic {@link OneTime} events. */
   public static class ExpandRecurringBillingEventsMapper
       extends Mapper<Recurring, DateTime, DateTime> {
@@ -155,100 +212,7 @@ public class ExpandRecurringBillingEventsAction implements Runnable {
       try {
         numBillingEventsSaved =
             tm().transactNew(
-                    () -> {
-                      ImmutableSet.Builder<OneTime> syntheticOneTimesBuilder =
-                          new ImmutableSet.Builder<>();
-                      final Registry tld =
-                          Registry.get(getTldFromDomainName(recurring.getTargetId()));
-
-                      // Determine the complete set of times at which this recurring event should
-                      // occur (up to and including the runtime of the mapreduce).
-                      Iterable<DateTime> eventTimes =
-                          recurring
-                              .getRecurrenceTimeOfYear()
-                              .getInstancesInRange(
-                                  Range.closed(
-                                      recurring.getEventTime(),
-                                      earliestOf(recurring.getRecurrenceEndTime(), executeTime)));
-
-                      // Convert these event times to billing times
-                      final ImmutableSet<DateTime> billingTimes =
-                          getBillingTimesInScope(eventTimes, cursorTime, executeTime, tld);
-
-                      Key<? extends EppResource> domainKey = recurring.getParentKey().getParent();
-                      Iterable<OneTime> oneTimesForDomain =
-                          ofy().load().type(OneTime.class).ancestor(domainKey);
-
-                      // Determine the billing times that already have OneTime events persisted.
-                      ImmutableSet<DateTime> existingBillingTimes =
-                          getExistingBillingTimes(oneTimesForDomain, recurring);
-
-                      ImmutableSet.Builder<HistoryEntry> historyEntriesBuilder =
-                          new ImmutableSet.Builder<>();
-                      // Create synthetic OneTime events for all billing times that do not yet have
-                      // an event persisted.
-                      for (DateTime billingTime : difference(billingTimes, existingBillingTimes)) {
-                        // Construct a new HistoryEntry that parents over the OneTime
-                        DomainHistory historyEntry =
-                            new DomainHistory.Builder()
-                                .setBySuperuser(false)
-                                .setClientId(recurring.getClientId())
-                                .setModificationTime(tm().getTransactionTime())
-                                // TODO (jianglai): modify this to use setDomain instead when
-                                //  converting this action to be SQL-aware.
-                                .setDomainRepoId(domainKey.getName())
-                                .setPeriod(Period.create(1, YEARS))
-                                .setReason(
-                                    "Domain autorenewal by ExpandRecurringBillingEventsAction")
-                                .setRequestedByRegistrar(false)
-                                .setType(DOMAIN_AUTORENEW)
-                                // Don't write a domain transaction record if the recurrence was
-                                // ended prior to the billing time (i.e. a domain was deleted
-                                // during the autorenew grace period).
-                                .setDomainTransactionRecords(
-                                    recurring.getRecurrenceEndTime().isBefore(billingTime)
-                                        ? ImmutableSet.of()
-                                        : ImmutableSet.of(
-                                            DomainTransactionRecord.create(
-                                                tld.getTldStr(),
-                                                // We report this when the autorenew grace period
-                                                // ends
-                                                billingTime,
-                                                TransactionReportField.netRenewsFieldFromYears(1),
-                                                1)))
-                                .build();
-                        historyEntriesBuilder.add(historyEntry);
-
-                        DateTime eventTime = billingTime.minus(tld.getAutoRenewGracePeriodLength());
-                        // Determine the cost for a one-year renewal.
-                        Money renewCost = getDomainRenewCost(recurring.getTargetId(), eventTime, 1);
-                        syntheticOneTimesBuilder.add(
-                            new OneTime.Builder()
-                                .setBillingTime(billingTime)
-                                .setClientId(recurring.getClientId())
-                                .setCost(renewCost)
-                                .setEventTime(eventTime)
-                                .setFlags(union(recurring.getFlags(), Flag.SYNTHETIC))
-                                .setParent(historyEntry)
-                                .setPeriodYears(1)
-                                .setReason(recurring.getReason())
-                                .setSyntheticCreationTime(executeTime)
-                                .setCancellationMatchingBillingEvent(recurring.createVKey())
-                                .setTargetId(recurring.getTargetId())
-                                .build());
-                      }
-                      Set<HistoryEntry> historyEntries = historyEntriesBuilder.build();
-                      Set<OneTime> syntheticOneTimes = syntheticOneTimesBuilder.build();
-                      if (!isDryRun) {
-                        ImmutableSet<ImmutableObject> entitiesToSave =
-                            new ImmutableSet.Builder<ImmutableObject>()
-                                .addAll(historyEntries)
-                                .addAll(syntheticOneTimes)
-                                .build();
-                        ofy().save().entities(entitiesToSave).now();
-                      }
-                      return syntheticOneTimes.size();
-                    });
+                    () -> expandBillingEvent(recurring, executeTime, cursorTime, isDryRun));
       } catch (Throwable t) {
         getContext().incrementCounter("error: " + t.getClass().getSimpleName());
         getContext().incrementCounter(ERROR_COUNTER);
@@ -260,45 +224,12 @@ public class ExpandRecurringBillingEventsAction implements Runnable {
       if (!isDryRun) {
         getContext().incrementCounter("Saved OneTime billing events", numBillingEventsSaved);
       } else {
-        getContext().incrementCounter(
-            "Generated OneTime billing events (dry run)", numBillingEventsSaved);
+        getContext()
+            .incrementCounter("Generated OneTime billing events (dry run)", numBillingEventsSaved);
       }
     }
-
-    /**
-     * Filters a set of {@link DateTime}s down to event times that are in scope for a particular
-     * mapreduce run, given the cursor time and the mapreduce execution time.
-     */
-    private ImmutableSet<DateTime> getBillingTimesInScope(
-        Iterable<DateTime> eventTimes,
-        DateTime cursorTime,
-        DateTime executeTime,
-        final Registry tld) {
-      return Streams.stream(eventTimes)
-          .map(eventTime -> eventTime.plus(tld.getAutoRenewGracePeriodLength()))
-          .filter(Range.closedOpen(cursorTime, executeTime))
-          .collect(toImmutableSet());
-    }
-
-    /**
-     * Determines an {@link ImmutableSet} of {@link DateTime}s that have already been persisted
-     * for a given recurring billing event.
-     */
-    private ImmutableSet<DateTime> getExistingBillingTimes(
-        Iterable<BillingEvent.OneTime> oneTimesForDomain,
-        final BillingEvent.Recurring recurringEvent) {
-      return Streams.stream(oneTimesForDomain)
-          .filter(
-              billingEvent ->
-                  recurringEvent
-                      .createVKey()
-                      .equals(billingEvent.getCancellationMatchingBillingEvent()))
-          .map(OneTime::getBillingTime)
-          .collect(toImmutableSet());
-    }
   }
 
-
   /**
    * "Reducer" to advance the cursor after all map jobs have been completed. The NullInput into the
    * mapper will cause the mapper to emit one timestamp pair (current cursor and execution time),
@@ -331,7 +262,8 @@ public class ExpandRecurringBillingEventsAction implements Runnable {
           isDryRun ? "(dry run) " : "", cursorTime, executionTime);
       tm().transact(
               () -> {
-                Cursor cursor = ofy().load().key(Cursor.createGlobalKey(RECURRING_BILLING)).now();
+                Cursor cursor =
+                    auditedOfy().load().key(Cursor.createGlobalKey(RECURRING_BILLING)).now();
                 DateTime currentCursorTime =
                     (cursor == null ? START_OF_TIME : cursor.getCursorTime());
                 if (!currentCursorTime.equals(expectedPersistedCursorTime)) {
@@ -346,4 +278,135 @@ public class ExpandRecurringBillingEventsAction implements Runnable {
               });
     }
   }
+
+  private static int expandBillingEvent(
+      Recurring recurring, DateTime executeTime, DateTime cursorTime, boolean isDryRun) {
+    ImmutableSet.Builder<OneTime> syntheticOneTimesBuilder = new ImmutableSet.Builder<>();
+    final Registry tld = Registry.get(getTldFromDomainName(recurring.getTargetId()));
+
+    // Determine the complete set of times at which this recurring event should
+    // occur (up to and including the runtime of the mapreduce).
+    Iterable<DateTime> eventTimes =
+        recurring
+            .getRecurrenceTimeOfYear()
+            .getInstancesInRange(
+                Range.closed(
+                    recurring.getEventTime(),
+                    earliestOf(recurring.getRecurrenceEndTime(), executeTime)));
+
+    // Convert these event times to billing times
+    final ImmutableSet<DateTime> billingTimes =
+        getBillingTimesInScope(eventTimes, cursorTime, executeTime, tld);
+
+    VKey<DomainBase> domainKey =
+        VKey.create(
+            DomainBase.class, recurring.getDomainRepoId(), recurring.getParentKey().getParent());
+    Iterable<OneTime> oneTimesForDomain;
+    if (tm().isOfy()) {
+      oneTimesForDomain = auditedOfy().load().type(OneTime.class).ancestor(domainKey.getOfyKey());
+    } else {
+      oneTimesForDomain =
+          tm().createQueryComposer(OneTime.class)
+              .where("domainRepoId", EQ, recurring.getDomainRepoId())
+              .list();
+    }
+
+    // Determine the billing times that already have OneTime events persisted.
+    ImmutableSet<DateTime> existingBillingTimes =
+        getExistingBillingTimes(oneTimesForDomain, recurring);
+
+    ImmutableSet.Builder<DomainHistory> historyEntriesBuilder = new ImmutableSet.Builder<>();
+    // Create synthetic OneTime events for all billing times that do not yet have
+    // an event persisted.
+    for (DateTime billingTime : difference(billingTimes, existingBillingTimes)) {
+      // Construct a new HistoryEntry that parents over the OneTime
+      DomainHistory historyEntry =
+          new DomainHistory.Builder()
+              .setBySuperuser(false)
+              .setClientId(recurring.getClientId())
+              .setModificationTime(tm().getTransactionTime())
+              .setDomain(tm().loadByKey(domainKey))
+              .setPeriod(Period.create(1, YEARS))
+              .setReason("Domain autorenewal by ExpandRecurringBillingEventsAction")
+              .setRequestedByRegistrar(false)
+              .setType(DOMAIN_AUTORENEW)
+              // Don't write a domain transaction record if the recurrence was
+              // ended prior to the billing time (i.e. a domain was deleted
+              // during the autorenew grace period).
+              .setDomainTransactionRecords(
+                  recurring.getRecurrenceEndTime().isBefore(billingTime)
+                      ? ImmutableSet.of()
+                      : ImmutableSet.of(
+                          DomainTransactionRecord.create(
+                              tld.getTldStr(),
+                              // We report this when the autorenew grace period
+                              // ends
+                              billingTime,
+                              TransactionReportField.netRenewsFieldFromYears(1),
+                              1)))
+              .build();
+      historyEntriesBuilder.add(historyEntry);
+
+      DateTime eventTime = billingTime.minus(tld.getAutoRenewGracePeriodLength());
+      // Determine the cost for a one-year renewal.
+      Money renewCost = getDomainRenewCost(recurring.getTargetId(), eventTime, 1);
+      syntheticOneTimesBuilder.add(
+          new OneTime.Builder()
+              .setBillingTime(billingTime)
+              .setClientId(recurring.getClientId())
+              .setCost(renewCost)
+              .setEventTime(eventTime)
+              .setFlags(union(recurring.getFlags(), Flag.SYNTHETIC))
+              .setParent(historyEntry)
+              .setPeriodYears(1)
+              .setReason(recurring.getReason())
+              .setSyntheticCreationTime(executeTime)
+              .setCancellationMatchingBillingEvent(recurring.createVKey())
+              .setTargetId(recurring.getTargetId())
+              .build());
+    }
+    Set<DomainHistory> historyEntries = historyEntriesBuilder.build();
+    Set<OneTime> syntheticOneTimes = syntheticOneTimesBuilder.build();
+    if (!isDryRun) {
+      ImmutableSet<ImmutableObject> entitiesToSave =
+          new ImmutableSet.Builder<ImmutableObject>()
+              .addAll(historyEntries)
+              .addAll(syntheticOneTimes)
+              .build();
+      tm().putAll(entitiesToSave);
+    }
+    return syntheticOneTimes.size();
+  }
+
+  /**
+   * Filters a set of {@link DateTime}s down to event times that are in scope for a particular
+   * mapreduce run, given the cursor time and the mapreduce execution time.
+   */
+  protected static ImmutableSet<DateTime> getBillingTimesInScope(
+      Iterable<DateTime> eventTimes,
+      DateTime cursorTime,
+      DateTime executeTime,
+      final Registry tld) {
+    return Streams.stream(eventTimes)
+        .map(eventTime -> eventTime.plus(tld.getAutoRenewGracePeriodLength()))
+        .filter(Range.closedOpen(cursorTime, executeTime))
+        .collect(toImmutableSet());
+  }
+
+  /**
+   * Determines an {@link ImmutableSet} of {@link DateTime}s that have already been persisted for a
+   * given recurring billing event.
+   */
+  private static ImmutableSet<DateTime> getExistingBillingTimes(
+      Iterable<BillingEvent.OneTime> oneTimesForDomain,
+      final BillingEvent.Recurring recurringEvent) {
+    return Streams.stream(oneTimesForDomain)
+        .filter(
+            billingEvent ->
+                recurringEvent
+                    .createVKey()
+                    .equals(billingEvent.getCancellationMatchingBillingEvent()))
+        .map(OneTime::getBillingTime)
+        .collect(toImmutableSet());
+  }
 }

core/src/main/java/google/registry/beam/invoicing/BillingEvent.java

@@ -251,7 +251,14 @@ public abstract class BillingEvent implements Serializable {
   InvoiceGroupingKey getInvoiceGroupingKey() {
     return new AutoValue_BillingEvent_InvoiceGroupingKey(
         billingTime().toLocalDate().withDayOfMonth(1).toString(),
-        billingTime().toLocalDate().withDayOfMonth(1).plusYears(years()).minusDays(1).toString(),
+        years() == 0
+            ? ""
+            : billingTime()
+                .toLocalDate()
+                .withDayOfMonth(1)
+                .plusYears(years())
+                .minusDays(1)
+                .toString(),
         billingId(),
         String.format("%s - %s", registrarId(), tld()),
         String.format("%s | TLD: %s | TERM: %d-year", action(), tld(), years()),

core/src/main/java/google/registry/config/RegistryConfig.java

@@ -1502,22 +1502,6 @@ public final class RegistryConfig {
     CONFIG_SETTINGS.get().cloudSql.replicateTransactions = replicateTransactions;
   }
 
-  /**
-   * Returns whether or not to replay commit logs to the SQL database after export to GCS.
-   *
-   * <p>If true, we will trigger the {@link google.registry.backup.ReplayCommitLogsToSqlAction}
-   * after the {@link google.registry.backup.ExportCommitLogDiffAction} to load the commit logs and
-   * replay them to SQL.
-   */
-  public static boolean getCloudSqlReplayCommitLogs() {
-    return CONFIG_SETTINGS.get().cloudSql.replayCommitLogs;
-  }
-
-  @VisibleForTesting
-  public static void overrideCloudSqlReplayCommitLogs(boolean replayCommitLogs) {
-    CONFIG_SETTINGS.get().cloudSql.replayCommitLogs = replayCommitLogs;
-  }
-
   /** Returns the roid suffix to be used for the roids of all contacts and hosts. */
   public static String getContactAndHostRoidSuffix() {
     return CONFIG_SETTINGS.get().registryPolicy.contactAndHostRoidSuffix;

core/src/main/java/google/registry/config/RegistryConfigSettings.java

@@ -127,7 +127,6 @@ public class RegistryConfigSettings {
     public String username;
     public String instanceConnectionName;
     public boolean replicateTransactions;
-    public boolean replayCommitLogs;
   }
 
   /** Configuration for Apache Beam (Cloud Dataflow). */

core/src/main/java/google/registry/config/files/default-config.yaml

@@ -230,8 +230,6 @@ cloudSql:
   # Set this to true to replicate cloud SQL transactions to datastore in the
   # background.
   replicateTransactions: false
-  # Set this to true to enable replay of commit logs to SQL
-  replayCommitLogs: false
 
 cloudDns:
   # Set both properties to null in Production.

core/src/main/java/google/registry/env/common/default/WEB-INF/datastore-indexes.xml

@@ -41,6 +41,11 @@
     <property name="nsHosts" direction="asc"/>
     <property name="deletionTime" direction="asc"/>
   </datastore-index>
+  <!-- For deleting expired not-previously-deleted domains. -->
+  <datastore-index kind="DomainBase" ancestor="false" source="manual">
+    <property name="deletionTime" direction="asc"/>
+    <property name="autorenewEndTime" direction="asc"/>
+  </datastore-index>
   <!-- For RDAP searches by linked nameserver. -->
   <datastore-index kind="DomainBase" ancestor="false" source="manual">
     <property name="nsHosts" direction="asc"/>

core/src/main/java/google/registry/keyring/kms/KmsKeyring.java

@@ -16,24 +16,12 @@ package google.registry.keyring.kms;
 
 import static com.google.common.base.CaseFormat.LOWER_HYPHEN;
 import static com.google.common.base.CaseFormat.UPPER_UNDERSCORE;
-import static com.google.common.base.Preconditions.checkState;
-import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
-import static google.registry.model.ofy.ObjectifyService.auditedOfy;
-import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
-import com.google.common.flogger.FluentLogger;
-import com.googlecode.objectify.Key;
-import google.registry.config.RegistryConfig.Config;
 import google.registry.keyring.api.KeySerializer;
 import google.registry.keyring.api.Keyring;
 import google.registry.keyring.api.KeyringException;
-import google.registry.model.server.KmsSecret;
-import google.registry.model.server.KmsSecretRevision;
-import google.registry.model.server.KmsSecretRevisionSqlDao;
 import google.registry.privileges.secretmanager.KeyringSecretStore;
 import java.io.IOException;
-import java.util.Arrays;
-import java.util.Optional;
 import javax.inject.Inject;
 import org.bouncycastle.openpgp.PGPException;
 import org.bouncycastle.openpgp.PGPKeyPair;
@@ -47,11 +35,9 @@ import org.bouncycastle.openpgp.PGPPublicKey;
  * @see <a href="https://cloud.google.com/kms/docs/">Google Cloud Key Management Service
  *     Documentation</a>
  */
-// TODO(2021-06-01): rename this class to SecretManagerKeyring and delete KmsSecretRevision
+// TODO(2021-07-01): rename this class to SecretManagerKeyring and delete KmsSecretRevision
 public class KmsKeyring implements Keyring {
 
-  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
-
   /** Key labels for private key secrets. */
   enum PrivateKeyLabel {
     BRDA_SIGNING_PRIVATE,
@@ -92,13 +78,10 @@ public class KmsKeyring implements Keyring {
     }
   }
 
-  private final KmsConnection kmsConnection;
   private final KeyringSecretStore secretStore;
 
   @Inject
-  KmsKeyring(
-      @Config("defaultKmsConnection") KmsConnection kmsConnection, KeyringSecretStore secretStore) {
-    this.kmsConnection = kmsConnection;
+  KmsKeyring(KeyringSecretStore secretStore) {
     this.secretStore = secretStore;
   }
 
@@ -201,44 +184,11 @@ public class KmsKeyring implements Keyring {
     return getKeyPair(keyLabel).getPrivateKey();
   }
 
-  private byte[] getDecryptedDataFromDatastore(String keyName) {
-    String encryptedData;
-    if (tm().isOfy()) {
-      KmsSecret secret =
-          auditedOfy().load().key(Key.create(getCrossTldKey(), KmsSecret.class, keyName)).now();
-      checkState(secret != null, "Requested secret '%s' does not exist.", keyName);
-      encryptedData = auditedOfy().load().key(secret.getLatestRevision()).now().getEncryptedValue();
-    } else {
-      Optional<KmsSecretRevision> revision =
-          tm().transact(() -> KmsSecretRevisionSqlDao.getLatestRevision(keyName));
-      checkState(revision.isPresent(), "Requested secret '%s' does not exist.", keyName);
-      encryptedData = revision.get().getEncryptedValue();
-    }
-
-    try {
-      return kmsConnection.decrypt(keyName, encryptedData);
-    } catch (Exception e) {
-      return new byte[0];
-    }
-  }
-
-  private byte[] getDataFromSecretStore(String keyName) {
+  private byte[] getDecryptedData(String keyName) {
     try {
       return secretStore.getSecret(keyName);
     } catch (Exception e) {
       throw new KeyringException("Failed to retrieve secret for " + keyName, e);
     }
   }
-
-  private byte[] getDecryptedData(String keyName) {
-    byte[] dsData = getDecryptedDataFromDatastore(keyName);
-    byte[] secretStoreData = getDataFromSecretStore(keyName);
-
-    if (Arrays.equals(dsData, secretStoreData)) {
-      logger.atInfo().log("Values for %s in Datastore and Secret Manager match.", keyName);
-      return secretStoreData;
-    }
-    logger.atWarning().log("Values for %s in Datastore and Secret Manager do not match.", keyName);
-    return secretStoreData;
-  }
 }

core/src/main/java/google/registry/keyring/kms/KmsUpdater.java

@@ -32,18 +32,13 @@ import static google.registry.keyring.kms.KmsKeyring.StringKeyLabel.MARKSDB_SMDR
 import static google.registry.keyring.kms.KmsKeyring.StringKeyLabel.RDE_SSH_CLIENT_PRIVATE_STRING;
 import static google.registry.keyring.kms.KmsKeyring.StringKeyLabel.RDE_SSH_CLIENT_PUBLIC_STRING;
 import static google.registry.keyring.kms.KmsKeyring.StringKeyLabel.SAFE_BROWSING_API_KEY;
-import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 
-import com.google.common.collect.ImmutableMap;
 import com.google.common.flogger.FluentLogger;
-import google.registry.config.RegistryConfig.Config;
 import google.registry.keyring.api.KeySerializer;
 import google.registry.keyring.kms.KmsKeyring.PrivateKeyLabel;
 import google.registry.keyring.kms.KmsKeyring.PublicKeyLabel;
 import google.registry.keyring.kms.KmsKeyring.StringKeyLabel;
-import google.registry.model.server.KmsSecret;
-import google.registry.model.server.KmsSecretRevision;
 import google.registry.privileges.secretmanager.KeyringSecretStore;
 import java.io.IOException;
 import java.util.HashMap;
@@ -62,14 +57,11 @@ import org.bouncycastle.openpgp.PGPPublicKey;
 public final class KmsUpdater {
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
 
-  private final KmsConnection kmsConnection;
   private final KeyringSecretStore secretStore;
   private final HashMap<String, byte[]> secretValues;
 
   @Inject
-  public KmsUpdater(
-      @Config("defaultKmsConnection") KmsConnection kmsConnection, KeyringSecretStore secretStore) {
-    this.kmsConnection = kmsConnection;
+  public KmsUpdater(KeyringSecretStore secretStore) {
     this.secretStore = secretStore;
 
     // Use LinkedHashMap to preserve insertion order on update() to simplify testing and debugging
@@ -150,24 +142,6 @@ public final class KmsUpdater {
               + "Please check the status of Secret Manager and re-run the command.",
           e);
     }
-
-    // TODO(2021-06-01): remove the writes to Datastore
-    persistEncryptedValues(encryptValues(secretValues));
-  }
-
-  /**
-   * Encrypts updated secrets using KMS. If the configured {@code KeyRing} or {@code CryptoKey}
-   * associated with a secret doesn't exist, they will first be created.
-   *
-   * @see google.registry.config.RegistryConfigSettings.Kms
-   */
-  private ImmutableMap<String, EncryptResponse> encryptValues(Map<String, byte[]> keyValues) {
-    ImmutableMap.Builder<String, EncryptResponse> encryptedValues = new ImmutableMap.Builder<>();
-    for (Map.Entry<String, byte[]> entry : keyValues.entrySet()) {
-      String secretName = entry.getKey();
-      encryptedValues.put(secretName, kmsConnection.encrypt(secretName, entry.getValue()));
-    }
-    return encryptedValues.build();
   }
 
   private KmsUpdater setString(String key, StringKeyLabel stringKeyLabel) {
@@ -195,32 +169,6 @@ public final class KmsUpdater {
     return this;
   }
 
-  /**
-   * Persists encrypted secrets to Datastore as {@link KmsSecretRevision} entities and makes them
-   * primary. {@link KmsSecret} entities point to the latest {@link KmsSecretRevision}.
-   *
-   * <p>The changes are committed transactionally; if an error occurs, all existing {@link
-   * KmsSecretRevision} entities will remain primary.
-   */
-  private static void persistEncryptedValues(
-      final ImmutableMap<String, EncryptResponse> encryptedValues) {
-    tm().transact(
-            () -> {
-              for (Map.Entry<String, EncryptResponse> entry : encryptedValues.entrySet()) {
-                String secretName = entry.getKey();
-                EncryptResponse revisionData = entry.getValue();
-
-                KmsSecretRevision secretRevision =
-                    new KmsSecretRevision.Builder()
-                        .setEncryptedValue(revisionData.ciphertext())
-                        .setKmsCryptoKeyVersionName(revisionData.cryptoKeyVersionName())
-                        .setParent(secretName)
-                        .build();
-                tm().putAll(secretRevision, KmsSecret.create(secretName, secretRevision));
-              }
-            });
-  }
-
   private void setSecret(String secretName, byte[] value) {
     checkArgument(!secretValues.containsKey(secretName), "Attempted to set %s twice", secretName);
     secretValues.put(secretName, value);

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

@@ -1,55 +0,0 @@
-// Copyright 2020 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;
-
-import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
-
-import com.google.common.flogger.FluentLogger;
-import google.registry.config.RegistryEnvironment;
-import google.registry.model.common.DatabaseTransitionSchedule;
-import google.registry.model.common.DatabaseTransitionSchedule.PrimaryDatabase;
-import google.registry.model.common.DatabaseTransitionSchedule.TransitionId;
-
-/** Utility methods related to migrating dual-read/dual-write entities. */
-public class DatabaseMigrationUtils {
-
-  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
-
-  /** Throws exceptions only in unit tests, otherwise only logs exceptions. */
-  public static void suppressExceptionUnlessInTest(Runnable work, String message) {
-    try {
-      work.run();
-    } catch (Exception e) {
-      if (RegistryEnvironment.get().equals(RegistryEnvironment.UNITTEST)) {
-        throw e;
-      }
-      logger.atWarning().withCause(e).log(message);
-    }
-  }
-
-  /** Gets the value for the database currently considered primary. */
-  public static PrimaryDatabase getPrimaryDatabase(TransitionId transitionId) {
-    return DatabaseTransitionSchedule.getCached(transitionId)
-        .map(DatabaseTransitionSchedule::getPrimaryDatabase)
-        .orElse(PrimaryDatabase.DATASTORE);
-  }
-
-  public static boolean isDatastore(TransitionId transitionId) {
-    return tm().transactNew(() -> DatabaseMigrationUtils.getPrimaryDatabase(transitionId))
-        .equals(PrimaryDatabase.DATASTORE);
-  }
-
-  private DatabaseMigrationUtils() {}
-}

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

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

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

@@ -155,9 +155,16 @@ public abstract class EppResource extends BackupGroupRoot implements Buildable {
     return repoId;
   }
 
-  /** This method exists solely to satisfy Hibernate. Use {@link Builder} instead. */
-  @SuppressWarnings("UnusedMethod")
-  private void setRepoId(String repoId) {
+  /**
+   * Sets the repository ID.
+   *
+   * <p>This should only be used for restoring the repo id of an object being loaded in a PostLoad
+   * method (effectively, when it is still under construction by Hibernate). In all other cases, the
+   * object should be regarded as immutable and changes should go through a Builder.
+   *
+   * <p>In addition to this special case use, this method must exist to satisfy Hibernate.
+   */
+  public void setRepoId(String repoId) {
     this.repoId = repoId;
   }
 

core/src/main/java/google/registry/model/billing/BillingEvent.java

@@ -324,13 +324,22 @@ public abstract class BillingEvent extends ImmutableObject
     DateTime syntheticCreationTime;
 
     /**
-     * For {@link Flag#SYNTHETIC} events, a {@link Key} to the {@link BillingEvent} from which this
-     * OneTime was created. This is needed in order to properly match billing events against {@link
-     * Cancellation}s.
+     * For {@link Flag#SYNTHETIC} events, a {@link Key} to the {@link Recurring} from which this
+     * {@link OneTime} was created. This is needed in order to properly match billing events against
+     * {@link Cancellation}s.
      */
     @Column(name = "cancellation_matching_billing_recurrence_id")
     VKey<Recurring> cancellationMatchingBillingEvent;
 
+    /**
+     * For {@link Flag#SYNTHETIC} events, the {@link DomainHistory} revision ID of the the {@link
+     * Recurring} from which this {@link OneTime} was created. This is needed in order to recreate
+     * the {@link VKey} when reading from SQL.
+     */
+    @Ignore
+    @Column(name = "recurrence_history_revision_id")
+    Long recurringEventHistoryRevisionId;
+
     /**
      * The {@link AllocationToken} used in the creation of this event, or null if one was not used.
      */
@@ -354,10 +363,14 @@ public abstract class BillingEvent extends ImmutableObject
       return syntheticCreationTime;
     }
 
-    public VKey<? extends BillingEvent> getCancellationMatchingBillingEvent() {
+    public VKey<Recurring> getCancellationMatchingBillingEvent() {
       return cancellationMatchingBillingEvent;
     }
 
+    public Long getRecurringEventHistoryRevisionId() {
+      return recurringEventHistoryRevisionId;
+    }
+
     public Optional<VKey<AllocationToken>> getAllocationToken() {
       return Optional.ofNullable(allocationToken);
     }
@@ -376,6 +389,28 @@ public abstract class BillingEvent extends ImmutableObject
       return new Builder(clone(this));
     }
 
+    @Override
+    void onLoad() {
+      super.onLoad();
+      if (cancellationMatchingBillingEvent != null) {
+        recurringEventHistoryRevisionId =
+            cancellationMatchingBillingEvent.getOfyKey().getParent().getId();
+      }
+    }
+
+    @Override
+    void postLoad() {
+      super.postLoad();
+      if (cancellationMatchingBillingEvent != null) {
+        cancellationMatchingBillingEvent =
+            cancellationMatchingBillingEvent.restoreOfy(
+                DomainBase.class,
+                domainRepoId,
+                DomainHistory.class,
+                recurringEventHistoryRevisionId);
+      }
+    }
+
     /** A builder for {@link OneTime} since it is immutable. */
     public static class Builder extends BillingEvent.Builder<OneTime, Builder> {
 
@@ -410,6 +445,8 @@ public abstract class BillingEvent extends ImmutableObject
       public Builder setCancellationMatchingBillingEvent(
           VKey<Recurring> cancellationMatchingBillingEvent) {
         getInstance().cancellationMatchingBillingEvent = cancellationMatchingBillingEvent;
+        getInstance().recurringEventHistoryRevisionId =
+            cancellationMatchingBillingEvent.getOfyKey().getParent().getId();
         return this;
       }
 
@@ -444,6 +481,11 @@ public abstract class BillingEvent extends ImmutableObject
                 == (instance.cancellationMatchingBillingEvent != null),
             "Cancellation matching billing event must be set if and only if the SYNTHETIC flag "
                 + "is set.");
+        checkState(
+            !instance.getFlags().contains(Flag.SYNTHETIC)
+                || (instance.cancellationMatchingBillingEvent.getOfyKey().getParent().getId()
+                    == instance.recurringEventHistoryRevisionId),
+            "Cancellation matching billing event and its history revision ID does not match.");
         return super.build();
       }
     }

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

@@ -0,0 +1,251 @@
+// 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 static google.registry.util.DateTimeUtils.START_OF_TIME;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.cache.CacheBuilder;
+import com.google.common.cache.CacheLoader;
+import com.google.common.cache.LoadingCache;
+import com.google.common.collect.ImmutableMultimap;
+import com.google.common.collect.ImmutableSortedMap;
+import com.googlecode.objectify.annotation.Embed;
+import com.googlecode.objectify.annotation.Entity;
+import com.googlecode.objectify.annotation.Mapify;
+import google.registry.model.annotations.InCrossTld;
+import google.registry.model.common.TimedTransitionProperty.TimeMapper;
+import google.registry.model.common.TimedTransitionProperty.TimedTransition;
+import google.registry.schema.replay.DatastoreOnlyEntity;
+import java.time.Duration;
+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 Datastore throughout the entire migration so as to have a single point
+ * of access.
+ */
+@Entity
+@InCrossTld
+public class DatabaseMigrationStateSchedule extends CrossTldSingleton
+    implements DatastoreOnlyEntity {
+
+  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_ONLY(PrimaryDatabase.DATASTORE, false, ReplayDirection.NO_REPLAY),
+    DATASTORE_PRIMARY(PrimaryDatabase.DATASTORE, false, ReplayDirection.DATASTORE_TO_SQL),
+    DATASTORE_PRIMARY_READ_ONLY(PrimaryDatabase.DATASTORE, true, ReplayDirection.DATASTORE_TO_SQL),
+    SQL_PRIMARY_READ_ONLY(PrimaryDatabase.CLOUD_SQL, true, ReplayDirection.SQL_TO_DATASTORE),
+    SQL_PRIMARY(PrimaryDatabase.CLOUD_SQL, false, ReplayDirection.SQL_TO_DATASTORE),
+    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;
+    }
+  }
+
+  @Embed
+  public static class MigrationStateTransition extends TimedTransition<MigrationState> {
+    private MigrationState migrationState;
+
+    @Override
+    protected 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.
+   */
+  private static final LoadingCache<
+          Class<DatabaseMigrationStateSchedule>,
+          TimedTransitionProperty<MigrationState, MigrationStateTransition>>
+      // Each instance should cache the migration schedule for five minutes before reloading
+      CACHE =
+          CacheBuilder.newBuilder()
+              .expireAfterWrite(Duration.ofMinutes(5))
+              .build(
+                  new CacheLoader<
+                      Class<DatabaseMigrationStateSchedule>,
+                      TimedTransitionProperty<MigrationState, MigrationStateTransition>>() {
+                    @Override
+                    public TimedTransitionProperty<MigrationState, MigrationStateTransition> load(
+                        Class<DatabaseMigrationStateSchedule> unused) {
+                      return 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_READ_ONLY)
+            .putAll(
+                MigrationState.DATASTORE_PRIMARY_READ_ONLY,
+                MigrationState.DATASTORE_ONLY,
+                MigrationState.DATASTORE_PRIMARY,
+                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).
+    for (MigrationState migrationState : MigrationState.values()) {
+      builder.put(migrationState, migrationState);
+    }
+    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
+  @Mapify(TimeMapper.class)
+  public TimedTransitionProperty<MigrationState, MigrationStateTransition> migrationTransitions =
+      TimedTransitionProperty.forMapify(
+          MigrationState.DATASTORE_ONLY, MigrationStateTransition.class);
+
+  // Required for Objectify initialization
+  private DatabaseMigrationStateSchedule() {}
+
+  @VisibleForTesting
+  DatabaseMigrationStateSchedule(
+      TimedTransitionProperty<MigrationState, MigrationStateTransition> migrationTransitions) {
+    this.migrationTransitions = migrationTransitions;
+  }
+
+  /** Sets and persists to Datastore the provided migration transition schedule. */
+  public static void set(ImmutableSortedMap<DateTime, MigrationState> migrationTransitionMap) {
+    ofyTm().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);
+    ofyTm().put(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.getUnchecked(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 Datastore, or the default if none exists. */
+  @VisibleForTesting
+  static TimedTransitionProperty<MigrationState, MigrationStateTransition> getUncached() {
+    return ofyTm()
+        .transact(
+            () ->
+                ofyTm()
+                    .loadSingleton(DatabaseMigrationStateSchedule.class)
+                    .map(s -> s.migrationTransitions)
+                    .orElse(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 time must also be valid.
+   */
+  private static void validateTransitionAtCurrentTime(
+      TimedTransitionProperty<MigrationState, MigrationStateTransition> newTransitions) {
+    MigrationState currentValue = getUncached().getValueAtTime(ofyTm().getTransactionTime());
+    MigrationState nextCurrentValue = newTransitions.getValueAtTime(ofyTm().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/model/common/DatabaseMigrationStateWrapper.java

@@ -1,122 +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.
-
-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/main/java/google/registry/model/common/DatabaseTransitionSchedule.java

@@ -1,166 +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.
-
-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.annotations.InCrossTld;
-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.model.tmch.ClaimsList;
-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
-@InCrossTld
-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 migration of {@link ClaimsList} entities. */
-    CLAIMS_LIST,
-    /** 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 schedule for all asynchronously-replayed entities, ones not dually-written. */
-    REPLAYED_ENTITIES,
-  }
-
-  /**
-   * 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<TransitionId, Optional<DatabaseTransitionSchedule>> CACHE =
-      CacheBuilder.newBuilder()
-          .expireAfterWrite(
-              java.time.Duration.ofMillis(getSingletonCacheRefreshDuration().getMillis()))
-          .build(
-              new CacheLoader<TransitionId, Optional<DatabaseTransitionSchedule>>() {
-                @Override
-                public Optional<DatabaseTransitionSchedule> load(TransitionId transitionId) {
-                  return DatabaseTransitionSchedule.get(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);
-  }
-
-  /** 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/model/contact/ContactHistory.java

@@ -115,7 +115,10 @@ public class ContactHistory extends HistoryEntry implements SqlEntity {
       contactBase = null;
     }
     if (contactBase != null && contactBase.getRepoId() == null) {
-      contactBase = contactBase.asBuilder().setRepoId(parent.getName()).build();
+      // contactBase hasn't been fully constructed yet, so it's ok to go in and mutate it.  Though
+      // the use of the Builder is not necessarily problematic in this case, this is still safer as
+      // the Builder can do things like comparisons that compute the hash code.
+      contactBase.setRepoId(parent.getName());
     }
   }
 

core/src/main/java/google/registry/model/domain/DomainHistory.java

@@ -258,7 +258,10 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
         domainContent = null;
       } else {
         if (domainContent.getRepoId() == null) {
-          domainContent = domainContent.asBuilder().setRepoId(parent.getName()).build();
+          // domainContent still hasn't been fully constructed yet, so it's ok to go in and mutate
+          // it.  In fact, we have to because going through the builder causes the hash codes of
+          // contained objects to be calculated prematurely.
+          domainContent.setRepoId(parent.getName());
         }
       }
     }

core/src/main/java/google/registry/model/host/HostHistory.java

@@ -116,7 +116,10 @@ public class HostHistory extends HistoryEntry implements SqlEntity {
       hostBase = null;
     }
     if (hostBase != null && hostBase.getRepoId() == null) {
-      hostBase = hostBase.asBuilder().setRepoId(parent.getName()).build();
+      // hostBase hasn't been fully constructed yet, so it's ok to go in and mutate it.  Though the
+      // use of the Builder is not necessarily problematic in this case, this is still safer as the
+      // Builder can do things like comparisons that compute the hash code.
+      hostBase.setRepoId(parent.getName());
     }
   }
 

core/src/main/java/google/registry/model/index/ForeignKeyIndex.java

@@ -228,8 +228,7 @@ public abstract class ForeignKeyIndex<E extends EppResource> extends BackupGroup
           tm().transact(
                   () ->
                       jpaTm()
-                          .getEntityManager()
-                          .createQuery(
+                          .query(
                               CriteriaQueryBuilder.create(clazz)
                                   .whereFieldIsIn(property, foreignKeys)
                                   .build())

core/src/main/java/google/registry/model/reporting/HistoryEntry.java

@@ -61,6 +61,7 @@ import javax.persistence.EnumType;
 import javax.persistence.Enumerated;
 import javax.persistence.MappedSuperclass;
 import javax.persistence.Transient;
+import org.apache.commons.lang3.BooleanUtils;
 import org.joda.time.DateTime;
 
 /**
@@ -435,7 +436,8 @@ public class HistoryEntry extends ImmutableObject implements Buildable, Datastor
       checkArgumentNotNull(getInstance().modificationTime, "Modification time must be specified");
       checkArgumentNotNull(getInstance().clientId, "Registrar ID must be specified");
       checkArgument(
-          !getInstance().type.equals(Type.SYNTHETIC) || !getInstance().requestedByRegistrar,
+          !getInstance().type.equals(Type.SYNTHETIC)
+              || BooleanUtils.isNotTrue(getInstance().requestedByRegistrar),
           "Synthetic history entries cannot be requested by a registrar");
       return super.build();
     }

core/src/main/java/google/registry/model/reporting/HistoryEntryDao.java

@@ -164,8 +164,7 @@ public class HistoryEntryDao {
   private static <T extends HistoryEntry> Stream<T> loadHistoryObjectFromSqlByRegistrars(
       Class<T> historyClass, ImmutableCollection<String> registrarIds) {
     return jpaTm()
-        .getEntityManager()
-        .createQuery(
+        .query(
             CriteriaQueryBuilder.create(historyClass)
                 .whereFieldIsIn("clientId", registrarIds)
                 .build())
@@ -184,11 +183,12 @@ public class HistoryEntryDao {
             .where("modificationTime", criteriaBuilder::greaterThanOrEqualTo, afterTime)
             .where("modificationTime", criteriaBuilder::lessThanOrEqualTo, beforeTime)
             .where(repoIdFieldName, criteriaBuilder::equal, parentKey.getSqlKey().toString())
+            .orderByAsc("id")
             .build();
 
     return ImmutableList.sortedCopyOf(
         Comparator.comparing(HistoryEntry::getModificationTime),
-        jpaTm().getEntityManager().createQuery(criteriaQuery).getResultList());
+        jpaTm().query(criteriaQuery).getResultList());
   }
 
   private static Class<? extends HistoryEntry> getHistoryClassFromParent(
@@ -215,8 +215,7 @@ public class HistoryEntryDao {
       Class<T> historyClass, DateTime afterTime, DateTime beforeTime) {
     CriteriaBuilder criteriaBuilder = jpaTm().getEntityManager().getCriteriaBuilder();
     return jpaTm()
-        .getEntityManager()
-        .createQuery(
+        .query(
             CriteriaQueryBuilder.create(historyClass)
                 .where("modificationTime", criteriaBuilder::greaterThanOrEqualTo, afterTime)
                 .where("modificationTime", criteriaBuilder::lessThanOrEqualTo, beforeTime)

core/src/main/java/google/registry/model/transfer/DomainTransferData.java

@@ -164,7 +164,6 @@ public class DomainTransferData extends TransferData<DomainTransferData.Builder>
       serverApproveEntities = null;
       postLoad();
     }
-    hashCode = null; // reset the hash code since we may have changed the entities
   }
 
   /**

core/src/main/java/google/registry/persistence/transaction/JpaTransactionManager.java

@@ -19,6 +19,7 @@ import java.util.function.Supplier;
 import javax.persistence.EntityManager;
 import javax.persistence.Query;
 import javax.persistence.TypedQuery;
+import javax.persistence.criteria.CriteriaQuery;
 
 /** Sub-interface of {@link TransactionManager} which defines JPA related methods. */
 public interface JpaTransactionManager extends TransactionManager {
@@ -34,11 +35,18 @@ public interface JpaTransactionManager extends TransactionManager {
    */
   <T> TypedQuery<T> query(String sqlString, Class<T> resultClass);
 
+  /** Creates a JPA SQU query for the given criteria query. */
+  <T> TypedQuery<T> query(CriteriaQuery<T> criteriaQuery);
+
   /**
-   * Creates a JPA SQL query for the given query string (which does not return results).
+   * Creates a JPA SQL query for the given query string.
    *
    * <p>This is a convenience method for the longer <code>
    * jpaTm().getEntityManager().createQuery(...)</code>.
+   *
+   * <p>Note that while this method can legally be used for queries that return results, <u>it
+   * should not be</u>, as it does not correctly detach entities as must be done for nomulus model
+   * objects.
    */
   Query query(String sqlString);
 

core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java

@@ -23,6 +23,7 @@ import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 import static java.util.AbstractMap.SimpleEntry;
 import static java.util.stream.Collectors.joining;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableCollection;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -43,8 +44,11 @@ import google.registry.persistence.VKey;
 import google.registry.util.Clock;
 import google.registry.util.Retrier;
 import google.registry.util.SystemSleeper;
+import java.lang.reflect.Array;
 import java.lang.reflect.Field;
+import java.util.Calendar;
 import java.util.Collections;
+import java.util.Date;
 import java.util.IdentityHashMap;
 import java.util.List;
 import java.util.Map;
@@ -58,9 +62,14 @@ import javax.annotation.Nullable;
 import javax.persistence.EntityManager;
 import javax.persistence.EntityManagerFactory;
 import javax.persistence.EntityTransaction;
+import javax.persistence.FlushModeType;
+import javax.persistence.LockModeType;
+import javax.persistence.Parameter;
 import javax.persistence.PersistenceException;
 import javax.persistence.Query;
+import javax.persistence.TemporalType;
 import javax.persistence.TypedQuery;
+import javax.persistence.criteria.CriteriaQuery;
 import javax.persistence.metamodel.EntityType;
 import javax.persistence.metamodel.SingularAttribute;
 import org.joda.time.DateTime;
@@ -116,7 +125,12 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
 
   @Override
   public <T> TypedQuery<T> query(String sqlString, Class<T> resultClass) {
-    return getEntityManager().createQuery(sqlString, resultClass);
+    return new DetachingTypedQuery(getEntityManager().createQuery(sqlString, resultClass));
+  }
+
+  @Override
+  public <T> TypedQuery<T> query(CriteriaQuery<T> criteriaQuery) {
+    return new DetachingTypedQuery(getEntityManager().createQuery(criteriaQuery));
   }
 
   @Override
@@ -496,7 +510,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
         elements.size() <= 1,
         "Expected at most one entity of type %s, found at least two",
         clazz.getSimpleName());
-    return elements.stream().findFirst();
+    return elements.stream().findFirst().map(this::detach);
   }
 
   private int internalDelete(VKey<?> key) {
@@ -665,6 +679,33 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     }
   }
 
+  @Nullable
+  private <T> T detachIfEntity(@Nullable T object) {
+    if (object == null) {
+      return null;
+    }
+
+    // Check if the object is an array, if so we'll want to recurse through the elements.
+    if (object.getClass().isArray()) {
+      for (int i = 0; i < Array.getLength(object); ++i) {
+        detachIfEntity(Array.get(object, i));
+      }
+      return object;
+    }
+
+    // Check to see if it is an entity (queries can return raw column values or counts, so this
+    // could be String, Long, ...).
+    try {
+      getEntityManager().getMetamodel().entity(object.getClass());
+    } catch (IllegalArgumentException e) {
+      // The object is not an entity.  Return without detaching.
+      return object;
+    }
+
+    // At this point, object must be an entity.
+    return detach(object);
+  }
+
   /** Detach the entity, suitable for use in Optional.map(). */
   @Nullable
   private <T> T detach(@Nullable T entity) {
@@ -758,6 +799,207 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     }
   }
 
+  /**
+   * Typed query wrapper that applies a transform to all result objects.
+   *
+   * <p>This is used to detach objects upon load.
+   */
+  @VisibleForTesting
+  class DetachingTypedQuery<T> implements TypedQuery<T> {
+
+    TypedQuery<T> delegate;
+
+    public DetachingTypedQuery(TypedQuery<T> delegate) {
+      this.delegate = delegate;
+    }
+
+    @Override
+    public List<T> getResultList() {
+      return delegate
+          .getResultStream()
+          .map(JpaTransactionManagerImpl.this::detachIfEntity)
+          .collect(toImmutableList());
+    }
+
+    @Override
+    public Stream<T> getResultStream() {
+      return delegate.getResultStream().map(JpaTransactionManagerImpl.this::detachIfEntity);
+    }
+
+    @Override
+    public T getSingleResult() {
+      return detachIfEntity(delegate.getSingleResult());
+    }
+
+    @Override
+    public TypedQuery<T> setMaxResults(int maxResults) {
+      delegate.setMaxResults(maxResults);
+      return this;
+    }
+
+    @Override
+    public TypedQuery<T> setFirstResult(int startPosition) {
+      delegate.setFirstResult(startPosition);
+      return this;
+    }
+
+    @Override
+    public TypedQuery<T> setHint(String hintName, Object value) {
+      delegate.setHint(hintName, value);
+      return this;
+    }
+
+    @Override
+    public <U> TypedQuery<T> setParameter(Parameter<U> param, U value) {
+      delegate.setParameter(param, value);
+      return this;
+    }
+
+    @Override
+    public TypedQuery<T> setParameter(
+        Parameter<Calendar> param, Calendar value, TemporalType temporalType) {
+      delegate.setParameter(param, value, temporalType);
+      return this;
+    }
+
+    @Override
+    public TypedQuery<T> setParameter(
+        Parameter<Date> param, Date value, TemporalType temporalType) {
+      delegate.setParameter(param, value, temporalType);
+      return this;
+    }
+
+    @Override
+    public TypedQuery<T> setParameter(String name, Object value) {
+      delegate.setParameter(name, value);
+      return this;
+    }
+
+    @Override
+    public TypedQuery<T> setParameter(String name, Calendar value, TemporalType temporalType) {
+      delegate.setParameter(name, value, temporalType);
+      return this;
+    }
+
+    @Override
+    public TypedQuery<T> setParameter(String name, Date value, TemporalType temporalType) {
+      delegate.setParameter(name, value, temporalType);
+      return this;
+    }
+
+    @Override
+    public TypedQuery<T> setParameter(int position, Object value) {
+      delegate.setParameter(position, value);
+      return this;
+    }
+
+    @Override
+    public TypedQuery<T> setParameter(int position, Calendar value, TemporalType temporalType) {
+      delegate.setParameter(position, value, temporalType);
+      return this;
+    }
+
+    @Override
+    public TypedQuery<T> setParameter(int position, Date value, TemporalType temporalType) {
+      delegate.setParameter(position, value, temporalType);
+      return this;
+    }
+
+    @Override
+    public TypedQuery<T> setFlushMode(FlushModeType flushMode) {
+      delegate.setFlushMode(flushMode);
+      return this;
+    }
+
+    @Override
+    public TypedQuery<T> setLockMode(LockModeType lockMode) {
+      delegate.setLockMode(lockMode);
+      return this;
+    }
+
+    // Query interface
+
+    @Override
+    public int executeUpdate() {
+      return delegate.executeUpdate();
+    }
+
+    @Override
+    public int getMaxResults() {
+      return delegate.getMaxResults();
+    }
+
+    @Override
+    public int getFirstResult() {
+      return delegate.getFirstResult();
+    }
+
+    @Override
+    public Map<String, Object> getHints() {
+      return delegate.getHints();
+    }
+
+    @Override
+    public Set<Parameter<?>> getParameters() {
+      return delegate.getParameters();
+    }
+
+    @Override
+    public Parameter<?> getParameter(String name) {
+      return delegate.getParameter(name);
+    }
+
+    @Override
+    public <U> Parameter<U> getParameter(String name, Class<U> type) {
+      return delegate.getParameter(name, type);
+    }
+
+    @Override
+    public Parameter<?> getParameter(int position) {
+      return delegate.getParameter(position);
+    }
+
+    @Override
+    public <U> Parameter<U> getParameter(int position, Class<U> type) {
+      return delegate.getParameter(position, type);
+    }
+
+    @Override
+    public boolean isBound(Parameter<?> param) {
+      return delegate.isBound(param);
+    }
+
+    @Override
+    public <U> U getParameterValue(Parameter<U> param) {
+      return delegate.getParameterValue(param);
+    }
+
+    @Override
+    public Object getParameterValue(String name) {
+      return delegate.getParameterValue(name);
+    }
+
+    @Override
+    public Object getParameterValue(int position) {
+      return delegate.getParameterValue(position);
+    }
+
+    @Override
+    public FlushModeType getFlushMode() {
+      return delegate.getFlushMode();
+    }
+
+    @Override
+    public LockModeType getLockMode() {
+      return delegate.getLockMode();
+    }
+
+    @Override
+    public <U> U unwrap(Class<U> cls) {
+      return delegate.unwrap(cls);
+    }
+  }
+
   private class JpaQueryComposerImpl<T> extends QueryComposer<T> {
 
     private static final int DEFAULT_FETCH_SIZE = 1000;

core/src/main/java/google/registry/rdap/RdapActionBase.java

@@ -17,7 +17,6 @@ package google.registry.rdap;
 import static com.google.common.base.Charsets.UTF_8;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.net.HttpHeaders.ACCESS_CONTROL_ALLOW_ORIGIN;
-import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.request.Actions.getPathForAction;
 import static google.registry.util.DomainNameUtils.canonicalizeDomainName;
 import static javax.servlet.http.HttpServletResponse.SC_BAD_REQUEST;
@@ -29,10 +28,7 @@ import com.google.common.net.MediaType;
 import com.google.gson.Gson;
 import com.google.gson.GsonBuilder;
 import google.registry.config.RegistryConfig.Config;
-import google.registry.model.DatabaseMigrationUtils;
 import google.registry.model.EppResource;
-import google.registry.model.common.DatabaseTransitionSchedule.PrimaryDatabase;
-import google.registry.model.common.DatabaseTransitionSchedule.TransitionId;
 import google.registry.model.registrar.Registrar;
 import google.registry.rdap.RdapMetrics.EndpointType;
 import google.registry.rdap.RdapObjectClasses.ErrorResponse;
@@ -261,10 +257,4 @@ public abstract class RdapActionBase implements Runnable {
     return rdapJsonFormatter.getRequestTime();
   }
 
-  static boolean isDatastore() {
-    return tm().transact(
-            () ->
-                DatabaseMigrationUtils.getPrimaryDatabase(TransitionId.REPLAYED_ENTITIES)
-                    .equals(PrimaryDatabase.DATASTORE));
-  }
 }

core/src/main/java/google/registry/rdap/RdapDomainSearchAction.java

@@ -19,6 +19,7 @@ import static google.registry.model.EppResourceUtils.loadByForeignKey;
 import static google.registry.model.index.ForeignKeyIndex.loadAndGetKey;
 import static google.registry.model.ofy.ObjectifyService.auditedOfy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.request.Action.Method.GET;
 import static google.registry.request.Action.Method.HEAD;
 import static google.registry.util.DateTimeUtils.END_OF_TIME;
@@ -204,7 +205,7 @@ public class RdapDomainSearchAction extends RdapSearchActionBase {
     // need it.
     int querySizeLimit = RESULT_SET_SIZE_SCALING_FACTOR * rdapResultSetMaxSize;
     RdapResultSet<DomainBase> resultSet;
-    if (isDatastore()) {
+    if (tm().isOfy()) {
       Query<DomainBase> query =
           auditedOfy()
               .load()
@@ -260,7 +261,7 @@ public class RdapDomainSearchAction extends RdapSearchActionBase {
     // Don't use queryItems, because it doesn't handle pending deletes.
     int querySizeLimit = RESULT_SET_SIZE_SCALING_FACTOR * rdapResultSetMaxSize;
     RdapResultSet<DomainBase> resultSet;
-    if (isDatastore()) {
+    if (tm().isOfy()) {
       Query<DomainBase> query = auditedOfy().load().type(DomainBase.class).filter("tld", tld);
       if (cursorString.isPresent()) {
         query = query.filter("fullyQualifiedDomainName >", cursorString.get());
@@ -337,7 +338,7 @@ public class RdapDomainSearchAction extends RdapSearchActionBase {
     // incomplete result set if a search asks for something like "ns*", but we need to enforce a
     // limit in order to avoid arbitrarily long-running queries.
     Optional<String> desiredRegistrar = getDesiredRegistrar();
-    if (isDatastore()) {
+    if (tm().isOfy()) {
       Query<HostResource> query =
           queryItems(
               HostResource.class,
@@ -472,7 +473,7 @@ public class RdapDomainSearchAction extends RdapSearchActionBase {
   private DomainSearchResponse searchByNameserverIp(final InetAddress inetAddress) {
     Optional<String> desiredRegistrar = getDesiredRegistrar();
     ImmutableSet<VKey<HostResource>> hostKeys;
-    if (isDatastore()) {
+    if (tm().isOfy()) {
       Query<HostResource> query =
           queryItems(
               HostResource.class,
@@ -544,7 +545,7 @@ public class RdapDomainSearchAction extends RdapSearchActionBase {
     int numHostKeysSearched = 0;
     for (List<VKey<HostResource>> chunk : Iterables.partition(hostKeys, 30)) {
       numHostKeysSearched += chunk.size();
-      if (isDatastore()) {
+      if (tm().isOfy()) {
         Query<DomainBase> query =
             auditedOfy()
                 .load()
@@ -588,8 +589,7 @@ public class RdapDomainSearchAction extends RdapSearchActionBase {
                               cursorString.get());
                     }
                     jpaTm()
-                        .getEntityManager()
-                        .createQuery(queryBuilder.build())
+                        .query(queryBuilder.build())
                         .getResultStream()
                         .filter(this::isAuthorized)
                         .forEach(

core/src/main/java/google/registry/rdap/RdapEntitySearchAction.java

@@ -262,7 +262,7 @@ public class RdapEntitySearchAction extends RdapSearchActionBase {
           || (cursorType == CursorType.REGISTRAR)) {
         resultSet = RdapResultSet.create(ImmutableList.of());
       } else {
-        if (isDatastore()) {
+        if (tm().isOfy()) {
           Query<ContactResource> query =
               queryItems(
                   ContactResource.class,
@@ -386,7 +386,7 @@ public class RdapEntitySearchAction extends RdapSearchActionBase {
       if (subtype == Subtype.REGISTRARS) {
         contactResultSet = RdapResultSet.create(ImmutableList.of());
       } else {
-        if (isDatastore()) {
+        if (tm().isOfy()) {
           contactResultSet =
               getMatchingResources(
                   queryItemsByKey(

core/src/main/java/google/registry/rdap/RdapNameserverSearchAction.java

@@ -16,6 +16,7 @@ package google.registry.rdap;
 
 import static google.registry.model.EppResourceUtils.loadByForeignKey;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.request.Action.Method.GET;
 import static google.registry.request.Action.Method.HEAD;
 import static google.registry.util.DateTimeUtils.END_OF_TIME;
@@ -220,7 +221,7 @@ public class RdapNameserverSearchAction extends RdapSearchActionBase {
   private NameserverSearchResponse searchByNameUsingPrefix(RdapSearchPattern partialStringQuery) {
     // Add 1 so we can detect truncation.
     int querySizeLimit = getStandardQuerySizeLimit();
-    if (isDatastore()) {
+    if (tm().isOfy()) {
       Query<HostResource> query =
           queryItems(
               HostResource.class,
@@ -254,7 +255,7 @@ public class RdapNameserverSearchAction extends RdapSearchActionBase {
     // Add 1 so we can detect truncation.
     int querySizeLimit = getStandardQuerySizeLimit();
     RdapResultSet<HostResource> rdapResultSet;
-    if (isDatastore()) {
+    if (tm().isOfy()) {
       Query<HostResource> query =
           queryItems(
               HostResource.class,

core/src/main/java/google/registry/rdap/RdapSearchActionBase.java

@@ -202,11 +202,7 @@ public abstract class RdapSearchActionBase extends RdapActionBase {
               desiredRegistrar.get());
     }
     List<T> queryResult =
-        jpaTm()
-            .getEntityManager()
-            .createQuery(builder.build())
-            .setMaxResults(querySizeLimit)
-            .getResultList();
+        jpaTm().query(builder.build()).setMaxResults(querySizeLimit).getResultList();
     if (checkForVisibility) {
       return filterResourcesByVisibility(queryResult, querySizeLimit);
     } else {

core/src/main/java/google/registry/schema/replay/ReplicateToDatastoreAction.java

@@ -22,12 +22,17 @@ import static google.registry.request.Action.Method.GET;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
 import com.google.common.flogger.FluentLogger;
+import google.registry.model.common.DatabaseMigrationStateSchedule;
+import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState;
+import google.registry.model.common.DatabaseMigrationStateSchedule.ReplayDirection;
 import google.registry.persistence.transaction.Transaction;
 import google.registry.persistence.transaction.TransactionEntity;
 import google.registry.request.Action;
 import google.registry.request.auth.Auth;
+import google.registry.util.Clock;
 import java.io.IOException;
 import java.util.List;
+import javax.inject.Inject;
 import javax.persistence.NoResultException;
 
 /** Cron task to replicate from Cloud SQL to datastore. */
@@ -48,6 +53,13 @@ class ReplicateToDatastoreAction implements Runnable {
    */
   public static final int BATCH_SIZE = 200;
 
+  private final Clock clock;
+
+  @Inject
+  public ReplicateToDatastoreAction(Clock clock) {
+    this.clock = clock;
+  }
+
   @VisibleForTesting
   List<TransactionEntity> getTransactionBatch() {
     // Get the next batch of transactions that we haven't replicated.
@@ -59,7 +71,8 @@ class ReplicateToDatastoreAction implements Runnable {
                   jpaTm()
                       .query(
                           "SELECT txn FROM TransactionEntity txn WHERE id >"
-                              + " :lastId ORDER BY id")
+                              + " :lastId ORDER BY id",
+                          TransactionEntity.class)
                       .setParameter("lastId", lastSqlTxnBeforeBatch.getTransactionId())
                       .setMaxResults(BATCH_SIZE)
                       .getResultList());
@@ -69,7 +82,7 @@ class ReplicateToDatastoreAction implements Runnable {
   }
 
   /**
-   * Apply a transaction to datastore, returns true if there was a fatal error and the batch should
+   * Apply a transaction to Datastore, returns true if there was a fatal error and the batch should
    * be aborted.
    */
   @VisibleForTesting
@@ -122,6 +135,13 @@ class ReplicateToDatastoreAction implements Runnable {
 
   @Override
   public void run() {
+    MigrationState state = DatabaseMigrationStateSchedule.getValueAtTime(clock.nowUtc());
+    if (!state.getReplayDirection().equals(ReplayDirection.SQL_TO_DATASTORE)) {
+      logger.atInfo().log(
+          String.format(
+              "Skipping ReplicateToDatastoreAction because we are in migration phase %s.", state));
+      return;
+    }
     // TODO(b/181758163): Deal with objects that don't exist in Cloud SQL, e.g. ForeignKeyIndex,
     // EppResourceIndex.
     logger.atInfo().log("Processing transaction replay batch Cloud SQL -> Cloud Datastore");

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

@@ -0,0 +1,34 @@
+// 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.Parameters;
+import google.registry.model.common.DatabaseMigrationStateSchedule;
+import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState;
+import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationStateTransition;
+import google.registry.model.common.TimedTransitionProperty;
+
+/** A command to check the current Registry 3.0 migration state of the database. */
+@Parameters(separators = " =", commandDescription = "Check current Registry 3.0 migration state")
+public class GetDatabaseMigrationStateCommand implements CommandWithRemoteApi {
+
+  @Override
+  public void run() throws Exception {
+    TimedTransitionProperty<MigrationState, MigrationStateTransition> migrationSchedule =
+        DatabaseMigrationStateSchedule.get();
+    System.out.println(
+        String.format("Current migration schedule: %s", migrationSchedule.toValueMap()));
+  }
+}

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

@@ -1,44 +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.
-
-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

@@ -69,7 +69,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_database_migration_state", GetDatabaseMigrationStateCommand.class)
           .put("get_domain", GetDomainCommand.class)
           .put("get_history_entries", GetHistoryEntriesCommand.class)
           .put("get_host", GetHostCommand.class)
@@ -111,7 +111,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_database_migration_state", SetDatabaseMigrationStateCommand.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/SetDatabaseMigrationStateCommand.java

@@ -0,0 +1,70 @@
+// 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.persistence.transaction.TransactionManagerFactory.ofyTm;
+
+import com.beust.jcommander.Parameter;
+import com.beust.jcommander.Parameters;
+import com.google.common.collect.ImmutableSortedMap;
+import google.registry.model.common.DatabaseMigrationStateSchedule;
+import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState;
+import google.registry.tools.params.TransitionListParameter.MigrationStateTransitions;
+import org.joda.time.DateTime;
+
+/** Command to set the Registry 3.0 database migration state schedule. */
+@Parameters(
+    separators = " =",
+    commandDescription = "Set the current database migration state schedule.")
+public class SetDatabaseMigrationStateCommand extends ConfirmingCommand
+    implements CommandWithRemoteApi {
+
+  private static final String WARNING_MESSAGE =
+      "Attempting to change the schedule with an effect that would take place within the next 10 "
+          + "minutes. The cache expiration duration is 5 minutes so this MAY BE DANGEROUS.\n";
+
+  @Parameter(
+      names = "--migration_schedule",
+      converter = MigrationStateTransitions.class,
+      validateWith = MigrationStateTransitions.class,
+      required = true,
+      description =
+          "Comma-delimited list of database transitions, of the form"
+              + " <time>=<migration-state>[,<time>=<migration-state>]*")
+  ImmutableSortedMap<DateTime, MigrationState> transitionSchedule;
+
+  @Override
+  protected String prompt() {
+    return ofyTm()
+        .transact(
+            () -> {
+              StringBuilder result = new StringBuilder();
+              DateTime now = ofyTm().getTransactionTime();
+              DateTime nextTransition = transitionSchedule.ceilingKey(now);
+              if (nextTransition != null && nextTransition.isBefore(now.plusMinutes(10))) {
+                result.append(WARNING_MESSAGE);
+              }
+              return result
+                  .append(String.format("Set new migration state schedule %s?", transitionSchedule))
+                  .toString();
+            });
+  }
+
+  @Override
+  protected String execute() {
+    ofyTm().transact(() -> DatabaseMigrationStateSchedule.set(transitionSchedule));
+    return String.format("Successfully set new migration state schedule %s", transitionSchedule);
+  }
+}

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

@@ -1,69 +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.
-
-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;
-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 org.joda.time.DateTime;
-
-/** Command to update {@link DatabaseTransitionSchedule}. */
-@Parameters(
-    separators = " =",
-    commandDescription = "Set the database transition schedule for transition id.")
-public class SetDatabaseTransitionScheduleCommand extends ConfirmingCommand
-    implements CommandWithRemoteApi {
-
-  @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",
-      required = true,
-      description = "Transition id string for the schedule being updated")
-  private TransitionId transitionId;
-
-  @Override
-  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));
-    ofyTm().transact(() -> ofyTm().put(newSchedule));
-    return String.format(
-        "Inserted new schedule %s for transition ID %s.", transitionSchedule, transitionId);
-  }
-}

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

@@ -1,23 +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.
-
-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,7 +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.common.DatabaseMigrationStateSchedule.MigrationState;
 import google.registry.model.domain.token.AllocationToken.TokenStatus;
 import google.registry.model.registry.Registry.TldState;
 import org.joda.money.Money;
@@ -74,11 +74,11 @@ public abstract class TransitionListParameter<V> extends KeyValueMapParameter<Da
     }
   }
 
-  /** Converter-validator for primary database transitions. */
-  public static class PrimaryDatabaseTransitions extends TransitionListParameter<PrimaryDatabase> {
+  /** Converter-validator for states of the Registry 3.0 database migration. */
+  public static class MigrationStateTransitions extends TransitionListParameter<MigrationState> {
     @Override
-    protected PrimaryDatabase parseValue(String value) {
-      return PrimaryDatabase.valueOf(value);
+    protected MigrationState parseValue(String value) {
+      return MigrationState.valueOf(value);
     }
   }
 }

core/src/main/resources/google/registry/beam/invoicing/sql/billing_events.sql

@@ -27,7 +27,7 @@ SELECT
   reason as action,
   targetId as domain,
   BillingEvent.domainRepoId as repositoryId,
-  periodYears as years,
+  IFNULL(periodYears, 0) as years,
   BillingEvent.currency AS currency,
   BillingEvent.amount as amount,
   -- We'll strip out non-useful flags downstream

core/src/test/java/google/registry/backup/ReplayCommitLogsToSqlActionTest.java

@@ -21,13 +21,16 @@ import static google.registry.backup.RestoreCommitLogsActionTest.GCS_BUCKET;
 import static google.registry.backup.RestoreCommitLogsActionTest.createCheckpoint;
 import static google.registry.backup.RestoreCommitLogsActionTest.saveDiffFile;
 import static google.registry.backup.RestoreCommitLogsActionTest.saveDiffFileNotToRestore;
+import static google.registry.model.common.DatabaseMigrationStateSchedule.DEFAULT_TRANSITION_MAP;
 import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
 import static google.registry.model.ofy.CommitLogBucket.getBucketKey;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatabaseHelper.createTld;
 import static google.registry.testing.DatabaseHelper.newDomainBase;
 import static google.registry.testing.DatabaseHelper.persistActiveContact;
+import static google.registry.util.DateTimeUtils.START_OF_TIME;
 import static javax.servlet.http.HttpServletResponse.SC_NO_CONTENT;
 import static javax.servlet.http.HttpServletResponse.SC_OK;
 import static org.mockito.ArgumentMatchers.any;
@@ -40,9 +43,11 @@ import com.google.appengine.tools.cloudstorage.GcsServiceFactory;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.truth.Truth8;
 import com.googlecode.objectify.Key;
-import google.registry.config.RegistryConfig;
+import google.registry.model.common.DatabaseMigrationStateSchedule;
+import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState;
 import google.registry.model.contact.ContactResource;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.GracePeriod;
@@ -122,11 +127,20 @@ public class ReplayCommitLogsToSqlActionTest {
     action.gcsService = gcsService;
     action.response = response;
     action.requestStatusChecker = requestStatusChecker;
+    action.clock = fakeClock;
     action.diffLister = new GcsDiffFileLister();
     action.diffLister.gcsService = gcsService;
     action.diffLister.gcsBucket = GCS_BUCKET;
     action.diffLister.executor = newDirectExecutorService();
-    RegistryConfig.overrideCloudSqlReplayCommitLogs(true);
+    ofyTm()
+        .transact(
+            () ->
+                DatabaseMigrationStateSchedule.set(
+                    ImmutableSortedMap.of(
+                        START_OF_TIME,
+                        MigrationState.DATASTORE_ONLY,
+                        START_OF_TIME.plusMinutes(1),
+                        MigrationState.DATASTORE_PRIMARY)));
     TestObject.beforeSqlSaveCallCount = 0;
     TestObject.beforeSqlDeleteCallCount = 0;
   }
@@ -329,7 +343,7 @@ public class ReplayCommitLogsToSqlActionTest {
     ContactResource contactWithEdit =
         contact.asBuilder().setEmailAddress("replay@example.tld").build();
     CommitLogMutation contactMutation =
-        tm().transact(() -> CommitLogMutation.create(manifestKey, contactWithEdit));
+        ofyTm().transact(() -> CommitLogMutation.create(manifestKey, contactWithEdit));
 
     jpaTm().transact(() -> SqlReplayCheckpoint.set(now.minusMinutes(1).minusMillis(1)));
 
@@ -421,11 +435,13 @@ public class ReplayCommitLogsToSqlActionTest {
 
   @Test
   void testFailure_notEnabled() {
-    RegistryConfig.overrideCloudSqlReplayCommitLogs(false);
+    ofyTm().transact(() -> DatabaseMigrationStateSchedule.set(DEFAULT_TRANSITION_MAP.toValueMap()));
     action.run();
     assertThat(response.getStatus()).isEqualTo(SC_NO_CONTENT);
     assertThat(response.getPayload())
-        .isEqualTo("ReplayCommitLogsToSqlAction was called but disabled in the config.");
+        .isEqualTo(
+            "Skipping ReplayCommitLogsToSqlAction because we are in migration phase"
+                + " DATASTORE_ONLY.");
   }
 
   @Test

core/src/test/java/google/registry/batch/DeleteExpiredDomainsActionTest.java

@@ -14,12 +14,14 @@
 
 package google.registry.batch;
 
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.eppcommon.StatusValue.PENDING_DELETE;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_CREATE;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 import static google.registry.testing.DatabaseHelper.createTld;
+import static google.registry.testing.DatabaseHelper.loadByEntity;
 import static google.registry.testing.DatabaseHelper.newDomainBase;
 import static google.registry.testing.DatabaseHelper.persistActiveDomain;
 import static google.registry.testing.DatabaseHelper.persistResource;
@@ -38,27 +40,35 @@ import google.registry.model.ofy.Ofy;
 import google.registry.model.poll.PollMessage;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.monitoring.whitebox.EppMetric;
+import google.registry.persistence.transaction.QueryComposer.Comparator;
 import google.registry.testing.AppEngineExtension;
+import google.registry.testing.DualDatabaseTest;
 import google.registry.testing.FakeClock;
 import google.registry.testing.FakeLockHandler;
 import google.registry.testing.FakeResponse;
 import google.registry.testing.InjectExtension;
+import google.registry.testing.TestOfyAndSql;
 import java.util.Optional;
 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 DeleteExpiredDomainsAction}. */
+@DualDatabaseTest
 class DeleteExpiredDomainsActionTest {
 
+  private final FakeClock clock = new FakeClock(DateTime.parse("2016-06-13T20:21:22Z"));
+
   @RegisterExtension
   public final AppEngineExtension appEngine =
-      AppEngineExtension.builder().withDatastoreAndCloudSql().withTaskQueue().build();
+      AppEngineExtension.builder()
+          .withDatastoreAndCloudSql()
+          .withClock(clock)
+          .withTaskQueue()
+          .build();
 
   @RegisterExtension public final InjectExtension inject = new InjectExtension();
 
-  private final FakeClock clock = new FakeClock(DateTime.parse("2016-06-13T20:21:22Z"));
   private final FakeResponse response = new FakeResponse();
   private DeleteExpiredDomainsAction action;
 
@@ -78,7 +88,7 @@ class DeleteExpiredDomainsActionTest {
             eppController, "NewRegistrar", clock, new FakeLockHandler(true), response);
   }
 
-  @Test
+  @TestOfyAndSql
   void test_deletesOnlyExpiredDomain() {
     // A normal, active autorenewing domain that shouldn't be touched.
     DomainBase activeDomain = persistActiveDomain("foo.tld");
@@ -105,7 +115,7 @@ class DeleteExpiredDomainsActionTest {
     // to operate on.)
     DomainBase pendingExpirationDomain = persistNonAutorenewingDomain("fizz.tld");
 
-    assertThat(tm().loadByEntity(pendingExpirationDomain).getStatusValues())
+    assertThat(loadByEntity(pendingExpirationDomain).getStatusValues())
         .doesNotContain(PENDING_DELETE);
     // action.run() does not use any test helper that can advance the fake clock. We manually
     // advance the clock to emulate the actual behavior. This works because the action only has
@@ -113,17 +123,17 @@ class DeleteExpiredDomainsActionTest {
     clock.advanceOneMilli();
     action.run();
 
-    DomainBase reloadedActiveDomain = tm().loadByEntity(activeDomain);
+    DomainBase reloadedActiveDomain = loadByEntity(activeDomain);
     assertThat(reloadedActiveDomain).isEqualTo(activeDomain);
     assertThat(reloadedActiveDomain.getStatusValues()).doesNotContain(PENDING_DELETE);
-    assertThat(tm().loadByEntity(alreadyDeletedDomain)).isEqualTo(alreadyDeletedDomain);
-    assertThat(tm().loadByEntity(notYetExpiredDomain)).isEqualTo(notYetExpiredDomain);
-    DomainBase reloadedExpiredDomain = tm().loadByEntity(pendingExpirationDomain);
+    assertThat(loadByEntity(alreadyDeletedDomain)).isEqualTo(alreadyDeletedDomain);
+    assertThat(loadByEntity(notYetExpiredDomain)).isEqualTo(notYetExpiredDomain);
+    DomainBase reloadedExpiredDomain = loadByEntity(pendingExpirationDomain);
     assertThat(reloadedExpiredDomain.getStatusValues()).contains(PENDING_DELETE);
     assertThat(reloadedExpiredDomain.getDeletionTime()).isEqualTo(clock.nowUtc().plusDays(35));
   }
 
-  @Test
+  @TestOfyAndSql
   void test_deletesThreeDomainsInOneRun() throws Exception {
     DomainBase domain1 = persistNonAutorenewingDomain("ecck1.tld");
     DomainBase domain2 = persistNonAutorenewingDomain("veee2.tld");
@@ -135,14 +145,14 @@ class DeleteExpiredDomainsActionTest {
     int maxRetries = 5;
     while (true) {
       ImmutableSet<String> matchingDomains =
-          ofy()
-              .load()
-              .type(DomainBase.class)
-              .filter("autorenewEndTime <=", clock.nowUtc())
-              .list()
-              .stream()
-              .map(DomainBase::getDomainName)
-              .collect(ImmutableSet.toImmutableSet());
+          transactIfJpaTm(
+              () ->
+                  tm()
+                      .createQueryComposer(DomainBase.class)
+                      .where("autorenewEndTime", Comparator.LTE, clock.nowUtc())
+                      .stream()
+                      .map(DomainBase::getDomainName)
+                      .collect(toImmutableSet()));
       if (matchingDomains.containsAll(ImmutableSet.of("ecck1.tld", "veee2.tld", "tarm3.tld"))) {
         break;
       }
@@ -158,9 +168,9 @@ class DeleteExpiredDomainsActionTest {
     action.run();
     clock.disableAutoIncrement();
 
-    assertThat(tm().loadByEntity(domain1).getStatusValues()).contains(PENDING_DELETE);
-    assertThat(tm().loadByEntity(domain2).getStatusValues()).contains(PENDING_DELETE);
-    assertThat(tm().loadByEntity(domain3).getStatusValues()).contains(PENDING_DELETE);
+    assertThat(loadByEntity(domain1).getStatusValues()).contains(PENDING_DELETE);
+    assertThat(loadByEntity(domain2).getStatusValues()).contains(PENDING_DELETE);
+    assertThat(loadByEntity(domain3).getStatusValues()).contains(PENDING_DELETE);
   }
 
   private DomainBase persistNonAutorenewingDomain(String domainName) {

core/src/test/java/google/registry/batch/ExpandRecurringBillingEventsActionTest.java

@@ -21,6 +21,7 @@ import static google.registry.model.ofy.ObjectifyService.auditedOfy;
 import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_AUTORENEW;
 import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_CREATE;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 import static google.registry.testing.DatabaseHelper.assertBillingEvents;
 import static google.registry.testing.DatabaseHelper.assertBillingEventsForResource;
 import static google.registry.testing.DatabaseHelper.createTld;
@@ -52,9 +53,14 @@ import google.registry.model.registry.Registry;
 import google.registry.model.reporting.DomainTransactionRecord;
 import google.registry.model.reporting.DomainTransactionRecord.TransactionReportField;
 import google.registry.model.reporting.HistoryEntry;
+import google.registry.testing.DualDatabaseTest;
 import google.registry.testing.FakeClock;
 import google.registry.testing.FakeResponse;
 import google.registry.testing.InjectExtension;
+import google.registry.testing.ReplayExtension;
+import google.registry.testing.TestOfyAndSql;
+import google.registry.testing.TestOfyOnly;
+import google.registry.testing.TestSqlOnly;
 import google.registry.testing.mapreduce.MapreduceTestCase;
 import java.util.ArrayList;
 import java.util.List;
@@ -62,17 +68,25 @@ import java.util.Optional;
 import org.joda.money.Money;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Order;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Unit tests for {@link ExpandRecurringBillingEventsAction}. */
+@DualDatabaseTest
 public class ExpandRecurringBillingEventsActionTest
     extends MapreduceTestCase<ExpandRecurringBillingEventsAction> {
 
-  @RegisterExtension public final InjectExtension inject = new InjectExtension();
+  private DateTime currentTestTime = DateTime.parse("1999-01-05T00:00:00Z");
+  private final FakeClock clock = new FakeClock(currentTestTime);
 
-  private final DateTime beginningOfTest = DateTime.parse("2000-10-02T00:00:00Z");
-  private final FakeClock clock = new FakeClock(beginningOfTest);
+  @Order(Order.DEFAULT - 1)
+  @RegisterExtension
+  public final InjectExtension inject =
+      new InjectExtension().withStaticFieldOverride(Ofy.class, "clock", clock);
+
+  @Order(Order.DEFAULT - 2)
+  @RegisterExtension
+  public final ReplayExtension replayExtension = ReplayExtension.createWithCompare(clock);
 
   private DomainBase domain;
   private DomainHistory historyEntry;
@@ -80,7 +94,6 @@ public class ExpandRecurringBillingEventsActionTest
 
   @BeforeEach
   void beforeEach() {
-    inject.setStaticField(Ofy.class, "clock", clock);
     action = new ExpandRecurringBillingEventsAction();
     action.mrRunner = makeDefaultRunner();
     action.clock = clock;
@@ -111,27 +124,37 @@ public class ExpandRecurringBillingEventsActionTest
             .setRecurrenceEndTime(END_OF_TIME)
             .setTargetId(domain.getDomainName())
             .build();
+    currentTestTime = clock.nowUtc();
+    clock.setTo(DateTime.parse("2000-10-02T00:00:00Z"));
   }
 
   private void saveCursor(final DateTime cursorTime) {
     tm().transact(() -> tm().put(Cursor.createGlobal(RECURRING_BILLING, cursorTime)));
   }
 
-  private void runMapreduce() throws Exception {
+  private void runAction() throws Exception {
     action.response = new FakeResponse();
     action.run();
+    // Need to save the current test time before running the mapreduce, which increments the clock.
+    // The execution time (e. g. transaction time) is captured when the action starts running so
+    // the passage of time afterward does not affect the timestamp stored in the billing events.
+    currentTestTime = clock.nowUtc();
     executeTasksUntilEmpty("mapreduce", clock);
     auditedOfy().clearSessionCache();
   }
 
   private void assertCursorAt(DateTime expectedCursorTime) {
-    Cursor cursor = auditedOfy().load().key(Cursor.createGlobalKey(RECURRING_BILLING)).now();
+    Cursor cursor =
+        transactIfJpaTm(() -> tm().loadByKey(Cursor.createGlobalVKey(RECURRING_BILLING)));
     assertThat(cursor).isNotNull();
     assertThat(cursor.getCursorTime()).isEqualTo(expectedCursorTime);
   }
 
   private void assertHistoryEntryMatches(
-      DomainBase domain, HistoryEntry actual, String clientId, DateTime billingTime,
+      DomainBase domain,
+      HistoryEntry actual,
+      String clientId,
+      DateTime billingTime,
       boolean shouldHaveTxRecord) {
     assertThat(actual.getBySuperuser()).isFalse();
     assertThat(actual.getClientId()).isEqualTo(clientId);
@@ -145,10 +168,7 @@ public class ExpandRecurringBillingEventsActionTest
       assertThat(actual.getDomainTransactionRecords())
           .containsExactly(
               DomainTransactionRecord.create(
-                  "tld",
-                  billingTime,
-                  TransactionReportField.NET_RENEWS_1_YR,
-                  1));
+                  "tld", billingTime, TransactionReportField.NET_RENEWS_1_YR, 1));
     } else {
       assertThat(actual.getDomainTransactionRecords()).isEmpty();
     }
@@ -163,28 +183,26 @@ public class ExpandRecurringBillingEventsActionTest
         .setFlags(ImmutableSet.of(Flag.AUTO_RENEW, Flag.SYNTHETIC))
         .setPeriodYears(1)
         .setReason(Reason.RENEW)
-        .setSyntheticCreationTime(beginningOfTest)
+        .setSyntheticCreationTime(currentTestTime)
         .setCancellationMatchingBillingEvent(recurring.createVKey())
         .setTargetId(domain.getDomainName());
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_expandSingleEvent() throws Exception {
     persistResource(recurring);
     action.cursorTimeParam = Optional.of(START_OF_TIME);
-    runMapreduce();
+    runAction();
     DomainHistory persistedEntry =
         getOnlyHistoryEntryOfType(domain, DOMAIN_AUTORENEW, DomainHistory.class);
     assertHistoryEntryMatches(
         domain, persistedEntry, "TheRegistrar", DateTime.parse("2000-02-19T00:00:00Z"), true);
-    BillingEvent.OneTime expected = defaultOneTimeBuilder()
-        .setParent(persistedEntry)
-        .build();
+    BillingEvent.OneTime expected = defaultOneTimeBuilder().setParent(persistedEntry).build();
     assertBillingEventsForResource(domain, expected, recurring);
-    assertCursorAt(beginningOfTest);
+    assertCursorAt(currentTestTime);
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_expandSingleEvent_deletedDomain() throws Exception {
     DateTime deletionTime = DateTime.parse("2000-08-01T00:00:00Z");
     DomainBase deletedDomain = persistDeletedDomain("deleted.tld", deletionTime);
@@ -209,7 +227,7 @@ public class ExpandRecurringBillingEventsActionTest
                 .setTargetId(deletedDomain.getDomainName())
                 .build());
     action.cursorTimeParam = Optional.of(START_OF_TIME);
-    runMapreduce();
+    runAction();
     DomainHistory persistedEntry =
         getOnlyHistoryEntryOfType(deletedDomain, DOMAIN_AUTORENEW, DomainHistory.class);
     assertHistoryEntryMatches(
@@ -224,69 +242,69 @@ public class ExpandRecurringBillingEventsActionTest
             .setTargetId(deletedDomain.getDomainName())
             .build();
     assertBillingEventsForResource(deletedDomain, expected, recurring);
-    assertCursorAt(beginningOfTest);
+    assertCursorAt(currentTestTime);
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_expandSingleEvent_idempotentForDuplicateRuns() throws Exception {
     persistResource(recurring);
     action.cursorTimeParam = Optional.of(START_OF_TIME);
-    runMapreduce();
+    runAction();
     DomainHistory persistedEntry =
         getOnlyHistoryEntryOfType(domain, DOMAIN_AUTORENEW, DomainHistory.class);
     assertHistoryEntryMatches(
         domain, persistedEntry, "TheRegistrar", DateTime.parse("2000-02-19T00:00:00Z"), true);
     BillingEvent.OneTime expected = defaultOneTimeBuilder().setParent(persistedEntry).build();
-    assertCursorAt(beginningOfTest);
+    assertCursorAt(currentTestTime);
     DateTime beginningOfSecondRun = clock.nowUtc();
     action.response = new FakeResponse();
-    runMapreduce();
+    runAction();
     assertCursorAt(beginningOfSecondRun);
     assertBillingEventsForResource(domain, expected, recurring);
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_expandSingleEvent_idempotentForExistingOneTime() throws Exception {
     persistResource(recurring);
-    BillingEvent.OneTime persisted = persistResource(defaultOneTimeBuilder()
-        .setParent(historyEntry)
-        .build());
+    BillingEvent.OneTime persisted =
+        persistResource(defaultOneTimeBuilder().setParent(historyEntry).build());
     action.cursorTimeParam = Optional.of(START_OF_TIME);
-    runMapreduce();
+    runAction();
     // No new history entries should be generated
     assertThat(getHistoryEntriesOfType(domain, DOMAIN_AUTORENEW)).isEmpty();
-    assertCursorAt(beginningOfTest);
+    assertCursorAt(currentTestTime);
     // No additional billing events should be generated
     assertBillingEventsForResource(domain, persisted, recurring);
   }
 
-  @Test
+  @TestSqlOnly
   void testSuccess_expandSingleEvent_notIdempotentForDifferentBillingTime() throws Exception {
     persistResource(recurring);
     action.cursorTimeParam = Optional.of(START_OF_TIME);
-    runMapreduce();
+    runAction();
     DomainHistory persistedEntry =
         getOnlyHistoryEntryOfType(domain, DOMAIN_AUTORENEW, DomainHistory.class);
     assertHistoryEntryMatches(
         domain, persistedEntry, "TheRegistrar", DateTime.parse("2000-02-19T00:00:00Z"), true);
     BillingEvent.OneTime expected = defaultOneTimeBuilder().setParent(persistedEntry).build();
     // Persist an otherwise identical billing event that differs only in billing time.
-    BillingEvent.OneTime persisted = persistResource(expected.asBuilder()
-        .setBillingTime(DateTime.parse("1999-02-19T00:00:00Z"))
-        .setEventTime(DateTime.parse("1999-01-05T00:00:00Z"))
-        .build());
-    assertCursorAt(beginningOfTest);
+    BillingEvent.OneTime persisted =
+        persistResource(
+            expected
+                .asBuilder()
+                .setBillingTime(DateTime.parse("1999-02-19T00:00:00Z"))
+                .setEventTime(DateTime.parse("1999-01-05T00:00:00Z"))
+                .build());
+    assertCursorAt(currentTestTime);
     assertBillingEventsForResource(domain, persisted, expected, recurring);
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_expandSingleEvent_notIdempotentForDifferentRecurring() throws Exception {
     persistResource(recurring);
-    BillingEvent.Recurring recurring2 = persistResource(recurring.asBuilder()
-        .setId(3L)
-        .build());
+    BillingEvent.Recurring recurring2 = persistResource(recurring.asBuilder().setId(3L).build());
     action.cursorTimeParam = Optional.of(START_OF_TIME);
-    runMapreduce();
+    runAction();
     List<DomainHistory> persistedEntries =
         getHistoryEntriesOfType(domain, DOMAIN_AUTORENEW, DomainHistory.class);
     for (HistoryEntry persistedEntry : persistedEntries) {
@@ -294,9 +312,8 @@ public class ExpandRecurringBillingEventsActionTest
           domain, persistedEntry, "TheRegistrar", DateTime.parse("2000-02-19T00:00:00Z"), true);
     }
     assertThat(persistedEntries).hasSize(2);
-    BillingEvent.OneTime expected = defaultOneTimeBuilder()
-        .setParent(persistedEntries.get(0))
-        .build();
+    BillingEvent.OneTime expected =
+        defaultOneTimeBuilder().setParent(persistedEntries.get(0)).build();
     // Persist an otherwise identical billing event that differs only in recurring event key.
     BillingEvent.OneTime persisted =
         expected
@@ -304,164 +321,166 @@ public class ExpandRecurringBillingEventsActionTest
             .setParent(persistedEntries.get(1))
             .setCancellationMatchingBillingEvent(recurring2.createVKey())
             .build();
-    assertCursorAt(beginningOfTest);
+    assertCursorAt(currentTestTime);
     assertBillingEventsForResource(domain, persisted, expected, recurring, recurring2);
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_ignoreRecurringBeforeWindow() throws Exception {
-    recurring = persistResource(recurring.asBuilder()
-        .setEventTime(DateTime.parse("1997-01-05T00:00:00Z"))
-        .setRecurrenceEndTime(DateTime.parse("1999-10-05T00:00:00Z"))
-        .build());
+    recurring =
+        persistResource(
+            recurring
+                .asBuilder()
+                .setEventTime(DateTime.parse("1997-01-05T00:00:00Z"))
+                .setRecurrenceEndTime(DateTime.parse("1999-10-05T00:00:00Z"))
+                .build());
     action.cursorTimeParam = Optional.of(DateTime.parse("2000-01-01T00:00:00Z"));
-    runMapreduce();
+    runAction();
     // No new history entries should be generated
     assertThat(getHistoryEntriesOfType(domain, DOMAIN_AUTORENEW)).isEmpty();
     assertBillingEventsForResource(domain, recurring);
-    assertCursorAt(beginningOfTest);
+    assertCursorAt(currentTestTime);
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_ignoreRecurringAfterWindow() throws Exception {
-    recurring = persistResource(recurring.asBuilder()
-        .setEventTime(clock.nowUtc().plusYears(2))
-        .build());
+    recurring =
+        persistResource(recurring.asBuilder().setEventTime(clock.nowUtc().plusYears(2)).build());
     action.cursorTimeParam = Optional.of(START_OF_TIME);
-    runMapreduce();
+    runAction();
     // No new history entries should be generated
     assertThat(getHistoryEntriesOfType(domain, DOMAIN_AUTORENEW)).isEmpty();
     assertBillingEventsForResource(domain, recurring);
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_expandSingleEvent_billingTimeAtCursorTime() throws Exception {
     persistResource(recurring);
     action.cursorTimeParam = Optional.of(DateTime.parse("2000-02-19T00:00:00Z"));
-    runMapreduce();
+    runAction();
     DomainHistory persistedEntry =
         getOnlyHistoryEntryOfType(domain, DOMAIN_AUTORENEW, DomainHistory.class);
     assertHistoryEntryMatches(
         domain, persistedEntry, "TheRegistrar", DateTime.parse("2000-02-19T00:00:00Z"), true);
     BillingEvent.OneTime expected = defaultOneTimeBuilder().setParent(persistedEntry).build();
     assertBillingEventsForResource(domain, expected, recurring);
-    assertCursorAt(beginningOfTest);
+    assertCursorAt(currentTestTime);
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_expandSingleEvent_cursorTimeBetweenEventAndBillingTime() throws Exception {
     persistResource(recurring);
     action.cursorTimeParam = Optional.of(DateTime.parse("2000-01-12T00:00:00Z"));
-    runMapreduce();
+    runAction();
     DomainHistory persistedEntry =
         getOnlyHistoryEntryOfType(domain, DOMAIN_AUTORENEW, DomainHistory.class);
     assertHistoryEntryMatches(
         domain, persistedEntry, "TheRegistrar", DateTime.parse("2000-02-19T00:00:00Z"), true);
     BillingEvent.OneTime expected = defaultOneTimeBuilder().setParent(persistedEntry).build();
     assertBillingEventsForResource(domain, expected, recurring);
-    assertCursorAt(beginningOfTest);
+    assertCursorAt(currentTestTime);
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_expandSingleEvent_billingTimeAtExecutionTime() throws Exception {
-    DateTime testTime = DateTime.parse("2000-02-19T00:00:00Z").minusMillis(1);
+    clock.setTo(currentTestTime);
     persistResource(recurring);
     action.cursorTimeParam = Optional.of(START_OF_TIME);
-    // Clock is advanced one milli in runMapreduce()
-    clock.setTo(testTime);
-    runMapreduce();
+    clock.setTo(DateTime.parse("2000-02-19T00:00:00Z"));
+    runAction();
     // No new history entries should be generated
     assertThat(getHistoryEntriesOfType(domain, DOMAIN_AUTORENEW)).isEmpty();
     // A candidate billing event is set to be billed exactly on 2/19/00 @ 00:00,
     // but these should not be generated as the interval is closed on cursorTime, open on
     // executeTime.
     assertBillingEventsForResource(domain, recurring);
-    assertCursorAt(testTime);
+    assertCursorAt(currentTestTime);
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_expandSingleEvent_multipleYearCreate() throws Exception {
-    DateTime testTime = beginningOfTest.plusYears(2);
     action.cursorTimeParam = Optional.of(recurring.getEventTime());
     recurring =
         persistResource(
             recurring.asBuilder().setEventTime(recurring.getEventTime().plusYears(2)).build());
-    clock.setTo(testTime);
-    runMapreduce();
+    clock.setTo(DateTime.parse("2002-10-02T00:00:00Z"));
+    runAction();
     DomainHistory persistedEntry =
         getOnlyHistoryEntryOfType(domain, DOMAIN_AUTORENEW, DomainHistory.class);
     assertHistoryEntryMatches(
         domain, persistedEntry, "TheRegistrar", DateTime.parse("2002-02-19T00:00:00Z"), true);
-    BillingEvent.OneTime expected = defaultOneTimeBuilder()
-        .setBillingTime(DateTime.parse("2002-02-19T00:00:00Z"))
-        .setEventTime(DateTime.parse("2002-01-05T00:00:00Z"))
-        .setParent(persistedEntry)
-        .setSyntheticCreationTime(testTime)
-        .build();
+    BillingEvent.OneTime expected =
+        defaultOneTimeBuilder()
+            .setBillingTime(DateTime.parse("2002-02-19T00:00:00Z"))
+            .setEventTime(DateTime.parse("2002-01-05T00:00:00Z"))
+            .setParent(persistedEntry)
+            .build();
     assertBillingEventsForResource(domain, expected, recurring);
-    assertCursorAt(testTime);
+    assertCursorAt(currentTestTime);
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_expandSingleEvent_withCursor() throws Exception {
     persistResource(recurring);
     saveCursor(START_OF_TIME);
-    runMapreduce();
+    runAction();
     DomainHistory persistedEntry =
         getOnlyHistoryEntryOfType(domain, DOMAIN_AUTORENEW, DomainHistory.class);
     assertHistoryEntryMatches(
         domain, persistedEntry, "TheRegistrar", DateTime.parse("2000-02-19T00:00:00Z"), true);
     BillingEvent.OneTime expected = defaultOneTimeBuilder().setParent(persistedEntry).build();
     assertBillingEventsForResource(domain, expected, recurring);
-    assertCursorAt(beginningOfTest);
+    assertCursorAt(currentTestTime);
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_expandSingleEvent_withCursorPastExpected() throws Exception {
     persistResource(recurring);
     // Simulate a quick second run of the mapreduce (this should be a no-op).
     saveCursor(clock.nowUtc().minusSeconds(1));
-    runMapreduce();
+    runAction();
     // No new history entries should be generated
     assertThat(getHistoryEntriesOfType(domain, DOMAIN_AUTORENEW)).isEmpty();
     assertBillingEventsForResource(domain, recurring);
-    assertCursorAt(beginningOfTest);
+    assertCursorAt(currentTestTime);
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_expandSingleEvent_recurrenceEndBeforeEvent() throws Exception {
     // This can occur when a domain is transferred or deleted before a domain comes up for renewal.
-    recurring = persistResource(recurring.asBuilder()
-        .setRecurrenceEndTime(recurring.getEventTime().minusDays(5))
-        .build());
+    recurring =
+        persistResource(
+            recurring
+                .asBuilder()
+                .setRecurrenceEndTime(recurring.getEventTime().minusDays(5))
+                .build());
     action.cursorTimeParam = Optional.of(START_OF_TIME);
-    runMapreduce();
+    runAction();
     // No new history entries should be generated
     assertThat(getHistoryEntriesOfType(domain, DOMAIN_AUTORENEW)).isEmpty();
     assertBillingEventsForResource(domain, recurring);
-    assertCursorAt(beginningOfTest);
+    assertCursorAt(currentTestTime);
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_expandSingleEvent_dryRun() throws Exception {
     persistResource(recurring);
     action.isDryRun = true;
     saveCursor(START_OF_TIME); // Need a saved cursor to verify that it didn't move.
-    runMapreduce();
+    runAction();
     // No new history entries should be generated
     assertThat(getHistoryEntriesOfType(domain, DOMAIN_AUTORENEW)).isEmpty();
     assertBillingEventsForResource(domain, recurring);
     assertCursorAt(START_OF_TIME); // Cursor doesn't move on a dry run.
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_expandSingleEvent_multipleYears() throws Exception {
-    DateTime testTime = clock.nowUtc().plusYears(5);
-    clock.setTo(testTime);
+    clock.setTo(clock.nowUtc().plusYears(5));
     List<BillingEvent> expectedEvents = new ArrayList<>();
     expectedEvents.add(persistResource(recurring));
     action.cursorTimeParam = Optional.of(START_OF_TIME);
-    runMapreduce();
+    runAction();
     List<DomainHistory> persistedEntries =
         getHistoryEntriesOfType(domain, DOMAIN_AUTORENEW, DomainHistory.class);
     assertThat(persistedEntries).hasSize(6);
@@ -470,29 +489,25 @@ public class ExpandRecurringBillingEventsActionTest
     // Expecting events for '00, '01, '02, '03, '04, '05.
     for (int year = 0; year < 6; year++) {
       assertHistoryEntryMatches(
-          domain,
-          persistedEntries.get(year),
-          "TheRegistrar",
-          billingDate.plusYears(year), true);
-      expectedEvents.add(defaultOneTimeBuilder()
-          .setBillingTime(billingDate.plusYears(year))
-          .setEventTime(eventDate.plusYears(year))
-          .setParent(persistedEntries.get(year))
-          .setSyntheticCreationTime(testTime)
-          .build());
+          domain, persistedEntries.get(year), "TheRegistrar", billingDate.plusYears(year), true);
+      expectedEvents.add(
+          defaultOneTimeBuilder()
+              .setBillingTime(billingDate.plusYears(year))
+              .setEventTime(eventDate.plusYears(year))
+              .setParent(persistedEntries.get(year))
+              .build());
     }
     assertBillingEventsForResource(domain, Iterables.toArray(expectedEvents, BillingEvent.class));
-    assertCursorAt(testTime);
+    assertCursorAt(currentTestTime);
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_expandSingleEvent_multipleYears_cursorInBetweenYears() throws Exception {
-    DateTime testTime = clock.nowUtc().plusYears(5);
-    clock.setTo(testTime);
+    clock.setTo(clock.nowUtc().plusYears(5));
     List<BillingEvent> expectedEvents = new ArrayList<>();
     expectedEvents.add(persistResource(recurring));
     saveCursor(DateTime.parse("2003-10-02T00:00:00Z"));
-    runMapreduce();
+    runAction();
     List<DomainHistory> persistedEntries =
         getHistoryEntriesOfType(domain, DOMAIN_AUTORENEW, DomainHistory.class);
     assertThat(persistedEntries).hasSize(2);
@@ -502,139 +517,153 @@ public class ExpandRecurringBillingEventsActionTest
     for (int year = 0; year < 2; year++) {
       assertHistoryEntryMatches(
           domain, persistedEntries.get(year), "TheRegistrar", billingDate.plusYears(year), true);
-      expectedEvents.add(defaultOneTimeBuilder()
-          .setBillingTime(billingDate.plusYears(year))
-          .setParent(persistedEntries.get(year))
-          .setEventTime(eventDate.plusYears(year))
-          .setSyntheticCreationTime(testTime)
-          .build());
+      expectedEvents.add(
+          defaultOneTimeBuilder()
+              .setBillingTime(billingDate.plusYears(year))
+              .setParent(persistedEntries.get(year))
+              .setEventTime(eventDate.plusYears(year))
+              .build());
     }
     assertBillingEventsForResource(domain, Iterables.toArray(expectedEvents, BillingEvent.class));
-    assertCursorAt(testTime);
+    assertCursorAt(currentTestTime);
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_singleEvent_beforeRenewal() throws Exception {
-    DateTime testTime = DateTime.parse("2000-01-04T00:00:00Z");
-    clock.setTo(testTime);
+    // Need to restore to the time before the clock was advanced so that the commit log's timestamp
+    // is not inverted when the clock is later reverted.
+    clock.setTo(currentTestTime);
     persistResource(recurring);
+    clock.setTo(DateTime.parse("2000-01-04T00:00:00Z"));
     action.cursorTimeParam = Optional.of(START_OF_TIME);
-    runMapreduce();
+    runAction();
     // No new history entries should be generated
     assertThat(getHistoryEntriesOfType(domain, DOMAIN_AUTORENEW)).isEmpty();
     assertBillingEventsForResource(domain, recurring);
-    assertCursorAt(testTime);
+    assertCursorAt(currentTestTime);
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_singleEvent_afterRecurrenceEnd_inAutorenewGracePeriod() throws Exception {
     // The domain creation date is 1999-01-05, and the first renewal date is thus 2000-01-05.
-    DateTime testTime = DateTime.parse("2001-02-06T00:00:00Z");
-    clock.setTo(testTime);
-    recurring = persistResource(recurring.asBuilder()
-        // The domain deletion date is 2000-01-29, which is within the 45 day autorenew grace period
-        // from the renewal date.
-        .setRecurrenceEndTime(DateTime.parse("2000-01-29T00:00:00Z"))
-        .setEventTime(domain.getCreationTime().plusYears(1))
-        .build());
+    clock.setTo(DateTime.parse("2001-02-06T00:00:00Z"));
+    recurring =
+        persistResource(
+            recurring
+                .asBuilder()
+                // The domain deletion date is 2000-01-29, which is within the 45 day autorenew
+                // grace period
+                // from the renewal date.
+                .setRecurrenceEndTime(DateTime.parse("2000-01-29T00:00:00Z"))
+                .setEventTime(domain.getCreationTime().plusYears(1))
+                .build());
     action.cursorTimeParam = Optional.of(START_OF_TIME);
-    runMapreduce();
+    runAction();
     DomainHistory persistedEntry =
         getOnlyHistoryEntryOfType(domain, DOMAIN_AUTORENEW, DomainHistory.class);
     assertHistoryEntryMatches(
         domain, persistedEntry, "TheRegistrar", DateTime.parse("2000-02-19T00:00:00Z"), false);
-    BillingEvent.OneTime expected = defaultOneTimeBuilder()
-        .setBillingTime(DateTime.parse("2000-02-19T00:00:00Z"))
-        .setParent(persistedEntry)
-        .setSyntheticCreationTime(testTime)
-        .build();
+    BillingEvent.OneTime expected =
+        defaultOneTimeBuilder()
+            .setBillingTime(DateTime.parse("2000-02-19T00:00:00Z"))
+            .setParent(persistedEntry)
+            .build();
     assertBillingEventsForResource(domain, recurring, expected);
-    assertCursorAt(testTime);
+    assertCursorAt(currentTestTime);
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_singleEvent_afterRecurrenceEnd_outsideAutorenewGracePeriod() throws Exception {
     // The domain creation date is 1999-01-05, and the first renewal date is thus 2000-01-05.
-    DateTime testTime = DateTime.parse("2001-02-06T00:00:00Z");
-    clock.setTo(testTime);
-    recurring = persistResource(recurring.asBuilder()
-        // The domain deletion date is 2000-04-05, which is not within the 45 day autorenew grace
-        // period from the renewal date.
-        .setRecurrenceEndTime(DateTime.parse("2000-04-05T00:00:00Z"))
-        .setEventTime(domain.getCreationTime().plusYears(1))
-        .build());
+    clock.setTo(DateTime.parse("2001-02-06T00:00:00Z"));
+    recurring =
+        persistResource(
+            recurring
+                .asBuilder()
+                // The domain deletion date is 2000-04-05, which is not within the 45 day autorenew
+                // grace
+                // period from the renewal date.
+                .setRecurrenceEndTime(DateTime.parse("2000-04-05T00:00:00Z"))
+                .setEventTime(domain.getCreationTime().plusYears(1))
+                .build());
     action.cursorTimeParam = Optional.of(START_OF_TIME);
-    runMapreduce();
+    runAction();
     DomainHistory persistedEntry =
         getOnlyHistoryEntryOfType(domain, DOMAIN_AUTORENEW, DomainHistory.class);
     assertHistoryEntryMatches(
         domain, persistedEntry, "TheRegistrar", DateTime.parse("2000-02-19T00:00:00Z"), true);
-    BillingEvent.OneTime expected = defaultOneTimeBuilder()
-        .setBillingTime(DateTime.parse("2000-02-19T00:00:00Z"))
-        .setParent(persistedEntry)
-        .setSyntheticCreationTime(testTime)
-        .build();
+    BillingEvent.OneTime expected =
+        defaultOneTimeBuilder()
+            .setBillingTime(DateTime.parse("2000-02-19T00:00:00Z"))
+            .setParent(persistedEntry)
+            .build();
     assertBillingEventsForResource(domain, recurring, expected);
-    assertCursorAt(testTime);
+    assertCursorAt(currentTestTime);
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_expandSingleEvent_billingTimeOnLeapYear() throws Exception {
     recurring =
         persistResource(
             recurring.asBuilder().setEventTime(DateTime.parse("2000-01-15T00:00:00Z")).build());
     action.cursorTimeParam = Optional.of(START_OF_TIME);
-    runMapreduce();
+    runAction();
     DomainHistory persistedEntry =
         getOnlyHistoryEntryOfType(domain, DOMAIN_AUTORENEW, DomainHistory.class);
     assertHistoryEntryMatches(
         domain, persistedEntry, "TheRegistrar", DateTime.parse("2000-02-29T00:00:00Z"), true);
-    BillingEvent.OneTime expected = defaultOneTimeBuilder()
-        .setBillingTime(DateTime.parse("2000-02-29T00:00:00Z"))
-        .setEventTime(DateTime.parse("2000-01-15T00:00:00Z"))
-        .setParent(persistedEntry)
-        .build();
+    BillingEvent.OneTime expected =
+        defaultOneTimeBuilder()
+            .setBillingTime(DateTime.parse("2000-02-29T00:00:00Z"))
+            .setEventTime(DateTime.parse("2000-01-15T00:00:00Z"))
+            .setParent(persistedEntry)
+            .build();
     assertBillingEventsForResource(domain, expected, recurring);
-    assertCursorAt(beginningOfTest);
+    assertCursorAt(currentTestTime);
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_expandSingleEvent_billingTimeNotOnLeapYear() throws Exception {
-    DateTime testTime = DateTime.parse("2001-12-01T00:00:00Z");
     recurring =
         persistResource(
             recurring.asBuilder().setEventTime(DateTime.parse("2001-01-15T00:00:00Z")).build());
     action.cursorTimeParam = Optional.of(START_OF_TIME);
-    clock.setTo(testTime);
-    runMapreduce();
+    clock.setTo(DateTime.parse("2001-12-01T00:00:00Z"));
+    runAction();
     DomainHistory persistedEntry =
         getOnlyHistoryEntryOfType(domain, DOMAIN_AUTORENEW, DomainHistory.class);
     assertHistoryEntryMatches(
         domain, persistedEntry, "TheRegistrar", DateTime.parse("2001-03-01T00:00:00Z"), true);
-    BillingEvent.OneTime expected = defaultOneTimeBuilder()
-        .setBillingTime(DateTime.parse("2001-03-01T00:00:00Z"))
-        .setEventTime(DateTime.parse("2001-01-15T00:00:00Z"))
-        .setParent(persistedEntry)
-        .setSyntheticCreationTime(testTime)
-        .build();
+    BillingEvent.OneTime expected =
+        defaultOneTimeBuilder()
+            .setBillingTime(DateTime.parse("2001-03-01T00:00:00Z"))
+            .setEventTime(DateTime.parse("2001-01-15T00:00:00Z"))
+            .setParent(persistedEntry)
+            .build();
     assertBillingEventsForResource(domain, expected, recurring);
-    assertCursorAt(testTime);
+    assertCursorAt(currentTestTime);
   }
 
-  @Test
+  @TestSqlOnly
   void testSuccess_expandMultipleEvents() throws Exception {
     persistResource(recurring);
-    BillingEvent.Recurring recurring2 = persistResource(recurring.asBuilder()
-        .setEventTime(recurring.getEventTime().plusMonths(3))
-        .setId(3L)
-        .build());
+    BillingEvent.Recurring recurring2 =
+        persistResource(
+            recurring
+                .asBuilder()
+                .setEventTime(recurring.getEventTime().plusMonths(3))
+                .setId(3L)
+                .build());
     action.cursorTimeParam = Optional.of(START_OF_TIME);
-    runMapreduce();
+    runAction();
     List<DomainHistory> persistedEntries =
         getHistoryEntriesOfType(domain, DOMAIN_AUTORENEW, DomainHistory.class);
     assertThat(persistedEntries).hasSize(2);
     assertHistoryEntryMatches(
-        domain, persistedEntries.get(0), "TheRegistrar", DateTime.parse("2000-02-19T00:00:00Z"),
+        domain,
+        persistedEntries.get(0),
+        "TheRegistrar",
+        DateTime.parse("2000-02-19T00:00:00Z"),
         true);
     BillingEvent.OneTime expected =
         defaultOneTimeBuilder()
@@ -642,7 +671,10 @@ public class ExpandRecurringBillingEventsActionTest
             .setCancellationMatchingBillingEvent(recurring.createVKey())
             .build();
     assertHistoryEntryMatches(
-        domain, persistedEntries.get(1), "TheRegistrar", DateTime.parse("2000-05-20T00:00:00Z"),
+        domain,
+        persistedEntries.get(1),
+        "TheRegistrar",
+        DateTime.parse("2000-05-20T00:00:00Z"),
         true);
     BillingEvent.OneTime expected2 =
         defaultOneTimeBuilder()
@@ -652,10 +684,10 @@ public class ExpandRecurringBillingEventsActionTest
             .setCancellationMatchingBillingEvent(recurring2.createVKey())
             .build();
     assertBillingEventsForResource(domain, expected, expected2, recurring, recurring2);
-    assertCursorAt(beginningOfTest);
+    assertCursorAt(currentTestTime);
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_premiumDomain() throws Exception {
     persistResource(
         Registry.get("tld")
@@ -664,86 +696,87 @@ public class ExpandRecurringBillingEventsActionTest
             .build());
     persistResource(recurring);
     action.cursorTimeParam = Optional.of(START_OF_TIME);
-    runMapreduce();
+    runAction();
     DomainHistory persistedEntry =
         getOnlyHistoryEntryOfType(domain, DOMAIN_AUTORENEW, DomainHistory.class);
     assertHistoryEntryMatches(
         domain, persistedEntry, "TheRegistrar", DateTime.parse("2000-02-19T00:00:00Z"), true);
-    BillingEvent.OneTime expected = defaultOneTimeBuilder()
-        .setParent(persistedEntry)
-        .setCost(Money.of(USD, 100))
-        .build();
+    BillingEvent.OneTime expected =
+        defaultOneTimeBuilder().setParent(persistedEntry).setCost(Money.of(USD, 100)).build();
     assertBillingEventsForResource(domain, expected, recurring);
-    assertCursorAt(beginningOfTest);
+    assertCursorAt(currentTestTime);
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_varyingRenewPrices() throws Exception {
-    DateTime testTime = beginningOfTest.plusYears(1);
+    clock.setTo(currentTestTime);
     persistResource(
         Registry.get("tld")
             .asBuilder()
             .setRenewBillingCostTransitions(
                 ImmutableSortedMap.of(
-                    START_OF_TIME, Money.of(USD, 8),
-                    DateTime.parse("2000-06-01T00:00:00Z"), Money.of(USD, 10)))
+                    START_OF_TIME,
+                    Money.of(USD, 8),
+                    DateTime.parse("2000-06-01T00:00:00Z"),
+                    Money.of(USD, 10)))
             .build());
-    clock.setTo(testTime);
+    clock.setTo(DateTime.parse("2001-10-02T00:00:00Z"));
     persistResource(recurring);
     action.cursorTimeParam = Optional.of(START_OF_TIME);
-    runMapreduce();
+    runAction();
     List<DomainHistory> persistedEntries =
         getHistoryEntriesOfType(domain, DOMAIN_AUTORENEW, DomainHistory.class);
     assertThat(persistedEntries).hasSize(2);
     DateTime eventDate = DateTime.parse("2000-01-05T00:00:00Z");
     DateTime billingDate = DateTime.parse("2000-02-19T00:00:00Z");
     assertHistoryEntryMatches(domain, persistedEntries.get(0), "TheRegistrar", billingDate, true);
-    BillingEvent.OneTime cheaper = defaultOneTimeBuilder()
-        .setBillingTime(billingDate)
-        .setEventTime(eventDate)
-        .setParent(persistedEntries.get(0))
-        .setCost(Money.of(USD, 8))
-        .setSyntheticCreationTime(testTime)
-        .build();
+    BillingEvent.OneTime cheaper =
+        defaultOneTimeBuilder()
+            .setBillingTime(billingDate)
+            .setEventTime(eventDate)
+            .setParent(persistedEntries.get(0))
+            .setCost(Money.of(USD, 8))
+            .build();
     assertHistoryEntryMatches(
         domain, persistedEntries.get(1), "TheRegistrar", billingDate.plusYears(1), true);
-    BillingEvent.OneTime expensive = cheaper.asBuilder()
-        .setCost(Money.of(USD, 10))
-        .setBillingTime(billingDate.plusYears(1))
-        .setEventTime(eventDate.plusYears(1))
-        .setParent(persistedEntries.get(1))
-        .build();
+    BillingEvent.OneTime expensive =
+        cheaper
+            .asBuilder()
+            .setCost(Money.of(USD, 10))
+            .setBillingTime(billingDate.plusYears(1))
+            .setEventTime(eventDate.plusYears(1))
+            .setParent(persistedEntries.get(1))
+            .build();
     assertBillingEventsForResource(domain, recurring, cheaper, expensive);
-    assertCursorAt(testTime);
+    assertCursorAt(currentTestTime);
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_cursorAfterExecutionTime() {
     action.cursorTimeParam = Optional.of(clock.nowUtc().plusYears(1));
-    IllegalArgumentException thrown =
-        assertThrows(IllegalArgumentException.class, this::runMapreduce);
+    IllegalArgumentException thrown = assertThrows(IllegalArgumentException.class, this::runAction);
     assertThat(thrown)
         .hasMessageThat()
         .contains("Cursor time must be earlier than execution time.");
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_cursorAtExecutionTime() {
     // The clock advances one milli on runMapreduce.
     action.cursorTimeParam = Optional.of(clock.nowUtc().plusMillis(1));
-    IllegalArgumentException thrown =
-        assertThrows(IllegalArgumentException.class, this::runMapreduce);
+    IllegalArgumentException thrown = assertThrows(IllegalArgumentException.class, this::runAction);
     assertThat(thrown)
         .hasMessageThat()
         .contains("Cursor time must be earlier than execution time.");
   }
 
-  @Test
+  @TestOfyOnly
   void testFailure_mapperException_doesNotMoveCursor() throws Exception {
     saveCursor(START_OF_TIME); // Need a saved cursor to verify that it didn't move.
+    clock.advanceOneMilli();
     // Set target to a TLD that doesn't exist.
     recurring = persistResource(recurring.asBuilder().setTargetId("domain.junk").build());
-    runMapreduce();
+    runAction();
     // No new history entries should be generated
     assertThat(getHistoryEntriesOfType(domain, DOMAIN_AUTORENEW)).isEmpty();
     assertBillingEvents(recurring); // only the bogus one in Datastore

core/src/test/java/google/registry/beam/invoicing/BillingEventTest.java

@@ -185,6 +185,19 @@ class BillingEventTest {
                 + "myRegistrar - test,3,RENEW | TLD: test | TERM: 5-year,20.50,USD,");
   }
 
+  @Test
+  void testConvertInvoiceGroupingKey_zeroYears_toCsv() {
+    GenericRecord record = schemaAndRecord.getRecord();
+    record.put("years", 0);
+    schemaAndRecord = new SchemaAndRecord(record, null);
+    BillingEvent event = BillingEvent.parseFromRecord(schemaAndRecord);
+    InvoiceGroupingKey invoiceKey = event.getInvoiceGroupingKey();
+    assertThat(invoiceKey.toCsv(3L))
+        .isEqualTo(
+            "2017-10-01,,12345-CRRHELLO,61.50,USD,10125,1,PURCHASE,"
+                + "myRegistrar - test,3,RENEW | TLD: test | TERM: 0-year,20.50,USD,");
+  }
+
   @Test
   void testInvoiceGroupingKeyCoder_deterministicSerialization() throws IOException {
     InvoiceGroupingKey invoiceKey =

core/src/test/java/google/registry/beam/invoicing/InvoicingPipelineTest.java

@@ -119,7 +119,37 @@ class InvoicingPipelineTest {
               1,
               "USD",
               0,
-              "SUNRISE ANCHOR_TENANT"));
+              "SUNRISE ANCHOR_TENANT"),
+          BillingEvent.create(
+              1,
+              ZonedDateTime.of(2017, 10, 4, 0, 0, 0, 0, ZoneId.of("UTC")),
+              ZonedDateTime.of(2017, 10, 4, 0, 0, 0, 0, ZoneId.of("UTC")),
+              "theRegistrar",
+              "234",
+              "",
+              "test",
+              "SERVER_STATUS",
+              "locked.test",
+              "REPO-ID",
+              0,
+              "USD",
+              0,
+              ""),
+          BillingEvent.create(
+              1,
+              ZonedDateTime.of(2017, 10, 4, 0, 0, 0, 0, ZoneId.of("UTC")),
+              ZonedDateTime.of(2017, 10, 4, 0, 0, 0, 0, ZoneId.of("UTC")),
+              "theRegistrar",
+              "234",
+              "",
+              "test",
+              "SERVER_STATUS",
+              "update-prohibited.test",
+              "REPO-ID",
+              0,
+              "USD",
+              20,
+              ""));
 
   private static final ImmutableMap<String, ImmutableList<String>> EXPECTED_DETAILED_REPORT_MAP =
       ImmutableMap.of(
@@ -128,7 +158,11 @@ class InvoicingPipelineTest {
               "1,2017-10-04 00:00:00 UTC,2017-10-04 00:00:00 UTC,theRegistrar,234,,"
                   + "test,RENEW,mydomain2.test,REPO-ID,3,USD,20.50,",
               "1,2017-10-04 00:00:00 UTC,2017-10-04 00:00:00 UTC,theRegistrar,234,,"
-                  + "test,RENEW,mydomain.test,REPO-ID,3,USD,20.50,"),
+                  + "test,RENEW,mydomain.test,REPO-ID,3,USD,20.50,",
+              "1,2017-10-04 00:00:00 UTC,2017-10-04 00:00:00 UTC,theRegistrar,234,,"
+                  + "test,SERVER_STATUS,update-prohibited.test,REPO-ID,0,USD,20.00,",
+              "1,2017-10-04 00:00:00 UTC,2017-10-04 00:00:00 UTC,theRegistrar,234,,"
+                  + "test,SERVER_STATUS,locked.test,REPO-ID,0,USD,0.00,"),
           "invoice_details_2017-10_theRegistrar_hello.csv",
           ImmutableList.of(
               "1,2017-10-02 00:00:00 UTC,2017-09-29 00:00:00 UTC,theRegistrar,234,,"
@@ -148,6 +182,8 @@ class InvoicingPipelineTest {
               + "RENEW | TLD: test | TERM: 3-year,20.50,USD,",
           "2017-10-01,2022-09-30,234,70.75,JPY,10125,1,PURCHASE,theRegistrar - hello,1,"
               + "CREATE | TLD: hello | TERM: 5-year,70.75,JPY,",
+          "2017-10-01,,234,20.00,USD,10125,1,PURCHASE,theRegistrar - test,1,"
+              + "SERVER_STATUS | TLD: test | TERM: 0-year,20.00,USD,",
           "2017-10-01,2018-09-30,456,20.50,USD,10125,1,PURCHASE,bestdomains - test,1,"
               + "RENEW | TLD: test | TERM: 1-year,20.50,USD,116688");
 

core/src/test/java/google/registry/keyring/kms/FakeKmsConnection.java

@@ -1,45 +0,0 @@
-// Copyright 2017 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.keyring.kms;
-
-import com.google.common.io.BaseEncoding;
-import org.bouncycastle.util.Arrays;
-
-class FakeKmsConnection implements KmsConnection {
-
-  FakeKmsConnection() {}
-
-  /**
-   * Returns a dummy {@link EncryptResponse}.
-   *
-   * <p>The "encrypted value" in the response is the provided value reversed and then base64-encoded
-   * and the name of the cryptoKeyVersion is {@code cryptoKeyName + "/foo"}.
-   */
-  @Override
-  public EncryptResponse encrypt(String cryptoKeyName, byte[] plaintext) {
-    return EncryptResponse.create(
-        BaseEncoding.base64().encode(Arrays.reverse(plaintext)), cryptoKeyName + "/foo");
-  }
-
-  /**
-   * Returns a "decrypted" plaintext.
-   *
-   * <p>The plaintext is the encodedCiphertext base64-decoded and then reversed.
-   */
-  @Override
-  public byte[] decrypt(String cryptoKeyName, String encodedCiphertext) {
-    return Arrays.reverse(BaseEncoding.base64().decode(encodedCiphertext));
-  }
-}

core/src/test/java/google/registry/keyring/kms/KmsKeyringTest.java

@@ -1,187 +0,0 @@
-// Copyright 2017 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.keyring.kms;
-
-import static com.google.common.truth.Truth.assertThat;
-import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
-
-import google.registry.keyring.api.KeySerializer;
-import google.registry.model.server.KmsSecret;
-import google.registry.model.server.KmsSecretRevision;
-import google.registry.privileges.secretmanager.FakeSecretManagerClient;
-import google.registry.privileges.secretmanager.KeyringSecretStore;
-import google.registry.testing.AppEngineExtension;
-import google.registry.testing.BouncyCastleProviderExtension;
-import google.registry.testing.DualDatabaseTest;
-import google.registry.testing.TestOfyAndSql;
-import org.bouncycastle.openpgp.PGPKeyPair;
-import org.bouncycastle.openpgp.PGPPrivateKey;
-import org.bouncycastle.openpgp.PGPPublicKey;
-import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.extension.RegisterExtension;
-
-/** Unit tests for {@link KmsKeyring}. */
-@DualDatabaseTest
-class KmsKeyringTest {
-
-  @RegisterExtension
-  final BouncyCastleProviderExtension bouncy = new BouncyCastleProviderExtension();
-
-  @RegisterExtension
-  final AppEngineExtension appEngine =
-      AppEngineExtension.builder().withDatastoreAndCloudSql().build();
-
-  private KmsKeyring keyring;
-  private KeyringSecretStore fakeSecretStore =
-      new KeyringSecretStore(new FakeSecretManagerClient());
-
-  @BeforeEach
-  void beforeEach() {
-    keyring = new KmsKeyring(new FakeKmsConnection(), fakeSecretStore);
-  }
-
-  @TestOfyAndSql
-  void test_getRdeSigningKey() throws Exception {
-    saveKeyPairSecret("rde-signing-public", "rde-signing-private");
-    PGPKeyPair rdeSigningKey = keyring.getRdeSigningKey();
-    assertThat(KeySerializer.serializeKeyPair(rdeSigningKey))
-        .isEqualTo(KeySerializer.serializeKeyPair(KmsTestHelper.getKeyPair()));
-  }
-
-  @TestOfyAndSql
-  void test_getRdeStagingEncryptionKey() throws Exception {
-    savePublicKeySecret("rde-staging-public");
-    PGPPublicKey rdeStagingEncryptionKey = keyring.getRdeStagingEncryptionKey();
-    assertThat(rdeStagingEncryptionKey.getFingerprint())
-        .isEqualTo(KmsTestHelper.getPublicKey().getFingerprint());
-  }
-
-  @TestOfyAndSql
-  void test_getRdeStagingDecryptionKey() throws Exception {
-    savePrivateKeySecret("rde-staging-private");
-    savePublicKeySecret("rde-staging-public");
-
-    PGPPrivateKey rdeStagingDecryptionKey = keyring.getRdeStagingDecryptionKey();
-    PGPPublicKey rdeStagingEncryptionKey = keyring.getRdeStagingEncryptionKey();
-    PGPKeyPair keyPair = new PGPKeyPair(rdeStagingEncryptionKey, rdeStagingDecryptionKey);
-
-    assertThat(KeySerializer.serializeKeyPair(keyPair))
-        .isEqualTo(KeySerializer.serializeKeyPair(KmsTestHelper.getKeyPair()));
-  }
-
-  @TestOfyAndSql
-  void test_getRdeReceiverKey() throws Exception {
-    savePublicKeySecret("rde-receiver-public");
-    PGPPublicKey rdeReceiverKey = keyring.getRdeReceiverKey();
-    assertThat(rdeReceiverKey.getFingerprint())
-        .isEqualTo(KmsTestHelper.getPublicKey().getFingerprint());
-  }
-
-  @TestOfyAndSql
-  void test_getBrdaSigningKey() throws Exception {
-    saveKeyPairSecret("brda-signing-public", "brda-signing-private");
-    PGPKeyPair brdaSigningKey = keyring.getBrdaSigningKey();
-    assertThat(KeySerializer.serializeKeyPair(brdaSigningKey))
-        .isEqualTo(KeySerializer.serializeKeyPair(KmsTestHelper.getKeyPair()));
-  }
-
-  @TestOfyAndSql
-  void test_getBrdaReceiverKey() throws Exception {
-    savePublicKeySecret("brda-receiver-public");
-    PGPPublicKey brdaReceiverKey = keyring.getBrdaReceiverKey();
-    assertThat(brdaReceiverKey.getFingerprint())
-        .isEqualTo(KmsTestHelper.getPublicKey().getFingerprint());
-  }
-
-  @TestOfyAndSql
-  void test_getRdeSshClientPublicKey() {
-    saveCleartextSecret("rde-ssh-client-public-string");
-    String rdeSshClientPublicKey = keyring.getRdeSshClientPublicKey();
-    assertThat(rdeSshClientPublicKey).isEqualTo("rde-ssh-client-public-stringmoo");
-  }
-
-  @TestOfyAndSql
-  void test_getRdeSshClientPrivateKey() {
-    saveCleartextSecret("rde-ssh-client-private-string");
-    String rdeSshClientPrivateKey = keyring.getRdeSshClientPrivateKey();
-    assertThat(rdeSshClientPrivateKey).isEqualTo("rde-ssh-client-private-stringmoo");
-  }
-
-  @TestOfyAndSql
-  void test_getIcannReportingPassword() {
-    saveCleartextSecret("icann-reporting-password-string");
-    String icannReportingPassword = keyring.getIcannReportingPassword();
-    assertThat(icannReportingPassword).isEqualTo("icann-reporting-password-stringmoo");
-  }
-
-  @TestOfyAndSql
-  void test_getMarksdbDnlLoginAndPassword() {
-    saveCleartextSecret("marksdb-dnl-login-string");
-    String marksdbDnlLoginAndPassword = keyring.getMarksdbDnlLoginAndPassword();
-    assertThat(marksdbDnlLoginAndPassword).isEqualTo("marksdb-dnl-login-stringmoo");
-  }
-
-  @TestOfyAndSql
-  void test_getMarksdbLordnPassword() {
-    saveCleartextSecret("marksdb-lordn-password-string");
-    String marksdbLordnPassword = keyring.getMarksdbLordnPassword();
-    assertThat(marksdbLordnPassword).isEqualTo("marksdb-lordn-password-stringmoo");
-  }
-
-  @TestOfyAndSql
-  void test_getMarksdbSmdrlLoginAndPassword() {
-    saveCleartextSecret("marksdb-smdrl-login-string");
-    String marksdbSmdrlLoginAndPassword = keyring.getMarksdbSmdrlLoginAndPassword();
-    assertThat(marksdbSmdrlLoginAndPassword).isEqualTo("marksdb-smdrl-login-stringmoo");
-  }
-
-  @TestOfyAndSql
-  void test_getJsonCredential() {
-    saveCleartextSecret("json-credential-string");
-    String jsonCredential = keyring.getJsonCredential();
-    assertThat(jsonCredential).isEqualTo("json-credential-stringmoo");
-  }
-
-  private void persistSecret(String secretName, byte[] secretValue) {
-    KmsConnection kmsConnection = new FakeKmsConnection();
-
-    KmsSecretRevision secretRevision =
-        new KmsSecretRevision.Builder()
-            .setEncryptedValue(kmsConnection.encrypt(secretName, secretValue).ciphertext())
-            .setKmsCryptoKeyVersionName(KmsTestHelper.DUMMY_CRYPTO_KEY_VERSION)
-            .setParent(secretName)
-            .build();
-    KmsSecret secret = KmsSecret.create(secretName, secretRevision);
-    tm().transact(() -> tm().putAll(secretRevision, secret));
-    fakeSecretStore.createOrUpdateSecret(secretName, secretValue);
-  }
-
-  private void saveCleartextSecret(String secretName) {
-    persistSecret(secretName, KeySerializer.serializeString(secretName + "moo"));
-  }
-
-  private void savePublicKeySecret(String publicKeyName) throws Exception {
-    persistSecret(publicKeyName, KeySerializer.serializePublicKey(KmsTestHelper.getPublicKey()));
-  }
-
-  private void savePrivateKeySecret(String privateKeyName) throws Exception {
-    persistSecret(privateKeyName, KeySerializer.serializeKeyPair(KmsTestHelper.getKeyPair()));
-  }
-
-  private void saveKeyPairSecret(String publicKeyName, String privateKeyName) throws Exception {
-    savePublicKeySecret(publicKeyName);
-    savePrivateKeySecret(privateKeyName);
-  }
-}

core/src/test/java/google/registry/keyring/kms/KmsTestHelper.java

@@ -26,7 +26,7 @@ import org.bouncycastle.openpgp.bc.BcPGPSecretKeyRing;
 import org.bouncycastle.openpgp.operator.bc.BcPBESecretKeyDecryptorBuilder;
 import org.bouncycastle.openpgp.operator.bc.BcPGPDigestCalculatorProvider;
 
-/** Stores dummy values for test use in {@link KmsUpdaterTest} and {@link KmsKeyringTest}. */
+/** Stores dummy values for test use in {@link KmsUpdaterTest}. */
 final class KmsTestHelper {
 
   static final String DUMMY_CRYPTO_KEY_VERSION = "cheeseburger";

core/src/test/java/google/registry/keyring/kms/KmsUpdaterTest.java

@@ -15,226 +15,185 @@
 package google.registry.keyring.kms;
 
 import static com.google.common.truth.Truth.assertThat;
-import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
-import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
-import com.googlecode.objectify.Key;
 import google.registry.keyring.api.KeySerializer;
-import google.registry.model.server.KmsSecret;
-import google.registry.model.server.KmsSecretRevision;
-import google.registry.model.server.KmsSecretRevisionSqlDao;
-import google.registry.persistence.VKey;
 import google.registry.privileges.secretmanager.FakeSecretManagerClient;
 import google.registry.privileges.secretmanager.KeyringSecretStore;
-import google.registry.testing.AppEngineExtension;
 import google.registry.testing.BouncyCastleProviderExtension;
-import google.registry.testing.DualDatabaseTest;
-import google.registry.testing.TestOfyAndSql;
 import java.io.IOException;
+import java.nio.charset.StandardCharsets;
 import org.bouncycastle.openpgp.PGPKeyPair;
 import org.bouncycastle.openpgp.PGPPublicKey;
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
-/** Unit tests for {@link KmsUpdater} */
-@DualDatabaseTest
+/** Unit tests for {@link KmsKeyring} and {@link KmsUpdater} */
+// TODO(2021-07-01): Rename this class along with KmsKeyring
 public class KmsUpdaterTest {
 
-  @RegisterExtension
-  public final AppEngineExtension appEngine =
-      AppEngineExtension.builder().withDatastoreAndCloudSql().build();
-
   @RegisterExtension
   public final BouncyCastleProviderExtension bouncy = new BouncyCastleProviderExtension();
 
+  private KeyringSecretStore secretStore;
   private KmsUpdater updater;
+  private KmsKeyring keyring;
 
   @BeforeEach
   void beforeEach() {
-    updater =
-        new KmsUpdater(
-            new FakeKmsConnection(), new KeyringSecretStore(new FakeSecretManagerClient()));
+    secretStore = new KeyringSecretStore(new FakeSecretManagerClient());
+    updater = new KmsUpdater(secretStore);
+    keyring = new KmsKeyring(secretStore);
   }
 
-  @TestOfyAndSql
-  void test_setMultipleSecrets() {
+  @Test
+  void setAndReadMultiple() {
+    String secretPrefix = "setAndReadMultiple_";
     updater
-        .setMarksdbDnlLoginAndPassword("value1")
-        .setIcannReportingPassword("value2")
-        .setJsonCredential("value3")
+        .setMarksdbDnlLoginAndPassword(secretPrefix + "marksdb")
+        .setIcannReportingPassword(secretPrefix + "icann")
+        .setJsonCredential(secretPrefix + "json")
         .update();
 
-    verifySecretAndSecretRevisionWritten(
-        "marksdb-dnl-login-string",
-        "marksdb-dnl-login-string/foo",
-        getCiphertext("value1"));
-    verifySecretAndSecretRevisionWritten(
-        "icann-reporting-password-string",
-        "icann-reporting-password-string/foo",
-        getCiphertext("value2"));
-    verifySecretAndSecretRevisionWritten(
-        "json-credential-string", "json-credential-string/foo", getCiphertext("value3"));
+    assertThat(keyring.getMarksdbDnlLoginAndPassword()).isEqualTo(secretPrefix + "marksdb");
+    assertThat(keyring.getIcannReportingPassword()).isEqualTo(secretPrefix + "icann");
+    assertThat(keyring.getJsonCredential()).isEqualTo(secretPrefix + "json");
+
+    verifyPersistedSecret("marksdb-dnl-login-string", secretPrefix + "marksdb");
+    verifyPersistedSecret("icann-reporting-password-string", secretPrefix + "icann");
+    verifyPersistedSecret("json-credential-string", secretPrefix + "json");
   }
 
-  @TestOfyAndSql
-  void test_setBrdaReceiverKey() throws Exception {
-    updater.setBrdaReceiverPublicKey(KmsTestHelper.getPublicKey()).update();
+  @Test
+  void brdaReceiverKey() throws Exception {
+    PGPPublicKey publicKey = KmsTestHelper.getPublicKey();
+    updater.setBrdaReceiverPublicKey(publicKey).update();
 
-    verifySecretAndSecretRevisionWritten(
-        "brda-receiver-public",
-        "brda-receiver-public/foo",
-        getCiphertext(KmsTestHelper.getPublicKey()));
+    assertThat(keyring.getBrdaReceiverKey().getFingerprint()).isEqualTo(publicKey.getFingerprint());
+    verifyPersistedSecret("brda-receiver-public", serializePublicKey(publicKey));
   }
 
-  @TestOfyAndSql
-  void test_setBrdaSigningKey() throws Exception {
-    updater.setBrdaSigningKey(KmsTestHelper.getKeyPair()).update();
+  @Test
+  void brdaSigningKey() throws Exception {
+    PGPKeyPair keyPair = KmsTestHelper.getKeyPair();
+    updater.setBrdaSigningKey(keyPair).update();
 
-    verifySecretAndSecretRevisionWritten(
-        "brda-signing-private",
-        "brda-signing-private/foo",
-        getCiphertext(KmsTestHelper.getKeyPair()));
-    verifySecretAndSecretRevisionWritten(
-        "brda-signing-public",
-        "brda-signing-public/foo",
-        getCiphertext(KmsTestHelper.getPublicKey()));
+    assertThat(serializeKeyPair(keyring.getBrdaSigningKey())).isEqualTo(serializeKeyPair(keyPair));
+    verifyPersistedSecret("brda-signing-private", serializeKeyPair(KmsTestHelper.getKeyPair()));
+    verifyPersistedSecret("brda-signing-public", serializePublicKey(KmsTestHelper.getPublicKey()));
   }
 
-  @TestOfyAndSql
-  void test_setIcannReportingPassword() {
-    updater.setIcannReportingPassword("value1").update();
+  @Test
+  void icannReportingPassword() {
+    String secret = "icannReportingPassword";
+    updater.setIcannReportingPassword(secret).update();
 
-    verifySecretAndSecretRevisionWritten(
-        "icann-reporting-password-string",
-        "icann-reporting-password-string/foo",
-        getCiphertext("value1"));
+    assertThat(keyring.getIcannReportingPassword()).isEqualTo(secret);
+    verifyPersistedSecret("icann-reporting-password-string", secret);
   }
 
-  @TestOfyAndSql
-  void test_setJsonCredential() {
-    updater.setJsonCredential("value1").update();
+  @Test
+  void jsonCredential() {
+    String secret = "jsonCredential";
+    updater.setJsonCredential(secret).update();
 
-    verifySecretAndSecretRevisionWritten(
-        "json-credential-string", "json-credential-string/foo", getCiphertext("value1"));
+    assertThat(keyring.getJsonCredential()).isEqualTo(secret);
+    verifyPersistedSecret("json-credential-string", secret);
   }
 
-  @TestOfyAndSql
-  void test_setMarksdbDnlLoginAndPassword() {
-    updater.setMarksdbDnlLoginAndPassword("value1").update();
+  @Test
+  void marksdbDnlLoginAndPassword() {
+    String secret = "marksdbDnlLoginAndPassword";
+    updater.setMarksdbDnlLoginAndPassword(secret).update();
 
-    verifySecretAndSecretRevisionWritten(
-        "marksdb-dnl-login-string", "marksdb-dnl-login-string/foo", getCiphertext("value1"));
+    assertThat(keyring.getMarksdbDnlLoginAndPassword()).isEqualTo(secret);
+    verifyPersistedSecret("marksdb-dnl-login-string", secret);
   }
 
-  @TestOfyAndSql
-  void test_setMarksdbLordnPassword() {
-    updater.setMarksdbLordnPassword("value1").update();
+  @Test
+  void marksdbLordnPassword() {
+    String secret = "marksdbLordnPassword";
+    updater.setMarksdbLordnPassword(secret).update();
 
-    verifySecretAndSecretRevisionWritten(
-        "marksdb-lordn-password-string",
-        "marksdb-lordn-password-string/foo",
-        getCiphertext("value1"));
+    assertThat(keyring.getMarksdbLordnPassword()).isEqualTo(secret);
+    verifyPersistedSecret("marksdb-lordn-password-string", secret);
   }
 
-  @TestOfyAndSql
-  void test_setMarksdbSmdrlLoginAndPassword() {
-    updater.setMarksdbSmdrlLoginAndPassword("value1").update();
+  @Test
+  void marksdbSmdrlLoginAndPassword() {
+    String secret = "marksdbSmdrlLoginAndPassword";
+    updater.setMarksdbSmdrlLoginAndPassword(secret).update();
 
-    verifySecretAndSecretRevisionWritten(
-        "marksdb-smdrl-login-string", "marksdb-smdrl-login-string/foo", getCiphertext("value1"));
+    assertThat(keyring.getMarksdbSmdrlLoginAndPassword()).isEqualTo(secret);
+    verifyPersistedSecret("marksdb-smdrl-login-string", secret);
   }
 
-  @TestOfyAndSql
-  void test_setRdeReceiverKey() throws Exception {
-    updater.setRdeReceiverPublicKey(KmsTestHelper.getPublicKey()).update();
+  @Test
+  void rdeReceiverKey() throws Exception {
+    PGPPublicKey publicKey = KmsTestHelper.getPublicKey();
+    updater.setRdeReceiverPublicKey(publicKey).update();
 
-    verifySecretAndSecretRevisionWritten(
-        "rde-receiver-public",
-        "rde-receiver-public/foo",
-        getCiphertext(
-            KeySerializer.serializePublicKey(KmsTestHelper.getPublicKey())));
+    assertThat(keyring.getRdeReceiverKey().getFingerprint()).isEqualTo(publicKey.getFingerprint());
+    verifyPersistedSecret("rde-receiver-public", serializePublicKey(KmsTestHelper.getPublicKey()));
   }
 
-  @TestOfyAndSql
-  void test_setRdeSigningKey() throws Exception {
-    updater.setRdeSigningKey(KmsTestHelper.getKeyPair()).update();
+  @Test
+  void rdeSigningKey() throws Exception {
+    PGPKeyPair keyPair = KmsTestHelper.getKeyPair();
+    updater.setRdeSigningKey(keyPair).update();
 
-    verifySecretAndSecretRevisionWritten(
-        "rde-signing-private",
-        "rde-signing-private/foo",
-        getCiphertext(KmsTestHelper.getKeyPair()));
-    verifySecretAndSecretRevisionWritten(
-        "rde-signing-public",
-        "rde-signing-public/foo",
-        getCiphertext(KmsTestHelper.getPublicKey()));
+    assertThat(serializeKeyPair(keyring.getRdeSigningKey())).isEqualTo(serializeKeyPair(keyPair));
+
+    verifyPersistedSecret("rde-signing-private", serializeKeyPair(keyPair));
+    verifyPersistedSecret("rde-signing-public", serializePublicKey(keyPair.getPublicKey()));
   }
 
-  @TestOfyAndSql
-  void test_setRdeSshClientPrivateKey() {
-    updater.setRdeSshClientPrivateKey("value1").update();
+  @Test
+  void rdeSshClientPrivateKey() {
+    String secret = "rdeSshClientPrivateKey";
+    updater.setRdeSshClientPrivateKey(secret).update();
 
-    verifySecretAndSecretRevisionWritten(
-        "rde-ssh-client-private-string",
-        "rde-ssh-client-private-string/foo",
-        getCiphertext("value1"));
+    assertThat(keyring.getRdeSshClientPrivateKey()).isEqualTo(secret);
+    verifyPersistedSecret("rde-ssh-client-private-string", secret);
   }
 
-  @TestOfyAndSql
-  void test_setRdeSshClientPublicKey() {
-    updater.setRdeSshClientPublicKey("value1").update();
+  @Test
+  void rdeSshClientPublicKey() {
+    String secret = "rdeSshClientPublicKey";
+    updater.setRdeSshClientPublicKey(secret).update();
 
-    verifySecretAndSecretRevisionWritten(
-        "rde-ssh-client-public-string",
-        "rde-ssh-client-public-string/foo",
-        getCiphertext("value1"));
+    assertThat(keyring.getRdeSshClientPublicKey()).isEqualTo(secret);
+    verifyPersistedSecret("rde-ssh-client-public-string", secret);
   }
 
-  @TestOfyAndSql
-  void test_setRdeStagingKey() throws Exception {
-    updater.setRdeStagingKey(KmsTestHelper.getKeyPair()).update();
+  @Test
+  void rdeStagingKey() throws Exception {
+    PGPKeyPair keyPair = KmsTestHelper.getKeyPair();
+    updater.setRdeStagingKey(keyPair).update();
 
-    verifySecretAndSecretRevisionWritten(
-        "rde-staging-private",
-        "rde-staging-private/foo",
-        getCiphertext(KmsTestHelper.getKeyPair()));
-    verifySecretAndSecretRevisionWritten(
-        "rde-staging-public",
-        "rde-staging-public/foo",
-        getCiphertext(KmsTestHelper.getPublicKey()));
+    assertThat(serializePublicKey(keyring.getRdeStagingEncryptionKey()))
+        .isEqualTo(serializePublicKey(keyPair.getPublicKey()));
+    // Since we do not have dedicated tools to compare private keys, we leverage key-pair
+    // serialization util to compare private keys.
+    assertThat(
+            serializeKeyPair(
+                new PGPKeyPair(
+                    keyring.getRdeStagingEncryptionKey(), keyring.getRdeStagingDecryptionKey())))
+        .isEqualTo(serializeKeyPair(keyPair));
+    verifyPersistedSecret("rde-staging-private", serializeKeyPair(keyPair));
+    verifyPersistedSecret("rde-staging-public", serializePublicKey(KmsTestHelper.getPublicKey()));
   }
 
-  private static void verifySecretAndSecretRevisionWritten(
-      String secretName, String expectedCryptoKeyVersionName, String expectedEncryptedValue) {
-    KmsSecretRevision secretRevision;
-    if (tm().isOfy()) {
-      KmsSecret secret =
-          tm().loadByKey(
-                  VKey.createOfy(
-                      KmsSecret.class, Key.create(getCrossTldKey(), KmsSecret.class, secretName)));
-      assertThat(secret).isNotNull();
-      secretRevision =
-          tm().loadByKey(VKey.createOfy(KmsSecretRevision.class, secret.getLatestRevision()));
-    } else {
-      secretRevision =
-          tm().transact(() -> KmsSecretRevisionSqlDao.getLatestRevision(secretName).get());
-    }
-    assertThat(secretRevision.getKmsCryptoKeyVersionName()).isEqualTo(expectedCryptoKeyVersionName);
-    assertThat(secretRevision.getEncryptedValue()).isEqualTo(expectedEncryptedValue);
+  private void verifyPersistedSecret(String secretName, String expectedPlainTextValue) {
+    assertThat(new String(secretStore.getSecret(secretName), StandardCharsets.UTF_8))
+        .isEqualTo(expectedPlainTextValue);
   }
 
-  private static String getCiphertext(byte[] plaintext) {
-    return new FakeKmsConnection().encrypt("blah", plaintext).ciphertext();
+  private static String serializePublicKey(PGPPublicKey publicKey) throws IOException {
+    return new String(KeySerializer.serializePublicKey(publicKey), StandardCharsets.UTF_8);
   }
 
-  private static String getCiphertext(String plaintext) {
-    return getCiphertext(KeySerializer.serializeString(plaintext));
-  }
-
-  private static String getCiphertext(PGPPublicKey publicKey) throws IOException {
-    return getCiphertext(KeySerializer.serializePublicKey(publicKey));
-  }
-
-  private static String getCiphertext(PGPKeyPair keyPair) throws Exception {
-    return getCiphertext(KeySerializer.serializeKeyPair(keyPair));
+  private static String serializeKeyPair(PGPKeyPair keyPair) throws Exception {
+    return new String(KeySerializer.serializeKeyPair(keyPair), StandardCharsets.UTF_8);
   }
 }

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

@@ -0,0 +1,173 @@
+// 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.DatabaseMigrationStateSchedule.MigrationState.DATASTORE_ONLY;
+import static google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState.DATASTORE_PRIMARY;
+import static google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState.DATASTORE_PRIMARY_READ_ONLY;
+import static google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState.SQL_ONLY;
+import static google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState.SQL_PRIMARY;
+import static google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState.SQL_PRIMARY_READ_ONLY;
+import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
+import static google.registry.util.DateTimeUtils.START_OF_TIME;
+import static org.junit.Assert.assertThrows;
+
+import com.google.common.collect.ImmutableSortedMap;
+import google.registry.model.EntityTestCase;
+import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState;
+import org.joda.time.DateTime;
+import org.joda.time.Duration;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+public class DatabaseMigrationStateScheduleTest extends EntityTestCase {
+
+  @BeforeEach
+  void beforeEach() {
+    fakeClock.setAutoIncrementByOneMilli();
+  }
+
+  @Test
+  void testEmpty_returnsDatastoreOnlyMap() {
+    assertThat(DatabaseMigrationStateSchedule.getUncached())
+        .isEqualTo(DatabaseMigrationStateSchedule.DEFAULT_TRANSITION_MAP);
+  }
+
+  @Test
+  void testValidTransitions() {
+    // First, verify that no-ops are safe
+    for (MigrationState migrationState : MigrationState.values()) {
+      runValidTransition(migrationState, migrationState);
+    }
+
+    // Next, the transitions that will actually cause a change
+    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_READ_ONLY);
+    runValidTransition(DATASTORE_PRIMARY_READ_ONLY, SQL_PRIMARY);
+
+    runValidTransition(SQL_PRIMARY_READ_ONLY, DATASTORE_PRIMARY_READ_ONLY);
+    runValidTransition(SQL_PRIMARY_READ_ONLY, SQL_PRIMARY);
+
+    runValidTransition(SQL_PRIMARY, SQL_PRIMARY_READ_ONLY);
+    runValidTransition(SQL_PRIMARY, SQL_ONLY);
+
+    runValidTransition(SQL_ONLY, SQL_PRIMARY);
+  }
+
+  @Test
+  void testInvalidTransitions() {
+    runInvalidTransition(DATASTORE_ONLY, DATASTORE_PRIMARY_READ_ONLY);
+    runInvalidTransition(DATASTORE_ONLY, SQL_PRIMARY_READ_ONLY);
+    runInvalidTransition(DATASTORE_ONLY, SQL_PRIMARY);
+    runInvalidTransition(DATASTORE_ONLY, SQL_ONLY);
+
+    runInvalidTransition(DATASTORE_PRIMARY, SQL_PRIMARY_READ_ONLY);
+    runInvalidTransition(DATASTORE_PRIMARY, SQL_PRIMARY);
+    runInvalidTransition(DATASTORE_PRIMARY, SQL_ONLY);
+
+    runInvalidTransition(DATASTORE_PRIMARY_READ_ONLY, SQL_ONLY);
+
+    runInvalidTransition(SQL_PRIMARY_READ_ONLY, DATASTORE_ONLY);
+    runInvalidTransition(SQL_PRIMARY_READ_ONLY, DATASTORE_PRIMARY);
+    runInvalidTransition(SQL_PRIMARY_READ_ONLY, SQL_ONLY);
+
+    runInvalidTransition(SQL_PRIMARY, DATASTORE_ONLY);
+    runInvalidTransition(SQL_PRIMARY, DATASTORE_PRIMARY);
+    runInvalidTransition(SQL_PRIMARY, DATASTORE_PRIMARY_READ_ONLY);
+
+    runInvalidTransition(SQL_ONLY, DATASTORE_ONLY);
+    runInvalidTransition(SQL_ONLY, DATASTORE_PRIMARY);
+    runInvalidTransition(SQL_ONLY, DATASTORE_PRIMARY_READ_ONLY);
+  }
+
+  @Test
+  void testFailure_newMapImpliesInvalidChangeNow() {
+    DateTime startTime = fakeClock.nowUtc();
+    fakeClock.advanceBy(Duration.standardHours(6));
+
+    // The new map is valid by itself, but not with the current state of DATASTORE_ONLY because the
+    // new map implies that the current state is DATASTORE_PRIMARY_READ_ONLY
+    ImmutableSortedMap<DateTime, MigrationState> nowInvalidMap =
+        ImmutableSortedMap.<DateTime, MigrationState>naturalOrder()
+            .put(START_OF_TIME, DATASTORE_ONLY)
+            .put(startTime.plusHours(1), DATASTORE_PRIMARY)
+            .put(startTime.plusHours(2), DATASTORE_PRIMARY_READ_ONLY)
+            .build();
+    assertThat(
+            assertThrows(
+                IllegalArgumentException.class,
+                () -> ofyTm().transact(() -> DatabaseMigrationStateSchedule.set(nowInvalidMap))))
+        .hasMessageThat()
+        .isEqualTo(
+            "Cannot transition from current state-as-of-now DATASTORE_ONLY "
+                + "to new state-as-of-now DATASTORE_PRIMARY_READ_ONLY");
+  }
+
+  @Test
+  void testFailure_notInTransaction() {
+    assertThat(
+            assertThrows(
+                IllegalStateException.class,
+                () ->
+                    DatabaseMigrationStateSchedule.set(
+                        DatabaseMigrationStateSchedule.DEFAULT_TRANSITION_MAP.toValueMap())))
+        .hasMessageThat()
+        .isEqualTo("Must be called in a transaction");
+  }
+
+  private void runValidTransition(MigrationState from, MigrationState to) {
+    ImmutableSortedMap<DateTime, MigrationState> transitions =
+        createMapEndingWithTransition(from, to);
+    ofyTm().transact(() -> DatabaseMigrationStateSchedule.set(transitions));
+    assertThat(DatabaseMigrationStateSchedule.getUncached().toValueMap())
+        .containsExactlyEntriesIn(transitions);
+  }
+
+  private void runInvalidTransition(MigrationState from, MigrationState to) {
+    ImmutableSortedMap<DateTime, MigrationState> transitions =
+        createMapEndingWithTransition(from, to);
+    assertThat(
+            assertThrows(
+                IllegalArgumentException.class,
+                () -> ofyTm().transact(() -> DatabaseMigrationStateSchedule.set(transitions))))
+        .hasMessageThat()
+        .isEqualTo(
+            String.format("validStateTransitions map cannot transition from %s to %s.", from, to));
+  }
+
+  // Create a transition map that is valid up to the "from" transition, then add the "to" transition
+  private ImmutableSortedMap<DateTime, MigrationState> createMapEndingWithTransition(
+      MigrationState from, MigrationState to) {
+    ImmutableSortedMap.Builder<DateTime, MigrationState> builder =
+        ImmutableSortedMap.naturalOrder();
+    builder.put(START_OF_TIME, DATASTORE_ONLY);
+    MigrationState[] allMigrationStates = MigrationState.values();
+    for (int i = 0; i < allMigrationStates.length; i++) {
+      builder.put(fakeClock.nowUtc().plusMinutes(i), allMigrationStates[i]);
+      if (allMigrationStates[i].equals(from)) {
+        break;
+      }
+    }
+    builder.put(fakeClock.nowUtc().plusDays(1), to);
+    return builder.build();
+  }
+}

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

@@ -1,111 +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.
-
-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/java/google/registry/model/common/DatabaseTransitionScheduleTest.java

@@ -1,79 +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.
-
-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/persistence/transaction/CriteriaQueryBuilderTest.java

@@ -59,8 +59,7 @@ class CriteriaQueryBuilderTest {
                 .transact(
                     () ->
                         jpaTm()
-                            .getEntityManager()
-                            .createQuery(
+                            .query(
                                 CriteriaQueryBuilder.create(CriteriaQueryBuilderTestEntity.class)
                                     .build())
                             .getResultList()))
@@ -77,10 +76,11 @@ class CriteriaQueryBuilderTest {
                   CriteriaQuery<CriteriaQueryBuilderTestEntity> query =
                       CriteriaQueryBuilder.create(CriteriaQueryBuilderTestEntity.class)
                           .where(
-                              "data", jpaTm().getEntityManager().getCriteriaBuilder()::equal,
+                              "data",
+                              jpaTm().getEntityManager().getCriteriaBuilder()::equal,
                               "zztz")
                           .build();
-                  return jpaTm().getEntityManager().createQuery(query).getResultList();
+                  return jpaTm().query(query).getResultList();
                 });
     assertThat(result).containsExactly(entity2);
   }
@@ -96,7 +96,7 @@ class CriteriaQueryBuilderTest {
                           .where(
                               "data", jpaTm().getEntityManager().getCriteriaBuilder()::like, "a%")
                           .build();
-                  return jpaTm().getEntityManager().createQuery(query).getResultList();
+                  return jpaTm().query(query).getResultList();
                 });
     assertThat(result).containsExactly(entity3);
   }
@@ -112,7 +112,7 @@ class CriteriaQueryBuilderTest {
                           .where(
                               "data", jpaTm().getEntityManager().getCriteriaBuilder()::like, "%a%")
                           .build();
-                  return jpaTm().getEntityManager().createQuery(query).getResultList();
+                  return jpaTm().query(query).getResultList();
                 });
     assertThat(result).containsExactly(entity1, entity3).inOrder();
   }
@@ -132,7 +132,7 @@ class CriteriaQueryBuilderTest {
                           .where(
                               "data", jpaTm().getEntityManager().getCriteriaBuilder()::like, "%t%")
                           .build();
-                  return jpaTm().getEntityManager().createQuery(query).getResultList();
+                  return jpaTm().query(query).getResultList();
                 });
     assertThat(result).containsExactly(entity1);
   }
@@ -147,7 +147,7 @@ class CriteriaQueryBuilderTest {
                       CriteriaQueryBuilder.create(CriteriaQueryBuilderTestEntity.class)
                           .whereFieldIsIn("data", ImmutableList.of("aaa", "bbb"))
                           .build();
-                  return jpaTm().getEntityManager().createQuery(query).getResultList();
+                  return jpaTm().query(query).getResultList();
                 });
     assertThat(result).containsExactly(entity3).inOrder();
   }
@@ -162,7 +162,7 @@ class CriteriaQueryBuilderTest {
                       CriteriaQueryBuilder.create(CriteriaQueryBuilderTestEntity.class)
                           .whereFieldIsIn("data", ImmutableList.of("aaa", "bbb", "data"))
                           .build();
-                  return jpaTm().getEntityManager().createQuery(query).getResultList();
+                  return jpaTm().query(query).getResultList();
                 });
     assertThat(result).containsExactly(entity1, entity3).inOrder();
   }
@@ -179,7 +179,7 @@ class CriteriaQueryBuilderTest {
                           .where(
                               "data", jpaTm().getEntityManager().getCriteriaBuilder()::like, "%a%")
                           .build();
-                  return jpaTm().getEntityManager().createQuery(query).getResultList();
+                  return jpaTm().query(query).getResultList();
                 });
     assertThat(result).containsExactly(entity3, entity1).inOrder();
   }
@@ -194,7 +194,7 @@ class CriteriaQueryBuilderTest {
                       CriteriaQueryBuilder.create(CriteriaQueryBuilderTestEntity.class)
                           .orderByDesc("data")
                           .build();
-                  return jpaTm().getEntityManager().createQuery(query).getResultList();
+                  return jpaTm().query(query).getResultList();
                 });
     assertThat(result).containsExactly(entity2, entity1, entity3).inOrder();
   }

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

@@ -508,6 +508,19 @@ class JpaTransactionManagerImplTest {
     assertThat(persisted).containsExactlyElementsIn(moreEntities);
   }
 
+  @Test
+  void loadSingleton_detaches() {
+    jpaTm().transact(() -> jpaTm().insert(theEntity));
+    jpaTm()
+        .transact(
+            () ->
+                assertThat(
+                    jpaTm()
+                        .getEntityManager()
+                        .contains(jpaTm().loadSingleton(TestEntity.class).get())))
+        .isFalse();
+  }
+
   @Test
   void delete_succeeds() {
     jpaTm().transact(() -> jpaTm().insert(theEntity));
@@ -572,6 +585,29 @@ class JpaTransactionManagerImplTest {
         .contains("Inserted/updated object reloaded: ");
   }
 
+  @Test
+  void cqQuery_detaches() {
+    jpaTm().transact(() -> jpaTm().insertAll(moreEntities));
+    jpaTm()
+        .transact(
+            () ->
+                assertThat(
+                        jpaTm()
+                            .getEntityManager()
+                            .contains(
+                                jpaTm()
+                                    .query(
+                                        CriteriaQueryBuilder.create(TestEntity.class)
+                                            .where(
+                                                "name",
+                                                jpaTm().getEntityManager().getCriteriaBuilder()
+                                                    ::equal,
+                                                "entity1")
+                                            .build())
+                                    .getSingleResult()))
+                    .isFalse());
+  }
+
   @Test
   void loadAfterPut_fails() {
     assertThat(
@@ -588,6 +624,36 @@ class JpaTransactionManagerImplTest {
         .contains("Inserted/updated object reloaded: ");
   }
 
+  @Test
+  void query_detachesResults() {
+    jpaTm().transact(() -> jpaTm().insertAll(moreEntities));
+    jpaTm()
+        .transact(
+            () ->
+                jpaTm().query("FROM TestEntity", TestEntity.class).getResultList().stream()
+                    .forEach(e -> assertThat(jpaTm().getEntityManager().contains(e)).isFalse()));
+    jpaTm()
+        .transact(
+            () ->
+                jpaTm()
+                    .query("FROM TestEntity", TestEntity.class)
+                    .getResultStream()
+                    .forEach(e -> assertThat(jpaTm().getEntityManager().contains(e)).isFalse()));
+
+    jpaTm()
+        .transact(
+            () ->
+                assertThat(
+                        jpaTm()
+                            .getEntityManager()
+                            .contains(
+                                jpaTm()
+                                    .query(
+                                        "FROM TestEntity WHERE name = 'entity1'", TestEntity.class)
+                                    .getSingleResult()))
+                    .isFalse());
+  }
+
   private void insertPerson(int age) {
     jpaTm()
         .getEntityManager()

core/src/test/java/google/registry/schema/replay/ReplicateToDatastoreActionTest.java

@@ -18,19 +18,28 @@ import static com.google.common.truth.Truth.assertThat;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
 import static google.registry.testing.LogsSubject.assertAboutLogs;
+import static google.registry.util.DateTimeUtils.START_OF_TIME;
 
+import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.testing.TestLogHandler;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Entity;
 import com.googlecode.objectify.annotation.Id;
 import google.registry.config.RegistryConfig;
 import google.registry.model.ImmutableObject;
+import google.registry.model.common.DatabaseMigrationStateSchedule;
+import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState;
+import google.registry.model.ofy.Ofy;
 import google.registry.persistence.VKey;
 import google.registry.persistence.transaction.TransactionEntity;
 import google.registry.testing.AppEngineExtension;
+import google.registry.testing.FakeClock;
+import google.registry.testing.InjectExtension;
 import java.util.List;
 import java.util.logging.Level;
 import java.util.logging.Logger;
+import org.joda.time.DateTime;
+import org.joda.time.Duration;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -38,25 +47,42 @@ import org.junit.jupiter.api.extension.RegisterExtension;
 
 public class ReplicateToDatastoreActionTest {
 
+  private final FakeClock fakeClock = new FakeClock(DateTime.parse("2000-01-01TZ"));
+
   @RegisterExtension
   public final AppEngineExtension appEngine =
       AppEngineExtension.builder()
           .withDatastoreAndCloudSql()
           .withOfyTestEntities(TestEntity.class)
           .withJpaUnitTestEntities(TestEntity.class)
+          .withClock(fakeClock)
           .build();
 
-  ReplicateToDatastoreAction task = new ReplicateToDatastoreAction();
+  @RegisterExtension final InjectExtension injectExtension = new InjectExtension();
 
-  TestLogHandler logHandler;
-
-  public ReplicateToDatastoreActionTest() {}
+  private final ReplicateToDatastoreAction task = new ReplicateToDatastoreAction(fakeClock);
+  private final TestLogHandler logHandler = new TestLogHandler();
 
   @BeforeEach
   public void setUp() {
+    injectExtension.setStaticField(Ofy.class, "clock", fakeClock);
     RegistryConfig.overrideCloudSqlReplicateTransactions(true);
-    logHandler = new TestLogHandler();
     Logger.getLogger(ReplicateToDatastoreAction.class.getCanonicalName()).addHandler(logHandler);
+    DateTime now = fakeClock.nowUtc();
+    ofyTm()
+        .transact(
+            () ->
+                DatabaseMigrationStateSchedule.set(
+                    ImmutableSortedMap.of(
+                        START_OF_TIME,
+                        MigrationState.DATASTORE_ONLY,
+                        now,
+                        MigrationState.DATASTORE_PRIMARY,
+                        now.plusHours(1),
+                        MigrationState.DATASTORE_PRIMARY_READ_ONLY,
+                        now.plusHours(2),
+                        MigrationState.SQL_PRIMARY)));
+    fakeClock.advanceBy(Duration.standardDays(1));
   }
 
   @AfterEach
@@ -170,6 +196,36 @@ public class ReplicateToDatastoreActionTest {
             "Missing transaction: last transaction id = -1, next available transaction = 1");
   }
 
+  @Test
+  void testNotInMigrationState_doesNothing() {
+    // set a schedule that backtracks the current status to DATASTORE_PRIMARY_READ_ONLY
+    DateTime now = fakeClock.nowUtc();
+    ofyTm()
+        .transact(
+            () ->
+                DatabaseMigrationStateSchedule.set(
+                    ImmutableSortedMap.<DateTime, MigrationState>naturalOrder()
+                        .put(START_OF_TIME, MigrationState.DATASTORE_ONLY)
+                        .put(START_OF_TIME.plusHours(1), MigrationState.DATASTORE_PRIMARY)
+                        .put(START_OF_TIME.plusHours(2), MigrationState.DATASTORE_PRIMARY_READ_ONLY)
+                        .put(START_OF_TIME.plusHours(3), MigrationState.SQL_PRIMARY)
+                        .put(now.plusHours(1), MigrationState.SQL_PRIMARY_READ_ONLY)
+                        .put(now.plusHours(2), MigrationState.DATASTORE_PRIMARY_READ_ONLY)
+                        .build()));
+    fakeClock.advanceBy(Duration.standardDays(1));
+
+    jpaTm().transact(() -> jpaTm().insert(new TestEntity("foo")));
+    task.run();
+    // Replication shouldn't have happened
+    assertThat(ofyTm().loadAllOf(TestEntity.class)).isEmpty();
+    assertAboutLogs()
+        .that(logHandler)
+        .hasLogAtLevelWithMessage(
+            Level.INFO,
+            "Skipping ReplicateToDatastoreAction because we are in migration phase "
+                + "DATASTORE_PRIMARY_READ_ONLY.");
+  }
+
   @Entity(name = "ReplicationTestEntity")
   @javax.persistence.Entity(name = "TestEntity")
   private static class TestEntity extends ImmutableObject {

core/src/test/java/google/registry/testing/AppEngineExtension.java

@@ -17,12 +17,10 @@ package google.registry.testing;
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.io.Files.asCharSink;
 import static com.google.common.truth.Truth.assertWithMessage;
-import static google.registry.model.ofy.ObjectifyService.auditedOfy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatabaseHelper.persistSimpleResources;
 import static google.registry.testing.DualDatabaseTestInvocationContextProvider.injectTmForDualDatabaseTest;
 import static google.registry.testing.DualDatabaseTestInvocationContextProvider.restoreTmAfterDualDatabaseTest;
-import static google.registry.util.DateTimeUtils.START_OF_TIME;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 import static google.registry.util.ResourceUtils.readResourceUtf8;
 import static java.nio.charset.StandardCharsets.UTF_8;
@@ -42,16 +40,10 @@ import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.collect.Sets;
 import com.google.common.io.Files;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.ObjectifyFilter;
-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.ObjectifyService;
 import google.registry.model.registrar.Registrar;
 import google.registry.model.registrar.Registrar.State;
@@ -402,18 +394,8 @@ public final class AppEngineExtension implements BeforeEachCallback, AfterEachCa
       if (withDatastore && !withoutCannedData) {
         loadInitialData();
       }
-    } else {
-      // If we're using SQL, set replayed entities to use SQL
-      DatabaseTransitionSchedule schedule =
-          DatabaseTransitionSchedule.create(
-              TransitionId.REPLAYED_ENTITIES,
-              TimedTransitionProperty.fromValueMap(
-                  ImmutableSortedMap.of(START_OF_TIME, PrimaryDatabase.CLOUD_SQL),
-                  PrimaryDatabaseTransition.class));
-      tm().transactNew(() -> auditedOfy().saveWithoutBackup().entity(schedule).now());
-      if (withCloudSql && !withJpaUnitTest && !withoutCannedData) {
-        loadInitialData();
-      }
+    } else if (withCloudSql && !withJpaUnitTest && !withoutCannedData) {
+      loadInitialData();
     }
   }
 

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

@@ -0,0 +1,65 @@
+// 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.model.common.DatabaseMigrationStateSchedule.DEFAULT_TRANSITION_MAP;
+import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
+import static google.registry.util.DateTimeUtils.START_OF_TIME;
+
+import com.google.common.collect.ImmutableSortedMap;
+import google.registry.model.common.DatabaseMigrationStateSchedule;
+import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState;
+import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.TestOfyAndSql;
+import org.joda.time.DateTime;
+import org.junit.jupiter.api.BeforeEach;
+
+/** Tests for {@link GetDatabaseMigrationStateCommand}. */
+@DualDatabaseTest
+public class GetDatabaseMigrationStateCommandTest
+    extends CommandTestCase<GetDatabaseMigrationStateCommand> {
+
+  @BeforeEach
+  void beforeEach() {
+    ofyTm().transact(() -> DatabaseMigrationStateSchedule.set(DEFAULT_TRANSITION_MAP.toValueMap()));
+  }
+
+  @TestOfyAndSql
+  void testInitial_returnsDatastoreOnly() throws Exception {
+    runCommand();
+    assertStdoutIs(
+        String.format("Current migration schedule: %s\n", DEFAULT_TRANSITION_MAP.toValueMap()));
+  }
+
+  @TestOfyAndSql
+  void testFullSchedule() throws Exception {
+    DateTime now = fakeClock.nowUtc();
+    ImmutableSortedMap<DateTime, MigrationState> transitions =
+        ImmutableSortedMap.of(
+            START_OF_TIME,
+            MigrationState.DATASTORE_ONLY,
+            now.plusHours(1),
+            MigrationState.DATASTORE_PRIMARY,
+            now.plusHours(2),
+            MigrationState.DATASTORE_PRIMARY_READ_ONLY,
+            now.plusHours(3),
+            MigrationState.SQL_PRIMARY,
+            now.plusHours(4),
+            MigrationState.SQL_ONLY);
+    ofyTm().transact(() -> DatabaseMigrationStateSchedule.set(transitions));
+    runCommand();
+    assertStdoutIs(String.format("Current migration schedule: %s\n", transitions));
+  }
+}

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

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

@@ -0,0 +1,165 @@
+// 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 com.google.common.truth.Truth8.assertThat;
+import static google.registry.model.common.DatabaseMigrationStateSchedule.DEFAULT_TRANSITION_MAP;
+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.DatabaseMigrationStateSchedule;
+import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState;
+import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.TestOfyAndSql;
+import org.joda.time.DateTime;
+
+/** Tests for {@link SetDatabaseMigrationStateCommand}. */
+@DualDatabaseTest
+public class SetDatabaseMigrationStateCommandTest
+    extends CommandTestCase<SetDatabaseMigrationStateCommand> {
+
+  @TestOfyAndSql
+  void testSuccess_setsBasicSchedule() throws Exception {
+    assertThat(DatabaseMigrationStateSchedule.get()).isEqualTo(DEFAULT_TRANSITION_MAP);
+    assertThat(ofyTm().transact(() -> ofyTm().loadSingleton(DatabaseMigrationStateSchedule.class)))
+        .isEmpty();
+    runCommandForced("--migration_schedule=1970-01-01T00:00:00.000Z=DATASTORE_ONLY");
+    // use a raw ofy call to check what's in the DB
+    ofyTm()
+        .transact(
+            () ->
+                assertThat(
+                        ofyTm()
+                            .loadSingleton(DatabaseMigrationStateSchedule.class)
+                            .get()
+                            .migrationTransitions)
+                    .isEqualTo(DEFAULT_TRANSITION_MAP));
+    assertThat(DatabaseMigrationStateSchedule.get()).isEqualTo(DEFAULT_TRANSITION_MAP);
+  }
+
+  @TestOfyAndSql
+  void testSuccess_fullSchedule() throws Exception {
+    DateTime now = fakeClock.nowUtc();
+    DateTime datastorePrimary = now.plusHours(1);
+    DateTime datastorePrimaryReadOnly = now.plusHours(2);
+    DateTime sqlPrimary = now.plusHours(3);
+    DateTime sqlOnly = now.plusHours(4);
+    runCommandForced(
+        String.format(
+            "--migration_schedule=%s=DATASTORE_ONLY,%s=DATASTORE_PRIMARY,"
+                + "%s=DATASTORE_PRIMARY_READ_ONLY,%s=SQL_PRIMARY,%s=SQL_ONLY",
+            START_OF_TIME, datastorePrimary, datastorePrimaryReadOnly, sqlPrimary, sqlOnly));
+    assertThat(DatabaseMigrationStateSchedule.get().toValueMap())
+        .containsExactlyEntriesIn(
+            ImmutableSortedMap.of(
+                START_OF_TIME,
+                MigrationState.DATASTORE_ONLY,
+                datastorePrimary,
+                MigrationState.DATASTORE_PRIMARY,
+                datastorePrimaryReadOnly,
+                MigrationState.DATASTORE_PRIMARY_READ_ONLY,
+                sqlPrimary,
+                MigrationState.SQL_PRIMARY,
+                sqlOnly,
+                MigrationState.SQL_ONLY));
+  }
+
+  @TestOfyAndSql
+  void testSuccess_warnsOnChangeSoon() throws Exception {
+    DateTime now = fakeClock.nowUtc();
+    runCommandForced(
+        String.format(
+            "--migration_schedule=%s=DATASTORE_ONLY,%s=DATASTORE_PRIMARY",
+            START_OF_TIME, now.plusMinutes(1)));
+    assertThat(DatabaseMigrationStateSchedule.get().toValueMap())
+        .containsExactlyEntriesIn(
+            ImmutableSortedMap.of(
+                START_OF_TIME,
+                MigrationState.DATASTORE_ONLY,
+                now.plusMinutes(1),
+                MigrationState.DATASTORE_PRIMARY));
+    assertInStdout("MAY BE DANGEROUS");
+  }
+
+  @TestOfyAndSql
+  void testSuccess_goesBackward() throws Exception {
+    DateTime now = fakeClock.nowUtc();
+    runCommandForced(
+        String.format(
+            "--migration_schedule=%s=DATASTORE_ONLY,%s=DATASTORE_PRIMARY,"
+                + "%s=DATASTORE_PRIMARY_READ_ONLY,%s=DATASTORE_PRIMARY",
+            START_OF_TIME, now.plusHours(1), now.plusHours(2), now.plusHours(3)));
+    assertThat(DatabaseMigrationStateSchedule.get().toValueMap())
+        .containsExactlyEntriesIn(
+            ImmutableSortedMap.of(
+                START_OF_TIME,
+                MigrationState.DATASTORE_ONLY,
+                now.plusHours(1),
+                MigrationState.DATASTORE_PRIMARY,
+                now.plusHours(2),
+                MigrationState.DATASTORE_PRIMARY_READ_ONLY,
+                now.plusHours(3),
+                MigrationState.DATASTORE_PRIMARY));
+  }
+
+  @TestOfyAndSql
+  void testFailure_invalidTransition() {
+    assertThat(
+            assertThrows(
+                IllegalArgumentException.class,
+                () ->
+                    runCommandForced(
+                        String.format(
+                            "--migration_schedule=%s=DATASTORE_ONLY,%s=DATASTORE_PRIMARY_READ_ONLY",
+                            START_OF_TIME, START_OF_TIME.plusHours(1)))))
+        .hasMessageThat()
+        .isEqualTo(
+            "validStateTransitions map cannot transition from DATASTORE_ONLY "
+                + "to DATASTORE_PRIMARY_READ_ONLY.");
+  }
+
+  @TestOfyAndSql
+  void testFailure_invalidTransitionFromOldToNew() {
+    // The map we pass in is valid by itself, but we can't go from DATASTORE_ONLY now to
+    // DATASTORE_PRIMARY_READ_ONLY now
+    DateTime now = fakeClock.nowUtc();
+    assertThat(
+            assertThrows(
+                IllegalArgumentException.class,
+                () ->
+                    runCommandForced(
+                        String.format(
+                            "--migration_schedule=%s=DATASTORE_ONLY,"
+                                + "%s=DATASTORE_PRIMARY,%s=DATASTORE_PRIMARY_READ_ONLY",
+                            START_OF_TIME, now.minusHours(2), now.minusHours(1)))))
+        .hasMessageThat()
+        .isEqualTo(
+            "Cannot transition from current state-as-of-now DATASTORE_ONLY "
+                + "to new state-as-of-now DATASTORE_PRIMARY_READ_ONLY");
+  }
+
+  @TestOfyAndSql
+  void testFailure_invalidParams() {
+    assertThrows(ParameterException.class, this::runCommandForced);
+    assertThrows(ParameterException.class, () -> runCommandForced("--migration_schedule=FOOBAR"));
+    assertThrows(
+        ParameterException.class,
+        () -> runCommandForced("--migration_schedule=1970-01-01T00:00:00.000Z=FOOBAR"));
+  }
+}

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

@@ -1,134 +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.
-
-package google.registry.tools;
-
-import static com.google.common.truth.Truth.assertThat;
-import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
-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 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;
-import org.junit.jupiter.api.Test;
-
-/** Unit tests for {@link SetDatabaseTransitionScheduleCommand}. */
-public class SetDatabaseTransitionScheduleCommandTest
-    extends CommandTestCase<SetDatabaseTransitionScheduleCommand> {
-
-  @BeforeEach
-  void setup() {
-    fakeClock.setTo(DateTime.parse("2020-12-01T00:00:00Z"));
-  }
-
-  @Test
-  void testFailure_noTransitionId() throws Exception {
-    ParameterException thrown =
-        assertThrows(
-            ParameterException.class,
-            () -> runCommandForced("--transition_schedule=2021-04-14T00:00:00.000Z=DATASTORE"));
-    assertThat(thrown).hasMessageThat().contains("--transition_id");
-  }
-
-  @Test
-  void testSuccess_currentScheduleIsEmpty() throws Exception {
-    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");
-    assertThat(
-            ofyTm()
-                .transact(
-                    () ->
-                        DatabaseTransitionSchedule.get(TransitionId.SIGNED_MARK_REVOCATION_LIST)
-                            .get()
-                            .getPrimaryDatabase()))
-        .isEqualTo(PrimaryDatabase.DATASTORE);
-    assertThat(command.prompt())
-        .isEqualTo(
-            "Insert new schedule {1970-01-01T00:00:00.000Z=DATASTORE} "
-                + "for transition ID SIGNED_MARK_REVOCATION_LIST?");
-  }
-
-  @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())
-        .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?");
-  }
-}

core/src/test/resources/google/registry/beam/invoicing/billing_events_test.sql

@@ -27,7 +27,7 @@ SELECT
   reason as action,
   targetId as domain,
   BillingEvent.domainRepoId as repositoryId,
-  periodYears as years,
+  IFNULL(periodYears, 0) as years,
   BillingEvent.currency AS currency,
   BillingEvent.amount as amount,
   -- We'll strip out non-useful flags downstream

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

@@ -2,8 +2,7 @@ AllocationToken
 Cancellation
 ContactResource
 Cursor
-DatabaseMigrationStateWrapper
-DatabaseTransitionSchedule
+DatabaseMigrationStateSchedule
 DomainBase
 EntityGroupRoot
 EppResourceIndex

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

@@ -1,8 +1,7 @@
 ClaimsList
 ClaimsListSingleton
 Cursor
-DatabaseMigrationStateWrapper
-DatabaseTransitionSchedule
+DatabaseMigrationStateSchedule
 KmsSecret
 KmsSecretRevision
 PremiumList

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

@@ -78,30 +78,21 @@ class google.registry.model.common.Cursor {
   google.registry.model.UpdateAutoTimestamp lastUpdateTime;
   org.joda.time.DateTime cursorTime;
 }
-class google.registry.model.common.DatabaseMigrationStateWrapper {
+class google.registry.model.common.DatabaseMigrationStateSchedule {
   @Id long id;
   @Parent com.googlecode.objectify.Key<google.registry.model.common.EntityGroupRoot> parent;
-  google.registry.model.common.DatabaseMigrationStateWrapper$MigrationState migrationState;
+  google.registry.model.common.TimedTransitionProperty<google.registry.model.common.DatabaseMigrationStateSchedule$MigrationState, google.registry.model.common.DatabaseMigrationStateSchedule$MigrationStateTransition> migrationTransitions;
 }
-enum google.registry.model.common.DatabaseMigrationStateWrapper$MigrationState {
+enum google.registry.model.common.DatabaseMigrationStateSchedule$MigrationState {
   DATASTORE_ONLY;
   DATASTORE_PRIMARY;
   DATASTORE_PRIMARY_READ_ONLY;
   SQL_ONLY;
   SQL_PRIMARY;
+  SQL_PRIMARY_READ_ONLY;
 }
-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;
+class google.registry.model.common.DatabaseMigrationStateSchedule$MigrationStateTransition {
+  google.registry.model.common.DatabaseMigrationStateSchedule$MigrationState migrationState;
   org.joda.time.DateTime transitionTime;
 }
 class google.registry.model.common.EntityGroupRoot {

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

@@ -94,3 +94,4 @@ V93__defer_all_fkeys.sql
 V94__rename_lock_scope.sql
 V95__add_contacts_indexes_on_domain.sql
 V96__rename_sql_checkpoint_fields.sql
+V97__add_recurrence_history_id_column_to_onetime.sql

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

@@ -0,0 +1,21 @@
+-- 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.
+
+ALTER TABLE "BillingEvent"
+    ADD COLUMN IF NOT EXISTS "recurrence_history_revision_id" INT8;
+
+ALTER TABLE "BillingEvent" ADD CONSTRAINT fk_billing_event_recurrence_history
+   FOREIGN KEY (domain_repo_id, recurrence_history_revision_id)
+   REFERENCES "DomainHistory"(domain_repo_id, history_revision_id)
+   DEFERRABLE INITIALLY DEFERRED;

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

@@ -63,6 +63,7 @@
         cost_amount numeric(19, 2),
         cost_currency text,
         period_years int4,
+        recurrence_history_revision_id int8,
         synthetic_creation_time timestamptz,
         primary key (billing_event_id)
     );

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

@@ -97,7 +97,8 @@ CREATE TABLE public."BillingEvent" (
     cost_amount numeric(19,2),
     cost_currency text,
     period_years integer,
-    synthetic_creation_time timestamp with time zone
+    synthetic_creation_time timestamp with time zone,
+    recurrence_history_revision_id bigint
 );
 
 
@@ -1985,6 +1986,14 @@ ALTER TABLE ONLY public."BillingEvent"
     ADD CONSTRAINT fk_billing_event_domain_history FOREIGN KEY (domain_repo_id, domain_history_revision_id) REFERENCES public."DomainHistory"(domain_repo_id, history_revision_id) DEFERRABLE INITIALLY DEFERRED;
 
 
+--
+-- Name: BillingEvent fk_billing_event_recurrence_history; Type: FK CONSTRAINT; Schema: public; Owner: -
+--
+
+ALTER TABLE ONLY public."BillingEvent"
+    ADD CONSTRAINT fk_billing_event_recurrence_history FOREIGN KEY (domain_repo_id, recurrence_history_revision_id) REFERENCES public."DomainHistory"(domain_repo_id, history_revision_id) DEFERRABLE INITIALLY DEFERRED;
+
+
 --
 -- Name: BillingEvent fk_billing_event_registrar_id; Type: FK CONSTRAINT; Schema: public; Owner: -
 --

release/cloudbuild-deploy.yaml

@@ -20,22 +20,14 @@
 # See https://github.com/spinnaker/spinnaker/issues/3028 for more information.
 steps:
 # Pull the credential for nomulus tool.
-- name: 'gcr.io/$PROJECT_ID/builder:latest'
-  args:
-  - gsutil
-  - cp
-  - gs://$PROJECT_ID-deploy/secrets/tool-credential.json.enc
-  - .
-# Decrypt the credential.
 - name: 'gcr.io/$PROJECT_ID/builder:latest'
   entrypoint: /bin/bash
   args:
   - -c
   - |
     set -e
-    cat tool-credential.json.enc | base64 -d | gcloud kms decrypt \
-      --ciphertext-file=- --plaintext-file=tool-credential.json \
-      --location=global --keyring=nomulus-tool-keyring --key=nomulus-tool-key
+    gcloud secrets versions access latest \
+      --secret nomulus-tool-cloudbuild-credential > tool-credential.json
 # Deploy the GAE config files.
 # First authorize the gcloud tool to use the credential json file, then
 # download and unzip the tarball that contains the relevant config files
@@ -53,8 +45,6 @@ steps:
     fi
     gsutil cp gs://$PROJECT_ID-deploy/${TAG_NAME}/${_ENV}.tar .
     tar -xvf ${_ENV}.tar
-    # Note that this currently does not work for google.com projects that
-    # we use due to b/137891685. External projects are likely to work.
     for filename in cron dispatch dos index queue; do
       gcloud -q --project $project_id app deploy \
         default/WEB-INF/appengine-generated/$filename.yaml

release/cloudbuild-schema-deploy.yaml

@@ -39,12 +39,9 @@ steps:
   - -c
   - |
     set -e
-    gsutil cp gs://$PROJECT_ID-deploy/secrets/tool-credential.json.enc - \
-      | base64 -d \
-      | gcloud kms decrypt \
-        --ciphertext-file=- \
-        --plaintext-file=/secrets/cloud_sql_credential.json \
-        --location=global --keyring=nomulus-tool-keyring --key=nomulus-tool-key
+    gcloud secrets versions access latest \
+      --secret nomulus-tool-cloudbuild-credential \
+      > /secrets/cloud_sql_credential.json
 # Fetch the Cloud SQL credential for schema_deployer
 - name: 'gcr.io/$PROJECT_ID/nomulus-tool:latest'
   volumes: