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

@@ -370,11 +370,10 @@ public class DeleteContactsAndHostsAction implements Runnable {
                       : "it was transferred prior to deletion");
 
       HistoryEntry historyEntry =
-          new HistoryEntry.Builder()
+          HistoryEntry.createBuilderForResource(resource)
               .setClientId(deletionRequest.requestingClientId())
               .setModificationTime(now)
               .setType(getHistoryEntryType(resource, deleteAllowed))
-              .setParent(deletionRequest.key())
               .build();
 
       PollMessage.OneTime pollMessage =
@@ -409,7 +408,9 @@ public class DeleteContactsAndHostsAction implements Runnable {
       } else {
         resourceToSave = resource.asBuilder().removeStatusValue(PENDING_DELETE).build();
       }
-      auditedOfy().save().<ImmutableObject>entities(resourceToSave, historyEntry, pollMessage);
+      auditedOfy()
+          .save()
+          .<ImmutableObject>entities(resourceToSave, historyEntry.asHistoryEntry(), pollMessage);
       return DeletionResult.create(
           deleteAllowed ? Type.DELETED : Type.NOT_DELETED, pollMessageText);
     }

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

@@ -44,11 +44,11 @@ import google.registry.mapreduce.MapreduceRunner;
 import google.registry.mapreduce.inputs.EppResourceInputs;
 import google.registry.model.EppResourceUtils;
 import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.DomainHistory;
 import google.registry.model.index.EppResourceIndex;
 import google.registry.model.index.ForeignKeyIndex;
 import google.registry.model.registry.Registry;
 import google.registry.model.registry.Registry.TldType;
-import google.registry.model.reporting.HistoryEntry;
 import google.registry.request.Action;
 import google.registry.request.Parameter;
 import google.registry.request.Response;
@@ -253,9 +253,9 @@ public class DeleteProberDataAction implements Runnable {
                         .setDeletionTime(tm().getTransactionTime())
                         .setStatusValues(null)
                         .build();
-                HistoryEntry historyEntry =
-                    new HistoryEntry.Builder()
-                        .setParent(domain)
+                DomainHistory historyEntry =
+                    new DomainHistory.Builder()
+                        .setDomain(domain)
                         .setType(DOMAIN_DELETE)
                         .setModificationTime(tm().getTransactionTime())
                         .setBySuperuser(true)
@@ -263,11 +263,9 @@ public class DeleteProberDataAction implements Runnable {
                         .setClientId(registryAdminClientId)
                         .build();
                 // Note that we don't bother handling grace periods, billing events, pending
-                // transfers,
-                // poll messages, or auto-renews because these will all be hard-deleted the next
-                // time the
-                // mapreduce runs anyway.
-                auditedOfy().save().entities(deletedDomain, historyEntry);
+                // transfers, poll messages, or auto-renews because these will all be hard-deleted
+                // the next time the mapreduce runs anyway.
+                tm().putAll(deletedDomain, historyEntry);
                 updateForeignKeyIndexDeletionTime(deletedDomain);
                 dnsQueue.addDomainRefreshTask(deletedDomain.getDomainName());
               });

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;
@@ -49,11 +50,12 @@ import google.registry.model.billing.BillingEvent.OneTime;
 import google.registry.model.billing.BillingEvent.Recurring;
 import google.registry.model.common.Cursor;
 import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.DomainHistory;
 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;
