Restore commit logs from other project (#1236)

* Restore commit logs from other project

Allow non-production projects to restore commit logs from another
project. This feature can be used to duplicate a realistic testing
environment.

An optional parameter is added that can override the default commit log
location.

Tested successfully in QA.

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 {
@@ -777,6 +780,10 @@ if (environment in ['alpha', 'crash']) {
           mainClass: 'google.registry.beam.invoicing.InvoicingPipeline',
           metaData: 'google/registry/beam/invoicing_pipeline_metadata.json'
       ],
+      [
+          mainClass: 'google.registry.beam.rde.RdePipeline',
+          metaData: 'google/registry/beam/rde_pipeline_metadata.json'
+      ],
   ]
   project.tasks.create("stage_beam_pipelines") {
     doLast {

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

@@ -18,8 +18,10 @@ import static com.google.appengine.api.ThreadManager.currentRequestThreadFactory
 import static com.google.common.util.concurrent.MoreExecutors.listeningDecorator;
 import static google.registry.backup.ExportCommitLogDiffAction.LOWER_CHECKPOINT_TIME_PARAM;
 import static google.registry.backup.ExportCommitLogDiffAction.UPPER_CHECKPOINT_TIME_PARAM;
+import static google.registry.backup.RestoreCommitLogsAction.BUCKET_OVERRIDE_PARAM;
 import static google.registry.backup.RestoreCommitLogsAction.FROM_TIME_PARAM;
 import static google.registry.backup.RestoreCommitLogsAction.TO_TIME_PARAM;
+import static google.registry.request.RequestParameters.extractOptionalParameter;
 import static google.registry.request.RequestParameters.extractRequiredDatetimeParameter;
 import static google.registry.request.RequestParameters.extractRequiredParameter;
 import static java.util.concurrent.Executors.newFixedThreadPool;
@@ -32,6 +34,7 @@ import google.registry.cron.CommitLogFanoutAction;
 import google.registry.request.HttpException.BadRequestException;
 import google.registry.request.Parameter;
 import java.lang.annotation.Documented;
+import java.util.Optional;
 import javax.inject.Qualifier;
 import javax.servlet.http.HttpServletRequest;
 import org.joda.time.DateTime;
@@ -75,6 +78,12 @@ public final class BackupModule {
     return extractRequiredDatetimeParameter(req, UPPER_CHECKPOINT_TIME_PARAM);
   }
 
+  @Provides
+  @Parameter(BUCKET_OVERRIDE_PARAM)
+  static Optional<String> provideBucketOverride(HttpServletRequest req) {
+    return extractOptionalParameter(req, BUCKET_OVERRIDE_PARAM);
+  }
+
   @Provides
   @Parameter(FROM_TIME_PARAM)
   static DateTime provideFromTime(HttpServletRequest req) {

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

@@ -32,7 +32,6 @@ import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.ListeningExecutorService;
 import google.registry.backup.BackupModule.Backups;
-import google.registry.config.RegistryConfig.Config;
 import java.io.IOException;
 import java.util.Iterator;
 import java.util.Map;
@@ -47,16 +46,16 @@ class GcsDiffFileLister {
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
 
   @Inject GcsService gcsService;
-  @Inject @Config("commitLogGcsBucket") String gcsBucket;
   @Inject @Backups ListeningExecutorService executor;
   @Inject GcsDiffFileLister() {}
 
   /**
    * Traverses the sequence of diff files backwards from checkpointTime and inserts the file
-   * metadata into "sequence".  Returns true if a complete sequence was discovered, false if one or
+   * metadata into "sequence". Returns true if a complete sequence was discovered, false if one or
    * more files are missing.
    */
   private boolean constructDiffSequence(
+      String gcsBucket,
       Map<DateTime, ListenableFuture<GcsFileMetadata>> upperBoundTimesToMetadata,
       DateTime fromTime,
       DateTime lastTime,
@@ -69,7 +68,7 @@ class GcsDiffFileLister {
       } else {
         String filename = DIFF_FILE_PREFIX + checkpointTime;
         logger.atInfo().log("Patching GCS list; discovered file: %s", filename);
-        metadata = getMetadata(filename);
+        metadata = getMetadata(gcsBucket, filename);
 
         // If we hit a gap, quit.
         if (metadata == null) {
@@ -87,7 +86,8 @@ class GcsDiffFileLister {
     return true;
   }
 
-  ImmutableList<GcsFileMetadata> listDiffFiles(DateTime fromTime, @Nullable DateTime toTime) {
+  ImmutableList<GcsFileMetadata> listDiffFiles(
+      String gcsBucket, DateTime fromTime, @Nullable DateTime toTime) {
     logger.atInfo().log("Requested restore from time: %s", fromTime);
     if (toTime != null) {
       logger.atInfo().log("  Until time: %s", toTime);
@@ -111,7 +111,8 @@ class GcsDiffFileLister {
       final String filename = listItems.next().getName();
       DateTime upperBoundTime = DateTime.parse(filename.substring(DIFF_FILE_PREFIX.length()));
       if (isInRange(upperBoundTime, fromTime, toTime)) {
-        upperBoundTimesToMetadata.put(upperBoundTime, executor.submit(() -> getMetadata(filename)));
+        upperBoundTimesToMetadata.put(
+            upperBoundTime, executor.submit(() -> getMetadata(gcsBucket, filename)));
         lastUpperBoundTime = latestOf(upperBoundTime, lastUpperBoundTime);
       }
     }
@@ -130,8 +131,9 @@ class GcsDiffFileLister {
     // may be missing files at the end).
     TreeMap<DateTime, GcsFileMetadata> sequence = new TreeMap<>();
     logger.atInfo().log("Restoring until: %s", lastUpperBoundTime);
-    boolean inconsistentFileSet = !constructDiffSequence(
-        upperBoundTimesToMetadata, fromTime, lastUpperBoundTime, sequence);
+    boolean inconsistentFileSet =
+        !constructDiffSequence(
+            gcsBucket, upperBoundTimesToMetadata, fromTime, lastUpperBoundTime, sequence);
 
     // Verify that all of the elements in the original set are represented in the sequence.  If we
     // find anything that's not represented, construct a sequence for it.
@@ -143,7 +145,7 @@ class GcsDiffFileLister {
           break;
         }
         if (!sequence.containsKey(key)) {
-          constructDiffSequence(upperBoundTimesToMetadata, fromTime, key, sequence);
+          constructDiffSequence(gcsBucket, upperBoundTimesToMetadata, fromTime, key, sequence);
           checkForMoreExtraDiffs = true;
           inconsistentFileSet = true;
           break;
@@ -175,7 +177,7 @@ class GcsDiffFileLister {
     return DateTime.parse(metadata.getOptions().getUserMetadata().get(LOWER_BOUND_CHECKPOINT));
   }
 
-  private GcsFileMetadata getMetadata(String filename) {
+  private GcsFileMetadata getMetadata(String gcsBucket, String filename) {
     try {
       return gcsService.getMetadata(new GcsFilename(gcsBucket, filename));
     } catch (IOException e) {

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

@@ -14,7 +14,9 @@
 
 package google.registry.backup;
 
+import static com.google.common.collect.ImmutableList.toImmutableList;
 import static google.registry.backup.ExportCommitLogDiffAction.DIFF_FILE_PREFIX;
+import static google.registry.backup.RestoreCommitLogsAction.DRY_RUN_PARAM;
 import static google.registry.model.ofy.EntityWritePriorities.getEntityPriority;
 import static google.registry.model.ofy.ObjectifyService.auditedOfy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
@@ -25,8 +27,10 @@ import com.google.appengine.api.datastore.Entity;
 import com.google.appengine.api.datastore.Key;
 import com.google.appengine.tools.cloudstorage.GcsFileMetadata;
 import com.google.appengine.tools.cloudstorage.GcsService;
+import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableList;
 import com.google.common.flogger.FluentLogger;
+import google.registry.config.RegistryConfig.Config;
 import google.registry.model.common.DatabaseMigrationStateSchedule;
 import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState;
 import google.registry.model.common.DatabaseMigrationStateSchedule.ReplayDirection;
@@ -35,6 +39,7 @@ 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.Parameter;
 import google.registry.request.Response;
 import google.registry.request.auth.Auth;
 import google.registry.schema.replay.DatastoreEntity;
@@ -64,10 +69,13 @@ 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;
@@ -75,6 +83,15 @@ public class ReplayCommitLogsToSqlAction implements Runnable {
   @Inject GcsDiffFileLister diffLister;
   @Inject Clock clock;
 
+  @Inject
+  @Config("commitLogGcsBucket")
+  String gcsBucket;
+
+  /** If true, will exit after logging the commit log files that would otherwise be replayed. */
+  @Inject
+  @Parameter(DRY_RUN_PARAM)
+  boolean dryRun;
+
   @Inject
   ReplayCommitLogsToSqlAction() {}
 
@@ -104,44 +121,91 @@ public class ReplayCommitLogsToSqlAction implements Runnable {
       return;
     }
     try {
-      replayFiles();
-      response.setStatus(HttpServletResponse.SC_OK);
-      logger.atInfo().log("ReplayCommitLogsToSqlAction completed successfully.");
+      logger.atInfo().log("Beginning replay of commit logs.");
+      ImmutableList<GcsFileMetadata> commitLogFiles = getFilesToReplay();
+      if (dryRun) {
+        response.setStatus(HttpServletResponse.SC_OK);
+        ImmutableList<String> filenames =
+            commitLogFiles.stream()
+                .limit(10)
+                .map(file -> file.getFilename().getObjectName())
+                .collect(toImmutableList());
+        String dryRunMessage =
+            "Running in dry-run mode; would have processed %d files. They are (limit 10):\n"
+                + Joiner.on('\n').join(filenames);
+        response.setPayload(dryRunMessage);
+        logger.atInfo().log(dryRunMessage);
+      } else {
+        replayFiles(commitLogFiles);
+        response.setStatus(HttpServletResponse.SC_OK);
+        String message = "ReplayCommitLogsToSqlAction completed successfully.";
+        response.setPayload(message);
+        logger.atInfo().log(message);
+      }
+    } catch (Throwable t) {
+      String message = "Errored out replaying files.";
+      logger.atSevere().withCause(t).log(message);
+      response.setStatus(HttpServletResponse.SC_INTERNAL_SERVER_ERROR);
+      response.setPayload(message);
     } finally {
       lock.ifPresent(Lock::release);
     }
   }
 
-  private void replayFiles() {
+  private ImmutableList<GcsFileMetadata> getFilesToReplay() {
     // Start at the first millisecond we haven't seen yet
     DateTime fromTime = jpaTm().transact(() -> SqlReplayCheckpoint.get().plusMillis(1));
+    logger.atInfo().log("Starting replay from: %s.", fromTime);
     // 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);
+        diffLister.listDiffFiles(gcsBucket, fromTime, /* current time */ null);
+    logger.atInfo().log("Found %d new commit log files to process.", commitLogFiles.size());
+    return commitLogFiles;
+  }
+
+  private void replayFiles(ImmutableList<GcsFileMetadata> commitLogFiles) {
+    DateTime replayTimeoutTime = clock.nowUtc().plus(REPLAY_TIMEOUT_DURATION);
+    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) {
+    logger.atInfo().log(
+        "Processing commit log file %s of size %d B.",
+        metadata.getFilename(), metadata.getLength());
     try (InputStream input =
         Channels.newInputStream(
             gcsService.openPrefetchingReadChannel(metadata.getFilename(), 0, BLOCK_SIZE))) {
       // Load and process the Datastore transactions one at a time
       ImmutableList<ImmutableList<VersionedEntity>> allTransactions =
           CommitLogImports.loadEntitiesByTransaction(input);
+      logger.atInfo().log(
+          "Replaying %d transactions from commit log file %s.",
+          allTransactions.size(), metadata.getFilename());
       allTransactions.forEach(this::replayTransaction);
       // if we succeeded, set the last-seen time
       DateTime checkpoint =
           DateTime.parse(
               metadata.getFilename().getObjectName().substring(DIFF_FILE_PREFIX.length()));
       SqlReplayCheckpoint.set(checkpoint);
-      logger.atInfo().log("Replayed %d transactions from commit log file.", allTransactions.size());
+      logger.atInfo().log(
+          "Replayed %d transactions from commit log file %s.",
+          allTransactions.size(), metadata.getFilename());
     } catch (IOException e) {
-      throw new RuntimeException(e);
+      throw new RuntimeException(
+          "Errored out while replaying commit log file " + metadata.getFilename(), e);
     }
   }
 

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

@@ -26,6 +26,7 @@ import com.google.appengine.api.datastore.EntityTranslator;
 import com.google.appengine.tools.cloudstorage.GcsFileMetadata;
 import com.google.appengine.tools.cloudstorage.GcsService;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Lists;
 import com.google.common.collect.PeekingIterator;
 import com.google.common.collect.Streams;
@@ -33,6 +34,7 @@ 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.Config;
 import google.registry.config.RegistryEnvironment;
 import google.registry.model.ImmutableObject;
 import google.registry.model.ofy.CommitLogBucket;
@@ -50,6 +52,7 @@ import java.nio.channels.Channels;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Stream;
 import javax.inject.Inject;
@@ -72,27 +75,41 @@ public class RestoreCommitLogsAction implements Runnable {
   static final String DRY_RUN_PARAM = "dryRun";
   static final String FROM_TIME_PARAM = "fromTime";
   static final String TO_TIME_PARAM = "toTime";
+  static final String BUCKET_OVERRIDE_PARAM = "gcsBucket";
+
+  private static final ImmutableSet<RegistryEnvironment> FORBIDDEN_ENVIRONMENTS =
+      ImmutableSet.of(RegistryEnvironment.PRODUCTION, RegistryEnvironment.SANDBOX);
 
   @Inject GcsService gcsService;
   @Inject @Parameter(DRY_RUN_PARAM) boolean dryRun;
   @Inject @Parameter(FROM_TIME_PARAM) DateTime fromTime;
   @Inject @Parameter(TO_TIME_PARAM) DateTime toTime;
+
+  @Inject
+  @Parameter(BUCKET_OVERRIDE_PARAM)
+  Optional<String> gcsBucketOverride;
+
   @Inject DatastoreService datastoreService;
   @Inject GcsDiffFileLister diffLister;
+
+  @Inject
+  @Config("commitLogGcsBucket")
+  String defaultGcsBucket;
+
   @Inject Retrier retrier;
   @Inject RestoreCommitLogsAction() {}
 
   @Override
   public void run() {
     checkArgument(
-        RegistryEnvironment.get() == RegistryEnvironment.ALPHA
-            || RegistryEnvironment.get() == RegistryEnvironment.CRASH
-            || RegistryEnvironment.get() == RegistryEnvironment.UNITTEST,
-        "DO NOT RUN ANYWHERE ELSE EXCEPT ALPHA, CRASH OR TESTS.");
+        !FORBIDDEN_ENVIRONMENTS.contains(RegistryEnvironment.get()),
+        "DO NOT RUN IN PRODUCTION OR SANDBOX.");
     if (dryRun) {
       logger.atInfo().log("Running in dryRun mode");
     }
-    List<GcsFileMetadata> diffFiles = diffLister.listDiffFiles(fromTime, toTime);
+    String gcsBucket = gcsBucketOverride.orElse(defaultGcsBucket);
+    logger.atInfo().log("Restoring from %s.", gcsBucket);
+    List<GcsFileMetadata> diffFiles = diffLister.listDiffFiles(gcsBucket, fromTime, toTime);
     if (diffFiles.isEmpty()) {
       logger.atInfo().log("Nothing to restore");
       return;
@@ -104,7 +121,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/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/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/rde/RdePipeline.java

@@ -0,0 +1,260 @@
+// 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.beam.rde;
+
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
+import static google.registry.model.EppResourceUtils.loadAtPointInTimeAsync;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableSetMultimap;
+import com.google.common.collect.Maps;
+import com.google.common.collect.Sets;
+import com.google.common.io.BaseEncoding;
+import google.registry.beam.common.RegistryJpaIO;
+import google.registry.model.EppResource;
+import google.registry.model.contact.ContactResource;
+import google.registry.model.domain.DomainBase;
+import google.registry.model.host.HostResource;
+import google.registry.model.rde.RdeMode;
+import google.registry.model.registrar.Registrar;
+import google.registry.model.registrar.Registrar.Type;
+import google.registry.persistence.VKey;
+import google.registry.rde.DepositFragment;
+import google.registry.rde.PendingDeposit;
+import google.registry.rde.PendingDeposit.PendingDepositCoder;
+import google.registry.rde.RdeFragmenter;
+import google.registry.rde.RdeMarshaller;
+import google.registry.xml.ValidationMode;
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Optional;
+import java.util.function.Supplier;
+import javax.persistence.Entity;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.PipelineResult;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.coders.SerializableCoder;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.transforms.FlatMapElements;
+import org.apache.beam.sdk.transforms.Flatten;
+import org.apache.beam.sdk.transforms.Reshuffle;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionList;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.sdk.values.TypeDescriptors;
+import org.joda.time.DateTime;
+
+/**
+ * Definition of a Dataflow Flex template, which generates RDE/BRDA deposits.
+ *
+ * <p>To stage this template locally, run the {@code stage_beam_pipeline.sh} shell script.
+ *
+ * <p>Then, you can run the staged template via the API client library, gCloud or a raw REST call.
+ *
+ * @see <a href="https://cloud.google.com/dataflow/docs/guides/templates/using-flex-templates">Using
+ *     Flex Templates</a>
+ */
+public class RdePipeline implements Serializable {
+
+  private final RdeMarshaller marshaller;
+  private final ImmutableSetMultimap<String, PendingDeposit> pendings;
+
+  // Registrars to be excluded from data escrow. Not including the sandbox-only OTE type so that
+  // if sneaks into production we would get an extra signal.
+  private static final ImmutableSet<Type> IGNORED_REGISTRAR_TYPES =
+      Sets.immutableEnumSet(Registrar.Type.MONITORING, Registrar.Type.TEST);
+
+  private static final String EPP_RESOURCE_QUERY =
+      "SELECT id FROM %entity% "
+          + "WHERE COALESCE(creationClientId, '') NOT LIKE 'prober-%' "
+          + "AND COALESCE(currentSponsorClientId, '') NOT LIKE 'prober-%' "
+          + "AND COALESCE(lastEppUpdateClientId, '') NOT LIKE 'prober-%'";
+
+  public static String createEppResourceQuery(Class<? extends EppResource> clazz) {
+    return EPP_RESOURCE_QUERY.replace("%entity%", clazz.getAnnotation(Entity.class).name())
+        + (clazz.equals(DomainBase.class) ? " AND tld in (:tlds)" : "");
+  }
+
+  RdePipeline(RdePipelineOptions options) throws IOException, ClassNotFoundException {
+    this.marshaller = new RdeMarshaller(ValidationMode.valueOf(options.getValidationMode()));
+    this.pendings = decodePendings(options.getPendings());
+  }
+
+  @VisibleForTesting
+  PipelineResult run(Pipeline pipeline) {
+    createFragments(pipeline);
+    return pipeline.run();
+  }
+
+  PipelineResult run() {
+    return run(Pipeline.create());
+  }
+
+  PCollection<KV<PendingDeposit, DepositFragment>> createFragments(Pipeline pipeline) {
+    PCollection<KV<PendingDeposit, DepositFragment>> fragments =
+        PCollectionList.of(processRegistrars(pipeline))
+            .and(processNonRegistrarEntities(pipeline, DomainBase.class))
+            .and(processNonRegistrarEntities(pipeline, ContactResource.class))
+            .and(processNonRegistrarEntities(pipeline, HostResource.class))
+            .apply(Flatten.pCollections())
+            .setCoder(
+                KvCoder.of(PendingDepositCoder.of(), SerializableCoder.of(DepositFragment.class)));
+    return fragments;
+  }
+
+  PCollection<KV<PendingDeposit, DepositFragment>> processRegistrars(Pipeline pipeline) {
+    return pipeline
+        .apply(
+            "Read all production Registrar entities",
+            RegistryJpaIO.read(
+                "SELECT clientIdentifier FROM Registrar WHERE type NOT IN (:types)",
+                ImmutableMap.of("types", IGNORED_REGISTRAR_TYPES),
+                String.class,
+                // TODO: consider adding coders for entities and pass them directly instead of using
+                // VKeys.
+                id -> VKey.createSql(Registrar.class, id)))
+        .apply(
+            "Marshal Registrar into DepositFragment",
+            FlatMapElements.into(
+                    TypeDescriptors.kvs(
+                        TypeDescriptor.of(PendingDeposit.class),
+                        TypeDescriptor.of(DepositFragment.class)))
+                .via(
+                    (VKey<Registrar> key) -> {
+                      Registrar registrar = jpaTm().transact(() -> jpaTm().loadByKey(key));
+                      DepositFragment fragment = marshaller.marshalRegistrar(registrar);
+                      return pendings.values().stream()
+                          .map(pending -> KV.of(pending, fragment))
+                          .collect(toImmutableSet());
+                    }));
+  }
+
+  @SuppressWarnings("deprecation") // Reshuffle is still recommended by Dataflow.
+  <T extends EppResource>
+      PCollection<KV<PendingDeposit, DepositFragment>> processNonRegistrarEntities(
+          Pipeline pipeline, Class<T> clazz) {
+    return createInputs(pipeline, clazz)
+        .apply("Marshal " + clazz.getSimpleName() + " into DepositFragment", mapToFragments(clazz))
+        .setCoder(KvCoder.of(PendingDepositCoder.of(), SerializableCoder.of(DepositFragment.class)))
+        .apply(
+            "Reshuffle KV<PendingDeposit, DepositFragment> of "
+                + clazz.getSimpleName()
+                + " to prevent fusion",
+            Reshuffle.of());
+  }
+
+  <T extends EppResource> PCollection<VKey<T>> createInputs(Pipeline pipeline, Class<T> clazz) {
+    return pipeline.apply(
+        "Read all production " + clazz.getSimpleName() + " entities",
+        RegistryJpaIO.read(
+            createEppResourceQuery(clazz),
+            clazz.equals(DomainBase.class)
+                ? ImmutableMap.of("tlds", pendings.keySet())
+                : ImmutableMap.of(),
+            String.class,
+            // TODO: consider adding coders for entities and pass them directly instead of using
+            // VKeys.
+            x -> VKey.create(clazz, x)));
+  }
+
+  <T extends EppResource>
+      FlatMapElements<VKey<T>, KV<PendingDeposit, DepositFragment>> mapToFragments(Class<T> clazz) {
+    return FlatMapElements.into(
+            TypeDescriptors.kvs(
+                TypeDescriptor.of(PendingDeposit.class), TypeDescriptor.of(DepositFragment.class)))
+        .via(
+            (VKey<T> key) -> {
+              T resource = jpaTm().transact(() -> jpaTm().loadByKey(key));
+              // The set of all TLDs to which this resource should be emitted.
+              ImmutableSet<String> tlds =
+                  clazz.equals(DomainBase.class)
+                      ? ImmutableSet.of(((DomainBase) resource).getTld())
+                      : pendings.keySet();
+              // Get the set of all point-in-time watermarks we need, to minimize rewinding.
+              ImmutableSet<DateTime> dates =
+                  tlds.stream()
+                      .map(pendings::get)
+                      .flatMap(ImmutableSet::stream)
+                      .map(PendingDeposit::watermark)
+                      .collect(toImmutableSet());
+              // Launch asynchronous fetches of point-in-time representations of resource.
+              ImmutableMap<DateTime, Supplier<EppResource>> resourceAtTimes =
+                  ImmutableMap.copyOf(
+                      Maps.asMap(dates, input -> loadAtPointInTimeAsync(resource, input)));
+              // Convert resource to an XML fragment for each watermark/mode pair lazily and cache
+              // the result.
+              RdeFragmenter fragmenter = new RdeFragmenter(resourceAtTimes, marshaller);
+              List<KV<PendingDeposit, DepositFragment>> results = new ArrayList<>();
+              for (String tld : tlds) {
+                for (PendingDeposit pending : pendings.get(tld)) {
+                  // Hosts and contacts don't get included in BRDA deposits.
+                  if (pending.mode() == RdeMode.THIN && !clazz.equals(DomainBase.class)) {
+                    continue;
+                  }
+                  Optional<DepositFragment> fragment =
+                      fragmenter.marshal(pending.watermark(), pending.mode());
+                  fragment.ifPresent(
+                      depositFragment -> results.add(KV.of(pending, depositFragment)));
+                }
+              }
+              return results;
+            });
+  }
+
+  /**
+   * Decodes the pipeline option extracted from the URL parameter sent by the pipeline launcher to
+   * the original TLD to pending deposit map.
+   */
+  @SuppressWarnings("unchecked")
+  static ImmutableSetMultimap<String, PendingDeposit> decodePendings(String encodedPending)
+      throws IOException, ClassNotFoundException {
+    try (ObjectInputStream ois =
+        new ObjectInputStream(
+            new ByteArrayInputStream(
+                BaseEncoding.base64Url().omitPadding().decode(encodedPending)))) {
+      return (ImmutableSetMultimap<String, PendingDeposit>) ois.readObject();
+    }
+  }
+
+  /**
+   * Encodes the TLD to pending deposit map in an URL safe string that is sent to the pipeline
+   * worker by the pipeline launcher as a pipeline option.
+   */
+  public static String encodePendings(ImmutableSetMultimap<String, PendingDeposit> pendings)
+      throws IOException {
+    try (ByteArrayOutputStream baos = new ByteArrayOutputStream()) {
+      ObjectOutputStream oos = new ObjectOutputStream(baos);
+      oos.writeObject(pendings);
+      oos.flush();
+      return BaseEncoding.base64Url().omitPadding().encode(baos.toByteArray());
+    }
+  }
+
+  public static void main(String[] args) throws IOException, ClassNotFoundException {
+    PipelineOptionsFactory.register(RdePipelineOptions.class);
+    RdePipelineOptions options = PipelineOptionsFactory.fromArgs(args).as(RdePipelineOptions.class);
+    new RdePipeline(options).run();
+  }
+}

core/src/main/java/google/registry/beam/rde/RdePipelineOptions.java

@@ -0,0 +1,32 @@
+// 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.beam.rde;
+
+import google.registry.beam.common.RegistryPipelineOptions;
+import org.apache.beam.sdk.options.Description;
+
+/** Custom options for running the spec11 pipeline. */
+public interface RdePipelineOptions extends RegistryPipelineOptions {
+
+  @Description("The Base64-encoded serialized map of TLDs to PendingDeposit.")
+  String getPendings();
+
+  void setPendings(String value);
+
+  @Description("The validation mode (LENIENT|STRICT) that the RDE marshaller uses.")
+  String getValidationMode();
+
+  void setValidationMode(String value);
+}

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

@@ -116,6 +116,7 @@ public class Spec11Pipeline implements Serializable {
             "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);

core/src/main/java/google/registry/dns/writer/clouddns/CloudDnsWriter.java

@@ -284,9 +284,10 @@ public class CloudDnsWriter extends BaseDnsWriter {
         .filter(hostName -> hostName.endsWith("." + domainName) && !hostName.equals(domainName));
   }
 
-  /** Mutate the zone with the provided {@code desiredRecords}. */
+  /** Mutate the zone with the provided map of hostnames to desired DNS records. */
   @VisibleForTesting
   void mutateZone(ImmutableMap<String, ImmutableSet<ResourceRecordSet>> desiredRecords) {
+    logger.atInfo().log("Updating DNS records for hostname(s) %s.", desiredRecords.keySet());
     // Fetch all existing records for names that this writer is trying to modify
     ImmutableSet.Builder<ResourceRecordSet> flattenedExistingRecords = new ImmutableSet.Builder<>();
 
@@ -313,7 +314,12 @@ public class CloudDnsWriter extends BaseDnsWriter {
     desiredRecords.values().forEach(flattenedDesiredRecords::addAll);
 
     // Delete all existing records and add back the desired records
-    updateResourceRecords(flattenedDesiredRecords.build(), flattenedExistingRecords.build());
+    try {
+      updateResourceRecords(flattenedDesiredRecords.build(), flattenedExistingRecords.build());
+    } catch (IOException e) {
+      throw new RuntimeException(
+          "Error updating DNS records for hostname(s) " + desiredRecords.keySet(), e);
+    }
   }
 
   /**
@@ -363,10 +369,12 @@ public class CloudDnsWriter extends BaseDnsWriter {
    *
    * @throws ZoneStateException if the operation could not be completely successfully because the
    *     records to delete do not exist, already exist or have been modified with different
-   *     attributes since being queried.
+   *     attributes since being queried. These errors will be retried.
+   * @throws IOException on non-retryable API errors, e.g. invalid request.
    */
   private void updateResourceRecords(
-      ImmutableSet<ResourceRecordSet> additions, ImmutableSet<ResourceRecordSet> deletions) {
+      ImmutableSet<ResourceRecordSet> additions, ImmutableSet<ResourceRecordSet> deletions)
+      throws IOException {
     // Find records that are both in additions and deletions, so we can remove them from both before
     // requesting the change. This is mostly for optimization reasons - not doing so doesn't affect
     // the result.
@@ -392,17 +400,15 @@ public class CloudDnsWriter extends BaseDnsWriter {
       GoogleJsonError err = e.getDetails();
       // We did something really wrong here, just give up and re-throw
       if (err == null || err.getErrors().size() > 1) {
-        throw new RuntimeException(e);
+        throw e;
       }
       String errorReason = err.getErrors().get(0).getReason();
 
       if (RETRYABLE_EXCEPTION_REASONS.contains(errorReason)) {
         throw new ZoneStateException(errorReason);
       } else {
-        throw new RuntimeException(e);
+        throw e;
       }
-    } catch (IOException e) {
-      throw new RuntimeException(e);
     }
   }
 

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

@@ -230,7 +230,7 @@
   </cron>
 
   <cron>
-    <url><![CDATA[/_dr/cron/fanout?queue=replay-commit-logs-to-sql&endpoint=/_dr/task/replayCommitLogsToSql&runInEmpty]]></url>
+    <url><![CDATA[/_dr/cron/fanout?queue=replay-commit-logs-to-sql&endpoint=/_dr/task/replayCommitLogsToSql&runInEmpty&dryRun=true]]></url>
     <description>
       Replays recent commit logs from Datastore to the SQL secondary backend.
     </description>

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

@@ -113,7 +113,7 @@ public class ExportPremiumTermsAction implements Runnable {
           "Skipping premium terms export for TLD %s because Drive folder isn't specified", tld);
       return Optional.of("Skipping export because no Drive folder is associated with this TLD");
     }
-    if (!registry.getPremiumList().isPresent()) {
+    if (!registry.getPremiumListName().isPresent()) {
       logger.atInfo().log("No premium terms to export for TLD %s", tld);
       return Optional.of("No premium lists configured");
     }
@@ -137,8 +137,8 @@ public class ExportPremiumTermsAction implements Runnable {
   }
 
   private String getFormattedPremiumTerms(Registry registry) {
-    checkState(registry.getPremiumList().isPresent(), "%s does not have a premium list", tld);
-    String premiumListName = registry.getPremiumList().get().getName();
+    checkState(registry.getPremiumListName().isPresent(), "%s does not have a premium list", tld);
+    String premiumListName = registry.getPremiumListName().get();
     checkState(
         PremiumListDao.getLatestRevision(premiumListName).isPresent(),
         "Could not load premium list for " + tld);

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

@@ -63,7 +63,7 @@ public class ExportReservedTermsAction implements Runnable {
     try {
       Registry registry = Registry.get(tld);
       String resultMsg;
-      if (registry.getReservedLists().isEmpty() && isNullOrEmpty(registry.getDriveFolderId())) {
+      if (registry.getReservedListNames().isEmpty() && isNullOrEmpty(registry.getDriveFolderId())) {
         resultMsg = "No reserved lists configured";
         logger.atInfo().log("No reserved terms to export for TLD %s", tld);
       } else if (registry.getDriveFolderId() == null) {

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

@@ -15,11 +15,11 @@
 package google.registry.export;
 
 import com.google.common.base.Joiner;
-import com.googlecode.objectify.Key;
 import google.registry.config.RegistryConfig.Config;
 import google.registry.model.registry.Registry;
 import google.registry.model.registry.label.ReservedList;
 import google.registry.model.registry.label.ReservedList.ReservedListEntry;
+import google.registry.model.registry.label.ReservedListDao;
 import java.util.Set;
 import java.util.TreeSet;
 import javax.inject.Inject;
@@ -39,8 +39,13 @@ public final class ExportUtils {
   public String exportReservedTerms(Registry registry) {
     StringBuilder termsBuilder = new StringBuilder(reservedTermsExportDisclaimer).append("\n");
     Set<String> reservedTerms = new TreeSet<>();
-    for (Key<ReservedList> key : registry.getReservedLists()) {
-      ReservedList reservedList = ReservedList.load(key).get();
+    for (String reservedListName : registry.getReservedListNames()) {
+      ReservedList reservedList =
+          ReservedListDao.getLatestRevision(reservedListName)
+              .orElseThrow(
+                  () ->
+                      new IllegalStateException(
+                          String.format("Reserved list %s does not exist", reservedListName)));
       if (reservedList.getShouldPublish()) {
         for (ReservedListEntry entry : reservedList.getReservedListEntries().values()) {
           reservedTerms.add(entry.getLabel());

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;
@@ -331,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(

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

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/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.auditedOfy;
+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,25 +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(
-                      auditedOfy()
-                          .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/EppResourceUtils.java

@@ -398,7 +398,7 @@ public final class EppResourceUtils {
                 .orElse(null);
     if (resourceAtPointInTime == null) {
       logger.atSevere().log(
-          "Couldn't load resource at % for key %s, falling back to resource %s.",
+          "Couldn't load resource at %s for key %s, falling back to resource %s.",
           timestamp, resource.createVKey(), resource);
       return resource;
     }

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/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

@@ -110,7 +110,8 @@ public class ContactHistory extends HistoryEntry implements SqlEntity {
 
   @Override
   public Optional<? extends EppResource> getResourceAtPointInTime() {
-    return getContactBase();
+    return getContactBase()
+        .map(contactBase -> new ContactResource.Builder().copyFrom(contactBase).build());
   }
 
   @PostLoad

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

@@ -79,5 +79,26 @@ public class ContactResource extends ContactBase
     private Builder(ContactResource instance) {
       super(instance);
     }
+
+    public Builder copyFrom(ContactBase contactBase) {
+      return this.setAuthInfo(contactBase.getAuthInfo())
+          .setContactId(contactBase.getContactId())
+          .setCreationClientId(contactBase.getCreationClientId())
+          .setCreationTime(contactBase.getCreationTime())
+          .setDeletionTime(contactBase.getDeletionTime())
+          .setDisclose(contactBase.getDisclose())
+          .setEmailAddress(contactBase.getEmailAddress())
+          .setFaxNumber(contactBase.getFaxNumber())
+          .setInternationalizedPostalInfo(contactBase.getInternationalizedPostalInfo())
+          .setLastTransferTime(contactBase.getLastTransferTime())
+          .setLastEppUpdateClientId(contactBase.getLastEppUpdateClientId())
+          .setLastEppUpdateTime(contactBase.getLastEppUpdateTime())
+          .setLocalizedPostalInfo(contactBase.getLocalizedPostalInfo())
+          .setPersistedCurrentSponsorClientId(contactBase.getPersistedCurrentSponsorClientId())
+          .setRepoId(contactBase.getRepoId())
+          .setStatusValues(contactBase.getStatusValues())
+          .setTransferData(contactBase.getTransferData())
+          .setVoiceNumber(contactBase.getVoiceNumber());
+    }
   }
 }

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

@@ -181,5 +181,34 @@ public class DomainBase extends DomainContent
     Builder(DomainBase instance) {
       super(instance);
     }
+
+    public Builder copyFrom(DomainContent domainContent) {
+      return this.setAuthInfo(domainContent.getAuthInfo())
+          .setAutorenewPollMessage(domainContent.getAutorenewPollMessage())
+          .setAutorenewBillingEvent(domainContent.getAutorenewBillingEvent())
+          .setAutorenewEndTime(domainContent.getAutorenewEndTime())
+          .setContacts(domainContent.getContacts())
+          .setCreationClientId(domainContent.getCreationClientId())
+          .setCreationTime(domainContent.getCreationTime())
+          .setDomainName(domainContent.getDomainName())
+          .setDeletePollMessage(domainContent.getDeletePollMessage())
+          .setDsData(domainContent.getDsData())
+          .setDeletionTime(domainContent.getDeletionTime())
+          .setGracePeriods(domainContent.getGracePeriods())
+          .setIdnTableName(domainContent.getIdnTableName())
+          .setLastTransferTime(domainContent.getLastTransferTime())
+          .setLaunchNotice(domainContent.getLaunchNotice())
+          .setLastEppUpdateClientId(domainContent.getLastEppUpdateClientId())
+          .setLastEppUpdateTime(domainContent.getLastEppUpdateTime())
+          .setNameservers(domainContent.getNameservers())
+          .setPersistedCurrentSponsorClientId(domainContent.getPersistedCurrentSponsorClientId())
+          .setRegistrant(domainContent.getRegistrant())
+          .setRegistrationExpirationTime(domainContent.getRegistrationExpirationTime())
+          .setRepoId(domainContent.getRepoId())
+          .setSmdId(domainContent.getSmdId())
+          .setSubordinateHosts(domainContent.getSubordinateHosts())
+          .setStatusValues(domainContent.getStatusValues())
+          .setTransferData(domainContent.getTransferData());
+    }
   }
 }

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

@@ -252,11 +252,18 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
 
   @Override
   public Optional<? extends EppResource> getResourceAtPointInTime() {
-    return getDomainContent();
+    return getDomainContent()
+        .map(domainContent -> new DomainBase.Builder().copyFrom(domainContent).build());
   }
 
   @PostLoad
   void postLoad() {
+    // TODO(b/188044616): Determine why Eager loading doesn't work here.
+    Hibernate.initialize(domainTransactionRecords);
+    Hibernate.initialize(nsHosts);
+    Hibernate.initialize(dsDataHistories);
+    Hibernate.initialize(gracePeriodHistories);
+
     if (domainContent != null) {
       domainContent.nsHosts = nullToEmptyImmutableCopy(nsHosts);
       domainContent.gracePeriods =
@@ -278,12 +285,6 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
         }
       }
     }
-
-    // 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);
@@ -210,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/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/HostBase.java

@@ -137,7 +137,7 @@ public class HostBase extends EppResource {
   }
 
   @Override
-  public Builder asBuilder() {
+  public Builder<? extends HostBase, ?> asBuilder() {
     return new Builder<>(clone(this));
   }
 

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

@@ -111,7 +111,7 @@ public class HostHistory extends HistoryEntry implements SqlEntity {
 
   @Override
   public Optional<? extends EppResource> getResourceAtPointInTime() {
-    return getHostBase();
+    return getHostBase().map(hostBase -> new HostResource.Builder().copyFrom(hostBase).build());
   }
 
   @PostLoad

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

@@ -63,5 +63,21 @@ public class HostResource extends HostBase
     private Builder(HostResource instance) {
       super(instance);
     }
+
+    public Builder copyFrom(HostBase hostBase) {
+      return this.setCreationClientId(hostBase.getCreationClientId())
+          .setCreationTime(hostBase.getCreationTime())
+          .setDeletionTime(hostBase.getDeletionTime())
+          .setHostName(hostBase.getHostName())
+          .setInetAddresses(hostBase.getInetAddresses())
+          .setLastTransferTime(hostBase.getLastTransferTime())
+          .setLastSuperordinateChange(hostBase.getLastSuperordinateChange())
+          .setLastEppUpdateClientId(hostBase.getLastEppUpdateClientId())
+          .setLastEppUpdateTime(hostBase.getLastEppUpdateTime())
+          .setPersistedCurrentSponsorClientId(hostBase.getPersistedCurrentSponsorClientId())
+          .setRepoId(hostBase.getRepoId())
+          .setSuperordinateDomain(hostBase.getSuperordinateDomain())
+          .setStatusValues(hostBase.getStatusValues());
+    }
   }
 }

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

@@ -279,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();

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,7 +46,6 @@ 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 #auditedOfy()}, should be used to access all persistable
@@ -57,12 +54,6 @@ import java.util.concurrent.atomic.AtomicLong;
  */
 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/pricing/StaticPremiumListPricingEngine.java

@@ -39,9 +39,7 @@ public final class StaticPremiumListPricingEngine implements PremiumPricingEngin
     String label = InternetDomainName.from(fullyQualifiedDomainName).parts().get(0);
     Registry registry = Registry.get(checkNotNull(tld, "tld"));
     Optional<Money> premiumPrice =
-        registry
-            .getPremiumList()
-            .flatMap(listKey -> PremiumListDao.getPremiumPrice(listKey.getName(), label));
+        registry.getPremiumListName().flatMap(pl -> PremiumListDao.getPremiumPrice(pl, label));
     return DomainPrices.create(
         premiumPrice.isPresent(),
         premiumPrice.orElse(registry.getStandardCreateCost()),

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

@@ -387,18 +387,32 @@ public class Registry extends ImmutableObject implements Buildable, DatastoreAnd
   @Column(nullable = false)
   CreateAutoTimestamp creationTime = CreateAutoTimestamp.create(null);
 
-  /** The set of reserved lists that are applicable to this registry. */
+  /** The set of reserved list names that are applicable to this registry. */
   @Column(name = "reserved_list_names")
-  Set<Key<ReservedList>> reservedLists;
+  Set<String> reservedListNames;
 
-  /** Retrieves an ImmutableSet of all ReservedLists associated with this tld. */
-  public ImmutableSet<Key<ReservedList>> getReservedLists() {
-    return nullToEmptyImmutableCopy(reservedLists);
+  /**
+   * 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<String> getReservedListNames() {
+    return nullToEmptyImmutableCopy(reservedListNames);
   }
 
-  /** The static {@link PremiumList} for this TLD, if there is one. */
+  /**
+   * 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)
@@ -606,8 +620,8 @@ public class Registry extends ImmutableObject implements Buildable, DatastoreAnd
     return anchorTenantAddGracePeriodLength;
   }
 
-  public Optional<Key<PremiumList>> getPremiumList() {
-    return Optional.ofNullable(premiumList);
+  public Optional<String> getPremiumListName() {
+    return Optional.ofNullable(premiumListName);
   }
 
   public CurrencyUnit getCurrency() {
@@ -878,22 +892,16 @@ 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) {
-      getInstance().premiumList = (premiumList == null) ? null : Key.create(premiumList);
-      return this;
-    }
-
-    @VisibleForTesting
-    public Builder setPremiumListKey(@Nullable Key<PremiumList> premiumList) {
-      getInstance().premiumList = premiumList;
+    public Builder setPremiumList(@Nullable PremiumList premiumList) {
+      getInstance().premiumListName = (premiumList == null) ? null : premiumList.getName();
       return this;
     }
 

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

@@ -171,6 +171,7 @@ public abstract class BaseDomainLabelList<T extends Comparable<?>, R extends Dom
         .collect(toImmutableSet());
   }
 
+  // TODO(b/193043636): Refactor this class to no longer use key references
   protected abstract boolean refersToKey(
       Registry registry, Key<? extends BaseDomainLabelList<?, ?>> key);
 

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.
    *
@@ -270,7 +283,7 @@ public final class PremiumList extends BaseDomainLabelList<Money, PremiumList.Pr
 
   @Override
   public boolean refersToKey(Registry registry, Key<? extends BaseDomainLabelList<?, ?>> key) {
-    return Objects.equals(registry.getPremiumList().orElse(null), key);
+    return Objects.equals(registry.getPremiumListName().orElse(null), key.getName());
   }
 
   @Override

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

@@ -20,7 +20,6 @@ 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;
@@ -187,7 +186,7 @@ public final class ReservedList
 
   @Override
   protected boolean refersToKey(Registry registry, Key<? extends BaseDomainLabelList<?, ?>> key) {
-    return registry.getReservedLists().contains(key);
+    return registry.getReservedListNames().contains(key.getName());
   }
 
   /** Determines whether the ReservedList is in use on any Registry */
@@ -236,11 +235,6 @@ public final class ReservedList
     return getFromCache(listName, cache);
   }
 
-  /** Loads a ReservedList from its Objectify key. */
-  public static Optional<ReservedList> load(Key<ReservedList> key) {
-    return get(key.getName());
-  }
-
   /**
    * Queries the set of all reserved lists associated with the specified TLD and returns the
    * reservation types of the label.
@@ -270,7 +264,7 @@ public final class ReservedList
         new ImmutableSet.Builder<>();
 
     // Loop through all reservation lists and add each of them.
-    for (ReservedList rl : loadReservedLists(registry.getReservedLists())) {
+    for (ReservedList rl : loadReservedLists(registry.getReservedListNames())) {
       if (rl.getReservedListEntries().containsKey(label)) {
         ReservedListEntry entry = rl.getReservedListEntries().get(label);
         entriesBuilder.add(entry);
@@ -285,17 +279,15 @@ public final class ReservedList
   }
 
   private static ImmutableSet<ReservedList> loadReservedLists(
-      ImmutableSet<Key<ReservedList>> reservedListKeys) {
-    return reservedListKeys
-        .stream()
+      ImmutableSet<String> reservedListNames) {
+    return reservedListNames.stream()
         .map(
-            (listKey) -> {
+            (listName) -> {
               try {
-                return cache.get(listKey.getName());
+                return cache.get(listName);
               } catch (ExecutionException e) {
                 throw new UncheckedExecutionException(
-                    String.format(
-                        "Could not load the reserved list '%s' from the cache", listKey.getName()),
+                    String.format("Could not load the reserved list '%s' from the cache", listName),
                     e);
               }
             })

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

@@ -17,7 +17,7 @@ package google.registry.model.tmch;
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
-import static google.registry.model.ofy.ObjectifyService.allocateId;
+import static google.registry.model.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;
@@ -36,6 +36,7 @@ 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;
@@ -45,6 +46,7 @@ import javax.persistence.Column;
 import javax.persistence.GeneratedValue;
 import javax.persistence.GenerationType;
 import javax.persistence.PostPersist;
+import javax.persistence.PostUpdate;
 import javax.persistence.PreRemove;
 import javax.persistence.Table;
 import javax.persistence.Transient;

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/persistence/BillingVKey.java

@@ -21,6 +21,7 @@ import google.registry.model.billing.BillingEvent.OneTime;
 import google.registry.model.billing.BillingEvent.Recurring;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.reporting.HistoryEntry;
+import java.io.Serializable;
 import javax.annotation.Nullable;
 import javax.persistence.AttributeOverride;
 import javax.persistence.AttributeOverrides;
@@ -46,7 +47,7 @@ public abstract class BillingVKey<K> extends EppHistoryVKey<K, DomainBase> {
   }
 
   @Override
-  public Object createSqlKey() {
+  public Serializable createSqlKey() {
     return billingId;
   }
 

core/src/main/java/google/registry/persistence/DomainHistoryVKey.java

@@ -21,6 +21,7 @@ import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.DomainHistory;
 import google.registry.model.domain.DomainHistory.DomainHistoryId;
 import google.registry.model.reporting.HistoryEntry;
+import java.io.Serializable;
 import javax.persistence.Embeddable;
 
 /** {@link VKey} for {@link HistoryEntry} which parent is {@link DomainBase}. */
@@ -35,7 +36,7 @@ public class DomainHistoryVKey extends EppHistoryVKey<HistoryEntry, DomainBase>
   }
 
   @Override
-  public Object createSqlKey() {
+  public Serializable createSqlKey() {
     return new DomainHistoryId(repoId, historyRevisionId);
   }
 

core/src/main/java/google/registry/persistence/EppHistoryVKey.java

@@ -101,7 +101,7 @@ public abstract class EppHistoryVKey<K, E extends EppResource> extends Immutable
     return VKey.create(vKeyType, createSqlKey(), createOfyKey());
   }
 
-  public abstract Object createSqlKey();
+  public abstract Serializable createSqlKey();
 
   public abstract Key<K> createOfyKey();
 }

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

@@ -37,7 +37,7 @@ public class VKey<T> extends ImmutableObject implements Serializable {
   private static final long serialVersionUID = -5291472863840231240L;
 
   // The SQL key for the referenced entity.
-  Object sqlKey;
+  Serializable sqlKey;
 
   // The objectify key for the referenced entity.
   Key<T> ofyKey;
@@ -46,7 +46,7 @@ public class VKey<T> extends ImmutableObject implements Serializable {
 
   VKey() {}
 
-  VKey(Class<? extends T> kind, Key<T> ofyKey, Object sqlKey) {
+  VKey(Class<? extends T> kind, Key<T> ofyKey, Serializable sqlKey) {
     this.kind = kind;
     this.ofyKey = ofyKey;
     this.sqlKey = sqlKey;
@@ -57,7 +57,7 @@ public class VKey<T> extends ImmutableObject implements Serializable {
    *
    * <p>Deprecated. Create symmetric keys with create() instead.
    */
-  public static <T> VKey<T> createSql(Class<T> kind, Object sqlKey) {
+  public static <T> VKey<T> createSql(Class<T> kind, Serializable sqlKey) {
     checkArgumentNotNull(kind, "kind must not be null");
     checkArgumentNotNull(sqlKey, "sqlKey must not be null");
     return new VKey<T>(kind, null, sqlKey);
@@ -71,7 +71,7 @@ public class VKey<T> extends ImmutableObject implements Serializable {
   }
 
   /** Creates a {@link VKey} which only contains both sql and ofy primary key. */
-  public static <T> VKey<T> create(Class<T> kind, Object sqlKey, Key<T> ofyKey) {
+  public static <T> VKey<T> create(Class<T> kind, Serializable sqlKey, Key<T> ofyKey) {
     checkArgumentNotNull(kind, "kind must not be null");
     checkArgumentNotNull(sqlKey, "sqlKey must not be null");
     checkArgumentNotNull(ofyKey, "ofyKey must not be null");
@@ -84,8 +84,8 @@ public class VKey<T> extends ImmutableObject implements Serializable {
    * <p>IMPORTANT USAGE NOTE: Datastore entities that are not roots of entity groups (i.e. those
    * that do not have a null parent in their Objectify keys) require the full entity group
    * inheritance chain to be specified and thus cannot use this create method. You need to use
-   * {@link #create(Class, Object, Key)} instead and pass in the full, valid parent field in the
-   * Datastore key.
+   * {@link #create(Class, Serializable, Key)} instead and pass in the full, valid parent field in
+   * the Datastore key.
    */
   public static <T> VKey<T> create(Class<T> kind, long id) {
     checkArgument(
@@ -102,8 +102,8 @@ public class VKey<T> extends ImmutableObject implements Serializable {
    * <p>IMPORTANT USAGE NOTE: Datastore entities that are not roots of entity groups (i.e. those
    * that do not have a null parent in their Objectify keys) require the full entity group
    * inheritance chain to be specified and thus cannot use this create method. You need to use
-   * {@link #create(Class, Object, Key)} instead and pass in the full, valid parent field in the
-   * Datastore key.
+   * {@link #create(Class, Serializable, Key)} instead and pass in the full, valid parent field in
+   * the Datastore key.
    */
   public static <T> VKey<T> create(Class<T> kind, String name) {
     checkArgument(
@@ -172,7 +172,7 @@ public class VKey<T> extends ImmutableObject implements Serializable {
       throw new IllegalArgumentException("Missing value for last key of type " + lastClass);
     }
 
-    Object sqlKey = getSqlKey();
+    Serializable sqlKey = getSqlKey();
     Key<T> ofyKey =
         sqlKey instanceof Long
             ? Key.create(lastKey, getKind(), (Long) sqlKey)
@@ -197,7 +197,7 @@ public class VKey<T> extends ImmutableObject implements Serializable {
   }
 
   /** Returns the SQL primary key. */
-  public Object getSqlKey() {
+  public Serializable getSqlKey() {
     checkState(sqlKey != null, "Attempting obtain a null SQL key.");
     return this.sqlKey;
   }

core/src/main/java/google/registry/persistence/converter/VKeyConverter.java

@@ -16,11 +16,13 @@ package google.registry.persistence.converter;
 
 import com.googlecode.objectify.Key;
 import google.registry.persistence.VKey;
+import java.io.Serializable;
 import javax.annotation.Nullable;
 import javax.persistence.AttributeConverter;
 
 /** Converts VKey to a string column. */
-public abstract class VKeyConverter<T, C> implements AttributeConverter<VKey<? extends T>, C> {
+public abstract class VKeyConverter<T, C extends Serializable>
+    implements AttributeConverter<VKey<? extends T>, C> {
   @Override
   @Nullable
   public C convertToDatabaseColumn(@Nullable VKey<? extends T> attribute) {

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

@@ -62,6 +62,13 @@ public class CriteriaQueryBuilder<T> {
     return where(root.get(fieldName).in(values));
   }
 
+  /**
+   * Adds a WHERE clause to the query specifying that a value must not be in the given collection.
+   */
+  public CriteriaQueryBuilder<T> whereFieldIsNotIn(String fieldName, Collection<?> values) {
+    return where(root.get(fieldName).in(values).not());
+  }
+
   /**
    * Adds a WHERE clause to the query specifying that a collection field must contain a particular
    * value.

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

@@ -50,6 +50,11 @@ public interface JpaTransactionManager extends TransactionManager {
    */
   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);
 
@@ -65,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

@@ -41,9 +41,12 @@ 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.io.Serializable;
 import java.lang.reflect.Array;
 import java.lang.reflect.Field;
 import java.util.Calendar;
@@ -152,6 +155,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()) {
@@ -162,7 +174,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();
@@ -194,7 +206,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();
@@ -478,7 +490,8 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
             loadByKey(
                 VKey.createSql(
                     possibleChild.getClass(),
-                    emf.getPersistenceUnitUtil().getIdentifier(possibleChild)));
+                    // Casting to Serializable is safe according to JPA (JSR 338 sec. 2.4).
+                    (Serializable) emf.getPersistenceUnitUtil().getIdentifier(possibleChild)));
     return returnValue;
   }
 
@@ -489,13 +502,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 getEntityManager()
         .createQuery(String.format("FROM %s", getEntityType(clazz).getName()), clazz)
         .getResultStream()
-        .map(this::detach)
-        .collect(toImmutableList());
+        .map(this::detach);
   }
 
   @Override
@@ -736,11 +753,11 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     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();
       }
     }
@@ -759,17 +776,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();
@@ -1052,11 +1075,7 @@ 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.");
-      }
+      JpaTransactionManager.setQueryFetchSize(query, fetchSize);
       return query.getResultStream().map(JpaTransactionManagerImpl.this::detach);
     }
 

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

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

@@ -22,6 +22,7 @@ import google.registry.persistence.VKey;
 import java.util.NoSuchElementException;
 import java.util.Optional;
 import java.util.function.Supplier;
+import java.util.stream.Stream;
 import org.joda.time.DateTime;
 
 /**
@@ -240,6 +241,15 @@ public interface TransactionManager {
    */
   <T> ImmutableList<T> loadByEntities(Iterable<T> entities);
 
+  /**
+   * Returns a list of all entities of the given type that exist in the database.
+   *
+   * <p>The resulting list is empty if there are no entities of this type. In Datastore mode, if the
+   * class is a member of the cross-TLD entity group (i.e. if it has the {@link InCrossTld}
+   * annotation, then the correct ancestor query will automatically be applied.
+   */
+  <T> ImmutableList<T> loadAllOf(Class<T> clazz);
+
   /**
    * Returns a stream of all entities of the given type that exist in the database.
    *
@@ -247,7 +257,7 @@ public interface TransactionManager {
    * the class is a member of the cross-TLD entity group (i.e. if it has the {@link InCrossTld}
    * annotation, then the correct ancestor query will automatically be applied.
    */
-  <T> ImmutableList<T> loadAllOf(Class<T> clazz);
+  <T> Stream<T> loadAllOfStream(Class<T> clazz);
 
   /**
    * Loads the only instance of this particular class, or empty if none exists.

core/src/main/java/google/registry/rde/PendingDeposit.java

@@ -17,8 +17,17 @@ package google.registry.rde;
 import com.google.auto.value.AutoValue;
 import google.registry.model.common.Cursor.CursorType;
 import google.registry.model.rde.RdeMode;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
 import java.io.Serializable;
 import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.AtomicCoder;
+import org.apache.beam.sdk.coders.BooleanCoder;
+import org.apache.beam.sdk.coders.NullableCoder;
+import org.apache.beam.sdk.coders.SerializableCoder;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.coders.VarIntCoder;
 import org.joda.time.DateTime;
 import org.joda.time.Duration;
 
@@ -69,17 +78,9 @@ public abstract class PendingDeposit implements Serializable {
   @Nullable
   public abstract Integer revision();
 
-  static PendingDeposit create(
+  public static PendingDeposit create(
       String tld, DateTime watermark, RdeMode mode, CursorType cursor, Duration interval) {
-    return new AutoValue_PendingDeposit(
-        false,
-        tld,
-        watermark,
-        mode,
-        cursor,
-        interval,
-        null,
-        null);
+    return new AutoValue_PendingDeposit(false, tld, watermark, mode, cursor, interval, null, null);
   }
 
   static PendingDeposit createInManualOperation(
@@ -89,15 +90,52 @@ public abstract class PendingDeposit implements Serializable {
       String directoryWithTrailingSlash,
       @Nullable Integer revision) {
     return new AutoValue_PendingDeposit(
-        true,
-        tld,
-        watermark,
-        mode,
-        null,
-        null,
-        directoryWithTrailingSlash,
-        revision);
+        true, tld, watermark, mode, null, null, directoryWithTrailingSlash, revision);
   }
 
   PendingDeposit() {}
+
+  /**
+   * A deterministic coder for {@link PendingDeposit} used during a GroupBy transform.
+   *
+   * <p>We cannot use a {@link SerializableCoder} directly because it does not guarantee
+   * determinism, which is required by GroupBy.
+   */
+  public static class PendingDepositCoder extends AtomicCoder<PendingDeposit> {
+
+    private PendingDepositCoder() {
+      super();
+    }
+
+    private static final PendingDepositCoder INSTANCE = new PendingDepositCoder();
+
+    public static PendingDepositCoder of() {
+      return INSTANCE;
+    }
+
+    @Override
+    public void encode(PendingDeposit value, OutputStream outStream) throws IOException {
+      BooleanCoder.of().encode(value.manual(), outStream);
+      StringUtf8Coder.of().encode(value.tld(), outStream);
+      SerializableCoder.of(DateTime.class).encode(value.watermark(), outStream);
+      SerializableCoder.of(RdeMode.class).encode(value.mode(), outStream);
+      NullableCoder.of(SerializableCoder.of(CursorType.class)).encode(value.cursor(), outStream);
+      NullableCoder.of(SerializableCoder.of(Duration.class)).encode(value.interval(), outStream);
+      NullableCoder.of(StringUtf8Coder.of()).encode(value.directoryWithTrailingSlash(), outStream);
+      NullableCoder.of(VarIntCoder.of()).encode(value.revision(), outStream);
+    }
+
+    @Override
+    public PendingDeposit decode(InputStream inStream) throws IOException {
+      return new AutoValue_PendingDeposit(
+          BooleanCoder.of().decode(inStream),
+          StringUtf8Coder.of().decode(inStream),
+          SerializableCoder.of(DateTime.class).decode(inStream),
+          SerializableCoder.of(RdeMode.class).decode(inStream),
+          NullableCoder.of(SerializableCoder.of(CursorType.class)).decode(inStream),
+          NullableCoder.of(SerializableCoder.of(Duration.class)).decode(inStream),
+          NullableCoder.of(StringUtf8Coder.of()).decode(inStream),
+          NullableCoder.of(VarIntCoder.of()).decode(inStream));
+    }
+  }
 }

core/src/main/java/google/registry/rde/RdeFragmenter.java

@@ -0,0 +1,109 @@
+// 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.rde;
+
+import static google.registry.model.EppResourceUtils.loadAtPointInTime;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+
+import com.google.auto.value.AutoValue;
+import com.google.common.collect.ImmutableMap;
+import google.registry.model.EppResource;
+import google.registry.model.contact.ContactResource;
+import google.registry.model.domain.DomainBase;
+import google.registry.model.host.HostResource;
+import google.registry.model.rde.RdeMode;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.function.Supplier;
+import org.joda.time.DateTime;
+
+/** Loading cache that turns a resource into XML for the various points in time and modes. */
+public class RdeFragmenter {
+  private final Map<WatermarkModePair, Optional<DepositFragment>> cache = new HashMap<>();
+  private final ImmutableMap<DateTime, Supplier<EppResource>> resourceAtTimes;
+  private final RdeMarshaller marshaller;
+
+  long cacheHits = 0;
+  long resourcesNotFound = 0;
+  long resourcesFound = 0;
+
+  public RdeFragmenter(
+      ImmutableMap<DateTime, Supplier<EppResource>> resourceAtTimes, RdeMarshaller marshaller) {
+    this.resourceAtTimes = resourceAtTimes;
+    this.marshaller = marshaller;
+  }
+
+  public Optional<DepositFragment> marshal(DateTime watermark, RdeMode mode) {
+    Optional<DepositFragment> result = cache.get(WatermarkModePair.create(watermark, mode));
+    if (result != null) {
+      cacheHits++;
+      return result;
+    }
+    EppResource resource = resourceAtTimes.get(watermark).get();
+    if (resource == null) {
+      result = Optional.empty();
+      cache.put(WatermarkModePair.create(watermark, RdeMode.FULL), result);
+      cache.put(WatermarkModePair.create(watermark, RdeMode.THIN), result);
+      resourcesNotFound++;
+      return result;
+    }
+    resourcesFound++;
+    if (resource instanceof DomainBase) {
+      result = Optional.of(marshaller.marshalDomain((DomainBase) resource, mode));
+      cache.put(WatermarkModePair.create(watermark, mode), result);
+      return result;
+    } else if (resource instanceof ContactResource) {
+      result = Optional.of(marshaller.marshalContact((ContactResource) resource));
+      cache.put(WatermarkModePair.create(watermark, RdeMode.FULL), result);
+      cache.put(WatermarkModePair.create(watermark, RdeMode.THIN), result);
+      return result;
+    } else if (resource instanceof HostResource) {
+      HostResource host = (HostResource) resource;
+      result =
+          Optional.of(
+              host.isSubordinate()
+                  ? marshaller.marshalSubordinateHost(
+                      host,
+                      // Note that loadAtPointInTime() does cloneProjectedAtTime(watermark) for
+                      // us.
+                      Objects.requireNonNull(
+                          loadAtPointInTime(
+                              tm().loadByKey(host.getSuperordinateDomain()), watermark)))
+                  : marshaller.marshalExternalHost(host));
+      cache.put(WatermarkModePair.create(watermark, RdeMode.FULL), result);
+      cache.put(WatermarkModePair.create(watermark, RdeMode.THIN), result);
+      return result;
+    } else {
+      throw new IllegalStateException(
+          String.format(
+              "Resource %s of type %s cannot be converted to XML.",
+              resource, resource.getClass().getSimpleName()));
+    }
+  }
+
+  /** Map key for {@link RdeFragmenter} cache. */
+  @AutoValue
+  abstract static class WatermarkModePair {
+    abstract DateTime watermark();
+
+    abstract RdeMode mode();
+
+    static WatermarkModePair create(DateTime watermark, RdeMode mode) {
+      return new AutoValue_RdeFragmenter_WatermarkModePair(watermark, mode);
+    }
+  }
+}

core/src/main/java/google/registry/rde/RdeStagingMapper.java

@@ -16,12 +16,10 @@ package google.registry.rde;
 
 import static com.google.common.base.Strings.nullToEmpty;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
-import static google.registry.model.EppResourceUtils.loadAtPointInTime;
 import static google.registry.model.EppResourceUtils.loadAtPointInTimeAsync;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.appengine.tools.mapreduce.Mapper;
-import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSetMultimap;
@@ -34,9 +32,6 @@ import google.registry.model.host.HostResource;
 import google.registry.model.rde.RdeMode;
 import google.registry.model.registrar.Registrar;
 import google.registry.xml.ValidationMode;
-import java.util.HashMap;
-import java.util.Map;
-import java.util.Objects;
 import java.util.Optional;
 import java.util.function.Supplier;
 import org.joda.time.DateTime;
@@ -129,7 +124,7 @@ public final class RdeStagingMapper extends Mapper<EppResource, PendingDeposit,
         ImmutableMap.copyOf(Maps.asMap(dates, input -> loadAtPointInTimeAsync(resource, input)));
 
     // Convert resource to an XML fragment for each watermark/mode pair lazily and cache the result.
-    Fragmenter fragmenter = new Fragmenter(resourceAtTimes);
+    RdeFragmenter fragmenter = new RdeFragmenter(resourceAtTimes, marshaller);
 
     // Emit resource as an XML fragment for all TLDs and modes pending deposit.
     long resourcesEmitted = 0;
@@ -157,74 +152,4 @@ public final class RdeStagingMapper extends Mapper<EppResource, PendingDeposit,
     // Avoid running out of memory.
     tm().clearSessionCache();
   }
-
-  /** Loading cache that turns a resource into XML for the various points in time and modes. */
-  private class Fragmenter {
-    private final Map<WatermarkModePair, Optional<DepositFragment>> cache = new HashMap<>();
-    private final ImmutableMap<DateTime, Supplier<EppResource>> resourceAtTimes;
-
-    long cacheHits = 0;
-    long resourcesNotFound = 0;
-    long resourcesFound = 0;
-
-    Fragmenter(ImmutableMap<DateTime, Supplier<EppResource>> resourceAtTimes) {
-      this.resourceAtTimes = resourceAtTimes;
-    }
-
-    Optional<DepositFragment> marshal(DateTime watermark, RdeMode mode) {
-      Optional<DepositFragment> result = cache.get(WatermarkModePair.create(watermark, mode));
-      if (result != null) {
-        cacheHits++;
-        return result;
-      }
-      EppResource resource = resourceAtTimes.get(watermark).get();
-      if (resource == null) {
-        result = Optional.empty();
-        cache.put(WatermarkModePair.create(watermark, RdeMode.FULL), result);
-        cache.put(WatermarkModePair.create(watermark, RdeMode.THIN), result);
-        resourcesNotFound++;
-        return result;
-      }
-      resourcesFound++;
-      if (resource instanceof DomainBase) {
-        result = Optional.of(marshaller.marshalDomain((DomainBase) resource, mode));
-        cache.put(WatermarkModePair.create(watermark, mode), result);
-        return result;
-      } else if (resource instanceof ContactResource) {
-        result = Optional.of(marshaller.marshalContact((ContactResource) resource));
-        cache.put(WatermarkModePair.create(watermark, RdeMode.FULL), result);
-        cache.put(WatermarkModePair.create(watermark, RdeMode.THIN), result);
-        return result;
-      } else if (resource instanceof HostResource) {
-        HostResource host = (HostResource) resource;
-        result =
-            Optional.of(
-                host.isSubordinate()
-                    ? marshaller.marshalSubordinateHost(
-                        host,
-                        // Note that loadAtPointInTime() does cloneProjectedAtTime(watermark) for
-                        // us.
-                        Objects.requireNonNull(
-                            loadAtPointInTime(
-                                tm().loadByKey(host.getSuperordinateDomain()), watermark)))
-                    : marshaller.marshalExternalHost(host));
-        cache.put(WatermarkModePair.create(watermark, RdeMode.FULL), result);
-        cache.put(WatermarkModePair.create(watermark, RdeMode.THIN), result);
-        return result;
-      } else {
-        throw new AssertionError(resource.toString());
-      }
-    }
-  }
-
-  /** Map key for {@link Fragmenter} cache. */
-  @AutoValue
-  abstract static class WatermarkModePair {
-    abstract DateTime watermark();
-    abstract RdeMode mode();
-
-    static WatermarkModePair create(DateTime watermark, RdeMode mode) {
-      return new AutoValue_RdeStagingMapper_WatermarkModePair(watermark, mode);
-    }
-  }
 }

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

@@ -42,7 +42,8 @@ import javax.persistence.NoResultException;
     method = GET,
     automaticallyPrintOk = true,
     auth = Auth.AUTH_INTERNAL_OR_ADMIN)
-class ReplicateToDatastoreAction implements Runnable {
+@VisibleForTesting
+public class ReplicateToDatastoreAction implements Runnable {
   public static final String PATH = "/_dr/cron/replicateToDatastore";
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
 
@@ -61,12 +62,12 @@ class ReplicateToDatastoreAction implements Runnable {
   }
 
   @VisibleForTesting
-  List<TransactionEntity> getTransactionBatch() {
+  public List<TransactionEntity> getTransactionBatch() {
     // Get the next batch of transactions that we haven't replicated.
     LastSqlTransaction lastSqlTxnBeforeBatch = ofyTm().transact(() -> LastSqlTransaction.load());
     try {
       return jpaTm()
-          .transact(
+          .transactWithoutBackup(
               () ->
                   jpaTm()
                       .query(
@@ -86,7 +87,7 @@ class ReplicateToDatastoreAction implements Runnable {
    * be aborted.
    */
   @VisibleForTesting
-  boolean applyTransaction(TransactionEntity txnEntity) {
+  public boolean applyTransaction(TransactionEntity txnEntity) {
     logger.atInfo().log("Applying a single transaction Cloud SQL -> Cloud Datastore");
     return ofyTm()
         .transact(

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

@@ -143,12 +143,7 @@ public class PremiumListDao {
     RevisionIdAndLabel revisionIdAndLabel =
         RevisionIdAndLabel.create(loadedList.getRevisionId(), label);
     try {
-      Optional<BigDecimal> price = premiumEntryCache.get(revisionIdAndLabel);
-      return price.map(
-          p ->
-              Money.of(
-                  loadedList.getCurrency(),
-                  p.setScale(loadedList.getCurrency().getDecimalPlaces())));
+      return premiumEntryCache.get(revisionIdAndLabel).map(loadedList::convertAmountToMoney);
     } catch (InvalidCacheLoadException | ExecutionException e) {
       throw new RuntimeException(
           String.format(

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

@@ -125,13 +125,7 @@ final class AckPollMessagesCommand implements CommandWithRemoteApi {
               if (!isNullOrEmpty(message)) {
                 query = query.where("msg", LIKE, "%" + message + "%");
               }
-
-              query.stream()
-                  // Detach it so that we can print out the old, non-acked version
-                  // (for autorenews, acking changes the next event time)
-                  // TODO(mmuller): remove after PR 1116 is merged.
-                  .peek(jpaTm().getEntityManager()::detach)
-                  .forEach(this::actOnPollMessage);
+              query.stream().forEach(this::actOnPollMessage);
             });
   }
 

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

@@ -17,7 +17,7 @@ package google.registry.tools;
 import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
 import com.google.template.soy.data.SoyMapData;
-import google.registry.tools.soy.DeleteDomainSoyInfo;
+import google.registry.tools.soy.DomainDeleteSoyInfo;
 
 /** A command to delete a domain via EPP. */
 @Parameters(separators = " =", commandDescription = "Delete domain")
@@ -60,7 +60,7 @@ final class DeleteDomainCommand extends MutatingEppToolCommand {
       // Immediate deletion is accomplished using the superuser extension.
       superuser = true;
     }
-    setSoyTemplate(DeleteDomainSoyInfo.getInstance(), DeleteDomainSoyInfo.DELETEDOMAIN);
+    setSoyTemplate(DomainDeleteSoyInfo.getInstance(), DomainDeleteSoyInfo.DELETEDOMAIN);
     addSoyRecord(clientId, new SoyMapData(
         "domainName", domainName,
         "immediately", immediately,

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

@@ -17,7 +17,7 @@ package google.registry.tools;
 import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
 import com.google.template.soy.data.SoyMapData;
-import google.registry.tools.soy.DeleteHostSoyInfo;
+import google.registry.tools.soy.HostDeleteSoyInfo;
 
 /** A command to delete a host via EPP. */
 @Parameters(separators = " =", commandDescription = "Delete host")
@@ -49,7 +49,7 @@ final class DeleteHostCommand extends MutatingEppToolCommand {
 
   @Override
   protected void initMutatingEppToolCommand() {
-    setSoyTemplate(DeleteHostSoyInfo.getInstance(), DeleteHostSoyInfo.DELETEHOST);
+    setSoyTemplate(HostDeleteSoyInfo.getInstance(), HostDeleteSoyInfo.DELETEHOST);
     addSoyRecord(clientId, new SoyMapData(
         "hostName", hostName,
         "reason", reason,

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

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

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

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

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

@@ -25,7 +25,7 @@ import com.beust.jcommander.Parameters;
 import com.google.common.base.Joiner;
 import com.google.template.soy.data.SoyMapData;
 import google.registry.model.domain.DomainBase;
-import google.registry.tools.soy.RenewDomainSoyInfo;
+import google.registry.tools.soy.DomainRenewSoyInfo;
 import google.registry.util.Clock;
 import java.util.List;
 import java.util.Optional;
@@ -68,7 +68,7 @@ final class RenewDomainCommand extends MutatingEppToolCommand {
       Optional<DomainBase> domainOptional =
           loadByForeignKey(DomainBase.class, domainName, now);
       checkArgumentPresent(domainOptional, "Domain '%s' does not exist or is deleted", domainName);
-      setSoyTemplate(RenewDomainSoyInfo.getInstance(), RenewDomainSoyInfo.RENEWDOMAIN);
+      setSoyTemplate(DomainRenewSoyInfo.getInstance(), DomainRenewSoyInfo.RENEWDOMAIN);
       DomainBase domain = domainOptional.get();
       addSoyRecord(
           isNullOrEmpty(clientId) ? domain.getCurrentSponsorClientId() : clientId,

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

@@ -15,7 +15,6 @@
 package google.registry.tools;
 
 import static com.google.common.base.Preconditions.checkArgument;
-import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static com.google.common.collect.Sets.difference;
 import static com.google.common.collect.Sets.intersection;
 import static com.google.common.collect.Sets.union;
@@ -26,7 +25,6 @@ import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Maps;
-import com.googlecode.objectify.Key;
 import google.registry.config.RegistryEnvironment;
 import google.registry.model.registry.Registry;
 import google.registry.model.registry.Registry.TldState;
@@ -113,7 +111,7 @@ class UpdateTldCommand extends CreateOrUpdateTldCommand {
   ImmutableSet<String> getReservedLists(Registry oldRegistry) {
     return formUpdatedList(
         "reserved lists",
-        oldRegistry.getReservedLists().stream().map(Key::getName).collect(toImmutableSet()),
+        oldRegistry.getReservedListNames(),
         reservedListNames,
         reservedListsAdd,
         reservedListsRemove);

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

@@ -0,0 +1,30 @@
+// Copyright 2021 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.tools.javascrap;
+
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+
+import com.beust.jcommander.Parameters;
+import google.registry.model.registry.Registry;
+import google.registry.tools.CommandWithRemoteApi;
+
+/** Scrap command to resave all Registry entities. */
+@Parameters(commandDescription = "Resave all TLDs")
+public class ResaveAllTldsCommand implements CommandWithRemoteApi {
+  @Override
+  public void run() throws Exception {
+    tm().transact(() -> tm().putAll(tm().loadAllOf(Registry.class)));
+  }
+}

core/src/main/resources/google/registry/beam/rde_pipeline_metadata.json

@@ -0,0 +1,23 @@
+{
+  "name": "RDE/BRDA Deposit Generation",
+  "description": "An Apache Beam pipeline generates RDE or BRDA deposits and deposits them to GCS with GhostRyde encryption.",
+  "parameters": [
+    {
+      "name": "pendings",
+      "label": "The pendings deposits to generate.",
+      "helpText": "A TLD to PendingDeposit map that is serialized and Base64 URL-safe encoded.",
+      "regexes": [
+        "A-Za-z0-9\\-_"
+      ]
+    },
+    {
+      "name": "validationMode",
+      "label": "How strict the marshaller validates the given EPP resources.",
+      "helpText": "If set to LENIENT the marshaller will not warn about missing data on the EPP resources.",
+      "is_optional": true,
+      "regexes": [
+        "^STRICT|LENIENT$"
+      ]
+    }
+  ]
+}

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

@@ -0,0 +1,87 @@
+// 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 static com.google.common.truth.Truth.assertThat;
+import static google.registry.model.ofy.ObjectifyService.auditedOfy;
+
+import com.google.appengine.api.datastore.DatastoreService;
+import com.google.appengine.api.datastore.DatastoreServiceFactory;
+import com.google.appengine.api.datastore.Entity;
+import com.google.appengine.api.datastore.EntityTranslator;
+import com.google.common.collect.ImmutableList;
+import com.google.common.io.Resources;
+import com.google.storage.onestore.v3.OnestoreEntity.EntityProto;
+import google.registry.model.ofy.CommitLogCheckpoint;
+import google.registry.testing.AppEngineExtension;
+import google.registry.testing.DatastoreEntityExtension;
+import java.io.InputStream;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Order;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+public class EntityImportsTest {
+
+  @RegisterExtension
+  @Order(value = 1)
+  final DatastoreEntityExtension datastoreEntityExtension = new DatastoreEntityExtension();
+
+  @RegisterExtension
+  final AppEngineExtension appEngine =
+      new AppEngineExtension.Builder().withDatastoreAndCloudSql().withoutCannedData().build();
+
+  private DatastoreService datastoreService;
+
+  @BeforeEach
+  void beforeEach() {
+    datastoreService = DatastoreServiceFactory.getDatastoreService();
+  }
+
+  @Test
+  void importCommitLogs_keysFixed() throws Exception {
+    // Input resource is a standard commit log file whose entities has "AppId_1" as appId. The key
+    // fixes can be verified by checking that the appId of an imported entity's key has been updated
+    // to 'test' (which is set by AppEngineExtension) and/or that after persistence the imported
+    // entity can be loaded by Objectify.
+    try (InputStream commitLogInputStream =
+        Resources.getResource("google/registry/backup/commitlog.data").openStream()) {
+      ImmutableList<Entity> entities =
+          loadEntityProtos(commitLogInputStream).stream()
+              .map(EntityImports::fixEntity)
+              .map(EntityTranslator::createFromPb)
+              .collect(ImmutableList.toImmutableList());
+      // Verifies that the original appId has been overwritten.
+      assertThat(entities.get(0).getKey().getAppId()).isEqualTo("test");
+      datastoreService.put(entities);
+      // Imported entity can be found by Ofy after appId conversion.
+      assertThat(auditedOfy().load().type(CommitLogCheckpoint.class).count()).isGreaterThan(0);
+    }
+  }
+
+  private static ImmutableList<EntityProto> loadEntityProtos(InputStream inputStream) {
+    ImmutableList.Builder<EntityProto> protosBuilder = new ImmutableList.Builder<>();
+    while (true) {
+      EntityProto proto = new EntityProto();
+      boolean parsed = proto.parseDelimitedFrom(inputStream);
+      if (parsed && proto.isInitialized()) {
+        protosBuilder.add(proto);
+      } else {
+        break;
+      }
+    }
+    return protosBuilder.build();
+  }
+}

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

@@ -65,7 +65,6 @@ public class GcsDiffFileListerTest {
   @BeforeEach
   void beforeEach() throws Exception {
     diffLister.gcsService = gcsService;
-    diffLister.gcsBucket = GCS_BUCKET;
     diffLister.executor = newDirectExecutorService();
     for (int i = 0; i < 5; i++) {
       gcsService.createOrReplace(
@@ -87,7 +86,7 @@ public class GcsDiffFileListerTest {
   }
 
   private Iterable<DateTime> listDiffFiles(DateTime fromTime, DateTime toTime) {
-    return extractTimesFromDiffFiles(diffLister.listDiffFiles(fromTime, toTime));
+    return extractTimesFromDiffFiles(diffLister.listDiffFiles(GCS_BUCKET, fromTime, toTime));
   }
 
   private void addGcsFile(int fileAge, int prevAge) throws IOException {

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

@@ -17,7 +17,6 @@ package google.registry.backup;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.util.concurrent.MoreExecutors.newDirectExecutorService;
-import static google.registry.backup.RestoreCommitLogsActionTest.GCS_BUCKET;
 import static google.registry.backup.RestoreCommitLogsActionTest.createCheckpoint;
 import static google.registry.backup.RestoreCommitLogsActionTest.saveDiffFile;
 import static google.registry.backup.RestoreCommitLogsActionTest.saveDiffFileNotToRestore;
@@ -128,9 +127,9 @@ public class ReplayCommitLogsToSqlActionTest {
     action.response = response;
     action.requestStatusChecker = requestStatusChecker;
     action.clock = fakeClock;
+    action.gcsBucket = "gcs bucket";
     action.diffLister = new GcsDiffFileLister();
     action.diffLister.gcsService = gcsService;
-    action.diffLister.gcsBucket = GCS_BUCKET;
     action.diffLister.executor = newDirectExecutorService();
     ofyTm()
         .transact(
@@ -208,6 +207,30 @@ public class ReplayCommitLogsToSqlActionTest {
     assertExpectedIds("previous to keep");
   }
 
+  @Test
+  void testReplay_dryRun() throws Exception {
+    action.dryRun = true;
+    DateTime now = fakeClock.nowUtc();
+    jpaTm().transact(() -> jpaTm().insertWithoutBackup(TestObject.create("previous to keep")));
+    Key<CommitLogBucket> bucketKey = getBucketKey(1);
+    Key<CommitLogManifest> manifestKey = CommitLogManifest.createKey(bucketKey, now);
+    saveDiffFileNotToRestore(gcsService, now.minusMinutes(2));
+    jpaTm().transact(() -> SqlReplayCheckpoint.set(now.minusMinutes(1).minusMillis(1)));
+    saveDiffFile(
+        gcsService,
+        createCheckpoint(now.minusMinutes(1)),
+        CommitLogManifest.create(bucketKey, now, null),
+        CommitLogMutation.create(manifestKey, TestObject.create("a")),
+        CommitLogMutation.create(manifestKey, TestObject.create("b")));
+
+    action.run();
+    assertThat(response.getStatus()).isEqualTo(SC_OK);
+    assertThat(response.getPayload())
+        .isEqualTo(
+            "Running in dry-run mode; would have processed %d files. They are (limit 10):\n"
+                + "commit_diff_until_1999-12-31T23:59:00.000Z");
+  }
+
   @Test
   void testReplay_manifestWithNoDeletions() throws Exception {
     DateTime now = fakeClock.nowUtc();
@@ -494,6 +517,8 @@ public class ReplayCommitLogsToSqlActionTest {
   private void runAndAssertSuccess(DateTime expectedCheckpointTime) {
     action.run();
     assertThat(response.getStatus()).isEqualTo(SC_OK);
+    assertThat(response.getPayload())
+        .isEqualTo("ReplayCommitLogsToSqlAction completed successfully.");
     assertThat(jpaTm().transact(SqlReplayCheckpoint::get)).isEqualTo(expectedCheckpointTime);
   }
 

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

@@ -34,9 +34,11 @@ import com.google.appengine.tools.cloudstorage.GcsService;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Lists;
+import com.google.common.io.Resources;
 import com.google.common.primitives.Longs;
 import com.googlecode.objectify.Key;
 import google.registry.model.ImmutableObject;
+import google.registry.model.domain.DomainBase;
 import google.registry.model.ofy.CommitLogBucket;
 import google.registry.model.ofy.CommitLogCheckpoint;
 import google.registry.model.ofy.CommitLogCheckpointRoot;
@@ -53,6 +55,7 @@ import java.nio.ByteBuffer;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
+import java.util.Optional;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -81,9 +84,10 @@ public class RestoreCommitLogsActionTest {
     action.datastoreService = DatastoreServiceFactory.getDatastoreService();
     action.fromTime = now.minusMillis(1);
     action.retrier = new Retrier(new FakeSleeper(new FakeClock()), 1);
+    action.defaultGcsBucket = GCS_BUCKET;
+    action.gcsBucketOverride = Optional.empty();
     action.diffLister = new GcsDiffFileLister();
     action.diffLister.gcsService = gcsService;
-    action.diffLister.gcsBucket = GCS_BUCKET;
     action.diffLister.executor = newDirectExecutorService();
   }
 
@@ -258,10 +262,40 @@ public class RestoreCommitLogsActionTest {
     assertInDatastore(CommitLogCheckpointRoot.create(now));
   }
 
+  @Test
+  void testRestore_fromOtherProject() throws IOException {
+    // Input resource is a standard commit log file whose entities has "AppId_1" as appId. Among the
+    // entities are CommitLogMutations that have an embedded DomainBase and a ContactResource, both
+    // having "AppId_1" as appId. This test verifies that the embedded entities are properly
+    // imported, in particular, the domain's 'registrant' key can be used by Objectify to load the
+    // contact.
+    saveDiffFile(
+        gcsService,
+        Resources.toByteArray(Resources.getResource("google/registry/backup/commitlog.data")),
+        now);
+    action.run();
+    auditedOfy().clearSessionCache();
+    List<DomainBase> domainBases = auditedOfy().load().type(DomainBase.class).list();
+    assertThat(domainBases).hasSize(1);
+    DomainBase domainBase = domainBases.get(0);
+    // If the registrant is found, then the key instance in domainBase is fixed.
+    assertThat(auditedOfy().load().key(domainBase.getRegistrant().getOfyKey()).now()).isNotNull();
+  }
+
   static CommitLogCheckpoint createCheckpoint(DateTime now) {
     return CommitLogCheckpoint.create(now, toMap(getBucketIds(), x -> now));
   }
 
+  static void saveDiffFile(GcsService gcsService, byte[] rawBytes, DateTime timestamp)
+      throws IOException {
+    gcsService.createOrReplace(
+        new GcsFilename(GCS_BUCKET, DIFF_FILE_PREFIX + timestamp),
+        new GcsFileOptions.Builder()
+            .addUserMetadata(LOWER_BOUND_CHECKPOINT, timestamp.minusMinutes(1).toString())
+            .build(),
+        ByteBuffer.wrap(rawBytes));
+  }
+
   static Iterable<ImmutableObject> saveDiffFile(
       GcsService gcsService, CommitLogCheckpoint checkpoint, ImmutableObject... entities)
       throws IOException {
@@ -271,12 +305,7 @@ public class RestoreCommitLogsActionTest {
     for (ImmutableObject entity : allEntities) {
       serializeEntity(entity, output);
     }
-    gcsService.createOrReplace(
-        new GcsFilename(GCS_BUCKET, DIFF_FILE_PREFIX + now),
-        new GcsFileOptions.Builder()
-            .addUserMetadata(LOWER_BOUND_CHECKPOINT, now.minusMinutes(1).toString())
-            .build(),
-        ByteBuffer.wrap(output.toByteArray()));
+    saveDiffFile(gcsService, output.toByteArray(), now);
     return allEntities;
   }