A better Datastore bulk delete tool (#947)

* A better Datastore bulk delete tool

Implements an alternative to GCP builtin tool with better performance
and usability.

CONTRIBUTORS

@@ -34,3 +34,5 @@ Guy Bensky <guyben@google.com>
 Weimin Yu <weiminyu@google.com>
 Shicong Huang <shicong@google.com>
 Gustav Brodman <gbrodman@google.com>
+Sarah Botwinick <sarahbot@google.com>
+Legina Chen <legina@google.com>

core/build.gradle

@@ -172,6 +172,8 @@ dependencies {
 
   compile deps['com.beust:jcommander']
   compile deps['com.google.api:gax']
+  compile deps['com.google.api.grpc:proto-google-cloud-datastore-v1']
+  compile deps['com.google.api.grpc:proto-google-common-protos']
   compile deps['com.google.api.grpc:proto-google-cloud-secretmanager-v1']
   compile deps['com.google.api-client:google-api-client']
   compile deps['com.google.api-client:google-api-client-appengine']
@@ -196,6 +198,8 @@ dependencies {
   compile deps['com.google.appengine:appengine-remote-api']
   compile deps['com.google.auth:google-auth-library-credentials']
   compile deps['com.google.auth:google-auth-library-oauth2-http']
+  compile deps['com.google.cloud.bigdataoss:util']
+  compile deps['com.google.cloud.datastore:datastore-v1-proto-client']
   compile deps['com.google.cloud.sql:jdbc-socket-factory-core']
   runtimeOnly deps['com.google.cloud.sql:postgres-socket-factory']
   compile deps['com.google.cloud:google-cloud-secretmanager']
@@ -600,7 +604,6 @@ task compileProdJS(type: JavaExec) {
   closureArgs << "--js=${jsDir}/soyutils_usegoog.js"
   closureArgs << "--js=${cssSourceDir}/registrar_bin.css.js"
   closureArgs << "--js=${jsSourceDir}/**.js"
-  // TODO(shicong) Verify the compiled JS file works in Alpha
   closureArgs << "--js=${externsDir}/json.js"
   closureArgs << "--js=${soySourceDir}/**.js"
   args closureArgs
@@ -737,6 +740,13 @@ project.tasks.create('initSqlPipeline', JavaExec) {
   }
 }
 
+// Caller must provide projectId, GCP region, runner, and the kinds to delete
+// (comma-separated kind names or '*' for all). E.g.:
+// nom_build :core:bulkDeleteDatastore --args="--project=domain-registry-crash \
+//     --region=us-central1 --runner=DataflowRunner --kindsToDelete=*"
+createToolTask(
+    'bulkDeleteDatastore', 'google.registry.beam.datastore.BulkDeletePipeline')
+
 project.tasks.create('generateSqlSchema', JavaExec) {
   classpath = sourceSets.nonprod.runtimeClasspath
   main = 'google.registry.tools.DevTool'
@@ -828,9 +838,8 @@ task buildToolImage(dependsOn: nomulus, type: Exec) {
   commandLine 'docker', 'build', '-t', 'nomulus-tool', '.'
 }
 
-task generateInitSqlPipelineGraph(type: Test) {
-  include "**/InitSqlPipelineGraphTest.*"
-  testNameIncludePatterns = ["**createPipeline_compareGraph"]
+task generateInitSqlPipelineGraph(type: FilteringTest) {
+  tests = ['InitSqlPipelineGraphTest.createPipeline_compareGraph']
   ignoreFailures = true
 }
 

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

@@ -85,7 +85,7 @@ public class ReplayCommitLogsToSqlAction implements Runnable {
     Optional<Lock> lock =
         Lock.acquire(
             this.getClass().getSimpleName(), null, LEASE_LENGTH, requestStatusChecker, false);
-    if (lock.isEmpty()) {
+    if (!lock.isPresent()) {
       String message = "Can't acquire SQL commit log replay lock, aborting.";
       logger.atSevere().log(message);
       // App Engine will retry on any non-2xx status code, which we don't want in this case.
@@ -182,7 +182,7 @@ public class ReplayCommitLogsToSqlAction implements Runnable {
   }
 
   private static int compareByWeight(VersionedEntity a, VersionedEntity b) {
-    return getEntityPriority(a.key().getKind(), a.getEntity().isEmpty())
-        - getEntityPriority(b.key().getKind(), b.getEntity().isEmpty());
+    return getEntityPriority(a.key().getKind(), !a.getEntity().isPresent())
+        - getEntityPriority(b.key().getKind(), !b.getEntity().isPresent());
   }
 }

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

@@ -345,7 +345,7 @@ public class DeleteContactsAndHostsAction implements Runnable {
       String resourceClientId = resource.getPersistedCurrentSponsorClientId();
       if (resource instanceof HostResource && ((HostResource) resource).isSubordinate()) {
         resourceClientId =
-            tm().load(((HostResource) resource).getSuperordinateDomain())
+            tm().loadByKey(((HostResource) resource).getSuperordinateDomain())
                 .cloneProjectedAtTime(now)
                 .getCurrentSponsorClientId();
       }
@@ -465,7 +465,7 @@ public class DeleteContactsAndHostsAction implements Runnable {
         if (host.isSubordinate()) {
           dnsQueue.addHostRefreshTask(host.getHostName());
           tm().put(
-                  tm().load(host.getSuperordinateDomain())
+                  tm().loadByKey(host.getSuperordinateDomain())
                       .asBuilder()
                       .removeSubordinateHost(host.getHostName())
                       .build());

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

@@ -0,0 +1,193 @@
+// 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.collect.ImmutableList.toImmutableList;
+import static com.google.common.net.MediaType.PLAIN_TEXT_UTF_8;
+import static google.registry.flows.FlowUtils.marshalWithLenientRetry;
+import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.util.DateTimeUtils.END_OF_TIME;
+import static google.registry.util.ResourceUtils.readResourceUtf8;
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static javax.servlet.http.HttpServletResponse.SC_INTERNAL_SERVER_ERROR;
+import static javax.servlet.http.HttpServletResponse.SC_NO_CONTENT;
+import static javax.servlet.http.HttpServletResponse.SC_OK;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.flogger.FluentLogger;
+import google.registry.config.RegistryConfig.Config;
+import google.registry.flows.EppController;
+import google.registry.flows.EppRequestSource;
+import google.registry.flows.PasswordOnlyTransportCredentials;
+import google.registry.flows.StatelessRequestSessionMetadata;
+import google.registry.model.domain.DomainBase;
+import google.registry.model.eppcommon.ProtocolDefinition;
+import google.registry.model.eppoutput.EppOutput;
+import google.registry.request.Action;
+import google.registry.request.Action.Method;
+import google.registry.request.Response;
+import google.registry.request.auth.Auth;
+import google.registry.request.lock.LockHandler;
+import google.registry.util.Clock;
+import java.util.Optional;
+import java.util.concurrent.Callable;
+import javax.inject.Inject;
+import org.joda.time.DateTime;
+import org.joda.time.Duration;
+
+/**
+ * An action that deletes all non-renewing domains whose expiration dates have now passed.
+ *
+ * <p>The registry runs on an autorenew domain model, so domains don't ever expire naturally; they
+ * are only ever autorenewed. However, in some situations (such as URS) we don't want this to
+ * happen. Thus, the domains are tagged as non-renewing and are deleted by the next daily invocation
+ * of this action once they are past the date at which they were to expire.
+ *
+ * <p>Note that this action works by running a superuser EPP domain delete command, and as a side
+ * effect of when domains are deleted (just past their expiration date), they are invariably in the
+ * autorenew grace period when this happens.
+ */
+@Action(
+    service = Action.Service.BACKEND,
+    path = DeleteExpiredDomainsAction.PATH,
+    auth = Auth.AUTH_INTERNAL_OR_ADMIN,
+    method = Method.POST)
+public class DeleteExpiredDomainsAction implements Runnable {
+
+  public static final String PATH = "/_dr/task/deleteExpiredDomains";
+  private static final String LOCK_NAME = "Delete expired domains";
+
+  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
+  private final EppController eppController;
+  private final String registryAdminClientId;
+  private final Clock clock;
+  private final LockHandler lockHandler;
+  private final Response response;
+  private final String deleteXmlTmpl;
+
+  @Inject
+  DeleteExpiredDomainsAction(
+      EppController eppController,
+      @Config("registryAdminClientId") String registryAdminClientId,
+      Clock clock,
+      LockHandler lockHandler,
+      Response response) {
+    this.eppController = eppController;
+    this.registryAdminClientId = registryAdminClientId;
+    this.clock = clock;
+    this.lockHandler = lockHandler;
+    this.response = response;
+    this.deleteXmlTmpl =
+        readResourceUtf8(DeleteExpiredDomainsAction.class, "delete_expired_domain.xml");
+  }
+
+  @Override
+  public void run() {
+    response.setContentType(PLAIN_TEXT_UTF_8);
+
+    Callable<Void> runner =
+        () -> {
+          try {
+            runLocked();
+            response.setStatus(SC_OK);
+          } catch (Exception e) {
+            response.setStatus(SC_INTERNAL_SERVER_ERROR);
+            response.setPayload("Encountered error; see GCP logs for full details.");
+          }
+          return null;
+        };
+
+    if (!lockHandler.executeWithLocks(runner, null, Duration.standardHours(1), LOCK_NAME)) {
+      // Send a 200-series status code to prevent this conflicting action from retrying.
+      response.setStatus(SC_NO_CONTENT);
+      response.setPayload("Could not acquire lock; already running?");
+    }
+  }
+
+  private void runLocked() {
+    DateTime runTime = clock.nowUtc();
+    logger.atInfo().log(
+        "Deleting non-renewing domains with autorenew end times up through %s.", runTime);
+
+    // Note: This query is (and must be) non-transactional, and thus, is only eventually consistent.
+    ImmutableList<DomainBase> domainsToDelete =
+        ofy().load().type(DomainBase.class).filter("autorenewEndTime <=", runTime).list().stream()
+            // Datastore can't do two inequalities in one query, so the second happens in-memory.
+            .filter(d -> d.getDeletionTime().isEqual(END_OF_TIME))
+            .collect(toImmutableList());
+    if (domainsToDelete.isEmpty()) {
+      logger.atInfo().log("Found 0 domains to delete.");
+      response.setPayload("Found 0 domains to delete.");
+      return;
+    }
+
+    logger.atInfo().log(
+        "Found %d domains to delete: %s.",
+        domainsToDelete.size(),
+        String.join(
+            ", ",
+            domainsToDelete.stream().map(DomainBase::getDomainName).collect(toImmutableList())));
+    domainsToDelete.forEach(this::runDomainDeleteFlow);
+    logger.atInfo().log("Finished deleting domains.");
+    response.setPayload("Finished deleting domains.");
+  }
+
+  private void runDomainDeleteFlow(DomainBase domain) {
+    logger.atInfo().log("Attempting to delete domain %s", domain.getDomainName());
+    // Create a new transaction that the flow's execution will be enlisted in that loads the domain
+    // transactionally. This way we can ensure that nothing else has modified the domain in question
+    // in the intervening period since the query above found it.
+    Optional<EppOutput> eppOutput =
+        tm().transact(
+                () -> {
+                  DomainBase transDomain = tm().loadByKey(domain.createVKey());
+                  if (!domain.getAutorenewEndTime().isPresent()
+                      || domain.getAutorenewEndTime().get().isAfter(tm().getTransactionTime())) {
+                    logger.atSevere().log(
+                        "Failed to delete domain %s because of its autorenew end time: %s.",
+                        transDomain.getDomainName(), transDomain.getAutorenewEndTime());
+                    return Optional.empty();
+                  } else if (domain.getDeletionTime().isBefore(END_OF_TIME)) {
+                    logger.atSevere().log(
+                        "Failed to delete domain %s because it was already deleted on %s.",
+                        transDomain.getDomainName(), transDomain.getDeletionTime());
+                    return Optional.empty();
+                  }
+                  return Optional.of(
+                      eppController.handleEppCommand(
+                          new StatelessRequestSessionMetadata(
+                              registryAdminClientId,
+                              ProtocolDefinition.getVisibleServiceExtensionUris()),
+                          new PasswordOnlyTransportCredentials(),
+                          EppRequestSource.BACKEND,
+                          false,
+                          true,
+                          deleteXmlTmpl
+                              .replace("%DOMAIN%", transDomain.getDomainName())
+                              .getBytes(UTF_8)));
+                });
+
+    if (eppOutput.isPresent()) {
+      if (eppOutput.get().isSuccess()) {
+        logger.atInfo().log("Successfully deleted domain %s", domain.getDomainName());
+      } else {
+        logger.atWarning().log(
+            "Failed to delete domain %s; EPP response:\n\n%s",
+            domain.getDomainName(), new String(marshalWithLenientRetry(eppOutput.get()), UTF_8));
+      }
+    }
+  }
+}

core/src/main/java/google/registry/batch/delete_expired_domain.xml

@@ -0,0 +1,17 @@
+<epp xmlns="urn:ietf:params:xml:ns:epp-1.0">
+  <command>
+    <delete>
+      <domain:delete
+       xmlns:domain="urn:ietf:params:xml:ns:domain-1.0">
+        <domain:name>%DOMAIN%</domain:name>
+      </domain:delete>
+    </delete>
+    <extension>
+      <metadata:metadata xmlns:metadata="urn:google:params:xml:ns:metadata-1.0">
+        <metadata:reason>Non-renewing domain has reached expiration date.</metadata:reason>
+        <metadata:requestedByRegistrar>false</metadata:requestedByRegistrar>
+      </metadata:metadata>
+    </extension>
+    <clTRID>ABC-12345</clTRID>
+  </command>
+</epp>

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

@@ -0,0 +1,330 @@
+// 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.datastore;
+
+import static com.google.common.base.Preconditions.checkState;
+import static org.apache.beam.sdk.values.TypeDescriptors.kvs;
+import static org.apache.beam.sdk.values.TypeDescriptors.strings;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Splitter;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableSortedSet;
+import com.google.common.flogger.FluentLogger;
+import com.google.datastore.v1.Entity;
+import java.util.Iterator;
+import java.util.Map;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.extensions.gcp.options.GcpOptions;
+import org.apache.beam.sdk.io.gcp.datastore.DatastoreIO;
+import org.apache.beam.sdk.options.Default;
+import org.apache.beam.sdk.options.Description;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.options.Validation;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.GroupByKey;
+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.Reshuffle;
+import org.apache.beam.sdk.transforms.View;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PBegin;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.PCollectionView;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
+
+/**
+ * A BEAM pipeline that deletes Datastore entities in bulk.
+ *
+ * <p>This pipeline provides an alternative to the <a
+ * href="https://cloud.google.com/datastore/docs/bulk-delete">GCP builtin template</a> that performs
+ * the same task. It solves the following performance and usability problems in the builtin
+ * template:
+ *
+ * <ul>
+ *   <li>When deleting all data (by using the {@code select __key__} or {@code select *} queries),
+ *       the builtin template cannot parallelize the query, therefore has to query with a single
+ *       worker.
+ *   <li>When deleting all data, the builtin template also attempts to delete Datastore internal
+ *       tables which would cause permission-denied errors, which in turn MAY cause the pipeline to
+ *       abort before all data has been deleted.
+ *   <li>With the builtin template, it is possible to delete multiple entity types in one pipeline
+ *       ONLY if the user can come up with a single literal query that covers all of them. This is
+ *       not the case with most Nomulus entity types.
+ * </ul>
+ *
+ * <p>A user of this pipeline must specify the types of entities to delete using the {@code
+ * --kindsToDelete} command line argument. To delete specific entity types, give a comma-separated
+ * string of their kind names; to delete all data, give {@code "*"}.
+ *
+ * <p>When deleting all data, it is recommended for the user to specify the number of user entity
+ * types in the Datastore using the {@code --numOfKindsHint} argument. If the default value for this
+ * parameter is too low, performance will suffer.
+ */
+public class BulkDeletePipeline {
+  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
+
+  // This tool is not for use in our critical projects.
+  private static final ImmutableSet<String> FORBIDDEN_PROJECTS =
+      ImmutableSet.of("domain-registry", "domain-registry-sandbox");
+
+  private final BulkDeletePipelineOptions options;
+
+  private final Pipeline pipeline;
+
+  BulkDeletePipeline(BulkDeletePipelineOptions options) {
+    this.options = options;
+    pipeline = Pipeline.create(options);
+  }
+
+  public void run() {
+    setupPipeline();
+    pipeline.run();
+  }
+
+  @SuppressWarnings("deprecation") // org.apache.beam.sdk.transforms.Reshuffle
+  private void setupPipeline() {
+    checkState(
+        !FORBIDDEN_PROJECTS.contains(options.getProject()),
+        "Bulk delete is forbidden in %s",
+        options.getProject());
+
+    // Pre-allocated tags to label entities by kind. In the case of delete-all, we must use a guess.
+    TupleTagList deletionTags;
+    PCollection<String> kindsToDelete;
+
+    if (options.getKindsToDelete().equals("*")) {
+      deletionTags = getDeletionTags(options.getNumOfKindsHint());
+      kindsToDelete =
+          pipeline.apply("DiscoverEntityKinds", discoverEntityKinds(options.getProject()));
+    } else {
+      ImmutableList<String> kindsToDeleteParam = parseKindsToDelete(options);
+      checkState(
+          !kindsToDeleteParam.contains("*"),
+          "The --kindsToDelete argument should not contain both '*' and other kinds.");
+      deletionTags = getDeletionTags(kindsToDeleteParam.size());
+      kindsToDelete = pipeline.apply("UseProvidedKinds", Create.of(kindsToDeleteParam));
+    }
+
+    // Map each kind to a tag. The "SplitByKind" stage below will group entities by kind using
+    // this mapping. In practice, this has been effective at avoiding entity group contentions.
+    PCollectionView<Map<String, TupleTag<Entity>>> kindToTagMapping =
+        mapKindsToDeletionTags(kindsToDelete, deletionTags).apply("GetKindsToTagMap", View.asMap());
+
+    PCollectionTuple entities =
+        kindsToDelete
+            .apply("GenerateQueries", ParDo.of(new GenerateQueries()))
+            .apply("ReadEntities", DatastoreV1.read().withProjectId(options.getProject()))
+            .apply(
+                "SplitByKind",
+                ParDo.of(new SplitEntities(kindToTagMapping))
+                    .withSideInputs(kindToTagMapping)
+                    .withOutputTags(getOneDeletionTag("placeholder"), deletionTags));
+
+    for (TupleTag<?> tag : deletionTags.getAll()) {
+      entities
+          .get((TupleTag<Entity>) tag)
+          // Reshuffle calls GroupByKey which is one way to trigger load rebalance in the pipeline.
+          // Using the deprecated "Reshuffle" for convenience given the short life of this tool.
+          .apply("RebalanceLoad", Reshuffle.viaRandomKey())
+          .apply(
+              "DeleteEntities_" + tag.getId(),
+              DatastoreIO.v1().deleteEntity().withProjectId(options.getProject()));
+    }
+  }
+
+  private static String toKeyOnlyQueryForKind(String kind) {
+    return "select __key__ from `" + kind + "`";
+  }
+
+  /**
+   * Returns a {@link TupleTag} that retains the generic type parameter and may be used in a
+   * multi-output {@link ParDo} (e.g. {@link SplitEntities}).
+   *
+   * <p>This method is NOT needed in tests when creating tags for assertions. Simply create them
+   * with {@code new TupleTag<Entity>(String)}.
+   */
+  @VisibleForTesting
+  static TupleTag<Entity> getOneDeletionTag(String id) {
+    // The trailing {} is needed to retain generic param type.
+    return new TupleTag<Entity>(id) {};
+  }
+
+  @VisibleForTesting
+  static ImmutableList<String> parseKindsToDelete(BulkDeletePipelineOptions options) {
+    return ImmutableList.copyOf(
+        Splitter.on(",").omitEmptyStrings().trimResults().split(options.getKindsToDelete().trim()));
+  }
+
+  /**
+   * Returns a list of {@code n} {@link TupleTag TupleTags} numbered from {@code 0} to {@code n-1}.
+   */
+  @VisibleForTesting
+  static TupleTagList getDeletionTags(int n) {
+    ImmutableList.Builder<TupleTag<?>> builder = new ImmutableList.Builder<>();
+    for (int i = 0; i < n; i++) {
+      builder.add(getOneDeletionTag(String.valueOf(i)));
+    }
+    return TupleTagList.of(builder.build());
+  }
+
+  /** Returns a {@link PTransform} that finds all entity kinds in Datastore. */
+  @VisibleForTesting
+  static PTransform<PBegin, PCollection<String>> discoverEntityKinds(String project) {
+    return new PTransform<PBegin, PCollection<String>>() {
+      @Override
+      public PCollection<String> expand(PBegin input) {
+        // Use the __kind__ table to discover entity kinds. Data in the more informational
+        // __Stat_Kind__ table may be up to 48-hour stale.
+        return input
+            .apply(
+                "LoadEntityMetaData",
+                DatastoreIO.v1()
+                    .read()
+                    .withProjectId(project)
+                    .withLiteralGqlQuery("select * from __kind__"))
+            .apply(
+                "GetKindNames",
+                ParDo.of(
+                    new DoFn<Entity, String>() {
+                      @ProcessElement
+                      public void processElement(
+                          @Element Entity entity, OutputReceiver<String> out) {
+                        String kind = entity.getKey().getPath(0).getName();
+                        if (kind.startsWith("_")) {
+                          return;
+                        }
+                        out.output(kind);
+                      }
+                    }));
+      }
+    };
+  }
+
+  @VisibleForTesting
+  static PCollection<KV<String, TupleTag<Entity>>> mapKindsToDeletionTags(
+      PCollection<String> kinds, TupleTagList tags) {
+    // The first two stages send all strings in the 'kinds' PCollection to one worker which
+    // performs the mapping in the last stage.
+    return kinds
+        .apply(
+            "AssignSingletonKeyToKinds",
+            MapElements.into(kvs(strings(), strings())).via(kind -> KV.of("", kind)))
+        .apply("GatherKindsIntoCollection", GroupByKey.create())
+        .apply("MapKindsToTag", ParDo.of(new MapKindsToTags(tags)));
+  }
+
+  /** Transforms each {@code kind} string into a Datastore query for that kind. */
+  @VisibleForTesting
+  static class GenerateQueries extends DoFn<String, String> {
+    @ProcessElement
+    public void processElement(@Element String kind, OutputReceiver<String> out) {
+      out.output(toKeyOnlyQueryForKind(kind));
+    }
+  }
+
+  private static class MapKindsToTags
+      extends DoFn<KV<String, Iterable<String>>, KV<String, TupleTag<Entity>>> {
+    private final TupleTagList tupleTags;
+
+    MapKindsToTags(TupleTagList tupleTags) {
+      this.tupleTags = tupleTags;
+    }
+
+    @ProcessElement
+    public void processElement(
+        @Element KV<String, Iterable<String>> kv,
+        OutputReceiver<KV<String, TupleTag<Entity>>> out) {
+      // Sort kinds so that mapping is deterministic.
+      ImmutableSortedSet<String> sortedKinds = ImmutableSortedSet.copyOf(kv.getValue());
+      Iterator<String> kinds = sortedKinds.iterator();
+      Iterator<TupleTag<?>> tags = tupleTags.getAll().iterator();
+
+      while (kinds.hasNext() && tags.hasNext()) {
+        out.output(KV.of(kinds.next(), (TupleTag<Entity>) tags.next()));
+      }
+
+      if (kinds.hasNext()) {
+        logger.atWarning().log(
+            "There are more kinds to delete (%s) than our estimate (%s). "
+                + "Performance may suffer.",
+            sortedKinds.size(), tupleTags.size());
+      }
+      // Round robin assignment so that mapping is deterministic
+      while (kinds.hasNext()) {
+        tags = tupleTags.getAll().iterator();
+        while (kinds.hasNext() && tags.hasNext()) {
+          out.output(KV.of(kinds.next(), (TupleTag<Entity>) tags.next()));
+        }
+      }
+    }
+  }
+
+  /**
+   * {@link DoFn} that splits one {@link PCollection} of mixed kinds into multiple single-kind
+   * {@code PCollections}.
+   */
+  @VisibleForTesting
+  static class SplitEntities extends DoFn<Entity, Entity> {
+    private final PCollectionView<Map<String, TupleTag<Entity>>> kindToTagMapping;
+
+    SplitEntities(PCollectionView<Map<String, TupleTag<Entity>>> kindToTagMapping) {
+      super();
+      this.kindToTagMapping = kindToTagMapping;
+    }
+
+    @ProcessElement
+    public void processElement(ProcessContext context) {
+      Entity entity = context.element();
+      com.google.datastore.v1.Key entityKey = entity.getKey();
+      String kind = entityKey.getPath(entityKey.getPathCount() - 1).getKind();
+      TupleTag<Entity> tag = context.sideInput(kindToTagMapping).get(kind);
+      context.output(tag, entity);
+    }
+  }
+
+  public static void main(String[] args) {
+    BulkDeletePipelineOptions options =
+        PipelineOptionsFactory.fromArgs(args).withValidation().as(BulkDeletePipelineOptions.class);
+    BulkDeletePipeline pipeline = new BulkDeletePipeline(options);
+    pipeline.run();
+    System.exit(0);
+  }
+
+  public interface BulkDeletePipelineOptions extends GcpOptions {
+
+    @Description(
+        "The Datastore KINDs to be deleted. The format may be:\n"
+            + "\t- The list of kinds to be deleted as a comma-separated string, or\n"
+            + "\t- '*', which causes all kinds to be deleted.")
+    @Validation.Required
+    String getKindsToDelete();
+
+    void setKindsToDelete(String kinds);
+
+    @Description(
+        "An estimate of the number of KINDs to be deleted. "
+            + "This is recommended if --kindsToDelete is '*' and the default value is too low.")
+    @Default.Integer(30)
+    int getNumOfKindsHint();
+
+    void setNumOfKindsHint(int numOfKindsHint);
+  }
+}

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

@@ -0,0 +1,765 @@
+// 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.
+
+// This class is adapted from the Apache BEAM SDK. The original license may
+// be found at <a href="https://github.com/apache/beam/blob/master/LICENSE">
+// this link</a>.
+
+package google.registry.beam.datastore;
+
+import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Verify.verify;
+import static com.google.datastore.v1.PropertyFilter.Operator.EQUAL;
+import static com.google.datastore.v1.PropertyOrder.Direction.DESCENDING;
+import static com.google.datastore.v1.QueryResultBatch.MoreResultsType.NOT_FINISHED;
+import static com.google.datastore.v1.client.DatastoreHelper.makeAndFilter;
+import static com.google.datastore.v1.client.DatastoreHelper.makeFilter;
+import static com.google.datastore.v1.client.DatastoreHelper.makeOrder;
+import static com.google.datastore.v1.client.DatastoreHelper.makeValue;
+
+import com.google.api.client.http.HttpRequestInitializer;
+import com.google.auth.Credentials;
+import com.google.auth.http.HttpCredentialsAdapter;
+import com.google.auto.value.AutoValue;
+import com.google.cloud.hadoop.util.ChainingHttpRequestInitializer;
+import com.google.common.base.Strings;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.flogger.FluentLogger;
+import com.google.datastore.v1.Entity;
+import com.google.datastore.v1.EntityResult;
+import com.google.datastore.v1.GqlQuery;
+import com.google.datastore.v1.PartitionId;
+import com.google.datastore.v1.Query;
+import com.google.datastore.v1.QueryResultBatch;
+import com.google.datastore.v1.RunQueryRequest;
+import com.google.datastore.v1.RunQueryResponse;
+import com.google.datastore.v1.client.Datastore;
+import com.google.datastore.v1.client.DatastoreException;
+import com.google.datastore.v1.client.DatastoreFactory;
+import com.google.datastore.v1.client.DatastoreHelper;
+import com.google.datastore.v1.client.DatastoreOptions;
+import com.google.datastore.v1.client.QuerySplitter;
+import com.google.protobuf.Int32Value;
+import com.google.rpc.Code;
+import java.io.Serializable;
+import java.util.List;
+import java.util.NoSuchElementException;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.extensions.gcp.options.GcpOptions;
+import org.apache.beam.sdk.extensions.gcp.util.RetryHttpRequestInitializer;
+import org.apache.beam.sdk.metrics.Counter;
+import org.apache.beam.sdk.metrics.Metrics;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.Reshuffle;
+import org.apache.beam.sdk.transforms.display.DisplayData;
+import org.apache.beam.sdk.transforms.display.HasDisplayData;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.BackOffUtils;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.sdk.util.Sleeper;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.joda.time.Duration;
+
+/**
+ * Contains an adaptation of {@link org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.Read}. See
+ * {@link MultiRead} for details.
+ */
+public class DatastoreV1 {
+
+  // A package-private constructor to prevent direct instantiation from outside of this package
+  DatastoreV1() {}
+
+  /**
+   * Non-retryable errors. See https://cloud.google.com/datastore/docs/concepts/errors#Error_Codes .
+   */
+  private static final ImmutableSet<Code> NON_RETRYABLE_ERRORS =
+      ImmutableSet.of(
+          Code.FAILED_PRECONDITION,
+          Code.INVALID_ARGUMENT,
+          Code.PERMISSION_DENIED,
+          Code.UNAUTHENTICATED);
+
+  /**
+   * Returns an empty {@link MultiRead} builder. Configure the source {@code projectId}, {@code
+   * query}, and optionally {@code namespace} and {@code numQuerySplits} using {@link
+   * MultiRead#withProjectId}, {@link MultiRead#withNamespace}, {@link
+   * MultiRead#withNumQuerySplits}.
+   */
+  public static MultiRead read() {
+    return new AutoValue_DatastoreV1_MultiRead.Builder().setNumQuerySplits(0).build();
+  }
+
+  /**
+   * A {@link PTransform} that executes every Cloud SQL queries in a {@link PCollection } and reads
+   * their result rows as {@code Entity} objects.
+   *
+   * <p>This class is adapted from {@link org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.Read}. It
+   * uses literal GQL queries in the input {@link PCollection} instead of a constant query provided
+   * to the builder. Only the {@link #expand} method is modified from the original. Everything else
+   * including comments have been copied verbatim.
+   */
+  @AutoValue
+  public abstract static class MultiRead
+      extends PTransform<PCollection<String>, PCollection<Entity>> {
+    private static final FluentLogger logger = FluentLogger.forEnclosingClass();
+
+    /** An upper bound on the number of splits for a query. */
+    public static final int NUM_QUERY_SPLITS_MAX = 50000;
+
+    /** A lower bound on the number of splits for a query. */
+    static final int NUM_QUERY_SPLITS_MIN = 12;
+
+    /** Default bundle size of 64MB. */
+    static final long DEFAULT_BUNDLE_SIZE_BYTES = 64L * 1024L * 1024L;
+
+    /**
+     * Maximum number of results to request per query.
+     *
+     * <p>Must be set, or it may result in an I/O error when querying Cloud Datastore.
+     */
+    static final int QUERY_BATCH_LIMIT = 500;
+
+    public abstract @Nullable String getProjectId();
+
+    public abstract @Nullable String getNamespace();
+
+    public abstract int getNumQuerySplits();
+
+    public abstract @Nullable String getLocalhost();
+
+    @Override
+    public abstract String toString();
+
+    abstract Builder toBuilder();
+
+    @AutoValue.Builder
+    abstract static class Builder {
+      abstract Builder setProjectId(String projectId);
+
+      abstract Builder setNamespace(String namespace);
+
+      abstract Builder setNumQuerySplits(int numQuerySplits);
+
+      abstract Builder setLocalhost(String localhost);
+
+      abstract MultiRead build();
+    }
+
+    /**
+     * Computes the number of splits to be performed on the given query by querying the estimated
+     * size from Cloud Datastore.
+     */
+    static int getEstimatedNumSplits(Datastore datastore, Query query, @Nullable String namespace) {
+      int numSplits;
+      try {
+        long estimatedSizeBytes = getEstimatedSizeBytes(datastore, query, namespace);
+        logger.atInfo().log("Estimated size bytes for the query is: %s", estimatedSizeBytes);
+        numSplits =
+            (int)
+                Math.min(
+                    NUM_QUERY_SPLITS_MAX,
+                    Math.round(((double) estimatedSizeBytes) / DEFAULT_BUNDLE_SIZE_BYTES));
+      } catch (Exception e) {
+        logger.atWarning().log("Failed the fetch estimatedSizeBytes for query: %s", query, e);
+        // Fallback in case estimated size is unavailable.
+        numSplits = NUM_QUERY_SPLITS_MIN;
+      }
+      return Math.max(numSplits, NUM_QUERY_SPLITS_MIN);
+    }
+
+    /**
+     * Cloud Datastore system tables with statistics are periodically updated. This method fetches
+     * the latest timestamp (in microseconds) of statistics update using the {@code __Stat_Total__}
+     * table.
+     */
+    private static long queryLatestStatisticsTimestamp(
+        Datastore datastore, @Nullable String namespace) throws DatastoreException {
+      Query.Builder query = Query.newBuilder();
+      // Note: namespace either being null or empty represents the default namespace, in which
+      // case we treat it as not provided by the user.
+      if (Strings.isNullOrEmpty(namespace)) {
+        query.addKindBuilder().setName("__Stat_Total__");
+      } else {
+        query.addKindBuilder().setName("__Stat_Ns_Total__");
+      }
+      query.addOrder(makeOrder("timestamp", DESCENDING));
+      query.setLimit(Int32Value.newBuilder().setValue(1));
+      RunQueryRequest request = makeRequest(query.build(), namespace);
+
+      RunQueryResponse response = datastore.runQuery(request);
+      QueryResultBatch batch = response.getBatch();
+      if (batch.getEntityResultsCount() == 0) {
+        throw new NoSuchElementException("Datastore total statistics unavailable");
+      }
+      Entity entity = batch.getEntityResults(0).getEntity();
+      return entity.getProperties().get("timestamp").getTimestampValue().getSeconds() * 1000000;
+    }
+
+    /** Retrieve latest table statistics for a given kind, namespace, and datastore. */
+    private static Entity getLatestTableStats(
+        String ourKind, @Nullable String namespace, Datastore datastore) throws DatastoreException {
+      long latestTimestamp = queryLatestStatisticsTimestamp(datastore, namespace);
+      logger.atInfo().log("Latest stats timestamp for kind %s is %s", ourKind, latestTimestamp);
+
+      Query.Builder queryBuilder = Query.newBuilder();
+      if (Strings.isNullOrEmpty(namespace)) {
+        queryBuilder.addKindBuilder().setName("__Stat_Kind__");
+      } else {
+        queryBuilder.addKindBuilder().setName("__Stat_Ns_Kind__");
+      }
+
+      queryBuilder.setFilter(
+          makeAndFilter(
+              makeFilter("kind_name", EQUAL, makeValue(ourKind).build()).build(),
+              makeFilter("timestamp", EQUAL, makeValue(latestTimestamp).build()).build()));
+
+      RunQueryRequest request = makeRequest(queryBuilder.build(), namespace);
+
+      long now = System.currentTimeMillis();
+      RunQueryResponse response = datastore.runQuery(request);
+      logger.atFine().log(
+          "Query for per-kind statistics took %sms", System.currentTimeMillis() - now);
+
+      QueryResultBatch batch = response.getBatch();
+      if (batch.getEntityResultsCount() == 0) {
+        throw new NoSuchElementException(
+            "Datastore statistics for kind " + ourKind + " unavailable");
+      }
+      return batch.getEntityResults(0).getEntity();
+    }
+
+    /**
+     * Get the estimated size of the data returned by the given query.
+     *
+     * <p>Cloud Datastore provides no way to get a good estimate of how large the result of a query
+     * entity kind being queried, using the __Stat_Kind__ system table, assuming exactly 1 kind is
+     * specified in the query.
+     *
+     * <p>See https://cloud.google.com/datastore/docs/concepts/stats.
+     */
+    static long getEstimatedSizeBytes(Datastore datastore, Query query, @Nullable String namespace)
+        throws DatastoreException {
+      String ourKind = query.getKind(0).getName();
+      Entity entity = getLatestTableStats(ourKind, namespace, datastore);
+      return entity.getProperties().get("entity_bytes").getIntegerValue();
+    }
+
+    private static PartitionId.Builder forNamespace(@Nullable String namespace) {
+      PartitionId.Builder partitionBuilder = PartitionId.newBuilder();
+      // Namespace either being null or empty represents the default namespace.
+      // Datastore Client libraries expect users to not set the namespace proto field in
+      // either of these cases.
+      if (!Strings.isNullOrEmpty(namespace)) {
+        partitionBuilder.setNamespaceId(namespace);
+      }
+      return partitionBuilder;
+    }
+
+    /** Builds a {@link RunQueryRequest} from the {@code query} and {@code namespace}. */
+    static RunQueryRequest makeRequest(Query query, @Nullable String namespace) {
+      return RunQueryRequest.newBuilder()
+          .setQuery(query)
+          .setPartitionId(forNamespace(namespace))
+          .build();
+    }
+
+    /** Builds a {@link RunQueryRequest} from the {@code GqlQuery} and {@code namespace}. */
+    private static RunQueryRequest makeRequest(GqlQuery gqlQuery, @Nullable String namespace) {
+      return RunQueryRequest.newBuilder()
+          .setGqlQuery(gqlQuery)
+          .setPartitionId(forNamespace(namespace))
+          .build();
+    }
+
+    /**
+     * A helper function to get the split queries, taking into account the optional {@code
+     * namespace}.
+     */
+    private static List<Query> splitQuery(
+        Query query,
+        @Nullable String namespace,
+        Datastore datastore,
+        QuerySplitter querySplitter,
+        int numSplits)
+        throws DatastoreException {
+      // If namespace is set, include it in the split request so splits are calculated accordingly.
+      return querySplitter.getSplits(query, forNamespace(namespace).build(), numSplits, datastore);
+    }
+
+    /**
+     * Translates a Cloud Datastore gql query string to {@link Query}.
+     *
+     * <p>Currently, the only way to translate a gql query string to a Query is to run the query
+     * against Cloud Datastore and extract the {@code Query} from the response. To prevent reading
+     * any data, we set the {@code LIMIT} to 0 but if the gql query already has a limit set, we
+     * catch the exception with {@code INVALID_ARGUMENT} error code and retry the translation
+     * without the zero limit.
+     *
+     * <p>Note: This may result in reading actual data from Cloud Datastore but the service has a
+     * cap on the number of entities returned for a single rpc request, so this should not be a
+     * problem in practice.
+     */
+    private static Query translateGqlQueryWithLimitCheck(
+        String gql, Datastore datastore, String namespace) throws DatastoreException {
+      String gqlQueryWithZeroLimit = gql + " LIMIT 0";
+      try {
+        Query translatedQuery = translateGqlQuery(gqlQueryWithZeroLimit, datastore, namespace);
+        // Clear the limit that we set.
+        return translatedQuery.toBuilder().clearLimit().build();
+      } catch (DatastoreException e) {
+        // Note: There is no specific error code or message to detect if the query already has a
+        // limit, so we just check for INVALID_ARGUMENT and assume that that the query might have
+        // a limit already set.
+        if (e.getCode() == Code.INVALID_ARGUMENT) {
+          logger.atWarning().log(
+              "Failed to translate Gql query '%s': %s", gqlQueryWithZeroLimit, e.getMessage());
+          logger.atWarning().log(
+              "User query might have a limit already set, so trying without zero limit");
+          // Retry without the zero limit.
+          return translateGqlQuery(gql, datastore, namespace);
+        } else {
+          throw e;
+        }
+      }
+    }
+
+    /** Translates a gql query string to {@link Query}. */
+    private static Query translateGqlQuery(String gql, Datastore datastore, String namespace)
+        throws DatastoreException {
+      logger.atInfo().log("Translating gql %s", gql);
+      GqlQuery gqlQuery = GqlQuery.newBuilder().setQueryString(gql).setAllowLiterals(true).build();
+      RunQueryRequest req = makeRequest(gqlQuery, namespace);
+      return datastore.runQuery(req).getQuery();
+    }
+
+    /**
+     * Returns a new {@link MultiRead} that reads from the Cloud Datastore for the specified
+     * project.
+     */
+    public MultiRead withProjectId(String projectId) {
+      checkArgument(projectId != null, "projectId can not be null");
+      return toBuilder().setProjectId(projectId).build();
+    }
+
+    /** Returns a new {@link MultiRead} that reads from the given namespace. */
+    public MultiRead withNamespace(String namespace) {
+      return toBuilder().setNamespace(namespace).build();
+    }
+
+    /**
+     * Returns a new {@link MultiRead} that reads by splitting the given {@code query} into {@code
+     * numQuerySplits}.
+     *
+     * <p>The semantics for the query splitting is defined below:
+     *
+     * <ul>
+     *   <li>Any value less than or equal to 0 will be ignored, and the number of splits will be
+     *       chosen dynamically at runtime based on the query data size.
+     *   <li>Any value greater than {@link MultiRead#NUM_QUERY_SPLITS_MAX} will be capped at {@code
+     *       NUM_QUERY_SPLITS_MAX}.
+     *   <li>If the {@code query} has a user limit set, then {@code numQuerySplits} will be ignored
+     *       and no split will be performed.
+     *   <li>Under certain cases Cloud Datastore is unable to split query to the requested number of
+     *       splits. In such cases we just use whatever the Cloud Datastore returns.
+     * </ul>
+     */
+    public MultiRead withNumQuerySplits(int numQuerySplits) {
+      return toBuilder()
+          .setNumQuerySplits(Math.min(Math.max(numQuerySplits, 0), NUM_QUERY_SPLITS_MAX))
+          .build();
+    }
+
+    /**
+     * Returns a new {@link MultiRead} that reads from a Datastore Emulator running at the given
+     * localhost address.
+     */
+    public MultiRead withLocalhost(String localhost) {
+      return toBuilder().setLocalhost(localhost).build();
+    }
+
+    /** Returns Number of entities available for reading. */
+    public long getNumEntities(
+        PipelineOptions options, String ourKind, @Nullable String namespace) {
+      try {
+        V1Options v1Options = V1Options.from(getProjectId(), getNamespace(), getLocalhost());
+        V1DatastoreFactory datastoreFactory = new V1DatastoreFactory();
+        Datastore datastore =
+            datastoreFactory.getDatastore(
+                options, v1Options.getProjectId(), v1Options.getLocalhost());
+
+        Entity entity = getLatestTableStats(ourKind, namespace, datastore);
+        return entity.getProperties().get("count").getIntegerValue();
+      } catch (Exception e) {
+        return -1;
+      }
+    }
+
+    @Override
+    public PCollection<Entity> expand(PCollection<String> gqlQueries) {
+      checkArgument(getProjectId() != null, "projectId cannot be null");
+
+      V1Options v1Options = V1Options.from(getProjectId(), getNamespace(), getLocalhost());
+
+      /*
+       * This composite transform involves the following steps:
+       *   1. Apply a {@link ParDo} that translates each query in {@code gqlQueries} into a {@code
+       *   query}.
+       *
+       *   2. A {@link ParDo} splits the resulting query into {@code numQuerySplits} and
+       *   assign each split query a unique {@code Integer} as the key. The resulting output is
+       *   of the type {@code PCollection<KV<Integer, Query>>}.
+       *
+       *   If the value of {@code numQuerySplits} is less than or equal to 0, then the number of
+       *   splits will be computed dynamically based on the size of the data for the {@code query}.
+       *
+       *   3. The resulting {@code PCollection} is sharded using a {@link GroupByKey} operation. The
+       *   queries are extracted from they {@code KV<Integer, Iterable<Query>>} and flattened to
+       *   output a {@code PCollection<Query>}.
+       *
+       *   4. In the third step, a {@code ParDo} reads entities for each query and outputs
+       *   a {@code PCollection<Entity>}.
+       */
+
+      PCollection<Query> inputQuery =
+          gqlQueries.apply(ParDo.of(new GqlQueryTranslateFn(v1Options)));
+
+      return inputQuery
+          .apply("Split", ParDo.of(new SplitQueryFn(v1Options, getNumQuerySplits())))
+          .apply("Reshuffle", Reshuffle.viaRandomKey())
+          .apply("Read", ParDo.of(new ReadFn(v1Options)));
+    }
+
+    @Override
+    public void populateDisplayData(DisplayData.Builder builder) {
+      super.populateDisplayData(builder);
+      builder
+          .addIfNotNull(DisplayData.item("projectId", getProjectId()).withLabel("ProjectId"))
+          .addIfNotNull(DisplayData.item("namespace", getNamespace()).withLabel("Namespace"));
+    }
+
+    private static class V1Options implements HasDisplayData, Serializable {
+      private final String project;
+      private final @Nullable String namespace;
+      private final @Nullable String localhost;
+
+      private V1Options(String project, @Nullable String namespace, @Nullable String localhost) {
+        this.project = project;
+        this.namespace = namespace;
+        this.localhost = localhost;
+      }
+
+      public static V1Options from(
+          String projectId, @Nullable String namespace, @Nullable String localhost) {
+        return new V1Options(projectId, namespace, localhost);
+      }
+
+      public String getProjectId() {
+        return project;
+      }
+
+      public @Nullable String getNamespace() {
+        return namespace;
+      }
+
+      public @Nullable String getLocalhost() {
+        return localhost;
+      }
+
+      @Override
+      public void populateDisplayData(DisplayData.Builder builder) {
+        builder
+            .addIfNotNull(DisplayData.item("projectId", getProjectId()).withLabel("ProjectId"))
+            .addIfNotNull(DisplayData.item("namespace", getNamespace()).withLabel("Namespace"));
+      }
+    }
+
+    /** A DoFn that translates a Cloud Datastore gql query string to {@code Query}. */
+    static class GqlQueryTranslateFn extends DoFn<String, Query> {
+      private final V1Options v1Options;
+      private transient Datastore datastore;
+      private final V1DatastoreFactory datastoreFactory;
+
+      GqlQueryTranslateFn(V1Options options) {
+        this(options, new V1DatastoreFactory());
+      }
+
+      GqlQueryTranslateFn(V1Options options, V1DatastoreFactory datastoreFactory) {
+        this.v1Options = options;
+        this.datastoreFactory = datastoreFactory;
+      }
+
+      @StartBundle
+      public void startBundle(StartBundleContext c) throws Exception {
+        datastore =
+            datastoreFactory.getDatastore(
+                c.getPipelineOptions(), v1Options.getProjectId(), v1Options.getLocalhost());
+      }
+
+      @ProcessElement
+      public void processElement(ProcessContext c) throws Exception {
+        String gqlQuery = c.element();
+        logger.atInfo().log("User query: '%s'", gqlQuery);
+        Query query =
+            translateGqlQueryWithLimitCheck(gqlQuery, datastore, v1Options.getNamespace());
+        logger.atInfo().log("User gql query translated to Query(%s)", query);
+        c.output(query);
+      }
+    }
+
+    /**
+     * A {@link DoFn} that splits a given query into multiple sub-queries, assigns them unique keys
+     * and outputs them as {@link KV}.
+     */
+    private static class SplitQueryFn extends DoFn<Query, Query> {
+      private final V1Options options;
+      // number of splits to make for a given query
+      private final int numSplits;
+
+      private final V1DatastoreFactory datastoreFactory;
+      // Datastore client
+      private transient Datastore datastore;
+      // Query splitter
+      private transient QuerySplitter querySplitter;
+
+      public SplitQueryFn(V1Options options, int numSplits) {
+        this(options, numSplits, new V1DatastoreFactory());
+      }
+
+      private SplitQueryFn(V1Options options, int numSplits, V1DatastoreFactory datastoreFactory) {
+        this.options = options;
+        this.numSplits = numSplits;
+        this.datastoreFactory = datastoreFactory;
+      }
+
+      @StartBundle
+      public void startBundle(StartBundleContext c) throws Exception {
+        datastore =
+            datastoreFactory.getDatastore(
+                c.getPipelineOptions(), options.getProjectId(), options.getLocalhost());
+        querySplitter = datastoreFactory.getQuerySplitter();
+      }
+
+      @ProcessElement
+      public void processElement(ProcessContext c) throws Exception {
+        Query query = c.element();
+
+        // If query has a user set limit, then do not split.
+        if (query.hasLimit()) {
+          c.output(query);
+          return;
+        }
+
+        int estimatedNumSplits;
+        // Compute the estimated numSplits if numSplits is not specified by the user.
+        if (numSplits <= 0) {
+          estimatedNumSplits = getEstimatedNumSplits(datastore, query, options.getNamespace());
+        } else {
+          estimatedNumSplits = numSplits;
+        }
+
+        logger.atInfo().log("Splitting the query into %s splits", estimatedNumSplits);
+        List<Query> querySplits;
+        try {
+          querySplits =
+              splitQuery(
+                  query, options.getNamespace(), datastore, querySplitter, estimatedNumSplits);
+        } catch (Exception e) {
+          logger.atWarning().log("Unable to parallelize the given query: %s", query, e);
+          querySplits = ImmutableList.of(query);
+        }
+
+        // assign unique keys to query splits.
+        for (Query subquery : querySplits) {
+          c.output(subquery);
+        }
+      }
+
+      @Override
+      public void populateDisplayData(DisplayData.Builder builder) {
+        super.populateDisplayData(builder);
+        builder.include("options", options);
+        if (numSplits > 0) {
+          builder.add(
+              DisplayData.item("numQuerySplits", numSplits)
+                  .withLabel("Requested number of Query splits"));
+        }
+      }
+    }
+
+    /** A {@link DoFn} that reads entities from Cloud Datastore for each query. */
+    private static class ReadFn extends DoFn<Query, Entity> {
+      private final V1Options options;
+      private final V1DatastoreFactory datastoreFactory;
+      // Datastore client
+      private transient Datastore datastore;
+      private final Counter rpcErrors = Metrics.counter(ReadFn.class, "datastoreRpcErrors");
+      private final Counter rpcSuccesses = Metrics.counter(ReadFn.class, "datastoreRpcSuccesses");
+      private static final int MAX_RETRIES = 5;
+      private static final FluentBackoff RUNQUERY_BACKOFF =
+          FluentBackoff.DEFAULT
+              .withMaxRetries(MAX_RETRIES)
+              .withInitialBackoff(Duration.standardSeconds(5));
+
+      public ReadFn(V1Options options) {
+        this(options, new V1DatastoreFactory());
+      }
+
+      private ReadFn(V1Options options, V1DatastoreFactory datastoreFactory) {
+        this.options = options;
+        this.datastoreFactory = datastoreFactory;
+      }
+
+      @StartBundle
+      public void startBundle(StartBundleContext c) throws Exception {
+        datastore =
+            datastoreFactory.getDatastore(
+                c.getPipelineOptions(), options.getProjectId(), options.getLocalhost());
+      }
+
+      private RunQueryResponse runQueryWithRetries(RunQueryRequest request) throws Exception {
+        Sleeper sleeper = Sleeper.DEFAULT;
+        BackOff backoff = RUNQUERY_BACKOFF.backoff();
+        while (true) {
+          try {
+            RunQueryResponse response = datastore.runQuery(request);
+            rpcSuccesses.inc();
+            return response;
+          } catch (DatastoreException exception) {
+            rpcErrors.inc();
+
+            if (NON_RETRYABLE_ERRORS.contains(exception.getCode())) {
+              throw exception;
+            }
+            if (!BackOffUtils.next(sleeper, backoff)) {
+              logger.atSevere().log("Aborting after %s retries.", MAX_RETRIES);
+              throw exception;
+            }
+          }
+        }
+      }
+
+      /** Read and output entities for the given query. */
+      @ProcessElement
+      public void processElement(ProcessContext context) throws Exception {
+        Query query = context.element();
+        String namespace = options.getNamespace();
+        int userLimit = query.hasLimit() ? query.getLimit().getValue() : Integer.MAX_VALUE;
+
+        boolean moreResults = true;
+        QueryResultBatch currentBatch = null;
+
+        while (moreResults) {
+          Query.Builder queryBuilder = query.toBuilder();
+          queryBuilder.setLimit(
+              Int32Value.newBuilder().setValue(Math.min(userLimit, QUERY_BATCH_LIMIT)));
+
+          if (currentBatch != null && !currentBatch.getEndCursor().isEmpty()) {
+            queryBuilder.setStartCursor(currentBatch.getEndCursor());
+          }
+
+          RunQueryRequest request = makeRequest(queryBuilder.build(), namespace);
+          RunQueryResponse response = runQueryWithRetries(request);
+
+          currentBatch = response.getBatch();
+
+          // MORE_RESULTS_AFTER_LIMIT is not implemented yet:
+          // https://groups.google.com/forum/#!topic/gcd-discuss/iNs6M1jA2Vw, so
+          // use result count to determine if more results might exist.
+          int numFetch = currentBatch.getEntityResultsCount();
+          if (query.hasLimit()) {
+            verify(
+                userLimit >= numFetch,
+                "Expected userLimit %s >= numFetch %s, because query limit %s must be <= userLimit",
+                userLimit,
+                numFetch,
+                query.getLimit());
+            userLimit -= numFetch;
+          }
+
+          // output all the entities from the current batch.
+          for (EntityResult entityResult : currentBatch.getEntityResultsList()) {
+            context.output(entityResult.getEntity());
+          }
+
+          // Check if we have more entities to be read.
+          moreResults =
+              // User-limit does not exist (so userLimit == MAX_VALUE) and/or has not been satisfied
+              (userLimit > 0)
+                  // All indications from the API are that there are/may be more results.
+                  && ((numFetch == QUERY_BATCH_LIMIT)
+                      || (currentBatch.getMoreResults() == NOT_FINISHED));
+        }
+      }
+
+      @Override
+      public void populateDisplayData(DisplayData.Builder builder) {
+        super.populateDisplayData(builder);
+        builder.include("options", options);
+      }
+    }
+  }
+
+  /**
+   * A wrapper factory class for Cloud Datastore singleton classes {@link DatastoreFactory} and
+   * {@link QuerySplitter}
+   *
+   * <p>{@link DatastoreFactory} and {@link QuerySplitter} are not java serializable, hence wrapping
+   * them under this class, which implements {@link Serializable}.
+   */
+  private static class V1DatastoreFactory implements Serializable {
+
+    /** Builds a Cloud Datastore client for the given pipeline options and project. */
+    public Datastore getDatastore(PipelineOptions pipelineOptions, String projectId) {
+      return getDatastore(pipelineOptions, projectId, null);
+    }
+
+    /**
+     * Builds a Cloud Datastore client for the given pipeline options, project and an optional
+     * locahost.
+     */
+    public Datastore getDatastore(
+        PipelineOptions pipelineOptions, String projectId, @Nullable String localhost) {
+      Credentials credential = pipelineOptions.as(GcpOptions.class).getGcpCredential();
+      HttpRequestInitializer initializer;
+      if (credential != null) {
+        initializer =
+            new ChainingHttpRequestInitializer(
+                new HttpCredentialsAdapter(credential), new RetryHttpRequestInitializer());
+      } else {
+        initializer = new RetryHttpRequestInitializer();
+      }
+
+      DatastoreOptions.Builder builder =
+          new DatastoreOptions.Builder().projectId(projectId).initializer(initializer);
+
+      if (localhost != null) {
+        builder.localHost(localhost);
+      } else {
+        builder.host("batch-datastore.googleapis.com");
+      }
+
+      return DatastoreFactory.get().create(builder.build());
+    }
+
+    /** Builds a Cloud Datastore {@link QuerySplitter}. */
+    public QuerySplitter getQuerySplitter() {
+      return DatastoreHelper.getQuerySplitter();
+    }
+  }
+}

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

@@ -30,6 +30,7 @@ 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;
@@ -57,6 +58,7 @@ public class BeamJpaModule {
 
   @Nullable private final String sqlAccessInfoFile;
   @Nullable private final String cloudKmsProjectId;
+  @Nullable private final TransactionIsolationLevel isolationOverride;
 
   /**
    * Constructs a new instance of {@link BeamJpaModule}.
@@ -73,10 +75,20 @@ public class BeamJpaModule {
    *     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) {
+  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). */
@@ -154,6 +166,13 @@ public class BeamJpaModule {
     return "nomulus-tool-keyring";
   }
 
+  @Provides
+  @Config("beamIsolationOverride")
+  @Nullable
+  TransactionIsolationLevel providesIsolationOverride() {
+    return isolationOverride;
+  }
+
   @Provides
   @Config("beamHibernateHikariMaximumPoolSize")
   static int getBeamHibernateHikariMaximumPoolSize() {

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

@@ -26,17 +26,15 @@ final class DomainBaseUtil {
   private DomainBaseUtil() {}
 
   /**
-   * Removes {@link google.registry.model.billing.BillingEvent.Recurring}, {@link
-   * google.registry.model.poll.PollMessage PollMessages} and {@link
-   * google.registry.model.host.HostResource name servers} from a Datastore {@link Entity} that
-   * represents an Ofy {@link google.registry.model.domain.DomainBase}. This breaks the cycle of
-   * foreign key constraints between these entity kinds, allowing {@code DomainBases} to be inserted
-   * into the SQL database. See {@link InitSqlPipeline} for a use case, where the full {@code
-   * DomainBases} are written again during the last stage of the pipeline.
+   * Removes properties that contain foreign keys from a Datastore {@link Entity} that represents an
+   * Ofy {@link google.registry.model.domain.DomainBase}. This breaks the cycle of foreign key
+   * constraints between entity kinds, allowing {@code DomainBases} to be inserted into the SQL
+   * database. See {@link InitSqlPipeline} for a use case, where the full {@code DomainBases} are
+   * written again during the last stage of the pipeline.
    *
    * <p>The returned object may be in bad state. Specifically, {@link
    * google.registry.model.eppcommon.StatusValue#INACTIVE} is not added after name servers are
-   * removed. This only impacts tests.
+   * removed. This only impacts tests that manipulate Datastore entities directly.
    *
    * <p>This operation is performed on an Datastore {@link Entity} instead of Ofy Java object
    * because Objectify requires access to a Datastore service when converting an Ofy object to a
@@ -70,6 +68,9 @@ final class DomainBaseUtil {
     domainBase.getProperties().keySet().stream()
         .filter(s -> s.startsWith("transferData."))
         .forEach(s -> clone.removeProperty(s));
+    domainBase.getProperties().keySet().stream()
+        .filter(s -> s.startsWith("gracePeriods."))
+        .forEach(s -> clone.removeProperty(s));
     return clone;
   }
 }

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

@@ -34,6 +34,7 @@ import google.registry.model.registrar.Registrar;
 import google.registry.model.registrar.RegistrarContact;
 import google.registry.model.registry.Registry;
 import google.registry.model.reporting.HistoryEntry;
+import google.registry.persistence.PersistenceModule.TransactionIsolationLevel;
 import google.registry.persistence.transaction.JpaTransactionManager;
 import java.io.Serializable;
 import java.util.Collection;
@@ -77,14 +78,22 @@ import org.joda.time.DateTime;
  *       HistoryEntry}.
  *   <li>{@link BillingEvent.OneTime}: references {@code Registrar}, {@code DomainBase}, {@code
  *       BillingEvent.Recurring}, {@code HistoryEntry} and {@code AllocationToken}.
- *   <li>{@link BillingEvent.Modification}: SQL model TBD. Will reference {@code Registrar}, {@code
- *       DomainBase} and {@code BillingEvent.OneTime}.
  *   <li>{@link BillingEvent.Cancellation}: references {@code Registrar}, {@code DomainBase}, {@code
  *       BillingEvent.Recurring}, {@code BillingEvent.OneTime}, and {@code HistoryEntry}.
  *   <li>{@link PollMessage}: references {@code Registrar}, {@code DomainBase}, {@code
  *       ContactResource}, {@code HostResource}, and {@code HistoryEntry}.
  *   <li>{@link DomainBase}, original copy from Datastore.
  * </ol>
+ *
+ * <p>This pipeline expects that the source Datastore has at least one entity in each of the types
+ * above. This assumption allows us to construct a simpler pipeline graph that can be visually
+ * examined, and is true in all intended use cases. However, tests must not violate this assumption
+ * when setting up data, otherwise they may run into foreign key constraint violations. The reason
+ * is that this pipeline uses the {@link Wait} transform to order the persistence by entity type.
+ * However, the wait is skipped if the target type has no data, resulting in subsequent entity types
+ * starting prematurely. E.g., if a Datastore has no {@code RegistrarContact} entities, the pipeline
+ * may start writing {@code DomainBase} entities before all {@code Registry}, {@code Registrar} and
+ * {@code ContactResource} entities have been persisted.
  */
 public class InitSqlPipeline implements Serializable {
 
@@ -93,24 +102,23 @@ public class InitSqlPipeline implements Serializable {
    * DomainBase}.
    */
   private static final ImmutableList<Class<?>> PHASE_ONE_ORDERED =
-      ImmutableList.of(Registry.class, Registrar.class, ContactResource.class);
+      ImmutableList.of(
+          Registry.class, Registrar.class, ContactResource.class, RegistrarContact.class);
 
   /**
    * Datastore kinds to be written to the SQL database after the cleansed version of {@link
    * DomainBase}.
-   *
-   * <p>The following entities are missing from the list:
-   *
-   * <ul>
-   *   <li>Those not modeled in JPA yet, e.g., {@code BillingEvent.Modification}.
-   *   <li>Those waiting for sanitation, e.g., {@code HistoryEntry}, which would have duplicate keys
-   *       after converting to SQL model.
-   *   <li>Those that have foreign key constraints on the above.
-   * </ul>
    */
-  // TODO(weiminyu): add more entities when available.
   private static final ImmutableList<Class<?>> PHASE_TWO_ORDERED =
-      ImmutableList.of(HostResource.class);
+      ImmutableList.of(
+          HostResource.class,
+          HistoryEntry.class,
+          AllocationToken.class,
+          BillingEvent.Recurring.class,
+          BillingEvent.OneTime.class,
+          BillingEvent.Cancellation.class,
+          PollMessage.class,
+          DomainBase.class);
 
   private final InitSqlPipelineOptions options;
 
@@ -226,7 +234,11 @@ public class InitSqlPipeline implements Serializable {
             transformId,
             options.getMaxConcurrentSqlWriters(),
             options.getSqlWriteBatchSize(),
-            new JpaSupplierFactory(credentialFileUrl, options.getCloudKmsProjectId(), jpaGetter)));
+            new JpaSupplierFactory(
+                credentialFileUrl,
+                options.getCloudKmsProjectId(),
+                jpaGetter,
+                TransactionIsolationLevel.TRANSACTION_READ_UNCOMMITTED)));
   }
 
   private static ImmutableList<String> toKindStrings(Collection<Class<?>> entityClasses) {

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

@@ -16,6 +16,7 @@ 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;
@@ -28,21 +29,32 @@ public class JpaSupplierFactory implements SerializableSupplier<JpaTransactionMa
   @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))
+            .beamJpaModule(
+                new BeamJpaModule(credentialFileUrl, cloudKmsProjectId, isolationLevelOverride))
             .build());
   }
 }

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

@@ -17,11 +17,9 @@ package google.registry.beam.initsql;
 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.base.Throwables.throwIfUnchecked;
 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.JpaRetries.isFailedTxnRetriable;
 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;
@@ -38,27 +36,34 @@ import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 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.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;
-import google.registry.util.SystemSleeper;
 import java.io.Serializable;
 import java.util.Collection;
 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;
 import org.apache.beam.sdk.io.Compression;
 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;
@@ -78,7 +83,6 @@ import org.apache.beam.sdk.values.TupleTag;
 import org.apache.beam.sdk.values.TupleTagList;
 import org.apache.beam.sdk.values.TypeDescriptor;
 import org.joda.time.DateTime;
-import org.joda.time.Duration;
 
 /**
  * {@link PTransform Pipeline transforms} used in pipelines that load from both Datastore export
@@ -288,7 +292,7 @@ public final class Transforms {
         maxWriters,
         batchSize,
         jpaSupplier,
-        (e) -> ofy().toPojo(e.getEntity().get()),
+        Transforms::convertVersionedEntityToSqlEntity,
         TypeDescriptor.of(VersionedEntity.class));
   }
 
@@ -329,11 +333,50 @@ public final class Transforms {
             .apply("Batch output by shard " + transformId, GroupIntoBatches.ofSize(batchSize))
             .apply(
                 "Write in batch for " + transformId,
-                ParDo.of(new SqlBatchWriter<T>(jpaSupplier, jpaConverter)));
+                ParDo.of(new SqlBatchWriter<T>(transformId, jpaSupplier, jpaConverter)));
       }
     };
   }
 
+  private static Key toOfyKey(Object ofyEntity) {
+    return Key.create(ofyEntity);
+  }
+
+  private static boolean isMigratable(Entity entity) {
+    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
+      // be deserialized.
+      return !Objects.equals(entity.getProperty("type"), "DOMAIN_APPLICATION_CREATE");
+    }
+    return true;
+  }
+
+  private static SqlEntity toSqlEntity(Object ofyEntity) {
+    if (ofyEntity instanceof HistoryEntry) {
+      HistoryEntry ofyHistory = (HistoryEntry) ofyEntity;
+      return (SqlEntity) ofyHistory.toChildHistoryEntity();
+    }
+    return ((DatastoreAndSqlEntity) ofyEntity).toSqlEntity().get();
+  }
+
+  /**
+   * Converts a {@link VersionedEntity} to an JPA entity for persistence.
+   *
+   * @return An object to be persisted to SQL, or null if the input is not to be migrated. (Not
+   *     using Optional in return because as a one-use method, we do not want to invest the effort
+   *     to make Optional work with BEAM)
+   */
+  @Nullable
+  private static Object convertVersionedEntityToSqlEntity(VersionedEntity dsEntity) {
+    return dsEntity
+        .getEntity()
+        .filter(Transforms::isMigratable)
+        .map(e -> ofy().toPojo(e))
+        .map(Transforms::toSqlEntity)
+        .orElse(null);
+  }
+
   /** Interface for serializable {@link Supplier suppliers}. */
   public interface SerializableSupplier<T> extends Supplier<T>, Serializable {}
 
@@ -428,22 +471,22 @@ public final class Transforms {
     private static int instanceCount = 0;
     private static JpaTransactionManager originalJpa;
 
+    private Counter counter;
+
     private final SerializableSupplier<JpaTransactionManager> jpaSupplier;
     private final SerializableFunction<T, Object> jpaConverter;
 
-    private transient SystemSleeper sleeper;
-
     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() {
-      sleeper = new SystemSleeper();
-
       try (AppEngineEnvironment env = new AppEngineEnvironment()) {
         ObjectifyService.initOfy();
       }
@@ -474,31 +517,29 @@ public final class Transforms {
         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());
-        retry(() -> jpaTm().transact(() -> jpaTm().putAll(ofyEntities)));
+        try {
+          jpaTm().transact(() -> jpaTm().putAll(ofyEntities));
+          counter.inc(ofyEntities.size());
+        } catch (RuntimeException e) {
+          processSingly(ofyEntities);
+        }
       }
     }
 
-    // TODO(b/160632289): Enhance Retrier and use it here.
-    private void retry(Runnable runnable) {
-      int maxAttempts = 5;
-      int initialDelayMillis = 100;
-      double jitterRatio = 0.2;
-
-      for (int attempt = 0; attempt < maxAttempts; attempt++) {
+    /**
+     * 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 {
-          runnable.run();
-          return;
-        } catch (Throwable throwable) {
-          if (!isFailedTxnRetriable(throwable)) {
-            throwIfUnchecked(throwable);
-            throw new RuntimeException(throwable);
-          }
-          int sleepMillis = (1 << attempt) * initialDelayMillis;
-          int jitter =
-              ThreadLocalRandom.current().nextInt((int) (sleepMillis * jitterRatio))
-                  - (int) (sleepMillis * jitterRatio / 2);
-          sleeper.sleepUninterruptibly(Duration.millis(sleepMillis + jitter));
+          jpaTm().transact(() -> jpaTm().put(ofyEntity));
+          counter.inc();
+        } catch (RuntimeException e) {
+          throw new RuntimeException(toOfyKey(ofyEntity).toString(), e);
         }
       }
     }

core/src/main/java/google/registry/env/common/default/WEB-INF/datastore-indexes.xml

@@ -25,11 +25,6 @@
     <property name="tld" direction="asc"/>
     <property name="creationTime" direction="desc"/>
   </datastore-index>
-  <!-- For finding non-autorenewing domains to be deleted. -->
-  <datastore-index kind="DomainBase" ancestor="false" source="manual">
-    <property name="autorenewEndTime" direction="asc"/>
-    <property name="deletionTime" direction="asc"/>
-  </datastore-index>
   <!-- For finding host resources by registrar. -->
   <datastore-index kind="HostResource" ancestor="false" source="manual">
     <property name="currentSponsorClientId" direction="asc"/>

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

@@ -22,5 +22,6 @@ public enum EppRequestSource {
   TLS,
   TOOL,
   CHECK_API,
-  UNIT_TEST
+  UNIT_TEST,
+  BACKEND
 }

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

@@ -104,11 +104,14 @@ public final class ExtensionManager {
         clientId, flowClass.getSimpleName(), undeclaredUris);
   }
 
+  private static final ImmutableSet<EppRequestSource> ALLOWED_METADATA_EPP_REQUEST_SOURCES =
+      ImmutableSet.of(EppRequestSource.TOOL, EppRequestSource.BACKEND);
+
   private void checkForRestrictedExtensions(
       ImmutableSet<Class<? extends CommandExtension>> suppliedExtensions)
       throws OnlyToolCanPassMetadataException, UnauthorizedForSuperuserExtensionException {
     if (suppliedExtensions.contains(MetadataExtension.class)
-        && !eppRequestSource.equals(EppRequestSource.TOOL)) {
+        && !ALLOWED_METADATA_EPP_REQUEST_SOURCES.contains(eppRequestSource)) {
       throw new OnlyToolCanPassMetadataException();
     }
     // Can't use suppliedExtension.contains() here because the SuperuserExtension has child classes.

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

@@ -15,8 +15,8 @@
 package google.registry.flows;
 
 import static com.google.common.collect.Sets.intersection;
+import static google.registry.model.EppResourceUtils.getLinkedDomainKeys;
 import static google.registry.model.EppResourceUtils.loadByForeignKey;
-import static google.registry.model.EppResourceUtils.queryForLinkedDomains;
 import static google.registry.model.index.ForeignKeyIndex.loadAndGetKey;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
@@ -94,14 +94,13 @@ public final class ResourceFlowUtils {
                    * actual reference then we can reliably fail. If we don't find any, we can't
                    * trust the query and need to do the full mapreduce.
                    */
-                  Iterable<Key<DomainBase>> keys =
-                      queryForLinkedDomains(fki.getResourceKey().getOfyKey(), now)
-                          .limit(FAILFAST_CHECK_COUNT)
-                          .keys();
+                  Iterable<VKey<DomainBase>> keys =
+                      getLinkedDomainKeys(fki.getResourceKey(), now, FAILFAST_CHECK_COUNT);
+
                   VKey<R> resourceVKey = fki.getResourceKey();
                   Predicate<DomainBase> predicate =
                       domain -> getPotentialReferences.apply(domain).contains(resourceVKey);
-                  return ofy().load().keys(keys).values().stream().anyMatch(predicate)
+                  return tm().loadByKeys(keys).values().stream().anyMatch(predicate)
                       ? new ResourceToDeleteIsReferencedException()
                       : null;
                 });
@@ -139,7 +138,7 @@ public final class ResourceFlowUtils {
       Class<R> clazz, String targetId, DateTime now, String clientId) throws EppException {
     VKey<R> key = loadAndGetKey(clazz, targetId, now);
     if (key != null) {
-      R resource = tm().load(key);
+      R resource = tm().loadByKey(key);
       // These are similar exceptions, but we can track them internally as log-based metrics.
       if (Objects.equals(clientId, resource.getPersistedCurrentSponsorClientId())) {
         throw new ResourceAlreadyExistsForThisClientException(targetId);

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

@@ -15,9 +15,8 @@
 package google.registry.flows;
 
 import static com.google.common.base.MoreObjects.toStringHelper;
-import static com.google.common.base.Strings.isNullOrEmpty;
 import static google.registry.request.RequestParameters.extractOptionalHeader;
-import static google.registry.request.RequestParameters.extractRequiredHeader;
+import static google.registry.util.X509Utils.loadCertificate;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
@@ -27,11 +26,18 @@ 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 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;
@@ -44,6 +50,9 @@ import javax.servlet.http.HttpServletRequest;
  * <dl>
  *   <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
@@ -56,17 +65,23 @@ public class TlsCredentials implements TransportCredentials {
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
 
   private final boolean requireSslCertificates;
-  private final String clientCertificateHash;
-  private final InetAddress clientInetAddr;
+  private final Optional<String> clientCertificateHash;
+  private final Optional<String> clientCertificate;
+  private final Optional<InetAddress> clientInetAddr;
+  private final CertificateChecker certificateChecker;
 
   @Inject
   public TlsCredentials(
       @Config("requireSslCertificates") boolean requireSslCertificates,
-      @Header("X-SSL-Certificate") String clientCertificateHash,
-      @Header("X-Forwarded-For") Optional<String> clientAddress) {
+      @Header("X-SSL-Certificate") Optional<String> clientCertificateHash,
+      @Header("X-SSL-Full-Certificate") Optional<String> clientCertificate,
+      @Header("X-Forwarded-For") Optional<String> clientAddress,
+      CertificateChecker certificateChecker) {
     this.requireSslCertificates = requireSslCertificates;
     this.clientCertificateHash = clientCertificateHash;
-    this.clientInetAddr = clientAddress.isPresent() ? parseInetAddress(clientAddress.get()) : null;
+    this.clientCertificate = clientCertificate;
+    this.clientInetAddr = clientAddress.map(TlsCredentials::parseInetAddress);
+    this.certificateChecker = certificateChecker;
   }
 
   static InetAddress parseInetAddress(String asciiAddr) {
@@ -93,14 +108,18 @@ public class TlsCredentials implements TransportCredentials {
     ImmutableList<CidrAddressBlock> ipAddressAllowList = registrar.getIpAddressAllowList();
     if (ipAddressAllowList.isEmpty()) {
       logger.atInfo().log(
-          "Skipping IP allow list check because %s doesn't have an IP allow list",
+          "Skipping IP allow list check because %s doesn't have an IP allow list.",
           registrar.getClientId());
       return;
     }
-    for (CidrAddressBlock cidrAddressBlock : ipAddressAllowList) {
-      if (cidrAddressBlock.contains(clientInetAddr)) {
-        // IP address is in allow list; return early.
-        return;
+    // In the rare unexpected case that the client inet address wasn't passed along at all, then
+    // by default deny access.
+    if (clientInetAddr.isPresent()) {
+      for (CidrAddressBlock cidrAddressBlock : ipAddressAllowList) {
+        if (cidrAddressBlock.contains(clientInetAddr.get())) {
+          // IP address is in allow list; return early.
+          return;
+        }
       }
     }
     logger.atInfo().log(
@@ -113,13 +132,82 @@ public class TlsCredentials implements TransportCredentials {
   /**
    * Verifies client SSL certificate is permitted to issue commands as {@code registrar}.
    *
-   * @throws MissingRegistrarCertificateException if frontend didn't send certificate hash header
+   * @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 {
-    if (isNullOrEmpty(registrar.getClientCertificateHash())
-        && isNullOrEmpty(registrar.getFailoverClientCertificateHash())) {
+    // 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) {
+          // throw exception in unit tests and Sandbox
+          if (RegistryEnvironment.get().equals(RegistryEnvironment.UNITTEST)
+              || RegistryEnvironment.get().equals(RegistryEnvironment.SANDBOX)) {
+            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 {
+    // Check the certificate hash as a failover
+    // TODO(sarahbot): Remove hash checks once certificate checks are working.
+    if (!registrar.getClientCertificateHash().isPresent()
+        && !registrar.getFailoverClientCertificateHash().isPresent()) {
       if (requireSslCertificates) {
         throw new RegistrarCertificateNotConfiguredException();
       } else {
@@ -128,14 +216,17 @@ public class TlsCredentials implements TransportCredentials {
         return;
       }
     }
-    if (isNullOrEmpty(clientCertificateHash)) {
-      logger.atInfo().log("Request did not include X-SSL-Certificate");
+    // Check that the request included the certificate hash
+    if (!clientCertificateHash.isPresent()) {
+      logger.atInfo().log(
+          "Request from registrar %s did not include X-SSL-Certificate.", registrar.getClientId());
       throw new MissingRegistrarCertificateException();
     }
+    // Check if the certificate hash is equal to the one on file for the registrar.
     if (!clientCertificateHash.equals(registrar.getClientCertificateHash())
         && !clientCertificateHash.equals(registrar.getFailoverClientCertificateHash())) {
       logger.atWarning().log(
-          "bad certificate hash (%s) for %s, wanted either %s or %s",
+          "Non-matching certificate hash (%s) for %s, wanted either %s or %s.",
           clientCertificateHash,
           registrar.getClientId(),
           registrar.getClientCertificateHash(),
@@ -151,24 +242,55 @@ 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("clientCertificateHash", clientCertificateHash)
-        .add("clientAddress", clientInetAddr)
+        .add("clientCertificate", clientCertificate.orElse(null))
+        .add("clientCertificateHash", clientCertificateHash.orElse(null))
+        .add("clientAddress", clientInetAddr.orElse(null))
         .toString();
   }
 
   /** Registrar certificate does not match stored certificate. */
   public static class BadRegistrarCertificateException extends AuthenticationErrorException {
-    public BadRegistrarCertificateException() {
+    BadRegistrarCertificateException() {
       super("Registrar certificate does not match stored certificate");
     }
   }
 
+  /** Registrar certificate contains the following security violations: ... */
+  public static class CertificateContainsSecurityViolationsException
+      extends AuthenticationErrorException {
+    InsecureCertificateException exception;
+
+    CertificateContainsSecurityViolationsException(InsecureCertificateException exception) {
+      super(
+          String.format(
+              "Registrar certificate contains the following security violations:\n%s",
+              exception.getMessage()));
+      this.exception = exception;
+    }
+  }
+
   /** Registrar certificate not present. */
   public static class MissingRegistrarCertificateException extends AuthenticationErrorException {
-    public MissingRegistrarCertificateException() {
+    MissingRegistrarCertificateException() {
       super("Registrar certificate not present");
     }
   }
@@ -176,14 +298,14 @@ public class TlsCredentials implements TransportCredentials {
   /** Registrar certificate is not configured. */
   public static class RegistrarCertificateNotConfiguredException
       extends AuthenticationErrorException {
-    public RegistrarCertificateNotConfiguredException() {
+    RegistrarCertificateNotConfiguredException() {
       super("Registrar certificate is not configured");
     }
   }
 
   /** Registrar IP address is not in stored allow list. */
   public static class BadRegistrarIpAddressException extends AuthenticationErrorException {
-    public BadRegistrarIpAddressException() {
+    BadRegistrarIpAddressException() {
       super("Registrar IP address is not in stored allow list");
     }
   }
@@ -191,10 +313,21 @@ public class TlsCredentials implements TransportCredentials {
   /** Dagger module for the EPP TLS endpoint. */
   @Module
   public static final class EppTlsModule {
+
     @Provides
     @Header("X-SSL-Certificate")
-    static String provideClientCertificateHash(HttpServletRequest req) {
-      return extractRequiredHeader(req, "X-SSL-Certificate");
+    static Optional<String> provideClientCertificateHash(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, "X-SSL-Certificate");
+    }
+
+    @Provides
+    @Header("X-SSL-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, "X-SSL-Full-Certificate");
     }
 
     @Provides

core/src/main/java/google/registry/flows/certs/CertificateChecker.java

@@ -89,14 +89,26 @@ public class CertificateChecker {
    * Checks the given certificate string for violations and throws an exception if any violations
    * exist.
    */
-  public void validateCertificate(String certificateString) {
-    ImmutableSet<CertificateViolation> violations = checkCertificate(certificateString);
+  public void validateCertificate(String certificateString) throws InsecureCertificateException {
+    handleCertViolations(checkCertificate(certificateString));
+  }
+
+  /**
+   * Checks the given certificate string for violations and throws an exception if any violations
+   * exist.
+   */
+  public void validateCertificate(X509Certificate certificate) throws InsecureCertificateException {
+    handleCertViolations(checkCertificate(certificate));
+  }
+
+  private void handleCertViolations(ImmutableSet<CertificateViolation> violations)
+      throws InsecureCertificateException {
     if (!violations.isEmpty()) {
       String displayMessages =
           violations.stream()
               .map(violation -> getViolationDisplayMessage(violation))
               .collect(Collectors.joining("\n"));
-      throw new IllegalArgumentException(displayMessages);
+      throw new InsecureCertificateException(violations, displayMessages);
     }
   }
 
@@ -258,4 +270,14 @@ public class CertificateChecker {
       return certificateChecker.getViolationDisplayMessage(this);
     }
   }
+
+  /** Exception to throw when a certificate has security violations. */
+  public static class InsecureCertificateException extends Exception {
+    ImmutableSet<CertificateViolation> violations;
+
+    InsecureCertificateException(ImmutableSet<CertificateViolation> violations, String message) {
+      super(message);
+      this.violations = violations;
+    }
+  }
 }

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

@@ -19,9 +19,9 @@ import static google.registry.flows.ResourceFlowUtils.verifyResourceDoesNotExist
 import static google.registry.flows.contact.ContactFlowUtils.validateAsciiPostalInfo;
 import static google.registry.flows.contact.ContactFlowUtils.validateContactAgainstPolicy;
 import static google.registry.model.EppResourceUtils.createRepoId;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
+import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
 import google.registry.config.RegistryConfig.Config;
 import google.registry.flows.EppException;
@@ -95,11 +95,12 @@ public final class ContactCreateFlow implements TransactionalFlow {
         .setModificationTime(now)
         .setXmlBytes(null)  // We don't want to store contact details in the history entry.
         .setParent(Key.create(newContact));
-    ofy().save().entities(
-        newContact,
-        historyBuilder.build(),
-        ForeignKeyIndex.create(newContact, newContact.getDeletionTime()),
-        EppResourceIndex.create(Key.create(newContact)));
+    tm().insertAll(
+            ImmutableSet.of(
+                newContact,
+                historyBuilder.build().toChildHistoryEntity(),
+                ForeignKeyIndex.create(newContact, newContact.getDeletionTime()),
+                EppResourceIndex.create(Key.create(newContact))));
     return responseBuilder
         .setResData(ContactCreateData.create(newContact.getContactId(), now))
         .build();

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

@@ -21,7 +21,6 @@ import static google.registry.flows.ResourceFlowUtils.verifyNoDisallowedStatuses
 import static google.registry.flows.ResourceFlowUtils.verifyOptionalAuthInfo;
 import static google.registry.flows.ResourceFlowUtils.verifyResourceOwnership;
 import static google.registry.model.eppoutput.Result.Code.SUCCESS_WITH_ACTION_PENDING;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.common.collect.ImmutableSet;
@@ -101,7 +100,8 @@ public final class ContactDeleteFlow implements TransactionalFlow {
         .setType(HistoryEntry.Type.CONTACT_PENDING_DELETE)
         .setModificationTime(now)
         .setParent(Key.create(existingContact));
-    ofy().save().<Object>entities(newContact, historyBuilder.build());
+    tm().insert(historyBuilder.build().toChildHistoryEntity());
+    tm().update(newContact);
     return responseBuilder.setResultFromCode(SUCCESS_WITH_ACTION_PENDING).build();
   }
 }

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

@@ -20,7 +20,6 @@ import static google.registry.flows.ResourceFlowUtils.verifyResourceOwnership;
 import static google.registry.model.EppResourceUtils.isLinked;
 
 import com.google.common.collect.ImmutableSet;
-import com.googlecode.objectify.Key;
 import google.registry.flows.EppException;
 import google.registry.flows.ExtensionManager;
 import google.registry.flows.Flow;
@@ -77,7 +76,7 @@ public final class ContactInfoFlow implements Flow {
         clientId.equals(contact.getCurrentSponsorClientId()) || authInfo.isPresent();
     ImmutableSet.Builder<StatusValue> statusValues = new ImmutableSet.Builder<>();
     statusValues.addAll(contact.getStatusValues());
-    if (isLinked(Key.create(contact), now)) {
+    if (isLinked(contact.createVKey(), now)) {
       statusValues.add(StatusValue.LINKED);
     }
     return responseBuilder

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

@@ -22,9 +22,9 @@ import static google.registry.flows.ResourceFlowUtils.verifyResourceOwnership;
 import static google.registry.flows.contact.ContactFlowUtils.createGainingTransferPollMessage;
 import static google.registry.flows.contact.ContactFlowUtils.createTransferResponse;
 import static google.registry.model.ResourceTransferUtils.approvePendingTransfer;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
+import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
 import google.registry.flows.EppException;
 import google.registry.flows.ExtensionManager;
@@ -94,7 +94,8 @@ public final class ContactTransferApproveFlow implements TransactionalFlow {
     // Create a poll message for the gaining client.
     PollMessage gainingPollMessage =
         createGainingTransferPollMessage(targetId, newContact.getTransferData(), historyEntry);
-    ofy().save().<Object>entities(newContact, historyEntry, gainingPollMessage);
+    tm().insertAll(ImmutableSet.of(historyEntry.toChildHistoryEntity(), gainingPollMessage));
+    tm().update(newContact);
     // Delete the billing event and poll messages that were written in case the transfer would have
     // been implicitly server approved.
     tm().delete(existingContact.getTransferData().getServerApproveEntities());

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

@@ -22,9 +22,9 @@ import static google.registry.flows.ResourceFlowUtils.verifyTransferInitiator;
 import static google.registry.flows.contact.ContactFlowUtils.createLosingTransferPollMessage;
 import static google.registry.flows.contact.ContactFlowUtils.createTransferResponse;
 import static google.registry.model.ResourceTransferUtils.denyPendingTransfer;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
+import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
 import google.registry.flows.EppException;
 import google.registry.flows.ExtensionManager;
@@ -90,7 +90,8 @@ public final class ContactTransferCancelFlow implements TransactionalFlow {
     // Create a poll message for the losing client.
     PollMessage losingPollMessage =
         createLosingTransferPollMessage(targetId, newContact.getTransferData(), historyEntry);
-    ofy().save().<Object>entities(newContact, historyEntry, losingPollMessage);
+    tm().insertAll(ImmutableSet.of(historyEntry.toChildHistoryEntity(), losingPollMessage));
+    tm().update(newContact);
     // Delete the billing event and poll messages that were written in case the transfer would have
     // been implicitly server approved.
     tm().delete(existingContact.getTransferData().getServerApproveEntities());

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

@@ -22,9 +22,9 @@ import static google.registry.flows.ResourceFlowUtils.verifyResourceOwnership;
 import static google.registry.flows.contact.ContactFlowUtils.createGainingTransferPollMessage;
 import static google.registry.flows.contact.ContactFlowUtils.createTransferResponse;
 import static google.registry.model.ResourceTransferUtils.denyPendingTransfer;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
+import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
 import google.registry.flows.EppException;
 import google.registry.flows.ExtensionManager;
@@ -87,7 +87,8 @@ public final class ContactTransferRejectFlow implements TransactionalFlow {
         .build();
     PollMessage gainingPollMessage =
         createGainingTransferPollMessage(targetId, newContact.getTransferData(), historyEntry);
-    ofy().save().<Object>entities(newContact, historyEntry, gainingPollMessage);
+    tm().insertAll(ImmutableSet.of(historyEntry.toChildHistoryEntity(), gainingPollMessage));
+    tm().update(newContact);
     // Delete the billing event and poll messages that were written in case the transfer would have
     // been implicitly server approved.
     tm().delete(existingContact.getTransferData().getServerApproveEntities());

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

@@ -23,7 +23,6 @@ import static google.registry.flows.contact.ContactFlowUtils.createGainingTransf
 import static google.registry.flows.contact.ContactFlowUtils.createLosingTransferPollMessage;
 import static google.registry.flows.contact.ContactFlowUtils.createTransferResponse;
 import static google.registry.model.eppoutput.Result.Code.SUCCESS_WITH_ACTION_PENDING;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.common.collect.ImmutableSet;
@@ -145,12 +144,13 @@ public final class ContactTransferRequestFlow implements TransactionalFlow {
         .setTransferData(pendingTransferData)
         .addStatusValue(StatusValue.PENDING_TRANSFER)
         .build();
-    ofy().save().<Object>entities(
-        newContact,
-        historyEntry,
-        requestPollMessage,
-        serverApproveGainingPollMessage,
-        serverApproveLosingPollMessage);
+    tm().update(newContact);
+    tm().insertAll(
+            ImmutableSet.of(
+                historyEntry.toChildHistoryEntity(),
+                requestPollMessage,
+                serverApproveGainingPollMessage,
+                serverApproveLosingPollMessage));
     return responseBuilder
         .setResultFromCode(SUCCESS_WITH_ACTION_PENDING)
         .setResData(createTransferResponse(targetId, newContact.getTransferData()))

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

@@ -24,7 +24,6 @@ import static google.registry.flows.ResourceFlowUtils.verifyOptionalAuthInfo;
 import static google.registry.flows.ResourceFlowUtils.verifyResourceOwnership;
 import static google.registry.flows.contact.ContactFlowUtils.validateAsciiPostalInfo;
 import static google.registry.flows.contact.ContactFlowUtils.validateContactAgainstPolicy;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.common.collect.ImmutableSet;
@@ -151,7 +150,8 @@ public final class ContactUpdateFlow implements TransactionalFlow {
     }
     validateAsciiPostalInfo(newContact.getInternationalizedPostalInfo());
     validateContactAgainstPolicy(newContact);
-    ofy().save().<Object>entities(newContact, historyBuilder.build());
+    tm().insert(historyBuilder.build().toChildHistoryEntity());
+    tm().update(newContact);
     return responseBuilder.build();
   }
 

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

@@ -225,7 +225,7 @@ public final class DomainDeleteFlow implements TransactionalFlow {
         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.
-          BillingEvent.OneTime oneTime = tm().load(gracePeriod.getOneTimeBillingEvent());
+          BillingEvent.OneTime oneTime = tm().loadByKey(gracePeriod.getOneTimeBillingEvent());
           newExpirationTime = newExpirationTime.minusYears(oneTime.getPeriodYears());
         } else if (gracePeriod.getRecurringBillingEvent() != null) {
           // Take 1 year off the registration if in the autorenew grace period (no need to load the
@@ -372,12 +372,12 @@ public final class DomainDeleteFlow implements TransactionalFlow {
   private Money getGracePeriodCost(GracePeriod gracePeriod, DateTime now) {
     if (gracePeriod.getType() == GracePeriodStatus.AUTO_RENEW) {
       DateTime autoRenewTime =
-          tm().load(checkNotNull(gracePeriod.getRecurringBillingEvent()))
+          tm().loadByKey(checkNotNull(gracePeriod.getRecurringBillingEvent()))
               .getRecurrenceTimeOfYear()
               .getLastInstanceBeforeOrAt(now);
       return getDomainRenewCost(targetId, autoRenewTime, 1);
     }
-    return tm().load(checkNotNull(gracePeriod.getOneTimeBillingEvent())).getCost();
+    return tm().loadByKey(checkNotNull(gracePeriod.getOneTimeBillingEvent())).getCost();
   }
 
   @Nullable

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

@@ -517,7 +517,7 @@ public class DomainFlowUtils {
    */
   public static void updateAutorenewRecurrenceEndTime(DomainBase domain, DateTime newEndTime) {
     Optional<PollMessage.Autorenew> autorenewPollMessage =
-        tm().maybeLoad(domain.getAutorenewPollMessage());
+        tm().loadByKeyIfPresent(domain.getAutorenewPollMessage());
 
     // Construct an updated autorenew poll message. If the autorenew poll message no longer exists,
     // create a new one at the same id. This can happen if a transfer was requested on a domain
@@ -542,7 +542,7 @@ public class DomainFlowUtils {
       ofy().save().entity(updatedAutorenewPollMessage);
     }
 
-    Recurring recurring = tm().load(domain.getAutorenewBillingEvent());
+    Recurring recurring = tm().loadByKey(domain.getAutorenewBillingEvent());
     ofy().save().entity(recurring.asBuilder().setRecurrenceEndTime(newEndTime).build());
   }
 
@@ -1022,7 +1022,7 @@ public class DomainFlowUtils {
     for (DesignatedContact contact : contacts) {
       builder.add(
           ForeignKeyedDesignatedContact.create(
-              contact.getType(), tm().load(contact.getContactKey()).getContactId()));
+              contact.getType(), tm().loadByKey(contact.getContactKey()).getContactId()));
     }
     return builder.build();
   }

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

@@ -101,8 +101,8 @@ public final class DomainInfoFlow implements Flow {
     flowCustomLogic.afterValidation(
         AfterValidationParameters.newBuilder().setDomain(domain).build());
     // Prefetch all referenced resources. Calling values() blocks until loading is done.
-    tm().load(domain.getNameservers());
-    tm().load(domain.getReferencedContacts());
+    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,7 +110,7 @@ public final class DomainInfoFlow implements Flow {
             .setFullyQualifiedDomainName(domain.getDomainName())
             .setRepoId(domain.getRepoId())
             .setCurrentSponsorClientId(domain.getCurrentSponsorClientId())
-            .setRegistrant(tm().load(domain.getRegistrant()).getContactId());
+            .setRegistrant(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()) {

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

@@ -153,8 +153,8 @@ public class AllocationTokenFlowUtils {
       throw new InvalidAllocationTokenException();
     }
     Optional<AllocationToken> maybeTokenEntity =
-        tm().maybeLoad(VKey.create(AllocationToken.class, token));
-    if (maybeTokenEntity.isEmpty()) {
+        tm().loadByKeyIfPresent(VKey.create(AllocationToken.class, token));
+    if (!maybeTokenEntity.isPresent()) {
       throw new InvalidAllocationTokenException();
     }
     if (maybeTokenEntity.get().isRedeemed()) {

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

@@ -21,7 +21,6 @@ import static google.registry.flows.ResourceFlowUtils.verifyNoDisallowedStatuses
 import static google.registry.flows.ResourceFlowUtils.verifyResourceOwnership;
 import static google.registry.flows.host.HostFlowUtils.validateHostName;
 import static google.registry.model.eppoutput.Result.Code.SUCCESS_WITH_ACTION_PENDING;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.common.collect.ImmutableSet;
@@ -96,7 +95,7 @@ public final class HostDeleteFlow implements TransactionalFlow {
       // the client id, needs to be read off of it.
       EppResource owningResource =
           existingHost.isSubordinate()
-              ? tm().load(existingHost.getSuperordinateDomain()).cloneProjectedAtTime(now)
+              ? tm().loadByKey(existingHost.getSuperordinateDomain()).cloneProjectedAtTime(now)
               : existingHost;
       verifyResourceOwnership(clientId, owningResource);
     }
@@ -108,7 +107,8 @@ public final class HostDeleteFlow implements TransactionalFlow {
         .setType(HistoryEntry.Type.HOST_PENDING_DELETE)
         .setModificationTime(now)
         .setParent(Key.create(existingHost));
-    ofy().save().<Object>entities(newHost, historyBuilder.build());
+    tm().insert(historyBuilder.build().toChildHistoryEntity());
+    tm().update(newHost);
     return responseBuilder.setResultFromCode(SUCCESS_WITH_ACTION_PENDING).build();
   }
 }

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

@@ -21,7 +21,6 @@ import static google.registry.model.EppResourceUtils.isLinked;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.common.collect.ImmutableSet;
-import com.googlecode.objectify.Key;
 import google.registry.flows.EppException;
 import google.registry.flows.ExtensionManager;
 import google.registry.flows.Flow;
@@ -68,7 +67,7 @@ public final class HostInfoFlow implements Flow {
     HostResource host = loadAndVerifyExistence(HostResource.class, targetId, now);
     ImmutableSet.Builder<StatusValue> statusValues = new ImmutableSet.Builder<>();
     statusValues.addAll(host.getStatusValues());
-    if (isLinked(Key.create(host), now)) {
+    if (isLinked(host.createVKey(), now)) {
       statusValues.add(StatusValue.LINKED);
     }
     HostInfoData.Builder hostInfoDataBuilder = HostInfoData.newBuilder();
@@ -77,7 +76,7 @@ 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().load(host.getSuperordinateDomain()).cloneProjectedAtTime(now);
+          tm().loadByKey(host.getSuperordinateDomain()).cloneProjectedAtTime(now);
       hostInfoDataBuilder
           .setCurrentSponsorClientId(superordinateDomain.getCurrentSponsorClientId())
           .setLastTransferTime(host.computeLastTransferTime(superordinateDomain));

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

@@ -27,7 +27,6 @@ import static google.registry.flows.host.HostFlowUtils.validateHostName;
 import static google.registry.flows.host.HostFlowUtils.verifySuperordinateDomainNotInPendingDelete;
 import static google.registry.flows.host.HostFlowUtils.verifySuperordinateDomainOwnership;
 import static google.registry.model.index.ForeignKeyIndex.loadAndGetKey;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.CollectionUtils.isNullOrEmpty;
 
@@ -139,7 +138,7 @@ public final class HostUpdateFlow implements TransactionalFlow {
     String newHostName = firstNonNull(suppliedNewHostName, oldHostName);
     DomainBase oldSuperordinateDomain =
         existingHost.isSubordinate()
-            ? tm().load(existingHost.getSuperordinateDomain()).cloneProjectedAtTime(now)
+            ? tm().loadByKey(existingHost.getSuperordinateDomain()).cloneProjectedAtTime(now)
             : null;
     // Note that lookupSuperordinateDomain calls cloneProjectedAtTime on the domain for us.
     Optional<DomainBase> newSuperordinateDomain =
@@ -191,23 +190,26 @@ public final class HostUpdateFlow implements TransactionalFlow {
             .setPersistedCurrentSponsorClientId(newPersistedClientId)
             .build();
     verifyHasIpsIffIsExternal(command, existingHost, newHost);
-    ImmutableSet.Builder<ImmutableObject> entitiesToSave = new ImmutableSet.Builder<>();
-    entitiesToSave.add(newHost);
+    ImmutableSet.Builder<ImmutableObject> entitiesToInsert = new ImmutableSet.Builder<>();
+    ImmutableSet.Builder<ImmutableObject> entitiesToUpdate = new ImmutableSet.Builder<>();
+    entitiesToUpdate.add(newHost);
     // Keep the {@link ForeignKeyIndex} for this host up to date.
     if (isHostRename) {
       // Update the foreign key for the old host name and save one for the new host name.
-      entitiesToSave.add(
-          ForeignKeyIndex.create(existingHost, now),
-          ForeignKeyIndex.create(newHost, newHost.getDeletionTime()));
+      entitiesToUpdate.add(ForeignKeyIndex.create(existingHost, now));
+      entitiesToUpdate.add(ForeignKeyIndex.create(newHost, newHost.getDeletionTime()));
       updateSuperordinateDomains(existingHost, newHost);
     }
     enqueueTasks(existingHost, newHost);
-    entitiesToSave.add(historyBuilder
-        .setType(HistoryEntry.Type.HOST_UPDATE)
-        .setModificationTime(now)
-        .setParent(Key.create(existingHost))
-        .build());
-    ofy().save().entities(entitiesToSave.build());
+    entitiesToInsert.add(
+        historyBuilder
+            .setType(HistoryEntry.Type.HOST_UPDATE)
+            .setModificationTime(now)
+            .setParent(Key.create(existingHost))
+            .build()
+            .toChildHistoryEntity());
+    tm().updateAll(entitiesToUpdate.build());
+    tm().insertAll(entitiesToInsert.build());
     return responseBuilder.build();
   }
 
@@ -286,7 +288,7 @@ public final class HostUpdateFlow implements TransactionalFlow {
         && Objects.equals(
             existingHost.getSuperordinateDomain(), newHost.getSuperordinateDomain())) {
       tm().put(
-              tm().load(existingHost.getSuperordinateDomain())
+              tm().loadByKey(existingHost.getSuperordinateDomain())
                   .asBuilder()
                   .removeSubordinateHost(existingHost.getHostName())
                   .addSubordinateHost(newHost.getHostName())
@@ -295,14 +297,14 @@ public final class HostUpdateFlow implements TransactionalFlow {
     }
     if (existingHost.isSubordinate()) {
       tm().put(
-              tm().load(existingHost.getSuperordinateDomain())
+              tm().loadByKey(existingHost.getSuperordinateDomain())
                   .asBuilder()
                   .removeSubordinateHost(existingHost.getHostName())
                   .build());
     }
     if (newHost.isSubordinate()) {
       tm().put(
-              tm().load(newHost.getSuperordinateDomain())
+              tm().loadByKey(newHost.getSuperordinateDomain())
                   .asBuilder()
                   .addSubordinateHost(newHost.getHostName())
                   .build());

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

@@ -148,7 +148,7 @@ public abstract class EppResource extends BackupGroupRoot implements Buildable {
    *
    * @see google.registry.model.translators.CommitLogRevisionsTranslatorFactory
    */
-  @Transient
+  @Transient @DoNotCompare
   ImmutableSortedMap<DateTime, Key<CommitLogManifest>> revisions = ImmutableSortedMap.of();
 
   public String getRepoId() {
@@ -360,13 +360,13 @@ public abstract class EppResource extends BackupGroupRoot implements Buildable {
 
         @Override
         public EppResource load(VKey<? extends EppResource> key) {
-          return tm().doTransactionless(() -> tm().load(key));
+          return tm().doTransactionless(() -> tm().loadByKey(key));
         }
 
         @Override
         public Map<VKey<? extends EppResource>, EppResource> loadAll(
             Iterable<? extends VKey<? extends EppResource>> keys) {
-          return tm().doTransactionless(() -> tm().load(keys));
+          return tm().doTransactionless(() -> tm().loadByKeys(keys));
         }
       };
 
@@ -406,7 +406,7 @@ public abstract class EppResource extends BackupGroupRoot implements Buildable {
   public static ImmutableMap<VKey<? extends EppResource>, EppResource> loadCached(
       Iterable<VKey<? extends EppResource>> keys) {
     if (!RegistryConfig.isEppResourceCachingEnabled()) {
-      return tm().load(keys);
+      return tm().loadByKeys(keys);
     }
     try {
       return cacheEppResources.getAll(keys);
@@ -423,7 +423,7 @@ public abstract class EppResource extends BackupGroupRoot implements Buildable {
    */
   public static <T extends EppResource> T loadCached(VKey<T> key) {
     if (!RegistryConfig.isEppResourceCachingEnabled()) {
-      return tm().load(key);
+      return tm().loadByKey(key);
     }
     try {
       // Safe to cast because loading a Key<T> returns an entity of type T.

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

@@ -17,6 +17,7 @@ package google.registry.model;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 import static google.registry.util.DateTimeUtils.isAtOrAfter;
@@ -28,7 +29,6 @@ import com.google.common.collect.ImmutableSet;
 import com.google.common.flogger.FluentLogger;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.Result;
-import com.googlecode.objectify.cmd.Query;
 import com.googlecode.objectify.util.ResultNow;
 import google.registry.config.RegistryConfig;
 import google.registry.model.EppResource.BuilderWithTransferData;
@@ -37,6 +37,7 @@ import google.registry.model.EppResource.ResourceWithTransferData;
 import google.registry.model.contact.ContactResource;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.eppcommon.StatusValue;
+import google.registry.model.host.HostResource;
 import google.registry.model.index.ForeignKeyIndex;
 import google.registry.model.ofy.CommitLogManifest;
 import google.registry.model.ofy.CommitLogMutation;
@@ -44,11 +45,13 @@ import google.registry.model.registry.Registry;
 import google.registry.model.transfer.DomainTransferData;
 import google.registry.model.transfer.TransferData;
 import google.registry.model.transfer.TransferStatus;
+import google.registry.persistence.VKey;
 import java.util.List;
 import java.util.Map.Entry;
 import java.util.Optional;
 import java.util.function.Function;
 import javax.annotation.Nullable;
+import javax.persistence.Query;
 import org.joda.time.DateTime;
 import org.joda.time.Interval;
 
@@ -57,6 +60,22 @@ public final class EppResourceUtils {
 
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
 
+  private static final String CONTACT_LINKED_DOMAIN_QUERY =
+      "SELECT repoId FROM Domain "
+          + "WHERE (adminContact = :fkRepoId "
+          + "OR billingContact = :fkRepoId "
+          + "OR techContact = :fkRepoId "
+          + "OR registrantContact = :fkRepoId) "
+          + "AND deletionTime > :now";
+
+  // We have to use the native SQL query here because DomainHost table doesn't have its entity
+  // class so we cannot reference its property like domainHost.hostRepoId in a JPQL query.
+  private static final String HOST_LINKED_DOMAIN_QUERY =
+      "SELECT d.repo_id FROM \"Domain\" d "
+          + "JOIN \"DomainHost\" dh ON dh.domain_repo_id = d.repo_id "
+          + "WHERE d.deletion_time > :now "
+          + "AND dh.host_repo_id = :fkRepoId";
+
   /** Returns the full domain repoId in the format HEX-TLD for the specified long id and tld. */
   public static String createDomainRepoId(long repoId, String tld) {
     return createRepoId(repoId, Registry.get(tld).getRoidSuffix());
@@ -144,7 +163,7 @@ public final class EppResourceUtils {
     T resource =
         useCache
             ? EppResource.loadCached(fki.getResourceKey())
-            : transactIfJpaTm(() -> tm().maybeLoad(fki.getResourceKey()).orElse(null));
+            : transactIfJpaTm(() -> tm().loadByKeyIfPresent(fki.getResourceKey()).orElse(null));
     if (resource == null || isAtOrAfter(now, resource.getDeletionTime())) {
       return Optional.empty();
     }
@@ -365,21 +384,63 @@ public final class EppResourceUtils {
   }
 
   /**
-   * Returns a query for domains or applications that reference a specified contact or host.
+   * Returns a set of {@link VKey} for domains that reference a specified contact or host.
    *
-   * <p>This is an eventually consistent query.
+   * <p>This is an eventually consistent query if used for Datastore.
    *
    * @param key the referent key
    * @param now the logical time of the check
+   * @param limit the maximum number of returned keys
    */
-  public static Query<DomainBase> queryForLinkedDomains(
-      Key<? extends EppResource> key, DateTime now) {
-    boolean isContactKey = key.getKind().equals(Key.getKind(ContactResource.class));
-    return ofy()
-        .load()
-        .type(DomainBase.class)
-        .filter(isContactKey ? "allContacts.contact" : "nsHosts", key)
-        .filter("deletionTime >", now);
+  public static ImmutableSet<VKey<DomainBase>> getLinkedDomainKeys(
+      VKey<? extends EppResource> key, DateTime now, int limit) {
+    checkArgument(
+        key.getKind().equals(ContactResource.class) || key.getKind().equals(HostResource.class),
+        "key must be either VKey<ContactResource> or VKey<HostResource>, but it is %s",
+        key);
+    boolean isContactKey = key.getKind().equals(ContactResource.class);
+    if (tm().isOfy()) {
+      return ofy()
+          .load()
+          .type(DomainBase.class)
+          .filter(isContactKey ? "allContacts.contact" : "nsHosts", key.getOfyKey())
+          .filter("deletionTime >", now)
+          .limit(limit)
+          .keys()
+          .list()
+          .stream()
+          .map(DomainBase::createVKey)
+          .collect(toImmutableSet());
+    } else {
+      return tm().transact(
+              () -> {
+                Query query;
+                if (isContactKey) {
+                  query =
+                      jpaTm()
+                          .getEntityManager()
+                          .createQuery(CONTACT_LINKED_DOMAIN_QUERY, String.class)
+                          .setParameter("fkRepoId", key)
+                          .setParameter("now", now);
+                } else {
+                  query =
+                      jpaTm()
+                          .getEntityManager()
+                          .createNativeQuery(HOST_LINKED_DOMAIN_QUERY)
+                          .setParameter("fkRepoId", key.getSqlKey())
+                          .setParameter("now", now.toDate());
+                }
+                return (ImmutableSet<VKey<DomainBase>>)
+                    query
+                        .setMaxResults(limit)
+                        .getResultStream()
+                        .map(
+                            repoId ->
+                                DomainBase.createVKey(
+                                    Key.create(DomainBase.class, (String) repoId)))
+                        .collect(toImmutableSet());
+              });
+    }
   }
 
   /**
@@ -390,8 +451,8 @@ public final class EppResourceUtils {
    * @param key the referent key
    * @param now the logical time of the check
    */
-  public static boolean isLinked(Key<? extends EppResource> key, DateTime now) {
-    return queryForLinkedDomains(key, now).limit(1).count() > 0;
+  public static boolean isLinked(VKey<? extends EppResource> key, DateTime now) {
+    return getLinkedDomainKeys(key, now, 1).size() > 0;
   }
 
   private EppResourceUtils() {}

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

@@ -54,9 +54,37 @@ public abstract class ImmutableObject implements Cloneable {
   @Target(FIELD)
   public @interface DoNotHydrate {}
 
-  @Ignore
-  @XmlTransient
-  Integer hashCode;
+  /**
+   * Indicates that the field should be ignored when comparing an object in the datastore to the
+   * corresponding object in Cloud SQL.
+   */
+  @Documented
+  @Retention(RUNTIME)
+  @Target(FIELD)
+  public @interface DoNotCompare {}
+
+  /**
+   * Indicates that the field stores a null value to indicate an empty set. This is also used in
+   * object comparison.
+   */
+  @Documented
+  @Retention(RUNTIME)
+  @Target(FIELD)
+  public @interface EmptySetToNull {}
+
+  /**
+   * Indicates that the field does not take part in the immutability contract.
+   *
+   * <p>Certain fields currently get modified by hibernate and there is nothing we can do about it.
+   * As well as violating immutability, this breaks hashing and equality comparisons, so we mark
+   * these fields with this annotation to exclude them from most operations.
+   */
+  @Documented
+  @Retention(RUNTIME)
+  @Target(FIELD)
+  public @interface Insignificant {}
+
+  @Ignore @XmlTransient protected Integer hashCode;
 
   private boolean equalsImmutableObject(ImmutableObject other) {
     return getClass().equals(other.getClass())
@@ -71,7 +99,14 @@ public abstract class ImmutableObject implements Cloneable {
    * <p>Isolated into a method so that derived classes can override it.
    */
   protected Map<Field, Object> getSignificantFields() {
-    return ModelUtils.getFieldValues(this);
+    // Can't use streams or ImmutableMap because we can have null values.
+    Map<Field, Object> result = new LinkedHashMap();
+    for (Map.Entry<Field, Object> entry : ModelUtils.getFieldValues(this).entrySet()) {
+      if (!entry.getKey().isAnnotationPresent(Insignificant.class)) {
+        result.put(entry.getKey(), entry.getValue());
+      }
+    }
+    return result;
   }
 
   @Override

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

@@ -46,6 +46,8 @@ import google.registry.model.domain.rgp.GracePeriodStatus;
 import google.registry.model.domain.token.AllocationToken;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.transfer.TransferData.TransferServerApproveEntity;
+import google.registry.persistence.BillingVKey.BillingEventVKey;
+import google.registry.persistence.BillingVKey.BillingRecurrenceVKey;
 import google.registry.persistence.VKey;
 import google.registry.persistence.WithLongVKey;
 import google.registry.schema.replay.DatastoreAndSqlEntity;
@@ -572,8 +574,7 @@ public abstract class BillingEvent extends ImmutableObject
      * <p>Although the type is {@link Key} the name "ref" is preserved for historical reasons.
      */
     @IgnoreSave(IfNull.class)
-    @Column(name = "billing_event_id")
-    VKey<BillingEvent.OneTime> refOneTime = null;
+    BillingEventVKey refOneTime = null;
 
     /**
      * The recurring billing event to cancel, or null for non-autorenew cancellations.
@@ -581,15 +582,14 @@ public abstract class BillingEvent extends ImmutableObject
      * <p>Although the type is {@link Key} the name "ref" is preserved for historical reasons.
      */
     @IgnoreSave(IfNull.class)
-    @Column(name = "billing_recurrence_id")
-    VKey<BillingEvent.Recurring> refRecurring = null;
+    BillingRecurrenceVKey refRecurring = null;
 
     public DateTime getBillingTime() {
       return billingTime;
     }
 
     public VKey<? extends BillingEvent> getEventKey() {
-      return firstNonNull(refOneTime, refRecurring);
+      return firstNonNull(refOneTime, refRecurring).createVKey();
     }
 
     /** The mapping from billable grace period types to originating billing event reasons. */
@@ -656,12 +656,12 @@ public abstract class BillingEvent extends ImmutableObject
       }
 
       public Builder setOneTimeEventKey(VKey<BillingEvent.OneTime> eventKey) {
-        getInstance().refOneTime = eventKey;
+        getInstance().refOneTime = BillingEventVKey.create(eventKey);
         return this;
       }
 
       public Builder setRecurringEventKey(VKey<BillingEvent.Recurring> eventKey) {
-        getInstance().refRecurring = eventKey;
+        getInstance().refRecurring = BillingRecurrenceVKey.create(eventKey);
         return this;
       }
 

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

@@ -14,6 +14,7 @@
 
 package google.registry.model.domain;
 
+import static com.google.common.base.MoreObjects.firstNonNull;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Strings.emptyToNull;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
@@ -65,6 +66,7 @@ import google.registry.model.transfer.DomainTransferData;
 import google.registry.model.transfer.TransferStatus;
 import google.registry.persistence.VKey;
 import google.registry.util.CollectionUtils;
+import google.registry.util.DateTimeUtils;
 import java.util.HashSet;
 import java.util.Objects;
 import java.util.Optional;
@@ -132,7 +134,7 @@ public class DomainContent extends EppResource
   @Index String tld;
 
   /** References to hosts that are the nameservers for the domain. */
-  @Index @Transient Set<VKey<HostResource>> nsHosts;
+  @EmptySetToNull @Index @Transient Set<VKey<HostResource>> nsHosts;
 
   /**
    * The union of the contacts visible via {@link #getContacts} and {@link #getRegistrant}.
@@ -283,9 +285,10 @@ public class DomainContent extends EppResource
   /**
    * When the domain's autorenewal status will expire.
    *
-   * <p>This will be null for the vast majority of domains because all domains autorenew
-   * indefinitely by default and autorenew can only be countermanded by administrators, typically
-   * for reasons of the URS process or termination of a registrar for nonpayment.
+   * <p>This will be {@link DateTimeUtils#END_OF_TIME} for the vast majority of domains because all
+   * domains autorenew indefinitely by default and autorenew can only be countermanded by
+   * administrators, typically for reasons of the URS process or termination of a registrar for
+   * nonpayment.
    *
    * <p>When a domain is scheduled to not autorenew, this field is set to the current value of its
    * {@link #registrationExpirationTime}, after which point the next invocation of a periodic
@@ -294,10 +297,16 @@ public class DomainContent extends EppResource
    * difference domains that have reached their life and must be deleted now, and domains that
    * happen to be in the autorenew grace period now but should be deleted in roughly a year.
    */
-  @Nullable @Index DateTime autorenewEndTime;
+  @Index DateTime autorenewEndTime;
 
   @OnLoad
   void load() {
+    // Back fill with correct END_OF_TIME sentinel value.
+    // TODO(mcilwain): Remove this once back-filling is complete.
+    if (autorenewEndTime == null) {
+      autorenewEndTime = END_OF_TIME;
+    }
+
     // Reconstitute all of the contacts so that they have VKeys.
     allContacts =
         allContacts.stream().map(DesignatedContact::reconstitute).collect(toImmutableSet());
@@ -319,6 +328,11 @@ public class DomainContent extends EppResource
     autorenewPollMessageHistoryId = getHistoryId(autorenewPollMessage);
     autorenewBillingEventHistoryId = getHistoryId(autorenewBillingEvent);
     deletePollMessageHistoryId = getHistoryId(deletePollMessage);
+
+    // Fix PollMessage VKeys.
+    autorenewPollMessage = PollMessage.Autorenew.convertVKey(autorenewPollMessage);
+    deletePollMessage = PollMessage.OneTime.convertVKey(deletePollMessage);
+
     dsData =
         nullToEmptyImmutableCopy(dsData).stream()
             .map(dsData -> dsData.cloneWithDomainRepoId(getRepoId()))
@@ -397,8 +411,20 @@ public class DomainContent extends EppResource
     return smdId;
   }
 
+  /**
+   * Returns the autorenew end time if there is one, otherwise empty.
+   *
+   * <p>Note that {@link DateTimeUtils#END_OF_TIME} is used as a sentinel value in the database
+   * representation to signify that autorenew doesn't end, and is mapped to empty here for the
+   * purposes of more legible business logic.
+   */
   public Optional<DateTime> getAutorenewEndTime() {
-    return Optional.ofNullable(autorenewEndTime);
+    // TODO(mcilwain): Remove null handling for autorenewEndTime once data migration away from null
+    //                 is complete.
+    return Optional.ofNullable(
+        (autorenewEndTime == null || autorenewEndTime.equals(END_OF_TIME))
+            ? null
+            : autorenewEndTime);
   }
 
   @Override
@@ -773,6 +799,8 @@ public class DomainContent extends EppResource
       } else { // There are nameservers, so make sure INACTIVE isn't there.
         removeStatusValue(StatusValue.INACTIVE);
       }
+      // If there is no autorenew end time, set it to END_OF_TIME.
+      instance.autorenewEndTime = firstNonNull(getInstance().autorenewEndTime, END_OF_TIME);
 
       checkArgumentNotNull(emptyToNull(instance.fullyQualifiedDomainName), "Missing domainName");
       if (instance.getRegistrant() == null
@@ -952,8 +980,15 @@ public class DomainContent extends EppResource
       return thisCastToDerived();
     }
 
+    /**
+     * Sets the autorenew end time, or clears it if empty is passed.
+     *
+     * <p>Note that {@link DateTimeUtils#END_OF_TIME} is used as a sentinel value in the database
+     * representation to signify that autorenew doesn't end, and is mapped to empty here for the
+     * purposes of more legible business logic.
+     */
     public B setAutorenewEndTime(Optional<DateTime> autorenewEndTime) {
-      getInstance().autorenewEndTime = autorenewEndTime.orElse(null);
+      getInstance().autorenewEndTime = autorenewEndTime.orElse(END_OF_TIME);
       return thisCastToDerived();
     }
 

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

@@ -104,6 +104,7 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
                 "domain_history_history_revision_id,domain_history_domain_repo_id,host_repo_id",
             unique = true),
       })
+  @ImmutableObject.EmptySetToNull
   @Column(name = "host_repo_id")
   Set<VKey<HostResource>> nsHosts;
 
@@ -180,7 +181,9 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
    * #getDomainTransactionRecords()}.
    */
   @Access(AccessType.PROPERTY)
-  @OneToMany(cascade = {CascadeType.ALL})
+  @OneToMany(
+      cascade = {CascadeType.ALL},
+      fetch = FetchType.EAGER)
   @JoinColumn(name = "historyRevisionId", referencedColumnName = "historyRevisionId")
   @JoinColumn(name = "domainRepoId", referencedColumnName = "domainRepoId")
   @SuppressWarnings("unused")

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

@@ -193,19 +193,6 @@ public class GracePeriod extends GracePeriodBase implements DatastoreAndSqlEntit
     return clone;
   }
 
-  /**
-   * Returns a clone of this {@link GracePeriod} with prepopulated {@link #gracePeriodId} generated
-   * by {@link ObjectifyService#allocateId()}.
-   *
-   * <p>TODO(shicong): Figure out how to generate the id only when the entity is used for Cloud SQL.
-   */
-  @VisibleForTesting
-  public GracePeriod cloneWithPrepopulatedId() {
-    GracePeriod clone = clone(this);
-    clone.gracePeriodId = ObjectifyService.allocateId();
-    return clone;
-  }
-
   /** Entity class to represent a historic {@link GracePeriod}. */
   @Entity(name = "GracePeriodHistory")
   @Table(indexes = @Index(columnList = "domainRepoId"))

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

@@ -17,15 +17,11 @@ package google.registry.model.domain;
 import com.googlecode.objectify.annotation.Embed;
 import com.googlecode.objectify.annotation.Ignore;
 import google.registry.model.ImmutableObject;
-import google.registry.model.ModelUtils;
 import google.registry.model.billing.BillingEvent;
 import google.registry.model.domain.rgp.GracePeriodStatus;
 import google.registry.persistence.BillingVKey.BillingEventVKey;
 import google.registry.persistence.BillingVKey.BillingRecurrenceVKey;
 import google.registry.persistence.VKey;
-import java.lang.reflect.Field;
-import java.util.LinkedHashMap;
-import java.util.Map;
 import javax.persistence.Access;
 import javax.persistence.AccessType;
 import javax.persistence.Column;
@@ -125,20 +121,4 @@ public class GracePeriodBase extends ImmutableObject {
   public VKey<BillingEvent.Recurring> getRecurringBillingEvent() {
     return billingEventRecurring == null ? null : billingEventRecurring.createVKey();
   }
-
-  /**
-   * Override {@link ImmutableObject#getSignificantFields()} to exclude "id", which breaks equality
-   * testing in the unit tests.
-   */
-  @Override
-  protected Map<Field, Object> getSignificantFields() {
-    // Can't use streams or ImmutableMap because we can have null values.
-    Map<Field, Object> result = new LinkedHashMap();
-    for (Map.Entry<Field, Object> entry : ModelUtils.getFieldValues(this).entrySet()) {
-      if (!entry.getKey().getName().equals("id")) {
-        result.put(entry.getKey(), entry.getValue());
-      }
-    }
-    return result;
-  }
 }

core/src/main/java/google/registry/model/domain/secdns/DelegationSignerData.java

@@ -35,7 +35,7 @@ import javax.xml.bind.annotation.XmlType;
  *
  * @see <a href="http://tools.ietf.org/html/rfc5910">RFC 5910</a>
  * @see <a href="http://tools.ietf.org/html/rfc4034">RFC 4034</a>
- *     <p>TODO(shicong): Rename this class to DomainDsData.
+ *     <p>TODO(b/177567432): Rename this class to DomainDsData.
  */
 @Embed
 @XmlType(name = "dsData")

core/src/main/java/google/registry/model/eppcommon/Address.java

@@ -35,6 +35,7 @@ import javax.persistence.Embeddable;
 import javax.persistence.MappedSuperclass;
 import javax.persistence.PostLoad;
 import javax.persistence.Transient;
+import javax.xml.bind.Unmarshaller;
 import javax.xml.bind.annotation.XmlElement;
 import javax.xml.bind.annotation.XmlTransient;
 import javax.xml.bind.annotation.adapters.CollapsedStringAdapter;
@@ -58,8 +59,8 @@ import javax.xml.bind.annotation.adapters.XmlJavaTypeAdapter;
 public class Address extends ImmutableObject implements Jsonifiable {
 
   /** The schema validation will enforce that this has 3 lines at most. */
-  // TODO(shicong): Remove this field after migration. We need to figure out how to generate same
-  // XML from streetLine[1,2,3].
+  // TODO(b/177569726): Remove this field after migration. We need to figure out how to generate
+  // same XML from streetLine[1,2,3].
   @XmlJavaTypeAdapter(NormalizedStringAdapter.class)
   @Transient
   List<String> street;
@@ -174,15 +175,10 @@ public class Address extends ImmutableObject implements Jsonifiable {
    * entity from Datastore.
    *
    * <p>This callback method is used by Objectify to set streetLine[1,2,3] fields as they are not
-   * persisted in the Datastore. TODO(shicong): Delete this method after database migration.
+   * persisted in the Datastore.
    */
   void onLoad(@AlsoLoad("street") List<String> street) {
-    if (street == null || street.size() == 0) {
-      return;
-    }
-    streetLine1 = street.get(0);
-    streetLine2 = street.size() >= 2 ? street.get(1) : null;
-    streetLine3 = street.size() >= 3 ? street.get(2) : null;
+    mapStreetListToIndividualFields(street);
   }
 
   /**
@@ -202,4 +198,23 @@ public class Address extends ImmutableObject implements Jsonifiable {
                 .filter(Objects::nonNull)
                 .collect(toImmutableList());
   }
+
+  /**
+   * Sets {@link #streetLine1}, {@link #streetLine2} and {@link #streetLine3} when the entity is
+   * reconstructed from XML message.
+   *
+   * <p>This is a callback function that JAXB invokes after unmarshalling the XML message.
+   */
+  void afterUnmarshal(Unmarshaller unmarshaller, Object parent) {
+    mapStreetListToIndividualFields(street);
+  }
+
+  private void mapStreetListToIndividualFields(List<String> street) {
+    if (street == null || street.size() == 0) {
+      return;
+    }
+    streetLine1 = street.get(0);
+    streetLine2 = street.size() >= 2 ? street.get(1) : null;
+    streetLine3 = street.size() >= 3 ? street.get(2) : null;
+  }
 }

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

@@ -16,6 +16,7 @@ package google.registry.model.ofy;
 
 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.model.ofy.ObjectifyService.ofy;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 
@@ -23,10 +24,13 @@ import com.google.common.base.Functions;
 import com.google.common.collect.ImmutableCollection;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Iterables;
 import com.google.common.collect.Streams;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.Result;
 import google.registry.model.contact.ContactHistory;
+import google.registry.model.domain.DomainHistory;
 import google.registry.model.host.HostHistory;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.persistence.VKey;
@@ -177,26 +181,13 @@ public class DatastoreTransactionManager implements TransactionManager {
   // VKey instead of by ofy Key.  But ideally, there should be one set of TransactionManager
   // interface tests that are applied to both the datastore and SQL implementations.
   @Override
-  public <T> Optional<T> maybeLoad(VKey<T> key) {
+  public <T> Optional<T> loadByKeyIfPresent(VKey<T> key) {
     return Optional.ofNullable(loadNullable(key));
   }
 
   @Override
-  public <T> T load(VKey<T> key) {
-    T result = loadNullable(key);
-    if (result == null) {
-      throw new NoSuchElementException(key.toString());
-    }
-    return result;
-  }
-
-  @Override
-  public <T> T load(T entity) {
-    return ofy().load().entity(entity).now();
-  }
-
-  @Override
-  public <T> ImmutableMap<VKey<? extends T>, T> load(Iterable<? extends VKey<? extends T>> keys) {
+  public <T> ImmutableMap<VKey<? extends T>, T> loadByKeysIfPresent(
+      Iterable<? extends VKey<? extends T>> keys) {
     // Keep track of the Key -> VKey mapping so we can translate them back.
     ImmutableMap<Key<T>, VKey<? extends T>> keyMap =
         StreamSupport.stream(keys.spliterator(), false)
@@ -211,13 +202,51 @@ public class DatastoreTransactionManager implements TransactionManager {
   }
 
   @Override
-  public <T> ImmutableList<T> loadAll(Class<T> clazz) {
-    return ImmutableList.copyOf(getOfy().load().type(clazz));
+  public <T> ImmutableList<T> loadByEntitiesIfPresent(Iterable<T> entities) {
+    return ImmutableList.copyOf(getOfy().load().entities(entities).values());
   }
 
   @Override
-  public <T> ImmutableList<T> loadAll(Iterable<T> entities) {
-    return ImmutableList.copyOf(getOfy().load().entities(entities).values());
+  public <T> T loadByKey(VKey<T> key) {
+    T result = loadNullable(key);
+    if (result == null) {
+      throw new NoSuchElementException(key.toString());
+    }
+    return result;
+  }
+
+  @Override
+  public <T> ImmutableMap<VKey<? extends T>, T> loadByKeys(
+      Iterable<? extends VKey<? extends T>> keys) {
+    ImmutableMap<VKey<? extends T>, T> result = loadByKeysIfPresent(keys);
+    ImmutableSet<? extends VKey<? extends T>> missingKeys =
+        Streams.stream(keys).filter(k -> !result.containsKey(k)).collect(toImmutableSet());
+    if (!missingKeys.isEmpty()) {
+      // Ofy ignores nonexistent keys but the method contract specifies to throw if nonexistent
+      throw new NoSuchElementException(
+          String.format("Failed to load nonexistent entities for keys: %s", missingKeys));
+    }
+    return result;
+  }
+
+  @Override
+  public <T> T loadByEntity(T entity) {
+    return ofy().load().entity(entity).now();
+  }
+
+  @Override
+  public <T> ImmutableList<T> loadByEntities(Iterable<T> entities) {
+    ImmutableList<T> result = loadByEntitiesIfPresent(entities);
+    if (result.size() != Iterables.size(entities)) {
+      throw new NoSuchElementException(
+          String.format("Attempted to load entities, some of which are missing: %s", entities));
+    }
+    return result;
+  }
+
+  @Override
+  public <T> ImmutableList<T> loadAllOf(Class<T> clazz) {
+    return ImmutableList.copyOf(getOfy().load().type(clazz));
   }
 
   @Override
@@ -264,6 +293,21 @@ public class DatastoreTransactionManager implements TransactionManager {
     getOfy().clearSessionCache();
   }
 
+  /**
+   * Executes the given {@link Result} instance synchronously if not in a transaction.
+   *
+   * <p>The {@link Result} instance contains a task that will be executed by Objectify
+   * asynchronously. If it is in a transaction, we don't need to execute the task immediately
+   * because it is guaranteed to be done by the end of the transaction. However, if it is not in a
+   * transaction, we need to execute it in case the following code expects that happens before
+   * themselves.
+   */
+  private void syncIfTransactionless(Result<?> result) {
+    if (!inTransaction()) {
+      result.now();
+    }
+  }
+
   /**
    * The following three methods exist due to the migration to Cloud SQL.
    *
@@ -281,34 +325,23 @@ public class DatastoreTransactionManager implements TransactionManager {
     syncIfTransactionless(getOfy().save().entity(entity));
   }
 
-  @SuppressWarnings("unchecked")
-  private <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 != null && HistoryEntry.class.isAssignableFrom(obj.getClass())) {
-      return (T) ((HistoryEntry) obj).toChildHistoryEntity();
-    }
-    return obj;
-  }
-
   @Nullable
   private <T> T loadNullable(VKey<T> key) {
     return toChildHistoryEntryIfPossible(getOfy().load().key(key.getOfyKey()).now());
   }
 
-  /**
-   * Executes the given {@link Result} instance synchronously if not in a transaction.
-   *
-   * <p>The {@link Result} instance contains a task that will be executed by Objectify
-   * asynchronously. If it is in a transaction, we don't need to execute the task immediately
-   * because it is guaranteed to be done by the end of the transaction. However, if it is not in a
-   * transaction, we need to execute it in case the following code expects that happens before
-   * themselves.
-   */
-  private void syncIfTransactionless(Result<?> result) {
-    if (!inTransaction()) {
-      result.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();
     }
+    return obj;
   }
 }

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

@@ -14,7 +14,18 @@
 
 package google.registry.model.ofy;
 
+import static google.registry.model.ofy.EntityWritePriorities.getEntityPriority;
+import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+
+import com.google.common.collect.ImmutableMap;
+import com.googlecode.objectify.Key;
 import google.registry.config.RegistryEnvironment;
+import google.registry.model.UpdateAutoTimestamp;
+import google.registry.persistence.VKey;
+import google.registry.schema.replay.DatastoreEntity;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.concurrent.ConcurrentLinkedQueue;
 
 /**
@@ -24,23 +35,79 @@ import java.util.concurrent.ConcurrentLinkedQueue;
  */
 public class ReplayQueue {
 
-  static ConcurrentLinkedQueue<TransactionInfo> queue =
-      new ConcurrentLinkedQueue<TransactionInfo>();
+  static ConcurrentLinkedQueue<ImmutableMap<Key<?>, Object>> queue =
+      new ConcurrentLinkedQueue<ImmutableMap<Key<?>, Object>>();
 
   static void addInTests(TransactionInfo info) {
     if (RegistryEnvironment.get() == RegistryEnvironment.UNITTEST) {
-      queue.add(info);
+      // Transform the entities to be persisted to the set of values as they were actually
+      // persisted.
+      ImmutableMap.Builder<Key<?>, Object> builder = new ImmutableMap.Builder<Key<?>, Object>();
+      for (ImmutableMap.Entry<Key<?>, Object> entry : info.getChanges().entrySet()) {
+        if (entry.getValue().equals(TransactionInfo.Delete.SENTINEL)) {
+          builder.put(entry.getKey(), entry.getValue());
+        } else {
+          // The value is an entity object that has not yet been persisted, and thus some of the
+          // special transformations that we do (notably the auto-timestamp transformations) have
+          // not been applied.  Converting the object to an entity and then back again performs
+          // those transformations so that we persist the same values to SQL that we have in
+          // Datastore.
+          builder.put(entry.getKey(), ofy().toPojo(ofy().toEntity(entry.getValue())));
+        }
+      }
+      queue.add(builder.build());
     }
   }
 
-  public static void replay() {
-    TransactionInfo info;
-    while ((info = queue.poll()) != null) {
-      info.saveToJpa();
+  /** Replay all transactions, return the set of keys that were replayed. */
+  public static ImmutableMap<Key<?>, Object> replay() {
+    // We can't use an ImmutableMap.Builder here, we need to be able to overwrite existing values
+    // and the builder doesn't support that.
+    Map<Key<?>, Object> result = new HashMap<Key<?>, Object>();
+    ImmutableMap<Key<?>, Object> changes;
+    while ((changes = queue.poll()) != null) {
+      saveToJpa(changes);
+      result.putAll(changes);
     }
+
+    return ImmutableMap.copyOf(result);
   }
 
   public static void clear() {
     queue.clear();
   }
+
+  /** Returns the priority of the entity type in the map entry. */
+  private static int getPriority(ImmutableMap.Entry<Key<?>, Object> entry) {
+    return getEntityPriority(
+        entry.getKey().getKind(), entry.getValue().equals(TransactionInfo.Delete.SENTINEL));
+  }
+
+  private static int compareByPriority(
+      ImmutableMap.Entry<Key<?>, Object> a, ImmutableMap.Entry<Key<?>, Object> b) {
+    return getPriority(a) - getPriority(b);
+  }
+
+  private static void saveToJpa(ImmutableMap<Key<?>, Object> changes) {
+    try (UpdateAutoTimestamp.DisableAutoUpdateResource disabler =
+        UpdateAutoTimestamp.disableAutoUpdate()) {
+      // Sort the changes into an order that will work for insertion into the database.
+      jpaTm()
+          .transact(
+              () -> {
+                changes.entrySet().stream()
+                    .sorted(ReplayQueue::compareByPriority)
+                    .forEach(
+                        entry -> {
+                          if (entry.getValue().equals(TransactionInfo.Delete.SENTINEL)) {
+                            jpaTm().delete(VKey.from(entry.getKey()));
+                          } else {
+                            ((DatastoreEntity) entry.getValue())
+                                .toSqlEntity()
+                                .ifPresent(jpaTm()::put);
+                          }
+                        });
+              });
+    }
+  }
 }

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

@@ -20,24 +20,20 @@ import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static com.google.common.collect.Maps.filterValues;
 import static com.google.common.collect.Maps.toMap;
 import static google.registry.model.ofy.CommitLogBucket.getArbitraryBucketId;
-import static google.registry.model.ofy.EntityWritePriorities.getEntityPriority;
 import static google.registry.model.ofy.ObjectifyService.ofy;
-import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
-import google.registry.persistence.VKey;
-import google.registry.schema.replay.DatastoreEntity;
 import java.util.Map;
 import org.joda.time.DateTime;
 
 /** Metadata for an {@link Ofy} transaction that saves commit logs. */
-class TransactionInfo {
+public class TransactionInfo {
 
   @VisibleForTesting
-  enum Delete {
+  public enum Delete {
     SENTINEL
   }
 
@@ -87,6 +83,10 @@ class TransactionInfo {
     return ImmutableSet.copyOf(changesBuilder.build().keySet());
   }
 
+  ImmutableMap<Key<?>, Object> getChanges() {
+    return changesBuilder.build();
+  }
+
   ImmutableSet<Key<?>> getDeletes() {
     return ImmutableSet.copyOf(
         filterValues(changesBuilder.build(), Delete.SENTINEL::equals).keySet());
@@ -100,35 +100,4 @@ class TransactionInfo {
         .filter(not(Delete.SENTINEL::equals))
         .collect(toImmutableSet());
   }
-
-  /** Returns the weight of the entity type in the map entry. */
-  @VisibleForTesting
-  static int getWeight(ImmutableMap.Entry<Key<?>, Object> entry) {
-    return getEntityPriority(entry.getKey().getKind(), entry.getValue().equals(Delete.SENTINEL));
-  }
-
-  private static int compareByWeight(
-      ImmutableMap.Entry<Key<?>, Object> a, ImmutableMap.Entry<Key<?>, Object> b) {
-    return getWeight(a) - getWeight(b);
-  }
-
-  void saveToJpa() {
-    // Sort the changes into an order that will work for insertion into the database.
-    jpaTm()
-        .transact(
-            () -> {
-              changesBuilder.build().entrySet().stream()
-                  .sorted(TransactionInfo::compareByWeight)
-                  .forEach(
-                      entry -> {
-                        if (entry.getValue().equals(Delete.SENTINEL)) {
-                          jpaTm().delete(VKey.from(entry.getKey()));
-                        } else {
-                          ((DatastoreEntity) entry.getValue())
-                              .toSqlEntity()
-                              .ifPresent(jpaTm()::put);
-                        }
-                      });
-            });
-  }
 }

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

@@ -16,6 +16,7 @@ package google.registry.model.poll;
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static google.registry.util.CollectionUtils.forceEmptyToNull;
+import static google.registry.util.CollectionUtils.isNullOrEmpty;
 import static google.registry.util.CollectionUtils.nullToEmpty;
 import static google.registry.util.DateTimeUtils.END_OF_TIME;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
@@ -52,6 +53,7 @@ import google.registry.persistence.WithLongVKey;
 import google.registry.schema.replay.DatastoreAndSqlEntity;
 import java.util.List;
 import java.util.Optional;
+import javax.annotation.Nullable;
 import javax.persistence.AttributeOverride;
 import javax.persistence.AttributeOverrides;
 import javax.persistence.Column;
@@ -185,6 +187,7 @@ public abstract class PollMessage extends ImmutableObject
   @Override
   public abstract VKey<? extends PollMessage> createVKey();
 
+  /** Static VKey factory method for use by VKeyTranslatorFactory. */
   public static VKey<PollMessage> createVKey(Key<PollMessage> key) {
     return VKey.create(PollMessage.class, key.getId(), key);
   }
@@ -289,7 +292,7 @@ public abstract class PollMessage extends ImmutableObject
 
     @Transient List<ContactTransferResponse> contactTransferResponses;
 
-    @Transient
+    @Transient @ImmutableObject.DoNotCompare
     List<DomainPendingActionNotificationResponse> domainPendingActionNotificationResponses;
 
     @Transient List<DomainTransferResponse> domainTransferResponses;
@@ -355,6 +358,11 @@ public abstract class PollMessage extends ImmutableObject
       return VKey.create(OneTime.class, getId(), Key.create(this));
     }
 
+    /** Converts an unspecialized VKey&lt;PollMessage&gt; to a VKey of the derived class. */
+    public static @Nullable VKey<OneTime> convertVKey(@Nullable VKey<OneTime> key) {
+      return key == null ? null : VKey.create(OneTime.class, key.getSqlKey(), key.getOfyKey());
+    }
+
     @Override
     public Builder asBuilder() {
       return new Builder(clone(this));
@@ -371,6 +379,47 @@ public abstract class PollMessage extends ImmutableObject
           .build();
     }
 
+    @Override
+    @OnLoad
+    void onLoad() {
+      super.onLoad();
+      if (!isNullOrEmpty(contactPendingActionNotificationResponses)) {
+        pendingActionNotificationResponse = contactPendingActionNotificationResponses.get(0);
+      }
+      if (!isNullOrEmpty(contactTransferResponses)) {
+        contactId = contactTransferResponses.get(0).getContactId();
+        transferResponse = contactTransferResponses.get(0);
+      }
+    }
+
+    @Override
+    @PostLoad
+    void postLoad() {
+      super.postLoad();
+      if (pendingActionNotificationResponse != null) {
+        contactPendingActionNotificationResponses =
+            ImmutableList.of(
+                ContactPendingActionNotificationResponse.create(
+                    pendingActionNotificationResponse.nameOrId.value,
+                    pendingActionNotificationResponse.getActionResult(),
+                    pendingActionNotificationResponse.getTrid(),
+                    pendingActionNotificationResponse.processedDate));
+      }
+      if (contactId != null && transferResponse != null) {
+        contactTransferResponses =
+            ImmutableList.of(
+                new ContactTransferResponse.Builder()
+                    .setContactId(contactId)
+                    .setGainingClientId(transferResponse.getGainingClientId())
+                    .setLosingClientId(transferResponse.getLosingClientId())
+                    .setTransferStatus(transferResponse.getTransferStatus())
+                    .setTransferRequestTime(transferResponse.getTransferRequestTime())
+                    .setPendingTransferExpirationTime(
+                        transferResponse.getPendingTransferExpirationTime())
+                    .build());
+      }
+    }
+
     /** A builder for {@link OneTime} since it is immutable. */
     public static class Builder extends PollMessage.Builder<OneTime, Builder> {
 
@@ -389,6 +438,10 @@ public abstract class PollMessage extends ImmutableObject
                     .filter(ContactPendingActionNotificationResponse.class::isInstance)
                     .map(ContactPendingActionNotificationResponse.class::cast)
                     .collect(toImmutableList()));
+        if (getInstance().contactPendingActionNotificationResponses != null) {
+          getInstance().pendingActionNotificationResponse =
+              getInstance().contactPendingActionNotificationResponses.get(0);
+        }
         getInstance().contactTransferResponses =
             forceEmptyToNull(
                 responseData
@@ -396,6 +449,11 @@ public abstract class PollMessage extends ImmutableObject
                     .filter(ContactTransferResponse.class::isInstance)
                     .map(ContactTransferResponse.class::cast)
                     .collect(toImmutableList()));
+        if (getInstance().contactTransferResponses != null) {
+          getInstance().contactId = getInstance().contactTransferResponses.get(0).getContactId();
+          getInstance().transferResponse = getInstance().contactTransferResponses.get(0);
+        }
+
         getInstance().domainPendingActionNotificationResponses =
             forceEmptyToNull(
                 responseData
@@ -456,6 +514,11 @@ public abstract class PollMessage extends ImmutableObject
       return VKey.create(Autorenew.class, getId(), Key.create(this));
     }
 
+    /** Converts an unspecialized VKey&lt;PollMessage&gt; to a VKey of the derived class. */
+    public static @Nullable VKey<Autorenew> convertVKey(VKey<Autorenew> key) {
+      return key == null ? null : VKey.create(Autorenew.class, key.getSqlKey(), key.getOfyKey());
+    }
+
     @Override
     public ImmutableList<ResponseData> getResponseData() {
       // Note that the event time is when the auto-renew occured, so the expiration time in the

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

@@ -97,7 +97,7 @@ 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().maybeLoad(VKey.create(RdeRevision.class, sqlKey, ofyKey));
+        tm().loadByKeyIfPresent(VKey.create(RdeRevision.class, sqlKey, ofyKey));
     return revisionOptional.map(rdeRevision -> rdeRevision.revision + 1).orElse(0);
   }
 
@@ -117,7 +117,7 @@ public final class RdeRevision extends BackupGroupRoot implements NonReplicatedE
     RdeRevisionId sqlKey = RdeRevisionId.create(tld, date.toLocalDate(), mode);
     Key<RdeRevision> ofyKey = Key.create(RdeRevision.class, triplet);
     Optional<RdeRevision> revisionOptional =
-        tm().maybeLoad(VKey.create(RdeRevision.class, sqlKey, ofyKey));
+        tm().loadByKeyIfPresent(VKey.create(RdeRevision.class, sqlKey, ofyKey));
     if (revision == 0) {
       revisionOptional.ifPresent(
           rdeRevision -> {

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

@@ -234,7 +234,7 @@ public class Registrar extends ImmutableObject
    * Unique registrar client id. Must conform to "clIDType" as defined in RFC5730.
    *
    * @see <a href="http://tools.ietf.org/html/rfc5730#section-4.2">Shared Structure Schema</a>
-   *     <p>TODO(shicong): Rename this field to clientId
+   *     <p>TODO(b/177568946): Rename this field to registrarId.
    */
   @Id
   @javax.persistence.Id
@@ -537,20 +537,24 @@ public class Registrar extends ImmutableObject
     return LIVE_STATES.contains(state) && PUBLICLY_VISIBLE_TYPES.contains(type);
   }
 
-  public String getClientCertificate() {
-    return clientCertificate;
+  /** Returns the client certificate string if it has been set, or empty otherwise. */
+  public Optional<String> getClientCertificate() {
+    return Optional.ofNullable(clientCertificate);
   }
 
-  public String getClientCertificateHash() {
-    return clientCertificateHash;
+  /** Returns the client certificate hash if it has been set, or empty otherwise. */
+  public Optional<String> getClientCertificateHash() {
+    return Optional.ofNullable(clientCertificateHash);
   }
 
-  public String getFailoverClientCertificate() {
-    return failoverClientCertificate;
+  /** Returns the failover client certificate string if it has been set, or empty otherwise. */
+  public Optional<String> getFailoverClientCertificate() {
+    return Optional.ofNullable(failoverClientCertificate);
   }
 
-  public String getFailoverClientCertificateHash() {
-    return failoverClientCertificateHash;
+  /** Returns the failover client certificate hash if it has been set, or empty otherwise. */
+  public Optional<String> getFailoverClientCertificateHash() {
+    return Optional.ofNullable(failoverClientCertificateHash);
   }
 
   public ImmutableList<CidrAddressBlock> getIpAddressAllowList() {
@@ -815,7 +819,8 @@ public class Registrar extends ImmutableObject
               .map(Registry::createVKey)
               .collect(toImmutableSet());
       Set<VKey<Registry>> missingTldKeys =
-          Sets.difference(newTldKeys, transactIfJpaTm(() -> tm().load(newTldKeys)).keySet());
+          Sets.difference(
+              newTldKeys, transactIfJpaTm(() -> tm().loadByKeysIfPresent(newTldKeys)).keySet());
       checkArgument(missingTldKeys.isEmpty(), "Trying to set nonexisting TLDs: %s", missingTldKeys);
       getInstance().allowedTlds = ImmutableSortedSet.copyOf(allowedTlds);
       return this;
@@ -983,7 +988,7 @@ public class Registrar extends ImmutableObject
   public static Iterable<Registrar> loadAll() {
     return tm().isOfy()
         ? ImmutableList.copyOf(ofy().load().type(Registrar.class).ancestor(getCrossTldKey()))
-        : tm().transact(() -> tm().loadAll(Registrar.class));
+        : tm().transact(() -> tm().loadAllOf(Registrar.class));
   }
 
   /** Loads all registrar entities using an in-memory cache. */
@@ -994,7 +999,7 @@ public class Registrar extends ImmutableObject
   /** Loads and returns a registrar entity by its client id directly from Datastore. */
   public static Optional<Registrar> loadByClientId(String clientId) {
     checkArgument(!Strings.isNullOrEmpty(clientId), "clientId must be specified");
-    return transactIfJpaTm(() -> tm().maybeLoad(createVKey(clientId)));
+    return transactIfJpaTm(() -> tm().loadByKeyIfPresent(createVKey(clientId)));
   }
 
   /**

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

@@ -69,7 +69,7 @@ public final class Registries {
                                   .stream()
                                   .map(Key::getName)
                                   .collect(toImmutableSet())
-                              : tm().loadAll(Registry.class).stream()
+                              : tm().loadAllOf(Registry.class).stream()
                                   .map(Registry::getTldStr)
                                   .collect(toImmutableSet());
                       return Registry.getAll(tlds).stream()

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

@@ -267,7 +267,7 @@ public class Registry extends ImmutableObject implements Buildable, DatastoreAnd
                 public Optional<Registry> load(final String tld) {
                   // Enter a transaction-less context briefly; we don't want to enroll every TLD in
                   // a transaction that might be wrapping this call.
-                  return tm().doTransactionless(() -> tm().maybeLoad(createVKey(tld)));
+                  return tm().doTransactionless(() -> tm().loadByKeyIfPresent(createVKey(tld)));
                 }
 
                 @Override
@@ -275,7 +275,7 @@ public class Registry extends ImmutableObject implements Buildable, DatastoreAnd
                   ImmutableMap<String, VKey<Registry>> keysMap =
                       toMap(ImmutableSet.copyOf(tlds), Registry::createVKey);
                   Map<VKey<? extends Registry>, Registry> entities =
-                      tm().doTransactionless(() -> tm().load(keysMap.values()));
+                      tm().doTransactionless(() -> tm().loadByKeys(keysMap.values()));
                   return Maps.transformEntries(
                       keysMap, (k, v) -> Optional.ofNullable(entities.getOrDefault(v, null)));
                 }

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

@@ -62,7 +62,7 @@ public class ReservedListDualWriteDao {
   public static Optional<ReservedList> getLatestRevision(String reservedListName) {
     Optional<ReservedList> maybeDatastoreList =
         ofyTm()
-            .maybeLoad(
+            .loadByKeyIfPresent(
                 VKey.createOfy(
                     ReservedList.class,
                     Key.create(getCrossTldKey(), ReservedList.class, reservedListName)));

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

@@ -49,7 +49,9 @@ public class DomainTransactionRecord extends ImmutableObject
 
   @Id
   @Ignore
+  @ImmutableObject.DoNotCompare
   @GeneratedValue(strategy = GenerationType.IDENTITY)
+  @ImmutableObject.Insignificant
   Long id;
 
   /** The TLD this record operates on. */

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

@@ -198,6 +198,7 @@ public class HistoryEntry extends ImmutableObject implements Buildable, Datastor
    * transaction counts (such as contact or host mutations).
    */
   @Transient // domain-specific
+  @ImmutableObject.EmptySetToNull
   protected Set<DomainTransactionRecord> domainTransactionRecords;
 
   public long getId() {

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

@@ -0,0 +1,144 @@
+// Copyright 2020 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.model.reporting;
+
+import static com.google.common.collect.ImmutableList.toImmutableList;
+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.END_OF_TIME;
+import static google.registry.util.DateTimeUtils.START_OF_TIME;
+
+import com.google.common.collect.Iterables;
+import google.registry.model.EppResource;
+import google.registry.model.contact.ContactHistory;
+import google.registry.model.contact.ContactResource;
+import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.DomainHistory;
+import google.registry.model.host.HostHistory;
+import google.registry.model.host.HostResource;
+import google.registry.persistence.VKey;
+import java.util.Comparator;
+import javax.persistence.EntityManager;
+import org.joda.time.DateTime;
+
+/**
+ * Retrieves {@link HistoryEntry} descendants (e.g. {@link DomainHistory}).
+ *
+ * <p>This class is configured to retrieve either from Datastore or SQL, depending on which database
+ * is currently considered the primary database.
+ */
+public class HistoryEntryDao {
+
+  /** Loads all history objects in the times specified, including all types. */
+  public static Iterable<? extends HistoryEntry> loadAllHistoryObjects(
+      DateTime afterTime, DateTime beforeTime) {
+    if (tm().isOfy()) {
+      return ofy()
+          .load()
+          .type(HistoryEntry.class)
+          .order("modificationTime")
+          .filter("modificationTime >=", afterTime)
+          .filter("modificationTime <=", beforeTime);
+    } else {
+      return jpaTm()
+          .transact(
+              () ->
+                  Iterables.concat(
+                      loadAllHistoryObjectsFromSql(ContactHistory.class, afterTime, beforeTime),
+                      loadAllHistoryObjectsFromSql(DomainHistory.class, afterTime, beforeTime),
+                      loadAllHistoryObjectsFromSql(HostHistory.class, afterTime, beforeTime)));
+    }
+  }
+
+  /** Loads all history objects corresponding to the given {@link EppResource}. */
+  public static Iterable<? extends HistoryEntry> loadHistoryObjectsForResource(
+      VKey<? extends EppResource> parentKey) {
+    return loadHistoryObjectsForResource(parentKey, START_OF_TIME, END_OF_TIME);
+  }
+
+  /** Loads all history objects in the time period specified for the given {@link EppResource}. */
+  public static Iterable<? extends HistoryEntry> loadHistoryObjectsForResource(
+      VKey<? extends EppResource> parentKey, DateTime afterTime, DateTime beforeTime) {
+    if (tm().isOfy()) {
+      return ofy()
+          .load()
+          .type(HistoryEntry.class)
+          .ancestor(parentKey.getOfyKey())
+          .order("modificationTime")
+          .filter("modificationTime >=", afterTime)
+          .filter("modificationTime <=", beforeTime);
+    } else {
+      return jpaTm()
+          .transact(() -> loadHistoryObjectsForResourceFromSql(parentKey, afterTime, beforeTime));
+    }
+  }
+
+  private static Iterable<? extends HistoryEntry> loadHistoryObjectsForResourceFromSql(
+      VKey<? extends EppResource> parentKey, DateTime afterTime, DateTime beforeTime) {
+    Class<? extends HistoryEntry> historyClass = getHistoryClassFromParent(parentKey.getKind());
+    String repoIdFieldName = getRepoIdFieldNameFromHistoryClass(historyClass);
+    EntityManager entityManager = jpaTm().getEntityManager();
+    String tableName = entityManager.getMetamodel().entity(historyClass).getName();
+    String queryString =
+        String.format(
+            "SELECT entry FROM %s entry WHERE entry.modificationTime >= :afterTime AND "
+                + "entry.modificationTime <= :beforeTime AND entry.%s = :parentKey",
+            tableName, repoIdFieldName);
+    return entityManager
+        .createQuery(queryString, historyClass)
+        .setParameter("afterTime", afterTime)
+        .setParameter("beforeTime", beforeTime)
+        .setParameter("parentKey", parentKey.getSqlKey().toString())
+        .getResultStream()
+        .sorted(Comparator.comparing(HistoryEntry::getModificationTime))
+        .collect(toImmutableList());
+  }
+
+  private static Class<? extends HistoryEntry> getHistoryClassFromParent(
+      Class<? extends EppResource> parent) {
+    if (parent.equals(ContactResource.class)) {
+      return ContactHistory.class;
+    } else if (parent.equals(DomainBase.class)) {
+      return DomainHistory.class;
+    } else if (parent.equals(HostResource.class)) {
+      return HostHistory.class;
+    }
+    throw new IllegalArgumentException(
+        String.format("Unknown history type for parent %s", parent.getName()));
+  }
+
+  private static String getRepoIdFieldNameFromHistoryClass(
+      Class<? extends HistoryEntry> historyClass) {
+    return historyClass.equals(ContactHistory.class)
+        ? "contactRepoId"
+        : historyClass.equals(DomainHistory.class) ? "domainRepoId" : "hostRepoId";
+  }
+
+  private static Iterable<? extends HistoryEntry> loadAllHistoryObjectsFromSql(
+      Class<? extends HistoryEntry> historyClass, DateTime afterTime, DateTime beforeTime) {
+    EntityManager entityManager = jpaTm().getEntityManager();
+    return entityManager
+        .createQuery(
+            String.format(
+                "SELECT entry FROM %s entry WHERE entry.modificationTime >= :afterTime AND "
+                    + "entry.modificationTime <= :beforeTime",
+                entityManager.getMetamodel().entity(historyClass).getName()),
+            historyClass)
+        .setParameter("afterTime", afterTime)
+        .setParameter("beforeTime", beforeTime)
+        .getResultList();
+  }
+}

core/src/main/java/google/registry/model/server/ServerSecret.java

@@ -77,7 +77,7 @@ public class ServerSecret extends CrossTldSingleton implements NonReplicatedEnti
               // transactionally create a new ServerSecret (once per app setup) if necessary.
               // return the ofy() result during Datastore-primary phase
               ServerSecret secret =
-                  ofyTm().maybeLoad(key).orElseGet(() -> create(UUID.randomUUID()));
+                  ofyTm().loadByKeyIfPresent(key).orElseGet(() -> create(UUID.randomUUID()));
               // During a dual-write period, write it to both Datastore and SQL
               // even if we didn't have to retrieve it from the DB
               ofyTm().transact(() -> ofyTm().putWithoutBackup(secret));

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

@@ -58,7 +58,7 @@ public class ClaimsListDao {
 
   /**
    * Returns the most recent revision of the {@link ClaimsListShard} in Cloud SQL, if it exists.
-   * TODO(shicong): Change this method to package level access after dual-read phase.
+   * TODO(b/177569979): Change this method to package level access after dual-read phase.
    * ClaimsListShard uses this method to retrieve claims list in Cloud SQL for the comparison, and
    * ClaimsListShard is not in this package.
    */

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

@@ -56,7 +56,7 @@ public final class TmchCrl extends CrossTldSingleton implements NonReplicatedEnt
         VKey.create(
             TmchCrl.class, SINGLETON_ID, Key.create(getCrossTldKey(), TmchCrl.class, SINGLETON_ID));
     // return the ofy() result during Datastore-primary phase
-    return ofyTm().transact(() -> ofyTm().maybeLoad(key));
+    return ofyTm().transact(() -> ofyTm().loadByKeyIfPresent(key));
   }
 
   /**

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

@@ -14,15 +14,25 @@
 
 package google.registry.model.transfer;
 
+import static com.google.common.collect.ImmutableList.toImmutableList;
+import static google.registry.util.CollectionUtils.isNullOrEmpty;
+import static google.registry.util.CollectionUtils.nullToEmpty;
 import static google.registry.util.CollectionUtils.nullToEmptyImmutableCopy;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
+import com.googlecode.objectify.Key;
+import com.googlecode.objectify.annotation.AlsoLoad;
 import com.googlecode.objectify.annotation.Ignore;
 import com.googlecode.objectify.annotation.IgnoreSave;
 import com.googlecode.objectify.condition.IfNull;
 import google.registry.model.Buildable;
 import google.registry.model.EppResource;
+import google.registry.model.contact.ContactResource;
+import google.registry.model.domain.DomainBase;
 import google.registry.model.eppcommon.Trid;
+import google.registry.model.poll.PollMessage;
+import google.registry.model.reporting.HistoryEntry;
 import google.registry.persistence.VKey;
 import google.registry.util.TypeUtils.TypeInstantiator;
 import java.util.Set;
@@ -32,6 +42,7 @@ import javax.persistence.AttributeOverrides;
 import javax.persistence.Column;
 import javax.persistence.Embedded;
 import javax.persistence.MappedSuperclass;
+import javax.persistence.PostLoad;
 import javax.persistence.Transient;
 
 /**
@@ -68,15 +79,27 @@ public abstract class TransferData<
   @IgnoreSave(IfNull.class)
   Set<VKey<? extends TransferServerApproveEntity>> serverApproveEntities;
 
-  // The following 3 fields are the replacement for serverApproveEntities in Cloud SQL.
-  // TODO(shicong): Add getter/setter for these 3 fields and use them in the application code.
   @Ignore
-  @Column(name = "transfer_gaining_poll_message_id")
-  Long gainingTransferPollMessageId;
+  @Column(name = "transfer_repo_id")
+  String repoId;
 
   @Ignore
-  @Column(name = "transfer_losing_poll_message_id")
-  Long losingTransferPollMessageId;
+  @Column(name = "transfer_history_entry_id")
+  Long historyEntryId;
+
+  // The pollMessageId1 and pollMessageId2 are used to store the IDs for gaining and losing poll
+  // messages in Cloud SQL, and they are added to replace the VKeys in serverApproveEntities.
+  // Although we can distinguish which is which when we construct the TransferData instance from
+  // the transfer request flow, when the instance is loaded from Datastore, we cannot make this
+  // distinction because they are just VKeys. Also, the only way we use serverApproveEntities is to
+  // just delete all the entities referenced by the VKeys, so we don't need to make the distinction.
+  @Ignore
+  @Column(name = "transfer_poll_message_id_1")
+  Long pollMessageId1;
+
+  @Ignore
+  @Column(name = "transfer_poll_message_id_2")
+  Long pollMessageId2;
 
   public abstract boolean isEmpty();
 
@@ -116,6 +139,83 @@ public abstract class TransferData<
     return newBuilder;
   }
 
+  void onLoad(
+      @AlsoLoad("serverApproveEntities")
+          Set<VKey<? extends TransferServerApproveEntity>> serverApproveEntities) {
+    mapServerApproveEntitiesToFields(serverApproveEntities, this);
+  }
+
+  @PostLoad
+  void postLoad() {
+    mapFieldsToServerApproveEntities();
+  }
+
+  /**
+   * Reconstructs serverApproveEntities set from the individual fields, e.g. repoId, historyEntryId,
+   * pollMessageId1.
+   */
+  void mapFieldsToServerApproveEntities() {
+    if (repoId == null) {
+      return;
+    }
+    Key<? extends EppResource> eppKey;
+    if (getClass().equals(DomainBase.class)) {
+      eppKey = Key.create(DomainBase.class, repoId);
+    } else {
+      eppKey = Key.create(ContactResource.class, repoId);
+    }
+    Key<HistoryEntry> historyEntryKey = Key.create(eppKey, HistoryEntry.class, historyEntryId);
+    ImmutableSet.Builder<VKey<? extends TransferServerApproveEntity>> entityKeysBuilder =
+        new ImmutableSet.Builder<>();
+    if (pollMessageId1 != null) {
+      Key<PollMessage> ofyKey = Key.create(historyEntryKey, PollMessage.class, pollMessageId1);
+      entityKeysBuilder.add(PollMessage.createVKey(ofyKey));
+    }
+    if (pollMessageId2 != null) {
+      Key<PollMessage> ofyKey = Key.create(historyEntryKey, PollMessage.class, pollMessageId2);
+      entityKeysBuilder.add(PollMessage.createVKey(ofyKey));
+    }
+    serverApproveEntities = entityKeysBuilder.build();
+  }
+
+  /** Maps serverApproveEntities set to the individual fields. */
+  static void mapServerApproveEntitiesToFields(
+      Set<VKey<? extends TransferServerApproveEntity>> serverApproveEntities,
+      TransferData transferData) {
+    if (isNullOrEmpty(serverApproveEntities)) {
+      transferData.historyEntryId = null;
+      transferData.repoId = null;
+      transferData.pollMessageId1 = null;
+      transferData.pollMessageId2 = null;
+      return;
+    }
+    // Each element in serverApproveEntities should have the exact same Key<HistoryEntry> as its
+    // parent. So, we can use any to set historyEntryId and repoId.
+    Key<?> key = serverApproveEntities.iterator().next().getOfyKey();
+    transferData.historyEntryId = key.getParent().getId();
+    transferData.repoId = key.getParent().getParent().getName();
+
+    ImmutableList<Long> sortedPollMessageIds = getSortedPollMessageIds(serverApproveEntities);
+    if (sortedPollMessageIds.size() >= 1) {
+      transferData.pollMessageId1 = sortedPollMessageIds.get(0);
+    }
+    if (sortedPollMessageIds.size() >= 2) {
+      transferData.pollMessageId2 = sortedPollMessageIds.get(1);
+    }
+  }
+
+  /**
+   * Gets poll message IDs from the given serverApproveEntities and sorted the IDs in natural order.
+   */
+  private static ImmutableList<Long> getSortedPollMessageIds(
+      Set<VKey<? extends TransferServerApproveEntity>> serverApproveEntities) {
+    return nullToEmpty(serverApproveEntities).stream()
+        .filter(vKey -> PollMessage.class.isAssignableFrom(vKey.getKind()))
+        .map(vKey -> (long) vKey.getSqlKey())
+        .sorted()
+        .collect(toImmutableList());
+  }
+
   /** Builder for {@link TransferData} because it is immutable. */
   public abstract static class Builder<T extends TransferData, B extends Builder<T, B>>
       extends BaseTransferObject.Builder<T, B> {
@@ -141,6 +241,7 @@ public abstract class TransferData<
 
     @Override
     public T build() {
+      mapServerApproveEntitiesToFields(getInstance().serverApproveEntities, getInstance());
       return super.build();
     }
   }

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

@@ -27,6 +27,7 @@ import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Maps;
 import com.google.common.flogger.FluentLogger;
+import dagger.BindsOptionalOf;
 import dagger.Module;
 import dagger.Provides;
 import google.registry.config.RegistryConfig.Config;
@@ -42,8 +43,12 @@ import google.registry.privileges.secretmanager.SqlUser.RobotUser;
 import google.registry.tools.AuthModule.CloudSqlClientCredential;
 import google.registry.util.Clock;
 import java.lang.annotation.Documented;
+import java.sql.Connection;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Optional;
+import javax.annotation.Nullable;
+import javax.inject.Provider;
 import javax.inject.Qualifier;
 import javax.inject.Singleton;
 import javax.persistence.EntityManagerFactory;
@@ -52,7 +57,7 @@ import org.hibernate.cfg.Environment;
 
 /** Dagger module class for the persistence layer. */
 @Module
-public class PersistenceModule {
+public abstract class PersistenceModule {
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
 
   // This name must be the same as the one defined in persistence.xml.
@@ -102,25 +107,48 @@ public class PersistenceModule {
       @Config("cloudSqlJdbcUrl") String jdbcUrl,
       @Config("cloudSqlInstanceConnectionName") String instanceConnectionName,
       @DefaultHibernateConfigs ImmutableMap<String, String> defaultConfigs) {
-    return createPartialSqlConfigs(jdbcUrl, instanceConnectionName, defaultConfigs);
+    return createPartialSqlConfigs(
+        jdbcUrl, instanceConnectionName, defaultConfigs, Optional.empty());
   }
 
+  /**
+   * Optionally overrides the isolation level in the config file.
+   *
+   * <p>The binding for {@link TransactionIsolationLevel} may be {@link Nullable}. As a result, it
+   * is a compile-time error to inject {@code Optional<TransactionIsolation>} (See {@link
+   * BindsOptionalOf} for more information). User should inject {@code
+   * Optional<Provider<TransactionIsolation>>} instead.
+   */
+  @BindsOptionalOf
+  @Config("beamIsolationOverride")
+  abstract TransactionIsolationLevel bindBeamIsolationOverride();
+
   @Provides
   @Singleton
   @BeamPipelineCloudSqlConfigs
   static ImmutableMap<String, String> provideBeamPipelineCloudSqlConfigs(
       @Config("beamCloudSqlJdbcUrl") String jdbcUrl,
       @Config("beamCloudSqlInstanceConnectionName") String instanceConnectionName,
-      @DefaultHibernateConfigs ImmutableMap<String, String> defaultConfigs) {
-    return createPartialSqlConfigs(jdbcUrl, instanceConnectionName, defaultConfigs);
+      @DefaultHibernateConfigs ImmutableMap<String, String> defaultConfigs,
+      @Config("beamIsolationOverride")
+          Optional<Provider<TransactionIsolationLevel>> isolationOverride) {
+    return createPartialSqlConfigs(
+        jdbcUrl, instanceConnectionName, defaultConfigs, isolationOverride);
   }
 
-  private static ImmutableMap<String, String> createPartialSqlConfigs(
-      String jdbcUrl, String instanceConnectionName, ImmutableMap<String, String> defaultConfigs) {
+  @VisibleForTesting
+  static ImmutableMap<String, String> createPartialSqlConfigs(
+      String jdbcUrl,
+      String instanceConnectionName,
+      ImmutableMap<String, String> defaultConfigs,
+      Optional<Provider<TransactionIsolationLevel>> isolationOverride) {
     HashMap<String, String> overrides = Maps.newHashMap(defaultConfigs);
     overrides.put(Environment.URL, jdbcUrl);
     overrides.put(HIKARI_DS_SOCKET_FACTORY, "com.google.cloud.sql.postgres.SocketFactory");
     overrides.put(HIKARI_DS_CLOUD_SQL_INSTANCE, instanceConnectionName);
+    isolationOverride
+        .map(Provider::get)
+        .ifPresent(override -> overrides.put(Environment.ISOLATION, override.name()));
     return ImmutableMap.copyOf(overrides);
   }
 
@@ -254,6 +282,37 @@ public class PersistenceModule {
     }
   }
 
+  /**
+   * Transaction isolation levels supported by Cloud SQL (mysql and postgresql).
+   *
+   * <p>Enum names may be used for property-based configuration, and must match the corresponding
+   * variable names in {@link Connection}.
+   */
+  public enum TransactionIsolationLevel {
+    TRANSACTION_READ_UNCOMMITTED,
+    TRANSACTION_READ_COMMITTED,
+    TRANSACTION_REPEATABLE_READ,
+    TRANSACTION_SERIALIZABLE;
+
+    private final int value;
+
+    TransactionIsolationLevel() {
+      try {
+        // name() is final in parent class (Enum.java), therefore safe to call in constructor.
+        value = Connection.class.getField(name()).getInt(null);
+      } catch (Exception e) {
+        throw new IllegalStateException(
+            String.format(
+                "%s Enum name %s has no matching public field in java.sql.Connection.",
+                getClass().getSimpleName(), name()));
+      }
+    }
+
+    public final int getValue() {
+      return value;
+    }
+  }
+
   /** Dagger qualifier for {@link JpaTransactionManager} used for App Engine application. */
   @Qualifier
   @Documented
@@ -292,8 +351,8 @@ public class PersistenceModule {
   @interface BeamPipelineCloudSqlConfigs {}
 
   /** Dagger qualifier for the default Hibernate configurations. */
-  // TODO(shicong): Change annotations in this class to none public or put them in a top level
-  //  package
+  // TODO(b/177568911): Change annotations in this class to none public or put them in a top level
+  // package.
   @Qualifier
   @Documented
   public @interface DefaultHibernateConfigs {}

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

@@ -20,7 +20,6 @@ import javax.persistence.Converter;
 
 /**
  * JPA {@link AttributeConverter} for storing/retrieving {@code List<CidrAddressBlock>} objects.
- * TODO(shicong): Investigate if we can have one converter for any List type
  */
 @Converter(autoApply = true)
 public class CidrAddressBlockListConverter extends StringListConverterBase<CidrAddressBlock> {

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

@@ -18,6 +18,7 @@ import static com.google.common.base.Preconditions.checkArgument;
 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.model.ofy.DatastoreTransactionManager.toChildHistoryEntryIfPossible;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 import static java.util.AbstractMap.SimpleEntry;
 import static java.util.stream.Collectors.joining;
@@ -34,13 +35,14 @@ import google.registry.model.index.EppResourceIndex;
 import google.registry.model.index.ForeignKeyIndex.ForeignKeyContactIndex;
 import google.registry.model.index.ForeignKeyIndex.ForeignKeyDomainIndex;
 import google.registry.model.index.ForeignKeyIndex.ForeignKeyHostIndex;
+import google.registry.model.ofy.DatastoreTransactionManager;
 import google.registry.persistence.JpaRetries;
 import google.registry.persistence.VKey;
 import google.registry.util.Clock;
 import google.registry.util.Retrier;
 import google.registry.util.SystemSleeper;
 import java.lang.reflect.Field;
-import java.util.Map.Entry;
+import java.util.Map;
 import java.util.NoSuchElementException;
 import java.util.Optional;
 import java.util.function.Supplier;
@@ -76,8 +78,9 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   // EntityManagerFactory is thread safe.
   private final EntityManagerFactory emf;
   private final Clock clock;
-  // TODO(shicong): Investigate alternatives for managing transaction information. ThreadLocal adds
-  //  an unnecessary restriction that each request has to be processed by one thread synchronously.
+  // TODO(b/177588434): Investigate alternatives for managing transaction information. ThreadLocal
+  // adds an unnecessary restriction that each request has to be processed by one thread
+  // synchronously.
   private final ThreadLocal<TransactionInfo> transactionInfo =
       ThreadLocal.withInitial(TransactionInfo::new);
 
@@ -115,8 +118,6 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
 
   @Override
   public <T> T transact(Supplier<T> work) {
-    // TODO(shicong): Investigate removing transactNew functionality after migration as it may
-    //  be same as this one.
     return retrier.callWithRetry(
         () -> {
           if (inTransaction()) {
@@ -196,6 +197,8 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
         });
   }
 
+  // TODO(b/177674699): Remove all transactNew methods as they are same as transact after the
+  // database migration.
   @Override
   public <T> T transactNew(Supplier<T> work) {
     return transact(work);
@@ -249,8 +252,10 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
       return;
     }
     assertInTransaction();
-    getEntityManager().persist(entity);
-    transactionInfo.get().addUpdate(entity);
+    // Necessary due to the changes in HistoryEntry representation during the migration to SQL
+    Object toPersist = toChildHistoryEntryIfPossible(entity);
+    getEntityManager().persist(toPersist);
+    transactionInfo.get().addUpdate(toPersist);
   }
 
   @Override
@@ -277,8 +282,10 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
       return;
     }
     assertInTransaction();
-    getEntityManager().merge(entity);
-    transactionInfo.get().addUpdate(entity);
+    // Necessary due to the changes in HistoryEntry representation during the migration to SQL
+    Object toPersist = toChildHistoryEntryIfPossible(entity);
+    getEntityManager().merge(toPersist);
+    transactionInfo.get().addUpdate(toPersist);
   }
 
   @Override
@@ -306,8 +313,10 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     }
     assertInTransaction();
     checkArgument(exists(entity), "Given entity does not exist");
-    getEntityManager().merge(entity);
-    transactionInfo.get().addUpdate(entity);
+    // Necessary due to the changes in HistoryEntry representation during the migration to SQL
+    Object toPersist = toChildHistoryEntryIfPossible(entity);
+    getEntityManager().merge(toPersist);
+    transactionInfo.get().addUpdate(toPersist);
   }
 
   @Override
@@ -338,6 +347,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   @Override
   public boolean exists(Object entity) {
     checkArgumentNotNull(entity, "entity must be specified");
+    entity = toChildHistoryEntryIfPossible(entity);
     EntityType<?> entityType = getEntityType(entity.getClass());
     ImmutableSet<EntityId> entityIds = getEntityIdsFromEntity(entityType, entity);
     return exists(entityType.getName(), entityIds);
@@ -356,33 +366,15 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   }
 
   @Override
-  public <T> Optional<T> maybeLoad(VKey<T> key) {
+  public <T> Optional<T> loadByKeyIfPresent(VKey<T> key) {
     checkArgumentNotNull(key, "key must be specified");
     assertInTransaction();
     return Optional.ofNullable(getEntityManager().find(key.getKind(), key.getSqlKey()));
   }
 
   @Override
-  public <T> T load(VKey<T> key) {
-    checkArgumentNotNull(key, "key must be specified");
-    assertInTransaction();
-    T result = getEntityManager().find(key.getKind(), key.getSqlKey());
-    if (result == null) {
-      throw new NoSuchElementException(key.toString());
-    }
-    return result;
-  }
-
-  @Override
-  public <T> T load(T entity) {
-    checkArgumentNotNull(entity, "entity must be specified");
-    assertInTransaction();
-    return (T)
-        load(VKey.createSql(entity.getClass(), emf.getPersistenceUnitUtil().getIdentifier(entity)));
-  }
-
-  @Override
-  public <T> ImmutableMap<VKey<? extends T>, T> load(Iterable<? extends VKey<? extends T>> keys) {
+  public <T> ImmutableMap<VKey<? extends T>, T> loadByKeysIfPresent(
+      Iterable<? extends VKey<? extends T>> keys) {
     checkArgumentNotNull(keys, "keys must be specified");
     assertInTransaction();
     return StreamSupport.stream(keys.spliterator(), false)
@@ -393,11 +385,64 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
                 new SimpleEntry<VKey<? extends T>, T>(
                     key, getEntityManager().find(key.getKind(), key.getSqlKey())))
         .filter(entry -> entry.getValue() != null)
-        .collect(toImmutableMap(Entry::getKey, Entry::getValue));
+        .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue));
   }
 
   @Override
-  public <T> ImmutableList<T> loadAll(Class<T> clazz) {
+  public <T> ImmutableList<T> loadByEntitiesIfPresent(Iterable<T> entities) {
+    return Streams.stream(entities)
+        .map(DatastoreTransactionManager::toChildHistoryEntryIfPossible)
+        .filter(this::exists)
+        .map(this::loadByEntity)
+        .collect(toImmutableList());
+  }
+
+  @Override
+  public <T> T loadByKey(VKey<T> key) {
+    checkArgumentNotNull(key, "key must be specified");
+    assertInTransaction();
+    T result = getEntityManager().find(key.getKind(), key.getSqlKey());
+    if (result == null) {
+      throw new NoSuchElementException(key.toString());
+    }
+    return result;
+  }
+
+  @Override
+  public <T> ImmutableMap<VKey<? extends T>, T> loadByKeys(
+      Iterable<? extends VKey<? extends T>> keys) {
+    ImmutableMap<VKey<? extends T>, T> existing = loadByKeysIfPresent(keys);
+    ImmutableSet<? extends VKey<? extends T>> missingKeys =
+        Streams.stream(keys).filter(k -> !existing.containsKey(k)).collect(toImmutableSet());
+    if (!missingKeys.isEmpty()) {
+      throw new NoSuchElementException(
+          String.format(
+              "Expected to find the following VKeys but they were missing: %s.", missingKeys));
+    }
+    return existing;
+  }
+
+  @Override
+  public <T> T loadByEntity(T entity) {
+    checkArgumentNotNull(entity, "entity must be specified");
+    assertInTransaction();
+    entity = toChildHistoryEntryIfPossible(entity);
+    // If the caller requested a HistoryEntry, load the corresponding *History class
+    T possibleChild = toChildHistoryEntryIfPossible(entity);
+    return (T)
+        loadByKey(
+            VKey.createSql(
+                possibleChild.getClass(),
+                emf.getPersistenceUnitUtil().getIdentifier(possibleChild)));
+  }
+
+  @Override
+  public <T> ImmutableList<T> loadByEntities(Iterable<T> entities) {
+    return Streams.stream(entities).map(this::loadByEntity).collect(toImmutableList());
+  }
+
+  @Override
+  public <T> ImmutableList<T> loadAllOf(Class<T> clazz) {
     checkArgumentNotNull(clazz, "clazz must be specified");
     assertInTransaction();
     return ImmutableList.copyOf(
@@ -408,11 +453,6 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
             .getResultList());
   }
 
-  @Override
-  public <T> ImmutableList<T> loadAll(Iterable<T> entities) {
-    return Streams.stream(entities).map(this::load).collect(toImmutableList());
-  }
-
   private int internalDelete(VKey<?> key) {
     checkArgumentNotNull(key, "key must be specified");
     assertInTransaction();
@@ -444,6 +484,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
       return;
     }
     assertInTransaction();
+    entity = toChildHistoryEntryIfPossible(entity);
     Object managedEntity = entity;
     if (!getEntityManager().contains(entity)) {
       managedEntity = getEntityManager().merge(entity);

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

@@ -29,17 +29,19 @@ import org.joda.time.DateTime;
  */
 public interface TransactionManager {
 
-  /** Returns {@code true} if the caller is in a transaction.
+  /**
+   * Returns {@code true} if the caller is in a transaction.
    *
-   *  <p>Note that this function is kept for backward compatibility. We will review the use case
-   *  later when adding the cloud sql implementation.
+   * <p>Note that this function is kept for backward compatibility. We will review the use case
+   * later when adding the cloud sql implementation.
    */
   boolean inTransaction();
 
-  /** Throws {@link IllegalStateException} if the caller is not in a transaction.
+  /**
+   * Throws {@link IllegalStateException} if the caller is not in a transaction.
    *
-   *  <p>Note that this function is kept for backward compatibility. We will review the use case
-   *  later when adding the cloud sql implementation.
+   * <p>Note that this function is kept for backward compatibility. We will review the use case
+   * later when adding the cloud sql implementation.
    */
   void assertInTransaction();
 
@@ -58,10 +60,11 @@ public interface TransactionManager {
    */
   <T> T transactNew(Supplier<T> work);
 
-  /** Pauses the current transaction (if any) and executes the work in a new transaction.
+  /**
+   * Pauses the current transaction (if any) and executes the work in a new transaction.
    *
-   *  <p>Note that this function is kept for backward compatibility. We will review the use case
-   *  later when adding the cloud sql implementation.
+   * <p>Note that this function is kept for backward compatibility. We will review the use case
+   * later when adding the cloud sql implementation.
    */
   void transactNew(Runnable work);
 
@@ -73,10 +76,11 @@ public interface TransactionManager {
    */
   <R> R transactNewReadOnly(Supplier<R> work);
 
-  /** Executes the work in a read-only transaction.
+  /**
+   * Executes the work in a read-only transaction.
    *
-   *  <p>Note that this function is kept for backward compatibility. We will review the use case
-   *  later when adding the cloud sql implementation.
+   * <p>Note that this function is kept for backward compatibility. We will review the use case
+   * later when adding the cloud sql implementation.
    */
   void transactNewReadOnly(Runnable work);
 
@@ -182,31 +186,60 @@ public interface TransactionManager {
   /** Returns whether the entity of given key exists. */
   <T> boolean exists(VKey<T> key);
 
-  /** Loads the entity by its id, returns empty if the entity doesn't exist. */
-  <T> Optional<T> maybeLoad(VKey<T> key);
-
-  /** Loads the entity by its id, throws NoSuchElementException if it doesn't exist. */
-  <T> T load(VKey<T> key);
+  /** Loads the entity by its key, returns empty if the entity doesn't exist. */
+  <T> Optional<T> loadByKeyIfPresent(VKey<T> key);
 
   /**
-   * Loads the given entity from the database, throws NoSuchElementException if it doesn't exist.
-   */
-  <T> T load(T entity);
-
-  /**
-   * Loads the set of entities by their key id.
+   * Loads the set of entities by their keys.
    *
-   * @throws NoSuchElementException if any of the keys are not found.
+   * <p>Nonexistent keys / entities are absent from the resulting map, but no {@link
+   * NoSuchElementException} will be thrown.
    */
-  <T> ImmutableMap<VKey<? extends T>, T> load(Iterable<? extends VKey<? extends T>> keys);
-
-  /** Loads all entities of the given type, returns empty if there is no such entity. */
-  <T> ImmutableList<T> loadAll(Class<T> clazz);
+  <T> ImmutableMap<VKey<? extends T>, T> loadByKeysIfPresent(
+      Iterable<? extends VKey<? extends T>> keys);
 
   /**
-   * Loads all given entities from the database, throws NoSuchElementException if it doesn't exist.
+   * Loads all given entities from the database if possible.
+   *
+   * <p>Nonexistent entities are absent from the resulting list, but no {@link
+   * NoSuchElementException} will be thrown.
    */
-  <T> ImmutableList<T> loadAll(Iterable<T> entities);
+  <T> ImmutableList<T> loadByEntitiesIfPresent(Iterable<T> entities);
+
+  /**
+   * Loads the entity by its key.
+   *
+   * @throws NoSuchElementException if this key does not correspond to an existing entity.
+   */
+  <T> T loadByKey(VKey<T> key);
+
+  /**
+   * Loads the set of entities by their keys.
+   *
+   * @throws NoSuchElementException if any of the keys do not correspond to an existing entity.
+   */
+  <T> ImmutableMap<VKey<? extends T>, T> loadByKeys(Iterable<? extends VKey<? extends T>> keys);
+
+  /**
+   * Loads the given entity from the database.
+   *
+   * @throws NoSuchElementException if the entity does not exist in the database.
+   */
+  <T> T loadByEntity(T entity);
+
+  /**
+   * Loads all given entities from the database.
+   *
+   * @throws NoSuchElementException if any of the entities do not exist in the database.
+   */
+  <T> ImmutableList<T> loadByEntities(Iterable<T> entities);
+
+  /**
+   * Returns a stream of all entities of the given type that exist in the database.
+   *
+   * <p>The resulting stream is empty if there are no entities of this type.
+   */
+  <T> ImmutableList<T> loadAllOf(Class<T> clazz);
 
   /** Deletes the entity by its id. */
   void delete(VKey<?> key);

core/src/main/java/google/registry/rdap/RdapJsonFormatter.java

@@ -52,6 +52,7 @@ import google.registry.model.registrar.Registrar;
 import google.registry.model.registrar.RegistrarAddress;
 import google.registry.model.registrar.RegistrarContact;
 import google.registry.model.reporting.HistoryEntry;
+import google.registry.model.reporting.HistoryEntryDao;
 import google.registry.persistence.VKey;
 import google.registry.rdap.RdapDataStructures.Event;
 import google.registry.rdap.RdapDataStructures.EventAction;
@@ -342,7 +343,7 @@ public class RdapJsonFormatter {
     // Kick off the database loads of the nameservers that we will need, so it can load
     // asynchronously while we load and process the contacts.
     ImmutableSet<HostResource> loadedHosts =
-        ImmutableSet.copyOf(tm().load(domainBase.getNameservers()).values());
+        ImmutableSet.copyOf(tm().loadByKeys(domainBase.getNameservers()).values());
     // Load the registrant and other contacts and add them to the data.
     Map<Key<ContactResource>, ContactResource> loadedContacts =
         ofy()
@@ -425,11 +426,11 @@ public class RdapJsonFormatter {
     if (outputDataType == OutputDataType.FULL) {
       ImmutableSet.Builder<StatusValue> statuses = new ImmutableSet.Builder<>();
       statuses.addAll(hostResource.getStatusValues());
-      if (isLinked(Key.create(hostResource), getRequestTime())) {
+      if (isLinked(hostResource.createVKey(), getRequestTime())) {
         statuses.add(StatusValue.LINKED);
       }
       if (hostResource.isSubordinate()
-          && tm().load(hostResource.getSuperordinateDomain())
+          && tm().loadByKey(hostResource.getSuperordinateDomain())
               .cloneProjectedAtTime(getRequestTime())
               .getStatusValues()
               .contains(StatusValue.PENDING_TRANSFER)) {
@@ -562,7 +563,7 @@ public class RdapJsonFormatter {
           .statusBuilder()
           .addAll(
               makeStatusValueList(
-                  isLinked(Key.create(contactResource), getRequestTime())
+                  isLinked(contactResource.createVKey(), getRequestTime())
                       ? union(contactResource.getStatusValues(), StatusValue.LINKED)
                       : contactResource.getStatusValues(),
                   false,
@@ -880,8 +881,9 @@ public class RdapJsonFormatter {
     // 2.3.2.3 An event of *eventAction* type *transfer*, with the last date and time that the
     // domain was transferred. The event of *eventAction* type *transfer* MUST be omitted if the
     // domain name has not been transferred since it was created.
-    for (HistoryEntry historyEntry :
-        ofy().load().type(HistoryEntry.class).ancestor(resource).order("modificationTime")) {
+    Iterable<? extends HistoryEntry> historyEntries =
+        HistoryEntryDao.loadHistoryObjectsForResource(resource.createVKey());
+    for (HistoryEntry historyEntry : historyEntries) {
       EventAction rdapEventAction =
           HISTORY_ENTRY_TYPE_TO_RDAP_EVENT_ACTION_MAP.get(historyEntry.getType());
       // Only save the historyEntries if this is a type we care about.
@@ -930,13 +932,9 @@ public class RdapJsonFormatter {
     return eventsBuilder.build();
   }
 
-  /**
-   * Creates an RDAP event object as defined by RFC 7483.
-   */
+  /** Creates an RDAP event object as defined by RFC 7483. */
   private static Event makeEvent(
-      EventAction eventAction,
-      @Nullable String eventActor,
-      DateTime eventDate) {
+      EventAction eventAction, @Nullable String eventActor, DateTime eventDate) {
     Event.Builder builder = Event.builder()
         .setEventAction(eventAction)
         .setEventDate(eventDate);

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

@@ -172,7 +172,7 @@ final class DomainBaseToXjcConverter {
         if (registrant == null) {
           logger.atWarning().log("Domain %s has no registrant contact.", domainName);
         } else {
-          ContactResource registrantContact = tm().load(registrant);
+          ContactResource registrantContact = tm().loadByKey(registrant);
           checkState(
               registrantContact != null,
               "Registrant contact %s on domain %s does not exist",
@@ -305,7 +305,7 @@ final class DomainBaseToXjcConverter {
         "Contact key for type %s is null on domain %s",
         model.getType(),
         domainName);
-    ContactResource contact = tm().load(model.getContactKey());
+    ContactResource contact = tm().loadByKey(model.getContactKey());
     checkState(
         contact != null,
         "Contact %s on domain %s does not exist",

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

@@ -203,7 +203,7 @@ public final class RdeStagingMapper extends Mapper<EppResource, PendingDeposit,
                         host,
                         // Note that loadAtPointInTime() does cloneProjectedAtTime(watermark) for
                         // us.
-                        loadAtPointInTime(tm().load(host.getSuperordinateDomain()), watermark)
+                        loadAtPointInTime(tm().loadByKey(host.getSuperordinateDomain()), watermark)
                             .now())
                     : marshaller.marshalExternalHost(host));
         cache.put(WatermarkModePair.create(watermark, RdeMode.FULL), result);

core/src/main/java/google/registry/request/RequestParameters.java

@@ -310,7 +310,7 @@ public final class RequestParameters {
    * @param name case insensitive header name
    */
   public static Optional<String> extractOptionalHeader(HttpServletRequest req, String name) {
-    return Optional.ofNullable(req.getHeader(name));
+    return Optional.ofNullable(emptyToNull(req.getHeader(name)));
   }
 
   private RequestParameters() {}

core/src/main/java/google/registry/schema/domain/RegistryLock.java

@@ -131,6 +131,7 @@ public final class RegistryLock extends ImmutableObject implements Buildable, Sq
   private boolean isSuperuser;
 
   /** The lock that undoes this lock, if this lock has been unlocked and the domain locked again. */
+  // TODO(b/176498743): Lazy loading on scalar field not supported by default. See bug for details.
   @OneToOne(fetch = FetchType.LAZY)
   @JoinColumn(name = "relockRevisionId", referencedColumnName = "revisionId")
   private RegistryLock relock;

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

@@ -75,8 +75,8 @@ final class DeleteAllocationTokensCommand extends UpdateOrDeleteAllocationTokens
     // Load the tokens in the same transaction as they are deleted to verify they weren't redeemed
     // since the query ran. This also filters out per-domain tokens if they're not to be deleted.
     ImmutableSet<VKey<AllocationToken>> tokensToDelete =
-        tm().load(batch).values().stream()
-            .filter(t -> withDomains || t.getDomainName().isEmpty())
+        tm().loadByKeys(batch).values().stream()
+            .filter(t -> withDomains || !t.getDomainName().isPresent())
             .filter(t -> SINGLE_USE.equals(t.getTokenType()))
             .filter(t -> !t.isRedeemed())
             .map(AllocationToken::createVKey)

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

@@ -263,7 +263,7 @@ class GenerateAllocationTokensCommand implements CommandWithRemoteApi {
       savedTokens = tokens;
     } else {
       transactIfJpaTm(() -> tm().transact(() -> tm().putAll(tokens)));
-      savedTokens = tm().transact(() -> tm().loadAll(tokens));
+      savedTokens = tm().transact(() -> tm().loadByEntities(tokens));
     }
     savedTokens.forEach(
         t -> System.out.println(SKIP_NULLS.join(t.getDomainName().orElse(null), t.getToken())));
@@ -293,7 +293,7 @@ class GenerateAllocationTokensCommand implements CommandWithRemoteApi {
             .collect(toImmutableSet());
     return transactIfJpaTm(
         () ->
-            tm().load(existingTokenKeys).values().stream()
+            tm().loadByKeysIfPresent(existingTokenKeys).values().stream()
                 .map(AllocationToken::getToken)
                 .collect(toImmutableSet()));
   }

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

@@ -49,7 +49,8 @@ final class GetAllocationTokenCommand implements CommandWithRemoteApi {
           tokens.stream()
               .map(t -> VKey.create(AllocationToken.class, t))
               .collect(toImmutableList());
-      tm().load(tokenKeys).forEach((k, v) -> builder.put(k.getSqlKey().toString(), v));
+      tm().loadByKeysIfPresent(tokenKeys)
+          .forEach((k, v) -> builder.put(k.getSqlKey().toString(), v));
     }
     ImmutableMap<String, AllocationToken> loadedTokens = builder.build();
     ImmutableMap<VKey<DomainBase>, DomainBase> domains = loadRedeemedDomains(loadedTokens.values());
@@ -58,7 +59,7 @@ final class GetAllocationTokenCommand implements CommandWithRemoteApi {
       if (loadedTokens.containsKey(token)) {
         AllocationToken loadedToken = loadedTokens.get(token);
         System.out.println(loadedToken.toString());
-        if (loadedToken.getRedemptionHistoryEntry().isEmpty()) {
+        if (!loadedToken.getRedemptionHistoryEntry().isPresent()) {
           System.out.printf("Token %s was not redeemed.\n", token);
         } else {
           Key<DomainBase> domainOfyKey =
@@ -88,14 +89,14 @@ final class GetAllocationTokenCommand implements CommandWithRemoteApi {
             .map(AllocationToken::getRedemptionHistoryEntry)
             .filter(Optional::isPresent)
             .map(Optional::get)
-            .map(key -> tm().load(key))
+            .map(key -> tm().loadByKey(key))
             .map(he -> (Key<DomainBase>) he.getParent())
             .map(key -> VKey.create(DomainBase.class, key.getName(), key))
             .collect(toImmutableList());
     ImmutableMap.Builder<VKey<DomainBase>, DomainBase> domainsBuilder =
         new ImmutableMap.Builder<>();
     for (List<VKey<DomainBase>> keys : Lists.partition(domainKeys, BATCH_SIZE)) {
-      tm().load(ImmutableList.copyOf(keys))
+      tm().loadByKeys(ImmutableList.copyOf(keys))
           .forEach((k, v) -> domainsBuilder.put((VKey<DomainBase>) k, v));
     }
     return domainsBuilder.build();

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

@@ -15,7 +15,6 @@
 package google.registry.tools;
 
 import static com.google.common.base.Preconditions.checkArgument;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.util.DateTimeUtils.END_OF_TIME;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
@@ -25,14 +24,16 @@ import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
 import google.registry.model.EppResource;
 import google.registry.model.reporting.HistoryEntry;
+import google.registry.model.reporting.HistoryEntryDao;
 import google.registry.persistence.VKey;
 import google.registry.tools.CommandUtilities.ResourceType;
 import google.registry.xml.XmlTransformer;
 import org.joda.time.DateTime;
 
 /** Command to show history entries. */
-@Parameters(separators = " =",
-  commandDescription = "Show history entries that occurred in a given time range")
+@Parameters(
+    separators = " =",
+    commandDescription = "Show history entries that occurred in a given time range")
 final class GetHistoryEntriesCommand implements CommandWithRemoteApi {
 
   @Parameter(
@@ -45,33 +46,26 @@ final class GetHistoryEntriesCommand implements CommandWithRemoteApi {
       description = "Only show history entries that occurred at or before this time")
   private DateTime before = END_OF_TIME;
 
-  @Parameter(
-    names = "--type",
-    description = "Resource type.")
+  @Parameter(names = "--type", description = "Resource type.")
   private ResourceType type;
 
-  @Parameter(
-    names = "--id",
-    description = "Foreign key of the resource.")
+  @Parameter(names = "--id", description = "Foreign key of the resource.")
   private String uniqueId;
 
   @Override
   public void run() {
-    VKey<? extends EppResource> parentKey = null;
+    Iterable<? extends HistoryEntry> historyEntries;
     if (type != null || uniqueId != null) {
       checkArgument(
           type != null && uniqueId != null,
           "If either of 'type' or 'id' are set then both must be");
-      parentKey = type.getKey(uniqueId, DateTime.now(UTC));
+      VKey<? extends EppResource> parentKey = type.getKey(uniqueId, DateTime.now(UTC));
       checkArgumentNotNull(parentKey, "Invalid resource ID");
+      historyEntries = HistoryEntryDao.loadHistoryObjectsForResource(parentKey, after, before);
+    } else {
+      historyEntries = HistoryEntryDao.loadAllHistoryObjects(after, before);
     }
-    for (HistoryEntry entry :
-        (parentKey == null
-                ? ofy().load().type(HistoryEntry.class)
-                : ofy().load().type(HistoryEntry.class).ancestor(parentKey.getOfyKey()))
-            .order("modificationTime")
-            .filter("modificationTime >=", after)
-            .filter("modificationTime <=", before)) {
+    for (HistoryEntry entry : historyEntries) {
       System.out.printf(
           "Client: %s\nTime: %s\nClient TRID: %s\nServer TRID: %s\n%s\n",
           entry.getClientId(),

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

@@ -168,6 +168,8 @@ interface RegistryToolComponent {
 
   void inject(ValidateEscrowDepositCommand command);
 
+  void inject(ValidateLoginCredentialsCommand command);
+
   void inject(WhoisQueryCommand command);
 
   AppEngineConnection appEngineConnection();

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

@@ -56,7 +56,7 @@ public final class ResaveEppResourceCommand extends MutatingCommand {
             uniqueId);
     // Load the resource directly to bypass running cloneProjectedAtTime() automatically, which can
     // cause stageEntityChange() to fail due to implicit projection changes.
-    EppResource resource = tm().load(resourceKey);
+    EppResource resource = tm().loadByKey(resourceKey);
     stageEntityChange(resource, resource);
   }
 }

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

@@ -162,7 +162,7 @@ final class UniformRapidSuspensionCommand extends MutatingEppToolCommand {
 
   private ImmutableSortedSet<String> getExistingNameservers(DomainBase domain) {
     ImmutableSortedSet.Builder<String> nameservers = ImmutableSortedSet.naturalOrder();
-    for (HostResource host : tm().load(domain.getNameservers()).values()) {
+    for (HostResource host : tm().loadByKeys(domain.getNameservers()).values()) {
       nameservers.add(host.getForeignKey());
     }
     return nameservers.build();

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

@@ -107,7 +107,7 @@ final class UpdateAllocationTokensCommand extends UpdateOrDeleteAllocationTokens
     tokensToSave =
         transactIfJpaTm(
             () ->
-                tm().load(getTokenKeys()).values().stream()
+                tm().loadByKeys(getTokenKeys()).values().stream()
                     .collect(toImmutableMap(Function.identity(), this::updateToken))
                     .entrySet()
                     .stream()

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

@@ -18,6 +18,7 @@ import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static google.registry.model.EppResourceUtils.loadByForeignKey;
 import static google.registry.model.domain.rgp.GracePeriodStatus.AUTO_RENEW;
+import static google.registry.model.eppcommon.StatusValue.PENDING_DELETE;
 import static google.registry.model.eppcommon.StatusValue.SERVER_UPDATE_PROHIBITED;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.PreconditionsUtils.checkArgumentPresent;
@@ -139,6 +140,11 @@ final class UpdateDomainCommand extends CreateOrUpdateDomainCommand {
               + " deleted at the end of its current registration period.")
   Boolean autorenews;
 
+  @Parameter(
+      names = {"--force_in_pending_delete"},
+      description = "Force a superuser update even on domains that are in pending delete")
+  boolean forceInPendingDelete;
+
   @Override
   protected void initMutatingEppToolCommand() {
     if (!nameservers.isEmpty()) {
@@ -186,6 +192,12 @@ final class UpdateDomainCommand extends CreateOrUpdateDomainCommand {
           "The domain '%s' has status SERVER_UPDATE_PROHIBITED. Verify that you are allowed "
               + "to make updates, and if so, use the domain_unlock command to enable updates.",
           domain);
+      checkArgument(
+          !domainBase.getStatusValues().contains(PENDING_DELETE) || forceInPendingDelete,
+          "The domain '%s' has status PENDING_DELETE. Verify that you really are intending to "
+              + "update a domain in pending delete (this is uncommon), and if so, pass the "
+              + "--force_in_pending_delete parameter to allow this update.",
+          domain);
 
       // Use TreeSets so that the results are always in the same order (this makes testing easier).
       Set<String> addAdminsThisDomain = new TreeSet<>(addAdmins);
@@ -329,7 +341,7 @@ final class UpdateDomainCommand extends CreateOrUpdateDomainCommand {
       DomainBase domainBase, final DesignatedContact.Type contactType) {
     return domainBase.getContacts().stream()
         .filter(contact -> contact.getType().equals(contactType))
-        .map(contact -> tm().load(contact.getContactKey()).getContactId())
+        .map(contact -> tm().loadByKey(contact.getContactKey()).getContactId())
         .collect(toImmutableSet());
   }
 }

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

@@ -67,7 +67,7 @@ abstract class UpdateOrDeleteAllocationTokensCommand extends ConfirmingCommand
       checkArgument(!prefix.isEmpty(), "Provided prefix should not be blank");
       return transactIfJpaTm(
           () ->
-              tm().loadAll(AllocationToken.class).stream()
+              tm().loadAllOf(AllocationToken.class).stream()
                   .filter(token -> token.getToken().startsWith(prefix))
                   .map(AllocationToken::createVKey)
                   .collect(toImmutableSet()));

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

@@ -18,6 +18,7 @@ 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 static google.registry.util.PreconditionsUtils.checkArgumentPresent;
+import static google.registry.util.X509Utils.encodeX509CertificateFromPemString;
 import static google.registry.util.X509Utils.getCertificateHash;
 import static google.registry.util.X509Utils.loadCertificate;
 import static java.nio.charset.StandardCharsets.US_ASCII;
@@ -25,12 +26,14 @@ import static java.nio.charset.StandardCharsets.US_ASCII;
 import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
 import google.registry.flows.TlsCredentials;
+import google.registry.flows.certs.CertificateChecker;
 import google.registry.model.registrar.Registrar;
 import google.registry.tools.params.PathParameter;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.Optional;
 import javax.annotation.Nullable;
+import javax.inject.Inject;
 
 /** A command to test registrar login credentials. */
 @Parameters(separators = " =", commandDescription = "Test registrar login credentials")
@@ -55,30 +58,41 @@ final class ValidateLoginCredentialsCommand implements CommandWithRemoteApi {
       validateWith = PathParameter.InputFile.class)
   private Path clientCertificatePath;
 
+  // TODO(sarahbot@): Remove this after hash fallback is removed
   @Nullable
   @Parameter(
       names = {"-h", "--cert_hash"},
       description = "Hash of the client certificate.")
   private String clientCertificateHash;
 
+  @Nullable
   @Parameter(
       names = {"-i", "--ip_address"},
       description = "Client ip address to pretend to use")
   private String clientIpAddress = "10.0.0.1";
 
+  @Inject CertificateChecker certificateChecker;
+
   @Override
   public void run() throws Exception {
     checkArgument(
         clientCertificatePath == null || isNullOrEmpty(clientCertificateHash),
         "Can't specify both --cert_hash and --cert_file");
+    String encodedCertificate = "";
     if (clientCertificatePath != null) {
-      clientCertificateHash = getCertificateHash(
-          loadCertificate(new String(Files.readAllBytes(clientCertificatePath), US_ASCII)));
+      String certificateString = new String(Files.readAllBytes(clientCertificatePath), US_ASCII);
+      encodedCertificate = encodeX509CertificateFromPemString(certificateString);
+      clientCertificateHash = getCertificateHash(loadCertificate(clientCertificatePath));
     }
     Registrar registrar =
         checkArgumentPresent(
             Registrar.loadByClientId(clientId), "Registrar %s not found", clientId);
-    new TlsCredentials(true, clientCertificateHash, Optional.of(clientIpAddress))
+    new TlsCredentials(
+            true,
+            Optional.ofNullable(clientCertificateHash),
+            Optional.ofNullable(encodedCertificate),
+            Optional.ofNullable(clientIpAddress),
+            certificateChecker)
         .validate(registrar, password);
     checkState(
         registrar.isLive(), "Registrar %s has non-live state: %s", clientId, registrar.getState());

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

@@ -16,20 +16,24 @@ package google.registry.tools.javascrap;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.ImmutableList.toImmutableList;
-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.tools.LockOrUnlockDomainCommand.REGISTRY_LOCK_STATUSES;
 
 import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
+import com.google.common.collect.ImmutableCollection;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
+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.model.domain.DomainBase;
 import google.registry.model.registry.RegistryLockDao;
 import google.registry.model.reporting.HistoryEntry;
+import google.registry.model.reporting.HistoryEntryDao;
+import google.registry.persistence.VKey;
 import google.registry.schema.domain.RegistryLock;
 import google.registry.tools.CommandWithRemoteApi;
 import google.registry.tools.ConfirmingCommand;
@@ -130,7 +134,7 @@ public class BackfillRegistryLocksCommand extends ConfirmingCommand
   private DateTime getLockCompletionTimestamp(DomainBase domainBase, DateTime now) {
     // Best-effort, if a domain was URS-locked we should use that time
     // If we can't find that, return now.
-    return ofy().load().type(HistoryEntry.class).ancestor(domainBase).list().stream()
+    return Streams.stream(HistoryEntryDao.loadHistoryObjectsForResource(domainBase.createVKey()))
         // sort by modification time descending so we get the most recent one if it was locked twice
         .sorted(Comparator.comparing(HistoryEntry::getModificationTime).reversed())
         .filter(entry -> entry.getReason().equals("Uniform Rapid Suspension"))
@@ -140,18 +144,14 @@ public class BackfillRegistryLocksCommand extends ConfirmingCommand
   }
 
   private ImmutableList<DomainBase> getLockedDomainsWithoutLocks(DateTime now) {
-    return ImmutableList.copyOf(
-        ofy()
-            .load()
-            .keys(
-                roids.stream()
-                    .map(roid -> Key.create(DomainBase.class, roid))
-                    .collect(toImmutableList()))
-            .values()
-            .stream()
-            .filter(d -> d.getDeletionTime().isAfter(now))
-            .filter(d -> d.getStatusValues().containsAll(REGISTRY_LOCK_STATUSES))
-            .filter(d -> !RegistryLockDao.getMostRecentByRepoId(d.getRepoId()).isPresent())
-            .collect(toImmutableList()));
+    ImmutableList<VKey<DomainBase>> domainKeys =
+        roids.stream().map(roid -> VKey.create(DomainBase.class, roid)).collect(toImmutableList());
+    ImmutableCollection<DomainBase> domains =
+        transactIfJpaTm(() -> tm().loadByKeys(domainKeys)).values();
+    return domains.stream()
+        .filter(d -> d.getDeletionTime().isAfter(now))
+        .filter(d -> d.getStatusValues().containsAll(REGISTRY_LOCK_STATUSES))
+        .filter(d -> RegistryLockDao.getMostRecentByRepoId(d.getRepoId()).isEmpty())
+        .collect(toImmutableList());
   }
 }

core/src/main/java/google/registry/tools/server/GenerateZoneFilesAction.java

@@ -214,7 +214,7 @@ public class GenerateZoneFilesAction implements Runnable, JsonActionRunner.JsonA
     private void emitForSubordinateHosts(DomainBase domain) {
       ImmutableSet<String> subordinateHosts = domain.getSubordinateHosts();
       if (!subordinateHosts.isEmpty()) {
-        for (HostResource unprojectedHost : tm().load(domain.getNameservers()).values()) {
+        for (HostResource unprojectedHost : tm().loadByKeys(domain.getNameservers()).values()) {
           HostResource host = loadAtPointInTime(unprojectedHost, exportTime).now();
           // A null means the host was deleted (or not created) at this time.
           if ((host != null) && subordinateHosts.contains(host.getHostName())) {
@@ -283,7 +283,7 @@ public class GenerateZoneFilesAction implements Runnable, JsonActionRunner.JsonA
       Duration dnsDefaultDsTtl) {
     StringBuilder result = new StringBuilder();
     String domainLabel = stripTld(domain.getDomainName(), domain.getTld());
-    for (HostResource nameserver : tm().load(domain.getNameservers()).values()) {
+    for (HostResource nameserver : tm().loadByKeys(domain.getNameservers()).values()) {
       result.append(
           String.format(
               NS_FORMAT,

core/src/main/java/google/registry/ui/server/registrar/RegistrarSettingsAction.java

@@ -39,6 +39,7 @@ import com.google.common.collect.Streams;
 import com.google.common.flogger.FluentLogger;
 import google.registry.config.RegistryEnvironment;
 import google.registry.flows.certs.CertificateChecker;
+import google.registry.flows.certs.CertificateChecker.InsecureCertificateException;
 import google.registry.model.registrar.Registrar;
 import google.registry.model.registrar.RegistrarContact;
 import google.registry.model.registrar.RegistrarContact.Type;
@@ -334,9 +335,14 @@ public class RegistrarSettingsAction implements Runnable, JsonActionRunner.JsonA
    * Returns true if the registrar should accept the new certificate. Returns false if the
    * certificate is already the one stored for the registrar.
    */
-  private boolean validateCertificate(String existingCertificate, String certificateString) {
-    if ((existingCertificate == null) || !existingCertificate.equals(certificateString)) {
+  private boolean validateCertificate(
+      Optional<String> existingCertificate, String certificateString) {
+    if (!existingCertificate.isPresent() || !existingCertificate.get().equals(certificateString)) {
+      try {
         certificateChecker.validateCertificate(certificateString);
+      } catch (InsecureCertificateException e) {
+        throw new IllegalArgumentException(e.getMessage());
+      }
       return true;
     }
     return false;

core/src/main/java/google/registry/whois/NameserverWhoisResponse.java

@@ -49,7 +49,7 @@ final class NameserverWhoisResponse extends WhoisResponseImpl {
       HostResource host = hosts.get(i);
       String clientId =
           host.isSubordinate()
-              ? tm().load(host.getSuperordinateDomain())
+              ? tm().loadByKey(host.getSuperordinateDomain())
                   .cloneProjectedAtTime(getTimestamp())
                   .getCurrentSponsorClientId()
               : host.getPersistedCurrentSponsorClientId();

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

@@ -243,7 +243,7 @@ public class ReplayCommitLogsToSqlActionTest {
         CommitLogMutation.create(manifestKey, TestObject.create("existing", "b")));
     action.run();
     TestObject fromDatabase =
-        jpaTm().transact(() -> jpaTm().load(VKey.createSql(TestObject.class, "existing")));
+        jpaTm().transact(() -> jpaTm().loadByKey(VKey.createSql(TestObject.class, "existing")));
     assertThat(fromDatabase.getField()).isEqualTo("b");
   }
 
@@ -276,7 +276,7 @@ public class ReplayCommitLogsToSqlActionTest {
     DomainBase domain = newDomainBase("example.tld");
     CommitLogMutation domainMutation =
         tm().transact(() -> CommitLogMutation.create(manifestKey, domain));
-    ContactResource contact = tm().transact(() -> tm().load(domain.getRegistrant()));
+    ContactResource contact = tm().transact(() -> tm().loadByKey(domain.getRegistrant()));
     CommitLogMutation contactMutation =
         tm().transact(() -> CommitLogMutation.create(manifestKey, contact));
 
@@ -348,7 +348,7 @@ public class ReplayCommitLogsToSqlActionTest {
     inOrder.verify(spy).delete(contact.createVKey());
     inOrder.verify(spy).put(putCaptor.capture());
     assertThat(putCaptor.getValue().getClass()).isEqualTo(ContactResource.class);
-    assertThat(jpaTm().transact(() -> jpaTm().load(contact.createVKey()).getEmailAddress()))
+    assertThat(jpaTm().transact(() -> jpaTm().loadByKey(contact.createVKey()).getEmailAddress()))
         .isEqualTo("replay@example.tld");
   }
 
@@ -450,7 +450,7 @@ public class ReplayCommitLogsToSqlActionTest {
         jpaTm()
             .transact(
                 () ->
-                    jpaTm().loadAll(TestObject.class).stream()
+                    jpaTm().loadAllOf(TestObject.class).stream()
                         .map(TestObject::getId)
                         .collect(toImmutableList()));
     assertThat(actualIds).containsExactlyElementsIn(expectedIds);

core/src/test/java/google/registry/batch/DeleteExpiredDomainsActionTest.java

@@ -0,0 +1,183 @@
+// 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.truth.Truth.assertThat;
+import static google.registry.model.eppcommon.StatusValue.PENDING_DELETE;
+import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_CREATE;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.testing.DatabaseHelper.createTld;
+import static google.registry.testing.DatabaseHelper.newDomainBase;
+import static google.registry.testing.DatabaseHelper.persistActiveDomain;
+import static google.registry.testing.DatabaseHelper.persistResource;
+import static google.registry.util.DateTimeUtils.END_OF_TIME;
+
+import com.google.common.collect.ImmutableSet;
+import google.registry.flows.DaggerEppTestComponent;
+import google.registry.flows.EppController;
+import google.registry.flows.EppTestComponent.FakesAndMocksModule;
+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.ofy.Ofy;
+import google.registry.model.poll.PollMessage;
+import google.registry.model.reporting.HistoryEntry;
+import google.registry.monitoring.whitebox.EppMetric;
+import google.registry.testing.AppEngineExtension;
+import google.registry.testing.FakeClock;
+import google.registry.testing.FakeLockHandler;
+import google.registry.testing.FakeResponse;
+import google.registry.testing.InjectExtension;
+import java.util.Optional;
+import org.joda.time.DateTime;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+/** Unit tests for {@link DeleteExpiredDomainsAction}. */
+class DeleteExpiredDomainsActionTest {
+
+  @RegisterExtension
+  public final AppEngineExtension appEngine =
+      AppEngineExtension.builder().withDatastoreAndCloudSql().withTaskQueue().build();
+
+  @RegisterExtension public final InjectExtension inject = new InjectExtension();
+
+  private final FakeClock clock = new FakeClock(DateTime.parse("2016-06-13T20:21:22Z"));
+  private final FakeResponse response = new FakeResponse();
+  private DeleteExpiredDomainsAction action;
+
+  @BeforeEach
+  void beforeEach() {
+    inject.setStaticField(Ofy.class, "clock", clock);
+    createTld("tld");
+    EppController eppController =
+        DaggerEppTestComponent.builder()
+            .fakesAndMocksModule(
+                FakesAndMocksModule.create(clock, EppMetric.builderForRequest(clock)))
+            .build()
+            .startRequest()
+            .eppController();
+    action =
+        new DeleteExpiredDomainsAction(
+            eppController, "NewRegistrar", clock, new FakeLockHandler(true), response);
+  }
+
+  @Test
+  void test_deletesOnlyExpiredDomain() {
+    // A normal, active autorenewing domain that shouldn't be touched.
+    DomainBase activeDomain = persistActiveDomain("foo.tld");
+    clock.advanceOneMilli();
+
+    // A non-autorenewing domain that is already pending delete and shouldn't be touched.
+    DomainBase alreadyDeletedDomain =
+        persistResource(
+            newDomainBase("bar.tld")
+                .asBuilder()
+                .setAutorenewEndTime(Optional.of(clock.nowUtc().minusDays(10)))
+                .setDeletionTime(clock.nowUtc().plusDays(17))
+                .build());
+    clock.advanceOneMilli();
+
+    // A non-autorenewing domain that hasn't reached its expiration time and shouldn't be touched.
+    DomainBase notYetExpiredDomain =
+        persistResource(
+            newDomainBase("baz.tld")
+                .asBuilder()
+                .setAutorenewEndTime(Optional.of(clock.nowUtc().plusDays(15)))
+                .build());
+    clock.advanceOneMilli();
+
+    // A non-autorenewing domain that is past its expiration time and should be deleted.
+    // (This is the only one that needs a full set of subsidiary resources, for the delete flow to
+    // to operate on.)
+    DomainBase pendingExpirationDomain = persistNonAutorenewingDomain("fizz.tld");
+
+    assertThat(tm().loadByEntity(pendingExpirationDomain).getStatusValues())
+        .doesNotContain(PENDING_DELETE);
+    action.run();
+
+    DomainBase reloadedActiveDomain = tm().loadByEntity(activeDomain);
+    assertThat(reloadedActiveDomain).isEqualTo(activeDomain);
+    assertThat(reloadedActiveDomain.getStatusValues()).doesNotContain(PENDING_DELETE);
+    assertThat(tm().loadByEntity(alreadyDeletedDomain)).isEqualTo(alreadyDeletedDomain);
+    assertThat(tm().loadByEntity(notYetExpiredDomain)).isEqualTo(notYetExpiredDomain);
+    DomainBase reloadedExpiredDomain = tm().loadByEntity(pendingExpirationDomain);
+    assertThat(reloadedExpiredDomain.getStatusValues()).contains(PENDING_DELETE);
+    assertThat(reloadedExpiredDomain.getDeletionTime()).isEqualTo(clock.nowUtc().plusDays(35));
+  }
+
+  @Test
+  void test_deletesThreeDomainsInOneRun() {
+    DomainBase domain1 = persistNonAutorenewingDomain("ecck1.tld");
+    DomainBase domain2 = persistNonAutorenewingDomain("veee2.tld");
+    DomainBase domain3 = persistNonAutorenewingDomain("tarm3.tld");
+
+    action.run();
+
+    assertThat(tm().loadByEntity(domain1).getStatusValues()).contains(PENDING_DELETE);
+    assertThat(tm().loadByEntity(domain2).getStatusValues()).contains(PENDING_DELETE);
+    assertThat(tm().loadByEntity(domain3).getStatusValues()).contains(PENDING_DELETE);
+  }
+
+  private DomainBase persistNonAutorenewingDomain(String domainName) {
+    DomainBase pendingExpirationDomain = persistActiveDomain(domainName);
+    HistoryEntry createHistoryEntry =
+        persistResource(
+            new HistoryEntry.Builder()
+                .setType(DOMAIN_CREATE)
+                .setParent(pendingExpirationDomain)
+                .setModificationTime(clock.nowUtc().minusMonths(9))
+                .build());
+    BillingEvent.Recurring autorenewBillingEvent =
+        persistResource(createAutorenewBillingEvent(createHistoryEntry).build());
+    PollMessage.Autorenew autorenewPollMessage =
+        persistResource(createAutorenewPollMessage(createHistoryEntry).build());
+    pendingExpirationDomain =
+        persistResource(
+            pendingExpirationDomain
+                .asBuilder()
+                .setAutorenewEndTime(Optional.of(clock.nowUtc().minusDays(10)))
+                .setAutorenewBillingEvent(autorenewBillingEvent.createVKey())
+                .setAutorenewPollMessage(autorenewPollMessage.createVKey())
+                .build());
+    clock.advanceOneMilli();
+
+    return pendingExpirationDomain;
+  }
+
+  private BillingEvent.Recurring.Builder createAutorenewBillingEvent(
+      HistoryEntry createHistoryEntry) {
+    return new BillingEvent.Recurring.Builder()
+        .setReason(Reason.RENEW)
+        .setFlags(ImmutableSet.of(Flag.AUTO_RENEW))
+        .setTargetId("fizz.tld")
+        .setClientId("TheRegistrar")
+        .setEventTime(clock.nowUtc().plusYears(1))
+        .setRecurrenceEndTime(END_OF_TIME)
+        .setParent(createHistoryEntry);
+  }
+
+  private PollMessage.Autorenew.Builder createAutorenewPollMessage(
+      HistoryEntry createHistoryEntry) {
+    return new PollMessage.Autorenew.Builder()
+        .setTargetId("fizz.tld")
+        .setClientId("TheRegistrar")
+        .setEventTime(clock.nowUtc().plusYears(1))
+        .setAutorenewEndTime(END_OF_TIME)
+        .setParent(createHistoryEntry);
+  }
+}

core/src/test/java/google/registry/beam/datastore/BulkDeletePipelineTest.java

@@ -0,0 +1,152 @@
+// 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.datastore;
+
+import static google.registry.beam.datastore.BulkDeletePipeline.discoverEntityKinds;
+import static google.registry.beam.datastore.BulkDeletePipeline.getDeletionTags;
+import static google.registry.beam.datastore.BulkDeletePipeline.getOneDeletionTag;
+
+import com.google.common.base.Verify;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Iterables;
+import com.google.datastore.v1.Entity;
+import com.google.datastore.v1.Key;
+import com.google.datastore.v1.Key.PathElement;
+import google.registry.beam.TestPipelineExtension;
+import google.registry.beam.datastore.BulkDeletePipeline.GenerateQueries;
+import google.registry.beam.datastore.BulkDeletePipeline.SplitEntities;
+import java.io.Serializable;
+import java.util.Map;
+import org.apache.beam.sdk.testing.PAssert;
+import org.apache.beam.sdk.transforms.Count;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.View;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.PCollectionView;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.condition.EnabledIfSystemProperty;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+/** Unit tests for {@link BulkDeletePipeline}. */
+class BulkDeletePipelineTest implements Serializable {
+
+  @RegisterExtension
+  final transient TestPipelineExtension testPipeline =
+      TestPipelineExtension.create().enableAbandonedNodeEnforcement(true);
+
+  @Test
+  void generateQueries() {
+    PCollection<String> queries =
+        testPipeline
+            .apply("InjectKinds", Create.of("A", "B"))
+            .apply("GenerateQueries", ParDo.of(new GenerateQueries()));
+
+    PAssert.that(queries).containsInAnyOrder("select __key__ from `A`", "select __key__ from `B`");
+    testPipeline.run();
+  }
+
+  @Test
+  void mapKindsToTags() {
+    TupleTagList tags = getDeletionTags(2);
+    PCollection<String> kinds = testPipeline.apply("InjectKinds", Create.of("A", "B"));
+    PCollection<KV<String, TupleTag<Entity>>> kindToTagMapping =
+        BulkDeletePipeline.mapKindsToDeletionTags(kinds, tags);
+    PAssert.thatMap(kindToTagMapping)
+        .isEqualTo(
+            ImmutableMap.of(
+                "A", new TupleTag<Entity>("0"),
+                "B", new TupleTag<Entity>("1")));
+    testPipeline.run();
+  }
+
+  @Test
+  void mapKindsToTags_fewerKindsThanTags() {
+    TupleTagList tags = getDeletionTags(3);
+    PCollection<String> kinds = testPipeline.apply("InjectKinds", Create.of("A", "B"));
+    PCollection<KV<String, TupleTag<Entity>>> kindToTagMapping =
+        BulkDeletePipeline.mapKindsToDeletionTags(kinds, tags);
+    PAssert.thatMap(kindToTagMapping)
+        .isEqualTo(
+            ImmutableMap.of(
+                "A", new TupleTag<Entity>("0"),
+                "B", new TupleTag<Entity>("1")));
+    testPipeline.run();
+  }
+
+  @Test
+  void mapKindsToTags_moreKindsThanTags() {
+    TupleTagList tags = getDeletionTags(2);
+    PCollection<String> kinds = testPipeline.apply("InjectKinds", Create.of("A", "B", "C"));
+    PCollection<KV<String, TupleTag<Entity>>> kindToTagMapping =
+        BulkDeletePipeline.mapKindsToDeletionTags(kinds, tags);
+    PAssert.thatMap(kindToTagMapping)
+        .isEqualTo(
+            ImmutableMap.of(
+                "A", new TupleTag<Entity>("0"),
+                "B", new TupleTag<Entity>("1"),
+                "C", new TupleTag<Entity>("0")));
+    testPipeline.run();
+  }
+
+  @Test
+  void splitEntitiesByKind() {
+    TupleTagList tags = getDeletionTags(2);
+    PCollection<String> kinds = testPipeline.apply("InjectKinds", Create.of("A", "B"));
+    PCollectionView<Map<String, TupleTag<Entity>>> kindToTagMapping =
+        BulkDeletePipeline.mapKindsToDeletionTags(kinds, tags).apply(View.asMap());
+    Entity entityA = createTestEntity("A", 1);
+    Entity entityB = createTestEntity("B", 2);
+    PCollection<Entity> entities =
+        testPipeline.apply("InjectEntities", Create.of(entityA, entityB));
+    PCollectionTuple allCollections =
+        entities.apply(
+            "SplitByKind",
+            ParDo.of(new SplitEntities(kindToTagMapping))
+                .withSideInputs(kindToTagMapping)
+                .withOutputTags(getOneDeletionTag("placeholder"), tags));
+    PAssert.that(allCollections.get((TupleTag<Entity>) tags.get(0))).containsInAnyOrder(entityA);
+    PAssert.that(allCollections.get((TupleTag<Entity>) tags.get(1))).containsInAnyOrder(entityB);
+    testPipeline.run();
+  }
+
+  private static Entity createTestEntity(String kind, long id) {
+    return Entity.newBuilder()
+        .setKey(Key.newBuilder().addPath(PathElement.newBuilder().setId(id).setKind(kind)))
+        .build();
+  }
+
+  @Test
+  @EnabledIfSystemProperty(named = "test.gcp_integration.env", matches = "\\S+")
+  void discoverKindsFromDatastore() {
+    String environmentName = System.getProperty("test.gcp_integration.env");
+    String project = "domain-registry-" + environmentName;
+
+    PCollection<String> kinds =
+        testPipeline.apply("DiscoverEntityKinds", discoverEntityKinds(project));
+
+    PAssert.that(kinds.apply(Count.globally()))
+        .satisfies(
+            longs -> {
+              Verify.verify(Iterables.size(longs) == 1 && Iterables.getFirst(longs, -1L) > 0);
+              return null;
+            });
+    testPipeline.run();
+  }
+}

core/src/test/java/google/registry/beam/initsql/DomainBaseUtilTest.java

@@ -179,6 +179,7 @@ public class DomainBaseUtilTest {
             .setNameservers(ImmutableSet.of())
             .setDeletePollMessage(null)
             .setTransferData(null)
+            .setGracePeriods(ImmutableSet.of())
             .build();
     DomainBase domainTransformedByUtil =
         (DomainBase) ofy().toPojo(DomainBaseUtil.removeBillingAndPollAndHosts(domainEntity));
@@ -199,6 +200,7 @@ public class DomainBaseUtilTest {
             .setNameservers(ImmutableSet.of())
             .setDeletePollMessage(null)
             .setTransferData(null)
+            .setGracePeriods(ImmutableSet.of())
             .build();
     Entity entityWithoutFkeys = tm().transact(() -> ofy().toEntity(domainWithoutFKeys));
     DomainBase domainTransformedByUtil =

core/src/test/java/google/registry/beam/initsql/InitSqlPipelineTest.java

@@ -17,9 +17,12 @@ package google.registry.beam.initsql;
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.ImmutableObjectSubject.assertAboutImmutableObjects;
 import static google.registry.model.ImmutableObjectSubject.immutableObjectCorrespondence;
+import static google.registry.model.domain.token.AllocationToken.TokenType.SINGLE_USE;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.testing.DatabaseHelper.newRegistry;
 import static google.registry.testing.DatabaseHelper.persistResource;
+import static google.registry.testing.DatabaseHelper.persistSimpleResource;
+import static google.registry.util.DateTimeUtils.END_OF_TIME;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
 
 import com.google.common.collect.ImmutableList;
@@ -27,7 +30,10 @@ import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
 import google.registry.backup.AppEngineEnvironment;
 import google.registry.beam.TestPipelineExtension;
+import google.registry.flows.domain.DomainFlowUtils;
 import google.registry.model.billing.BillingEvent;
+import google.registry.model.billing.BillingEvent.Flag;
+import google.registry.model.billing.BillingEvent.Reason;
 import google.registry.model.contact.ContactResource;
 import google.registry.model.domain.DesignatedContact;
 import google.registry.model.domain.DomainAuthInfo;
@@ -36,6 +42,7 @@ import google.registry.model.domain.GracePeriod;
 import google.registry.model.domain.launch.LaunchNotice;
 import google.registry.model.domain.rgp.GracePeriodStatus;
 import google.registry.model.domain.secdns.DelegationSignerData;
+import google.registry.model.domain.token.AllocationToken;
 import google.registry.model.eppcommon.AuthInfo.PasswordAuth;
 import google.registry.model.eppcommon.StatusValue;
 import google.registry.model.eppcommon.Trid;
@@ -43,6 +50,7 @@ import google.registry.model.host.HostResource;
 import google.registry.model.ofy.Ofy;
 import google.registry.model.poll.PollMessage;
 import google.registry.model.registrar.Registrar;
+import google.registry.model.registrar.RegistrarContact;
 import google.registry.model.registry.Registry;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.transfer.DomainTransferData;
@@ -58,6 +66,7 @@ import java.io.File;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.joda.money.Money;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Order;
@@ -69,14 +78,24 @@ import org.junit.jupiter.api.io.TempDir;
 class InitSqlPipelineTest {
   private static final DateTime START_TIME = DateTime.parse("2000-01-01T00:00:00.0Z");
 
+  /**
+   * All kinds of entities to be set up in the Datastore. Must contain all kinds known to {@link
+   * InitSqlPipeline}.
+   */
   private static final ImmutableList<Class<?>> ALL_KINDS =
       ImmutableList.of(
           Registry.class,
           Registrar.class,
           ContactResource.class,
-          HostResource.class,
+          RegistrarContact.class,
           DomainBase.class,
-          HistoryEntry.class);
+          HostResource.class,
+          HistoryEntry.class,
+          AllocationToken.class,
+          BillingEvent.Recurring.class,
+          BillingEvent.OneTime.class,
+          BillingEvent.Cancellation.class,
+          PollMessage.class);
 
   private transient FakeClock fakeClock = new FakeClock(START_TIME);
 
@@ -152,16 +171,75 @@ class InitSqlPipelineTest {
                   .setCreationClientId(registrar1.getClientId())
                   .setPersistedCurrentSponsorClientId(registrar1.getClientId())
                   .build());
-      historyEntry = persistResource(new HistoryEntry.Builder().setParent(domainKey).build());
+      persistSimpleResource(
+          new RegistrarContact.Builder()
+              .setParent(registrar1)
+              .setName("John Abused")
+              .setEmailAddress("johnabuse@example.com")
+              .setVisibleInWhoisAsAdmin(true)
+              .setVisibleInWhoisAsTech(false)
+              .setPhoneNumber("+1.2125551213")
+              .setFaxNumber("+1.2125551213")
+              .setTypes(ImmutableSet.of(RegistrarContact.Type.ABUSE, RegistrarContact.Type.ADMIN))
+              .build());
+      historyEntry =
+          persistResource(
+              new HistoryEntry.Builder()
+                  .setParent(domainKey)
+                  .setModificationTime(fakeClock.nowUtc())
+                  .setType(HistoryEntry.Type.DOMAIN_CREATE)
+                  .build());
+      persistResource(
+          new AllocationToken.Builder().setToken("abc123").setTokenType(SINGLE_USE).build());
       Key<HistoryEntry> historyEntryKey = Key.create(historyEntry);
-      Key<BillingEvent.OneTime> oneTimeBillKey =
-          Key.create(historyEntryKey, BillingEvent.OneTime.class, 1);
-      VKey<BillingEvent.Recurring> recurringBillKey =
-          VKey.from(Key.create(historyEntryKey, BillingEvent.Recurring.class, 2));
-      VKey<PollMessage.Autorenew> autorenewPollKey =
-          VKey.from(Key.create(historyEntryKey, PollMessage.Autorenew.class, 3));
-      VKey<PollMessage.OneTime> onetimePollKey =
-          VKey.from(Key.create(historyEntryKey, PollMessage.OneTime.class, 1));
+      BillingEvent.OneTime onetimeBillEvent =
+          new BillingEvent.OneTime.Builder()
+              .setId(1)
+              .setReason(Reason.RENEW)
+              .setTargetId("example.com")
+              .setClientId("TheRegistrar")
+              .setCost(Money.parse("USD 44.00"))
+              .setPeriodYears(4)
+              .setEventTime(fakeClock.nowUtc())
+              .setBillingTime(fakeClock.nowUtc())
+              .setParent(historyEntryKey)
+              .build();
+      persistResource(onetimeBillEvent);
+      Key<BillingEvent.OneTime> oneTimeBillKey = Key.create(onetimeBillEvent);
+      BillingEvent.Recurring recurringBillEvent =
+          new BillingEvent.Recurring.Builder()
+              .setId(2)
+              .setReason(Reason.RENEW)
+              .setFlags(ImmutableSet.of(Flag.AUTO_RENEW))
+              .setTargetId("example.com")
+              .setClientId("TheRegistrar")
+              .setEventTime(fakeClock.nowUtc())
+              .setRecurrenceEndTime(END_OF_TIME)
+              .setParent(historyEntryKey)
+              .build();
+      persistResource(recurringBillEvent);
+      VKey<BillingEvent.Recurring> recurringBillKey = recurringBillEvent.createVKey();
+      PollMessage.Autorenew autorenewPollMessage =
+          new PollMessage.Autorenew.Builder()
+              .setId(3L)
+              .setTargetId("example.com")
+              .setClientId("TheRegistrar")
+              .setEventTime(fakeClock.nowUtc())
+              .setMsg("Domain was auto-renewed.")
+              .setParent(historyEntry)
+              .build();
+      persistResource(autorenewPollMessage);
+      VKey<PollMessage.Autorenew> autorenewPollKey = autorenewPollMessage.createVKey();
+      PollMessage.OneTime oneTimePollMessage =
+          new PollMessage.OneTime.Builder()
+              .setId(1L)
+              .setParent(historyEntry)
+              .setEventTime(fakeClock.nowUtc())
+              .setClientId("TheRegistrar")
+              .setMsg(DomainFlowUtils.COLLISION_MESSAGE)
+              .build();
+      persistResource(oneTimePollMessage);
+      VKey<PollMessage.OneTime> onetimePollKey = oneTimePollMessage.createVKey();
       domain =
           persistResource(
               new DomainBase.Builder()
@@ -220,6 +298,16 @@ class InitSqlPipelineTest {
                           "TheRegistrar",
                           null))
                   .build());
+      persistResource(
+          new BillingEvent.Cancellation.Builder()
+              .setReason(Reason.RENEW)
+              .setTargetId(domain.getDomainName())
+              .setClientId(domain.getCurrentSponsorClientId())
+              .setEventTime(fakeClock.nowUtc())
+              .setBillingTime(fakeClock.nowUtc())
+              .setRecurringEventKey(recurringBillEvent.createVKey())
+              .setParent(historyEntryKey)
+              .build());
       exportDir = store.export(exportRootDir.getAbsolutePath(), ALL_KINDS, ImmutableSet.of());
       commitLogDir = Files.createDirectory(tmpDir.resolve("commits")).toFile();
     }
@@ -241,14 +329,14 @@ class InitSqlPipelineTest {
     initSqlPipeline.run().waitUntilFinish();
     try (AppEngineEnvironment env = new AppEngineEnvironment("test")) {
       assertHostResourceEquals(
-          jpaTm().transact(() -> jpaTm().load(hostResource.createVKey())), hostResource);
-      assertThat(jpaTm().transact(() -> jpaTm().loadAll(Registrar.class)))
+          jpaTm().transact(() -> jpaTm().loadByKey(hostResource.createVKey())), hostResource);
+      assertThat(jpaTm().transact(() -> jpaTm().loadAllOf(Registrar.class)))
           .comparingElementsUsing(immutableObjectCorrespondence("lastUpdateTime"))
           .containsExactly(registrar1, registrar2);
-      assertThat(jpaTm().transact(() -> jpaTm().loadAll(ContactResource.class)))
+      assertThat(jpaTm().transact(() -> jpaTm().loadAllOf(ContactResource.class)))
           .comparingElementsUsing(immutableObjectCorrespondence("revisions", "updateTimestamp"))
           .containsExactly(contact1, contact2);
-      assertCleansedDomainEquals(jpaTm().transact(() -> jpaTm().load(domain.createVKey())), domain);
+      assertDomainEquals(jpaTm().transact(() -> jpaTm().loadByKey(domain.createVKey())), domain);
     }
   }
 
@@ -260,29 +348,26 @@ class InitSqlPipelineTest {
         .isEqualTo(expected.getSuperordinateDomain().getSqlKey());
   }
 
-  private static void assertCleansedDomainEquals(DomainBase actual, DomainBase expected) {
+  private static void assertDomainEquals(DomainBase actual, DomainBase expected) {
     assertAboutImmutableObjects()
         .that(actual)
         .isEqualExceptFields(
             expected,
-            "adminContact",
-            "registrantContact",
-            "gracePeriods",
-            "dsData",
-            "allContacts",
             "revisions",
             "updateTimestamp",
-            "autorenewBillingEvent",
-            "autorenewBillingEventHistoryId",
             "autorenewPollMessage",
-            "autorenewPollMessageHistoryId",
             "deletePollMessage",
-            "deletePollMessageHistoryId",
             "nsHosts",
+            "gracePeriods",
             "transferData");
     assertThat(actual.getAdminContact().getSqlKey())
         .isEqualTo(expected.getAdminContact().getSqlKey());
     assertThat(actual.getRegistrant().getSqlKey()).isEqualTo(expected.getRegistrant().getSqlKey());
-    // TODO(weiminyu): compare gracePeriods, allContacts and dsData, when SQL model supports them.
+    assertThat(actual.getNsHosts()).isEqualTo(expected.getNsHosts());
+    assertThat(actual.getAutorenewPollMessage().getOfyKey())
+        .isEqualTo(expected.getAutorenewPollMessage().getOfyKey());
+    assertThat(actual.getDeletePollMessage().getOfyKey())
+        .isEqualTo(expected.getDeletePollMessage().getOfyKey());
+    // TODO(weiminyu): check gracePeriods and transferData when it is easier to do
   }
 }

core/src/test/java/google/registry/beam/initsql/InitSqlTestUtils.java

@@ -92,16 +92,10 @@ public final class InitSqlTestUtils {
    * VersionedEntities} contains exactly the same elements in the {@code expected} array.
    *
    * <p>This method makes assertions in the pipeline and only use {@link PAssert} on the result.
-   * This has two advantages over {@code PAssert}:
-   *
-   * <ul>
-   *   <li>It supports assertions on 'containsExactlyElementsIn', which is not available in {@code
-   *       PAssert}.
-   *   <li>It supports assertions on Objectify entities, which {@code PAssert} cannot not do.
-   *       Compared with PAssert-compatible options like {@code google.registry.tools.EntityWrapper}
-   *       or {@link EntityProto}, Objectify entities in Java give better-formatted error messages
-   *       when assertions fail.
-   * </ul>
+   * This way it supports assertions on Objectify entities, which {@code PAssert} cannot do ( since
+   * we have not implemented Coders for them). Compared with PAssert-compatible options like {@code
+   * google.registry.tools.EntityWrapper} or {@link EntityProto}, Objectify entities in Java give
+   * better-formatted error messages when assertions fail.
    *
    * <p>Each {@code expected} {@link KV key-value pair} refers to a versioned state of an Ofy
    * entity. The {@link KV#getKey key} is the timestamp, while the {@link KV#getValue value} is

core/src/test/java/google/registry/beam/initsql/WriteToSqlTest.java

@@ -119,7 +119,7 @@ class WriteToSqlTest implements Serializable {
                         .localDbJpaTransactionManager()));
     testPipeline.run().waitUntilFinish();
 
-    ImmutableList<?> sqlContacts = jpaTm().transact(() -> jpaTm().loadAll(ContactResource.class));
+    ImmutableList<?> sqlContacts = jpaTm().transact(() -> jpaTm().loadAllOf(ContactResource.class));
     assertThat(sqlContacts)
         .comparingElementsUsing(immutableObjectCorrespondence("revisions", "updateTimestamp"))
         .containsExactlyElementsIn(

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

@@ -16,16 +16,36 @@ package google.registry.flows;
 
 import static google.registry.testing.DatabaseHelper.loadRegistrar;
 import static google.registry.testing.DatabaseHelper.persistResource;
+import static google.registry.testing.LogsSubject.assertAboutLogs;
+import static google.registry.util.DateTimeUtils.START_OF_TIME;
+import static google.registry.util.X509Utils.encodeX509Certificate;
+import static google.registry.util.X509Utils.encodeX509CertificateFromPemString;
 import static org.joda.time.DateTimeZone.UTC;
 
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableSortedMap;
+import com.google.common.testing.TestLogHandler;
+import google.registry.config.RegistryEnvironment;
+import google.registry.flows.certs.CertificateChecker;
 import google.registry.testing.AppEngineExtension;
 import google.registry.testing.CertificateSamples;
+import google.registry.testing.SystemPropertyExtension;
+import google.registry.util.SelfSignedCaCertificate;
+import java.io.StringWriter;
+import java.security.cert.CertificateException;
+import java.security.cert.X509Certificate;
 import java.util.Optional;
+import java.util.logging.Level;
+import java.util.logging.Logger;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Order;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
+import org.testcontainers.shaded.org.bouncycastle.openssl.jcajce.JcaMiscPEMGenerator;
+import org.testcontainers.shaded.org.bouncycastle.util.io.pem.PemObjectGenerator;
+import org.testcontainers.shaded.org.bouncycastle.util.io.pem.PemWriter;
 
 /** Test logging in with TLS credentials. */
 class EppLoginTlsTest extends EppTestCase {
@@ -34,17 +54,40 @@ class EppLoginTlsTest extends EppTestCase {
   final AppEngineExtension appEngine =
       AppEngineExtension.builder().withDatastoreAndCloudSql().build();
 
-  void setClientCertificateHash(String clientCertificateHash) {
+  @RegisterExtension
+  @Order(value = Integer.MAX_VALUE)
+  final SystemPropertyExtension systemPropertyExtension = new SystemPropertyExtension();
+
+  private final CertificateChecker certificateChecker =
+      new CertificateChecker(
+          ImmutableSortedMap.of(START_OF_TIME, 825, DateTime.parse("2020-09-01T00:00:00Z"), 398),
+          30,
+          2048,
+          ImmutableSet.of("secp256r1", "secp384r1"),
+          clock);
+
+  private final Logger loggerToIntercept =
+      Logger.getLogger(TlsCredentials.class.getCanonicalName());
+  private final TestLogHandler handler = new TestLogHandler();
+
+  private String encodedCertString;
+
+  void setCredentials(String clientCertificateHash, String clientCertificate) {
     setTransportCredentials(
-        new TlsCredentials(true, clientCertificateHash, Optional.of("192.168.1.100:54321")));
+        new TlsCredentials(
+            true,
+            Optional.ofNullable(clientCertificateHash),
+            Optional.ofNullable(clientCertificate),
+            Optional.of("192.168.1.100:54321"),
+            certificateChecker));
   }
 
   @BeforeEach
-  void beforeEach() {
+  void beforeEach() throws CertificateException {
     persistResource(
         loadRegistrar("NewRegistrar")
             .asBuilder()
-            .setClientCertificate(CertificateSamples.SAMPLE_CERT, DateTime.now(UTC))
+            .setClientCertificate(CertificateSamples.SAMPLE_CERT3, DateTime.now(UTC))
             .build());
     // Set a cert for the second registrar, or else any cert will be allowed for login.
     persistResource(
@@ -52,18 +95,20 @@ class EppLoginTlsTest extends EppTestCase {
             .asBuilder()
             .setClientCertificate(CertificateSamples.SAMPLE_CERT2, DateTime.now(UTC))
             .build());
+    loggerToIntercept.addHandler(handler);
+    encodedCertString = encodeX509CertificateFromPemString(CertificateSamples.SAMPLE_CERT3);
   }
 
   @Test
   void testLoginLogout() throws Exception {
-    setClientCertificateHash(CertificateSamples.SAMPLE_CERT_HASH);
+    setCredentials(null, encodedCertString);
     assertThatLoginSucceeds("NewRegistrar", "foo-BAR2");
     assertThatLogoutSucceeds();
   }
 
   @Test
   void testLogin_wrongPasswordFails() throws Exception {
-    setClientCertificateHash(CertificateSamples.SAMPLE_CERT_HASH);
+    setCredentials(CertificateSamples.SAMPLE_CERT3_HASH, encodedCertString);
     // For TLS login, we also check the epp xml password.
     assertThatLogin("NewRegistrar", "incorrect")
         .hasResponse(
@@ -73,7 +118,7 @@ class EppLoginTlsTest extends EppTestCase {
 
   @Test
   void testMultiLogin() throws Exception {
-    setClientCertificateHash(CertificateSamples.SAMPLE_CERT_HASH);
+    setCredentials(CertificateSamples.SAMPLE_CERT3_HASH, encodedCertString);
     assertThatLoginSucceeds("NewRegistrar", "foo-BAR2");
     assertThatLogoutSucceeds();
     assertThatLoginSucceeds("NewRegistrar", "foo-BAR2");
@@ -87,7 +132,7 @@ class EppLoginTlsTest extends EppTestCase {
 
   @Test
   void testNonAuthedLogin_fails() throws Exception {
-    setClientCertificateHash(CertificateSamples.SAMPLE_CERT_HASH);
+    setCredentials(CertificateSamples.SAMPLE_CERT3_HASH, encodedCertString);
     assertThatLogin("TheRegistrar", "password2")
         .hasResponse(
             "response_error.xml",
@@ -97,7 +142,7 @@ class EppLoginTlsTest extends EppTestCase {
 
   @Test
   void testBadCertificate_failsBadCertificate2200() throws Exception {
-    setClientCertificateHash("laffo");
+    setCredentials("laffo", "cert");
     assertThatLogin("NewRegistrar", "foo-BAR2")
         .hasResponse(
             "response_error.xml",
@@ -107,7 +152,7 @@ class EppLoginTlsTest extends EppTestCase {
 
   @Test
   void testGfeDidntProvideClientCertificate_failsMissingCertificate2200() throws Exception {
-    setClientCertificateHash("");
+    setCredentials(null, null);
     assertThatLogin("NewRegistrar", "foo-BAR2")
         .hasResponse(
             "response_error.xml",
@@ -116,12 +161,12 @@ class EppLoginTlsTest extends EppTestCase {
 
   @Test
   void testGoodPrimaryCertificate() throws Exception {
-    setClientCertificateHash(CertificateSamples.SAMPLE_CERT_HASH);
+    setCredentials(null, encodedCertString);
     DateTime now = DateTime.now(UTC);
     persistResource(
         loadRegistrar("NewRegistrar")
             .asBuilder()
-            .setClientCertificate(CertificateSamples.SAMPLE_CERT, now)
+            .setClientCertificate(CertificateSamples.SAMPLE_CERT3, now)
             .setFailoverClientCertificate(CertificateSamples.SAMPLE_CERT2, now)
             .build());
     assertThatLoginSucceeds("NewRegistrar", "foo-BAR2");
@@ -129,33 +174,33 @@ class EppLoginTlsTest extends EppTestCase {
 
   @Test
   void testGoodFailoverCertificate() throws Exception {
-    setClientCertificateHash(CertificateSamples.SAMPLE_CERT2_HASH);
+    setCredentials(null, encodedCertString);
     DateTime now = DateTime.now(UTC);
     persistResource(
         loadRegistrar("NewRegistrar")
             .asBuilder()
             .setClientCertificate(CertificateSamples.SAMPLE_CERT, now)
-            .setFailoverClientCertificate(CertificateSamples.SAMPLE_CERT2, now)
+            .setFailoverClientCertificate(CertificateSamples.SAMPLE_CERT3, now)
             .build());
     assertThatLoginSucceeds("NewRegistrar", "foo-BAR2");
   }
 
   @Test
   void testMissingPrimaryCertificateButHasFailover_usesFailover() throws Exception {
-    setClientCertificateHash(CertificateSamples.SAMPLE_CERT2_HASH);
+    setCredentials(null, encodedCertString);
     DateTime now = DateTime.now(UTC);
     persistResource(
         loadRegistrar("NewRegistrar")
             .asBuilder()
             .setClientCertificate(null, now)
-            .setFailoverClientCertificate(CertificateSamples.SAMPLE_CERT2, now)
+            .setFailoverClientCertificate(CertificateSamples.SAMPLE_CERT3, now)
             .build());
     assertThatLoginSucceeds("NewRegistrar", "foo-BAR2");
   }
 
   @Test
   void testRegistrarHasNoCertificatesOnFile_fails() throws Exception {
-    setClientCertificateHash("laffo");
+    setCredentials("laffo", "cert");
     DateTime now = DateTime.now(UTC);
     persistResource(
         loadRegistrar("NewRegistrar")
@@ -168,4 +213,145 @@ class EppLoginTlsTest extends EppTestCase {
             "response_error.xml",
             ImmutableMap.of("CODE", "2200", "MSG", "Registrar certificate is not configured"));
   }
+
+  @Test
+  void testCertificateDoesNotMeetRequirements_fails() throws Exception {
+    String proxyEncoded = encodeX509CertificateFromPemString(CertificateSamples.SAMPLE_CERT);
+    // SAMPLE_CERT has a validity period that is too long
+    setCredentials(CertificateSamples.SAMPLE_CERT_HASH, proxyEncoded);
+    persistResource(
+        loadRegistrar("NewRegistrar")
+            .asBuilder()
+            .setClientCertificate(CertificateSamples.SAMPLE_CERT, clock.nowUtc())
+            .setFailoverClientCertificate(CertificateSamples.SAMPLE_CERT2, clock.nowUtc())
+            .build());
+    assertThatLogin("NewRegistrar", "foo-BAR2")
+        .hasResponse(
+            "response_error.xml",
+            ImmutableMap.of(
+                "CODE",
+                "2200",
+                "MSG",
+                "Registrar certificate contains the following security violations:\n"
+                    + "Certificate validity period is too long; it must be less than or equal to"
+                    + " 398 days."));
+  }
+
+  @Test
+  void testCertificateDoesNotMeetMultipleRequirements_fails() throws Exception {
+    X509Certificate certificate =
+        SelfSignedCaCertificate.create(
+                "test", clock.nowUtc().plusDays(100), clock.nowUtc().plusDays(5000))
+            .cert();
+
+    StringWriter sw = new StringWriter();
+    try (PemWriter pw = new PemWriter(sw)) {
+      PemObjectGenerator generator = new JcaMiscPEMGenerator(certificate);
+      pw.writeObject(generator);
+    }
+
+    String proxyEncoded = encodeX509Certificate(certificate);
+
+    // SAMPLE_CERT has a validity period that is too long
+    setCredentials(null, proxyEncoded);
+    persistResource(
+        loadRegistrar("NewRegistrar")
+            .asBuilder()
+            .setClientCertificate(sw.toString(), clock.nowUtc())
+            .setFailoverClientCertificate(CertificateSamples.SAMPLE_CERT2, clock.nowUtc())
+            .build());
+    assertThatLogin("NewRegistrar", "foo-BAR2")
+        .hasResponse(
+            "response_error.xml",
+            ImmutableMap.of(
+                "CODE",
+                "2200",
+                "MSG",
+                "Registrar certificate contains the following security violations:\n"
+                    + "Certificate is expired.\n"
+                    + "Certificate validity period is too long; it must be less than or equal to"
+                    + " 398 days."));
+  }
+
+  @Test
+  // TODO(sarahbot@): Remove this test once requirements are enforced in production
+  void testCertificateDoesNotMeetRequirementsInProduction_succeeds() throws Exception {
+    RegistryEnvironment.PRODUCTION.setup(systemPropertyExtension);
+    // SAMPLE_CERT has a validity period that is too long
+    String proxyEncoded = encodeX509CertificateFromPemString(CertificateSamples.SAMPLE_CERT);
+    setCredentials(null, proxyEncoded);
+    persistResource(
+        loadRegistrar("NewRegistrar")
+            .asBuilder()
+            .setClientCertificate(CertificateSamples.SAMPLE_CERT, clock.nowUtc())
+            .setFailoverClientCertificate(CertificateSamples.SAMPLE_CERT2, clock.nowUtc())
+            .build());
+    // Even though the certificate contains security violations, the login will succeed in
+    // production
+    assertThatLoginSucceeds("NewRegistrar", "foo-BAR2");
+    assertAboutLogs()
+        .that(handler)
+        .hasLogAtLevelWithMessage(
+            Level.WARNING,
+            "Registrar certificate used for NewRegistrar does not meet certificate requirements:"
+                + " Certificate validity period is too long; it must be less than or equal to 398"
+                + " days.");
+  }
+
+  @Test
+  void testRegistrarCertificateContainsExtraMetadata_succeeds() throws Exception {
+    String certPem =
+        String.format(
+            "Bag Attributes\n"
+                + "    localKeyID: 1F 1C 3A 3A 4C 03 EC C4 BC 7A C3 21 A9 F2 13 66 21 B8 7B 26 \n"
+                + "subject=/C=US/ST=New York/L=New"
+                + " York/O=Test/OU=ABC/CN=tester.test/emailAddress=test-certificate@test.test\n"
+                + "issuer=/C=US/ST=NY/L=NYC/O=ABC/OU=TEST CA/CN=TEST"
+                + " CA/emailAddress=testing@test.test\n"
+                + "%s",
+            CertificateSamples.SAMPLE_CERT3);
+
+    setCredentials(null, encodeX509CertificateFromPemString(certPem));
+    DateTime now = DateTime.now(UTC);
+    persistResource(
+        loadRegistrar("NewRegistrar")
+            .asBuilder()
+            .setClientCertificate(certPem, now)
+            .setFailoverClientCertificate(CertificateSamples.SAMPLE_CERT2, now)
+            .build());
+    assertThatLoginSucceeds("NewRegistrar", "foo-BAR2");
+  }
+
+  @Test
+  void testRegistrarCertificateContainsExtraMetadataAndViolations_fails() throws Exception {
+    String certPem =
+        String.format(
+            "Bag Attributes\n"
+                + "    localKeyID: 1F 1C 3A 3A 4C 03 EC C4 BC 7A C3 21 A9 F2 13 66 21 B8 7B 26 \n"
+                + "subject=/C=US/ST=New York/L=New"
+                + " York/O=Test/OU=ABC/CN=tester.test/emailAddress=test-certificate@test.test\n"
+                + "issuer=/C=US/ST=NY/L=NYC/O=ABC/OU=TEST CA/CN=TEST"
+                + " CA/emailAddress=testing@test.test\n"
+                + "%s",
+            CertificateSamples.SAMPLE_CERT);
+
+    setCredentials(null, encodeX509CertificateFromPemString(certPem));
+    DateTime now = DateTime.now(UTC);
+    persistResource(
+        loadRegistrar("NewRegistrar")
+            .asBuilder()
+            .setClientCertificate(certPem, now)
+            .setFailoverClientCertificate(CertificateSamples.SAMPLE_CERT2, now)
+            .build());
+    assertThatLogin("NewRegistrar", "foo-BAR2")
+        .hasResponse(
+            "response_error.xml",
+            ImmutableMap.of(
+                "CODE",
+                "2200",
+                "MSG",
+                "Registrar certificate contains the following security violations:\n"
+                    + "Certificate validity period is too long; it must be less than or equal to"
+                    + " 398 days."));
+  }
 }

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

@@ -45,12 +45,8 @@ import javax.inject.Singleton;
 
 /** Dagger component for running EPP tests. */
 @Singleton
-@Component(
-    modules = {
-        ConfigModule.class,
-        EppTestComponent.FakesAndMocksModule.class
-    })
-interface EppTestComponent {
+@Component(modules = {ConfigModule.class, EppTestComponent.FakesAndMocksModule.class})
+public interface EppTestComponent {
 
   RequestComponent startRequest();
 

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

@@ -19,6 +19,7 @@ import static com.google.common.truth.Truth.assertWithMessage;
 import static com.google.common.truth.Truth8.assertThat;
 import static google.registry.testing.TestDataHelper.loadFile;
 import static google.registry.testing.TestLogHandlerUtils.findFirstLogMessageByPrefix;
+import static google.registry.util.DateTimeUtils.START_OF_TIME;
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
@@ -27,8 +28,10 @@ import static org.mockito.Mockito.verify;
 import com.google.common.base.Joiner;
 import com.google.common.base.Splitter;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.flogger.LoggerConfig;
 import com.google.common.testing.TestLogHandler;
+import google.registry.flows.certs.CertificateChecker;
 import google.registry.model.eppcommon.Trid;
 import google.registry.model.eppoutput.EppOutput.ResponseOrGreeting;
 import google.registry.model.eppoutput.EppResponse;
@@ -38,6 +41,7 @@ import google.registry.testing.FakeClock;
 import google.registry.testing.FakeHttpSession;
 import java.util.List;
 import java.util.Optional;
+import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
@@ -54,6 +58,16 @@ class FlowRunnerTest {
 
   private final TestLogHandler handler = new TestLogHandler();
 
+  protected final FakeClock clock = new FakeClock();
+
+  private final CertificateChecker certificateChecker =
+      new CertificateChecker(
+          ImmutableSortedMap.of(START_OF_TIME, 825, DateTime.parse("2020-09-01T00:00:00Z"), 398),
+          30,
+          2048,
+          ImmutableSet.of("secp256r1", "secp384r1"),
+          clock);
+
   static class TestCommandFlow implements Flow {
     @Override
     public ResponseOrGreeting run() {
@@ -137,10 +151,18 @@ class FlowRunnerTest {
 
   @Test
   void testRun_loggingStatement_tlsCredentials() throws Exception {
-    flowRunner.credentials = new TlsCredentials(true, "abc123def", Optional.of("127.0.0.1"));
+    flowRunner.credentials =
+        new TlsCredentials(
+            true,
+            Optional.of("abc123def"),
+            Optional.of("cert046F5A3"),
+            Optional.of("127.0.0.1"),
+            certificateChecker);
     flowRunner.run(eppMetricBuilder);
     assertThat(Splitter.on("\n\t").split(findFirstLogMessageByPrefix(handler, "EPP Command\n\t")))
-        .contains("TlsCredentials{clientCertificateHash=abc123def, clientAddress=/127.0.0.1}");
+        .contains(
+            "TlsCredentials{clientCertificate=cert046F5A3, clientCertificateHash=abc123def,"
+                + " clientAddress=/127.0.0.1}");
   }
 
   @Test

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

@@ -203,7 +203,7 @@ public abstract class FlowTestCase<F extends Flow> {
     assertWithMessage("Billing event is present for grace period: " + gracePeriod)
         .that(gracePeriod.hasBillingEvent())
         .isTrue();
-    return tm().load(
+    return tm().loadByKey(
             firstNonNull(
                 gracePeriod.getOneTimeBillingEvent(), gracePeriod.getRecurringBillingEvent()));
   }

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

@@ -87,7 +87,8 @@ public abstract class ResourceFlowTestCase<F extends Flow, R extends EppResource
     // Force the session to be cleared.
     tm().clearSessionCache();
     @SuppressWarnings("unchecked")
-    T refreshedResource = (T) transactIfJpaTm(() -> tm().load(resource)).cloneProjectedAtTime(now);
+    T refreshedResource =
+        (T) transactIfJpaTm(() -> tm().loadByEntity(resource)).cloneProjectedAtTime(now);
     return refreshedResource;
   }
 

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

@@ -14,17 +14,25 @@
 
 package google.registry.flows;
 
-import static com.google.common.truth.Truth.assertThat;
+import static com.google.common.truth.Truth8.assertThat;
+import static google.registry.testing.CertificateSamples.SAMPLE_CERT;
 import static google.registry.testing.DatabaseHelper.loadRegistrar;
 import static google.registry.testing.DatabaseHelper.persistResource;
-import static org.joda.time.DateTimeZone.UTC;
+import static google.registry.util.DateTimeUtils.START_OF_TIME;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableSortedMap;
+import google.registry.flows.TlsCredentials.BadRegistrarIpAddressException;
+import google.registry.flows.TlsCredentials.MissingRegistrarCertificateException;
+import google.registry.flows.TlsCredentials.RegistrarCertificateNotConfiguredException;
+import google.registry.flows.certs.CertificateChecker;
 import google.registry.model.registrar.Registrar;
-import google.registry.request.HttpException.BadRequestException;
 import google.registry.testing.AppEngineExtension;
+import google.registry.testing.FakeClock;
+import google.registry.util.CidrAddressBlock;
 import java.util.Optional;
 import javax.servlet.http.HttpServletRequest;
 import org.joda.time.DateTime;
@@ -38,31 +46,114 @@ final class TlsCredentialsTest {
   final AppEngineExtension appEngine =
       AppEngineExtension.builder().withDatastoreAndCloudSql().build();
 
+  protected final FakeClock clock = new FakeClock();
+
+  private final CertificateChecker certificateChecker =
+      new CertificateChecker(
+          ImmutableSortedMap.of(START_OF_TIME, 825, DateTime.parse("2020-09-01T00:00:00Z"), 398),
+          30,
+          2048,
+          ImmutableSet.of("secp256r1", "secp384r1"),
+          clock);
+
   @Test
   void testProvideClientCertificateHash() {
     HttpServletRequest req = mock(HttpServletRequest.class);
     when(req.getHeader("X-SSL-Certificate")).thenReturn("data");
-    assertThat(TlsCredentials.EppTlsModule.provideClientCertificateHash(req)).isEqualTo("data");
+    assertThat(TlsCredentials.EppTlsModule.provideClientCertificateHash(req)).hasValue("data");
   }
 
   @Test
-  void testProvideClientCertificateHash_missing() {
-    HttpServletRequest req = mock(HttpServletRequest.class);
-    BadRequestException thrown =
-        assertThrows(
-            BadRequestException.class,
-            () -> TlsCredentials.EppTlsModule.provideClientCertificateHash(req));
-    assertThat(thrown).hasMessageThat().contains("Missing header: X-SSL-Certificate");
+  void testClientCertificateAndHash_missing() {
+    TlsCredentials tls =
+        new TlsCredentials(
+            true,
+            Optional.empty(),
+            Optional.empty(),
+            Optional.of("192.168.1.1"),
+            certificateChecker);
+    persistResource(
+        loadRegistrar("TheRegistrar")
+            .asBuilder()
+            .setClientCertificate(SAMPLE_CERT, clock.nowUtc())
+            .build());
+    assertThrows(
+        MissingRegistrarCertificateException.class,
+        () -> tls.validateCertificate(Registrar.loadByClientId("TheRegistrar").get()));
+  }
+
+  @Test
+  void test_missingIpAddress_doesntAllowAccess() {
+    TlsCredentials tls =
+        new TlsCredentials(
+            false, Optional.of("certHash"), Optional.empty(), Optional.empty(), certificateChecker);
+    persistResource(
+        loadRegistrar("TheRegistrar")
+            .asBuilder()
+            .setClientCertificate(SAMPLE_CERT, clock.nowUtc())
+            .setIpAddressAllowList(ImmutableSet.of(CidrAddressBlock.create("3.5.8.13")))
+            .build());
+    assertThrows(
+        BadRegistrarIpAddressException.class,
+        () -> tls.validate(Registrar.loadByClientId("TheRegistrar").get(), "password"));
   }
 
   @Test
   void test_validateCertificate_canBeConfiguredToBypassCertHashes() throws Exception {
-    TlsCredentials tls = new TlsCredentials(false, "certHash", Optional.of("192.168.1.1"));
+    TlsCredentials tls =
+        new TlsCredentials(
+            false,
+            Optional.of("certHash"),
+            Optional.of("cert"),
+            Optional.of("192.168.1.1"),
+            certificateChecker);
     persistResource(
         loadRegistrar("TheRegistrar")
             .asBuilder()
-            .setClientCertificate(null, DateTime.now(UTC))
-            .setFailoverClientCertificate(null, DateTime.now(UTC))
+            .setClientCertificate(null, clock.nowUtc())
+            .setFailoverClientCertificate(null, clock.nowUtc())
+            .build());
+    // This would throw a RegistrarCertificateNotConfiguredException if cert hashes were not
+    // bypassed
+    tls.validateCertificate(Registrar.loadByClientId("TheRegistrar").get());
+  }
+
+  @Test
+  void testProvideClientCertificate() {
+    HttpServletRequest req = mock(HttpServletRequest.class);
+    when(req.getHeader("X-SSL-Full-Certificate")).thenReturn("data");
+    assertThat(TlsCredentials.EppTlsModule.provideClientCertificate(req)).hasValue("data");
+  }
+
+  @Test
+  void testClientCertificate_notConfigured() {
+    TlsCredentials tls =
+        new TlsCredentials(
+            true,
+            Optional.of("hash"),
+            Optional.of(SAMPLE_CERT),
+            Optional.of("192.168.1.1"),
+            certificateChecker);
+    persistResource(loadRegistrar("TheRegistrar").asBuilder().build());
+    assertThrows(
+        RegistrarCertificateNotConfiguredException.class,
+        () -> tls.validateCertificate(Registrar.loadByClientId("TheRegistrar").get()));
+  }
+
+  @Test
+  void test_validateCertificate_canBeConfiguredToBypassCerts() throws Exception {
+    TlsCredentials tls =
+        new TlsCredentials(
+            false,
+            Optional.of("certHash"),
+            Optional.of("cert"),
+            Optional.of("192.168.1.1"),
+            certificateChecker);
+    persistResource(
+        loadRegistrar("TheRegistrar")
+            .asBuilder()
+            .setClientCertificate(null, clock.nowUtc())
+            .setFailoverClientCertificate(null, clock.nowUtc())
             .build());
     // This would throw a RegistrarCertificateNotConfiguredException if cert hashes wren't bypassed.
     tls.validateCertificate(Registrar.loadByClientId("TheRegistrar").get());

core/src/test/java/google/registry/flows/contact/ContactCheckFlowTest.java

@@ -24,16 +24,18 @@ import google.registry.flows.EppException;
 import google.registry.flows.ResourceCheckFlowTestCase;
 import google.registry.flows.exceptions.TooManyResourceChecksException;
 import google.registry.model.contact.ContactResource;
-import org.junit.jupiter.api.Test;
+import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.TestOfyAndSql;
 
 /** Unit tests for {@link ContactCheckFlow}. */
+@DualDatabaseTest
 class ContactCheckFlowTest extends ResourceCheckFlowTestCase<ContactCheckFlow, ContactResource> {
 
   ContactCheckFlowTest() {
     setEppInput("contact_check.xml");
   }
 
-  @Test
+  @TestOfyAndSql
   void testNothingExists() throws Exception {
     // These ids come from the check xml.
     doCheckTest(
@@ -42,7 +44,7 @@ class ContactCheckFlowTest extends ResourceCheckFlowTestCase<ContactCheckFlow, C
         create(true, "8013sah", null));
   }
 
-  @Test
+  @TestOfyAndSql
   void testOneExists() throws Exception {
     persistActiveContact("sh8013");
     // These ids come from the check xml.
@@ -52,7 +54,7 @@ class ContactCheckFlowTest extends ResourceCheckFlowTestCase<ContactCheckFlow, C
         create(true, "8013sah", null));
   }
 
-  @Test
+  @TestOfyAndSql
   void testOneExistsButWasDeleted() throws Exception {
     persistDeletedContact("sh8013", clock.nowUtc().minusDays(1));
     // These ids come from the check xml.
@@ -62,27 +64,27 @@ class ContactCheckFlowTest extends ResourceCheckFlowTestCase<ContactCheckFlow, C
         create(true, "8013sah", null));
   }
 
-  @Test
+  @TestOfyAndSql
   void testXmlMatches() throws Exception {
     persistActiveContact("sah8013");
     runFlowAssertResponse(loadFile("contact_check_response.xml"));
   }
 
-  @Test
+  @TestOfyAndSql
   void test50IdsAllowed() throws Exception {
     // Make sure we don't have a regression that reduces the number of allowed checks.
     setEppInput("contact_check_50.xml");
     runFlow();
   }
 
-  @Test
+  @TestOfyAndSql
   void testTooManyIds() {
     setEppInput("contact_check_51.xml");
     EppException thrown = assertThrows(TooManyResourceChecksException.class, this::runFlow);
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testIcannActivityReportField_getsLogged() throws Exception {
     runFlow();
     assertIcannReportingActivityFieldLogged("srs-cont-check");

core/src/test/java/google/registry/flows/contact/ContactCreateFlowTest.java

@@ -15,6 +15,7 @@
 package google.registry.flows.contact;
 
 import static com.google.common.truth.Truth.assertThat;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.ContactResourceSubject.assertAboutContacts;
 import static google.registry.testing.DatabaseHelper.assertNoBillingEvents;
 import static google.registry.testing.DatabaseHelper.newContactResource;
@@ -31,10 +32,12 @@ import google.registry.flows.contact.ContactFlowUtils.DeclineContactDisclosureFi
 import google.registry.flows.exceptions.ResourceAlreadyExistsForThisClientException;
 import google.registry.flows.exceptions.ResourceCreateContentionException;
 import google.registry.model.contact.ContactResource;
+import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.TestOfyAndSql;
 import org.joda.time.DateTime;
-import org.junit.jupiter.api.Test;
 
 /** Unit tests for {@link ContactCreateFlow}. */
+@DualDatabaseTest
 class ContactCreateFlowTest extends ResourceFlowTestCase<ContactCreateFlow, ContactResource> {
 
   ContactCreateFlowTest() {
@@ -51,27 +54,29 @@ class ContactCreateFlowTest extends ResourceFlowTestCase<ContactCreateFlow, Cont
         .hasOnlyOneHistoryEntryWhich()
         .hasNoXml();
     assertNoBillingEvents();
-    assertEppResourceIndexEntityFor(reloadResourceByForeignKey());
+    if (tm().isOfy()) {
+      assertEppResourceIndexEntityFor(reloadResourceByForeignKey());
+    }
   }
 
-  @Test
+  @TestOfyAndSql
   void testDryRun() throws Exception {
     dryRunFlowAssertResponse(loadFile("contact_create_response.xml"));
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_neverExisted() throws Exception {
     doSuccessfulTest();
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_existedButWasDeleted() throws Exception {
     persistDeletedContact(getUniqueIdFromCommand(), clock.nowUtc().minusDays(1));
     clock.advanceOneMilli();
     doSuccessfulTest();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_alreadyExists() throws Exception {
     persistActiveContact(getUniqueIdFromCommand());
     ResourceAlreadyExistsForThisClientException thrown =
@@ -83,7 +88,7 @@ class ContactCreateFlowTest extends ResourceFlowTestCase<ContactCreateFlow, Cont
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_resourceContention() throws Exception {
     String targetId = getUniqueIdFromCommand();
     persistResource(
@@ -99,13 +104,13 @@ class ContactCreateFlowTest extends ResourceFlowTestCase<ContactCreateFlow, Cont
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_nonAsciiInLocAddress() throws Exception {
     setEppInput("contact_create_hebrew_loc.xml");
     doSuccessfulTest();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_nonAsciiInIntAddress() {
     setEppInput("contact_create_hebrew_int.xml");
     EppException thrown =
@@ -113,7 +118,7 @@ class ContactCreateFlowTest extends ResourceFlowTestCase<ContactCreateFlow, Cont
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_declineDisclosure() {
     setEppInput("contact_create_decline_disclosure.xml");
     EppException thrown =
@@ -121,7 +126,7 @@ class ContactCreateFlowTest extends ResourceFlowTestCase<ContactCreateFlow, Cont
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testIcannActivityReportField_getsLogged() throws Exception {
     runFlow();
     assertIcannReportingActivityFieldLogged("srs-cont-create");

core/src/test/java/google/registry/flows/contact/ContactDeleteFlowTest.java

@@ -37,10 +37,12 @@ import google.registry.model.contact.ContactResource;
 import google.registry.model.eppcommon.StatusValue;
 import google.registry.model.eppcommon.Trid;
 import google.registry.model.reporting.HistoryEntry;
+import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.TestOfyAndSql;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
 
 /** Unit tests for {@link ContactDeleteFlow}. */
+@DualDatabaseTest
 class ContactDeleteFlowTest extends ResourceFlowTestCase<ContactDeleteFlow, ContactResource> {
 
   @BeforeEach
@@ -48,13 +50,13 @@ class ContactDeleteFlowTest extends ResourceFlowTestCase<ContactDeleteFlow, Cont
     setEppInput("contact_delete.xml");
   }
 
-  @Test
+  @TestOfyAndSql
   void testDryRun() throws Exception {
     persistActiveContact(getUniqueIdFromCommand());
     dryRunFlowAssertResponse(loadFile("contact_delete_response.xml"));
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess() throws Exception {
     persistActiveContact(getUniqueIdFromCommand());
     clock.advanceOneMilli();
@@ -71,7 +73,7 @@ class ContactDeleteFlowTest extends ResourceFlowTestCase<ContactDeleteFlow, Cont
     assertNoBillingEvents();
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_clTridNotSpecified() throws Exception {
     setEppInput("contact_delete_no_cltrid.xml");
     persistActiveContact(getUniqueIdFromCommand());
@@ -89,7 +91,7 @@ class ContactDeleteFlowTest extends ResourceFlowTestCase<ContactDeleteFlow, Cont
     assertNoBillingEvents();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_neverExisted() throws Exception {
     ResourceDoesNotExistException thrown =
         assertThrows(ResourceDoesNotExistException.class, this::runFlow);
@@ -97,7 +99,7 @@ class ContactDeleteFlowTest extends ResourceFlowTestCase<ContactDeleteFlow, Cont
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_existedButWasDeleted() throws Exception {
     persistDeletedContact(getUniqueIdFromCommand(), clock.nowUtc().minusDays(1));
     ResourceDoesNotExistException thrown =
@@ -106,19 +108,19 @@ class ContactDeleteFlowTest extends ResourceFlowTestCase<ContactDeleteFlow, Cont
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_existedButWasClientDeleteProhibited() throws Exception {
     doFailingStatusTest(
         StatusValue.CLIENT_DELETE_PROHIBITED, ResourceStatusProhibitsOperationException.class);
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_existedButWasServerDeleteProhibited() throws Exception {
     doFailingStatusTest(
         StatusValue.SERVER_DELETE_PROHIBITED, ResourceStatusProhibitsOperationException.class);
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_existedButWasPendingDelete() throws Exception {
     doFailingStatusTest(
         StatusValue.PENDING_DELETE, ResourceStatusProhibitsOperationException.class);
@@ -135,7 +137,7 @@ class ContactDeleteFlowTest extends ResourceFlowTestCase<ContactDeleteFlow, Cont
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_unauthorizedClient() throws Exception {
     sessionMetadata.setClientId("NewRegistrar");
     persistActiveContact(getUniqueIdFromCommand());
@@ -143,7 +145,7 @@ class ContactDeleteFlowTest extends ResourceFlowTestCase<ContactDeleteFlow, Cont
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_superuserUnauthorizedClient() throws Exception {
     sessionMetadata.setClientId("NewRegistrar");
     persistActiveContact(getUniqueIdFromCommand());
@@ -161,7 +163,7 @@ class ContactDeleteFlowTest extends ResourceFlowTestCase<ContactDeleteFlow, Cont
     assertNoBillingEvents();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_failfastWhenLinkedToDomain() throws Exception {
     createTld("tld");
     persistResource(
@@ -170,7 +172,7 @@ class ContactDeleteFlowTest extends ResourceFlowTestCase<ContactDeleteFlow, Cont
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_failfastWhenLinkedToApplication() throws Exception {
     createTld("tld");
     persistResource(
@@ -179,7 +181,7 @@ class ContactDeleteFlowTest extends ResourceFlowTestCase<ContactDeleteFlow, Cont
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testIcannActivityReportField_getsLogged() throws Exception {
     persistActiveContact(getUniqueIdFromCommand());
     clock.advanceOneMilli();