Retry flaky tests for ReplicateToDatastoreAction (#1226)

* Retry flaky tests for ReplicateToDatastoreAction

The occassional failures seem to be caused by the test Datastore.

.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

@@ -79,6 +79,9 @@ def fragileTestPatterns = [
     // Changes cache timeouts and for some reason appears to have contention
     // with other tests.
     "google/registry/whois/WhoisCommandFactoryTest.*",
+    // Currently changes a global configuration parameter that for some reason
+    // results in timestamp inversions for other tests.  TODO(mmuller): fix.
+    "google/registry/flows/host/HostInfoFlowTest.*",
 ] + dockerIncompatibleTestPatterns
 
 sourceSets {
@@ -317,6 +320,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/BackupUtils.java

@@ -56,12 +56,16 @@ public class BackupUtils {
    *
    * <p>The iterator reads from the stream on demand, and as such will fail if the stream is closed.
    */
-  public static Iterator<ImmutableObject> createDeserializingIterator(final InputStream input) {
+  public static Iterator<ImmutableObject> createDeserializingIterator(
+      final InputStream input, boolean withAppIdOverride) {
     return new AbstractIterator<ImmutableObject>() {
       @Override
       protected ImmutableObject computeNext() {
         EntityProto proto = new EntityProto();
         if (proto.parseDelimitedFrom(input)) { // False means end of stream; other errors throw.
+          if (withAppIdOverride) {
+            proto = EntityImports.fixEntity(proto);
+          }
           return auditedOfy().load().fromEntity(EntityTranslator.createFromPb(proto));
         }
         return endOfData();
@@ -70,6 +74,7 @@ public class BackupUtils {
   }
 
   public static ImmutableList<ImmutableObject> deserializeEntities(byte[] bytes) {
-    return ImmutableList.copyOf(createDeserializingIterator(new ByteArrayInputStream(bytes)));
+    return ImmutableList.copyOf(
+        createDeserializingIterator(new ByteArrayInputStream(bytes), false));
   }
 }

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

@@ -59,7 +59,7 @@ public final class CommitLogImports {
       InputStream inputStream) {
     try (AppEngineEnvironment appEngineEnvironment = new AppEngineEnvironment();
         InputStream input = new BufferedInputStream(inputStream)) {
-      Iterator<ImmutableObject> commitLogs = createDeserializingIterator(input);
+      Iterator<ImmutableObject> commitLogs = createDeserializingIterator(input, false);
       checkState(commitLogs.hasNext());
       checkState(commitLogs.next() instanceof CommitLogCheckpoint);
 

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

@@ -0,0 +1,115 @@
+// 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.backup;
+
+import com.google.apphosting.api.ApiProxy;
+import com.google.storage.onestore.v3.OnestoreEntity;
+import com.google.storage.onestore.v3.OnestoreEntity.EntityProto;
+import com.google.storage.onestore.v3.OnestoreEntity.Path;
+import com.google.storage.onestore.v3.OnestoreEntity.Property.Meaning;
+import com.google.storage.onestore.v3.OnestoreEntity.PropertyValue.ReferenceValue;
+import java.nio.charset.StandardCharsets;
+import java.util.Objects;
+
+/** Utilities for handling imported Datastore entities. */
+public class EntityImports {
+
+  /**
+   * Transitively sets the {@code appId} of all keys in a foreign entity to that of the current
+   * system.
+   */
+  public static EntityProto fixEntity(EntityProto entityProto) {
+    String currentAappId = ApiProxy.getCurrentEnvironment().getAppId();
+    if (Objects.equals(currentAappId, entityProto.getKey().getApp())) {
+      return entityProto;
+    }
+    return fixEntity(entityProto, currentAappId);
+  }
+
+  private static EntityProto fixEntity(EntityProto entityProto, String appId) {
+    if (entityProto.hasKey()) {
+      fixKey(entityProto, appId);
+    }
+
+    for (OnestoreEntity.Property property : entityProto.mutablePropertys()) {
+      fixProperty(property, appId);
+    }
+
+    for (OnestoreEntity.Property property : entityProto.mutableRawPropertys()) {
+      fixProperty(property, appId);
+    }
+
+    // CommitLogMutation embeds an entity as bytes, which needs additional fixes.
+    if (isCommitLogMutation(entityProto)) {
+      fixMutationEntityProtoBytes(entityProto, appId);
+    }
+    return entityProto;
+  }
+
+  private static boolean isCommitLogMutation(EntityProto entityProto) {
+    if (!entityProto.hasKey()) {
+      return false;
+    }
+    Path path = entityProto.getKey().getPath();
+    if (path.elementSize() == 0) {
+      return false;
+    }
+    return Objects.equals(
+        path.getElement(path.elementSize() - 1).getType(StandardCharsets.UTF_8),
+        "CommitLogMutation");
+  }
+
+  private static void fixMutationEntityProtoBytes(EntityProto entityProto, String appId) {
+    for (OnestoreEntity.Property property : entityProto.mutableRawPropertys()) {
+      if (Objects.equals(property.getName(), "entityProtoBytes")) {
+        OnestoreEntity.PropertyValue value = property.getValue();
+        EntityProto fixedProto =
+            fixEntity(bytesToEntityProto(value.getStringValueAsBytes()), appId);
+        value.setStringValueAsBytes(fixedProto.toByteArray());
+        return;
+      }
+    }
+  }
+
+  private static void fixKey(EntityProto entityProto, String appId) {
+    entityProto.getMutableKey().setApp(appId);
+  }
+
+  private static void fixKey(ReferenceValue referenceValue, String appId) {
+    referenceValue.setApp(appId);
+  }
+
+  private static void fixProperty(OnestoreEntity.Property property, String appId) {
+    OnestoreEntity.PropertyValue value = property.getMutableValue();
+    if (value.hasReferenceValue()) {
+      fixKey(value.getMutableReferenceValue(), appId);
+      return;
+    }
+    if (property.getMeaningEnum().equals(Meaning.ENTITY_PROTO)) {
+      EntityProto embeddedProto = bytesToEntityProto(value.getStringValueAsBytes());
+      fixEntity(embeddedProto, appId);
+      value.setStringValueAsBytes(embeddedProto.toByteArray());
+    }
+  }
+
+  private static EntityProto bytesToEntityProto(byte[] bytes) {
+    EntityProto entityProto = new EntityProto();
+    boolean isParsed = entityProto.parseFrom(bytes);
+    if (!isParsed) {
+      throw new IllegalStateException("Failed to parse raw bytes as EntityProto.");
+    }
+    return entityProto;
+  }
+}

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,31 +57,38 @@ 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 {
 
   static final String PATH = "/_dr/task/replayCommitLogsToSql";
 
+  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
   private static final int BLOCK_SIZE =
       1024 * 1024; // Buffer 1mb at a time, for no particular reason.
-  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
   private static final Duration LEASE_LENGTH = standardHours(1);
+  // Stop / pause where we are if we've been replaying for more than five minutes to avoid GAE
+  // request timeouts
+  private static final Duration REPLAY_TIMEOUT_DURATION = Duration.standardMinutes(5);
 
   @Inject GcsService gcsService;
   @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);
@@ -105,17 +116,26 @@ public class ReplayCommitLogsToSqlAction implements Runnable {
   }
 
   private void replayFiles() {
+    DateTime replayTimeoutTime = clock.nowUtc().plus(REPLAY_TIMEOUT_DURATION);
     // Start at the first millisecond we haven't seen yet
     DateTime fromTime = jpaTm().transact(() -> SqlReplayCheckpoint.get().plusMillis(1));
     // If there's an inconsistent file set, this will throw IllegalStateException and the job
     // will try later -- this is likely because an export hasn't finished yet.
     ImmutableList<GcsFileMetadata> commitLogFiles =
         diffLister.listDiffFiles(fromTime, /* current time */ null);
+    int processedFiles = 0;
     for (GcsFileMetadata metadata : commitLogFiles) {
       // One transaction per GCS file
       jpaTm().transact(() -> processFile(metadata));
+      processedFiles++;
+      if (clock.nowUtc().isAfter(replayTimeoutTime)) {
+        logger.atInfo().log(
+            "Reached max execution time after replaying %d files, leaving %d files for next run.",
+            processedFiles, commitLogFiles.size() - processedFiles);
+        return;
+      }
     }
-    logger.atInfo().log("Replayed %d commit log files to SQL successfully.", commitLogFiles.size());
+    logger.atInfo().log("Replayed %d commit log files to SQL successfully.", processedFiles);
   }
 
   private void processFile(GcsFileMetadata metadata) {

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

@@ -104,7 +104,7 @@ public class RestoreCommitLogsAction implements Runnable {
       try (InputStream input = Channels.newInputStream(
           gcsService.openPrefetchingReadChannel(metadata.getFilename(), 0, BLOCK_SIZE))) {
         PeekingIterator<ImmutableObject> commitLogs =
-            peekingIterator(createDeserializingIterator(input));
+            peekingIterator(createDeserializingIterator(input, true));
         lastCheckpoint = (CommitLogCheckpoint) commitLogs.next();
         saveOfy(ImmutableList.of(lastCheckpoint));  // Save the checkpoint itself.
         while (commitLogs.hasNext()) {

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

@@ -15,12 +15,14 @@
 package google.registry.batch;
 
 import static com.google.common.base.Preconditions.checkState;
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static google.registry.config.RegistryEnvironment.PRODUCTION;
 import static google.registry.mapreduce.MapreduceRunner.PARAM_DRY_RUN;
 import static google.registry.mapreduce.inputs.EppResourceInputs.createEntityInput;
 import static google.registry.model.ofy.ObjectifyService.auditedOfy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.request.Action.Method.POST;
+import static google.registry.util.DateTimeUtils.END_OF_TIME;
 
 import com.google.appengine.tools.mapreduce.Mapper;
 import com.google.common.collect.ImmutableList;
@@ -28,16 +30,24 @@ import com.google.common.collect.ImmutableSet;
 import com.google.common.flogger.FluentLogger;
 import com.googlecode.objectify.Key;
 import google.registry.config.RegistryEnvironment;
+import google.registry.flows.poll.PollFlowUtils;
 import google.registry.mapreduce.MapreduceRunner;
 import google.registry.model.EppResource;
+import google.registry.model.EppResourceUtils;
 import google.registry.model.contact.ContactResource;
+import google.registry.model.domain.DomainBase;
 import google.registry.model.host.HostResource;
 import google.registry.model.index.EppResourceIndex;
 import google.registry.model.index.ForeignKeyIndex;
+import google.registry.model.poll.PollMessage;
+import google.registry.model.reporting.HistoryEntry;
+import google.registry.model.reporting.HistoryEntryDao;
+import google.registry.persistence.VKey;
 import google.registry.request.Action;
 import google.registry.request.Parameter;
 import google.registry.request.Response;
 import google.registry.request.auth.Auth;
+import google.registry.util.Clock;
 import java.util.List;
 import javax.inject.Inject;
 
@@ -46,8 +56,8 @@ import javax.inject.Inject;
  * the associated ForeignKey and EppResourceIndex entities.
  *
  * <p>This only deletes contacts and hosts, NOT domains. To delete domains, use {@link
- * DeleteLoadTestDataAction} and pass it the TLD(s) that the load test domains were created on. Note
- * that DeleteLoadTestDataAction is safe enough to run in production whereas this mapreduce is not,
+ * DeleteProberDataAction} and pass it the TLD(s) that the load test domains were created on. Note
+ * that DeleteProberDataAction is safe enough to run in production whereas this mapreduce is not,
  * but this one does not need to be runnable in production because load testing isn't run against
  * production.
  */
@@ -68,15 +78,22 @@ public class DeleteLoadTestDataAction implements Runnable {
    */
   private static final ImmutableSet<String> LOAD_TEST_REGISTRARS = ImmutableSet.of("proxy");
 
-  @Inject
-  @Parameter(PARAM_DRY_RUN)
-  boolean isDryRun;
-
-  @Inject MapreduceRunner mrRunner;
-  @Inject Response response;
+  private final boolean isDryRun;
+  private final MapreduceRunner mrRunner;
+  private final Response response;
+  private final Clock clock;
 
   @Inject
-  DeleteLoadTestDataAction() {}
+  DeleteLoadTestDataAction(
+      @Parameter(PARAM_DRY_RUN) boolean isDryRun,
+      MapreduceRunner mrRunner,
+      Response response,
+      Clock clock) {
+    this.isDryRun = isDryRun;
+    this.mrRunner = mrRunner;
+    this.response = response;
+    this.clock = clock;
+  }
 
   @Override
   public void run() {
@@ -87,14 +104,85 @@ public class DeleteLoadTestDataAction implements Runnable {
         !RegistryEnvironment.get().equals(PRODUCTION),
         "This mapreduce is not safe to run on PRODUCTION.");
 
-    mrRunner
-        .setJobName("Delete load test data")
-        .setModuleName("backend")
-        .runMapOnly(
-            new DeleteLoadTestDataMapper(isDryRun),
-            ImmutableList.of(
-                createEntityInput(ContactResource.class), createEntityInput(HostResource.class)))
-        .sendLinkToMapreduceConsole(response);
+    if (tm().isOfy()) {
+      mrRunner
+          .setJobName("Delete load test data")
+          .setModuleName("backend")
+          .runMapOnly(
+              new DeleteLoadTestDataMapper(isDryRun),
+              ImmutableList.of(
+                  createEntityInput(ContactResource.class), createEntityInput(HostResource.class)))
+          .sendLinkToMapreduceConsole(response);
+    } else {
+      tm().transact(
+              () -> {
+                LOAD_TEST_REGISTRARS.forEach(this::deletePollMessages);
+                tm().loadAllOfStream(ContactResource.class).forEach(this::deleteContact);
+                tm().loadAllOfStream(HostResource.class).forEach(this::deleteHost);
+              });
+    }
+  }
+
+  private void deletePollMessages(String registrarId) {
+    ImmutableList<PollMessage> pollMessages =
+        PollFlowUtils.createPollMessageQuery(registrarId, END_OF_TIME).list();
+    if (isDryRun) {
+      logger.atInfo().log(
+          "Would delete %d poll messages for registrar %s.", pollMessages.size(), registrarId);
+    } else {
+      pollMessages.forEach(tm()::delete);
+    }
+  }
+
+  private void deleteContact(ContactResource contact) {
+    if (!LOAD_TEST_REGISTRARS.contains(contact.getPersistedCurrentSponsorClientId())) {
+      return;
+    }
+    // We cannot remove contacts from domains in the general case, so we cannot delete contacts
+    // that are linked to domains (since it would break the foreign keys)
+    if (EppResourceUtils.isLinked(contact.createVKey(), clock.nowUtc())) {
+      logger.atWarning().log(
+          "Cannot delete contact with repo ID %s since it is referenced from a domain",
+          contact.getRepoId());
+      return;
+    }
+    deleteResource(contact);
+  }
+
+  private void deleteHost(HostResource host) {
+    if (!LOAD_TEST_REGISTRARS.contains(host.getPersistedCurrentSponsorClientId())) {
+      return;
+    }
+    VKey<HostResource> hostVKey = host.createVKey();
+    // We can remove hosts from linked domains, so we should do so then delete the hosts
+    ImmutableSet<VKey<DomainBase>> linkedDomains =
+        EppResourceUtils.getLinkedDomainKeys(hostVKey, clock.nowUtc(), null);
+    tm().loadByKeys(linkedDomains)
+        .values()
+        .forEach(
+            domain -> {
+              ImmutableSet<VKey<HostResource>> remainingHosts =
+                  domain.getNsHosts().stream()
+                      .filter(vkey -> !vkey.equals(hostVKey))
+                      .collect(toImmutableSet());
+              tm().put(domain.asBuilder().setNameservers(remainingHosts).build());
+            });
+    deleteResource(host);
+  }
+
+  private void deleteResource(EppResource eppResource) {
+    // In SQL, the only objects parented on the resource are poll messages (deleted above) and
+    // history objects.
+    ImmutableList<HistoryEntry> historyObjects =
+        HistoryEntryDao.loadHistoryObjectsForResource(eppResource.createVKey());
+    if (isDryRun) {
+      logger.atInfo().log(
+          "Would delete repo ID %s along with %d history objects",
+          eppResource.getRepoId(), historyObjects.size());
+    } else {
+      historyObjects.forEach(tm()::delete);
+      tm().delete(eppResource);
+    }
   }
 
   /** Provides the map method that runs for each existing contact and host entity. */

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

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

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

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

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

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

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

@@ -18,17 +18,23 @@ import static com.google.common.base.Verify.verify;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 
 import google.registry.backup.AppEngineEnvironment;
+import google.registry.model.contact.ContactResource;
+import google.registry.persistence.transaction.CriteriaQueryBuilder;
 import google.registry.persistence.transaction.JpaTransactionManager;
 import java.io.Serializable;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.metrics.Counter;
 import org.apache.beam.sdk.metrics.Metrics;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
-import org.apache.beam.sdk.transforms.Create;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.ParDo;
 
-/** Toy pipeline that demonstrates how to use {@link JpaTransactionManager} in BEAM pipelines. */
+/**
+ * Toy pipeline that demonstrates how to use {@link JpaTransactionManager} in BEAM pipelines.
+ *
+ * <p>This pipeline may also be used as an integration test for {@link RegistryJpaIO.Read} in a
+ * project with realistic data.
+ */
 public class JpaDemoPipeline implements Serializable {
 
   public static void main(String[] args) {
@@ -38,23 +44,16 @@ public class JpaDemoPipeline implements Serializable {
 
     Pipeline pipeline = Pipeline.create(options);
     pipeline
-        .apply("Start", Create.of((Void) null))
         .apply(
-            "Generate Elements",
-            ParDo.of(
-                new DoFn<Void, Void>() {
-                  @ProcessElement
-                  public void processElement(OutputReceiver<Void> output) {
-                    for (int i = 0; i < 500; i++) {
-                      output.output(null);
-                    }
-                  }
-                }))
+            "Read contacts",
+            RegistryJpaIO.read(
+                () -> CriteriaQueryBuilder.create(ContactResource.class).build(),
+                ContactResource::getRepoId))
         .apply(
-            "Make Query",
+            "Count Contacts",
             ParDo.of(
-                new DoFn<Void, Void>() {
-                  private Counter counter = Metrics.counter("Demo", "Read");
+                new DoFn<String, Void>() {
+                  private Counter counter = Metrics.counter("Contacts", "Read");
 
                   @ProcessElement
                   public void processElement() {

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

@@ -15,6 +15,7 @@
 package google.registry.flows;
 
 import static com.google.common.base.Preconditions.checkState;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.common.base.Strings;
 import dagger.Module;
@@ -222,6 +223,7 @@ public class FlowModule {
           String clientId,
           EppInput eppInput) {
     builder
+        .setModificationTime(tm().getTransactionTime())
         .setTrid(trid)
         .setXmlBytes(inputXmlBytes)
         .setBySuperuser(isSuperuser)

core/src/main/java/google/registry/flows/FlowUtils.java

@@ -15,6 +15,7 @@
 package google.registry.flows;
 
 import static com.google.common.base.Preconditions.checkState;
+import static google.registry.model.IdService.allocateId;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.xml.ValidationMode.LENIENT;
 import static google.registry.xml.ValidationMode.STRICT;
@@ -33,7 +34,6 @@ import google.registry.model.eppcommon.EppXmlTransformer;
 import google.registry.model.eppinput.EppInput.WrongProtocolVersionException;
 import google.registry.model.eppoutput.EppOutput;
 import google.registry.model.host.InetAddressAdapter.IpVersionMismatchException;
-import google.registry.model.ofy.ObjectifyService;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.translators.CurrencyUnitAdapter.UnknownCurrencyException;
 import google.registry.xml.XmlException;
@@ -105,7 +105,7 @@ public final class FlowUtils {
 
   public static <H extends HistoryEntry> Key<H> createHistoryKey(
       EppResource parent, Class<H> clazz) {
-    return Key.create(Key.create(parent), clazz, ObjectifyService.allocateId());
+    return Key.create(Key.create(parent), clazz, allocateId());
   }
 
   /** Registrar is not logged in. */

core/src/main/java/google/registry/flows/contact/ContactCreateFlow.java

@@ -19,6 +19,7 @@ import static google.registry.flows.ResourceFlowUtils.verifyResourceDoesNotExist
 import static google.registry.flows.contact.ContactFlowUtils.validateAsciiPostalInfo;
 import static google.registry.flows.contact.ContactFlowUtils.validateContactAgainstPolicy;
 import static google.registry.model.EppResourceUtils.createRepoId;
+import static google.registry.model.IdService.allocateId;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.common.collect.ImmutableSet;
@@ -41,7 +42,6 @@ import google.registry.model.eppoutput.CreateData.ContactCreateData;
 import google.registry.model.eppoutput.EppResponse;
 import google.registry.model.index.EppResourceIndex;
 import google.registry.model.index.ForeignKeyIndex;
-import google.registry.model.ofy.ObjectifyService;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.reporting.IcannReportingTypes.ActivityReportField;
 import javax.inject.Inject;
@@ -81,7 +81,7 @@ public final class ContactCreateFlow implements TransactionalFlow {
             .setAuthInfo(command.getAuthInfo())
             .setCreationClientId(clientId)
             .setPersistedCurrentSponsorClientId(clientId)
-            .setRepoId(createRepoId(ObjectifyService.allocateId(), roidSuffix))
+            .setRepoId(createRepoId(allocateId(), roidSuffix))
             .setFaxNumber(command.getFax())
             .setVoiceNumber(command.getVoice())
             .setDisclose(command.getDisclose())
@@ -93,7 +93,6 @@ public final class ContactCreateFlow implements TransactionalFlow {
     validateContactAgainstPolicy(newContact);
     historyBuilder
         .setType(HistoryEntry.Type.CONTACT_CREATE)
-        .setModificationTime(now)
         .setXmlBytes(null) // We don't want to store contact details in the history entry.
         .setContact(newContact);
     tm().insertAll(

core/src/main/java/google/registry/flows/contact/ContactDeleteFlow.java

@@ -122,11 +122,7 @@ public final class ContactDeleteFlow implements TransactionalFlow {
       resultCode = SUCCESS;
     }
     ContactHistory contactHistory =
-        historyBuilder
-            .setType(historyEntryType)
-            .setModificationTime(now)
-            .setContact(newContact)
-            .build();
+        historyBuilder.setType(historyEntryType).setContact(newContact).build();
     if (!tm().isOfy()) {
       handlePendingTransferOnDelete(existingContact, newContact, now, contactHistory);
     }

core/src/main/java/google/registry/flows/contact/ContactTransferApproveFlow.java

@@ -22,6 +22,7 @@ import static google.registry.flows.ResourceFlowUtils.verifyResourceOwnership;
 import static google.registry.flows.contact.ContactFlowUtils.createGainingTransferPollMessage;
 import static google.registry.flows.contact.ContactFlowUtils.createTransferResponse;
 import static google.registry.model.ResourceTransferUtils.approvePendingTransfer;
+import static google.registry.model.reporting.HistoryEntry.Type.CONTACT_TRANSFER_APPROVE;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.common.collect.ImmutableSet;
@@ -39,7 +40,6 @@ import google.registry.model.eppcommon.AuthInfo;
 import google.registry.model.eppinput.ResourceCommand;
 import google.registry.model.eppoutput.EppResponse;
 import google.registry.model.poll.PollMessage;
-import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.reporting.IcannReportingTypes.ActivityReportField;
 import google.registry.model.transfer.TransferStatus;
 import java.util.Optional;
@@ -88,11 +88,7 @@ public final class ContactTransferApproveFlow implements TransactionalFlow {
     ContactResource newContact =
         approvePendingTransfer(existingContact, TransferStatus.CLIENT_APPROVED, now);
     ContactHistory contactHistory =
-        historyBuilder
-            .setType(HistoryEntry.Type.CONTACT_TRANSFER_APPROVE)
-            .setModificationTime(now)
-            .setContact(newContact)
-            .build();
+        historyBuilder.setType(CONTACT_TRANSFER_APPROVE).setContact(newContact).build();
     // Create a poll message for the gaining client.
     PollMessage gainingPollMessage =
         createGainingTransferPollMessage(

core/src/main/java/google/registry/flows/contact/ContactTransferCancelFlow.java

@@ -22,6 +22,7 @@ import static google.registry.flows.ResourceFlowUtils.verifyTransferInitiator;
 import static google.registry.flows.contact.ContactFlowUtils.createLosingTransferPollMessage;
 import static google.registry.flows.contact.ContactFlowUtils.createTransferResponse;
 import static google.registry.model.ResourceTransferUtils.denyPendingTransfer;
+import static google.registry.model.reporting.HistoryEntry.Type.CONTACT_TRANSFER_CANCEL;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.common.collect.ImmutableSet;
@@ -39,7 +40,6 @@ import google.registry.model.eppcommon.AuthInfo;
 import google.registry.model.eppinput.ResourceCommand;
 import google.registry.model.eppoutput.EppResponse;
 import google.registry.model.poll.PollMessage;
-import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.reporting.IcannReportingTypes.ActivityReportField;
 import google.registry.model.transfer.TransferStatus;
 import java.util.Optional;
@@ -84,11 +84,7 @@ public final class ContactTransferCancelFlow implements TransactionalFlow {
     ContactResource newContact =
         denyPendingTransfer(existingContact, TransferStatus.CLIENT_CANCELLED, now, clientId);
     ContactHistory contactHistory =
-        historyBuilder
-            .setType(HistoryEntry.Type.CONTACT_TRANSFER_CANCEL)
-            .setModificationTime(now)
-            .setContact(newContact)
-            .build();
+        historyBuilder.setType(CONTACT_TRANSFER_CANCEL).setContact(newContact).build();
     // Create a poll message for the losing client.
     PollMessage losingPollMessage =
         createLosingTransferPollMessage(

core/src/main/java/google/registry/flows/contact/ContactTransferRejectFlow.java

@@ -22,6 +22,7 @@ import static google.registry.flows.ResourceFlowUtils.verifyResourceOwnership;
 import static google.registry.flows.contact.ContactFlowUtils.createGainingTransferPollMessage;
 import static google.registry.flows.contact.ContactFlowUtils.createTransferResponse;
 import static google.registry.model.ResourceTransferUtils.denyPendingTransfer;
+import static google.registry.model.reporting.HistoryEntry.Type.CONTACT_TRANSFER_REJECT;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.common.collect.ImmutableSet;
@@ -38,7 +39,6 @@ import google.registry.model.domain.metadata.MetadataExtension;
 import google.registry.model.eppcommon.AuthInfo;
 import google.registry.model.eppoutput.EppResponse;
 import google.registry.model.poll.PollMessage;
-import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.reporting.IcannReportingTypes.ActivityReportField;
 import google.registry.model.transfer.TransferStatus;
 import java.util.Optional;
@@ -82,11 +82,7 @@ public final class ContactTransferRejectFlow implements TransactionalFlow {
     ContactResource newContact =
         denyPendingTransfer(existingContact, TransferStatus.CLIENT_REJECTED, now, clientId);
     ContactHistory contactHistory =
-        historyBuilder
-            .setType(HistoryEntry.Type.CONTACT_TRANSFER_REJECT)
-            .setModificationTime(now)
-            .setContact(newContact)
-            .build();
+        historyBuilder.setType(CONTACT_TRANSFER_REJECT).setContact(newContact).build();
     PollMessage gainingPollMessage =
         createGainingTransferPollMessage(
             targetId, newContact.getTransferData(), now, Key.create(contactHistory));

core/src/main/java/google/registry/flows/contact/ContactTransferRequestFlow.java

@@ -24,6 +24,7 @@ import static google.registry.flows.contact.ContactFlowUtils.createGainingTransf
 import static google.registry.flows.contact.ContactFlowUtils.createLosingTransferPollMessage;
 import static google.registry.flows.contact.ContactFlowUtils.createTransferResponse;
 import static google.registry.model.eppoutput.Result.Code.SUCCESS_WITH_ACTION_PENDING;
+import static google.registry.model.reporting.HistoryEntry.Type.CONTACT_TRANSFER_REQUEST;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.common.collect.ImmutableSet;
@@ -45,7 +46,6 @@ import google.registry.model.eppcommon.StatusValue;
 import google.registry.model.eppcommon.Trid;
 import google.registry.model.eppoutput.EppResponse;
 import google.registry.model.poll.PollMessage;
-import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.reporting.IcannReportingTypes.ActivityReportField;
 import google.registry.model.transfer.ContactTransferData;
 import google.registry.model.transfer.TransferStatus;
@@ -120,10 +120,7 @@ public final class ContactTransferRequestFlow implements TransactionalFlow {
             .setTransferStatus(TransferStatus.SERVER_APPROVED)
             .build();
     Key<ContactHistory> contactHistoryKey = createHistoryKey(existingContact, ContactHistory.class);
-    historyBuilder
-        .setId(contactHistoryKey.getId())
-        .setType(HistoryEntry.Type.CONTACT_TRANSFER_REQUEST)
-        .setModificationTime(now);
+    historyBuilder.setId(contactHistoryKey.getId()).setType(CONTACT_TRANSFER_REQUEST);
     // If the transfer is server approved, this message will be sent to the losing registrar. */
     PollMessage serverApproveLosingPollMessage =
         createLosingTransferPollMessage(targetId, serverApproveTransferData, contactHistoryKey);

core/src/main/java/google/registry/flows/contact/ContactUpdateFlow.java

@@ -24,6 +24,7 @@ import static google.registry.flows.ResourceFlowUtils.verifyOptionalAuthInfo;
 import static google.registry.flows.ResourceFlowUtils.verifyResourceOwnership;
 import static google.registry.flows.contact.ContactFlowUtils.validateAsciiPostalInfo;
 import static google.registry.flows.contact.ContactFlowUtils.validateContactAgainstPolicy;
+import static google.registry.model.reporting.HistoryEntry.Type.CONTACT_UPDATE;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.common.collect.ImmutableSet;
@@ -45,7 +46,6 @@ import google.registry.model.eppcommon.AuthInfo;
 import google.registry.model.eppcommon.StatusValue;
 import google.registry.model.eppinput.ResourceCommand;
 import google.registry.model.eppoutput.EppResponse;
-import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.reporting.IcannReportingTypes.ActivityReportField;
 import java.util.Optional;
 import javax.annotation.Nullable;
@@ -146,8 +146,7 @@ public final class ContactUpdateFlow implements TransactionalFlow {
     validateAsciiPostalInfo(newContact.getInternationalizedPostalInfo());
     validateContactAgainstPolicy(newContact);
     historyBuilder
-        .setType(HistoryEntry.Type.CONTACT_UPDATE)
-        .setModificationTime(now)
+        .setType(CONTACT_UPDATE)
         .setXmlBytes(null) // We don't want to store contact details in the history entry.
         .setContact(newContact);
     tm().insert(historyBuilder.build());

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

@@ -42,11 +42,13 @@ import static google.registry.flows.domain.DomainFlowUtils.verifyPremiumNameIsNo
 import static google.registry.flows.domain.DomainFlowUtils.verifyRegistrarIsActive;
 import static google.registry.flows.domain.DomainFlowUtils.verifyUnitIsYears;
 import static google.registry.model.EppResourceUtils.createDomainRepoId;
+import static google.registry.model.IdService.allocateId;
 import static google.registry.model.eppcommon.StatusValue.SERVER_HOLD;
 import static google.registry.model.registry.Registry.TldState.GENERAL_AVAILABILITY;
 import static google.registry.model.registry.Registry.TldState.QUIET_PERIOD;
 import static google.registry.model.registry.Registry.TldState.START_DATE_SUNRISE;
 import static google.registry.model.registry.label.ReservationType.NAME_COLLISION;
+import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_CREATE;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.DateTimeUtils.END_OF_TIME;
 import static google.registry.util.DateTimeUtils.leapSafeAddYears;
@@ -99,7 +101,6 @@ import google.registry.model.eppoutput.CreateData.DomainCreateData;
 import google.registry.model.eppoutput.EppResponse;
 import google.registry.model.index.EppResourceIndex;
 import google.registry.model.index.ForeignKeyIndex;
-import google.registry.model.ofy.ObjectifyService;
 import google.registry.model.poll.PendingActionNotificationResponse.DomainPendingActionNotificationResponse;
 import google.registry.model.poll.PollMessage;
 import google.registry.model.poll.PollMessage.Autorenew;
@@ -302,12 +303,9 @@ public class DomainCreateFlow implements TransactionalFlow {
     Optional<SecDnsCreateExtension> secDnsCreate =
         validateSecDnsExtension(eppInput.getSingleExtension(SecDnsCreateExtension.class));
     DateTime registrationExpirationTime = leapSafeAddYears(now, years);
-    String repoId = createDomainRepoId(ObjectifyService.allocateId(), registry.getTldStr());
+    String repoId = createDomainRepoId(allocateId(), registry.getTldStr());
     Key<DomainHistory> domainHistoryKey =
-        Key.create(
-            Key.create(DomainBase.class, repoId),
-            DomainHistory.class,
-            ObjectifyService.allocateId());
+        Key.create(Key.create(DomainBase.class, repoId), DomainHistory.class, allocateId());
     historyBuilder.setId(domainHistoryKey.getId());
     // Bill for the create.
     BillingEvent.OneTime createBillingEvent =
@@ -498,12 +496,7 @@ public class DomainCreateFlow implements TransactionalFlow {
                       TransactionReportField.netAddsFieldFromYears(period.getValue()),
                       1)));
     }
-    return historyBuilder
-        .setType(HistoryEntry.Type.DOMAIN_CREATE)
-        .setPeriod(period)
-        .setModificationTime(now)
-        .setDomain(domain)
-        .build();
+    return historyBuilder.setType(DOMAIN_CREATE).setPeriod(period).setDomain(domain).build();
   }
 
   private BillingEvent.OneTime createOneTimeBillingEvent(

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

@@ -34,6 +34,7 @@ import static google.registry.model.eppoutput.Result.Code.SUCCESS;
 import static google.registry.model.eppoutput.Result.Code.SUCCESS_WITH_ACTION_PENDING;
 import static google.registry.model.reporting.DomainTransactionRecord.TransactionReportField.ADD_FIELDS;
 import static google.registry.model.reporting.DomainTransactionRecord.TransactionReportField.RENEW_FIELDS;
+import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_DELETE;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.pricing.PricingEngineProxy.getDomainRenewCost;
 import static google.registry.util.CollectionUtils.nullToEmpty;
@@ -88,7 +89,6 @@ import google.registry.model.registry.Registry;
 import google.registry.model.registry.Registry.TldType;
 import google.registry.model.reporting.DomainTransactionRecord;
 import google.registry.model.reporting.DomainTransactionRecord.TransactionReportField;
-import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.reporting.IcannReportingTypes.ActivityReportField;
 import google.registry.model.transfer.TransferStatus;
 import java.util.Collections;
@@ -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
@@ -328,11 +331,7 @@ public final class DomainDeleteFlow implements TransactionalFlow {
                       : TransactionReportField.DELETED_DOMAINS_NOGRACE,
                   1)));
     }
-    return historyBuilder
-        .setType(HistoryEntry.Type.DOMAIN_DELETE)
-        .setModificationTime(now)
-        .setDomain(domain)
-        .build();
+    return historyBuilder.setType(DOMAIN_DELETE).setDomain(domain).build();
   }
 
   private PollMessage.OneTime createDeletePollMessage(
@@ -377,7 +376,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 +384,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 +397,12 @@ public final class DomainDeleteFlow implements TransactionalFlow {
     return ImmutableList.of(feeResponseBuilder.setCredits(credits).build());
   }
 
-  private Money getGracePeriodCost(GracePeriod gracePeriod, DateTime now) {
+  private Money getGracePeriodCost(
+      BillingEvent.Recurring recurringBillingEvent, GracePeriod gracePeriod, DateTime now) {
     if (gracePeriod.getType() == GracePeriodStatus.AUTO_RENEW) {
+      // If we updated the autorenew billing event, reuse it.
       DateTime autoRenewTime =
-          tm().loadByKey(checkNotNull(gracePeriod.getRecurringBillingEvent()))
-              .getRecurrenceTimeOfYear()
-              .getLastInstanceBeforeOrAt(now);
+          recurringBillingEvent.getRecurrenceTimeOfYear().getLastInstanceBeforeOrAt(now);
       return getDomainRenewCost(targetId, autoRenewTime, 1);
     }
     return tm().loadByKey(checkNotNull(gracePeriod.getOneTimeBillingEvent())).getCost();

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

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

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

@@ -29,6 +29,7 @@ import static google.registry.flows.domain.DomainFlowUtils.validateFeeChallenge;
 import static google.registry.flows.domain.DomainFlowUtils.validateRegistrationPeriod;
 import static google.registry.flows.domain.DomainFlowUtils.verifyRegistrarIsActive;
 import static google.registry.flows.domain.DomainFlowUtils.verifyUnitIsYears;
+import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_RENEW;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.DateTimeUtils.leapSafeAddYears;
 
@@ -73,7 +74,6 @@ import google.registry.model.poll.PollMessage;
 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.model.reporting.IcannReportingTypes.ActivityReportField;
 import java.util.Optional;
 import javax.inject.Inject;
@@ -230,9 +230,8 @@ public final class DomainRenewFlow implements TransactionalFlow {
   private DomainHistory buildDomainHistory(
       DomainBase newDomain, DateTime now, Period period, Duration renewGracePeriod) {
     return historyBuilder
-        .setType(HistoryEntry.Type.DOMAIN_RENEW)
+        .setType(DOMAIN_RENEW)
         .setPeriod(period)
-        .setModificationTime(now)
         .setDomain(newDomain)
         .setDomainTransactionRecords(
             ImmutableSet.of(

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

@@ -27,6 +27,7 @@ import static google.registry.flows.domain.DomainFlowUtils.verifyNotReserved;
 import static google.registry.flows.domain.DomainFlowUtils.verifyPremiumNameIsNotBlocked;
 import static google.registry.flows.domain.DomainFlowUtils.verifyRegistrarIsActive;
 import static google.registry.model.ResourceTransferUtils.updateForeignKeyIndexDeletionTime;
+import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_RESTORE;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.DateTimeUtils.END_OF_TIME;
 
@@ -66,7 +67,6 @@ import google.registry.model.poll.PollMessage;
 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.model.reporting.IcannReportingTypes.ActivityReportField;
 import java.util.Optional;
 import javax.inject.Inject;
@@ -188,8 +188,7 @@ public final class DomainRestoreRequestFlow implements TransactionalFlow  {
 
   private DomainHistory buildDomainHistory(DomainBase newDomain, DateTime now) {
     return historyBuilder
-        .setType(HistoryEntry.Type.DOMAIN_RESTORE)
-        .setModificationTime(now)
+        .setType(DOMAIN_RESTORE)
         .setDomain(newDomain)
         .setDomainTransactionRecords(
             ImmutableSet.of(

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

@@ -29,6 +29,7 @@ import static google.registry.flows.domain.DomainTransferUtils.createGainingTran
 import static google.registry.flows.domain.DomainTransferUtils.createTransferResponse;
 import static google.registry.model.ResourceTransferUtils.approvePendingTransfer;
 import static google.registry.model.reporting.DomainTransactionRecord.TransactionReportField.TRANSFER_SUCCESSFUL;
+import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_TRANSFER_APPROVE;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.pricing.PricingEngineProxy.getDomainRenewCost;
 import static google.registry.util.CollectionUtils.union;
@@ -58,7 +59,6 @@ import google.registry.model.eppoutput.EppResponse;
 import google.registry.model.poll.PollMessage;
 import google.registry.model.registry.Registry;
 import google.registry.model.reporting.DomainTransactionRecord;
-import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.reporting.IcannReportingTypes.ActivityReportField;
 import google.registry.model.transfer.DomainTransferData;
 import google.registry.model.transfer.TransferStatus;
@@ -240,8 +240,7 @@ public final class DomainTransferApproveFlow implements TransactionalFlow {
             registry.getAutomaticTransferLength().plus(registry.getTransferGracePeriodLength()),
             ImmutableSet.of(TRANSFER_SUCCESSFUL));
     return historyBuilder
-        .setType(HistoryEntry.Type.DOMAIN_TRANSFER_APPROVE)
-        .setModificationTime(now)
+        .setType(DOMAIN_TRANSFER_APPROVE)
         .setOtherClientId(gainingClientId)
         .setDomain(newDomain)
         .setDomainTransactionRecords(

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

@@ -27,6 +27,7 @@ import static google.registry.flows.domain.DomainTransferUtils.createLosingTrans
 import static google.registry.flows.domain.DomainTransferUtils.createTransferResponse;
 import static google.registry.model.ResourceTransferUtils.denyPendingTransfer;
 import static google.registry.model.reporting.DomainTransactionRecord.TransactionReportField.TRANSFER_SUCCESSFUL;
+import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_TRANSFER_CANCEL;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.DateTimeUtils.END_OF_TIME;
 
@@ -46,7 +47,6 @@ import google.registry.model.eppcommon.AuthInfo;
 import google.registry.model.eppoutput.EppResponse;
 import google.registry.model.registry.Registry;
 import google.registry.model.reporting.DomainTransactionRecord;
-import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.reporting.IcannReportingTypes.ActivityReportField;
 import google.registry.model.transfer.TransferStatus;
 import java.util.Optional;
@@ -130,8 +130,7 @@ public final class DomainTransferCancelFlow implements TransactionalFlow {
             registry.getAutomaticTransferLength().plus(registry.getTransferGracePeriodLength()),
             ImmutableSet.of(TRANSFER_SUCCESSFUL));
     return historyBuilder
-        .setType(HistoryEntry.Type.DOMAIN_TRANSFER_CANCEL)
-        .setModificationTime(now)
+        .setType(DOMAIN_TRANSFER_CANCEL)
         .setDomain(newDomain)
         .setDomainTransactionRecords(cancelingRecords)
         .build();

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

@@ -28,6 +28,7 @@ import static google.registry.flows.domain.DomainTransferUtils.createTransferRes
 import static google.registry.model.ResourceTransferUtils.denyPendingTransfer;
 import static google.registry.model.reporting.DomainTransactionRecord.TransactionReportField.TRANSFER_NACKED;
 import static google.registry.model.reporting.DomainTransactionRecord.TransactionReportField.TRANSFER_SUCCESSFUL;
+import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_TRANSFER_REJECT;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.CollectionUtils.union;
 import static google.registry.util.DateTimeUtils.END_OF_TIME;
@@ -48,7 +49,6 @@ import google.registry.model.eppcommon.AuthInfo;
 import google.registry.model.eppoutput.EppResponse;
 import google.registry.model.registry.Registry;
 import google.registry.model.reporting.DomainTransactionRecord;
-import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.reporting.IcannReportingTypes.ActivityReportField;
 import google.registry.model.transfer.TransferStatus;
 import java.util.Optional;
@@ -131,8 +131,7 @@ public final class DomainTransferRejectFlow implements TransactionalFlow {
             registry.getAutomaticTransferLength().plus(registry.getTransferGracePeriodLength()),
             ImmutableSet.of(TRANSFER_SUCCESSFUL));
     return historyBuilder
-        .setType(HistoryEntry.Type.DOMAIN_TRANSFER_REJECT)
-        .setModificationTime(now)
+        .setType(DOMAIN_TRANSFER_REJECT)
         .setDomainTransactionRecords(
             union(
                 cancelingRecords,

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

@@ -32,6 +32,7 @@ import static google.registry.flows.domain.DomainTransferUtils.createPendingTran
 import static google.registry.flows.domain.DomainTransferUtils.createTransferResponse;
 import static google.registry.flows.domain.DomainTransferUtils.createTransferServerApproveEntities;
 import static google.registry.model.eppoutput.Result.Code.SUCCESS_WITH_ACTION_PENDING;
+import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_TRANSFER_REQUEST;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.common.collect.ImmutableList;
@@ -68,7 +69,6 @@ import google.registry.model.poll.PollMessage;
 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.model.reporting.IcannReportingTypes.ActivityReportField;
 import google.registry.model.transfer.DomainTransferData;
 import google.registry.model.transfer.TransferData.TransferServerApproveEntity;
@@ -315,9 +315,8 @@ public final class DomainTransferRequestFlow implements TransactionalFlow {
   private DomainHistory buildDomainHistory(
       DomainBase newDomain, Registry registry, DateTime now, Period period) {
     return historyBuilder
-        .setType(HistoryEntry.Type.DOMAIN_TRANSFER_REQUEST)
+        .setType(DOMAIN_TRANSFER_REQUEST)
         .setPeriod(period)
-        .setModificationTime(now)
         .setDomain(newDomain)
         .setDomainTransactionRecords(
             ImmutableSet.of(

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

@@ -39,6 +39,7 @@ import static google.registry.flows.domain.DomainFlowUtils.validateRegistrantAll
 import static google.registry.flows.domain.DomainFlowUtils.validateRequiredContactsPresent;
 import static google.registry.flows.domain.DomainFlowUtils.verifyClientUpdateNotProhibited;
 import static google.registry.flows.domain.DomainFlowUtils.verifyNotInPendingDelete;
+import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_UPDATE;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.common.collect.ImmutableSet;
@@ -76,7 +77,6 @@ import google.registry.model.eppinput.EppInput;
 import google.registry.model.eppinput.ResourceCommand;
 import google.registry.model.eppoutput.EppResponse;
 import google.registry.model.registry.Registry;
-import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.reporting.IcannReportingTypes.ActivityReportField;
 import java.util.Optional;
 import javax.inject.Inject;
@@ -166,7 +166,8 @@ public final class DomainUpdateFlow implements TransactionalFlow {
     flowCustomLogic.afterValidation(
         AfterValidationParameters.newBuilder().setExistingDomain(existingDomain).build());
     DomainBase newDomain = performUpdate(command, existingDomain, now);
-    DomainHistory domainHistory = buildDomainHistory(newDomain, now);
+    DomainHistory domainHistory =
+        historyBuilder.setType(DOMAIN_UPDATE).setDomain(newDomain).build();
     validateNewState(newDomain);
     dnsQueue.addDomainRefreshTask(targetId);
     ImmutableSet.Builder<ImmutableObject> entitiesToSave = new ImmutableSet.Builder<>();
@@ -217,14 +218,6 @@ public final class DomainUpdateFlow implements TransactionalFlow {
         tld, add.getNameserverFullyQualifiedHostNames());
   }
 
-  private DomainHistory buildDomainHistory(DomainBase newDomain, DateTime now) {
-    return historyBuilder
-        .setType(HistoryEntry.Type.DOMAIN_UPDATE)
-        .setModificationTime(now)
-        .setDomain(newDomain)
-        .build();
-  }
-
   private DomainBase performUpdate(Update command, DomainBase domain, DateTime now)
       throws EppException {
     AddRemove add = command.getInnerAdd();
@@ -290,10 +283,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/flows/host/HostCreateFlow.java

@@ -21,6 +21,8 @@ import static google.registry.flows.host.HostFlowUtils.validateHostName;
 import static google.registry.flows.host.HostFlowUtils.verifySuperordinateDomainNotInPendingDelete;
 import static google.registry.flows.host.HostFlowUtils.verifySuperordinateDomainOwnership;
 import static google.registry.model.EppResourceUtils.createRepoId;
+import static google.registry.model.IdService.allocateId;
+import static google.registry.model.reporting.HistoryEntry.Type.HOST_CREATE;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.CollectionUtils.isNullOrEmpty;
 
@@ -49,8 +51,6 @@ import google.registry.model.host.HostHistory;
 import google.registry.model.host.HostResource;
 import google.registry.model.index.EppResourceIndex;
 import google.registry.model.index.ForeignKeyIndex;
-import google.registry.model.ofy.ObjectifyService;
-import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.reporting.IcannReportingTypes.ActivityReportField;
 import java.util.Optional;
 import javax.inject.Inject;
@@ -126,10 +126,10 @@ public final class HostCreateFlow implements TransactionalFlow {
             .setPersistedCurrentSponsorClientId(clientId)
             .setHostName(targetId)
             .setInetAddresses(command.getInetAddresses())
-            .setRepoId(createRepoId(ObjectifyService.allocateId(), roidSuffix))
+            .setRepoId(createRepoId(allocateId(), roidSuffix))
             .setSuperordinateDomain(superordinateDomain.map(DomainBase::createVKey).orElse(null))
             .build();
-    historyBuilder.setType(HistoryEntry.Type.HOST_CREATE).setModificationTime(now).setHost(newHost);
+    historyBuilder.setType(HOST_CREATE).setHost(newHost);
     ImmutableSet<ImmutableObject> entitiesToSave =
         ImmutableSet.of(
             newHost,

core/src/main/java/google/registry/flows/host/HostDeleteFlow.java

@@ -130,7 +130,7 @@ public final class HostDeleteFlow implements TransactionalFlow {
       historyEntryType = Type.HOST_DELETE;
       resultCode = SUCCESS;
     }
-    historyBuilder.setType(historyEntryType).setModificationTime(now).setHost(newHost);
+    historyBuilder.setType(historyEntryType).setHost(newHost);
     tm().insert(historyBuilder.build());
     tm().update(newHost);
     return responseBuilder.setResultFromCode(resultCode).build();

core/src/main/java/google/registry/flows/host/HostUpdateFlow.java

@@ -27,6 +27,7 @@ import static google.registry.flows.host.HostFlowUtils.validateHostName;
 import static google.registry.flows.host.HostFlowUtils.verifySuperordinateDomainNotInPendingDelete;
 import static google.registry.flows.host.HostFlowUtils.verifySuperordinateDomainOwnership;
 import static google.registry.model.index.ForeignKeyIndex.loadAndGetKey;
+import static google.registry.model.reporting.HistoryEntry.Type.HOST_UPDATE;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.CollectionUtils.isNullOrEmpty;
 
@@ -57,7 +58,6 @@ import google.registry.model.host.HostCommand.Update.Change;
 import google.registry.model.host.HostHistory;
 import google.registry.model.host.HostResource;
 import google.registry.model.index.ForeignKeyIndex;
-import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.reporting.IcannReportingTypes.ActivityReportField;
 import google.registry.persistence.VKey;
 import java.util.Objects;
@@ -201,12 +201,7 @@ public final class HostUpdateFlow implements TransactionalFlow {
       updateSuperordinateDomains(existingHost, newHost);
     }
     enqueueTasks(existingHost, newHost);
-    entitiesToInsert.add(
-        historyBuilder
-            .setType(HistoryEntry.Type.HOST_UPDATE)
-            .setModificationTime(now)
-            .setHost(newHost)
-            .build());
+    entitiesToInsert.add(historyBuilder.setType(HOST_UPDATE).setHost(newHost).build());
     tm().updateAll(entitiesToUpdate.build());
     tm().insertAll(entitiesToInsert.build());
     return responseBuilder.build();

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,9 +16,10 @@ 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.IdService.allocateId;
+import static google.registry.model.ModelUtils.getAllFields;
 
-import google.registry.model.ofy.ObjectifyService;
+import com.googlecode.objectify.annotation.Id;
 import google.registry.util.TypeUtils.TypeInstantiator;
 import java.lang.reflect.Field;
 import java.util.Optional;
@@ -54,19 +55,18 @@ public interface Buildable {
     /** Build the instance. */
     public S build() {
       try {
-        // 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());
-        } catch (Exception e) {
-          // Expected if the class is not registered with Objectify.
-        }
+        // If this object has a Long or long Objectify @Id field that is not set, set it now. For
+        // any entity it has one and only one @Id field in its class hierarchy.
+        Field idField =
+            getAllFields(instance.getClass()).values().stream()
+                .filter(field -> field.isAnnotationPresent(Id.class))
+                .findFirst()
+                .orElse(null);
         if (idField != null
             && !idField.getType().equals(String.class)
             && Optional.ofNullable((Long) ModelUtils.getFieldValue(instance, idField))
                 .orElse(0L) == 0) {
-          ModelUtils.setFieldValue(instance, idField, ObjectifyService.allocateId());
+          ModelUtils.setFieldValue(instance, idField, allocateId());
         }
         return instance;
       } finally {

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

@@ -1,55 +0,0 @@
-// Copyright 2020 The Nomulus Authors. All Rights Reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package google.registry.model;
-
-import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
-
-import com.google.common.flogger.FluentLogger;
-import google.registry.config.RegistryEnvironment;
-import google.registry.model.common.DatabaseTransitionSchedule;
-import google.registry.model.common.DatabaseTransitionSchedule.PrimaryDatabase;
-import google.registry.model.common.DatabaseTransitionSchedule.TransitionId;
-
-/** Utility methods related to migrating dual-read/dual-write entities. */
-public class DatabaseMigrationUtils {
-
-  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
-
-  /** Throws exceptions only in unit tests, otherwise only logs exceptions. */
-  public static void suppressExceptionUnlessInTest(Runnable work, String message) {
-    try {
-      work.run();
-    } catch (Exception e) {
-      if (RegistryEnvironment.get().equals(RegistryEnvironment.UNITTEST)) {
-        throw e;
-      }
-      logger.atWarning().withCause(e).log(message);
-    }
-  }
-
-  /** Gets the value for the database currently considered primary. */
-  public static PrimaryDatabase getPrimaryDatabase(TransitionId transitionId) {
-    return DatabaseTransitionSchedule.getCached(transitionId)
-        .map(DatabaseTransitionSchedule::getPrimaryDatabase)
-        .orElse(PrimaryDatabase.DATASTORE);
-  }
-
-  public static boolean isDatastore(TransitionId transitionId) {
-    return tm().transactNew(() -> DatabaseMigrationUtils.getPrimaryDatabase(transitionId))
-        .equals(PrimaryDatabase.DATASTORE);
-  }
-
-  private DatabaseMigrationUtils() {}
-}

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

@@ -17,8 +17,7 @@ package google.registry.model;
 import com.google.common.collect.ImmutableSet;
 import google.registry.model.billing.BillingEvent;
 import google.registry.model.common.Cursor;
-import google.registry.model.common.DatabaseMigrationStateWrapper;
-import google.registry.model.common.DatabaseTransitionSchedule;
+import google.registry.model.common.DatabaseMigrationStateSchedule;
 import google.registry.model.common.EntityGroupRoot;
 import google.registry.model.common.GaeUserIdConverter;
 import google.registry.model.contact.ContactHistory;
@@ -76,8 +75,7 @@ public final class EntityClasses {
           ContactHistory.class,
           ContactResource.class,
           Cursor.class,
-          DatabaseMigrationStateWrapper.class,
-          DatabaseTransitionSchedule.class,
+          DatabaseMigrationStateSchedule.class,
           DomainBase.class,
           DomainHistory.class,
           EntityGroupRoot.class,

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

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

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

@@ -16,10 +16,11 @@ package google.registry.model;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
-import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.model.ofy.ObjectifyService.auditedOfy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
+import static google.registry.util.DateTimeUtils.START_OF_TIME;
 import static google.registry.util.DateTimeUtils.isAtOrAfter;
 import static google.registry.util.DateTimeUtils.isBeforeOrAt;
 import static google.registry.util.DateTimeUtils.latestOf;
@@ -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/IdService.java

@@ -0,0 +1,61 @@
+// 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;
+
+import static com.google.common.base.Preconditions.checkState;
+
+import com.google.appengine.api.datastore.DatastoreServiceFactory;
+import com.google.common.annotations.VisibleForTesting;
+import google.registry.config.RegistryEnvironment;
+import java.util.concurrent.atomic.AtomicLong;
+
+/**
+ * Allocates a globally unique {@link Long} number to use as an Ofy {@code @Id}.
+ *
+ * <p>In non-test environments the Id is generated by Datastore, whereas in tests it's from an
+ * atomic long number that's incremented every time this method is called.
+ */
+public final class IdService {
+
+  /**
+   * A placeholder String passed into DatastoreService.allocateIds that ensures that all ids are
+   * initialized from the same id pool.
+   */
+  private static final String APP_WIDE_ALLOCATION_KIND = "common";
+
+  /** Counts of used ids for use in unit tests. Outside tests this is never used. */
+  private static final AtomicLong nextTestId = new AtomicLong(1); // ids cannot be zero
+
+  /** Allocates an id. */
+  public static long allocateId() {
+    return RegistryEnvironment.UNITTEST.equals(RegistryEnvironment.get())
+        ? nextTestId.getAndIncrement()
+        : DatastoreServiceFactory.getDatastoreService()
+            .allocateIds(APP_WIDE_ALLOCATION_KIND, 1)
+            .iterator()
+            .next()
+            .getId();
+  }
+
+  /** Resets the global test id counter (i.e. sets the next id to 1). */
+  @VisibleForTesting
+  public static void resetNextTestId() {
+    checkState(
+        RegistryEnvironment.UNITTEST.equals(RegistryEnvironment.get()),
+        "Can't call resetTestIdCounts() from RegistryEnvironment.%s",
+        RegistryEnvironment.get());
+    nextTestId.set(1); // ids cannot be zero
+  }
+}

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

@@ -1,122 +0,0 @@
-// Copyright 2021 The Nomulus Authors. All Rights Reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package google.registry.model.common;
-
-import static com.google.common.base.Preconditions.checkArgument;
-import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
-
-import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.ImmutableSet;
-import com.googlecode.objectify.annotation.Entity;
-import google.registry.model.annotations.InCrossTld;
-import google.registry.model.common.DatabaseTransitionSchedule.PrimaryDatabase;
-import google.registry.schema.replay.DatastoreOnlyEntity;
-
-/**
- * A wrapper object representing the current stage of the Registry 3.0 Cloud SQL migration.
- *
- * <p>The entity is stored in Datastore throughout the entire migration so as to have a single point
- * of access (avoiding a two-phase commit problem).
- */
-@Entity
-@InCrossTld
-public class DatabaseMigrationStateWrapper extends CrossTldSingleton
-    implements DatastoreOnlyEntity {
-
-  /**
-   * The current phase of the migration plus information about which database to use and whether or
-   * not the phase is read-only.
-   */
-  public enum MigrationState {
-    DATASTORE_ONLY(PrimaryDatabase.DATASTORE, false),
-    DATASTORE_PRIMARY(PrimaryDatabase.DATASTORE, false),
-    DATASTORE_PRIMARY_READ_ONLY(PrimaryDatabase.DATASTORE, true),
-    SQL_PRIMARY(PrimaryDatabase.CLOUD_SQL, false),
-    SQL_ONLY(PrimaryDatabase.CLOUD_SQL, false);
-
-    private final PrimaryDatabase primaryDatabase;
-    private final boolean readOnly;
-
-    public PrimaryDatabase getPrimaryDatabase() {
-      return primaryDatabase;
-    }
-
-    public boolean isReadOnly() {
-      return readOnly;
-    }
-
-    MigrationState(PrimaryDatabase primaryDatabase, boolean readOnly) {
-      this.primaryDatabase = primaryDatabase;
-      this.readOnly = readOnly;
-    }
-  }
-
-  private MigrationState migrationState;
-
-  // Required for Objectify initialization
-  private DatabaseMigrationStateWrapper() {}
-
-  DatabaseMigrationStateWrapper(MigrationState migrationState) {
-    this.migrationState = migrationState;
-  }
-
-  // The valid state transitions. Basically, at state N, state N+1 is valid as well as all previous
-  // states, with one type of exception: when in either of the SQL states, we can only move back
-  // one step so that we can make sure that any modifications have been replayed back to Datastore.
-  private static final ImmutableMap<MigrationState, ImmutableSet<MigrationState>>
-      VALID_STATE_TRANSITIONS =
-          ImmutableMap.of(
-              MigrationState.DATASTORE_ONLY,
-              ImmutableSet.of(MigrationState.DATASTORE_PRIMARY),
-              MigrationState.DATASTORE_PRIMARY,
-              ImmutableSet.of(
-                  MigrationState.DATASTORE_ONLY, MigrationState.DATASTORE_PRIMARY_READ_ONLY),
-              MigrationState.DATASTORE_PRIMARY_READ_ONLY,
-              ImmutableSet.of(
-                  MigrationState.DATASTORE_ONLY,
-                  MigrationState.DATASTORE_PRIMARY,
-                  MigrationState.SQL_PRIMARY),
-              MigrationState.SQL_PRIMARY,
-              ImmutableSet.of(MigrationState.DATASTORE_PRIMARY_READ_ONLY, MigrationState.SQL_ONLY),
-              MigrationState.SQL_ONLY,
-              ImmutableSet.of(MigrationState.SQL_PRIMARY));
-
-  private static boolean isValidStateTransition(MigrationState from, MigrationState to) {
-    return VALID_STATE_TRANSITIONS.get(from).contains(to);
-  }
-
-  /** Sets and persists to Datastore the current database migration state. */
-  public static void set(MigrationState newState) {
-    MigrationState currentState = get();
-    checkArgument(
-        isValidStateTransition(currentState, newState),
-        "Moving from migration state %s to %s is not a valid transition",
-        currentState,
-        newState);
-    DatabaseMigrationStateWrapper wrapper = new DatabaseMigrationStateWrapper(newState);
-    ofyTm().transact(() -> ofyTm().put(wrapper));
-  }
-
-  /** Retrieves the current state of the migration (or DATASTORE_ONLY if it hasn't started). */
-  public static MigrationState get() {
-    return ofyTm()
-        .transact(
-            () ->
-                ofyTm()
-                    .loadSingleton(DatabaseMigrationStateWrapper.class)
-                    .map(s -> s.migrationState)
-                    .orElse(MigrationState.DATASTORE_ONLY));
-  }
-}

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

@@ -1,166 +0,0 @@
-// Copyright 2021 The Nomulus Authors. All Rights Reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package google.registry.model.common;
-
-import static com.google.common.base.Preconditions.checkNotNull;
-import static google.registry.config.RegistryConfig.getSingletonCacheRefreshDuration;
-import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
-import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
-import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
-
-import com.google.common.cache.CacheBuilder;
-import com.google.common.cache.CacheLoader;
-import com.google.common.cache.LoadingCache;
-import com.google.common.collect.ImmutableSortedMap;
-import com.googlecode.objectify.Key;
-import com.googlecode.objectify.annotation.Embed;
-import com.googlecode.objectify.annotation.Entity;
-import com.googlecode.objectify.annotation.Id;
-import com.googlecode.objectify.annotation.Mapify;
-import com.googlecode.objectify.annotation.Parent;
-import google.registry.model.ImmutableObject;
-import google.registry.model.UpdateAutoTimestamp;
-import google.registry.model.annotations.InCrossTld;
-import google.registry.model.common.TimedTransitionProperty.TimeMapper;
-import google.registry.model.common.TimedTransitionProperty.TimedTransition;
-import google.registry.model.registry.label.PremiumList;
-import google.registry.model.registry.label.ReservedList;
-import google.registry.model.smd.SignedMarkRevocationList;
-import google.registry.model.tmch.ClaimsList;
-import google.registry.persistence.VKey;
-import google.registry.schema.replay.DatastoreOnlyEntity;
-import java.util.Optional;
-import javax.annotation.concurrent.Immutable;
-import org.joda.time.DateTime;
-
-@Entity
-@Immutable
-@InCrossTld
-public class DatabaseTransitionSchedule extends ImmutableObject implements DatastoreOnlyEntity {
-
-  /**
-   * The name of the database to be treated as the primary database. The first entry in the schedule
-   * will always be Datastore.
-   */
-  public enum PrimaryDatabase {
-    CLOUD_SQL,
-    DATASTORE
-  }
-
-  /** The id of the transition schedule. */
-  public enum TransitionId {
-    /** The schedule for migration of {@link ClaimsList} entities. */
-    CLAIMS_LIST,
-    /** The schedule for the migration of {@link PremiumList} and {@link ReservedList}. */
-    DOMAIN_LABEL_LISTS,
-    /** The schedule for the migration of the {@link SignedMarkRevocationList} entity. */
-    SIGNED_MARK_REVOCATION_LIST,
-    /** The schedule for all asynchronously-replayed entities, ones not dually-written. */
-    REPLAYED_ENTITIES,
-  }
-
-  /**
-   * The transition to a specified primary database at a specific point in time, for use in a
-   * TimedTransitionProperty.
-   */
-  @Embed
-  public static class PrimaryDatabaseTransition extends TimedTransition<PrimaryDatabase> {
-    private PrimaryDatabase primaryDatabase;
-
-    @Override
-    protected PrimaryDatabase getValue() {
-      return primaryDatabase;
-    }
-
-    @Override
-    protected void setValue(PrimaryDatabase primaryDatabase) {
-      this.primaryDatabase = primaryDatabase;
-    }
-  }
-
-  @Parent Key<EntityGroupRoot> parent = getCrossTldKey();
-
-  @Id String transitionId;
-
-  /** An automatically managed timestamp of when this schedule was last written to Datastore. */
-  UpdateAutoTimestamp lastUpdateTime = UpdateAutoTimestamp.create(null);
-
-  /** A property that tracks the primary database for a dual-read/dual-write database migration. */
-  @Mapify(TimeMapper.class)
-  TimedTransitionProperty<PrimaryDatabase, PrimaryDatabaseTransition> databaseTransitions =
-      TimedTransitionProperty.forMapify(PrimaryDatabase.DATASTORE, PrimaryDatabaseTransition.class);
-
-  /** A cache that loads the {@link DatabaseTransitionSchedule} for a given id. */
-  private static final LoadingCache<TransitionId, Optional<DatabaseTransitionSchedule>> CACHE =
-      CacheBuilder.newBuilder()
-          .expireAfterWrite(
-              java.time.Duration.ofMillis(getSingletonCacheRefreshDuration().getMillis()))
-          .build(
-              new CacheLoader<TransitionId, Optional<DatabaseTransitionSchedule>>() {
-                @Override
-                public Optional<DatabaseTransitionSchedule> load(TransitionId transitionId) {
-                  return DatabaseTransitionSchedule.get(transitionId);
-                }
-              });
-
-  public static DatabaseTransitionSchedule create(
-      TransitionId transitionId,
-      TimedTransitionProperty<PrimaryDatabase, PrimaryDatabaseTransition> databaseTransitions) {
-    checkNotNull(transitionId, "Id cannot be null");
-    checkNotNull(databaseTransitions, "databaseTransitions cannot be null");
-    databaseTransitions.checkValidity();
-    DatabaseTransitionSchedule instance = new DatabaseTransitionSchedule();
-    instance.transitionId = transitionId.name();
-    instance.databaseTransitions = databaseTransitions;
-    return instance;
-  }
-
-  /** Returns the database that is indicated as primary at the given time. */
-  public PrimaryDatabase getPrimaryDatabase() {
-    return databaseTransitions.getValueAtTime(tm().getTransactionTime());
-  }
-
-  /** Returns the database transitions as a map of start time to primary database. */
-  public ImmutableSortedMap<DateTime, PrimaryDatabase> getDatabaseTransitions() {
-    return databaseTransitions.toValueMap();
-  }
-
-  /**
-   * Returns the current cached schedule for the given id.
-   *
-   * <p>WARNING: The schedule returned by this method could be up to 10 minutes out of date.
-   */
-  public static Optional<DatabaseTransitionSchedule> getCached(TransitionId id) {
-    return CACHE.getUnchecked(id);
-  }
-
-  /** Returns the schedule for a given id. */
-  public static Optional<DatabaseTransitionSchedule> get(TransitionId transitionId) {
-    VKey<DatabaseTransitionSchedule> key =
-        VKey.create(
-            DatabaseTransitionSchedule.class,
-            transitionId,
-            Key.create(getCrossTldKey(), DatabaseTransitionSchedule.class, transitionId.name()));
-
-    return ofyTm().transact(() -> ofyTm().loadByKeyIfPresent(key));
-  }
-
-  @Override
-  public String toString() {
-    return String.format(
-        "%s(last updated at %s): %s",
-        transitionId, lastUpdateTime.getTimestamp(), databaseTransitions.toValueMap());
-  }
-}

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

@@ -14,11 +14,13 @@
 
 package google.registry.model.common;
 
+import com.google.apphosting.api.ApiProxy;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Entity;
 import com.googlecode.objectify.annotation.Id;
 import google.registry.model.BackupGroupRoot;
 import google.registry.schema.replay.DatastoreOnlyEntity;
+import javax.annotation.Nullable;
 
 /**
  * The root key for the entity group which is known as the cross-tld entity group for historical
@@ -42,7 +44,13 @@ public class EntityGroupRoot extends BackupGroupRoot implements DatastoreOnlyEnt
   private String id;
 
   /** The root key for cross-tld resources such as registrars. */
-  public static Key<EntityGroupRoot> getCrossTldKey() {
-    return Key.create(EntityGroupRoot.class, "cross-tld");
+  public static @Nullable Key<EntityGroupRoot> getCrossTldKey() {
+    // If we cannot get a current environment, calling Key.create() will fail. Instead we return a
+    // null in cases where this key is not actually needed (for example when loading an entity from
+    // SQL) to initialize an object, to avoid having to register a DatastoreEntityExtension in
+    // tests.
+    return ApiProxy.getCurrentEnvironment() == null
+        ? null
+        : Key.create(EntityGroupRoot.class, "cross-tld");
   }
 }

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

@@ -15,7 +15,7 @@
 package google.registry.model.common;
 
 import static com.google.common.base.Preconditions.checkState;
-import static google.registry.model.ofy.ObjectifyService.allocateId;
+import static google.registry.model.IdService.allocateId;
 import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
 
 import com.google.appengine.api.users.User;

core/src/main/java/google/registry/model/contact/ContactHistory.java

@@ -18,6 +18,7 @@ import static google.registry.persistence.transaction.TransactionManagerFactory.
 
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.EntitySubclass;
+import google.registry.model.EppResource;
 import google.registry.model.ImmutableObject;
 import google.registry.model.contact.ContactHistory.ContactHistoryId;
 import google.registry.model.reporting.HistoryEntry;
@@ -107,6 +108,11 @@ public class ContactHistory extends HistoryEntry implements SqlEntity {
     return (VKey<ContactHistory>) createVKey(Key.create(this));
   }
 
+  @Override
+  public Optional<? extends EppResource> getResourceAtPointInTime() {
+    return getContactBase();
+  }
+
   @PostLoad
   void postLoad() {
     // Normally Hibernate would see that the contact fields are all null and would fill contactBase
@@ -115,7 +121,10 @@ public class ContactHistory extends HistoryEntry implements SqlEntity {
       contactBase = null;
     }
     if (contactBase != null && contactBase.getRepoId() == null) {
-      contactBase = contactBase.asBuilder().setRepoId(parent.getName()).build();
+      // contactBase hasn't been fully constructed yet, so it's ok to go in and mutate it.  Though
+      // the use of the Builder is not necessarily problematic in this case, this is still safer as
+      // the Builder can do things like comparisons that compute the hash code.
+      contactBase.setRepoId(parent.getName());
     }
   }
 

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

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

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

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

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

@@ -21,9 +21,11 @@ import static google.registry.util.CollectionUtils.nullToEmptyImmutableCopy;
 import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.EntitySubclass;
+import google.registry.model.EppResource;
 import google.registry.model.ImmutableObject;
 import google.registry.model.domain.DomainHistory.DomainHistoryId;
 import google.registry.model.domain.GracePeriod.GracePeriodHistory;
+import google.registry.model.domain.secdns.DelegationSignerData;
 import google.registry.model.domain.secdns.DomainDsDataHistory;
 import google.registry.model.host.HostResource;
 import google.registry.model.reporting.DomainTransactionRecord;
@@ -54,6 +56,7 @@ import javax.persistence.JoinTable;
 import javax.persistence.OneToMany;
 import javax.persistence.PostLoad;
 import javax.persistence.Table;
+import org.hibernate.Hibernate;
 
 /**
  * A persisted history entry representing an EPP modification to a domain.
@@ -247,20 +250,40 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
     return (VKey<DomainHistory>) createVKey(Key.create(this));
   }
 
+  @Override
+  public Optional<? extends EppResource> getResourceAtPointInTime() {
+    return getDomainContent();
+  }
+
   @PostLoad
   void postLoad() {
     if (domainContent != null) {
       domainContent.nsHosts = nullToEmptyImmutableCopy(nsHosts);
+      domainContent.gracePeriods =
+          gracePeriodHistories.stream()
+              .map(GracePeriod::createFromHistory)
+              .collect(toImmutableSet());
+      domainContent.dsData =
+          dsDataHistories.stream().map(DelegationSignerData::create).collect(toImmutableSet());
       // Normally Hibernate would see that the domain fields are all null and would fill
       // domainContent with a null object. Unfortunately, the updateTimestamp is never null in SQL.
       if (domainContent.getDomainName() == null) {
         domainContent = null;
       } else {
         if (domainContent.getRepoId() == null) {
-          domainContent = domainContent.asBuilder().setRepoId(parent.getName()).build();
+          // domainContent still hasn't been fully constructed yet, so it's ok to go in and mutate
+          // it.  In fact, we have to because going through the builder causes the hash codes of
+          // contained objects to be calculated prematurely.
+          domainContent.setRepoId(parent.getName());
         }
       }
     }
+
+    // TODO(b/188044616): Determine why Eager loading doesn't work here.
+    Hibernate.initialize(domainTransactionRecords);
+    Hibernate.initialize(nsHosts);
+    Hibernate.initialize(dsDataHistories);
+    Hibernate.initialize(gracePeriodHistories);
   }
 
   // In Datastore, save as a HistoryEntry object regardless of this object's type

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

@@ -15,6 +15,7 @@
 package google.registry.model.domain;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static google.registry.model.IdService.allocateId;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 
 import com.google.common.annotations.VisibleForTesting;
@@ -22,7 +23,6 @@ import com.googlecode.objectify.annotation.Embed;
 import google.registry.model.billing.BillingEvent;
 import google.registry.model.billing.BillingEvent.Recurring;
 import google.registry.model.domain.rgp.GracePeriodStatus;
-import google.registry.model.ofy.ObjectifyService;
 import google.registry.persistence.BillingVKey.BillingEventVKey;
 import google.registry.persistence.BillingVKey.BillingRecurrenceVKey;
 import google.registry.persistence.VKey;
@@ -68,7 +68,7 @@ public class GracePeriod extends GracePeriodBase implements DatastoreAndSqlEntit
         (billingEventRecurring != null) == GracePeriodStatus.AUTO_RENEW.equals(type),
         "Recurring billing events must be present on (and only on) autorenew grace periods");
     GracePeriod instance = new GracePeriod();
-    instance.gracePeriodId = gracePeriodId == null ? ObjectifyService.allocateId() : gracePeriodId;
+    instance.gracePeriodId = gracePeriodId == null ? allocateId() : gracePeriodId;
     instance.type = checkArgumentNotNull(type);
     instance.domainRepoId = checkArgumentNotNull(domainRepoId);
     instance.expirationTime = checkArgumentNotNull(expirationTime);
@@ -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,
@@ -199,7 +210,7 @@ public class GracePeriod extends GracePeriodBase implements DatastoreAndSqlEntit
 
     static GracePeriodHistory createFrom(long historyRevisionId, GracePeriod gracePeriod) {
       GracePeriodHistory instance = new GracePeriodHistory();
-      instance.gracePeriodHistoryRevisionId = ObjectifyService.allocateId();
+      instance.gracePeriodHistoryRevisionId = allocateId();
       instance.domainHistoryRevisionId = historyRevisionId;
       instance.gracePeriodId = gracePeriod.gracePeriodId;
       instance.type = gracePeriod.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/domain/secdns/DomainDsDataHistory.java

@@ -14,8 +14,9 @@
 
 package google.registry.model.domain.secdns;
 
+import static google.registry.model.IdService.allocateId;
+
 import google.registry.model.domain.DomainHistory;
-import google.registry.model.ofy.ObjectifyService;
 import google.registry.schema.replay.SqlOnlyEntity;
 import javax.persistence.Access;
 import javax.persistence.AccessType;
@@ -48,7 +49,7 @@ public class DomainDsDataHistory extends DomainDsDataBase implements SqlOnlyEnti
     instance.algorithm = dsData.getAlgorithm();
     instance.digestType = dsData.getDigestType();
     instance.digest = dsData.getDigest();
-    instance.dsDataHistoryRevisionId = ObjectifyService.allocateId();
+    instance.dsDataHistoryRevisionId = allocateId();
     return instance;
   }
 

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

@@ -18,6 +18,7 @@ import static google.registry.persistence.transaction.TransactionManagerFactory.
 
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.EntitySubclass;
+import google.registry.model.EppResource;
 import google.registry.model.ImmutableObject;
 import google.registry.model.host.HostHistory.HostHistoryId;
 import google.registry.model.reporting.HistoryEntry;
@@ -108,6 +109,11 @@ public class HostHistory extends HistoryEntry implements SqlEntity {
     return (VKey<HostHistory>) createVKey(Key.create(this));
   }
 
+  @Override
+  public Optional<? extends EppResource> getResourceAtPointInTime() {
+    return getHostBase();
+  }
+
   @PostLoad
   void postLoad() {
     // Normally Hibernate would see that the host fields are all null and would fill hostBase
@@ -116,7 +122,10 @@ public class HostHistory extends HistoryEntry implements SqlEntity {
       hostBase = null;
     }
     if (hostBase != null && hostBase.getRepoId() == null) {
-      hostBase = hostBase.asBuilder().setRepoId(parent.getName()).build();
+      // hostBase hasn't been fully constructed yet, so it's ok to go in and mutate it.  Though the
+      // use of the Builder is not necessarily problematic in this case, this is still safer as the
+      // Builder can do things like comparisons that compute the hash code.
+      hostBase.setRepoId(parent.getName());
     }
   }
 

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

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

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

@@ -222,9 +222,17 @@ public class DatastoreTransactionManager implements TransactionManager {
                 entry -> keyMap.get(entry.getKey()), entry -> toSqlEntity(entry.getValue())));
   }
 
+  @SuppressWarnings("unchecked")
   @Override
   public <T> ImmutableList<T> loadByEntitiesIfPresent(Iterable<T> entities) {
-    return ImmutableList.copyOf(getOfy().load().entities(entities).values());
+    return getOfy()
+        .load()
+        .entities(toDatastoreEntities(ImmutableList.copyOf(entities)))
+        .values()
+        .stream()
+        .map(DatastoreTransactionManager::toSqlEntity)
+        .map(entity -> (T) entity)
+        .collect(toImmutableList());
   }
 
   @Override
@@ -250,6 +258,7 @@ public class DatastoreTransactionManager implements TransactionManager {
     return result;
   }
 
+  @SuppressWarnings("unchecked")
   @Override
   public <T> T loadByEntity(T entity) {
     return (T) toSqlEntity(auditedOfy().load().entity(toDatastoreEntity(entity)).now());
@@ -270,6 +279,11 @@ public class DatastoreTransactionManager implements TransactionManager {
     return ImmutableList.copyOf(getPossibleAncestorQuery(clazz));
   }
 
+  @Override
+  public <T> Stream<T> loadAllOfStream(Class<T> clazz) {
+    return Streams.stream(getPossibleAncestorQuery(clazz));
+  }
+
   @Override
   public <T> Optional<T> loadSingleton(Class<T> clazz) {
     List<T> elements = getPossibleAncestorQuery(clazz).limit(2).list();
@@ -297,8 +311,9 @@ public class DatastoreTransactionManager implements TransactionManager {
   }
 
   @Override
-  public void delete(Object entity) {
+  public <T> T delete(T entity) {
     syncIfTransactionless(getOfy().delete().entity(toDatastoreEntity(entity)));
+    return entity;
   }
 
   @Override
@@ -470,7 +485,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

@@ -21,8 +21,6 @@ import static google.registry.util.TypeUtils.hasAnnotation;
 
 import com.google.appengine.api.datastore.AsyncDatastoreService;
 import com.google.appengine.api.datastore.DatastoreServiceConfig;
-import com.google.appengine.api.datastore.DatastoreServiceFactory;
-import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Streams;
@@ -48,21 +46,14 @@ import google.registry.model.translators.InetAddressTranslatorFactory;
 import google.registry.model.translators.ReadableInstantUtcTranslatorFactory;
 import google.registry.model.translators.UpdateAutoTimestampTranslatorFactory;
 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.
  */
 public class ObjectifyService {
 
-  /**
-   * A placeholder String passed into DatastoreService.allocateIds that ensures that all ids are
-   * initialized from the same id pool.
-   */
-  public static final String APP_WIDE_ALLOCATION_KIND = "common";
-
   /** A singleton instance of our Ofy wrapper. */
   private static final Ofy OFY = new Ofy(null);
 
@@ -101,21 +92,24 @@ public class ObjectifyService {
   private static void initOfyOnce() {
     // Set an ObjectifyFactory that uses our extended ObjectifyImpl.
     // The "false" argument means that we are not using the v5-style Objectify embedded entities.
-    com.googlecode.objectify.ObjectifyService.setFactory(new ObjectifyFactory(false) {
-      @Override
-      public Objectify begin() {
-        return new SessionKeyExposingObjectify(this);
-      }
+    com.googlecode.objectify.ObjectifyService.setFactory(
+        new ObjectifyFactory(false) {
+          @Override
+          public Objectify begin() {
+            return new SessionKeyExposingObjectify(this);
+          }
 
-      @Override
-      protected AsyncDatastoreService createRawAsyncDatastoreService(DatastoreServiceConfig cfg) {
-        // In the unit test environment, wrap the Datastore service in a proxy that can be used to
-        // examine the number of requests sent to Datastore.
-        AsyncDatastoreService service = super.createRawAsyncDatastoreService(cfg);
-        return RegistryEnvironment.get().equals(RegistryEnvironment.UNITTEST)
-            ? new RequestCapturingAsyncDatastoreService(service)
-            : service;
-      }});
+          @Override
+          protected AsyncDatastoreService createRawAsyncDatastoreService(
+              DatastoreServiceConfig cfg) {
+            // In the unit test environment, wrap the Datastore service in a proxy that can be used
+            // to examine the number of requests sent to Datastore.
+            AsyncDatastoreService service = super.createRawAsyncDatastoreService(cfg);
+            return RegistryEnvironment.get().equals(RegistryEnvironment.UNITTEST)
+                ? new RequestCapturingAsyncDatastoreService(service)
+                : service;
+          }
+        });
 
     // Translators must be registered before any entities can be registered.
     registerTranslators();
@@ -168,9 +162,11 @@ public class ObjectifyService {
       } else if (clazz.isAnnotationPresent(EntitySubclass.class)) {
         // Ensure that any @EntitySubclass classes have also had their parent @Entity registered,
         // which Objectify nominally requires but doesn't enforce in 4.x (though it may in 5.x).
-        checkState(registered,
+        checkState(
+            registered,
             "No base entity for kind '%s' registered yet, cannot register new @EntitySubclass %s",
-            kind, clazz.getCanonicalName());
+            kind,
+            clazz.getCanonicalName());
       }
       com.googlecode.objectify.ObjectifyService.register(clazz);
       // Autogenerated ids make the commit log code very difficult since we won't always be able
@@ -186,27 +182,4 @@ public class ObjectifyService {
       }
     }
   }
-
-  /** Counts of used ids for use in unit tests. Outside tests this is never used. */
-  private static final AtomicLong nextTestId = new AtomicLong(1); // ids cannot be zero
-
-  /** Allocates an id. */
-  public static long allocateId() {
-    return RegistryEnvironment.UNITTEST.equals(RegistryEnvironment.get())
-      ? nextTestId.getAndIncrement()
-      : DatastoreServiceFactory.getDatastoreService()
-          .allocateIds(APP_WIDE_ALLOCATION_KIND, 1)
-          .iterator()
-          .next()
-          .getId();
-  }
-
-  /** Resets the global test id counter (i.e. sets the next id to 1). */
-  @VisibleForTesting
-  public static void resetNextTestId() {
-    checkState(RegistryEnvironment.UNITTEST.equals(RegistryEnvironment.get()),
-        "Can't call resetTestIdCounts() from RegistryEnvironment.%s",
-        RegistryEnvironment.get());
-    nextTestId.set(1); // ids cannot be zero
-  }
 }

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

@@ -19,7 +19,7 @@ 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.model.IdService.allocateId;
 import static google.registry.persistence.transaction.QueryComposer.Comparator.EQ;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 
@@ -44,6 +44,7 @@ import google.registry.schema.tld.PremiumListDao;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.math.BigDecimal;
+import java.math.RoundingMode;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -53,6 +54,7 @@ import javax.persistence.Column;
 import javax.persistence.Index;
 import javax.persistence.PostLoad;
 import javax.persistence.PostPersist;
+import javax.persistence.PostUpdate;
 import javax.persistence.PrePersist;
 import javax.persistence.PreRemove;
 import javax.persistence.Table;
@@ -182,11 +184,22 @@ public final class PremiumList extends BaseDomainLabelList<Money, PremiumList.Pr
               .createQueryComposer(PremiumEntry.class)
               .where("revisionId", EQ, revisionId)
               .stream()
-              .collect(toImmutableMap(PremiumEntry::getDomainLabel, PremiumEntry::getPrice));
+              .collect(
+                  toImmutableMap(
+                      PremiumEntry::getDomainLabel,
+                      // Set the correct amount of precision for the premium list's currency.
+                      entry -> convertAmountToMoney(entry.getPrice()).getAmount()));
     }
     return labelsToPrices;
   }
 
+  /**
+   * Converts a raw {@link BigDecimal} amount to a {@link Money} by applying the list's currency.
+   */
+  public Money convertAmountToMoney(BigDecimal amount) {
+    return Money.of(currency, amount.setScale(currency.getDecimalPlaces(), RoundingMode.HALF_EVEN));
+  }
+
   /**
    * Returns a Bloom filter to determine whether a label might be premium, or is definitely not.
    *

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

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

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

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

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

@@ -61,6 +61,7 @@ import javax.persistence.EnumType;
 import javax.persistence.Enumerated;
 import javax.persistence.MappedSuperclass;
 import javax.persistence.Transient;
+import org.apache.commons.lang3.BooleanUtils;
 import org.joda.time.DateTime;
 
 /**
@@ -289,6 +290,17 @@ public class HistoryEntry extends ImmutableObject implements Buildable, Datastor
     return nullToEmptyImmutableCopy(domainTransactionRecords);
   }
 
+  /**
+   * Throws an error when attempting to retrieve the EppResource at this point in time.
+   *
+   * <p>Subclasses must override this to return the resource; it is non-abstract for legacy reasons
+   * and objects created prior to the Registry 3.0 migration.
+   */
+  public Optional<? extends EppResource> getResourceAtPointInTime() {
+    throw new UnsupportedOperationException(
+        "Raw HistoryEntry objects do not store the resource at that point in time.");
+  }
+
   /** This method exists solely to satisfy Hibernate. Use the {@link Builder} instead. */
   @SuppressWarnings("UnusedMethod")
   private void setPeriod(Period period) {
@@ -435,7 +447,8 @@ public class HistoryEntry extends ImmutableObject implements Buildable, Datastor
       checkArgumentNotNull(getInstance().modificationTime, "Modification time must be specified");
       checkArgumentNotNull(getInstance().clientId, "Registrar ID must be specified");
       checkArgument(
-          !getInstance().type.equals(Type.SYNTHETIC) || !getInstance().requestedByRegistrar,
+          !getInstance().type.equals(Type.SYNTHETIC)
+              || BooleanUtils.isNotTrue(getInstance().requestedByRegistrar),
           "Synthetic history entries cannot be requested by a registrar");
       return super.build();
     }

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

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

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

@@ -0,0 +1,56 @@
+// Copyright 2021 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.model.tmch;
+
+import google.registry.model.ImmutableObject;
+import google.registry.schema.replay.NonReplicatedEntity;
+import java.io.Serializable;
+import javax.persistence.Column;
+import javax.persistence.Entity;
+import javax.persistence.Id;
+
+/**
+ * Claims entry record, used by ClaimsList for persistence.
+ *
+ * <p>It would be preferable to have this nested in {@link ClaimsList}, but for some reason
+ * hibernate won't generate this into the schema in this case. We may not care, as we only use the
+ * generated schema for informational purposes and persistence against the actual schema seems to
+ * work.
+ */
+@Entity(name = "ClaimsEntry")
+class ClaimsEntry extends ImmutableObject implements NonReplicatedEntity, Serializable {
+  @Id private Long revisionId;
+  @Id private String domainLabel;
+
+  @Column(nullable = false)
+  private String claimKey;
+
+  /** Default constructor for Hibernate. */
+  ClaimsEntry() {}
+
+  ClaimsEntry(Long revisionId, String domainLabel, String claimKey) {
+    this.revisionId = revisionId;
+    this.domainLabel = domainLabel;
+    this.claimKey = claimKey;
+  }
+
+  String getDomainLabel() {
+    return domainLabel;
+  }
+
+  String getClaimKey() {
+    return claimKey;
+  }
+}

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

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

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

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

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

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

core/src/main/java/google/registry/model/translators/CommitLogRevisionsTranslatorFactory.java

@@ -17,13 +17,14 @@ package google.registry.model.translators;
 import static com.google.common.base.MoreObjects.firstNonNull;
 import static google.registry.config.RegistryConfig.getCommitLogDatastoreRetention;
 import static google.registry.model.ofy.ObjectifyService.auditedOfy;
-import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
 
 import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.collect.Ordering;
 import com.googlecode.objectify.Key;
 import google.registry.model.ofy.CommitLogManifest;
+import google.registry.persistence.transaction.Transaction;
 import org.joda.time.DateTime;
 
 /**
@@ -58,12 +59,20 @@ public final class CommitLogRevisionsTranslatorFactory
    * <p>We store a maximum of one entry per day. It will be the last transaction that happened on
    * that day.
    *
+   * <p>In serialization mode, this method just returns "revisions" without modification.
+   *
    * @see google.registry.config.RegistryConfig#getCommitLogDatastoreRetention()
    */
   @Override
   ImmutableSortedMap<DateTime, Key<CommitLogManifest>> transformBeforeSave(
       ImmutableSortedMap<DateTime, Key<CommitLogManifest>> revisions) {
-    DateTime now = tm().getTransactionTime();
+
+    // Don't do anything if we're just doing object serialization.
+    if (Transaction.inSerializationMode()) {
+      return revisions;
+    }
+
+    DateTime now = ofyTm().getTransactionTime();
     DateTime threshold = now.minus(getCommitLogDatastoreRetention());
     DateTime preThresholdTime = firstNonNull(revisions.floorKey(threshold), START_OF_TIME);
     return new ImmutableSortedMap.Builder<DateTime, Key<CommitLogManifest>>(Ordering.natural())

core/src/main/java/google/registry/model/translators/CreateAutoTimestampTranslatorFactory.java

@@ -15,10 +15,11 @@
 package google.registry.model.translators;
 
 import static com.google.common.base.MoreObjects.firstNonNull;
-import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
 import static org.joda.time.DateTimeZone.UTC;
 
 import google.registry.model.CreateAutoTimestamp;
+import google.registry.persistence.transaction.Transaction;
 import java.util.Date;
 import org.joda.time.DateTime;
 
@@ -46,7 +47,14 @@ public class CreateAutoTimestampTranslatorFactory
       /** Save a timestamp, setting it to the current time if it did not have a previous value. */
       @Override
       public Date saveValue(CreateAutoTimestamp pojoValue) {
-        return firstNonNull(pojoValue.getTimestamp(), tm().getTransactionTime()).toDate();
-      }};
+
+        // Don't do this if we're in the course of transaction serialization.
+        if (Transaction.inSerializationMode()) {
+          return pojoValue.getTimestamp() == null ? null : pojoValue.getTimestamp().toDate();
+        }
+
+        return firstNonNull(pojoValue.getTimestamp(), ofyTm().getTransactionTime()).toDate();
+      }
+    };
   }
 }

core/src/main/java/google/registry/model/translators/UpdateAutoTimestampTranslatorFactory.java

@@ -14,10 +14,11 @@
 
 package google.registry.model.translators;
 
-import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
 import static org.joda.time.DateTimeZone.UTC;
 
 import google.registry.model.UpdateAutoTimestamp;
+import google.registry.persistence.transaction.Transaction;
 import java.util.Date;
 import org.joda.time.DateTime;
 
@@ -46,8 +47,14 @@ public class UpdateAutoTimestampTranslatorFactory
       /** Save a timestamp, setting it to the current time. */
       @Override
       public Date saveValue(UpdateAutoTimestamp pojoValue) {
+
+        // Don't do this if we're in the course of transaction serialization.
+        if (Transaction.inSerializationMode()) {
+          return pojoValue.getTimestamp() == null ? null : pojoValue.getTimestamp().toDate();
+        }
+
         return UpdateAutoTimestamp.autoUpdateEnabled()
-            ? tm().getTransactionTime().toDate()
+            ? ofyTm().getTransactionTime().toDate()
             : pojoValue.getTimestamp().toDate();
       }
     };

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

@@ -74,6 +74,15 @@ public abstract class PersistenceModule {
   public static final String HIKARI_DS_CLOUD_SQL_INSTANCE =
       "hibernate.hikari.dataSource.cloudSqlInstance";
 
+  /**
+   * Postgresql-specific: driver default fetch size is 0, which disables streaming result sets. Here
+   * we set a small default geared toward Nomulus server transactions. Large queries can override
+   * the defaults using {@link JpaTransactionManager#setQueryFetchSize}.
+   */
+  public static final String JDBC_FETCH_SIZE = "hibernate.jdbc.fetch_size";
+
+  private static final int DEFAULT_SERVER_FETCH_SIZE = 20;
+
   @VisibleForTesting
   @Provides
   @DefaultHibernateConfigs
@@ -100,6 +109,7 @@ public abstract class PersistenceModule {
     properties.put(HIKARI_MAXIMUM_POOL_SIZE, getHibernateHikariMaximumPoolSize());
     properties.put(HIKARI_IDLE_TIMEOUT, getHibernateHikariIdleTimeout());
     properties.put(Environment.DIALECT, NomulusPostgreSQLDialect.class.getName());
+    properties.put(JDBC_FETCH_SIZE, Integer.toString(DEFAULT_SERVER_FETCH_SIZE));
     return properties.build();
   }
 

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,14 +35,26 @@ 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);
 
+  /**
+   * Execute the work in a transaction without recording the transaction for replay to datastore.
+   */
+  <T> T transactWithoutBackup(Supplier<T> work);
+
   /** Executes the work in a transaction with no retries and returns the result. */
   <T> T transactNoRetry(Supplier<T> work);
 
@@ -57,4 +70,12 @@ public interface JpaTransactionManager extends TransactionManager {
    * <p>The errorprone check forbids injection of {@link java.io.Closeable} resources.
    */
   void teardown();
+
+  /**
+   * Sets the JDBC driver fetch size for the {@code query}. This overrides the default
+   * configuration.
+   */
+  static Query setQueryFetchSize(Query query, int fetchSize) {
+    return query.setHint("org.hibernate.fetchSize", fetchSize);
+  }
 }

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;
@@ -40,23 +41,37 @@ import google.registry.model.server.KmsSecret;
 import google.registry.model.tmch.ClaimsList.ClaimsListSingleton;
 import google.registry.persistence.JpaRetries;
 import google.registry.persistence.VKey;
+import google.registry.schema.replay.NonReplicatedEntity;
+import google.registry.schema.replay.SqlOnlyEntity;
 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 +127,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
@@ -134,6 +154,15 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
 
   @Override
   public <T> T transact(Supplier<T> work) {
+    return transact(work, true);
+  }
+
+  @Override
+  public <T> T transactWithoutBackup(Supplier<T> work) {
+    return transact(work, false);
+  }
+
+  private <T> T transact(Supplier<T> work, boolean withBackup) {
     return retrier.callWithRetry(
         () -> {
           if (inTransaction()) {
@@ -144,7 +173,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
           EntityTransaction txn = txnInfo.entityManager.getTransaction();
           try {
             txn.begin();
-            txnInfo.start(clock);
+            txnInfo.start(clock, withBackup);
             T result = work.get();
             txnInfo.recordTransaction();
             txn.commit();
@@ -176,7 +205,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     EntityTransaction txn = txnInfo.entityManager.getTransaction();
     try {
       txn.begin();
-      txnInfo.start(clock);
+      txnInfo.start(clock, true);
       T result = work.get();
       txnInfo.recordTransaction();
       txn.commit();
@@ -269,8 +298,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 +327,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 +366,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     checkArgument(exists(entity), "Given entity does not exist");
     // Necessary due to the changes in HistoryEntry representation during the migration to SQL
     Object toPersist = toSqlEntity(entity);
-    getEntityManager().merge(toPersist);
-    transactionInfo.get().addUpdate(toPersist);
+    transactionInfo.get().updateObject(toPersist);
   }
 
   @Override
@@ -397,7 +423,8 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   public <T> Optional<T> loadByKeyIfPresent(VKey<T> key) {
     checkArgumentNotNull(key, "key must be specified");
     assertInTransaction();
-    return Optional.ofNullable(getEntityManager().find(key.getKind(), key.getSqlKey()));
+    return Optional.ofNullable(getEntityManager().find(key.getKind(), key.getSqlKey()))
+        .map(this::detach);
   }
 
   @Override
@@ -411,7 +438,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
         .map(
             key ->
                 new SimpleEntry<VKey<? extends T>, T>(
-                    key, getEntityManager().find(key.getKind(), key.getSqlKey())))
+                    key, detach(getEntityManager().find(key.getKind(), key.getSqlKey()))))
         .filter(entry -> entry.getValue() != null)
         .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue));
   }
@@ -433,7 +460,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     if (result == null) {
       throw new NoSuchElementException(key.toString());
     }
-    return result;
+    return detach(result);
   }
 
   @Override
@@ -456,11 +483,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
@@ -470,14 +500,17 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
 
   @Override
   public <T> ImmutableList<T> loadAllOf(Class<T> clazz) {
+    return loadAllOfStream(clazz).collect(toImmutableList());
+  }
+
+  @Override
+  public <T> Stream<T> loadAllOfStream(Class<T> clazz) {
     checkArgumentNotNull(clazz, "clazz must be specified");
     assertInTransaction();
-    return ImmutableList.copyOf(
-        getEntityManager()
-            .createQuery(
-                String.format("SELECT entity FROM %s entity", getEntityType(clazz).getName()),
-                clazz)
-            .getResultList());
+    return getEntityManager()
+        .createQuery(String.format("FROM %s", getEntityType(clazz).getName()), clazz)
+        .getResultStream()
+        .map(this::detach);
   }
 
   @Override
@@ -492,7 +525,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
         elements.size() <= 1,
         "Expected at most one entity of type %s, found at least two",
         clazz.getSimpleName());
-    return elements.stream().findFirst();
+    return elements.stream().findFirst().map(this::detach);
   }
 
   private int internalDelete(VKey<?> key) {
@@ -524,18 +557,21 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   }
 
   @Override
-  public void delete(Object entity) {
+  public <T> T delete(T entity) {
     checkArgumentNotNull(entity, "entity must be specified");
     if (isEntityOfIgnoredClass(entity)) {
-      return;
+      return entity;
     }
     assertInTransaction();
     entity = toSqlEntity(entity);
-    Object managedEntity = entity;
+    T managedEntity = entity;
     if (!getEntityManager().contains(entity)) {
+      // We don't add the entity to "objectsToSave": once deleted, the object should never be
+      // returned as a result of the query or lookup.
       managedEntity = getEntityManager().merge(entity);
     }
     getEntityManager().remove(managedEntity);
+    return managedEntity;
   }
 
   @Override
@@ -555,7 +591,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
 
   @Override
   public <T> QueryComposer<T> createQueryComposer(Class<T> entity) {
-    return new JpaQueryComposerImpl<T>(entity, getEntityManager());
+    return new JpaQueryComposerImpl<T>(entity);
   }
 
   @Override
@@ -658,6 +694,48 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     }
   }
 
+  @Nullable
+  private <T> T detachIfEntity(@Nullable T object) {
+    if (object == null) {
+      return null;
+    }
+
+    // Check if the object is an array, if so we'll want to recurse through the elements.
+    if (object.getClass().isArray()) {
+      for (int i = 0; i < Array.getLength(object); ++i) {
+        detachIfEntity(Array.get(object, i));
+      }
+      return object;
+    }
+
+    // Check to see if it is an entity (queries can return raw column values or counts, so this
+    // could be String, Long, ...).
+    try {
+      getEntityManager().getMetamodel().entity(object.getClass());
+    } catch (IllegalArgumentException e) {
+      // The object is not an entity.  Return without detaching.
+      return object;
+    }
+
+    // At this point, object must be an entity.
+    return detach(object);
+  }
+
+  /** Detach the entity, suitable for use in Optional.map(). */
+  @Nullable
+  private <T> T detach(@Nullable T entity) {
+    if (entity != null) {
+
+      // If the entity was previously persisted or merged, we have to throw an exception.
+      if (transactionInfo.get().willSave(entity)) {
+        throw new IllegalStateException("Inserted/updated object reloaded: " + entity);
+      }
+
+      getEntityManager().detach(entity);
+    }
+    return entity;
+  }
+
   private static class TransactionInfo {
     EntityManager entityManager;
     boolean inTransaction = false;
@@ -666,12 +744,18 @@ 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) {
+    private void start(Clock clock, boolean withBackup) {
       checkArgumentNotNull(clock);
       inTransaction = true;
       transactionTime = clock.nowUtc();
-      if (RegistryConfig.getCloudSqlReplicateTransactions()) {
+      if (withBackup && RegistryConfig.getCloudSqlReplicateTransactions()) {
         contentsBuilder = new Transaction.Builder();
       }
     }
@@ -680,6 +764,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.
@@ -689,17 +774,23 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     }
 
     private void addUpdate(Object entity) {
-      if (contentsBuilder != null) {
+      if (contentsBuilder != null && shouldReplicate(entity.getClass())) {
         contentsBuilder.addUpdate(entity);
       }
     }
 
     private void addDelete(VKey<?> key) {
-      if (contentsBuilder != null) {
+      if (contentsBuilder != null && shouldReplicate(key.getKind())) {
         contentsBuilder.addDelete(key);
       }
     }
 
+    /** Returns true if the entity class should be replicated from SQL to datastore. */
+    private boolean shouldReplicate(Class<?> entityClass) {
+      return !NonReplicatedEntity.class.isAssignableFrom(entityClass)
+          && !SqlOnlyEntity.class.isAssignableFrom(entityClass);
+    }
+
     private void recordTransaction() {
       if (contentsBuilder != null) {
         Transaction persistedTxn = contentsBuilder.build();
@@ -708,25 +799,241 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
         }
       }
     }
+
+    /** Does the full "update" on an object including all internal housekeeping. */
+    private void updateObject(Object object) {
+      Object merged = entityManager.merge(object);
+      objectsToSave.add(merged);
+      addUpdate(object);
+    }
+
+    /** Does the full "insert" on a new object including all internal housekeeping. */
+    private void insertObject(Object object) {
+      entityManager.persist(object);
+      objectsToSave.add(object);
+      addUpdate(object);
+    }
+
+    /** Returns true if the object has been persisted/merged and will be saved on commit. */
+    private boolean willSave(Object object) {
+      return objectsToSave.contains(object);
+    }
   }
 
-  private static class JpaQueryComposerImpl<T> extends QueryComposer<T> {
+  /**
+   * Typed query wrapper that applies a transform to all result objects.
+   *
+   * <p>This is used to detach objects upon load.
+   */
+  @VisibleForTesting
+  class DetachingTypedQuery<T> implements TypedQuery<T> {
+
+    TypedQuery<T> delegate;
+
+    public DetachingTypedQuery(TypedQuery<T> delegate) {
+      this.delegate = delegate;
+    }
+
+    @Override
+    public List<T> getResultList() {
+      return delegate
+          .getResultStream()
+          .map(JpaTransactionManagerImpl.this::detachIfEntity)
+          .collect(toImmutableList());
+    }
+
+    @Override
+    public Stream<T> getResultStream() {
+      return delegate.getResultStream().map(JpaTransactionManagerImpl.this::detachIfEntity);
+    }
+
+    @Override
+    public T getSingleResult() {
+      return detachIfEntity(delegate.getSingleResult());
+    }
+
+    @Override
+    public TypedQuery<T> setMaxResults(int maxResults) {
+      delegate.setMaxResults(maxResults);
+      return this;
+    }
+
+    @Override
+    public TypedQuery<T> setFirstResult(int startPosition) {
+      delegate.setFirstResult(startPosition);
+      return this;
+    }
+
+    @Override
+    public TypedQuery<T> setHint(String hintName, Object value) {
+      delegate.setHint(hintName, value);
+      return this;
+    }
+
+    @Override
+    public <U> TypedQuery<T> setParameter(Parameter<U> param, U value) {
+      delegate.setParameter(param, value);
+      return this;
+    }
+
+    @Override
+    public TypedQuery<T> setParameter(
+        Parameter<Calendar> param, Calendar value, TemporalType temporalType) {
+      delegate.setParameter(param, value, temporalType);
+      return this;
+    }
+
+    @Override
+    public TypedQuery<T> setParameter(
+        Parameter<Date> param, Date value, TemporalType temporalType) {
+      delegate.setParameter(param, value, temporalType);
+      return this;
+    }
+
+    @Override
+    public TypedQuery<T> setParameter(String name, Object value) {
+      delegate.setParameter(name, value);
+      return this;
+    }
+
+    @Override
+    public TypedQuery<T> setParameter(String name, Calendar value, TemporalType temporalType) {
+      delegate.setParameter(name, value, temporalType);
+      return this;
+    }
+
+    @Override
+    public TypedQuery<T> setParameter(String name, Date value, TemporalType temporalType) {
+      delegate.setParameter(name, value, temporalType);
+      return this;
+    }
+
+    @Override
+    public TypedQuery<T> setParameter(int position, Object value) {
+      delegate.setParameter(position, value);
+      return this;
+    }
+
+    @Override
+    public TypedQuery<T> setParameter(int position, Calendar value, TemporalType temporalType) {
+      delegate.setParameter(position, value, temporalType);
+      return this;
+    }
+
+    @Override
+    public TypedQuery<T> setParameter(int position, Date value, TemporalType temporalType) {
+      delegate.setParameter(position, value, temporalType);
+      return this;
+    }
+
+    @Override
+    public TypedQuery<T> setFlushMode(FlushModeType flushMode) {
+      delegate.setFlushMode(flushMode);
+      return this;
+    }
+
+    @Override
+    public TypedQuery<T> setLockMode(LockModeType lockMode) {
+      delegate.setLockMode(lockMode);
+      return this;
+    }
+
+    // Query interface
+
+    @Override
+    public int executeUpdate() {
+      return delegate.executeUpdate();
+    }
+
+    @Override
+    public int getMaxResults() {
+      return delegate.getMaxResults();
+    }
+
+    @Override
+    public int getFirstResult() {
+      return delegate.getFirstResult();
+    }
+
+    @Override
+    public Map<String, Object> getHints() {
+      return delegate.getHints();
+    }
+
+    @Override
+    public Set<Parameter<?>> getParameters() {
+      return delegate.getParameters();
+    }
+
+    @Override
+    public Parameter<?> getParameter(String name) {
+      return delegate.getParameter(name);
+    }
+
+    @Override
+    public <U> Parameter<U> getParameter(String name, Class<U> type) {
+      return delegate.getParameter(name, type);
+    }
+
+    @Override
+    public Parameter<?> getParameter(int position) {
+      return delegate.getParameter(position);
+    }
+
+    @Override
+    public <U> Parameter<U> getParameter(int position, Class<U> type) {
+      return delegate.getParameter(position, type);
+    }
+
+    @Override
+    public boolean isBound(Parameter<?> param) {
+      return delegate.isBound(param);
+    }
+
+    @Override
+    public <U> U getParameterValue(Parameter<U> param) {
+      return delegate.getParameterValue(param);
+    }
+
+    @Override
+    public Object getParameterValue(String name) {
+      return delegate.getParameterValue(name);
+    }
+
+    @Override
+    public Object getParameterValue(int position) {
+      return delegate.getParameterValue(position);
+    }
+
+    @Override
+    public FlushModeType getFlushMode() {
+      return delegate.getFlushMode();
+    }
+
+    @Override
+    public LockModeType getLockMode() {
+      return delegate.getLockMode();
+    }
+
+    @Override
+    public <U> U unwrap(Class<U> cls) {
+      return delegate.unwrap(cls);
+    }
+  }
+
+  private class JpaQueryComposerImpl<T> extends QueryComposer<T> {
 
     private static final int DEFAULT_FETCH_SIZE = 1000;
 
-    EntityManager em;
-
     private int fetchSize = DEFAULT_FETCH_SIZE;
 
-    private boolean autoDetachOnLoad = true;
-
-    JpaQueryComposerImpl(Class<T> entityClass, EntityManager em) {
+    JpaQueryComposerImpl(Class<T> entityClass) {
       super(entityClass);
-      this.em = em;
     }
 
     private TypedQuery<T> buildQuery() {
-      CriteriaQueryBuilder<T> queryBuilder = CriteriaQueryBuilder.create(em, entityClass);
+      CriteriaQueryBuilder<T> queryBuilder =
+          CriteriaQueryBuilder.create(getEntityManager(), entityClass);
       return addCriteria(queryBuilder);
     }
 
@@ -739,13 +1046,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
         queryBuilder.orderByAsc(orderBy);
       }
 
-      return em.createQuery(queryBuilder.build());
-    }
-
-    @Override
-    public QueryComposer<T> withAutoDetachOnLoad(boolean autoDetachOnLoad) {
-      this.autoDetachOnLoad = autoDetachOnLoad;
-      return this;
+      return getEntityManager().createQuery(queryBuilder.build());
     }
 
     @Override
@@ -758,12 +1059,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
@@ -772,32 +1073,22 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
         logger.atWarning().log("Query result streaming is not enabled.");
       }
       TypedQuery<T> query = buildQuery();
-      if (query instanceof org.hibernate.query.Query) {
-        ((org.hibernate.query.Query) query).setFetchSize(fetchSize);
-      } else {
-        logger.atWarning().log("Query implemention does not support result streaming.");
-      }
-      return query.getResultStream().map(this::maybeDetachEntity);
+      JpaTransactionManager.setQueryFetchSize(query, fetchSize);
+      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/Transaction.java

@@ -18,7 +18,6 @@ import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
 import static google.registry.model.ofy.ObjectifyService.auditedOfy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
-import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.appengine.api.datastore.Entity;
 import com.google.appengine.api.datastore.EntityTranslator;
@@ -51,11 +50,17 @@ public class Transaction extends ImmutableObject implements Buildable {
   // unique and inherently informative.
   private static final int VERSION_ID = 20200604;
 
+  // Keep a per-thread flag to keep track of whether we're serializing an entity for a transaction.
+  // This is used by internal translators to avoid doing things that are dependent on being in a
+  // datastore transaction and alter the persisted representation of the entity.
+  private static ThreadLocal<Boolean> inSerializationMode = ThreadLocal.withInitial(() -> false);
+
   private transient ImmutableList<Mutation> mutations;
 
   /** Write the entire transaction to the datastore in a datastore transaction. */
   public void writeToDatastore() {
-    tm().transact(
+    ofyTm()
+        .transact(
             () -> {
               for (Mutation mutation : mutations) {
                 mutation.writeToDatastore();
@@ -75,8 +80,13 @@ public class Transaction extends ImmutableObject implements Buildable {
 
       // Write all of the mutations, preceded by their count.
       out.writeInt(mutations.size());
-      for (Mutation mutation : mutations) {
-        mutation.serializeTo(out);
+      try {
+        inSerializationMode.set(true);
+        for (Mutation mutation : mutations) {
+          mutation.serializeTo(out);
+        }
+      } finally {
+        inSerializationMode.set(false);
       }
 
       out.close();
@@ -114,6 +124,17 @@ public class Transaction extends ImmutableObject implements Buildable {
     return mutations.isEmpty();
   }
 
+  /**
+   * Returns true if we are serializing a transaction in the current thread.
+   *
+   * <p>This should be checked by any Ofy translators prior to making any changes to an entity's
+   * state representation based on the assumption that we are currently pseristing the entity to
+   * datastore.
+   */
+  public static boolean inSerializationMode() {
+    return inSerializationMode.get();
+  }
+
   @Override
   public Builder asBuilder() {
     return new Builder(clone(this));