Add a dry-run option to commit-log-replay action and use it in Sandbox (#1234)

* Set payload response in happy path of ReplayCommitLogsToSqlAction

I suspect this may be the reason the logs are missing on the happy path (when it
runs successfully), but are visible on the exception paths (which do set the
payload response). I don't think App Engine likes it when a Web request
terminates without a response.

This also adds more logging and error handling.

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/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/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,6 +27,7 @@ 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.model.common.DatabaseMigrationStateSchedule;
@@ -35,6 +38,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 +68,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 +82,11 @@ public class ReplayCommitLogsToSqlAction implements Runnable {
   @Inject GcsDiffFileLister diffLister;
   @Inject Clock clock;
 
+  /** 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 +116,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);
+    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

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

core/src/main/java/google/registry/batch/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/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/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/registry/Registry.java

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

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

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

core/src/main/java/google/registry/model/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/PersistenceModule.java

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

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

@@ -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,6 +41,8 @@ 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;
@@ -152,6 +154,15 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
 
   @Override
   public <T> T transact(Supplier<T> work) {
+    return transact(work, true);
+  }
+
+  @Override
+  public <T> T transactWithoutBackup(Supplier<T> work) {
+    return transact(work, false);
+  }
+
+  private <T> T transact(Supplier<T> work, boolean withBackup) {
     return retrier.callWithRetry(
         () -> {
           if (inTransaction()) {
@@ -162,7 +173,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
           EntityTransaction txn = txnInfo.entityManager.getTransaction();
           try {
             txn.begin();
-            txnInfo.start(clock);
+            txnInfo.start(clock, withBackup);
             T result = work.get();
             txnInfo.recordTransaction();
             txn.commit();
@@ -194,7 +205,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     EntityTransaction txn = txnInfo.entityManager.getTransaction();
     try {
       txn.begin();
-      txnInfo.start(clock);
+      txnInfo.start(clock, true);
       T result = work.get();
       txnInfo.recordTransaction();
       txn.commit();
@@ -489,13 +500,17 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
 
   @Override
   public <T> ImmutableList<T> loadAllOf(Class<T> clazz) {
+    return loadAllOfStream(clazz).collect(toImmutableList());
+  }
+
+  @Override
+  public <T> Stream<T> loadAllOfStream(Class<T> clazz) {
     checkArgumentNotNull(clazz, "clazz must be specified");
     assertInTransaction();
     return getEntityManager()
         .createQuery(String.format("FROM %s", getEntityType(clazz).getName()), clazz)
         .getResultStream()
-        .map(this::detach)
-        .collect(toImmutableList());
+        .map(this::detach);
   }
 
   @Override
@@ -736,11 +751,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 +774,23 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     }
 
     private void addUpdate(Object entity) {
-      if (contentsBuilder != null) {
+      if (contentsBuilder != null && shouldReplicate(entity.getClass())) {
         contentsBuilder.addUpdate(entity);
       }
     }
 
     private void addDelete(VKey<?> key) {
-      if (contentsBuilder != null) {
+      if (contentsBuilder != null && shouldReplicate(key.getKind())) {
         contentsBuilder.addDelete(key);
       }
     }
 
+    /** Returns true if the entity class should be replicated from SQL to datastore. */
+    private boolean shouldReplicate(Class<?> entityClass) {
+      return !NonReplicatedEntity.class.isAssignableFrom(entityClass)
+          && !SqlOnlyEntity.class.isAssignableFrom(entityClass);
+    }
+
     private void recordTransaction() {
       if (contentsBuilder != null) {
         Transaction persistedTxn = contentsBuilder.build();
@@ -1052,11 +1073,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/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/ReplayCommitLogsToSqlActionTest.java

@@ -208,6 +208,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 +518,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;
@@ -258,10 +260,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 +303,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;
   }
 

core/src/test/java/google/registry/beam/common/RegistryJpaReadTest.java

@@ -21,6 +21,7 @@ import static google.registry.util.DateTimeUtils.END_OF_TIME;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import google.registry.beam.TestPipelineExtension;
 import google.registry.beam.common.RegistryJpaIO.Read;
@@ -37,14 +38,13 @@ import google.registry.model.eppcommon.StatusValue;
 import google.registry.model.registrar.Registrar;
 import google.registry.model.registry.Registry;
 import google.registry.model.transfer.ContactTransferData;
+import google.registry.persistence.transaction.CriteriaQueryBuilder;
 import google.registry.persistence.transaction.JpaTestRules;
 import google.registry.persistence.transaction.JpaTestRules.JpaIntegrationTestExtension;
-import google.registry.persistence.transaction.JpaTransactionManager;
 import google.registry.testing.AppEngineExtension;
 import google.registry.testing.DatabaseHelper;
 import google.registry.testing.DatastoreEntityExtension;
 import google.registry.testing.FakeClock;
-import google.registry.testing.SystemPropertyExtension;
 import org.apache.beam.sdk.testing.PAssert;
 import org.apache.beam.sdk.values.PCollection;
 import org.joda.time.DateTime;
@@ -64,14 +64,6 @@ public class RegistryJpaReadTest {
   @Order(Order.DEFAULT - 1)
   final transient DatastoreEntityExtension datastore = new DatastoreEntityExtension();
 
-  // The pipeline runner on Kokoro sometimes mistakes the platform as appengine, resulting in
-  // a null thread factory. The cause is unknown but it may be due to the interaction with
-  // the DatastoreEntityExtension above. To work around the problem, we explicitly unset the
-  // relevant property before test starts.
-  @RegisterExtension
-  final transient SystemPropertyExtension systemPropertyExtension =
-      new SystemPropertyExtension().setProperty("com.google.appengine.runtime.environment", null);
-
   @RegisterExtension
   final transient JpaIntegrationTestExtension database =
       new JpaTestRules.Builder().withClock(fakeClock).buildIntegrationTestRule();
@@ -98,10 +90,10 @@ public class RegistryJpaReadTest {
   }
 
   @Test
-  void readWithQueryComposer() {
+  void readWithCriteriaQuery() {
     Read<ContactResource, String> read =
         RegistryJpaIO.read(
-            (JpaTransactionManager jpaTm) -> jpaTm.createQueryComposer(ContactResource.class),
+            () -> CriteriaQueryBuilder.create(ContactResource.class).build(),
             ContactBase::getContactId);
     PCollection<String> repoIds = testPipeline.apply(read);
 
@@ -115,19 +107,57 @@ public class RegistryJpaReadTest {
     Read<Object[], String> read =
         RegistryJpaIO.read(
             "select d, r.emailAddress from Domain d join Registrar r on"
-                + " d.currentSponsorClientId = r.clientIdentifier where r.type = 'REAL'"
+                + " d.currentSponsorClientId = r.clientIdentifier where r.type = :type"
                 + " and d.deletionTime > now()",
-            RegistryJpaReadTest::parseRow);
+            ImmutableMap.of("type", Registrar.Type.REAL),
+            false,
+            (Object[] row) -> {
+              DomainBase domainBase = (DomainBase) row[0];
+              String emailAddress = (String) row[1];
+              return domainBase.getRepoId() + "-" + emailAddress;
+            });
     PCollection<String> joinedStrings = testPipeline.apply(read);
 
     PAssert.that(joinedStrings).containsInAnyOrder("4-COM-me@google.com");
     testPipeline.run();
   }
 
-  private static String parseRow(Object[] row) {
-    DomainBase domainBase = (DomainBase) row[0];
-    String emailAddress = (String) row[1];
-    return domainBase.getRepoId() + "-" + emailAddress;
+  @Test
+  void readWithStringNativeQuery() {
+    setupForJoinQuery();
+    Read<Object[], String> read =
+        RegistryJpaIO.read(
+            "select d.repo_id, r.email_address from \"Domain\" d join \"Registrar\" r on"
+                + " d.current_sponsor_registrar_id = r.registrar_id where r.type = :type"
+                + " and d.deletion_time > now()",
+            ImmutableMap.of("type", "REAL"),
+            true,
+            (Object[] row) -> {
+              String repoId = (String) row[0];
+              String emailAddress = (String) row[1];
+              return repoId + "-" + emailAddress;
+            });
+    PCollection<String> joinedStrings = testPipeline.apply(read);
+
+    PAssert.that(joinedStrings).containsInAnyOrder("4-COM-me@google.com");
+    testPipeline.run();
+  }
+
+  @Test
+  void readWithStringTypedQuery() {
+    setupForJoinQuery();
+    Read<DomainBase, String> read =
+        RegistryJpaIO.read(
+            "select d from Domain d join Registrar r on"
+                + " d.currentSponsorClientId = r.clientIdentifier where r.type = :type"
+                + " and d.deletionTime > now()",
+            ImmutableMap.of("type", Registrar.Type.REAL),
+            DomainBase.class,
+            DomainBase::getRepoId);
+    PCollection<String> repoIds = testPipeline.apply(read);
+
+    PAssert.that(repoIds).containsInAnyOrder("4-COM");
+    testPipeline.run();
   }
 
   private void setupForJoinQuery() {

core/src/test/java/google/registry/beam/rde/RdePipelineTest.java

@@ -0,0 +1,257 @@
+// 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.base.Preconditions.checkState;
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
+import static com.google.common.truth.Truth.assertThat;
+import static google.registry.beam.rde.RdePipeline.decodePendings;
+import static google.registry.beam.rde.RdePipeline.encodePendings;
+import static google.registry.model.common.Cursor.CursorType.RDE_STAGING;
+import static google.registry.model.rde.RdeMode.FULL;
+import static google.registry.model.rde.RdeMode.THIN;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.setTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.rde.RdeResourceType.CONTACT;
+import static google.registry.rde.RdeResourceType.DOMAIN;
+import static google.registry.rde.RdeResourceType.HOST;
+import static google.registry.rde.RdeResourceType.REGISTRAR;
+import static google.registry.testing.AppEngineExtension.loadInitialData;
+import static google.registry.testing.DatabaseHelper.createTld;
+import static google.registry.testing.DatabaseHelper.newHostResource;
+import static google.registry.testing.DatabaseHelper.persistActiveContact;
+import static google.registry.testing.DatabaseHelper.persistActiveDomain;
+import static google.registry.testing.DatabaseHelper.persistDeletedDomain;
+import static google.registry.testing.DatabaseHelper.persistEppResource;
+import static google.registry.testing.DatabaseHelper.persistNewRegistrar;
+import static org.joda.time.Duration.standardDays;
+
+import com.google.common.base.Splitter;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSetMultimap;
+import com.google.common.collect.Iterables;
+import com.google.common.collect.Streams;
+import google.registry.beam.TestPipelineExtension;
+import google.registry.model.host.HostResource;
+import google.registry.model.registrar.Registrar;
+import google.registry.model.registrar.Registrar.State;
+import google.registry.persistence.transaction.JpaTestRules;
+import google.registry.persistence.transaction.JpaTestRules.JpaIntegrationTestExtension;
+import google.registry.persistence.transaction.TransactionManager;
+import google.registry.rde.DepositFragment;
+import google.registry.rde.PendingDeposit;
+import google.registry.rde.RdeResourceType;
+import google.registry.testing.DatastoreEntityExtension;
+import google.registry.testing.FakeClock;
+import java.util.function.Function;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+import java.util.stream.Stream;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.testing.PAssert;
+import org.apache.beam.sdk.transforms.GroupByKey;
+import org.apache.beam.sdk.values.KV;
+import org.joda.time.DateTime;
+import org.joda.time.Duration;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Order;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+/** Unit tests for {@link RdePipeline}. */
+public class RdePipelineTest {
+
+  private static final String REGISTRAR_NAME_PATTERN =
+      "<rdeRegistrar:name>(.*)</rdeRegistrar:name>";
+
+  private static final String DOMAIN_NAME_PATTERN = "<rdeDomain:name>(.*)</rdeDomain:name>";
+
+  private static final String CONTACT_ID_PATTERN = "<rdeContact:id>(.*)</rdeContact:id>";
+
+  private static final String HOST_NAME_PATTERN = "<rdeHost:name>(.*)</rdeHost:name>";
+
+  private static final ImmutableSetMultimap<String, PendingDeposit> PENDINGS =
+      ImmutableSetMultimap.of(
+          "pal",
+          PendingDeposit.create(
+              "pal", DateTime.parse("2000-01-01TZ"), FULL, RDE_STAGING, standardDays(1)),
+          "pal",
+          PendingDeposit.create(
+              "pal", DateTime.parse("2000-01-01TZ"), THIN, RDE_STAGING, standardDays(1)),
+          "fun",
+          PendingDeposit.create(
+              "fun", DateTime.parse("2000-01-01TZ"), FULL, RDE_STAGING, standardDays(1)));
+
+  // This is the default creation time for test data.
+  private final FakeClock clock = new FakeClock(DateTime.parse("1999-12-31TZ"));
+
+  // The pipeline runs in a different thread, which needs to be masqueraded as a GAE thread as well.
+  @RegisterExtension
+  @Order(Order.DEFAULT - 1)
+  final DatastoreEntityExtension datastore = new DatastoreEntityExtension().allThreads(true);
+
+  @RegisterExtension
+  final JpaIntegrationTestExtension database =
+      new JpaTestRules.Builder().withClock(clock).buildIntegrationTestRule();
+
+  @RegisterExtension
+  final TestPipelineExtension pipeline =
+      TestPipelineExtension.create().enableAbandonedNodeEnforcement(true);
+
+  private final RdePipelineOptions options =
+      PipelineOptionsFactory.create().as(RdePipelineOptions.class);
+
+  private RdePipeline rdePipeline;
+
+  private TransactionManager originalTm;
+
+  @BeforeEach
+  void beforeEach() throws Exception {
+    originalTm = tm();
+    setTm(jpaTm());
+    loadInitialData();
+
+    // Two real registrars have been created by loadInitialData(), named "New Registrar" and "The
+    // Registrar". Create one included registrar (external_monitoring) and two excluded ones.
+    Registrar monitoringRegistrar =
+        persistNewRegistrar("monitoring", "monitoring", Registrar.Type.MONITORING, null);
+    Registrar testRegistrar = persistNewRegistrar("test", "test", Registrar.Type.TEST, null);
+    Registrar externalMonitoringRegistrar =
+        persistNewRegistrar(
+            "externalmonitor", "external_monitoring", Registrar.Type.EXTERNAL_MONITORING, 9997L);
+    // Set Registrar states which are required for reporting.
+    tm().transact(
+            () ->
+                tm().putAll(
+                        ImmutableList.of(
+                            externalMonitoringRegistrar.asBuilder().setState(State.ACTIVE).build(),
+                            testRegistrar.asBuilder().setState(State.ACTIVE).build(),
+                            monitoringRegistrar.asBuilder().setState(State.ACTIVE).build())));
+
+    createTld("pal");
+    createTld("fun");
+    createTld("cat");
+
+    // Also persists a "contact1234" contact in the process.
+    persistActiveDomain("hello.pal");
+    persistActiveDomain("kitty.fun");
+    // Should not appear anywhere.
+    persistActiveDomain("lol.cat");
+    persistDeletedDomain("deleted.pal", DateTime.parse("1999-12-30TZ"));
+
+    persistActiveContact("contact456");
+
+    HostResource host = persistEppResource(newHostResource("old.host.test"));
+    // Set the clock to 2000-01-02, the updated host should NOT show up in RDE.
+    clock.advanceBy(Duration.standardDays(2));
+    persistEppResource(host.asBuilder().setHostName("new.host.test").build());
+
+    options.setPendings(encodePendings(PENDINGS));
+    // The EPP resources created in tests do not have all the fields populated, using a STRICT
+    // validation mode will result in a lot of warnings during marshalling.
+    options.setValidationMode("LENIENT");
+    rdePipeline = new RdePipeline(options);
+  }
+
+  @AfterEach
+  void afterEach() {
+    setTm(originalTm);
+  }
+
+  @Test
+  void testSuccess_encodeAndDecodePendingsMap() throws Exception {
+    String encodedString = encodePendings(PENDINGS);
+    assertThat(decodePendings(encodedString)).isEqualTo(PENDINGS);
+  }
+
+  @Test
+  void testSuccess_createFragments() {
+    PAssert.that(
+            rdePipeline
+                .createFragments(pipeline)
+                .apply("Group by PendingDeposit", GroupByKey.create()))
+        .satisfies(
+            kvs -> {
+              kvs.forEach(
+                  kv -> {
+                    // Registrar fragments.
+                    assertThat(
+                            getFragmentForType(kv, REGISTRAR)
+                                .map(getXmlElement(REGISTRAR_NAME_PATTERN))
+                                .collect(toImmutableSet()))
+                        // The same registrars are attached to all the pending deposits.
+                        .containsExactly("New Registrar", "The Registrar", "external_monitoring");
+                    // Domain fragments.
+                    assertThat(
+                            getFragmentForType(kv, DOMAIN)
+                                .map(getXmlElement(DOMAIN_NAME_PATTERN))
+                                .anyMatch(
+                                    domain ->
+                                        // Deleted domain should not be included
+                                        domain.equals("deleted.pal")
+                                            // Only domains on the pending deposit's tld should
+                                            // appear.
+                                            || !kv.getKey()
+                                                .tld()
+                                                .equals(
+                                                    Iterables.get(
+                                                        Splitter.on('.').split(domain), 1))))
+                        .isFalse();
+                    if (kv.getKey().mode().equals(FULL)) {
+                      // Contact fragments.
+                      assertThat(
+                              getFragmentForType(kv, CONTACT)
+                                  .map(getXmlElement(CONTACT_ID_PATTERN))
+                                  .collect(toImmutableSet()))
+                          // The same contacts are attached too all pending deposits.
+                          .containsExactly("contact1234", "contact456");
+                      // Host fragments.
+                      assertThat(
+                              getFragmentForType(kv, HOST)
+                                  .map(getXmlElement(HOST_NAME_PATTERN))
+                                  .collect(toImmutableSet()))
+                          // Should load the resource before update.
+                          .containsExactly("old.host.test");
+                    } else {
+                      // BRDA does not contain contact or hosts.
+                      assertThat(
+                              Streams.stream(kv.getValue())
+                                  .anyMatch(
+                                      fragment ->
+                                          fragment.type().equals(CONTACT)
+                                              || fragment.type().equals(HOST)))
+                          .isFalse();
+                    }
+                  });
+              return null;
+            });
+    pipeline.run().waitUntilFinish();
+  }
+
+  private static Function<DepositFragment, String> getXmlElement(String pattern) {
+    return (fragment) -> {
+      Matcher matcher = Pattern.compile(pattern).matcher(fragment.xml());
+      checkState(matcher.find(), "Missing %s in xml.", pattern);
+      return matcher.group(1);
+    };
+  }
+
+  private static Stream<DepositFragment> getFragmentForType(
+      KV<PendingDeposit, Iterable<DepositFragment>> kv, RdeResourceType type) {
+    return Streams.stream(kv.getValue()).filter(fragment -> fragment.type().equals(type));
+  }
+}

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

@@ -18,6 +18,7 @@ import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.ImmutableObjectSubject.immutableObjectCorrespondence;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.setTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.AppEngineExtension.makeRegistrar1;
 import static google.registry.testing.DatabaseHelper.createTld;
@@ -49,7 +50,6 @@ import google.registry.model.reporting.Spec11ThreatMatchDao;
 import google.registry.persistence.transaction.JpaTestRules;
 import google.registry.persistence.transaction.JpaTestRules.JpaIntegrationTestExtension;
 import google.registry.persistence.transaction.TransactionManager;
-import google.registry.persistence.transaction.TransactionManagerFactory;
 import google.registry.testing.DatastoreEntityExtension;
 import google.registry.testing.FakeClock;
 import google.registry.testing.FakeSleeper;
@@ -71,7 +71,6 @@ import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.DateTime;
 import org.joda.time.LocalDate;
 import org.json.JSONObject;
-import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Order;
 import org.junit.jupiter.api.Test;
@@ -113,8 +112,6 @@ class Spec11PipelineTest {
           ThreatMatch.create("THREAT_TYPE_UNSPECIFIED", "no-eamil.com"),
           ThreatMatch.create("UNWANTED_SOFTWARE", "anti-anti-anti-virus.dev"));
 
-  // This extension is only needed because Spec11ThreatMatch uses Ofy to generate the ID. Can be
-  // removed after the SQL migration.
   @RegisterExtension
   @Order(Order.DEFAULT - 1)
   final transient DatastoreEntityExtension datastore = new DatastoreEntityExtension();
@@ -136,12 +133,9 @@ class Spec11PipelineTest {
   private PCollection<KV<Subdomain, ThreatMatch>> threatMatches;
 
   ImmutableSet<Spec11ThreatMatch> sqlThreatMatches;
-  TransactionManager tm;
 
   @BeforeEach
   void beforeEach() throws Exception {
-    tm = tm();
-    TransactionManagerFactory.setTm(jpaTm());
     reportingBucketUrl = Files.createDirectory(tmpDir.resolve(REPORTING_BUCKET_URL)).toFile();
     options.setDate(DATE);
     options.setSafeBrowsingApiKey(SAFE_BROWSING_API_KEY);
@@ -196,11 +190,6 @@ class Spec11PipelineTest {
                 .build());
   }
 
-  @AfterEach
-  void afterEach() {
-    TransactionManagerFactory.setTm(tm);
-  }
-
   @Test
   void testSuccess_fullSqlPipeline() throws Exception {
     setupCloudSql();
@@ -241,6 +230,8 @@ class Spec11PipelineTest {
   }
 
   private void setupCloudSql() {
+    TransactionManager originalTm = tm();
+    setTm(jpaTm());
     persistNewRegistrar("TheRegistrar");
     persistNewRegistrar("NewRegistrar");
     Registrar registrar1 =
@@ -280,6 +271,7 @@ class Spec11PipelineTest {
     persistResource(createDomain("no-email.com", "2A4BA9BBC-COM", registrar2, contact2));
     persistResource(
         createDomain("anti-anti-anti-virus.dev", "555666888-DEV", registrar3, contact3));
+    setTm(originalTm);
   }
 
   private void verifySaveToGcs() throws Exception {

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

@@ -145,7 +145,7 @@ class EppPointInTimeTest {
     tm().clearSessionCache();
     assertAboutImmutableObjects()
         .that(loadAtPointInTime(latest, timeAtCreate.plusDays(1)))
-        .hasFieldsEqualTo(domainAfterCreate);
+        .isEqualExceptFields(domainAfterCreate, "updateTimestamp");
 
     tm().clearSessionCache();
     if (tm().isOfy()) {
@@ -159,18 +159,18 @@ class EppPointInTimeTest {
       // second update occurred one millisecond later.
       assertAboutImmutableObjects()
           .that(loadAtPointInTime(latest, timeAtFirstUpdate))
-          .hasFieldsEqualTo(domainAfterFirstUpdate);
+          .isEqualExceptFields(domainAfterFirstUpdate, "updateTimestamp");
     }
 
     tm().clearSessionCache();
     assertAboutImmutableObjects()
         .that(loadAtPointInTime(latest, timeAtSecondUpdate))
-        .hasFieldsEqualTo(domainAfterSecondUpdate);
+        .isEqualExceptFields(domainAfterSecondUpdate, "updateTimestamp");
 
     tm().clearSessionCache();
     assertAboutImmutableObjects()
         .that(loadAtPointInTime(latest, timeAtSecondUpdate.plusDays(1)))
-        .hasFieldsEqualTo(domainAfterSecondUpdate);
+        .isEqualExceptFields(domainAfterSecondUpdate, "updateTimestamp");
 
     // Deletion time has millisecond granularity due to isActive() check.
     tm().clearSessionCache();

core/src/test/java/google/registry/flows/host/HostInfoFlowTest.java

@@ -51,7 +51,7 @@ class HostInfoFlowTest extends ResourceFlowTestCase<HostInfoFlow, HostResource>
 
   @Order(value = Order.DEFAULT - 2)
   @RegisterExtension
-  final ReplayExtension replayExtension = ReplayExtension.createWithCompare(clock);
+  final ReplayExtension replayExtension = ReplayExtension.createWithDoubleReplay(clock);
 
   HostInfoFlowTest() {
     setEppInput("host_info.xml", ImmutableMap.of("HOSTNAME", "ns1.example.tld"));

core/src/test/java/google/registry/model/contact/ContactResourceTest.java

@@ -122,6 +122,13 @@ public class ContactResourceTest extends EntityTestCase {
     contactResource = persistResource(cloneAndSetAutoTimestamps(originalContact));
   }
 
+  @Test
+  void testContactBaseToContactResource() {
+    ImmutableObjectSubject.assertAboutImmutableObjects()
+        .that(new ContactResource.Builder().copyFrom(contactResource).build())
+        .isEqualExceptFields(contactResource, "updateTimestamp", "revisions");
+  }
+
   @Test
   void testCloudSqlPersistence_failWhenViolateForeignKeyConstraint() {
     assertThrowForeignKeyViolation(() -> jpaTm().transact(() -> jpaTm().insert(originalContact)));

core/src/test/java/google/registry/model/domain/DomainBaseTest.java

@@ -39,6 +39,7 @@ import com.google.common.collect.Streams;
 import com.googlecode.objectify.Key;
 import google.registry.model.EntityTestCase;
 import google.registry.model.ImmutableObject;
+import google.registry.model.ImmutableObjectSubject;
 import google.registry.model.billing.BillingEvent;
 import google.registry.model.billing.BillingEvent.Reason;
 import google.registry.model.contact.ContactResource;
@@ -171,6 +172,13 @@ public class DomainBaseTest extends EntityTestCase {
                     .build()));
   }
 
+  @Test
+  void testDomainContentToDomainBase() {
+    ImmutableObjectSubject.assertAboutImmutableObjects()
+        .that(new DomainBase.Builder().copyFrom(domain).build())
+        .isEqualExceptFields(domain, "updateTimestamp", "revisions");
+  }
+
   @Test
   void testPersistence() {
     // Note that this only verifies that the value stored under the foreign key is the same as that

core/src/test/java/google/registry/model/host/HostResourceTest.java

@@ -31,6 +31,7 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.net.InetAddresses;
 import google.registry.model.EntityTestCase;
+import google.registry.model.ImmutableObjectSubject;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.eppcommon.StatusValue;
 import google.registry.model.eppcommon.Trid;
@@ -89,6 +90,13 @@ class HostResourceTest extends EntityTestCase {
                     .build()));
   }
 
+  @TestOfyAndSql
+  void testHostBaseToHostResource() {
+    ImmutableObjectSubject.assertAboutImmutableObjects()
+        .that(new HostResource.Builder().copyFrom(host).build())
+        .isEqualExceptFields(host, "updateTimestamp", "revisions");
+  }
+
   @TestOfyAndSql
   void testPersistence() {
     HostResource newHost = host.asBuilder().setRepoId("NEWHOST").build();

core/src/test/java/google/registry/model/registry/label/PremiumListTest.java

@@ -20,6 +20,8 @@ import static google.registry.testing.DatabaseHelper.createTld;
 import static google.registry.testing.DatabaseHelper.persistPremiumList;
 import static google.registry.testing.DatabaseHelper.persistReservedList;
 import static google.registry.testing.DatabaseHelper.persistResource;
+import static org.joda.money.CurrencyUnit.JPY;
+import static org.joda.money.CurrencyUnit.USD;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
 import com.google.common.collect.ImmutableList;
@@ -28,6 +30,7 @@ import google.registry.model.registry.Registry;
 import google.registry.model.registry.label.PremiumList.PremiumListEntry;
 import google.registry.schema.tld.PremiumListDao;
 import google.registry.testing.AppEngineExtension;
+import java.math.BigDecimal;
 import org.joda.money.Money;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -125,4 +128,30 @@ public class PremiumListTest {
                     .build());
     assertThat(e).hasMessageThat().contains("must be in puny-coded, lower-case form");
   }
+
+  @Test
+  void testConvertAmountToMoney_USD() {
+    PremiumList premiumList = new PremiumList.Builder().setName("foo").setCurrency(USD).build();
+    assertThat(premiumList.convertAmountToMoney(new BigDecimal("20.000")))
+        .isEqualTo(Money.of(USD, new BigDecimal("20.00")));
+    assertThat(premiumList.convertAmountToMoney(new BigDecimal("37")))
+        .isEqualTo(Money.of(USD, new BigDecimal("37.00")));
+    assertThat(premiumList.convertAmountToMoney(new BigDecimal("42.5")))
+        .isEqualTo(Money.of(USD, new BigDecimal("42.50")));
+    assertThat(premiumList.convertAmountToMoney(new BigDecimal("15.678")))
+        .isEqualTo(Money.of(USD, new BigDecimal("15.68")));
+  }
+
+  @Test
+  void testConvertAmountToMoney_JPY() {
+    PremiumList premiumList = new PremiumList.Builder().setName("foo").setCurrency(JPY).build();
+    assertThat(premiumList.convertAmountToMoney(new BigDecimal("20.000")))
+        .isEqualTo(Money.of(JPY, new BigDecimal("20")));
+    assertThat(premiumList.convertAmountToMoney(new BigDecimal("37")))
+        .isEqualTo(Money.of(JPY, new BigDecimal("37")));
+    assertThat(premiumList.convertAmountToMoney(new BigDecimal("42.5")))
+        .isEqualTo(Money.of(JPY, new BigDecimal("42")));
+    assertThat(premiumList.convertAmountToMoney(new BigDecimal("15.678")))
+        .isEqualTo(Money.of(JPY, new BigDecimal("16")));
+  }
 }

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

@@ -46,6 +46,7 @@ import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junitpioneer.jupiter.RetryingTest;
 
 public class ReplicateToDatastoreActionTest {
 
@@ -97,7 +98,7 @@ public class ReplicateToDatastoreActionTest {
     RegistryConfig.overrideCloudSqlReplicateTransactions(false);
   }
 
-  @Test
+  @RetryingTest(4)
   public void testReplication() {
     TestEntity foo = new TestEntity("foo");
     TestEntity bar = new TestEntity("bar");
@@ -127,7 +128,7 @@ public class ReplicateToDatastoreActionTest {
     assertThat(ofyTm().transact(() -> ofyTm().loadByKey(baz.key()))).isEqualTo(baz);
   }
 
-  @Test
+  @RetryingTest(4)
   public void testReplayFromLastTxn() {
     TestEntity foo = new TestEntity("foo");
     TestEntity bar = new TestEntity("bar");
@@ -149,7 +150,7 @@ public class ReplicateToDatastoreActionTest {
     assertThat(ofyTm().transact(() -> ofyTm().loadByKeyIfPresent(foo.key()).isPresent())).isFalse();
   }
 
-  @Test
+  @RetryingTest(4)
   public void testUnintentionalConcurrency() {
     TestEntity foo = new TestEntity("foo");
     TestEntity bar = new TestEntity("bar");
@@ -184,7 +185,7 @@ public class ReplicateToDatastoreActionTest {
             Level.WARNING, "Ignoring transaction 1, which appears to have already been applied.");
   }
 
-  @Test
+  @RetryingTest(4)
   public void testMissingTransactions() {
     // Write a transaction (should have a transaction id of 1).
     TestEntity foo = new TestEntity("foo");

core/src/test/java/google/registry/schema/tld/PremiumListDaoTest.java

@@ -60,25 +60,24 @@ public class PremiumListDaoTest {
           .withPremiumListsCache(standardDays(1))
           .build();
 
-  private ImmutableMap<String, BigDecimal> testPrices;
+  private static final ImmutableMap<String, BigDecimal> TEST_PRICES =
+      ImmutableMap.of(
+          "silver",
+          BigDecimal.valueOf(10.23),
+          "gold",
+          BigDecimal.valueOf(1305.47),
+          "palladium",
+          BigDecimal.valueOf(1552.78));
 
   private PremiumList testList;
 
   @BeforeEach
   void beforeEach() {
-    testPrices =
-        ImmutableMap.of(
-            "silver",
-            BigDecimal.valueOf(10.23),
-            "gold",
-            BigDecimal.valueOf(1305.47),
-            "palladium",
-            BigDecimal.valueOf(1552.78));
     testList =
         new PremiumList.Builder()
             .setName("testname")
             .setCurrency(USD)
-            .setLabelsToPrices(testPrices)
+            .setLabelsToPrices(TEST_PRICES)
             .setCreationTime(fakeClock.nowUtc())
             .build();
   }
@@ -92,7 +91,7 @@ public class PremiumListDaoTest {
               Optional<PremiumList> persistedListOpt = PremiumListDao.getLatestRevision("testname");
               assertThat(persistedListOpt).isPresent();
               PremiumList persistedList = persistedListOpt.get();
-              assertThat(persistedList.getLabelsToPrices()).containsExactlyEntriesIn(testPrices);
+              assertThat(persistedList.getLabelsToPrices()).containsExactlyEntriesIn(TEST_PRICES);
               assertThat(persistedList.getCreationTime()).isEqualTo(fakeClock.nowUtc());
             });
   }
@@ -155,15 +154,15 @@ public class PremiumListDaoTest {
     PremiumListDao.save(
         new PremiumList.Builder()
             .setName("list1")
-            .setCurrency(JPY)
+            .setCurrency(USD)
             .setLabelsToPrices(ImmutableMap.of("wrong", BigDecimal.valueOf(1000.50)))
             .setCreationTime(fakeClock.nowUtc())
             .build());
     PremiumListDao.save(
         new PremiumList.Builder()
             .setName("list1")
-            .setCurrency(JPY)
-            .setLabelsToPrices(testPrices)
+            .setCurrency(USD)
+            .setLabelsToPrices(TEST_PRICES)
             .setCreationTime(fakeClock.nowUtc())
             .build());
     jpaTm()
@@ -172,9 +171,33 @@ public class PremiumListDaoTest {
               Optional<PremiumList> persistedList = PremiumListDao.getLatestRevision("list1");
               assertThat(persistedList).isPresent();
               assertThat(persistedList.get().getName()).isEqualTo("list1");
-              assertThat(persistedList.get().getCurrency()).isEqualTo(JPY);
+              assertThat(persistedList.get().getCurrency()).isEqualTo(USD);
               assertThat(persistedList.get().getLabelsToPrices())
-                  .containsExactlyEntriesIn(testPrices);
+                  .containsExactlyEntriesIn(TEST_PRICES);
+            });
+  }
+
+  @Test
+  void getLabelsToPrices_worksForJpy() {
+    PremiumListDao.save(
+        new PremiumList.Builder()
+            .setName("list1")
+            .setCurrency(JPY)
+            .setLabelsToPrices(TEST_PRICES)
+            .setCreationTime(fakeClock.nowUtc())
+            .build());
+    jpaTm()
+        .transact(
+            () -> {
+              PremiumList premiumList = PremiumListDao.getLatestRevision("list1").get();
+              assertThat(premiumList.getLabelsToPrices())
+                  .containsExactly(
+                      "silver",
+                      BigDecimal.valueOf(10),
+                      "gold",
+                      BigDecimal.valueOf(1305),
+                      "palladium",
+                      BigDecimal.valueOf(1553));
             });
   }
 
@@ -193,7 +216,7 @@ public class PremiumListDaoTest {
         new PremiumList.Builder()
             .setName("premlist")
             .setCurrency(USD)
-            .setLabelsToPrices(testPrices)
+            .setLabelsToPrices(TEST_PRICES)
             .setCreationTime(fakeClock.nowUtc())
             .build());
     assertThat(PremiumListDao.getPremiumPrice("premlist", "silver")).hasValue(Money.of(USD, 10.23));

core/src/test/java/google/registry/testing/AppEngineExtension.java

@@ -44,6 +44,7 @@ import com.google.common.collect.Sets;
 import com.google.common.io.Files;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.ObjectifyFilter;
+import google.registry.model.IdService;
 import google.registry.model.ofy.ObjectifyService;
 import google.registry.model.registrar.Registrar;
 import google.registry.model.registrar.Registrar.State;
@@ -466,7 +467,7 @@ public final class AppEngineExtension implements BeforeEachCallback, AfterEachCa
     if (withDatastore) {
       ObjectifyService.initOfy();
       // Reset id allocation in ObjectifyService so that ids are deterministic in tests.
-      ObjectifyService.resetNextTestId();
+      IdService.resetNextTestId();
       this.ofyTestEntities.forEach(AppEngineExtension::register);
     }
   }
@@ -476,7 +477,7 @@ public final class AppEngineExtension implements BeforeEachCallback, AfterEachCa
   public void afterEach(ExtensionContext context) throws Exception {
     checkArgumentNotNull(context, "The ExtensionContext must not be null");
     try {
-      // If there is a replay extension, we'll want to call its replayToSql() method.
+      // If there is a replay extension, we'll want to call its replay() method.
       //
       // We have to provide this hook here for ReplayExtension instead of relying on
       // ReplayExtension's afterEach() method because of ordering and the conflation of environment
@@ -491,7 +492,7 @@ public final class AppEngineExtension implements BeforeEachCallback, AfterEachCa
           (ReplayExtension)
               context.getStore(ExtensionContext.Namespace.GLOBAL).get(ReplayExtension.class);
       if (replayer != null) {
-        replayer.replayToSql();
+        replayer.replay();
       }
 
       if (withCloudSql) {

core/src/test/java/google/registry/testing/DatabaseHelper.java

@@ -28,10 +28,10 @@ import static google.registry.config.RegistryConfig.getContactAndHostRoidSuffix;
 import static google.registry.config.RegistryConfig.getContactAutomaticTransferLength;
 import static google.registry.model.EppResourceUtils.createDomainRepoId;
 import static google.registry.model.EppResourceUtils.createRepoId;
+import static google.registry.model.IdService.allocateId;
 import static google.registry.model.ImmutableObjectSubject.assertAboutImmutableObjects;
 import static google.registry.model.ImmutableObjectSubject.immutableObjectCorrespondence;
 import static google.registry.model.ResourceTransferUtils.createTransferResponse;
-import static google.registry.model.ofy.ObjectifyService.allocateId;
 import static google.registry.model.ofy.ObjectifyService.auditedOfy;
 import static google.registry.model.registry.Registry.TldState.GENERAL_AVAILABILITY;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
@@ -90,7 +90,6 @@ import google.registry.model.host.HostResource;
 import google.registry.model.index.EppResourceIndex;
 import google.registry.model.index.EppResourceIndexBucket;
 import google.registry.model.index.ForeignKeyIndex;
-import google.registry.model.ofy.ObjectifyService;
 import google.registry.model.poll.PollMessage;
 import google.registry.model.pricing.StaticPremiumListPricingEngine;
 import google.registry.model.registrar.Registrar;
@@ -957,12 +956,12 @@ public class DatabaseHelper {
 
   /** Returns a newly allocated, globally unique domain repoId of the format HEX-TLD. */
   public static String generateNewDomainRoid(String tld) {
-    return createDomainRepoId(ObjectifyService.allocateId(), tld);
+    return createDomainRepoId(allocateId(), tld);
   }
 
   /** Returns a newly allocated, globally unique contact/host repoId of the format HEX_TLD-ROID. */
   public static String generateNewContactHostRoid() {
-    return createRepoId(ObjectifyService.allocateId(), getContactAndHostRoidSuffix());
+    return createRepoId(allocateId(), getContactAndHostRoidSuffix());
   }
 
   /**
@@ -1086,8 +1085,7 @@ public class DatabaseHelper {
                           .setClientId(resource.getCreationClientId())
                           .setType(getHistoryEntryType(resource))
                           .setModificationTime(tm().getTransactionTime())
-                          .build()
-                          .toChildHistoryEntity());
+                          .build());
               ofyTmOrDoNothing(
                   () -> tm().put(ForeignKeyIndex.create(resource, resource.getDeletionTime())));
             });

core/src/test/java/google/registry/testing/DatastoreEntityExtension.java

@@ -14,8 +14,12 @@
 
 package google.registry.testing;
 
+import static google.registry.model.ofy.ObjectifyService.auditedOfy;
+
 import com.google.apphosting.api.ApiProxy;
 import com.google.apphosting.api.ApiProxy.Environment;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
 import java.util.Map;
 import org.junit.jupiter.api.extension.AfterEachCallback;
 import org.junit.jupiter.api.extension.BeforeEachCallback;
@@ -23,8 +27,12 @@ import org.junit.jupiter.api.extension.ExtensionContext;
 import org.testcontainers.shaded.com.google.common.collect.ImmutableMap;
 
 /**
- * Allows instantiation of Datastore {@code Entity entities} without the heavyweight {@code
- * AppEngineRule}.
+ * Allows instantiation of Datastore {@code Entity}s without the heavyweight {@link
+ * AppEngineExtension}.
+ *
+ * <p>When an Ofy key is created, by calling the various Key.create() methods, whether the current
+ * executing thread is a GAE thread is checked, which this extension masquerades. This happens
+ * frequently when an entity which has the key of another entity as a field is instantiated.
  *
  * <p>When used together with {@code JpaIntegrationWithCoverageExtension} or @{@code
  * TestPipelineExtension}, this extension must be registered first. For consistency's sake, it is
@@ -38,15 +46,49 @@ public class DatastoreEntityExtension implements BeforeEachCallback, AfterEachCa
 
   private static final Environment PLACEHOLDER_ENV = new PlaceholderEnvironment();
 
-  @Override
-  public void beforeEach(ExtensionContext context) {
-    ApiProxy.setEnvironmentForCurrentThread(PLACEHOLDER_ENV);
+  private boolean allThreads = false;
+
+  /**
+   * Whether all threads should be masqueraded as GAE threads.
+   *
+   * <p>This is particularly useful when new threads are spawned during a test. For example when
+   * testing Beam pipelines, the test pipeline runs the transforms in separate threads than the test
+   * thread. If Ofy keys are created in transforms, this value needs to be set to true.
+   *
+   * <p>Warning: by setting this value to true, any thread spawned by the current JVM will be have
+   * the thread local property set to the placeholder value during the execution of the current
+   * test, including those running other tests in parallel. This may or may not cause an issue when
+   * other tests have {@link AppEngineExtension} registered, that creates a much more fully
+   * functional GAE test environment. Consider moving tests using this extension to {@code
+   * outcastTest} if necessary.
+   */
+  public DatastoreEntityExtension allThreads(boolean value) {
+    allThreads = value;
+    return this;
   }
 
   @Override
-  public void afterEach(ExtensionContext context) {
+  public void beforeEach(ExtensionContext context) {
+    ApiProxy.setEnvironmentForCurrentThread(PLACEHOLDER_ENV);
+    // In order to create keys for entities they must be registered with Ofy. Calling this method
+    // will load the ObjectifyService class, whose static initialization block registers all Ofy
+    // entities.
+    auditedOfy();
+    if (allThreads) {
+      ApiProxy.setEnvironmentFactory(() -> PLACEHOLDER_ENV);
+    }
+  }
+
+  @Override
+  public void afterEach(ExtensionContext context)
+      throws NoSuchMethodException, InvocationTargetException, IllegalAccessException {
     // Clear the cached instance.
-    ApiProxy.setEnvironmentForCurrentThread(null);
+    ApiProxy.clearEnvironmentForCurrentThread();
+    if (allThreads) {
+      Method method = ApiProxy.class.getDeclaredMethod("clearEnvironmentFactory");
+      method.setAccessible(true);
+      method.invoke(null);
+    }
   }
 
   private static final class PlaceholderEnvironment implements Environment {

core/src/test/java/google/registry/testing/DualDatabaseTestInvocationContextProvider.java

@@ -51,15 +51,23 @@ class DualDatabaseTestInvocationContextProvider implements TestTemplateInvocatio
     return true;
   }
 
+  /**
+   * Returns true if "context" is an objectify unit test.
+   *
+   * <p>Provided to allow ReplayExtension to make this determination.
+   */
+  static boolean inOfyContext(ExtensionContext context) {
+    return (DatabaseType) context.getStore(NAMESPACE).get(INJECTED_TM_SUPPLIER_KEY)
+        == DatabaseType.OFY;
+  }
+
   @Override
   public Stream<TestTemplateInvocationContext> provideTestTemplateInvocationContexts(
       ExtensionContext context) {
     TestTemplateInvocationContext ofyContext =
-        createInvocationContext(
-            context.getDisplayName() + " with Datastore", TransactionManagerFactory::ofyTm);
+        createInvocationContext(context.getDisplayName() + " with Datastore", DatabaseType.OFY);
     TestTemplateInvocationContext sqlContext =
-        createInvocationContext(
-            context.getDisplayName() + " with PostgreSQL", TransactionManagerFactory::jpaTm);
+        createInvocationContext(context.getDisplayName() + " with PostgreSQL", DatabaseType.JPA);
     Method testMethod = context.getTestMethod().orElseThrow(IllegalStateException::new);
     if (testMethod.isAnnotationPresent(TestOfyAndSql.class)) {
       return Stream.of(ofyContext, sqlContext);
@@ -74,7 +82,7 @@ class DualDatabaseTestInvocationContextProvider implements TestTemplateInvocatio
   }
 
   private TestTemplateInvocationContext createInvocationContext(
-      String name, Supplier<? extends TransactionManager> tmSupplier) {
+      String name, DatabaseType databaseType) {
     return new TestTemplateInvocationContext() {
       @Override
       public String getDisplayName(int invocationIndex) {
@@ -83,17 +91,17 @@ class DualDatabaseTestInvocationContextProvider implements TestTemplateInvocatio
 
       @Override
       public List<Extension> getAdditionalExtensions() {
-        return ImmutableList.of(new DatabaseSwitchInvocationContext(tmSupplier));
+        return ImmutableList.of(new DatabaseSwitchInvocationContext(databaseType));
       }
     };
   }
 
   private static class DatabaseSwitchInvocationContext implements TestInstancePostProcessor {
 
-    private Supplier<? extends TransactionManager> tmSupplier;
+    private DatabaseType databaseType;
 
-    private DatabaseSwitchInvocationContext(Supplier<? extends TransactionManager> tmSupplier) {
-      this.tmSupplier = tmSupplier;
+    private DatabaseSwitchInvocationContext(DatabaseType databaseType) {
+      this.databaseType = databaseType;
     }
 
     @Override
@@ -113,7 +121,7 @@ class DualDatabaseTestInvocationContextProvider implements TestTemplateInvocatio
         throw new IllegalStateException(
             "AppEngineExtension in @DualDatabaseTest test must set withDatastoreAndCloudSql()");
       }
-      context.getStore(NAMESPACE).put(INJECTED_TM_SUPPLIER_KEY, tmSupplier);
+      context.getStore(NAMESPACE).put(INJECTED_TM_SUPPLIER_KEY, databaseType);
     }
 
     private static ImmutableList<Field> getAppEngineExtensionFields(Class<?> clazz) {
@@ -144,10 +152,9 @@ class DualDatabaseTestInvocationContextProvider implements TestTemplateInvocatio
                 }
               });
       context.getStore(NAMESPACE).put(ORIGINAL_TM_KEY, tm());
-      Supplier<? extends TransactionManager> tmSupplier =
-          (Supplier<? extends TransactionManager>)
-              context.getStore(NAMESPACE).get(INJECTED_TM_SUPPLIER_KEY);
-      TransactionManagerFactory.setTm(tmSupplier.get());
+      DatabaseType databaseType =
+          (DatabaseType) context.getStore(NAMESPACE).get(INJECTED_TM_SUPPLIER_KEY);
+      TransactionManagerFactory.setTm(databaseType.getTm());
     }
   }
 
@@ -171,4 +178,20 @@ class DualDatabaseTestInvocationContextProvider implements TestTemplateInvocatio
     return testInstance.getClass().isAnnotationPresent(DualDatabaseTest.class)
         && isDeclaredTestMethod;
   }
+
+  private enum DatabaseType {
+    JPA(TransactionManagerFactory::jpaTm),
+    OFY(TransactionManagerFactory::ofyTm);
+
+    @SuppressWarnings("Immutable") // Supplier is immutable, but not annotated as such.
+    private final Supplier<? extends TransactionManager> supplier;
+
+    DatabaseType(Supplier<? extends TransactionManager> supplier) {
+      this.supplier = supplier;
+    }
+
+    TransactionManager getTm() {
+      return supplier.get();
+    }
+  }
 }

core/src/test/java/google/registry/testing/ReplayExtension.java

@@ -22,13 +22,17 @@ import com.google.common.base.Suppliers;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
+import google.registry.config.RegistryConfig;
 import google.registry.model.ImmutableObject;
 import google.registry.model.ofy.CommitLogBucket;
 import google.registry.model.ofy.ReplayQueue;
 import google.registry.model.ofy.TransactionInfo;
 import google.registry.persistence.VKey;
+import google.registry.persistence.transaction.TransactionEntity;
 import google.registry.schema.replay.DatastoreEntity;
+import google.registry.schema.replay.ReplicateToDatastoreAction;
 import java.util.Optional;
+import javax.annotation.Nullable;
 import org.junit.jupiter.api.extension.AfterEachCallback;
 import org.junit.jupiter.api.extension.BeforeEachCallback;
 import org.junit.jupiter.api.extension.ExtensionContext;
@@ -48,19 +52,27 @@ public class ReplayExtension implements BeforeEachCallback, AfterEachCallback {
 
   FakeClock clock;
   boolean compare;
+  boolean replayed = false;
+  boolean inOfyContext;
   InjectExtension injectExtension = new InjectExtension();
+  @Nullable ReplicateToDatastoreAction sqlToDsReplicator;
 
-  private ReplayExtension(FakeClock clock, boolean compare) {
+  private ReplayExtension(
+      FakeClock clock, boolean compare, @Nullable ReplicateToDatastoreAction sqlToDsReplicator) {
     this.clock = clock;
     this.compare = compare;
+    this.sqlToDsReplicator = sqlToDsReplicator;
   }
 
   public static ReplayExtension createWithCompare(FakeClock clock) {
-    return new ReplayExtension(clock, true);
+    return new ReplayExtension(clock, true, null);
   }
 
-  public static ReplayExtension createWithoutCompare(FakeClock clock) {
-    return new ReplayExtension(clock, false);
+  /**
+   * Create a replay extension that replays from SQL to cloud datastore when running in SQL mode.
+   */
+  public static ReplayExtension createWithDoubleReplay(FakeClock clock) {
+    return new ReplayExtension(clock, true, new ReplicateToDatastoreAction(clock));
   }
 
   @Override
@@ -74,16 +86,27 @@ public class ReplayExtension implements BeforeEachCallback, AfterEachCallback {
     DatabaseHelper.setClock(clock);
     DatabaseHelper.setAlwaysSaveWithBackup(true);
     ReplayQueue.clear();
+
+    // When running in JPA mode with double replay enabled, enable JPA transaction replication.
+    // Note that we can't just use isOfy() here because this extension gets run before the dual-test
+    // transaction manager gets injected.
+    inOfyContext = DualDatabaseTestInvocationContextProvider.inOfyContext(context);
+    if (sqlToDsReplicator != null && !inOfyContext) {
+      RegistryConfig.overrideCloudSqlReplicateTransactions(true);
+    }
+
     context.getStore(ExtensionContext.Namespace.GLOBAL).put(ReplayExtension.class, this);
   }
 
   @Override
   public void afterEach(ExtensionContext context) {
     // This ensures that we do the replay even if we're not called from AppEngineExtension.  It
-    // should be safe to call replayToSql() twice, as the replay queue should be empty the second
-    // time.
-    replayToSql();
+    // is safe to call replay() twice, as the method ensures idempotence.
+    replay();
     injectExtension.afterEach(context);
+    if (sqlToDsReplicator != null) {
+      RegistryConfig.overrideCloudSqlReplicateTransactions(false);
+    }
   }
 
   private static ImmutableSet<String> NON_REPLICATED_TYPES =
@@ -104,7 +127,26 @@ public class ReplayExtension implements BeforeEachCallback, AfterEachCallback {
           "ForeignKeyContactIndex",
           "ForeignKeyDomainIndex");
 
-  public void replayToSql() {
+  public void replay() {
+    if (!replayed) {
+      if (inOfyContext) {
+        replayToSql();
+      } else {
+        // Disable database backups.  For unknown reason, if we don't do this we get residual commit
+        // log entries that cause timestamp inversions in other tests.
+        DatabaseHelper.setAlwaysSaveWithBackup(false);
+
+        // Do the ofy replay.
+        replayToOfy();
+
+        // Clean out anything that ends up in the replay queue.
+        ReplayQueue.clear();
+      }
+      replayed = true;
+    }
+  }
+
+  private void replayToSql() {
     DatabaseHelper.setAlwaysSaveWithBackup(false);
     ImmutableMap<Key<?>, Object> changes = ReplayQueue.replay();
 
@@ -139,4 +181,18 @@ public class ReplayExtension implements BeforeEachCallback, AfterEachCallback {
       }
     }
   }
+
+  private void replayToOfy() {
+    if (sqlToDsReplicator == null) {
+      return;
+    }
+
+    // TODO(mmuller): Verify that all entities are the same across both databases.
+    for (TransactionEntity txn : sqlToDsReplicator.getTransactionBatch()) {
+      if (sqlToDsReplicator.applyTransaction(txn)) {
+        break;
+      }
+      clock.advanceOneMilli();
+    }
+  }
 }