Add Cloud SQL read to Spec11Pipeline (#1173)

* Add Cloud SQL read to Spec11Pipeline

* Add database option

* Add database parameter

* Add a test of the full pipeline

* Use DatabaseHelper in tests

* restore the original tm

* More test fixes

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

@@ -66,6 +66,8 @@ import org.joda.time.Duration;
     service = Action.Service.BACKEND,
     path = "/_dr/task/deleteOldCommitLogs",
     auth = Auth.AUTH_INTERNAL_OR_ADMIN)
+// No longer needed in SQL. Subject to future removal.
+@Deprecated
 public final class DeleteOldCommitLogsAction implements Runnable {
 
   private static final int NUM_MAP_SHARDS = 20;

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,88 @@ 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 eventTime <= :executeTime "
+                                  + "AND eventTime < recurrenceEndTime "
+                                  + "ORDER BY id ASC",
+                              Recurring.class)
+                          .setParameter("executeTime", executeTime)
+                          // 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 +213,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 +225,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 +263,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 +279,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/batch/RefreshDnsOnHostRenameAction.java

@@ -23,6 +23,7 @@ import static google.registry.batch.AsyncTaskEnqueuer.PARAM_REQUESTED_TIME;
 import static google.registry.batch.AsyncTaskEnqueuer.QUEUE_ASYNC_HOST_RENAME;
 import static google.registry.batch.AsyncTaskMetrics.OperationType.DNS_REFRESH;
 import static google.registry.mapreduce.inputs.EppResourceInputs.createEntityInput;
+import static google.registry.model.EppResourceUtils.getLinkedDomainKeys;
 import static google.registry.model.EppResourceUtils.isActive;
 import static google.registry.model.EppResourceUtils.isDeleted;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
@@ -69,6 +70,7 @@ import java.util.logging.Level;
 import javax.annotation.Nullable;
 import javax.inject.Inject;
 import javax.inject.Named;
+import org.apache.http.HttpStatus;
 import org.joda.time.DateTime;
 import org.joda.time.Duration;
 
@@ -86,6 +88,8 @@ public class RefreshDnsOnHostRenameAction implements Runnable {
   @Inject Clock clock;
   @Inject MapreduceRunner mrRunner;
   @Inject @Named(QUEUE_ASYNC_HOST_RENAME) Queue pullQueue;
+
+  @Inject DnsQueue dnsQueue;
   @Inject RequestStatusChecker requestStatusChecker;
   @Inject Response response;
   @Inject Retrier retrier;
@@ -153,7 +157,39 @@ public class RefreshDnsOnHostRenameAction implements Runnable {
     } else {
       logger.atInfo().log(
           "Processing asynchronous DNS refresh for renamed hosts: %s", hostKeys.build());
-      runMapreduce(refreshRequests, lock);
+      if (tm().isOfy()) {
+        runMapreduce(refreshRequests, lock);
+      } else {
+        try {
+          refreshRequests.stream()
+              .flatMap(
+                  request ->
+                      getLinkedDomainKeys(request.hostKey(), request.lastUpdateTime(), null)
+                          .stream())
+              .distinct()
+              .map(domainKey -> tm().transact(() -> tm().loadByKey(domainKey).getDomainName()))
+              .forEach(
+                  domainName -> {
+                    retrier.callWithRetry(
+                        () -> dnsQueue.addDomainRefreshTask(domainName),
+                        TransientFailureException.class);
+                    logger.atInfo().log("Enqueued DNS refresh for domain %s.", domainName);
+                  });
+          deleteTasksWithRetry(
+              refreshRequests,
+              getQueue(QUEUE_ASYNC_HOST_RENAME),
+              asyncTaskMetrics,
+              retrier,
+              OperationResult.SUCCESS);
+        } catch (Throwable t) {
+          String message = "Error refreshing DNS on host rename.";
+          logger.atSevere().withCause(t).log(message);
+          response.setPayload(message);
+          response.setStatus(HttpStatus.SC_INTERNAL_SERVER_ERROR);
+        } finally {
+          lock.get().release();
+        }
+      }
     }
   }
 

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

@@ -54,6 +54,8 @@ import javax.inject.Inject;
     service = Action.Service.BACKEND,
     path = "/_dr/task/resaveAllEppResources",
     auth = Auth.AUTH_INTERNAL_OR_ADMIN)
