Fix typo in header name in Client Certificate Provider (#946)

* Fix typo in header name

* fix test

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/backup/ReplayCommitLogsToSqlAction.java

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

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

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

core/src/main/java/google/registry/batch/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;
@@ -27,7 +25,10 @@ import com.google.common.net.InetAddresses;
 import dagger.Module;
 import dagger.Provides;
 import google.registry.config.RegistryConfig.Config;
+import google.registry.config.RegistryEnvironment;
 import google.registry.flows.EppException.AuthenticationErrorException;
+import google.registry.flows.certs.CertificateChecker;
+import google.registry.flows.certs.CertificateChecker.InsecureCertificateException;
 import google.registry.model.registrar.Registrar;
 import google.registry.request.Header;
 import google.registry.util.CidrAddressBlock;
@@ -44,6 +45,9 @@ import javax.servlet.http.HttpServletRequest;
  * <dl>
  *   <dt>X-SSL-Certificate
  *   <dd>This field should contain a base64 encoded digest of the client's TLS certificate. It is
+ *       used only if the validation of the full certificate fails.
+ *   <dt>X-SSL-Full-Certificate
+ *   <dd>This field should contain a base64 encoding of the client's TLS certificate. It is
  *       validated during an EPP login command against a known good value that is transmitted out of
  *       band.
  *   <dt>X-Forwarded-For
@@ -56,17 +60,23 @@ public class TlsCredentials implements TransportCredentials {
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
 
   private final boolean requireSslCertificates;
-  private final String clientCertificateHash;
-  private final InetAddress clientInetAddr;
+  private final Optional<String> clientCertificateHash;
+  private final Optional<String> clientCertificate;
+  private final Optional<InetAddress> clientInetAddr;
+  private final CertificateChecker certificateChecker;
 
   @Inject
   public TlsCredentials(
       @Config("requireSslCertificates") boolean requireSslCertificates,
-      @Header("X-SSL-Certificate") String clientCertificateHash,
-      @Header("X-Forwarded-For") Optional<String> clientAddress) {
+      @Header("X-SSL-Certificate") Optional<String> clientCertificateHash,
+      @Header("X-SSL-Full-Certificate") Optional<String> clientCertificate,
+      @Header("X-Forwarded-For") Optional<String> clientAddress,
+      CertificateChecker certificateChecker) {
     this.requireSslCertificates = requireSslCertificates;
     this.clientCertificateHash = clientCertificateHash;
-    this.clientInetAddr = clientAddress.isPresent() ? parseInetAddress(clientAddress.get()) : null;
+    this.clientCertificate = clientCertificate;
+    this.clientInetAddr = clientAddress.map(TlsCredentials::parseInetAddress);
+    this.certificateChecker = certificateChecker;
   }
 
   static InetAddress parseInetAddress(String asciiAddr) {
@@ -93,14 +103,18 @@ public class TlsCredentials implements TransportCredentials {
     ImmutableList<CidrAddressBlock> ipAddressAllowList = registrar.getIpAddressAllowList();
     if (ipAddressAllowList.isEmpty()) {
       logger.atInfo().log(
-          "Skipping IP allow list check because %s doesn't have an IP allow list",
+          "Skipping IP allow list check because %s doesn't have an IP allow list.",
           registrar.getClientId());
       return;
     }
-    for (CidrAddressBlock cidrAddressBlock : ipAddressAllowList) {
-      if (cidrAddressBlock.contains(clientInetAddr)) {
-        // IP address is in allow list; return early.
-        return;
+    // In the rare unexpected case that the client inet address wasn't passed along at all, then
+    // by default deny access.
+    if (clientInetAddr.isPresent()) {
+      for (CidrAddressBlock cidrAddressBlock : ipAddressAllowList) {
+        if (cidrAddressBlock.contains(clientInetAddr.get())) {
+          // IP address is in allow list; return early.
+          return;
+        }
       }
     }
     logger.atInfo().log(
@@ -113,13 +127,65 @@ public class TlsCredentials implements TransportCredentials {
   /**
    * Verifies client SSL certificate is permitted to issue commands as {@code registrar}.
    *
-   * @throws MissingRegistrarCertificateException if frontend didn't send certificate hash header
+   * @throws MissingRegistrarCertificateException if frontend didn't send certificate header
    * @throws BadRegistrarCertificateException if registrar requires certificate and it didn't match
    */
   @VisibleForTesting
   void validateCertificate(Registrar registrar) throws AuthenticationErrorException {
-    if (isNullOrEmpty(registrar.getClientCertificateHash())
-        && isNullOrEmpty(registrar.getFailoverClientCertificateHash())) {
+    // Check that certificate is present in registrar object
+    if (!registrar.getClientCertificate().isPresent()
+        && !registrar.getFailoverClientCertificate().isPresent()) {
+      // Log an error and validate using certificate hash instead
+      // TODO(sarahbot): throw a RegistrarCertificateNotConfiguredException once hash is no longer
+      // used as failover
+      logger.atWarning().log(
+          "There is no certificate configured for registrar %s.", registrar.getClientId());
+    } else if (!clientCertificate.isPresent()) {
+      // Check that the request included the full certificate
+      // Log an error and validate using certificate hash instead
+      // TODO(sarahbot): throw a MissingRegistrarCertificateException once hash is no longer used as
+      // failover
+      logger.atWarning().log(
+          "Request from registrar %s did not include X-SSL-Full-Certificate.",
+          registrar.getClientId());
+    } else {
+      // Check if the certificate is equal to the one on file for the registrar.
+      if (clientCertificate.equals(registrar.getClientCertificate())
+          || clientCertificate.equals(registrar.getFailoverClientCertificate())) {
+        // Check certificate for any requirement violations
+        // TODO(Sarahbot@): Throw exceptions instead of just logging once requirement enforcement
+        // begins
+        try {
+          certificateChecker.validateCertificate(clientCertificate.get());
+        } catch (InsecureCertificateException e) {
+          // throw exception in unit tests and Sandbox
+          if (RegistryEnvironment.get().equals(RegistryEnvironment.UNITTEST)
+              || RegistryEnvironment.get().equals(RegistryEnvironment.SANDBOX)) {
+            throw new CertificateContainsSecurityViolationsException(e);
+          }
+          logger.atWarning().log(
+              "Registrar certificate used for %s does not meet certificate requirements: %s",
+              registrar.getClientId(), e.getMessage());
+        } catch (Exception e) {
+          logger.atWarning().log(
+              "Error validating certificate for %s: %s", registrar.getClientId(), e);
+        }
+        // successfully validated, return here since hash validation is not necessary
+        return;
+      }
+      // Log an error and validate using certificate hash instead
+      // TODO(sarahbot): throw a BadRegistrarCertificateException once hash is no longer used as
+      // failover
+      logger.atWarning().log("Non-matching certificate for registrar %s.", registrar.getClientId());
+    }
+    validateCertificateHash(registrar);
+  }
+
+  private void validateCertificateHash(Registrar registrar) throws AuthenticationErrorException {
+    // Check the certificate hash as a failover
+    // TODO(sarahbot): Remove hash checks once certificate checks are working.
+    if (!registrar.getClientCertificateHash().isPresent()
+        && !registrar.getFailoverClientCertificateHash().isPresent()) {
       if (requireSslCertificates) {
         throw new RegistrarCertificateNotConfiguredException();
       } else {
@@ -128,14 +194,17 @@ public class TlsCredentials implements TransportCredentials {
         return;
       }
     }
-    if (isNullOrEmpty(clientCertificateHash)) {
-      logger.atInfo().log("Request did not include X-SSL-Certificate");
+    // Check that the request included the certificate hash
+    if (!clientCertificateHash.isPresent()) {
+      logger.atInfo().log(
+          "Request from registrar %s did not include X-SSL-Certificate.", registrar.getClientId());
       throw new MissingRegistrarCertificateException();
     }
+    // Check if the certificate hash is equal to the one on file for the registrar.
     if (!clientCertificateHash.equals(registrar.getClientCertificateHash())
         && !clientCertificateHash.equals(registrar.getFailoverClientCertificateHash())) {
       logger.atWarning().log(
-          "bad certificate hash (%s) for %s, wanted either %s or %s",
+          "Non-matching certificate hash (%s) for %s, wanted either %s or %s.",
           clientCertificateHash,
           registrar.getClientId(),
           registrar.getClientCertificateHash(),
@@ -154,21 +223,36 @@ public class TlsCredentials implements TransportCredentials {
   @Override
   public String toString() {
     return toStringHelper(getClass())
-        .add("clientCertificateHash", clientCertificateHash)
-        .add("clientAddress", clientInetAddr)
+        .add("clientCertificate", clientCertificate.orElse(null))
+        .add("clientCertificateHash", clientCertificateHash.orElse(null))
+        .add("clientAddress", clientInetAddr.orElse(null))
         .toString();
   }
 
   /** Registrar certificate does not match stored certificate. */
   public static class BadRegistrarCertificateException extends AuthenticationErrorException {
-    public BadRegistrarCertificateException() {
+    BadRegistrarCertificateException() {
       super("Registrar certificate does not match stored certificate");
     }
   }
 
+  /** Registrar certificate contains the following security violations: ... */
+  public static class CertificateContainsSecurityViolationsException
+      extends AuthenticationErrorException {
+    InsecureCertificateException exception;
+
+    CertificateContainsSecurityViolationsException(InsecureCertificateException exception) {
+      super(
+          String.format(
+              "Registrar certificate contains the following security violations:\n%s",
+              exception.getMessage()));
+      this.exception = exception;
+    }
+  }
+
   /** Registrar certificate not present. */
   public static class MissingRegistrarCertificateException extends AuthenticationErrorException {
-    public MissingRegistrarCertificateException() {
+    MissingRegistrarCertificateException() {
       super("Registrar certificate not present");
     }
   }
@@ -176,14 +260,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 +275,21 @@ public class TlsCredentials implements TransportCredentials {
   /** Dagger module for the EPP TLS endpoint. */
   @Module
   public static final class EppTlsModule {
+
     @Provides
     @Header("X-SSL-Certificate")
-    static String provideClientCertificateHash(HttpServletRequest req) {
-      return extractRequiredHeader(req, "X-SSL-Certificate");
+    static Optional<String> provideClientCertificateHash(HttpServletRequest req) {
+      // Note: This header is actually required, we just want to handle its absence explicitly
+      // by throwing an EPP exception rather than a generic Bad Request exception.
+      return extractOptionalHeader(req, "X-SSL-Certificate");
+    }
+
+    @Provides
+    @Header("X-SSL-Full-Certificate")
+    static Optional<String> provideClientCertificate(HttpServletRequest req) {
+      // Note: This header is actually required, we just want to handle its absence explicitly
+      // by throwing an EPP exception rather than a generic Bad Request exception.
+      return extractOptionalHeader(req, "X-SSL-Full-Certificate");
     }
 
     @Provides

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

@@ -110,7 +110,6 @@ import google.registry.model.reporting.DomainTransactionRecord;
 import google.registry.model.reporting.DomainTransactionRecord.TransactionReportField;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.reporting.IcannReportingTypes.ActivityReportField;
-import google.registry.persistence.DomainHistoryVKey;
 import google.registry.persistence.VKey;
 import google.registry.tmch.LordnTaskUtils;
 import java.util.Optional;
@@ -372,7 +371,7 @@ public class DomainCreateFlow implements TransactionalFlow {
         && TokenType.SINGLE_USE.equals(allocationToken.get().getTokenType())) {
       entitiesToSave.add(
           allocationTokenFlowUtils.redeemToken(
-              allocationToken.get(), DomainHistoryVKey.create(Key.create(historyEntry))));
+              allocationToken.get(), HistoryEntry.createVKey(Key.create(historyEntry))));
     }
     enqueueTasks(newDomain, hasSignedMarks, hasClaimsNotice);
 

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;
@@ -32,7 +31,8 @@ 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.model.registry.Registry;
-import google.registry.persistence.DomainHistoryVKey;
+import google.registry.model.reporting.HistoryEntry;
+import google.registry.persistence.VKey;
 import java.util.List;
 import java.util.Optional;
 import javax.inject.Inject;
@@ -107,7 +107,7 @@ public class AllocationTokenFlowUtils {
 
   /** Redeems a SINGLE_USE {@link AllocationToken}, returning the redeemed copy. */
   public AllocationToken redeemToken(
-      AllocationToken token, DomainHistoryVKey redemptionHistoryEntry) {
+      AllocationToken token, VKey<? extends HistoryEntry> redemptionHistoryEntry) {
     checkArgument(
         TokenType.SINGLE_USE.equals(token.getTokenType()),
         "Only SINGLE_USE tokens can be marked as redeemed");
@@ -152,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.isPresent()) {
       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

@@ -23,6 +23,8 @@ import google.registry.model.billing.BillingEvent;
 import google.registry.model.billing.BillingEvent.Recurring;
 import google.registry.model.domain.rgp.GracePeriodStatus;
 import google.registry.model.ofy.ObjectifyService;
+import google.registry.persistence.BillingVKey.BillingEventVKey;
+import google.registry.persistence.BillingVKey.BillingRecurrenceVKey;
 import google.registry.persistence.VKey;
 import google.registry.schema.replay.DatastoreAndSqlEntity;
 import javax.annotation.Nullable;
@@ -82,10 +84,8 @@ public class GracePeriod extends GracePeriodBase implements DatastoreAndSqlEntit
     instance.domainRepoId = checkArgumentNotNull(domainRepoId);
     instance.expirationTime = checkArgumentNotNull(expirationTime);
     instance.clientId = checkArgumentNotNull(clientId);
-    instance.billingEventOneTime = billingEventOneTime;
-    instance.billingEventOneTimeHistoryId = DomainBase.getHistoryId(billingEventOneTime);
-    instance.billingEventRecurring = billingEventRecurring;
-    instance.billingEventRecurringHistoryId = DomainBase.getHistoryId(billingEventRecurring);
+    instance.billingEventOneTime = BillingEventVKey.create(billingEventOneTime);
+    instance.billingEventRecurring = BillingRecurrenceVKey.create(billingEventRecurring);
     return instance;
   }
 
@@ -178,7 +178,6 @@ public class GracePeriod extends GracePeriodBase implements DatastoreAndSqlEntit
   public GracePeriod cloneAfterOfyLoad(String domainRepoId) {
     GracePeriod clone = clone(this);
     clone.domainRepoId = checkArgumentNotNull(domainRepoId);
-    clone.restoreHistoryIds();
     return clone;
   }
 
@@ -190,20 +189,7 @@ public class GracePeriod extends GracePeriodBase implements DatastoreAndSqlEntit
    */
   public GracePeriod cloneWithRecurringBillingEvent(VKey<BillingEvent.Recurring> recurring) {
     GracePeriod clone = clone(this);
-    clone.billingEventRecurring = recurring;
-    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();
+    clone.billingEventRecurring = BillingRecurrenceVKey.create(recurring);
     return clone;
   }
 
@@ -232,9 +218,7 @@ public class GracePeriod extends GracePeriodBase implements DatastoreAndSqlEntit
       instance.expirationTime = gracePeriod.expirationTime;
       instance.clientId = gracePeriod.clientId;
       instance.billingEventOneTime = gracePeriod.billingEventOneTime;
-      instance.billingEventOneTimeHistoryId = gracePeriod.billingEventOneTimeHistoryId;
       instance.billingEventRecurring = gracePeriod.billingEventRecurring;
-      instance.billingEventRecurringHistoryId = gracePeriod.billingEventRecurringHistoryId;
       return instance;
     }
   }

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

@@ -14,18 +14,14 @@
 
 package google.registry.model.domain;
 
-import com.googlecode.objectify.Key;
 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.billing.BillingEvent.OneTime;
 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;
@@ -68,24 +64,16 @@ public class GracePeriodBase extends ImmutableObject {
    * billingEventRecurring}) or for redemption grace periods (since deletes have no cost).
    */
   // NB: Would @IgnoreSave(IfNull.class), but not allowed for @Embed collections.
-  @Column(name = "billing_event_id")
-  VKey<OneTime> billingEventOneTime = null;
-
-  @Ignore
-  @Column(name = "billing_event_history_id")
-  Long billingEventOneTimeHistoryId;
+  @Access(AccessType.FIELD)
+  BillingEventVKey billingEventOneTime = null;
 
   /**
    * The recurring billing event corresponding to the action that triggered this grace period, if
    * applicable - i.e. if the action was an autorenew - or null in all other cases.
    */
   // NB: Would @IgnoreSave(IfNull.class), but not allowed for @Embed collections.
-  @Column(name = "billing_recurrence_id")
-  VKey<BillingEvent.Recurring> billingEventRecurring = null;
-
-  @Ignore
-  @Column(name = "billing_recurrence_history_id")
-  Long billingEventRecurringHistoryId;
+  @Access(AccessType.FIELD)
+  BillingRecurrenceVKey billingEventRecurring = null;
 
   public long getGracePeriodId() {
     return gracePeriodId;
@@ -123,8 +111,7 @@ public class GracePeriodBase extends ImmutableObject {
    * period is not AUTO_RENEW.
    */
   public VKey<BillingEvent.OneTime> getOneTimeBillingEvent() {
-    restoreOfyKeys();
-    return billingEventOneTime;
+    return billingEventOneTime == null ? null : billingEventOneTime.createVKey();
   }
 
   /**
@@ -132,63 +119,6 @@ public class GracePeriodBase extends ImmutableObject {
    * period is AUTO_RENEW.
    */
   public VKey<BillingEvent.Recurring> getRecurringBillingEvent() {
-    restoreOfyKeys();
-    return billingEventRecurring;
-  }
-
-  /**
-   * Restores history ids for composite VKeys after a load from datastore.
-   *
-   * <p>For use by DomainContent.load() ONLY.
-   */
-  protected void restoreHistoryIds() {
-    billingEventOneTimeHistoryId = DomainBase.getHistoryId(billingEventOneTime);
-    billingEventRecurringHistoryId = DomainBase.getHistoryId(billingEventRecurring);
-  }
-
-  /**
-   * Override {@link ImmutableObject#getSignificantFields()} to exclude "id", which breaks equality
-   * testing in the unit tests.
-   */
-  @Override
-  protected Map<Field, Object> getSignificantFields() {
-    restoreOfyKeys();
-    // 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;
-  }
-
-  /**
-   * Restores Ofy keys in the billing events.
-   *
-   * <p>This must be called by all methods that access the one time or recurring billing event keys.
-   * When the billing event keys are loaded from SQL, they are loaded as asymmetric keys because the
-   * database columns that we load them from do not contain all of the information necessary to
-   * reconsitute the Ofy side of the key. In other cases, we restore the Ofy key during the
-   * hibernate {@link javax.persistence.PostLoad} method from the other fields of the object, but we
-   * have been unable to make this work with hibernate's internal persistence model in this case
-   * because the {@link GracePeriod}'s hash code is evaluated prior to these calls, and would be
-   * invalidated by changing the fields.
-   */
-  private final synchronized void restoreOfyKeys() {
-    if (billingEventOneTime != null && !billingEventOneTime.maybeGetOfyKey().isPresent()) {
-      billingEventOneTime =
-          DomainBase.restoreOfyFrom(
-              Key.create(DomainBase.class, domainRepoId),
-              billingEventOneTime,
-              billingEventOneTimeHistoryId);
-    }
-    if (billingEventRecurring != null && !billingEventRecurring.maybeGetOfyKey().isPresent()) {
-      billingEventRecurring =
-          DomainBase.restoreOfyFrom(
-              Key.create(DomainBase.class, domainRepoId),
-              billingEventRecurring,
-              billingEventRecurringHistoryId);
-    }
+    return billingEventRecurring == null ? null : billingEventRecurring.createVKey();
   }
 }

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/domain/token/AllocationToken.java

@@ -111,9 +111,9 @@ public class AllocationToken extends BackupGroupRoot implements Buildable, Datas
   @Nullable
   @Index
   @AttributeOverrides({
-    @AttributeOverride(name = "domainRepoId", column = @Column(name = "redemption_domain_repo_id")),
+    @AttributeOverride(name = "repoId", column = @Column(name = "redemption_domain_repo_id")),
     @AttributeOverride(
-        name = "domainHistoryId",
+        name = "historyRevisionId",
         column = @Column(name = "redemption_domain_history_id"))
   })
   DomainHistoryVKey redemptionHistoryEntry;
@@ -192,8 +192,9 @@ public class AllocationToken extends BackupGroupRoot implements Buildable, Datas
     return token;
   }
 
-  public Optional<VKey<HistoryEntry>> getRedemptionHistoryEntry() {
-    return Optional.ofNullable(redemptionHistoryEntry);
+  public Optional<VKey<? extends HistoryEntry>> getRedemptionHistoryEntry() {
+    return Optional.ofNullable(
+        redemptionHistoryEntry == null ? null : redemptionHistoryEntry.createDomainHistoryVKey());
   }
 
   public boolean isRedeemed() {
@@ -291,9 +292,10 @@ public class AllocationToken extends BackupGroupRoot implements Buildable, Datas
       return this;
     }
 
-    public Builder setRedemptionHistoryEntry(DomainHistoryVKey redemptionHistoryEntry) {
+    public Builder setRedemptionHistoryEntry(VKey<? extends HistoryEntry> redemptionHistoryEntry) {
+      checkArgumentNotNull(redemptionHistoryEntry, "Redemption history entry must not be null");
       getInstance().redemptionHistoryEntry =
-          checkArgumentNotNull(redemptionHistoryEntry, "Redemption history entry must not be null");
+          DomainHistoryVKey.create(redemptionHistoryEntry.getOfyKey());
       return this;
     }
 

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

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

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

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

@@ -42,8 +42,8 @@ import google.registry.model.translators.CidrAddressBlockTranslatorFactory;
 import google.registry.model.translators.CommitLogRevisionsTranslatorFactory;
 import google.registry.model.translators.CreateAutoTimestampTranslatorFactory;
 import google.registry.model.translators.CurrencyUnitTranslatorFactory;
-import google.registry.model.translators.DomainHistoryVKeyTranslatorFactory;
 import google.registry.model.translators.DurationTranslatorFactory;
+import google.registry.model.translators.EppHistoryVKeyTranslatorFactory;
 import google.registry.model.translators.InetAddressTranslatorFactory;
 import google.registry.model.translators.ReadableInstantUtcTranslatorFactory;
 import google.registry.model.translators.UpdateAutoTimestampTranslatorFactory;
@@ -128,7 +128,7 @@ public class ObjectifyService {
             new CreateAutoTimestampTranslatorFactory(),
             new CurrencyUnitTranslatorFactory(),
             new DurationTranslatorFactory(),
-            new DomainHistoryVKeyTranslatorFactory(),
+            new EppHistoryVKeyTranslatorFactory(),
             new InetAddressTranslatorFactory(),
             new MoneyStringTranslatorFactory(),
             new ReadableInstantUtcTranslatorFactory(),

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

core/src/main/java/google/registry/model/translators/DomainHistoryVKeyTranslatorFactory.java

@@ -1,48 +0,0 @@
-// Copyright 2020 The Nomulus Authors. All Rights Reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package google.registry.model.translators;
-
-import com.google.appengine.api.datastore.Key;
-import google.registry.persistence.DomainHistoryVKey;
-import javax.annotation.Nullable;
-
-/** Translator factory for {@link DomainHistoryVKey}. */
-public class DomainHistoryVKeyTranslatorFactory
-    extends AbstractSimpleTranslatorFactory<DomainHistoryVKey, Key> {
-
-  public DomainHistoryVKeyTranslatorFactory() {
-    super(DomainHistoryVKey.class);
-  }
-
-  @Override
-  SimpleTranslator<DomainHistoryVKey, Key> createTranslator() {
-    return new SimpleTranslator<DomainHistoryVKey, Key>() {
-
-      @Nullable
-      @Override
-      public DomainHistoryVKey loadValue(@Nullable Key datastoreValue) {
-        return datastoreValue == null
-            ? null
-            : DomainHistoryVKey.create(com.googlecode.objectify.Key.create(datastoreValue));
-      }
-
-      @Nullable
-      @Override
-      public Key saveValue(@Nullable DomainHistoryVKey pojoValue) {
-        return pojoValue == null ? null : pojoValue.getOfyKey().getRaw();
-      }
-    };
-  }
-}

core/src/main/java/google/registry/model/translators/EppHistoryVKeyTranslatorFactory.java

@@ -0,0 +1,110 @@
+// Copyright 2020 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.model.translators;
+
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
+import static java.util.function.Function.identity;
+
+import com.google.appengine.api.datastore.Key;
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import google.registry.persistence.BillingVKey.BillingEventVKey;
+import google.registry.persistence.BillingVKey.BillingRecurrenceVKey;
+import google.registry.persistence.DomainHistoryVKey;
+import google.registry.persistence.EppHistoryVKey;
+import java.lang.reflect.Constructor;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import javax.annotation.Nullable;
+
+/** Translator factory for {@link EppHistoryVKey}. */
+public class EppHistoryVKeyTranslatorFactory
+    extends AbstractSimpleTranslatorFactory<EppHistoryVKey, Key> {
+
+  public EppHistoryVKeyTranslatorFactory() {
+    super(EppHistoryVKey.class);
+  }
+
+  // This map is used when we need to convert the raw Datastore key to its VKey instance. We have
+  // one dedicated VKey class, e.g. DomainHistoryVKey, for each such kind of entity, and we need
+  // a way to map the raw Datastore key to its VKey class. So, we use the kind path as the key of
+  // the map, and the kind path is created by concatenating all the kind strings in a raw Datastore
+  // key, e.g. the map key for ContactPollMessageVKey is "ContactResource/HistoryEntry/PollMessage".
+  @VisibleForTesting
+  static final ImmutableMap<String, Class<? extends EppHistoryVKey>> kindPathToVKeyClass =
+      ImmutableSet.of(DomainHistoryVKey.class, BillingEventVKey.class, BillingRecurrenceVKey.class)
+          .stream()
+          .collect(toImmutableMap(EppHistoryVKeyTranslatorFactory::getKindPath, identity()));
+
+  /**
+   * Gets the kind path string for the given {@link Class}.
+   *
+   * <p>This method calls the getKindPath method on an instance of the given {@link Class} to get
+   * the kind path string.
+   */
+  private static String getKindPath(Class<? extends EppHistoryVKey> clazz) {
+    try {
+      Constructor<?> constructor = clazz.getDeclaredConstructor();
+      constructor.setAccessible(true);
+      Object instance = constructor.newInstance();
+      Method getKindPathMethod = EppHistoryVKey.class.getDeclaredMethod("getKindPath");
+      getKindPathMethod.setAccessible(true);
+      return (String) getKindPathMethod.invoke(instance);
+    } catch (Throwable t) {
+      throw new IllegalStateException(t);
+    }
+  }
+
+  @Override
+  SimpleTranslator<EppHistoryVKey, Key> createTranslator() {
+    return new SimpleTranslator<EppHistoryVKey, Key>() {
+
+      @Nullable
+      @Override
+      public EppHistoryVKey loadValue(@Nullable Key datastoreValue) {
+        if (datastoreValue == null) {
+          return null;
+        } else {
+          com.googlecode.objectify.Key<?> ofyKey =
+              com.googlecode.objectify.Key.create(datastoreValue);
+          String kindPath = EppHistoryVKey.createKindPath(ofyKey);
+          if (kindPathToVKeyClass.containsKey(kindPath)) {
+            Class<? extends EppHistoryVKey> vKeyClass = kindPathToVKeyClass.get(kindPath);
+            try {
+              Method createVKeyMethod =
+                  vKeyClass.getDeclaredMethod("create", com.googlecode.objectify.Key.class);
+              return (EppHistoryVKey) createVKeyMethod.invoke(null, ofyKey);
+            } catch (NoSuchMethodException e) {
+              throw new IllegalStateException(
+                  "Missing static method create(com.googlecode.objectify.Key) on " + vKeyClass);
+            } catch (IllegalAccessException | InvocationTargetException e) {
+              throw new IllegalStateException("Error invoking createVKey on " + vKeyClass, e);
+            }
+          } else {
+            throw new IllegalStateException(
+                "Missing EppHistoryVKey implementation for kind path: " + kindPath);
+          }
+        }
+      }
+
+      @Nullable
+      @Override
+      public Key saveValue(@Nullable EppHistoryVKey pojoValue) {
+        return pojoValue == null ? null : pojoValue.createOfyKey().getRaw();
+      }
+    };
+  }
+}

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

@@ -0,0 +1,134 @@
+// Copyright 2020 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.persistence;
+
+
+import com.googlecode.objectify.Key;
+import google.registry.model.billing.BillingEvent;
+import google.registry.model.billing.BillingEvent.OneTime;
+import google.registry.model.billing.BillingEvent.Recurring;
+import google.registry.model.domain.DomainBase;
+import google.registry.model.reporting.HistoryEntry;
+import javax.annotation.Nullable;
+import javax.persistence.AttributeOverride;
+import javax.persistence.AttributeOverrides;
+import javax.persistence.Column;
+import javax.persistence.Embeddable;
+import javax.persistence.MappedSuperclass;
+
+/** Base class for {@link BillingEvent}'s {@link VKey}. */
+@MappedSuperclass
+public abstract class BillingVKey<K> extends EppHistoryVKey<K, DomainBase> {
+  Long billingId;
+
+  // Hibernate requires a default constructor.
+  BillingVKey() {}
+
+  BillingVKey(String repoId, long historyRevisionId, long billingId) {
+    super(repoId, historyRevisionId);
+    this.billingId = billingId;
+  }
+
+  Key<HistoryEntry> createHistoryEntryKey() {
+    return Key.create(Key.create(DomainBase.class, repoId), HistoryEntry.class, historyRevisionId);
+  }
+
+  @Override
+  public Object createSqlKey() {
+    return billingId;
+  }
+
+  /** VKey class for {@link BillingEvent.OneTime} that belongs to a {@link DomainBase} entity. */
+  @Embeddable
+  @AttributeOverrides({
+    @AttributeOverride(name = "repoId", column = @Column(name = "billing_event_domain_repo_id")),
+    @AttributeOverride(
+        name = "historyRevisionId",
+        column = @Column(name = "billing_event_history_id")),
+    @AttributeOverride(name = "billingId", column = @Column(name = "billing_event_id"))
+  })
+  public static class BillingEventVKey extends BillingVKey<OneTime> {
+
+    // Hibernate requires this default constructor
+    private BillingEventVKey() {}
+
+    private BillingEventVKey(String repoId, long historyRevisionId, long billingEventId) {
+      super(repoId, historyRevisionId, billingEventId);
+    }
+
+    @Override
+    public Key<OneTime> createOfyKey() {
+      return Key.create(createHistoryEntryKey(), BillingEvent.OneTime.class, billingId);
+    }
+
+    /** Creates a {@link BillingEventVKey} instance from the given {@link Key} instance. */
+    public static BillingEventVKey create(@Nullable Key<BillingEvent.OneTime> ofyKey) {
+      if (ofyKey == null) {
+        return null;
+      }
+      long billingEventId = ofyKey.getId();
+      long historyRevisionId = ofyKey.getParent().getId();
+      String repoId = ofyKey.getParent().getParent().getName();
+      return new BillingEventVKey(repoId, historyRevisionId, billingEventId);
+    }
+
+    /** Creates a {@link BillingEventVKey} instance from the given {@link VKey} instance. */
+    public static BillingEventVKey create(@Nullable VKey<BillingEvent.OneTime> vKey) {
+      return vKey == null ? null : create(vKey.getOfyKey());
+    }
+  }
+
+  /** VKey class for {@link BillingEvent.Recurring} that belongs to a {@link DomainBase} entity. */
+  @Embeddable
+  @AttributeOverrides({
+    @AttributeOverride(
+        name = "repoId",
+        column = @Column(name = "billing_recurrence_domain_repo_id")),
+    @AttributeOverride(
+        name = "historyRevisionId",
+        column = @Column(name = "billing_recurrence_history_id")),
+    @AttributeOverride(name = "billingId", column = @Column(name = "billing_recurrence_id"))
+  })
+  public static class BillingRecurrenceVKey extends BillingVKey<Recurring> {
+
+    // Hibernate requires this default constructor
+    private BillingRecurrenceVKey() {}
+
+    private BillingRecurrenceVKey(String repoId, long historyRevisionId, long billingEventId) {
+      super(repoId, historyRevisionId, billingEventId);
+    }
+
+    @Override
+    public Key<Recurring> createOfyKey() {
+      return Key.create(createHistoryEntryKey(), BillingEvent.Recurring.class, billingId);
+    }
+
+    /** Creates a {@link BillingRecurrenceVKey} instance from the given {@link Key} instance. */
+    public static BillingRecurrenceVKey create(@Nullable Key<BillingEvent.Recurring> ofyKey) {
+      if (ofyKey == null) {
+        return null;
+      }
+      long billingEventId = ofyKey.getId();
+      long historyRevisionId = ofyKey.getParent().getId();
+      String repoId = ofyKey.getParent().getParent().getName();
+      return new BillingRecurrenceVKey(repoId, historyRevisionId, billingEventId);
+    }
+
+    /** Creates a {@link BillingRecurrenceVKey} instance from the given {@link VKey} instance. */
+    public static BillingRecurrenceVKey create(@Nullable VKey<BillingEvent.Recurring> vKey) {
+      return vKey == null ? null : create(vKey.getOfyKey());
+    }
+  }
+}

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

@@ -18,45 +18,44 @@ import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 
 import com.googlecode.objectify.Key;
 import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.DomainHistory;
 import google.registry.model.domain.DomainHistory.DomainHistoryId;
 import google.registry.model.reporting.HistoryEntry;
 import javax.persistence.Embeddable;
-import javax.persistence.PostLoad;
 
 /** {@link VKey} for {@link HistoryEntry} which parent is {@link DomainBase}. */
 @Embeddable
-public class DomainHistoryVKey extends VKey<HistoryEntry> {
-
-  private String domainRepoId;
-
-  private Long domainHistoryId;
+public class DomainHistoryVKey extends EppHistoryVKey<HistoryEntry, DomainBase> {
 
   // Hibernate requires a default constructor
   private DomainHistoryVKey() {}
 
-  private DomainHistoryVKey(String domainRepoId, long domainHistoryId) {
-    initWith(domainRepoId, domainHistoryId);
+  private DomainHistoryVKey(String repoId, long historyRevisionId) {
+    super(repoId, historyRevisionId);
   }
 
-  @PostLoad
-  void postLoad() {
-    initWith(domainRepoId, domainHistoryId);
+  @Override
+  public Object createSqlKey() {
+    return new DomainHistoryId(repoId, historyRevisionId);
   }
 
-  private void initWith(String domainRepoId, long domainHistoryId) {
-    this.kind = HistoryEntry.class;
-    this.ofyKey =
-        Key.create(Key.create(DomainBase.class, domainRepoId), HistoryEntry.class, domainHistoryId);
-    this.sqlKey = new DomainHistoryId(domainRepoId, domainHistoryId);
-    this.domainRepoId = domainRepoId;
-    this.domainHistoryId = domainHistoryId;
+  @Override
+  public Key<HistoryEntry> createOfyKey() {
+    return Key.create(Key.create(DomainBase.class, repoId), HistoryEntry.class, historyRevisionId);
   }
 
   /** Creates {@link DomainHistoryVKey} from the given {@link Key} instance. */
-  public static DomainHistoryVKey create(Key<HistoryEntry> ofyKey) {
+  public static DomainHistoryVKey create(Key<? extends HistoryEntry> ofyKey) {
     checkArgumentNotNull(ofyKey, "ofyKey must be specified");
-    String domainRepoId = ofyKey.getParent().getName();
-    long domainHistoryId = ofyKey.getId();
-    return new DomainHistoryVKey(domainRepoId, domainHistoryId);
+    String repoId = ofyKey.getParent().getName();
+    long historyRevisionId = ofyKey.getId();
+    return new DomainHistoryVKey(repoId, historyRevisionId);
+  }
+
+  public VKey<? extends HistoryEntry> createDomainHistoryVKey() {
+    return VKey.create(
+        DomainHistory.class,
+        createSqlKey(),
+        Key.create(Key.create(DomainBase.class, repoId), DomainHistory.class, historyRevisionId));
   }
 }

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

@@ -0,0 +1,107 @@
+// Copyright 2020 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.persistence;
+
+import com.google.common.base.Joiner;
+import com.googlecode.objectify.Key;
+import google.registry.model.EppResource;
+import google.registry.model.ImmutableObject;
+import google.registry.model.reporting.HistoryEntry;
+import google.registry.util.TypeUtils.TypeInstantiator;
+import java.io.Serializable;
+import javax.annotation.Nullable;
+import javax.persistence.Access;
+import javax.persistence.AccessType;
+import javax.persistence.MappedSuperclass;
+
+/**
+ * Base class for {@link VKey} which ofyKey has a {@link HistoryEntry} key as its parent and a key
+ * for EPP resource as its grandparent.
+ *
+ * <p>For such a {@link VKey}, we need to provide two type parameters to indicate the type of {@link
+ * VKey} itself and the type of EPP resource respectively.
+ *
+ * @param <K> type of the {@link VKey}
+ * @param <E> type of the EPP resource that the key belongs to
+ */
+@MappedSuperclass
+@Access(AccessType.FIELD)
+public abstract class EppHistoryVKey<K, E extends EppResource> extends ImmutableObject
+    implements Serializable {
+
+  private static final long serialVersionUID = -3906580677709539818L;
+
+  String repoId;
+
+  Long historyRevisionId;
+
+  // Hibernate requires a default constructor.
+  EppHistoryVKey() {}
+
+  EppHistoryVKey(String repoId, long historyRevisionId) {
+    this.repoId = repoId;
+    this.historyRevisionId = historyRevisionId;
+  }
+
+  /**
+   * Returns the kind path for the ofyKey in this instance.
+   *
+   * <p>This method is only used reflectively by {@link EppHistoryVKeyTranslatorFactory} to get the
+   * kind path for a given {@link EppHistoryVKey} instance so it is marked as a private method.
+   *
+   * @see #createKindPath(Key)
+   */
+  @SuppressWarnings("unused")
+  private String getKindPath() {
+    String eppKind = Key.getKind(new TypeInstantiator<E>(getClass()) {}.getExactType());
+    String keyKind = Key.getKind(new TypeInstantiator<K>(getClass()) {}.getExactType());
+    if (keyKind.equals(Key.getKind(HistoryEntry.class))) {
+      return createKindPath(eppKind, keyKind);
+    } else {
+      return createKindPath(eppKind, Key.getKind(HistoryEntry.class), keyKind);
+    }
+  }
+
+  /**
+   * Creates the kind path for the given ofyKey}.
+   *
+   * <p>The kind path is a string including all kind names(delimited by slash) of a hierarchical
+   * {@link Key}, e.g., the kind path for BillingEvent.OneTime is "DomainBase/HistoryEntry/OneTime".
+   */
+  @Nullable
+  public static String createKindPath(@Nullable Key<?> ofyKey) {
+    if (ofyKey == null) {
+      return null;
+    } else if (ofyKey.getParent() == null) {
+      return ofyKey.getKind();
+    } else {
+      return createKindPath(createKindPath(ofyKey.getParent()), ofyKey.getKind());
+    }
+  }
+
+  private static String createKindPath(String... kinds) {
+    return Joiner.on("/").join(kinds);
+  }
+
+  /** Creates a {@link VKey} from this instance. */
+  public VKey<K> createVKey() {
+    Class<K> vKeyType = new TypeInstantiator<K>(getClass()) {}.getExactType();
+    return VKey.create(vKeyType, createSqlKey(), createOfyKey());
+  }
+
+  public abstract Object createSqlKey();
+
+  public abstract Key<K> createOfyKey();
+}

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

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

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

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

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

@@ -18,6 +18,7 @@ import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
+import static google.registry.model.ofy.DatastoreTransactionManager.toChildHistoryEntryIfPossible;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 import static java.util.AbstractMap.SimpleEntry;
 import static java.util.stream.Collectors.joining;
@@ -29,13 +30,19 @@ 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.model.ofy.DatastoreTransactionManager;
 import google.registry.persistence.JpaRetries;
 import google.registry.persistence.VKey;
 import google.registry.util.Clock;
 import google.registry.util.Retrier;
 import google.registry.util.SystemSleeper;
 import java.lang.reflect.Field;
-import java.util.Map.Entry;
+import java.util.Map;
 import java.util.NoSuchElementException;
 import java.util.Optional;
 import java.util.function.Supplier;
@@ -56,11 +63,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);
 
@@ -98,8 +118,6 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
 
   @Override
   public <T> T transact(Supplier<T> work) {
-    // TODO(shicong): Investigate removing transactNew functionality after migration as it may
-    //  be same as this one.
     return retrier.callWithRetry(
         () -> {
           if (inTransaction()) {
@@ -179,6 +197,8 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
         });
   }
 
+  // TODO(b/177674699): Remove all transactNew methods as they are same as transact after the
+  // database migration.
   @Override
   public <T> T transactNew(Supplier<T> work) {
     return transact(work);
@@ -228,9 +248,14 @@ 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);
+    // Necessary due to the changes in HistoryEntry representation during the migration to SQL
+    Object toPersist = toChildHistoryEntryIfPossible(entity);
+    getEntityManager().persist(toPersist);
+    transactionInfo.get().addUpdate(toPersist);
   }
 
   @Override
@@ -253,9 +278,14 @@ 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);
+    // Necessary due to the changes in HistoryEntry representation during the migration to SQL
+    Object toPersist = toChildHistoryEntryIfPossible(entity);
+    getEntityManager().merge(toPersist);
+    transactionInfo.get().addUpdate(toPersist);
   }
 
   @Override
@@ -278,10 +308,15 @@ 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);
-    transactionInfo.get().addUpdate(entity);
+    // Necessary due to the changes in HistoryEntry representation during the migration to SQL
+    Object toPersist = toChildHistoryEntryIfPossible(entity);
+    getEntityManager().merge(toPersist);
+    transactionInfo.get().addUpdate(toPersist);
   }
 
   @Override
@@ -312,6 +347,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   @Override
   public boolean exists(Object entity) {
     checkArgumentNotNull(entity, "entity must be specified");
+    entity = toChildHistoryEntryIfPossible(entity);
     EntityType<?> entityType = getEntityType(entity.getClass());
     ImmutableSet<EntityId> entityIds = getEntityIdsFromEntity(entityType, entity);
     return exists(entityType.getName(), entityIds);
@@ -330,33 +366,15 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   }
 
   @Override
-  public <T> Optional<T> maybeLoad(VKey<T> key) {
+  public <T> Optional<T> loadByKeyIfPresent(VKey<T> key) {
     checkArgumentNotNull(key, "key must be specified");
     assertInTransaction();
     return Optional.ofNullable(getEntityManager().find(key.getKind(), key.getSqlKey()));
   }
 
   @Override
-  public <T> T load(VKey<T> key) {
-    checkArgumentNotNull(key, "key must be specified");
-    assertInTransaction();
-    T result = getEntityManager().find(key.getKind(), key.getSqlKey());
-    if (result == null) {
-      throw new NoSuchElementException(key.toString());
-    }
-    return result;
-  }
-
-  @Override
-  public <T> T load(T entity) {
-    checkArgumentNotNull(entity, "entity must be specified");
-    assertInTransaction();
-    return (T)
-        load(VKey.createSql(entity.getClass(), emf.getPersistenceUnitUtil().getIdentifier(entity)));
-  }
-
-  @Override
-  public <T> ImmutableMap<VKey<? extends T>, T> load(Iterable<? extends VKey<? extends T>> keys) {
+  public <T> ImmutableMap<VKey<? extends T>, T> loadByKeysIfPresent(
+      Iterable<? extends VKey<? extends T>> keys) {
     checkArgumentNotNull(keys, "keys must be specified");
     assertInTransaction();
     return StreamSupport.stream(keys.spliterator(), false)
@@ -367,11 +385,64 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
                 new SimpleEntry<VKey<? extends T>, T>(
                     key, getEntityManager().find(key.getKind(), key.getSqlKey())))
         .filter(entry -> entry.getValue() != null)
-        .collect(toImmutableMap(Entry::getKey, Entry::getValue));
+        .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue));
   }
 
   @Override
-  public <T> ImmutableList<T> loadAll(Class<T> clazz) {
+  public <T> ImmutableList<T> loadByEntitiesIfPresent(Iterable<T> entities) {
+    return Streams.stream(entities)
+        .map(DatastoreTransactionManager::toChildHistoryEntryIfPossible)
+        .filter(this::exists)
+        .map(this::loadByEntity)
+        .collect(toImmutableList());
+  }
+
+  @Override
+  public <T> T loadByKey(VKey<T> key) {
+    checkArgumentNotNull(key, "key must be specified");
+    assertInTransaction();
+    T result = getEntityManager().find(key.getKind(), key.getSqlKey());
+    if (result == null) {
+      throw new NoSuchElementException(key.toString());
+    }
+    return result;
+  }
+
+  @Override
+  public <T> ImmutableMap<VKey<? extends T>, T> loadByKeys(
+      Iterable<? extends VKey<? extends T>> keys) {
+    ImmutableMap<VKey<? extends T>, T> existing = loadByKeysIfPresent(keys);
+    ImmutableSet<? extends VKey<? extends T>> missingKeys =
+        Streams.stream(keys).filter(k -> !existing.containsKey(k)).collect(toImmutableSet());
+    if (!missingKeys.isEmpty()) {
+      throw new NoSuchElementException(
+          String.format(
+              "Expected to find the following VKeys but they were missing: %s.", missingKeys));
+    }
+    return existing;
+  }
+
+  @Override
+  public <T> T loadByEntity(T entity) {
+    checkArgumentNotNull(entity, "entity must be specified");
+    assertInTransaction();
+    entity = toChildHistoryEntryIfPossible(entity);
+    // If the caller requested a HistoryEntry, load the corresponding *History class
+    T possibleChild = toChildHistoryEntryIfPossible(entity);
+    return (T)
+        loadByKey(
+            VKey.createSql(
+                possibleChild.getClass(),
+                emf.getPersistenceUnitUtil().getIdentifier(possibleChild)));
+  }
+
+  @Override
+  public <T> ImmutableList<T> loadByEntities(Iterable<T> entities) {
+    return Streams.stream(entities).map(this::loadByEntity).collect(toImmutableList());
+  }
+
+  @Override
+  public <T> ImmutableList<T> loadAllOf(Class<T> clazz) {
     checkArgumentNotNull(clazz, "clazz must be specified");
     assertInTransaction();
     return ImmutableList.copyOf(
@@ -382,11 +453,6 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
             .getResultList());
   }
 
-  @Override
-  public <T> ImmutableList<T> loadAll(Iterable<T> entities) {
-    return Streams.stream(entities).map(this::load).collect(toImmutableList());
-  }
-
   private int internalDelete(VKey<?> key) {
     checkArgumentNotNull(key, "key must be specified");
     assertInTransaction();
@@ -414,7 +480,11 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   @Override
   public void delete(Object entity) {
     checkArgumentNotNull(entity, "entity must be specified");
+    if (isEntityOfIgnoredClass(entity)) {
+      return;
+    }
     assertInTransaction();
+    entity = toChildHistoryEntryIfPossible(entity);
     Object managedEntity = entity;
     if (!getEntityManager().contains(entity)) {
       managedEntity = getEntityManager().merge(entity);
@@ -464,6 +534,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

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

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

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

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

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

core/src/main/java/google/registry/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()
+    ImmutableSet<VKey<AllocationToken>> tokensToDelete =
+        tm().loadByKeys(batch).values().stream()
             .filter(t -> withDomains || !t.getDomainName().isPresent())
             .filter(t -> SINGLE_USE.equals(t.getTokenType()))
             .filter(t -> !t.isRedeemed())
+            .map(AllocationToken::createVKey)
             .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,12 +45,15 @@ 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)) {
@@ -59,8 +62,10 @@ final class GetAllocationTokenCommand implements CommandWithRemoteApi {
         if (!loadedToken.getRedemptionHistoryEntry().isPresent()) {
           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/GetHistoryEntriesCommand.java

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

core/src/test/java/google/registry/batch/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 =