Add a GCB job to delete stopped GAE versions (#1128)

Unfortunately, much of the time there's a bit of a circular dependency
in the object creation, e.g. the Domain object stores references to the
billing events which store references to the history object which
contains the Domain object. As a result, we allocate the history
object's ID before creating it, so that it can be referenced in the
other objects that store that reference, e.g. billing events.

In addition, we add a utility copyFrom method in HistoryEntry.Builder to
avoid unnecessary ID allocations.

build.gradle

@@ -24,7 +24,7 @@ buildscript {
   }
 
   dependencies {
-    classpath 'com.google.cloud.tools:appengine-gradle-plugin:2.0.1'
+    classpath 'com.google.cloud.tools:appengine-gradle-plugin:2.4.1'
     classpath 'net.ltgt.gradle:gradle-errorprone-plugin:0.6.1'
     classpath 'org.sonatype.aether:aether-api:1.13.1'
     classpath 'org.sonatype.aether:aether-impl:1.13.1'

buildSrc/src/main/java/google/registry/gradle/plugin/ProjectData.java

@@ -91,7 +91,7 @@ abstract class ProjectData {
       /** The task was actually run and has finished successfully. */
       SUCCESS,
       /** The task was up-to-date and successful, and hence didn't need to run again. */
-      UP_TO_DATE;
+      UP_TO_DATE
     }
 
     abstract String uniqueName();

common/src/testing/java/google/registry/testing/FakeClock.java

@@ -35,6 +35,8 @@ public final class FakeClock implements Clock {
   // threads should see a consistent flow.
   private final AtomicLong currentTimeMillis = new AtomicLong();
 
+  private volatile long autoIncrementStepMs;
+
   /** Creates a FakeClock that starts at START_OF_TIME. */
   public FakeClock() {
     this(START_OF_TIME);
@@ -48,7 +50,21 @@ public final class FakeClock implements Clock {
   /** Returns the current time. */
   @Override
   public DateTime nowUtc() {
-    return new DateTime(currentTimeMillis.get(), UTC);
+    return new DateTime(currentTimeMillis.addAndGet(autoIncrementStepMs), UTC);
+  }
+
+  /**
+   * Sets the increment applied to the clock whenever it is queried. The increment is zero by
+   * default: the clock is left unchanged when queried.
+   *
+   * <p>Passing a duration of zero to this method effectively unsets the auto increment mode.
+   *
+   * @param autoIncrementStep the new auto increment duration
+   * @return this
+   */
+  public FakeClock setAutoIncrementStep(ReadableDuration autoIncrementStep) {
+    this.autoIncrementStepMs = autoIncrementStep.getMillis();
+    return this;
   }
 
   /** Advances clock by one millisecond. */

config/nom_build.py

@@ -256,6 +256,7 @@ GRADLE_FLAGS = [
                'Specify a task to be excluded from execution.',
                True),
 ]
+
 def generate_gradle_properties() -> str:
     """Returns the expected contents of gradle.properties."""
     out = io.StringIO()
@@ -270,7 +271,7 @@ def generate_gradle_properties() -> str:
 def get_root() -> str:
     """Returns the root of the nomulus build tree."""
     cur_dir = os.getcwd()
-    if not os.path.exists(os.path.join(cur_dir, '.git')) or \
+    if not os.path.exists(os.path.join(cur_dir, 'buildSrc')) or \
        not os.path.exists(os.path.join(cur_dir, 'core')) or \
        not os.path.exists(os.path.join(cur_dir, 'gradle.properties')):
         raise Exception('You must run this script from the root directory')

core/build.gradle

@@ -330,6 +330,7 @@ dependencies {
   testCompile deps['org.junit.platform:junit-platform-suite-api']
   testCompile deps['org.mockito:mockito-core']
   testCompile deps['org.mockito:mockito-junit-jupiter']
+  testCompile 'org.checkerframework:checker-qual:3.9.1'
   runtime deps['org.postgresql:postgresql']
 
   // Indirect dependency found by undeclared-dependency check. Such
@@ -423,7 +424,7 @@ task jaxbToJava {
       }
     }
     execInBash(
-        'find . -name *.java -exec sed -i /\\*\\ \\<p\\>\\$/d {} +',
+        "find . -name *.java -exec sed -i -e '/" + /\* <p>$/ + "/d' {} +",
         generatedDir)
   }
 }
@@ -805,6 +806,14 @@ if (environment in ['alpha', 'crash']) {
           mainClass: 'google.registry.beam.datastore.BulkDeleteDatastorePipeline',
           metaData: 'google/registry/beam/bulk_delete_datastore_pipeline_metadata.json'
       ],
+      [
+          mainClass: 'google.registry.beam.spec11.Spec11Pipeline',
+          metaData: 'google/registry/beam/spec11_pipeline_metadata.json'
+      ],
+      [
+          mainClass: 'google.registry.beam.invoicing.InvoicingPipeline',
+          metaData: 'google/registry/beam/invoicing_pipeline_metadata.json'
+      ],
   ]
   project.tasks.create("stage_beam_pipelines") {
     doLast {

core/gradle/dependency-locks/compile.lockfile

@@ -253,10 +253,10 @@ org.postgresql:postgresql:42.2.18
 org.rnorth.duct-tape:duct-tape:1.0.8
 org.rnorth.visible-assertions:visible-assertions:2.1.2
 org.slf4j:slf4j-api:1.7.30
-org.testcontainers:database-commons:1.15.1
-org.testcontainers:jdbc:1.15.1
-org.testcontainers:postgresql:1.15.1
-org.testcontainers:testcontainers:1.15.1
+org.testcontainers:database-commons:1.15.2
+org.testcontainers:jdbc:1.15.2
+org.testcontainers:postgresql:1.15.2
+org.testcontainers:testcontainers:1.15.2
 org.threeten:threetenbp:1.5.0
 org.tukaani:xz:1.5
 org.w3c.css:sac:1.3

core/gradle/dependency-locks/compileClasspath.lockfile

@@ -246,10 +246,10 @@ org.postgresql:postgresql:42.2.18
 org.rnorth.duct-tape:duct-tape:1.0.8
 org.rnorth.visible-assertions:visible-assertions:2.1.2
 org.slf4j:slf4j-api:1.7.30
-org.testcontainers:database-commons:1.15.1
-org.testcontainers:jdbc:1.15.1
-org.testcontainers:postgresql:1.15.1
-org.testcontainers:testcontainers:1.15.1
+org.testcontainers:database-commons:1.15.2
+org.testcontainers:jdbc:1.15.2
+org.testcontainers:postgresql:1.15.2
+org.testcontainers:testcontainers:1.15.2
 org.threeten:threetenbp:1.5.0
 org.tukaani:xz:1.5
 org.w3c.css:sac:1.3

core/gradle/dependency-locks/default.lockfile

@@ -267,10 +267,10 @@ org.slf4j:jcl-over-slf4j:1.7.30
 org.slf4j:jul-to-slf4j:1.7.30
 org.slf4j:slf4j-api:1.7.30
 org.slf4j:slf4j-jdk14:1.7.28
-org.testcontainers:database-commons:1.15.1
-org.testcontainers:jdbc:1.15.1
-org.testcontainers:postgresql:1.15.1
-org.testcontainers:testcontainers:1.15.1
+org.testcontainers:database-commons:1.15.2
+org.testcontainers:jdbc:1.15.2
+org.testcontainers:postgresql:1.15.2
+org.testcontainers:testcontainers:1.15.2
 org.threeten:threetenbp:1.5.0
 org.tukaani:xz:1.5
 org.w3c.css:sac:1.3

core/gradle/dependency-locks/deploy_jar.lockfile

@@ -266,10 +266,10 @@ org.slf4j:jcl-over-slf4j:1.7.30
 org.slf4j:jul-to-slf4j:1.7.30
 org.slf4j:slf4j-api:1.7.30
 org.slf4j:slf4j-jdk14:1.7.28
-org.testcontainers:database-commons:1.15.1
-org.testcontainers:jdbc:1.15.1
-org.testcontainers:postgresql:1.15.1
-org.testcontainers:testcontainers:1.15.1
+org.testcontainers:database-commons:1.15.2
+org.testcontainers:jdbc:1.15.2
+org.testcontainers:postgresql:1.15.2
+org.testcontainers:testcontainers:1.15.2
 org.threeten:threetenbp:1.5.0
 org.tukaani:xz:1.5
 org.w3c.css:sac:1.3

core/gradle/dependency-locks/nonprodCompile.lockfile

@@ -253,10 +253,10 @@ org.postgresql:postgresql:42.2.18
 org.rnorth.duct-tape:duct-tape:1.0.8
 org.rnorth.visible-assertions:visible-assertions:2.1.2
 org.slf4j:slf4j-api:1.7.30
-org.testcontainers:database-commons:1.15.1
-org.testcontainers:jdbc:1.15.1
-org.testcontainers:postgresql:1.15.1
-org.testcontainers:testcontainers:1.15.1
+org.testcontainers:database-commons:1.15.2
+org.testcontainers:jdbc:1.15.2
+org.testcontainers:postgresql:1.15.2
+org.testcontainers:testcontainers:1.15.2
 org.threeten:threetenbp:1.5.0
 org.tukaani:xz:1.5
 org.w3c.css:sac:1.3

core/gradle/dependency-locks/nonprodCompileClasspath.lockfile

@@ -247,10 +247,10 @@ org.postgresql:postgresql:42.2.18
 org.rnorth.duct-tape:duct-tape:1.0.8
 org.rnorth.visible-assertions:visible-assertions:2.1.2
 org.slf4j:slf4j-api:1.7.30
-org.testcontainers:database-commons:1.15.1
-org.testcontainers:jdbc:1.15.1
-org.testcontainers:postgresql:1.15.1
-org.testcontainers:testcontainers:1.15.1
+org.testcontainers:database-commons:1.15.2
+org.testcontainers:jdbc:1.15.2
+org.testcontainers:postgresql:1.15.2
+org.testcontainers:testcontainers:1.15.2
 org.threeten:threetenbp:1.5.0
 org.tukaani:xz:1.5
 org.w3c.css:sac:1.3

core/gradle/dependency-locks/nonprodRuntime.lockfile

@@ -265,10 +265,10 @@ org.rnorth.visible-assertions:visible-assertions:2.1.2
 org.slf4j:jcl-over-slf4j:1.7.30
 org.slf4j:jul-to-slf4j:1.7.30
 org.slf4j:slf4j-api:1.7.30
-org.testcontainers:database-commons:1.15.1
-org.testcontainers:jdbc:1.15.1
-org.testcontainers:postgresql:1.15.1
-org.testcontainers:testcontainers:1.15.1
+org.testcontainers:database-commons:1.15.2
+org.testcontainers:jdbc:1.15.2
+org.testcontainers:postgresql:1.15.2
+org.testcontainers:testcontainers:1.15.2
 org.threeten:threetenbp:1.5.0
 org.tukaani:xz:1.5
 org.w3c.css:sac:1.3

core/gradle/dependency-locks/nonprodRuntimeClasspath.lockfile

@@ -265,10 +265,10 @@ org.rnorth.visible-assertions:visible-assertions:2.1.2
 org.slf4j:jcl-over-slf4j:1.7.30
 org.slf4j:jul-to-slf4j:1.7.30
 org.slf4j:slf4j-api:1.7.30
-org.testcontainers:database-commons:1.15.1
-org.testcontainers:jdbc:1.15.1
-org.testcontainers:postgresql:1.15.1
-org.testcontainers:testcontainers:1.15.1
+org.testcontainers:database-commons:1.15.2
+org.testcontainers:jdbc:1.15.2
+org.testcontainers:postgresql:1.15.2
+org.testcontainers:testcontainers:1.15.2
 org.threeten:threetenbp:1.5.0
 org.tukaani:xz:1.5
 org.w3c.css:sac:1.3

core/gradle/dependency-locks/runtime.lockfile

@@ -265,10 +265,10 @@ org.rnorth.visible-assertions:visible-assertions:2.1.2
 org.slf4j:jcl-over-slf4j:1.7.30
 org.slf4j:jul-to-slf4j:1.7.30
 org.slf4j:slf4j-api:1.7.30
-org.testcontainers:database-commons:1.15.1
-org.testcontainers:jdbc:1.15.1
-org.testcontainers:postgresql:1.15.1
-org.testcontainers:testcontainers:1.15.1
+org.testcontainers:database-commons:1.15.2
+org.testcontainers:jdbc:1.15.2
+org.testcontainers:postgresql:1.15.2
+org.testcontainers:testcontainers:1.15.2
 org.threeten:threetenbp:1.5.0
 org.tukaani:xz:1.5
 org.w3c.css:sac:1.3

core/gradle/dependency-locks/runtimeClasspath.lockfile

@@ -266,10 +266,10 @@ org.slf4j:jcl-over-slf4j:1.7.30
 org.slf4j:jul-to-slf4j:1.7.30
 org.slf4j:slf4j-api:1.7.30
 org.slf4j:slf4j-jdk14:1.7.28
-org.testcontainers:database-commons:1.15.1
-org.testcontainers:jdbc:1.15.1
-org.testcontainers:postgresql:1.15.1
-org.testcontainers:testcontainers:1.15.1
+org.testcontainers:database-commons:1.15.2
+org.testcontainers:jdbc:1.15.2
+org.testcontainers:postgresql:1.15.2
+org.testcontainers:testcontainers:1.15.2
 org.threeten:threetenbp:1.5.0
 org.tukaani:xz:1.5
 org.w3c.css:sac:1.3

core/gradle/dependency-locks/testCompile.lockfile

@@ -300,12 +300,12 @@ org.seleniumhq.selenium:selenium-remote-driver:3.141.59
 org.seleniumhq.selenium:selenium-safari-driver:3.141.59
 org.seleniumhq.selenium:selenium-support:3.141.59
 org.slf4j:slf4j-api:1.7.30
-org.testcontainers:database-commons:1.15.1
-org.testcontainers:jdbc:1.15.1
-org.testcontainers:junit-jupiter:1.15.1
-org.testcontainers:postgresql:1.15.1
-org.testcontainers:selenium:1.15.1
-org.testcontainers:testcontainers:1.15.1
+org.testcontainers:database-commons:1.15.2
+org.testcontainers:jdbc:1.15.2
+org.testcontainers:junit-jupiter:1.15.2
+org.testcontainers:postgresql:1.15.2
+org.testcontainers:selenium:1.15.2
+org.testcontainers:testcontainers:1.15.2
 org.threeten:threetenbp:1.5.0
 org.tukaani:xz:1.5
 org.w3c.css:sac:1.3

core/gradle/dependency-locks/testCompileClasspath.lockfile

@@ -294,12 +294,12 @@ org.seleniumhq.selenium:selenium-remote-driver:3.141.59
 org.seleniumhq.selenium:selenium-safari-driver:3.141.59
 org.seleniumhq.selenium:selenium-support:3.141.59
 org.slf4j:slf4j-api:1.7.30
-org.testcontainers:database-commons:1.15.1
-org.testcontainers:jdbc:1.15.1
-org.testcontainers:junit-jupiter:1.15.1
-org.testcontainers:postgresql:1.15.1
-org.testcontainers:selenium:1.15.1
-org.testcontainers:testcontainers:1.15.1
+org.testcontainers:database-commons:1.15.2
+org.testcontainers:jdbc:1.15.2
+org.testcontainers:junit-jupiter:1.15.2
+org.testcontainers:postgresql:1.15.2
+org.testcontainers:selenium:1.15.2
+org.testcontainers:testcontainers:1.15.2
 org.threeten:threetenbp:1.5.0
 org.tukaani:xz:1.5
 org.w3c.css:sac:1.3

core/gradle/dependency-locks/testRuntime.lockfile

@@ -313,12 +313,12 @@ org.seleniumhq.selenium:selenium-support:3.141.59
 org.slf4j:jcl-over-slf4j:1.7.30
 org.slf4j:jul-to-slf4j:1.7.30
 org.slf4j:slf4j-api:1.7.30
-org.testcontainers:database-commons:1.15.1
-org.testcontainers:jdbc:1.15.1
-org.testcontainers:junit-jupiter:1.15.1
-org.testcontainers:postgresql:1.15.1
-org.testcontainers:selenium:1.15.1
-org.testcontainers:testcontainers:1.15.1
+org.testcontainers:database-commons:1.15.2
+org.testcontainers:jdbc:1.15.2
+org.testcontainers:junit-jupiter:1.15.2
+org.testcontainers:postgresql:1.15.2
+org.testcontainers:selenium:1.15.2
+org.testcontainers:testcontainers:1.15.2
 org.threeten:threetenbp:1.5.0
 org.tukaani:xz:1.5
 org.w3c.css:sac:1.3

core/gradle/dependency-locks/testRuntimeClasspath.lockfile

@@ -314,12 +314,12 @@ org.slf4j:jcl-over-slf4j:1.7.30
 org.slf4j:jul-to-slf4j:1.7.30
 org.slf4j:slf4j-api:1.7.30
 org.slf4j:slf4j-jdk14:1.7.28
-org.testcontainers:database-commons:1.15.1
-org.testcontainers:jdbc:1.15.1
-org.testcontainers:junit-jupiter:1.15.1
-org.testcontainers:postgresql:1.15.1
-org.testcontainers:selenium:1.15.1
-org.testcontainers:testcontainers:1.15.1
+org.testcontainers:database-commons:1.15.2
+org.testcontainers:jdbc:1.15.2
+org.testcontainers:junit-jupiter:1.15.2
+org.testcontainers:postgresql:1.15.2
+org.testcontainers:selenium:1.15.2
+org.testcontainers:testcontainers:1.15.2
 org.threeten:threetenbp:1.5.0
 org.tukaani:xz:1.5
 org.w3c.css:sac:1.3

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

@@ -154,7 +154,13 @@ public class ReplayCommitLogsToSqlAction implements Runnable {
     Object ofyPojo = ofy().toPojo(entity);
     if (ofyPojo instanceof DatastoreEntity) {
       DatastoreEntity datastoreEntity = (DatastoreEntity) ofyPojo;
-      datastoreEntity.toSqlEntity().ifPresent(jpaTm()::put);
+      datastoreEntity
+          .toSqlEntity()
+          .ifPresent(
+              sqlEntity -> {
+                ReplaySpecializer.beforeSqlSave(sqlEntity);
+                jpaTm().put(sqlEntity);
+              });
     } else {
       // this should never happen, but we shouldn't fail on it
       logger.atSevere().log(

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

@@ -57,8 +57,6 @@ public final class AsyncTaskEnqueuer {
   public static final String QUEUE_ASYNC_DELETE = "async-delete-pull";
   public static final String QUEUE_ASYNC_HOST_RENAME = "async-host-rename-pull";
 
-  public static final String PATH_RESAVE_ENTITY = "/_dr/task/resaveEntity";
-
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
   private static final Duration MAX_ASYNC_ETA = Duration.standardDays(30);
 
@@ -112,7 +110,7 @@ public final class AsyncTaskEnqueuer {
     logger.atInfo().log("Enqueuing async re-save of %s to run at %s.", entityKey, whenToResave);
     String backendHostname = appEngineServiceUtils.getServiceHostname("backend");
     TaskOptions task =
-        TaskOptions.Builder.withUrl(PATH_RESAVE_ENTITY)
+        TaskOptions.Builder.withUrl(ResaveEntityAction.PATH)
             .method(Method.POST)
             .header("Host", backendHostname)
             .countdownMillis(etaDuration.getMillis())

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

@@ -25,7 +25,7 @@ import static google.registry.batch.AsyncTaskMetrics.OperationType.DNS_REFRESH;
 import static google.registry.mapreduce.inputs.EppResourceInputs.createEntityInput;
 import static google.registry.model.EppResourceUtils.isActive;
 import static google.registry.model.EppResourceUtils.isDeleted;
-import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.DateTimeUtils.latestOf;
 import static java.util.concurrent.TimeUnit.DAYS;
 import static java.util.concurrent.TimeUnit.SECONDS;
@@ -44,7 +44,6 @@ import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.flogger.FluentLogger;
-import com.googlecode.objectify.Key;
 import google.registry.batch.AsyncTaskMetrics.OperationResult;
 import google.registry.dns.DnsQueue;
 import google.registry.mapreduce.MapreduceRunner;
@@ -64,6 +63,7 @@ import google.registry.util.SystemClock;
 import java.io.Serializable;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.NoSuchElementException;
 import java.util.Optional;
 import java.util.logging.Level;
 import javax.annotation.Nullable;
@@ -123,7 +123,7 @@ public class RefreshDnsOnHostRenameAction implements Runnable {
     }
 
     ImmutableList.Builder<DnsRefreshRequest> requestsBuilder = new ImmutableList.Builder<>();
-    ImmutableList.Builder<Key<HostResource>> hostKeys = new ImmutableList.Builder<>();
+    ImmutableList.Builder<VKey<HostResource>> hostKeys = new ImmutableList.Builder<>();
     final List<DnsRefreshRequest> requestsToDelete = new ArrayList<>();
 
     for (TaskHandle task : tasks) {
@@ -204,10 +204,10 @@ public class RefreshDnsOnHostRenameAction implements Runnable {
         emit(true, true);
         return;
       }
-      Key<HostResource> referencingHostKey = null;
+      VKey<HostResource> referencingHostKey = null;
       for (DnsRefreshRequest request : refreshRequests) {
         if (isActive(domain, request.lastUpdateTime())
-            && domain.getNameservers().contains(VKey.from(request.hostKey()))) {
+            && domain.getNameservers().contains(request.hostKey())) {
           referencingHostKey = request.hostKey();
           break;
         }
@@ -293,7 +293,8 @@ public class RefreshDnsOnHostRenameAction implements Runnable {
 
     private static final long serialVersionUID = 1772812852271288622L;
 
-    abstract Key<HostResource> hostKey();
+    abstract VKey<HostResource> hostKey();
+
     abstract DateTime lastUpdateTime();
     abstract DateTime requestedTime();
     abstract boolean isRefreshNeeded();
@@ -301,7 +302,8 @@ public class RefreshDnsOnHostRenameAction implements Runnable {
 
     @AutoValue.Builder
     abstract static class Builder {
-      abstract Builder setHostKey(Key<HostResource> hostKey);
+      abstract Builder setHostKey(VKey<HostResource> hostKey);
+
       abstract Builder setLastUpdateTime(DateTime lastUpdateTime);
       abstract Builder setRequestedTime(DateTime requestedTime);
       abstract Builder setIsRefreshNeeded(boolean isRefreshNeeded);
@@ -314,10 +316,12 @@ public class RefreshDnsOnHostRenameAction implements Runnable {
      */
     static DnsRefreshRequest createFromTask(TaskHandle task, DateTime now) throws Exception {
       ImmutableMap<String, String> params = ImmutableMap.copyOf(task.extractParams());
-      Key<HostResource> hostKey =
-          Key.create(checkNotNull(params.get(PARAM_HOST_KEY), "Host to refresh not specified"));
+      VKey<HostResource> hostKey =
+          VKey.fromWebsafeKey(
+              checkNotNull(params.get(PARAM_HOST_KEY), "Host to refresh not specified"));
       HostResource host =
-          checkNotNull(ofy().load().key(hostKey).now(), "Host to refresh doesn't exist");
+          tm().transact(() -> tm().loadByKeyIfPresent(hostKey))
+              .orElseThrow(() -> new NoSuchElementException("Host to refresh doesn't exist"));
       boolean isHostDeleted =
           isDeleted(host, latestOf(now, host.getUpdateTimestamp().getTimestamp()));
       if (isHostDeleted) {

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

@@ -16,7 +16,6 @@ package google.registry.batch;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.request.Action.Method.POST;
@@ -33,6 +32,7 @@ import google.registry.model.eppcommon.StatusValue;
 import google.registry.model.registrar.Registrar;
 import google.registry.model.registrar.RegistrarContact;
 import google.registry.model.registry.RegistryLockDao;
+import google.registry.persistence.VKey;
 import google.registry.request.Action;
 import google.registry.request.Parameter;
 import google.registry.request.Response;
@@ -125,6 +125,7 @@ public class RelockDomainAction implements Runnable {
     response.setContentType(MediaType.PLAIN_TEXT_UTF_8);
 
     // nb: DomainLockUtils relies on the JPA transaction being the outermost transaction
+    // if we have Datastore as the primary DB (if SQL is the primary DB, it's irrelevant)
     jpaTm().transact(() -> tm().transact(this::relockDomain));
   }
 
@@ -139,12 +140,8 @@ public class RelockDomainAction implements Runnable {
                       new IllegalArgumentException(
                           String.format("Unknown revision ID %d", oldUnlockRevisionId)));
       domain =
-          ofy()
-              .load()
-              .type(DomainBase.class)
-              .id(oldLock.getRepoId())
-              .now()
-              .cloneProjectedAtTime(jpaTm().getTransactionTime());
+          tm().loadByKey(VKey.create(DomainBase.class, oldLock.getRepoId()))
+              .cloneProjectedAtTime(tm().getTransactionTime());
     } catch (Throwable t) {
       handleTransientFailure(Optional.ofNullable(oldLock), t);
       return;

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

@@ -17,7 +17,6 @@ package google.registry.batch;
 import static google.registry.batch.AsyncTaskEnqueuer.PARAM_REQUESTED_TIME;
 import static google.registry.batch.AsyncTaskEnqueuer.PARAM_RESAVE_TIMES;
 import static google.registry.batch.AsyncTaskEnqueuer.PARAM_RESOURCE_KEY;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.common.collect.ImmutableSet;
@@ -26,6 +25,7 @@ import com.google.common.flogger.FluentLogger;
 import com.googlecode.objectify.Key;
 import google.registry.model.EppResource;
 import google.registry.model.ImmutableObject;
+import google.registry.persistence.VKey;
 import google.registry.request.Action;
 import google.registry.request.Action.Method;
 import google.registry.request.Parameter;
@@ -74,16 +74,17 @@ public class ResaveEntityAction implements Runnable {
   public void run() {
     logger.atInfo().log(
         "Re-saving entity %s which was enqueued at %s.", resourceKey, requestedTime);
-    tm().transact(() -> {
-      ImmutableObject entity = ofy().load().key(resourceKey).now();
-      ofy().save().entity(
-          (entity instanceof EppResource)
-              ? ((EppResource) entity).cloneProjectedAtTime(tm().getTransactionTime()) : entity
-      );
-      if (!resaveTimes.isEmpty()) {
-        asyncTaskEnqueuer.enqueueAsyncResave(entity, requestedTime, resaveTimes);
-      }
-    });
+    tm().transact(
+            () -> {
+              ImmutableObject entity = tm().loadByKey(VKey.from(resourceKey));
+              tm().put(
+                      (entity instanceof EppResource)
+                          ? ((EppResource) entity).cloneProjectedAtTime(tm().getTransactionTime())
+                          : entity);
+              if (!resaveTimes.isEmpty()) {
+                asyncTaskEnqueuer.enqueueAsyncResave(entity, requestedTime, resaveTimes);
+              }
+            });
     response.setPayload("Entity re-saved.");
   }
 }

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

@@ -19,6 +19,7 @@ import static javax.servlet.http.HttpServletResponse.SC_FORBIDDEN;
 import static javax.servlet.http.HttpServletResponse.SC_INTERNAL_SERVER_ERROR;
 import static javax.servlet.http.HttpServletResponse.SC_OK;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.flogger.FluentLogger;
 import google.registry.config.RegistryConfig.Config;
@@ -28,10 +29,11 @@ import google.registry.request.Response;
 import google.registry.request.auth.Auth;
 import google.registry.util.Retrier;
 import java.sql.Connection;
+import java.sql.ResultSet;
+import java.sql.SQLException;
 import java.sql.Statement;
 import java.util.function.Supplier;
 import javax.inject.Inject;
-import org.flywaydb.core.api.FlywayException;
 
 /**
  * Wipes out all Cloud SQL data in a Nomulus GCP environment.
@@ -80,13 +82,13 @@ public class WipeOutCloudSqlAction implements Runnable {
     try {
       retrier.callWithRetry(
           () -> {
-            try (Connection conn = connectionSupplier.get();
-                Statement statement = conn.createStatement()) {
-              statement.execute("drop owned by schema_deployer;");
+            try (Connection conn = connectionSupplier.get()) {
+              dropAllTables(conn, listTables(conn));
+              dropAllSequences(conn, listSequences(conn));
             }
             return null;
           },
-          e -> !(e instanceof FlywayException));
+          e -> !(e instanceof SQLException));
       response.setStatus(SC_OK);
       response.setPayload("Wiped out Cloud SQL in " + projectId);
     } catch (RuntimeException e) {
@@ -95,4 +97,69 @@ public class WipeOutCloudSqlAction implements Runnable {
       response.setPayload("Failed to wipe out Cloud SQL in " + projectId);
     }
   }
+
+  /** Returns a list of all tables in the public schema of a Postgresql database. */
+  static ImmutableList<String> listTables(Connection connection) throws SQLException {
+    try (ResultSet resultSet =
+        connection.getMetaData().getTables(null, null, null, new String[] {"TABLE"})) {
+      ImmutableList.Builder<String> tables = new ImmutableList.Builder<>();
+      while (resultSet.next()) {
+        String schema = resultSet.getString("TABLE_SCHEM");
+        if (schema == null || !schema.equalsIgnoreCase("public")) {
+          continue;
+        }
+        String tableName = resultSet.getString("TABLE_NAME");
+        tables.add("public.\"" + tableName + "\"");
+      }
+      return tables.build();
+    }
+  }
+
+  static void dropAllTables(Connection conn, ImmutableList<String> tables) throws SQLException {
+    if (tables.isEmpty()) {
+      return;
+    }
+
+    try (Statement statement = conn.createStatement()) {
+      for (String table : tables) {
+        statement.addBatch(String.format("DROP TABLE IF EXISTS %s CASCADE;", table));
+      }
+      for (int code : statement.executeBatch()) {
+        if (code == Statement.EXECUTE_FAILED) {
+          throw new RuntimeException("Failed to drop some tables. Please check.");
+        }
+      }
+    }
+  }
+
+  /** Returns a list of all sequences in a Postgresql database. */
+  static ImmutableList<String> listSequences(Connection conn) throws SQLException {
+    try (Statement statement = conn.createStatement();
+        ResultSet resultSet =
+            statement.executeQuery("SELECT c.relname FROM pg_class c WHERE c.relkind = 'S';")) {
+      ImmutableList.Builder<String> sequences = new ImmutableList.Builder<>();
+      while (resultSet.next()) {
+        sequences.add('\"' + resultSet.getString(1) + '\"');
+      }
+      return sequences.build();
+    }
+  }
+
+  static void dropAllSequences(Connection conn, ImmutableList<String> sequences)
+      throws SQLException {
+    if (sequences.isEmpty()) {
+      return;
+    }
+
+    try (Statement statement = conn.createStatement()) {
+      for (String sequence : sequences) {
+        statement.addBatch(String.format("DROP SEQUENCE IF EXISTS %s CASCADE;", sequence));
+      }
+      for (int code : statement.executeBatch()) {
+        if (code == Statement.EXECUTE_FAILED) {
+          throw new RuntimeException("Failed to drop some sequences. Please check.");
+        }
+      }
+    }
+  }
 }

core/src/main/java/google/registry/batch/WipeoutDatastoreAction.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.batch;
+
+import static com.google.common.net.MediaType.PLAIN_TEXT_UTF_8;
+import static google.registry.beam.BeamUtils.createJobName;
+import static javax.servlet.http.HttpServletResponse.SC_FORBIDDEN;
+import static javax.servlet.http.HttpServletResponse.SC_INTERNAL_SERVER_ERROR;
+import static javax.servlet.http.HttpServletResponse.SC_OK;
+
+import com.google.api.services.dataflow.Dataflow;
+import com.google.api.services.dataflow.model.LaunchFlexTemplateParameter;
+import com.google.api.services.dataflow.model.LaunchFlexTemplateRequest;
+import com.google.api.services.dataflow.model.LaunchFlexTemplateResponse;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.flogger.FluentLogger;
+import google.registry.config.RegistryConfig.Config;
+import google.registry.request.Action;
+import google.registry.request.Response;
+import google.registry.request.auth.Auth;
+import google.registry.util.Clock;
+import javax.inject.Inject;
+
+/**
+ * Wipes out all Cloud Datastore data in a Nomulus GCP environment.
+ *
+ * <p>This class is created for the QA environment, where migration testing with production data
+ * will happen. A regularly scheduled wipeout is a prerequisite to using production data there.
+ */
+@Action(
+    service = Action.Service.BACKEND,
+    path = "/_dr/task/wipeOutDatastore",
+    auth = Auth.AUTH_INTERNAL_OR_ADMIN)
+public class WipeoutDatastoreAction implements Runnable {
+  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
+
+  private static final String PIPELINE_NAME = "bulk_delete_datastore_pipeline";
+
+  // As a short-lived class, hardcode allowed projects here instead of using config files.
+  private static final ImmutableSet<String> ALLOWED_PROJECTS =
+      ImmutableSet.of("domain-registry-qa");
+
+  private final String projectId;
+  private final String jobRegion;
+  private final Response response;
+  private final Dataflow dataflow;
+  private final String stagingBucketUrl;
+  private final Clock clock;
+
+  @Inject
+  WipeoutDatastoreAction(
+      @Config("projectId") String projectId,
+      @Config("defaultJobRegion") String jobRegion,
+      @Config("beamStagingBucketUrl") String stagingBucketUrl,
+      Clock clock,
+      Response response,
+      Dataflow dataflow) {
+    this.projectId = projectId;
+    this.jobRegion = jobRegion;
+    this.stagingBucketUrl = stagingBucketUrl;
+    this.clock = clock;
+    this.response = response;
+    this.dataflow = dataflow;
+  }
+
+  @Override
+  public void run() {
+    response.setContentType(PLAIN_TEXT_UTF_8);
+
+    if (!ALLOWED_PROJECTS.contains(projectId)) {
+      response.setStatus(SC_FORBIDDEN);
+      response.setPayload("Wipeout is not allowed in " + projectId);
+      return;
+    }
+
+    try {
+      LaunchFlexTemplateParameter parameters =
+          new LaunchFlexTemplateParameter()
+              .setJobName(createJobName("bulk-delete-datastore-", clock))
+              .setContainerSpecGcsPath(
+                  String.format("%s/%s_metadata.json", stagingBucketUrl, PIPELINE_NAME))
+              .setParameters(ImmutableMap.of("kindsToDelete", "*"));
+      LaunchFlexTemplateResponse launchResponse =
+          dataflow
+              .projects()
+              .locations()
+              .flexTemplates()
+              .launch(
+                  projectId,
+                  jobRegion,
+                  new LaunchFlexTemplateRequest().setLaunchParameter(parameters))
+              .execute();
+      response.setStatus(SC_OK);
+      response.setPayload("Launched " + launchResponse.getJob().getName());
+    } catch (Exception e) {
+      String msg = String.format("Failed to launch %s.", PIPELINE_NAME);
+      logger.atSevere().withCause(e).log(msg);
+      response.setStatus(SC_INTERNAL_SERVER_ERROR);
+      response.setPayload(msg);
+    }
+  }
+}

core/src/main/java/google/registry/beam/BeamUtils.java

@@ -14,10 +14,14 @@
 
 package google.registry.beam;
 
+import static com.google.common.base.Preconditions.checkArgument;
+
 import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableList;
 import com.google.common.io.Resources;
+import google.registry.util.Clock;
 import google.registry.util.ResourceUtils;
+import java.util.regex.Pattern;
 import org.apache.avro.generic.GenericRecord;
 import org.apache.beam.sdk.io.gcp.bigquery.SchemaAndRecord;
 
@@ -41,8 +45,7 @@ public class BeamUtils {
       ImmutableList<String> fieldNames, SchemaAndRecord schemaAndRecord) {
     GenericRecord record = schemaAndRecord.getRecord();
     ImmutableList<String> nullFields =
-        fieldNames
-            .stream()
+        fieldNames.stream()
             .filter(fieldName -> record.get(fieldName) == null)
             .collect(ImmutableList.toImmutableList());
     String missingFieldList = Joiner.on(", ").join(nullFields);
@@ -61,4 +64,19 @@ public class BeamUtils {
   public static String getQueryFromFile(Class<?> clazz, String filename) {
     return ResourceUtils.readResourceUtf8(Resources.getResource(clazz, "sql/" + filename));
   }
+
+  /** Creates a beam job name and validates that it conforms to the requirements. */
+  public static String createJobName(String prefix, Clock clock) {
+    // Flex template job name must be unique and consists of only characters [-a-z0-9], starting
+    // with a letter and ending with a letter or number. So we replace the "T" and "Z" in ISO 8601
+    // with lowercase letters.
+    String jobName =
+        String.format("%s-%s", prefix, clock.nowUtc().toString("yyyy-MM-dd't'HH-mm-ss'z'"));
+    checkArgument(
+        Pattern.compile("^[a-z][-a-z0-9]*[a-z0-9]*").matcher(jobName).matches(),
+        "The job name %s is illegal, it consists of only characters [-a-z0-9], "
+            + "starting with a letter and ending with a letter or number,",
+        jobName);
+    return jobName;
+  }
 }

core/src/main/java/google/registry/beam/datastore/BulkDeleteDatastorePipeline.java

@@ -87,20 +87,18 @@ public class BulkDeleteDatastorePipeline {
 
   private final BulkDeletePipelineOptions options;
 
-  private final Pipeline pipeline;
-
   BulkDeleteDatastorePipeline(BulkDeletePipelineOptions options) {
     this.options = options;
-    pipeline = Pipeline.create(options);
   }
 
   public void run() {
-    setupPipeline();
+    Pipeline pipeline = Pipeline.create(options);
+    setupPipeline(pipeline);
     pipeline.run();
   }
 
   @SuppressWarnings("deprecation") // org.apache.beam.sdk.transforms.Reshuffle
-  private void setupPipeline() {
+  private void setupPipeline(Pipeline pipeline) {
     checkState(
         !FORBIDDEN_PROJECTS.contains(options.getProject()),
         "Bulk delete is forbidden in %s",

core/src/main/java/google/registry/beam/datastore/DatastoreV1.java

@@ -505,7 +505,7 @@ public class DatastoreV1 {
       }
 
       @StartBundle
-      public void startBundle(StartBundleContext c) throws Exception {
+      public void startBundle(StartBundleContext c) {
         datastore =
             datastoreFactory.getDatastore(
                 c.getPipelineOptions(), v1Options.getProjectId(), v1Options.getLocalhost());
@@ -548,7 +548,7 @@ public class DatastoreV1 {
       }
 
       @StartBundle
-      public void startBundle(StartBundleContext c) throws Exception {
+      public void startBundle(StartBundleContext c) {
         datastore =
             datastoreFactory.getDatastore(
                 c.getPipelineOptions(), options.getProjectId(), options.getLocalhost());
@@ -556,7 +556,7 @@ public class DatastoreV1 {
       }
 
       @ProcessElement
-      public void processElement(ProcessContext c) throws Exception {
+      public void processElement(ProcessContext c) {
         Query query = c.element();
 
         // If query has a user set limit, then do not split.
@@ -626,7 +626,7 @@ public class DatastoreV1 {
       }
 
       @StartBundle
-      public void startBundle(StartBundleContext c) throws Exception {
+      public void startBundle(StartBundleContext c) {
         datastore =
             datastoreFactory.getDatastore(
                 c.getPipelineOptions(), options.getProjectId(), options.getLocalhost());

core/src/main/java/google/registry/beam/initsql/BackupPaths.java

@@ -93,7 +93,7 @@ public final class BackupPaths {
     checkArgument(!isNullOrEmpty(exportDir), "Null or empty exportDir.");
     checkArgument(!isNullOrEmpty(kind), "Null or empty kind.");
     checkArgument(shard >= 0, "Negative shard %s not allowed.", shard);
-    return String.format(EXPORT_PATTERN_TEMPLATE, exportDir, kind, Integer.toString(shard));
+    return String.format(EXPORT_PATTERN_TEMPLATE, exportDir, kind, shard);
   }
 
   /** Returns an {@link ImmutableList} of regex patterns that match all CommitLog files. */

core/src/main/java/google/registry/beam/initsql/BeamJpaModule.java

@@ -1,201 +0,0 @@
-// Copyright 2020 The Nomulus Authors. All Rights Reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package google.registry.beam.initsql;
-
-import static com.google.common.base.Preconditions.checkArgument;
-import static com.google.common.base.Preconditions.checkState;
-import static com.google.common.base.Strings.isNullOrEmpty;
-
-import com.google.common.base.Splitter;
-import dagger.Component;
-import dagger.Lazy;
-import dagger.Module;
-import dagger.Provides;
-import google.registry.config.CredentialModule;
-import google.registry.config.RegistryConfig.Config;
-import google.registry.config.RegistryConfig.ConfigModule;
-import google.registry.keyring.kms.KmsModule;
-import google.registry.persistence.PersistenceModule;
-import google.registry.persistence.PersistenceModule.JdbcJpaTm;
-import google.registry.persistence.PersistenceModule.SocketFactoryJpaTm;
-import google.registry.persistence.PersistenceModule.TransactionIsolationLevel;
-import google.registry.persistence.transaction.JpaTransactionManager;
-import google.registry.privileges.secretmanager.SecretManagerModule;
-import google.registry.util.UtilsModule;
-import java.io.BufferedReader;
-import java.io.IOException;
-import java.io.InputStreamReader;
-import java.nio.channels.Channels;
-import java.nio.charset.StandardCharsets;
-import java.util.List;
-import javax.annotation.Nullable;
-import javax.inject.Singleton;
-import org.apache.beam.sdk.io.FileSystems;
-import org.apache.beam.sdk.io.fs.ResourceId;
-
-/**
- * Provides bindings for {@link JpaTransactionManager} to Cloud SQL.
- *
- * <p>This module is intended for use in BEAM pipelines, and uses a BEAM utility to access GCS like
- * a regular file system.
- */
-@Module
-public class BeamJpaModule {
-
-  private static final String GCS_SCHEME = "gs://";
-
-  @Nullable private final String sqlAccessInfoFile;
-  @Nullable private final String cloudKmsProjectId;
-  @Nullable private final TransactionIsolationLevel isolationOverride;
-
-  /**
-   * Constructs a new instance of {@link BeamJpaModule}.
-   *
-   * <p>Note: it is an unfortunately necessary antipattern to check for the validity of
-   * sqlAccessInfoFile in {@link #provideCloudSqlAccessInfo} rather than in the constructor.
-   * Unfortunately, this is a restriction imposed upon us by Dagger. Specifically, because we use
-   * this in at least one 1 {@link google.registry.tools.RegistryTool} command(s), it must be
-   * instantiated in {@code google.registry.tools.RegistryToolComponent} for all possible commands;
-   * Dagger doesn't permit it to ever be null. For the vast majority of commands, it will never be
-   * used (so a null credential file path is fine in those cases).
-   *
-   * @param sqlAccessInfoFile the path to a Cloud SQL credential file. This must refer to either a
-   *     real encrypted file on GCS as returned by {@link
-   *     BackupPaths#getCloudSQLCredentialFilePatterns} or an unencrypted file on local filesystem
-   *     with credentials to a test database.
-   * @param cloudKmsProjectId the GCP project where the credential decryption key can be found
-   * @param isolationOverride the desired Transaction Isolation level for all JDBC connections
-   */
-  public BeamJpaModule(
-      @Nullable String sqlAccessInfoFile,
-      @Nullable String cloudKmsProjectId,
-      @Nullable TransactionIsolationLevel isolationOverride) {
-    this.sqlAccessInfoFile = sqlAccessInfoFile;
-    this.cloudKmsProjectId = cloudKmsProjectId;
-    this.isolationOverride = isolationOverride;
-  }
-
-  public BeamJpaModule(@Nullable String sqlAccessInfoFile, @Nullable String cloudKmsProjectId) {
-    this(sqlAccessInfoFile, cloudKmsProjectId, null);
-  }
-
-  /** Returns true if the credential file is on GCS (and therefore expected to be encrypted). */
-  private boolean isCloudSqlCredential() {
-    return sqlAccessInfoFile.startsWith(GCS_SCHEME);
-  }
-
-  @Provides
-  @Singleton
-  SqlAccessInfo provideCloudSqlAccessInfo(Lazy<CloudSqlCredentialDecryptor> lazyDecryptor) {
-    checkArgument(!isNullOrEmpty(sqlAccessInfoFile), "Null or empty credentialFilePath");
-    String line = readOnlyLineFromCredentialFile();
-    if (isCloudSqlCredential()) {
-      line = lazyDecryptor.get().decrypt(line);
-    }
-    // See ./BackupPaths.java for explanation of the line format.
-    List<String> parts = Splitter.on(' ').splitToList(line.trim());
-    checkState(parts.size() == 3, "Expecting three phrases in %s", line);
-    if (isCloudSqlCredential()) {
-      return SqlAccessInfo.createCloudSqlAccessInfo(parts.get(0), parts.get(1), parts.get(2));
-    } else {
-      return SqlAccessInfo.createLocalSqlAccessInfo(parts.get(0), parts.get(1), parts.get(2));
-    }
-  }
-
-  String readOnlyLineFromCredentialFile() {
-    try {
-      ResourceId resourceId = FileSystems.matchSingleFileSpec(sqlAccessInfoFile).resourceId();
-      try (BufferedReader reader =
-          new BufferedReader(
-              new InputStreamReader(
-                  Channels.newInputStream(FileSystems.open(resourceId)), StandardCharsets.UTF_8))) {
-        return reader.readLine();
-      }
-    } catch (IOException e) {
-      throw new RuntimeException(e);
-    }
-  }
-
-  @Provides
-  @Config("beamCloudSqlJdbcUrl")
-  String provideJdbcUrl(SqlAccessInfo sqlAccessInfo) {
-    return sqlAccessInfo.jdbcUrl();
-  }
-
-  @Provides
-  @Config("beamCloudSqlInstanceConnectionName")
-  String provideSqlInstanceName(SqlAccessInfo sqlAccessInfo) {
-    return sqlAccessInfo
-        .cloudSqlInstanceName()
-        .orElseThrow(() -> new IllegalStateException("Cloud SQL not provisioned."));
-  }
-
-  @Provides
-  @Config("beamCloudSqlUsername")
-  String provideSqlUsername(SqlAccessInfo sqlAccessInfo) {
-    return sqlAccessInfo.user();
-  }
-
-  @Provides
-  @Config("beamCloudSqlPassword")
-  String provideSqlPassword(SqlAccessInfo sqlAccessInfo) {
-    return sqlAccessInfo.password();
-  }
-
-  @Provides
-  @Config("beamCloudKmsProjectId")
-  String kmsProjectId() {
-    return cloudKmsProjectId;
-  }
-
-  @Provides
-  @Config("beamCloudKmsKeyRing")
-  static String keyRingName() {
-    return "nomulus-tool-keyring";
-  }
-
-  @Provides
-  @Config("beamIsolationOverride")
-  @Nullable
-  TransactionIsolationLevel providesIsolationOverride() {
-    return isolationOverride;
-  }
-
-  @Provides
-  @Config("beamHibernateHikariMaximumPoolSize")
-  static int getBeamHibernateHikariMaximumPoolSize() {
-    // TODO(weiminyu): make this configurable. Should be equal to number of cores.
-    return 4;
-  }
-
-  @Singleton
-  @Component(
-      modules = {
-        ConfigModule.class,
-        CredentialModule.class,
-        BeamJpaModule.class,
-        KmsModule.class,
-        PersistenceModule.class,
-        SecretManagerModule.class,
-        UtilsModule.class
-      })
-  public interface JpaTransactionManagerComponent {
-    @SocketFactoryJpaTm
-    JpaTransactionManager cloudSqlJpaTransactionManager();
-
-    @JdbcJpaTm
-    JpaTransactionManager localDbJpaTransactionManager();
-  }
-}

core/src/main/java/google/registry/beam/initsql/InitSqlPipeline.java

@@ -120,26 +120,22 @@ public class InitSqlPipeline implements Serializable {
 
   private final InitSqlPipelineOptions options;
 
-  private final Pipeline pipeline;
-
   InitSqlPipeline(InitSqlPipelineOptions options) {
     this.options = options;
-    pipeline = Pipeline.create(options);
+  }
+
+  PipelineResult run() {
+    return run(Pipeline.create(options));
   }
 
   @VisibleForTesting
-  InitSqlPipeline(InitSqlPipelineOptions options, Pipeline pipeline) {
-    this.options = options;
-    this.pipeline = pipeline;
-  }
-
-  public PipelineResult run() {
-    setupPipeline();
+  PipelineResult run(Pipeline pipeline) {
+    setupPipeline(pipeline);
     return pipeline.run();
   }
 
   @VisibleForTesting
-  void setupPipeline() {
+  void setupPipeline(Pipeline pipeline) {
     options.setIsolationOverride(TransactionIsolationLevel.TRANSACTION_READ_UNCOMMITTED);
     PCollectionTuple datastoreSnapshot =
         pipeline.apply(

core/src/main/java/google/registry/beam/initsql/JpaSupplierFactory.java

@@ -1,60 +0,0 @@
-// Copyright 2020 The Nomulus Authors. All Rights Reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package google.registry.beam.initsql;
-
-import google.registry.beam.initsql.BeamJpaModule.JpaTransactionManagerComponent;
-import google.registry.beam.initsql.Transforms.SerializableSupplier;
-import google.registry.persistence.PersistenceModule.TransactionIsolationLevel;
-import google.registry.persistence.transaction.JpaTransactionManager;
-import javax.annotation.Nullable;
-import org.apache.beam.sdk.transforms.SerializableFunction;
-
-public class JpaSupplierFactory implements SerializableSupplier<JpaTransactionManager> {
-
-  private static final long serialVersionUID = 1L;
-
-  private final String credentialFileUrl;
-  @Nullable private final String cloudKmsProjectId;
-  private final SerializableFunction<JpaTransactionManagerComponent, JpaTransactionManager>
-      jpaGetter;
-  @Nullable private final TransactionIsolationLevel isolationLevelOverride;
-
-  public JpaSupplierFactory(
-      String credentialFileUrl,
-      @Nullable String cloudKmsProjectId,
-      SerializableFunction<JpaTransactionManagerComponent, JpaTransactionManager> jpaGetter) {
-    this(credentialFileUrl, cloudKmsProjectId, jpaGetter, null);
-  }
-
-  public JpaSupplierFactory(
-      String credentialFileUrl,
-      @Nullable String cloudKmsProjectId,
-      SerializableFunction<JpaTransactionManagerComponent, JpaTransactionManager> jpaGetter,
-      @Nullable TransactionIsolationLevel isolationLevelOverride) {
-    this.credentialFileUrl = credentialFileUrl;
-    this.cloudKmsProjectId = cloudKmsProjectId;
-    this.jpaGetter = jpaGetter;
-    this.isolationLevelOverride = isolationLevelOverride;
-  }
-
-  @Override
-  public JpaTransactionManager get() {
-    return jpaGetter.apply(
-        DaggerBeamJpaModule_JpaTransactionManagerComponent.builder()
-            .beamJpaModule(
-                new BeamJpaModule(credentialFileUrl, cloudKmsProjectId, isolationLevelOverride))
-            .build());
-  }
-}

core/src/main/java/google/registry/beam/initsql/Transforms.java

@@ -20,12 +20,9 @@ import static com.google.common.base.Preconditions.checkState;
 import static google.registry.beam.initsql.BackupPaths.getCommitLogTimestamp;
 import static google.registry.beam.initsql.BackupPaths.getExportFilePatterns;
 import static google.registry.model.ofy.ObjectifyService.ofy;
-import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
-import static google.registry.persistence.transaction.TransactionManagerFactory.setJpaTm;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
 import static google.registry.util.DateTimeUtils.isBeforeOrAt;
 import static java.util.Comparator.comparing;
-import static org.apache.beam.sdk.values.TypeDescriptors.integers;
 import static org.apache.beam.sdk.values.TypeDescriptors.kvs;
 import static org.apache.beam.sdk.values.TypeDescriptors.strings;
 
@@ -35,15 +32,14 @@ import com.google.appengine.api.datastore.EntityTranslator;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Streams;
-import com.googlecode.objectify.Key;
-import google.registry.backup.AppEngineEnvironment;
 import google.registry.backup.CommitLogImports;
 import google.registry.backup.VersionedEntity;
+import google.registry.model.billing.BillingEvent.Flag;
+import google.registry.model.billing.BillingEvent.Reason;
 import google.registry.model.domain.DomainBase;
-import google.registry.model.ofy.ObjectifyService;
 import google.registry.model.reporting.HistoryEntry;
-import google.registry.persistence.transaction.JpaTransactionManager;
 import google.registry.schema.replay.DatastoreAndSqlEntity;
 import google.registry.schema.replay.SqlEntity;
 import google.registry.tools.LevelDbLogReader;
@@ -53,7 +49,6 @@ import java.util.Iterator;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
-import java.util.concurrent.ThreadLocalRandom;
 import java.util.function.Supplier;
 import javax.annotation.Nullable;
 import org.apache.beam.sdk.coders.StringUtf8Coder;
@@ -62,18 +57,14 @@ import org.apache.beam.sdk.io.FileIO;
 import org.apache.beam.sdk.io.FileIO.ReadableFile;
 import org.apache.beam.sdk.io.fs.EmptyMatchTreatment;
 import org.apache.beam.sdk.io.fs.MatchResult.Metadata;
-import org.apache.beam.sdk.metrics.Counter;
-import org.apache.beam.sdk.metrics.Metrics;
 import org.apache.beam.sdk.transforms.Create;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.Flatten;
 import org.apache.beam.sdk.transforms.GroupByKey;
-import org.apache.beam.sdk.transforms.GroupIntoBatches;
 import org.apache.beam.sdk.transforms.MapElements;
 import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.transforms.ParDo;
 import org.apache.beam.sdk.transforms.ProcessFunction;
-import org.apache.beam.sdk.transforms.SerializableFunction;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.PBegin;
 import org.apache.beam.sdk.values.PCollection;
@@ -268,81 +259,58 @@ public final class Transforms {
                     .iterator()));
   }
 
-  /**
-   * Returns a {@link PTransform} that writes a {@link PCollection} of {@link VersionedEntity}s to a
-   * SQL database. and outputs an empty {@code PCollection<Void>}. This allows other operations to
-   * {@link org.apache.beam.sdk.transforms.Wait wait} for the completion of this transform.
-   *
-   * <p>Errors are handled according to the pipeline runner's default policy. As part of a one-time
-   * job, we will not add features unless proven necessary.
-   *
-   * @param transformId a unique ID for an instance of the returned transform
-   * @param maxWriters the max number of concurrent writes to SQL, which also determines the max
-   *     number of connection pools created
-   * @param batchSize the number of entities to write in each operation
-   * @param jpaSupplier supplier of a {@link JpaTransactionManager}
-   */
-  public static PTransform<PCollection<VersionedEntity>, PCollection<Void>> writeToSql(
-      String transformId,
-      int maxWriters,
-      int batchSize,
-      SerializableSupplier<JpaTransactionManager> jpaSupplier) {
-    return writeToSql(
-        transformId,
-        maxWriters,
-        batchSize,
-        jpaSupplier,
-        Transforms::convertVersionedEntityToSqlEntity,
-        TypeDescriptor.of(VersionedEntity.class));
-  }
+  // Production data repair configs go below. See b/185954992.
 
-  /**
-   * Returns a {@link PTransform} that writes a {@link PCollection} of entities to a SQL database.
-   * and outputs an empty {@code PCollection<Void>}. This allows other operations to {@link
-   * org.apache.beam.sdk.transforms.Wait wait} for the completion of this transform.
-   *
-   * <p>The converter and type descriptor are generics so that we can convert any type of entity to
-   * an object to be placed in SQL.
-   *
-   * <p>Errors are handled according to the pipeline runner's default policy. As part of a one-time
-   * job, we will not add features unless proven necessary.
-   *
-   * @param transformId a unique ID for an instance of the returned transform
-   * @param maxWriters the max number of concurrent writes to SQL, which also determines the max
-   *     number of connection pools created
-   * @param batchSize the number of entities to write in each operation
-   * @param jpaSupplier supplier of a {@link JpaTransactionManager}
-   * @param jpaConverter the function that converts the input object to a JPA entity
-   * @param objectDescriptor the type descriptor of the input object
-   */
-  public static <T> PTransform<PCollection<T>, PCollection<Void>> writeToSql(
-      String transformId,
-      int maxWriters,
-      int batchSize,
-      SerializableSupplier<JpaTransactionManager> jpaSupplier,
-      SerializableFunction<T, Object> jpaConverter,
-      TypeDescriptor<T> objectDescriptor) {
-    return new PTransform<PCollection<T>, PCollection<Void>>() {
-      @Override
-      public PCollection<Void> expand(PCollection<T> input) {
-        return input
-            .apply(
-                "Shard data for " + transformId,
-                MapElements.into(kvs(integers(), objectDescriptor))
-                    .via(ve -> KV.of(ThreadLocalRandom.current().nextInt(maxWriters), ve)))
-            .apply("Batch output by shard " + transformId, GroupIntoBatches.ofSize(batchSize))
-            .apply(
-                "Write in batch for " + transformId,
-                ParDo.of(new SqlBatchWriter<T>(transformId, jpaSupplier, jpaConverter)));
-      }
-    };
-  }
+  // Prober domains in bad state, without associated contacts, hosts, billings, and history.
+  // They can be safely ignored.
+  private static final ImmutableSet<String> IGNORED_DOMAINS =
+      ImmutableSet.of("6AF6D2-IQCANT", "2-IQANYT");
 
-  private static Key toOfyKey(Object ofyEntity) {
-    return Key.create(ofyEntity);
-  }
+  // Prober hosts referencing phantom registrars. They and their associated history entries can be
+  // safely ignored.
+  private static final ImmutableSet<String> IGNORED_HOSTS =
+      ImmutableSet.of(
+          "4E21_WJ0TEST-GOOGLE",
+          "4E21_WJ1TEST-GOOGLE",
+          "4E21_WJ2TEST-GOOGLE",
+          "4E21_WJ3TEST-GOOGLE");
+
+  // Prober contacts referencing phantom registrars. They and their associated history entries can
+  // be safely ignored.
+  private static final ImmutableSet IGNORED_CONTACTS =
+      ImmutableSet.of(
+          "1_WJ0TEST-GOOGLE", "1_WJ1TEST-GOOGLE", "1_WJ2TEST-GOOGLE", "1_WJ3TEST-GOOGLE");
 
   private static boolean isMigratable(Entity entity) {
+    // Checks specific to production data. See b/185954992 for details.
+    // The names of these bad entities in production do not conflict with other environments. For
+    // simplicities sake we apply them regardless of the source of the data.
+    if (entity.getKind().equals("DomainBase")
+        && IGNORED_DOMAINS.contains(entity.getKey().getName())) {
+      return false;
+    }
+    if (entity.getKind().equals("ContactResource")) {
+      String roid = entity.getKey().getName();
+      return !IGNORED_CONTACTS.contains(roid);
+    }
+    if (entity.getKind().equals("HostResource")) {
+      String roid = entity.getKey().getName();
+      return !IGNORED_HOSTS.contains(roid);
+    }
+    if (entity.getKind().equals("HistoryEntry")) {
+      // Remove production bad data: History of the contacts to be ignored:
+      com.google.appengine.api.datastore.Key parentKey = entity.getKey().getParent();
+      if (parentKey.getKind().equals("ContactResource")) {
+        String contactRoid = parentKey.getName();
+        return !IGNORED_CONTACTS.contains(contactRoid);
+      }
+      if (parentKey.getKind().equals("HostResource")) {
+        String hostRoid = parentKey.getName();
+        return !IGNORED_HOSTS.contains(hostRoid);
+      }
+    }
+    // End of production-specific checks.
+
     if (entity.getKind().equals("HistoryEntry")) {
       // DOMAIN_APPLICATION_CREATE is deprecated type and should not be migrated.
       // The Enum name DOMAIN_APPLICATION_CREATE no longer exists in Java and cannot
@@ -352,6 +320,18 @@ public final class Transforms {
     return true;
   }
 
+  private static Entity repairBadData(Entity entity) {
+    if (entity.getKind().equals("Cancellation")
+        && Objects.equals(entity.getProperty("reason"), "AUTO_RENEW")) {
+      // AUTO_RENEW has been moved from 'reason' to flags. Change reason to RENEW and add the
+      // AUTO_RENEW flag. Note: all affected entities have empty flags so we can simply assign
+      // instead of append. See b/185954992.
+      entity.setUnindexedProperty("reason", Reason.RENEW.name());
+      entity.setUnindexedProperty("flags", ImmutableList.of(Flag.AUTO_RENEW.name()));
+    }
+    return entity;
+  }
+
   private static SqlEntity toSqlEntity(Object ofyEntity) {
     if (ofyEntity instanceof HistoryEntry) {
       HistoryEntry ofyHistory = (HistoryEntry) ofyEntity;
@@ -372,6 +352,7 @@ public final class Transforms {
     return dsEntity
         .getEntity()
         .filter(Transforms::isMigratable)
+        .map(Transforms::repairBadData)
         .map(e -> ofy().toPojo(e))
         .map(Transforms::toSqlEntity)
         .orElse(null);
@@ -458,93 +439,6 @@ public final class Transforms {
     }
   }
 
-  /**
-   * Writes a batch of entities to a SQL database.
-   *
-   * <p>Note that an arbitrary number of instances of this class may be created and freed in
-   * arbitrary order in a single JVM. Due to the tech debt that forced us to use a static variable
-   * to hold the {@code JpaTransactionManager} instance, we must ensure that JpaTransactionManager
-   * is not changed or torn down while being used by some instance.
-   */
-  private static class SqlBatchWriter<T> extends DoFn<KV<Integer, Iterable<T>>, Void> {
-
-    private static int instanceCount = 0;
-    private static JpaTransactionManager originalJpa;
-
-    private Counter counter;
-
-    private final SerializableSupplier<JpaTransactionManager> jpaSupplier;
-    private final SerializableFunction<T, Object> jpaConverter;
-
-    SqlBatchWriter(
-        String type,
-        SerializableSupplier<JpaTransactionManager> jpaSupplier,
-        SerializableFunction<T, Object> jpaConverter) {
-      counter = Metrics.counter("SQL_WRITE", type);
-      this.jpaSupplier = jpaSupplier;
-      this.jpaConverter = jpaConverter;
-    }
-
-    @Setup
-    public void setup() {
-      try (AppEngineEnvironment env = new AppEngineEnvironment()) {
-        ObjectifyService.initOfy();
-      }
-
-      synchronized (SqlBatchWriter.class) {
-        if (instanceCount == 0) {
-          originalJpa = jpaTm();
-          setJpaTm(jpaSupplier);
-        }
-        instanceCount++;
-      }
-    }
-
-    @Teardown
-    public void teardown() {
-      synchronized (SqlBatchWriter.class) {
-        instanceCount--;
-        if (instanceCount == 0) {
-          jpaTm().teardown();
-          setJpaTm(() -> originalJpa);
-        }
-      }
-    }
-
-    @ProcessElement
-    public void processElement(@Element KV<Integer, Iterable<T>> kv) {
-      try (AppEngineEnvironment env = new AppEngineEnvironment()) {
-        ImmutableList<Object> ofyEntities =
-            Streams.stream(kv.getValue())
-                .map(this.jpaConverter::apply)
-                // TODO(b/177340730): post migration delete the line below.
-                .filter(Objects::nonNull)
-                .collect(ImmutableList.toImmutableList());
-        try {
-          jpaTm().transact(() -> jpaTm().putAll(ofyEntities));
-          counter.inc(ofyEntities.size());
-        } catch (RuntimeException e) {
-          processSingly(ofyEntities);
-        }
-      }
-    }
-
-    /**
-     * Writes entities in a failed batch one by one to identify the first bad entity and throws a
-     * {@link RuntimeException} on it.
-     */
-    private void processSingly(ImmutableList<Object> ofyEntities) {
-      for (Object ofyEntity : ofyEntities) {
-        try {
-          jpaTm().transact(() -> jpaTm().put(ofyEntity));
-          counter.inc();
-        } catch (RuntimeException e) {
-          throw new RuntimeException(toOfyKey(ofyEntity).toString(), e);
-        }
-      }
-    }
-  }
-
   /**
    * Removes BillingEvents, {@link google.registry.model.poll.PollMessage PollMessages} and {@link
    * google.registry.model.host.HostResource} from a {@link DomainBase}. These are circular foreign

core/src/main/java/google/registry/beam/invoicing/BillingEvent.java

@@ -260,6 +260,11 @@ public abstract class BillingEvent implements Serializable {
         poNumber());
   }
 
+  /** Returns the grouping key for this {@code BillingEvent}, to generate the detailed report. */
+  String getDetailedReportGroupingKey() {
+    return String.format("%s_%s", registrarId(), tld());
+  }
+
   /** Key for each {@code BillingEvent}, when aggregating for the overall invoice. */
   @AutoValue
   abstract static class InvoiceGroupingKey implements Serializable {

core/src/main/java/google/registry/beam/invoicing/InvoicingPipeline.java

@@ -14,28 +14,27 @@
 
 package google.registry.beam.invoicing;
 
-import com.google.auth.oauth2.GoogleCredentials;
+import static google.registry.beam.BeamUtils.getQueryFromFile;
+import static org.apache.beam.sdk.values.TypeDescriptors.strings;
+
 import google.registry.beam.invoicing.BillingEvent.InvoiceGroupingKey;
 import google.registry.beam.invoicing.BillingEvent.InvoiceGroupingKey.InvoiceGroupingKeyCoder;
-import google.registry.config.CredentialModule.LocalCredential;
-import google.registry.config.RegistryConfig.Config;
 import google.registry.reporting.billing.BillingModule;
-import google.registry.reporting.billing.GenerateInvoicesAction;
-import google.registry.util.GoogleCredentialsBundle;
+import google.registry.util.SqlTemplate;
 import java.io.Serializable;
-import javax.inject.Inject;
-import org.apache.beam.runners.dataflow.DataflowRunner;
-import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions;
+import java.time.LocalDateTime;
+import java.time.LocalTime;
+import java.time.YearMonth;
+import java.time.format.DateTimeFormatter;
 import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.PipelineResult;
 import org.apache.beam.sdk.coders.SerializableCoder;
-import org.apache.beam.sdk.io.DefaultFilenamePolicy.Params;
-import org.apache.beam.sdk.io.FileBasedSink;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.io.FileIO;
 import org.apache.beam.sdk.io.TextIO;
 import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO;
-import org.apache.beam.sdk.options.Description;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
-import org.apache.beam.sdk.options.ValueProvider;
-import org.apache.beam.sdk.options.ValueProvider.NestedValueProvider;
+import org.apache.beam.sdk.transforms.Contextful;
 import org.apache.beam.sdk.transforms.Count;
 import org.apache.beam.sdk.transforms.Filter;
 import org.apache.beam.sdk.transforms.MapElements;
@@ -43,107 +42,48 @@ import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.sdk.values.TypeDescriptor;
-import org.apache.beam.sdk.values.TypeDescriptors;
 
 /**
- * Definition of a Dataflow pipeline template, which generates a given month's invoices.
+ * Definition of a Dataflow Flex pipeline template, which generates a given month's invoices.
  *
- * <p>To stage this template on GCS, run the {@link
- * google.registry.tools.DeployInvoicingPipelineCommand} Nomulus command.
+ * <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.
- * For an example using the API client library, see {@link GenerateInvoicesAction}.
  *
- * @see <a href="https://cloud.google.com/dataflow/docs/templates/overview">Dataflow Templates</a>
+ * @see <a href="https://cloud.google.com/dataflow/docs/guides/templates/using-flex-templates">Using
+ *     Flex Templates</a>
  */
 public class InvoicingPipeline implements Serializable {
 
-  private final String projectId;
-  private final String beamJobRegion;
-  private final String beamBucketUrl;
-  private final String invoiceTemplateUrl;
-  private final String beamStagingUrl;
-  private final String billingBucketUrl;
-  private final String invoiceFilePrefix;
-  private final GoogleCredentials googleCredentials;
+  private static final DateTimeFormatter TIMESTAMP_FORMATTER =
+      DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss.SSSSSS");
 
-  @Inject
-  public InvoicingPipeline(
-      @Config("projectId") String projectId,
-      @Config("defaultJobRegion") String beamJobRegion,
-      @Config("apacheBeamBucketUrl") String beamBucketUrl,
-      @Config("invoiceTemplateUrl") String invoiceTemplateUrl,
-      @Config("beamStagingUrl") String beamStagingUrl,
-      @Config("billingBucketUrl") String billingBucketUrl,
-      @Config("invoiceFilePrefix") String invoiceFilePrefix,
-      @LocalCredential GoogleCredentialsBundle googleCredentialsBundle) {
-    this.projectId = projectId;
-    this.beamJobRegion = beamJobRegion;
-    this.beamBucketUrl = beamBucketUrl;
-    this.invoiceTemplateUrl = invoiceTemplateUrl;
-    this.beamStagingUrl = beamStagingUrl;
-    this.billingBucketUrl = billingBucketUrl;
-    this.invoiceFilePrefix = invoiceFilePrefix;
-    this.googleCredentials = googleCredentialsBundle.getGoogleCredentials();
+  private final InvoicingPipelineOptions options;
+
+  InvoicingPipeline(InvoicingPipelineOptions options) {
+    this.options = options;
   }
 
-  /** Custom options for running the invoicing pipeline. */
-  public interface InvoicingPipelineOptions extends DataflowPipelineOptions {
-    /** Returns the yearMonth we're generating invoices for, in yyyy-MM format. */
-    @Description("The yearMonth we generate invoices for, in yyyy-MM format.")
-    ValueProvider<String> getYearMonth();
-    /**
-     * Sets the yearMonth we generate invoices for.
-     *
-     * <p>This is implicitly set when executing the Dataflow template, by specifying the 'yearMonth
-     * parameter.
-     */
-    void setYearMonth(ValueProvider<String> value);
+  PipelineResult run() {
+    Pipeline pipeline = Pipeline.create(options);
+    setupPipeline(pipeline);
+    return pipeline.run();
   }
 
-  /** Deploys the invoicing pipeline as a template on GCS, for a given projectID and GCS bucket. */
-  public void deploy() {
-    // We can't store options as a member variable due to serialization concerns.
-    InvoicingPipelineOptions options = PipelineOptionsFactory.as(InvoicingPipelineOptions.class);
-    options.setProject(projectId);
-    options.setRegion(beamJobRegion);
-    options.setRunner(DataflowRunner.class);
-    // This causes p.run() to stage the pipeline as a template on GCS, as opposed to running it.
-    options.setTemplateLocation(invoiceTemplateUrl);
-    options.setStagingLocation(beamStagingUrl);
-    // This credential is used when Dataflow deploys the template to GCS in target GCP project.
-    // So, make sure the credential has write permission to GCS in that project.
-    options.setGcpCredential(googleCredentials);
-
-    Pipeline p = Pipeline.create(options);
-
+  void setupPipeline(Pipeline pipeline) {
     PCollection<BillingEvent> billingEvents =
-        p.apply(
+        pipeline.apply(
             "Read BillingEvents from Bigquery",
             BigQueryIO.read(BillingEvent::parseFromRecord)
-                .fromQuery(InvoicingUtils.makeQueryProvider(options.getYearMonth(), projectId))
+                .fromQuery(makeQuery(options.getYearMonth(), options.getProject()))
                 .withCoder(SerializableCoder.of(BillingEvent.class))
                 .usingStandardSql()
                 .withoutValidation()
                 .withTemplateCompatibility());
-    applyTerminalTransforms(billingEvents, options.getYearMonth());
-    p.run();
-  }
 
-  /**
-   * Applies output transforms to the {@code BillingEvent} source collection.
-   *
-   * <p>This is factored out purely to facilitate testing.
-   */
-  void applyTerminalTransforms(
-      PCollection<BillingEvent> billingEvents, ValueProvider<String> yearMonthProvider) {
-    billingEvents
-        .apply("Generate overall invoice rows", new GenerateInvoiceRows())
-        .apply("Write overall invoice to CSV", writeInvoice(yearMonthProvider));
+    saveInvoiceCsv(billingEvents, options);
 
-    billingEvents.apply(
-        "Write detail reports to separate CSVs keyed by registrarId_tld pair",
-        writeDetailReports(yearMonthProvider));
+    saveDetailedCsv(billingEvents, options);
   }
 
   /** Transform that converts a {@code BillingEvent} into an invoice CSV row. */
@@ -156,49 +96,85 @@ public class InvoicingPipeline implements Serializable {
               "Map to invoicing key",
               MapElements.into(TypeDescriptor.of(InvoiceGroupingKey.class))
                   .via(BillingEvent::getInvoiceGroupingKey))
-          .apply(Filter.by((InvoiceGroupingKey key) -> key.unitPrice() != 0))
+          .apply(
+              "Filter out free events", Filter.by((InvoiceGroupingKey key) -> key.unitPrice() != 0))
           .setCoder(new InvoiceGroupingKeyCoder())
           .apply("Count occurrences", Count.perElement())
           .apply(
               "Format as CSVs",
-              MapElements.into(TypeDescriptors.strings())
+              MapElements.into(strings())
                   .via((KV<InvoiceGroupingKey, Long> kv) -> kv.getKey().toCsv(kv.getValue())));
     }
   }
 
-  /** Returns an IO transform that writes the overall invoice to a single CSV file. */
-  private TextIO.Write writeInvoice(ValueProvider<String> yearMonthProvider) {
-    return TextIO.write()
-        .to(
-            NestedValueProvider.of(
-                yearMonthProvider,
-                yearMonth ->
+  /** Saves the billing events to a single overall invoice CSV file. */
+  static void saveInvoiceCsv(
+      PCollection<BillingEvent> billingEvents, InvoicingPipelineOptions options) {
+    billingEvents
+        .apply("Generate overall invoice rows", new GenerateInvoiceRows())
+        .apply(
+            "Write overall invoice to CSV",
+            TextIO.write()
+                .to(
                     String.format(
                         "%s/%s/%s/%s-%s",
-                        billingBucketUrl,
+                        options.getBillingBucketUrl(),
                         BillingModule.INVOICES_DIRECTORY,
-                        yearMonth,
-                        invoiceFilePrefix,
-                        yearMonth)))
-        .withHeader(InvoiceGroupingKey.invoiceHeader())
-        .withoutSharding()
-        .withSuffix(".csv");
+                        options.getYearMonth(),
+                        options.getInvoiceFilePrefix(),
+                        options.getYearMonth()))
+                .withHeader(InvoiceGroupingKey.invoiceHeader())
+                .withoutSharding()
+                .withSuffix(".csv"));
   }
 
-  /** Returns an IO transform that writes detail reports to registrar-tld keyed CSV files. */
-  private TextIO.TypedWrite<BillingEvent, Params> writeDetailReports(
-      ValueProvider<String> yearMonthProvider) {
-    return TextIO.<BillingEvent>writeCustomType()
-        .to(
-            InvoicingUtils.makeDestinationFunction(
-                String.format("%s/%s", billingBucketUrl, BillingModule.INVOICES_DIRECTORY),
-                yearMonthProvider),
-            InvoicingUtils.makeEmptyDestinationParams(billingBucketUrl + "/errors"))
-        .withFormatFunction(BillingEvent::toCsv)
-        .withoutSharding()
-        .withTempDirectory(
-            FileBasedSink.convertToFileResourceIfPossible(beamBucketUrl + "/temporary"))
-        .withHeader(BillingEvent.getHeader())
-        .withSuffix(".csv");
+  /** Saves the billing events to detailed report CSV files keyed by registrar-tld pairs. */
+  static void saveDetailedCsv(
+      PCollection<BillingEvent> billingEvents, InvoicingPipelineOptions options) {
+    String yearMonth = options.getYearMonth();
+    billingEvents.apply(
+        "Write detailed report for each registrar-tld pair",
+        FileIO.<String, BillingEvent>writeDynamic()
+            .to(
+                String.format(
+                    "%s/%s/%s",
+                    options.getBillingBucketUrl(), BillingModule.INVOICES_DIRECTORY, yearMonth))
+            .by(BillingEvent::getDetailedReportGroupingKey)
+            .withNumShards(1)
+            .withDestinationCoder(StringUtf8Coder.of())
+            .withNaming(
+                key ->
+                    (window, pane, numShards, shardIndex, compression) ->
+                        String.format(
+                            "%s_%s_%s.csv", BillingModule.DETAIL_REPORT_PREFIX, yearMonth, key))
+            .via(
+                Contextful.fn(BillingEvent::toCsv),
+                TextIO.sink().withHeader(BillingEvent.getHeader())));
+  }
+
+  /** Create the Bigquery query for a given project and yearMonth at runtime. */
+  static String makeQuery(String yearMonth, String projectId) {
+    // Get the timestamp endpoints capturing the entire month with microsecond precision
+    YearMonth reportingMonth = YearMonth.parse(yearMonth);
+    LocalDateTime firstMoment = reportingMonth.atDay(1).atTime(LocalTime.MIDNIGHT);
+    LocalDateTime lastMoment = reportingMonth.atEndOfMonth().atTime(LocalTime.MAX);
+    // Construct the month's query by filling in the billing_events.sql template
+    return SqlTemplate.create(getQueryFromFile(InvoicingPipeline.class, "billing_events.sql"))
+        .put("FIRST_TIMESTAMP_OF_MONTH", firstMoment.format(TIMESTAMP_FORMATTER))
+        .put("LAST_TIMESTAMP_OF_MONTH", lastMoment.format(TIMESTAMP_FORMATTER))
+        .put("PROJECT_ID", projectId)
+        .put("DATASTORE_EXPORT_DATA_SET", "latest_datastore_export")
+        .put("ONETIME_TABLE", "OneTime")
+        .put("REGISTRY_TABLE", "Registry")
+        .put("REGISTRAR_TABLE", "Registrar")
+        .put("CANCELLATION_TABLE", "Cancellation")
+        .build();
+  }
+
+  public static void main(String[] args) {
+    PipelineOptionsFactory.register(InvoicingPipelineOptions.class);
+    InvoicingPipelineOptions options =
+        PipelineOptionsFactory.fromArgs(args).withValidation().as(InvoicingPipelineOptions.class);
+    new InvoicingPipeline(options).run();
   }
 }

core/src/main/java/google/registry/beam/invoicing/InvoicingPipelineOptions.java

@@ -0,0 +1,37 @@
+// 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.invoicing;
+
+import google.registry.beam.common.RegistryPipelineOptions;
+import org.apache.beam.sdk.options.Description;
+
+/** Custom options for running the invoicing pipeline. */
+public interface InvoicingPipelineOptions extends RegistryPipelineOptions {
+
+  @Description("The year and month we generate invoices for, in yyyy-MM format.")
+  String getYearMonth();
+
+  void setYearMonth(String value);
+
+  @Description("Filename prefix for the invoice CSV file.")
+  String getInvoiceFilePrefix();
+
+  void setInvoiceFilePrefix(String value);
+
+  @Description("The GCS bucket URL for invoices and detailed  reports to be uploaded.")
+  String getBillingBucketUrl();
+
+  void setBillingBucketUrl(String value);
+}

core/src/main/java/google/registry/beam/invoicing/InvoicingUtils.java

@@ -1,106 +0,0 @@
-// Copyright 2018 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.invoicing;
-
-import static google.registry.beam.BeamUtils.getQueryFromFile;
-
-import google.registry.util.SqlTemplate;
-import java.time.LocalDateTime;
-import java.time.LocalTime;
-import java.time.YearMonth;
-import java.time.format.DateTimeFormatter;
-import org.apache.beam.sdk.io.DefaultFilenamePolicy.Params;
-import org.apache.beam.sdk.io.FileBasedSink;
-import org.apache.beam.sdk.options.ValueProvider;
-import org.apache.beam.sdk.options.ValueProvider.NestedValueProvider;
-import org.apache.beam.sdk.transforms.SerializableFunction;
-
-/** Pipeline helper functions used to generate invoices from instances of {@link BillingEvent}. */
-public class InvoicingUtils {
-
-  private InvoicingUtils() {}
-
-  private static final DateTimeFormatter TIMESTAMP_FORMATTER =
-      DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss.SSSSSS");
-
-  /**
-   * Returns a function mapping from {@code BillingEvent} to filename {@code Params}.
-   *
-   * <p>Beam uses this to determine which file a given {@code BillingEvent} should get placed into.
-   *
-   * @param outputBucket the GCS bucket we're outputting reports to
-   * @param yearMonthProvider a runtime provider for the yyyy-MM we're generating the invoice for
-   */
-  static SerializableFunction<BillingEvent, Params> makeDestinationFunction(
-      String outputBucket, ValueProvider<String> yearMonthProvider) {
-    return billingEvent ->
-        new Params()
-            .withShardTemplate("")
-            .withSuffix(".csv")
-            .withBaseFilename(
-                NestedValueProvider.of(
-                    yearMonthProvider,
-                    yearMonth ->
-                        FileBasedSink.convertToFileResourceIfPossible(
-                            String.format(
-                                "%s/%s/%s",
-                                outputBucket, yearMonth, billingEvent.toFilename(yearMonth)))));
-  }
-
-  /**
-   * Returns the default filename parameters for an unmappable {@code BillingEvent}.
-   *
-   * <p>The "failed" file should only be populated when an error occurs, which warrants further
-   * investigation.
-   */
-  static Params makeEmptyDestinationParams(String outputBucket) {
-    return new Params()
-        .withBaseFilename(
-            FileBasedSink.convertToFileResourceIfPossible(
-                String.format("%s/%s", outputBucket, "FAILURES")));
-  }
-
-  /**
-   * Returns a provider that creates a Bigquery query for a given project and yearMonth at runtime.
-   *
-   * <p>We only know yearMonth at runtime, so this provider fills in the {@code
-   * sql/billing_events.sql} template at runtime.
-   *
-   * @param yearMonthProvider a runtime provider that returns which month we're invoicing for.
-   * @param projectId the projectId we're generating invoicing for.
-   */
-  static ValueProvider<String> makeQueryProvider(
-      ValueProvider<String> yearMonthProvider, String projectId) {
-    return NestedValueProvider.of(
-        yearMonthProvider,
-        (yearMonth) -> {
-          // Get the timestamp endpoints capturing the entire month with microsecond precision
-          YearMonth reportingMonth = YearMonth.parse(yearMonth);
-          LocalDateTime firstMoment = reportingMonth.atDay(1).atTime(LocalTime.MIDNIGHT);
-          LocalDateTime lastMoment = reportingMonth.atEndOfMonth().atTime(LocalTime.MAX);
-          // Construct the month's query by filling in the billing_events.sql template
-          return SqlTemplate.create(getQueryFromFile(InvoicingPipeline.class, "billing_events.sql"))
-              .put("FIRST_TIMESTAMP_OF_MONTH", firstMoment.format(TIMESTAMP_FORMATTER))
-              .put("LAST_TIMESTAMP_OF_MONTH", lastMoment.format(TIMESTAMP_FORMATTER))
-              .put("PROJECT_ID", projectId)
-              .put("DATASTORE_EXPORT_DATA_SET", "latest_datastore_export")
-              .put("ONETIME_TABLE", "OneTime")
-              .put("REGISTRY_TABLE", "Registry")
-              .put("REGISTRAR_TABLE", "Registrar")
-              .put("CANCELLATION_TABLE", "Cancellation")
-              .build();
-        });
-  }
-}

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

@@ -14,7 +14,6 @@
 
 package google.registry.beam.spec11;
 
-
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.apache.http.HttpStatus.SC_OK;
 
@@ -30,7 +29,6 @@ import java.net.URISyntaxException;
 import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.function.Supplier;
-import org.apache.beam.sdk.options.ValueProvider;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
 import org.apache.beam.sdk.values.KV;
@@ -73,7 +71,7 @@ public class SafeBrowsingTransforms {
     private static final int BATCH_SIZE = 490;
 
     /** Provides the SafeBrowsing API key at runtime. */
-    private final ValueProvider<String> apiKeyProvider;
+    private final String apiKey;
 
     /**
      * Maps a subdomain's {@code fullyQualifiedDomainName} to its corresponding {@link Subdomain} to
@@ -93,20 +91,18 @@ public class SafeBrowsingTransforms {
     private final Retrier retrier;
 
     /**
-     * Constructs a {@link EvaluateSafeBrowsingFn} that gets its API key from the given provider.
+     * Constructs a {@link EvaluateSafeBrowsingFn} with a given API key.
      *
      * <p>We need to dual-cast the closeableHttpClientSupplier lambda because all {@code DoFn}
      * member variables need to be serializable. The (Supplier & Serializable) dual cast is safe
      * because class methods are generally serializable, especially a static function such as {@link
      * HttpClients#createDefault()}.
-     *
-     * @param apiKeyProvider provides the SafeBrowsing API key from {@code KMS} at runtime
      */
     @SuppressWarnings("unchecked")
-    EvaluateSafeBrowsingFn(ValueProvider<String> apiKeyProvider, Retrier retrier) {
-      this.apiKeyProvider = apiKeyProvider;
+    EvaluateSafeBrowsingFn(String apiKey, Retrier retrier) {
+      this.apiKey = apiKey;
       this.retrier = retrier;
-      this.closeableHttpClientSupplier = (Supplier & Serializable) HttpClients::createDefault;
+      closeableHttpClientSupplier = (Supplier & Serializable) HttpClients::createDefault;
     }
 
     /**
@@ -117,12 +113,10 @@ public class SafeBrowsingTransforms {
      */
     @VisibleForTesting
     EvaluateSafeBrowsingFn(
-        ValueProvider<String> apiKeyProvider,
-        Retrier retrier,
-        Supplier<CloseableHttpClient> clientSupplier) {
-      this.apiKeyProvider = apiKeyProvider;
+        String apiKey, Retrier retrier, Supplier<CloseableHttpClient> clientSupplier) {
+      this.apiKey = apiKey;
       this.retrier = retrier;
-      this.closeableHttpClientSupplier = clientSupplier;
+      closeableHttpClientSupplier = clientSupplier;
     }
 
     /** Evaluates any buffered {@link Subdomain} objects upon completing the bundle. */
@@ -159,7 +153,7 @@ public class SafeBrowsingTransforms {
       try {
         URIBuilder uriBuilder = new URIBuilder(SAFE_BROWSING_URL);
         // Add the API key param
-        uriBuilder.addParameter("key", apiKeyProvider.get());
+        uriBuilder.addParameter("key", apiKey);
 
         HttpPost httpPost = new HttpPost(uriBuilder.build());
         httpPost.addHeader(HTTP.CONTENT_TYPE, ContentType.APPLICATION_JSON.toString());
@@ -175,7 +169,7 @@ public class SafeBrowsingTransforms {
               }
             },
             IOException.class);
-      } catch (URISyntaxException | JSONException  e) {
+      } catch (URISyntaxException | JSONException e) {
         // Fail the pipeline on a parsing exception- this indicates the API likely changed.
         throw new RuntimeException("Caught parsing exception, failing pipeline.", e);
       } finally {
@@ -239,7 +233,9 @@ public class SafeBrowsingTransforms {
             String url = match.getJSONObject("threat").getString("url");
             Subdomain subdomain = subdomainBuffer.get(url);
             resultBuilder.add(
-                KV.of(subdomain, ThreatMatch.create(match, subdomain.domainName())));
+                KV.of(
+                    subdomain,
+                    ThreatMatch.create(match.getString("threatType"), subdomain.domainName())));
           }
         }
       }

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

@@ -17,32 +17,27 @@ package google.registry.beam.spec11;
 import static com.google.common.base.Preconditions.checkArgument;
 import static google.registry.beam.BeamUtils.getQueryFromFile;
 
-import com.google.auth.oauth2.GoogleCredentials;
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableSet;
-import google.registry.beam.initsql.Transforms;
-import google.registry.beam.initsql.Transforms.SerializableSupplier;
+import dagger.Component;
+import dagger.Module;
+import dagger.Provides;
+import google.registry.beam.common.RegistryJpaIO;
 import google.registry.beam.spec11.SafeBrowsingTransforms.EvaluateSafeBrowsingFn;
-import google.registry.config.CredentialModule.LocalCredential;
-import google.registry.config.RegistryConfig.Config;
+import google.registry.config.RegistryConfig.ConfigModule;
 import google.registry.model.reporting.Spec11ThreatMatch;
 import google.registry.model.reporting.Spec11ThreatMatch.ThreatType;
-import google.registry.persistence.transaction.JpaTransactionManager;
-import google.registry.util.GoogleCredentialsBundle;
 import google.registry.util.Retrier;
 import google.registry.util.SqlTemplate;
+import google.registry.util.UtilsModule;
 import java.io.Serializable;
-import javax.inject.Inject;
-import org.apache.beam.runners.dataflow.DataflowRunner;
-import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions;
+import javax.inject.Singleton;
 import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.PipelineResult;
 import org.apache.beam.sdk.coders.SerializableCoder;
 import org.apache.beam.sdk.io.TextIO;
 import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO;
-import org.apache.beam.sdk.options.Description;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
-import org.apache.beam.sdk.options.ValueProvider;
-import org.apache.beam.sdk.options.ValueProvider.NestedValueProvider;
 import org.apache.beam.sdk.transforms.GroupByKey;
 import org.apache.beam.sdk.transforms.MapElements;
 import org.apache.beam.sdk.transforms.ParDo;
@@ -58,21 +53,20 @@ import org.json.JSONException;
 import org.json.JSONObject;
 
 /**
- * Definition of a Dataflow pipeline template, which generates a given month's spec11 report.
+ * Definition of a Dataflow Flex template, which generates a given month's spec11 report.
  *
- * <p>To stage this template on GCS, run the {@link
- * google.registry.tools.DeploySpec11PipelineCommand} Nomulus command.
+ * <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/templates/overview">Dataflow Templates</a>
+ * @see <a href="https://cloud.google.com/dataflow/docs/guides/templates/using-flex-templates">Using
+ *     Flex Templates</a>
  */
 public class Spec11Pipeline implements Serializable {
 
   /**
    * Returns the subdirectory spec11 reports reside in for a given local date in yyyy-MM-dd format.
    *
-   * @see google.registry.beam.spec11.Spec11Pipeline
    * @see google.registry.reporting.spec11.Spec11EmailUtils
    */
   public static String getSpec11ReportFilePath(LocalDate localDate) {
@@ -87,84 +81,28 @@ public class Spec11Pipeline implements Serializable {
   /** The JSON object field into which we put the threat match array for Spec11 reports. */
   public static final String THREAT_MATCHES_FIELD = "threatMatches";
 
-  private final String projectId;
-  private final String beamJobRegion;
-  private final String beamStagingUrl;
-  private final String spec11TemplateUrl;
-  private final String reportingBucketUrl;
-  private final GoogleCredentials googleCredentials;
-  private final Retrier retrier;
-  private final SerializableSupplier<JpaTransactionManager> jpaSupplierFactory;
+  private final Spec11PipelineOptions options;
+  private final EvaluateSafeBrowsingFn safeBrowsingFn;
 
-  @Inject
-  public Spec11Pipeline(
-      @Config("projectId") String projectId,
-      @Config("defaultJobRegion") String beamJobRegion,
-      @Config("beamStagingUrl") String beamStagingUrl,
-      @Config("spec11TemplateUrl") String spec11TemplateUrl,
-      @Config("reportingBucketUrl") String reportingBucketUrl,
-      SerializableSupplier<JpaTransactionManager> jpaSupplierFactory,
-      @LocalCredential GoogleCredentialsBundle googleCredentialsBundle,
-      Retrier retrier) {
-    this.projectId = projectId;
-    this.beamJobRegion = beamJobRegion;
-    this.beamStagingUrl = beamStagingUrl;
-    this.spec11TemplateUrl = spec11TemplateUrl;
-    this.reportingBucketUrl = reportingBucketUrl;
-    this.jpaSupplierFactory = jpaSupplierFactory;
-    this.googleCredentials = googleCredentialsBundle.getGoogleCredentials();
-    this.retrier = retrier;
+  Spec11Pipeline(Spec11PipelineOptions options, EvaluateSafeBrowsingFn safeBrowsingFn) {
+    this.options = options;
+    this.safeBrowsingFn = safeBrowsingFn;
   }
 
-  /** Custom options for running the spec11 pipeline. */
-  public interface Spec11PipelineOptions extends DataflowPipelineOptions {
-    /** Returns the local date we're generating the report for, in yyyy-MM-dd format. */
-    @Description("The local date we generate the report for, in yyyy-MM-dd format.")
-    ValueProvider<String> getDate();
-
-    /**
-     * Sets the local date we generate invoices for.
-     *
-     * <p>This is implicitly set when executing the Dataflow template, by specifying the "date"
-     * parameter.
-     */
-    void setDate(ValueProvider<String> value);
-
-    /** Returns the SafeBrowsing API key we use to evaluate subdomain health. */
-    @Description("The API key we use to access the SafeBrowsing API.")
-    ValueProvider<String> getSafeBrowsingApiKey();
-
-    /**
-     * Sets the SafeBrowsing API key we use.
-     *
-     * <p>This is implicitly set when executing the Dataflow template, by specifying the
-     * "safeBrowsingApiKey" parameter.
-     */
-    void setSafeBrowsingApiKey(ValueProvider<String> value);
+  PipelineResult run() {
+    Pipeline pipeline = Pipeline.create(options);
+    setupPipeline(pipeline);
+    return pipeline.run();
   }
 
-  /** Deploys the spec11 pipeline as a template on GCS. */
-  public void deploy() {
-    // We can't store options as a member variable due to serialization concerns.
-    Spec11PipelineOptions options = PipelineOptionsFactory.as(Spec11PipelineOptions.class);
-    options.setProject(projectId);
-    options.setRegion(beamJobRegion);
-    options.setRunner(DataflowRunner.class);
-    // This causes p.run() to stage the pipeline as a template on GCS, as opposed to running it.
-    options.setTemplateLocation(spec11TemplateUrl);
-    options.setStagingLocation(beamStagingUrl);
-    // This credential is used when Dataflow deploys the template to GCS in target GCP project.
-    // So, make sure the credential has write permission to GCS in that project.
-    options.setGcpCredential(googleCredentials);
-
-    Pipeline p = Pipeline.create(options);
+  void setupPipeline(Pipeline pipeline) {
     PCollection<Subdomain> domains =
-        p.apply(
+        pipeline.apply(
             "Read active domains from BigQuery",
             BigQueryIO.read(Subdomain::parseFromRecord)
                 .fromQuery(
                     SqlTemplate.create(getQueryFromFile(Spec11Pipeline.class, "subdomains.sql"))
-                        .put("PROJECT_ID", projectId)
+                        .put("PROJECT_ID", options.getProject())
                         .put("DATASTORE_EXPORT_DATASET", "latest_datastore_export")
                         .put("REGISTRAR_TABLE", "Registrar")
                         .put("DOMAIN_BASE_TABLE", "DomainBase")
@@ -174,48 +112,40 @@ public class Spec11Pipeline implements Serializable {
                 .withoutValidation()
                 .withTemplateCompatibility());
 
-    evaluateUrlHealth(
-        domains,
-        new EvaluateSafeBrowsingFn(options.getSafeBrowsingApiKey(), retrier),
-        options.getDate());
-    p.run();
+    PCollection<KV<Subdomain, ThreatMatch>> threatMatches =
+        domains.apply("Run through SafeBrowsing API", ParDo.of(safeBrowsingFn));
+
+    saveToSql(threatMatches, options);
+    saveToGcs(threatMatches, options);
   }
 
-  /**
-   * Evaluate each {@link Subdomain} URL via the SafeBrowsing API.
-   *
-   * <p>This is factored out to facilitate testing.
-   */
-  void evaluateUrlHealth(
-      PCollection<Subdomain> domains,
-      EvaluateSafeBrowsingFn evaluateSafeBrowsingFn,
-      ValueProvider<String> dateProvider) {
-    PCollection<KV<Subdomain, ThreatMatch>> subdomainsSql =
-        domains.apply("Run through SafeBrowsing API", ParDo.of(evaluateSafeBrowsingFn));
-    TypeDescriptor<KV<Subdomain, ThreatMatch>> descriptor =
-        new TypeDescriptor<KV<Subdomain, ThreatMatch>>() {};
-    subdomainsSql.apply(
-        Transforms.writeToSql(
-            "Spec11ThreatMatch",
-            4,
-            4,
-            jpaSupplierFactory,
-            (kv) -> {
-              Subdomain subdomain = kv.getKey();
-              return new Spec11ThreatMatch.Builder()
-                  .setThreatTypes(ImmutableSet.of(ThreatType.valueOf(kv.getValue().threatType())))
-                  .setCheckDate(LocalDate.parse(dateProvider.get(), ISODateTimeFormat.date()))
-                  .setDomainName(subdomain.domainName())
-                  .setDomainRepoId(subdomain.domainRepoId())
-                  .setRegistrarId(subdomain.registrarId())
-                  .build();
-            },
-            descriptor));
+  static void saveToSql(
+      PCollection<KV<Subdomain, ThreatMatch>> threatMatches, Spec11PipelineOptions options) {
+    String transformId = "Spec11 Threat Matches";
+    LocalDate date = LocalDate.parse(options.getDate(), ISODateTimeFormat.date());
+    threatMatches.apply(
+        "Write to Sql: " + transformId,
+        RegistryJpaIO.<KV<Subdomain, ThreatMatch>>write()
+            .withName(transformId)
+            .withBatchSize(options.getSqlWriteBatchSize())
+            .withShards(options.getSqlWriteShards())
+            .withJpaConverter(
+                (kv) -> {
+                  Subdomain subdomain = kv.getKey();
+                  return new Spec11ThreatMatch.Builder()
+                      .setThreatTypes(
+                          ImmutableSet.of(ThreatType.valueOf(kv.getValue().threatType())))
+                      .setCheckDate(date)
+                      .setDomainName(subdomain.domainName())
+                      .setDomainRepoId(subdomain.domainRepoId())
+                      .setRegistrarId(subdomain.registrarId())
+                      .build();
+                }));
+  }
 
-    /* Store ThreatMatch objects in JSON. */
-    PCollection<KV<Subdomain, ThreatMatch>> subdomainsJson =
-        domains.apply("Run through SafeBrowsingAPI", ParDo.of(evaluateSafeBrowsingFn));
-    subdomainsJson
+  static void saveToGcs(
+      PCollection<KV<Subdomain, ThreatMatch>> threatMatches, Spec11PipelineOptions options) {
+    threatMatches
         .apply(
             "Map registrar ID to email/ThreatMatch pair",
             MapElements.into(
@@ -260,17 +190,54 @@ public class Spec11Pipeline implements Serializable {
             "Output to text file",
             TextIO.write()
                 .to(
-                    NestedValueProvider.of(
-                        dateProvider,
-                        date ->
-                            String.format(
-                                "%s/%s",
-                                reportingBucketUrl,
-                                getSpec11ReportFilePath(LocalDate.parse(date)))))
+                    String.format(
+                        "%s/%s",
+                        options.getReportingBucketUrl(),
+                        getSpec11ReportFilePath(LocalDate.parse(options.getDate()))))
                 .withoutSharding()
                 .withHeader("Map from registrar email / name to detected subdomain threats:"));
   }
 
+  public static void main(String[] args) {
+    PipelineOptionsFactory.register(Spec11PipelineOptions.class);
+    DaggerSpec11Pipeline_Spec11PipelineComponent.builder()
+        .spec11PipelineModule(new Spec11PipelineModule(args))
+        .build()
+        .spec11Pipeline()
+        .run();
+  }
+
+  @Module
+  static class Spec11PipelineModule {
+    private final String[] args;
+
+    Spec11PipelineModule(String[] args) {
+      this.args = args;
+    }
+
+    @Provides
+    Spec11PipelineOptions provideOptions() {
+      return PipelineOptionsFactory.fromArgs(args).withValidation().as(Spec11PipelineOptions.class);
+    }
+
+    @Provides
+    EvaluateSafeBrowsingFn provideSafeBrowsingFn(Spec11PipelineOptions options, Retrier retrier) {
+      return new EvaluateSafeBrowsingFn(options.getSafeBrowsingApiKey(), retrier);
+    }
+
+    @Provides
+    Spec11Pipeline providePipeline(
+        Spec11PipelineOptions options, EvaluateSafeBrowsingFn safeBrowsingFn) {
+      return new Spec11Pipeline(options, safeBrowsingFn);
+    }
+  }
+
+  @Component(modules = {Spec11PipelineModule.class, UtilsModule.class, ConfigModule.class})
+  @Singleton
+  interface Spec11PipelineComponent {
+    Spec11Pipeline spec11Pipeline();
+  }
+
   @AutoValue
   abstract static class EmailAndThreatMatch implements Serializable {
 

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

@@ -0,0 +1,37 @@
+// 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.spec11;
+
+import google.registry.beam.common.RegistryPipelineOptions;
+import org.apache.beam.sdk.options.Description;
+
+/** Custom options for running the spec11 pipeline. */
+public interface Spec11PipelineOptions extends RegistryPipelineOptions {
+
+  @Description("The local date we generate the report for, in yyyy-MM-dd format.")
+  String getDate();
+
+  void setDate(String value);
+
+  @Description("The API key we use to access the SafeBrowsing API.")
+  String getSafeBrowsingApiKey();
+
+  void setSafeBrowsingApiKey(String value);
+
+  @Description("The GCS bucket URL for Spec11 reports to be uploaded.")
+  String getReportingBucketUrl();
+
+  void setReportingBucketUrl(String value);
+}

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

@@ -15,6 +15,7 @@
 package google.registry.beam.spec11;
 
 import com.google.auto.value.AutoValue;
+import com.google.common.annotations.VisibleForTesting;
 import java.io.Serializable;
 import org.json.JSONException;
 import org.json.JSONObject;
@@ -31,16 +32,9 @@ public abstract class ThreatMatch implements Serializable {
   /** Returns the fully qualified domain name [SLD].[TLD] of the matched threat. */
   public abstract String fullyQualifiedDomainName();
 
-  /**
-   * Constructs a {@link ThreatMatch} by parsing a {@code SafeBrowsing API} response {@link
-   * JSONObject}.
-   *
-   * @throws JSONException when encountering parse errors in the response format
-   */
-  static ThreatMatch create(JSONObject threatMatchJSON, String fullyQualifiedDomainName)
-      throws JSONException {
-    return new AutoValue_ThreatMatch(
-        threatMatchJSON.getString(THREAT_TYPE_FIELD), fullyQualifiedDomainName);
+  @VisibleForTesting
+  static ThreatMatch create(String threatType, String fullyQualifiedDomainName) {
+    return new AutoValue_ThreatMatch(threatType, fullyQualifiedDomainName);
   }
 
   /** Returns a {@link JSONObject} representing a subset of this object's data. */

core/src/main/java/google/registry/config/RegistryConfig.java

@@ -384,31 +384,12 @@ public final class RegistryConfig {
       return Duration.standardHours(1);
     }
 
-    /**
-     * Number of sharded entity group roots used for performing strongly consistent scans.
-     *
-     * <p><b>Warning:</b> This number may increase but never decrease.
-     *
-     * @see google.registry.model.index.EppResourceIndex
-     */
-    @Provides
-    @Config("eppResourceIndexBucketCount")
-    public static int provideEppResourceIndexBucketCount(RegistryConfigSettings config) {
-      return config.datastore.eppResourceIndexBucketsNum;
-    }
-
     @Provides
     @Config("cloudSqlJdbcUrl")
     public static String providesCloudSqlJdbcUrl(RegistryConfigSettings config) {
       return config.cloudSql.jdbcUrl;
     }
 
-    @Provides
-    @Config("cloudSqlUsername")
-    public static String providesCloudSqlUsername(RegistryConfigSettings config) {
-      return config.cloudSql.username;
-    }
-
     @Provides
     @Config("cloudSqlInstanceConnectionName")
     public static String providesCloudSqlInstanceConnectionName(RegistryConfigSettings config) {
@@ -570,53 +551,6 @@ public final class RegistryConfig {
       return config.gSuite.outgoingEmailDisplayName;
     }
 
-    /**
-     * Returns the name of the GCS bucket for storing Beam templates and results.
-     *
-     * @see google.registry.reporting.billing.GenerateInvoicesAction
-     */
-    @Provides
-    @Config("apacheBeamBucket")
-    public static String provideApacheBeamBucket(@Config("projectId") String projectId) {
-      return projectId + "-beam";
-    }
-
-    /**
-     * Returns the URL of the GCS location for storing Apache Beam related objects.
-     *
-     * @see google.registry.reporting.billing.GenerateInvoicesAction
-     */
-    @Provides
-    @Config("apacheBeamBucketUrl")
-    public static String provideApacheBeamBucketUrl(@Config("apacheBeamBucket") String beamBucket) {
-      return "gs://" + beamBucket;
-    }
-
-    /**
-     * Returns the URL of the GCS location for storing the monthly invoicing Beam template.
-     *
-     * @see google.registry.reporting.billing.GenerateInvoicesAction
-     * @see google.registry.beam.invoicing.InvoicingPipeline
-     */
-    @Provides
-    @Config("invoiceTemplateUrl")
-    public static String provideInvoiceTemplateUrl(
-        @Config("apacheBeamBucketUrl") String beamBucketUrl) {
-      return beamBucketUrl + "/templates/invoicing";
-    }
-
-    /**
-     * Returns the URL of the GCS location for storing the monthly spec11 Beam template.
-     *
-     * @see google.registry.beam.spec11.Spec11Pipeline
-     */
-    @Provides
-    @Config("spec11TemplateUrl")
-    public static String provideSpec11TemplateUrl(
-        @Config("apacheBeamBucketUrl") String beamBucketUrl) {
-      return beamBucketUrl + "/templates/spec11";
-    }
-
     /**
      * Returns whether an SSL certificate hash is required to log in via EPP and run flows.
      *
@@ -640,18 +574,6 @@ public final class RegistryConfig {
       return config.beam.defaultJobRegion;
     }
 
-    /**
-     * Returns the default job zone to run Apache Beam (Cloud Dataflow) jobs in.
-     *
-     * @see google.registry.reporting.billing.GenerateInvoicesAction
-     * @see google.registry.reporting.spec11.GenerateSpec11ReportAction
-     */
-    @Provides
-    @Config("defaultJobZone")
-    public static String provideDefaultJobZone(RegistryConfigSettings config) {
-      return config.beam.defaultJobZone;
-    }
-
     /** Returns the GCS bucket URL with all staged BEAM flex templates. */
     @Provides
     @Config("beamStagingBucketUrl")
@@ -659,19 +581,6 @@ public final class RegistryConfig {
       return config.beam.stagingBucketUrl;
     }
 
-    /**
-     * Returns the URL of the GCS location we store jar dependencies for beam pipelines.
-     *
-     * @see google.registry.beam.invoicing.InvoicingPipeline
-     * @see google.registry.beam.spec11.Spec11Pipeline
-     */
-    @Provides
-    @Config("beamStagingUrl")
-    public static String provideInvoiceStagingUrl(
-        @Config("apacheBeamBucketUrl") String beamBucketUrl) {
-      return beamBucketUrl + "/staging";
-    }
-
     /**
      * Returns the Google Cloud Storage bucket for Spec11 and ICANN transaction and activity reports
      * to be uploaded.
@@ -1233,14 +1142,6 @@ public final class RegistryConfig {
       return formatComments(config.registryPolicy.reservedTermsExportDisclaimer);
     }
 
-    /** Returns the clientId of the registrar used by the {@code CheckApiServlet}. */
-    // TODO(b/80417678): remove this once CheckApiAction no longer uses this id.
-    @Provides
-    @Config("checkApiServletRegistrarClientId")
-    public static String provideCheckApiServletRegistrarClientId(RegistryConfigSettings config) {
-      return config.registryPolicy.checkApiServletClientId;
-    }
-
     /**
      * Returns the clientId of the registrar that admins are automatically logged in as if they
      * aren't otherwise associated with one.
@@ -1342,12 +1243,6 @@ public final class RegistryConfig {
       return config.registryTool.clientSecret;
     }
 
-    @Provides
-    @Config("toolsCloudSqlUsername")
-    public static String providesToolsCloudSqlUsername(RegistryConfigSettings config) {
-      return config.registryTool.username;
-    }
-
     @Provides
     @Config("rdapTos")
     public static ImmutableList<String> provideRdapTos(RegistryConfigSettings config) {

core/src/main/java/google/registry/config/RegistryConfigSettings.java

@@ -123,6 +123,7 @@ public class RegistryConfigSettings {
   /** Configuration for Cloud SQL. */
   public static class CloudSql {
     public String jdbcUrl;
+    // TODO(05012021): remove username field after it is removed from all yaml files.
     public String username;
     public String instanceConnectionName;
     public boolean replicateTransactions;
@@ -132,7 +133,6 @@ public class RegistryConfigSettings {
   /** Configuration for Apache Beam (Cloud Dataflow). */
   public static class Beam {
     public String defaultJobRegion;
-    public String defaultJobZone;
     public String stagingBucketUrl;
   }
 
@@ -221,6 +221,7 @@ public class RegistryConfigSettings {
   public static class RegistryTool {
     public String clientId;
     public String clientSecret;
+    // TODO(05012021): remove username field after it is removed from all yaml files.
     public String username;
   }
 

core/src/main/java/google/registry/config/files/default-config.yaml

@@ -225,8 +225,6 @@ cloudSql:
   # If jdbcUrl in this file is moved elsewhere, be sure to move this notice
   # with it until the change is applied.
   jdbcUrl: jdbc:postgresql://localhost
-  # Username for the database user.
-  username: username
   # This name is used by Cloud SQL when connecting to the database.
   instanceConnectionName: project-id:region:instance-id
   # Set this to true to replicate cloud SQL transactions to datastore in the
@@ -422,9 +420,6 @@ misc:
 beam:
   # The default region to run Apache Beam (Cloud Dataflow) jobs in.
   defaultJobRegion: us-east1
-  # The default zone to run Apache Beam (Cloud Dataflow) jobs in.
-  # TODO(weiminyu): consider dropping zone config. No obvious needs for this.
-  defaultJobZone: us-east1-c
   stagingBucketUrl: gcs-bucket-with-staged-templates
 
 keyring:
@@ -447,7 +442,6 @@ registryTool:
   clientId: YOUR_CLIENT_ID
   # OAuth client secret used by the tool.
   clientSecret: YOUR_CLIENT_SECRET
-  username: toolusername
 
 # Configuration options for checking SSL certificates.
 sslCertificateValidation:

core/src/main/java/google/registry/env/common/backend/WEB-INF/web.xml

@@ -385,6 +385,12 @@
     <url-pattern>/_dr/task/wipeOutCloudSql</url-pattern>
   </servlet-mapping>
 
+  <!-- Action to wipeout Cloud Datastore data -->
+  <servlet-mapping>
+    <servlet-name>backend-servlet</servlet-name>
+    <url-pattern>/_dr/task/wipeOutDatastore</url-pattern>
+  </servlet-mapping>
+
   <!-- Security config -->
   <security-constraint>
     <web-resource-collection>

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

@@ -91,4 +91,13 @@
     <target>backend</target>
   </cron>
 
+  <cron>
+    <url><![CDATA[/_dr/task/wipeOutDatastore]]></url>
+    <description>
+      This job runs an action that deletes all data in Cloud Datastore.
+    </description>
+    <schedule>every saturday 03:07</schedule>
+    <target>backend</target>
+  </cron>
+
 </cronentries>

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

@@ -16,7 +16,6 @@ package google.registry.export;
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.request.Action.Method.POST;
 import static google.registry.util.CollectionUtils.nullToEmpty;
@@ -164,7 +163,7 @@ public final class SyncGroupMembersAction implements Runnable {
         registrarsToSave.add(result.getKey().asBuilder().setContactsRequireSyncing(false).build());
       }
     }
-    tm().transactNew(() -> ofy().save().entities(registrarsToSave.build()));
+    tm().transactNew(() -> tm().updateAll(registrarsToSave.build()));
     return errors;
   }
 

core/src/main/java/google/registry/export/sheet/SyncRegistrarsSheet.java

@@ -17,7 +17,6 @@ package google.registry.export.sheet;
 import static com.google.common.base.MoreObjects.firstNonNull;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static google.registry.model.common.Cursor.CursorType.SYNC_REGISTRAR_SHEET;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.registrar.RegistrarContact.Type.ABUSE;
 import static google.registry.model.registrar.RegistrarContact.Type.ADMIN;
 import static google.registry.model.registrar.RegistrarContact.Type.BILLING;
@@ -26,6 +25,7 @@ import static google.registry.model.registrar.RegistrarContact.Type.MARKETING;
 import static google.registry.model.registrar.RegistrarContact.Type.TECH;
 import static google.registry.model.registrar.RegistrarContact.Type.WHOIS;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
 
 import com.google.common.base.Joiner;
@@ -40,6 +40,7 @@ import google.registry.model.registrar.RegistrarContact;
 import google.registry.util.Clock;
 import google.registry.util.DateTimeUtils;
 import java.io.IOException;
+import java.util.Optional;
 import java.util.function.Predicate;
 import javax.annotation.Nullable;
 import javax.inject.Inject;
@@ -61,8 +62,10 @@ class SyncRegistrarsSheet {
    * successfully completed, as measured by a cursor.
    */
   boolean wereRegistrarsModified() {
-    Cursor cursor = ofy().load().key(Cursor.createGlobalKey(SYNC_REGISTRAR_SHEET)).now();
-    DateTime lastUpdateTime = (cursor == null) ? START_OF_TIME : cursor.getCursorTime();
+    Optional<Cursor> cursor =
+        transactIfJpaTm(
+            () -> tm().loadByKeyIfPresent(Cursor.createGlobalVKey(SYNC_REGISTRAR_SHEET)));
+    DateTime lastUpdateTime = !cursor.isPresent() ? START_OF_TIME : cursor.get().getCursorTime();
     for (Registrar registrar : Registrar.loadAllCached()) {
       if (DateTimeUtils.isAtOrAfter(registrar.getLastUpdateTime(), lastUpdateTime)) {
         return true;

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

@@ -20,6 +20,7 @@ import com.google.common.base.Strings;
 import dagger.Module;
 import dagger.Provides;
 import google.registry.flows.picker.FlowPicker;
+import google.registry.model.domain.DomainHistory;
 import google.registry.model.domain.metadata.MetadataExtension;
 import google.registry.model.eppcommon.AuthInfo;
 import google.registry.model.eppcommon.Trid;
@@ -239,6 +240,12 @@ public class FlowModule {
     return historyBuilder;
   }
 
+  @Provides
+  static DomainHistory.Builder provideDomainHistoryBuilder(
+      HistoryEntry.Builder historyEntryBuilder) {
+    return new DomainHistory.Builder().copyFrom(historyEntryBuilder);
+  }
+
   /**
    * Provides a partially filled in {@link EppResponse} builder.
    *

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

@@ -15,22 +15,26 @@
 package google.registry.flows;
 
 import static com.google.common.base.Preconditions.checkState;
-import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.xml.ValidationMode.LENIENT;
 import static google.registry.xml.ValidationMode.STRICT;
 import static java.nio.charset.StandardCharsets.UTF_8;
 
 import com.google.common.base.Throwables;
 import com.google.common.flogger.FluentLogger;
+import com.googlecode.objectify.Key;
 import google.registry.flows.EppException.CommandUseErrorException;
 import google.registry.flows.EppException.ParameterValueRangeErrorException;
 import google.registry.flows.EppException.SyntaxErrorException;
 import google.registry.flows.EppException.UnimplementedProtocolVersionException;
 import google.registry.flows.custom.EntityChanges;
+import google.registry.model.EppResource;
 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;
 import java.util.List;
@@ -51,8 +55,8 @@ public final class FlowUtils {
 
   /** Persists the saves and deletes in an {@link EntityChanges} to Datastore. */
   public static void persistEntityChanges(EntityChanges entityChanges) {
-    ofy().save().entities(entityChanges.getSaves());
-    ofy().delete().keys(entityChanges.getDeletes());
+    tm().putAll(entityChanges.getSaves());
+    tm().delete(entityChanges.getDeletes());
   }
 
   /**
@@ -99,6 +103,11 @@ 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());
+  }
+
   /** Registrar is not logged in. */
   public static class NotLoggedInException extends CommandUseErrorException {
     public NotLoggedInException() {

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

@@ -16,7 +16,6 @@ package google.registry.flows;
 
 import static com.google.common.base.MoreObjects.toStringHelper;
 import static google.registry.request.RequestParameters.extractOptionalHeader;
-import static google.registry.util.X509Utils.loadCertificate;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
@@ -26,24 +25,17 @@ import com.google.common.net.InetAddresses;
 import dagger.Module;
 import dagger.Provides;
 import google.registry.config.RegistryConfig.Config;
-import google.registry.config.RegistryEnvironment;
 import google.registry.flows.EppException.AuthenticationErrorException;
 import google.registry.flows.certs.CertificateChecker;
 import google.registry.flows.certs.CertificateChecker.InsecureCertificateException;
 import google.registry.model.registrar.Registrar;
 import google.registry.request.Header;
 import google.registry.util.CidrAddressBlock;
-import google.registry.util.Clock;
 import google.registry.util.ProxyHttpHeaders;
-import java.io.ByteArrayInputStream;
 import java.net.InetAddress;
-import java.security.cert.CertificateException;
-import java.security.cert.X509Certificate;
-import java.util.Base64;
 import java.util.Optional;
 import javax.inject.Inject;
 import javax.servlet.http.HttpServletRequest;
-import org.joda.time.DateTime;
 
 /**
  * Container and validation for TLS certificate and IP-allow-listing.
@@ -54,10 +46,6 @@ import org.joda.time.DateTime;
  *   <dt>X-SSL-Certificate
  *   <dd>This field should contain a base64 encoded digest of the client's TLS certificate. It is
  *       used only if the validation of the full certificate fails.
- *   <dt>X-SSL-Full-Certificate
- *   <dd>This field should contain a base64 encoding of the client's TLS certificate. It is
- *       validated during an EPP login command against a known good value that is transmitted out of
- *       band.
  *   <dt>X-Forwarded-For
  *   <dd>This field should contain the host and port of the connecting client. It is validated
  *       during an EPP login command against an IP allow list that is transmitted out of band.
@@ -66,30 +54,22 @@ import org.joda.time.DateTime;
 public class TlsCredentials implements TransportCredentials {
 
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
-  private static final DateTime CERT_ENFORCEMENT_START_TIME =
-      DateTime.parse("2021-03-01T16:00:00Z");
 
   private final boolean requireSslCertificates;
   private final Optional<String> clientCertificateHash;
-  private final Optional<String> clientCertificate;
   private final Optional<InetAddress> clientInetAddr;
   private final CertificateChecker certificateChecker;
-  private final Clock clock;
 
   @Inject
   public TlsCredentials(
       @Config("requireSslCertificates") boolean requireSslCertificates,
       @Header(ProxyHttpHeaders.CERTIFICATE_HASH) Optional<String> clientCertificateHash,
-      @Header(ProxyHttpHeaders.FULL_CERTIFICATE) Optional<String> clientCertificate,
       @Header(ProxyHttpHeaders.IP_ADDRESS) Optional<String> clientAddress,
-      CertificateChecker certificateChecker,
-      Clock clock) {
+      CertificateChecker certificateChecker) {
     this.requireSslCertificates = requireSslCertificates;
     this.clientCertificateHash = clientCertificateHash;
-    this.clientCertificate = clientCertificate;
     this.clientInetAddr = clientAddress.map(TlsCredentials::parseInetAddress);
     this.certificateChecker = certificateChecker;
-    this.clock = clock;
   }
 
   static InetAddress parseInetAddress(String asciiAddr) {
@@ -103,7 +83,7 @@ public class TlsCredentials implements TransportCredentials {
   @Override
   public void validate(Registrar registrar, String password) throws AuthenticationErrorException {
     validateIp(registrar);
-    validateCertificate(registrar);
+    validateCertificateHash(registrar);
     validatePassword(registrar, password);
   }
 
@@ -137,89 +117,8 @@ public class TlsCredentials implements TransportCredentials {
     throw new BadRegistrarIpAddressException();
   }
 
-  /**
-   * Verifies client SSL certificate is permitted to issue commands as {@code registrar}.
-   *
-   * @throws MissingRegistrarCertificateException if frontend didn't send certificate header
-   * @throws BadRegistrarCertificateException if registrar requires certificate and it didn't match
-   */
   @VisibleForTesting
-  void validateCertificate(Registrar registrar) throws AuthenticationErrorException {
-    // Check that certificate is present in registrar object
-    if (!registrar.getClientCertificate().isPresent()
-        && !registrar.getFailoverClientCertificate().isPresent()) {
-      // Log an error and validate using certificate hash instead
-      // TODO(sarahbot): throw a RegistrarCertificateNotConfiguredException once hash is no longer
-      // used as failover
-      logger.atWarning().log(
-          "There is no certificate configured for registrar %s.", registrar.getClientId());
-    } else if (!clientCertificate.isPresent()) {
-      // Check that the request included the full certificate
-      // Log an error and validate using certificate hash instead
-      // TODO(sarahbot): throw a MissingRegistrarCertificateException once hash is no longer used as
-      // failover
-      logger.atWarning().log(
-          "Request from registrar %s did not include X-SSL-Full-Certificate.",
-          registrar.getClientId());
-    } else {
-      X509Certificate passedCert;
-      Optional<X509Certificate> storedCert;
-      Optional<X509Certificate> storedFailoverCert;
-
-      try {
-        storedCert = deserializePemCert(registrar.getClientCertificate());
-        storedFailoverCert = deserializePemCert(registrar.getFailoverClientCertificate());
-        passedCert = decodeCertString(clientCertificate.get());
-      } catch (Exception e) {
-        // TODO(Sarahbot@): remove this catch once we know it's working
-        logger.atWarning().log(
-            "Error converting certificate string to certificate for %s: %s",
-            registrar.getClientId(), e);
-        validateCertificateHash(registrar);
-        return;
-      }
-
-      // Check if the certificate is equal to the one on file for the registrar.
-      if (passedCert.equals(storedCert.orElse(null))
-          || passedCert.equals(storedFailoverCert.orElse(null))) {
-        // Check certificate for any requirement violations
-        // TODO(Sarahbot@): Throw exceptions instead of just logging once requirement enforcement
-        // begins
-        try {
-          certificateChecker.validateCertificate(passedCert);
-        } catch (InsecureCertificateException e) {
-          // TODO(Sarahbot@): Remove this if statement after March 1. After March 1, exception
-          // should be thrown in all environments.
-          // throw exception in unit tests and Sandbox
-          if (RegistryEnvironment.get().equals(RegistryEnvironment.UNITTEST)
-              || RegistryEnvironment.get().equals(RegistryEnvironment.SANDBOX)
-              || clock.nowUtc().isAfter(CERT_ENFORCEMENT_START_TIME)) {
-            throw new CertificateContainsSecurityViolationsException(e);
-          }
-          logger.atWarning().log(
-              "Registrar certificate used for %s does not meet certificate requirements: %s",
-              registrar.getClientId(), e.getMessage());
-        } catch (Exception e) {
-          logger.atWarning().log(
-              "Error validating certificate for %s: %s", registrar.getClientId(), e);
-        }
-        // successfully validated, return here since hash validation is not necessary
-        return;
-      }
-      // Log an error and validate using certificate hash instead
-      // TODO(sarahbot): throw a BadRegistrarCertificateException once hash is no longer used as
-      // failover
-      logger.atWarning().log("Non-matching certificate for registrar %s.", registrar.getClientId());
-    }
-    validateCertificateHash(registrar);
-  }
-
-  private void validateCertificateHash(Registrar registrar) throws AuthenticationErrorException {
-    logger.atWarning().log(
-        "Error validating certificate for %s, attempting to validate using certificate hash.",
-        registrar.getClientId());
-    // Check the certificate hash as a failover
-    // TODO(sarahbot): Remove hash checks once certificate checks are working.
+  void validateCertificateHash(Registrar registrar) throws AuthenticationErrorException {
     if (!registrar.getClientCertificateHash().isPresent()
         && !registrar.getFailoverClientCertificateHash().isPresent()) {
       if (requireSslCertificates) {
@@ -247,6 +146,20 @@ public class TlsCredentials implements TransportCredentials {
           registrar.getFailoverClientCertificateHash());
       throw new BadRegistrarCertificateException();
     }
+    if (requireSslCertificates) {
+      String passedCert =
+          clientCertificateHash.equals(registrar.getClientCertificateHash())
+              ? registrar.getClientCertificate().get()
+              : registrar.getFailoverClientCertificate().get();
+      try {
+        certificateChecker.validateCertificate(passedCert);
+      } catch (InsecureCertificateException e) {
+        logger.atWarning().log(
+            "Registrar certificate used for %s does not meet certificate requirements: %s",
+            registrar.getClientId(), e.getMessage());
+        throw new CertificateContainsSecurityViolationsException(e);
+      }
+    }
   }
 
   private void validatePassword(Registrar registrar, String password)
@@ -256,26 +169,9 @@ public class TlsCredentials implements TransportCredentials {
     }
   }
 
-  // Converts a PEM formatted certificate string into an X509Certificate
-  private Optional<X509Certificate> deserializePemCert(Optional<String> certificateString)
-      throws CertificateException {
-    if (certificateString.isPresent()) {
-      return Optional.of(loadCertificate(certificateString.get()));
-    }
-    return Optional.empty();
-  }
-
-  // Decodes the string representation of an encoded certificate back into an X509Certificate
-  private X509Certificate decodeCertString(String encodedCertString) throws CertificateException {
-    byte decodedCert[] = Base64.getDecoder().decode(encodedCertString);
-    ByteArrayInputStream inputStream = new ByteArrayInputStream(decodedCert);
-    return loadCertificate(inputStream);
-  }
-
   @Override
   public String toString() {
     return toStringHelper(getClass())
-        .add("clientCertificate", clientCertificate.orElse(null))
         .add("clientCertificateHash", clientCertificateHash.orElse(null))
         .add("clientAddress", clientInetAddr.orElse(null))
         .toString();
@@ -336,14 +232,6 @@ public class TlsCredentials implements TransportCredentials {
       return extractOptionalHeader(req, ProxyHttpHeaders.CERTIFICATE_HASH);
     }
 
-    @Provides
-    @Header(ProxyHttpHeaders.FULL_CERTIFICATE)
-    static Optional<String> provideClientCertificate(HttpServletRequest req) {
-      // Note: This header is actually required, we just want to handle its absence explicitly
-      // by throwing an EPP exception rather than a generic Bad Request exception.
-      return extractOptionalHeader(req, ProxyHttpHeaders.FULL_CERTIFICATE);
-    }
-
     @Provides
     @Header(ProxyHttpHeaders.IP_ADDRESS)
     static Optional<String> provideIpAddress(HttpServletRequest req) {

core/src/main/java/google/registry/flows/custom/EntityChanges.java

@@ -16,8 +16,8 @@ package google.registry.flows.custom;
 
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableSet;
-import com.googlecode.objectify.Key;
 import google.registry.model.ImmutableObject;
+import google.registry.persistence.VKey;
 
 /** A wrapper class that encapsulates Datastore entities to both save and delete. */
 @AutoValue
@@ -25,7 +25,7 @@ public abstract class EntityChanges {
 
   public abstract ImmutableSet<ImmutableObject> getSaves();
 
-  public abstract ImmutableSet<Key<ImmutableObject>> getDeletes();
+  public abstract ImmutableSet<VKey<ImmutableObject>> getDeletes();
 
   public static Builder newBuilder() {
     // Default both entities to save and entities to delete to empty sets, so that the build()
@@ -48,11 +48,11 @@ public abstract class EntityChanges {
       return this;
     }
 
-    public abstract Builder setDeletes(ImmutableSet<Key<ImmutableObject>> entitiesToDelete);
+    public abstract Builder setDeletes(ImmutableSet<VKey<ImmutableObject>> entitiesToDelete);
 
-    public abstract ImmutableSet.Builder<Key<ImmutableObject>> deletesBuilder();
+    public abstract ImmutableSet.Builder<VKey<ImmutableObject>> deletesBuilder();
 
-    public Builder addDelete(Key<ImmutableObject> entityToDelete) {
+    public Builder addDelete(VKey<ImmutableObject> entityToDelete) {
       deletesBuilder().add(entityToDelete);
       return this;
     }

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

@@ -80,6 +80,7 @@ import google.registry.model.billing.BillingEvent.Recurring;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.DomainCommand;
 import google.registry.model.domain.DomainCommand.Create;
+import google.registry.model.domain.DomainHistory;
 import google.registry.model.domain.GracePeriod;
 import google.registry.model.domain.Period;
 import google.registry.model.domain.fee.FeeCreateCommandExtension;
@@ -96,7 +97,6 @@ import google.registry.model.eppinput.EppInput;
 import google.registry.model.eppinput.ResourceCommand;
 import google.registry.model.eppoutput.CreateData.DomainCreateData;
 import google.registry.model.eppoutput.EppResponse;
-import google.registry.model.host.HostResource;
 import google.registry.model.index.EppResourceIndex;
 import google.registry.model.index.ForeignKeyIndex;
 import google.registry.model.ofy.ObjectifyService;
@@ -106,11 +106,11 @@ import google.registry.model.poll.PollMessage.Autorenew;
 import google.registry.model.registry.Registry;
 import google.registry.model.registry.Registry.TldState;
 import google.registry.model.registry.Registry.TldType;
+import google.registry.model.registry.label.ReservationType;
 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.persistence.VKey;
 import google.registry.tmch.LordnTaskUtils;
 import java.util.Optional;
 import javax.inject.Inject;
@@ -206,7 +206,7 @@ public class DomainCreateFlow implements TransactionalFlow {
   @Inject @ClientId String clientId;
   @Inject @TargetId String targetId;
   @Inject @Superuser boolean isSuperuser;
-  @Inject HistoryEntry.Builder historyBuilder;
+  @Inject DomainHistory.Builder historyBuilder;
   @Inject EppResponse.Builder responseBuilder;
   @Inject AllocationTokenFlowUtils allocationTokenFlowUtils;
   @Inject DomainCreateFlowCustomLogic flowCustomLogic;
@@ -301,10 +301,14 @@ public class DomainCreateFlow implements TransactionalFlow {
     validateFeeChallenge(targetId, now, feeCreate, feesAndCredits);
     Optional<SecDnsCreateExtension> secDnsCreate =
         validateSecDnsExtension(eppInput.getSingleExtension(SecDnsCreateExtension.class));
-    String repoId = createDomainRepoId(ObjectifyService.allocateId(), registry.getTldStr());
     DateTime registrationExpirationTime = leapSafeAddYears(now, years);
-    HistoryEntry historyEntry = buildHistoryEntry(
-        repoId, registry, now, period, registry.getAddGracePeriodLength());
+    String repoId = createDomainRepoId(ObjectifyService.allocateId(), registry.getTldStr());
+    Key<DomainHistory> domainHistoryKey =
+        Key.create(
+            Key.create(DomainBase.class, repoId),
+            DomainHistory.class,
+            ObjectifyService.allocateId());
+    historyBuilder.setId(domainHistoryKey.getId());
     // Bill for the create.
     BillingEvent.OneTime createBillingEvent =
         createOneTimeBillingEvent(
@@ -314,33 +318,27 @@ public class DomainCreateFlow implements TransactionalFlow {
             isReserved(domainName, isSunriseCreate),
             years,
             feesAndCredits,
-            historyEntry,
+            domainHistoryKey,
             allocationToken,
             now);
     // Create a new autorenew billing event and poll message starting at the expiration time.
     BillingEvent.Recurring autorenewBillingEvent =
-        createAutorenewBillingEvent(historyEntry, registrationExpirationTime);
+        createAutorenewBillingEvent(domainHistoryKey, registrationExpirationTime);
     PollMessage.Autorenew autorenewPollMessage =
-        createAutorenewPollMessage(historyEntry, registrationExpirationTime);
+        createAutorenewPollMessage(domainHistoryKey, registrationExpirationTime);
     ImmutableSet.Builder<ImmutableObject> entitiesToSave = new ImmutableSet.Builder<>();
-    entitiesToSave.add(
-        historyEntry,
-        createBillingEvent,
-        autorenewBillingEvent,
-        autorenewPollMessage);
+    entitiesToSave.add(createBillingEvent, autorenewBillingEvent, autorenewPollMessage);
     // Bill for EAP cost, if any.
     if (!feesAndCredits.getEapCost().isZero()) {
       entitiesToSave.add(createEapBillingEvent(feesAndCredits, createBillingEvent));
     }
 
-    ImmutableSet.Builder<StatusValue> statuses = new ImmutableSet.Builder<>();
-    if (getReservationTypes(domainName).contains(NAME_COLLISION)) {
-      statuses.add(SERVER_HOLD);
-      entitiesToSave.add(
-          createNameCollisionOneTimePollMessage(targetId, historyEntry, clientId, now));
-    }
-
-    DomainBase newDomain =
+    ImmutableSet<ReservationType> reservationTypes = getReservationTypes(domainName);
+    ImmutableSet<StatusValue> statuses =
+        reservationTypes.contains(NAME_COLLISION)
+            ? ImmutableSet.of(SERVER_HOLD)
+            : ImmutableSet.of();
+    DomainBase domain =
         new DomainBase.Builder()
             .setCreationClientId(clientId)
             .setPersistedCurrentSponsorClientId(clientId)
@@ -351,35 +349,39 @@ public class DomainCreateFlow implements TransactionalFlow {
             .setAutorenewPollMessage(autorenewPollMessage.createVKey())
             .setLaunchNotice(hasClaimsNotice ? launchCreate.get().getNotice() : null)
             .setSmdId(signedMarkId)
-            .setDsData(secDnsCreate.isPresent() ? secDnsCreate.get().getDsData() : null)
+            .setDsData(secDnsCreate.map(SecDnsCreateExtension::getDsData).orElse(null))
             .setRegistrant(command.getRegistrant())
             .setAuthInfo(command.getAuthInfo())
             .setDomainName(targetId)
-            .setNameservers(
-                (ImmutableSet<VKey<HostResource>>)
-                    command.getNameservers().stream().collect(toImmutableSet()))
-            .setStatusValues(statuses.build())
+            .setNameservers(command.getNameservers().stream().collect(toImmutableSet()))
+            .setStatusValues(statuses)
             .setContacts(command.getContacts())
             .addGracePeriod(
                 GracePeriod.forBillingEvent(GracePeriodStatus.ADD, repoId, createBillingEvent))
             .build();
+    DomainHistory domainHistory =
+        buildDomainHistory(domain, registry, now, period, registry.getAddGracePeriodLength());
+    if (reservationTypes.contains(NAME_COLLISION)) {
+      entitiesToSave.add(
+          createNameCollisionOneTimePollMessage(targetId, domainHistory, clientId, now));
+    }
     entitiesToSave.add(
-        newDomain,
-        ForeignKeyIndex.create(newDomain, newDomain.getDeletionTime()),
-        EppResourceIndex.create(Key.create(newDomain)));
+        domain,
+        domainHistory,
+        ForeignKeyIndex.create(domain, domain.getDeletionTime()),
+        EppResourceIndex.create(Key.create(domain)));
     if (allocationToken.isPresent()
         && TokenType.SINGLE_USE.equals(allocationToken.get().getTokenType())) {
       entitiesToSave.add(
-          allocationTokenFlowUtils.redeemToken(
-              allocationToken.get(), HistoryEntry.createVKey(Key.create(historyEntry))));
+          allocationTokenFlowUtils.redeemToken(allocationToken.get(), domainHistory.createVKey()));
     }
-    enqueueTasks(newDomain, hasSignedMarks, hasClaimsNotice);
+    enqueueTasks(domain, hasSignedMarks, hasClaimsNotice);
 
     EntityChanges entityChanges =
         flowCustomLogic.beforeSave(
             DomainCreateFlowCustomLogic.BeforeSaveParameters.newBuilder()
-                .setNewDomain(newDomain)
-                .setHistoryEntry(historyEntry)
+                .setNewDomain(domain)
+                .setHistoryEntry(domainHistory)
                 .setEntityChanges(
                     EntityChanges.newBuilder().setSaves(entitiesToSave.build()).build())
                 .setYears(years)
@@ -483,8 +485,8 @@ public class DomainCreateFlow implements TransactionalFlow {
             : null);
   }
 
-  private HistoryEntry buildHistoryEntry(
-      String repoId, Registry registry, DateTime now, Period period, Duration addGracePeriod) {
+  private DomainHistory buildDomainHistory(
+      DomainBase domain, Registry registry, DateTime now, Period period, Duration addGracePeriod) {
     // We ignore prober transactions
     if (registry.getTldType() == TldType.REAL) {
       historyBuilder
@@ -500,7 +502,7 @@ public class DomainCreateFlow implements TransactionalFlow {
         .setType(HistoryEntry.Type.DOMAIN_CREATE)
         .setPeriod(period)
         .setModificationTime(now)
-        .setParent(Key.create(DomainBase.class, repoId))
+        .setDomainContent(domain)
         .build();
   }
 
@@ -511,7 +513,7 @@ public class DomainCreateFlow implements TransactionalFlow {
       boolean isReserved,
       int years,
       FeesAndCredits feesAndCredits,
-      HistoryEntry historyEntry,
+      Key<DomainHistory> domainHistoryKey,
       Optional<AllocationToken> allocationToken,
       DateTime now) {
     ImmutableSet.Builder<Flag> flagsBuilder = new ImmutableSet.Builder<>();
@@ -540,12 +542,12 @@ public class DomainCreateFlow implements TransactionalFlow {
                     ? registry.getAnchorTenantAddGracePeriodLength()
                     : registry.getAddGracePeriodLength()))
         .setFlags(flagsBuilder.build())
-        .setParent(historyEntry)
+        .setParent(domainHistoryKey)
         .build();
   }
 
   private Recurring createAutorenewBillingEvent(
-      HistoryEntry historyEntry, DateTime registrationExpirationTime) {
+      Key<DomainHistory> domainHistoryKey, DateTime registrationExpirationTime) {
     return new BillingEvent.Recurring.Builder()
         .setReason(Reason.RENEW)
         .setFlags(ImmutableSet.of(Flag.AUTO_RENEW))
@@ -553,18 +555,18 @@ public class DomainCreateFlow implements TransactionalFlow {
         .setClientId(clientId)
         .setEventTime(registrationExpirationTime)
         .setRecurrenceEndTime(END_OF_TIME)
-        .setParent(historyEntry)
+        .setParent(domainHistoryKey)
         .build();
   }
 
   private Autorenew createAutorenewPollMessage(
-      HistoryEntry historyEntry, DateTime registrationExpirationTime) {
+      Key<DomainHistory> domainHistoryKey, DateTime registrationExpirationTime) {
     return new PollMessage.Autorenew.Builder()
         .setTargetId(targetId)
         .setClientId(clientId)
         .setEventTime(registrationExpirationTime)
         .setMsg("Domain was auto-renewed.")
-        .setParent(historyEntry)
+        .setParentKey(domainHistoryKey)
         .build();
   }
 

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

@@ -16,6 +16,7 @@ package google.registry.flows.domain;
 
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Strings.isNullOrEmpty;
+import static google.registry.flows.FlowUtils.createHistoryKey;
 import static google.registry.flows.FlowUtils.persistEntityChanges;
 import static google.registry.flows.FlowUtils.validateClientIsLoggedIn;
 import static google.registry.flows.ResourceFlowUtils.loadAndVerifyExistence;
@@ -63,6 +64,7 @@ import google.registry.flows.custom.EntityChanges;
 import google.registry.model.ImmutableObject;
 import google.registry.model.billing.BillingEvent;
 import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.DomainHistory;
 import google.registry.model.domain.GracePeriod;
 import google.registry.model.domain.fee.BaseFee.FeeType;
 import google.registry.model.domain.fee.Credit;
@@ -125,7 +127,7 @@ public final class DomainDeleteFlow implements TransactionalFlow {
   @Inject @ClientId String clientId;
   @Inject @TargetId String targetId;
   @Inject @Superuser boolean isSuperuser;
-  @Inject HistoryEntry.Builder historyBuilder;
+  @Inject DomainHistory.Builder historyBuilder;
   @Inject DnsQueue dnsQueue;
   @Inject Trid trid;
   @Inject AsyncTaskEnqueuer asyncTaskEnqueuer;
@@ -177,8 +179,8 @@ public final class DomainDeleteFlow implements TransactionalFlow {
             ? Duration.ZERO
             // By default, this should be 30 days of grace, and 5 days of pending delete.
             : redemptionGracePeriodLength.plus(pendingDeleteLength);
-    HistoryEntry historyEntry =
-        buildHistoryEntry(existingDomain, registry, now, durationUntilDelete, inAddGracePeriod);
+    Key<DomainHistory> domainHistoryKey = createHistoryKey(existingDomain, DomainHistory.class);
+    historyBuilder.setId(domainHistoryKey.getId());
     DateTime deletionTime = now.plus(durationUntilDelete);
     if (durationUntilDelete.equals(Duration.ZERO)) {
       builder.setDeletionTime(now).setStatusValues(null);
@@ -208,20 +210,28 @@ public final class DomainDeleteFlow implements TransactionalFlow {
     // Enqueue the deletion poll message if the delete is asynchronous or if requested by a
     // superuser (i.e. the registrar didn't request this delete and thus should be notified even if
     // it is synchronous).
-    if (!durationUntilDelete.equals(Duration.ZERO) || isSuperuser) {
+    if (durationUntilDelete.isLongerThan(Duration.ZERO) || isSuperuser) {
       PollMessage.OneTime deletePollMessage =
-          createDeletePollMessage(existingDomain, historyEntry, deletionTime);
+          createDeletePollMessage(existingDomain, domainHistoryKey, deletionTime);
       entitiesToSave.add(deletePollMessage);
       builder.setDeletePollMessage(deletePollMessage.createVKey());
     }
 
+    // Send a second poll message immediately if the domain is being deleted asynchronously by a
+    // registrar other than the sponsoring registrar (which will necessarily be a superuser).
+    if (durationUntilDelete.isLongerThan(Duration.ZERO)
+        && !clientId.equals(existingDomain.getPersistedCurrentSponsorClientId())) {
+      entitiesToSave.add(
+          createImmediateDeletePollMessage(existingDomain, domainHistoryKey, now, deletionTime));
+    }
+
     // Cancel any grace periods that were still active, and set the expiration time accordingly.
     DateTime newExpirationTime = existingDomain.getRegistrationExpirationTime();
     for (GracePeriod gracePeriod : existingDomain.getGracePeriods()) {
       // No cancellation is written if the grace period was not for a billable event.
       if (gracePeriod.hasBillingEvent()) {
         entitiesToSave.add(
-            BillingEvent.Cancellation.forGracePeriod(gracePeriod, historyEntry, targetId));
+            BillingEvent.Cancellation.forGracePeriod(gracePeriod, now, domainHistoryKey, targetId));
         if (gracePeriod.getOneTimeBillingEvent() != null) {
           // Take the amount of amount of registration time being refunded off the expiration time.
           // This can be either add grace periods or renew grace periods.
@@ -237,8 +247,10 @@ public final class DomainDeleteFlow implements TransactionalFlow {
     builder.setRegistrationExpirationTime(newExpirationTime);
 
     DomainBase newDomain = builder.build();
+    DomainHistory domainHistory =
+        buildDomainHistory(newDomain, registry, now, durationUntilDelete, inAddGracePeriod);
     updateForeignKeyIndexDeletionTime(newDomain);
-    handlePendingTransferOnDelete(existingDomain, newDomain, now, historyEntry);
+    handlePendingTransferOnDelete(existingDomain, newDomain, now, domainHistory);
     // Close the autorenew billing event and poll message. This may delete the poll message.
     updateAutorenewRecurrenceEndTime(existingDomain, now);
     // If there's a pending transfer, the gaining client's autorenew billing
@@ -246,15 +258,16 @@ public final class DomainDeleteFlow implements TransactionalFlow {
     // ResourceDeleteFlow since it's listed in serverApproveEntities.
     dnsQueue.addDomainRefreshTask(existingDomain.getDomainName());
 
-    entitiesToSave.add(newDomain, historyEntry);
-    EntityChanges entityChanges = flowCustomLogic.beforeSave(
-        BeforeSaveParameters.newBuilder()
-            .setExistingDomain(existingDomain)
-            .setNewDomain(newDomain)
-            .setHistoryEntry(historyEntry)
-            .setEntityChanges(EntityChanges.newBuilder().setSaves(entitiesToSave.build()).build())
-            .build());
-    persistEntityChanges(entityChanges);
+    entitiesToSave.add(newDomain, domainHistory);
+    EntityChanges entityChanges =
+        flowCustomLogic.beforeSave(
+            BeforeSaveParameters.newBuilder()
+                .setExistingDomain(existingDomain)
+                .setNewDomain(newDomain)
+                .setHistoryEntry(domainHistory)
+                .setEntityChanges(
+                    EntityChanges.newBuilder().setSaves(entitiesToSave.build()).build())
+                .build());
     BeforeResponseReturnData responseData =
         flowCustomLogic.beforeResponse(
             BeforeResponseParameters.newBuilder()
@@ -264,6 +277,7 @@ public final class DomainDeleteFlow implements TransactionalFlow {
                         : SUCCESS)
                 .setResponseExtensions(getResponseExtensions(existingDomain, now))
                 .build());
+    persistEntityChanges(entityChanges);
     return responseBuilder
         .setResultFromCode(responseData.resultCode())
         .setExtensions(responseData.responseExtensions())
@@ -284,8 +298,8 @@ public final class DomainDeleteFlow implements TransactionalFlow {
     }
   }
 
-  private HistoryEntry buildHistoryEntry(
-      DomainBase existingResource,
+  private DomainHistory buildDomainHistory(
+      DomainBase domain,
       Registry registry,
       DateTime now,
       Duration durationUntilDelete,
@@ -299,31 +313,30 @@ public final class DomainDeleteFlow implements TransactionalFlow {
               registry.getRenewGracePeriodLength()));
       ImmutableSet<DomainTransactionRecord> cancelledRecords =
           createCancelingRecords(
-              existingResource,
+              domain,
               now,
               maxGracePeriod,
               Sets.immutableEnumSet(Sets.union(ADD_FIELDS, RENEW_FIELDS)));
-      historyBuilder
-          .setDomainTransactionRecords(
-              union(
-                  cancelledRecords,
-                  DomainTransactionRecord.create(
-                      existingResource.getTld(),
-                      now.plus(durationUntilDelete),
-                      inAddGracePeriod
-                          ? TransactionReportField.DELETED_DOMAINS_GRACE
-                          : TransactionReportField.DELETED_DOMAINS_NOGRACE,
-                      1)));
+      historyBuilder.setDomainTransactionRecords(
+          union(
+              cancelledRecords,
+              DomainTransactionRecord.create(
+                  domain.getTld(),
+                  now.plus(durationUntilDelete),
+                  inAddGracePeriod
+                      ? TransactionReportField.DELETED_DOMAINS_GRACE
+                      : TransactionReportField.DELETED_DOMAINS_NOGRACE,
+                  1)));
     }
     return historyBuilder
         .setType(HistoryEntry.Type.DOMAIN_DELETE)
         .setModificationTime(now)
-        .setParent(Key.create(existingResource))
+        .setDomainContent(domain)
         .build();
   }
 
   private PollMessage.OneTime createDeletePollMessage(
-      DomainBase existingDomain, HistoryEntry historyEntry, DateTime deletionTime) {
+      DomainBase existingDomain, Key<DomainHistory> domainHistoryKey, DateTime deletionTime) {
     Optional<MetadataExtension> metadataExtension =
         eppInput.getSingleExtension(MetadataExtension.class);
     boolean hasMetadataMessage =
@@ -342,7 +355,23 @@ public final class DomainDeleteFlow implements TransactionalFlow {
             ImmutableList.of(
                 DomainPendingActionNotificationResponse.create(
                     existingDomain.getDomainName(), true, trid, deletionTime)))
-        .setParent(historyEntry)
+        .setParentKey(domainHistoryKey)
+        .build();
+  }
+
+  private PollMessage.OneTime createImmediateDeletePollMessage(
+      DomainBase existingDomain,
+      Key<DomainHistory> domainHistoryKey,
+      DateTime now,
+      DateTime deletionTime) {
+    return new PollMessage.OneTime.Builder()
+        .setClientId(existingDomain.getPersistedCurrentSponsorClientId())
+        .setEventTime(now)
+        .setParentKey(domainHistoryKey)
+        .setMsg(
+            String.format(
+                "Domain %s was deleted by registry administrator with final deletion effective: %s",
+                existingDomain.getDomainName(), deletionTime))
         .build();
   }
 

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

@@ -22,6 +22,7 @@ import static google.registry.flows.domain.DomainFlowUtils.handleFeeRequest;
 import static google.registry.flows.domain.DomainFlowUtils.loadForeignKeyedDesignatedContacts;
 import static google.registry.model.EppResourceUtils.loadByForeignKey;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
@@ -100,9 +101,11 @@ public final class DomainInfoFlow implements Flow {
     verifyOptionalAuthInfo(authInfo, domain);
     flowCustomLogic.afterValidation(
         AfterValidationParameters.newBuilder().setDomain(domain).build());
-    // Prefetch all referenced resources. Calling values() blocks until loading is done.
-    tm().loadByKeys(domain.getNameservers());
-    tm().loadByKeys(domain.getReferencedContacts());
+    // In ofy, refetch all referenced resources.
+    if (tm().isOfy()) {
+      tm().loadByKeys(domain.getNameservers());
+      tm().loadByKeys(domain.getReferencedContacts());
+    }
     // Registrars can only see a few fields on unauthorized domains.
     // This is a policy decision that is left up to us by the rfcs.
     DomainInfoData.Builder infoBuilder =
@@ -110,14 +113,16 @@ public final class DomainInfoFlow implements Flow {
             .setFullyQualifiedDomainName(domain.getDomainName())
             .setRepoId(domain.getRepoId())
             .setCurrentSponsorClientId(domain.getCurrentSponsorClientId())
-            .setRegistrant(tm().loadByKey(domain.getRegistrant()).getContactId());
+            .setRegistrant(
+                transactIfJpaTm(() -> tm().loadByKey(domain.getRegistrant())).getContactId());
     // If authInfo is non-null, then the caller is authorized to see the full information since we
     // will have already verified the authInfo is valid.
     if (clientId.equals(domain.getCurrentSponsorClientId()) || authInfo.isPresent()) {
       HostsRequest hostsRequest = ((Info) resourceCommand).getHostsRequest();
       infoBuilder
           .setStatusValues(domain.getStatusValues())
-          .setContacts(loadForeignKeyedDesignatedContacts(domain.getContacts()))
+          .setContacts(
+              transactIfJpaTm(() -> loadForeignKeyedDesignatedContacts(domain.getContacts())))
           .setNameservers(hostsRequest.requestDelegated() ? domain.loadNameserverHostNames() : null)
           .setSubordinateHosts(
               hostsRequest.requestSubordinate() ? domain.getSubordinateHosts() : null)

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

@@ -14,6 +14,7 @@
 
 package google.registry.flows.domain;
 
+import static google.registry.flows.FlowUtils.createHistoryKey;
 import static google.registry.flows.FlowUtils.persistEntityChanges;
 import static google.registry.flows.FlowUtils.validateClientIsLoggedIn;
 import static google.registry.flows.ResourceFlowUtils.loadAndVerifyExistence;
@@ -33,6 +34,7 @@ import static google.registry.util.DateTimeUtils.leapSafeAddYears;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
+import com.googlecode.objectify.Key;
 import google.registry.flows.EppException;
 import google.registry.flows.EppException.ParameterValueRangeErrorException;
 import google.registry.flows.ExtensionManager;
@@ -52,6 +54,7 @@ import google.registry.model.billing.BillingEvent.OneTime;
 import google.registry.model.billing.BillingEvent.Reason;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.DomainCommand.Renew;
+import google.registry.model.domain.DomainHistory;
 import google.registry.model.domain.DomainRenewData;
 import google.registry.model.domain.GracePeriod;
 import google.registry.model.domain.Period;
@@ -123,7 +126,7 @@ public final class DomainRenewFlow implements TransactionalFlow {
   @Inject @ClientId String clientId;
   @Inject @TargetId String targetId;
   @Inject @Superuser boolean isSuperuser;
-  @Inject HistoryEntry.Builder historyBuilder;
+  @Inject DomainHistory.Builder historyBuilder;
   @Inject EppResponse.Builder responseBuilder;
   @Inject DomainRenewFlowCustomLogic flowCustomLogic;
   @Inject DomainPricingLogic pricingLogic;
@@ -156,22 +159,23 @@ public final class DomainRenewFlow implements TransactionalFlow {
             .setNow(now)
             .setYears(years)
             .build());
-    Registry registry = Registry.get(existingDomain.getTld());
-    HistoryEntry historyEntry = buildHistoryEntry(
-        existingDomain, now, command.getPeriod(), registry.getRenewGracePeriodLength());
+    Key<DomainHistory> domainHistoryKey = createHistoryKey(existingDomain, DomainHistory.class);
+    historyBuilder.setId(domainHistoryKey.getId());
     String tld = existingDomain.getTld();
     // Bill for this explicit renew itself.
     BillingEvent.OneTime explicitRenewEvent =
-        createRenewBillingEvent(tld, feesAndCredits.getTotalCost(), years, historyEntry, now);
+        createRenewBillingEvent(tld, feesAndCredits.getTotalCost(), years, domainHistoryKey, now);
     // Create a new autorenew billing event and poll message starting at the new expiration time.
-    BillingEvent.Recurring newAutorenewEvent = newAutorenewBillingEvent(existingDomain)
-        .setEventTime(newExpirationTime)
-        .setParent(historyEntry)
-        .build();
-    PollMessage.Autorenew newAutorenewPollMessage = newAutorenewPollMessage(existingDomain)
-        .setEventTime(newExpirationTime)
-        .setParent(historyEntry)
-        .build();
+    BillingEvent.Recurring newAutorenewEvent =
+        newAutorenewBillingEvent(existingDomain)
+            .setEventTime(newExpirationTime)
+            .setParent(domainHistoryKey)
+            .build();
+    PollMessage.Autorenew newAutorenewPollMessage =
+        newAutorenewPollMessage(existingDomain)
+            .setEventTime(newExpirationTime)
+            .setParentKey(domainHistoryKey)
+            .build();
     // End the old autorenew billing event and poll message now. This may delete the poll message.
     updateAutorenewRecurrenceEndTime(existingDomain, now);
     DomainBase newDomain =
@@ -186,6 +190,10 @@ public final class DomainRenewFlow implements TransactionalFlow {
                 GracePeriod.forBillingEvent(
                     GracePeriodStatus.RENEW, existingDomain.getRepoId(), explicitRenewEvent))
             .build();
+    Registry registry = Registry.get(existingDomain.getTld());
+    DomainHistory domainHistory =
+        buildDomainHistory(
+            newDomain, now, command.getPeriod(), registry.getRenewGracePeriodLength());
     EntityChanges entityChanges =
         flowCustomLogic.beforeSave(
             BeforeSaveParameters.newBuilder()
@@ -193,19 +201,18 @@ public final class DomainRenewFlow implements TransactionalFlow {
                 .setNewDomain(newDomain)
                 .setNow(now)
                 .setYears(years)
-                .setHistoryEntry(historyEntry)
+                .setHistoryEntry(domainHistory)
                 .setEntityChanges(
                     EntityChanges.newBuilder()
                         .setSaves(
                             ImmutableSet.of(
                                 newDomain,
-                                historyEntry,
+                                domainHistory,
                                 explicitRenewEvent,
                                 newAutorenewEvent,
                                 newAutorenewPollMessage))
                         .build())
                 .build());
-    persistEntityChanges(entityChanges);
     BeforeResponseReturnData responseData =
         flowCustomLogic.beforeResponse(
             BeforeResponseParameters.newBuilder()
@@ -213,23 +220,24 @@ public final class DomainRenewFlow implements TransactionalFlow {
                 .setResData(DomainRenewData.create(targetId, newExpirationTime))
                 .setResponseExtensions(createResponseExtensions(feesAndCredits, feeRenew))
                 .build());
+    persistEntityChanges(entityChanges);
     return responseBuilder
         .setResData(responseData.resData())
         .setExtensions(responseData.responseExtensions())
         .build();
   }
 
-  private HistoryEntry buildHistoryEntry(
-      DomainBase existingDomain, DateTime now, Period period, Duration renewGracePeriod) {
+  private DomainHistory buildDomainHistory(
+      DomainBase newDomain, DateTime now, Period period, Duration renewGracePeriod) {
     return historyBuilder
         .setType(HistoryEntry.Type.DOMAIN_RENEW)
         .setPeriod(period)
         .setModificationTime(now)
-        .setParent(existingDomain)
+        .setDomainContent(newDomain)
         .setDomainTransactionRecords(
             ImmutableSet.of(
                 DomainTransactionRecord.create(
-                    existingDomain.getTld(),
+                    newDomain.getTld(),
                     now.plus(renewGracePeriod),
                     TransactionReportField.netRenewsFieldFromYears(period.getValue()),
                     1)))
@@ -255,7 +263,7 @@ public final class DomainRenewFlow implements TransactionalFlow {
   }
 
   private OneTime createRenewBillingEvent(
-      String tld, Money renewCost, int years, HistoryEntry historyEntry, DateTime now) {
+      String tld, Money renewCost, int years, Key<DomainHistory> domainHistoryKey, DateTime now) {
     return new BillingEvent.OneTime.Builder()
         .setReason(Reason.RENEW)
         .setTargetId(targetId)
@@ -264,7 +272,7 @@ public final class DomainRenewFlow implements TransactionalFlow {
         .setCost(renewCost)
         .setEventTime(now)
         .setBillingTime(now.plus(Registry.get(tld).getRenewGracePeriodLength()))
-        .setParent(historyEntry)
+        .setParent(domainHistoryKey)
         .build();
   }
 

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

@@ -14,6 +14,7 @@
 
 package google.registry.flows.domain;
 
+import static google.registry.flows.FlowUtils.createHistoryKey;
 import static google.registry.flows.FlowUtils.validateClientIsLoggedIn;
 import static google.registry.flows.ResourceFlowUtils.loadAndVerifyExistence;
 import static google.registry.flows.ResourceFlowUtils.verifyOptionalAuthInfo;
@@ -49,6 +50,7 @@ import google.registry.model.billing.BillingEvent.OneTime;
 import google.registry.model.billing.BillingEvent.Reason;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.DomainCommand.Update;
+import google.registry.model.domain.DomainHistory;
 import google.registry.model.domain.fee.BaseFee.FeeType;
 import google.registry.model.domain.fee.Fee;
 import google.registry.model.domain.fee.FeeTransformResponseExtension;
@@ -117,7 +119,7 @@ public final class DomainRestoreRequestFlow implements TransactionalFlow  {
   @Inject @ClientId String clientId;
   @Inject @TargetId String targetId;
   @Inject @Superuser boolean isSuperuser;
-  @Inject HistoryEntry.Builder historyBuilder;
+  @Inject DomainHistory.Builder historyBuilder;
   @Inject DnsQueue dnsQueue;
   @Inject EppResponse.Builder responseBuilder;
   @Inject DomainPricingLogic pricingLogic;
@@ -142,7 +144,8 @@ public final class DomainRestoreRequestFlow implements TransactionalFlow  {
     Optional<FeeUpdateCommandExtension> feeUpdate =
         eppInput.getSingleExtension(FeeUpdateCommandExtension.class);
     verifyRestoreAllowed(command, existingDomain, feeUpdate, feesAndCredits, now);
-    HistoryEntry historyEntry = buildHistoryEntry(existingDomain, now);
+    Key<DomainHistory> domainHistoryKey = createHistoryKey(existingDomain, DomainHistory.class);
+    historyBuilder.setId(domainHistoryKey.getId());
     ImmutableSet.Builder<ImmutableObject> entitiesToSave = new ImmutableSet.Builder<>();
 
     DateTime newExpirationTime =
@@ -150,29 +153,31 @@ public final class DomainRestoreRequestFlow implements TransactionalFlow  {
     // Restore the expiration time on the deleted domain, except if that's already passed, then add
     // a year and bill for it immediately, with no grace period.
     if (isExpired) {
-      entitiesToSave.add(createRenewBillingEvent(historyEntry, feesAndCredits.getRenewCost(), now));
+      entitiesToSave.add(
+          createRenewBillingEvent(domainHistoryKey, feesAndCredits.getRenewCost(), now));
     }
     // Always bill for the restore itself.
     entitiesToSave.add(
-        createRestoreBillingEvent(historyEntry, feesAndCredits.getRestoreCost(), now));
+        createRestoreBillingEvent(domainHistoryKey, feesAndCredits.getRestoreCost(), now));
 
     BillingEvent.Recurring autorenewEvent =
         newAutorenewBillingEvent(existingDomain)
             .setEventTime(newExpirationTime)
             .setRecurrenceEndTime(END_OF_TIME)
-            .setParent(historyEntry)
+            .setParent(domainHistoryKey)
             .build();
     PollMessage.Autorenew autorenewPollMessage =
         newAutorenewPollMessage(existingDomain)
             .setEventTime(newExpirationTime)
             .setAutorenewEndTime(END_OF_TIME)
-            .setParent(historyEntry)
+            .setParentKey(domainHistoryKey)
             .build();
     DomainBase newDomain =
         performRestore(
             existingDomain, newExpirationTime, autorenewEvent, autorenewPollMessage, now, clientId);
     updateForeignKeyIndexDeletionTime(newDomain);
-    entitiesToSave.add(newDomain, historyEntry, autorenewEvent, autorenewPollMessage);
+    DomainHistory domainHistory = buildDomainHistory(newDomain, now);
+    entitiesToSave.add(newDomain, domainHistory, autorenewEvent, autorenewPollMessage);
     tm().putAll(entitiesToSave.build());
     tm().delete(existingDomain.getDeletePollMessage());
     dnsQueue.addDomainRefreshTask(existingDomain.getDomainName());
@@ -181,15 +186,15 @@ public final class DomainRestoreRequestFlow implements TransactionalFlow  {
         .build();
   }
 
-  private HistoryEntry buildHistoryEntry(DomainBase existingDomain, DateTime now) {
+  private DomainHistory buildDomainHistory(DomainBase newDomain, DateTime now) {
     return historyBuilder
         .setType(HistoryEntry.Type.DOMAIN_RESTORE)
         .setModificationTime(now)
-        .setParent(Key.create(existingDomain))
+        .setDomainContent(newDomain)
         .setDomainTransactionRecords(
             ImmutableSet.of(
                 DomainTransactionRecord.create(
-                    existingDomain.getTld(), now, TransactionReportField.RESTORED_DOMAINS, 1)))
+                    newDomain.getTld(), now, TransactionReportField.RESTORED_DOMAINS, 1)))
         .build();
   }
 
@@ -242,20 +247,19 @@ public final class DomainRestoreRequestFlow implements TransactionalFlow  {
   }
 
   private OneTime createRenewBillingEvent(
-      HistoryEntry historyEntry, Money renewCost, DateTime now) {
-    return prepareBillingEvent(historyEntry, renewCost, now)
-        .setReason(Reason.RENEW)
-        .build();
+      Key<DomainHistory> domainHistoryKey, Money renewCost, DateTime now) {
+    return prepareBillingEvent(domainHistoryKey, renewCost, now).setReason(Reason.RENEW).build();
   }
 
   private BillingEvent.OneTime createRestoreBillingEvent(
-      HistoryEntry historyEntry, Money restoreCost, DateTime now) {
-    return prepareBillingEvent(historyEntry, restoreCost, now)
+      Key<DomainHistory> domainHistoryKey, Money restoreCost, DateTime now) {
+    return prepareBillingEvent(domainHistoryKey, restoreCost, now)
         .setReason(Reason.RESTORE)
         .build();
   }
 
-  private OneTime.Builder prepareBillingEvent(HistoryEntry historyEntry, Money cost, DateTime now) {
+  private OneTime.Builder prepareBillingEvent(
+      Key<DomainHistory> domainHistoryKey, Money cost, DateTime now) {
     return new BillingEvent.OneTime.Builder()
         .setTargetId(targetId)
         .setClientId(clientId)
@@ -263,7 +267,7 @@ public final class DomainRestoreRequestFlow implements TransactionalFlow  {
         .setBillingTime(now)
         .setPeriodYears(1)
         .setCost(cost)
-        .setParent(historyEntry);
+        .setParent(domainHistoryKey);
   }
 
   private static ImmutableList<FeeTransformResponseExtension> createResponseExtensions(

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

@@ -15,6 +15,7 @@
 package google.registry.flows.domain;
 
 import static com.google.common.collect.Iterables.getOnlyElement;
+import static google.registry.flows.FlowUtils.createHistoryKey;
 import static google.registry.flows.FlowUtils.validateClientIsLoggedIn;
 import static google.registry.flows.ResourceFlowUtils.computeExDateForApprovalTime;
 import static google.registry.flows.ResourceFlowUtils.loadAndVerifyExistence;
@@ -27,13 +28,13 @@ import static google.registry.flows.domain.DomainFlowUtils.updateAutorenewRecurr
 import static google.registry.flows.domain.DomainTransferUtils.createGainingTransferPollMessage;
 import static google.registry.flows.domain.DomainTransferUtils.createTransferResponse;
 import static google.registry.model.ResourceTransferUtils.approvePendingTransfer;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.reporting.DomainTransactionRecord.TransactionReportField.TRANSFER_SUCCESSFUL;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.pricing.PricingEngineProxy.getDomainRenewCost;
 import static google.registry.util.CollectionUtils.union;
 import static google.registry.util.DateTimeUtils.END_OF_TIME;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
 import google.registry.flows.EppException;
@@ -48,6 +49,7 @@ import google.registry.model.billing.BillingEvent;
 import google.registry.model.billing.BillingEvent.Flag;
 import google.registry.model.billing.BillingEvent.Reason;
 import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.DomainHistory;
 import google.registry.model.domain.GracePeriod;
 import google.registry.model.domain.metadata.MetadataExtension;
 import google.registry.model.domain.rgp.GracePeriodStatus;
@@ -90,7 +92,7 @@ public final class DomainTransferApproveFlow implements TransactionalFlow {
   @Inject @ClientId String clientId;
   @Inject @TargetId String targetId;
   @Inject @Superuser boolean isSuperuser;
-  @Inject HistoryEntry.Builder historyBuilder;
+  @Inject DomainHistory.Builder historyBuilder;
   @Inject EppResponse.Builder responseBuilder;
   @Inject DomainTransferApproveFlow() {}
 
@@ -114,10 +116,10 @@ public final class DomainTransferApproveFlow implements TransactionalFlow {
     }
     DomainTransferData transferData = existingDomain.getTransferData();
     String gainingClientId = transferData.getGainingClientId();
-    Registry registry = Registry.get(existingDomain.getTld());
-    HistoryEntry historyEntry = buildHistoryEntry(existingDomain, registry, now, gainingClientId);
     // Create a transfer billing event for 1 year, unless the superuser extension was used to set
     // the transfer period to zero. There is not a transfer cost if the transfer period is zero.
+    Key<DomainHistory> domainHistoryKey = createHistoryKey(existingDomain, DomainHistory.class);
+    historyBuilder.setId(domainHistoryKey.getId());
     Optional<BillingEvent.OneTime> billingEvent =
         (transferData.getTransferPeriod().getValue() == 0)
             ? Optional.empty()
@@ -130,8 +132,9 @@ public final class DomainTransferApproveFlow implements TransactionalFlow {
                     .setCost(getDomainRenewCost(targetId, transferData.getTransferRequestTime(), 1))
                     .setEventTime(now)
                     .setBillingTime(now.plus(Registry.get(tld).getTransferGracePeriodLength()))
-                    .setParent(historyEntry)
+                    .setParent(domainHistoryKey)
                     .build());
+    ImmutableList.Builder<ImmutableObject> entitiesToSave = new ImmutableList.Builder<>();
     // If we are within an autorenew grace period, cancel the autorenew billing event and don't
     // increase the registration time, since the transfer subsumes the autorenew's extra year.
     GracePeriod autorenewGrace =
@@ -143,8 +146,9 @@ public final class DomainTransferApproveFlow implements TransactionalFlow {
       // then the gaining registrar is not charged for the one year renewal and the losing registrar
       // still needs to be charged for the auto-renew.
       if (billingEvent.isPresent()) {
-        ofy().save().entity(
-            BillingEvent.Cancellation.forGracePeriod(autorenewGrace, historyEntry, targetId));
+        entitiesToSave.add(
+            BillingEvent.Cancellation.forGracePeriod(
+                autorenewGrace, now, domainHistoryKey, targetId));
       }
     }
     // Close the old autorenew event and poll message at the transfer time (aka now). This may end
@@ -153,24 +157,26 @@ public final class DomainTransferApproveFlow implements TransactionalFlow {
     DateTime newExpirationTime =
         computeExDateForApprovalTime(existingDomain, now, transferData.getTransferPeriod());
     // Create a new autorenew event starting at the expiration time.
-    BillingEvent.Recurring autorenewEvent = new BillingEvent.Recurring.Builder()
-        .setReason(Reason.RENEW)
-        .setFlags(ImmutableSet.of(Flag.AUTO_RENEW))
-        .setTargetId(targetId)
-        .setClientId(gainingClientId)
-        .setEventTime(newExpirationTime)
-        .setRecurrenceEndTime(END_OF_TIME)
-        .setParent(historyEntry)
-        .build();
+    BillingEvent.Recurring autorenewEvent =
+        new BillingEvent.Recurring.Builder()
+            .setReason(Reason.RENEW)
+            .setFlags(ImmutableSet.of(Flag.AUTO_RENEW))
+            .setTargetId(targetId)
+            .setClientId(gainingClientId)
+            .setEventTime(newExpirationTime)
+            .setRecurrenceEndTime(END_OF_TIME)
+            .setParent(domainHistoryKey)
+            .build();
     // Create a new autorenew poll message.
-    PollMessage.Autorenew gainingClientAutorenewPollMessage = new PollMessage.Autorenew.Builder()
-        .setTargetId(targetId)
-        .setClientId(gainingClientId)
-        .setEventTime(newExpirationTime)
-        .setAutorenewEndTime(END_OF_TIME)
-        .setMsg("Domain was auto-renewed.")
-        .setParent(historyEntry)
-        .build();
+    PollMessage.Autorenew gainingClientAutorenewPollMessage =
+        new PollMessage.Autorenew.Builder()
+            .setTargetId(targetId)
+            .setClientId(gainingClientId)
+            .setEventTime(newExpirationTime)
+            .setAutorenewEndTime(END_OF_TIME)
+            .setMsg("Domain was auto-renewed.")
+            .setParentKey(domainHistoryKey)
+            .build();
     // Construct the post-transfer domain.
     DomainBase partiallyApprovedDomain =
         approvePendingTransfer(existingDomain, TransferStatus.CLIENT_APPROVED, now);
@@ -190,31 +196,32 @@ public final class DomainTransferApproveFlow implements TransactionalFlow {
             .setAutorenewPollMessage(gainingClientAutorenewPollMessage.createVKey())
             // Remove all the old grace periods and add a new one for the transfer.
             .setGracePeriods(
-                billingEvent.isPresent()
-                    ? ImmutableSet.of(
-                        GracePeriod.forBillingEvent(
-                            GracePeriodStatus.TRANSFER,
-                            existingDomain.getRepoId(),
-                            billingEvent.get()))
-                    : ImmutableSet.of())
+                billingEvent
+                    .map(
+                        oneTime ->
+                            ImmutableSet.of(
+                                GracePeriod.forBillingEvent(
+                                    GracePeriodStatus.TRANSFER,
+                                    existingDomain.getRepoId(),
+                                    oneTime)))
+                    .orElseGet(ImmutableSet::of))
             .setLastEppUpdateTime(now)
             .setLastEppUpdateClientId(clientId)
             .build();
+    Registry registry = Registry.get(existingDomain.getTld());
+    DomainHistory domainHistory = buildDomainHistory(newDomain, registry, now, gainingClientId);
     // Create a poll message for the gaining client.
-    PollMessage gainingClientPollMessage = createGainingTransferPollMessage(
-        targetId,
-        newDomain.getTransferData(),
-        newExpirationTime,
-        historyEntry);
-    ImmutableSet.Builder<ImmutableObject> entitiesToSave = new ImmutableSet.Builder<>();
+    PollMessage gainingClientPollMessage =
+        createGainingTransferPollMessage(
+            targetId, newDomain.getTransferData(), newExpirationTime, now, domainHistoryKey);
+    billingEvent.ifPresent(entitiesToSave::add);
     entitiesToSave.add(
-        newDomain,
-        historyEntry,
         autorenewEvent,
         gainingClientPollMessage,
-        gainingClientAutorenewPollMessage);
-    billingEvent.ifPresent(entitiesToSave::add);
-    ofy().save().entities(entitiesToSave.build());
+        gainingClientAutorenewPollMessage,
+        newDomain,
+        domainHistory);
+    tm().putAll(entitiesToSave.build());
     // Delete the billing event and poll messages that were written in case the transfer would have
     // been implicitly server approved.
     tm().delete(existingDomain.getTransferData().getServerApproveEntities());
@@ -224,11 +231,11 @@ public final class DomainTransferApproveFlow implements TransactionalFlow {
         .build();
   }
 
-  private HistoryEntry buildHistoryEntry(
-      DomainBase existingDomain, Registry registry, DateTime now, String gainingClientId) {
+  private DomainHistory buildDomainHistory(
+      DomainBase newDomain, Registry registry, DateTime now, String gainingClientId) {
     ImmutableSet<DomainTransactionRecord> cancelingRecords =
         createCancelingRecords(
-            existingDomain,
+            newDomain,
             now,
             registry.getAutomaticTransferLength().plus(registry.getTransferGracePeriodLength()),
             ImmutableSet.of(TRANSFER_SUCCESSFUL));
@@ -236,15 +243,15 @@ public final class DomainTransferApproveFlow implements TransactionalFlow {
         .setType(HistoryEntry.Type.DOMAIN_TRANSFER_APPROVE)
         .setModificationTime(now)
         .setOtherClientId(gainingClientId)
-        .setParent(Key.create(existingDomain))
+        .setDomainContent(newDomain)
         .setDomainTransactionRecords(
             union(
                 cancelingRecords,
                 DomainTransactionRecord.create(
-                        existingDomain.getTld(),
-                        now.plus(registry.getTransferGracePeriodLength()),
-                        TRANSFER_SUCCESSFUL,
-                        1)))
+                    newDomain.getTld(),
+                    now.plus(registry.getTransferGracePeriodLength()),
+                    TRANSFER_SUCCESSFUL,
+                    1)))
         .build();
   }
 }

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

@@ -14,6 +14,7 @@
 
 package google.registry.flows.domain;
 
+import static google.registry.flows.FlowUtils.createHistoryKey;
 import static google.registry.flows.FlowUtils.validateClientIsLoggedIn;
 import static google.registry.flows.ResourceFlowUtils.loadAndVerifyExistence;
 import static google.registry.flows.ResourceFlowUtils.verifyHasPendingTransfer;
@@ -25,7 +26,6 @@ import static google.registry.flows.domain.DomainFlowUtils.updateAutorenewRecurr
 import static google.registry.flows.domain.DomainTransferUtils.createLosingTransferPollMessage;
 import static google.registry.flows.domain.DomainTransferUtils.createTransferResponse;
 import static google.registry.model.ResourceTransferUtils.denyPendingTransfer;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.reporting.DomainTransactionRecord.TransactionReportField.TRANSFER_SUCCESSFUL;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.DateTimeUtils.END_OF_TIME;
@@ -39,8 +39,8 @@ import google.registry.flows.FlowModule.Superuser;
 import google.registry.flows.FlowModule.TargetId;
 import google.registry.flows.TransactionalFlow;
 import google.registry.flows.annotations.ReportingSpec;
-import google.registry.model.ImmutableObject;
 import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.DomainHistory;
 import google.registry.model.domain.metadata.MetadataExtension;
 import google.registry.model.eppcommon.AuthInfo;
 import google.registry.model.eppoutput.EppResponse;
@@ -79,7 +79,7 @@ public final class DomainTransferCancelFlow implements TransactionalFlow {
   @Inject @ClientId String clientId;
   @Inject @TargetId String targetId;
   @Inject @Superuser boolean isSuperuser;
-  @Inject HistoryEntry.Builder historyBuilder;
+  @Inject DomainHistory.Builder historyBuilder;
   @Inject EppResponse.Builder responseBuilder;
   @Inject DomainTransferCancelFlow() {}
 
@@ -97,14 +97,20 @@ public final class DomainTransferCancelFlow implements TransactionalFlow {
       checkAllowedAccessToTld(clientId, existingDomain.getTld());
     }
     Registry registry = Registry.get(existingDomain.getTld());
-    HistoryEntry historyEntry = buildHistoryEntry(existingDomain, registry, now);
+
+    Key<DomainHistory> domainHistoryKey = createHistoryKey(existingDomain, DomainHistory.class);
+    historyBuilder
+        .setId(domainHistoryKey.getId())
+        .setOtherClientId(existingDomain.getTransferData().getLosingClientId());
+
     DomainBase newDomain =
         denyPendingTransfer(existingDomain, TransferStatus.CLIENT_CANCELLED, now, clientId);
-    ofy().save().<ImmutableObject>entities(
-        newDomain,
-        historyEntry,
-        createLosingTransferPollMessage(
-            targetId, newDomain.getTransferData(), null, historyEntry));
+    DomainHistory domainHistory = buildDomainHistory(newDomain, registry, now);
+    tm().putAll(
+            newDomain,
+            domainHistory,
+            createLosingTransferPollMessage(
+                targetId, newDomain.getTransferData(), null, domainHistoryKey));
     // Reopen the autorenew event and poll message that we closed for the implicit transfer. This
     // may recreate the autorenew poll message if it was deleted when the transfer request was made.
     updateAutorenewRecurrenceEndTime(existingDomain, END_OF_TIME);
@@ -116,19 +122,17 @@ public final class DomainTransferCancelFlow implements TransactionalFlow {
         .build();
   }
 
-  private HistoryEntry buildHistoryEntry(
-      DomainBase existingDomain, Registry registry, DateTime now) {
+  private DomainHistory buildDomainHistory(DomainBase newDomain, Registry registry, DateTime now) {
     ImmutableSet<DomainTransactionRecord> cancelingRecords =
         createCancelingRecords(
-            existingDomain,
+            newDomain,
             now,
             registry.getAutomaticTransferLength().plus(registry.getTransferGracePeriodLength()),
             ImmutableSet.of(TRANSFER_SUCCESSFUL));
     return historyBuilder
         .setType(HistoryEntry.Type.DOMAIN_TRANSFER_CANCEL)
-        .setOtherClientId(existingDomain.getTransferData().getLosingClientId())
         .setModificationTime(now)
-        .setParent(Key.create(existingDomain))
+        .setDomainContent(newDomain)
         .setDomainTransactionRecords(cancelingRecords)
         .build();
   }

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

@@ -14,6 +14,7 @@
 
 package google.registry.flows.domain;
 
+import static google.registry.flows.FlowUtils.createHistoryKey;
 import static google.registry.flows.FlowUtils.validateClientIsLoggedIn;
 import static google.registry.flows.ResourceFlowUtils.loadAndVerifyExistence;
 import static google.registry.flows.ResourceFlowUtils.verifyHasPendingTransfer;
@@ -41,6 +42,7 @@ import google.registry.flows.FlowModule.TargetId;
 import google.registry.flows.TransactionalFlow;
 import google.registry.flows.annotations.ReportingSpec;
 import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.DomainHistory;
 import google.registry.model.domain.metadata.MetadataExtension;
 import google.registry.model.eppcommon.AuthInfo;
 import google.registry.model.eppoutput.EppResponse;
@@ -79,7 +81,7 @@ public final class DomainTransferRejectFlow implements TransactionalFlow {
   @Inject @ClientId String clientId;
   @Inject @TargetId String targetId;
   @Inject @Superuser boolean isSuperuser;
-  @Inject HistoryEntry.Builder historyBuilder;
+  @Inject DomainHistory.Builder historyBuilder;
   @Inject EppResponse.Builder responseBuilder;
   @Inject DomainTransferRejectFlow() {}
 
@@ -91,7 +93,11 @@ public final class DomainTransferRejectFlow implements TransactionalFlow {
     DateTime now = tm().getTransactionTime();
     DomainBase existingDomain = loadAndVerifyExistence(DomainBase.class, targetId, now);
     Registry registry = Registry.get(existingDomain.getTld());
-    HistoryEntry historyEntry = buildHistoryEntry(existingDomain, registry, now);
+    Key<DomainHistory> domainHistoryKey = createHistoryKey(existingDomain, DomainHistory.class);
+    historyBuilder
+        .setId(domainHistoryKey.getId())
+        .setOtherClientId(existingDomain.getTransferData().getGainingClientId());
+
     verifyOptionalAuthInfo(authInfo, existingDomain);
     verifyHasPendingTransfer(existingDomain);
     verifyResourceOwnership(clientId, existingDomain);
@@ -100,11 +106,12 @@ public final class DomainTransferRejectFlow implements TransactionalFlow {
     }
     DomainBase newDomain =
         denyPendingTransfer(existingDomain, TransferStatus.CLIENT_REJECTED, now, clientId);
+    DomainHistory domainHistory = buildDomainHistory(newDomain, registry, now);
     tm().putAll(
             newDomain,
-            historyEntry,
+            domainHistory,
             createGainingTransferPollMessage(
-                targetId, newDomain.getTransferData(), null, historyEntry));
+                targetId, newDomain.getTransferData(), null, now, domainHistoryKey));
     // Reopen the autorenew event and poll message that we closed for the implicit transfer. This
     // may end up recreating the poll message if it was deleted upon the transfer request.
     updateAutorenewRecurrenceEndTime(existingDomain, END_OF_TIME);
@@ -116,24 +123,21 @@ public final class DomainTransferRejectFlow implements TransactionalFlow {
         .build();
   }
 
-  private HistoryEntry buildHistoryEntry(
-      DomainBase existingDomain, Registry registry, DateTime now) {
+  private DomainHistory buildDomainHistory(DomainBase newDomain, Registry registry, DateTime now) {
     ImmutableSet<DomainTransactionRecord> cancelingRecords =
         createCancelingRecords(
-            existingDomain,
+            newDomain,
             now,
             registry.getAutomaticTransferLength().plus(registry.getTransferGracePeriodLength()),
             ImmutableSet.of(TRANSFER_SUCCESSFUL));
     return historyBuilder
         .setType(HistoryEntry.Type.DOMAIN_TRANSFER_REJECT)
         .setModificationTime(now)
-        .setOtherClientId(existingDomain.getTransferData().getGainingClientId())
-        .setParent(Key.create(existingDomain))
         .setDomainTransactionRecords(
             union(
                 cancelingRecords,
-                DomainTransactionRecord.create(
-                    existingDomain.getTld(), now, TRANSFER_NACKED, 1)))
+                DomainTransactionRecord.create(newDomain.getTld(), now, TRANSFER_NACKED, 1)))
+        .setDomainContent(newDomain)
         .build();
   }
 }

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

@@ -14,6 +14,7 @@
 
 package google.registry.flows.domain;
 
+import static google.registry.flows.FlowUtils.createHistoryKey;
 import static google.registry.flows.FlowUtils.validateClientIsLoggedIn;
 import static google.registry.flows.ResourceFlowUtils.computeExDateForApprovalTime;
 import static google.registry.flows.ResourceFlowUtils.loadAndVerifyExistence;
@@ -51,6 +52,7 @@ import google.registry.flows.exceptions.TransferPeriodMustBeOneYearException;
 import google.registry.flows.exceptions.TransferPeriodZeroAndFeeTransferExtensionException;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.DomainCommand.Transfer;
+import google.registry.model.domain.DomainHistory;
 import google.registry.model.domain.Period;
 import google.registry.model.domain.fee.FeeTransferCommandExtension;
 import google.registry.model.domain.fee.FeeTransformResponseExtension;
@@ -126,7 +128,7 @@ public final class DomainTransferRequestFlow implements TransactionalFlow {
   @Inject @ClientId String gainingClientId;
   @Inject @TargetId String targetId;
   @Inject @Superuser boolean isSuperuser;
-  @Inject HistoryEntry.Builder historyBuilder;
+  @Inject DomainHistory.Builder historyBuilder;
   @Inject Trid trid;
   @Inject AsyncTaskEnqueuer asyncTaskEnqueuer;
   @Inject EppResponse.Builder responseBuilder;
@@ -169,7 +171,10 @@ public final class DomainTransferRequestFlow implements TransactionalFlow {
     if (feesAndCredits.isPresent()) {
       validateFeeChallenge(targetId, now, feeTransfer, feesAndCredits.get());
     }
-    HistoryEntry historyEntry = buildHistoryEntry(existingDomain, registry, now, period);
+    Key<DomainHistory> domainHistoryKey = createHistoryKey(existingDomain, DomainHistory.class);
+    historyBuilder
+        .setId(domainHistoryKey.getId())
+        .setOtherClientId(existingDomain.getCurrentSponsorClientId());
     DateTime automaticTransferTime =
         superuserExtension.isPresent()
             ? now.plusDays(superuserExtension.get().getAutomaticTransferLength())
@@ -190,7 +195,7 @@ public final class DomainTransferRequestFlow implements TransactionalFlow {
         createTransferServerApproveEntities(
             automaticTransferTime,
             serverApproveNewExpirationTime,
-            historyEntry,
+            domainHistoryKey,
             existingDomain,
             trid,
             gainingClientId,
@@ -209,9 +214,12 @@ public final class DomainTransferRequestFlow implements TransactionalFlow {
             serverApproveEntities,
             period);
     // Create a poll message to notify the losing registrar that a transfer was requested.
-    PollMessage requestPollMessage = createLosingTransferPollMessage(
-        targetId, pendingTransferData, serverApproveNewExpirationTime, historyEntry)
-            .asBuilder().setEventTime(now).build();
+    PollMessage requestPollMessage =
+        createLosingTransferPollMessage(
+                targetId, pendingTransferData, serverApproveNewExpirationTime, domainHistoryKey)
+            .asBuilder()
+            .setEventTime(now)
+            .build();
     // End the old autorenew event and poll message at the implicit transfer time. This may delete
     // the poll message if it has no events left. Note that if the automatic transfer succeeds, then
     // cloneProjectedAtTime() will replace these old autorenew entities with the server approve ones
@@ -225,10 +233,12 @@ public final class DomainTransferRequestFlow implements TransactionalFlow {
             .setLastEppUpdateTime(now)
             .setLastEppUpdateClientId(gainingClientId)
             .build();
+    DomainHistory domainHistory = buildDomainHistory(newDomain, registry, now, period);
+
     asyncTaskEnqueuer.enqueueAsyncResave(newDomain, now, automaticTransferTime);
     tm().putAll(
             new ImmutableSet.Builder<>()
-                .add(newDomain, historyEntry, requestPollMessage)
+                .add(newDomain, domainHistory, requestPollMessage)
                 .addAll(serverApproveEntities)
                 .build());
     return responseBuilder
@@ -302,14 +312,13 @@ public final class DomainTransferRequestFlow implements TransactionalFlow {
     }
   }
 
-  private HistoryEntry buildHistoryEntry(
-      DomainBase existingDomain, Registry registry, DateTime now, Period period) {
+  private DomainHistory buildDomainHistory(
+      DomainBase newDomain, Registry registry, DateTime now, Period period) {
     return historyBuilder
         .setType(HistoryEntry.Type.DOMAIN_TRANSFER_REQUEST)
-        .setOtherClientId(existingDomain.getCurrentSponsorClientId())
         .setPeriod(period)
         .setModificationTime(now)
-        .setParent(Key.create(existingDomain))
+        .setDomainContent(newDomain)
         .setDomainTransactionRecords(
             ImmutableSet.of(
                 DomainTransactionRecord.create(

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

@@ -20,10 +20,12 @@ import static google.registry.util.DateTimeUtils.END_OF_TIME;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
+import com.googlecode.objectify.Key;
 import google.registry.model.billing.BillingEvent;
 import google.registry.model.billing.BillingEvent.Flag;
 import google.registry.model.billing.BillingEvent.Reason;
 import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.DomainHistory;
 import google.registry.model.domain.GracePeriod;
 import google.registry.model.domain.Period;
 import google.registry.model.domain.rgp.GracePeriodStatus;
@@ -31,7 +33,6 @@ import google.registry.model.eppcommon.Trid;
 import google.registry.model.poll.PendingActionNotificationResponse.DomainPendingActionNotificationResponse;
 import google.registry.model.poll.PollMessage;
 import google.registry.model.registry.Registry;
-import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.transfer.DomainTransferData;
 import google.registry.model.transfer.TransferData;
 import google.registry.model.transfer.TransferData.TransferServerApproveEntity;
@@ -90,20 +91,21 @@ public final class DomainTransferUtils {
    * Returns a set of entities created speculatively in anticipation of a server approval.
    *
    * <p>This set consists of:
+   *
    * <ul>
-   *    <li>The one-time billing event charging the gaining registrar for the transfer
-   *    <li>A cancellation of an autorenew charge for the losing registrar, if the autorenew grace
-   *        period will apply at transfer time
-   *    <li>A new post-transfer autorenew billing event for the domain (and gaining registrar)
-   *    <li>A new post-transfer autorenew poll message for the domain (and gaining registrar)
-   *    <li>A poll message for the gaining registrar
-   *    <li>A poll message for the losing registrar
+   *   <li>The one-time billing event charging the gaining registrar for the transfer
+   *   <li>A cancellation of an autorenew charge for the losing registrar, if the autorenew grace
+   *       period will apply at transfer time
+   *   <li>A new post-transfer autorenew billing event for the domain (and gaining registrar)
+   *   <li>A new post-transfer autorenew poll message for the domain (and gaining registrar)
+   *   <li>A poll message for the gaining registrar
+   *   <li>A poll message for the losing registrar
    * </ul>
    */
   public static ImmutableSet<TransferServerApproveEntity> createTransferServerApproveEntities(
       DateTime automaticTransferTime,
       DateTime serverApproveNewExpirationTime,
-      HistoryEntry historyEntry,
+      Key<DomainHistory> domainHistoryKey,
       DomainBase existingDomain,
       Trid trid,
       String gainingClientId,
@@ -123,32 +125,39 @@ public final class DomainTransferUtils {
             .build();
     Registry registry = Registry.get(existingDomain.getTld());
     ImmutableSet.Builder<TransferServerApproveEntity> builder = new ImmutableSet.Builder<>();
-    if (transferCost.isPresent()) {
-      builder.add(
-          createTransferBillingEvent(
-              automaticTransferTime,
-              historyEntry,
-              targetId,
-              gainingClientId,
-              registry,
-              transferCost.get()));
-    }
+    transferCost.ifPresent(
+        cost ->
+            builder.add(
+                createTransferBillingEvent(
+                    automaticTransferTime,
+                    domainHistoryKey,
+                    targetId,
+                    gainingClientId,
+                    registry,
+                    cost)));
     createOptionalAutorenewCancellation(
-            automaticTransferTime, historyEntry, targetId, existingDomain, transferCost)
+            automaticTransferTime, now, domainHistoryKey, targetId, existingDomain, transferCost)
         .ifPresent(builder::add);
     return builder
         .add(
             createGainingClientAutorenewEvent(
-                serverApproveNewExpirationTime, historyEntry, targetId, gainingClientId))
+                serverApproveNewExpirationTime, domainHistoryKey, targetId, gainingClientId))
         .add(
             createGainingClientAutorenewPollMessage(
-                serverApproveNewExpirationTime, historyEntry, targetId, gainingClientId))
+                serverApproveNewExpirationTime, domainHistoryKey, targetId, gainingClientId))
         .add(
             createGainingTransferPollMessage(
-                targetId, serverApproveTransferData, serverApproveNewExpirationTime, historyEntry))
+                targetId,
+                serverApproveTransferData,
+                serverApproveNewExpirationTime,
+                now,
+                domainHistoryKey))
         .add(
             createLosingTransferPollMessage(
-                targetId, serverApproveTransferData, serverApproveNewExpirationTime, historyEntry))
+                targetId,
+                serverApproveTransferData,
+                serverApproveNewExpirationTime,
+                domainHistoryKey))
         .build();
   }
 
@@ -157,19 +166,21 @@ public final class DomainTransferUtils {
       String targetId,
       TransferData transferData,
       @Nullable DateTime extendedRegistrationExpirationTime,
-      HistoryEntry historyEntry) {
+      DateTime now,
+      Key<DomainHistory> domainHistoryKey) {
     return new PollMessage.OneTime.Builder()
         .setClientId(transferData.getGainingClientId())
         .setEventTime(transferData.getPendingTransferExpirationTime())
         .setMsg(transferData.getTransferStatus().getMessage())
-        .setResponseData(ImmutableList.of(
-            createTransferResponse(targetId, transferData, extendedRegistrationExpirationTime),
-            DomainPendingActionNotificationResponse.create(
-                  targetId,
-                  transferData.getTransferStatus().isApproved(),
-                  transferData.getTransferRequestTrid(),
-                  historyEntry.getModificationTime())))
-        .setParent(historyEntry)
+        .setResponseData(
+            ImmutableList.of(
+                createTransferResponse(targetId, transferData, extendedRegistrationExpirationTime),
+                DomainPendingActionNotificationResponse.create(
+                    targetId,
+                    transferData.getTransferStatus().isApproved(),
+                    transferData.getTransferRequestTrid(),
+                    now)))
+        .setParentKey(domainHistoryKey)
         .build();
   }
 
@@ -178,14 +189,15 @@ public final class DomainTransferUtils {
       String targetId,
       TransferData transferData,
       @Nullable DateTime extendedRegistrationExpirationTime,
-      HistoryEntry historyEntry) {
+      Key<DomainHistory> domainHistoryKey) {
     return new PollMessage.OneTime.Builder()
         .setClientId(transferData.getLosingClientId())
         .setEventTime(transferData.getPendingTransferExpirationTime())
         .setMsg(transferData.getTransferStatus().getMessage())
-        .setResponseData(ImmutableList.of(
-            createTransferResponse(targetId, transferData, extendedRegistrationExpirationTime)))
-        .setParent(historyEntry)
+        .setResponseData(
+            ImmutableList.of(
+                createTransferResponse(targetId, transferData, extendedRegistrationExpirationTime)))
+        .setParentKey(domainHistoryKey)
         .build();
   }
 
@@ -207,7 +219,7 @@ public final class DomainTransferUtils {
 
   private static PollMessage.Autorenew createGainingClientAutorenewPollMessage(
       DateTime serverApproveNewExpirationTime,
-      HistoryEntry historyEntry,
+      Key<DomainHistory> domainHistoryKey,
       String targetId,
       String gainingClientId) {
     return new PollMessage.Autorenew.Builder()
@@ -216,13 +228,13 @@ public final class DomainTransferUtils {
         .setEventTime(serverApproveNewExpirationTime)
         .setAutorenewEndTime(END_OF_TIME)
         .setMsg("Domain was auto-renewed.")
-        .setParent(historyEntry)
+        .setParentKey(domainHistoryKey)
         .build();
   }
 
   private static BillingEvent.Recurring createGainingClientAutorenewEvent(
       DateTime serverApproveNewExpirationTime,
-      HistoryEntry historyEntry,
+      Key<DomainHistory> domainHistoryKey,
       String targetId,
       String gainingClientId) {
     return new BillingEvent.Recurring.Builder()
@@ -232,7 +244,7 @@ public final class DomainTransferUtils {
         .setClientId(gainingClientId)
         .setEventTime(serverApproveNewExpirationTime)
         .setRecurrenceEndTime(END_OF_TIME)
-        .setParent(historyEntry)
+        .setParent(domainHistoryKey)
         .build();
   }
 
@@ -254,7 +266,8 @@ public final class DomainTransferUtils {
    */
   private static Optional<BillingEvent.Cancellation> createOptionalAutorenewCancellation(
       DateTime automaticTransferTime,
-      HistoryEntry historyEntry,
+      DateTime now,
+      Key<DomainHistory> domainHistoryKey,
       String targetId,
       DomainBase existingDomain,
       Optional<Money> transferCost) {
@@ -265,7 +278,8 @@ public final class DomainTransferUtils {
             domainAtTransferTime.getGracePeriodsOfType(GracePeriodStatus.AUTO_RENEW), null);
     if (autorenewGracePeriod != null && transferCost.isPresent()) {
       return Optional.of(
-          BillingEvent.Cancellation.forGracePeriod(autorenewGracePeriod, historyEntry, targetId)
+          BillingEvent.Cancellation.forGracePeriod(
+                  autorenewGracePeriod, now, domainHistoryKey, targetId)
               .asBuilder()
               .setEventTime(automaticTransferTime)
               .build());
@@ -275,7 +289,7 @@ public final class DomainTransferUtils {
 
   private static BillingEvent.OneTime createTransferBillingEvent(
       DateTime automaticTransferTime,
-      HistoryEntry historyEntry,
+      Key<DomainHistory> domainHistoryKey,
       String targetId,
       String gainingClientId,
       Registry registry,
@@ -288,7 +302,7 @@ public final class DomainTransferUtils {
         .setPeriodYears(1)
         .setEventTime(automaticTransferTime)
         .setBillingTime(automaticTransferTime.plus(registry.getTransferGracePeriodLength()))
-        .setParent(historyEntry)
+        .setParent(domainHistoryKey)
         .build();
   }
 

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

@@ -43,7 +43,6 @@ import static google.registry.persistence.transaction.TransactionManagerFactory.
 
 import com.google.common.collect.ImmutableSet;
 import com.google.common.net.InternetDomainName;
-import com.googlecode.objectify.Key;
 import google.registry.dns.DnsQueue;
 import google.registry.flows.EppException;
 import google.registry.flows.ExtensionManager;
@@ -65,6 +64,7 @@ import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.DomainCommand.Update;
 import google.registry.model.domain.DomainCommand.Update.AddRemove;
 import google.registry.model.domain.DomainCommand.Update.Change;
+import google.registry.model.domain.DomainHistory;
 import google.registry.model.domain.fee.FeeUpdateCommandExtension;
 import google.registry.model.domain.metadata.MetadataExtension;
 import google.registry.model.domain.secdns.DelegationSignerData;
@@ -142,7 +142,7 @@ public final class DomainUpdateFlow implements TransactionalFlow {
   @Inject @ClientId String clientId;
   @Inject @TargetId String targetId;
   @Inject @Superuser boolean isSuperuser;
-  @Inject HistoryEntry.Builder historyBuilder;
+  @Inject DomainHistory.Builder historyBuilder;
   @Inject DnsQueue dnsQueue;
   @Inject EppResponse.Builder responseBuilder;
   @Inject DomainUpdateFlowCustomLogic flowCustomLogic;
@@ -165,19 +165,19 @@ public final class DomainUpdateFlow implements TransactionalFlow {
     verifyUpdateAllowed(command, existingDomain, now);
     flowCustomLogic.afterValidation(
         AfterValidationParameters.newBuilder().setExistingDomain(existingDomain).build());
-    HistoryEntry historyEntry = buildHistoryEntry(existingDomain, now);
     DomainBase newDomain = performUpdate(command, existingDomain, now);
+    DomainHistory domainHistory = buildDomainHistory(newDomain, now);
     validateNewState(newDomain);
     dnsQueue.addDomainRefreshTask(targetId);
     ImmutableSet.Builder<ImmutableObject> entitiesToSave = new ImmutableSet.Builder<>();
-    entitiesToSave.add(newDomain, historyEntry);
+    entitiesToSave.add(newDomain, domainHistory);
     Optional<BillingEvent.OneTime> statusUpdateBillingEvent =
-        createBillingEventForStatusUpdates(existingDomain, newDomain, historyEntry, now);
+        createBillingEventForStatusUpdates(existingDomain, newDomain, domainHistory, now);
     statusUpdateBillingEvent.ifPresent(entitiesToSave::add);
     EntityChanges entityChanges =
         flowCustomLogic.beforeSave(
             BeforeSaveParameters.newBuilder()
-                .setHistoryEntry(historyEntry)
+                .setHistoryEntry(domainHistory)
                 .setNewDomain(newDomain)
                 .setExistingDomain(existingDomain)
                 .setEntityChanges(
@@ -217,11 +217,11 @@ public final class DomainUpdateFlow implements TransactionalFlow {
         tld, add.getNameserverFullyQualifiedHostNames());
   }
 
-  private HistoryEntry buildHistoryEntry(DomainBase existingDomain, DateTime now) {
+  private DomainHistory buildDomainHistory(DomainBase newDomain, DateTime now) {
     return historyBuilder
         .setType(HistoryEntry.Type.DOMAIN_UPDATE)
         .setModificationTime(now)
-        .setParent(Key.create(existingDomain))
+        .setDomainContent(newDomain)
         .build();
   }
 

core/src/main/java/google/registry/flows/domain/token/AllocationTokenFlowUtils.java

@@ -16,6 +16,7 @@ package google.registry.flows.domain.token;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 
 import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
@@ -153,7 +154,7 @@ public class AllocationTokenFlowUtils {
       throw new InvalidAllocationTokenException();
     }
     Optional<AllocationToken> maybeTokenEntity =
-        tm().loadByKeyIfPresent(VKey.create(AllocationToken.class, token));
+        transactIfJpaTm(() -> tm().loadByKeyIfPresent(VKey.create(AllocationToken.class, token)));
     if (!maybeTokenEntity.isPresent()) {
       throw new InvalidAllocationTokenException();
     }

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

@@ -19,6 +19,7 @@ import static google.registry.flows.ResourceFlowUtils.loadAndVerifyExistence;
 import static google.registry.flows.host.HostFlowUtils.validateHostName;
 import static google.registry.model.EppResourceUtils.isLinked;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 
 import com.google.common.collect.ImmutableSet;
 import google.registry.flows.EppException;
@@ -76,7 +77,8 @@ public final class HostInfoFlow implements Flow {
     // there is no superordinate domain, the host's own values for these fields will be correct.
     if (host.isSubordinate()) {
       DomainBase superordinateDomain =
-          tm().loadByKey(host.getSuperordinateDomain()).cloneProjectedAtTime(now);
+          transactIfJpaTm(
+              () -> tm().loadByKey(host.getSuperordinateDomain()).cloneProjectedAtTime(now));
       hostInfoDataBuilder
           .setCurrentSponsorClientId(superordinateDomain.getCurrentSponsorClientId())
           .setLastTransferTime(host.computeLastTransferTime(superordinateDomain));

core/src/main/java/google/registry/flows/poll/PollAckFlow.java

@@ -16,15 +16,13 @@ package google.registry.flows.poll;
 
 import static google.registry.flows.FlowUtils.validateClientIsLoggedIn;
 import static google.registry.flows.poll.PollFlowUtils.ackPollMessage;
-import static google.registry.flows.poll.PollFlowUtils.getPollMessagesQuery;
+import static google.registry.flows.poll.PollFlowUtils.getPollMessageCount;
 import static google.registry.model.eppoutput.Result.Code.SUCCESS_WITH_NO_MESSAGES;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.poll.PollMessageExternalKeyConverter.makePollMessageExternalId;
 import static google.registry.model.poll.PollMessageExternalKeyConverter.parsePollMessageExternalId;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.DateTimeUtils.isBeforeOrAt;
 
-import com.googlecode.objectify.Key;
 import google.registry.flows.EppException;
 import google.registry.flows.EppException.AuthorizationErrorException;
 import google.registry.flows.EppException.ObjectDoesNotExistException;
@@ -39,6 +37,8 @@ import google.registry.model.poll.MessageQueueInfo;
 import google.registry.model.poll.PollMessage;
 import google.registry.model.poll.PollMessageExternalKeyConverter;
 import google.registry.model.poll.PollMessageExternalKeyConverter.PollMessageExternalKeyParseException;
+import google.registry.persistence.VKey;
+import java.util.Optional;
 import javax.inject.Inject;
 import org.joda.time.DateTime;
 
@@ -71,7 +71,7 @@ public class PollAckFlow implements TransactionalFlow {
       throw new MissingMessageIdException();
     }
 
-    Key<PollMessage> pollMessageKey;
+    VKey<PollMessage> pollMessageKey;
     // Try parsing the messageId, and throw an exception if it's invalid.
     try {
       pollMessageKey = parsePollMessageExternalId(messageId);
@@ -84,12 +84,13 @@ public class PollAckFlow implements TransactionalFlow {
     // Load the message to be acked. If a message is queued to be delivered in the future, we treat
     // it as if it doesn't exist yet. Same for if the message ID year isn't the same as the actual
     // poll message's event time (that means they're passing in an old already-acked ID).
-    PollMessage pollMessage = ofy().load().key(pollMessageKey).now();
-    if (pollMessage == null
-        || !isBeforeOrAt(pollMessage.getEventTime(), now)
-        || !makePollMessageExternalId(pollMessage).equals(messageId)) {
+    Optional<PollMessage> maybePollMessage = tm().loadByKeyIfPresent(pollMessageKey);
+    if (!maybePollMessage.isPresent()
+        || !isBeforeOrAt(maybePollMessage.get().getEventTime(), now)
+        || !makePollMessageExternalId(maybePollMessage.get()).equals(messageId)) {
       throw new MessageDoesNotExistException(messageId);
     }
+    PollMessage pollMessage = maybePollMessage.get();
 
     // Make sure this client is authorized to ack this message. It could be that the message is
     // supposed to go to a different registrar.
@@ -106,8 +107,11 @@ public class PollAckFlow implements TransactionalFlow {
     // acked, then we return a special status code indicating that. Note that the query will
     // include the message being acked.
 
-    int messageCount = tm().doTransactionless(() -> getPollMessagesQuery(clientId, now).count());
-    if (!includeAckedMessageInCount) {
+    int messageCount = tm().doTransactionless(() -> getPollMessageCount(clientId, now));
+    // Within the same transaction, Datastore will not reflect the updated count (potentially
+    // reduced by one thanks to the acked poll message). SQL will, however, so we shouldn't reduce
+    // the count in the SQL case.
+    if (!includeAckedMessageInCount && tm().isOfy()) {
       messageCount--;
     }
     if (messageCount <= 0) {

core/src/main/java/google/registry/flows/poll/PollFlowUtils.java

@@ -16,25 +16,56 @@ package google.registry.flows.poll;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.DateTimeUtils.isBeforeOrAt;
 
 import com.googlecode.objectify.cmd.Query;
 import google.registry.model.poll.PollMessage;
+import java.util.Optional;
 import org.joda.time.DateTime;
 
 /** Static utility functions for poll flows. */
 public final class PollFlowUtils {
 
-  private PollFlowUtils() {}
+  public static final String SQL_POLL_MESSAGE_QUERY =
+      "FROM PollMessage WHERE clientId = :registrarId AND eventTime <= :now ORDER BY eventTime ASC";
+  private static final String SQL_POLL_MESSAGE_COUNT_QUERY =
+      "SELECT COUNT(*) FROM PollMessage WHERE clientId = :registrarId AND eventTime <= :now";
 
-  /** Returns a query for poll messages for the logged in registrar which are not in the future. */
-  public static Query<PollMessage> getPollMessagesQuery(String clientId, DateTime now) {
-    return ofy().load()
-        .type(PollMessage.class)
-        .filter("clientId", clientId)
-        .filter("eventTime <=", now.toDate())
-        .order("eventTime");
+  /** Returns the number of poll messages for the given registrar that are not in the future. */
+  public static int getPollMessageCount(String registrarId, DateTime now) {
+    if (tm().isOfy()) {
+      return datastorePollMessageQuery(registrarId, now).count();
+    } else {
+      return jpaTm()
+          .transact(
+              () ->
+                  jpaTm()
+                      .query(SQL_POLL_MESSAGE_COUNT_QUERY, Long.class)
+                      .setParameter("registrarId", registrarId)
+                      .setParameter("now", now)
+                      .getSingleResult()
+                      .intValue());
+    }
+  }
+
+  /** Returns the first (by event time) poll message not in the future for this registrar. */
+  public static Optional<PollMessage> getFirstPollMessage(String registrarId, DateTime now) {
+    if (tm().isOfy()) {
+      return Optional.ofNullable(datastorePollMessageQuery(registrarId, now).first().now());
+    } else {
+      return jpaTm()
+          .transact(
+              () ->
+                  jpaTm()
+                      .query(SQL_POLL_MESSAGE_QUERY, PollMessage.class)
+                      .setParameter("registrarId", registrarId)
+                      .setParameter("now", now)
+                      .setMaxResults(1)
+                      .getResultStream()
+                      .findFirst());
+    }
   }
 
   /**
@@ -74,4 +105,16 @@ public final class PollFlowUtils {
     }
     return includeAckedMessageInCount;
   }
+
+  /** A Datastore query for poll messages from the given registrar that are not in the future. */
+  public static Query<PollMessage> datastorePollMessageQuery(String registrarId, DateTime now) {
+    return ofy()
+        .load()
+        .type(PollMessage.class)
+        .filter("clientId", registrarId)
+        .filter("eventTime <=", now.toDate())
+        .order("eventTime");
+  }
+
+  private PollFlowUtils() {}
 }

core/src/main/java/google/registry/flows/poll/PollRequestFlow.java

@@ -15,7 +15,8 @@
 package google.registry.flows.poll;
 
 import static google.registry.flows.FlowUtils.validateClientIsLoggedIn;
-import static google.registry.flows.poll.PollFlowUtils.getPollMessagesQuery;
+import static google.registry.flows.poll.PollFlowUtils.getFirstPollMessage;
+import static google.registry.flows.poll.PollFlowUtils.getPollMessageCount;
 import static google.registry.model.eppoutput.Result.Code.SUCCESS_WITH_ACK_MESSAGE;
 import static google.registry.model.eppoutput.Result.Code.SUCCESS_WITH_NO_MESSAGES;
 import static google.registry.model.poll.PollMessageExternalKeyConverter.makePollMessageExternalId;
@@ -31,6 +32,7 @@ import google.registry.model.poll.MessageQueueInfo;
 import google.registry.model.poll.PollMessage;
 import google.registry.model.poll.PollMessageExternalKeyConverter;
 import google.registry.util.Clock;
+import java.util.Optional;
 import javax.inject.Inject;
 import org.joda.time.DateTime;
 
@@ -63,18 +65,20 @@ public class PollRequestFlow implements Flow {
     }
     // Return the oldest message from the queue.
     DateTime now = clock.nowUtc();
-    PollMessage pollMessage = getPollMessagesQuery(clientId, now).first().now();
-    if (pollMessage == null) {
+    Optional<PollMessage> maybePollMessage = getFirstPollMessage(clientId, now);
+    if (!maybePollMessage.isPresent()) {
       return responseBuilder.setResultFromCode(SUCCESS_WITH_NO_MESSAGES).build();
     }
+    PollMessage pollMessage = maybePollMessage.get();
     return responseBuilder
         .setResultFromCode(SUCCESS_WITH_ACK_MESSAGE)
-        .setMessageQueueInfo(new MessageQueueInfo.Builder()
-            .setQueueDate(pollMessage.getEventTime())
-            .setMsg(pollMessage.getMsg())
-            .setQueueLength(getPollMessagesQuery(clientId, now).count())
-            .setMessageId(makePollMessageExternalId(pollMessage))
-            .build())
+        .setMessageQueueInfo(
+            new MessageQueueInfo.Builder()
+                .setQueueDate(pollMessage.getEventTime())
+                .setMsg(pollMessage.getMsg())
+                .setQueueLength(getPollMessageCount(clientId, now))
+                .setMessageId(makePollMessageExternalId(pollMessage))
+                .build())
         .setMultipleResData(pollMessage.getResponseData())
         .build();
   }

core/src/main/java/google/registry/keyring/api/InMemoryKeyring.java

@@ -39,8 +39,6 @@ public final class InMemoryKeyring implements Keyring {
   private final String marksdbLordnPassword;
   private final String marksdbSmdrlLoginAndPassword;
   private final String jsonCredential;
-  private final String cloudSqlPassword;
-  private final String toolsCloudSqlPassword;
 
   public InMemoryKeyring(
       PGPKeyPair rdeStagingKey,
@@ -83,8 +81,6 @@ public final class InMemoryKeyring implements Keyring {
     this.marksdbSmdrlLoginAndPassword =
         checkNotNull(marksdbSmdrlLoginAndPassword, "marksdbSmdrlLoginAndPassword");
     this.jsonCredential = checkNotNull(jsonCredential, "jsonCredential");
-    this.cloudSqlPassword = checkNotNull(cloudSqlPassword, "cloudSqlPassword");
-    this.toolsCloudSqlPassword = checkNotNull(toolsCloudSqlPassword, "toolsCloudSqlPassword");
   }
 
   @Override
@@ -157,16 +153,6 @@ public final class InMemoryKeyring implements Keyring {
     return jsonCredential;
   }
 
-  @Override
-  public String getCloudSqlPassword() {
-    return cloudSqlPassword;
-  }
-
-  @Override
-  public String getToolsCloudSqlPassword() {
-    return toolsCloudSqlPassword;
-  }
-
   /** Does nothing. */
   @Override
   public void close() {}

core/src/main/java/google/registry/keyring/api/KeyModule.java

@@ -36,18 +36,6 @@ public final class KeyModule {
     String value();
   }
 
-  @Provides
-  @Key("cloudSqlPassword")
-  static String providesCloudSqlPassword(Keyring keyring) {
-    return keyring.getCloudSqlPassword();
-  }
-
-  @Provides
-  @Key("toolsCloudSqlPassword")
-  static String providesToolsCloudSqlPassword(Keyring keyring) {
-    return keyring.getToolsCloudSqlPassword();
-  }
-
   @Provides
   @Key("brdaReceiverKey")
   static PGPPublicKey provideBrdaReceiverKey(Keyring keyring) {

core/src/main/java/google/registry/keyring/api/Keyring.java

@@ -28,12 +28,6 @@ import org.bouncycastle.openpgp.PGPPublicKey;
 @ThreadSafe
 public interface Keyring extends AutoCloseable {
 
-  /** Returns the password which is used by App Engine to connect to the Cloud SQL database. */
-  String getCloudSqlPassword();
-
-  /** Returns the password which is used by nomulus tool to connect to the Cloud SQL database. */
-  String getToolsCloudSqlPassword();
-
   /**
    * Returns the key which should be used to sign RDE deposits being uploaded to a third-party.
    *

core/src/main/java/google/registry/keyring/kms/KmsKeyring.java

@@ -19,14 +19,25 @@ import static com.google.common.base.CaseFormat.UPPER_UNDERSCORE;
 import static com.google.common.base.Preconditions.checkState;
 import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
 import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Streams;
+import com.google.common.flogger.FluentLogger;
 import com.googlecode.objectify.Key;
 import google.registry.config.RegistryConfig.Config;
 import google.registry.keyring.api.KeySerializer;
 import google.registry.keyring.api.Keyring;
 import google.registry.keyring.api.KeyringException;
 import google.registry.model.server.KmsSecret;
+import google.registry.model.server.KmsSecretRevision;
+import google.registry.model.server.KmsSecretRevisionSqlDao;
+import google.registry.privileges.secretmanager.KeyringSecretStore;
 import java.io.IOException;
+import java.util.Arrays;
+import java.util.Optional;
+import java.util.stream.Stream;
 import javax.inject.Inject;
 import org.bouncycastle.openpgp.PGPException;
 import org.bouncycastle.openpgp.PGPKeyPair;
@@ -42,6 +53,8 @@ import org.bouncycastle.openpgp.PGPPublicKey;
  */
 public class KmsKeyring implements Keyring {
 
+  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
+
   /** Key labels for private key secrets. */
   enum PrivateKeyLabel {
     BRDA_SIGNING_PRIVATE,
@@ -68,7 +81,6 @@ public class KmsKeyring implements Keyring {
 
   /** Key labels for string secrets. */
   enum StringKeyLabel {
-    CLOUD_SQL_PASSWORD_STRING,
     SAFE_BROWSING_API_KEY,
     ICANN_REPORTING_PASSWORD_STRING,
     JSON_CREDENTIAL_STRING,
@@ -76,8 +88,7 @@ public class KmsKeyring implements Keyring {
     MARKSDB_LORDN_PASSWORD_STRING,
     MARKSDB_SMDRL_LOGIN_STRING,
     RDE_SSH_CLIENT_PRIVATE_STRING,
-    RDE_SSH_CLIENT_PUBLIC_STRING,
-    TOOLS_CLOUD_SQL_PASSWORD_STRING;
+    RDE_SSH_CLIENT_PUBLIC_STRING;
 
     String getLabel() {
       return UPPER_UNDERSCORE.to(LOWER_HYPHEN, name());
@@ -85,20 +96,13 @@ public class KmsKeyring implements Keyring {
   }
 
   private final KmsConnection kmsConnection;
+  private final KeyringSecretStore secretStore;
 
   @Inject
-  KmsKeyring(@Config("defaultKmsConnection") KmsConnection kmsConnection) {
+  KmsKeyring(
+      @Config("defaultKmsConnection") KmsConnection kmsConnection, KeyringSecretStore secretStore) {
     this.kmsConnection = kmsConnection;
-  }
-
-  @Override
-  public String getCloudSqlPassword() {
-    return getString(StringKeyLabel.CLOUD_SQL_PASSWORD_STRING);
-  }
-
-  @Override
-  public String getToolsCloudSqlPassword() {
-    return getString(StringKeyLabel.TOOLS_CLOUD_SQL_PASSWORD_STRING);
+    this.secretStore = secretStore;
   }
 
   @Override
@@ -200,17 +204,83 @@ public class KmsKeyring implements Keyring {
     return getKeyPair(keyLabel).getPrivateKey();
   }
 
-  private byte[] getDecryptedData(String keyName) {
-    KmsSecret secret =
-        ofy().load().key(Key.create(getCrossTldKey(), KmsSecret.class, keyName)).now();
-    checkState(secret != null, "Requested secret '%s' does not exist.", keyName);
-    String encryptedData = ofy().load().key(secret.getLatestRevision()).now().getEncryptedValue();
+  private byte[] getDecryptedDataFromDatastore(String keyName) {
+    String encryptedData;
+    if (tm().isOfy()) {
+      KmsSecret secret =
+          ofy().load().key(Key.create(getCrossTldKey(), KmsSecret.class, keyName)).now();
+      checkState(secret != null, "Requested secret '%s' does not exist.", keyName);
+      encryptedData = ofy().load().key(secret.getLatestRevision()).now().getEncryptedValue();
+    } else {
+      Optional<KmsSecretRevision> revision =
+          tm().transact(() -> KmsSecretRevisionSqlDao.getLatestRevision(keyName));
+      checkState(revision.isPresent(), "Requested secret '%s' does not exist.", keyName);
+      encryptedData = revision.get().getEncryptedValue();
+    }
 
     try {
-      return kmsConnection.decrypt(secret.getName(), encryptedData);
+      return kmsConnection.decrypt(keyName, encryptedData);
     } catch (Exception e) {
       throw new KeyringException(
           String.format("CloudKMS decrypt operation failed for secret %s", keyName), e);
     }
   }
+
+  private byte[] getDataFromSecretStore(String keyName) {
+    try {
+      return secretStore.getSecret(keyName);
+    } catch (Exception e) {
+      return new byte[0];
+    }
+  }
+
+  private byte[] getDecryptedData(String keyName) {
+    byte[] dsData = getDecryptedDataFromDatastore(keyName);
+    byte[] secretStoreData = getDataFromSecretStore(keyName);
+
+    if (Arrays.equals(dsData, secretStoreData)) {
+      logger.atInfo().log("Values for %s in Datastore and Secret Manager match.", keyName);
+      return secretStoreData;
+    }
+    logger.atWarning().log("Values for %s in Datastore and Secret Manager do not match.", keyName);
+    return dsData;
+  }
+
+  /**
+   * Generates the tasks to migrate secrets from Datastore to Secret Manager.
+   *
+   * <p>The keys in the returned {@link ImmutableMap} are the names of the secrets that need
+   * migration. The values in the map are {@link Runnable Runnables} that copy secret data from
+   * Datastore to Secret Manager for their corresponding keys. Only secrets that are absent in
+   * Secret Manager or have inconsistent values are included in the returned map.
+   */
+  public ImmutableMap<String, Runnable> migrationPlan() {
+
+    ImmutableMap.Builder<String, Runnable> tasks = new ImmutableMap.Builder<>();
+
+    ImmutableList<String> labels =
+        Streams.concat(
+                Stream.of(PrivateKeyLabel.values()).map(PrivateKeyLabel::getLabel),
+                Stream.of(PublicKeyLabel.values()).map(PublicKeyLabel::getLabel),
+                Stream.of(StringKeyLabel.values()).map(StringKeyLabel::getLabel))
+            .collect(ImmutableList.toImmutableList());
+
+    for (String keyName : labels) {
+      byte[] dsData;
+      try {
+        dsData = getDecryptedDataFromDatastore(keyName);
+      } catch (IllegalStateException e) {
+        logger.atWarning().log("Cannot load %s from Datastore. Skipping...", keyName);
+        continue;
+      }
+      byte[] secretStoreData = getDataFromSecretStore(keyName);
+      if (Arrays.equals(dsData, secretStoreData)) {
+        logger.atInfo().log("%s is already up to date.\n", keyName);
+        continue;
+      }
+      logger.atInfo().log("%s needs to be migrated.\n", keyName);
+      tasks.put(keyName, () -> secretStore.createOrUpdateSecret(keyName, dsData));
+    }
+    return tasks.build();
+  }
 }

core/src/main/java/google/registry/keyring/kms/KmsUpdater.java

@@ -24,7 +24,6 @@ import static google.registry.keyring.kms.KmsKeyring.PublicKeyLabel.BRDA_SIGNING
 import static google.registry.keyring.kms.KmsKeyring.PublicKeyLabel.RDE_RECEIVER_PUBLIC;
 import static google.registry.keyring.kms.KmsKeyring.PublicKeyLabel.RDE_SIGNING_PUBLIC;
 import static google.registry.keyring.kms.KmsKeyring.PublicKeyLabel.RDE_STAGING_PUBLIC;
-import static google.registry.keyring.kms.KmsKeyring.StringKeyLabel.CLOUD_SQL_PASSWORD_STRING;
 import static google.registry.keyring.kms.KmsKeyring.StringKeyLabel.ICANN_REPORTING_PASSWORD_STRING;
 import static google.registry.keyring.kms.KmsKeyring.StringKeyLabel.JSON_CREDENTIAL_STRING;
 import static google.registry.keyring.kms.KmsKeyring.StringKeyLabel.MARKSDB_DNL_LOGIN_STRING;
@@ -33,8 +32,6 @@ import static google.registry.keyring.kms.KmsKeyring.StringKeyLabel.MARKSDB_SMDR
 import static google.registry.keyring.kms.KmsKeyring.StringKeyLabel.RDE_SSH_CLIENT_PRIVATE_STRING;
 import static google.registry.keyring.kms.KmsKeyring.StringKeyLabel.RDE_SSH_CLIENT_PUBLIC_STRING;
 import static google.registry.keyring.kms.KmsKeyring.StringKeyLabel.SAFE_BROWSING_API_KEY;
-import static google.registry.keyring.kms.KmsKeyring.StringKeyLabel.TOOLS_CLOUD_SQL_PASSWORD_STRING;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 
@@ -46,6 +43,7 @@ import google.registry.keyring.kms.KmsKeyring.PublicKeyLabel;
 import google.registry.keyring.kms.KmsKeyring.StringKeyLabel;
 import google.registry.model.server.KmsSecret;
 import google.registry.model.server.KmsSecretRevision;
+import google.registry.privileges.secretmanager.KeyringSecretStore;
 import java.io.IOException;
 import java.util.HashMap;
 import java.util.LinkedHashMap;
@@ -62,20 +60,19 @@ import org.bouncycastle.openpgp.PGPPublicKey;
 public final class KmsUpdater {
 
   private final KmsConnection kmsConnection;
+  private final KeyringSecretStore secretStore;
   private final HashMap<String, byte[]> secretValues;
 
   @Inject
-  public KmsUpdater(@Config("defaultKmsConnection") KmsConnection kmsConnection) {
+  public KmsUpdater(
+      @Config("defaultKmsConnection") KmsConnection kmsConnection, KeyringSecretStore secretStore) {
     this.kmsConnection = kmsConnection;
+    this.secretStore = secretStore;
 
     // Use LinkedHashMap to preserve insertion order on update() to simplify testing and debugging
     this.secretValues = new LinkedHashMap<>();
   }
 
-  public KmsUpdater setCloudSqlPassword(String password) {
-    return setString(password, CLOUD_SQL_PASSWORD_STRING);
-  }
-
   public KmsUpdater setRdeSigningKey(PGPKeyPair keyPair) throws IOException, PGPException {
     return setKeyPair(keyPair, RDE_SIGNING_PRIVATE, RDE_SIGNING_PUBLIC);
   }
@@ -108,10 +105,6 @@ public final class KmsUpdater {
     return setString(apiKey, SAFE_BROWSING_API_KEY);
   }
 
-  public KmsUpdater setToolsCloudSqlPassword(String password) {
-    return setString(password, TOOLS_CLOUD_SQL_PASSWORD_STRING);
-  }
-
   public KmsUpdater setIcannReportingPassword(String password) {
     return setString(password, ICANN_REPORTING_PASSWORD_STRING);
   }
@@ -143,6 +136,19 @@ public final class KmsUpdater {
     checkState(!secretValues.isEmpty(), "At least one Keyring value must be persisted");
 
     persistEncryptedValues(encryptValues(secretValues));
+
+    // Errors when writing to secret store can be thrown to the top, since writes are always
+    // executed by a human user using the UpdateKmsKeyringCommand.
+    try {
+      secretValues
+          .entrySet()
+          .forEach(e -> secretStore.createOrUpdateSecret(e.getKey(), e.getValue()));
+    } catch (RuntimeException e) {
+      throw new RuntimeException(
+          "Failed to persist secrets to Secret Manager. "
+              + "Please check the status of Secret Manager and re-run the command.",
+          e);
+    }
   }
 
   /**
@@ -194,8 +200,7 @@ public final class KmsUpdater {
    */
   private static void persistEncryptedValues(
       final ImmutableMap<String, EncryptResponse> encryptedValues) {
-    tm()
-        .transact(
+    tm().transact(
             () -> {
               for (Map.Entry<String, EncryptResponse> entry : encryptedValues.entrySet()) {
                 String secretName = entry.getKey();
@@ -207,7 +212,7 @@ public final class KmsUpdater {
                         .setKmsCryptoKeyVersionName(revisionData.cryptoKeyVersionName())
                         .setParent(secretName)
                         .build();
-                ofy().save().entities(secretRevision, KmsSecret.create(secretName, secretRevision));
+                tm().putAll(secretRevision, KmsSecret.create(secretName, secretRevision));
               }
             });
   }

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

@@ -16,6 +16,8 @@ package google.registry.model;
 
 import javax.persistence.Access;
 import javax.persistence.AccessType;
+import javax.persistence.AttributeOverride;
+import javax.persistence.Column;
 import javax.persistence.MappedSuperclass;
 import javax.xml.bind.annotation.XmlTransient;
 
@@ -40,6 +42,7 @@ public abstract class BackupGroupRoot extends ImmutableObject {
   // Prevents subclasses from unexpectedly accessing as property (e.g., HostResource), which would
   // require an unnecessary non-private setter method.
   @Access(AccessType.FIELD)
+  @AttributeOverride(name = "lastUpdateTime", column = @Column(name = "updateTimestamp"))
   UpdateAutoTimestamp updateTimestamp = UpdateAutoTimestamp.create(null);
 
   /** Get the {@link UpdateAutoTimestamp} for this entity. */

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

@@ -47,7 +47,6 @@ import google.registry.model.server.KmsSecret;
 import google.registry.model.server.KmsSecretRevision;
 import google.registry.model.server.Lock;
 import google.registry.model.server.ServerSecret;
-import google.registry.model.smd.SignedMarkRevocationList;
 import google.registry.model.tmch.ClaimsListShard;
 import google.registry.model.tmch.ClaimsListShard.ClaimsListRevision;
 import google.registry.model.tmch.ClaimsListShard.ClaimsListSingleton;
@@ -105,7 +104,6 @@ public final class EntityClasses {
           Registry.class,
           ReservedList.class,
           ServerSecret.class,
-          SignedMarkRevocationList.class,
           TmchCrl.class);
 
   private EntityClasses() {}

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

@@ -194,30 +194,6 @@ public final class EppResourceUtils {
     return ForeignKeyIndex.load(clazz, uniqueIds, now).keySet();
   }
 
-  /**
-   * Loads resources that match some filter and that have {@link EppResource#deletionTime} that is
-   * not before "now".
-   *
-   * <p>This is an eventually consistent query.
-   *
-   * @param clazz the resource type to load
-   * @param now the logical time of the check
-   * @param filterDefinition the filter to apply when loading resources
-   * @param filterValue the acceptable value for the filter
-   */
-  public static <T extends EppResource> Iterable<T> queryNotDeleted(
-      Class<T> clazz, DateTime now, String filterDefinition, Object filterValue) {
-    return ofy()
-        .load()
-        .type(clazz)
-        .filter(filterDefinition, filterValue)
-        .filter("deletionTime >", now.toDate())
-        .list()
-        .stream()
-        .map(EppResourceUtils.transformAtTime(now))
-        .collect(toImmutableSet());
-  }
-
   /**
    * Returns a Function that transforms an EppResource to the given DateTime, suitable for use with
    * Iterables.transform() over a collection of EppResources.

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

@@ -17,6 +17,7 @@ package google.registry.model;
 import static com.google.common.collect.Iterables.transform;
 import static com.google.common.collect.Maps.transformValues;
 import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static java.lang.annotation.ElementType.FIELD;
 import static java.lang.annotation.RetentionPolicy.RUNTIME;
 import static java.util.stream.Collectors.toCollection;
@@ -187,7 +188,10 @@ public abstract class ImmutableObject implements Cloneable {
   /** Helper function to recursively hydrate an ImmutableObject. */
   private static Object hydrate(Object value) {
     if (value instanceof Key) {
-      return hydrate(ofy().load().key((Key<?>) value).now());
+      if (tm().isOfy()) {
+        return hydrate(ofy().load().key((Key<?>) value).now());
+      }
+      return value;
     } else if (value instanceof Map) {
       return transformValues((Map<?, ?>) value, ImmutableObject::hydrate);
     } else if (value instanceof Collection) {

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

@@ -18,7 +18,6 @@ import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.collect.ImmutableList.toImmutableList;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.registry.Registry.TldState.GENERAL_AVAILABILITY;
 import static google.registry.model.registry.Registry.TldState.START_DATE_SUNRISE;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
@@ -41,10 +40,10 @@ import google.registry.model.registry.Registry;
 import google.registry.model.registry.Registry.TldState;
 import google.registry.model.registry.label.PremiumList;
 import google.registry.model.registry.label.PremiumListDualDao;
+import google.registry.persistence.VKey;
 import google.registry.util.CidrAddressBlock;
 import java.util.Collection;
 import java.util.Optional;
-import java.util.Set;
 import java.util.function.Function;
 import java.util.regex.Pattern;
 import org.joda.money.CurrencyUnit;
@@ -225,15 +224,14 @@ public final class OteAccountBuilder {
   }
 
   /**
-   * Persists all the OT&E entities to datastore.
+   * Persists all the OT&E entities to the database.
    *
    * @return map from the new clientIds created to the new TLDs they have access to. Can be used to
    *     go over all the newly created Registrars / Registries / RegistrarContacts if any
    *     post-creation work is needed.
    */
   public ImmutableMap<String, String> buildAndPersist() {
-    // save all the entitiesl in a single transaction
-    tm().transact(this::saveAllEntities);
+    saveAllEntities();
     return clientIdToTld;
   }
 
@@ -246,30 +244,38 @@ public final class OteAccountBuilder {
 
   /** Saves all the OT&amp;E entities we created. */
   private void saveAllEntities() {
-    tm().assertInTransaction();
-
     // use ImmutableObject instead of Registry so that the Key generation doesn't break
-    ImmutableList<ImmutableObject> registries = ImmutableList.of(sunriseTld, gaTld, eapTld);
+    ImmutableList<Registry> registries = ImmutableList.of(sunriseTld, gaTld, eapTld);
     ImmutableList<RegistrarContact> contacts = contactsBuilder.build();
 
-    if (!replaceExisting) {
-      ImmutableList<Key<ImmutableObject>> keys =
-          Streams.concat(registries.stream(), registrars.stream(), contacts.stream())
-              .map(Key::create)
-              .collect(toImmutableList());
-      Set<Key<ImmutableObject>> existingKeys = ofy().load().keys(keys).keySet();
-      checkState(
-          existingKeys.isEmpty(),
-          "Found existing object(s) conflicting with OT&E objects: %s",
-          existingKeys);
-    }
-    // Save the Registries (TLDs) first
-    ofy().save().entities(registries).now();
-    // Now we can set the allowedTlds for the registrars
-    registrars = registrars.stream().map(this::addAllowedTld).collect(toImmutableList());
-    // and we can save the registrars and contacts!
-    ofy().save().entities(registrars);
-    ofy().save().entities(contacts);
+    tm().transact(
+            () -> {
+              if (!replaceExisting) {
+                ImmutableList<VKey<? extends ImmutableObject>> keys =
+                    Streams.concat(
+                            registries.stream()
+                                .map(registry -> Registry.createVKey(registry.getTldStr())),
+                            registrars.stream().map(Registrar::createVKey),
+                            contacts.stream().map(RegistrarContact::createVKey))
+                        .collect(toImmutableList());
+                ImmutableMap<VKey<? extends ImmutableObject>, ImmutableObject> existingObjects =
+                    tm().loadByKeysIfPresent(keys);
+                checkState(
+                    existingObjects.isEmpty(),
+                    "Found existing object(s) conflicting with OT&E objects: %s",
+                    existingObjects.keySet());
+              }
+              // Save the Registries (TLDs) first
+              tm().putAll(registries);
+            });
+    // Now we can set the allowedTlds for the registrars in a new transaction
+    tm().transact(
+            () -> {
+              registrars = registrars.stream().map(this::addAllowedTld).collect(toImmutableList());
+              // and we can save the registrars and contacts!
+              tm().putAll(registrars);
+              tm().putAll(contacts);
+            });
   }
 
   private Registrar addAllowedTld(Registrar registrar) {

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

@@ -17,7 +17,6 @@ package google.registry.model;
 import static com.google.common.base.Predicates.equalTo;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static google.registry.model.eppcommon.EppXmlTransformer.unmarshal;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.util.CollectionUtils.isNullOrEmpty;
 import static google.registry.util.DomainNameUtils.ACE_PREFIX;
 
@@ -28,7 +27,6 @@ import com.google.common.collect.ImmutableCollection;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Multiset;
 import com.googlecode.objectify.Key;
-import com.googlecode.objectify.cmd.Query;
 import google.registry.model.domain.DomainCommand;
 import google.registry.model.domain.fee.FeeCreateCommandExtension;
 import google.registry.model.domain.launch.LaunchCreateExtension;
@@ -39,6 +37,7 @@ import google.registry.model.eppinput.EppInput.ResourceCommandWrapper;
 import google.registry.model.host.HostCommand;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.reporting.HistoryEntry.Type;
+import google.registry.model.reporting.HistoryEntryDao;
 import google.registry.xml.XmlException;
 import java.util.Arrays;
 import java.util.EnumSet;
@@ -196,16 +195,10 @@ public class OteStats {
    * <p>Stops when it notices that all tests have passed.
    */
   private OteStats recordRegistrarHistory(String registrarName) {
-    ImmutableCollection<String> clientIds =
+    ImmutableCollection<String> registrarIds =
         OteAccountBuilder.createClientIdToTldMap(registrarName).keySet();
 
-    Query<HistoryEntry> query =
-        ofy()
-            .load()
-            .type(HistoryEntry.class)
-            .filter("clientId in", clientIds)
-            .order("modificationTime");
-    for (HistoryEntry historyEntry : query) {
+    for (HistoryEntry historyEntry : HistoryEntryDao.loadHistoryObjectsByRegistrars(registrarIds)) {
       try {
         record(historyEntry);
       } catch (XmlException e) {

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

@@ -17,7 +17,6 @@ package google.registry.model;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Preconditions.checkState;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.common.collect.ImmutableList;
@@ -107,14 +106,14 @@ public final class ResourceTransferUtils {
   /** Update the relevant {@link ForeignKeyIndex} to cache the new deletion time. */
   public static <R extends EppResource> void updateForeignKeyIndexDeletionTime(R resource) {
     if (resource instanceof ForeignKeyedEppResource) {
-      ofy().save().entity(ForeignKeyIndex.create(resource, resource.getDeletionTime()));
+      tm().insert(ForeignKeyIndex.create(resource, resource.getDeletionTime()));
     }
   }
 
   /** If there is a transfer out, delete the server-approve entities and enqueue a poll message. */
   public static <R extends EppResource & ResourceWithTransferData>
       void handlePendingTransferOnDelete(
-            R resource, R newResource, DateTime now, HistoryEntry historyEntry) {
+          R resource, R newResource, DateTime now, HistoryEntry historyEntry) {
     if (resource.getStatusValues().contains(StatusValue.PENDING_TRANSFER)) {
       TransferData oldTransferData = resource.getTransferData();
       tm().delete(oldTransferData.getServerApproveEntities());

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

@@ -14,11 +14,22 @@
 
 package google.registry.model;
 
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
 
+import com.googlecode.objectify.annotation.Ignore;
+import com.googlecode.objectify.annotation.OnLoad;
 import google.registry.model.translators.UpdateAutoTimestampTranslatorFactory;
+import google.registry.util.DateTimeUtils;
+import java.time.ZonedDateTime;
 import java.util.Optional;
 import javax.annotation.Nullable;
+import javax.persistence.Column;
+import javax.persistence.Embeddable;
+import javax.persistence.PostLoad;
+import javax.persistence.PrePersist;
+import javax.persistence.PreUpdate;
+import javax.persistence.Transient;
 import org.joda.time.DateTime;
 
 /**
@@ -26,14 +37,45 @@ import org.joda.time.DateTime;
  *
  * @see UpdateAutoTimestampTranslatorFactory
  */
+@Embeddable
 public class UpdateAutoTimestamp extends ImmutableObject {
 
-  // When set to true, database converters/translators should do tha auto update.  When set to
+  // When set to true, database converters/translators should do the auto update.  When set to
   // false, auto update should be suspended (this exists to allow us to preserve the original value
   // during a replay).
   private static ThreadLocal<Boolean> autoUpdateEnabled = ThreadLocal.withInitial(() -> true);
 
-  DateTime timestamp;
+  @Transient DateTime timestamp;
+
+  @Ignore
+  @Column(nullable = false)
+  ZonedDateTime lastUpdateTime;
+
+  // Unfortunately, we cannot use the @UpdateTimestamp annotation on "lastUpdateTime" in this class
+  // because Hibernate does not allow it to be used on @Embeddable classes, see
+  // https://hibernate.atlassian.net/browse/HHH-13235. This is a workaround.
+  @PrePersist
+  @PreUpdate
+  void setTimestamp() {
+    if (autoUpdateEnabled() || lastUpdateTime == null) {
+      timestamp = jpaTm().getTransactionTime();
+      lastUpdateTime = DateTimeUtils.toZonedDateTime(timestamp);
+    }
+  }
+
+  @OnLoad
+  void onLoad() {
+    if (timestamp != null) {
+      lastUpdateTime = DateTimeUtils.toZonedDateTime(timestamp);
+    }
+  }
+
+  @PostLoad
+  void postLoad() {
+    if (lastUpdateTime != null) {
+      timestamp = DateTimeUtils.toJodaDateTime(lastUpdateTime);
+    }
+  }
 
   /** Returns the timestamp, or {@code START_OF_TIME} if it's null. */
   public DateTime getTimestamp() {
@@ -43,6 +85,7 @@ public class UpdateAutoTimestamp extends ImmutableObject {
   public static UpdateAutoTimestamp create(@Nullable DateTime timestamp) {
     UpdateAutoTimestamp instance = new UpdateAutoTimestamp();
     instance.timestamp = timestamp;
+    instance.lastUpdateTime = timestamp == null ? null : DateTimeUtils.toZonedDateTime(timestamp);
     return instance;
   }
 

core/src/main/java/google/registry/model/billing/BillingEvent.java

@@ -18,7 +18,8 @@ import static com.google.common.base.MoreObjects.firstNonNull;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Preconditions.checkState;
-import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 import static google.registry.util.CollectionUtils.forceEmptyToNull;
 import static google.registry.util.CollectionUtils.nullToEmptyImmutableCopy;
 import static google.registry.util.DateTimeUtils.END_OF_TIME;
@@ -41,6 +42,7 @@ import google.registry.model.ImmutableObject;
 import google.registry.model.annotations.ReportedOn;
 import google.registry.model.common.TimeOfYear;
 import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.DomainHistory;
 import google.registry.model.domain.GracePeriod;
 import google.registry.model.domain.rgp.GracePeriodStatus;
 import google.registry.model.domain.token.AllocationToken;
@@ -113,7 +115,7 @@ public abstract class BillingEvent extends ImmutableObject
   /** Entity id. */
   @Id @javax.persistence.Id Long id;
 
-  @Parent @DoNotHydrate @Transient Key<HistoryEntry> parent;
+  @Parent @DoNotHydrate @Transient Key<? extends HistoryEntry> parent;
 
   /** The registrar to bill. */
   @Index
@@ -190,7 +192,7 @@ public abstract class BillingEvent extends ImmutableObject
     return targetId;
   }
 
-  public Key<HistoryEntry> getParentKey() {
+  public Key<? extends HistoryEntry> getParentKey() {
     return parent;
   }
 
@@ -257,7 +259,7 @@ public abstract class BillingEvent extends ImmutableObject
       return thisCastToDerived();
     }
 
-    public B setParent(Key<HistoryEntry> parentKey) {
+    public B setParent(Key<? extends HistoryEntry> parentKey) {
       getInstance().parent = parentKey;
       return thisCastToDerived();
     }
@@ -601,23 +603,27 @@ public abstract class BillingEvent extends ImmutableObject
             GracePeriodStatus.TRANSFER, Reason.TRANSFER);
 
     /**
-     * Creates a cancellation billing event (parented on the provided history entry, and with the
-     * history entry's event time) that will cancel out the provided grace period's billing event,
+     * Creates a cancellation billing event (parented on the provided history key, and with the
+     * corresponding event time) that will cancel out the provided grace period's billing event,
      * using the supplied targetId and deriving other metadata (clientId, billing time, and the
      * cancellation reason) from the grace period.
      */
     public static BillingEvent.Cancellation forGracePeriod(
-        GracePeriod gracePeriod, HistoryEntry historyEntry, String targetId) {
+        GracePeriod gracePeriod,
+        DateTime eventTime,
+        Key<DomainHistory> domainHistoryKey,
+        String targetId) {
       checkArgument(gracePeriod.hasBillingEvent(),
           "Cannot create cancellation for grace period without billing event");
-      BillingEvent.Cancellation.Builder builder = new BillingEvent.Cancellation.Builder()
-          .setReason(checkNotNull(GRACE_PERIOD_TO_REASON.get(gracePeriod.getType())))
-          .setTargetId(targetId)
-          .setClientId(gracePeriod.getClientId())
-          .setEventTime(historyEntry.getModificationTime())
-          // The charge being cancelled will take place at the grace period's expiration time.
-          .setBillingTime(gracePeriod.getExpirationTime())
-          .setParent(historyEntry);
+      BillingEvent.Cancellation.Builder builder =
+          new BillingEvent.Cancellation.Builder()
+              .setReason(checkNotNull(GRACE_PERIOD_TO_REASON.get(gracePeriod.getType())))
+              .setTargetId(targetId)
+              .setClientId(gracePeriod.getClientId())
+              .setEventTime(eventTime)
+              // The charge being cancelled will take place at the grace period's expiration time.
+              .setBillingTime(gracePeriod.getExpirationTime())
+              .setParent(domainHistoryKey);
       // Set the grace period's billing event using the appropriate Cancellation builder method.
       if (gracePeriod.getOneTimeBillingEvent() != null) {
         builder.setOneTimeEventKey(gracePeriod.getOneTimeBillingEvent());
@@ -772,10 +778,10 @@ public abstract class BillingEvent extends ImmutableObject
         Modification instance = getInstance();
         checkNotNull(instance.reason);
         checkNotNull(instance.eventRef);
-        BillingEvent.OneTime billingEvent = ofy().load().key(instance.eventRef).now();
-        checkArgument(Objects.equals(
-            instance.cost.getCurrencyUnit(),
-            billingEvent.cost.getCurrencyUnit()),
+        BillingEvent.OneTime billingEvent =
+            transactIfJpaTm(() -> tm().loadByKey(VKey.from(instance.eventRef)));
+        checkArgument(
+            Objects.equals(instance.cost.getCurrencyUnit(), billingEvent.cost.getCurrencyUnit()),
             "Referenced billing event is in a different currency");
         return super.build();
       }

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

@@ -31,7 +31,7 @@ public abstract class CrossTldSingleton extends ImmutableObject {
 
   public static final long SINGLETON_ID = 1; // There is always exactly one of these.
 
-  @Id @Transient long id = SINGLETON_ID;
+  @Id @javax.persistence.Id long id = SINGLETON_ID;
 
   @Transient @Parent Key<EntityGroupRoot> parent = getCrossTldKey();
 }

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

@@ -36,6 +36,7 @@ import google.registry.persistence.VKey;
 import google.registry.schema.replay.DatastoreAndSqlEntity;
 import java.io.Serializable;
 import java.util.List;
+import java.util.Optional;
 import javax.persistence.Column;
 import javax.persistence.EnumType;
 import javax.persistence.Enumerated;
@@ -280,8 +281,8 @@ public class Cursor extends ImmutableObject implements DatastoreAndSqlEntity {
   /**
    * Returns the current time for a given cursor, or {@code START_OF_TIME} if the cursor is null.
    */
-  public static DateTime getCursorTimeOrStartOfTime(Cursor cursor) {
-    return cursor != null ? cursor.getCursorTime() : START_OF_TIME;
+  public static DateTime getCursorTimeOrStartOfTime(Optional<Cursor> cursor) {
+    return cursor.map(Cursor::getCursorTime).orElse(START_OF_TIME);
   }
 
   public DateTime getCursorTime() {

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

@@ -16,8 +16,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.ofy.ObjectifyService.ofy;
-import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
 
 import com.google.appengine.api.users.User;
 import com.google.common.base.Splitter;
@@ -56,13 +55,22 @@ public class GaeUserIdConverter extends ImmutableObject {
     try {
       // Perform these operations in a transactionless context to avoid enlisting in some outer
       // transaction (if any).
-      tm().doTransactionless(() -> ofy().saveWithoutBackup().entity(gaeUserIdConverter).now());
+      ofyTm()
+          .doTransactionless(
+              () -> {
+                ofyTm().putWithoutBackup(gaeUserIdConverter);
+                return null;
+              });
 
       // The read must be done in its own transaction to avoid reading from the session cache.
-      return tm()
-          .transactNew(() -> ofy().load().entity(gaeUserIdConverter).safe().user.getUserId());
+      return ofyTm().transactNew(() -> ofyTm().loadByEntity(gaeUserIdConverter).user.getUserId());
     } finally {
-      tm().doTransactionless(() -> ofy().deleteWithoutBackup().entity(gaeUserIdConverter).now());
+      ofyTm()
+          .doTransactionless(
+              () -> {
+                ofyTm().deleteWithoutBackup(gaeUserIdConverter);
+                return null;
+              });
     }
   }
 }

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

@@ -14,6 +14,8 @@
 
 package google.registry.model.contact;
 
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.EntitySubclass;
 import google.registry.model.ImmutableObject;
@@ -54,6 +56,7 @@ import javax.persistence.PostLoad;
 public class ContactHistory extends HistoryEntry implements SqlEntity {
 
   // Store ContactBase instead of ContactResource so we don't pick up its @Id
+  // Nullable for the sake of pre-Registry-3.0 history objects
   @Nullable ContactBase contactBase;
 
   @Id
@@ -114,6 +117,12 @@ public class ContactHistory extends HistoryEntry implements SqlEntity {
     return Optional.of(asHistoryEntry());
   }
 
+  // Used to fill out the contactBase field during asynchronous replay
+  public static void beforeSqlSave(ContactHistory contactHistory) {
+    contactHistory.contactBase =
+        jpaTm().loadByKey(VKey.createSql(ContactResource.class, contactHistory.getContactRepoId()));
+  }
+
   /** Class to represent the composite primary key of {@link ContactHistory} entity. */
   public static class ContactHistoryId extends ImmutableObject implements Serializable {
 
@@ -185,9 +194,13 @@ public class ContactHistory extends HistoryEntry implements SqlEntity {
       super(instance);
     }
 
-    public Builder setContactBase(ContactBase contactBase) {
+    public Builder setContactBase(@Nullable ContactBase contactBase) {
+      // Nullable for the sake of pre-Registry-3.0 history objects
+      if (contactBase == null) {
+        return this;
+      }
       getInstance().contactBase = contactBase;
-      return this;
+      return super.setParent(contactBase);
     }
 
     public Builder setContactRepoId(String contactRepoId) {

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

@@ -23,8 +23,9 @@ import static com.google.common.collect.Sets.difference;
 import static com.google.common.collect.Sets.intersection;
 import static google.registry.model.EppResourceUtils.projectResourceOntoBuilderAtTime;
 import static google.registry.model.EppResourceUtils.setAutomaticTransferSuccessProperties;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 import static google.registry.util.CollectionUtils.forceEmptyToNull;
 import static google.registry.util.CollectionUtils.nullToEmpty;
 import static google.registry.util.CollectionUtils.nullToEmptyImmutableCopy;
@@ -337,7 +338,7 @@ public class DomainContent extends EppResource
 
   @PostLoad
   @SuppressWarnings("UnusedMethod")
-  private final void postLoad() {
+  private void postLoad() {
     // Reconstitute the contact list.
     ImmutableSet.Builder<DesignatedContact> contactsBuilder = new ImmutableSet.Builder<>();
 
@@ -670,13 +671,11 @@ public class DomainContent extends EppResource
 
   /** Loads and returns the fully qualified host names of all linked nameservers. */
   public ImmutableSortedSet<String> loadNameserverHostNames() {
-    return ofy()
-        .load()
-        .keys(getNameservers().stream().map(VKey::getOfyKey).collect(toImmutableSet()))
-        .values()
-        .stream()
-        .map(HostResource::getHostName)
-        .collect(toImmutableSortedSet(Ordering.natural()));
+    return transactIfJpaTm(
+        () ->
+            tm().loadByKeys(getNameservers()).values().stream()
+                .map(HostResource::getHostName)
+                .collect(toImmutableSortedSet(Ordering.natural())));
   }
 
   /** A key to the registrant who registered this domain. */

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

@@ -15,6 +15,7 @@
 package google.registry.model.domain;
 
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.util.CollectionUtils.nullToEmptyImmutableCopy;
 
 import com.google.common.collect.ImmutableSet;
@@ -32,6 +33,7 @@ import google.registry.persistence.VKey;
 import google.registry.schema.replay.DatastoreEntity;
 import google.registry.schema.replay.SqlEntity;
 import java.io.Serializable;
+import java.util.HashSet;
 import java.util.Optional;
 import java.util.Set;
 import javax.annotation.Nullable;
@@ -75,6 +77,7 @@ import javax.persistence.Table;
 public class DomainHistory extends HistoryEntry implements SqlEntity {
 
   // Store DomainContent instead of DomainBase so we don't pick up its @Id
+  // Nullable for the sake of pre-Registry-3.0 history objects
   @Nullable DomainContent domainContent;
 
   @Id
@@ -125,7 +128,8 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
         insertable = false,
         updatable = false)
   })
-  Set<DomainDsDataHistory> dsDataHistories = ImmutableSet.of();
+  // HashSet rather than ImmutableSet so that Hibernate can fill them out lazily on request
+  Set<DomainDsDataHistory> dsDataHistories = new HashSet<>();
 
   @Ignore
   @OneToMany(
@@ -144,7 +148,8 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
         insertable = false,
         updatable = false)
   })
-  Set<GracePeriodHistory> gracePeriodHistories = ImmutableSet.of();
+  // HashSet rather than ImmutableSet so that Hibernate can fill them out lazily on request
+  Set<GracePeriodHistory> gracePeriodHistories = new HashSet<>();
 
   @Override
   @Nullable
@@ -263,6 +268,12 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
     return Optional.of(asHistoryEntry());
   }
 
+  // Used to fill out the domainContent field during asynchronous replay
+  public static void beforeSqlSave(DomainHistory domainHistory) {
+    domainHistory.domainContent =
+        jpaTm().loadByKey(VKey.createSql(DomainBase.class, domainHistory.getDomainRepoId()));
+  }
+
   /** Class to represent the composite primary key of {@link DomainHistory} entity. */
   public static class DomainHistoryId extends ImmutableObject implements Serializable {
 
@@ -334,9 +345,13 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
       super(instance);
     }
 
-    public Builder setDomainContent(DomainContent domainContent) {
+    public Builder setDomainContent(@Nullable DomainContent domainContent) {
+      // Nullable for the sake of pre-Registry-3.0 history objects
+      if (domainContent == null) {
+        return this;
+      }
       getInstance().domainContent = domainContent;
-      return this;
+      return super.setParent(domainContent);
     }
 
     public Builder setDomainRepoId(String domainRepoId) {

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

@@ -38,7 +38,7 @@ import org.joda.time.DateTime;
 public class GracePeriodBase extends ImmutableObject {
 
   /** Unique id required for hibernate representation. */
-  @Transient Long gracePeriodId;
+  @Transient long gracePeriodId;
 
   /** Repository id for the domain which this grace period belongs to. */
   @Ignore

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

@@ -14,6 +14,8 @@
 
 package google.registry.model.host;
 
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.EntitySubclass;
 import google.registry.model.ImmutableObject;
@@ -55,6 +57,7 @@ import javax.persistence.PostLoad;
 public class HostHistory extends HistoryEntry implements SqlEntity {
 
   // Store HostBase instead of HostResource so we don't pick up its @Id
+  // Nullable for the sake of pre-Registry-3.0 history objects
   @Nullable HostBase hostBase;
 
   @Id
@@ -115,6 +118,12 @@ public class HostHistory extends HistoryEntry implements SqlEntity {
     return Optional.of(asHistoryEntry());
   }
 
+  // Used to fill out the hostBase field during asynchronous replay
+  public static void beforeSqlSave(HostHistory hostHistory) {
+    hostHistory.hostBase =
+        jpaTm().loadByKey(VKey.createSql(HostResource.class, hostHistory.getHostRepoId()));
+  }
+
   /** Class to represent the composite primary key of {@link HostHistory} entity. */
   public static class HostHistoryId extends ImmutableObject implements Serializable {
 
@@ -186,9 +195,13 @@ public class HostHistory extends HistoryEntry implements SqlEntity {
       super(instance);
     }
 
-    public Builder setHostBase(HostBase hostBase) {
+    public Builder setHostBase(@Nullable HostBase hostBase) {
+      // Nullable for the sake of pre-Registry-3.0 history objects
+      if (hostBase == null) {
+        return this;
+      }
       getInstance().hostBase = hostBase;
-      return this;
+      return super.setParent(hostBase);
     }
 
     public Builder setHostRepoId(String hostRepoId) {

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

@@ -16,17 +16,20 @@ package google.registry.model.index;
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static google.registry.config.RegistryConfig.getEppResourceCachingDuration;
 import static google.registry.config.RegistryConfig.getEppResourceMaxCachedEntries;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.util.CollectionUtils.entriesToImmutableMap;
 import static google.registry.util.TypeUtils.instantiate;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.cache.CacheBuilder;
 import com.google.common.cache.CacheLoader;
 import com.google.common.cache.LoadingCache;
+import com.google.common.collect.ImmutableBiMap;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
@@ -45,10 +48,11 @@ import google.registry.model.contact.ContactResource;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.host.HostResource;
 import google.registry.persistence.VKey;
+import google.registry.persistence.transaction.CriteriaQueryBuilder;
 import google.registry.schema.replay.DatastoreOnlyEntity;
 import google.registry.util.NonFinalForTesting;
+import java.util.Collection;
 import java.util.Comparator;
-import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.ExecutionException;
@@ -81,10 +85,10 @@ public abstract class ForeignKeyIndex<E extends EppResource> extends BackupGroup
   public static class ForeignKeyHostIndex extends ForeignKeyIndex<HostResource>
       implements DatastoreOnlyEntity {}
 
-  private static final ImmutableMap<
+  private static final ImmutableBiMap<
           Class<? extends EppResource>, Class<? extends ForeignKeyIndex<?>>>
       RESOURCE_CLASS_TO_FKI_CLASS =
-          ImmutableMap.of(
+          ImmutableBiMap.of(
               ContactResource.class, ForeignKeyContactIndex.class,
               DomainBase.class, ForeignKeyDomainIndex.class,
               HostResource.class, ForeignKeyHostIndex.class);
@@ -184,71 +188,100 @@ public abstract class ForeignKeyIndex<E extends EppResource> extends BackupGroup
   }
 
   /**
-   * Load a list of {@link ForeignKeyIndex} instances by class and id strings that are active at or
+   * Load a map of {@link ForeignKeyIndex} instances by class and id strings that are active at or
    * after the specified moment in time.
    *
    * <p>The returned map will omit any keys for which the {@link ForeignKeyIndex} doesn't exist or
    * has been soft deleted.
    */
   public static <E extends EppResource> ImmutableMap<String, ForeignKeyIndex<E>> load(
-      Class<E> clazz, Iterable<String> foreignKeys, final DateTime now) {
+      Class<E> clazz, Collection<String> foreignKeys, final DateTime now) {
+    return loadIndexesFromStore(clazz, foreignKeys, true).entrySet().stream()
+        .filter(e -> now.isBefore(e.getValue().getDeletionTime()))
+        .collect(entriesToImmutableMap());
+  }
+
+  /**
+   * Helper method to load all of the most recent {@link ForeignKeyIndex}es for the given foreign
+   * keys, regardless of whether or not they have been soft-deleted.
+   *
+   * <p>Used by both the cached (w/o deletion check) and the non-cached (with deletion check) calls.
+   *
+   * <p>Note that in the cached case, we wish to run this outside of any transaction because we may
+   * be loading many entities, going over the Datastore limit on the number of enrolled entity
+   * groups per transaction (25). If we require consistency, however, we must use a transaction.
+   *
+   * @param inTransaction whether or not to use an Objectify transaction
+   */
+  private static <E extends EppResource>
+      ImmutableMap<String, ForeignKeyIndex<E>> loadIndexesFromStore(
+          Class<E> clazz, Collection<String> foreignKeys, boolean inTransaction) {
     if (tm().isOfy()) {
-      return ofy().load().type(mapToFkiClass(clazz)).ids(foreignKeys).entrySet().stream()
-          .filter(e -> now.isBefore(e.getValue().deletionTime))
-          .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue));
+      Class<ForeignKeyIndex<E>> fkiClass = mapToFkiClass(clazz);
+      return ImmutableMap.copyOf(
+          inTransaction
+              ? ofy().load().type(fkiClass).ids(foreignKeys)
+              : tm().doTransactionless(() -> ofy().load().type(fkiClass).ids(foreignKeys)));
     } else {
       String property = RESOURCE_CLASS_TO_FKI_PROPERTY.get(clazz);
-      List<E> entities =
+      ImmutableList<ForeignKeyIndex<E>> indexes =
           tm().transact(
-                  () -> {
-                    String entityName =
-                        jpaTm().getEntityManager().getMetamodel().entity(clazz).getName();
-                    return jpaTm()
-                        .query(
-                            String.format(
-                                "FROM %s WHERE %s IN :propertyValue and deletionTime > :now ",
-                                entityName, property),
-                            clazz)
-                        .setParameter("propertyValue", foreignKeys)
-                        .setParameter("now", now)
-                        .getResultList();
-                  });
+                  () ->
+                      jpaTm()
+                          .getEntityManager()
+                          .createQuery(
+                              CriteriaQueryBuilder.create(clazz)
+                                  .whereFieldIsIn(property, foreignKeys)
+                                  .build())
+                          .getResultStream()
+                          .map(e -> ForeignKeyIndex.create(e, e.getDeletionTime()))
+                          .collect(toImmutableList()));
       // We need to find and return the entities with the maximum deletionTime for each foreign key.
-      return Multimaps.index(entities, EppResource::getForeignKey).asMap().entrySet().stream()
+      return Multimaps.index(indexes, ForeignKeyIndex::getForeignKey).asMap().entrySet().stream()
           .map(
               entry ->
                   Maps.immutableEntry(
                       entry.getKey(),
                       entry.getValue().stream()
-                          .max(Comparator.comparing(EppResource::getDeletionTime))
+                          .max(Comparator.comparing(ForeignKeyIndex::getDeletionTime))
                           .get()))
-          .collect(
-              toImmutableMap(
-                  Map.Entry::getKey,
-                  entry -> create(entry.getValue(), entry.getValue().getDeletionTime())));
+          .collect(entriesToImmutableMap());
     }
   }
 
-  static final CacheLoader<Key<ForeignKeyIndex<?>>, Optional<ForeignKeyIndex<?>>> CACHE_LOADER =
-      new CacheLoader<Key<ForeignKeyIndex<?>>, Optional<ForeignKeyIndex<?>>>() {
+  static final CacheLoader<VKey<ForeignKeyIndex<?>>, Optional<ForeignKeyIndex<?>>> CACHE_LOADER =
+      new CacheLoader<VKey<ForeignKeyIndex<?>>, Optional<ForeignKeyIndex<?>>>() {
 
         @Override
-        public Optional<ForeignKeyIndex<?>> load(Key<ForeignKeyIndex<?>> key) {
-          return Optional.ofNullable(tm().doTransactionless(() -> ofy().load().key(key).now()));
+        public Optional<ForeignKeyIndex<?>> load(VKey<ForeignKeyIndex<?>> key) {
+          String foreignKey = key.getSqlKey().toString();
+          return Optional.ofNullable(
+              loadIndexesFromStore(
+                      RESOURCE_CLASS_TO_FKI_CLASS.inverse().get(key.getKind()),
+                      ImmutableSet.of(foreignKey),
+                      false)
+                  .get(foreignKey));
         }
 
         @Override
-        public Map<Key<ForeignKeyIndex<?>>, Optional<ForeignKeyIndex<?>>> loadAll(
-            Iterable<? extends Key<ForeignKeyIndex<?>>> keys) {
-          ImmutableSet<Key<ForeignKeyIndex<?>>> typedKeys = ImmutableSet.copyOf(keys);
-          Map<Key<ForeignKeyIndex<?>>, ForeignKeyIndex<?>> existingFkis =
-              tm().doTransactionless(() -> ofy().load().keys(typedKeys));
+        public Map<VKey<ForeignKeyIndex<?>>, Optional<ForeignKeyIndex<?>>> loadAll(
+            Iterable<? extends VKey<ForeignKeyIndex<?>>> keys) {
+          if (!keys.iterator().hasNext()) {
+            return ImmutableMap.of();
+          }
+          Class<? extends EppResource> resourceClass =
+              RESOURCE_CLASS_TO_FKI_CLASS.inverse().get(keys.iterator().next().getKind());
+          ImmutableSet<String> foreignKeys =
+              Streams.stream(keys).map(v -> v.getSqlKey().toString()).collect(toImmutableSet());
+          ImmutableSet<VKey<ForeignKeyIndex<?>>> typedKeys = ImmutableSet.copyOf(keys);
+          ImmutableMap<String, ? extends ForeignKeyIndex<? extends EppResource>> existingFkis =
+              loadIndexesFromStore(resourceClass, foreignKeys, false);
           // ofy() omits keys that don't have values in Datastore, so re-add them in
           // here with Optional.empty() values.
           return Maps.asMap(
               typedKeys,
-              (Key<ForeignKeyIndex<?>> key) ->
-                  Optional.ofNullable(existingFkis.getOrDefault(key, null)));
+              (VKey<ForeignKeyIndex<?>> key) ->
+                  Optional.ofNullable(existingFkis.getOrDefault(key.getSqlKey().toString(), null)));
         }
       };
 
@@ -266,10 +299,10 @@ public abstract class ForeignKeyIndex<E extends EppResource> extends BackupGroup
    * given IDs (blah) don't exist."
    */
   @NonFinalForTesting
-  private static LoadingCache<Key<ForeignKeyIndex<?>>, Optional<ForeignKeyIndex<?>>>
+  private static LoadingCache<VKey<ForeignKeyIndex<?>>, Optional<ForeignKeyIndex<?>>>
       cacheForeignKeyIndexes = createForeignKeyIndexesCache(getEppResourceCachingDuration());
 
-  private static LoadingCache<Key<ForeignKeyIndex<?>>, Optional<ForeignKeyIndex<?>>>
+  private static LoadingCache<VKey<ForeignKeyIndex<?>>, Optional<ForeignKeyIndex<?>>>
       createForeignKeyIndexesCache(Duration expiry) {
     return CacheBuilder.newBuilder()
         .expireAfterWrite(java.time.Duration.ofMillis(expiry.getMillis()))
@@ -294,25 +327,28 @@ public abstract class ForeignKeyIndex<E extends EppResource> extends BackupGroup
    * reasons, and are OK with the trade-offs in loss of transactional consistency.
    */
   public static <E extends EppResource> ImmutableMap<String, ForeignKeyIndex<E>> loadCached(
-      Class<E> clazz, Iterable<String> foreignKeys, final DateTime now) {
+      Class<E> clazz, Collection<String> foreignKeys, final DateTime now) {
     if (!RegistryConfig.isEppResourceCachingEnabled()) {
       return tm().doTransactionless(() -> load(clazz, foreignKeys, now));
     }
-    ImmutableList<Key<ForeignKeyIndex<?>>> fkiKeys =
+    Class<? extends ForeignKeyIndex<?>> fkiClass = mapToFkiClass(clazz);
+    // Safe to cast VKey<FKI<E>> to VKey<FKI<?>>
+    @SuppressWarnings("unchecked")
+    ImmutableList<VKey<ForeignKeyIndex<?>>> fkiVKeys =
         Streams.stream(foreignKeys)
-            .map(fk -> Key.<ForeignKeyIndex<?>>create(mapToFkiClass(clazz), fk))
+            .map(fk -> (VKey<ForeignKeyIndex<?>>) VKey.create(fkiClass, fk))
             .collect(toImmutableList());
     try {
       // This cast is safe because when we loaded ForeignKeyIndexes above we used type clazz, which
       // is scoped to E.
       @SuppressWarnings("unchecked")
       ImmutableMap<String, ForeignKeyIndex<E>> fkisFromCache =
-          cacheForeignKeyIndexes.getAll(fkiKeys).entrySet().stream()
+          cacheForeignKeyIndexes.getAll(fkiVKeys).entrySet().stream()
               .filter(entry -> entry.getValue().isPresent())
               .filter(entry -> now.isBefore(entry.getValue().get().getDeletionTime()))
               .collect(
                   toImmutableMap(
-                      entry -> entry.getKey().getName(),
+                      entry -> entry.getKey().getSqlKey().toString(),
                       entry -> (ForeignKeyIndex<E>) entry.getValue().get()));
       return fkisFromCache;
     } catch (ExecutionException e) {

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

@@ -37,12 +37,19 @@ import google.registry.model.domain.DomainHistory;
 import google.registry.model.host.HostHistory;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.persistence.VKey;
+import google.registry.persistence.transaction.QueryComposer;
 import google.registry.persistence.transaction.TransactionManager;
+import google.registry.schema.replay.DatastoreEntity;
+import google.registry.schema.replay.SqlEntity;
+import java.util.List;
 import java.util.NoSuchElementException;
 import java.util.Optional;
 import java.util.function.Supplier;
+import java.util.stream.Stream;
 import java.util.stream.StreamSupport;
 import javax.annotation.Nullable;
+import javax.persistence.NoResultException;
+import javax.persistence.NonUniqueResultException;
 import org.joda.time.DateTime;
 
 /** Datastore implementation of {@link TransactionManager}. */
@@ -136,22 +143,23 @@ public class DatastoreTransactionManager implements TransactionManager {
 
   @Override
   public void putAll(Object... entities) {
-    syncIfTransactionless(getOfy().save().entities(entities));
+    syncIfTransactionless(
+        getOfy().save().entities(toDatastoreEntities(ImmutableList.copyOf(entities))));
   }
 
   @Override
   public void putAll(ImmutableCollection<?> entities) {
-    syncIfTransactionless(getOfy().save().entities(entities));
+    syncIfTransactionless(getOfy().save().entities(toDatastoreEntities(entities)));
   }
 
   @Override
   public void putWithoutBackup(Object entity) {
-    syncIfTransactionless(getOfy().saveWithoutBackup().entities(entity));
+    syncIfTransactionless(getOfy().saveWithoutBackup().entities(toDatastoreEntity(entity)));
   }
 
   @Override
   public void putAllWithoutBackup(ImmutableCollection<?> entities) {
-    syncIfTransactionless(getOfy().saveWithoutBackup().entities(entities));
+    syncIfTransactionless(getOfy().saveWithoutBackup().entities(toDatastoreEntities(entities)));
   }
 
   @Override
@@ -176,7 +184,7 @@ public class DatastoreTransactionManager implements TransactionManager {
 
   @Override
   public boolean exists(Object entity) {
-    return getOfy().load().key(Key.create(entity)).now() != null;
+    return getOfy().load().key(Key.create(toDatastoreEntity(entity))).now() != null;
   }
 
   @Override
@@ -205,8 +213,7 @@ public class DatastoreTransactionManager implements TransactionManager {
     return getOfy().load().keys(keyMap.keySet()).entrySet().stream()
         .collect(
             toImmutableMap(
-                entry -> keyMap.get(entry.getKey()),
-                entry -> toChildHistoryEntryIfPossible(entry.getValue())));
+                entry -> keyMap.get(entry.getKey()), entry -> toSqlEntity(entry.getValue())));
   }
 
   @Override
@@ -239,7 +246,7 @@ public class DatastoreTransactionManager implements TransactionManager {
 
   @Override
   public <T> T loadByEntity(T entity) {
-    return ofy().load().entity(entity).now();
+    return (T) toSqlEntity(ofy().load().entity(toDatastoreEntity(entity)).now());
   }
 
   @Override
@@ -281,7 +288,7 @@ public class DatastoreTransactionManager implements TransactionManager {
 
   @Override
   public void delete(Object entity) {
-    syncIfTransactionless(getOfy().delete().entity(entity));
+    syncIfTransactionless(getOfy().delete().entity(toDatastoreEntity(entity)));
   }
 
   @Override
@@ -299,7 +306,12 @@ public class DatastoreTransactionManager implements TransactionManager {
 
   @Override
   public void deleteWithoutBackup(Object entity) {
-    syncIfTransactionless(getOfy().deleteWithoutBackup().entity(entity));
+    syncIfTransactionless(getOfy().deleteWithoutBackup().entity(toDatastoreEntity(entity)));
+  }
+
+  @Override
+  public <T> QueryComposer<T> createQueryComposer(Class<T> entity) {
+    return new DatastoreQueryComposerImpl(entity);
   }
 
   @Override
@@ -338,29 +350,104 @@ public class DatastoreTransactionManager implements TransactionManager {
    */
   private void saveEntity(Object entity) {
     checkArgumentNotNull(entity, "entity must be specified");
-    if (entity instanceof HistoryEntry) {
-      entity = ((HistoryEntry) entity).asHistoryEntry();
-    }
-    syncIfTransactionless(getOfy().save().entity(entity));
+    syncIfTransactionless(getOfy().save().entity(toDatastoreEntity(entity)));
   }
 
   @Nullable
   private <T> T loadNullable(VKey<T> key) {
-    return toChildHistoryEntryIfPossible(getOfy().load().key(key.getOfyKey()).now());
+    return toSqlEntity(getOfy().load().key(key.getOfyKey()).now());
   }
 
-  /** Converts a nonnull {@link HistoryEntry} to the child format, e.g. {@link DomainHistory} */
-  @SuppressWarnings("unchecked")
-  public static <T> T toChildHistoryEntryIfPossible(@Nullable T obj) {
-    // NB: The Key of the object in question may not necessarily be the resulting class that we
-    // wish to have. Because all *History classes are @EntitySubclasses, their Keys will have type
-    // HistoryEntry -- even if you create them based off the *History class.
-    if (obj instanceof HistoryEntry
-        && !(obj instanceof ContactHistory)
-        && !(obj instanceof DomainHistory)
-        && !(obj instanceof HostHistory)) {
-      return (T) ((HistoryEntry) obj).toChildHistoryEntity();
+  /**
+   * Converts a possible {@link SqlEntity} to a {@link DatastoreEntity}.
+   *
+   * <p>One example is that this would convert a {@link DomainHistory} to a {@link HistoryEntry}.
+   */
+  private static Object toDatastoreEntity(@Nullable Object obj) {
+    if (obj instanceof SqlEntity) {
+      Optional<DatastoreEntity> possibleDatastoreEntity = ((SqlEntity) obj).toDatastoreEntity();
+      if (possibleDatastoreEntity.isPresent()) {
+        return possibleDatastoreEntity.get();
+      }
     }
     return obj;
   }
+
+  /** Converts many possible {@link SqlEntity} objects to {@link DatastoreEntity} objects. */
+  private static ImmutableList<Object> toDatastoreEntities(ImmutableCollection<?> collection) {
+    return collection.stream()
+        .map(DatastoreTransactionManager::toDatastoreEntity)
+        .collect(toImmutableList());
+  }
+
+  /**
+   * Converts an object to the corresponding {@link SqlEntity} if necessary and possible.
+   *
+   * <p>This should be used when returning objects from Datastore to make sure they reflect the most
+   * recent type of the object in question.
+   */
+  @SuppressWarnings("unchecked")
+  public static <T> T toSqlEntity(@Nullable T obj) {
+    // NB: The Key of the object in question may not necessarily be the resulting class that we
+    // wish to have. For example, because all *History classes are @EntitySubclasses, their Keys
+    // will have type HistoryEntry -- even if you create them based off the *History class.
+    if (obj instanceof DatastoreEntity && !(obj instanceof SqlEntity)) {
+      Optional<SqlEntity> possibleSqlEntity = ((DatastoreEntity) obj).toSqlEntity();
+      if (possibleSqlEntity.isPresent()) {
+        return (T) possibleSqlEntity.get();
+      }
+    }
+    return obj;
+  }
+
+  private static class DatastoreQueryComposerImpl<T> extends QueryComposer<T> {
+    DatastoreQueryComposerImpl(Class<T> entityClass) {
+      super(entityClass);
+    }
+
+    Query<T> buildQuery() {
+      Query<T> result = ofy().load().type(entityClass);
+      for (WhereClause pred : predicates) {
+        result = result.filter(pred.fieldName + pred.comparator.getDatastoreString(), pred.value);
+      }
+
+      if (orderBy != null) {
+        result = result.order(orderBy);
+      }
+
+      return result;
+    }
+
+    @Override
+    public Optional<T> first() {
+      return Optional.ofNullable(buildQuery().first().now());
+    }
+
+    @Override
+    public T getSingleResult() {
+      List<T> results = buildQuery().limit(2).list();
+      if (results.size() == 0) {
+        // The exception text here is the same as what we get for JPA queries.
+        throw new NoResultException("No entity found for query");
+      } else if (results.size() > 1) {
+        throw new NonUniqueResultException("More than one result found for getSingleResult query");
+      }
+      return results.get(0);
+    }
+
+    @Override
+    public Stream<T> stream() {
+      return Streams.stream(buildQuery());
+    }
+
+    @Override
+    public long count() {
+      return buildQuery().count();
+    }
+
+    @Override
+    public List<T> list() {
+      return buildQuery().list();
+    }
+  }
 }

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

@@ -41,10 +41,11 @@ public class EntityWritePriorities {
    */
   static final ImmutableMap<String, Integer> CLASS_PRIORITIES =
       ImmutableMap.of(
-          "ContactResource", -15,
-          "HistoryEntry", -10,
           "AllocationToken", -9,
-          "DomainBase", 10);
+          "ContactResource", 8,
+          "HostResource", 9,
+          "DomainBase", 10,
+          "HistoryEntry", 20);
 
   // The beginning of the range of priority numbers reserved for delete.  This must be greater than
   // any of the values in CLASS_PRIORITIES by enough overhead to accommodate  any negative values in

core/src/main/java/google/registry/model/poll/PollMessage.java

@@ -106,7 +106,7 @@ public abstract class PollMessage extends ImmutableObject
   @Column(name = "poll_message_id")
   Long id;
 
-  @Parent @DoNotHydrate @Transient Key<HistoryEntry> parent;
+  @Parent @DoNotHydrate @Transient Key<? extends HistoryEntry> parent;
 
   /** The registrar that this poll message will be delivered to. */
   @Index
@@ -134,7 +134,7 @@ public abstract class PollMessage extends ImmutableObject
 
   @Ignore Long hostHistoryRevisionId;
 
-  public Key<HistoryEntry> getParentKey() {
+  public Key<? extends HistoryEntry> getParentKey() {
     return parent;
   }
 
@@ -239,7 +239,7 @@ public abstract class PollMessage extends ImmutableObject
       return thisCastToDerived();
     }
 
-    public B setParentKey(Key<HistoryEntry> parentKey) {
+    public B setParentKey(Key<? extends HistoryEntry> parentKey) {
       getInstance().parent = parentKey;
       return thisCastToDerived();
     }
@@ -398,6 +398,7 @@ public abstract class PollMessage extends ImmutableObject
       }
       if (!isNullOrEmpty(domainPendingActionNotificationResponses)) {
         pendingActionNotificationResponse = domainPendingActionNotificationResponses.get(0);
+        fullyQualifiedDomainName = pendingActionNotificationResponse.nameOrId.value;
       }
       if (!isNullOrEmpty(domainTransferResponses)) {
         fullyQualifiedDomainName = domainTransferResponses.get(0).getFullyQualifiedDomainName();
@@ -414,21 +415,23 @@ public abstract class PollMessage extends ImmutableObject
       // Take the SQL-specific fields and map them to the Objectify-specific fields, if applicable
       if (pendingActionNotificationResponse != null) {
         if (contactId != null) {
-          contactPendingActionNotificationResponses =
-              ImmutableList.of(
-                  ContactPendingActionNotificationResponse.create(
-                      pendingActionNotificationResponse.nameOrId.value,
-                      pendingActionNotificationResponse.getActionResult(),
-                      pendingActionNotificationResponse.getTrid(),
-                      pendingActionNotificationResponse.processedDate));
+          ContactPendingActionNotificationResponse contactPendingResponse =
+              ContactPendingActionNotificationResponse.create(
+                  pendingActionNotificationResponse.nameOrId.value,
+                  pendingActionNotificationResponse.getActionResult(),
+                  pendingActionNotificationResponse.getTrid(),
+                  pendingActionNotificationResponse.processedDate);
+          pendingActionNotificationResponse = contactPendingResponse;
+          contactPendingActionNotificationResponses = ImmutableList.of(contactPendingResponse);
         } else if (fullyQualifiedDomainName != null) {
-          domainPendingActionNotificationResponses =
-              ImmutableList.of(
-                  DomainPendingActionNotificationResponse.create(
-                      pendingActionNotificationResponse.nameOrId.value,
-                      pendingActionNotificationResponse.getActionResult(),
-                      pendingActionNotificationResponse.getTrid(),
-                      pendingActionNotificationResponse.processedDate));
+          DomainPendingActionNotificationResponse domainPendingResponse =
+              DomainPendingActionNotificationResponse.create(
+                  pendingActionNotificationResponse.nameOrId.value,
+                  pendingActionNotificationResponse.getActionResult(),
+                  pendingActionNotificationResponse.getTrid(),
+                  pendingActionNotificationResponse.processedDate);
+          pendingActionNotificationResponse = domainPendingResponse;
+          domainPendingActionNotificationResponses = ImmutableList.of(domainPendingResponse);
         }
       }
       if (transferResponse != null) {
@@ -474,38 +477,35 @@ public abstract class PollMessage extends ImmutableObject
       }
 
       public Builder setResponseData(ImmutableList<? extends ResponseData> responseData) {
-        getInstance().contactPendingActionNotificationResponses =
+        OneTime instance = getInstance();
+        instance.contactPendingActionNotificationResponses =
             forceEmptyToNull(
-                responseData
-                    .stream()
+                responseData.stream()
                     .filter(ContactPendingActionNotificationResponse.class::isInstance)
                     .map(ContactPendingActionNotificationResponse.class::cast)
                     .collect(toImmutableList()));
 
-        getInstance().contactTransferResponses =
+        instance.contactTransferResponses =
             forceEmptyToNull(
-                responseData
-                    .stream()
+                responseData.stream()
                     .filter(ContactTransferResponse.class::isInstance)
                     .map(ContactTransferResponse.class::cast)
                     .collect(toImmutableList()));
 
-        getInstance().domainPendingActionNotificationResponses =
+        instance.domainPendingActionNotificationResponses =
             forceEmptyToNull(
-                responseData
-                    .stream()
+                responseData.stream()
                     .filter(DomainPendingActionNotificationResponse.class::isInstance)
                     .map(DomainPendingActionNotificationResponse.class::cast)
                     .collect(toImmutableList()));
-        getInstance().domainTransferResponses =
+        instance.domainTransferResponses =
             forceEmptyToNull(
-                responseData
-                    .stream()
+                responseData.stream()
                     .filter(DomainTransferResponse.class::isInstance)
                     .map(DomainTransferResponse.class::cast)
                     .collect(toImmutableList()));
 
-        getInstance().hostPendingActionNotificationResponses =
+        instance.hostPendingActionNotificationResponses =
             forceEmptyToNull(
                 responseData.stream()
                     .filter(HostPendingActionNotificationResponse.class::isInstance)
@@ -513,26 +513,30 @@ public abstract class PollMessage extends ImmutableObject
                     .collect(toImmutableList()));
 
         // Set the generic pending-action field as appropriate
-        if (getInstance().contactPendingActionNotificationResponses != null) {
-          getInstance().pendingActionNotificationResponse =
-              getInstance().contactPendingActionNotificationResponses.get(0);
-        } else if (getInstance().domainPendingActionNotificationResponses != null) {
-          getInstance().pendingActionNotificationResponse =
-              getInstance().domainPendingActionNotificationResponses.get(0);
-        } else if (getInstance().hostPendingActionNotificationResponses != null) {
-          getInstance().pendingActionNotificationResponse =
-              getInstance().hostPendingActionNotificationResponses.get(0);
+        if (instance.contactPendingActionNotificationResponses != null) {
+          instance.pendingActionNotificationResponse =
+              instance.contactPendingActionNotificationResponses.get(0);
+          instance.contactId =
+              instance.contactPendingActionNotificationResponses.get(0).nameOrId.value;
+        } else if (instance.domainPendingActionNotificationResponses != null) {
+          instance.pendingActionNotificationResponse =
+              instance.domainPendingActionNotificationResponses.get(0);
+          instance.fullyQualifiedDomainName =
+              instance.domainPendingActionNotificationResponses.get(0).nameOrId.value;
+        } else if (instance.hostPendingActionNotificationResponses != null) {
+          instance.pendingActionNotificationResponse =
+              instance.hostPendingActionNotificationResponses.get(0);
         }
         // Set the generic transfer response field as appropriate
-        if (getInstance().contactTransferResponses != null) {
-          getInstance().contactId = getInstance().contactTransferResponses.get(0).getContactId();
-          getInstance().transferResponse = getInstance().contactTransferResponses.get(0);
-        } else if (getInstance().domainTransferResponses != null) {
-          getInstance().fullyQualifiedDomainName =
-              getInstance().domainTransferResponses.get(0).getFullyQualifiedDomainName();
-          getInstance().transferResponse = getInstance().domainTransferResponses.get(0);
-          getInstance().extendedRegistrationExpirationTime =
-              getInstance().domainTransferResponses.get(0).getExtendedRegistrationExpirationTime();
+        if (instance.contactTransferResponses != null) {
+          instance.contactId = getInstance().contactTransferResponses.get(0).getContactId();
+          instance.transferResponse = getInstance().contactTransferResponses.get(0);
+        } else if (instance.domainTransferResponses != null) {
+          instance.fullyQualifiedDomainName =
+              instance.domainTransferResponses.get(0).getFullyQualifiedDomainName();
+          instance.transferResponse = getInstance().domainTransferResponses.get(0);
+          instance.extendedRegistrationExpirationTime =
+              instance.domainTransferResponses.get(0).getExtendedRegistrationExpirationTime();
         }
         return this;
       }

core/src/main/java/google/registry/model/poll/PollMessageExternalKeyConverter.java

@@ -24,6 +24,7 @@ import google.registry.model.contact.ContactResource;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.host.HostResource;
 import google.registry.model.reporting.HistoryEntry;
+import google.registry.persistence.VKey;
 import java.util.List;
 
 /**
@@ -78,14 +79,14 @@ public class PollMessageExternalKeyConverter {
   /**
    * Returns an Objectify Key to a PollMessage corresponding with the external ID.
    *
-   * <p>Note that the year field that is included at the end of the poll message isn't actually
-   * used for anything; it exists solely to create unique externally visible IDs for autorenews. We
-   * thus ignore it (for now) for backwards compatibility reasons, so that registrars can still ACK
+   * <p>Note that the year field that is included at the end of the poll message isn't actually used
+   * for anything; it exists solely to create unique externally visible IDs for autorenews. We thus
+   * ignore it (for now) for backwards compatibility reasons, so that registrars can still ACK
    * existing poll message IDs they may have lying around.
    *
    * @throws PollMessageExternalKeyParseException if the external key has an invalid format.
    */
-  public static Key<PollMessage> parsePollMessageExternalId(String externalKey) {
+  public static VKey<PollMessage> parsePollMessageExternalId(String externalKey) {
     List<String> idComponents = Splitter.on('-').splitToList(externalKey);
     if (idComponents.size() != 6) {
       throw new PollMessageExternalKeyParseException();
@@ -96,16 +97,17 @@ public class PollMessageExternalKeyConverter {
       if (resourceClazz == null) {
         throw new PollMessageExternalKeyParseException();
       }
-      return Key.create(
+      return VKey.from(
           Key.create(
               Key.create(
-                  null,
-                  resourceClazz,
-                  String.format("%s-%s", idComponents.get(1), idComponents.get(2))),
-              HistoryEntry.class,
-              Long.parseLong(idComponents.get(3))),
-          PollMessage.class,
-          Long.parseLong(idComponents.get(4)));
+                  Key.create(
+                      null,
+                      resourceClazz,
+                      String.format("%s-%s", idComponents.get(1), idComponents.get(2))),
+                  HistoryEntry.class,
+                  Long.parseLong(idComponents.get(3))),
+              PollMessage.class,
+              Long.parseLong(idComponents.get(4))));
       // Note that idComponents.get(5) is entirely ignored; we never use the year field internally.
     } catch (NumberFormatException e) {
       throw new PollMessageExternalKeyParseException();

core/src/main/java/google/registry/model/rde/RdeRevision.java

@@ -17,6 +17,7 @@ package google.registry.model.rde;
 import static com.google.common.base.Preconditions.checkArgument;
 import static google.registry.model.rde.RdeNamingUtils.makePartialName;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 
 import com.google.common.base.VerifyException;
 import com.googlecode.objectify.Key;
@@ -97,10 +98,20 @@ public final class RdeRevision extends BackupGroupRoot implements NonReplicatedE
     RdeRevisionId sqlKey = RdeRevisionId.create(tld, date.toLocalDate(), mode);
     Key<RdeRevision> ofyKey = Key.create(RdeRevision.class, id);
     Optional<RdeRevision> revisionOptional =
-        tm().loadByKeyIfPresent(VKey.create(RdeRevision.class, sqlKey, ofyKey));
+        transactIfJpaTm(
+            () -> tm().loadByKeyIfPresent(VKey.create(RdeRevision.class, sqlKey, ofyKey)));
     return revisionOptional.map(rdeRevision -> rdeRevision.revision + 1).orElse(0);
   }
 
+  /** Returns the latest revision of the report already generated for the given triplet. */
+  public static Optional<Integer> getCurrentRevision(String tld, DateTime date, RdeMode mode) {
+    int nextRevision = getNextRevision(tld, date, mode);
+    if (nextRevision == 0) {
+      return Optional.empty();
+    }
+    return Optional.of(nextRevision - 1);
+  }
+
   /**
    * Sets the revision ID for a given triplet.
    *

core/src/main/java/google/registry/model/registrar/RegistrarContact.java

@@ -23,6 +23,7 @@ import static com.google.common.io.BaseEncoding.base64;
 import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.registrar.Registrar.checkValidEmail;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.CollectionUtils.nullToEmptyImmutableSortedCopy;
 import static google.registry.util.PasswordUtils.SALT_SUPPLIER;
@@ -200,17 +201,34 @@ public class RegistrarContact extends ImmutableObject
    * relevant Registrar entity with the {@link Registrar#contactsRequireSyncing} field set to true.
    */
   public static void updateContacts(
-      final Registrar registrar, final Set<RegistrarContact> contacts) {
+      final Registrar registrar, final ImmutableSet<RegistrarContact> contacts) {
     tm().transact(
             () -> {
-              ofy()
-                  .delete()
-                  .keys(
-                      difference(
-                          ImmutableSet.copyOf(
-                              ofy().load().type(RegistrarContact.class).ancestor(registrar).keys()),
-                          contacts.stream().map(Key::create).collect(toImmutableSet())));
-              ofy().save().entities(contacts);
+              if (tm().isOfy()) {
+                ImmutableSet<Key<RegistrarContact>> existingKeys =
+                    ImmutableSet.copyOf(
+                        ofy().load().type(RegistrarContact.class).ancestor(registrar).keys());
+                tm().delete(
+                        difference(
+                                existingKeys,
+                                contacts.stream().map(Key::create).collect(toImmutableSet()))
+                            .stream()
+                            .map(key -> VKey.createOfy(RegistrarContact.class, key))
+                            .collect(toImmutableSet()));
+              } else {
+                ImmutableSet<String> emailAddressesToKeep =
+                    contacts.stream()
+                        .map(RegistrarContact::getEmailAddress)
+                        .collect(toImmutableSet());
+                jpaTm()
+                    .query(
+                        "DELETE FROM RegistrarPoc WHERE registrarId = :registrarId AND "
+                            + "emailAddress NOT IN :emailAddressesToKeep")
+                    .setParameter("registrarId", registrar.getClientId())
+                    .setParameter("emailAddressesToKeep", emailAddressesToKeep)
+                    .executeUpdate();
+              }
+              tm().putAll(contacts);
             });
   }
 

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

@@ -106,6 +106,17 @@ public final class Registries {
     return tld;
   }
 
+  /**
+   * Pass-through check that the TLD exists, otherwise throw using the given error format message.
+   *
+   * <p>The specified TLD will be passed to the format message string.
+   */
+  public static String assertTldExists(String tld, String fmtMessage) {
+    String message = String.format(fmtMessage, tld);
+    checkArgument(getTlds().contains(checkArgumentNotNull(emptyToNull(tld), message)), message);
+    return tld;
+  }
+
   /** Pass-through check that every TLD in the given iterable exists, otherwise throw an IAE. */
   public static Iterable<String> assertTldsExist(Iterable<String> tlds) {
     for (String tld : tlds) {

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

@@ -51,10 +51,8 @@ public final class RegistryLockDao {
     return ImmutableList.copyOf(
         jpaTm()
             .query(
-                "SELECT lock FROM RegistryLock lock"
-                    + " WHERE lock.registrarId = :registrarId"
-                    + " AND lock.unlockCompletionTimestamp IS NULL"
-                    + " ORDER BY lock.domainName ASC",
+                "SELECT lock FROM RegistryLock lock WHERE lock.registrarId = :registrarId"
+                    + " AND lock.unlockCompletionTime IS NULL ORDER BY lock.domainName ASC",
                 RegistryLock.class)
             .setParameter("registrarId", registrarId)
             .getResultList());
@@ -89,9 +87,8 @@ public final class RegistryLockDao {
     return jpaTm()
         .query(
             "SELECT lock FROM RegistryLock lock WHERE lock.repoId = :repoId AND"
-                + " lock.lockCompletionTimestamp IS NOT NULL AND"
-                + " lock.unlockCompletionTimestamp IS NULL ORDER BY lock.revisionId"
-                + " DESC",
+                + " lock.lockCompletionTime IS NOT NULL AND lock.unlockCompletionTime IS NULL"
+                + " ORDER BY lock.revisionId DESC",
             RegistryLock.class)
         .setParameter("repoId", repoId)
         .setMaxResults(1)
@@ -110,8 +107,7 @@ public final class RegistryLockDao {
     return jpaTm()
         .query(
             "SELECT lock FROM RegistryLock lock WHERE lock.repoId = :repoId AND"
-                + " lock.unlockCompletionTimestamp IS NOT NULL ORDER BY lock.revisionId"
-                + " DESC",
+                + " lock.unlockCompletionTime IS NOT NULL ORDER BY lock.revisionId DESC",
             RegistryLock.class)
         .setParameter("repoId", repoId)
         .setMaxResults(1)