+// No longer needed in SQL. Subject to future removal.
+@Deprecated
 public class ResaveAllEppResourcesAction implements Runnable {
 
   @Inject MapreduceRunner mrRunner;

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/beam/spec11/Spec11Pipeline.java

@@ -23,8 +23,10 @@ import dagger.Component;
 import dagger.Module;
 import dagger.Provides;
 import google.registry.beam.common.RegistryJpaIO;
+import google.registry.beam.common.RegistryJpaIO.Read;
 import google.registry.beam.spec11.SafeBrowsingTransforms.EvaluateSafeBrowsingFn;
 import google.registry.config.RegistryConfig.ConfigModule;
+import google.registry.model.domain.DomainBase;
 import google.registry.model.reporting.Spec11ThreatMatch;
 import google.registry.model.reporting.Spec11ThreatMatch.ThreatType;
 import google.registry.util.Retrier;
@@ -97,20 +99,9 @@ public class Spec11Pipeline implements Serializable {
 
   void setupPipeline(Pipeline pipeline) {
     PCollection<Subdomain> domains =
-        pipeline.apply(
-            "Read active domains from BigQuery",
-            BigQueryIO.read(Subdomain::parseFromRecord)
-                .fromQuery(
-                    SqlTemplate.create(getQueryFromFile(Spec11Pipeline.class, "subdomains.sql"))
-                        .put("PROJECT_ID", options.getProject())
-                        .put("DATASTORE_EXPORT_DATASET", "latest_datastore_export")
-                        .put("REGISTRAR_TABLE", "Registrar")
-                        .put("DOMAIN_BASE_TABLE", "DomainBase")
-                        .build())
-                .withCoder(SerializableCoder.of(Subdomain.class))
-                .usingStandardSql()
-                .withoutValidation()
-                .withTemplateCompatibility());
+        options.getDatabase().equals("DATASTORE")
+            ? readFromBigQuery(options, pipeline)
+            : readFromCloudSql(pipeline);
 
     PCollection<KV<Subdomain, ThreatMatch>> threatMatches =
         domains.apply("Run through SafeBrowsing API", ParDo.of(safeBrowsingFn));
@@ -119,6 +110,47 @@ public class Spec11Pipeline implements Serializable {
     saveToGcs(threatMatches, options);
   }
 
+  static PCollection<Subdomain> readFromCloudSql(Pipeline pipeline) {
+    Read<Object[], Subdomain> read =
+        RegistryJpaIO.read(
+            "select d, r.emailAddress from Domain d join Registrar r on"
+                + " d.currentSponsorClientId = r.clientIdentifier where r.type = 'REAL'"
+                + " and d.deletionTime > now()",
+            Spec11Pipeline::parseRow);
+
+    return pipeline.apply("Read active domains from Cloud SQL", read);
+  }
+
+  static PCollection<Subdomain> readFromBigQuery(Spec11PipelineOptions options, Pipeline pipeline) {
+    return pipeline.apply(
+        "Read active domains from BigQuery",
+        BigQueryIO.read(Subdomain::parseFromRecord)
+            .fromQuery(
+                SqlTemplate.create(getQueryFromFile(Spec11Pipeline.class, "subdomains.sql"))
+                    .put("PROJECT_ID", options.getProject())
+                    .put("DATASTORE_EXPORT_DATASET", "latest_datastore_export")
+                    .put("REGISTRAR_TABLE", "Registrar")
+                    .put("DOMAIN_BASE_TABLE", "DomainBase")
+                    .build())
+            .withCoder(SerializableCoder.of(Subdomain.class))
+            .usingStandardSql()
+            .withoutValidation()
+            .withTemplateCompatibility());
+  }
+
+  private static Subdomain parseRow(Object[] row) {
+    DomainBase domainBase = (DomainBase) row[0];
+    String emailAddress = (String) row[1];
+    if (emailAddress == null) {
+      emailAddress = "";
+    }
+    return Subdomain.create(
+        domainBase.getDomainName(),
+        domainBase.getRepoId(),
+        domainBase.getCurrentSponsorClientId(),
+        emailAddress);
+  }
+
   static void saveToSql(
       PCollection<KV<Subdomain, ThreatMatch>> threatMatches, Spec11PipelineOptions options) {
     String transformId = "Spec11 Threat Matches";

core/src/main/java/google/registry/beam/spec11/Spec11PipelineOptions.java

@@ -34,4 +34,9 @@ public interface Spec11PipelineOptions extends RegistryPipelineOptions {
   String getReportingBucketUrl();
 
   void setReportingBucketUrl(String value);
+
+  @Description("The database to read data from.")
+  String getDatabase();
+
+  void setDatabase(String value);
 }

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/alpha/default/WEB-INF/cron.xml

@@ -199,6 +199,15 @@
     <target>backend</target>
   </cron>
 
+  <cron>
+    <url><![CDATA[/_dr/task/replayCommitLogsToSql]]></url>
+    <description>
+      Replays recent commit logs from Datastore to the SQL secondary backend.
+    </description>
+    <schedule>every 3 minutes</schedule>
+    <target>backend</target>
+  </cron>
+
   <cron>
     <url><![CDATA[/_dr/cron/readDnsQueue?jitterSeconds=45]]></url>
     <description>

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/env/crash/default/WEB-INF/cron.xml

@@ -200,4 +200,12 @@
     <target>backend</target>
   </cron>
 
+  <cron>
+    <url><![CDATA[/_dr/task/replayCommitLogsToSql]]></url>
+    <description>
+      Replays recent commit logs from Datastore to the SQL secondary backend.
+    </description>
+    <schedule>every 3 minutes</schedule>
+    <target>backend</target>
+  </cron>
 </cronentries>

core/src/main/java/google/registry/env/qa/default/WEB-INF/cron.xml

@@ -100,4 +100,12 @@
     <target>backend</target>
   </cron>
 
+  <cron>
+    <url><![CDATA[/_dr/task/replayCommitLogsToSql]]></url>
+    <description>
+      Replays recent commit logs from Datastore to the SQL secondary backend.
+    </description>
+    <schedule>every 3 minutes</schedule>
+    <target>backend</target>
+  </cron>
 </cronentries>

core/src/main/java/google/registry/env/sandbox/default/WEB-INF/cron.xml

@@ -229,4 +229,12 @@
     <target>backend</target>
   </cron>
 
+  <cron>
+    <url><![CDATA[/_dr/task/replayCommitLogsToSql]]></url>
+    <description>
+      Replays recent commit logs from Datastore to the SQL secondary backend.
+    </description>
+    <schedule>every 3 minutes</schedule>
+    <target>backend</target>
+  </cron>
 </cronentries>

core/src/main/java/google/registry/export/ExportDomainListsAction.java

@@ -19,9 +19,10 @@ import static com.google.common.base.Verify.verifyNotNull;
 import static google.registry.mapreduce.inputs.EppResourceInputs.createEntityInput;
 import static google.registry.model.EppResourceUtils.isActive;
 import static google.registry.model.registry.Registries.getTldsOfType;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.request.Action.Method.POST;
 import static java.nio.charset.StandardCharsets.UTF_8;
-import static org.joda.time.DateTimeZone.UTC;
 
 import com.google.appengine.tools.cloudstorage.GcsFilename;
 import com.google.appengine.tools.cloudstorage.RetryParams;
@@ -45,12 +46,14 @@ import google.registry.request.Action;
 import google.registry.request.Response;
 import google.registry.request.auth.Auth;
 import google.registry.storage.drive.DriveConnection;
+import google.registry.util.Clock;
 import google.registry.util.NonFinalForTesting;
 import java.io.IOException;
 import java.io.ObjectInputStream;
 import java.io.OutputStream;
 import java.io.OutputStreamWriter;
 import java.io.Writer;
+import java.util.List;
 import java.util.function.Supplier;
 import javax.inject.Inject;
 import org.joda.time.DateTime;
@@ -70,9 +73,13 @@ public class ExportDomainListsAction implements Runnable {
 
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
   private static final int MAX_NUM_REDUCE_SHARDS = 100;
+  public static final String REGISTERED_DOMAINS_FILENAME = "registered_domains.txt";
 
   @Inject MapreduceRunner mrRunner;
   @Inject Response response;
+  @Inject Clock clock;
+  @Inject DriveConnection driveConnection;
+
   @Inject @Config("domainListsGcsBucket") String gcsBucket;
   @Inject @Config("gcsBufferSize") int gcsBufferSize;
   @Inject ExportDomainListsAction() {}
@@ -81,15 +88,99 @@ public class ExportDomainListsAction implements Runnable {
   public void run() {
     ImmutableSet<String> realTlds = getTldsOfType(TldType.REAL);
     logger.atInfo().log("Exporting domain lists for tlds %s", realTlds);
-    mrRunner
-        .setJobName("Export domain lists")
-        .setModuleName("backend")
-        .setDefaultReduceShards(Math.min(realTlds.size(), MAX_NUM_REDUCE_SHARDS))
-        .runMapreduce(
-            new ExportDomainListsMapper(DateTime.now(UTC), realTlds),
-            new ExportDomainListsReducer(gcsBucket, gcsBufferSize),
-            ImmutableList.of(createEntityInput(DomainBase.class)))
-        .sendLinkToMapreduceConsole(response);
+    if (tm().isOfy()) {
+      mrRunner
+          .setJobName("Export domain lists")
+          .setModuleName("backend")
+          .setDefaultReduceShards(Math.min(realTlds.size(), MAX_NUM_REDUCE_SHARDS))
+          .runMapreduce(
+              new ExportDomainListsMapper(clock.nowUtc(), realTlds),
+              new ExportDomainListsReducer(gcsBucket, gcsBufferSize),
+              ImmutableList.of(createEntityInput(DomainBase.class)))
+          .sendLinkToMapreduceConsole(response);
+    } else {
+      realTlds.forEach(
+          tld -> {
+            List<String> domains =
+                tm().transact(
+                        () ->
+                            // Note that if we had "creationTime <= :now" in the condition (not
+                            // necessary as there is no pending creation, the order of deletionTime
+                            // and creationTime in the query would have been significant and it
+                            // should come after deletionTime. When Hibernate substitutes "now" it
+                            // will first validate that the **first** field that is to be compared
+                            // with it (deletionTime) is assignable from the substituted Java object
+                            // (click.nowUtc()). Since creationTime is a CreateAutoTimestamp, if it
+                            // comes first, we will need to substitute "now" with
+                            // CreateAutoTimestamp.create(clock.nowUtc()). This might look a bit
+                            // strange as the Java object type is clearly incompatible between the
+                            // two fields deletionTime (DateTime) and creationTime, yet they are
+                            // compared with the same "now". It is actually OK because in the end
+                            // Hibernate converts everything to SQL types (and Java field names to
+                            // SQL column names) to run the query. Both CreateAutoTimestamp and
+                            // DateTime are persisted as timestamp_z in SQL. It is only the
+                            // validation that compares the Java types, and only with the first
+                            // field that compares with the substituted value.
+                            jpaTm()
+                                .query(
+                                    "SELECT fullyQualifiedDomainName FROM Domain "
+                                        + "WHERE tld = :tld "
+                                        + "AND deletionTime > :now "
+                                        + "ORDER by fullyQualifiedDomainName ASC",
+                                    String.class)
+                                .setParameter("tld", tld)
+                                .setParameter("now", clock.nowUtc())
+                                .getResultList());
+            String domainsList = Joiner.on("\n").join(domains);
+            logger.atInfo().log(
+                "Exporting %d domains for TLD %s to GCS and Drive.", domains.size(), tld);
+            exportToGcs(tld, domainsList, gcsBucket, gcsBufferSize);
+            exportToDrive(tld, domainsList, driveConnection);
+          });
+    }
+  }
+
+  protected static boolean exportToDrive(
+      String tld, String domains, DriveConnection driveConnection) {
+    verifyNotNull(driveConnection, "Expecting non-null driveConnection");
+    try {
+      Registry registry = Registry.get(tld);
+      if (registry.getDriveFolderId() == null) {
+        logger.atInfo().log(
+            "Skipping registered domains export for TLD %s because Drive folder isn't specified",
+            tld);
+      } else {
+        String resultMsg =
+            driveConnection.createOrUpdateFile(
+                REGISTERED_DOMAINS_FILENAME,
+                MediaType.PLAIN_TEXT_UTF_8,
+                registry.getDriveFolderId(),
+                domains.getBytes(UTF_8));
+        logger.atInfo().log(
+            "Exporting registered domains succeeded for TLD %s, response was: %s", tld, resultMsg);
+      }
+    } catch (Throwable e) {
+      logger.atSevere().withCause(e).log(
+          "Error exporting registered domains for TLD %s to Drive, skipping...", tld);
+      return false;
+    }
+    return true;
+  }
+
+  protected static boolean exportToGcs(
+      String tld, String domains, String gcsBucket, int gcsBufferSize) {
+    GcsFilename filename = new GcsFilename(gcsBucket, tld + ".txt");
+    GcsUtils cloudStorage =
+        new GcsUtils(createGcsService(RetryParams.getDefaultInstance()), gcsBufferSize);
+    try (OutputStream gcsOutput = cloudStorage.openOutputStream(filename);
+        Writer osWriter = new OutputStreamWriter(gcsOutput, UTF_8)) {
+      osWriter.write(domains);
+    } catch (Throwable e) {
+      logger.atSevere().withCause(e).log(
+          "Error exporting registered domains for TLD %s to GCS, skipping...", tld);
+      return false;
+    }
+    return true;
   }
 
   static class ExportDomainListsMapper extends Mapper<DomainBase, String, String> {
@@ -122,9 +213,6 @@ public class ExportDomainListsAction implements Runnable {
     private static Supplier<DriveConnection> driveConnectionSupplier =
         Suppliers.memoize(() -> DaggerDriveModule_DriveComponent.create().driveConnection());
 
-    static final String REGISTERED_DOMAINS_FILENAME = "registered_domains.txt";
-    static final MediaType EXPORT_MIME_TYPE = MediaType.PLAIN_TEXT_UTF_8;
-
     private final String gcsBucket;
     private final int gcsBufferSize;
 
@@ -147,53 +235,21 @@ public class ExportDomainListsAction implements Runnable {
       driveConnection = driveConnectionSupplier.get();
     }
 
-    private void exportToDrive(String tld, String domains) {
-      verifyNotNull(driveConnection, "expecting non-null driveConnection");
-      try {
-        Registry registry = Registry.get(tld);
-        if (registry.getDriveFolderId() == null) {
-          logger.atInfo().log(
-              "Skipping registered domains export for TLD %s because Drive folder isn't specified",
-              tld);
-        } else {
-          String resultMsg =
-              driveConnection.createOrUpdateFile(
-                  REGISTERED_DOMAINS_FILENAME,
-                  EXPORT_MIME_TYPE,
-                  registry.getDriveFolderId(),
-                  domains.getBytes(UTF_8));
-          logger.atInfo().log(
-              "Exporting registered domains succeeded for TLD %s, response was: %s",
-              tld, resultMsg);
-        }
-      } catch (Throwable e) {
-        logger.atSevere().withCause(e).log(
-            "Error exporting registered domains for TLD %s to Drive", tld);
-      }
-      getContext().incrementCounter("domain lists written out to Drive");
-    }
-
-    private void exportToGcs(String tld, String domains) {
-      GcsFilename filename = new GcsFilename(gcsBucket, tld + ".txt");
-      GcsUtils cloudStorage =
-          new GcsUtils(createGcsService(RetryParams.getDefaultInstance()), gcsBufferSize);
-      try (OutputStream gcsOutput = cloudStorage.openOutputStream(filename);
-          Writer osWriter = new OutputStreamWriter(gcsOutput, UTF_8)) {
-        osWriter.write(domains);
-      } catch (IOException e) {
-        logger.atSevere().withCause(e).log(
-            "Error exporting registered domains for TLD %s to GCS.", tld);
-      }
-      getContext().incrementCounter("domain lists written out to GCS");
-    }
-
     @Override
     public void reduce(String tld, ReducerInput<String> fqdns) {
       ImmutableList<String> domains = ImmutableList.sortedCopyOf(() -> fqdns);
       String domainsList = Joiner.on('\n').join(domains);
       logger.atInfo().log("Exporting %d domains for TLD %s to GCS and Drive.", domains.size(), tld);
-      exportToGcs(tld, domainsList);
-      exportToDrive(tld, domainsList);
+      if (exportToGcs(tld, domainsList, gcsBucket, gcsBufferSize)) {
+        getContext().incrementCounter("domain lists successful written out to GCS");
+      } else {
+        getContext().incrementCounter("domain lists failed to write out to GCS");
+      }
+      if (exportToDrive(tld, domainsList, driveConnection)) {
+        getContext().incrementCounter("domain lists successfully written out to Drive");
+      } else {
+        getContext().incrementCounter("domain lists failed to write out to Drive");
+      }
     }
 
     @VisibleForTesting

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/EppResourceUtils.java

@@ -16,10 +16,11 @@ package google.registry.model;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
-import static google.registry.model.ofy.ObjectifyService.ofy;
+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.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
+import static google.registry.util.DateTimeUtils.START_OF_TIME;
 import static google.registry.util.DateTimeUtils.isAtOrAfter;
 import static google.registry.util.DateTimeUtils.isBeforeOrAt;
 import static google.registry.util.DateTimeUtils.latestOf;
@@ -42,10 +43,13 @@ import google.registry.model.index.ForeignKeyIndex;
 import google.registry.model.ofy.CommitLogManifest;
 import google.registry.model.ofy.CommitLogMutation;
 import google.registry.model.registry.Registry;
+import google.registry.model.reporting.HistoryEntry;
+import google.registry.model.reporting.HistoryEntryDao;
 import google.registry.model.transfer.DomainTransferData;
 import google.registry.model.transfer.TransferData;
 import google.registry.model.transfer.TransferStatus;
 import google.registry.persistence.VKey;
+import java.util.Comparator;
 import java.util.List;
 import java.util.Map.Entry;
 import java.util.Optional;
@@ -265,26 +269,43 @@ public final class EppResourceUtils {
    * Rewinds an {@link EppResource} object to a given point in time.
    *
    * <p>This method costs nothing if {@code resource} is already current. Otherwise it needs to
-   * perform a single asynchronous key fetch operation.
+   * perform a single fetch operation.
    *
    * <p><b>Warning:</b> A resource can only be rolled backwards in time, not forwards; therefore
    * {@code resource} should be whatever's currently in Datastore.
    *
-   * <p><b>Warning:</b> Revisions are granular to 24-hour periods. It's recommended that
-   * {@code timestamp} be set to midnight. Otherwise you must take into consideration that under
-   * certain circumstances, a resource might be restored to a revision on the previous day, even if
-   * there were revisions made earlier on the same date as {@code timestamp}; however, a resource
-   * will never be restored to a revision occurring after {@code timestamp}. This behavior is due to
-   * the way {@link google.registry.model.translators.CommitLogRevisionsTranslatorFactory
+   * <p><b>Warning:</b> In Datastore, revisions are granular to 24-hour periods. It's recommended
+   * that {@code timestamp} be set to midnight. If you don't use midnight, you must take into
+   * consideration that under certain circumstances, a resource might be restored to a revision on
+   * the previous day, even if there were revisions made earlier on the same date as {@code
+   * timestamp}; however, a resource will never be restored to a revision occurring after {@code
+   * timestamp}. This behavior is due to the way {@link
+   * google.registry.model.translators.CommitLogRevisionsTranslatorFactory
    * CommitLogRevisionsTranslatorFactory} manages the {@link EppResource#revisions} field. Please
    * note however that the creation and deletion times of a resource are granular to the
    * millisecond.
    *
+   * <p>Example: a resource in Datastore has three revisions A, B, and C
+   *
+   * <ul>
+   *   <li>A: Day 0, 1pm
+   *   <li>B: Day 1, 1pm
+   *   <li>C: Day 1, 3pm
+   * </ul>
+   *
+   * <p>If one requests the resource as of day 1 at 2pm, we will return revision A because as far as
+   * the commit logs are concerned, revision C completely overwrites the existence of revision B.
+   *
+   * <p>When using the SQL backend (post-Registry-3.0-migration) this restriction goes away and
+   * objects can be restored to any revision.
+   *
+   * <p>TODO(b/177567432): Once Datastore is completely removed, remove the Result wrapping.
+   *
    * @return an asynchronous operation returning resource at {@code timestamp} or {@code null} if
    *     resource is deleted or not yet created
    */
-  public static <T extends EppResource>
-      Result<T> loadAtPointInTime(final T resource, final DateTime timestamp) {
+  public static <T extends EppResource> Result<T> loadAtPointInTime(
+      final T resource, final DateTime timestamp) {
     // If we're before the resource creation time, don't try to find a "most recent revision".
     if (timestamp.isBefore(resource.getCreationTime())) {
       return new ResultNow<>(null);
@@ -299,7 +320,8 @@ public final class EppResourceUtils {
             : loadMostRecentRevisionAtTime(resource, timestamp);
     return () -> {
       T loadedResource = loadResult.now();
-      return (loadedResource == null) ? null
+      return (loadedResource == null)
+          ? null
           : (isActive(loadedResource, timestamp)
               ? cloneProjectedAtTime(loadedResource, timestamp)
               : null);
@@ -307,26 +329,43 @@ public final class EppResourceUtils {
   }
 
   /**
-   * Returns an asynchronous result holding the most recent Datastore revision of a given
-   * EppResource before or at the provided timestamp using the EppResource revisions map, falling
-   * back to using the earliest revision or the resource as-is if there are no revisions.
+   * Returns an asynchronous result holding the most recent revision of a given EppResource before
+   * or at the provided timestamp, falling back to using the resource as-is if there are no
+   * revisions.
    *
    * @see #loadAtPointInTime(EppResource, DateTime)
    */
   private static <T extends EppResource> Result<T> loadMostRecentRevisionAtTime(
       final T resource, final DateTime timestamp) {
+    if (tm().isOfy()) {
+      return loadMostRecentRevisionAtTimeDatastore(resource, timestamp);
+    } else {
+      return loadMostRecentRevisionAtTimeSql(resource, timestamp);
+    }
+  }
+
+  /**
+   * Returns an asynchronous result holding the most recent Datastore revision of a given
+   * EppResource before or at the provided timestamp using the EppResource revisions map, falling
+   * back to using the resource as-is if there are no revisions.
+   *
+   * @see #loadAtPointInTime(EppResource, DateTime)
+   */
+  private static <T extends EppResource> Result<T> loadMostRecentRevisionAtTimeDatastore(
+      final T resource, final DateTime timestamp) {
     final Key<T> resourceKey = Key.create(resource);
-    final Key<CommitLogManifest> revision = findMostRecentRevisionAtTime(resource, timestamp);
+    final Key<CommitLogManifest> revision =
+        findMostRecentDatastoreRevisionAtTime(resource, timestamp);
     if (revision == null) {
       logger.atSevere().log("No revision found for %s, falling back to resource.", resourceKey);
       return new ResultNow<>(resource);
     }
     final Result<CommitLogMutation> mutationResult =
-        ofy().load().key(CommitLogMutation.createKey(revision, resourceKey));
+        auditedOfy().load().key(CommitLogMutation.createKey(revision, resourceKey));
     return () -> {
       CommitLogMutation mutation = mutationResult.now();
       if (mutation != null) {
-        return ofy().load().fromEntity(mutation.getEntity());
+        return auditedOfy().load().fromEntity(mutation.getEntity());
       }
       logger.atSevere().log(
           "Couldn't load mutation for revision at %s for %s, falling back to resource."
@@ -336,9 +375,37 @@ public final class EppResourceUtils {
     };
   }
 
+  /**
+   * Returns an asynchronous result holding the most recent SQL revision of a given EppResource
+   * before or at the provided timestamp using *History objects, falling back to using the resource
+   * as-is if there are no revisions.
+   *
+   * @see #loadAtPointInTime(EppResource, DateTime)
+   */
+  @SuppressWarnings("unchecked")
+  private static <T extends EppResource> Result<T> loadMostRecentRevisionAtTimeSql(
+      T resource, DateTime timestamp) {
+    T resourceAtPointInTime =
+        (T)
+            HistoryEntryDao.loadHistoryObjectsForResource(
+                    resource.createVKey(), START_OF_TIME, timestamp)
+                .stream()
+                .max(Comparator.comparing(HistoryEntry::getModificationTime))
+                .flatMap(HistoryEntry::getResourceAtPointInTime)
+                .orElse(null);
+    if (resourceAtPointInTime == null) {
+      logger.atSevere().log(
+          "Couldn't load resource at % for key %s, falling back to resource %s.",
+          timestamp, resource.createVKey(), resource);
+      return new ResultNow<>(resource);
+    }
+    return new ResultNow<>(resourceAtPointInTime);
+  }
+
   @Nullable
-  private static <T extends EppResource> Key<CommitLogManifest>
-      findMostRecentRevisionAtTime(final T resource, final DateTime timestamp) {
+  private static <T extends EppResource>
+      Key<CommitLogManifest> findMostRecentDatastoreRevisionAtTime(
+          final T resource, final DateTime timestamp) {
     final Key<T> resourceKey = Key.create(resource);
     Entry<?, Key<CommitLogManifest>> revision = resource.getRevisions().floorEntry(timestamp);
     if (revision != null) {
@@ -366,27 +433,26 @@ public final class EppResourceUtils {
    *
    * @param key the referent key
    * @param now the logical time of the check
-   * @param limit the maximum number of returned keys
+   * @param limit the maximum number of returned keys, unlimited if null
    */
   public static ImmutableSet<VKey<DomainBase>> getLinkedDomainKeys(
-      VKey<? extends EppResource> key, DateTime now, int limit) {
+      VKey<? extends EppResource> key, DateTime now, @Nullable Integer limit) {
     checkArgument(
         key.getKind().equals(ContactResource.class) || key.getKind().equals(HostResource.class),
         "key must be either VKey<ContactResource> or VKey<HostResource>, but it is %s",
         key);
     boolean isContactKey = key.getKind().equals(ContactResource.class);
     if (tm().isOfy()) {
-      return ofy()
-          .load()
-          .type(DomainBase.class)
-          .filter(isContactKey ? "allContacts.contact" : "nsHosts", key.getOfyKey())
-          .filter("deletionTime >", now)
-          .limit(limit)
-          .keys()
-          .list()
-          .stream()
-          .map(DomainBase::createVKey)
-          .collect(toImmutableSet());
+      com.googlecode.objectify.cmd.Query<DomainBase> query =
+          auditedOfy()
+              .load()
+              .type(DomainBase.class)
+              .filter(isContactKey ? "allContacts.contact" : "nsHosts", key.getOfyKey())
+              .filter("deletionTime >", now);
+      if (limit != null) {
+        query.limit(limit);
+      }
+      return query.keys().list().stream().map(DomainBase::createVKey).collect(toImmutableSet());
     } else {
       return tm().transact(
               () -> {
@@ -405,11 +471,13 @@ public final class EppResourceUtils {
                           .setParameter("fkRepoId", key.getSqlKey())
                           .setParameter("now", now.toDate());
                 }
+                if (limit != null) {
+                  query.setMaxResults(limit);
+                }
                 @SuppressWarnings("unchecked")
                 ImmutableSet<VKey<DomainBase>> domainBaseKeySet =
                     (ImmutableSet<VKey<DomainBase>>)
                         query
-                            .setMaxResults(limit)
                             .getResultStream()
                             .map(
                                 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,252 @@
+// 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.
+   */
+  @VisibleForTesting
+  public 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

@@ -18,6 +18,7 @@ import static google.registry.persistence.transaction.TransactionManagerFactory.
 
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.EntitySubclass;
+import google.registry.model.EppResource;
 import google.registry.model.ImmutableObject;
 import google.registry.model.contact.ContactHistory.ContactHistoryId;
 import google.registry.model.reporting.HistoryEntry;
@@ -107,6 +108,11 @@ public class ContactHistory extends HistoryEntry implements SqlEntity {
     return (VKey<ContactHistory>) createVKey(Key.create(this));
   }
 
+  @Override
+  public Optional<? extends EppResource> getResourceAtPointInTime() {
+    return getContactBase();
+  }
+
   @PostLoad
   void postLoad() {
     // Normally Hibernate would see that the contact fields are all null and would fill contactBase
@@ -115,7 +121,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

@@ -21,9 +21,11 @@ import static google.registry.util.CollectionUtils.nullToEmptyImmutableCopy;
 import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.EntitySubclass;
+import google.registry.model.EppResource;
 import google.registry.model.ImmutableObject;
 import google.registry.model.domain.DomainHistory.DomainHistoryId;
 import google.registry.model.domain.GracePeriod.GracePeriodHistory;
+import google.registry.model.domain.secdns.DelegationSignerData;
 import google.registry.model.domain.secdns.DomainDsDataHistory;
 import google.registry.model.host.HostResource;
 import google.registry.model.reporting.DomainTransactionRecord;
@@ -248,17 +250,31 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
     return (VKey<DomainHistory>) createVKey(Key.create(this));
   }
 
+  @Override
+  public Optional<? extends EppResource> getResourceAtPointInTime() {
+    return getDomainContent();
+  }
+
   @PostLoad
   void postLoad() {
     if (domainContent != null) {
       domainContent.nsHosts = nullToEmptyImmutableCopy(nsHosts);
+      domainContent.gracePeriods =
+          gracePeriodHistories.stream()
+              .map(GracePeriod::createFromHistory)
+              .collect(toImmutableSet());
+      domainContent.dsData =
+          dsDataHistories.stream().map(DelegationSignerData::create).collect(toImmutableSet());
       // Normally Hibernate would see that the domain fields are all null and would fill
       // domainContent with a null object. Unfortunately, the updateTimestamp is never null in SQL.
       if (domainContent.getDomainName() == null) {
         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/domain/GracePeriod.java

@@ -115,6 +115,17 @@ public class GracePeriod extends GracePeriodBase implements DatastoreAndSqlEntit
         type, domainRepoId, expirationTime, clientId, billingEventOneTime, null, gracePeriodId);
   }
 
+  public static GracePeriod createFromHistory(GracePeriodHistory history) {
+    return createInternal(
+        history.type,
+        history.domainRepoId,
+        history.expirationTime,
+        history.clientId,
+        history.billingEventOneTime == null ? null : history.billingEventOneTime.createVKey(),
+        history.billingEventRecurring == null ? null : history.billingEventRecurring.createVKey(),
+        history.gracePeriodId);
+  }
+
   /** Creates a GracePeriod for a Recurring billing event. */
   public static GracePeriod createForRecurring(
       GracePeriodStatus type,

core/src/main/java/google/registry/model/domain/secdns/DelegationSignerData.java

@@ -114,6 +114,15 @@ public class DelegationSignerData extends DomainDsDataBase {
     return create(keyTag, algorithm, digestType, DatatypeConverter.parseHexBinary(digestAsHex));
   }
 
+  public static DelegationSignerData create(DomainDsDataHistory history) {
+    return create(
+        history.keyTag,
+        history.algorithm,
+        history.digestType,
+        history.digest,
+        history.domainRepoId);
+  }
+
   /** Class to represent the composite primary key of {@link DelegationSignerData} entity. */
   static class DomainDsDataId extends ImmutableObject implements Serializable {
 

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

@@ -18,6 +18,7 @@ import static google.registry.persistence.transaction.TransactionManagerFactory.
 
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.EntitySubclass;
+import google.registry.model.EppResource;
 import google.registry.model.ImmutableObject;
 import google.registry.model.host.HostHistory.HostHistoryId;
 import google.registry.model.reporting.HistoryEntry;
@@ -108,6 +109,11 @@ public class HostHistory extends HistoryEntry implements SqlEntity {
     return (VKey<HostHistory>) createVKey(Key.create(this));
   }
 
+  @Override
+  public Optional<? extends EppResource> getResourceAtPointInTime() {
+    return getHostBase();
+  }
+
   @PostLoad
   void postLoad() {
     // Normally Hibernate would see that the host fields are all null and would fill hostBase
@@ -116,7 +122,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/ofy/DatastoreTransactionManager.java

@@ -480,7 +480,13 @@ public class DatastoreTransactionManager implements TransactionManager {
 
     @Override
     public long count() {
-      return buildQuery().count();
+      // Objectify provides a count() function, but unfortunately that doesn't work if there are
+      // more than 1000 (the default response page size?) entries in the result set.  We also use
+      // chunkAll() here as it provides a nice performance boost.
+      //
+      // There is some information on this issue on SO, see:
+      // https://stackoverflow.com/questions/751124/how-does-one-get-a-count-of-rows-in-a-datastore-model-in-google-app-engine
+      return Iterables.size(buildQuery().chunkAll().keys());
     }
 
     @Override

core/src/main/java/google/registry/model/ofy/ObjectifyService.java

@@ -51,7 +51,7 @@ import google.registry.model.translators.VKeyTranslatorFactory;
 import java.util.concurrent.atomic.AtomicLong;
 
 /**
- * An instance of Ofy, obtained via {@code #ofy()}, should be used to access all persistable
+ * An instance of Ofy, obtained via {@code #auditedOfy()}, should be used to access all persistable
  * objects. The class contains a static initializer to call factory().register(...) on all
  * persistable objects in this package.
  */

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;
 
 /**
@@ -289,6 +290,17 @@ public class HistoryEntry extends ImmutableObject implements Buildable, Datastor
     return nullToEmptyImmutableCopy(domainTransactionRecords);
   }
 
+  /**
+   * Throws an error when attempting to retrieve the EppResource at this point in time.
+   *
+   * <p>Subclasses must override this to return the resource; it is non-abstract for legacy reasons
+   * and objects created prior to the Registry 3.0 migration.
+   */
+  public Optional<? extends EppResource> getResourceAtPointInTime() {
+    throw new UnsupportedOperationException(
+        "Raw HistoryEntry objects do not store the resource at that point in time.");
+  }
+
   /** This method exists solely to satisfy Hibernate. Use the {@link Builder} instead. */
   @SuppressWarnings("UnusedMethod")
   private void setPeriod(Period period) {
@@ -435,7 +447,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/reporting/ReportingModule.java

@@ -14,6 +14,7 @@
 
 package google.registry.reporting;
 
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.request.RequestParameters.extractOptionalParameter;
 import static google.registry.request.RequestParameters.extractRequiredParameter;
 
@@ -55,6 +56,9 @@ public class ReportingModule {
   /** The request parameter specifying the jobId for a running Dataflow pipeline. */
   public static final String PARAM_JOB_ID = "jobId";
 
+  /** The request parameter for specifying which database reporting actions should read from. */
+  public static final String DATABASE = "database";
+
   /** Provides the Cloud Dataflow jobId for a pipeline. */
   @Provides
   @Parameter(PARAM_JOB_ID)
@@ -62,6 +66,14 @@ public class ReportingModule {
     return extractRequiredParameter(req, PARAM_JOB_ID);
   }
 
+  /** Provides the database for the pipeline to read from. */
+  @Provides
+  @Parameter(DATABASE)
+  static String provideDatabase(HttpServletRequest req) {
+    Optional<String> optionalDatabase = extractOptionalParameter(req, DATABASE);
+    return optionalDatabase.orElse(tm().isOfy() ? "DATASTORE" : "CLOUD_SQL");
+  }
+
   /** Extracts an optional YearMonth in yyyy-MM format from the request. */
   @Provides
   @Parameter(PARAM_YEAR_MONTH)

core/src/main/java/google/registry/reporting/spec11/GenerateSpec11ReportAction.java

@@ -15,6 +15,8 @@
 package google.registry.reporting.spec11;
 
 import static google.registry.beam.BeamUtils.createJobName;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.reporting.ReportingModule.DATABASE;
 import static google.registry.reporting.ReportingModule.PARAM_DATE;
 import static google.registry.reporting.ReportingUtils.enqueueBeamReportingTask;
 import static google.registry.request.Action.Method.POST;
@@ -69,6 +71,7 @@ public class GenerateSpec11ReportAction implements Runnable {
   private final Clock clock;
   private final Response response;
   private final Dataflow dataflow;
+  private final String database;
 
   @Inject
   GenerateSpec11ReportAction(
@@ -78,15 +81,20 @@ public class GenerateSpec11ReportAction implements Runnable {
       @Config("reportingBucketUrl") String reportingBucketUrl,
       @Key("safeBrowsingAPIKey") String apiKey,
       @Parameter(PARAM_DATE) LocalDate date,
+      @Parameter(DATABASE) String database,
       Clock clock,
       Response response,
       Dataflow dataflow) {
     this.projectId = projectId;
     this.jobRegion = jobRegion;
     this.stagingBucketUrl = stagingBucketUrl;
+    if (tm().isOfy() && database.equals("CLOUD_SQL")) {
+      reportingBucketUrl = reportingBucketUrl.concat("-sql");
+    }
     this.reportingBucketUrl = reportingBucketUrl;
     this.apiKey = apiKey;
     this.date = date;
+    this.database = database;
     this.clock = clock;
     this.response = response;
     this.dataflow = dataflow;
@@ -105,6 +113,8 @@ public class GenerateSpec11ReportAction implements Runnable {
                   ImmutableMap.of(
                       "safeBrowsingApiKey",
                       apiKey,
+                      "database",
+                      database,
                       ReportingModule.PARAM_DATE,
                       date.toString(),
                       "reportingBucketUrl",

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/CreateDomainCommand.java

@@ -80,8 +80,7 @@ final class CreateDomainCommand extends CreateOrUpdateDomainCommand
       }
 
       setSoyTemplate(DomainCreateSoyInfo.getInstance(), DomainCreateSoyInfo.DOMAINCREATE);
-      addSoyRecord(
-          clientId,
+      SoyMapData soyMapData =
           new SoyMapData(
               "domain", domain,
               "period", period,
@@ -92,7 +91,12 @@ final class CreateDomainCommand extends CreateOrUpdateDomainCommand
               "password", password,
               "currency", currency,
               "price", cost,
-              "dsRecords", DsRecord.convertToSoy(dsRecords)));
+              "dsRecords", DsRecord.convertToSoy(dsRecords),
+              "reason", reason);
+      if (requestedByRegistrar != null) {
+        soyMapData.put("requestedByRegistrar", requestedByRegistrar.toString());
+      }
+      addSoyRecord(clientId, soyMapData);
     }
   }
 }

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

@@ -78,6 +78,17 @@ abstract class CreateOrUpdateDomainCommand extends MutatingEppToolCommand {
       converter = DsRecord.Converter.class)
   List<DsRecord> dsRecords = new ArrayList<>();
 
+  @Parameter(
+      names = {"--reason"},
+      description = "Reason for the change.")
+  String reason;
+
+  @Parameter(
+      names = {"--registrar_request"},
+      description = "Whether the change was requested by a registrar.",
+      arity = 1)
+  Boolean requestedByRegistrar;
+
   Set<String> domains;
 
   @Override

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/UpdateDomainCommand.java

@@ -315,10 +315,14 @@ final class UpdateDomainCommand extends CreateOrUpdateDomainCommand {
               "secdns", secDns,
               "addDsRecords", DsRecord.convertToSoy(addDsRecords),
               "removeDsRecords", DsRecord.convertToSoy(removeDsRecords),
-              "removeAllDsRecords", clearDsRecords);
+              "removeAllDsRecords", clearDsRecords,
+              "reason", reason);
       if (autorenews != null) {
         soyMapData.put("autorenews", autorenews.toString());
       }
+      if (requestedByRegistrar != null) {
+        soyMapData.put("requestedByRegistrar", requestedByRegistrar.toString());
+      }
       addSoyRecord(clientId, soyMapData);
     }
 

core/src/main/java/google/registry/tools/javascrap/CreateSyntheticHistoryEntriesAction.java

@@ -14,7 +14,7 @@
 
 package google.registry.tools.javascrap;
 
-import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.model.ofy.ObjectifyService.auditedOfy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.appengine.tools.mapreduce.Mapper;
@@ -111,7 +111,7 @@ public class CreateSyntheticHistoryEntriesAction implements Runnable {
     public final void map(final Key<EppResource> resourceKey) {
       tm().transact(
               () -> {
-                EppResource eppResource = ofy().load().key(resourceKey).now();
+                EppResource eppResource = auditedOfy().load().key(resourceKey).now();
                 tm().put(
                         HistoryEntry.createBuilderForResource(eppResource)
                             .setClientId(registryAdminRegistrarId)

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/java/google/registry/tools/server/KillAllCommitLogsAction.java

@@ -49,6 +49,8 @@ import javax.inject.Inject;
     path = "/_dr/task/killAllCommitLogs",
     method = POST,
     auth = Auth.AUTH_INTERNAL_OR_ADMIN)
+// No longer needed in SQL. Subject to future removal.
+@Deprecated
 public class KillAllCommitLogsAction implements Runnable {
 
   @Inject MapreduceRunner mrRunner;

core/src/main/java/google/registry/tools/server/RefreshDnsForAllDomainsAction.java

@@ -18,6 +18,8 @@ import static com.google.common.base.Preconditions.checkArgument;
 import static google.registry.mapreduce.inputs.EppResourceInputs.createEntityInput;
 import static google.registry.model.EppResourceUtils.isActive;
 import static google.registry.model.registry.Registries.assertTldsExist;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.request.RequestParameters.PARAM_TLDS;
 
 import com.google.appengine.tools.mapreduce.Mapper;
@@ -32,11 +34,12 @@ import google.registry.request.Action;
 import google.registry.request.Parameter;
 import google.registry.request.Response;
 import google.registry.request.auth.Auth;
+import google.registry.util.Clock;
 import google.registry.util.NonFinalForTesting;
 import java.util.Random;
 import javax.inject.Inject;
+import org.apache.http.HttpStatus;
 import org.joda.time.DateTime;
-import org.joda.time.DateTimeZone;
 import org.joda.time.Duration;
 
 /**
@@ -74,6 +77,8 @@ public class RefreshDnsForAllDomainsAction implements Runnable {
   @Parameter("smearMinutes")
   int smearMinutes;
 
+  @Inject DnsQueue dnsQueue;
+  @Inject Clock clock;
   @Inject Random random;
 
   @Inject
@@ -83,14 +88,41 @@ public class RefreshDnsForAllDomainsAction implements Runnable {
   public void run() {
     assertTldsExist(tlds);
     checkArgument(smearMinutes > 0, "Must specify a positive number of smear minutes");
-    mrRunner
-        .setJobName("Refresh DNS for all domains")
-        .setModuleName("tools")
-        .setDefaultMapShards(10)
-        .runMapOnly(
-            new RefreshDnsForAllDomainsActionMapper(tlds, smearMinutes, random),
-            ImmutableList.of(createEntityInput(DomainBase.class)))
-        .sendLinkToMapreduceConsole(response);
+    if (tm().isOfy()) {
+      mrRunner
+          .setJobName("Refresh DNS for all domains")
+          .setModuleName("tools")
+          .setDefaultMapShards(10)
+          .runMapOnly(
+              new RefreshDnsForAllDomainsActionMapper(tlds, smearMinutes, random, clock.nowUtc()),
+              ImmutableList.of(createEntityInput(DomainBase.class)))
+          .sendLinkToMapreduceConsole(response);
+    } else {
+      tm().transact(
+              () ->
+                  jpaTm()
+                      .query(
+                          "SELECT fullyQualifiedDomainName FROM Domain "
+                              + "WHERE tld IN (:tlds) "
+                              + "AND deletionTime > :now",
+                          String.class)
+                      .setParameter("tlds", tlds)
+                      .setParameter("now", clock.nowUtc())
+                      .getResultStream()
+                      .forEach(
+                          domainName -> {
+                            try {
+                              // Smear the task execution time over the next N minutes.
+                              dnsQueue.addDomainRefreshTask(
+                                  domainName,
+                                  Duration.standardMinutes(random.nextInt(smearMinutes)));
+                            } catch (Throwable t) {
+                              logger.atSevere().withCause(t).log(
+                                  "Error while enqueuing DNS refresh for domain %s", domainName);
+                              response.setStatus(HttpStatus.SC_INTERNAL_SERVER_ERROR);
+                            }
+                          }));
+    }
   }
 
   /** Mapper to refresh DNS for all active domain resources. */
@@ -104,19 +136,21 @@ public class RefreshDnsForAllDomainsAction implements Runnable {
     private final ImmutableSet<String> tlds;
     private final int smearMinutes;
     private final Random random;
+    private final DateTime now;
 
     RefreshDnsForAllDomainsActionMapper(
-        ImmutableSet<String> tlds, int smearMinutes, Random random) {
+        ImmutableSet<String> tlds, int smearMinutes, Random random, DateTime now) {
       this.tlds = tlds;
       this.smearMinutes = smearMinutes;
       this.random = random;
+      this.now = now;
     }
 
     @Override
     public void map(final DomainBase domain) {
       String domainName = domain.getDomainName();
       if (tlds.contains(domain.getTld())) {
-        if (isActive(domain, DateTime.now(DateTimeZone.UTC))) {
+        if (isActive(domain, now)) {
           try {
             // Smear the task execution time over the next N minutes.
             dnsQueue.addDomainRefreshTask(
@@ -124,7 +158,7 @@ public class RefreshDnsForAllDomainsAction implements Runnable {
             getContext().incrementCounter("active domains refreshed");
           } catch (Throwable t) {
             logger.atSevere().withCause(t).log(
-                "Error while refreshing DNS for domain %s", domainName);
+                "Error while enqueuing DNS refresh for domain %s", domainName);
             getContext().incrementCounter("active domains errored");
           }
         } else {

core/src/main/java/google/registry/tools/server/ResaveAllHistoryEntriesAction.java

@@ -42,6 +42,8 @@ import javax.inject.Inject;
     service = Action.Service.TOOLS,
     path = "/_dr/task/resaveAllHistoryEntries",
     auth = Auth.AUTH_INTERNAL_OR_ADMIN)
+// No longer needed in SQL. Subject to future removal.
+@Deprecated
 public class ResaveAllHistoryEntriesAction implements Runnable {
 
   @Inject MapreduceRunner mrRunner;

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/main/resources/google/registry/beam/spec11_pipeline_metadata.json

@@ -61,6 +61,14 @@
       "regexes": [
         "^gs:\\/\\/[^\\n\\r]+$"
       ]
+    },
+    {
+      "name": "database",
+      "label": "Database to read from.",
+      "helpText": "DATASTORE or CLOUD_SQL.",
+      "regexes": [
+        "^DATASTORE|CLOUD_SQL$"
+      ]
     }
   ]
 }

core/src/main/resources/google/registry/tools/soy/DomainCreate.soy

@@ -27,6 +27,8 @@
   {@param? currency: string}
   {@param? price: string}
   {@param dsRecords: list<[keyTag:int, alg:int, digestType:int, digest:string]>}
+  {@param? reason: string}
+  {@param? requestedByRegistrar: string}
 
   <?xml version="1.0" encoding="UTF-8" standalone="yes"?>
   <epp xmlns="urn:ietf:params:xml:ns:epp-1.0">
@@ -54,7 +56,7 @@
           </domain:authInfo>
         </domain:create>
       </create>
-      {if length($dsRecords) > 0 or $price != null}
+      {if length($dsRecords) > 0 or $price != null or $reason or $requestedByRegistrar}
         <extension>
           {if $price != null}
             <fee:create xmlns:fee="urn:ietf:params:xml:ns:fee-0.12">
@@ -74,6 +76,16 @@
               {/for}
             </secDNS:create>
           {/if}
+          {if $reason or $requestedByRegistrar}
+            <metadata:metadata xmlns:metadata="urn:google:params:xml:ns:metadata-1.0">
+            {if $reason}
+              <metadata:reason>{$reason}</metadata:reason>
+            {/if}
+            {if $requestedByRegistrar}
+              <metadata:requestedByRegistrar>{$requestedByRegistrar}</metadata:requestedByRegistrar>
+            {/if}
+            </metadata:metadata>
+          {/if}
         </extension>
       {/if}
       <clTRID>RegistryTool</clTRID>

core/src/main/resources/google/registry/tools/soy/DomainUpdate.soy

@@ -36,6 +36,8 @@
   {@param removeDsRecords: list<[keyTag:int, alg:int, digestType:int, digest:string]>}
   {@param removeAllDsRecords: bool}
   {@param? autorenews: string}
+  {@param? reason: string}
+  {@param? requestedByRegistrar: string}
 
   <?xml version="1.0" encoding="UTF-8" standalone="yes"?>
   <epp xmlns="urn:ietf:params:xml:ns:epp-1.0">
@@ -97,7 +99,7 @@
           {/if}
         </domain:update>
       </update>
-      {if $secdns or $autorenews}
+      {if $secdns or $autorenews or $reason or $requestedByRegistrar}
         <extension>
           {if $secdns}
             <secDNS:update xmlns:secDNS="urn:ietf:params:xml:ns:secDNS-1.1">
@@ -137,6 +139,16 @@
               <superuser:autorenews>{$autorenews}</superuser:autorenews>
             </superuser:domainUpdate>
           {/if}
+          {if $reason or $requestedByRegistrar}
+            <metadata:metadata xmlns:metadata="urn:google:params:xml:ns:metadata-1.0">
+            {if $reason}
+              <metadata:reason>{$reason}</metadata:reason>
+            {/if}
+            {if $requestedByRegistrar}
+              <metadata:requestedByRegistrar>{$requestedByRegistrar}</metadata:requestedByRegistrar>
+            {/if}
+            </metadata:metadata>
+          {/if}
         </extension>
       {/if}
       <clTRID>RegistryTool</clTRID>

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

@@ -22,7 +22,7 @@ import static com.google.common.collect.Iterables.concat;
 import static com.google.common.collect.Lists.partition;
 import static google.registry.backup.BackupUtils.serializeEntity;
 import static google.registry.model.ofy.CommitLogBucket.getBucketKey;
-import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.model.ofy.ObjectifyService.auditedOfy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
 import static google.registry.util.DateTimeUtils.isAtOrAfter;
@@ -80,7 +80,7 @@ public final class CommitLogExports {
   public static CommitLogCheckpoint computeCheckpoint(Clock clock) {
     CommitLogCheckpointStrategy strategy = new CommitLogCheckpointStrategy();
     strategy.clock = clock;
-    strategy.ofy = ofy();
+    strategy.ofy = auditedOfy();
 
     CommitLogCheckpoint checkpoint = strategy.computeCheckpoint();
     tm().transact(
@@ -90,7 +90,7 @@ public final class CommitLogExports {
                   checkpoint.getCheckpointTime().isAfter(lastWrittenTime),
                   "Newer checkpoint already written at time: %s",
                   lastWrittenTime);
-              ofy()
+              auditedOfy()
                   .saveWithoutBackup()
                   .entities(
                       checkpoint, CommitLogCheckpointRoot.create(checkpoint.getCheckpointTime()));
@@ -135,17 +135,17 @@ public final class CommitLogExports {
       // asynchronously load the entities for the next one.
       List<List<Key<CommitLogManifest>>> keyChunks = partition(sortedKeys, EXPORT_DIFF_BATCH_SIZE);
       // Objectify's map return type is asynchronous. Calling .values() will block until it loads.
-      Map<?, CommitLogManifest> nextChunkToExport = ofy().load().keys(keyChunks.get(0));
+      Map<?, CommitLogManifest> nextChunkToExport = auditedOfy().load().keys(keyChunks.get(0));
       for (int i = 0; i < keyChunks.size(); i++) {
         // Force the async load to finish.
         Collection<CommitLogManifest> chunkValues = nextChunkToExport.values();
         // Since there is no hard bound on how much data this might be, take care not to let the
         // Objectify session cache fill up and potentially run out of memory. This is the only safe
         // point to do this since at this point there is no async load in progress.
-        ofy().clearSessionCache();
+        auditedOfy().clearSessionCache();
         // Kick off the next async load, which can happen in parallel to the current GCS export.
         if (i + 1 < keyChunks.size()) {
-          nextChunkToExport = ofy().load().keys(keyChunks.get(i + 1));
+          nextChunkToExport = auditedOfy().load().keys(keyChunks.get(i + 1));
         }
         exportChunk(commitLogStream, chunkValues);
       }
@@ -205,7 +205,7 @@ public final class CommitLogExports {
       return ImmutableSet.of();
     }
     Key<CommitLogBucket> bucketKey = getBucketKey(bucketNum);
-    return ofy()
+    return auditedOfy()
         .load()
         .type(CommitLogManifest.class)
         .ancestor(bucketKey)
@@ -222,7 +222,7 @@ public final class CommitLogExports {
         new ImmutableList.Builder<>();
     for (CommitLogManifest manifest : chunk) {
       entities.add(ImmutableList.of(manifest));
-      entities.add(ofy().load().type(CommitLogMutation.class).ancestor(manifest));
+      entities.add(auditedOfy().load().type(CommitLogMutation.class).ancestor(manifest));
     }
     for (ImmutableObject entity : concat(entities.build())) {
       serializeEntity(entity, gcsStream);

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,13 @@ 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.mapreduce.MapreduceTestCase;
 import java.util.ArrayList;
 import java.util.List;
@@ -62,17 +67,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 +93,6 @@ public class ExpandRecurringBillingEventsActionTest
 
   @BeforeEach
   void beforeEach() {
-    inject.setStaticField(Ofy.class, "clock", clock);
     action = new ExpandRecurringBillingEventsAction();
     action.mrRunner = makeDefaultRunner();
     action.clock = clock;
@@ -111,27 +123,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 +167,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 +182,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 +226,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 +241,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
+  @TestOfyAndSql
   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 +311,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 +320,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 +488,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 +516,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
+  @TestOfyAndSql
   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 +670,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 +683,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 +695,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/batch/RefreshDnsOnHostRenameActionTest.java

@@ -17,6 +17,8 @@ package google.registry.batch;
 import static com.google.appengine.api.taskqueue.QueueFactory.getQueue;
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth8.assertThat;
+import static google.registry.batch.AsyncTaskEnqueuer.PARAM_HOST_KEY;
+import static google.registry.batch.AsyncTaskEnqueuer.PARAM_REQUESTED_TIME;
 import static google.registry.batch.AsyncTaskEnqueuer.QUEUE_ASYNC_HOST_RENAME;
 import static google.registry.batch.AsyncTaskMetrics.OperationType.DNS_REFRESH;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
@@ -45,6 +47,7 @@ import static org.mockito.Mockito.when;
 import com.googlecode.objectify.Key;
 import google.registry.batch.AsyncTaskMetrics.OperationResult;
 import google.registry.batch.RefreshDnsOnHostRenameAction.RefreshDnsOnHostRenameReducer;
+import google.registry.dns.DnsQueue;
 import google.registry.model.host.HostResource;
 import google.registry.model.server.Lock;
 import google.registry.testing.DualDatabaseTest;
@@ -55,6 +58,7 @@ import google.registry.testing.InjectExtension;
 import google.registry.testing.TaskQueueHelper.TaskMatcher;
 import google.registry.testing.TestOfyAndSql;
 import google.registry.testing.TestOfyOnly;
+import google.registry.testing.TestSqlOnly;
 import google.registry.testing.mapreduce.MapreduceTestCase;
 import google.registry.util.AppEngineServiceUtils;
 import google.registry.util.RequestStatusChecker;
@@ -62,6 +66,7 @@ import google.registry.util.Retrier;
 import google.registry.util.Sleeper;
 import google.registry.util.SystemSleeper;
 import java.util.Optional;
+import org.apache.http.HttpStatus;
 import org.joda.time.DateTime;
 import org.joda.time.Duration;
 import org.junit.jupiter.api.BeforeEach;
@@ -94,6 +99,7 @@ public class RefreshDnsOnHostRenameActionTest
     action.clock = clock;
     action.mrRunner = makeDefaultRunner();
     action.pullQueue = getQueue(QUEUE_ASYNC_HOST_RENAME);
+    action.dnsQueue = DnsQueue.createForTesting(clock);
     action.requestStatusChecker = requestStatusChecker;
     action.response = fakeResponse;
     action.retrier = new Retrier(new FakeSleeper(clock), 1);
@@ -102,7 +108,7 @@ public class RefreshDnsOnHostRenameActionTest
         .thenThrow(new AssertionError("Should not be called"));
   }
 
-  private void runMapreduce() throws Exception {
+  private void runAction() throws Exception {
     clock.advanceOneMilli();
     // Use hard sleeps to ensure that the tasks are enqueued properly and will be leased.
     Sleeper sleeper = new SystemSleeper();
@@ -123,10 +129,32 @@ public class RefreshDnsOnHostRenameActionTest
     tm().clearSessionCache();
   }
 
-  // TODO(b/181662306) None of the map reduce tests work with SQL since our map-reduce setup is
-  // inherently Datastore oriented, but this is a bigger task.
+  @TestSqlOnly
+  void testFailure_dnsUpdateEnqueueFailed() throws Exception {
+    HostResource host = persistActiveHost("ns1.example.tld");
+    persistResource(newDomainBase("example.tld", host));
+    persistResource(newDomainBase("otherexample.tld", host));
+    persistResource(newDomainBase("untouched.tld", persistActiveHost("ns2.example.tld")));
+    DateTime timeEnqueued = clock.nowUtc();
+    enqueuer.enqueueAsyncDnsRefresh(host, timeEnqueued);
+    DnsQueue mockedQueue = mock(DnsQueue.class);
+    action.dnsQueue = mockedQueue;
+    when(mockedQueue.addDomainRefreshTask(anyString()))
+        .thenThrow(new RuntimeException("Cannot enqueue task."));
+    runAction();
+    assertNoDnsTasksEnqueued();
+    assertTasksEnqueued(
+        QUEUE_ASYNC_HOST_RENAME,
+        new TaskMatcher()
+            .param(PARAM_HOST_KEY, host.createVKey().getOfyKey().getString())
+            .param(PARAM_REQUESTED_TIME, timeEnqueued.toString()));
+    verify(action.asyncTaskMetrics).recordDnsRefreshBatchSize(1L);
+    verifyNoMoreInteractions(action.asyncTaskMetrics);
+    assertThat(fakeResponse.getStatus()).isEqualTo(HttpStatus.SC_INTERNAL_SERVER_ERROR);
+    assertThat(acquireLock()).isPresent();
+  }
 
-  @TestOfyOnly
+  @TestOfyAndSql
   void testSuccess_dnsUpdateEnqueued() throws Exception {
     HostResource host = persistActiveHost("ns1.example.tld");
     persistResource(newDomainBase("example.tld", host));
@@ -134,21 +162,22 @@ public class RefreshDnsOnHostRenameActionTest
     persistResource(newDomainBase("untouched.tld", persistActiveHost("ns2.example.tld")));
     DateTime timeEnqueued = clock.nowUtc();
     enqueuer.enqueueAsyncDnsRefresh(host, timeEnqueued);
-    runMapreduce();
+    runAction();
     assertDnsTasksEnqueued("example.tld", "otherexample.tld");
     assertNoTasksEnqueued(QUEUE_ASYNC_HOST_RENAME);
     verify(action.asyncTaskMetrics).recordDnsRefreshBatchSize(1L);
     verify(action.asyncTaskMetrics)
         .recordAsyncFlowResult(DNS_REFRESH, OperationResult.SUCCESS, timeEnqueued);
     verifyNoMoreInteractions(action.asyncTaskMetrics);
+    assertThat(acquireLock()).isPresent();
   }
 
-  @TestOfyOnly
+  @TestOfyAndSql
   void testSuccess_multipleHostsProcessedInBatch() throws Exception {
     HostResource host1 = persistActiveHost("ns1.example.tld");
     HostResource host2 = persistActiveHost("ns2.example.tld");
     HostResource host3 = persistActiveHost("ns3.example.tld");
-    persistResource(newDomainBase("example1.tld", host1));
+    persistResource(newDomainBase("example1.tld", host1, host2));
     persistResource(newDomainBase("example2.tld", host2));
     persistResource(newDomainBase("example3.tld", host3));
     DateTime timeEnqueued = clock.nowUtc();
@@ -156,7 +185,7 @@ public class RefreshDnsOnHostRenameActionTest
     enqueuer.enqueueAsyncDnsRefresh(host1, timeEnqueued);
     enqueuer.enqueueAsyncDnsRefresh(host2, timeEnqueued);
     enqueuer.enqueueAsyncDnsRefresh(host3, laterTimeEnqueued);
-    runMapreduce();
+    runAction();
     assertDnsTasksEnqueued("example1.tld", "example2.tld", "example3.tld");
     assertNoTasksEnqueued(QUEUE_ASYNC_HOST_RENAME);
     verify(action.asyncTaskMetrics).recordDnsRefreshBatchSize(3L);
@@ -165,21 +194,23 @@ public class RefreshDnsOnHostRenameActionTest
     verify(action.asyncTaskMetrics)
         .recordAsyncFlowResult(DNS_REFRESH, OperationResult.SUCCESS, laterTimeEnqueued);
     verifyNoMoreInteractions(action.asyncTaskMetrics);
+    assertThat(acquireLock()).isPresent();
   }
 
-  @TestOfyOnly
+  @TestOfyAndSql
   void testSuccess_deletedHost_doesntTriggerDnsRefresh() throws Exception {
     HostResource host = persistDeletedHost("ns11.fakesss.tld", clock.nowUtc().minusDays(4));
     persistResource(newDomainBase("example1.tld", host));
     DateTime timeEnqueued = clock.nowUtc();
     enqueuer.enqueueAsyncDnsRefresh(host, timeEnqueued);
-    runMapreduce();
+    runAction();
     assertNoDnsTasksEnqueued();
     assertNoTasksEnqueued(QUEUE_ASYNC_HOST_RENAME);
     verify(action.asyncTaskMetrics).recordDnsRefreshBatchSize(1L);
     verify(action.asyncTaskMetrics)
         .recordAsyncFlowResult(DNS_REFRESH, OperationResult.STALE, timeEnqueued);
     verifyNoMoreInteractions(action.asyncTaskMetrics);
+    assertThat(acquireLock()).isPresent();
   }
 
   @TestOfyAndSql
@@ -191,9 +222,10 @@ public class RefreshDnsOnHostRenameActionTest
             .setDeletionTime(START_OF_TIME)
             .build());
     enqueuer.enqueueAsyncDnsRefresh(renamedHost, clock.nowUtc());
-    runMapreduce();
+    runAction();
     assertNoDnsTasksEnqueued();
     assertNoTasksEnqueued(QUEUE_ASYNC_HOST_RENAME);
+    assertThat(acquireLock()).isPresent();
   }
 
   @TestOfyAndSql
@@ -217,9 +249,10 @@ public class RefreshDnsOnHostRenameActionTest
     enqueueMapreduceOnly();
     assertThat(fakeResponse.getPayload()).isEqualTo("Can't acquire lock; aborting.");
     assertNoDnsTasksEnqueued();
+    assertThat(acquireLock()).isEmpty();
   }
 
-  @TestOfyAndSql
+  @TestOfyOnly
   void test_mapreduceHasWorkToDo_lockIsAcquired() {
     HostResource host = persistActiveHost("ns1.example.tld");
     enqueuer.enqueueAsyncDnsRefresh(host, clock.nowUtc());

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/beam/spec11/SafeBrowsingTransformsTest.java

@@ -164,7 +164,7 @@ class SafeBrowsingTransformsTest {
    * A serializable {@link Answer} that returns a mock HTTP response based on the HTTP request's
    * content.
    */
-  private static class HttpResponder implements Answer<CloseableHttpResponse>, Serializable {
+  static class HttpResponder implements Answer<CloseableHttpResponse>, Serializable {
     @Override
     public CloseableHttpResponse answer(InvocationOnMock invocation) throws Throwable {
       return getMockResponse(

core/src/test/java/google/registry/beam/spec11/Spec11PipelineTest.java

@@ -18,33 +18,60 @@ import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.ImmutableObjectSubject.immutableObjectCorrespondence;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.testing.AppEngineExtension.makeRegistrar1;
+import static google.registry.testing.DatabaseHelper.createTld;
+import static google.registry.testing.DatabaseHelper.persistActiveContact;
+import static google.registry.testing.DatabaseHelper.persistNewRegistrar;
+import static google.registry.testing.DatabaseHelper.persistResource;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+import static org.mockito.Mockito.withSettings;
 
+import com.google.common.base.Suppliers;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Streams;
 import com.google.common.truth.Correspondence;
 import com.google.common.truth.Correspondence.BinaryPredicate;
 import google.registry.beam.TestPipelineExtension;
+import google.registry.beam.spec11.SafeBrowsingTransforms.EvaluateSafeBrowsingFn;
+import google.registry.beam.spec11.SafeBrowsingTransformsTest.HttpResponder;
+import google.registry.model.contact.ContactResource;
+import google.registry.model.domain.DomainAuthInfo;
+import google.registry.model.domain.DomainBase;
+import google.registry.model.eppcommon.AuthInfo.PasswordAuth;
+import google.registry.model.registrar.Registrar;
 import google.registry.model.reporting.Spec11ThreatMatch;
 import google.registry.model.reporting.Spec11ThreatMatch.ThreatType;
 import google.registry.model.reporting.Spec11ThreatMatchDao;
 import google.registry.persistence.transaction.JpaTestRules;
 import google.registry.persistence.transaction.JpaTestRules.JpaIntegrationTestExtension;
+import google.registry.persistence.transaction.TransactionManager;
+import google.registry.persistence.transaction.TransactionManagerFactory;
 import google.registry.testing.DatastoreEntityExtension;
 import google.registry.testing.FakeClock;
+import google.registry.testing.FakeSleeper;
 import google.registry.util.ResourceUtils;
+import google.registry.util.Retrier;
 import java.io.File;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import org.apache.beam.sdk.coders.KvCoder;
 import org.apache.beam.sdk.coders.SerializableCoder;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.testing.PAssert;
 import org.apache.beam.sdk.transforms.Create;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.PCollection;
+import org.apache.http.client.methods.HttpPost;
+import org.apache.http.impl.client.CloseableHttpClient;
 import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.DateTime;
 import org.joda.time.LocalDate;
 import org.json.JSONObject;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Order;
 import org.junit.jupiter.api.Test;
@@ -60,9 +87,14 @@ import org.junit.jupiter.api.io.TempDir;
  */
 class Spec11PipelineTest {
 
+  private static final DateTime START_TIME = DateTime.parse("2020-01-27T00:00:00.0Z");
+  private final FakeClock fakeClock = new FakeClock(START_TIME);
+
   private static final String DATE = "2020-01-27";
   private static final String SAFE_BROWSING_API_KEY = "api-key";
   private static final String REPORTING_BUCKET_URL = "reporting_bucket";
+  private final CloseableHttpClient mockHttpClient =
+      mock(CloseableHttpClient.class, withSettings().serializable());
 
   private static final ImmutableList<Subdomain> SUBDOMAINS =
       ImmutableList.of(
@@ -103,12 +135,18 @@ class Spec11PipelineTest {
   private File reportingBucketUrl;
   private PCollection<KV<Subdomain, ThreatMatch>> threatMatches;
 
+  ImmutableSet<Spec11ThreatMatch> sqlThreatMatches;
+  TransactionManager tm;
+
   @BeforeEach
   void beforeEach() throws Exception {
+    tm = tm();
+    TransactionManagerFactory.setTm(jpaTm());
     reportingBucketUrl = Files.createDirectory(tmpDir.resolve(REPORTING_BUCKET_URL)).toFile();
     options.setDate(DATE);
     options.setSafeBrowsingApiKey(SAFE_BROWSING_API_KEY);
     options.setReportingBucketUrl(reportingBucketUrl.getAbsolutePath());
+    options.setDatabase("DATASTORE");
     threatMatches =
         pipeline.apply(
             Create.of(
@@ -118,11 +156,8 @@ class Spec11PipelineTest {
                     KvCoder.of(
                         SerializableCoder.of(Subdomain.class),
                         SerializableCoder.of(ThreatMatch.class))));
-  }
 
-  @Test
-  void testSuccess_saveToSql() {
-    ImmutableSet<Spec11ThreatMatch> sqlThreatMatches =
+    sqlThreatMatches =
         ImmutableSet.of(
             new Spec11ThreatMatch.Builder()
                 .setDomainName("111.com")
@@ -159,25 +194,98 @@ class Spec11PipelineTest {
                 .setCheckDate(new LocalDate(2020, 1, 27))
                 .setThreatTypes(ImmutableSet.of(ThreatType.UNWANTED_SOFTWARE))
                 .build());
+  }
+
+  @AfterEach
+  void afterEach() {
+    TransactionManagerFactory.setTm(tm);
+  }
+
+  @Test
+  void testSuccess_fullSqlPipeline() throws Exception {
+    setupCloudSql();
+    options.setDatabase("CLOUD_SQL");
+    EvaluateSafeBrowsingFn safeBrowsingFn =
+        new EvaluateSafeBrowsingFn(
+            SAFE_BROWSING_API_KEY,
+            new Retrier(new FakeSleeper(new FakeClock()), 1),
+            Suppliers.ofInstance(mockHttpClient));
+    when(mockHttpClient.execute(any(HttpPost.class))).thenAnswer(new HttpResponder());
+    Spec11Pipeline spec11Pipeline = new Spec11Pipeline(options, safeBrowsingFn);
+    spec11Pipeline.setupPipeline(pipeline);
+    pipeline.run(options).waitUntilFinish();
+    verifySaveToGcs();
+    verifySaveToCloudSql();
+  }
+
+  @Test
+  void testSuccess_saveToSql() {
     Spec11Pipeline.saveToSql(threatMatches, options);
     pipeline.run().waitUntilFinish();
-    assertThat(
-            jpaTm()
-                .transact(
-                    () ->
-                        Spec11ThreatMatchDao.loadEntriesByDate(
-                            jpaTm(), new LocalDate(2020, 1, 27))))
-        .comparingElementsUsing(immutableObjectCorrespondence("id"))
-        .containsExactlyElementsIn(sqlThreatMatches);
+    verifySaveToCloudSql();
   }
 
   @Test
   void testSuccess_saveToGcs() throws Exception {
+    Spec11Pipeline.saveToGcs(threatMatches, options);
+    pipeline.run().waitUntilFinish();
+    verifySaveToGcs();
+  }
+
+  @Test
+  void testSuccess_readFromCloudSql() throws Exception {
+    setupCloudSql();
+    PCollection<Subdomain> subdomains = Spec11Pipeline.readFromCloudSql(pipeline);
+    PAssert.that(subdomains).containsInAnyOrder(SUBDOMAINS);
+    pipeline.run().waitUntilFinish();
+  }
+
+  private void setupCloudSql() {
+    persistNewRegistrar("TheRegistrar");
+    persistNewRegistrar("NewRegistrar");
+    Registrar registrar1 =
+        persistResource(
+            makeRegistrar1()
+                .asBuilder()
+                .setClientId("hello-registrar")
+                .setEmailAddress("email@hello.net")
+                .build());
+    Registrar registrar2 =
+        persistResource(
+            makeRegistrar1()
+                .asBuilder()
+                .setClientId("kitty-registrar")
+                .setEmailAddress("contact@kit.ty")
+                .build());
+    Registrar registrar3 =
+        persistResource(
+            makeRegistrar1()
+                .asBuilder()
+                .setClientId("cool-registrar")
+                .setEmailAddress("cool@aid.net")
+                .build());
+
+    createTld("com");
+    createTld("net");
+    createTld("bank");
+    createTld("dev");
+
+    ContactResource contact1 = persistActiveContact(registrar1.getClientId());
+    ContactResource contact2 = persistActiveContact(registrar2.getClientId());
+    ContactResource contact3 = persistActiveContact(registrar3.getClientId());
+
+    persistResource(createDomain("111.com", "123456789-COM", registrar1, contact1));
+    persistResource(createDomain("party-night.net", "2244AABBC-NET", registrar2, contact2));
+    persistResource(createDomain("bitcoin.bank", "1C3D5E7F9-BANK", registrar1, contact1));
+    persistResource(createDomain("no-email.com", "2A4BA9BBC-COM", registrar2, contact2));
+    persistResource(
+        createDomain("anti-anti-anti-virus.dev", "555666888-DEV", registrar3, contact3));
+  }
+
+  private void verifySaveToGcs() throws Exception {
     ImmutableList<String> expectedFileContents =
         ImmutableList.copyOf(
             ResourceUtils.readResourceUtf8(this.getClass(), "test_output.txt").split("\n"));
-    Spec11Pipeline.saveToGcs(threatMatches, options);
-    pipeline.run().waitUntilFinish();
     ImmutableList<String> resultFileContents = resultFileContents();
     assertThat(resultFileContents.size()).isEqualTo(expectedFileContents.size());
     assertThat(resultFileContents.get(0)).isEqualTo(expectedFileContents.get(0));
@@ -188,6 +296,34 @@ class Spec11PipelineTest {
         .containsExactlyElementsIn(expectedFileContents.subList(1, expectedFileContents.size()));
   }
 
+  private void verifySaveToCloudSql() {
+    jpaTm()
+        .transact(
+            () -> {
+              ImmutableList<Spec11ThreatMatch> sqlThreatMatches =
+                  Spec11ThreatMatchDao.loadEntriesByDate(jpaTm(), new LocalDate(2020, 1, 27));
+              assertThat(sqlThreatMatches)
+                  .comparingElementsUsing(immutableObjectCorrespondence("id"))
+                  .containsExactlyElementsIn(sqlThreatMatches);
+            });
+  }
+
+  private DomainBase createDomain(
+      String domainName, String repoId, Registrar registrar, ContactResource contact) {
+    return new DomainBase.Builder()
+        .setDomainName(domainName)
+        .setRepoId(repoId)
+        .setCreationClientId(registrar.getClientId())
+        .setLastEppUpdateTime(fakeClock.nowUtc())
+        .setLastEppUpdateClientId(registrar.getClientId())
+        .setLastTransferTime(fakeClock.nowUtc())
+        .setRegistrant(contact.createVKey())
+        .setPersistedCurrentSponsorClientId(registrar.getClientId())
+        .setRegistrationExpirationTime(fakeClock.nowUtc().plusYears(1))
+        .setAuthInfo(DomainAuthInfo.create(PasswordAuth.create("password")))
+        .build();
+  }
+
   /** Returns the text contents of a file under the beamBucket/results directory. */
   private ImmutableList<String> resultFileContents() throws Exception {
     File resultFile =

core/src/test/java/google/registry/export/ExportDomainListsActionTest.java

@@ -16,8 +16,7 @@ package google.registry.export;
 
 import static com.google.appengine.tools.cloudstorage.GcsServiceFactory.createGcsService;
 import static com.google.common.truth.Truth.assertThat;
-import static google.registry.export.ExportDomainListsAction.ExportDomainListsReducer.EXPORT_MIME_TYPE;
-import static google.registry.export.ExportDomainListsAction.ExportDomainListsReducer.REGISTERED_DOMAINS_FILENAME;
+import static google.registry.export.ExportDomainListsAction.REGISTERED_DOMAINS_FILENAME;
 import static google.registry.testing.DatabaseHelper.createTld;
 import static google.registry.testing.DatabaseHelper.persistActiveDomain;
 import static google.registry.testing.DatabaseHelper.persistDeletedDomain;
@@ -34,25 +33,40 @@ import com.google.appengine.tools.cloudstorage.GcsFilename;
 import com.google.appengine.tools.cloudstorage.GcsService;
 import com.google.appengine.tools.cloudstorage.ListOptions;
 import com.google.appengine.tools.cloudstorage.ListResult;
+import com.google.common.net.MediaType;
 import google.registry.export.ExportDomainListsAction.ExportDomainListsReducer;
+import google.registry.model.ofy.Ofy;
 import google.registry.model.registry.Registry;
 import google.registry.model.registry.Registry.TldType;
 import google.registry.storage.drive.DriveConnection;
+import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.FakeClock;
 import google.registry.testing.FakeResponse;
+import google.registry.testing.InjectExtension;
+import google.registry.testing.TestOfyAndSql;
+import google.registry.testing.TestOfyOnly;
 import google.registry.testing.mapreduce.MapreduceTestCase;
 import java.io.FileNotFoundException;
 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;
 import org.mockito.ArgumentCaptor;
 
 /** Unit tests for {@link ExportDomainListsAction}. */
+@DualDatabaseTest
 class ExportDomainListsActionTest extends MapreduceTestCase<ExportDomainListsAction> {
 
   private GcsService gcsService;
   private DriveConnection driveConnection = mock(DriveConnection.class);
   private ArgumentCaptor<byte[]> bytesExportedToDrive = ArgumentCaptor.forClass(byte[].class);
   private final FakeResponse response = new FakeResponse();
+  private final FakeClock clock = new FakeClock(DateTime.parse("2020-02-02T02:02:02Z"));
+
+  @Order(Order.DEFAULT - 1)
+  @RegisterExtension
+  public final InjectExtension inject =
+      new InjectExtension().withStaticFieldOverride(Ofy.class, "clock", clock);
 
   @BeforeEach
   void beforeEach() {
@@ -68,10 +82,12 @@ class ExportDomainListsActionTest extends MapreduceTestCase<ExportDomainListsAct
     action.response = response;
     action.gcsBucket = "outputbucket";
     action.gcsBufferSize = 500;
+    action.clock = clock;
+    action.driveConnection = driveConnection;
     gcsService = createGcsService();
   }
 
-  private void runMapreduce() throws Exception {
+  private void runAction() throws Exception {
     action.run();
     executeTasksUntilEmpty("mapreduce");
   }
@@ -80,27 +96,27 @@ class ExportDomainListsActionTest extends MapreduceTestCase<ExportDomainListsAct
     verify(driveConnection)
         .createOrUpdateFile(
             eq(REGISTERED_DOMAINS_FILENAME),
-            eq(EXPORT_MIME_TYPE),
+            eq(MediaType.PLAIN_TEXT_UTF_8),
             eq(folderId),
             bytesExportedToDrive.capture());
     assertThat(new String(bytesExportedToDrive.getValue(), UTF_8)).isEqualTo(domains);
   }
 
-  @Test
+  @TestOfyOnly
   void test_writesLinkToMapreduceConsoleToResponse() throws Exception {
-    runMapreduce();
+    runAction();
     assertThat(response.getPayload())
         .startsWith(
             "Mapreduce console: https://backend-dot-projectid.appspot.com"
                 + "/_ah/pipeline/status.html?root=");
   }
 
-  @Test
+  @TestOfyAndSql
   void test_outputsOnlyActiveDomains() throws Exception {
     persistActiveDomain("onetwo.tld");
     persistActiveDomain("rudnitzky.tld");
     persistDeletedDomain("mortuary.tld", DateTime.parse("2001-03-14T10:11:12Z"));
-    runMapreduce();
+    runAction();
     GcsFilename existingFile = new GcsFilename("outputbucket", "tld.txt");
     String tlds = new String(readGcsFile(gcsService, existingFile), UTF_8);
     // Check that it only contains the active domains, not the dead one.
@@ -109,12 +125,12 @@ class ExportDomainListsActionTest extends MapreduceTestCase<ExportDomainListsAct
     verifyNoMoreInteractions(driveConnection);
   }
 
-  @Test
+  @TestOfyAndSql
   void test_outputsOnlyDomainsOnRealTlds() throws Exception {
     persistActiveDomain("onetwo.tld");
     persistActiveDomain("rudnitzky.tld");
     persistActiveDomain("wontgo.testtld");
-    runMapreduce();
+    runAction();
     GcsFilename existingFile = new GcsFilename("outputbucket", "tld.txt");
     String tlds = new String(readGcsFile(gcsService, existingFile), UTF_8).trim();
     // Check that it only contains the domains on the real TLD, and not the test one.
@@ -130,7 +146,7 @@ class ExportDomainListsActionTest extends MapreduceTestCase<ExportDomainListsAct
     verifyNoMoreInteractions(driveConnection);
   }
 
-  @Test
+  @TestOfyAndSql
   void test_outputsDomainsFromDifferentTldsToMultipleFiles() throws Exception {
     createTld("tldtwo");
     persistResource(Registry.get("tldtwo").asBuilder().setDriveFolderId("hooray").build());
@@ -143,7 +159,7 @@ class ExportDomainListsActionTest extends MapreduceTestCase<ExportDomainListsAct
     persistActiveDomain("santa.tldtwo");
     persistActiveDomain("buddy.tldtwo");
     persistActiveDomain("cupid.tldthree");
-    runMapreduce();
+    runAction();
     GcsFilename firstTldFile = new GcsFilename("outputbucket", "tld.txt");
     String tlds = new String(readGcsFile(gcsService, firstTldFile), UTF_8).trim();
     assertThat(tlds).isEqualTo("dasher.tld\nprancer.tld");

core/src/test/java/google/registry/flows/EppPointInTimeTest.java

@@ -16,8 +16,11 @@ package google.registry.flows;
 
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.EppResourceUtils.loadAtPointInTime;
-import static google.registry.model.ofy.ObjectifyService.auditedOfy;
+import static google.registry.model.ImmutableObjectSubject.assertAboutImmutableObjects;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatabaseHelper.createTld;
+import static google.registry.testing.DatabaseHelper.loadAllOf;
+import static google.registry.testing.DatabaseHelper.loadByEntity;
 import static google.registry.testing.DatabaseHelper.persistActiveContact;
 import static google.registry.testing.DatabaseHelper.persistActiveHost;
 import static java.nio.charset.StandardCharsets.UTF_8;
@@ -25,31 +28,38 @@ import static org.joda.time.DateTimeZone.UTC;
 import static org.joda.time.Duration.standardDays;
 
 import com.google.common.collect.ImmutableMap;
-import com.googlecode.objectify.Key;
+import com.google.common.collect.Iterables;
 import google.registry.flows.EppTestComponent.FakesAndMocksModule;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.ofy.Ofy;
 import google.registry.monitoring.whitebox.EppMetric;
 import google.registry.testing.AppEngineExtension;
+import google.registry.testing.DualDatabaseTest;
 import google.registry.testing.EppLoader;
 import google.registry.testing.FakeClock;
 import google.registry.testing.FakeHttpSession;
 import google.registry.testing.InjectExtension;
+import google.registry.testing.TestOfyAndSql;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
-/** Test that domain flows create the commit logs needed to reload at points in the past. */
-class EppCommitLogsTest {
+/** Test that we can reload EPP resources as they were in the past. */
+@DualDatabaseTest
+class EppPointInTimeTest {
+
+  private final FakeClock clock = new FakeClock(DateTime.now(UTC));
 
   @RegisterExtension
   final AppEngineExtension appEngine =
-      AppEngineExtension.builder().withDatastoreAndCloudSql().withTaskQueue().build();
+      AppEngineExtension.builder()
+          .withDatastoreAndCloudSql()
+          .withClock(clock)
+          .withTaskQueue()
+          .build();
 
   @RegisterExtension final InjectExtension inject = new InjectExtension();
 
-  private final FakeClock clock = new FakeClock(DateTime.now(UTC));
   private EppLoader eppLoader;
 
   @BeforeEach
@@ -81,7 +91,7 @@ class EppCommitLogsTest {
         .run(EppMetric.builder());
   }
 
-  @Test
+  @TestOfyAndSql
   void testLoadAtPointInTime() throws Exception {
     clock.setTo(DateTime.parse("1984-12-18T12:30Z")); // not midnight
 
@@ -95,64 +105,75 @@ class EppCommitLogsTest {
     clock.setTo(timeAtCreate);
     eppLoader = new EppLoader(this, "domain_create.xml", ImmutableMap.of("DOMAIN", "example.tld"));
     runFlow();
-    auditedOfy().clearSessionCache();
-    Key<DomainBase> key = Key.create(auditedOfy().load().type(DomainBase.class).first().now());
-    DomainBase domainAfterCreate = auditedOfy().load().key(key).now();
+    tm().clearSessionCache();
+    DomainBase domainAfterCreate = Iterables.getOnlyElement(loadAllOf(DomainBase.class));
     assertThat(domainAfterCreate.getDomainName()).isEqualTo("example.tld");
 
     clock.advanceBy(standardDays(2));
     DateTime timeAtFirstUpdate = clock.nowUtc();
     eppLoader = new EppLoader(this, "domain_update_dsdata_add.xml");
     runFlow();
-    auditedOfy().clearSessionCache();
+    tm().clearSessionCache();
 
-    DomainBase domainAfterFirstUpdate = auditedOfy().load().key(key).now();
+    DomainBase domainAfterFirstUpdate = loadByEntity(domainAfterCreate);
     assertThat(domainAfterCreate).isNotEqualTo(domainAfterFirstUpdate);
 
     clock.advanceOneMilli(); // same day as first update
     DateTime timeAtSecondUpdate = clock.nowUtc();
     eppLoader = new EppLoader(this, "domain_update_dsdata_rem.xml");
     runFlow();
-    auditedOfy().clearSessionCache();
-    DomainBase domainAfterSecondUpdate = auditedOfy().load().key(key).now();
+    tm().clearSessionCache();
+    DomainBase domainAfterSecondUpdate = loadByEntity(domainAfterCreate);
 
     clock.advanceBy(standardDays(2));
     DateTime timeAtDelete = clock.nowUtc(); // before 'add' grace period ends
     eppLoader = new EppLoader(this, "domain_delete.xml", ImmutableMap.of("DOMAIN", "example.tld"));
     runFlow();
-    auditedOfy().clearSessionCache();
+    tm().clearSessionCache();
 
     assertThat(domainAfterFirstUpdate).isNotEqualTo(domainAfterSecondUpdate);
 
     // Point-in-time can only rewind an object from the current version, not roll forward.
-    DomainBase latest = auditedOfy().load().key(key).now();
+    DomainBase latest = loadByEntity(domainAfterCreate);
 
     // Creation time has millisecond granularity due to isActive() check.
-    auditedOfy().clearSessionCache();
+    tm().clearSessionCache();
     assertThat(loadAtPointInTime(latest, timeAtCreate.minusMillis(1)).now()).isNull();
     assertThat(loadAtPointInTime(latest, timeAtCreate).now()).isNotNull();
     assertThat(loadAtPointInTime(latest, timeAtCreate.plusMillis(1)).now()).isNotNull();
 
-    auditedOfy().clearSessionCache();
-    assertThat(loadAtPointInTime(latest, timeAtCreate.plusDays(1)).now())
-        .isEqualTo(domainAfterCreate);
+    tm().clearSessionCache();
+    assertAboutImmutableObjects()
+        .that(loadAtPointInTime(latest, timeAtCreate.plusDays(1)).now())
+        .hasFieldsEqualTo(domainAfterCreate);
 
-    // Both updates happened on the same day. Since the revisions field has day granularity, the
-    // key to the first update should have been overwritten by the second, and its timestamp rolled
-    // forward. So we have to fall back to the last revision before midnight.
-    auditedOfy().clearSessionCache();
-    assertThat(loadAtPointInTime(latest, timeAtFirstUpdate).now()).isEqualTo(domainAfterCreate);
+    tm().clearSessionCache();
+    if (tm().isOfy()) {
+      // Both updates happened on the same day. Since the revisions field has day granularity in
+      // Datastore, the key to the first update should have been overwritten by the second, and its
+      // timestamp rolled forward. So we have to fall back to the last revision before midnight.
+      assertThat(loadAtPointInTime(latest, timeAtFirstUpdate).now()).isEqualTo(domainAfterCreate);
+    } else {
+      // In SQL, however, we are not limited by the day granularity, so when we request the object
+      // at timeAtFirstUpdate we should receive the object at that first update, even though the
+      // second update occurred one millisecond later.
+      assertAboutImmutableObjects()
+          .that(loadAtPointInTime(latest, timeAtFirstUpdate).now())
+          .hasFieldsEqualTo(domainAfterFirstUpdate);
+    }
 
-    auditedOfy().clearSessionCache();
-    assertThat(loadAtPointInTime(latest, timeAtSecondUpdate).now())
-        .isEqualTo(domainAfterSecondUpdate);
+    tm().clearSessionCache();
+    assertAboutImmutableObjects()
+        .that(loadAtPointInTime(latest, timeAtSecondUpdate).now())
+        .hasFieldsEqualTo(domainAfterSecondUpdate);
 
-    auditedOfy().clearSessionCache();
-    assertThat(loadAtPointInTime(latest, timeAtSecondUpdate.plusDays(1)).now())
-        .isEqualTo(domainAfterSecondUpdate);
+    tm().clearSessionCache();
+    assertAboutImmutableObjects()
+        .that(loadAtPointInTime(latest, timeAtSecondUpdate.plusDays(1)).now())
+        .hasFieldsEqualTo(domainAfterSecondUpdate);
 
     // Deletion time has millisecond granularity due to isActive() check.
-    auditedOfy().clearSessionCache();
+    tm().clearSessionCache();
     assertThat(loadAtPointInTime(latest, timeAtDelete.minusMillis(1)).now()).isNotNull();
     assertThat(loadAtPointInTime(latest, timeAtDelete).now()).isNull();
     assertThat(loadAtPointInTime(latest, timeAtDelete.plusMillis(1)).now()).isNull();

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/EppResourceUtilsTest.java

@@ -41,14 +41,18 @@ import org.junit.jupiter.api.extension.RegisterExtension;
 @DualDatabaseTest
 class EppResourceUtilsTest {
 
+  private final FakeClock clock = new FakeClock(DateTime.now(UTC));
+
   @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.now(UTC));
-
   @BeforeEach
   void beforeEach() {
     createTld("tld");

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/model/index/EppResourceIndexTest.java

@@ -16,7 +16,7 @@ package google.registry.model.index;
 
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.config.RegistryConfig.getEppResourceIndexBucketCount;
-import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.model.ofy.ObjectifyService.auditedOfy;
 import static google.registry.testing.DatabaseHelper.createTld;
 import static google.registry.testing.DatabaseHelper.persistActiveContact;
 import static google.registry.testing.DatabaseHelper.persistResource;
@@ -44,7 +44,7 @@ class EppResourceIndexTest extends EntityTestCase {
   @Test
   void testPersistence() {
     EppResourceIndex loadedIndex = Iterables.getOnlyElement(getEppResourceIndexObjects());
-    assertThat(ofy().load().key(loadedIndex.reference).now()).isEqualTo(contact);
+    assertThat(auditedOfy().load().key(loadedIndex.reference).now()).isEqualTo(contact);
   }
 
   @Test
@@ -56,7 +56,7 @@ class EppResourceIndexTest extends EntityTestCase {
   void testIdempotentOnUpdate() {
     contact = persistResource(contact.asBuilder().setEmailAddress("abc@def.fake").build());
     EppResourceIndex loadedIndex = Iterables.getOnlyElement(getEppResourceIndexObjects());
-    assertThat(ofy().load().key(loadedIndex.reference).now()).isEqualTo(contact);
+    assertThat(auditedOfy().load().key(loadedIndex.reference).now()).isEqualTo(contact);
   }
 
   /**
@@ -65,9 +65,11 @@ class EppResourceIndexTest extends EntityTestCase {
   private static ImmutableList<EppResourceIndex> getEppResourceIndexObjects() {
     ImmutableList.Builder<EppResourceIndex> indexEntities = new ImmutableList.Builder<>();
     for (int i = 0; i < getEppResourceIndexBucketCount(); i++) {
-      indexEntities.addAll(ofy().load()
-          .type(EppResourceIndex.class)
-          .ancestor(Key.create(EppResourceIndexBucket.class, i + 1)));
+      indexEntities.addAll(
+          auditedOfy()
+              .load()
+              .type(EppResourceIndex.class)
+              .ancestor(Key.create(EppResourceIndexBucket.class, i + 1)));
     }
     return indexEntities.build();
   }

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/reporting/spec11/GenerateSpec11ReportActionTest.java

@@ -50,6 +50,7 @@ class GenerateSpec11ReportActionTest extends BeamActionTestBase {
             "gs://reporting-project/reporting-bucket/",
             "api_key/a",
             clock.nowUtc().toLocalDate(),
+            "DATASTORE",
             clock,
             response,
             dataflow);
@@ -71,6 +72,7 @@ class GenerateSpec11ReportActionTest extends BeamActionTestBase {
             "gs://reporting-project/reporting-bucket/",
             "api_key/a",
             clock.nowUtc().toLocalDate(),
+            "DATASTORE",
             clock,
             response,
             dataflow);

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

@@ -18,19 +18,30 @@ 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.base.Suppliers;
+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.CommitLogBucket;
+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,29 +49,51 @@ 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);
+    // Use a single bucket to expose timestamp inversion problems.
+    injectExtension.setStaticField(
+        CommitLogBucket.class, "bucketIdSupplier", Suppliers.ofInstance(1));
+    fakeClock.setAutoIncrementByOneMilli();
     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
   public void tearDown() {
+    fakeClock.disableAutoIncrement();
     RegistryConfig.overrideCloudSqlReplicateTransactions(false);
   }
 
@@ -170,6 +203,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/testing/DatabaseHelper.java

@@ -19,6 +19,7 @@ import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.base.Suppliers.memoize;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static com.google.common.collect.Iterables.toArray;
 import static com.google.common.collect.MoreCollectors.onlyElement;
 import static com.google.common.truth.Truth.assertThat;
@@ -110,6 +111,7 @@ import google.registry.model.transfer.TransferStatus;
 import google.registry.persistence.VKey;
 import google.registry.schema.tld.PremiumListDao;
 import google.registry.tmch.LordnTaskUtils;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
 import java.util.NoSuchElementException;
@@ -182,11 +184,10 @@ public class DatabaseHelper {
         domainName, generateNewDomainRoid(getTldFromDomainName(domainName)), contact);
   }
 
-  public static DomainBase newDomainBase(String domainName, HostResource host) {
-    return newDomainBase(domainName)
-        .asBuilder()
-        .setNameservers(ImmutableSet.of(host.createVKey()))
-        .build();
+  public static DomainBase newDomainBase(String domainName, HostResource... hosts) {
+    ImmutableSet<VKey<HostResource>> hostKeys =
+        Arrays.stream(hosts).map(HostResource::createVKey).collect(toImmutableSet());
+    return newDomainBase(domainName).asBuilder().setNameservers(hostKeys).build();
   }
 
   public static DomainBase newDomainBase(

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

@@ -162,6 +162,20 @@ class CreateDomainCommandTest extends EppToolCommandTestCase<CreateDomainCommand
             + "sending total cost for 1 year(s) of USD 877.00.");
   }
 
+  @Test
+  void testSuccess_reasonAndRegistrarRequest() throws Exception {
+    createTld("tld");
+    runCommandForced(
+        "--client=NewRegistrar",
+        "--registrant=crr-admin",
+        "--admins=crr-admin",
+        "--techs=crr-tech",
+        "--reason=\"Creating test domain\"",
+        "--registrar_request=false",
+        "example.tld");
+    eppVerifier.verifySent("domain_create_metadata.xml");
+  }
+
   @Test
   void testFailure_duplicateDomains() {
     IllegalArgumentException thrown =
@@ -175,7 +189,7 @@ class CreateDomainCommandTest extends EppToolCommandTestCase<CreateDomainCommand
                     "--techs=crr-tech",
                     "example.tld",
                     "example.tld"));
-    assertThat(thrown).hasMessageThat().contains("Duplicate arguments found: \'example.tld\'");
+    assertThat(thrown).hasMessageThat().contains("Duplicate arguments found: 'example.tld'");
   }
 
   @Test

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,178 @@
+// 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;
+import org.junit.jupiter.api.BeforeEach;
+
+/** Tests for {@link SetDatabaseMigrationStateCommand}. */
+@DualDatabaseTest
+public class SetDatabaseMigrationStateCommandTest
+    extends CommandTestCase<SetDatabaseMigrationStateCommand> {
+
+  @BeforeEach
+  void beforeEach() {
+    // clear out any static state that may have been persisted
+    ofyTm()
+        .transact(
+            () ->
+                ofyTm()
+                    .loadSingleton(DatabaseMigrationStateSchedule.class)
+                    .ifPresent(ofyTm()::delete));
+    DatabaseMigrationStateSchedule.CACHE.invalidateAll();
+  }
+
+  @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/java/google/registry/tools/UpdateDomainCommandTest.java

@@ -197,6 +197,18 @@ class UpdateDomainCommandTest extends EppToolCommandTestCase<UpdateDomainCommand
     eppVerifier.verifySent("domain_update_change.xml");
   }
 
+  @TestOfyAndSql
+  void testSuccess_change_reasonAndRegistrarRequest() throws Exception {
+    runCommandForced(
+        "--client=NewRegistrar",
+        "--registrant=crr-admin",
+        "--password=2fooBAR",
+        "--reason=\"Testing domain update\"",
+        "--registrar_request=false",
+        "example.tld");
+    eppVerifier.verifySent("domain_update_change_metadata.xml");
+  }
+
   @TestOfyAndSql
   void testSuccess_setNameservers() throws Exception {
     HostResource host1 = persistActiveHost("ns1.zdns.google");

core/src/test/java/google/registry/tools/server/KillAllCommitLogsActionTest.java

@@ -21,7 +21,6 @@ import static com.google.common.collect.Iterables.filter;
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
 import static google.registry.model.ofy.ObjectifyService.auditedOfy;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.testing.DatabaseHelper.createTld;
 import static google.registry.testing.DatabaseHelper.newContactResource;
 import static google.registry.testing.DatabaseHelper.persistResource;
@@ -78,7 +77,9 @@ class KillAllCommitLogsActionTest extends MapreduceTestCase<KillAllCommitLogsAct
             START_OF_TIME.plusDays(1),
             ImmutableMap.of(1, bucketTime, 2, bucketTime, 3, bucketTime)));
     for (Class<?> clazz : AFFECTED_TYPES) {
-      assertWithMessage("entities of type " + clazz).that(ofy().load().type(clazz)).isNotEmpty();
+      assertWithMessage("entities of type " + clazz)
+          .that(auditedOfy().load().type(clazz))
+          .isNotEmpty();
     }
     ImmutableList<?> otherStuff =
         Streams.stream(auditedOfy().load())