@@ -91,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> {
@@ -154,98 +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
-                        HistoryEntry historyEntry =
-                            new HistoryEntry.Builder()
-                                .setBySuperuser(false)
-                                .setClientId(recurring.getClientId())
-                                .setModificationTime(tm().getTransactionTime())
-                                .setParent(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<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);
@@ -257,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),
@@ -328,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)) {
@@ -343,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/FlowModule.java

@@ -213,50 +213,78 @@ public class FlowModule {
     return Strings.nullToEmpty(((Poll) eppInput.getCommandWrapper().getCommand()).getMessageId());
   }
 
+  private static <B extends HistoryEntry.Builder<? extends HistoryEntry, ?>>
+      B makeHistoryEntryBuilder(
+          B builder,
+          Trid trid,
+          byte[] inputXmlBytes,
+          boolean isSuperuser,
+          String clientId,
+          EppInput eppInput) {
+    builder
+        .setTrid(trid)
+        .setXmlBytes(inputXmlBytes)
+        .setBySuperuser(isSuperuser)
+        .setClientId(clientId);
+    Optional<MetadataExtension> metadataExtension =
+        eppInput.getSingleExtension(MetadataExtension.class);
+    metadataExtension.ifPresent(
+        extension ->
+            builder
+                .setReason(extension.getReason())
+                .setRequestedByRegistrar(extension.getRequestedByRegistrar()));
+    return builder;
+  }
+
   /**
-   * Provides a partially filled in {@link HistoryEntry} builder.
+   * Provides a partially filled in {@link ContactHistory.Builder}
    *
    * <p>This is not marked with {@link FlowScope} so that each retry gets a fresh one. Otherwise,
    * the fact that the builder is one-use would cause NPEs.
    */
   @Provides
-  static HistoryEntry.Builder provideHistoryEntryBuilder(
+  static ContactHistory.Builder provideContactHistoryBuilder(
       Trid trid,
       @InputXml byte[] inputXmlBytes,
       @Superuser boolean isSuperuser,
       @ClientId String clientId,
       EppInput eppInput) {
-    HistoryEntry.Builder historyBuilder =
-        new HistoryEntry.Builder()
-            .setTrid(trid)
-            .setXmlBytes(inputXmlBytes)
-            .setBySuperuser(isSuperuser)
-            .setClientId(clientId);
-    Optional<MetadataExtension> metadataExtension =
-        eppInput.getSingleExtension(MetadataExtension.class);
-    metadataExtension.ifPresent(
-        extension ->
-            historyBuilder
-                .setReason(extension.getReason())
-                .setRequestedByRegistrar(extension.getRequestedByRegistrar()));
-    return historyBuilder;
+    return makeHistoryEntryBuilder(
+        new ContactHistory.Builder(), trid, inputXmlBytes, isSuperuser, clientId, eppInput);
   }
 
+  /**
+   * Provides a partially filled in {@link HostHistory.Builder}
+   *
+   * <p>This is not marked with {@link FlowScope} so that each retry gets a fresh one. Otherwise,
+   * the fact that the builder is one-use would cause NPEs.
+   */
   @Provides
-  static ContactHistory.Builder provideContactHistoryBuilder(
-      HistoryEntry.Builder historyEntryBuilder) {
-    return new ContactHistory.Builder().copyFrom(historyEntryBuilder);
+  static HostHistory.Builder provideHostHistoryBuilder(
+      Trid trid,
+      @InputXml byte[] inputXmlBytes,
+      @Superuser boolean isSuperuser,
+      @ClientId String clientId,
+      EppInput eppInput) {
+    return makeHistoryEntryBuilder(
+        new HostHistory.Builder(), trid, inputXmlBytes, isSuperuser, clientId, eppInput);
   }
 
+  /**
+   * Provides a partially filled in {@link DomainHistory.Builder}
+   *
+   * <p>This is not marked with {@link FlowScope} so that each retry gets a fresh one. Otherwise,
+   * the fact that the builder is one-use would cause NPEs.
+   */
   @Provides
   static DomainHistory.Builder provideDomainHistoryBuilder(
-      HistoryEntry.Builder historyEntryBuilder) {
-    return new DomainHistory.Builder().copyFrom(historyEntryBuilder);
-  }
-
-  @Provides
-  static HostHistory.Builder provideHostHistoryBuilder(HistoryEntry.Builder historyEntryBuilder) {
-    return new HostHistory.Builder().copyFrom(historyEntryBuilder);
+      Trid trid,
+      @InputXml byte[] inputXmlBytes,
+      @Superuser boolean isSuperuser,
+      @ClientId String clientId,
+      EppInput eppInput) {
+    return makeHistoryEntryBuilder(
+        new DomainHistory.Builder(), trid, inputXmlBytes, isSuperuser, clientId, eppInput);
   }
 
   /**

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

@@ -45,7 +45,7 @@ import google.registry.model.eppinput.ResourceCommand;
 import google.registry.model.eppoutput.EppResponse;
 import google.registry.model.registry.Registry;
 import google.registry.model.reporting.IcannReportingTypes.ActivityReportField;
-import google.registry.model.tmch.ClaimsListDualDatabaseDao;
+import google.registry.model.tmch.ClaimsListDao;
 import google.registry.util.Clock;
 import java.util.HashSet;
 import java.util.Optional;
@@ -104,8 +104,7 @@ public final class DomainClaimsCheckFlow implements Flow {
           verifyClaimsPeriodNotEnded(registry, now);
         }
       }
-      Optional<String> claimKey =
-          ClaimsListDualDatabaseDao.get().getClaimKey(parsedDomain.parts().get(0));
+      Optional<String> claimKey = ClaimsListDao.get().getClaimKey(parsedDomain.parts().get(0));
       launchChecksBuilder.add(
           LaunchCheck.create(
               LaunchCheckName.create(claimKey.isPresent(), domainName), claimKey.orElse(null)));

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

@@ -126,7 +126,7 @@ import google.registry.model.registry.label.ReservedList;
 import google.registry.model.reporting.DomainTransactionRecord;
 import google.registry.model.reporting.DomainTransactionRecord.TransactionReportField;
 import google.registry.model.reporting.HistoryEntry;
-import google.registry.model.tmch.ClaimsListDualDatabaseDao;
+import google.registry.model.tmch.ClaimsListDao;
 import google.registry.persistence.VKey;
 import google.registry.tldconfig.idn.IdnLabelValidator;
 import google.registry.util.Idn;
@@ -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;
   }
 
   /**
@@ -990,8 +997,7 @@ public class DomainFlowUtils {
   static void verifyClaimsNoticeIfAndOnlyIfNeeded(
       InternetDomainName domainName, boolean hasSignedMarks, boolean hasClaimsNotice)
       throws EppException {
-    boolean isInClaimsList =
-        ClaimsListDualDatabaseDao.get().getClaimKey(domainName.parts().get(0)).isPresent();
+    boolean isInClaimsList = ClaimsListDao.get().getClaimKey(domainName.parts().get(0)).isPresent();
     if (hasClaimsNotice && !isInClaimsList) {
       throw new UnexpectedClaimsNoticeException(domainName.toString());
     }

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;
@@ -47,9 +47,9 @@ import google.registry.model.server.KmsSecret;
 import google.registry.model.server.KmsSecretRevision;
 import google.registry.model.server.Lock;
 import google.registry.model.server.ServerSecret;
-import google.registry.model.tmch.ClaimsListShard;
-import google.registry.model.tmch.ClaimsListShard.ClaimsListRevision;
-import google.registry.model.tmch.ClaimsListShard.ClaimsListSingleton;
+import google.registry.model.tmch.ClaimsList;
+import google.registry.model.tmch.ClaimsList.ClaimsListRevision;
+import google.registry.model.tmch.ClaimsList.ClaimsListSingleton;
 import google.registry.model.tmch.TmchCrl;
 import google.registry.schema.replay.LastSqlTransaction;
 
@@ -64,7 +64,7 @@ public final class EntityClasses {
           BillingEvent.Modification.class,
           BillingEvent.OneTime.class,
           BillingEvent.Recurring.class,
-          ClaimsListShard.class,
+          ClaimsList.class,
           ClaimsListRevision.class,
           ClaimsListSingleton.class,
           CommitLogBucket.class,
@@ -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.ClaimsListShard;
-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 ClaimsListShard} 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;
@@ -41,6 +42,11 @@ import javax.persistence.PostLoad;
  * <p>In addition to the general history fields (e.g. action time, registrar ID) we also persist a
  * copy of the contact entity at this point in time. We persist a raw {@link ContactBase} so that
  * the foreign-keyed fields in that class can refer to this object.
+ *
+ * <p>This class is only marked as a Datastore entity subclass and registered with Objectify so that
+ * when building it its ID can be auto-populated by Objectify. It is converted to its superclass
+ * {@link HistoryEntry} when persisted to Datastore using {@link
+ * google.registry.persistence.transaction.TransactionManager}.
  */
 @Entity
 @javax.persistence.Table(
@@ -102,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
@@ -110,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,10 +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 com.googlecode.objectify.annotation.Ignore;
+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;
@@ -55,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.
@@ -62,6 +64,11 @@ import javax.persistence.Table;
  * <p>In addition to the general history fields (e.g. action time, registrar ID) we also persist a
  * copy of the domain entity at this point in time. We persist a raw {@link DomainContent} so that
  * the foreign-keyed fields in that class can refer to this object.
+ *
+ * <p>This class is only marked as a Datastore entity subclass and registered with Objectify so that
+ * when building it its ID can be auto-populated by Objectify. It is converted to its superclass
+ * {@link HistoryEntry} when persisted to Datastore using {@link
+ * google.registry.persistence.transaction.TransactionManager}.
  */
 @Entity
 @Table(
@@ -97,7 +104,6 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
   // We could have reused domainContent.nsHosts here, but Hibernate throws a weird exception after
   // we change to use a composite primary key.
   // TODO(b/166776754): Investigate if we can reuse domainContent.nsHosts for storing host keys.
-  @Ignore
   @ElementCollection
   @JoinTable(
       name = "DomainHistoryHost",
@@ -111,7 +117,6 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
   @Column(name = "host_repo_id")
   Set<VKey<HostResource>> nsHosts;
 
-  @Ignore
   @OneToMany(
       cascade = {CascadeType.ALL},
       fetch = FetchType.EAGER,
@@ -131,7 +136,6 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
   // HashSet rather than ImmutableSet so that Hibernate can fill them out lazily on request
   Set<DomainDsDataHistory> dsDataHistories = new HashSet<>();
 
-  @Ignore
   @OneToMany(
       cascade = {CascadeType.ALL},
       fetch = FetchType.EAGER,
@@ -246,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;
@@ -41,6 +42,11 @@ import javax.persistence.PostLoad;
  * <p>In addition to the general history fields (e.g. action time, registrar ID) we also persist a
  * copy of the host entity at this point in time. We persist a raw {@link HostBase} so that the
  * foreign-keyed fields in that class can refer to this object.
+ *
+ * <p>This class is only marked as a Datastore entity subclass and registered with Objectify so that
+ * when building it its ID can be auto-populated by Objectify. It is converted to its superclass
+ * {@link HistoryEntry} when persisted to Datastore using {@link
+ * google.registry.persistence.transaction.TransactionManager}.
  */
 @Entity
 @javax.persistence.Table(
@@ -103,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
@@ -111,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

@@ -173,6 +173,11 @@ public class DatastoreTransactionManager implements TransactionManager {
     putAll(entities);
   }
 
+  @Override
+  public void updateAll(Object... entities) {
+    updateAll(ImmutableList.of(entities));
+  }
+
   @Override
   public void updateWithoutBackup(Object entity) {
     putWithoutBackup(entity);
@@ -217,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
@@ -245,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());
@@ -262,13 +276,17 @@ public class DatastoreTransactionManager implements TransactionManager {
 
   @Override
   public <T> ImmutableList<T> loadAllOf(Class<T> clazz) {
-    Query<T> query = getOfy().load().type(clazz);
-    // If the entity is in the cross-TLD entity group, then we can take advantage of an ancestor
-    // query to give us strong transactional consistency.
-    if (clazz.isAnnotationPresent(InCrossTld.class)) {
-      query = query.ancestor(getCrossTldKey());
-    }
-    return ImmutableList.copyOf(query);
+    return ImmutableList.copyOf(getPossibleAncestorQuery(clazz));
+  }
+
+  @Override
+  public <T> Optional<T> loadSingleton(Class<T> clazz) {
+    List<T> elements = getPossibleAncestorQuery(clazz).limit(2).list();
+    checkArgument(
+        elements.size() <= 1,
+        "Expected at most one entity of type %s, found at least two",
+        clazz.getSimpleName());
+    return elements.stream().findFirst();
   }
 
   @Override
@@ -288,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
@@ -401,6 +420,17 @@ public class DatastoreTransactionManager implements TransactionManager {
     return obj;
   }
 
+  /** A query for returning any/all results of an object, with an ancestor if possible. */
+  private <T> Query<T> getPossibleAncestorQuery(Class<T> clazz) {
+    Query<T> query = getOfy().load().type(clazz);
+    // If the entity is in the cross-TLD entity group, then we can take advantage of an ancestor
+    // query to give us strong transactional consistency.
+    if (clazz.isAnnotationPresent(InCrossTld.class)) {
+      query = query.ancestor(getCrossTldKey());
+    }
+    return query;
+  }
+
   private static class DatastoreQueryComposerImpl<T> extends QueryComposer<T> {
 
     DatastoreQueryComposerImpl(Class<T> entityClass) {
@@ -450,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/PremiumList.java

@@ -16,9 +16,12 @@ package google.registry.model.registry.label;
 
 import static com.google.common.base.Charsets.US_ASCII;
 import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static com.google.common.hash.Funnels.stringFunnel;
 import static com.google.common.hash.Funnels.unencodedCharsFunnel;
 import static google.registry.model.ofy.ObjectifyService.allocateId;
+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.base.Splitter;
@@ -36,6 +39,7 @@ import google.registry.model.annotations.ReportedOn;
 import google.registry.model.registry.Registry;
 import google.registry.schema.replay.DatastoreOnlyEntity;
 import google.registry.schema.replay.NonReplicatedEntity;
+import google.registry.schema.tld.PremiumEntry;
 import google.registry.schema.tld.PremiumListDao;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
@@ -45,17 +49,14 @@ import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
 import javax.annotation.Nullable;
-import javax.persistence.CollectionTable;
 import javax.persistence.Column;
-import javax.persistence.ElementCollection;
 import javax.persistence.Index;
-import javax.persistence.JoinColumn;
-import javax.persistence.MapKeyColumn;
 import javax.persistence.PostLoad;
+import javax.persistence.PostPersist;
 import javax.persistence.PrePersist;
+import javax.persistence.PreRemove;
 import javax.persistence.Table;
 import javax.persistence.Transient;
-import org.hibernate.LazyInitializationException;
 import org.joda.money.CurrencyUnit;
 import org.joda.money.Money;
 
@@ -81,14 +82,14 @@ public final class PremiumList extends BaseDomainLabelList<Money, PremiumList.Pr
   @Column(nullable = false)
   CurrencyUnit currency;
 
-  @Ignore
-  @ElementCollection
-  @CollectionTable(
-      name = "PremiumEntry",
-      joinColumns = @JoinColumn(name = "revisionId", referencedColumnName = "revisionId"))
-  @MapKeyColumn(name = "domainLabel")
-  @Column(name = "price", nullable = false)
-  Map<String, BigDecimal> labelsToPrices;
+  /**
+   * Mapping from unqualified domain names to their prices.
+   *
+   * <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.
+   */
+  @Ignore @ImmutableObject.Insignificant @Transient ImmutableMap<String, BigDecimal> labelsToPrices;
 
   @Ignore
   @Column(nullable = false)
@@ -170,14 +171,20 @@ public final class PremiumList extends BaseDomainLabelList<Money, PremiumList.Pr
   /**
    * Returns a {@link Map} of domain labels to prices.
    *
-   * <p>Note that this is lazily loaded and thus will throw a {@link LazyInitializationException} if
-   * used outside the transaction in which the given entity was loaded. You generally should not be
-   * using this anyway as it's inefficient to load all of the PremiumEntry rows if you don't need
-   * them. To check prices, use {@link PremiumListDao#getPremiumPrice} instead.
+   * <p>Note that this is lazily loaded and thus must be called inside a transaction. You generally
+   * should not be using this anyway as it's inefficient to load all of the PremiumEntry rows if you
+   * don't need them. To check prices, use {@link PremiumListDao#getPremiumPrice} instead.
    */
-  @Nullable
-  public ImmutableMap<String, BigDecimal> getLabelsToPrices() {
-    return labelsToPrices == null ? null : ImmutableMap.copyOf(labelsToPrices);
+  public synchronized ImmutableMap<String, BigDecimal> getLabelsToPrices() {
+    if (labelsToPrices == null) {
+      labelsToPrices =
+          jpaTm()
+              .createQueryComposer(PremiumEntry.class)
+              .where("revisionId", EQ, revisionId)
+              .stream()
+              .collect(toImmutableMap(PremiumEntry::getDomainLabel, PremiumEntry::getPrice));
+    }
+    return labelsToPrices;
   }
 
   /**
@@ -320,4 +327,32 @@ public final class PremiumList extends BaseDomainLabelList<Money, PremiumList.Pr
   void postLoad() {
     creationTime = lastUpdateTime;
   }
+
+  @PreRemove
+  void preRemove() {
+    jpaTm()
+        .query("DELETE FROM PremiumEntry WHERE revision_id = :revisionId")
+        .setParameter("revisionId", revisionId)
+        .executeUpdate();
+  }
+
+  /**
+   * Hibernate hook called on the insert of a new PremiumList. Stores the associated {@link
+   * PremiumEntry}'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 the price map is loaded, persist it too.
+    if (labelsToPrices != null) {
+      labelsToPrices.entrySet().stream()
+          .forEach(
+              entry ->
+                  jpaTm()
+                      .insert(PremiumEntry.create(revisionId, entry.getValue(), entry.getKey())));
+    }
+  }
 }

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

@@ -32,14 +32,17 @@ import google.registry.model.Buildable;
 import google.registry.model.EppResource;
 import google.registry.model.ImmutableObject;
 import google.registry.model.annotations.ReportedOn;
+import google.registry.model.contact.ContactBase;
 import google.registry.model.contact.ContactHistory;
 import google.registry.model.contact.ContactHistory.ContactHistoryId;
 import google.registry.model.contact.ContactResource;
 import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.DomainContent;
 import google.registry.model.domain.DomainHistory;
 import google.registry.model.domain.DomainHistory.DomainHistoryId;
 import google.registry.model.domain.Period;
 import google.registry.model.eppcommon.Trid;
+import google.registry.model.host.HostBase;
 import google.registry.model.host.HostHistory;
 import google.registry.model.host.HostHistory.HostHistoryId;
 import google.registry.model.host.HostResource;
@@ -58,9 +61,24 @@ 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;
 
-/** A record of an EPP command that mutated a resource. */
+/**
+ * A record of an EPP command that mutated a resource.
+ *
+ * <p>Due to historical reasons this class is persisted only to Datastore. It has three subclasses
+ * that include the parent resource itself which are persisted to Cloud SQL. During migration this
+ * class cannot be made abstract in order for the class to be persisted and loaded to and from
+ * Datastore. However it should never be used directly in the Java code itself. When it is loaded
+ * from Datastore it should be converted to a subclass for handling and when a new history entry is
+ * built it should always be a subclass, which is automatically converted to HistoryEntry when
+ * persisting to Datastore.
+ *
+ * <p>Some care has been taken to make it close to impossible to use this class directly, but the
+ * user should still exercise caution. After the migration is complete this class will be made
+ * abstract.
+ */
 @ReportedOn
 @Entity
 @MappedSuperclass
@@ -203,6 +221,12 @@ public class HistoryEntry extends ImmutableObject implements Buildable, Datastor
   @ImmutableObject.EmptySetToNull
   protected Set<DomainTransactionRecord> domainTransactionRecords;
 
+  // Make it impossible to instantiate a HistoryEntry explicitly. One should only instantiate a
+  // subtype of HistoryEntry.
+  protected HistoryEntry() {
+    super();
+  }
+
   public long getId() {
     // For some reason, Hibernate throws NPE during some initialization phase if we don't deal with
     // the null case. Setting the id to 0L when it is null should be fine because 0L for primitive
@@ -266,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) {
@@ -285,13 +320,50 @@ public class HistoryEntry extends ImmutableObject implements Buildable, Datastor
         domainTransactionRecords == null ? null : ImmutableSet.copyOf(domainTransactionRecords);
   }
 
+  /**
+   * Throws an error when trying to get a builder from a bare {@link HistoryEntry}.
+   *
+   * <p>This method only exists to satisfy the requirement that the {@link HistoryEntry} is NOT
+   * abstract, it should never be called directly and all three of the subclass of {@link
+   * HistoryEntry} implements it.
+   */
   @Override
-  public Builder asBuilder() {
-    return new Builder(clone(this));
+  public Builder<? extends HistoryEntry, ?> asBuilder() {
+    throw new UnsupportedOperationException(
+        "You should never attempt to build a HistoryEntry from a raw HistoryEntry. A raw "
+            + "HistoryEntry should only exist internally when persisting to datastore. If you need "
+            + "to build from a raw HistoryEntry, use "
+            + "{Contact,Host,Domain}History.Builder.copyFrom(HistoryEntry) instead.");
   }
 
+  /**
+   * Clones and returns a {@code HistoryEntry} objec
+   *
+   * <p>This is useful when converting a subclass to the base class to persist to Datastore.
+   */
   public HistoryEntry asHistoryEntry() {
-    return new Builder().copyFrom(this).build();
+    HistoryEntry historyEntry = new HistoryEntry();
+    copy(this, historyEntry);
+    return historyEntry;
+  }
+
+  protected static void copy(HistoryEntry src, HistoryEntry dst) {
+    dst.id = src.id;
+    dst.parent = src.parent;
+    dst.type = src.type;
+    dst.period = src.period;
+    dst.xmlBytes = src.xmlBytes;
+    dst.modificationTime = src.modificationTime;
+    dst.clientId = src.clientId;
+    dst.otherClientId = src.otherClientId;
+    dst.trid = src.trid;
+    dst.bySuperuser = src.bySuperuser;
+    dst.reason = src.reason;
+    dst.requestedByRegistrar = src.requestedByRegistrar;
+    dst.domainTransactionRecords =
+        src.domainTransactionRecords == null
+            ? null
+            : ImmutableSet.copyOf(src.domainTransactionRecords);
   }
 
   @SuppressWarnings("unchecked")
@@ -349,33 +421,18 @@ public class HistoryEntry extends ImmutableObject implements Buildable, Datastor
   }
 
   /** A builder for {@link HistoryEntry} since it is immutable */
-  public static class Builder<T extends HistoryEntry, B extends Builder<?, ?>>
+  public abstract static class Builder<T extends HistoryEntry, B extends Builder<?, ?>>
       extends GenericBuilder<T, B> {
-    public Builder() {}
+    protected Builder() {}
 
-    public Builder(T instance) {
+    protected Builder(T instance) {
       super(instance);
     }
 
     // Used to fill out the fields in this object from an object which may not be exactly the same
     // as the class T, where both classes still subclass HistoryEntry
     public B copyFrom(HistoryEntry historyEntry) {
-      setId(historyEntry.id);
-      setParent(historyEntry.parent);
-      setType(historyEntry.type);
-      setPeriod(historyEntry.period);
-      setXmlBytes(historyEntry.xmlBytes);
-      setModificationTime(historyEntry.modificationTime);
-      setClientId(historyEntry.clientId);
-      setOtherClientId(historyEntry.otherClientId);
-      setTrid(historyEntry.trid);
-      setBySuperuser(historyEntry.bySuperuser);
-      setReason(historyEntry.reason);
-      setRequestedByRegistrar(historyEntry.requestedByRegistrar);
-      setDomainTransactionRecords(
-          historyEntry.domainTransactionRecords == null
-              ? null
-              : ImmutableSet.copyOf(historyEntry.domainTransactionRecords));
+      copy(historyEntry, getInstance());
       return thisCastToDerived();
     }
 
@@ -390,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();
     }
@@ -400,13 +458,13 @@ public class HistoryEntry extends ImmutableObject implements Buildable, Datastor
       return thisCastToDerived();
     }
 
-    public B setParent(EppResource parent) {
+    protected B setParent(EppResource parent) {
       getInstance().parent = Key.create(parent);
       return thisCastToDerived();
     }
 
     // Until we move completely to SQL, override this in subclasses (e.g. HostHistory) to set VKeys
-    public B setParent(Key<? extends EppResource> parent) {
+    protected B setParent(Key<? extends EppResource> parent) {
       getInstance().parent = parent;
       return thisCastToDerived();
     }
@@ -467,4 +525,19 @@ public class HistoryEntry extends ImmutableObject implements Buildable, Datastor
       return thisCastToDerived();
     }
   }
+
+  public static <E extends EppResource>
+      HistoryEntry.Builder<? extends HistoryEntry, ?> createBuilderForResource(E parent) {
+    if (parent instanceof DomainContent) {
+      return new DomainHistory.Builder().setDomain((DomainContent) parent);
+    } else if (parent instanceof ContactBase) {
+      return new ContactHistory.Builder().setContact((ContactBase) parent);
+    } else if (parent instanceof HostBase) {
+      return new HostHistory.Builder().setHost((HostBase) parent);
+    } else {
+      throw new IllegalStateException(
+          String.format(
+              "Class %s does not have an associated HistoryEntry", parent.getClass().getName()));
+    }
+  }
 }

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/server/ServerSecret.java

@@ -14,7 +14,6 @@
 
 package google.registry.model.server;
 
-import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.common.annotations.VisibleForTesting;
@@ -22,7 +21,6 @@ import com.google.common.cache.CacheBuilder;
 import com.google.common.cache.CacheLoader;
 import com.google.common.cache.LoadingCache;
 import com.google.common.primitives.Longs;
-import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Entity;
 import com.googlecode.objectify.annotation.Ignore;
 import com.googlecode.objectify.annotation.OnLoad;
@@ -30,7 +28,6 @@ import com.googlecode.objectify.annotation.Unindex;
 import google.registry.model.annotations.NotBackedUp;
 import google.registry.model.annotations.NotBackedUp.Reason;
 import google.registry.model.common.CrossTldSingleton;
-import google.registry.persistence.VKey;
 import google.registry.schema.replay.NonReplicatedEntity;
 import java.nio.ByteBuffer;
 import java.util.Optional;
@@ -65,14 +62,9 @@ public class ServerSecret extends CrossTldSingleton implements NonReplicatedEnti
               });
 
   private static ServerSecret retrieveAndSaveSecret() {
-    VKey<ServerSecret> vkey =
-        VKey.create(
-            ServerSecret.class,
-            SINGLETON_ID,
-            Key.create(getCrossTldKey(), ServerSecret.class, SINGLETON_ID));
     if (tm().isOfy()) {
       // Attempt a quick load if we're in ofy first to short-circuit sans transaction
-      Optional<ServerSecret> secretWithoutTransaction = tm().loadByKeyIfPresent(vkey);
+      Optional<ServerSecret> secretWithoutTransaction = tm().loadSingleton(ServerSecret.class);
       if (secretWithoutTransaction.isPresent()) {
         return secretWithoutTransaction.get();
       }
@@ -81,7 +73,7 @@ public class ServerSecret extends CrossTldSingleton implements NonReplicatedEnti
             () -> {
               // Make sure we're in a transaction and attempt to load any existing secret, then
               // create it if it's absent.
-              Optional<ServerSecret> secret = tm().loadByKeyIfPresent(vkey);
+              Optional<ServerSecret> secret = tm().loadSingleton(ServerSecret.class);
               if (!secret.isPresent()) {
                 secret = Optional.of(create(UUID.randomUUID()));
                 tm().insertWithoutBackup(secret.get());

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

@@ -0,0 +1,279 @@
+// Copyright 2017 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.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.Entity;
+import com.googlecode.objectify.annotation.Id;
+import com.googlecode.objectify.annotation.Ignore;
+import com.googlecode.objectify.annotation.Parent;
+import google.registry.model.CreateAutoTimestamp;
+import google.registry.model.ImmutableObject;
+import google.registry.model.annotations.InCrossTld;
+import google.registry.model.annotations.NotBackedUp;
+import google.registry.model.annotations.NotBackedUp.Reason;
+import google.registry.model.annotations.VirtualEntity;
+import google.registry.model.common.CrossTldSingleton;
+import google.registry.schema.replay.DatastoreOnlyEntity;
+import google.registry.schema.replay.NonReplicatedEntity;
+import java.util.Map;
+import java.util.Optional;
+import javax.annotation.Nullable;
+import javax.persistence.Column;
+import javax.persistence.GeneratedValue;
+import javax.persistence.GenerationType;
+import javax.persistence.PostPersist;
+import javax.persistence.PreRemove;
+import javax.persistence.Table;
+import javax.persistence.Transient;
+import org.joda.time.DateTime;
+
+/**
+ * A list of TMCH claims labels and their associated claims keys.
+ *
+ * <p>Note that the primary key of this entity is {@link #revisionId}, which is auto-generated by
+ * the database. So, if a retry of insertion happens after the previous attempt unexpectedly
+ * succeeds, we will end up with having two exact same claims list with only different {@link
+ * #revisionId}. However, this is not an actual problem because we only use the claims list with
+ * highest {@link #revisionId}.
+ *
+ * <p>TODO(b/162007765): Remove Datastore related fields and methods.
+ */
+@Entity
+@NotBackedUp(reason = Reason.EXTERNALLY_SOURCED)
+@javax.persistence.Entity(name = "ClaimsList")
+@Table
+@InCrossTld
+public class ClaimsList extends ImmutableObject implements NonReplicatedEntity {
+
+  @Transient @Id long id;
+
+  @Transient @Parent Key<ClaimsListRevision> parent;
+
+  @Ignore
+  @javax.persistence.Id
+  @GeneratedValue(strategy = GenerationType.IDENTITY)
+  Long revisionId;
+
+  @Ignore
+  @Column(nullable = false)
+  CreateAutoTimestamp creationTimestamp = CreateAutoTimestamp.create(null);
+
+  /**
+   * When the claims list was last updated.
+   *
+   * <p>Note that the value of this field is parsed from the claims list file(See this <a
+   * href="https://tools.ietf.org/html/draft-lozano-tmch-func-spec-08#section-6.1">RFC</>), it is
+   * the DNL List creation datetime from the rfc. Since this field has been used by Datastore, we
+   * cannot change its name until we finish the migration.
+   *
+   * <p>TODO(b/177567432): Rename this field to tmdbGenerationTime.
+   */
+  @Column(name = "tmdb_generation_time", nullable = false)
+  DateTime creationTime;
+
+  /**
+   * 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() {
+    checkState(
+        revisionId != null, "revisionId is null because it is not persisted in the database");
+    return revisionId;
+  }
+
+  /**
+   * Returns the time when the external TMDB service generated this revision of the claims list.
+   *
+   * @see <a href="https://tools.ietf.org/html/draft-lozano-tmch-func-spec-08#section-6.1">DNL List
+   *     creation datetime</a>
+   */
+  public DateTime getTmdbGenerationTime() {
+    return creationTime;
+  }
+
+  /** Returns the creation time of this claims list. */
+  public DateTime getCreationTimestamp() {
+    return creationTimestamp.getTimestamp();
+  }
+
+  /**
+   * 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) {
+    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.
+   *
+   * <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() {
+    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.
+   *
+   * <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, ImmutableMap<String, String> labelsToKeys) {
+    ClaimsList instance = new ClaimsList();
+    instance.id = allocateId();
+    instance.creationTime = checkNotNull(tmdbGenerationTime);
+    instance.labelsToKeys = checkNotNull(labelsToKeys);
+    return instance;
+  }
+
+  /** Virtual parent entity for claims list shards of a specific revision. */
+  @Entity
+  @VirtualEntity
+  public static class ClaimsListRevision extends ImmutableObject implements DatastoreOnlyEntity {
+    @Parent Key<ClaimsListSingleton> parent;
+
+    @Id long versionId;
+
+    @VisibleForTesting
+    public static Key<ClaimsListRevision> createKey(ClaimsListSingleton singleton) {
+      ClaimsListRevision revision = new ClaimsListRevision();
+      revision.versionId = allocateId();
+      revision.parent = Key.create(singleton);
+      return Key.create(revision);
+    }
+
+    @VisibleForTesting
+    public static Key<ClaimsListRevision> createKey() {
+      return createKey(new ClaimsListSingleton());
+    }
+  }
+
+  /**
+   * Serves as the coordinating claims list singleton linking to the {@link ClaimsListRevision} that
+   * is live.
+   */
+  @Entity
+  @NotBackedUp(reason = Reason.EXTERNALLY_SOURCED)
+  public static class ClaimsListSingleton extends CrossTldSingleton implements DatastoreOnlyEntity {
+    Key<ClaimsListRevision> activeRevision;
+
+    static ClaimsListSingleton create(Key<ClaimsListRevision> revision) {
+      ClaimsListSingleton instance = new ClaimsListSingleton();
+      instance.activeRevision = revision;
+      return instance;
+    }
+
+    @VisibleForTesting
+    public void setActiveRevision(Key<ClaimsListRevision> revision) {
+      activeRevision = revision;
+    }
+  }
+
+  /**
+   * Returns the current ClaimsListRevision if there is one, or null if no claims list revisions
+   * have ever been persisted yet.
+   */
+  @Nullable
+  public static Key<ClaimsListRevision> getCurrentRevision() {
+    ClaimsListSingleton singleton = auditedOfy().load().entity(new ClaimsListSingleton()).now();
+    return singleton == null ? null : singleton.activeRevision;
+  }
+
+  /** Exception when trying to directly save a {@link ClaimsList} without sharding. */
+  public static class UnshardedSaveException extends RuntimeException {}
+}

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

@@ -14,23 +14,25 @@
 
 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;
 
-import java.util.Optional;
+import com.google.common.collect.ImmutableMap;
 
-/** Data access object for {@link ClaimsListShard}. */
-public class ClaimsListSqlDao {
+/** Data access object for {@link ClaimsList}. */
+public class ClaimsListDao {
 
-  /** Saves the given {@link ClaimsListShard} to Cloud SQL. */
-  static void save(ClaimsListShard claimsList) {
+  /** Saves the given {@link ClaimsList} to Cloud SQL. */
+  public static void save(ClaimsList claimsList) {
     jpaTm().transact(() -> jpaTm().insert(claimsList));
   }
 
   /**
-   * Returns the most recent revision of the {@link ClaimsListShard} in SQL or an empty list if it
+   * Returns the most recent revision of the {@link ClaimsList} in SQL or an empty list if it
    * doesn't exist.
    */
-  static Optional<ClaimsListShard> get() {
+  public static ClaimsList get() {
     return jpaTm()
         .transact(
             () -> {
@@ -39,15 +41,12 @@ public class ClaimsListSqlDao {
                       .query("SELECT MAX(revisionId) FROM ClaimsList", Long.class)
                       .getSingleResult();
               return jpaTm()
-                  .query(
-                      "FROM ClaimsList cl LEFT JOIN FETCH cl.labelsToKeys WHERE cl.revisionId ="
-                          + " :revisionId",
-                      ClaimsListShard.class)
-                  .setParameter("revisionId", revisionId)
-                  .getResultStream()
-                  .findFirst();
-            });
+                  .createQueryComposer(ClaimsList.class)
+                  .where("revisionId", EQ, revisionId)
+                  .first();
+            })
+        .orElse(ClaimsList.create(START_OF_TIME, ImmutableMap.of()));
   }
 
-  private ClaimsListSqlDao() {}
+  private ClaimsListDao() {}
 }

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

@@ -1,128 +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.tmch;
-
-import static google.registry.config.RegistryConfig.getDomainLabelListCacheDuration;
-import static google.registry.model.CacheUtils.tryMemoizeWithExpiration;
-import static google.registry.model.DatabaseMigrationUtils.suppressExceptionUnlessInTest;
-import static google.registry.util.DateTimeUtils.START_OF_TIME;
-
-import com.google.common.base.Supplier;
-import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.MapDifference;
-import com.google.common.collect.Maps;
-import google.registry.util.NonFinalForTesting;
-import java.util.Optional;
-
-/**
- * DAO for {@link ClaimsListShard} objects that handles the branching paths for SQL and Datastore.
- *
- * <p>For write actions, this class will perform the action against Cloud SQL then, after that
- * success or failure, against Datastore. If Datastore fails, an error is logged (but not thrown).
- *
- * <p>For read actions, we will log if the two databases have different values (or if the retrieval
- * from Datastore fails).
- */
-public class ClaimsListDualDatabaseDao {
-
-  /** In-memory cache for claims list. */
-  @NonFinalForTesting
-  private static Supplier<ClaimsListShard> claimsListCache =
-      tryMemoizeWithExpiration(
-          getDomainLabelListCacheDuration(), ClaimsListDualDatabaseDao::getUncached);
-
-  /**
-   * Saves the given {@link ClaimsListShard} to both the primary and secondary databases, logging
-   * and skipping errors in Datastore.
-   */
-  public static void save(ClaimsListShard claimsList) {
-      ClaimsListSqlDao.save(claimsList);
-      suppressExceptionUnlessInTest(
-          claimsList::saveToDatastore, "Error saving ClaimsListShard to Datastore.");
-  }
-
-  /** Returns the most recent revision of the {@link ClaimsListShard}, from cache. */
-  public static ClaimsListShard get() {
-    return claimsListCache.get();
-  }
-
-  /** Retrieves and compares the latest revision from the databases. */
-  private static ClaimsListShard getUncached() {
-    Optional<ClaimsListShard> cloudSqlResult = ClaimsListSqlDao.get();
-    suppressExceptionUnlessInTest(
-        () -> {
-          Optional<ClaimsListShard> datastoreResult = ClaimsListShard.getFromDatastore();
-          compareClaimsLists(cloudSqlResult, datastoreResult);
-        },
-        "Error loading ClaimsListShard from Datastore.");
-    return cloudSqlResult.orElse(ClaimsListShard.create(START_OF_TIME, ImmutableMap.of()));
-  }
-
-  private static void compareClaimsLists(
-      Optional<ClaimsListShard> maybeCloudSql, Optional<ClaimsListShard> maybeDatastore) {
-    if (maybeCloudSql.isPresent() && !maybeDatastore.isPresent()) {
-      throw new IllegalStateException("Claims list found in Cloud SQL but not in Datastore.");
-    }
-    if (!maybeCloudSql.isPresent() && maybeDatastore.isPresent()) {
-      throw new IllegalStateException("Claims list found in Datastore but not in Cloud SQL.");
-    }
-    if (!maybeCloudSql.isPresent()) {
-      return;
-    }
-    ClaimsListShard sqlList = maybeCloudSql.get();
-    ClaimsListShard datastoreList = maybeDatastore.get();
-    MapDifference<String, String> diff =
-        Maps.difference(sqlList.labelsToKeys, datastoreList.getLabelsToKeys());
-    if (!diff.areEqual()) {
-      if (diff.entriesDiffering().size()
-              + diff.entriesOnlyOnRight().size()
-              + diff.entriesOnlyOnLeft().size()
-          > 10) {
-        throw new IllegalStateException(
-            String.format(
-                "Unequal claims lists detected, Datastore list with revision id %d has %d"
-                    + " different records than the current Cloud SQL list.",
-                datastoreList.getRevisionId(), diff.entriesDiffering().size()));
-      } else {
-        StringBuilder diffMessage = new StringBuilder("Unequal claims lists detected:\n");
-        diff.entriesDiffering()
-            .forEach(
-                (label, valueDiff) ->
-                    diffMessage.append(
-                        String.format(
-                            "Domain label %s has key %s in Cloud SQL and key %s "
-                                + "in Datastore.\n",
-                            label, valueDiff.leftValue(), valueDiff.rightValue())));
-        diff.entriesOnlyOnLeft()
-            .forEach(
-                (label, valueDiff) ->
-                    diffMessage.append(
-                        String.format(
-                            "Domain label %s with key %s only appears in Cloud SQL.\n",
-                            label, valueDiff)));
-        diff.entriesOnlyOnRight()
-            .forEach(
-                (label, valueDiff) ->
-                    diffMessage.append(
-                        String.format(
-                            "Domain label %s with key %s only appears in Datastore.\n",
-                            label, valueDiff)));
-        throw new IllegalStateException(diffMessage.toString());
-      }
-    }
-  }
-
-  private ClaimsListDualDatabaseDao() {}
-}

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

@@ -1,359 +0,0 @@
-// Copyright 2017 The Nomulus Authors. All Rights Reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package google.registry.model.tmch;
-
-import static com.google.common.base.Preconditions.checkNotNull;
-import static com.google.common.base.Preconditions.checkState;
-import static com.google.common.base.Throwables.throwIfUnchecked;
-import static com.google.common.base.Verify.verify;
-import static google.registry.model.ofy.ObjectifyService.allocateId;
-import static google.registry.model.ofy.ObjectifyService.auditedOfy;
-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.collect.ImmutableMap;
-import com.google.common.util.concurrent.UncheckedExecutionException;
-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;
-import com.googlecode.objectify.annotation.OnSave;
-import com.googlecode.objectify.annotation.Parent;
-import google.registry.model.CreateAutoTimestamp;
-import google.registry.model.ImmutableObject;
-import google.registry.model.annotations.InCrossTld;
-import google.registry.model.annotations.NotBackedUp;
-import google.registry.model.annotations.NotBackedUp.Reason;
-import google.registry.model.annotations.VirtualEntity;
-import google.registry.model.common.CrossTldSingleton;
-import google.registry.schema.replay.DatastoreOnlyEntity;
-import google.registry.schema.replay.NonReplicatedEntity;
-import google.registry.util.CollectionUtils;
-import google.registry.util.Concurrent;
-import google.registry.util.Retrier;
-import google.registry.util.SystemSleeper;
-import java.util.HashMap;
-import java.util.List;
-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.Table;
-import javax.persistence.Transient;
-import org.joda.time.DateTime;
-
-/**
- * A list of TMCH claims labels and their associated claims keys.
- *
- * <p>The claims list is actually sharded into multiple {@link ClaimsListShard} entities to work
- * around the Datastore limitation of 1M max size per entity. However, when calling {@link
- * #getFromDatastore} all of the shards are recombined into one {@link ClaimsListShard} object.
- *
- * <p>ClaimsList shards are tied to a specific revision and are persisted individually, then the
- * entire claims list is atomically shifted over to using the new shards by persisting the new
- * revision object and updating the {@link ClaimsListSingleton} pointing to it. This bypasses the
- * 10MB per transaction limit.
- *
- * <p>Therefore, it is never OK to save an instance of this class directly to Datastore. Instead you
- * must use the {@link #saveToDatastore} method to do it for you.
- *
- * <p>Note that the primary key of this entity is {@link #revisionId}, which is auto-generated by
- * the database. So, if a retry of insertion happens after the previous attempt unexpectedly
- * succeeds, we will end up with having two exact same claims list with only different {@link
- * #revisionId}. However, this is not an actual problem because we only use the claims list with
- * highest {@link #revisionId}.
- *
- * <p>TODO(b/162007765): Rename the class to ClaimsList and remove Datastore related fields and
- * methods.
- */
-@Entity
-@NotBackedUp(reason = Reason.EXTERNALLY_SOURCED)
-@javax.persistence.Entity(name = "ClaimsList")
-@Table
-@InCrossTld
-public class ClaimsListShard extends ImmutableObject implements NonReplicatedEntity {
-
-  /** The number of claims list entries to store per shard. */
-  private static final int SHARD_SIZE = 10000;
-
-  @Transient @Id long id;
-
-  @Transient @Parent Key<ClaimsListRevision> parent;
-
-  @Ignore
-  @javax.persistence.Id
-  @GeneratedValue(strategy = GenerationType.IDENTITY)
-  Long revisionId;
-
-  @Ignore
-  @Column(nullable = false)
-  CreateAutoTimestamp creationTimestamp = CreateAutoTimestamp.create(null);
-
-  /**
-   * When the claims list was last updated.
-   *
-   * <p>Note that the value of this field is parsed from the claims list file(See this <a
-   * href="https://tools.ietf.org/html/draft-lozano-tmch-func-spec-08#section-6.1">RFC</>), it is
-   * the DNL List creation datetime from the rfc. Since this field has been used by Datastore, we
-   * cannot change its name until we finish the migration.
-   *
-   * <p>TODO(b/177567432): Rename this field to tmdbGenerationTime.
-   */
-  @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;
-
-  /** Indicates that this is a shard rather than a "full" list. */
-  @Ignore @Transient boolean isShard = false;
-
-  private static final Retrier LOADER_RETRIER = new Retrier(new SystemSleeper(), 2);
-
-  private static Optional<ClaimsListShard> loadClaimsListShard() {
-    // Find the most recent revision.
-    Key<ClaimsListRevision> revisionKey = getCurrentRevision();
-    if (revisionKey == null) {
-      return Optional.empty();
-    }
-
-    Map<String, String> combinedLabelsToKeys = new HashMap<>();
-    DateTime creationTime = START_OF_TIME;
-    // Grab all of the keys for the shards that belong to the current revision.
-    final List<Key<ClaimsListShard>> shardKeys =
-        auditedOfy().load().type(ClaimsListShard.class).ancestor(revisionKey).keys().list();
-
-    List<ClaimsListShard> shards;
-    try {
-      // Load all of the shards concurrently, each in a separate transaction.
-      shards =
-          Concurrent.transform(
-              shardKeys,
-              key ->
-                  ofyTm()
-                      .transactNewReadOnly(
-                          () -> {
-                            ClaimsListShard claimsListShard = auditedOfy().load().key(key).now();
-                            checkState(
-                                claimsListShard != null,
-                                "Key not found when loading claims list shards.");
-                            return claimsListShard;
-                          }));
-    } catch (UncheckedExecutionException e) {
-      // We retry on IllegalStateException. However, there's a checkState inside the
-      // Concurrent.transform, so if it's thrown it'll be wrapped in an
-      // UncheckedExecutionException. We want to unwrap it so it's caught by the retrier.
-      if (e.getCause() != null) {
-        throwIfUnchecked(e.getCause());
-      }
-      throw e;
-    }
-
-    // Combine the shards together and return the concatenated ClaimsList.
-    if (!shards.isEmpty()) {
-      creationTime = shards.get(0).creationTime;
-      for (ClaimsListShard shard : shards) {
-        combinedLabelsToKeys.putAll(shard.labelsToKeys);
-        checkState(
-            creationTime.equals(shard.creationTime),
-            "Inconsistent claims list shard creation times.");
-      }
-    }
-
-    return Optional.of(create(creationTime, ImmutableMap.copyOf(combinedLabelsToKeys)));
-  }
-
-  /** Returns the revision id of this claims list, or throws exception if it is null. */
-  public Long getRevisionId() {
-    checkState(
-        revisionId != null, "revisionId is null because it is not persisted in the database");
-    return revisionId;
-  }
-
-  /**
-   * Returns the time when the external TMDB service generated this revision of the claims list.
-   *
-   * @see <a href="https://tools.ietf.org/html/draft-lozano-tmch-func-spec-08#section-6.1">DNL List
-   *     creation datetime</a>
-   */
-  public DateTime getTmdbGenerationTime() {
-    return creationTime;
-  }
-
-  /** Returns the creation time of this claims list. */
-  public DateTime getCreationTimestamp() {
-    return creationTimestamp.getTimestamp();
-  }
-
-  /** Returns the claim key for a given domain if there is one, empty otherwise. */
-  public Optional<String> getClaimKey(String label) {
-    return Optional.ofNullable(labelsToKeys.get(label));
-  }
-
-  /** Returns an {@link Map} mapping domain label to its lookup key. */
-  public ImmutableMap<String, String> getLabelsToKeys() {
-    return ImmutableMap.copyOf(labelsToKeys);
-  }
-
-  /** Returns the number of claims. */
-  public int size() {
-    return labelsToKeys.size();
-  }
-
-  /**
-   * Save the Claims list to Datastore by writing the new shards in a series of transactions,
-   * switching over to using them atomically, then deleting the old ones.
-   */
-  void saveToDatastore() {
-    saveToDatastore(SHARD_SIZE);
-  }
-
-  @VisibleForTesting
-  void saveToDatastore(int shardSize) {
-    // Figure out what the next versionId should be based on which ones already exist.
-    final Key<ClaimsListRevision> oldRevision = getCurrentRevision();
-    final Key<ClaimsListRevision> parentKey = ClaimsListRevision.createKey();
-
-    // Save the ClaimsList shards in separate transactions.
-    Concurrent.transform(
-        CollectionUtils.partitionMap(labelsToKeys, shardSize),
-        (final ImmutableMap<String, String> labelsToKeysShard) ->
-            ofyTm()
-                .transact(
-                    () -> {
-                      ClaimsListShard shard = create(creationTime, labelsToKeysShard);
-                      shard.isShard = true;
-                      shard.parent = parentKey;
-                      auditedOfy().saveWithoutBackup().entity(shard);
-                      return shard;
-                    }));
-
-    // Persist the new revision, thus causing the newly created shards to go live.
-    ofyTm()
-        .transact(
-            () -> {
-              verify(
-                  (getCurrentRevision() == null && oldRevision == null)
-                      || getCurrentRevision().equals(oldRevision),
-                  "Registries' ClaimsList was updated by someone else while attempting to update.");
-              auditedOfy().saveWithoutBackup().entity(ClaimsListSingleton.create(parentKey));
-              // Delete the old ClaimsListShard entities.
-              if (oldRevision != null) {
-                auditedOfy()
-                    .deleteWithoutBackup()
-                    .keys(
-                        auditedOfy()
-                            .load()
-                            .type(ClaimsListShard.class)
-                            .ancestor(oldRevision)
-                            .keys());
-              }
-            });
-  }
-
-  public static ClaimsListShard create(
-      DateTime tmdbGenerationTime, Map<String, String> labelsToKeys) {
-    ClaimsListShard instance = new ClaimsListShard();
-    instance.id = allocateId();
-    instance.creationTime = checkNotNull(tmdbGenerationTime);
-    instance.labelsToKeys = checkNotNull(labelsToKeys);
-    return instance;
-  }
-
-  /** Return a single logical instance that combines all Datastore shards. */
-  static Optional<ClaimsListShard> getFromDatastore() {
-    return LOADER_RETRIER.callWithRetry(
-        ClaimsListShard::loadClaimsListShard, IllegalStateException.class);
-  }
-
-  /** As a safety mechanism, fail if someone tries to save this class directly. */
-  @OnSave
-  void disallowUnshardedSaves() {
-    if (!isShard) {
-      throw new UnshardedSaveException();
-    }
-  }
-
-  /** Virtual parent entity for claims list shards of a specific revision. */
-  @Entity
-  @VirtualEntity
-  public static class ClaimsListRevision extends ImmutableObject implements DatastoreOnlyEntity {
-    @Parent Key<ClaimsListSingleton> parent;
-
-    @Id long versionId;
-
-    @VisibleForTesting
-    public static Key<ClaimsListRevision> createKey(ClaimsListSingleton singleton) {
-      ClaimsListRevision revision = new ClaimsListRevision();
-      revision.versionId = allocateId();
-      revision.parent = Key.create(singleton);
-      return Key.create(revision);
-    }
-
-    @VisibleForTesting
-    public static Key<ClaimsListRevision> createKey() {
-      return createKey(new ClaimsListSingleton());
-    }
-  }
-
-  /**
-   * Serves as the coordinating claims list singleton linking to the {@link ClaimsListRevision} that
-   * is live.
-   */
-  @Entity
-  @NotBackedUp(reason = Reason.EXTERNALLY_SOURCED)
-  public static class ClaimsListSingleton extends CrossTldSingleton implements DatastoreOnlyEntity {
-    Key<ClaimsListRevision> activeRevision;
-
-    static ClaimsListSingleton create(Key<ClaimsListRevision> revision) {
-      ClaimsListSingleton instance = new ClaimsListSingleton();
-      instance.activeRevision = revision;
-      return instance;
-    }
-
-    @VisibleForTesting
-    public void setActiveRevision(Key<ClaimsListRevision> revision) {
-      activeRevision = revision;
-    }
-  }
-
-  /**
-   * Returns the current ClaimsListRevision if there is one, or null if no claims list revisions
-   * have ever been persisted yet.
-   */
-  @Nullable
-  public static Key<ClaimsListRevision> getCurrentRevision() {
-    ClaimsListSingleton singleton = auditedOfy().load().entity(new ClaimsListSingleton()).now();
-    return singleton == null ? null : singleton.activeRevision;
-  }
-
-  /** Exception when trying to directly save a {@link ClaimsListShard} without sharding. */
-  public static class UnshardedSaveException extends RuntimeException {}
-}

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

@@ -15,17 +15,14 @@
 package google.registry.model.tmch;
 
 import static com.google.common.base.Preconditions.checkNotNull;
-import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
 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 com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Entity;
 import google.registry.model.annotations.NotBackedUp;
 import google.registry.model.annotations.NotBackedUp.Reason;
 import google.registry.model.common.CrossTldSingleton;
-import google.registry.persistence.VKey;
 import google.registry.schema.replay.NonReplicatedEntity;
 import java.util.Optional;
 import javax.annotation.concurrent.Immutable;
@@ -50,13 +47,7 @@ public final class TmchCrl extends CrossTldSingleton implements NonReplicatedEnt
 
   /** Returns the singleton instance of this entity, without memoization. */
   public static Optional<TmchCrl> get() {
-    return tm().transact(
-            () ->
-                tm().loadByKeyIfPresent(
-                        VKey.create(
-                            TmchCrl.class,
-                            SINGLETON_ID,
-                            Key.create(getCrossTldKey(), TmchCrl.class, SINGLETON_ID))));
+    return tm().transact(() -> tm().loadSingleton(TmchCrl.class));
   }
 
   /**

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;
@@ -37,26 +38,38 @@ import google.registry.model.index.ForeignKeyIndex.ForeignKeyDomainIndex;
 import google.registry.model.index.ForeignKeyIndex.ForeignKeyHostIndex;
 import google.registry.model.ofy.DatastoreTransactionManager;
 import google.registry.model.server.KmsSecret;
-import google.registry.model.tmch.ClaimsListShard.ClaimsListSingleton;
+import google.registry.model.tmch.ClaimsList.ClaimsListSingleton;
 import google.registry.persistence.JpaRetries;
 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
@@ -350,6 +365,11 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     entities.forEach(this::update);
   }
 
+  @Override
+  public void updateAll(Object... entities) {
+    updateAll(ImmutableList.of(entities));
+  }
+
   @Override
   public void updateWithoutBackup(Object entity) {
     update(entity);
@@ -392,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
@@ -406,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));
   }
@@ -428,7 +449,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     if (result == null) {
       throw new NoSuchElementException(key.toString());
     }
-    return result;
+    return detach(result);
   }
 
   @Override
@@ -451,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
@@ -467,12 +491,26 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   public <T> ImmutableList<T> loadAllOf(Class<T> clazz) {
     checkArgumentNotNull(clazz, "clazz must be specified");
     assertInTransaction();
-    return ImmutableList.copyOf(
+    return getEntityManager()
+        .createQuery(String.format("FROM %s", getEntityType(clazz).getName()), clazz)
+        .getResultStream()
+        .map(this::detach)
+        .collect(toImmutableList());
+  }
+
+  @Override
+  public <T> Optional<T> loadSingleton(Class<T> clazz) {
+    assertInTransaction();
+    List<T> elements =
         getEntityManager()
-            .createQuery(
-                String.format("SELECT entity FROM %s entity", getEntityType(clazz).getName()),
-                clazz)
-            .getResultList());
+            .createQuery(String.format("FROM %s", getEntityType(clazz).getName()), clazz)
+            .setMaxResults(2)
+            .getResultList();
+    checkArgument(
+        elements.size() <= 1,
+        "Expected at most one entity of type %s, found at least two",
+        clazz.getSimpleName());
+    return elements.stream().findFirst().map(this::detach);
   }
 
   private int internalDelete(VKey<?> key) {
@@ -504,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
@@ -535,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
@@ -638,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;
@@ -646,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);
@@ -660,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.
@@ -688,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);
     }
 
@@ -719,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
@@ -738,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
@@ -757,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

@@ -159,6 +159,9 @@ public interface TransactionManager {
   /** Updates all entities in the database, throws exception if any entity does not exist. */
   void updateAll(ImmutableCollection<?> entities);
 
+  /** Updates all entities in the database, throws exception if any entity does not exist. */
+  void updateAll(Object... entities);
+
   /**
    * Updates an entity in the database without writing commit logs if the underlying database is
    * Datastore.
@@ -246,14 +249,27 @@ public interface TransactionManager {
    */
   <T> ImmutableList<T> loadAllOf(Class<T> clazz);
 
+  /**
+   * Loads the only instance of this particular class, or empty if none exists.
+   *
+   * <p>Throws an exception if there is more than one element in the table.
+   */
+  <T> Optional<T> loadSingleton(Class<T> clazz);
+
   /** Deletes the entity by its id. */
   void delete(VKey<?> key);
 
   /** 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/schema/replay/SqlReplayCheckpoint.java

@@ -14,26 +14,24 @@
 
 package google.registry.schema.replay;
 
-import static google.registry.model.common.CrossTldSingleton.SINGLETON_ID;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
 
+import google.registry.model.common.CrossTldSingleton;
+import javax.persistence.Column;
 import javax.persistence.Entity;
-import javax.persistence.Id;
 import org.joda.time.DateTime;
 
 @Entity
-public class SqlReplayCheckpoint implements SqlOnlyEntity {
-
-  // Hibernate doesn't allow us to have a converted DateTime as our primary key so we need this
-  @Id private long revisionId = SINGLETON_ID;
+public class SqlReplayCheckpoint extends CrossTldSingleton implements SqlOnlyEntity {
 
+  @Column(nullable = false)
   private DateTime lastReplayTime;
 
   public static DateTime get() {
     jpaTm().assertInTransaction();
-    return jpaTm().loadAllOf(SqlReplayCheckpoint.class).stream()
-        .findFirst()
+    return jpaTm()
+        .loadSingleton(SqlReplayCheckpoint.class)
         .map(checkpoint -> checkpoint.lastReplayTime)
         .orElse(START_OF_TIME);
   }

core/src/main/java/google/registry/schema/tld/PremiumEntry.java

@@ -52,8 +52,9 @@ public class PremiumEntry extends ImmutableObject implements Serializable, SqlOn
     return domainLabel;
   }
 
-  public static PremiumEntry create(BigDecimal price, String domainLabel) {
+  public static PremiumEntry create(long revisionId, BigDecimal price, String domainLabel) {
     PremiumEntry result = new PremiumEntry();
+    result.revisionId = revisionId;
     result.price = price;
     result.domainLabel = domainLabel;
     return result;

core/src/main/java/google/registry/tmch/ClaimsListParser.java

@@ -18,7 +18,7 @@ import static com.google.common.base.Preconditions.checkArgument;
 
 import com.google.common.base.Splitter;
 import com.google.common.collect.ImmutableMap;
-import google.registry.model.tmch.ClaimsListShard;
+import google.registry.model.tmch.ClaimsList;
 import java.util.List;
 import org.joda.time.DateTime;
 
@@ -34,11 +34,11 @@ import org.joda.time.DateTime;
 public class ClaimsListParser {
 
   /**
-   * Converts the lines from the DNL CSV file into a {@link ClaimsListShard} object.
+   * Converts the lines from the DNL CSV file into a {@link ClaimsList} object.
    *
    * <p>Please note that this does <b>not</b> insert the object into Datastore.
    */
-  public static ClaimsListShard parse(List<String> lines) {
+  public static ClaimsList parse(List<String> lines) {
     ImmutableMap.Builder<String, String> builder = new ImmutableMap.Builder<>();
 
     // First line: <version>,<DNL List creation datetime>
@@ -74,6 +74,6 @@ public class ClaimsListParser {
       builder.put(label, lookupKey);
     }
 
-    return ClaimsListShard.create(creationTime, builder.build());
+    return ClaimsList.create(creationTime, builder.build());
   }
 }

core/src/main/java/google/registry/tmch/TmchDnlAction.java

@@ -18,8 +18,8 @@ import static google.registry.request.Action.Method.POST;
 
 import com.google.common.flogger.FluentLogger;
 import google.registry.keyring.api.KeyModule.Key;
-import google.registry.model.tmch.ClaimsListDualDatabaseDao;
-import google.registry.model.tmch.ClaimsListShard;
+import google.registry.model.tmch.ClaimsList;
+import google.registry.model.tmch.ClaimsListDao;
 import google.registry.request.Action;
 import google.registry.request.auth.Auth;
 import java.io.IOException;
@@ -55,8 +55,8 @@ public final class TmchDnlAction implements Runnable {
     } catch (SignatureException | IOException | PGPException e) {
       throw new RuntimeException(e);
     }
-    ClaimsListShard claims = ClaimsListParser.parse(lines);
-    ClaimsListDualDatabaseDao.save(claims);
+    ClaimsList claims = ClaimsListParser.parse(lines);
+    ClaimsListDao.save(claims);
     logger.atInfo().log(
         "Inserted %,d claims into the DB(s), created at %s",
         claims.size(), claims.getTmdbGenerationTime());

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

@@ -22,7 +22,6 @@ import static google.registry.tools.LockOrUnlockDomainCommand.REGISTRY_LOCK_STAT
 
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Sets;
-import com.googlecode.objectify.Key;
 import google.registry.batch.AsyncTaskEnqueuer;
 import google.registry.config.RegistryConfig.Config;
 import google.registry.model.billing.BillingEvent;
@@ -370,7 +369,7 @@ public final class DomainLockUtils {
             .setRequestedByRegistrar(!lock.isSuperuser())
             .setType(HistoryEntry.Type.DOMAIN_UPDATE)
             .setModificationTime(now)
-            .setParent(Key.create(domain))
+            .setDomain(domain)
             .setReason(reason)
             .build();
     tm().update(domain);

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

@@ -21,8 +21,8 @@ import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
 import com.google.common.base.Joiner;
 import com.google.common.io.Files;
-import google.registry.model.tmch.ClaimsListDualDatabaseDao;
-import google.registry.model.tmch.ClaimsListShard;
+import google.registry.model.tmch.ClaimsList;
+import google.registry.model.tmch.ClaimsListDao;
 import google.registry.tools.params.PathParameter;
 import java.nio.file.Path;
 import java.nio.file.Paths;
@@ -44,7 +44,7 @@ final class GetClaimsListCommand implements CommandWithRemoteApi {
 
   @Override
   public void run() throws Exception {
-    ClaimsListShard cl = checkNotNull(ClaimsListDualDatabaseDao.get(), "Couldn't load ClaimsList");
+    ClaimsList cl = checkNotNull(ClaimsListDao.get(), "Couldn't load ClaimsList");
     String csv = Joiner.on('\n').withKeyValueSeparator(",").join(cl.getLabelsToKeys()) + "\n";
     Files.asCharSink(output.toFile(), UTF_8).write(csv);
   }

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

@@ -185,7 +185,7 @@ class UnrenewDomainCommand extends ConfirmingCommand implements CommandWithRemot
         leapSafeSubtractYears(domain.getRegistrationExpirationTime(), period);
     DomainHistory domainHistory =
         new DomainHistory.Builder()
-            .setParent(domain)
+            .setDomain(domain)
             .setModificationTime(now)
             .setBySuperuser(true)
             .setType(Type.SYNTHETIC)

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);
     }