Make RegistryJpaIO use CriteriaQuery intead of QueryComposer (#1209)

QueryComposer could be used when the transaction manager is not
determined (i. e. it supports both ofy and sql), but this also imposes
limits on what you can do with it. For example it does not support IN
operator in the where clause.

Since QueryComposer itself creates a CriteriaQuery for JPA TM it make
sense to have RegistryJpaIO take a CriteriaQuery directly as it only
uses JPA.

Also add some more helper methods to use native queries and typed
queires, and fix some generic type warnings.

<!-- Reviewable:start -->
---
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/google/nomulus/1209)
<!-- Reviewable:end -->

.gitignore

@@ -4,6 +4,7 @@
 ######################################################################
 # Java Ignores
 
+gjf.out
 *.class
 
 # Mobile Tools for Java (J2ME)

config/presubmits.py

@@ -202,6 +202,8 @@ PRESUBMITS = {
         "java",
         # ActivityReportingQueryBuilder deals with Dremel queries
         {"src/test", "ActivityReportingQueryBuilder.java",
+         # This class contains helper method to make queries in Beam.
+         "RegistryJpaIO.java",
          # TODO(b/179158393): Remove everything below, which should be done
          # using Criteria
          "ForeignKeyIndex.java",

core/build.gradle

@@ -317,6 +317,7 @@ dependencies {
   testCompile deps['com.google.monitoring-client:contrib']
   testCompile deps['com.google.truth:truth']
   testCompile deps['com.google.truth.extensions:truth-java8-extension']
+  testCompile deps['org.checkerframework:checker-qual']
   testCompile deps['org.hamcrest:hamcrest']
   testCompile deps['org.hamcrest:hamcrest-core']
   testCompile deps['org.hamcrest:hamcrest-library']

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,11 +27,14 @@ 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;
 import google.registry.request.Action;
+import google.registry.request.Action.Method;
 import google.registry.request.Response;
 import google.registry.request.auth.Auth;
 import google.registry.schema.replay.DatastoreEntity;
@@ -39,6 +42,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;
@@ -53,7 +57,7 @@ import org.joda.time.Duration;
 @Action(
     service = Action.Service.BACKEND,
     path = ReplayCommitLogsToSqlAction.PATH,
-    method = Action.Method.POST,
+    method = Method.POST,
     automaticallyPrintOk = true,
     auth = Auth.AUTH_INTERNAL_OR_ADMIN)
 public class ReplayCommitLogsToSqlAction implements Runnable {
@@ -69,15 +73,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/common/RegistryJpaIO.java

@@ -21,14 +21,15 @@ import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Streams;
 import google.registry.backup.AppEngineEnvironment;
-import google.registry.beam.common.RegistryQuery.QueryComposerFactory;
-import google.registry.beam.common.RegistryQuery.RegistryQueryFactory;
+import google.registry.beam.common.RegistryQuery.CriteriaQuerySupplier;
 import google.registry.model.ofy.ObjectifyService;
 import google.registry.persistence.transaction.JpaTransactionManager;
 import google.registry.persistence.transaction.TransactionManagerFactory;
 import java.io.Serializable;
+import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.ThreadLocalRandom;
+import javax.annotation.Nullable;
 import javax.persistence.criteria.CriteriaQuery;
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.coders.SerializableCoder;
@@ -60,13 +61,18 @@ public final class RegistryJpaIO {
 
   private RegistryJpaIO() {}
 
-  public static <R> Read<R, R> read(QueryComposerFactory<R> queryFactory) {
-    return Read.<R, R>builder().queryFactory(queryFactory).build();
+  public static <R> Read<R, R> read(CriteriaQuerySupplier<R> query) {
+    return read(query, x -> x);
   }
 
   public static <R, T> Read<R, T> read(
-      QueryComposerFactory<R> queryFactory, SerializableFunction<R, T> resultMapper) {
-    return Read.<R, T>builder().queryFactory(queryFactory).resultMapper(resultMapper).build();
+      CriteriaQuerySupplier<R> query, SerializableFunction<R, T> resultMapper) {
+    return Read.<R, T>builder().criteriaQuery(query).resultMapper(resultMapper).build();
+  }
+
+  public static <R, T> Read<R, T> read(
+      String sql, boolean nativeQuery, SerializableFunction<R, T> resultMapper) {
+    return read(sql, null, nativeQuery, resultMapper);
   }
 
   /**
@@ -74,8 +80,39 @@ public final class RegistryJpaIO {
    *
    * <p>User should take care to prevent sql-injection attacks.
    */
-  public static <R, T> Read<R, T> read(String jpql, SerializableFunction<R, T> resultMapper) {
-    return Read.<R, T>builder().jpqlQueryFactory(jpql).resultMapper(resultMapper).build();
+  public static <R, T> Read<R, T> read(
+      String sql,
+      @Nullable Map<String, Object> parameter,
+      boolean nativeQuery,
+      SerializableFunction<R, T> resultMapper) {
+    Read.Builder<R, T> builder = Read.builder();
+    if (nativeQuery) {
+      builder.nativeQuery(sql, parameter);
+    } else {
+      builder.jpqlQuery(sql, parameter);
+    }
+    return builder.resultMapper(resultMapper).build();
+  }
+
+  public static <R, T> Read<R, T> read(
+      String jpql, Class<R> clazz, SerializableFunction<R, T> resultMapper) {
+    return read(jpql, null, clazz, resultMapper);
+  }
+
+  /**
+   * Returns a {@link Read} connector based on the given {@code jpql} typed query string.
+   *
+   * <p>User should take care to prevent sql-injection attacks.
+   */
+  public static <R, T> Read<R, T> read(
+      String jpql,
+      @Nullable Map<String, Object> parameter,
+      Class<R> clazz,
+      SerializableFunction<R, T> resultMapper) {
+    return Read.<R, T>builder()
+        .jpqlQuery(jpql, clazz, parameter)
+        .resultMapper(resultMapper)
+        .build();
   }
 
   public static <T> Write<T> write() {
@@ -94,7 +131,7 @@ public final class RegistryJpaIO {
 
     abstract String name();
 
-    abstract RegistryQueryFactory<R> queryFactory();
+    abstract RegistryQuery<R> query();
 
     abstract SerializableFunction<R, T> resultMapper();
 
@@ -107,9 +144,7 @@ public final class RegistryJpaIO {
     public PCollection<T> expand(PBegin input) {
       return input
           .apply("Starting " + name(), Create.of((Void) null))
-          .apply(
-              "Run query for " + name(),
-              ParDo.of(new QueryRunner<>(queryFactory(), resultMapper())))
+          .apply("Run query for " + name(), ParDo.of(new QueryRunner<>(query(), resultMapper())))
           .setCoder(coder())
           .apply("Reshuffle", Reshuffle.viaRandomKey());
     }
@@ -127,9 +162,8 @@ public final class RegistryJpaIO {
     }
 
     static <R, T> Builder<R, T> builder() {
-      return new AutoValue_RegistryJpaIO_Read.Builder()
+      return new AutoValue_RegistryJpaIO_Read.Builder<R, T>()
           .name(DEFAULT_NAME)
-          .resultMapper(x -> x)
           .coder(SerializableCoder.of(Serializable.class));
     }
 
@@ -138,7 +172,7 @@ public final class RegistryJpaIO {
 
       abstract Builder<R, T> name(String name);
 
-      abstract Builder<R, T> queryFactory(RegistryQueryFactory<R> queryFactory);
+      abstract Builder<R, T> query(RegistryQuery<R> query);
 
       abstract Builder<R, T> resultMapper(SerializableFunction<R, T> mapper);
 
@@ -146,21 +180,29 @@ public final class RegistryJpaIO {
 
       abstract Read<R, T> build();
 
-      Builder<R, T> queryFactory(QueryComposerFactory<R> queryFactory) {
-        return queryFactory(RegistryQuery.createQueryFactory(queryFactory));
+      Builder<R, T> criteriaQuery(CriteriaQuerySupplier<R> criteriaQuery) {
+        return query(RegistryQuery.createQuery(criteriaQuery));
       }
 
-      Builder<R, T> jpqlQueryFactory(String jpql) {
-        return queryFactory(RegistryQuery.createQueryFactory(jpql));
+      Builder<R, T> nativeQuery(String sql, Map<String, Object> parameters) {
+        return query(RegistryQuery.createQuery(sql, parameters, true));
+      }
+
+      Builder<R, T> jpqlQuery(String jpql, Map<String, Object> parameters) {
+        return query(RegistryQuery.createQuery(jpql, parameters, false));
+      }
+
+      Builder<R, T> jpqlQuery(String jpql, Class<R> clazz, Map<String, Object> parameters) {
+        return query(RegistryQuery.createQuery(jpql, parameters, clazz));
       }
     }
 
     static class QueryRunner<R, T> extends DoFn<Void, T> {
-      private final RegistryQueryFactory<R> queryFactory;
+      private final RegistryQuery<R> query;
       private final SerializableFunction<R, T> resultMapper;
 
-      QueryRunner(RegistryQueryFactory<R> queryFactory, SerializableFunction<R, T> resultMapper) {
-        this.queryFactory = queryFactory;
+      QueryRunner(RegistryQuery<R> query, SerializableFunction<R, T> resultMapper) {
+        this.query = query;
         this.resultMapper = resultMapper;
       }
 
@@ -172,10 +214,7 @@ public final class RegistryJpaIO {
           // TODO(b/187210388): JpaTransactionManager should support non-transactional query.
           jpaTm()
               .transactNoRetry(
-                  () ->
-                      queryFactory.apply(jpaTm()).stream()
-                          .map(resultMapper::apply)
-                          .forEach(outputReceiver::output));
+                  () -> query.stream().map(resultMapper::apply).forEach(outputReceiver::output));
         }
       }
     }

core/src/main/java/google/registry/beam/common/RegistryQuery.java

@@ -14,48 +14,75 @@
 
 package google.registry.beam.common;
 
-import google.registry.persistence.transaction.JpaTransactionManager;
-import google.registry.persistence.transaction.QueryComposer;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+
+import java.io.Serializable;
+import java.util.Map;
+import java.util.function.Supplier;
 import java.util.stream.Stream;
+import javax.annotation.Nullable;
 import javax.persistence.EntityManager;
 import javax.persistence.Query;
-import org.apache.beam.sdk.transforms.SerializableFunction;
+import javax.persistence.TypedQuery;
+import javax.persistence.criteria.CriteriaQuery;
 
 /** Interface for query instances used by {@link RegistryJpaIO.Read}. */
-public interface RegistryQuery<T> {
+public interface RegistryQuery<T> extends Serializable {
   Stream<T> stream();
 
-  /** Factory for {@link RegistryQuery}. */
-  interface RegistryQueryFactory<T>
-      extends SerializableFunction<JpaTransactionManager, RegistryQuery<T>> {}
-
-  // TODO(mmuller): Consider detached JpaQueryComposer that works with any JpaTransactionManager
-  // instance, i.e., change composer.buildQuery() to composer.buildQuery(JpaTransactionManager).
-  // This way QueryComposer becomes reusable and serializable (at least with Hibernate), and this
-  // interface would no longer be necessary.
-  interface QueryComposerFactory<T>
-      extends SerializableFunction<JpaTransactionManager, QueryComposer<T>> {}
+  interface CriteriaQuerySupplier<T> extends Supplier<CriteriaQuery<T>>, Serializable {}
 
   /**
-   * Returns a {@link RegistryQueryFactory} that creates a JPQL query from constant text.
+   * Returns a {@link RegistryQuery} that creates a string query from constant text.
    *
+   * @param nativeQuery whether the given string is to be interpreted as a native query or JPQL.
+   * @param parameters parameters to be substituted in the query.
    * @param <T> Type of each row in the result set, {@link Object} in single-select queries, and
    *     {@code Object[]} in multi-select queries.
    */
-  @SuppressWarnings("unchecked") // query.getResultStream: jpa api uses raw type
-  static <T> RegistryQueryFactory<T> createQueryFactory(String jpql) {
-    return (JpaTransactionManager jpa) ->
-        () -> {
-          EntityManager entityManager = jpa.getEntityManager();
-          Query query = entityManager.createQuery(jpql);
-          return query.getResultStream().map(e -> detach(entityManager, e));
-        };
+  static <T> RegistryQuery<T> createQuery(
+      String sql, @Nullable Map<String, Object> parameters, boolean nativeQuery) {
+    return () -> {
+      EntityManager entityManager = jpaTm().getEntityManager();
+      Query query =
+          nativeQuery ? entityManager.createNativeQuery(sql) : entityManager.createQuery(sql);
+      if (parameters != null) {
+        parameters.forEach(query::setParameter);
+      }
+      @SuppressWarnings("unchecked")
+      Stream<T> resultStream = query.getResultStream();
+      return nativeQuery ? resultStream : resultStream.map(e -> detach(entityManager, e));
+    };
   }
 
-  static <T> RegistryQueryFactory<T> createQueryFactory(
-      QueryComposerFactory<T> queryComposerFactory) {
-    return (JpaTransactionManager jpa) ->
-        () -> queryComposerFactory.apply(jpa).withAutoDetachOnLoad(true).stream();
+  /**
+   * Returns a {@link RegistryQuery} that creates a typed JPQL query from constant text.
+   *
+   * @param parameters parameters to be substituted in the query.
+   * @param <T> Type of each row in the result set.
+   */
+  static <T> RegistryQuery<T> createQuery(
+      String jpql, @Nullable Map<String, Object> parameters, Class<T> clazz) {
+    return () -> {
+      TypedQuery<T> query = jpaTm().query(jpql, clazz);
+      if (parameters != null) {
+        parameters.forEach(query::setParameter);
+      }
+      return query.getResultStream();
+    };
+  }
+
+  /**
+   * Returns a {@link RegistryQuery} from a {@link CriteriaQuery} supplier.
+   *
+   * <p>A serializable supplier is needed in because {@link CriteriaQuery} itself must be created
+   * within a transaction, and we are not in a transaction yet when this function is called to set
+   * up the pipeline.
+   *
+   * @param <T> Type of each row in the result set.
+   */
+  static <T> RegistryQuery<T> createQuery(CriteriaQuerySupplier<T> criteriaQuery) {
+    return () -> jpaTm().query(criteriaQuery.get()).getResultStream();
   }
 
   /**

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,48 @@ 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()",
+            false,
+            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/cron/fanout?queue=replay-commit-logs-to-sql&endpoint=/_dr/task/replayCommitLogsToSql&runInEmpty]]></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/backend/WEB-INF/web.xml

@@ -237,6 +237,12 @@
     <url-pattern>/_dr/task/killCommitLogs</url-pattern>
   </servlet-mapping>
 
+  <!-- Replays Datastore commit logs to SQL. -->
+  <servlet-mapping>
+    <servlet-name>backend-servlet</servlet-name>
+    <url-pattern>/_dr/task/replayCommitLogsToSql</url-pattern>
+  </servlet-mapping>
+
   <!-- MapReduce servlet. -->
   <servlet>
     <servlet-name>mapreduce</servlet-name>

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/cron/fanout?queue=replay-commit-logs-to-sql&endpoint=/_dr/task/replayCommitLogsToSql&runInEmpty]]></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

@@ -1,6 +1,17 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <cronentries>
 
+  <cron>
+    <url>/_dr/task/rdeStaging</url>
+    <description>
+      This job generates a full RDE escrow deposit as a single gigantic XML document
+      and streams it to cloud storage. When this job has finished successfully, it'll
+      launch a separate task that uploads the deposit file to Iron Mountain via SFTP.
+    </description>
+    <schedule>every day 00:07</schedule>
+    <target>backend</target>
+  </cron>
+
   <cron>
     <url><![CDATA[/_dr/cron/readDnsQueue?jitterSeconds=45]]></url>
     <description>
@@ -100,4 +111,21 @@
     <target>backend</target>
   </cron>
 
+  <cron>
+    <url><![CDATA[/_dr/cron/fanout?queue=replay-commit-logs-to-sql&endpoint=/_dr/task/replayCommitLogsToSql&runInEmpty]]></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/commitLogCheckpoint]]></url>
+    <description>
+      This job checkpoints the commit log buckets and exports the diff since last checkpoint to GCS.
+    </description>
+    <schedule>every 3 minutes synchronized</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/cron/fanout?queue=replay-commit-logs-to-sql&endpoint=/_dr/task/replayCommitLogsToSql&runInEmpty]]></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/flows/domain/DomainDeleteFlow.java

@@ -251,8 +251,10 @@ public final class DomainDeleteFlow implements TransactionalFlow {
         buildDomainHistory(newDomain, registry, now, durationUntilDelete, inAddGracePeriod);
     updateForeignKeyIndexDeletionTime(newDomain);
     handlePendingTransferOnDelete(existingDomain, newDomain, now, domainHistory);
-    // Close the autorenew billing event and poll message. This may delete the poll message.
-    updateAutorenewRecurrenceEndTime(existingDomain, now);
+    // Close the autorenew billing event and poll message. This may delete the poll message.  Store
+    // the updated recurring billing event, we'll need it later and can't reload it.
+    BillingEvent.Recurring recurringBillingEvent =
+        updateAutorenewRecurrenceEndTime(existingDomain, now);
     // If there's a pending transfer, the gaining client's autorenew billing
     // event and poll message will already have been deleted in
     // ResourceDeleteFlow since it's listed in serverApproveEntities.
@@ -275,7 +277,8 @@ public final class DomainDeleteFlow implements TransactionalFlow {
                     newDomain.getDeletionTime().isAfter(now)
                         ? SUCCESS_WITH_ACTION_PENDING
                         : SUCCESS)
-                .setResponseExtensions(getResponseExtensions(existingDomain, now))
+                .setResponseExtensions(
+                    getResponseExtensions(recurringBillingEvent, existingDomain, now))
                 .build());
     persistEntityChanges(entityChanges);
     return responseBuilder
@@ -377,7 +380,7 @@ public final class DomainDeleteFlow implements TransactionalFlow {
 
   @Nullable
   private ImmutableList<FeeTransformResponseExtension> getResponseExtensions(
-      DomainBase existingDomain, DateTime now) {
+      BillingEvent.Recurring recurringBillingEvent, DomainBase existingDomain, DateTime now) {
     FeeTransformResponseExtension.Builder feeResponseBuilder = getDeleteResponseBuilder();
     if (feeResponseBuilder == null) {
       return ImmutableList.of();
@@ -385,7 +388,7 @@ public final class DomainDeleteFlow implements TransactionalFlow {
     ImmutableList.Builder<Credit> creditsBuilder = new ImmutableList.Builder<>();
     for (GracePeriod gracePeriod : existingDomain.getGracePeriods()) {
       if (gracePeriod.hasBillingEvent()) {
-        Money cost = getGracePeriodCost(gracePeriod, now);
+        Money cost = getGracePeriodCost(recurringBillingEvent, gracePeriod, now);
         creditsBuilder.add(Credit.create(
             cost.negated().getAmount(), FeeType.CREDIT, gracePeriod.getType().getXmlName()));
         feeResponseBuilder.setCurrency(checkNotNull(cost.getCurrencyUnit()));
@@ -398,12 +401,12 @@ public final class DomainDeleteFlow implements TransactionalFlow {
     return ImmutableList.of(feeResponseBuilder.setCredits(credits).build());
   }
 
-  private Money getGracePeriodCost(GracePeriod gracePeriod, DateTime now) {
+  private Money getGracePeriodCost(
+      BillingEvent.Recurring recurringBillingEvent, GracePeriod gracePeriod, DateTime now) {
     if (gracePeriod.getType() == GracePeriodStatus.AUTO_RENEW) {
+      // If we updated the autorenew billing event, reuse it.
       DateTime autoRenewTime =
-          tm().loadByKey(checkNotNull(gracePeriod.getRecurringBillingEvent()))
-              .getRecurrenceTimeOfYear()
-              .getLastInstanceBeforeOrAt(now);
+          recurringBillingEvent.getRecurrenceTimeOfYear().getLastInstanceBeforeOrAt(now);
       return getDomainRenewCost(targetId, autoRenewTime, 1);
     }
     return tm().loadByKey(checkNotNull(gracePeriod.getOneTimeBillingEvent())).getCost();

core/src/main/java/google/registry/flows/domain/DomainFlowUtils.java

@@ -517,8 +517,10 @@ public class DomainFlowUtils {
    * <p>This may end up deleting the poll message (if closing the message interval) or recreating it
    * (if opening the message interval). This may cause an autorenew billing event to have an end
    * time earlier than its event time (i.e. if it's being ended before it was ever triggered).
+   *
+   * <p>Returns the new autorenew recurring billing event.
    */
-  public static void updateAutorenewRecurrenceEndTime(DomainBase domain, DateTime newEndTime) {
+  public static Recurring updateAutorenewRecurrenceEndTime(DomainBase domain, DateTime newEndTime) {
     Optional<PollMessage.Autorenew> autorenewPollMessage =
         tm().loadByKeyIfPresent(domain.getAutorenewPollMessage());
 
@@ -545,8 +547,13 @@ public class DomainFlowUtils {
       tm().put(updatedAutorenewPollMessage);
     }
 
-    Recurring recurring = tm().loadByKey(domain.getAutorenewBillingEvent());
-    tm().put(recurring.asBuilder().setRecurrenceEndTime(newEndTime).build());
+    Recurring recurring =
+        tm().loadByKey(domain.getAutorenewBillingEvent())
+            .asBuilder()
+            .setRecurrenceEndTime(newEndTime)
+            .build();
+    tm().put(recurring);
+    return recurring;
   }
 
   /**

core/src/main/java/google/registry/flows/domain/DomainUpdateFlow.java

@@ -290,10 +290,7 @@ public final class DomainUpdateFlow implements TransactionalFlow {
 
   /** Some status updates cost money. Bill only once no matter how many of them are changed. */
   private Optional<BillingEvent.OneTime> createBillingEventForStatusUpdates(
-      DomainBase existingDomain,
-      DomainBase newDomain,
-      HistoryEntry historyEntry,
-      DateTime now) {
+      DomainBase existingDomain, DomainBase newDomain, DomainHistory historyEntry, DateTime now) {
     Optional<MetadataExtension> metadataExtension =
         eppInput.getSingleExtension(MetadataExtension.class);
     if (metadataExtension.isPresent() && metadataExtension.get().getRequestedByRegistrar()) {

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

@@ -16,7 +16,7 @@ package google.registry.model;
 
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Preconditions.checkState;
-import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.model.ofy.ObjectifyService.auditedOfy;
 
 import google.registry.model.ofy.ObjectifyService;
 import google.registry.util.TypeUtils.TypeInstantiator;
@@ -57,8 +57,14 @@ public interface Buildable {
         // If this object has a Long or long Objectify @Id field that is not set, set it now.
         Field idField = null;
         try {
-          idField = ModelUtils.getAllFields(instance.getClass()).get(
-              ofy().factory().getMetadata(instance.getClass()).getKeyMetadata().getIdFieldName());
+          idField =
+              ModelUtils.getAllFields(instance.getClass())
+                  .get(
+                      auditedOfy()
+                          .factory()
+                          .getMetadata(instance.getClass())
+                          .getKeyMetadata()
+                          .getIdFieldName());
         } catch (Exception e) {
           // Expected if the class is not registered with Objectify.
         }

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,7 +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.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;
@@ -75,7 +75,7 @@ public final class EntityClasses {
           ContactHistory.class,
           ContactResource.class,
           Cursor.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;
@@ -28,8 +29,6 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.flogger.FluentLogger;
 import com.googlecode.objectify.Key;
-import com.googlecode.objectify.Result;
-import com.googlecode.objectify.util.ResultNow;
 import google.registry.config.RegistryConfig;
 import google.registry.model.EppResource.BuilderWithTransferData;
 import google.registry.model.EppResource.ForeignKeyedEppResource;
@@ -42,14 +41,18 @@ 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;
 import java.util.function.Function;
+import java.util.function.Supplier;
 import javax.annotation.Nullable;
 import javax.persistence.Query;
 import org.joda.time.DateTime;
@@ -265,80 +268,147 @@ 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.
    *
-   * @return an asynchronous operation returning resource at {@code timestamp} or {@code null} if
-   *     resource is deleted or not yet created
+   * <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.
+   *
+   * @return the 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> 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);
+      return null;
     }
     // If the resource was not modified after the requested time, then use it as-is, otherwise find
-    // the most recent revision asynchronously, and return an async result that wraps that revision
-    // and returns it projected forward to exactly the desired timestamp, or null if the resource is
-    // deleted at that timestamp.
-    final Result<T> loadResult =
+    // the most recent revision and project it forward to exactly the desired timestamp, or null if
+    // the resource is deleted at that timestamp.
+    T loadedResource =
         isAtOrAfter(timestamp, resource.getUpdateTimestamp().getTimestamp())
-            ? new ResultNow<>(resource)
+            ? resource
             : loadMostRecentRevisionAtTime(resource, timestamp);
-    return () -> {
-      T loadedResource = loadResult.now();
-      return (loadedResource == null) ? null
-          : (isActive(loadedResource, timestamp)
-              ? cloneProjectedAtTime(loadedResource, timestamp)
-              : null);
-    };
+    return (loadedResource == null)
+        ? null
+        : (isActive(loadedResource, timestamp)
+            ? cloneProjectedAtTime(loadedResource, timestamp)
+            : null);
   }
 
   /**
-   * 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.
+   * Rewinds an {@link EppResource} object to a given point in time.
+   *
+   * <p>This method costs nothing if {@code resource} is already current. Otherwise it returns an
+   * async operation that performs a single fetch operation.
+   *
+   * @return an asynchronous operation returning resource at {@code timestamp} or {@code null} if
+   *     resource is deleted or not yet created
+   * @see #loadAtPointInTime(EppResource, DateTime)
+   */
+  public static <T extends EppResource> Supplier<T> loadAtPointInTimeAsync(
+      final T resource, final DateTime timestamp) {
+    return () -> loadAtPointInTime(resource, timestamp);
+  }
+
+  /**
+   * Returns 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(
+  private static <T extends EppResource> T loadMostRecentRevisionAtTime(
+      final T resource, final DateTime timestamp) {
+    if (tm().isOfy()) {
+      return loadMostRecentRevisionAtTimeDatastore(resource, timestamp);
+    } else {
+      return loadMostRecentRevisionAtTimeSql(resource, timestamp);
+    }
+  }
+
+  /**
+   * Returns 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 #loadAtPointInTimeAsync(EppResource, DateTime)
+   */
+  private static <T extends EppResource> 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));
-    return () -> {
-      CommitLogMutation mutation = mutationResult.now();
-      if (mutation != null) {
-        return ofy().load().fromEntity(mutation.getEntity());
-      }
-      logger.atSevere().log(
-          "Couldn't load mutation for revision at %s for %s, falling back to resource."
-              + " Revision: %s",
-          timestamp, resourceKey, revision);
       return resource;
-    };
+    }
+    final CommitLogMutation mutation =
+        auditedOfy().load().key(CommitLogMutation.createKey(revision, resourceKey)).now();
+    if (mutation != null) {
+      return auditedOfy().load().fromEntity(mutation.getEntity());
+    }
+    logger.atSevere().log(
+        "Couldn't load mutation for revision at %s for %s, falling back to resource."
+            + " Revision: %s",
+        timestamp, resourceKey, revision);
+    return resource;
+  }
+
+  /**
+   * Returns 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 #loadAtPointInTimeAsync(EppResource, DateTime)
+   */
+  private static <T extends EppResource> T loadMostRecentRevisionAtTimeSql(
+      T resource, DateTime timestamp) {
+    @SuppressWarnings("unchecked")
+    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 resource;
+    }
+    return 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 +436,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 +474,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

@@ -46,7 +46,6 @@ import google.registry.model.domain.DomainHistory;
 import google.registry.model.domain.GracePeriod;
 import google.registry.model.domain.rgp.GracePeriodStatus;
 import google.registry.model.domain.token.AllocationToken;
-import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.transfer.TransferData.TransferServerApproveEntity;
 import google.registry.persistence.BillingVKey.BillingEventVKey;
 import google.registry.persistence.BillingVKey.BillingRecurrenceVKey;
@@ -115,7 +114,7 @@ public abstract class BillingEvent extends ImmutableObject
   /** Entity id. */
   @Id @javax.persistence.Id Long id;
 
-  @Parent @DoNotHydrate @Transient Key<? extends HistoryEntry> parent;
+  @Parent @DoNotHydrate @Transient Key<DomainHistory> parent;
 
   /** The registrar to bill. */
   @Index
@@ -154,7 +153,7 @@ public abstract class BillingEvent extends ImmutableObject
     parent =
         Key.create(
             Key.create(DomainBase.class, domainRepoId),
-            HistoryEntry.class,
+            DomainHistory.class,
             domainHistoryRevisionId);
   }
 
@@ -192,7 +191,7 @@ public abstract class BillingEvent extends ImmutableObject
     return targetId;
   }
 
-  public Key<? extends HistoryEntry> getParentKey() {
+  public Key<DomainHistory> getParentKey() {
     return parent;
   }
 
@@ -254,12 +253,12 @@ public abstract class BillingEvent extends ImmutableObject
       return thisCastToDerived();
     }
 
-    public B setParent(HistoryEntry parent) {
+    public B setParent(DomainHistory parent) {
       getInstance().parent = Key.create(parent);
       return thisCastToDerived();
     }
 
-    public B setParent(Key<? extends HistoryEntry> parentKey) {
+    public B setParent(Key<DomainHistory> parentKey) {
       getInstance().parent = parentKey;
       return thisCastToDerived();
     }
@@ -325,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.
      */
@@ -355,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);
     }
@@ -377,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> {
 
@@ -411,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;
       }
 
@@ -445,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();
       }
     }
@@ -735,7 +776,7 @@ public abstract class BillingEvent extends ImmutableObject
      * because it is needed by one-off scrap tools that need to make billing adjustments.
      */
     public static Modification createRefundFor(
-        OneTime billingEvent, HistoryEntry historyEntry, String description) {
+        OneTime billingEvent, DomainHistory historyEntry, String description) {
       return new Builder()
           .setClientId(billingEvent.getClientId())
           .setFlags(billingEvent.getFlags())

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

@@ -17,7 +17,6 @@ package google.registry.model.common;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
 import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
 
 import com.google.common.base.Splitter;
@@ -210,12 +209,7 @@ public class Cursor extends ImmutableObject implements DatastoreAndSqlEntity {
   private static void checkValidCursorTypeForScope(
       CursorType cursorType, Key<? extends ImmutableObject> scope) {
     checkArgument(
-        cursorType
-            .getScopeClass()
-            .equals(
-                scope.equals(getCrossTldKey())
-                    ? EntityGroupRoot.class
-                    : ofy().factory().getMetadata(scope).getEntityClass()),
+        cursorType.getScopeClass().getSimpleName().equals(scope.getKind()),
         "Class required for cursor does not match scope class");
   }
 

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

@@ -36,7 +36,9 @@ import javax.persistence.Index;
 import javax.persistence.JoinColumn;
 import javax.persistence.JoinTable;
 import javax.persistence.OneToMany;
+import javax.persistence.PostLoad;
 import javax.persistence.Table;
+import org.hibernate.Hibernate;
 import org.joda.time.DateTime;
 
 /**
@@ -78,6 +80,8 @@ public class DomainBase extends DomainContent
     return super.getRepoId();
   }
 
+  // It seems like this should be FetchType.EAGER, but for some reason when we do that we get a lazy
+  // load error during the load of a domain.
   @ElementCollection
   @JoinTable(
       name = "DomainHost",
@@ -139,6 +143,16 @@ public class DomainBase extends DomainContent
     return dsData;
   }
 
+  /** Post-load method to eager load the collections. */
+  @PostLoad
+  @Override
+  protected void postLoad() {
+    super.postLoad();
+    // TODO(b/188044616): Determine why Eager loading doesn't work here.
+    Hibernate.initialize(dsData);
+    Hibernate.initialize(gracePeriods);
+  }
+
   @Override
   public VKey<DomainBase> createVKey() {
     return VKey.create(DomainBase.class, getRepoId(), Key.create(this));

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

@@ -338,8 +338,7 @@ public class DomainContent extends EppResource
   }
 
   @PostLoad
-  @SuppressWarnings("UnusedMethod")
-  private void postLoad() {
+  protected void postLoad() {
     // Reconstitute the contact list.
     ImmutableSet.Builder<DesignatedContact> contactsBuilder = new ImmutableSet.Builder<>();
 

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;
@@ -54,6 +56,7 @@ import javax.persistence.JoinTable;
 import javax.persistence.OneToMany;
 import javax.persistence.PostLoad;
 import javax.persistence.Table;
+import org.hibernate.Hibernate;
 
 /**
  * A persisted history entry representing an EPP modification to a domain.
@@ -247,20 +250,40 @@ 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());
         }
       }
     }
+
+    // TODO(b/188044616): Determine why Eager loading doesn't work here.
+    Hibernate.initialize(domainTransactionRecords);
+    Hibernate.initialize(nsHosts);
+    Hibernate.initialize(dsDataHistories);
+    Hibernate.initialize(gracePeriodHistories);
   }
 
   // In Datastore, save as a HistoryEntry object regardless of this object's type

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

@@ -222,9 +222,17 @@ public class DatastoreTransactionManager implements TransactionManager {
                 entry -> keyMap.get(entry.getKey()), entry -> toSqlEntity(entry.getValue())));
   }
 
+  @SuppressWarnings("unchecked")
   @Override
   public <T> ImmutableList<T> loadByEntitiesIfPresent(Iterable<T> entities) {
-    return ImmutableList.copyOf(getOfy().load().entities(entities).values());
+    return getOfy()
+        .load()
+        .entities(toDatastoreEntities(ImmutableList.copyOf(entities)))
+        .values()
+        .stream()
+        .map(DatastoreTransactionManager::toSqlEntity)
+        .map(entity -> (T) entity)
+        .collect(toImmutableList());
   }
 
   @Override
@@ -250,6 +258,7 @@ public class DatastoreTransactionManager implements TransactionManager {
     return result;
   }
 
+  @SuppressWarnings("unchecked")
   @Override
   public <T> T loadByEntity(T entity) {
     return (T) toSqlEntity(auditedOfy().load().entity(toDatastoreEntity(entity)).now());
@@ -297,8 +306,9 @@ public class DatastoreTransactionManager implements TransactionManager {
   }
 
   @Override
-  public void delete(Object entity) {
+  public <T> T delete(T entity) {
     syncIfTransactionless(getOfy().delete().entity(toDatastoreEntity(entity)));
+    return entity;
   }
 
   @Override
@@ -470,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/poll/PollMessageExternalKeyConverter.java

@@ -14,10 +14,10 @@
 
 package google.registry.model.poll;
 
-import static google.registry.model.ofy.ObjectifyService.ofy;
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
 
 import com.google.common.base.Splitter;
-import com.google.common.collect.ImmutableBiMap;
+import com.google.common.collect.ImmutableMap;
 import com.googlecode.objectify.Key;
 import google.registry.model.EppResource;
 import google.registry.model.contact.ContactResource;
@@ -26,6 +26,7 @@ import google.registry.model.host.HostResource;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.persistence.VKey;
 import java.util.List;
+import java.util.Map;
 
 /**
  * A converter between external key strings for {@link PollMessage}s (i.e. what registrars use to
@@ -49,24 +50,23 @@ public class PollMessageExternalKeyConverter {
   /** An exception thrown when an external key cannot be parsed. */
   public static class PollMessageExternalKeyParseException extends RuntimeException {}
 
-  /**
-   * A map of IDs used in external keys corresponding to which EppResource class the poll message
-   * belongs to.
-   */
-  public static final ImmutableBiMap<Class<? extends EppResource>, Long> EXTERNAL_KEY_CLASS_ID_MAP =
-      ImmutableBiMap.of(
-          DomainBase.class, 1L,
-          ContactResource.class, 2L,
-          HostResource.class, 3L);
+  /** Maps that detail the correspondence between EppResource classes and external IDs. */
+  private static final ImmutableMap<Long, Class<? extends EppResource>> ID_TO_CLASS_MAP =
+      ImmutableMap.of(
+          1L, DomainBase.class,
+          2L, ContactResource.class,
+          3L, HostResource.class);
+
+  private static final ImmutableMap<String, Long> KEY_KIND_TO_ID_MAP =
+      ID_TO_CLASS_MAP.entrySet().stream()
+          .collect(toImmutableMap(entry -> entry.getValue().getSimpleName(), Map.Entry::getKey));
 
   /** Returns an external poll message ID for the given poll message. */
   public static String makePollMessageExternalId(PollMessage pollMessage) {
     @SuppressWarnings("unchecked")
     Key<EppResource> ancestorResource =
         (Key<EppResource>) (Key<?>) pollMessage.getParentKey().getParent();
-    long externalKeyClassId =
-        EXTERNAL_KEY_CLASS_ID_MAP.get(
-            ofy().factory().getMetadata(ancestorResource.getKind()).getEntityClass());
+    long externalKeyClassId = KEY_KIND_TO_ID_MAP.get(ancestorResource.getKind());
     return String.format(
         "%d-%s-%d-%d-%d",
         externalKeyClassId,
@@ -92,8 +92,7 @@ public class PollMessageExternalKeyConverter {
       throw new PollMessageExternalKeyParseException();
     }
     try {
-      Class<?> resourceClazz =
-          EXTERNAL_KEY_CLASS_ID_MAP.inverse().get(Long.parseLong(idComponents.get(0)));
+      Class<?> resourceClazz = ID_TO_CLASS_MAP.get(Long.parseLong(idComponents.get(0)));
       if (resourceClazz == null) {
         throw new PollMessageExternalKeyParseException();
       }

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

@@ -47,6 +47,7 @@ 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.OnLoad;
 import com.googlecode.objectify.annotation.OnSave;
 import com.googlecode.objectify.annotation.Parent;
 import google.registry.model.Buildable;
@@ -111,6 +112,26 @@ public class Registry extends ImmutableObject implements Buildable, DatastoreAnd
   @PostLoad
   void postLoad() {
     tldStr = tldStrId;
+    // TODO(sarahbot@): Remove the rest of this method after this data migration is complete
+    if (premiumListName != null) {
+      premiumList = Key.create(getCrossTldKey(), PremiumList.class, premiumListName);
+    }
+    if (reservedListNames != null) {
+      reservedLists =
+          reservedListNames.stream()
+              .map(name -> Key.create(getCrossTldKey(), ReservedList.class, name))
+              .collect(toImmutableSet());
+    }
+  }
+
+  // TODO(sarahbot@): Remove this method after this data migration is complete
+  @OnLoad
+  void onLoad() {
+    if (reservedLists != null) {
+      reservedListNames =
+          reservedLists.stream().map(key -> key.getName()).collect(toImmutableSet());
+    }
+    premiumListName = premiumList == null ? null : premiumList.getName();
   }
 
   /** The suffix that identifies roids as belonging to this specific tld, e.g. -HOW for .how. */
@@ -388,17 +409,37 @@ public class Registry extends ImmutableObject implements Buildable, DatastoreAnd
   CreateAutoTimestamp creationTime = CreateAutoTimestamp.create(null);
 
   /** The set of reserved lists that are applicable to this registry. */
-  @Column(name = "reserved_list_names")
-  Set<Key<ReservedList>> reservedLists;
+  @Transient Set<Key<ReservedList>> reservedLists;
 
-  /** Retrieves an ImmutableSet of all ReservedLists associated with this tld. */
+  /** The set of reserved list names that are applicable to this registry. */
+  @Column(name = "reserved_list_names")
+  Set<String> reservedListNames;
+
+  /**
+   * Retrieves an ImmutableSet of all ReservedLists associated with this TLD.
+   *
+   * <p>This set contains only the names of the list and not a reference to the lists. Updates to a
+   * reserved list in Cloud SQL are saved as a new ReservedList entity. When using the ReservedList
+   * for a registry, the database should be queried for the entity with this name that has the
+   * largest revision ID.
+   */
   public ImmutableSet<Key<ReservedList>> getReservedLists() {
     return nullToEmptyImmutableCopy(reservedLists);
   }
 
   /** The static {@link PremiumList} for this TLD, if there is one. */
+  @Transient Key<PremiumList> premiumList;
+
+  /**
+   * The name of the {@link PremiumList} for this TLD, if there is one.
+   *
+   * <p>This is only the name of the list and not a reference to the list. Updates to the premium
+   * list in Cloud SQL are saved as a new PremiumList entity. When using the PremiumList for a
+   * registry, the database should be queried for the entity with this name that has the largest
+   * revision ID.
+   */
   @Column(name = "premium_list_name", nullable = true)
-  Key<PremiumList> premiumList;
+  String premiumListName;
 
   /** Should RDE upload a nightly escrow deposit for this TLD? */
   @Column(nullable = false)
@@ -879,21 +920,26 @@ public class Registry extends ImmutableObject implements Buildable, DatastoreAnd
     public Builder setReservedLists(Set<ReservedList> reservedLists) {
       checkArgumentNotNull(reservedLists, "reservedLists must not be null");
       ImmutableSet.Builder<Key<ReservedList>> builder = new ImmutableSet.Builder<>();
+      ImmutableSet.Builder<String> nameBuilder = new ImmutableSet.Builder<>();
       for (ReservedList reservedList : reservedLists) {
         builder.add(Key.create(reservedList));
+        nameBuilder.add(reservedList.getName());
       }
       getInstance().reservedLists = builder.build();
+      getInstance().reservedListNames = nameBuilder.build();
       return this;
     }
 
-    public Builder setPremiumList(PremiumList premiumList) {
+    public Builder setPremiumList(@Nullable PremiumList premiumList) {
       getInstance().premiumList = (premiumList == null) ? null : Key.create(premiumList);
+      getInstance().premiumListName = (premiumList == null) ? null : premiumList.getName();
       return this;
     }
 
     @VisibleForTesting
     public Builder setPremiumListKey(@Nullable Key<PremiumList> premiumList) {
       getInstance().premiumList = premiumList;
+      getInstance().premiumListName = (premiumList == null) ? null : premiumList.getName();
       return this;
     }
 

core/src/main/java/google/registry/model/registry/label/BaseDomainLabelList.java

@@ -38,7 +38,6 @@ import google.registry.model.ImmutableObject;
 import google.registry.model.annotations.InCrossTld;
 import google.registry.model.common.EntityGroupRoot;
 import google.registry.model.registry.Registry;
-import google.registry.model.registry.label.ReservedList.ReservedListEntry;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -56,8 +55,8 @@ import org.joda.time.DateTime;
  * Base class for {@link ReservedList} and {@link PremiumList} objects stored in Datastore.
  *
  * @param <T> The type of the root value being listed, e.g. {@link ReservationType}.
- * @param <R> The type of domain label entry being listed, e.g. {@link ReservedListEntry} (note,
- *     must subclass {@link DomainLabelEntry}.
+ * @param <R> The type of domain label entry being listed, e.g. {@link
+ *     ReservedList.ReservedListEntry} (note, must subclass {@link DomainLabelEntry}.
  */
 @MappedSuperclass
 @InCrossTld

core/src/main/java/google/registry/model/registry/label/DomainLabelEntry.java

@@ -36,7 +36,8 @@ public abstract class DomainLabelEntry<T extends Comparable<?>, D extends Domain
     extends ImmutableObject implements Comparable<D> {
 
   @Id
-  @Column(name = "domain_label", insertable = false, updatable = false)
+  @javax.persistence.Id
+  @Column(name = "domain_label", nullable = false)
   String label;
 
   String comment;

core/src/main/java/google/registry/model/registry/label/ReservedList.java

@@ -17,9 +17,13 @@ package google.registry.model.registry.label;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Strings.isNullOrEmpty;
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static google.registry.config.RegistryConfig.getDomainLabelListCacheDuration;
+import static google.registry.model.ImmutableObject.Insignificant;
 import static google.registry.model.registry.label.ReservationType.FULLY_BLOCKED;
+import static google.registry.persistence.transaction.QueryComposer.Comparator.EQ;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.util.CollectionUtils.nullToEmpty;
 import static org.joda.time.DateTimeZone.UTC;
 
@@ -40,19 +44,19 @@ import google.registry.model.annotations.ReportedOn;
 import google.registry.model.registry.Registry;
 import google.registry.model.registry.label.DomainLabelMetrics.MetricsReservedListMatch;
 import google.registry.schema.replay.NonReplicatedEntity;
+import java.io.Serializable;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.ExecutionException;
 import javax.annotation.Nullable;
-import javax.persistence.CollectionTable;
 import javax.persistence.Column;
-import javax.persistence.ElementCollection;
-import javax.persistence.Embeddable;
+import javax.persistence.Id;
 import javax.persistence.Index;
-import javax.persistence.JoinColumn;
-import javax.persistence.MapKeyColumn;
+import javax.persistence.PostPersist;
+import javax.persistence.PreRemove;
 import javax.persistence.Table;
+import javax.persistence.Transient;
 import org.joda.time.DateTime;
 
 /**
@@ -71,25 +75,60 @@ public final class ReservedList
     extends BaseDomainLabelList<ReservationType, ReservedList.ReservedListEntry>
     implements NonReplicatedEntity {
 
+  /**
+   * Mapping from domain name to its reserved list info.
+   *
+   * <p>This field requires special treatment since we want to lazy load it. We have to remove it
+   * from the immutability contract so we can modify it after construction and we have to handle the
+   * database processing on our own so we can detach it after load.
+   */
   @Mapify(ReservedListEntry.LabelMapper.class)
-  @ElementCollection
-  @CollectionTable(
-      name = "ReservedEntry",
-      joinColumns = @JoinColumn(name = "revisionId", referencedColumnName = "revisionId"))
-  @MapKeyColumn(name = "domain_label")
+  @Insignificant
+  @Transient
   Map<String, ReservedListEntry> reservedListMap;
 
   @Column(nullable = false)
   boolean shouldPublish = true;
 
+  @PreRemove
+  void preRemove() {
+    jpaTm()
+        .query("DELETE FROM ReservedEntry WHERE revision_id = :revisionId")
+        .setParameter("revisionId", revisionId)
+        .executeUpdate();
+  }
+
+  /**
+   * Hibernate hook called on the insert of a new ReservedList. Stores the associated {@link
+   * ReservedEntry}'s.
+   *
+   * <p>We need to persist the list entries, but only on the initial insert (not on update) since
+   * the entries themselves never get changed, so we only annotate it with {@link PostPersist}, not
+   * {@link PostUpdate}.
+   */
+  @PostPersist
+  void postPersist() {
+    if (reservedListMap != null) {
+      reservedListMap.values().stream()
+          .forEach(
+              entry -> {
+                // We can safely change the revision id since it's "Insignificant".
+                entry.revisionId = revisionId;
+                jpaTm().insert(entry);
+              });
+    }
+  }
+
   /**
    * A reserved list entry entity, persisted to Datastore, that represents a single label and its
    * reservation type.
    */
   @Embed
-  @Embeddable
+  @javax.persistence.Entity(name = "ReservedEntry")
   public static class ReservedListEntry extends DomainLabelEntry<ReservationType, ReservedListEntry>
-      implements Buildable {
+      implements Buildable, NonReplicatedEntity, Serializable {
+
+    @Insignificant @Id Long revisionId;
 
     @Column(nullable = false)
     ReservationType reservationType;
@@ -164,8 +203,24 @@ public final class ReservedList
     return shouldPublish;
   }
 
-  /** Returns a {@link Map} of domain labels to {@link ReservedListEntry}. */
-  public ImmutableMap<String, ReservedListEntry> getReservedListEntries() {
+  /**
+   * Returns a {@link Map} of domain labels to {@link ReservedListEntry}.
+   *
+   * <p>Note that this involves a database fetch of a potentially large number of elements and
+   * should be avoided unless necessary.
+   */
+  public synchronized ImmutableMap<String, ReservedListEntry> getReservedListEntries() {
+    if (reservedListMap == null) {
+      reservedListMap =
+          jpaTm()
+              .transact(
+                  () ->
+                      jpaTm()
+                          .createQueryComposer(ReservedListEntry.class)
+                          .where("revisionId", EQ, revisionId)
+                          .stream()
+                          .collect(toImmutableMap(ReservedListEntry::getLabel, e -> e)));
+    }
     return ImmutableMap.copyOf(nullToEmpty(reservedListMap));
   }
 

core/src/main/java/google/registry/model/registry/label/ReservedListDao.java

@@ -52,9 +52,8 @@ public class ReservedListDao {
             () ->
                 jpaTm()
                     .query(
-                        "FROM ReservedList rl LEFT JOIN FETCH rl.reservedListMap WHERE"
-                            + " rl.revisionId IN (SELECT MAX(revisionId) FROM ReservedList subrl"
-                            + " WHERE subrl.name = :name)",
+                        "FROM ReservedList WHERE revisionId IN "
+                            + "(SELECT MAX(revisionId) FROM ReservedList WHERE name = :name)",
                         ReservedList.class)
                     .setParameter("name", reservedListName)
                     .getResultStream()

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

@@ -14,6 +14,7 @@
 
 package google.registry.model.reporting;
 
+import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static google.registry.model.ofy.ObjectifyService.auditedOfy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
@@ -49,7 +50,7 @@ import org.joda.time.DateTime;
 public class HistoryEntryDao {
 
   /** Loads all history objects in the times specified, including all types. */
-  public static ImmutableList<? extends HistoryEntry> loadAllHistoryObjects(
+  public static ImmutableList<HistoryEntry> loadAllHistoryObjects(
       DateTime afterTime, DateTime beforeTime) {
     if (tm().isOfy()) {
       return Streams.stream(
@@ -77,13 +78,22 @@ public class HistoryEntryDao {
   }
 
   /** Loads all history objects corresponding to the given {@link EppResource}. */
-  public static ImmutableList<? extends HistoryEntry> loadHistoryObjectsForResource(
+  public static ImmutableList<HistoryEntry> loadHistoryObjectsForResource(
       VKey<? extends EppResource> parentKey) {
     return loadHistoryObjectsForResource(parentKey, START_OF_TIME, END_OF_TIME);
   }
 
+  /**
+   * Loads all history objects corresponding to the given {@link EppResource} and casted to the
+   * appropriate subclass.
+   */
+  public static <T extends HistoryEntry> ImmutableList<T> loadHistoryObjectsForResource(
+      VKey<? extends EppResource> parentKey, Class<T> subclazz) {
+    return loadHistoryObjectsForResource(parentKey, START_OF_TIME, END_OF_TIME, subclazz);
+  }
+
   /** Loads all history objects in the time period specified for the given {@link EppResource}. */
-  public static ImmutableList<? extends HistoryEntry> loadHistoryObjectsForResource(
+  public static ImmutableList<HistoryEntry> loadHistoryObjectsForResource(
       VKey<? extends EppResource> parentKey, DateTime afterTime, DateTime beforeTime) {
     if (tm().isOfy()) {
       return Streams.stream(
@@ -102,8 +112,35 @@ public class HistoryEntryDao {
     }
   }
 
+  /**
+   * Loads all history objects in the time period specified for the given {@link EppResource} and
+   * casted to the appropriate subclass.
+   *
+   * <p>Note that the subclass must be explicitly provided because we need the compile time
+   * information of T to return an {@code ImmutableList<T>}, even though at runtime we can call
+   * {@link #getHistoryClassFromParent(Class)} to obtain it, which we also did to confirm that the
+   * provided subclass is indeed correct.
+   */
+  public static <T extends HistoryEntry> ImmutableList<T> loadHistoryObjectsForResource(
+      VKey<? extends EppResource> parentKey,
+      DateTime afterTime,
+      DateTime beforeTime,
+      Class<T> subclazz) {
+    Class<? extends HistoryEntry> expectedSubclazz = getHistoryClassFromParent(parentKey.getKind());
+    checkArgument(
+        subclazz.equals(expectedSubclazz),
+        "The supplied HistoryEntry subclass %s is incompatible with the EppResource %s, "
+            + "use %s instead",
+        subclazz.getSimpleName(),
+        parentKey.getKind().getSimpleName(),
+        expectedSubclazz.getSimpleName());
+    return loadHistoryObjectsForResource(parentKey, afterTime, beforeTime).stream()
+        .map(subclazz::cast)
+        .collect(toImmutableList());
+  }
+
   /** Loads all history objects from all time from the given registrars. */
-  public static Iterable<? extends HistoryEntry> loadHistoryObjectsByRegistrars(
+  public static Iterable<HistoryEntry> loadHistoryObjectsByRegistrars(
       ImmutableCollection<String> registrarIds) {
     if (tm().isOfy()) {
       return auditedOfy()
@@ -124,18 +161,17 @@ public class HistoryEntryDao {
     }
   }
 
-  private static Stream<? extends HistoryEntry> loadHistoryObjectFromSqlByRegistrars(
-      Class<? extends HistoryEntry> historyClass, ImmutableCollection<String> registrarIds) {
+  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())
         .getResultStream();
   }
 
-  private static ImmutableList<? extends HistoryEntry> loadHistoryObjectsForResourceFromSql(
+  private static ImmutableList<HistoryEntry> loadHistoryObjectsForResourceFromSql(
       VKey<? extends EppResource> parentKey, DateTime afterTime, DateTime beforeTime) {
     // The class we're searching from is based on which parent type (e.g. Domain) we have
     Class<? extends HistoryEntry> historyClass = getHistoryClassFromParent(parentKey.getKind());
@@ -147,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(
@@ -174,12 +211,11 @@ public class HistoryEntryDao {
         : historyClass.equals(DomainHistory.class) ? "domainRepoId" : "hostRepoId";
   }
 
-  private static List<? extends HistoryEntry> loadAllHistoryObjectsFromSql(
-      Class<? extends HistoryEntry> historyClass, DateTime afterTime, DateTime beforeTime) {
+  private static <T extends HistoryEntry> List<T> loadAllHistoryObjectsFromSql(
+      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/tmch/ClaimsEntry.java

@@ -0,0 +1,56 @@
+// 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.tmch;
+
+import google.registry.model.ImmutableObject;
+import google.registry.schema.replay.NonReplicatedEntity;
+import java.io.Serializable;
+import javax.persistence.Column;
+import javax.persistence.Entity;
+import javax.persistence.Id;
+
+/**
+ * Claims entry record, used by ClaimsList for persistence.
+ *
+ * <p>It would be preferable to have this nested in {@link ClaimsList}, but for some reason
+ * hibernate won't generate this into the schema in this case. We may not care, as we only use the
+ * generated schema for informational purposes and persistence against the actual schema seems to
+ * work.
+ */
+@Entity(name = "ClaimsEntry")
+class ClaimsEntry extends ImmutableObject implements NonReplicatedEntity, Serializable {
+  @Id private Long revisionId;
+  @Id private String domainLabel;
+
+  @Column(nullable = false)
+  private String claimKey;
+
+  /** Default constructor for Hibernate. */
+  ClaimsEntry() {}
+
+  ClaimsEntry(Long revisionId, String domainLabel, String claimKey) {
+    this.revisionId = revisionId;
+    this.domainLabel = domainLabel;
+    this.claimKey = claimKey;
+  }
+
+  String getDomainLabel() {
+    return domainLabel;
+  }
+
+  String getClaimKey() {
+    return claimKey;
+  }
+}

core/src/main/java/google/registry/model/tmch/ClaimsList.java

@@ -16,13 +16,15 @@ package google.registry.model.tmch;
 
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Preconditions.checkState;
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static google.registry.model.ofy.ObjectifyService.allocateId;
 import static google.registry.model.ofy.ObjectifyService.auditedOfy;
+import static google.registry.persistence.transaction.QueryComposer.Comparator.EQ;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableMap;
 import com.googlecode.objectify.Key;
-import com.googlecode.objectify.annotation.EmbedMap;
 import com.googlecode.objectify.annotation.Entity;
 import com.googlecode.objectify.annotation.Id;
 import com.googlecode.objectify.annotation.Ignore;
@@ -39,13 +41,11 @@ import google.registry.schema.replay.NonReplicatedEntity;
 import java.util.Map;
 import java.util.Optional;
 import javax.annotation.Nullable;
-import javax.persistence.CollectionTable;
 import javax.persistence.Column;
-import javax.persistence.ElementCollection;
 import javax.persistence.GeneratedValue;
 import javax.persistence.GenerationType;
-import javax.persistence.JoinColumn;
-import javax.persistence.MapKeyColumn;
+import javax.persistence.PostPersist;
+import javax.persistence.PreRemove;
 import javax.persistence.Table;
 import javax.persistence.Transient;
 import org.joda.time.DateTime;
@@ -94,15 +94,40 @@ public class ClaimsList extends ImmutableObject implements NonReplicatedEntity {
   @Column(name = "tmdb_generation_time", nullable = false)
   DateTime creationTime;
 
-  /** A map from labels to claims keys. */
-  @EmbedMap
-  @ElementCollection
-  @CollectionTable(
-      name = "ClaimsEntry",
-      joinColumns = @JoinColumn(name = "revisionId", referencedColumnName = "revisionId"))
-  @MapKeyColumn(name = "domainLabel", nullable = false)
-  @Column(name = "claimKey", nullable = false)
-  Map<String, String> labelsToKeys;
+  /**
+   * A map from labels to claims keys.
+   *
+   * <p>This field requires special treatment since we want to lazy load it. We have to remove it
+   * from the immutability contract so we can modify it after construction and we have to handle the
+   * database processing on our own so we can detach it after load.
+   */
+  @Insignificant @Transient ImmutableMap<String, String> labelsToKeys;
+
+  @PreRemove
+  void preRemove() {
+    jpaTm()
+        .query("DELETE FROM ClaimsEntry WHERE revision_id = :revisionId")
+        .setParameter("revisionId", revisionId)
+        .executeUpdate();
+  }
+
+  /**
+   * Hibernate hook called on the insert of a new ReservedList. Stores the associated {@link
+   * ReservedListEntry}'s.
+   *
+   * <p>We need to persist the list entries, but only on the initial insert (not on update) since
+   * the entries themselves never get changed, so we only annotate it with {@link PostPersist}, not
+   * {@link PostUpdate}.
+   */
+  @PostPersist
+  void postPersist() {
+    if (labelsToKeys != null) {
+      labelsToKeys.entrySet().stream()
+          .forEach(
+              entry ->
+                  jpaTm().insert(new ClaimsEntry(revisionId, entry.getKey(), entry.getValue())));
+    }
+  }
 
   /** Returns the revision id of this claims list, or throws exception if it is null. */
   public Long getRevisionId() {
@@ -126,22 +151,69 @@ public class ClaimsList extends ImmutableObject implements NonReplicatedEntity {
     return creationTimestamp.getTimestamp();
   }
 
-  /** Returns the claim key for a given domain if there is one, empty otherwise. */
+  /**
+   * Returns the claim key for a given domain if there is one, empty otherwise.
+   *
+   * <p>Note that this may do a database query. For checking multiple keys against the claims list
+   * it may be more efficient to use {@link #getLabelsToKeys()} first, as this will prefetch all
+   * entries and cache them locally.
+   */
   public Optional<String> getClaimKey(String label) {
-    return Optional.ofNullable(labelsToKeys.get(label));
+    if (labelsToKeys != null) {
+      return Optional.ofNullable(labelsToKeys.get(label));
+    }
+    return jpaTm()
+        .transact(
+            () ->
+                jpaTm()
+                    .createQueryComposer(ClaimsEntry.class)
+                    .where("revisionId", EQ, revisionId)
+                    .where("domainLabel", EQ, label)
+                    .first()
+                    .map(ClaimsEntry::getClaimKey));
   }
 
-  /** Returns an {@link Map} mapping domain label to its lookup key. */
+  /**
+   * Returns an {@link Map} mapping domain label to its lookup key.
+   *
+   * <p>Note that this involves a database fetch of a potentially large number of elements and
+   * should be avoided unless necessary.
+   */
   public ImmutableMap<String, String> getLabelsToKeys() {
-    return ImmutableMap.copyOf(labelsToKeys);
+    if (labelsToKeys == null) {
+      labelsToKeys =
+          jpaTm()
+              .transact(
+                  () ->
+                      jpaTm()
+                          .createQueryComposer(ClaimsEntry.class)
+                          .where("revisionId", EQ, revisionId)
+                          .stream()
+                          .collect(
+                              toImmutableMap(
+                                  ClaimsEntry::getDomainLabel, ClaimsEntry::getClaimKey)));
+    }
+    return labelsToKeys;
   }
 
-  /** Returns the number of claims. */
-  public int size() {
+  /**
+   * Returns the number of claims.
+   *
+   * <p>Note that this will perform a database "count" query if the label to key map has not been
+   * previously cached by calling {@link #getLabelsToKeys()}.
+   */
+  public long size() {
+    if (labelsToKeys == null) {
+      return jpaTm()
+          .createQueryComposer(ClaimsEntry.class)
+          .where("revisionId", EQ, revisionId)
+          .count();
+    }
     return labelsToKeys.size();
   }
 
-  public static ClaimsList create(DateTime tmdbGenerationTime, Map<String, String> labelsToKeys) {
+  public static ClaimsList create(
+      DateTime tmdbGenerationTime, ImmutableMap<String, String> labelsToKeys) {
     ClaimsList instance = new ClaimsList();
     instance.id = allocateId();
     instance.creationTime = checkNotNull(tmdbGenerationTime);

core/src/main/java/google/registry/model/tmch/ClaimsListDao.java

@@ -14,6 +14,7 @@
 
 package google.registry.model.tmch;
 
+import static google.registry.persistence.transaction.QueryComposer.Comparator.EQ;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
 
@@ -40,13 +41,9 @@ public class ClaimsListDao {
                       .query("SELECT MAX(revisionId) FROM ClaimsList", Long.class)
                       .getSingleResult();
               return jpaTm()
-                  .query(
-                      "FROM ClaimsList cl LEFT JOIN FETCH cl.labelsToKeys WHERE cl.revisionId ="
-                          + " :revisionId",
-                      ClaimsList.class)
-                  .setParameter("revisionId", revisionId)
-                  .getResultStream()
-                  .findFirst();
+                  .createQueryComposer(ClaimsList.class)
+                  .where("revisionId", EQ, revisionId)
+                  .first();
             })
         .orElse(ClaimsList.create(START_OF_TIME, ImmutableMap.of()));
   }

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/module/backend/BackendRequestComponent.java

@@ -20,6 +20,7 @@ import google.registry.backup.BackupModule;
 import google.registry.backup.CommitLogCheckpointAction;
 import google.registry.backup.DeleteOldCommitLogsAction;
 import google.registry.backup.ExportCommitLogDiffAction;
+import google.registry.backup.ReplayCommitLogsToSqlAction;
 import google.registry.batch.BatchModule;
 import google.registry.batch.DeleteContactsAndHostsAction;
 import google.registry.batch.DeleteExpiredDomainsAction;
@@ -186,6 +187,8 @@ interface BackendRequestComponent {
 
   RelockDomainAction relockDomainAction();
 
+  ReplayCommitLogsToSqlAction replayCommitLogsToSqlAction();
+
   ResaveAllEppResourcesAction resaveAllEppResourcesAction();
 
   ResaveEntityAction resaveEntityAction();

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,20 +44,32 @@ 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;
 import java.util.NoSuchElementException;
 import java.util.Optional;
+import java.util.Set;
 import java.util.function.Supplier;
 import java.util.stream.Stream;
 import java.util.stream.StreamSupport;
+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;
@@ -112,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
@@ -269,8 +287,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     assertInTransaction();
     // Necessary due to the changes in HistoryEntry representation during the migration to SQL
     Object toPersist = toSqlEntity(entity);
-    getEntityManager().persist(toPersist);
-    transactionInfo.get().addUpdate(toPersist);
+    transactionInfo.get().insertObject(toPersist);
   }
 
   @Override
@@ -299,8 +316,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     assertInTransaction();
     // Necessary due to the changes in HistoryEntry representation during the migration to SQL
     Object toPersist = toSqlEntity(entity);
-    getEntityManager().merge(toPersist);
-    transactionInfo.get().addUpdate(toPersist);
+    transactionInfo.get().updateObject(toPersist);
   }
 
   @Override
@@ -339,8 +355,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     checkArgument(exists(entity), "Given entity does not exist");
     // Necessary due to the changes in HistoryEntry representation during the migration to SQL
     Object toPersist = toSqlEntity(entity);
-    getEntityManager().merge(toPersist);
-    transactionInfo.get().addUpdate(toPersist);
+    transactionInfo.get().updateObject(toPersist);
   }
 
   @Override
@@ -397,7 +412,8 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   public <T> Optional<T> loadByKeyIfPresent(VKey<T> key) {
     checkArgumentNotNull(key, "key must be specified");
     assertInTransaction();
-    return Optional.ofNullable(getEntityManager().find(key.getKind(), key.getSqlKey()));
+    return Optional.ofNullable(getEntityManager().find(key.getKind(), key.getSqlKey()))
+        .map(this::detach);
   }
 
   @Override
@@ -411,7 +427,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
         .map(
             key ->
                 new SimpleEntry<VKey<? extends T>, T>(
-                    key, getEntityManager().find(key.getKind(), key.getSqlKey())))
+                    key, detach(getEntityManager().find(key.getKind(), key.getSqlKey()))))
         .filter(entry -> entry.getValue() != null)
         .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue));
   }
@@ -433,7 +449,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     if (result == null) {
       throw new NoSuchElementException(key.toString());
     }
-    return result;
+    return detach(result);
   }
 
   @Override
@@ -456,11 +472,14 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     assertInTransaction();
     // If the caller requested a HistoryEntry, load the corresponding *History class
     T possibleChild = toSqlEntity(entity);
-    return (T)
-        loadByKey(
-            VKey.createSql(
-                possibleChild.getClass(),
-                emf.getPersistenceUnitUtil().getIdentifier(possibleChild)));
+    @SuppressWarnings("unchecked")
+    T returnValue =
+        (T)
+            loadByKey(
+                VKey.createSql(
+                    possibleChild.getClass(),
+                    emf.getPersistenceUnitUtil().getIdentifier(possibleChild)));
+    return returnValue;
   }
 
   @Override
@@ -472,12 +491,11 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   public <T> ImmutableList<T> loadAllOf(Class<T> clazz) {
     checkArgumentNotNull(clazz, "clazz must be specified");
     assertInTransaction();
-    return ImmutableList.copyOf(
-        getEntityManager()
-            .createQuery(
-                String.format("SELECT entity FROM %s entity", getEntityType(clazz).getName()),
-                clazz)
-            .getResultList());
+    return getEntityManager()
+        .createQuery(String.format("FROM %s", getEntityType(clazz).getName()), clazz)
+        .getResultStream()
+        .map(this::detach)
+        .collect(toImmutableList());
   }
 
   @Override
@@ -492,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) {
@@ -524,18 +542,21 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   }
 
   @Override
-  public void delete(Object entity) {
+  public <T> T delete(T entity) {
     checkArgumentNotNull(entity, "entity must be specified");
     if (isEntityOfIgnoredClass(entity)) {
-      return;
+      return entity;
     }
     assertInTransaction();
     entity = toSqlEntity(entity);
-    Object managedEntity = entity;
+    T managedEntity = entity;
     if (!getEntityManager().contains(entity)) {
+      // We don't add the entity to "objectsToSave": once deleted, the object should never be
+      // returned as a result of the query or lookup.
       managedEntity = getEntityManager().merge(entity);
     }
     getEntityManager().remove(managedEntity);
+    return managedEntity;
   }
 
   @Override
@@ -555,7 +576,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
 
   @Override
   public <T> QueryComposer<T> createQueryComposer(Class<T> entity) {
-    return new JpaQueryComposerImpl<T>(entity, getEntityManager());
+    return new JpaQueryComposerImpl<T>(entity);
   }
 
   @Override
@@ -658,6 +679,48 @@ 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) {
+    if (entity != null) {
+
+      // If the entity was previously persisted or merged, we have to throw an exception.
+      if (transactionInfo.get().willSave(entity)) {
+        throw new IllegalStateException("Inserted/updated object reloaded: " + entity);
+      }
+
+      getEntityManager().detach(entity);
+    }
+    return entity;
+  }
+
   private static class TransactionInfo {
     EntityManager entityManager;
     boolean inTransaction = false;
@@ -666,6 +729,12 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     // Serializable representation of the transaction to be persisted in the Transaction table.
     Transaction.Builder contentsBuilder;
 
+    // The set of entity objects that have been either persisted (via insert()) or merged (via
+    // put()/update()).  If the entity manager returns these as a result of a find() or query
+    // operation, we can not detach them -- detaching removes them from the transaction and causes
+    // them to not be saved to the database -- so we throw an exception instead.
+    Set<Object> objectsToSave = Collections.newSetFromMap(new IdentityHashMap<Object, Boolean>());
+
     /** Start a new transaction. */
     private void start(Clock clock) {
       checkArgumentNotNull(clock);
@@ -680,6 +749,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
       inTransaction = false;
       transactionTime = null;
       contentsBuilder = null;
+      objectsToSave = Collections.newSetFromMap(new IdentityHashMap<Object, Boolean>());
       if (entityManager != null) {
         // Close this EntityManager just let the connection pool be able to reuse it, it doesn't
         // close the underlying database connection.
@@ -708,25 +778,241 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
         }
       }
     }
+
+    /** Does the full "update" on an object including all internal housekeeping. */
+    private void updateObject(Object object) {
+      Object merged = entityManager.merge(object);
+      objectsToSave.add(merged);
+      addUpdate(object);
+    }
+
+    /** Does the full "insert" on a new object including all internal housekeeping. */
+    private void insertObject(Object object) {
+      entityManager.persist(object);
+      objectsToSave.add(object);
+      addUpdate(object);
+    }
+
+    /** Returns true if the object has been persisted/merged and will be saved on commit. */
+    private boolean willSave(Object object) {
+      return objectsToSave.contains(object);
+    }
   }
 
-  private static class JpaQueryComposerImpl<T> extends QueryComposer<T> {
+  /**
+   * 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;
 
-    EntityManager em;
-
     private int fetchSize = DEFAULT_FETCH_SIZE;
 
-    private boolean autoDetachOnLoad = true;
-
-    JpaQueryComposerImpl(Class<T> entityClass, EntityManager em) {
+    JpaQueryComposerImpl(Class<T> entityClass) {
       super(entityClass);
-      this.em = em;
     }
 
     private TypedQuery<T> buildQuery() {
-      CriteriaQueryBuilder<T> queryBuilder = CriteriaQueryBuilder.create(em, entityClass);
+      CriteriaQueryBuilder<T> queryBuilder =
+          CriteriaQueryBuilder.create(getEntityManager(), entityClass);
       return addCriteria(queryBuilder);
     }
 
@@ -739,13 +1025,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
         queryBuilder.orderByAsc(orderBy);
       }
 
-      return em.createQuery(queryBuilder.build());
-    }
-
-    @Override
-    public QueryComposer<T> withAutoDetachOnLoad(boolean autoDetachOnLoad) {
-      this.autoDetachOnLoad = autoDetachOnLoad;
-      return this;
+      return getEntityManager().createQuery(queryBuilder.build());
     }
 
     @Override
@@ -758,12 +1038,12 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     @Override
     public Optional<T> first() {
       List<T> results = buildQuery().setMaxResults(1).getResultList();
-      return results.size() > 0 ? Optional.of(maybeDetachEntity(results.get(0))) : Optional.empty();
+      return results.size() > 0 ? Optional.of(detach(results.get(0))) : Optional.empty();
     }
 
     @Override
     public T getSingleResult() {
-      return maybeDetachEntity(buildQuery().getSingleResult());
+      return detach(buildQuery().getSingleResult());
     }
 
     @Override
@@ -777,27 +1057,21 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
       } else {
         logger.atWarning().log("Query implemention does not support result streaming.");
       }
-      return query.getResultStream().map(this::maybeDetachEntity);
+      return query.getResultStream().map(JpaTransactionManagerImpl.this::detach);
     }
 
     @Override
     public long count() {
-      CriteriaQueryBuilder<Long> queryBuilder = CriteriaQueryBuilder.createCount(em, entityClass);
+      CriteriaQueryBuilder<Long> queryBuilder =
+          CriteriaQueryBuilder.createCount(getEntityManager(), entityClass);
       return addCriteria(queryBuilder).getSingleResult();
     }
 
     @Override
     public ImmutableList<T> list() {
       return buildQuery().getResultList().stream()
-          .map(this::maybeDetachEntity)
+          .map(JpaTransactionManagerImpl.this::detach)
           .collect(ImmutableList.toImmutableList());
     }
-
-    private T maybeDetachEntity(T entity) {
-      if (autoDetachOnLoad) {
-        em.detach(entity);
-      }
-      return entity;
-    }
   }
 }

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

@@ -93,16 +93,6 @@ public abstract class QueryComposer<T> {
     return this;
   }
 
-  /**
-   * Specifies if JPA entities should be automatically detached from the persistence context after
-   * loading. The default behavior is auto-detach.
-   *
-   * <p>This configuration has no effect on Datastore queries.
-   */
-  public QueryComposer<T> withAutoDetachOnLoad(boolean autoDetachOnLoad) {
-    return this;
-  }
-
   /** Returns the first result of the query or an empty optional if there is none. */
   public abstract Optional<T> first();
 

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

@@ -262,8 +262,14 @@ public interface TransactionManager {
   /** Deletes the set of entities by their key id. */
   void delete(Iterable<? extends VKey<?>> keys);
 
-  /** Deletes the given entity from the database. */
-  void delete(Object entity);
+  /**
+   * Deletes the given entity from the database.
+   *
+   * <p>This returns the deleted entity, which may not necessarily be the same as the original
+   * entity passed in, as it may be a) converted to a different type of object more appropriate to
+   * the database type or b) merged with an object managed by the database entity manager.
+   */
+  <T> T delete(T entity);
 
   /**
    * Deletes the entity by its id without writing commit logs if the underlying database is

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/rde/RdeStagingMapper.java

@@ -17,6 +17,7 @@ package google.registry.rde;
 import static com.google.common.base.Strings.nullToEmpty;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static google.registry.model.EppResourceUtils.loadAtPointInTime;
+import static google.registry.model.EppResourceUtils.loadAtPointInTimeAsync;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.appengine.tools.mapreduce.Mapper;
@@ -26,7 +27,6 @@ import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSetMultimap;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Sets;
-import com.googlecode.objectify.Result;
 import google.registry.model.EppResource;
 import google.registry.model.contact.ContactResource;
 import google.registry.model.domain.DomainBase;
@@ -36,7 +36,9 @@ import google.registry.model.registrar.Registrar;
 import google.registry.xml.ValidationMode;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
+import java.util.function.Supplier;
 import org.joda.time.DateTime;
 
 /** Mapper for {@link RdeStagingAction}. */
@@ -123,8 +125,8 @@ public final class RdeStagingMapper extends Mapper<EppResource, PendingDeposit,
             .collect(toImmutableSet());
 
     // Launch asynchronous fetches of point-in-time representations of resource.
-    ImmutableMap<DateTime, Result<EppResource>> resourceAtTimes =
-        ImmutableMap.copyOf(Maps.asMap(dates, input -> loadAtPointInTime(resource, input)));
+    ImmutableMap<DateTime, Supplier<EppResource>> resourceAtTimes =
+        ImmutableMap.copyOf(Maps.asMap(dates, input -> loadAtPointInTimeAsync(resource, input)));
 
     // Convert resource to an XML fragment for each watermark/mode pair lazily and cache the result.
     Fragmenter fragmenter = new Fragmenter(resourceAtTimes);
@@ -159,13 +161,13 @@ public final class RdeStagingMapper extends Mapper<EppResource, PendingDeposit,
   /** Loading cache that turns a resource into XML for the various points in time and modes. */
   private class Fragmenter {
     private final Map<WatermarkModePair, Optional<DepositFragment>> cache = new HashMap<>();
-    private final ImmutableMap<DateTime, Result<EppResource>> resourceAtTimes;
+    private final ImmutableMap<DateTime, Supplier<EppResource>> resourceAtTimes;
 
     long cacheHits = 0;
     long resourcesNotFound = 0;
     long resourcesFound = 0;
 
-    Fragmenter(ImmutableMap<DateTime, Result<EppResource>> resourceAtTimes) {
+    Fragmenter(ImmutableMap<DateTime, Supplier<EppResource>> resourceAtTimes) {
       this.resourceAtTimes = resourceAtTimes;
     }
 
@@ -175,7 +177,7 @@ public final class RdeStagingMapper extends Mapper<EppResource, PendingDeposit,
         cacheHits++;
         return result;
       }
-      EppResource resource = resourceAtTimes.get(watermark).now();
+      EppResource resource = resourceAtTimes.get(watermark).get();
       if (resource == null) {
         result = Optional.empty();
         cache.put(WatermarkModePair.create(watermark, RdeMode.FULL), result);
@@ -202,8 +204,9 @@ public final class RdeStagingMapper extends Mapper<EppResource, PendingDeposit,
                         host,
                         // Note that loadAtPointInTime() does cloneProjectedAtTime(watermark) for
                         // us.
-                        loadAtPointInTime(tm().loadByKey(host.getSuperordinateDomain()), watermark)
-                            .now())
+                        Objects.requireNonNull(
+                            loadAtPointInTime(
+                                tm().loadByKey(host.getSuperordinateDomain()), watermark)))
                     : marshaller.marshalExternalHost(host));
         cache.put(WatermarkModePair.create(watermark, RdeMode.FULL), result);
         cache.put(WatermarkModePair.create(watermark, RdeMode.THIN), result);

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

@@ -30,6 +30,7 @@ import google.registry.model.registry.label.ReservedList;
 import google.registry.persistence.VKey;
 import java.nio.file.Files;
 import java.util.List;
+import java.util.stream.Collectors;
 import org.joda.time.DateTime;
 
 /** Command to create a {@link ReservedList}. */
@@ -73,6 +74,30 @@ final class CreateReservedListCommand extends CreateOrUpdateReservedListCommand
         null, reservedList, VKey.createOfy(ReservedList.class, Key.create(reservedList)));
   }
 
+  @Override
+  protected String prompt() {
+    return getChangedEntities().isEmpty()
+        ? "No entity changes to apply."
+        : getChangedEntities().stream()
+            .map(
+                entity -> {
+                  if (entity instanceof ReservedList) {
+                    // Format the entries of the reserved list as well.
+                    String entries =
+                        ((ReservedList) entity)
+                            .getReservedListEntries().entrySet().stream()
+                                .map(
+                                    entry ->
+                                        String.format("%s=%s", entry.getKey(), entry.getValue()))
+                                .collect(Collectors.joining(", "));
+                    return String.format("%s\nreservedListMap={%s}\n", entity, entries);
+                  } else {
+                    return entity.toString();
+                  }
+                })
+            .collect(Collectors.joining("\n"));
+  }
+
   private static void validateListName(String name) {
     List<String> nameParts = Splitter.on('_').splitToList(name);
     checkArgument(nameParts.size() == 2, INVALID_FORMAT_ERROR_MESSAGE);

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

@@ -69,7 +69,6 @@ final class GenerateLordnCommand implements CommandWithRemoteApi {
                 .createQueryComposer(DomainBase.class)
                 .where("tld", Comparator.EQ, tld)
                 .orderBy("repoId")
-                .withAutoDetachOnLoad(false)
                 .stream()
                 .forEach(domain -> processDomain(claimsCsv, sunriseCsv, domain)));
     ImmutableList<String> claimsRows = claimsCsv.build();

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

@@ -20,6 +20,7 @@ import google.registry.tools.javascrap.BackfillSpec11ThreatMatchesCommand;
 import google.registry.tools.javascrap.DeleteContactByRoidCommand;
 import google.registry.tools.javascrap.PopulateNullRegistrarFieldsCommand;
 import google.registry.tools.javascrap.RemoveIpAddressCommand;
+import google.registry.tools.javascrap.ResaveAllTldsCommand;
 
 /** Container class to create and run remote commands against a Datastore instance. */
 public final class RegistryTool {
@@ -69,7 +70,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)
@@ -106,12 +107,13 @@ public final class RegistryTool {
           .put("remove_ip_address", RemoveIpAddressCommand.class)
           .put("remove_registry_one_key", RemoveRegistryOneKeyCommand.class)
           .put("renew_domain", RenewDomainCommand.class)
+          .put("resave_all_tlds", ResaveAllTldsCommand.class)
           .put("resave_entities", ResaveEntitiesCommand.class)
           .put("resave_environment_entities", ResaveEnvironmentEntitiesCommand.class)
           .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/javascrap/ResaveAllTldsCommand.java

@@ -12,12 +12,19 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package google.registry.tools.params;
+package google.registry.tools.javascrap;
 
-import google.registry.model.common.DatabaseTransitionSchedule.TransitionId;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
-/**
- * {@link TransitionId} CLI parameter converter/validator. Required to support multi-value
- * TransitionId parameters.
- */
-public final class TransitionIdParameter extends EnumParameter<TransitionId> {}
+import com.beust.jcommander.Parameters;
+import google.registry.model.registry.Registry;
+import google.registry.tools.CommandWithRemoteApi;
+
+/** Scrap command to resave all Registry entities. */
+@Parameters(commandDescription = "Resave all TLDs")
+public class ResaveAllTldsCommand implements CommandWithRemoteApi {
+  @Override
+  public void run() throws Exception {
+    tm().transact(() -> tm().putAll(tm().loadAllOf(Registry.class)));
+  }
+}

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

@@ -198,7 +198,7 @@ public class GenerateZoneFilesAction implements Runnable, JsonActionRunner.JsonA
     private void mapDomain(DomainBase domain) {
       // Domains never change their tld, so we can check if it's from the wrong tld right away.
       if (tlds.contains(domain.getTld())) {
-        domain = loadAtPointInTime(domain, exportTime).now();
+        domain = loadAtPointInTime(domain, exportTime);
         // A null means the domain was deleted (or not created) at this time.
         if (domain != null && domain.shouldPublishToDns()) {
           String stanza = domainStanza(domain, exportTime, dnsDefaultNsTtl, dnsDefaultDsTtl);
@@ -215,7 +215,7 @@ public class GenerateZoneFilesAction implements Runnable, JsonActionRunner.JsonA
       ImmutableSet<String> subordinateHosts = domain.getSubordinateHosts();
       if (!subordinateHosts.isEmpty()) {
         for (HostResource unprojectedHost : tm().loadByKeys(domain.getNameservers()).values()) {
-          HostResource host = loadAtPointInTime(unprojectedHost, exportTime).now();
+          HostResource host = loadAtPointInTime(unprojectedHost, exportTime);
           // A null means the host was deleted (or not created) at this time.
           if ((host != null) && subordinateHosts.contains(host.getHostName())) {
             String stanza = hostStanza(host, dnsDefaultATtl, domain.getTld());
@@ -290,7 +290,7 @@ public class GenerateZoneFilesAction implements Runnable, JsonActionRunner.JsonA
               domainLabel,
               dnsDefaultNsTtl.getStandardSeconds(),
               // Load the nameservers at the export time in case they've been renamed or deleted.
-              loadAtPointInTime(nameserver, exportTime).now().getHostName()));
+              loadAtPointInTime(nameserver, exportTime).getHostName()));
     }
     for (DelegationSignerData dsData : domain.getDsData()) {
       result.append(

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

@@ -19,7 +19,6 @@ import static google.registry.persistence.transaction.TransactionManagerFactory.
 import static google.registry.request.Action.Method.GET;
 import static google.registry.request.Action.Method.POST;
 
-import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import google.registry.model.registry.label.PremiumList;
 import google.registry.request.Action;
@@ -58,16 +57,6 @@ public final class ListPremiumListsAction extends ListObjectsAction<PremiumList>
                     .map(PremiumListDao::getLatestRevision)
                     .filter(Optional::isPresent)
                     .map(Optional::get)
-                    .peek(list -> list.getLabelsToPrices())
                     .collect(toImmutableSortedSet(Comparator.comparing(PremiumList::getName))));
   }
-
-  /**
-   * Provide a field override for labelsToPrices, since it is an {@code Insignificant} field and
-   * doesn't get returned from {@link google.registry.model.ImmutableObject#toDiffableFieldMap}.
-   */
-  @Override
-  public ImmutableMap<String, String> getFieldOverrides(PremiumList list) {
-    return ImmutableMap.of("labelsToPrices", list.getLabelsToPrices().toString());
-  }
 }

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/META-INF/persistence.xml

@@ -61,6 +61,7 @@
     <class>google.registry.model.registrar.RegistrarContact</class>
     <class>google.registry.model.registry.label.PremiumList</class>
     <class>google.registry.model.registry.label.ReservedList</class>
+    <class>google.registry.model.registry.label.ReservedList$ReservedListEntry</class>
     <class>google.registry.model.registry.Registry</class>
     <class>google.registry.model.reporting.DomainTransactionRecord</class>
     <class>google.registry.model.reporting.Spec11ThreatMatch</class>
@@ -69,6 +70,7 @@
     <class>google.registry.model.server.ServerSecret</class>
     <class>google.registry.model.smd.SignedMarkRevocationList</class>
     <class>google.registry.model.tmch.ClaimsList</class>
+    <class>google.registry.model.tmch.ClaimsEntry</class>
     <class>google.registry.model.tmch.TmchCrl</class>
     <class>google.registry.persistence.transaction.TransactionEntity</class>
     <class>google.registry.schema.domain.RegistryLock</class>

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,14 +168,14 @@ 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) {
     DomainBase pendingExpirationDomain = persistActiveDomain(domainName);
-    HistoryEntry createHistoryEntry =
+    DomainHistory createHistoryEntry =
         persistResource(
             new DomainHistory.Builder()
                 .setType(DOMAIN_CREATE)
@@ -190,7 +200,7 @@ class DeleteExpiredDomainsActionTest {
   }
 
   private BillingEvent.Recurring.Builder createAutorenewBillingEvent(
-      HistoryEntry createHistoryEntry) {
+      DomainHistory createHistoryEntry) {
     return new BillingEvent.Recurring.Builder()
         .setReason(Reason.RENEW)
         .setFlags(ImmutableSet.of(Flag.AUTO_RENEW))