Resolve TODOs under Shicong's name (#930)

* Add object comparison to replay tests

Allow optional object comparison in the replay test extension and enable it
for the DomainCreateFlow test.

To faciliate this, add two new field annotations to ImmutableObject:
DoNotCompare, to be used for fields that are not relevant to either database,
and Insignificant, to be used for fields that are mutated after they have been
accessed and therefore violate immutability (there is currently only one of
these, however we might discover more in the course of adding more comparisons
to the replay test.

* Revert commented out premium price error log

* Added static create methods for ReplayExtension

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

@@ -600,7 +600,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
@@ -828,9 +827,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/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/ResaveAllEppResourcesAction.java

@@ -14,6 +14,7 @@
 
 package google.registry.batch;
 
+import static google.registry.mapreduce.MapreduceRunner.PARAM_FAST;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
@@ -24,6 +25,7 @@ import google.registry.mapreduce.MapreduceRunner;
 import google.registry.mapreduce.inputs.EppResourceInputs;
 import google.registry.model.EppResource;
 import google.registry.request.Action;
+import google.registry.request.Parameter;
 import google.registry.request.Response;
 import google.registry.request.auth.Auth;
 import javax.inject.Inject;
@@ -39,6 +41,14 @@ import javax.inject.Inject;
  * <p>Because there are no auth settings in the {@link Action} annotation, this command can only be
  * run internally, or by pretending to be internal by setting the X-AppEngine-QueueName header,
  * which only admin users can do.
+ *
+ * <p>If the <code>?fast=true</code> querystring parameter is passed, then entities that are not
+ * changed by {@link EppResource#cloneProjectedAtTime} will not be re-saved. This helps prevent
+ * mutation load on the DB and has the beneficial side effect of writing out smaller commit logs.
+ * Note that this does NOT pick up mutations caused by migrations using the {@link
+ * com.googlecode.objectify.annotation.OnLoad} annotation, so if you are running a one-off schema
+ * migration, do not use fast mode. Fast mode defaults to false for this reason, but is used by the
+ * monthly invocation of the mapreduce.
  */
 @Action(
     service = Action.Service.BACKEND,
@@ -48,7 +58,13 @@ public class ResaveAllEppResourcesAction implements Runnable {
 
   @Inject MapreduceRunner mrRunner;
   @Inject Response response;
-  @Inject ResaveAllEppResourcesAction() {}
+
+  @Inject
+  @Parameter(PARAM_FAST)
+  boolean isFast;
+
+  @Inject
+  ResaveAllEppResourcesAction() {}
 
   /**
    * The number of shards to run the map-only mapreduce on.
@@ -66,7 +82,7 @@ public class ResaveAllEppResourcesAction implements Runnable {
         .setModuleName("backend")
         .setDefaultMapShards(NUM_SHARDS)
         .runMapOnly(
-            new ResaveAllEppResourcesActionMapper(),
+            new ResaveAllEppResourcesActionMapper(isFast),
             ImmutableList.of(EppResourceInputs.createKeyInput(EppResource.class)))
         .sendLinkToMapreduceConsole(response);
   }
@@ -76,23 +92,33 @@ public class ResaveAllEppResourcesAction implements Runnable {
       extends Mapper<Key<EppResource>, Void, Void> {
 
     private static final long serialVersionUID = -7721628665138087001L;
-    public ResaveAllEppResourcesActionMapper() {}
+
+    private final boolean isFast;
+
+    ResaveAllEppResourcesActionMapper(boolean isFast) {
+      this.isFast = isFast;
+    }
 
     @Override
     public final void map(final Key<EppResource> resourceKey) {
-      tm()
-          .transact(
-              () -> {
-                EppResource projectedResource =
-                    ofy()
-                        .load()
-                        .key(resourceKey)
-                        .now()
-                        .cloneProjectedAtTime(tm().getTransactionTime());
-                ofy().save().entity(projectedResource).now();
-              });
-      getContext().incrementCounter(String.format("%s entities re-saved", resourceKey.getKind()));
+      boolean resaved =
+          tm().transact(
+                  () -> {
+                    EppResource originalResource = ofy().load().key(resourceKey).now();
+                    EppResource projectedResource =
+                        originalResource.cloneProjectedAtTime(tm().getTransactionTime());
+                    if (isFast && originalResource.equals(projectedResource)) {
+                      return false;
+                    } else {
+                      ofy().save().entity(projectedResource).now();
+                      return true;
+                    }
+                  });
+      getContext()
+          .incrementCounter(
+              String.format(
+                  "%s entities %s",
+                  resourceKey.getKind(), resaved ? "re-saved" : "with no changes skipped"));
     }
   }
 }
-

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,28 +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.ofy.Ofy;
+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;
@@ -79,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
@@ -289,7 +292,7 @@ public final class Transforms {
         maxWriters,
         batchSize,
         jpaSupplier,
-        (e) -> ofy().toPojo(e.getEntity().get()),
+        Transforms::convertVersionedEntityToSqlEntity,
         TypeDescriptor.of(VersionedEntity.class));
   }
 
@@ -330,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 {}
 
@@ -429,26 +471,24 @@ 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 Ofy ofy;
-    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();
-        ofy = ObjectifyService.ofy();
       }
 
       synchronized (SqlBatchWriter.class) {
@@ -477,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/alpha/default/WEB-INF/cron.xml

@@ -80,7 +80,7 @@
   </cron>
 
   <cron>
-    <url><![CDATA[/_dr/task/resaveAllEppResources]]></url>
+    <url><![CDATA[/_dr/task/resaveAllEppResources?fast=true]]></url>
     <description>
       This job resaves all our resources, projected in time to "now".
       It is needed for "deleteOldCommitLogs" to work correctly.

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

@@ -103,7 +103,7 @@
   </cron>
 
   <cron>
-    <url><![CDATA[/_dr/task/resaveAllEppResources]]></url>
+    <url><![CDATA[/_dr/task/resaveAllEppResources?fast=true]]></url>
     <description>
       This job resaves all our resources, projected in time to "now".
       It is needed for "deleteOldCommitLogs" to work correctly.

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

@@ -21,7 +21,7 @@
   </cron>
 
   <cron>
-    <url><![CDATA[/_dr/task/resaveAllEppResources]]></url>
+    <url><![CDATA[/_dr/task/resaveAllEppResources?fast=true]]></url>
     <description>
       This job resaves all our resources, projected in time to "now".
       It is needed for "deleteOldCommitLogs" to work correctly.

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

@@ -87,7 +87,7 @@
   </cron>
 
   <cron>
-    <url><![CDATA[/_dr/task/resaveAllEppResources]]></url>
+    <url><![CDATA[/_dr/task/resaveAllEppResources?fast=true]]></url>
     <description>
       This job resaves all our resources, projected in time to "now".
       It is needed for "deleteOldCommitLogs" to work correctly.

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,7 @@
 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 com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
@@ -56,17 +54,17 @@ 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<InetAddress> clientInetAddr;
 
   @Inject
   public TlsCredentials(
       @Config("requireSslCertificates") boolean requireSslCertificates,
-      @Header("X-SSL-Certificate") String clientCertificateHash,
+      @Header("X-SSL-Certificate") Optional<String> clientCertificateHash,
       @Header("X-Forwarded-For") Optional<String> clientAddress) {
     this.requireSslCertificates = requireSslCertificates;
     this.clientCertificateHash = clientCertificateHash;
-    this.clientInetAddr = clientAddress.isPresent() ? parseInetAddress(clientAddress.get()) : null;
+    this.clientInetAddr = clientAddress.map(TlsCredentials::parseInetAddress);
   }
 
   static InetAddress parseInetAddress(String asciiAddr) {
@@ -97,10 +95,14 @@ public class TlsCredentials implements TransportCredentials {
           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(
@@ -118,8 +120,8 @@ public class TlsCredentials implements TransportCredentials {
    */
   @VisibleForTesting
   void validateCertificate(Registrar registrar) throws AuthenticationErrorException {
-    if (isNullOrEmpty(registrar.getClientCertificateHash())
-        && isNullOrEmpty(registrar.getFailoverClientCertificateHash())) {
+    if (!registrar.getClientCertificateHash().isPresent()
+        && !registrar.getFailoverClientCertificateHash().isPresent()) {
       if (requireSslCertificates) {
         throw new RegistrarCertificateNotConfiguredException();
       } else {
@@ -128,7 +130,7 @@ public class TlsCredentials implements TransportCredentials {
         return;
       }
     }
-    if (isNullOrEmpty(clientCertificateHash)) {
+    if (!clientCertificateHash.isPresent()) {
       logger.atInfo().log("Request did not include X-SSL-Certificate");
       throw new MissingRegistrarCertificateException();
     }
@@ -154,21 +156,21 @@ public class TlsCredentials implements TransportCredentials {
   @Override
   public String toString() {
     return toStringHelper(getClass())
-        .add("clientCertificateHash", clientCertificateHash)
-        .add("clientAddress", clientInetAddr)
+        .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 not present. */
   public static class MissingRegistrarCertificateException extends AuthenticationErrorException {
-    public MissingRegistrarCertificateException() {
+    MissingRegistrarCertificateException() {
       super("Registrar certificate not present");
     }
   }
@@ -176,14 +178,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 +193,13 @@ 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

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

@@ -15,14 +15,13 @@
 package google.registry.flows.domain.token;
 
 import static com.google.common.base.Preconditions.checkArgument;
-import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Maps;
 import com.google.common.net.InternetDomainName;
-import com.googlecode.objectify.Key;
 import google.registry.flows.EppException;
 import google.registry.flows.EppException.AssociationProhibitsOperationException;
 import google.registry.flows.EppException.AuthorizationErrorException;
@@ -153,14 +152,15 @@ public class AllocationTokenFlowUtils {
       // See https://tools.ietf.org/html/draft-ietf-regext-allocation-token-04#section-2.1
       throw new InvalidAllocationTokenException();
     }
-    AllocationToken tokenEntity = ofy().load().key(Key.create(AllocationToken.class, token)).now();
-    if (tokenEntity == null) {
+    Optional<AllocationToken> maybeTokenEntity =
+        tm().loadByKeyIfPresent(VKey.create(AllocationToken.class, token));
+    if (maybeTokenEntity.isEmpty()) {
       throw new InvalidAllocationTokenException();
     }
-    if (tokenEntity.isRedeemed()) {
+    if (maybeTokenEntity.get().isRedeemed()) {
       throw new AlreadyRedeemedAllocationTokenException();
     }
-    return tokenEntity;
+    return maybeTokenEntity.get();
   }
 
   // Note: exception messages should be <= 32 characters long for domain check results

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

@@ -21,10 +21,8 @@ import static google.registry.flows.host.HostFlowUtils.validateHostName;
 import static google.registry.flows.host.HostFlowUtils.verifySuperordinateDomainNotInPendingDelete;
 import static google.registry.flows.host.HostFlowUtils.verifySuperordinateDomainOwnership;
 import static google.registry.model.EppResourceUtils.createRepoId;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.CollectionUtils.isNullOrEmpty;
-import static google.registry.util.CollectionUtils.union;
 
 import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
@@ -137,13 +135,11 @@ public final class HostCreateFlow implements TransactionalFlow {
     ImmutableSet<ImmutableObject> entitiesToSave =
         ImmutableSet.of(
             newHost,
-            historyBuilder.build(),
+            historyBuilder.build().toChildHistoryEntity(),
             ForeignKeyIndex.create(newHost, newHost.getDeletionTime()),
             EppResourceIndex.create(Key.create(newHost)));
     if (superordinateDomain.isPresent()) {
-      entitiesToSave =
-          union(
-              entitiesToSave,
+      tm().update(
               superordinateDomain
                   .get()
                   .asBuilder()
@@ -153,7 +149,7 @@ public final class HostCreateFlow implements TransactionalFlow {
       // they are only written as NS records from the referencing domain.
       dnsQueue.addHostRefreshTask(targetId);
     }
-    ofy().save().entities(entitiesToSave);
+    tm().insertAll(entitiesToSave);
     return responseBuilder.setResData(HostCreateData.create(targetId, now)).build();
   }
 

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/mapreduce/MapreduceModule.java

@@ -15,6 +15,7 @@
 package google.registry.mapreduce;
 
 import static google.registry.mapreduce.MapreduceRunner.PARAM_DRY_RUN;
+import static google.registry.mapreduce.MapreduceRunner.PARAM_FAST;
 import static google.registry.mapreduce.MapreduceRunner.PARAM_MAP_SHARDS;
 import static google.registry.mapreduce.MapreduceRunner.PARAM_REDUCE_SHARDS;
 import static google.registry.request.RequestParameters.extractBooleanParameter;
@@ -36,6 +37,12 @@ public final class MapreduceModule {
     return extractBooleanParameter(req, PARAM_DRY_RUN);
   }
 
+  @Provides
+  @Parameter(PARAM_FAST)
+  static boolean provideIsFast(HttpServletRequest req) {
+    return extractBooleanParameter(req, PARAM_FAST);
+  }
+
   @Provides
   @Parameter(PARAM_MAP_SHARDS)
   static Optional<Integer> provideMapShards(HttpServletRequest req) {

core/src/main/java/google/registry/mapreduce/MapreduceRunner.java

@@ -55,6 +55,7 @@ public class MapreduceRunner {
   public static final String PARAM_DRY_RUN = "dryRun";
   public static final String PARAM_MAP_SHARDS = "mapShards";
   public static final String PARAM_REDUCE_SHARDS = "reduceShards";
+  public static final String PARAM_FAST = "fast";
 
   private static final String BASE_URL = "/_dr/mapreduce/";
   private static final String QUEUE_NAME = "mapreduce";

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,7 +17,9 @@ 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;
 import static google.registry.util.DateTimeUtils.isBeforeOrAt;
 import static google.registry.util.DateTimeUtils.latestOf;
@@ -27,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;
@@ -36,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;
@@ -43,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;
 
@@ -56,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());
@@ -135,7 +155,7 @@ public final class EppResourceUtils {
         useCache
             ? ForeignKeyIndex.loadCached(clazz, ImmutableList.of(foreignKey), now)
                 .getOrDefault(foreignKey, null)
-            : ofy().load().type(ForeignKeyIndex.mapToFkiClass(clazz)).id(foreignKey).now();
+            : ForeignKeyIndex.load(clazz, foreignKey, now);
     // The value of fki.getResourceKey() might be null for hard-deleted prober data.
     if (fki == null || isAtOrAfter(now, fki.getDeletionTime()) || fki.getResourceKey() == null) {
       return Optional.empty();
@@ -143,7 +163,7 @@ public final class EppResourceUtils {
     T resource =
         useCache
             ? EppResource.loadCached(fki.getResourceKey())
-            : tm().maybeLoad(fki.getResourceKey()).orElse(null);
+            : transactIfJpaTm(() -> tm().loadByKeyIfPresent(fki.getResourceKey()).orElse(null));
     if (resource == null || isAtOrAfter(now, resource.getDeletionTime())) {
       return Optional.empty();
     }
@@ -364,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());
+              });
+    }
   }
 
   /**
@@ -389,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/UpdateAutoTimestamp.java

@@ -22,7 +22,7 @@ import javax.annotation.Nullable;
 import org.joda.time.DateTime;
 
 /**
- * A timestamp that auto-updates on each save to Datastore.
+ * A timestamp that auto-updates on each save to Datastore/Cloud SQL.
  *
  * @see UpdateAutoTimestampTranslatorFactory
  */
@@ -31,7 +31,7 @@ public class UpdateAutoTimestamp extends ImmutableObject {
   // When set to true, database converters/translators should do tha auto update.  When set to
   // false, auto update should be suspended (this exists to allow us to preserve the original value
   // during a replay).
-  static ThreadLocal<Boolean> autoUpdateEnabled = ThreadLocal.withInitial(() -> true);
+  private static ThreadLocal<Boolean> autoUpdateEnabled = ThreadLocal.withInitial(() -> true);
 
   DateTime timestamp;
 

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

@@ -75,7 +75,9 @@ public class DomainBase extends DomainContent
   }
 
   @ElementCollection
-  @JoinTable(name = "DomainHost")
+  @JoinTable(
+      name = "DomainHost",
+      indexes = {@Index(columnList = "domain_repo_id,host_repo_id", unique = true)})
   @Access(AccessType.PROPERTY)
   @Column(name = "host_repo_id")
   public Set<VKey<HostResource>> getNsHosts() {

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

@@ -50,6 +50,7 @@ import com.googlecode.objectify.condition.IfNull;
 import google.registry.flows.ResourceFlowUtils;
 import google.registry.model.EppResource;
 import google.registry.model.EppResource.ResourceWithTransferData;
+import google.registry.model.ImmutableObject.EmptySetToNull;
 import google.registry.model.billing.BillingEvent;
 import google.registry.model.common.EntityGroupRoot;
 import google.registry.model.contact.ContactResource;
@@ -132,7 +133,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}.
@@ -319,6 +320,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()))

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

@@ -96,7 +96,15 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
   // TODO(b/166776754): Investigate if we can reuse domainContent.nsHosts for storing host keys.
   @Ignore
   @ElementCollection
-  @JoinTable(name = "DomainHistoryHost")
+  @JoinTable(
+      name = "DomainHistoryHost",
+      indexes = {
+        @Index(
+            columnList =
+                "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;
 
@@ -173,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;
@@ -177,12 +178,7 @@ public class Address extends ImmutableObject implements Jsonifiable {
    * persisted in the Datastore. TODO(shicong): Delete this method after database migration.
    */
   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/index/ForeignKeyIndex.java

@@ -15,9 +15,11 @@
 package google.registry.model.index;
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static google.registry.config.RegistryConfig.getEppResourceCachingDuration;
 import static google.registry.config.RegistryConfig.getEppResourceMaxCachedEntries;
 import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.TypeUtils.instantiate;
 
@@ -29,6 +31,7 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Maps;
+import com.google.common.collect.Multimaps;
 import com.google.common.collect.Streams;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Entity;
@@ -44,6 +47,8 @@ import google.registry.model.host.HostResource;
 import google.registry.persistence.VKey;
 import google.registry.schema.replay.DatastoreOnlyEntity;
 import google.registry.util.NonFinalForTesting;
+import java.util.Comparator;
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.ExecutionException;
@@ -76,13 +81,21 @@ public abstract class ForeignKeyIndex<E extends EppResource> extends BackupGroup
   public static class ForeignKeyHostIndex extends ForeignKeyIndex<HostResource>
       implements DatastoreOnlyEntity {}
 
-  static final ImmutableMap<Class<? extends EppResource>, Class<? extends ForeignKeyIndex<?>>>
+  private static final ImmutableMap<
+          Class<? extends EppResource>, Class<? extends ForeignKeyIndex<?>>>
       RESOURCE_CLASS_TO_FKI_CLASS =
           ImmutableMap.of(
               ContactResource.class, ForeignKeyContactIndex.class,
               DomainBase.class, ForeignKeyDomainIndex.class,
               HostResource.class, ForeignKeyHostIndex.class);
 
+  private static final ImmutableMap<Class<? extends EppResource>, String>
+      RESOURCE_CLASS_TO_FKI_PROPERTY =
+          ImmutableMap.of(
+              ContactResource.class, "contactId",
+              DomainBase.class, "fullyQualifiedDomainName",
+              HostResource.class, "fullyQualifiedHostName");
+
   @Id String foreignKey;
 
   /**
@@ -179,9 +192,42 @@ public abstract class ForeignKeyIndex<E extends EppResource> extends BackupGroup
    */
   public static <E extends EppResource> ImmutableMap<String, ForeignKeyIndex<E>> load(
       Class<E> clazz, Iterable<String> foreignKeys, final DateTime now) {
-    return ofy().load().type(mapToFkiClass(clazz)).ids(foreignKeys).entrySet().stream()
-        .filter(e -> now.isBefore(e.getValue().deletionTime))
-        .collect(ImmutableMap.toImmutableMap(Map.Entry::getKey, Map.Entry::getValue));
+    if (tm().isOfy()) {
+      return ofy().load().type(mapToFkiClass(clazz)).ids(foreignKeys).entrySet().stream()
+          .filter(e -> now.isBefore(e.getValue().deletionTime))
+          .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue));
+    } else {
+      String property = RESOURCE_CLASS_TO_FKI_PROPERTY.get(clazz);
+      List<E> entities =
+          tm().transact(
+                  () -> {
+                    String entityName =
+                        jpaTm().getEntityManager().getMetamodel().entity(clazz).getName();
+                    return jpaTm()
+                        .getEntityManager()
+                        .createQuery(
+                            String.format(
+                                "FROM %s WHERE %s IN :propertyValue and deletionTime > :now ",
+                                entityName, property),
+                            clazz)
+                        .setParameter("propertyValue", foreignKeys)
+                        .setParameter("now", now)
+                        .getResultList();
+                  });
+      // We need to find and return the entities with the maximum deletionTime for each foreign key.
+      return Multimaps.index(entities, EppResource::getForeignKey).asMap().entrySet().stream()
+          .map(
+              entry ->
+                  Maps.immutableEntry(
+                      entry.getKey(),
+                      entry.getValue().stream()
+                          .max(Comparator.comparing(EppResource::getDeletionTime))
+                          .get()))
+          .collect(
+              toImmutableMap(
+                  Map.Entry::getKey,
+                  entry -> create(entry.getValue(), entry.getValue().getDeletionTime())));
+    }
   }
 
   static final CacheLoader<Key<ForeignKeyIndex<?>>, Optional<ForeignKeyIndex<?>>> CACHE_LOADER =
@@ -266,7 +312,7 @@ public abstract class ForeignKeyIndex<E extends EppResource> extends BackupGroup
               .filter(entry -> entry.getValue().isPresent())
               .filter(entry -> now.isBefore(entry.getValue().get().getDeletionTime()))
               .collect(
-                  ImmutableMap.toImmutableMap(
+                  toImmutableMap(
                       entry -> entry.getKey().getName(),
                       entry -> (ForeignKeyIndex<E>) entry.getValue().get()));
       return fkisFromCache;

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,6 +24,8 @@ 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;
@@ -177,26 +180,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 +201,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

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

@@ -52,6 +52,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 +186,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 +291,7 @@ public abstract class PollMessage extends ImmutableObject
 
     @Transient List<ContactTransferResponse> contactTransferResponses;
 
-    @Transient
+    @Transient @ImmutableObject.DoNotCompare
     List<DomainPendingActionNotificationResponse> domainPendingActionNotificationResponses;
 
     @Transient List<DomainTransferResponse> domainTransferResponses;
@@ -355,6 +357,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));
@@ -456,6 +463,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/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

@@ -69,7 +69,7 @@ public abstract class TransferData<
   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.
+  // TODO(b/177589157): Make transfer flows work with Cloud SQL.
   @Ignore
   @Column(name = "transfer_gaining_poll_message_id")
   Long gainingTransferPollMessageId;

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

@@ -29,13 +29,18 @@ import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Streams;
 import com.google.common.flogger.FluentLogger;
 import google.registry.config.RegistryConfig;
+import google.registry.model.ImmutableObject;
+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.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;
@@ -56,11 +61,24 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
   private static final Retrier retrier = new Retrier(new SystemSleeper(), 3);
 
+  // The entity of classes in this set will be simply ignored when passed to modification
+  // operations, i.e. insert, put, update and delete. This is to help maintain a single code path
+  // when we switch from ofy() to tm() for the database migration as we don't need have a condition
+  // to exclude the Datastore specific entities when the underlying tm() is jpaTm().
+  // TODO(b/176108270): Remove this property after database migration.
+  private static final ImmutableSet<Class<? extends ImmutableObject>> IGNORED_ENTITY_CLASSES =
+      ImmutableSet.of(
+          EppResourceIndex.class,
+          ForeignKeyContactIndex.class,
+          ForeignKeyDomainIndex.class,
+          ForeignKeyHostIndex.class);
+
   // 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);
 
@@ -228,6 +246,9 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   @Override
   public void insert(Object entity) {
     checkArgumentNotNull(entity, "entity must be specified");
+    if (isEntityOfIgnoredClass(entity)) {
+      return;
+    }
     assertInTransaction();
     getEntityManager().persist(entity);
     transactionInfo.get().addUpdate(entity);
@@ -253,6 +274,9 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   @Override
   public void put(Object entity) {
     checkArgumentNotNull(entity, "entity must be specified");
+    if (isEntityOfIgnoredClass(entity)) {
+      return;
+    }
     assertInTransaction();
     getEntityManager().merge(entity);
     transactionInfo.get().addUpdate(entity);
@@ -278,6 +302,9 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   @Override
   public void update(Object entity) {
     checkArgumentNotNull(entity, "entity must be specified");
+    if (isEntityOfIgnoredClass(entity)) {
+      return;
+    }
     assertInTransaction();
     checkArgument(exists(entity), "Given entity does not exist");
     getEntityManager().merge(entity);
@@ -330,33 +357,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)
@@ -367,11 +376,58 @@ 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)
+        .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();
+    return (T)
+        loadByKey(
+            VKey.createSql(entity.getClass(), emf.getPersistenceUnitUtil().getIdentifier(entity)));
+  }
+
+  @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(
@@ -382,11 +438,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();
@@ -414,6 +465,9 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   @Override
   public void delete(Object entity) {
     checkArgumentNotNull(entity, "entity must be specified");
+    if (isEntityOfIgnoredClass(entity)) {
+      return;
+    }
     assertInTransaction();
     Object managedEntity = entity;
     if (!getEntityManager().contains(entity)) {
@@ -464,6 +518,10 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     }
   }
 
+  private static boolean isEntityOfIgnoredClass(Object entity) {
+    return IGNORED_ENTITY_CLASSES.contains(entity.getClass());
+  }
+
   private static ImmutableSet<EntityId> getEntityIdsFromEntity(
       EntityType<?> entityType, Object entity) {
     if (entityType.hasSingleIdAttribute()) {

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

@@ -342,7 +342,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 +425,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 +562,7 @@ public class RdapJsonFormatter {
           .statusBuilder()
           .addAll(
               makeStatusValueList(
-                  isLinked(Key.create(contactResource), getRequestTime())
+                  isLinked(contactResource.createVKey(), getRequestTime())
                       ? union(contactResource.getStatusValues(), StatusValue.LINKED)
                       : contactResource.getStatusValues(),
                   false,

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/reporting/spec11/Spec11RegistrarThreatMatchesParser.java

@@ -67,7 +67,6 @@ public class Spec11RegistrarThreatMatchesParser {
     if (!gcsUtils.existsAndNotEmpty(spec11ReportFilename)) {
       return ImmutableSet.of();
     }
-    ImmutableSet.Builder<RegistrarThreatMatches> builder = ImmutableSet.builder();
     try (InputStream in = gcsUtils.openInputStream(spec11ReportFilename);
         InputStreamReader isr = new InputStreamReader(in, UTF_8)) {
       // Skip the header at line 0

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

@@ -14,19 +14,19 @@
 
 package google.registry.tools;
 
+import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static com.google.common.collect.Iterables.partition;
 import static com.google.common.collect.Streams.stream;
 import static google.registry.model.domain.token.AllocationToken.TokenType.SINGLE_USE;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
 import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableSet;
-import com.googlecode.objectify.Key;
 import google.registry.model.domain.token.AllocationToken;
+import google.registry.persistence.VKey;
 import java.util.List;
 
 /**
@@ -48,7 +48,7 @@ final class DeleteAllocationTokensCommand extends UpdateOrDeleteAllocationTokens
   private static final int BATCH_SIZE = 20;
   private static final Joiner JOINER = Joiner.on(", ");
 
-  private ImmutableSet<Key<AllocationToken>> tokensToDelete;
+  private ImmutableSet<VKey<AllocationToken>> tokensToDelete;
 
   @Override
   public void init() {
@@ -71,26 +71,24 @@ final class DeleteAllocationTokensCommand extends UpdateOrDeleteAllocationTokens
   }
 
   /** Deletes a (filtered) batch of AllocationTokens and returns how many were deleted. */
-  private long deleteBatch(List<Key<AllocationToken>> batch) {
+  private long deleteBatch(List<VKey<AllocationToken>> batch) {
     // 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<AllocationToken> tokensToDelete =
-        ofy().load().keys(batch).values().stream()
-            .filter(t -> withDomains || !t.getDomainName().isPresent())
+    ImmutableSet<VKey<AllocationToken>> tokensToDelete =
+        tm().loadByKeys(batch).values().stream()
+            .filter(t -> withDomains || t.getDomainName().isEmpty())
             .filter(t -> SINGLE_USE.equals(t.getTokenType()))
             .filter(t -> !t.isRedeemed())
+            .map(AllocationToken::createVKey)
             .collect(toImmutableSet());
     if (!dryRun) {
-      ofy().delete().entities(tokensToDelete);
+      tm().delete(tokensToDelete);
     }
     System.out.printf(
         "%s tokens: %s\n",
         dryRun ? "Would delete" : "Deleted",
         JOINER.join(
-            tokensToDelete.stream()
-                .map(AllocationToken::getToken)
-                .sorted()
-                .collect(toImmutableSet())));
+            tokensToDelete.stream().map(VKey::getSqlKey).sorted().collect(toImmutableList())));
     return tokensToDelete.size();
   }
 }

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

@@ -20,8 +20,8 @@ import static com.google.common.collect.Queues.newArrayDeque;
 import static com.google.common.collect.Sets.difference;
 import static google.registry.model.domain.token.AllocationToken.TokenType.SINGLE_USE;
 import static google.registry.model.domain.token.AllocationToken.TokenType.UNLIMITED_USE;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 import static google.registry.util.CollectionUtils.nullToEmpty;
 import static google.registry.util.StringGenerator.DEFAULT_PASSWORD_LENGTH;
 import static java.nio.charset.StandardCharsets.UTF_8;
@@ -38,10 +38,10 @@ import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Streams;
 import com.google.common.io.Files;
-import com.googlecode.objectify.Key;
 import google.registry.model.domain.token.AllocationToken;
 import google.registry.model.domain.token.AllocationToken.TokenStatus;
 import google.registry.model.domain.token.AllocationToken.TokenType;
+import google.registry.persistence.VKey;
 import google.registry.tools.params.TransitionListParameter.TokenStatusTransitions;
 import google.registry.util.CollectionUtils;
 import google.registry.util.NonFinalForTesting;
@@ -258,8 +258,13 @@ class GenerateAllocationTokensCommand implements CommandWithRemoteApi {
 
   @VisibleForTesting
   int saveTokens(final ImmutableSet<AllocationToken> tokens) {
-    Collection<AllocationToken> savedTokens =
-        dryRun ? tokens : tm().transact(() -> ofy().save().entities(tokens).now().values());
+    Collection<AllocationToken> savedTokens;
+    if (dryRun) {
+      savedTokens = tokens;
+    } else {
+      transactIfJpaTm(() -> tm().transact(() -> tm().putAll(tokens)));
+      savedTokens = tm().transact(() -> tm().loadByEntities(tokens));
+    }
     savedTokens.forEach(
         t -> System.out.println(SKIP_NULLS.join(t.getDomainName().orElse(null), t.getToken())));
     return savedTokens.size();
@@ -282,12 +287,14 @@ class GenerateAllocationTokensCommand implements CommandWithRemoteApi {
   }
 
   private ImmutableSet<String> getExistingTokenStrings(ImmutableSet<String> candidates) {
-    ImmutableSet<Key<AllocationToken>> existingTokenKeys =
+    ImmutableSet<VKey<AllocationToken>> existingTokenKeys =
         candidates.stream()
-            .map(input -> Key.create(AllocationToken.class, input))
+            .map(input -> VKey.create(AllocationToken.class, input))
             .collect(toImmutableSet());
-    return ofy().load().keys(existingTokenKeys).values().stream()
-        .map(AllocationToken::getToken)
-        .collect(toImmutableSet());
+    return transactIfJpaTm(
+        () ->
+            tm().loadByKeysIfPresent(existingTokenKeys).values().stream()
+                .map(AllocationToken::getToken)
+                .collect(toImmutableSet()));
   }
 }

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

@@ -15,7 +15,6 @@
 package google.registry.tools;
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.beust.jcommander.Parameter;
@@ -26,6 +25,7 @@ import com.google.common.collect.Lists;
 import com.googlecode.objectify.Key;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.token.AllocationToken;
+import google.registry.persistence.VKey;
 import java.util.Collection;
 import java.util.List;
 import java.util.Optional;
@@ -45,22 +45,27 @@ final class GetAllocationTokenCommand implements CommandWithRemoteApi {
   public void run() {
     ImmutableMap.Builder<String, AllocationToken> builder = new ImmutableMap.Builder<>();
     for (List<String> tokens : Lists.partition(mainParameters, BATCH_SIZE)) {
-      ImmutableList<Key<AllocationToken>> tokenKeys =
-          tokens.stream().map(t -> Key.create(AllocationToken.class, t)).collect(toImmutableList());
-      ofy().load().keys(tokenKeys).forEach((k, v) -> builder.put(k.getName(), v));
+      ImmutableList<VKey<AllocationToken>> tokenKeys =
+          tokens.stream()
+              .map(t -> VKey.create(AllocationToken.class, t))
+              .collect(toImmutableList());
+      tm().loadByKeysIfPresent(tokenKeys)
+          .forEach((k, v) -> builder.put(k.getSqlKey().toString(), v));
     }
     ImmutableMap<String, AllocationToken> loadedTokens = builder.build();
-    ImmutableMap<Key<DomainBase>, DomainBase> domains = loadRedeemedDomains(loadedTokens.values());
+    ImmutableMap<VKey<DomainBase>, DomainBase> domains = loadRedeemedDomains(loadedTokens.values());
 
     for (String token : mainParameters) {
       if (loadedTokens.containsKey(token)) {
         AllocationToken loadedToken = loadedTokens.get(token);
         System.out.println(loadedToken.toString());
-        if (!loadedToken.getRedemptionHistoryEntry().isPresent()) {
+        if (loadedToken.getRedemptionHistoryEntry().isEmpty()) {
           System.out.printf("Token %s was not redeemed.\n", token);
         } else {
+          Key<DomainBase> domainOfyKey =
+              loadedToken.getRedemptionHistoryEntry().get().getOfyKey().getParent();
           DomainBase domain =
-              domains.get(loadedToken.getRedemptionHistoryEntry().get().getOfyKey().getParent());
+              domains.get(VKey.create(DomainBase.class, domainOfyKey.getName(), domainOfyKey));
           if (domain == null) {
             System.out.printf("ERROR: Token %s was redeemed but domain can't be loaded.\n", token);
           } else {
@@ -76,19 +81,23 @@ final class GetAllocationTokenCommand implements CommandWithRemoteApi {
     }
   }
 
-  private static ImmutableMap<Key<DomainBase>, DomainBase> loadRedeemedDomains(
+  @SuppressWarnings("unchecked")
+  private static ImmutableMap<VKey<DomainBase>, DomainBase> loadRedeemedDomains(
       Collection<AllocationToken> tokens) {
-    ImmutableList<Key<DomainBase>> domainKeys =
+    ImmutableList<VKey<DomainBase>> domainKeys =
         tokens.stream()
             .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<Key<DomainBase>, DomainBase> domainsBuilder = new ImmutableMap.Builder<>();
-    for (List<Key<DomainBase>> keys : Lists.partition(domainKeys, BATCH_SIZE)) {
-      domainsBuilder.putAll(ofy().load().keys(keys));
+    ImmutableMap.Builder<VKey<DomainBase>, DomainBase> domainsBuilder =
+        new ImmutableMap.Builder<>();
+    for (List<VKey<DomainBase>> keys : Lists.partition(domainKeys, BATCH_SIZE)) {
+      tm().loadByKeys(ImmutableList.copyOf(keys))
+          .forEach((k, v) -> domainsBuilder.put((VKey<DomainBase>) k, v));
     }
     return domainsBuilder.build();
   }

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

@@ -18,8 +18,8 @@ import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static com.google.common.collect.Iterables.partition;
 import static com.google.common.collect.Streams.stream;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 
 import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
@@ -105,13 +105,17 @@ final class UpdateAllocationTokensCommand extends UpdateOrDeleteAllocationTokens
     }
 
     tokensToSave =
-        ofy().load().keys(getTokenKeys()).values().stream()
-            .collect(toImmutableMap(Function.identity(), this::updateToken))
-            .entrySet()
-            .stream()
-            .filter(entry -> !entry.getKey().equals(entry.getValue())) // only update changed tokens
-            .map(Map.Entry::getValue)
-            .collect(toImmutableSet());
+        transactIfJpaTm(
+            () ->
+                tm().loadByKeys(getTokenKeys()).values().stream()
+                    .collect(toImmutableMap(Function.identity(), this::updateToken))
+                    .entrySet()
+                    .stream()
+                    .filter(
+                        entry ->
+                            !entry.getKey().equals(entry.getValue())) // only update changed tokens
+                    .map(Map.Entry::getValue)
+                    .collect(toImmutableSet()));
   }
 
   @Override
@@ -123,7 +127,7 @@ final class UpdateAllocationTokensCommand extends UpdateOrDeleteAllocationTokens
   protected String execute() {
     long numUpdated =
         stream(partition(tokensToSave, BATCH_SIZE))
-            .mapToLong(batch -> tm().transact(() -> saveBatch(batch)))
+            .mapToLong(batch -> tm().transact(() -> saveBatch(ImmutableList.copyOf(batch))))
             .sum();
     return String.format("Updated %d tokens in total.", numUpdated);
   }
@@ -141,9 +145,9 @@ final class UpdateAllocationTokensCommand extends UpdateOrDeleteAllocationTokens
     return builder.build();
   }
 
-  private long saveBatch(List<AllocationToken> batch) {
+  private long saveBatch(ImmutableList<AllocationToken> batch) {
     if (!dryRun) {
-      ofy().save().entities(batch);
+      tm().putAll(batch);
     }
     System.out.printf(
         "%s tokens: %s\n",

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

@@ -329,7 +329,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

@@ -15,13 +15,15 @@
 package google.registry.tools;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
-import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 
 import com.beust.jcommander.Parameter;
 import com.google.common.collect.ImmutableSet;
-import com.googlecode.objectify.Key;
 import google.registry.model.domain.token.AllocationToken;
+import google.registry.persistence.VKey;
 import java.util.List;
 
 /** Shared base class for commands to update or delete allocation tokens. */
@@ -47,19 +49,28 @@ abstract class UpdateOrDeleteAllocationTokensCommand extends ConfirmingCommand
       description = "Do not actually update or delete the tokens; defaults to false")
   protected boolean dryRun;
 
-  protected ImmutableSet<Key<AllocationToken>> getTokenKeys() {
+  protected ImmutableSet<VKey<AllocationToken>> getTokenKeys() {
     checkArgument(
         tokens == null ^ prefix == null,
         "Must provide one of --tokens or --prefix, not both / neither");
     if (tokens != null) {
-      return tokens.stream()
-          .map(token -> Key.create(AllocationToken.class, token))
-          .collect(toImmutableSet());
+      ImmutableSet<VKey<AllocationToken>> keys =
+          tokens.stream()
+              .map(token -> VKey.create(AllocationToken.class, token))
+              .collect(toImmutableSet());
+      ImmutableSet<VKey<AllocationToken>> nonexistentKeys =
+          transactIfJpaTm(
+              () -> keys.stream().filter(key -> !tm().exists(key)).collect(toImmutableSet()));
+      checkState(nonexistentKeys.isEmpty(), "Tokens with keys %s did not exist.", nonexistentKeys);
+      return keys;
     } else {
       checkArgument(!prefix.isEmpty(), "Provided prefix should not be blank");
-      return ofy().load().type(AllocationToken.class).keys().list().stream()
-          .filter(key -> key.getName().startsWith(prefix))
-          .collect(toImmutableSet());
+      return transactIfJpaTm(
+          () ->
+              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

@@ -61,6 +61,7 @@ final class ValidateLoginCredentialsCommand implements CommandWithRemoteApi {
       description = "Hash of the client certificate.")
   private String clientCertificateHash;
 
+  @Nullable
   @Parameter(
       names = {"-i", "--ip_address"},
       description = "Client ip address to pretend to use")
@@ -78,7 +79,8 @@ final class ValidateLoginCredentialsCommand implements CommandWithRemoteApi {
     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(clientIpAddress))
         .validate(registrar, password);
     checkState(
         registrar.isLive(), "Registrar %s has non-live state: %s", clientId, registrar.getState());

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

@@ -334,8 +334,9 @@ 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)) {
         certificateChecker.validateCertificate(certificateString);
       return true;
     }

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

@@ -55,6 +55,19 @@ class ResaveAllEppResourcesActionTest extends MapreduceTestCase<ResaveAllEppReso
         .isGreaterThan(creationTime);
   }
 
+  @Test
+  void test_fastMode_doesNotResaveEntityWithNoChanges() throws Exception {
+    ContactResource contact = persistActiveContact("test123");
+    DateTime creationTime = contact.getUpdateTimestamp().getTimestamp();
+    assertThat(ofy().load().entity(contact).now().getUpdateTimestamp().getTimestamp())
+        .isEqualTo(creationTime);
+    ofy().clearSessionCache();
+    action.isFast = true;
+    runMapreduce();
+    assertThat(ofy().load().entity(contact).now().getUpdateTimestamp().getTimestamp())
+        .isEqualTo(creationTime);
+  }
+
   @Test
   void test_mapreduceResolvesPendingTransfer() throws Exception {
     DateTime now = DateTime.now(UTC);

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

@@ -36,7 +36,8 @@ class EppLoginTlsTest extends EppTestCase {
 
   void setClientCertificateHash(String clientCertificateHash) {
     setTransportCredentials(
-        new TlsCredentials(true, clientCertificateHash, Optional.of("192.168.1.100:54321")));
+        new TlsCredentials(
+            true, Optional.ofNullable(clientCertificateHash), Optional.of("192.168.1.100:54321")));
   }
 
   @BeforeEach
@@ -107,7 +108,7 @@ class EppLoginTlsTest extends EppTestCase {
 
   @Test
   void testGfeDidntProvideClientCertificate_failsMissingCertificate2200() throws Exception {
-    setClientCertificateHash("");
+    setClientCertificateHash(null);
     assertThatLogin("NewRegistrar", "foo-BAR2")
         .hasResponse(
             "response_error.xml",

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

@@ -137,7 +137,8 @@ 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("127.0.0.1"));
     flowRunner.run(eppMetricBuilder);
     assertThat(Splitter.on("\n\t").split(findFirstLogMessageByPrefix(handler, "EPP Command\n\t")))
         .contains("TlsCredentials{clientCertificateHash=abc123def, clientAddress=/127.0.0.1}");

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

@@ -100,7 +100,11 @@ public abstract class FlowTestCase<F extends Flow> {
 
   @RegisterExtension
   final AppEngineExtension appEngine =
-      AppEngineExtension.builder().withDatastoreAndCloudSql().withTaskQueue().build();
+      AppEngineExtension.builder()
+          .withClock(clock)
+          .withDatastoreAndCloudSql()
+          .withTaskQueue()
+          .build();
 
   @BeforeEach
   public void beforeEachFlowTestCase() {
@@ -199,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()));
   }
@@ -288,7 +292,7 @@ public abstract class FlowTestCase<F extends Flow> {
           e);
     }
     // Clear the cache so that we don't see stale results in tests.
-    ofy().clearSessionCache();
+    tm().clearSessionCache();
     return output;
   }
 

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

@@ -19,6 +19,8 @@ import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.EppResourceUtils.loadByForeignKey;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.tmch.ClaimsListShardTest.createTestClaimsListShard;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 import static google.registry.testing.EppExceptionSubject.assertAboutEppExceptions;
 import static google.registry.testing.LogsSubject.assertAboutLogs;
 import static google.registry.testing.TaskQueueHelper.assertTasksEnqueued;
@@ -72,7 +74,7 @@ public abstract class ResourceFlowTestCase<F extends Flow, R extends EppResource
   protected R reloadResourceByForeignKey(DateTime now) throws Exception {
     // Force the session to be cleared so that when we read it back, we read from Datastore and not
     // from the transaction's session cache.
-    ofy().clearSessionCache();
+    tm().clearSessionCache();
     return loadByForeignKey(getResourceClass(), getUniqueIdFromCommand(), now).orElse(null);
   }
 
@@ -83,9 +85,10 @@ public abstract class ResourceFlowTestCase<F extends Flow, R extends EppResource
 
   protected <T extends EppResource> T reloadResourceAndCloneAtTime(T resource, DateTime now) {
     // Force the session to be cleared.
-    ofy().clearSessionCache();
+    tm().clearSessionCache();
     @SuppressWarnings("unchecked")
-    T refreshedResource = (T) ofy().load().entity(resource).now().cloneProjectedAtTime(now);
+    T refreshedResource =
+        (T) transactIfJpaTm(() -> tm().loadByEntity(resource)).cloneProjectedAtTime(now);
     return refreshedResource;
   }
 

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

@@ -14,7 +14,8 @@
 
 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;
@@ -22,9 +23,12 @@ 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 google.registry.flows.TlsCredentials.BadRegistrarIpAddressException;
+import google.registry.flows.TlsCredentials.MissingRegistrarCertificateException;
 import google.registry.model.registrar.Registrar;
-import google.registry.request.HttpException.BadRequestException;
 import google.registry.testing.AppEngineExtension;
+import google.registry.util.CidrAddressBlock;
 import java.util.Optional;
 import javax.servlet.http.HttpServletRequest;
 import org.joda.time.DateTime;
@@ -42,22 +46,40 @@ final class TlsCredentialsTest {
   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 testClientCertificateHash_missing() {
+    TlsCredentials tls = new TlsCredentials(true, Optional.empty(), Optional.of("192.168.1.1"));
+    persistResource(
+        loadRegistrar("TheRegistrar")
+            .asBuilder()
+            .setClientCertificate(SAMPLE_CERT, DateTime.now(UTC))
+            .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());
+    persistResource(
+        loadRegistrar("TheRegistrar")
+            .asBuilder()
+            .setClientCertificate(SAMPLE_CERT, DateTime.now(UTC))
+            .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("192.168.1.1"));
     persistResource(
         loadRegistrar("TheRegistrar")
             .asBuilder()

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

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

@@ -38,10 +38,12 @@ import google.registry.model.contact.PostalInfo.Type;
 import google.registry.model.eppcommon.AuthInfo.PasswordAuth;
 import google.registry.model.eppcommon.PresenceMarker;
 import google.registry.model.eppcommon.StatusValue;
+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 ContactInfoFlow}. */
+@DualDatabaseTest
 class ContactInfoFlowTest extends ResourceFlowTestCase<ContactInfoFlow, ContactResource> {
 
   ContactInfoFlowTest() {
@@ -94,7 +96,7 @@ class ContactInfoFlowTest extends ResourceFlowTestCase<ContactInfoFlow, ContactR
     return contact;
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess() throws Exception {
     persistContactResource(true);
     // Check that the persisted contact info was returned.
@@ -107,7 +109,7 @@ class ContactInfoFlowTest extends ResourceFlowTestCase<ContactInfoFlow, ContactR
     assertNoBillingEvents();
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_linked() throws Exception {
     createTld("foobar");
     persistResource(newDomainBase("example.foobar", persistContactResource(true)));
@@ -121,7 +123,7 @@ class ContactInfoFlowTest extends ResourceFlowTestCase<ContactInfoFlow, ContactR
     assertNoBillingEvents();
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_owningRegistrarWithoutAuthInfo_seesAuthInfo() throws Exception {
     setEppInput("contact_info_no_authinfo.xml");
     persistContactResource(true);
@@ -135,7 +137,7 @@ class ContactInfoFlowTest extends ResourceFlowTestCase<ContactInfoFlow, ContactR
     assertNoBillingEvents();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_otherRegistrar_notAuthorized() throws Exception {
     setClientIdForFlow("NewRegistrar");
     persistContactResource(true);
@@ -145,7 +147,7 @@ class ContactInfoFlowTest extends ResourceFlowTestCase<ContactInfoFlow, ContactR
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_otherRegistrarWithoutAuthInfoAsSuperuser_doesNotSeeAuthInfo() throws Exception {
     setClientIdForFlow("NewRegistrar");
     setEppInput("contact_info_no_authinfo.xml");
@@ -162,7 +164,7 @@ class ContactInfoFlowTest extends ResourceFlowTestCase<ContactInfoFlow, ContactR
     assertNoBillingEvents();
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_otherRegistrarWithAuthInfoAsSuperuser_seesAuthInfo() throws Exception {
     setClientIdForFlow("NewRegistrar");
     persistContactResource(true);
@@ -178,7 +180,7 @@ class ContactInfoFlowTest extends ResourceFlowTestCase<ContactInfoFlow, ContactR
     assertNoBillingEvents();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_neverExisted() throws Exception {
     ResourceDoesNotExistException thrown =
         assertThrows(ResourceDoesNotExistException.class, this::runFlow);
@@ -186,7 +188,7 @@ class ContactInfoFlowTest extends ResourceFlowTestCase<ContactInfoFlow, ContactR
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_existedButWasDeleted() throws Exception {
     persistContactResource(false);
     ResourceDoesNotExistException thrown =
@@ -195,7 +197,7 @@ class ContactInfoFlowTest extends ResourceFlowTestCase<ContactInfoFlow, ContactR
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testIcannActivityReportField_getsLogged() throws Exception {
     persistContactResource(true);
     runFlow();

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

@@ -41,9 +41,11 @@ import google.registry.model.contact.ContactResource;
 import google.registry.model.contact.PostalInfo;
 import google.registry.model.contact.PostalInfo.Type;
 import google.registry.model.eppcommon.StatusValue;
-import org.junit.jupiter.api.Test;
+import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.TestOfyAndSql;
 
 /** Unit tests for {@link ContactUpdateFlow}. */
+@DualDatabaseTest
 class ContactUpdateFlowTest extends ResourceFlowTestCase<ContactUpdateFlow, ContactResource> {
 
   ContactUpdateFlowTest() {
@@ -63,18 +65,18 @@ class ContactUpdateFlowTest extends ResourceFlowTestCase<ContactUpdateFlow, Cont
     assertNoBillingEvents();
   }
 
-  @Test
+  @TestOfyAndSql
   void testDryRun() throws Exception {
     persistActiveContact(getUniqueIdFromCommand());
     dryRunFlowAssertResponse(loadFile("generic_success_response.xml"));
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess() throws Exception {
     doSuccessfulTest();
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_updatingInternationalizedPostalInfoDeletesLocalized() throws Exception {
     ContactResource contact =
         persistResource(
@@ -112,7 +114,7 @@ class ContactUpdateFlowTest extends ResourceFlowTestCase<ContactUpdateFlow, Cont
             .build());
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_updatingLocalizedPostalInfoDeletesInternationalized() throws Exception {
     setEppInput("contact_update_localized.xml");
     ContactResource contact =
@@ -151,7 +153,7 @@ class ContactUpdateFlowTest extends ResourceFlowTestCase<ContactUpdateFlow, Cont
             .build());
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_partialPostalInfoUpdate() throws Exception {
     setEppInput("contact_update_partial_postalinfo.xml");
     persistResource(
@@ -187,7 +189,7 @@ class ContactUpdateFlowTest extends ResourceFlowTestCase<ContactUpdateFlow, Cont
             .build());
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_updateOnePostalInfo_touchOtherPostalInfoPreservesIt() throws Exception {
     setEppInput("contact_update_partial_postalinfo_preserve_int.xml");
     persistResource(
@@ -253,7 +255,7 @@ class ContactUpdateFlowTest extends ResourceFlowTestCase<ContactUpdateFlow, Cont
                 .build());
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_neverExisted() throws Exception {
     ResourceDoesNotExistException thrown =
         assertThrows(ResourceDoesNotExistException.class, this::runFlow);
@@ -261,7 +263,7 @@ class ContactUpdateFlowTest extends ResourceFlowTestCase<ContactUpdateFlow, Cont
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_existedButWasDeleted() throws Exception {
     persistDeletedContact(getUniqueIdFromCommand(), clock.nowUtc().minusDays(1));
     ResourceDoesNotExistException thrown =
@@ -270,7 +272,7 @@ class ContactUpdateFlowTest extends ResourceFlowTestCase<ContactUpdateFlow, Cont
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_statusValueNotClientSettable() throws Exception {
     setEppInput("contact_update_prohibited_status.xml");
     persistActiveContact(getUniqueIdFromCommand());
@@ -278,7 +280,7 @@ class ContactUpdateFlowTest extends ResourceFlowTestCase<ContactUpdateFlow, Cont
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_superuserStatusValueNotClientSettable() throws Exception {
     setEppInput("contact_update_prohibited_status.xml");
     persistActiveContact(getUniqueIdFromCommand());
@@ -287,7 +289,7 @@ class ContactUpdateFlowTest extends ResourceFlowTestCase<ContactUpdateFlow, Cont
         CommitMode.LIVE, UserPrivileges.SUPERUSER, loadFile("generic_success_response.xml"));
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_unauthorizedClient() throws Exception {
     sessionMetadata.setClientId("NewRegistrar");
     persistActiveContact(getUniqueIdFromCommand());
@@ -295,7 +297,7 @@ class ContactUpdateFlowTest extends ResourceFlowTestCase<ContactUpdateFlow, Cont
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_superuserUnauthorizedClient() throws Exception {
     sessionMetadata.setClientId("NewRegistrar");
     persistActiveContact(getUniqueIdFromCommand());
@@ -304,7 +306,7 @@ class ContactUpdateFlowTest extends ResourceFlowTestCase<ContactUpdateFlow, Cont
         CommitMode.LIVE, UserPrivileges.SUPERUSER, loadFile("generic_success_response.xml"));
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_clientUpdateProhibited_removed() throws Exception {
     setEppInput("contact_update_remove_client_update_prohibited.xml");
     persistResource(
@@ -318,7 +320,7 @@ class ContactUpdateFlowTest extends ResourceFlowTestCase<ContactUpdateFlow, Cont
         .doesNotHaveStatusValue(StatusValue.CLIENT_UPDATE_PROHIBITED);
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_superuserClientUpdateProhibited_notRemoved() throws Exception {
     setEppInput("contact_update_prohibited_status.xml");
     persistResource(
@@ -336,7 +338,7 @@ class ContactUpdateFlowTest extends ResourceFlowTestCase<ContactUpdateFlow, Cont
         .hasStatusValue(StatusValue.SERVER_DELETE_PROHIBITED);
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_clientUpdateProhibited_notRemoved() throws Exception {
     persistResource(
         newContactResource(getUniqueIdFromCommand())
@@ -348,7 +350,7 @@ class ContactUpdateFlowTest extends ResourceFlowTestCase<ContactUpdateFlow, Cont
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_serverUpdateProhibited() throws Exception {
     persistResource(
         newContactResource(getUniqueIdFromCommand())
@@ -361,7 +363,7 @@ class ContactUpdateFlowTest extends ResourceFlowTestCase<ContactUpdateFlow, Cont
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_pendingDeleteProhibited() throws Exception {
     persistResource(
         newContactResource(getUniqueIdFromCommand())
@@ -374,13 +376,13 @@ class ContactUpdateFlowTest extends ResourceFlowTestCase<ContactUpdateFlow, Cont
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_nonAsciiInLocAddress() throws Exception {
     setEppInput("contact_update_hebrew_loc.xml");
     doSuccessfulTest();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_nonAsciiInIntAddress() throws Exception {
     setEppInput("contact_update_hebrew_int.xml");
     persistActiveContact(getUniqueIdFromCommand());
@@ -389,7 +391,7 @@ class ContactUpdateFlowTest extends ResourceFlowTestCase<ContactUpdateFlow, Cont
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_declineDisclosure() throws Exception {
     setEppInput("contact_update_decline_disclosure.xml");
     persistActiveContact(getUniqueIdFromCommand());
@@ -398,7 +400,7 @@ class ContactUpdateFlowTest extends ResourceFlowTestCase<ContactUpdateFlow, Cont
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_addRemoveSameValue() throws Exception {
     setEppInput("contact_update_add_remove_same.xml");
     persistActiveContact(getUniqueIdFromCommand());
@@ -406,7 +408,7 @@ class ContactUpdateFlowTest extends ResourceFlowTestCase<ContactUpdateFlow, Cont
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testIcannActivityReportField_getsLogged() throws Exception {
     persistActiveContact(getUniqueIdFromCommand());
     clock.advanceOneMilli();

core/src/test/java/google/registry/flows/domain/DomainCreateFlowTest.java

@@ -184,7 +184,7 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
 
   @Order(value = Order.DEFAULT - 2)
   @RegisterExtension
-  final ReplayExtension replayExtension = new ReplayExtension(clock);
+  final ReplayExtension replayExtension = ReplayExtension.createWithCompare(clock);
 
   DomainCreateFlowTest() {
     setEppInput("domain_create.xml", ImmutableMap.of("DOMAIN", "example.tld"));
@@ -265,7 +265,8 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
     HistoryEntry historyEntry = getHistoryEntries(domain).get(0);
     assertAboutDomains()
         .that(domain)
-        .hasRegistrationExpirationTime(tm().load(domain.getAutorenewBillingEvent()).getEventTime())
+        .hasRegistrationExpirationTime(
+            tm().loadByKey(domain.getAutorenewBillingEvent()).getEventTime())
         .and()
         .hasOnlyOneHistoryEntryWhich()
         .hasType(HistoryEntry.Type.DOMAIN_CREATE)

core/src/test/java/google/registry/flows/domain/DomainDeleteFlowTest.java

@@ -112,7 +112,7 @@ class DomainDeleteFlowTest extends ResourceFlowTestCase<DomainDeleteFlow, Domain
 
   @Order(value = Order.DEFAULT - 2)
   @RegisterExtension
-  final ReplayExtension replayExtension = new ReplayExtension(clock);
+  final ReplayExtension replayExtension = ReplayExtension.createWithoutCompare(clock);
 
   private DomainBase domain;
   private HistoryEntry earlierHistoryEntry;
@@ -475,7 +475,7 @@ class DomainDeleteFlowTest extends ResourceFlowTestCase<DomainDeleteFlow, Domain
     assertThat(getPollMessages("TheRegistrar", deletionTime)).hasSize(1);
     assertThat(domain.getDeletePollMessage().getOfyKey())
         .isEqualTo(getOnlyPollMessage("TheRegistrar").createVKey().getOfyKey());
-    PollMessage.OneTime deletePollMessage = tm().load(domain.getDeletePollMessage());
+    PollMessage.OneTime deletePollMessage = tm().loadByKey(domain.getDeletePollMessage());
     assertThat(deletePollMessage.getMsg()).isEqualTo(expectedMessage);
   }
 
@@ -509,7 +509,7 @@ class DomainDeleteFlowTest extends ResourceFlowTestCase<DomainDeleteFlow, Domain
     // Modify the autorenew poll message so that it has unacked messages in the past. This should
     // prevent it from being deleted when the domain is deleted.
     persistResource(
-        tm().load(reloadResourceByForeignKey().getAutorenewPollMessage())
+        tm().loadByKey(reloadResourceByForeignKey().getAutorenewPollMessage())
             .asBuilder()
             .setEventTime(A_MONTH_FROM_NOW.minusYears(3))
             .build());

core/src/test/java/google/registry/flows/domain/DomainRenewFlowTest.java

@@ -108,7 +108,7 @@ class DomainRenewFlowTest extends ResourceFlowTestCase<DomainRenewFlow, DomainBa
 
   @Order(value = Order.DEFAULT - 2)
   @RegisterExtension
-  final ReplayExtension replayExtension = new ReplayExtension(clock);
+  final ReplayExtension replayExtension = ReplayExtension.createWithoutCompare(clock);
 
   @BeforeEach
   void initDomainTest() {
@@ -197,7 +197,7 @@ class DomainRenewFlowTest extends ResourceFlowTestCase<DomainRenewFlow, DomainBa
     DomainBase domain = reloadResourceByForeignKey();
     HistoryEntry historyEntryDomainRenew =
         getOnlyHistoryEntryOfType(domain, HistoryEntry.Type.DOMAIN_RENEW);
-    assertThat(tm().load(domain.getAutorenewBillingEvent()).getEventTime())
+    assertThat(tm().loadByKey(domain.getAutorenewBillingEvent()).getEventTime())
         .isEqualTo(newExpiration);
     assertAboutDomains()
         .that(domain)
@@ -494,7 +494,7 @@ class DomainRenewFlowTest extends ResourceFlowTestCase<DomainRenewFlow, DomainBa
     persistDomain();
     // Modify the autorenew poll message so that it has an undelivered message in the past.
     persistResource(
-        tm().load(reloadResourceByForeignKey().getAutorenewPollMessage())
+        tm().loadByKey(reloadResourceByForeignKey().getAutorenewPollMessage())
             .asBuilder()
             .setEventTime(expirationTime.minusYears(1))
             .build());

core/src/test/java/google/registry/flows/domain/DomainRestoreRequestFlowTest.java

@@ -150,7 +150,7 @@ class DomainRestoreRequestFlowTest
     DomainBase domain = reloadResourceByForeignKey();
     HistoryEntry historyEntryDomainRestore =
         getOnlyHistoryEntryOfType(domain, HistoryEntry.Type.DOMAIN_RESTORE);
-    assertThat(tm().load(domain.getAutorenewBillingEvent()).getEventTime())
+    assertThat(tm().loadByKey(domain.getAutorenewBillingEvent()).getEventTime())
         .isEqualTo(expirationTime);
     assertAboutDomains()
         .that(domain)
@@ -218,7 +218,7 @@ class DomainRestoreRequestFlowTest
     DomainBase domain = reloadResourceByForeignKey();
     HistoryEntry historyEntryDomainRestore =
         getOnlyHistoryEntryOfType(domain, HistoryEntry.Type.DOMAIN_RESTORE);
-    assertThat(tm().load(domain.getAutorenewBillingEvent()).getEventTime())
+    assertThat(tm().loadByKey(domain.getAutorenewBillingEvent()).getEventTime())
         .isEqualTo(newExpirationTime);
     assertAboutDomains()
         .that(domain)

core/src/test/java/google/registry/flows/domain/DomainTransferApproveFlowTest.java

@@ -191,7 +191,7 @@ class DomainTransferApproveFlowTest
     assertAboutHistoryEntries().that(historyEntryTransferApproved).hasOtherClientId("NewRegistrar");
     assertTransferApproved(domain, originalTransferData);
     assertAboutDomains().that(domain).hasRegistrationExpirationTime(expectedExpirationTime);
-    assertThat(tm().load(domain.getAutorenewBillingEvent()).getEventTime())
+    assertThat(tm().loadByKey(domain.getAutorenewBillingEvent()).getEventTime())
         .isEqualTo(expectedExpirationTime);
     // The poll message (in the future) to the losing registrar for implicit ack should be gone.
     assertThat(getPollMessages(domain, "TheRegistrar", clock.nowUtc().plusMonths(1))).isEmpty();

core/src/test/java/google/registry/flows/domain/DomainTransferRequestFlowTest.java

@@ -291,13 +291,13 @@ class DomainTransferRequestFlowTest
     // Assert that the domain's TransferData server-approve billing events match the above.
     if (expectTransferBillingEvent) {
       assertBillingEventsEqual(
-          tm().load(domain.getTransferData().getServerApproveBillingEvent()),
+          tm().loadByKey(domain.getTransferData().getServerApproveBillingEvent()),
           optionalTransferBillingEvent.get());
     } else {
       assertThat(domain.getTransferData().getServerApproveBillingEvent()).isNull();
     }
     assertBillingEventsEqual(
-        tm().load(domain.getTransferData().getServerApproveAutorenewEvent()),
+        tm().loadByKey(domain.getTransferData().getServerApproveAutorenewEvent()),
         gainingClientAutorenew);
     // Assert that the full set of server-approve billing events is exactly the extra ones plus
     // the transfer billing event (if present) and the gaining client autorenew.
@@ -318,7 +318,7 @@ class DomainTransferRequestFlowTest
             BillingEvent.class),
         Sets.union(expectedServeApproveBillingEvents, extraBillingEvents));
     // The domain's autorenew billing event should still point to the losing client's event.
-    BillingEvent.Recurring domainAutorenewEvent = tm().load(domain.getAutorenewBillingEvent());
+    BillingEvent.Recurring domainAutorenewEvent = tm().loadByKey(domain.getAutorenewBillingEvent());
     assertThat(domainAutorenewEvent.getClientId()).isEqualTo("TheRegistrar");
     assertThat(domainAutorenewEvent.getRecurrenceEndTime()).isEqualTo(implicitTransferTime);
     // The original grace periods should remain untouched.
@@ -414,7 +414,7 @@ class DomainTransferRequestFlowTest
 
     // Assert that the poll messages show up in the TransferData server approve entities.
     assertPollMessagesEqual(
-        tm().load(domain.getTransferData().getServerApproveAutorenewPollMessage()),
+        tm().loadByKey(domain.getTransferData().getServerApproveAutorenewPollMessage()),
         autorenewPollMessage);
     // Assert that the full set of server-approve poll messages is exactly the server approve
     // OneTime messages to gaining and losing registrars plus the gaining client autorenew.
@@ -446,7 +446,8 @@ class DomainTransferRequestFlowTest
         .hasLastEppUpdateTime(implicitTransferTime)
         .and()
         .hasLastEppUpdateClientId("NewRegistrar");
-    assertThat(tm().load(domainAfterAutomaticTransfer.getAutorenewBillingEvent()).getEventTime())
+    assertThat(
+            tm().loadByKey(domainAfterAutomaticTransfer.getAutorenewBillingEvent()).getEventTime())
         .isEqualTo(expectedExpirationTime);
     // And after the expected grace time, the grace period should be gone.
     DomainBase afterGracePeriod =