Remove a couple unused variables (#913)

* Remove a couple unused variables

config/presubmits.py

@@ -65,7 +65,7 @@ class PresubmitCheck:
     for pattern in self.skipped_patterns:
       if pattern in file:
         return False
-    with open(file, "r") as f:
+    with open(file, "r", encoding='utf8') as f:
       file_content = f.read()
       matches = re.match(self.regex, file_content, re.DOTALL)
       if self.regex_type == FORBIDDEN:
@@ -241,7 +241,7 @@ def verify_flyway_index():
 
   # Remove the sequence numbers and compare against the index file contents.
   files = [filename[1] for filename in sorted(files)]
-  with open('db/src/main/resources/sql/flyway.txt') as index:
+  with open('db/src/main/resources/sql/flyway.txt', encoding='utf8') as index:
     indexed_files = index.read().splitlines()
   if files != indexed_files:
     unindexed = set(files) - set(indexed_files)

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

@@ -31,6 +31,7 @@ import google.registry.persistence.PersistenceModule;
 import google.registry.persistence.PersistenceModule.JdbcJpaTm;
 import google.registry.persistence.PersistenceModule.SocketFactoryJpaTm;
 import google.registry.persistence.transaction.JpaTransactionManager;
+import google.registry.privileges.secretmanager.SecretManagerModule;
 import google.registry.util.UtilsModule;
 import java.io.BufferedReader;
 import java.io.IOException;
@@ -168,6 +169,7 @@ public class BeamJpaModule {
         BeamJpaModule.class,
         KmsModule.class,
         PersistenceModule.class,
+        SecretManagerModule.class,
         UtilsModule.class
       })
   public interface JpaTransactionManagerComponent {

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

@@ -43,7 +43,6 @@ 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.persistence.transaction.JpaTransactionManager;
 import google.registry.tools.LevelDbLogReader;
 import google.registry.util.SystemSleeper;
@@ -432,7 +431,6 @@ public final class Transforms {
     private final SerializableSupplier<JpaTransactionManager> jpaSupplier;
     private final SerializableFunction<T, Object> jpaConverter;
 
-    private transient Ofy ofy;
     private transient SystemSleeper sleeper;
 
     SqlBatchWriter(
@@ -448,7 +446,6 @@ public final class Transforms {
 
       try (AppEngineEnvironment env = new AppEngineEnvironment()) {
         ObjectifyService.initOfy();
-        ofy = ObjectifyService.ofy();
       }
 
       synchronized (SqlBatchWriter.class) {

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/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/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().maybeLoad(VKey.create(AllocationToken.class, token));
+    if (maybeTokenEntity.isEmpty()) {
       throw new InvalidAllocationTokenException();
     }
-    if (tokenEntity.isRedeemed()) {
+    if (maybeTokenEntity.get().isRedeemed()) {
       throw new AlreadyRedeemedAllocationTokenException();
     }
-    return tokenEntity;
+    return maybeTokenEntity.get();
   }
 
   // Note: exception messages should be <= 32 characters long for domain check results

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

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

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

@@ -0,0 +1,38 @@
+// 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;
+
+import com.google.common.flogger.FluentLogger;
+import google.registry.config.RegistryEnvironment;
+
+/** Utility methods related to migrating dual-read/dual-write entities. */
+public class DatabaseMigrationUtils {
+
+  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
+
+  /** Throws exceptions only in unit tests, otherwise only logs exceptions. */
+  public static void suppressExceptionUnlessInTest(Runnable work, String message) {
+    try {
+      work.run();
+    } catch (Exception e) {
+      if (RegistryEnvironment.get().equals(RegistryEnvironment.UNITTEST)) {
+        throw e;
+      }
+      logger.atWarning().withCause(e).log(message);
+    }
+  }
+
+  private DatabaseMigrationUtils() {}
+}

core/src/main/java/google/registry/model/EppResourceUtils.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.ofy.ObjectifyService.ofy;
 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;
@@ -135,7 +136,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 +144,7 @@ public final class EppResourceUtils {
     T resource =
         useCache
             ? EppResource.loadCached(fki.getResourceKey())
-            : tm().maybeLoad(fki.getResourceKey()).orElse(null);
+            : transactIfJpaTm(() -> tm().maybeLoad(fki.getResourceKey()).orElse(null));
     if (resource == null || isAtOrAfter(now, resource.getDeletionTime())) {
       return Optional.empty();
     }

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

@@ -22,12 +22,17 @@ 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
  */
 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).
+  private static ThreadLocal<Boolean> autoUpdateEnabled = ThreadLocal.withInitial(() -> true);
+
   DateTime timestamp;
 
   /** Returns the timestamp, or {@code START_OF_TIME} if it's null. */
@@ -40,4 +45,30 @@ public class UpdateAutoTimestamp extends ImmutableObject {
     instance.timestamp = timestamp;
     return instance;
   }
+
+  // TODO(b/175610935): Remove the auto-update disabling code below after migration.
+
+  /** Class to allow us to safely disable auto-update in a try-with-resources block. */
+  public static class DisableAutoUpdateResource implements AutoCloseable {
+    DisableAutoUpdateResource() {
+      autoUpdateEnabled.set(false);
+    }
+
+    @Override
+    public void close() {
+      autoUpdateEnabled.set(true);
+    }
+  }
+
+  /**
+   * Resturns a resource that disables auto-updates on all {@link UpdateAutoTimestamp}s in the
+   * current thread, suitable for use with in a try-with-resources block.
+   */
+  public static DisableAutoUpdateResource disableAutoUpdate() {
+    return new DisableAutoUpdateResource();
+  }
+
+  public static boolean autoUpdateEnabled() {
+    return autoUpdateEnabled.get();
+  }
 }

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

@@ -96,7 +96,14 @@ 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),
+      })
   @Column(name = "host_repo_id")
   Set<VKey<HostResource>> nsHosts;
 

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,7 +189,7 @@ public class GracePeriod extends GracePeriodBase implements DatastoreAndSqlEntit
    */
   public GracePeriod cloneWithRecurringBillingEvent(VKey<BillingEvent.Recurring> recurring) {
     GracePeriod clone = clone(this);
-    clone.billingEventRecurring = recurring;
+    clone.billingEventRecurring = BillingRecurrenceVKey.create(recurring);
     return clone;
   }
 
@@ -232,9 +231,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,14 +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;
@@ -68,24 +68,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 +115,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,18 +123,7 @@ 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);
+    return billingEventRecurring == null ? null : billingEventRecurring.createVKey();
   }
 
   /**
@@ -152,7 +132,6 @@ public class GracePeriodBase extends ImmutableObject {
    */
   @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()) {
@@ -162,33 +141,4 @@ public class GracePeriodBase extends ImmutableObject {
     }
     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);
-    }
-  }
 }

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

@@ -21,10 +21,10 @@ import javax.persistence.TemporalType;
 import org.joda.time.LocalDate;
 
 /**
- * Data access object for {@link google.registry.model.reporting.Spec11ThreatMatch}.
+ * Data access object for {@link Spec11ThreatMatch}.
  *
- * <p>A JpaTransactionManager is passed into each static method because they are called from a BEAM
- * pipeline and we don't know where it's coming from.
+ * <p>The transaction manager is passed as a parameter because this could be called either from a
+ * BEAM pipeline or standard non-BEAM code.
  */
 public class Spec11ThreatMatchDao {
 

core/src/main/java/google/registry/model/smd/SignedMarkRevocationList.java

@@ -19,6 +19,7 @@ import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.Iterables.isEmpty;
 import static google.registry.model.CacheUtils.memoizeWithShortExpiration;
+import static google.registry.model.DatabaseMigrationUtils.suppressExceptionUnlessInTest;
 import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
 import static google.registry.model.ofy.ObjectifyService.allocateId;
 import static google.registry.model.ofy.ObjectifyService.ofy;
@@ -32,7 +33,6 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.MapDifference;
 import com.google.common.collect.Maps;
-import com.google.common.flogger.FluentLogger;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.EmbedMap;
 import com.googlecode.objectify.annotation.Entity;
@@ -82,8 +82,6 @@ import org.joda.time.DateTime;
 @NotBackedUp(reason = Reason.EXTERNALLY_SOURCED)
 public class SignedMarkRevocationList extends ImmutableObject implements NonReplicatedEntity {
 
-  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
-
   @VisibleForTesting static final int SHARD_SIZE = 10000;
 
   /** Common ancestor for queries. */
@@ -121,12 +119,11 @@ public class SignedMarkRevocationList extends ImmutableObject implements NonRepl
       memoizeWithShortExpiration(
           () -> {
             SignedMarkRevocationList datastoreList = loadFromDatastore();
-            // Also load the list from Cloud SQL, compare the two lists, and log if different.
-            try {
-              loadAndCompareCloudSqlList(datastoreList);
-            } catch (Throwable t) {
-              logger.atSevere().withCause(t).log("Error comparing signed mark revocation lists.");
-            }
+            suppressExceptionUnlessInTest(
+                () -> {
+                  loadAndCompareCloudSqlList(datastoreList);
+                },
+                "Error comparing signed mark revocation lists.");
             return datastoreList;
           });
 
@@ -229,11 +226,12 @@ public class SignedMarkRevocationList extends ImmutableObject implements NonRepl
           Maps.difference(datastoreList.revokes, cloudSqlList.revokes);
       if (!diff.areEqual()) {
         if (diff.entriesDiffering().size() > 10) {
-          logger.atWarning().log(
+          String message =
               String.format(
                   "Unequal SM revocation lists detected, Cloud SQL list with revision id %d has %d"
                       + " different records than the current Datastore list.",
-                  cloudSqlList.revisionId, diff.entriesDiffering().size()));
+                  cloudSqlList.revisionId, diff.entriesDiffering().size());
+          throw new RuntimeException(message);
         } else {
           StringBuilder diffMessage = new StringBuilder("Unequal SM revocation lists detected:\n");
           diff.entriesDiffering()
@@ -243,11 +241,13 @@ public class SignedMarkRevocationList extends ImmutableObject implements NonRepl
                           String.format(
                               "SMD %s has key %s in Datastore and key %s in Cloud SQL.\n",
                               label, valueDiff.leftValue(), valueDiff.rightValue())));
-          logger.atWarning().log(diffMessage.toString());
+          throw new RuntimeException(diffMessage.toString());
         }
       }
     } else {
-      logger.atWarning().log("Signed mark revocation list in Cloud SQL is empty.");
+      if (datastoreList.size() != 0) {
+        throw new RuntimeException("Signed mark revocation list in Cloud SQL is empty.");
+      }
     }
   }
 

core/src/main/java/google/registry/model/smd/SignedMarkRevocationListDao.java

@@ -15,6 +15,7 @@
 package google.registry.model.smd;
 
 import static google.registry.model.CacheUtils.memoizeWithShortExpiration;
+import static google.registry.model.DatabaseMigrationUtils.suppressExceptionUnlessInTest;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 
 import com.google.common.base.Supplier;
@@ -60,14 +61,14 @@ public class SignedMarkRevocationListDao {
    * the authoritative database.
    */
   static void trySave(SignedMarkRevocationList signedMarkRevocationList) {
-    try {
-      SignedMarkRevocationListDao.save(signedMarkRevocationList);
-      logger.atInfo().log(
-          "Inserted %,d signed mark revocations into Cloud SQL",
-          signedMarkRevocationList.revokes.size());
-    } catch (Throwable e) {
-      logger.atSevere().withCause(e).log("Error inserting signed mark revocations into Cloud SQL");
-    }
+    suppressExceptionUnlessInTest(
+        () -> {
+          SignedMarkRevocationListDao.save(signedMarkRevocationList);
+          logger.atInfo().log(
+              "Inserted %,d signed mark revocations into Cloud SQL.",
+              signedMarkRevocationList.revokes.size());
+        },
+        "Error inserting signed mark revocations into Cloud SQL.");
   }
 
   private static void save(SignedMarkRevocationList signedMarkRevocationList) {

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/model/translators/UpdateAutoTimestampTranslatorFactory.java

@@ -46,7 +46,10 @@ public class UpdateAutoTimestampTranslatorFactory
       /** Save a timestamp, setting it to the current time. */
       @Override
       public Date saveValue(UpdateAutoTimestamp pojoValue) {
-        return tm().getTransactionTime().toDate();
-      }};
+        return UpdateAutoTimestamp.autoUpdateEnabled()
+            ? tm().getTransactionTime().toDate()
+            : pojoValue.getTimestamp().toDate();
+      }
+    };
   }
 }

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

@@ -20,6 +20,7 @@ import google.registry.config.RegistryConfig.ConfigModule;
 import google.registry.keyring.kms.KmsModule;
 import google.registry.persistence.PersistenceModule.AppEngineJpaTm;
 import google.registry.persistence.transaction.JpaTransactionManager;
+import google.registry.privileges.secretmanager.SecretManagerModule;
 import google.registry.util.UtilsModule;
 import javax.inject.Singleton;
 import javax.persistence.EntityManagerFactory;
@@ -32,6 +33,7 @@ import javax.persistence.EntityManagerFactory;
       CredentialModule.class,
       KmsModule.class,
       PersistenceModule.class,
+      SecretManagerModule.class,
       UtilsModule.class
     })
 public interface PersistenceComponent {

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

@@ -26,6 +26,7 @@ import com.google.api.client.auth.oauth2.Credential;
 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.Module;
 import dagger.Provides;
 import google.registry.config.RegistryConfig.Config;
@@ -33,6 +34,11 @@ import google.registry.keyring.kms.KmsKeyring;
 import google.registry.persistence.transaction.CloudSqlCredentialSupplier;
 import google.registry.persistence.transaction.JpaTransactionManager;
 import google.registry.persistence.transaction.JpaTransactionManagerImpl;
+import google.registry.privileges.secretmanager.SqlCredential;
+import google.registry.privileges.secretmanager.SqlCredentialStore;
+import google.registry.privileges.secretmanager.SqlUser;
+import google.registry.privileges.secretmanager.SqlUser.RobotId;
+import google.registry.privileges.secretmanager.SqlUser.RobotUser;
 import google.registry.tools.AuthModule.CloudSqlClientCredential;
 import google.registry.util.Clock;
 import java.lang.annotation.Documented;
@@ -47,6 +53,8 @@ import org.hibernate.cfg.Environment;
 /** Dagger module class for the persistence layer. */
 @Module
 public class PersistenceModule {
+  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
+
   // This name must be the same as the one defined in persistence.xml.
   public static final String PERSISTENCE_UNIT_NAME = "nomulus";
   public static final String HIKARI_CONNECTION_TIMEOUT = "hibernate.hikari.connectionTimeout";
@@ -122,11 +130,17 @@ public class PersistenceModule {
   static JpaTransactionManager provideAppEngineJpaTm(
       @Config("cloudSqlUsername") String username,
       KmsKeyring kmsKeyring,
+      SqlCredentialStore credentialStore,
       @PartialCloudSqlConfigs ImmutableMap<String, String> cloudSqlConfigs,
       Clock clock) {
     HashMap<String, String> overrides = Maps.newHashMap(cloudSqlConfigs);
     overrides.put(Environment.USER, username);
     overrides.put(Environment.PASS, kmsKeyring.getCloudSqlPassword());
+    validateCredentialStore(
+        credentialStore,
+        new RobotUser(RobotId.NOMULUS),
+        overrides.get(Environment.USER),
+        overrides.get(Environment.PASS));
     return new JpaTransactionManagerImpl(create(overrides), clock);
   }
 
@@ -136,6 +150,7 @@ public class PersistenceModule {
   static JpaTransactionManager provideNomulusToolJpaTm(
       @Config("toolsCloudSqlUsername") String username,
       KmsKeyring kmsKeyring,
+      SqlCredentialStore credentialStore,
       @PartialCloudSqlConfigs ImmutableMap<String, String> cloudSqlConfigs,
       @CloudSqlClientCredential Credential credential,
       Clock clock) {
@@ -143,6 +158,11 @@ public class PersistenceModule {
     HashMap<String, String> overrides = Maps.newHashMap(cloudSqlConfigs);
     overrides.put(Environment.USER, username);
     overrides.put(Environment.PASS, kmsKeyring.getToolsCloudSqlPassword());
+    validateCredentialStore(
+        credentialStore,
+        new RobotUser(RobotId.TOOL),
+        overrides.get(Environment.USER),
+        overrides.get(Environment.PASS));
     return new JpaTransactionManagerImpl(create(overrides), clock);
   }
 
@@ -150,6 +170,7 @@ public class PersistenceModule {
   @Singleton
   @SocketFactoryJpaTm
   static JpaTransactionManager provideSocketFactoryJpaTm(
+      SqlCredentialStore credentialStore,
       @Config("beamCloudSqlUsername") String username,
       @Config("beamCloudSqlPassword") String password,
       @Config("beamHibernateHikariMaximumPoolSize") int hikariMaximumPoolSize,
@@ -159,6 +180,12 @@ public class PersistenceModule {
     overrides.put(Environment.USER, username);
     overrides.put(Environment.PASS, password);
     overrides.put(HIKARI_MAXIMUM_POOL_SIZE, String.valueOf(hikariMaximumPoolSize));
+    // TODO(b/175700623): consider assigning different logins to pipelines
+    validateCredentialStore(
+        credentialStore,
+        new RobotUser(RobotId.NOMULUS),
+        overrides.get(Environment.USER),
+        overrides.get(Environment.PASS));
     return new JpaTransactionManagerImpl(create(overrides), clock);
   }
 
@@ -203,6 +230,30 @@ public class PersistenceModule {
     return emf;
   }
 
+  /** Verifies that the credential from the Secret Manager matches the one currently in use.
+   *
+   * <p>This is a helper for the transition to the Secret Manager, and will be removed once data
+   * and permissions are properly set up for all projects.
+   **/
+  private static void validateCredentialStore(
+      SqlCredentialStore credentialStore, SqlUser sqlUser, String login, String password) {
+    try {
+      SqlCredential credential = credentialStore.getCredential(sqlUser);
+      if (!credential.login().equals(login)) {
+        logger.atWarning().log(
+            "Wrong login for %s. Expecting %s, found %s.",
+            sqlUser.geUserName(), login, credential.login());
+        return;
+      }
+      if (!credential.password().equals(password)) {
+        logger.atWarning().log("Wrong password for %s.", sqlUser.geUserName());
+      }
+      logger.atWarning().log("Credentials in the kerying and the secret manager match.");
+    } catch (Throwable e) {
+      logger.atWarning().log(e.getMessage());
+    }
+  }
+
   /** Dagger qualifier for {@link JpaTransactionManager} used for App Engine application. */
   @Qualifier
   @Documented

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

@@ -31,7 +31,14 @@ public class UpdateAutoTimestampConverter
 
   @Override
   public Timestamp convertToDatabaseColumn(UpdateAutoTimestamp entity) {
-    return Timestamp.from(DateTimeUtils.toZonedDateTime(jpaTm().getTransactionTime()).toInstant());
+    return Timestamp.from(
+        DateTimeUtils.toZonedDateTime(
+                UpdateAutoTimestamp.autoUpdateEnabled()
+                        || entity == null
+                        || entity.getTimestamp() == null
+                    ? jpaTm().getTransactionTime()
+                    : entity.getTimestamp())
+            .toInstant());
   }
 
   @Override

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

@@ -29,6 +29,11 @@ import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Streams;
 import com.google.common.flogger.FluentLogger;
 import google.registry.config.RegistryConfig;
+import google.registry.model.ImmutableObject;
+import google.registry.model.index.EppResourceIndex;
+import google.registry.model.index.ForeignKeyIndex.ForeignKeyContactIndex;
+import google.registry.model.index.ForeignKeyIndex.ForeignKeyDomainIndex;
+import google.registry.model.index.ForeignKeyIndex.ForeignKeyHostIndex;
 import google.registry.persistence.JpaRetries;
 import google.registry.persistence.VKey;
 import google.registry.util.Clock;
@@ -56,6 +61,18 @@ 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;
@@ -228,6 +245,9 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   @Override
   public void insert(Object entity) {
     checkArgumentNotNull(entity, "entity must be specified");
+    if (isEntityOfIgnoredClass(entity)) {
+      return;
+    }
     assertInTransaction();
     getEntityManager().persist(entity);
     transactionInfo.get().addUpdate(entity);
@@ -253,6 +273,9 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   @Override
   public void put(Object entity) {
     checkArgumentNotNull(entity, "entity must be specified");
+    if (isEntityOfIgnoredClass(entity)) {
+      return;
+    }
     assertInTransaction();
     getEntityManager().merge(entity);
     transactionInfo.get().addUpdate(entity);
@@ -278,6 +301,9 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   @Override
   public void update(Object entity) {
     checkArgumentNotNull(entity, "entity must be specified");
+    if (isEntityOfIgnoredClass(entity)) {
+      return;
+    }
     assertInTransaction();
     checkArgument(exists(entity), "Given entity does not exist");
     getEntityManager().merge(entity);
@@ -414,6 +440,9 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   @Override
   public void delete(Object entity) {
     checkArgumentNotNull(entity, "entity must be specified");
+    if (isEntityOfIgnoredClass(entity)) {
+      return;
+    }
     assertInTransaction();
     Object managedEntity = entity;
     if (!getEntityManager().contains(entity)) {
@@ -464,6 +493,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/privileges/secretmanager/SecretManagerModule.java

@@ -16,13 +16,13 @@ package google.registry.privileges.secretmanager;
 
 
 import com.google.cloud.secretmanager.v1.SecretManagerServiceClient;
-import dagger.Component;
+import com.google.cloud.secretmanager.v1.SecretManagerServiceSettings;
 import dagger.Module;
 import dagger.Provides;
+import google.registry.config.CredentialModule.DefaultCredential;
 import google.registry.config.RegistryConfig.Config;
-import google.registry.config.RegistryConfig.ConfigModule;
+import google.registry.util.GoogleCredentialsBundle;
 import google.registry.util.Retrier;
-import google.registry.util.UtilsModule;
 import java.io.IOException;
 import javax.inject.Singleton;
 
@@ -32,20 +32,29 @@ public abstract class SecretManagerModule {
 
   @Provides
   @Singleton
-  static SecretManagerClient provideSecretManagerClient(
-      @Config("projectId") String project, Retrier retrier) {
+  static SecretManagerServiceSettings provideSecretManagerSetting(
+      @DefaultCredential GoogleCredentialsBundle credentialsBundle) {
     try {
-      SecretManagerServiceClient stub = SecretManagerServiceClient.create();
+      return SecretManagerServiceSettings.newBuilder()
+          .setCredentialsProvider(() -> credentialsBundle.getGoogleCredentials())
+          .build();
+    } catch (IOException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  @Provides
+  @Singleton
+  static SecretManagerClient provideSecretManagerClient(
+      SecretManagerServiceSettings serviceSettings,
+      @Config("projectId") String project,
+      Retrier retrier) {
+    try {
+      SecretManagerServiceClient stub = SecretManagerServiceClient.create(serviceSettings);
       Runtime.getRuntime().addShutdownHook(new Thread(stub::close));
       return new SecretManagerClientImpl(project, stub, retrier);
     } catch (IOException e) {
       throw new RuntimeException(e);
     }
   }
-
-  @Singleton
-  @Component(modules = {ConfigModule.class, SecretManagerModule.class, UtilsModule.class})
-  public interface SecretManagerComponent {
-    SecretManagerClient secretManagerClient();
-  }
 }

core/src/main/java/google/registry/privileges/secretmanager/SqlUser.java

@@ -48,7 +48,12 @@ public abstract class SqlUser {
 
   /** Enumerates the {@link RobotUser RobotUsers} in the system. */
   public enum RobotId {
-    NOMULUS;
+    NOMULUS,
+    /**
+     * Credential for RegistryTool. This is temporary, and will be removed when tool users are
+     * assigned their personal credentials.
+     */
+    TOOL;
   }
 
   /** Information of a RobotUser for privilege management purposes. */

core/src/main/java/google/registry/reporting/spec11/Spec11RegistrarThreatMatchesParser.java

@@ -14,9 +14,11 @@
 
 package google.registry.reporting.spec11;
 
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static java.nio.charset.StandardCharsets.UTF_8;
 
 import com.google.appengine.tools.cloudstorage.GcsFilename;
+import com.google.common.base.Splitter;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.flogger.FluentLogger;
@@ -59,6 +61,24 @@ public class Spec11RegistrarThreatMatchesParser {
     return getFromFile(getGcsFilename(date));
   }
 
+  /** Returns registrar:set-of-threat-match pairings from the file, or empty if it doesn't exist. */
+  public ImmutableSet<RegistrarThreatMatches> getFromFile(GcsFilename spec11ReportFilename)
+      throws IOException {
+    if (!gcsUtils.existsAndNotEmpty(spec11ReportFilename)) {
+      return ImmutableSet.of();
+    }
+    try (InputStream in = gcsUtils.openInputStream(spec11ReportFilename);
+        InputStreamReader isr = new InputStreamReader(in, UTF_8)) {
+      // Skip the header at line 0
+      return Splitter.on("\n")
+          .omitEmptyStrings()
+          .splitToStream(CharStreams.toString(isr))
+          .skip(1)
+          .map(this::parseRegistrarThreatMatch)
+          .collect(toImmutableSet());
+    }
+  }
+
   public Optional<LocalDate> getPreviousDateWithMatches(LocalDate date) {
     LocalDate yesterday = date.minusDays(1);
     GcsFilename gcsFilename = getGcsFilename(yesterday);
@@ -82,20 +102,6 @@ public class Spec11RegistrarThreatMatchesParser {
     return new GcsFilename(reportingBucket, Spec11Pipeline.getSpec11ReportFilePath(localDate));
   }
 
-  private ImmutableSet<RegistrarThreatMatches> getFromFile(GcsFilename spec11ReportFilename)
-      throws IOException, JSONException {
-    ImmutableSet.Builder<RegistrarThreatMatches> builder = ImmutableSet.builder();
-    try (InputStream in = gcsUtils.openInputStream(spec11ReportFilename)) {
-      ImmutableList<String> reportLines =
-          ImmutableList.copyOf(CharStreams.toString(new InputStreamReader(in, UTF_8)).split("\n"));
-      // Iterate from 1 to size() to skip the header at line 0.
-      for (int i = 1; i < reportLines.size(); i++) {
-        builder.add(parseRegistrarThreatMatch(reportLines.get(i)));
-      }
-      return builder.build();
-    }
-  }
-
   private RegistrarThreatMatches parseRegistrarThreatMatch(String line) throws JSONException {
     JSONObject reportJSON = new JSONObject(line);
     String clientId = reportJSON.getString(Spec11Pipeline.REGISTRAR_CLIENT_ID_FIELD);

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

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

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

@@ -20,8 +20,8 @@ import static com.google.common.collect.Queues.newArrayDeque;
 import static com.google.common.collect.Sets.difference;
 import static google.registry.model.domain.token.AllocationToken.TokenType.SINGLE_USE;
 import static google.registry.model.domain.token.AllocationToken.TokenType.UNLIMITED_USE;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 import static google.registry.util.CollectionUtils.nullToEmpty;
 import static google.registry.util.StringGenerator.DEFAULT_PASSWORD_LENGTH;
 import static java.nio.charset.StandardCharsets.UTF_8;
@@ -38,10 +38,10 @@ import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Streams;
 import com.google.common.io.Files;
-import com.googlecode.objectify.Key;
 import google.registry.model.domain.token.AllocationToken;
 import google.registry.model.domain.token.AllocationToken.TokenStatus;
 import google.registry.model.domain.token.AllocationToken.TokenType;
+import google.registry.persistence.VKey;
 import google.registry.tools.params.TransitionListParameter.TokenStatusTransitions;
 import google.registry.util.CollectionUtils;
 import google.registry.util.NonFinalForTesting;
@@ -258,8 +258,13 @@ class GenerateAllocationTokensCommand implements CommandWithRemoteApi {
 
   @VisibleForTesting
   int saveTokens(final ImmutableSet<AllocationToken> tokens) {
-    Collection<AllocationToken> savedTokens =
-        dryRun ? tokens : tm().transact(() -> ofy().save().entities(tokens).now().values());
+    Collection<AllocationToken> savedTokens;
+    if (dryRun) {
+      savedTokens = tokens;
+    } else {
+      transactIfJpaTm(() -> tm().transact(() -> tm().putAll(tokens)));
+      savedTokens = tm().transact(() -> tm().loadAll(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().load(existingTokenKeys).values().stream()
+                .map(AllocationToken::getToken)
+                .collect(toImmutableSet()));
   }
 }

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

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

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

@@ -16,6 +16,7 @@ package google.registry.tools;
 
 import com.google.common.collect.ImmutableMap;
 import google.registry.tools.javascrap.BackfillRegistryLocksCommand;
+import google.registry.tools.javascrap.BackfillSpec11ThreatMatchesCommand;
 import google.registry.tools.javascrap.PopulateNullRegistrarFieldsCommand;
 import google.registry.tools.javascrap.RemoveIpAddressCommand;
 
@@ -32,6 +33,7 @@ public final class RegistryTool {
       new ImmutableMap.Builder<String, Class<? extends Command>>()
           .put("ack_poll_messages", AckPollMessagesCommand.class)
           .put("backfill_registry_locks", BackfillRegistryLocksCommand.class)
+          .put("backfill_spec11_threat_matches", BackfillSpec11ThreatMatchesCommand.class)
           .put("canonicalize_labels", CanonicalizeLabelsCommand.class)
           .put("check_domain", CheckDomainCommand.class)
           .put("check_domain_claims", CheckDomainClaimsCommand.class)
@@ -108,6 +110,7 @@ public final class RegistryTool {
           .put("save_sql_credential", SaveSqlCredentialCommand.class)
           .put("send_escrow_report_to_icann", SendEscrowReportToIcannCommand.class)
           .put("set_num_instances", SetNumInstancesCommand.class)
+          .put("set_sql_replay_checkpoint", SetSqlReplayCheckpointCommand.class)
           .put("setup_ote", SetupOteCommand.class)
           .put("uniform_rapid_suspension", UniformRapidSuspensionCommand.class)
           .put("unlock_domain", UnlockDomainCommand.class)

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

@@ -148,6 +148,8 @@ interface RegistryToolComponent {
 
   void inject(SetNumInstancesCommand command);
 
+  void inject(SetSqlReplayCheckpointCommand command);
+
   void inject(SetupOteCommand command);
 
   void inject(UnlockDomainCommand command);

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

@@ -0,0 +1,48 @@
+// 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.tools;
+
+import static com.google.common.base.Preconditions.checkArgument;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+
+import com.beust.jcommander.Parameter;
+import com.beust.jcommander.Parameters;
+import com.google.common.collect.Iterables;
+import google.registry.schema.replay.SqlReplayCheckpoint;
+import java.util.List;
+import org.joda.time.DateTime;
+
+/** Command to set {@link SqlReplayCheckpoint} to a particular, post-initial-population time. */
+@Parameters(separators = " =", commandDescription = "Set SqlReplayCheckpoint to a particular time")
+public class SetSqlReplayCheckpointCommand extends ConfirmingCommand
+    implements CommandWithRemoteApi {
+
+  @Parameter(description = "Time to which SqlReplayCheckpoint will be set", required = true)
+  List<DateTime> mainParameters;
+
+  @Override
+  protected String prompt() {
+    checkArgument(mainParameters.size() == 1, "Must provide exactly one DateTime to set");
+    return String.format(
+        "Set SqlReplayCheckpoint to %s?", Iterables.getOnlyElement(mainParameters));
+  }
+
+  @Override
+  protected String execute() {
+    DateTime dateTime = Iterables.getOnlyElement(mainParameters);
+    jpaTm().transact(() -> SqlReplayCheckpoint.set(dateTime));
+    return String.format("Set SqlReplayCheckpoint time to %s", dateTime);
+  }
+}

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().load(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/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().loadAll(AllocationToken.class).stream()
+                  .filter(token -> token.getToken().startsWith(prefix))
+                  .map(AllocationToken::createVKey)
+                  .collect(toImmutableSet()));
     }
   }
 }

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

@@ -0,0 +1,216 @@
+// 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.tools.javascrap;
+
+import static com.google.common.base.Preconditions.checkState;
+import static com.google.common.collect.ImmutableListMultimap.flatteningToImmutableListMultimap;
+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 com.beust.jcommander.Parameter;
+import com.beust.jcommander.Parameters;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableListMultimap;
+import com.google.common.collect.ImmutableSet;
+import google.registry.beam.spec11.ThreatMatch;
+import google.registry.model.domain.DomainBase;
+import google.registry.model.reporting.Spec11ThreatMatch;
+import google.registry.model.reporting.Spec11ThreatMatch.ThreatType;
+import google.registry.model.reporting.Spec11ThreatMatchDao;
+import google.registry.reporting.spec11.RegistrarThreatMatches;
+import google.registry.reporting.spec11.Spec11RegistrarThreatMatchesParser;
+import google.registry.tools.CommandWithRemoteApi;
+import google.registry.tools.ConfirmingCommand;
+import google.registry.util.Clock;
+import java.io.IOException;
+import java.util.Comparator;
+import java.util.List;
+import java.util.function.Function;
+import javax.inject.Inject;
+import org.joda.time.LocalDate;
+
+/**
+ * Scrap tool to backfill {@link Spec11ThreatMatch} objects from prior days.
+ *
+ * <p>This will load the previously-existing Spec11 files from GCS (looking back to 2019-01-01 (a
+ * rough estimate of when we started using this format) and convert those RegistrarThreatMatches
+ * objects into the new Spec11ThreatMatch format. It will then insert these entries into SQL.
+ *
+ * <p>Note that the script will attempt to find the corresponding {@link DomainBase} object for each
+ * domain name on the day of the scan. It will fail if it cannot find a corresponding domain object,
+ * or if the domain objects were not active at the time of the scan.
+ */
+@Parameters(
+    commandDescription =
+        "Backfills Spec11 threat match entries from the old and deprecated GCS JSON files to the "
+            + "Cloud SQL database.")
+public class BackfillSpec11ThreatMatchesCommand extends ConfirmingCommand
+    implements CommandWithRemoteApi {
+
+  private static final LocalDate START_DATE = new LocalDate(2019, 1, 1);
+
+  @Parameter(
+      names = {"-o", "--overwrite_existing_dates"},
+      description =
+          "Whether the command will overwrite data that already exists for dates that exist in the "
+              + "GCS bucket. Defaults to false.")
+  private boolean overrideExistingDates;
+
+  @Inject Spec11RegistrarThreatMatchesParser threatMatchesParser;
+  // Inject the clock for testing purposes
+  @Inject Clock clock;
+
+  @Override
+  protected String prompt() {
+    return String.format("Backfill Spec11 results from %d files?", getDatesToBackfill().size());
+  }
+
+  @Override
+  protected String execute() {
+    ImmutableList<LocalDate> dates = getDatesToBackfill();
+    ImmutableListMultimap.Builder<LocalDate, RegistrarThreatMatches> threatMatchesBuilder =
+        new ImmutableListMultimap.Builder<>();
+    for (LocalDate date : dates) {
+      try {
+        // It's OK if the file doesn't exist for a particular date; the result will be empty.
+        threatMatchesBuilder.putAll(date, threatMatchesParser.getRegistrarThreatMatches(date));
+      } catch (IOException e) {
+        throw new RuntimeException(
+            String.format("Error parsing through file with date %s.", date), e);
+      }
+    }
+    ImmutableListMultimap<LocalDate, RegistrarThreatMatches> threatMatches =
+        threatMatchesBuilder.build();
+    // Look up all possible DomainBases for these domain names, any of which can be in the past
+    ImmutableListMultimap<String, DomainBase> domainsByDomainName =
+        getDomainsByDomainName(threatMatches);
+
+    // For each date, convert all threat matches with the proper domain repo ID
+    int totalNumThreats = 0;
+    for (LocalDate date : threatMatches.keySet()) {
+      ImmutableList.Builder<Spec11ThreatMatch> spec11ThreatsBuilder = new ImmutableList.Builder<>();
+      for (RegistrarThreatMatches rtm : threatMatches.get(date)) {
+        rtm.threatMatches().stream()
+            .map(
+                threatMatch ->
+                    threatMatchToCloudSqlObject(
+                        threatMatch, date, rtm.clientId(), domainsByDomainName))
+            .forEach(spec11ThreatsBuilder::add);
+      }
+      ImmutableList<Spec11ThreatMatch> spec11Threats = spec11ThreatsBuilder.build();
+      jpaTm()
+          .transact(
+              () -> {
+                Spec11ThreatMatchDao.deleteEntriesByDate(jpaTm(), date);
+                jpaTm().putAll(spec11Threats);
+              });
+      totalNumThreats += spec11Threats.size();
+    }
+    return String.format(
+        "Successfully parsed through %d files with %d threats.", dates.size(), totalNumThreats);
+  }
+
+  /** Returns a per-domain list of possible DomainBase objects, starting with the most recent. */
+  private ImmutableListMultimap<String, DomainBase> getDomainsByDomainName(
+      ImmutableListMultimap<LocalDate, RegistrarThreatMatches> threatMatchesByDate) {
+    return threatMatchesByDate.values().stream()
+        .map(RegistrarThreatMatches::threatMatches)
+        .flatMap(ImmutableList::stream)
+        .map(ThreatMatch::fullyQualifiedDomainName)
+        .distinct()
+        .collect(
+            flatteningToImmutableListMultimap(
+                Function.identity(),
+                (domainName) -> {
+                  List<DomainBase> domains =
+                      ofy()
+                          .load()
+                          .type(DomainBase.class)
+                          .filter("fullyQualifiedDomainName", domainName)
+                          .list();
+                  domains.sort(Comparator.comparing(DomainBase::getCreationTime).reversed());
+                  checkState(
+                      !domains.isEmpty(),
+                      "Domain name %s had no associated DomainBase objects.",
+                      domainName);
+                  return domains.stream();
+                }));
+  }
+
+  /** Converts the previous {@link ThreatMatch} object to {@link Spec11ThreatMatch}. */
+  private Spec11ThreatMatch threatMatchToCloudSqlObject(
+      ThreatMatch threatMatch,
+      LocalDate date,
+      String registrarId,
+      ImmutableListMultimap<String, DomainBase> domainsByDomainName) {
+    DomainBase domain =
+        findDomainAsOfDateOrThrow(
+            threatMatch.fullyQualifiedDomainName(), date, domainsByDomainName);
+    return new Spec11ThreatMatch.Builder()
+        .setThreatTypes(ImmutableSet.of(ThreatType.valueOf(threatMatch.threatType())))
+        .setCheckDate(date)
+        .setRegistrarId(registrarId)
+        .setDomainName(threatMatch.fullyQualifiedDomainName())
+        .setDomainRepoId(domain.getRepoId())
+        .build();
+  }
+
+  /** Returns the DomainBase object as of the particular date, which is likely in the past. */
+  private DomainBase findDomainAsOfDateOrThrow(
+      String domainName,
+      LocalDate date,
+      ImmutableListMultimap<String, DomainBase> domainsByDomainName) {
+    ImmutableList<DomainBase> domains = domainsByDomainName.get(domainName);
+    for (DomainBase domain : domains) {
+      // We only know the date (not datetime) of the threat scan, so we approximate
+      LocalDate creationDate = domain.getCreationTime().toLocalDate();
+      LocalDate deletionDate = domain.getDeletionTime().toLocalDate();
+      if (!date.isBefore(creationDate) && !date.isAfter(deletionDate)) {
+        return domain;
+      }
+    }
+    throw new IllegalStateException(
+        String.format("Could not find a DomainBase valid for %s on day %s.", domainName, date));
+  }
+
+  /** Returns the list of dates between {@link #START_DATE} and now (UTC), inclusive. */
+  private ImmutableList<LocalDate> getDatesToBackfill() {
+    ImmutableSet<LocalDate> datesToSkip =
+        overrideExistingDates ? ImmutableSet.of() : getExistingDates();
+    ImmutableList.Builder<LocalDate> result = new ImmutableList.Builder<>();
+    LocalDate endDate = clock.nowUtc().toLocalDate();
+    for (LocalDate currentDate = START_DATE;
+        !currentDate.isAfter(endDate);
+        currentDate = currentDate.plusDays(1)) {
+      if (!datesToSkip.contains(currentDate)) {
+        result.add(currentDate);
+      }
+    }
+    return result.build();
+  }
+
+  private ImmutableSet<LocalDate> getExistingDates() {
+    return jpaTm()
+        .transact(
+            () ->
+                jpaTm()
+                    .getEntityManager()
+                    .createQuery(
+                        "SELECT DISTINCT stm.checkDate FROM Spec11ThreatMatch stm", LocalDate.class)
+                    .getResultStream()
+                    .collect(toImmutableSet()));
+  }
+}

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/flows/FlowTestCase.java

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

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

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

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

@@ -71,7 +71,6 @@ import google.registry.model.registry.Registry;
 import google.registry.model.registry.Registry.TldState;
 import google.registry.model.registry.label.ReservedList;
 import google.registry.model.reporting.HistoryEntry;
-import google.registry.persistence.DomainHistoryVKey;
 import org.joda.money.CurrencyUnit;
 import org.joda.money.Money;
 import org.joda.time.DateTime;
@@ -169,7 +168,7 @@ class DomainCheckFlowTest extends ResourceCheckFlowTestCase<DomainCheckFlow, Dom
         new AllocationToken.Builder()
             .setToken("abc123")
             .setTokenType(SINGLE_USE)
-            .setRedemptionHistoryEntry(DomainHistoryVKey.create(historyEntryKey))
+            .setRedemptionHistoryEntry(HistoryEntry.createVKey(historyEntryKey))
             .build());
     doCheckTest(
         create(false, "example1.tld", "In use"),

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

@@ -162,7 +162,6 @@ import google.registry.model.reporting.DomainTransactionRecord;
 import google.registry.model.reporting.DomainTransactionRecord.TransactionReportField;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.monitoring.whitebox.EppMetric;
-import google.registry.persistence.DomainHistoryVKey;
 import google.registry.testing.ReplayExtension;
 import google.registry.testing.TaskQueueHelper.TaskMatcher;
 import java.math.BigDecimal;
@@ -505,7 +504,7 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
         new AllocationToken.Builder()
             .setToken("abc123")
             .setTokenType(SINGLE_USE)
-            .setRedemptionHistoryEntry(DomainHistoryVKey.create(historyEntryKey))
+            .setRedemptionHistoryEntry(HistoryEntry.createVKey(historyEntryKey))
             .build());
     clock.advanceOneMilli();
     EppException thrown =
@@ -528,7 +527,7 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
     HistoryEntry historyEntry =
         ofy().load().type(HistoryEntry.class).ancestor(reloadResourceByForeignKey()).first().now();
     assertThat(ofy().load().entity(token).now().getRedemptionHistoryEntry())
-        .hasValue(DomainHistoryVKey.create(Key.create(historyEntry)));
+        .hasValue(HistoryEntry.createVKey(Key.create(historyEntry)));
   }
 
   @Test
@@ -1275,7 +1274,7 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
     assertThat(reloadedToken.isRedeemed()).isTrue();
     assertThat(reloadedToken.getRedemptionHistoryEntry())
         .hasValue(
-            DomainHistoryVKey.create(
+            HistoryEntry.createVKey(
                 Key.create(getHistoryEntries(reloadResourceByForeignKey()).get(0))));
   }
 

core/src/test/java/google/registry/flows/domain/token/AllocationTokenFlowUtilsTest.java

@@ -49,7 +49,6 @@ import google.registry.model.domain.token.AllocationToken;
 import google.registry.model.domain.token.AllocationToken.TokenStatus;
 import google.registry.model.registry.Registry;
 import google.registry.model.reporting.HistoryEntry;
-import google.registry.persistence.DomainHistoryVKey;
 import google.registry.testing.AppEngineExtension;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
@@ -198,7 +197,7 @@ class AllocationTokenFlowUtilsTest {
         new AllocationToken.Builder()
             .setToken("tokeN")
             .setTokenType(SINGLE_USE)
-            .setRedemptionHistoryEntry(DomainHistoryVKey.create(historyEntryKey))
+            .setRedemptionHistoryEntry(HistoryEntry.createVKey(historyEntryKey))
             .build());
     assertThat(
             flowUtils

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

@@ -24,16 +24,18 @@ import google.registry.flows.EppException;
 import google.registry.flows.ResourceCheckFlowTestCase;
 import google.registry.flows.exceptions.TooManyResourceChecksException;
 import google.registry.model.host.HostResource;
-import org.junit.jupiter.api.Test;
+import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.TestOfyAndSql;
 
 /** Unit tests for {@link HostCheckFlow}. */
+@DualDatabaseTest
 class HostCheckFlowTest extends ResourceCheckFlowTestCase<HostCheckFlow, HostResource> {
 
   HostCheckFlowTest() {
     setEppInput("host_check.xml");
   }
 
-  @Test
+  @TestOfyAndSql
   void testNothingExists() throws Exception {
     // These ids come from the check xml.
     doCheckTest(
@@ -42,7 +44,7 @@ class HostCheckFlowTest extends ResourceCheckFlowTestCase<HostCheckFlow, HostRes
         create(true, "ns3.example.tld", null));
   }
 
-  @Test
+  @TestOfyAndSql
   void testOneExists() throws Exception {
     persistActiveHost("ns1.example.tld");
     // These ids come from the check xml.
@@ -52,7 +54,7 @@ class HostCheckFlowTest extends ResourceCheckFlowTestCase<HostCheckFlow, HostRes
         create(true, "ns3.example.tld", null));
   }
 
-  @Test
+  @TestOfyAndSql
   void testOneExistsButWasDeleted() throws Exception {
     persistDeletedHost("ns1.example.tld", clock.nowUtc().minusDays(1));
     // These ids come from the check xml.
@@ -62,27 +64,27 @@ class HostCheckFlowTest extends ResourceCheckFlowTestCase<HostCheckFlow, HostRes
         create(true, "ns3.example.tld", null));
   }
 
-  @Test
+  @TestOfyAndSql
   void testXmlMatches() throws Exception {
     persistActiveHost("ns2.example.tld");
     runFlowAssertResponse(loadFile("host_check_response.xml"));
   }
 
-  @Test
+  @TestOfyAndSql
   void test50IdsAllowed() throws Exception {
     // Make sure we don't have a regression that reduces the number of allowed checks.
     setEppInput("host_check_50.xml");
     runFlow();
   }
 
-  @Test
+  @TestOfyAndSql
   void testTooManyIds() {
     setEppInput("host_check_51.xml");
     EppException thrown = assertThrows(TooManyResourceChecksException.class, this::runFlow);
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testIcannActivityReportField_getsLogged() throws Exception {
     runFlow();
     assertIcannReportingActivityFieldLogged("srs-host-check");

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

@@ -16,6 +16,7 @@ package google.registry.flows.host;
 
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.EppResourceUtils.loadByForeignKey;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatabaseHelper.assertNoBillingEvents;
 import static google.registry.testing.DatabaseHelper.createTld;
 import static google.registry.testing.DatabaseHelper.createTlds;
@@ -53,10 +54,12 @@ import google.registry.model.domain.DomainBase;
 import google.registry.model.eppcommon.StatusValue;
 import google.registry.model.host.HostResource;
 import google.registry.model.reporting.HistoryEntry;
+import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.TestOfyAndSql;
 import org.joda.time.DateTime;
-import org.junit.jupiter.api.Test;
 
 /** Unit tests for {@link HostCreateFlow}. */
+@DualDatabaseTest
 class HostCreateFlowTest extends ResourceFlowTestCase<HostCreateFlow, HostResource> {
 
   private void setEppHostCreateInput(String hostName, String hostAddrs) {
@@ -90,7 +93,9 @@ class HostCreateFlowTest extends ResourceFlowTestCase<HostCreateFlow, HostResour
         .hasOnlyOneHistoryEntryWhich()
         .hasType(HistoryEntry.Type.HOST_CREATE);
     assertNoBillingEvents();
-    assertEppResourceIndexEntityFor(reloadResourceByForeignKey());
+    if (tm().isOfy()) {
+      assertEppResourceIndexEntityFor(reloadResourceByForeignKey());
+    }
   }
 
   private void doSuccessfulInternalTest(String tld) throws Exception {
@@ -100,19 +105,19 @@ class HostCreateFlowTest extends ResourceFlowTestCase<HostCreateFlow, HostResour
     doSuccessfulTest();
   }
 
-  @Test
+  @TestOfyAndSql
   void testDryRun() throws Exception {
     dryRunFlowAssertResponse(loadFile("host_create_response.xml"));
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_externalNeverExisted() throws Exception {
     doSuccessfulTest();
     assertAboutHosts().that(reloadResourceByForeignKey()).hasSuperordinateDomain(null);
     assertNoDnsTasksEnqueued();
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_internalNeverExisted() throws Exception {
     doSuccessfulInternalTest("tld");
     HostResource host = reloadResourceByForeignKey();
@@ -123,7 +128,7 @@ class HostCreateFlowTest extends ResourceFlowTestCase<HostCreateFlow, HostResour
     assertDnsTasksEnqueued("ns1.example.tld");
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_multipartTLDsAndInvalidHost() {
     createTlds("bar.tld", "tld");
 
@@ -132,7 +137,7 @@ class HostCreateFlowTest extends ResourceFlowTestCase<HostCreateFlow, HostResour
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_externalExistedButWasDeleted() throws Exception {
     persistDeletedHost(getUniqueIdFromCommand(), clock.nowUtc().minusDays(1));
     doSuccessfulTest();
@@ -140,7 +145,7 @@ class HostCreateFlowTest extends ResourceFlowTestCase<HostCreateFlow, HostResour
     assertNoDnsTasksEnqueued();
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_internalExistedButWasDeleted() throws Exception {
     persistDeletedHost(getUniqueIdFromCommand(), clock.nowUtc().minusDays(1));
     doSuccessfulInternalTest("tld");
@@ -152,7 +157,7 @@ class HostCreateFlowTest extends ResourceFlowTestCase<HostCreateFlow, HostResour
     assertDnsTasksEnqueued("ns1.example.tld");
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_subordinateNeedsIps() {
     setEppHostCreateInput("ns1.example.tld", null);
     createTld("tld");
@@ -161,7 +166,7 @@ class HostCreateFlowTest extends ResourceFlowTestCase<HostCreateFlow, HostResour
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_externalMustNotHaveIps() {
     setEppHostCreateInputWithIps("ns1.example.external");
     createTld("tld");
@@ -170,7 +175,7 @@ class HostCreateFlowTest extends ResourceFlowTestCase<HostCreateFlow, HostResour
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_superordinateMissing() {
     setEppHostCreateInput("ns1.example.tld", null);
     createTld("tld");
@@ -179,7 +184,7 @@ class HostCreateFlowTest extends ResourceFlowTestCase<HostCreateFlow, HostResour
     assertThat(thrown).hasMessageThat().contains("(example.tld)");
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_superordinateInPendingDelete() {
     setEppHostCreateInputWithIps("ns1.example.tld");
     createTld("tld");
@@ -197,7 +202,7 @@ class HostCreateFlowTest extends ResourceFlowTestCase<HostCreateFlow, HostResour
         .contains("Superordinate domain for this hostname is in pending delete");
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_alreadyExists() throws Exception {
     setEppHostCreateInput("ns1.example.tld", null);
     persistActiveHost(getUniqueIdFromCommand());
@@ -209,7 +214,7 @@ class HostCreateFlowTest extends ResourceFlowTestCase<HostCreateFlow, HostResour
             String.format("Object with given ID (%s) already exists", getUniqueIdFromCommand()));
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_resourceContention() throws Exception {
     setEppHostCreateInput("ns1.example.tld", null);
     String targetId = getUniqueIdFromCommand();
@@ -226,14 +231,14 @@ class HostCreateFlowTest extends ResourceFlowTestCase<HostCreateFlow, HostResour
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_nonLowerCaseHostname() {
     setEppHostCreateInput("ns1.EXAMPLE.tld", null);
     EppException thrown = assertThrows(HostNameNotLowerCaseException.class, this::runFlow);
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_nonPunyCodedHostname() {
     setEppHostCreateInput("ns1.çauçalito.みんな", null);
     HostNameNotPunyCodedException thrown =
@@ -241,21 +246,21 @@ class HostCreateFlowTest extends ResourceFlowTestCase<HostCreateFlow, HostResour
     assertThat(thrown).hasMessageThat().contains("expected ns1.xn--aualito-txac.xn--q9jyb4c");
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_nonCanonicalHostname() {
     setEppHostCreateInput("ns1.example.tld.", null);
     EppException thrown = assertThrows(HostNameNotNormalizedException.class, this::runFlow);
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_longHostName() {
     setEppHostCreateInputWithIps("a" + Strings.repeat(".labelpart", 25) + ".tld");
     EppException thrown = assertThrows(HostNameTooLongException.class, this::runFlow);
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_ip4AddressWithIp6Declaration() {
     setEppHostCreateInput(
         "ns1.example.tld",
@@ -272,37 +277,37 @@ class HostCreateFlowTest extends ResourceFlowTestCase<HostCreateFlow, HostResour
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_badCharacter() {
     doFailingHostNameTest("foo bar", InvalidHostNameException.class);
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_tooShallowPublicSuffix() {
     doFailingHostNameTest("example.tld", HostNameTooShallowException.class);
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_tooShallowCcTld() {
     doFailingHostNameTest("foo.co.uk", HostNameTooShallowException.class);
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_barePublicSuffix() {
     doFailingHostNameTest("com", HostNameTooShallowException.class);
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_bareCcTld() {
     doFailingHostNameTest("co.uk", HostNameTooShallowException.class);
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_tooShallowNewTld() {
     doFailingHostNameTest("example.lol", HostNameTooShallowException.class);
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_ccTldInBailiwick() {
     createTld("co.uk");
     setEppHostCreateInputWithIps("foo.co.uk");
@@ -310,7 +315,7 @@ class HostCreateFlowTest extends ResourceFlowTestCase<HostCreateFlow, HostResour
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
-  @Test
+  @TestOfyAndSql
   void testIcannActivityReportField_getsLogged() throws Exception {
     runFlow();
     assertIcannReportingActivityFieldLogged("srs-host-create");

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

@@ -26,66 +26,68 @@ import google.registry.flows.host.HostFlowUtils.HostNameTooLongException;
 import google.registry.flows.host.HostFlowUtils.HostNameTooShallowException;
 import google.registry.flows.host.HostFlowUtils.InvalidHostNameException;
 import google.registry.testing.AppEngineExtension;
-import org.junit.jupiter.api.Test;
+import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.TestOfyAndSql;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Unit tests for {@link HostFlowUtils}. */
+@DualDatabaseTest
 class HostFlowUtilsTest {
 
   @RegisterExtension
   final AppEngineExtension appEngine =
       AppEngineExtension.builder().withDatastoreAndCloudSql().build();
 
-  @Test
+  @TestOfyAndSql
   void test_validExternalHostName_validates() throws Exception {
     assertThat(validateHostName("host.example.com").toString()).isEqualTo("host.example.com");
   }
 
-  @Test
+  @TestOfyAndSql
   void test_validExternalHostNameOnRegistrySuffixList_validates() throws Exception {
     assertThat(validateHostName("host.blogspot.com").toString()).isEqualTo("host.blogspot.com");
   }
 
-  @Test
+  @TestOfyAndSql
   void test_validExternalHostNameOnRegistrySuffixList_multipartTLD_validates() throws Exception {
     assertThat(validateHostName("ns1.host.co.uk").toString()).isEqualTo("ns1.host.co.uk");
   }
 
-  @Test
+  @TestOfyAndSql
   void test_validExternalHostNameOnRegistrySuffixList_multipartTLD_tooShallow() {
     assertThrows(
         HostNameTooShallowException.class, () -> validateHostName("host.co.uk").toString());
   }
 
-  @Test
+  @TestOfyAndSql
   void test_validateHostName_hostNameTooLong() {
     assertThrows(
         HostNameTooLongException.class,
         () -> validateHostName(Strings.repeat("na", 200) + ".wat.man"));
   }
 
-  @Test
+  @TestOfyAndSql
   void test_validateHostName_hostNameNotLowerCase() {
     assertThrows(HostNameNotLowerCaseException.class, () -> validateHostName("NA.CAPS.TLD"));
   }
 
-  @Test
+  @TestOfyAndSql
   void test_validateHostName_hostNameNotPunyCoded() {
     assertThrows(
         HostNameNotPunyCodedException.class, () -> validateHostName("motörhead.death.metal"));
   }
 
-  @Test
+  @TestOfyAndSql
   void test_validateHostName_hostNameNotNormalized() {
     assertThrows(HostNameNotNormalizedException.class, () -> validateHostName("root.node.yeah."));
   }
 
-  @Test
+  @TestOfyAndSql
   void test_validateHostName_hostNameHasLeadingHyphen() {
     assertThrows(InvalidHostNameException.class, () -> validateHostName("-giga.mega.tld"));
   }
 
-  @Test
+  @TestOfyAndSql
   void test_validateHostName_hostNameTooShallow() {
     assertThrows(HostNameTooShallowException.class, () -> validateHostName("domain.tld"));
   }

core/src/test/java/google/registry/model/EntityTestCase.java

@@ -28,6 +28,7 @@ import com.googlecode.objectify.annotation.Parent;
 import com.googlecode.objectify.annotation.Serialize;
 import com.googlecode.objectify.cmd.Query;
 import google.registry.model.ofy.Ofy;
+import google.registry.persistence.EppHistoryVKey;
 import google.registry.persistence.VKey;
 import google.registry.testing.AppEngineExtension;
 import google.registry.testing.FakeClock;
@@ -173,7 +174,8 @@ public abstract class EntityTestCase {
           }
           // Descend into persisted ImmutableObject classes, but not anything else.
           if (ImmutableObject.class.isAssignableFrom(fieldClass)
-              && !VKey.class.isAssignableFrom(fieldClass)) {
+              && !VKey.class.isAssignableFrom(fieldClass)
+              && !EppHistoryVKey.class.isAssignableFrom(fieldClass)) {
             getAllPotentiallyIndexedFieldPaths(fieldClass).stream()
                 .map(subfield -> field.getName() + "." + subfield)
                 .distinct()

core/src/test/java/google/registry/model/EppResourceUtilsTest.java

@@ -18,6 +18,7 @@ import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.EppResourceUtils.loadAtPointInTime;
 import static google.registry.testing.DatabaseHelper.createTld;
 import static google.registry.testing.DatabaseHelper.newHostResource;
+import static google.registry.testing.DatabaseHelper.persistNewRegistrars;
 import static google.registry.testing.DatabaseHelper.persistResource;
 import static google.registry.testing.DatabaseHelper.persistResourceWithCommitLog;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
@@ -26,16 +27,19 @@ import static org.joda.time.DateTimeZone.UTC;
 import google.registry.model.host.HostResource;
 import google.registry.model.ofy.Ofy;
 import google.registry.testing.AppEngineExtension;
+import google.registry.testing.DualDatabaseTest;
 import google.registry.testing.FakeClock;
 import google.registry.testing.InjectExtension;
+import google.registry.testing.TestOfyAndSql;
+import google.registry.testing.TestOfyOnly;
 import org.joda.time.DateTime;
 import org.joda.time.Duration;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Tests for {@link EppResourceUtils}. */
-public class EppResourceUtilsTest {
+@DualDatabaseTest
+class EppResourceUtilsTest {
 
   @RegisterExtension
   public final AppEngineExtension appEngine =
@@ -51,7 +55,7 @@ public class EppResourceUtilsTest {
     inject.setStaticField(Ofy.class, "clock", clock);
   }
 
-  @Test
+  @TestOfyAndSql
   void testLoadAtPointInTime_beforeCreated_returnsNull() {
     clock.advanceOneMilli();
     // Don't save a commit log, we shouldn't need one.
@@ -62,7 +66,7 @@ public class EppResourceUtilsTest {
     assertThat(loadAtPointInTime(host, clock.nowUtc().minus(Duration.millis(1))).now()).isNull();
   }
 
-  @Test
+  @TestOfyAndSql
   void testLoadAtPointInTime_atOrAfterLastAutoUpdateTime_returnsResource() {
     clock.advanceOneMilli();
     // Don't save a commit log, we shouldn't need one.
@@ -73,8 +77,9 @@ public class EppResourceUtilsTest {
     assertThat(loadAtPointInTime(host, clock.nowUtc()).now()).isEqualTo(host);
   }
 
-  @Test
+  @TestOfyOnly
   void testLoadAtPointInTime_usingIntactRevisionHistory_returnsMutationValue() {
+    persistNewRegistrars("OLD", "NEW");
     clock.advanceOneMilli();
     // Save resource with a commit log that we can read in later as a revisions map value.
     HostResource oldHost = persistResourceWithCommitLog(
@@ -94,7 +99,7 @@ public class EppResourceUtilsTest {
         .isEqualTo(oldHost);
   }
 
-  @Test
+  @TestOfyOnly
   void testLoadAtPointInTime_brokenRevisionHistory_returnsResourceAsIs() {
     // Don't save a commit log since we want to test the handling of a broken revisions key.
     HostResource oldHost = persistResource(
@@ -114,7 +119,7 @@ public class EppResourceUtilsTest {
     assertThat(loadAtPointInTime(host, clock.nowUtc().minusMillis(1)).now()).isEqualTo(host);
   }
 
-  @Test
+  @TestOfyOnly
   void testLoadAtPointInTime_fallback_returnsMutationValueForOldestRevision() {
     clock.advanceOneMilli();
     // Save a commit log that we can fall back to.
@@ -136,7 +141,7 @@ public class EppResourceUtilsTest {
         .isEqualTo(oldHost);
   }
 
-  @Test
+  @TestOfyOnly
   void testLoadAtPointInTime_ultimateFallback_onlyOneRevision_returnsCurrentResource() {
     clock.advanceOneMilli();
     // Don't save a commit log; we want to test that we load from the current resource.
@@ -151,7 +156,7 @@ public class EppResourceUtilsTest {
     assertThat(loadAtPointInTime(host, clock.nowUtc().minusMillis(1)).now()).isEqualTo(host);
   }
 
-  @Test
+  @TestOfyOnly
   void testLoadAtPointInTime_moreThanThirtyDaysInPast_historyIsPurged() {
     clock.advanceOneMilli();
     HostResource host =

core/src/test/java/google/registry/model/UpdateAutoTimestampTest.java

@@ -15,64 +15,106 @@
 package google.registry.model;
 
 import static com.google.common.truth.Truth.assertThat;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static org.joda.time.DateTimeZone.UTC;
 
+import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Entity;
+import com.googlecode.objectify.annotation.Ignore;
 import google.registry.model.common.CrossTldSingleton;
+import google.registry.persistence.VKey;
 import google.registry.schema.replay.EntityTest.EntityForTesting;
 import google.registry.testing.AppEngineExtension;
+import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.FakeClock;
+import google.registry.testing.TestOfyAndSql;
 import org.joda.time.DateTime;
-import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Unit tests for {@link UpdateAutoTimestamp}. */
+@DualDatabaseTest
 public class UpdateAutoTimestampTest {
 
+  FakeClock clock = new FakeClock();
+
   @RegisterExtension
   public final AppEngineExtension appEngine =
       AppEngineExtension.builder()
           .withDatastoreAndCloudSql()
+          .withJpaUnitTestEntities(UpdateAutoTimestampTestObject.class)
           .withOfyTestEntities(UpdateAutoTimestampTestObject.class)
+          .withClock(clock)
           .build();
 
   /** Timestamped class. */
   @Entity(name = "UatTestEntity")
+  @javax.persistence.Entity
   @EntityForTesting
   public static class UpdateAutoTimestampTestObject extends CrossTldSingleton {
+    @Ignore @javax.persistence.Id long id = SINGLETON_ID;
     UpdateAutoTimestamp updateTime = UpdateAutoTimestamp.create(null);
   }
 
   private UpdateAutoTimestampTestObject reload() {
-    return ofy().load().entity(new UpdateAutoTimestampTestObject()).now();
+    return tm().transact(
+            () ->
+                tm().load(
+                        VKey.create(
+                            UpdateAutoTimestampTestObject.class,
+                            1L,
+                            Key.create(new UpdateAutoTimestampTestObject()))));
   }
 
-  @Test
+  @TestOfyAndSql
   void testSaveSetsTime() {
     DateTime transactionTime =
         tm().transact(
                 () -> {
                   UpdateAutoTimestampTestObject object = new UpdateAutoTimestampTestObject();
                   assertThat(object.updateTime.timestamp).isNull();
-                  ofy().save().entity(object);
+                  tm().insert(object);
                   return tm().getTransactionTime();
                 });
-    ofy().clearSessionCache();
+    tm().clearSessionCache();
     assertThat(reload().updateTime.timestamp).isEqualTo(transactionTime);
   }
 
-  @Test
+  @TestOfyAndSql
+  void testDisabledUpdates() throws Exception {
+    DateTime initialTime =
+        tm().transact(
+                () -> {
+                  tm().insert(new UpdateAutoTimestampTestObject());
+                  return tm().getTransactionTime();
+                });
+
+    UpdateAutoTimestampTestObject object = reload();
+    clock.advanceOneMilli();
+
+    try (UpdateAutoTimestamp.DisableAutoUpdateResource disabler =
+        new UpdateAutoTimestamp.DisableAutoUpdateResource()) {
+      DateTime secondTransactionTime =
+          tm().transact(
+                  () -> {
+                    tm().put(object);
+                    return tm().getTransactionTime();
+                  });
+      assertThat(secondTransactionTime).isGreaterThan(initialTime);
+    }
+    assertThat(reload().updateTime.timestamp).isEqualTo(initialTime);
+  }
+
+  @TestOfyAndSql
   void testResavingOverwritesOriginalTime() {
     DateTime transactionTime =
         tm().transact(
                 () -> {
                   UpdateAutoTimestampTestObject object = new UpdateAutoTimestampTestObject();
                   object.updateTime = UpdateAutoTimestamp.create(DateTime.now(UTC).minusDays(1));
-                  ofy().save().entity(object);
+                  tm().insert(object);
                   return tm().getTransactionTime();
                 });
-    ofy().clearSessionCache();
+    tm().clearSessionCache();
     assertThat(reload().updateTime.timestamp).isEqualTo(transactionTime);
   }
 }

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

@@ -633,14 +633,13 @@ public class DomainBaseSqlTest {
                           "4-COM",
                           END_OF_TIME,
                           "registrar1",
-                          createLegacyVKey(
-                              BillingEvent.OneTime.class, oneTimeBillingEvent.getId())),
+                          oneTimeBillingEvent.createVKey()),
                       GracePeriod.createForRecurring(
                           GracePeriodStatus.AUTO_RENEW,
                           "4-COM",
                           END_OF_TIME,
                           "registrar1",
-                          createLegacyVKey(BillingEvent.Recurring.class, billEvent.getId())));
+                          billEvent.createVKey()));
 
               jpaTm().insert(contact);
               jpaTm().insert(contact2);

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

@@ -875,16 +875,13 @@ public class DomainBaseTest extends EntityTestCase {
     ImmutableSet<BillEventInfo> historyIds =
         domain.getGracePeriods().stream()
             .map(
-                gp ->
-                    new BillEventInfo(
-                        gp.getRecurringBillingEvent(), gp.billingEventRecurringHistoryId,
-                        gp.getOneTimeBillingEvent(), gp.billingEventOneTimeHistoryId))
+                gp -> new BillEventInfo(gp.getRecurringBillingEvent(), gp.getOneTimeBillingEvent()))
             .collect(toImmutableSet());
     assertThat(historyIds)
         .isEqualTo(
             ImmutableSet.of(
-                new BillEventInfo(null, null, oneTimeBillKey, historyEntryKey.getId()),
-                new BillEventInfo(recurringBillKey, historyEntryKey.getId(), null, null)));
+                new BillEventInfo(null, oneTimeBillKey),
+                new BillEventInfo(recurringBillKey, null)));
   }
 
   static class BillEventInfo extends ImmutableObject {
@@ -895,13 +892,9 @@ public class DomainBaseTest extends EntityTestCase {
 
     BillEventInfo(
         VKey<BillingEvent.Recurring> billingEventRecurring,
-        Long billingEventRecurringHistoryId,
-        VKey<BillingEvent.OneTime> billingEventOneTime,
-        Long billingEventOneTimeHistoryId) {
+        VKey<BillingEvent.OneTime> billingEventOneTime) {
       this.billingEventRecurring = billingEventRecurring;
-      this.billingEventRecurringHistoryId = billingEventRecurringHistoryId;
       this.billingEventOneTime = billingEventOneTime;
-      this.billingEventOneTimeHistoryId = billingEventOneTimeHistoryId;
     }
   }
 }

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

@@ -63,11 +63,11 @@ public class GracePeriodTest {
     recurringKey =
         VKey.create(
             Recurring.class,
-            12345,
+            12345L,
             Key.create(
                 Key.create(Key.create(DomainBase.class, "1-TEST"), HistoryEntry.class, 343L),
                 Recurring.class,
-                12345));
+                12345L));
   }
 
   @Test
@@ -80,8 +80,6 @@ public class GracePeriodTest {
     assertThat(gracePeriod.getClientId()).isEqualTo("TheRegistrar");
     assertThat(gracePeriod.getExpirationTime()).isEqualTo(now.plusDays(1));
     assertThat(gracePeriod.hasBillingEvent()).isTrue();
-    assertThat(gracePeriod.billingEventOneTimeHistoryId).isEqualTo(12345L);
-    assertThat(gracePeriod.billingEventRecurringHistoryId).isNull();
   }
 
   @Test
@@ -96,8 +94,6 @@ public class GracePeriodTest {
     assertThat(gracePeriod.getClientId()).isEqualTo("TheRegistrar");
     assertThat(gracePeriod.getExpirationTime()).isEqualTo(now.plusDays(1));
     assertThat(gracePeriod.hasBillingEvent()).isTrue();
-    assertThat(gracePeriod.billingEventOneTimeHistoryId).isNull();
-    assertThat(gracePeriod.billingEventRecurringHistoryId).isEqualTo(343L);
   }
 
   @Test

core/src/test/java/google/registry/model/domain/token/AllocationTokenTest.java

@@ -16,15 +16,14 @@ package google.registry.model.domain.token;
 
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth8.assertThat;
-import static google.registry.model.ImmutableObjectSubject.assertAboutImmutableObjects;
 import static google.registry.model.domain.token.AllocationToken.TokenStatus.CANCELLED;
 import static google.registry.model.domain.token.AllocationToken.TokenStatus.ENDED;
 import static google.registry.model.domain.token.AllocationToken.TokenStatus.NOT_STARTED;
 import static google.registry.model.domain.token.AllocationToken.TokenStatus.VALID;
 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.jpaTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 import static google.registry.testing.DatabaseHelper.createTld;
 import static google.registry.testing.DatabaseHelper.persistActiveDomain;
 import static google.registry.testing.DatabaseHelper.persistResource;
@@ -40,13 +39,14 @@ import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.token.AllocationToken.TokenStatus;
 import google.registry.model.domain.token.AllocationToken.TokenType;
 import google.registry.model.reporting.HistoryEntry;
-import google.registry.persistence.DomainHistoryVKey;
-import google.registry.persistence.VKey;
+import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.TestOfyAndSql;
+import google.registry.testing.TestOfyOnly;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
 
 /** Unit tests for {@link AllocationToken}. */
+@DualDatabaseTest
 public class AllocationTokenTest extends EntityTestCase {
 
   public AllocationTokenTest() {
@@ -58,7 +58,7 @@ public class AllocationTokenTest extends EntityTestCase {
     createTld("foo");
   }
 
-  @Test
+  @TestOfyAndSql
   void testPersistence() {
     AllocationToken unlimitedUseToken =
         persistResource(
@@ -78,7 +78,7 @@ public class AllocationTokenTest extends EntityTestCase {
                         .put(DateTime.now(UTC).plusWeeks(8), TokenStatus.ENDED)
                         .build())
                 .build());
-    assertThat(ofy().load().entity(unlimitedUseToken).now()).isEqualTo(unlimitedUseToken);
+    assertThat(transactIfJpaTm(() -> tm().load(unlimitedUseToken))).isEqualTo(unlimitedUseToken);
 
     DomainBase domain = persistActiveDomain("example.foo");
     Key<HistoryEntry> historyEntryKey = Key.create(Key.create(domain), HistoryEntry.class, 1);
@@ -86,37 +86,15 @@ public class AllocationTokenTest extends EntityTestCase {
         persistResource(
             new AllocationToken.Builder()
                 .setToken("abc123Single")
-                .setRedemptionHistoryEntry(DomainHistoryVKey.create(historyEntryKey))
+                .setRedemptionHistoryEntry(HistoryEntry.createVKey(historyEntryKey))
                 .setDomainName("example.foo")
                 .setCreationTimeForTest(DateTime.parse("2010-11-12T05:00:00Z"))
                 .setTokenType(SINGLE_USE)
                 .build());
-    assertThat(ofy().load().entity(singleUseToken).now()).isEqualTo(singleUseToken);
-
-    jpaTm()
-        .transact(
-            () -> {
-              jpaTm().insert(unlimitedUseToken);
-              jpaTm().insert(singleUseToken);
-            });
-    jpaTm()
-        .transact(
-            () -> {
-              assertAboutImmutableObjects()
-                  .that(jpaTm().load(VKey.createSql(AllocationToken.class, "abc123Unlimited")))
-                  .isEqualExceptFields(
-                      unlimitedUseToken,
-                      "creationTime",
-                      "updateTimestamp",
-                      "redemptionHistoryEntry");
-              assertAboutImmutableObjects()
-                  .that(jpaTm().load(VKey.createSql(AllocationToken.class, "abc123Single")))
-                  .isEqualExceptFields(
-                      singleUseToken, "creationTime", "updateTimestamp", "redemptionHistoryEntry");
-            });
+    assertThat(transactIfJpaTm(() -> tm().load(singleUseToken))).isEqualTo(singleUseToken);
   }
 
-  @Test
+  @TestOfyOnly
   void testIndexing() throws Exception {
     DomainBase domain = persistActiveDomain("blahdomain.foo");
     Key<HistoryEntry> historyEntryKey = Key.create(Key.create(domain), HistoryEntry.class, 1);
@@ -125,7 +103,7 @@ public class AllocationTokenTest extends EntityTestCase {
             new AllocationToken.Builder()
                 .setToken("abc123")
                 .setTokenType(SINGLE_USE)
-                .setRedemptionHistoryEntry(DomainHistoryVKey.create(historyEntryKey))
+                .setRedemptionHistoryEntry(HistoryEntry.createVKey(historyEntryKey))
                 .setDomainName("blahdomain.foo")
                 .setCreationTimeForTest(DateTime.parse("2010-11-12T05:00:00Z"))
                 .build()),
@@ -134,7 +112,7 @@ public class AllocationTokenTest extends EntityTestCase {
         "domainName");
   }
 
-  @Test
+  @TestOfyAndSql
   void testCreationTime_autoPopulates() {
     AllocationToken tokenBeforePersisting =
         new AllocationToken.Builder().setToken("abc123").setTokenType(SINGLE_USE).build();
@@ -143,7 +121,7 @@ public class AllocationTokenTest extends EntityTestCase {
     assertThat(tokenAfterPersisting.getCreationTime()).hasValue(fakeClock.nowUtc());
   }
 
-  @Test
+  @TestOfyAndSql
   void testSetCreationTime_cantCallMoreThanOnce() {
     AllocationToken.Builder builder =
         new AllocationToken.Builder()
@@ -157,7 +135,7 @@ public class AllocationTokenTest extends EntityTestCase {
     assertThat(thrown).hasMessageThat().isEqualTo("Creation time can only be set once");
   }
 
-  @Test
+  @TestOfyAndSql
   void testSetToken_cantCallMoreThanOnce() {
     AllocationToken.Builder builder = new AllocationToken.Builder().setToken("foobar");
     IllegalStateException thrown =
@@ -165,7 +143,7 @@ public class AllocationTokenTest extends EntityTestCase {
     assertThat(thrown).hasMessageThat().isEqualTo("Token can only be set once");
   }
 
-  @Test
+  @TestOfyAndSql
   void testSetTokenType_cantCallMoreThanOnce() {
     AllocationToken.Builder builder =
         new AllocationToken.Builder().setTokenType(TokenType.UNLIMITED_USE);
@@ -174,7 +152,7 @@ public class AllocationTokenTest extends EntityTestCase {
     assertThat(thrown).hasMessageThat().isEqualTo("Token type can only be set once");
   }
 
-  @Test
+  @TestOfyAndSql
   void testBuild_DomainNameWithLessThanTwoParts() {
     IllegalArgumentException thrown =
         assertThrows(
@@ -192,7 +170,7 @@ public class AllocationTokenTest extends EntityTestCase {
     assertThat(thrown).hasMessageThat().isEqualTo("Invalid domain name: example");
   }
 
-  @Test
+  @TestOfyAndSql
   void testBuild_invalidTld() {
     IllegalArgumentException thrown =
         assertThrows(
@@ -210,7 +188,7 @@ public class AllocationTokenTest extends EntityTestCase {
     assertThat(thrown).hasMessageThat().isEqualTo("Invalid domain name: example.nosuchtld");
   }
 
-  @Test
+  @TestOfyAndSql
   void testBuild_domainNameOnlyOnSingleUse() {
     AllocationToken.Builder builder =
         new AllocationToken.Builder()
@@ -223,7 +201,7 @@ public class AllocationTokenTest extends EntityTestCase {
         .isEqualTo("Domain name can only be specified for SINGLE_USE tokens");
   }
 
-  @Test
+  @TestOfyAndSql
   void testBuild_redemptionHistoryEntryOnlyInSingleUse() {
     DomainBase domain = persistActiveDomain("blahdomain.foo");
     Key<HistoryEntry> historyEntryKey = Key.create(Key.create(domain), HistoryEntry.class, 1);
@@ -231,14 +209,14 @@ public class AllocationTokenTest extends EntityTestCase {
         new AllocationToken.Builder()
             .setToken("foobar")
             .setTokenType(TokenType.UNLIMITED_USE)
-            .setRedemptionHistoryEntry(DomainHistoryVKey.create(historyEntryKey));
+            .setRedemptionHistoryEntry(HistoryEntry.createVKey(historyEntryKey));
     IllegalArgumentException thrown = assertThrows(IllegalArgumentException.class, builder::build);
     assertThat(thrown)
         .hasMessageThat()
         .isEqualTo("Redemption history entry can only be specified for SINGLE_USE tokens");
   }
 
-  @Test
+  @TestOfyAndSql
   void testSetTransitions_notStartOfTime() {
     IllegalArgumentException thrown =
         assertThrows(
@@ -256,7 +234,7 @@ public class AllocationTokenTest extends EntityTestCase {
         .isEqualTo("tokenStatusTransitions map must start at START_OF_TIME.");
   }
 
-  @Test
+  @TestOfyAndSql
   void testSetTransitions_badInitialValue() {
     IllegalArgumentException thrown =
         assertThrows(
@@ -273,14 +251,14 @@ public class AllocationTokenTest extends EntityTestCase {
         .isEqualTo("tokenStatusTransitions must start with NOT_STARTED");
   }
 
-  @Test
+  @TestOfyAndSql
   void testSetTransitions_invalidInitialTransitions() {
     // NOT_STARTED can only go to VALID or CANCELLED
     assertBadInitialTransition(NOT_STARTED);
     assertBadInitialTransition(ENDED);
   }
 
-  @Test
+  @TestOfyAndSql
   void testSetTransitions_badTransitionsFromValid() {
     // VALID can only go to ENDED or CANCELLED
     assertBadTransition(
@@ -301,14 +279,14 @@ public class AllocationTokenTest extends EntityTestCase {
         NOT_STARTED);
   }
 
-  @Test
+  @TestOfyAndSql
   void testSetTransitions_terminalTransitions() {
     // both ENDED and CANCELLED are terminal
     assertTerminal(ENDED);
     assertTerminal(CANCELLED);
   }
 
-  @Test
+  @TestOfyAndSql
   void testSetDiscountFractionTooHigh() {
     IllegalArgumentException thrown =
         assertThrows(
@@ -319,7 +297,7 @@ public class AllocationTokenTest extends EntityTestCase {
         .isEqualTo("Discount fraction must be between 0 and 1 inclusive");
   }
 
-  @Test
+  @TestOfyAndSql
   void testSetDiscountFractionTooLow() {
     IllegalArgumentException thrown =
         assertThrows(
@@ -330,7 +308,7 @@ public class AllocationTokenTest extends EntityTestCase {
         .isEqualTo("Discount fraction must be between 0 and 1 inclusive");
   }
 
-  @Test
+  @TestOfyAndSql
   void testSetDiscountYearsTooHigh() {
     IllegalArgumentException thrown =
         assertThrows(
@@ -341,7 +319,7 @@ public class AllocationTokenTest extends EntityTestCase {
         .isEqualTo("Discount years must be between 1 and 10 inclusive");
   }
 
-  @Test
+  @TestOfyAndSql
   void testSetDiscountYearsTooLow() {
     IllegalArgumentException thrown =
         assertThrows(
@@ -352,7 +330,7 @@ public class AllocationTokenTest extends EntityTestCase {
         .isEqualTo("Discount years must be between 1 and 10 inclusive");
   }
 
-  @Test
+  @TestOfyAndSql
   void testBuild_noTokenType() {
     IllegalArgumentException thrown =
         assertThrows(
@@ -361,7 +339,7 @@ public class AllocationTokenTest extends EntityTestCase {
     assertThat(thrown).hasMessageThat().isEqualTo("Token type must be specified");
   }
 
-  @Test
+  @TestOfyAndSql
   void testBuild_noToken() {
     IllegalArgumentException thrown =
         assertThrows(
@@ -370,7 +348,7 @@ public class AllocationTokenTest extends EntityTestCase {
     assertThat(thrown).hasMessageThat().isEqualTo("Token must not be null or empty");
   }
 
-  @Test
+  @TestOfyAndSql
   void testBuild_emptyToken() {
     IllegalArgumentException thrown =
         assertThrows(
@@ -379,7 +357,7 @@ public class AllocationTokenTest extends EntityTestCase {
     assertThat(thrown).hasMessageThat().isEqualTo("Token must not be blank");
   }
 
-  @Test
+  @TestOfyAndSql
   void testBuild_discountPremiumsRequiresDiscountFraction() {
     IllegalArgumentException thrown =
         assertThrows(
@@ -395,7 +373,7 @@ public class AllocationTokenTest extends EntityTestCase {
         .isEqualTo("Discount premiums can only be specified along with a discount fraction");
   }
 
-  @Test
+  @TestOfyAndSql
   void testBuild_discountYearsRequiresDiscountFraction() {
     IllegalArgumentException thrown =
         assertThrows(

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

@@ -17,13 +17,17 @@ package google.registry.model.host;
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth8.assertThat;
 import static google.registry.model.EppResourceUtils.loadByForeignKey;
+import static google.registry.model.ImmutableObjectSubject.immutableObjectCorrespondence;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatabaseHelper.cloneAndSetAutoTimestamps;
 import static google.registry.testing.DatabaseHelper.createTld;
 import static google.registry.testing.DatabaseHelper.newDomainBase;
+import static google.registry.testing.DatabaseHelper.persistNewRegistrars;
 import static google.registry.testing.DatabaseHelper.persistResource;
 import static google.registry.testing.HostResourceSubject.assertAboutHosts;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.net.InetAddresses;
 import google.registry.model.EntityTestCase;
@@ -32,11 +36,14 @@ import google.registry.model.eppcommon.StatusValue;
 import google.registry.model.eppcommon.Trid;
 import google.registry.model.transfer.DomainTransferData;
 import google.registry.model.transfer.TransferStatus;
+import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.TestOfyAndSql;
+import google.registry.testing.TestOfyOnly;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
 
 /** Unit tests for {@link HostResource}. */
+@DualDatabaseTest
 class HostResourceTest extends EntityTestCase {
 
   private final DateTime day3 = fakeClock.nowUtc();
@@ -49,6 +56,7 @@ class HostResourceTest extends EntityTestCase {
   @BeforeEach
   void setUp() {
     createTld("com");
+    persistNewRegistrars("gaining", "losing", "thisRegistrar", "thatRegistrar");
     // Set up a new persisted registrar entity.
     domain =
         persistResource(
@@ -71,9 +79,9 @@ class HostResourceTest extends EntityTestCase {
                 new HostResource.Builder()
                     .setRepoId("DEADBEEF-COM")
                     .setHostName("ns1.example.com")
-                    .setCreationClientId("a registrar")
+                    .setCreationClientId("thisRegistrar")
                     .setLastEppUpdateTime(fakeClock.nowUtc())
-                    .setLastEppUpdateClientId("another registrar")
+                    .setLastEppUpdateClientId("thatRegistrar")
                     .setLastTransferTime(fakeClock.nowUtc())
                     .setInetAddresses(ImmutableSet.of(InetAddresses.forString("127.0.0.1")))
                     .setStatusValues(ImmutableSet.of(StatusValue.OK))
@@ -81,13 +89,22 @@ class HostResourceTest extends EntityTestCase {
                     .build()));
   }
 
-  @Test
+  @TestOfyAndSql
   void testPersistence() {
+    HostResource newHost = host.asBuilder().setRepoId("NEWHOST").build();
+    tm().transact(() -> tm().insert(newHost));
+    assertThat(ImmutableList.of(tm().transact(() -> tm().load(newHost.createVKey()))))
+        .comparingElementsUsing(immutableObjectCorrespondence("revisions"))
+        .containsExactly(newHost);
+  }
+
+  @TestOfyOnly
+  void testLoadingByForeignKey() {
     assertThat(loadByForeignKey(HostResource.class, host.getForeignKey(), fakeClock.nowUtc()))
         .hasValue(host);
   }
 
-  @Test
+  @TestOfyOnly
   void testIndexing() throws Exception {
     // Clone it and save it before running the indexing test so that its transferData fields are
     // populated from the superordinate domain.
@@ -100,7 +117,7 @@ class HostResourceTest extends EntityTestCase {
         "currentSponsorClientId");
   }
 
-  @Test
+  @TestOfyAndSql
   void testEmptyStringsBecomeNull() {
     assertThat(
             new HostResource.Builder()
@@ -122,7 +139,7 @@ class HostResourceTest extends EntityTestCase {
         .isNotNull();
   }
 
-  @Test
+  @TestOfyAndSql
   void testEmptySetsBecomeNull() {
     assertThat(new HostResource.Builder().setInetAddresses(null).build().inetAddresses).isNull();
     assertThat(new HostResource.Builder().setInetAddresses(ImmutableSet.of()).build().inetAddresses)
@@ -135,7 +152,7 @@ class HostResourceTest extends EntityTestCase {
         .isNotNull();
   }
 
-  @Test
+  @TestOfyAndSql
   void testImplicitStatusValues() {
     // OK is implicit if there's no other statuses.
     assertAboutHosts()
@@ -157,13 +174,13 @@ class HostResourceTest extends EntityTestCase {
         .hasExactlyStatusValues(StatusValue.CLIENT_HOLD);
   }
 
-  @Test
+  @TestOfyAndSql
   void testToHydratedString_notCircular() {
     // If there are circular references, this will overflow the stack.
     host.toHydratedString();
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_uppercaseHostName() {
     IllegalArgumentException thrown =
         assertThrows(
@@ -173,7 +190,7 @@ class HostResourceTest extends EntityTestCase {
         .contains("Host name must be in puny-coded, lower-case form");
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_utf8HostName() {
     IllegalArgumentException thrown =
         assertThrows(
@@ -183,14 +200,14 @@ class HostResourceTest extends EntityTestCase {
         .contains("Host name must be in puny-coded, lower-case form");
   }
 
-  @Test
+  @TestOfyAndSql
   void testComputeLastTransferTime_hostNeverSwitchedDomains_domainWasNeverTransferred() {
     domain = domain.asBuilder().setLastTransferTime(null).build();
     host = host.asBuilder().setLastTransferTime(null).setLastSuperordinateChange(null).build();
     assertThat(host.computeLastTransferTime(domain)).isNull();
   }
 
-  @Test
+  @TestOfyAndSql
   void testComputeLastTransferTime_hostNeverSwitchedDomains_domainWasTransferred() {
     // Host was created on Day 1.
     // Domain was transferred on Day 2.
@@ -205,7 +222,7 @@ class HostResourceTest extends EntityTestCase {
     assertThat(host.computeLastTransferTime(domain)).isEqualTo(day2);
   }
 
-  @Test
+  @TestOfyAndSql
   void testComputeLastTransferTime_hostCreatedAfterDomainWasTransferred() {
     // Domain was transferred on Day 1.
     // Host was created subordinate to domain on Day 2.
@@ -217,9 +234,9 @@ class HostResourceTest extends EntityTestCase {
                     .setCreationTime(day2)
                     .setRepoId("DEADBEEF-COM")
                     .setHostName("ns1.example.com")
-                    .setCreationClientId("a registrar")
+                    .setCreationClientId("thisRegistrar")
                     .setLastEppUpdateTime(fakeClock.nowUtc())
-                    .setLastEppUpdateClientId("another registrar")
+                    .setLastEppUpdateClientId("thatRegistrar")
                     .setInetAddresses(ImmutableSet.of(InetAddresses.forString("127.0.0.1")))
                     .setStatusValues(ImmutableSet.of(StatusValue.OK))
                     .setSuperordinateDomain(domain.createVKey())
@@ -227,7 +244,7 @@ class HostResourceTest extends EntityTestCase {
     assertThat(host.computeLastTransferTime(domain)).isNull();
   }
 
-  @Test
+  @TestOfyAndSql
   void testComputeLastTransferTime_hostWasTransferred_domainWasNeverTransferred() {
     // Host was transferred on Day 1.
     // Host was made subordinate to domain on Day 2.
@@ -237,7 +254,7 @@ class HostResourceTest extends EntityTestCase {
     assertThat(host.computeLastTransferTime(domain)).isEqualTo(day1);
   }
 
-  @Test
+  @TestOfyAndSql
   void testComputeLastTransferTime_domainWasTransferredBeforeHostBecameSubordinate() {
     // Host was transferred on Day 1.
     // Domain was transferred on Day 2.
@@ -247,7 +264,7 @@ class HostResourceTest extends EntityTestCase {
     assertThat(host.computeLastTransferTime(domain)).isEqualTo(day1);
   }
 
-  @Test
+  @TestOfyAndSql
   void testComputeLastTransferTime_domainWasTransferredAfterHostBecameSubordinate() {
     // Host was transferred on Day 1.
     // Host was made subordinate to domain on Day 2.

core/src/test/java/google/registry/model/index/ForeignKeyIndexTest.java

@@ -20,6 +20,7 @@ import static google.registry.model.EppResourceUtils.loadByForeignKey;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatabaseHelper.createTld;
 import static google.registry.testing.DatabaseHelper.deleteResource;
+import static google.registry.testing.DatabaseHelper.newDomainBase;
 import static google.registry.testing.DatabaseHelper.persistActiveContact;
 import static google.registry.testing.DatabaseHelper.persistActiveHost;
 import static google.registry.testing.DatabaseHelper.persistDeletedHost;
@@ -29,16 +30,21 @@ import static google.registry.util.DateTimeUtils.END_OF_TIME;
 import com.google.common.collect.ImmutableList;
 import google.registry.model.EntityTestCase;
 import google.registry.model.contact.ContactResource;
+import google.registry.model.domain.DomainBase;
 import google.registry.model.host.HostResource;
 import google.registry.model.index.ForeignKeyIndex.ForeignKeyHostIndex;
+import google.registry.testing.DualDatabaseTest;
 import google.registry.testing.TestCacheExtension;
+import google.registry.testing.TestOfyAndSql;
+import google.registry.testing.TestOfyOnly;
+import google.registry.testing.TestSqlOnly;
 import org.joda.time.Duration;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Unit tests for {@link ForeignKeyIndex}. */
-public class ForeignKeyIndexTest extends EntityTestCase {
+@DualDatabaseTest
+class ForeignKeyIndexTest extends EntityTestCase {
 
   @RegisterExtension
   public final TestCacheExtension testCacheExtension =
@@ -49,7 +55,17 @@ public class ForeignKeyIndexTest extends EntityTestCase {
     createTld("com");
   }
 
-  @Test
+  @TestSqlOnly
+  void testModifyForeignKeyIndex_notThrowExceptionInSql() {
+    DomainBase domainBase = newDomainBase("test.com");
+    ForeignKeyIndex<DomainBase> fki = ForeignKeyIndex.create(domainBase, fakeClock.nowUtc());
+    tm().transact(() -> tm().insert(fki));
+    tm().transact(() -> tm().put(fki));
+    tm().transact(() -> tm().delete(fki));
+    tm().transact(() -> tm().update(fki));
+  }
+
+  @TestOfyOnly
   void testPersistence() {
     // Persist a host and implicitly persist a ForeignKeyIndex for it.
     HostResource host = persistActiveHost("ns1.example.com");
@@ -59,7 +75,7 @@ public class ForeignKeyIndexTest extends EntityTestCase {
     assertThat(fki.getDeletionTime()).isEqualTo(END_OF_TIME);
   }
 
-  @Test
+  @TestOfyOnly
   void testIndexing() throws Exception {
     // Persist a host and implicitly persist a ForeignKeyIndex for it.
     persistActiveHost("ns1.example.com");
@@ -68,38 +84,50 @@ public class ForeignKeyIndexTest extends EntityTestCase {
         "deletionTime");
   }
 
-  @Test
+  @TestOfyAndSql
   void testLoadForNonexistentForeignKey_returnsNull() {
     assertThat(ForeignKeyIndex.load(HostResource.class, "ns1.example.com", fakeClock.nowUtc()))
         .isNull();
   }
 
-  @Test
+  @TestOfyAndSql
   void testLoadForDeletedForeignKey_returnsNull() {
     HostResource host = persistActiveHost("ns1.example.com");
-    persistResource(ForeignKeyIndex.create(host, fakeClock.nowUtc().minusDays(1)));
+    if (tm().isOfy()) {
+      persistResource(ForeignKeyIndex.create(host, fakeClock.nowUtc().minusDays(1)));
+    } else {
+      persistResource(host.asBuilder().setDeletionTime(fakeClock.nowUtc().minusDays(1)).build());
+    }
     assertThat(ForeignKeyIndex.load(HostResource.class, "ns1.example.com", fakeClock.nowUtc()))
         .isNull();
   }
 
-  @Test
+  @TestOfyAndSql
   void testLoad_newerKeyHasBeenSoftDeleted() {
     HostResource host1 = persistActiveHost("ns1.example.com");
     fakeClock.advanceOneMilli();
-    ForeignKeyHostIndex fki = new ForeignKeyHostIndex();
-    fki.foreignKey = "ns1.example.com";
-    fki.topReference = host1.createVKey();
-    fki.deletionTime = fakeClock.nowUtc();
-    persistResource(fki);
+    if (tm().isOfy()) {
+      ForeignKeyHostIndex fki = new ForeignKeyHostIndex();
+      fki.foreignKey = "ns1.example.com";
+      fki.topReference = host1.createVKey();
+      fki.deletionTime = fakeClock.nowUtc();
+      persistResource(fki);
+    } else {
+      persistResource(host1.asBuilder().setDeletionTime(fakeClock.nowUtc()).build());
+    }
     assertThat(ForeignKeyIndex.load(HostResource.class, "ns1.example.com", fakeClock.nowUtc()))
         .isNull();
   }
 
-  @Test
+  @TestOfyAndSql
   void testBatchLoad_skipsDeletedAndNonexistent() {
     persistActiveHost("ns1.example.com");
     HostResource host = persistActiveHost("ns2.example.com");
-    persistResource(ForeignKeyIndex.create(host, fakeClock.nowUtc().minusDays(1)));
+    if (tm().isOfy()) {
+      persistResource(ForeignKeyIndex.create(host, fakeClock.nowUtc().minusDays(1)));
+    } else {
+      persistResource(host.asBuilder().setDeletionTime(fakeClock.nowUtc().minusDays(1)).build());
+    }
     assertThat(
             ForeignKeyIndex.load(
                     HostResource.class,
@@ -109,7 +137,7 @@ public class ForeignKeyIndexTest extends EntityTestCase {
         .containsExactly("ns1.example.com");
   }
 
-  @Test
+  @TestOfyAndSql
   void testDeadCodeThatDeletedScrapCommandsReference() {
     persistActiveHost("omg");
     assertThat(ForeignKeyIndex.load(HostResource.class, "omg", fakeClock.nowUtc()).getForeignKey())
@@ -124,7 +152,7 @@ public class ForeignKeyIndexTest extends EntityTestCase {
     return ForeignKeyIndex.load(ContactResource.class, contactId, fakeClock.nowUtc());
   }
 
-  @Test
+  @TestOfyOnly
   void test_loadCached_cachesNonexistenceOfHosts() {
     assertThat(
             ForeignKeyIndex.loadCached(
@@ -144,7 +172,7 @@ public class ForeignKeyIndexTest extends EntityTestCase {
         .containsExactly("ns4.example.com", loadHostFki("ns4.example.com"));
   }
 
-  @Test
+  @TestOfyOnly
   void test_loadCached_cachesExistenceOfHosts() {
     HostResource host1 = persistActiveHost("ns1.example.com");
     HostResource host2 = persistActiveHost("ns2.example.com");
@@ -172,7 +200,7 @@ public class ForeignKeyIndexTest extends EntityTestCase {
             "ns3.example.com", loadHostFki("ns3.example.com"));
   }
 
-  @Test
+  @TestOfyOnly
   void test_loadCached_doesntSeeHostChangesWhileCacheIsValid() {
     HostResource originalHost = persistActiveHost("ns1.example.com");
     ForeignKeyIndex<HostResource> originalFki = loadHostFki("ns1.example.com");
@@ -195,7 +223,7 @@ public class ForeignKeyIndexTest extends EntityTestCase {
         .containsExactly("ns1.example.com", originalFki);
   }
 
-  @Test
+  @TestOfyOnly
   void test_loadCached_filtersOutSoftDeletedHosts() {
     persistActiveHost("ns1.example.com");
     persistDeletedHost("ns2.example.com", fakeClock.nowUtc().minusDays(1));
@@ -207,7 +235,7 @@ public class ForeignKeyIndexTest extends EntityTestCase {
         .containsExactly("ns1.example.com", loadHostFki("ns1.example.com"));
   }
 
-  @Test
+  @TestOfyOnly
   void test_loadCached_cachesContactFkis() {
     persistActiveContact("contactid1");
     ForeignKeyIndex<ContactResource> fki1 = loadContactFki("contactid1");

core/src/test/java/google/registry/model/smd/SignedMarkRevocationListDaoTest.java

@@ -50,21 +50,6 @@ public class SignedMarkRevocationListDaoTest {
     assertAboutImmutableObjects().that(fromDb).isEqualExceptFields(list);
   }
 
-  @Test
-  void trySave_failureIsSwallowed() {
-    SignedMarkRevocationList list =
-        SignedMarkRevocationList.create(
-            fakeClock.nowUtc(), ImmutableMap.of("mark", fakeClock.nowUtc().minusHours(1)));
-    SignedMarkRevocationListDao.trySave(list);
-    SignedMarkRevocationList fromDb = SignedMarkRevocationListDao.getLatestRevision().get();
-    assertAboutImmutableObjects().that(fromDb).isEqualExceptFields(list);
-
-    // This should throw an exception, which is swallowed and nothing changed
-    SignedMarkRevocationListDao.trySave(list);
-    SignedMarkRevocationList secondFromDb = SignedMarkRevocationListDao.getLatestRevision().get();
-    assertAboutImmutableObjects().that(secondFromDb).isEqualExceptFields(fromDb);
-  }
-
   @Test
   void testRetrieval_notPresent() {
     assertThat(SignedMarkRevocationListDao.getLatestRevision().isPresent()).isFalse();

core/src/test/java/google/registry/model/smd/SignedMarkRevocationListTest.java

@@ -18,6 +18,7 @@ import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.ImmutableObjectSubject.assertAboutImmutableObjects;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.smd.SignedMarkRevocationList.SHARD_SIZE;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
 import static org.joda.time.Duration.standardDays;
@@ -138,6 +139,33 @@ public class SignedMarkRevocationListTest {
         .isEqualTo(DateTime.parse("2000-01-01T00:00:00Z"));
   }
 
+  @Test
+  void test_getCreationTime_missingInCloudSQL() {
+    clock.setTo(DateTime.parse("2000-01-01T00:00:00Z"));
+    createSaveGetHelper(1);
+    jpaTm().transact(() -> jpaTm().delete(SignedMarkRevocationListDao.getLatestRevision().get()));
+    RuntimeException thrown =
+        assertThrows(RuntimeException.class, () -> SignedMarkRevocationList.get());
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo("Signed mark revocation list in Cloud SQL is empty.");
+  }
+
+  @Test
+  void test_getCreationTime_unequalListsInDatabases() {
+    clock.setTo(DateTime.parse("2000-01-01T00:00:00Z"));
+    createSaveGetHelper(1);
+    ImmutableMap.Builder<String, DateTime> revokes = new ImmutableMap.Builder<>();
+    for (int i = 0; i < 3; i++) {
+      revokes.put(Integer.toString(i), clock.nowUtc());
+    }
+    SignedMarkRevocationListDao.trySave(
+        SignedMarkRevocationList.create(clock.nowUtc(), revokes.build()));
+    RuntimeException thrown =
+        assertThrows(RuntimeException.class, () -> SignedMarkRevocationList.get());
+    assertThat(thrown).hasMessageThat().contains("Unequal SM revocation lists detected:");
+  }
+
   @Test
   void test_isSmdRevoked_present() {
     final int rows = SHARD_SIZE + 1;

core/src/test/java/google/registry/model/translators/EppHistoryVKeyTranslatorFactoryTest.java

@@ -0,0 +1,36 @@
+// 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 google.registry.model.translators.EppHistoryVKeyTranslatorFactory.kindPathToVKeyClass;
+import static org.junit.jupiter.api.Assertions.fail;
+
+import org.junit.jupiter.api.Test;
+
+/** Unit test for {@link EppHistoryVKeyTranslatorFactory}. */
+class EppHistoryVKeyTranslatorFactoryTest {
+
+  @Test
+  void assertAllVKeyClassesHavingCreateFromOfyKeyMethod() {
+    kindPathToVKeyClass.forEach(
+        (kindPath, vKeyClass) -> {
+          try {
+            vKeyClass.getDeclaredMethod("create", com.googlecode.objectify.Key.class);
+          } catch (NoSuchMethodException e) {
+            fail("Missing static method create(com.googlecode.objectify.Key) on " + vKeyClass, e);
+          }
+        });
+  }
+}

core/src/test/java/google/registry/persistence/BillingVKeyTest.java

@@ -0,0 +1,116 @@
+// 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 static com.google.common.truth.Truth.assertThat;
+import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+
+import com.googlecode.objectify.Key;
+import com.googlecode.objectify.annotation.Entity;
+import com.googlecode.objectify.annotation.Id;
+import com.googlecode.objectify.annotation.Parent;
+import google.registry.model.ImmutableObject;
+import google.registry.model.billing.BillingEvent;
+import google.registry.model.common.EntityGroupRoot;
+import google.registry.model.domain.DomainBase;
+import google.registry.model.reporting.HistoryEntry;
+import google.registry.persistence.BillingVKey.BillingEventVKey;
+import google.registry.persistence.BillingVKey.BillingRecurrenceVKey;
+import google.registry.schema.replay.EntityTest.EntityForTesting;
+import google.registry.testing.AppEngineExtension;
+import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.TestOfyAndSql;
+import javax.persistence.Transient;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+/** Unit test for {@link BillingVKey}. */
+@DualDatabaseTest
+class BillingVKeyTest {
+  @RegisterExtension
+  final AppEngineExtension appEngine =
+      AppEngineExtension.builder()
+          .withDatastoreAndCloudSql()
+          .withOfyTestEntities(BillingVKeyTestEntity.class)
+          .withJpaUnitTestEntities(BillingVKeyTestEntity.class)
+          .build();
+
+  @TestOfyAndSql
+  void testRestoreSymmetricVKey() {
+    Key<HistoryEntry> domainHistoryKey =
+        Key.create(Key.create(DomainBase.class, "domainRepoId"), HistoryEntry.class, 10L);
+
+    Key<BillingEvent.OneTime> oneTimeOfyKey =
+        Key.create(domainHistoryKey, BillingEvent.OneTime.class, 100L);
+    VKey<BillingEvent.OneTime> oneTimeVKey =
+        VKey.create(BillingEvent.OneTime.class, 100L, oneTimeOfyKey);
+
+    Key<BillingEvent.Recurring> recurringOfyKey =
+        Key.create(domainHistoryKey, BillingEvent.Recurring.class, 200L);
+    VKey<BillingEvent.Recurring> recurringVKey =
+        VKey.create(BillingEvent.Recurring.class, 200L, recurringOfyKey);
+
+    BillingVKeyTestEntity original = new BillingVKeyTestEntity(oneTimeVKey, recurringVKey);
+    tm().transact(() -> tm().insert(original));
+    BillingVKeyTestEntity persisted = tm().transact(() -> tm().load(original.createVKey()));
+
+    assertThat(persisted).isEqualTo(original);
+    assertThat(persisted.getBillingEventVKey()).isEqualTo(oneTimeVKey);
+    assertThat(persisted.getBillingRecurrenceVKey()).isEqualTo(recurringVKey);
+  }
+
+  @TestOfyAndSql
+  void testHandleNullVKeyCorrectly() {
+    BillingVKeyTestEntity original = new BillingVKeyTestEntity(null, null);
+    tm().transact(() -> tm().insert(original));
+    BillingVKeyTestEntity persisted = tm().transact(() -> tm().load(original.createVKey()));
+
+    assertThat(persisted).isEqualTo(original);
+  }
+
+  @EntityForTesting
+  @Entity
+  @javax.persistence.Entity
+  private static class BillingVKeyTestEntity extends ImmutableObject {
+    @Transient @Parent Key<EntityGroupRoot> parent = getCrossTldKey();
+
+    @Id @javax.persistence.Id String id = "id";
+
+    BillingEventVKey billingEventVKey;
+
+    BillingRecurrenceVKey billingRecurrenceVKey;
+
+    BillingVKeyTestEntity() {}
+
+    BillingVKeyTestEntity(
+        VKey<BillingEvent.OneTime> onetime, VKey<BillingEvent.Recurring> recurring) {
+      this.billingEventVKey = BillingEventVKey.create(onetime);
+      this.billingRecurrenceVKey = BillingRecurrenceVKey.create(recurring);
+    }
+
+    VKey<BillingEvent.OneTime> getBillingEventVKey() {
+      return billingEventVKey.createVKey();
+    }
+
+    VKey<BillingEvent.Recurring> getBillingRecurrenceVKey() {
+      return billingRecurrenceVKey.createVKey();
+    }
+
+    VKey<BillingVKeyTestEntity> createVKey() {
+      return VKey.create(
+          BillingVKeyTestEntity.class, id, Key.create(parent, BillingVKeyTestEntity.class, id));
+    }
+  }
+}

core/src/test/java/google/registry/persistence/DomainHistoryVKeyTest.java

@@ -56,12 +56,17 @@ class DomainHistoryVKeyTest {
     TestEntity persisted = tm().transact(() -> tm().load(original.createVKey()));
     assertThat(persisted).isEqualTo(original);
     // Double check that the persisted.domainHistoryVKey is a symmetric VKey
-    assertThat(persisted.domainHistoryVKey.getKind()).isEqualTo(HistoryEntry.class);
-    assertThat(persisted.domainHistoryVKey.getOfyKey())
+    assertThat(persisted.domainHistoryVKey.createOfyKey())
         .isEqualTo(
             Key.create(Key.create(DomainBase.class, "domainRepoId"), HistoryEntry.class, 10L));
-    assertThat(persisted.domainHistoryVKey.getSqlKey())
+    assertThat(persisted.domainHistoryVKey.createSqlKey())
         .isEqualTo(new DomainHistoryId("domainRepoId", 10L));
+    assertThat(persisted.domainHistoryVKey.createVKey())
+        .isEqualTo(
+            VKey.create(
+                HistoryEntry.class,
+                new DomainHistoryId("domainRepoId", 10L),
+                Key.create(Key.create(DomainBase.class, "domainRepoId"), HistoryEntry.class, 10L)));
   }
 
   @TestOfyAndSql
@@ -69,9 +74,12 @@ class DomainHistoryVKeyTest {
     Key<HistoryEntry> ofyKey =
         Key.create(Key.create(DomainBase.class, "domainRepoId"), HistoryEntry.class, 10L);
     DomainHistoryVKey domainHistoryVKey = DomainHistoryVKey.create(ofyKey);
-    assertThat(domainHistoryVKey.getKind()).isEqualTo(HistoryEntry.class);
-    assertThat(domainHistoryVKey.getOfyKey()).isEqualTo(ofyKey);
-    assertThat(domainHistoryVKey.getSqlKey()).isEqualTo(new DomainHistoryId("domainRepoId", 10L));
+    assertThat(domainHistoryVKey.createOfyKey()).isEqualTo(ofyKey);
+    assertThat(domainHistoryVKey.createSqlKey())
+        .isEqualTo(new DomainHistoryId("domainRepoId", 10L));
+    assertThat(domainHistoryVKey.createVKey())
+        .isEqualTo(
+            VKey.create(HistoryEntry.class, new DomainHistoryId("domainRepoId", 10L), ofyKey));
   }
 
   @EntityForTesting

core/src/test/java/google/registry/privileges/secretmanager/SecretManagerClientTest.java

@@ -17,6 +17,8 @@ package google.registry.privileges.secretmanager;
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
+import com.google.auth.oauth2.GoogleCredentials;
+import com.google.cloud.secretmanager.v1.SecretManagerServiceSettings;
 import com.google.cloud.secretmanager.v1.SecretVersion.State;
 import google.registry.privileges.secretmanager.SecretManagerClient.NoSuchSecretResourceException;
 import google.registry.privileges.secretmanager.SecretManagerClient.SecretAlreadyExistsException;
@@ -54,11 +56,14 @@ public class SecretManagerClientTest {
   private String secretId;
 
   @BeforeAll
-  static void beforeAll() {
+  static void beforeAll() throws IOException {
     String environmentName = System.getProperty("test.gcp_integration.env");
     if (environmentName != null) {
       secretManagerClient =
           SecretManagerModule.provideSecretManagerClient(
+              SecretManagerServiceSettings.newBuilder()
+                  .setCredentialsProvider(() -> GoogleCredentials.getApplicationDefault())
+                  .build(),
               String.format("domain-registry-%s", environmentName),
               new Retrier(new SystemSleeper(), 1));
       isUnitTest = false;

core/src/test/java/google/registry/reporting/spec11/Spec11RegistrarThreatMatchesParserTest.java

@@ -34,7 +34,7 @@ import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
 /** Unit tests for {@link Spec11RegistrarThreatMatchesParser}. */
-class Spec11RegistrarThreatMatchesParserTest {
+public class Spec11RegistrarThreatMatchesParserTest {
 
   private static final String TODAY = "2018-07-21";
   private static final String YESTERDAY = "2018-07-20";
@@ -66,6 +66,12 @@ class Spec11RegistrarThreatMatchesParserTest {
     assertThat(parser.getPreviousDateWithMatches(LocalDate.parse(TODAY))).isEmpty();
   }
 
+  @Test
+  void testNonexistent_returnsEmpty() throws Exception {
+    assertThat(parser.getRegistrarThreatMatches(LocalDate.parse(YESTERDAY).minusYears(1)))
+        .isEmpty();
+  }
+
   @Test
   void testFindPrevious_olderThanYesterdayFound() throws Exception {
     setupFile("spec11_fake_report_previous_day", "2018-07-14");
@@ -95,7 +101,7 @@ class Spec11RegistrarThreatMatchesParserTest {
   }
 
   /** The expected contents of the sample spec11 report file */
-  static ImmutableSet<RegistrarThreatMatches> sampleThreatMatches() throws Exception {
+  public static ImmutableSet<RegistrarThreatMatches> sampleThreatMatches() throws Exception {
     return ImmutableSet.of(getMatchA(), getMatchB());
   }
 

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

@@ -454,10 +454,6 @@ public final class AppEngineExtension implements BeforeEachCallback, AfterEachCa
           .setEnvIsAdmin(userInfo.isAdmin());
     }
 
-    if (clock != null) {
-      helper.setClock(() -> clock.nowUtc().getMillis());
-    }
-
     if (withLocalModules) {
       helper.setEnvInstance("0");
     }

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

@@ -750,6 +750,20 @@ public class DatabaseHelper {
             .build());
   }
 
+  /** Persists and returns a {@link Registrar} with the specified registrarId. */
+  public static Registrar persistNewRegistrar(String registrarId) {
+    return persistNewRegistrar(registrarId, registrarId + " name", Registrar.Type.REAL, 100L);
+  }
+
+  /** Persists and returns a list of {@link Registrar}s with the specified registrarIds. */
+  public static ImmutableList<Registrar> persistNewRegistrars(String... registrarIds) {
+    ImmutableList.Builder<Registrar> newRegistrars = new ImmutableList.Builder<>();
+    for (String registrarId : registrarIds) {
+      newRegistrars.add(persistNewRegistrar(registrarId));
+    }
+    return newRegistrars.build();
+  }
+
   private static Iterable<BillingEvent> getBillingEvents() {
     return transactIfJpaTm(
         () ->

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

@@ -18,6 +18,7 @@ import static com.google.common.collect.ImmutableList.toImmutableList;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import google.registry.persistence.transaction.TransactionManager;
 import google.registry.persistence.transaction.TransactionManagerFactory;
 import java.lang.reflect.Field;
@@ -160,6 +161,14 @@ class DualDatabaseTestInvocationContextProvider implements TestTemplateInvocatio
 
   private static boolean isDualDatabaseTest(ExtensionContext context) {
     Object testInstance = context.getTestInstance().orElseThrow(RuntimeException::new);
-    return testInstance.getClass().isAnnotationPresent(DualDatabaseTest.class);
+    // If the test method is declared in its parent class,
+    // e.g. google.registry.flows.ResourceFlowTestCase.testRequiresLogin,
+    // we don't consider it is a DualDatabaseTest. This is because there may exist some subclasses
+    // that have not been migrated to DualDatabaseTest.
+    boolean isDeclaredTestMethod =
+        ImmutableSet.copyOf(testInstance.getClass().getDeclaredMethods())
+            .contains(context.getTestMethod().orElseThrow(RuntimeException::new));
+    return testInstance.getClass().isAnnotationPresent(DualDatabaseTest.class)
+        && isDeclaredTestMethod;
   }
 }

core/src/test/java/google/registry/tools/CommandTestCase.java

@@ -17,7 +17,8 @@ package google.registry.tools;
 import static com.google.common.collect.Iterables.concat;
 import static com.google.common.collect.Iterables.toArray;
 import static com.google.common.truth.Truth.assertThat;
-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 java.nio.charset.StandardCharsets.UTF_8;
 import static org.joda.time.DateTimeZone.UTC;
 
@@ -116,7 +117,7 @@ public abstract class CommandTestCase<C extends Command> {
     } finally {
       // Clear the session cache so that subsequent reads for verification purposes hit Datastore.
       // This primarily matters for AutoTimestamp fields, which otherwise won't have updated values.
-      ofy().clearSessionCache();
+      tm().clearSessionCache();
       // Reset back to UNITTEST environment.
       RegistryToolEnvironment.UNITTEST.setup(systemPropertyExtension);
     }
@@ -185,12 +186,12 @@ public abstract class CommandTestCase<C extends Command> {
 
   /** Reloads the given resource from Datastore. */
   <T> T reloadResource(T resource) {
-    return ofy().load().entity(resource).now();
+    return transactIfJpaTm(() -> tm().load(resource));
   }
 
   /** Returns count of all poll messages in Datastore. */
   int getPollMessageCount() {
-    return ofy().load().type(PollMessage.class).count();
+    return transactIfJpaTm(() -> tm().loadAll(PollMessage.class).size());
   }
 
   /**

core/src/test/java/google/registry/tools/DeleteAllocationTokensCommandTest.java

@@ -16,24 +16,28 @@ package google.registry.tools;
 
 import static com.google.common.truth.Truth.assertThat;
 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 static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 import static google.registry.testing.DatabaseHelper.createTlds;
 import static google.registry.testing.DatabaseHelper.persistActiveDomain;
 import static google.registry.testing.DatabaseHelper.persistResource;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.token.AllocationToken;
 import google.registry.model.domain.token.AllocationToken.TokenType;
 import google.registry.model.reporting.HistoryEntry;
-import google.registry.persistence.DomainHistoryVKey;
-import java.util.Collection;
+import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.TestOfyAndSql;
+import java.util.Arrays;
 import javax.annotation.Nullable;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
 
 /** Unit tests for {@link DeleteAllocationTokensCommand}. */
+@DualDatabaseTest
 class DeleteAllocationTokensCommandTest extends CommandTestCase<DeleteAllocationTokensCommand> {
 
   private AllocationToken preRed1;
@@ -54,38 +58,38 @@ class DeleteAllocationTokensCommandTest extends CommandTestCase<DeleteAllocation
     othrNot = persistToken("asdgfho7HASDS", null, false);
   }
 
-  @Test
+  @TestOfyAndSql
   void test_deleteOnlyUnredeemedTokensWithPrefix() throws Exception {
     runCommandForced("--prefix", "prefix");
-    assertThat(reloadTokens(preNot1, preNot2)).isEmpty();
+    assertNonexistent(preNot1, preNot2);
     assertThat(reloadTokens(preRed1, preRed2, othrRed, othrNot))
         .containsExactly(preRed1, preRed2, othrRed, othrNot);
   }
 
-  @Test
+  @TestOfyAndSql
   void test_deleteSingleAllocationToken() throws Exception {
     runCommandForced("--prefix", "asdgfho7HASDS");
-    assertThat(reloadTokens(othrNot)).isEmpty();
+    assertNonexistent(othrNot);
     assertThat(reloadTokens(preRed1, preRed2, preNot1, preNot2, othrRed))
         .containsExactly(preRed1, preRed2, preNot1, preNot2, othrRed);
   }
 
-  @Test
+  @TestOfyAndSql
   void test_deleteParticularTokens() throws Exception {
     runCommandForced("--tokens", "prefix2978204,asdgfho7HASDS");
-    assertThat(reloadTokens(preNot1, othrNot)).isEmpty();
+    assertNonexistent(preNot1, othrNot);
     assertThat(reloadTokens(preRed1, preRed2, preNot2, othrRed))
         .containsExactly(preRed1, preRed2, preNot2, othrRed);
   }
 
-  @Test
+  @TestOfyAndSql
   void test_deleteTokensWithNonExistentPrefix_doesNothing() throws Exception {
     runCommandForced("--prefix", "nonexistent");
     assertThat(reloadTokens(preRed1, preRed2, preNot1, preNot2, othrRed, othrNot))
         .containsExactly(preRed1, preRed2, preNot1, preNot2, othrRed, othrNot);
   }
 
-  @Test
+  @TestOfyAndSql
   void test_dryRun_deletesNothing() throws Exception {
     runCommandForced("--prefix", "prefix", "--dry_run");
     assertThat(reloadTokens(preRed1, preRed2, preNot1, preNot2, othrRed, othrNot))
@@ -93,27 +97,27 @@ class DeleteAllocationTokensCommandTest extends CommandTestCase<DeleteAllocation
     assertInStdout("Would delete tokens: prefix2978204, prefix8ZZZhs8");
   }
 
-  @Test
+  @TestOfyAndSql
   void test_defaultOptions_doesntDeletePerDomainTokens() throws Exception {
     AllocationToken preDom1 = persistToken("prefixasdfg897as", "foo.bar", false);
     AllocationToken preDom2 = persistToken("prefix98HAZXadbn", "foo.bar", true);
     runCommandForced("--prefix", "prefix");
-    assertThat(reloadTokens(preNot1, preNot2)).isEmpty();
+    assertNonexistent(preNot1, preNot2);
     assertThat(reloadTokens(preRed1, preRed2, preDom1, preDom2, othrRed, othrNot))
         .containsExactly(preRed1, preRed2, preDom1, preDom2, othrRed, othrNot);
   }
 
-  @Test
+  @TestOfyAndSql
   void test_withDomains_doesDeletePerDomainTokens() throws Exception {
     AllocationToken preDom1 = persistToken("prefixasdfg897as", "foo.bar", false);
     AllocationToken preDom2 = persistToken("prefix98HAZXadbn", "foo.bar", true);
     runCommandForced("--prefix", "prefix", "--with_domains");
-    assertThat(reloadTokens(preNot1, preNot2, preDom1)).isEmpty();
+    assertNonexistent(preNot1, preNot2, preDom1);
     assertThat(reloadTokens(preRed1, preRed2, preDom2, othrRed, othrNot))
         .containsExactly(preRed1, preRed2, preDom2, othrRed, othrNot);
   }
 
-  @Test
+  @TestOfyAndSql
   void testSkipUnlimitedUseTokens() throws Exception {
     AllocationToken unlimitedUseToken =
         persistResource(
@@ -125,17 +129,18 @@ class DeleteAllocationTokensCommandTest extends CommandTestCase<DeleteAllocation
     assertThat(reloadTokens(unlimitedUseToken)).containsExactly(unlimitedUseToken);
   }
 
-  @Test
+  @TestOfyAndSql
   void test_batching() throws Exception {
     for (int i = 0; i < 50; i++) {
       persistToken(String.format("batch%2d", i), null, i % 2 == 0);
     }
-    assertThat(ofy().load().type(AllocationToken.class).count()).isEqualTo(56);
+    assertThat(transactIfJpaTm(() -> tm().loadAll(AllocationToken.class).size())).isEqualTo(56);
     runCommandForced("--prefix", "batch");
-    assertThat(ofy().load().type(AllocationToken.class).count()).isEqualTo(56 - 25);
+    assertThat(transactIfJpaTm(() -> tm().loadAll(AllocationToken.class).size()))
+        .isEqualTo(56 - 25);
   }
 
-  @Test
+  @TestOfyAndSql
   void test_prefixIsRequired() {
     IllegalArgumentException thrown =
         assertThrows(IllegalArgumentException.class, this::runCommandForced);
@@ -144,7 +149,7 @@ class DeleteAllocationTokensCommandTest extends CommandTestCase<DeleteAllocation
         .isEqualTo("Must provide one of --tokens or --prefix, not both / neither");
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_bothPrefixAndTokens() {
     IllegalArgumentException thrown =
         assertThrows(
@@ -155,7 +160,7 @@ class DeleteAllocationTokensCommandTest extends CommandTestCase<DeleteAllocation
         .isEqualTo("Must provide one of --tokens or --prefix, not both / neither");
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_emptyPrefix() {
     IllegalArgumentException thrown =
         assertThrows(IllegalArgumentException.class, () -> runCommandForced("--prefix", ""));
@@ -173,12 +178,16 @@ class DeleteAllocationTokensCommandTest extends CommandTestCase<DeleteAllocation
       String domainToPersist = domainName != null ? domainName : "example.foo";
       DomainBase domain = persistActiveDomain(domainToPersist);
       Key<HistoryEntry> historyEntryKey = Key.create(Key.create(domain), HistoryEntry.class, 1051L);
-      builder.setRedemptionHistoryEntry(DomainHistoryVKey.create(historyEntryKey));
+      builder.setRedemptionHistoryEntry(HistoryEntry.createVKey(historyEntryKey));
     }
     return persistResource(builder.build());
   }
 
-  private static Collection<AllocationToken> reloadTokens(AllocationToken... tokens) {
-    return ofy().load().entities(tokens).values();
+  private static ImmutableList<AllocationToken> reloadTokens(AllocationToken... tokens) {
+    return transactIfJpaTm(() -> tm().loadAll(ImmutableSet.copyOf(tokens)));
+  }
+
+  private static void assertNonexistent(AllocationToken... tokens) {
+    Arrays.stream(tokens).forEach(t -> transactIfJpaTm(() -> assertThat(tm().exists(t)).isFalse()));
   }
 }

core/src/test/java/google/registry/tools/GenerateAllocationTokensCommandTest.java

@@ -17,7 +17,8 @@ package google.registry.tools;
 import static com.google.common.truth.Truth.assertThat;
 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.testing.DatabaseHelper.assertAllocationTokens;
 import static google.registry.testing.DatabaseHelper.createTld;
 import static google.registry.testing.DatabaseHelper.persistResource;
@@ -39,11 +40,14 @@ import com.google.common.collect.Iterables;
 import com.google.common.io.Files;
 import google.registry.model.domain.token.AllocationToken;
 import google.registry.model.domain.token.AllocationToken.TokenStatus;
-import google.registry.persistence.DomainHistoryVKey;
+import google.registry.model.reporting.HistoryEntry;
+import google.registry.persistence.VKey;
 import google.registry.testing.DeterministicStringGenerator;
 import google.registry.testing.DeterministicStringGenerator.Rule;
+import google.registry.testing.DualDatabaseTest;
 import google.registry.testing.FakeClock;
 import google.registry.testing.FakeSleeper;
+import google.registry.testing.TestOfyAndSql;
 import google.registry.util.Retrier;
 import google.registry.util.StringGenerator.Alphabets;
 import java.io.File;
@@ -51,10 +55,10 @@ import java.util.Collection;
 import javax.annotation.Nullable;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
 import org.mockito.ArgumentMatchers;
 
 /** Unit tests for {@link GenerateAllocationTokensCommand}. */
+@DualDatabaseTest
 class GenerateAllocationTokensCommandTest extends CommandTestCase<GenerateAllocationTokensCommand> {
 
   @BeforeEach
@@ -64,14 +68,14 @@ class GenerateAllocationTokensCommandTest extends CommandTestCase<GenerateAlloca
         new Retrier(new FakeSleeper(new FakeClock(DateTime.parse("2000-01-01TZ"))), 3);
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_oneToken() throws Exception {
     runCommand("--prefix", "blah", "--number", "1", "--length", "9");
     assertAllocationTokens(createToken("blah123456789", null, null));
     assertInStdout("blah123456789");
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_threeTokens() throws Exception {
     runCommand("--prefix", "foo", "--number", "3", "--length", "10");
     assertAllocationTokens(
@@ -81,14 +85,14 @@ class GenerateAllocationTokensCommandTest extends CommandTestCase<GenerateAlloca
     assertInStdout("foo123456789A\nfooBCDEFGHJKL\nfooMNPQRSTUVW");
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_defaults() throws Exception {
     runCommand("--number", "1");
     assertAllocationTokens(createToken("123456789ABCDEFG", null, null));
     assertInStdout("123456789ABCDEFG");
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_retry() throws Exception {
     command = spy(command);
     RemoteApiException fakeException = new RemoteApiException("foo", "foo", "foo", new Exception());
@@ -103,7 +107,7 @@ class GenerateAllocationTokensCommandTest extends CommandTestCase<GenerateAlloca
     verify(command, times(3)).saveTokens(ArgumentMatchers.any());
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_tokenCollision() throws Exception {
     AllocationToken existingToken =
         persistResource(
@@ -116,24 +120,24 @@ class GenerateAllocationTokensCommandTest extends CommandTestCase<GenerateAlloca
     assertInStdout("DEADBEEFDEFGHJKLMNPQ");
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_dryRun_outputsButDoesntSave() throws Exception {
     runCommand("--prefix", "foo", "--number", "2", "--length", "10", "--dry_run");
     assertAllocationTokens();
     assertInStdout("foo123456789A\nfooBCDEFGHJKL");
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_largeNumberOfTokens() throws Exception {
     command.stringGenerator =
         new DeterministicStringGenerator(Alphabets.BASE_58, Rule.PREPEND_COUNTER);
     runCommand("--prefix", "ooo", "--number", "100", "--length", "16");
     // The deterministic string generator makes it too much hassle to assert about each token, so
     // just assert total number.
-    assertThat(ofy().load().type(AllocationToken.class).count()).isEqualTo(100);
+    assertThat(transactIfJpaTm(() -> tm().loadAll(AllocationToken.class).size())).isEqualTo(100);
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_domainNames() throws Exception {
     createTld("tld");
     File domainNamesFile = tmpDir.resolve("domain_names.txt").toFile();
@@ -147,7 +151,7 @@ class GenerateAllocationTokensCommandTest extends CommandTestCase<GenerateAlloca
         "foo1.tld,123456789ABCDEFG\nboo2.tld,HJKLMNPQRSTUVWXY\nbaz9.tld,Zabcdefghijkmnop");
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_promotionToken() throws Exception {
     DateTime promoStart = DateTime.now(UTC);
     DateTime promoEnd = promoStart.plusMonths(1);
@@ -181,24 +185,24 @@ class GenerateAllocationTokensCommandTest extends CommandTestCase<GenerateAlloca
             .build());
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_specifyTokens() throws Exception {
     runCommand("--tokens", "foobar,foobaz");
     assertAllocationTokens(createToken("foobar", null, null), createToken("foobaz", null, null));
     assertInStdout("foobar", "foobaz");
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_specifyManyTokens() throws Exception {
     command.stringGenerator =
         new DeterministicStringGenerator(Alphabets.BASE_58, Rule.PREPEND_COUNTER);
     Collection<String> sampleTokens = command.stringGenerator.createStrings(13, 100);
     runCommand("--tokens", Joiner.on(",").join(sampleTokens));
     assertInStdout(Iterables.toArray(sampleTokens, String.class));
-    assertThat(ofy().load().type(AllocationToken.class).count()).isEqualTo(100);
+    assertThat(transactIfJpaTm(() -> tm().loadAll(AllocationToken.class).size())).isEqualTo(100);
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_mustSpecifyNumberOfTokensOrDomainsFile() {
     IllegalArgumentException thrown =
         assertThrows(IllegalArgumentException.class, () -> runCommand("--prefix", "FEET"));
@@ -207,7 +211,7 @@ class GenerateAllocationTokensCommandTest extends CommandTestCase<GenerateAlloca
         .isEqualTo("Must specify exactly one of '--number', '--domain_names_file', and '--tokens'");
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_mustNotSpecifyBothNumberOfTokensAndDomainsFile() {
     IllegalArgumentException thrown =
         assertThrows(
@@ -222,7 +226,7 @@ class GenerateAllocationTokensCommandTest extends CommandTestCase<GenerateAlloca
         .isEqualTo("Must specify exactly one of '--number', '--domain_names_file', and '--tokens'");
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_mustNotSpecifyBothNumberOfTokensAndTokenStrings() {
     IllegalArgumentException thrown =
         assertThrows(
@@ -237,7 +241,7 @@ class GenerateAllocationTokensCommandTest extends CommandTestCase<GenerateAlloca
         .isEqualTo("Must specify exactly one of '--number', '--domain_names_file', and '--tokens'");
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_mustNotSpecifyBothTokenStringsAndDomainsFile() {
     IllegalArgumentException thrown =
         assertThrows(
@@ -252,7 +256,7 @@ class GenerateAllocationTokensCommandTest extends CommandTestCase<GenerateAlloca
         .isEqualTo("Must specify exactly one of '--number', '--domain_names_file', and '--tokens'");
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_specifiesAlreadyExistingToken() throws Exception {
     runCommand("--tokens", "foobar");
     beforeEachCommandTestCase(); // reset the command variables
@@ -263,7 +267,7 @@ class GenerateAllocationTokensCommandTest extends CommandTestCase<GenerateAlloca
         .isEqualTo("Cannot create specified tokens; the following tokens already exist: [foobar]");
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_invalidTokenType() {
     ParameterException thrown =
         assertThrows(
@@ -274,7 +278,7 @@ class GenerateAllocationTokensCommandTest extends CommandTestCase<GenerateAlloca
         .isEqualTo("Invalid value for -t parameter. Allowed values:[SINGLE_USE, UNLIMITED_USE]");
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_invalidTokenStatusTransition() {
     assertThat(
             assertThrows(
@@ -289,7 +293,7 @@ class GenerateAllocationTokensCommandTest extends CommandTestCase<GenerateAlloca
         .isInstanceOf(IllegalArgumentException.class);
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_lengthOfZero() {
     IllegalArgumentException thrown =
         assertThrows(
@@ -301,7 +305,7 @@ class GenerateAllocationTokensCommandTest extends CommandTestCase<GenerateAlloca
             "Token length should not be 0. To generate exact tokens, use the --tokens parameter.");
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_unlimitedUseMustHaveTransitions() {
     assertThat(
             assertThrows(
@@ -313,7 +317,7 @@ class GenerateAllocationTokensCommandTest extends CommandTestCase<GenerateAlloca
 
   private AllocationToken createToken(
       String token,
-      @Nullable DomainHistoryVKey redemptionHistoryEntry,
+      @Nullable VKey<? extends HistoryEntry> redemptionHistoryEntry,
       @Nullable String domainName) {
     AllocationToken.Builder builder =
         new AllocationToken.Builder().setToken(token).setTokenType(SINGLE_USE);

core/src/test/java/google/registry/tools/GetAllocationTokenCommandTest.java

@@ -28,7 +28,7 @@ import com.google.common.collect.ImmutableList;
 import com.googlecode.objectify.Key;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.token.AllocationToken;
-import google.registry.persistence.DomainHistoryVKey;
+import google.registry.model.reporting.HistoryEntry;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.Test;
 
@@ -85,7 +85,7 @@ class GetAllocationTokenCommandTest extends CommandTestCase<GetAllocationTokenCo
                 .setTokenType(SINGLE_USE)
                 .setDomainName("fqqdn.tld")
                 .setRedemptionHistoryEntry(
-                    DomainHistoryVKey.create(Key.create(createHistoryEntryForEppResource(domain))))
+                    HistoryEntry.createVKey(Key.create(createHistoryEntryForEppResource(domain))))
                 .build());
     runCommand("foo");
     assertInStdout(

core/src/test/java/google/registry/tools/SetSqlReplayCheckpointCommandTest.java

@@ -0,0 +1,42 @@
+// 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.tools;
+
+import static com.google.common.truth.Truth.assertThat;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.util.DateTimeUtils.START_OF_TIME;
+import static org.junit.Assert.assertThrows;
+
+import google.registry.schema.replay.SqlReplayCheckpoint;
+import org.joda.time.DateTime;
+import org.junit.jupiter.api.Test;
+
+class SetSqlReplayCheckpointCommandTest extends CommandTestCase<SetSqlReplayCheckpointCommand> {
+
+  @Test
+  void testSuccess() throws Exception {
+    assertThat(jpaTm().transact(SqlReplayCheckpoint::get)).isEqualTo(START_OF_TIME);
+    DateTime timeToSet = DateTime.parse("2000-06-06T22:00:00.0Z");
+    runCommandForced(timeToSet.toString());
+    assertThat(jpaTm().transact(SqlReplayCheckpoint::get)).isEqualTo(timeToSet);
+  }
+
+  @Test
+  void testFailure_multipleParams() throws Exception {
+    DateTime one = DateTime.parse("2000-06-06T22:00:00.0Z");
+    DateTime two = DateTime.parse("2001-06-06T22:00:00.0Z");
+    assertThrows(IllegalArgumentException.class, () -> runCommand(one.toString(), two.toString()));
+  }
+}

core/src/test/java/google/registry/tools/UpdateAllocationTokensCommandTest.java

@@ -30,12 +30,14 @@ import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
 import google.registry.model.domain.token.AllocationToken;
 import google.registry.model.domain.token.AllocationToken.TokenStatus;
+import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.TestOfyAndSql;
 import org.joda.time.DateTime;
-import org.junit.jupiter.api.Test;
 
+@DualDatabaseTest
 class UpdateAllocationTokensCommandTest extends CommandTestCase<UpdateAllocationTokensCommand> {
 
-  @Test
+  @TestOfyAndSql
   void testUpdateTlds_setTlds() throws Exception {
     AllocationToken token =
         persistResource(builderWithPromo().setAllowedTlds(ImmutableSet.of("toRemove")).build());
@@ -43,7 +45,7 @@ class UpdateAllocationTokensCommandTest extends CommandTestCase<UpdateAllocation
     assertThat(reloadResource(token).getAllowedTlds()).containsExactly("tld", "example");
   }
 
-  @Test
+  @TestOfyAndSql
   void testUpdateTlds_clearTlds() throws Exception {
     AllocationToken token =
         persistResource(builderWithPromo().setAllowedTlds(ImmutableSet.of("toRemove")).build());
@@ -51,7 +53,7 @@ class UpdateAllocationTokensCommandTest extends CommandTestCase<UpdateAllocation
     assertThat(reloadResource(token).getAllowedTlds()).isEmpty();
   }
 
-  @Test
+  @TestOfyAndSql
   void testUpdateClientIds_setClientIds() throws Exception {
     AllocationToken token =
         persistResource(
@@ -61,7 +63,7 @@ class UpdateAllocationTokensCommandTest extends CommandTestCase<UpdateAllocation
         .containsExactly("clientone", "clienttwo");
   }
 
-  @Test
+  @TestOfyAndSql
   void testUpdateClientIds_clearClientIds() throws Exception {
     AllocationToken token =
         persistResource(
@@ -70,14 +72,14 @@ class UpdateAllocationTokensCommandTest extends CommandTestCase<UpdateAllocation
     assertThat(reloadResource(token).getAllowedRegistrarIds()).isEmpty();
   }
 
-  @Test
+  @TestOfyAndSql
   void testUpdateDiscountFraction() throws Exception {
     AllocationToken token = persistResource(builderWithPromo().setDiscountFraction(0.5).build());
     runCommandForced("--prefix", "token", "--discount_fraction", "0.15");
     assertThat(reloadResource(token).getDiscountFraction()).isEqualTo(0.15);
   }
 
-  @Test
+  @TestOfyAndSql
   void testUpdateDiscountPremiums() throws Exception {
     AllocationToken token =
         persistResource(
@@ -88,14 +90,14 @@ class UpdateAllocationTokensCommandTest extends CommandTestCase<UpdateAllocation
     assertThat(reloadResource(token).shouldDiscountPremiums()).isFalse();
   }
 
-  @Test
+  @TestOfyAndSql
   void testUpdateDiscountYears() throws Exception {
     AllocationToken token = persistResource(builderWithPromo().setDiscountFraction(0.5).build());
     runCommandForced("--prefix", "token", "--discount_years", "4");
     assertThat(reloadResource(token).getDiscountYears()).isEqualTo(4);
   }
 
-  @Test
+  @TestOfyAndSql
   void testUpdateStatusTransitions() throws Exception {
     DateTime now = DateTime.now(UTC);
     AllocationToken token = persistResource(builderWithPromo().build());
@@ -110,7 +112,7 @@ class UpdateAllocationTokensCommandTest extends CommandTestCase<UpdateAllocation
         .containsExactly(START_OF_TIME, NOT_STARTED, now.minusDays(1), VALID, now, CANCELLED);
   }
 
-  @Test
+  @TestOfyAndSql
   void testUpdateStatusTransitions_badTransitions() {
     DateTime now = DateTime.now(UTC);
     persistResource(builderWithPromo().build());
@@ -130,7 +132,7 @@ class UpdateAllocationTokensCommandTest extends CommandTestCase<UpdateAllocation
         .isEqualTo("tokenStatusTransitions map cannot transition from NOT_STARTED to ENDED.");
   }
 
-  @Test
+  @TestOfyAndSql
   void testUpdate_onlyWithPrefix() throws Exception {
     AllocationToken token =
         persistResource(builderWithPromo().setAllowedTlds(ImmutableSet.of("tld")).build());
@@ -146,7 +148,7 @@ class UpdateAllocationTokensCommandTest extends CommandTestCase<UpdateAllocation
     assertThat(reloadResource(otherToken).getAllowedTlds()).isEmpty();
   }
 
-  @Test
+  @TestOfyAndSql
   void testUpdate_onlyTokensProvided() throws Exception {
     AllocationToken firstToken =
         persistResource(builderWithPromo().setAllowedTlds(ImmutableSet.of("tld")).build());
@@ -170,7 +172,7 @@ class UpdateAllocationTokensCommandTest extends CommandTestCase<UpdateAllocation
     assertThat(reloadResource(thirdToken).getAllowedTlds()).isEmpty();
   }
 
-  @Test
+  @TestOfyAndSql
   void testDoNothing() throws Exception {
     AllocationToken token =
         persistResource(
@@ -186,7 +188,7 @@ class UpdateAllocationTokensCommandTest extends CommandTestCase<UpdateAllocation
     assertThat(reloaded.getDiscountFraction()).isEqualTo(token.getDiscountFraction());
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_bothTokensAndPrefix() {
     assertThat(
             assertThrows(
@@ -196,7 +198,7 @@ class UpdateAllocationTokensCommandTest extends CommandTestCase<UpdateAllocation
         .isEqualTo("Must provide one of --tokens or --prefix, not both / neither");
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_neitherTokensNorPrefix() {
     assertThat(
             assertThrows(
@@ -205,7 +207,7 @@ class UpdateAllocationTokensCommandTest extends CommandTestCase<UpdateAllocation
         .isEqualTo("Must provide one of --tokens or --prefix, not both / neither");
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_emptyPrefix() {
     IllegalArgumentException thrown =
         assertThrows(IllegalArgumentException.class, () -> runCommandForced("--prefix", ""));

core/src/test/java/google/registry/tools/javascrap/BackfillSpec11ThreatMatchesCommandTest.java

@@ -0,0 +1,271 @@
+// 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.tools.javascrap;
+
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
+import static com.google.common.truth.Truth.assertThat;
+import static google.registry.model.ImmutableObjectSubject.assertAboutImmutableObjects;
+import static google.registry.model.ImmutableObjectSubject.immutableObjectCorrespondence;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.reporting.spec11.Spec11RegistrarThreatMatchesParserTest.sampleThreatMatches;
+import static google.registry.testing.DatabaseHelper.createTld;
+import static google.registry.testing.DatabaseHelper.deleteResource;
+import static google.registry.testing.DatabaseHelper.newDomainBase;
+import static google.registry.testing.DatabaseHelper.persistActiveDomain;
+import static google.registry.testing.DatabaseHelper.persistResource;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Iterables;
+import google.registry.model.domain.DomainBase;
+import google.registry.model.reporting.Spec11ThreatMatch;
+import google.registry.model.reporting.Spec11ThreatMatch.ThreatType;
+import google.registry.reporting.spec11.Spec11RegistrarThreatMatchesParser;
+import google.registry.tools.CommandTestCase;
+import java.io.IOException;
+import org.joda.time.DateTime;
+import org.joda.time.LocalDate;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+/** Tests for {@link BackfillSpec11ThreatMatchesCommand}. */
+public class BackfillSpec11ThreatMatchesCommandTest
+    extends CommandTestCase<BackfillSpec11ThreatMatchesCommand> {
+
+  private static final LocalDate CURRENT_DATE = DateTime.parse("2020-11-22").toLocalDate();
+  private final Spec11RegistrarThreatMatchesParser threatMatchesParser =
+      mock(Spec11RegistrarThreatMatchesParser.class);
+
+  private DomainBase domainA;
+
+  @BeforeEach
+  void beforeEach() throws Exception {
+    createTld("com");
+    domainA = persistActiveDomain("a.com");
+    persistActiveDomain("b.com");
+    persistActiveDomain("c.com");
+    fakeClock.setTo(CURRENT_DATE.toDateTimeAtStartOfDay());
+    command.threatMatchesParser = threatMatchesParser;
+    command.clock = fakeClock;
+    when(threatMatchesParser.getRegistrarThreatMatches(any(LocalDate.class)))
+        .thenReturn(ImmutableSet.of());
+  }
+
+  @Test
+  void testSuccess_singleFile() throws Exception {
+    when(threatMatchesParser.getRegistrarThreatMatches(CURRENT_DATE))
+        .thenReturn(sampleThreatMatches());
+    runCommandForced();
+    assertInStdout("Backfill Spec11 results from 692 files?");
+    assertInStdout("Successfully parsed through 692 files with 3 threats.");
+    verifyExactlyThreeEntriesInDbFromLastDay();
+  }
+
+  @Test
+  void testSuccess_sameDomain_multipleDays() throws Exception {
+    // If the same domains show up on multiple days, there should be multiple entries for them
+    when(threatMatchesParser.getRegistrarThreatMatches(CURRENT_DATE))
+        .thenReturn(sampleThreatMatches());
+    when(threatMatchesParser.getRegistrarThreatMatches(LocalDate.parse("2019-01-01")))
+        .thenReturn(sampleThreatMatches());
+    runCommandForced();
+    assertInStdout("Backfill Spec11 results from 692 files?");
+    assertInStdout("Successfully parsed through 692 files with 6 threats.");
+    jpaTm()
+        .transact(
+            () -> {
+              ImmutableList<Spec11ThreatMatch> threatMatches =
+                  jpaTm().loadAll(Spec11ThreatMatch.class);
+              assertThat(threatMatches).hasSize(6);
+              assertThat(
+                      threatMatches.stream()
+                          .map(Spec11ThreatMatch::getDomainName)
+                          .collect(toImmutableSet()))
+                  .containsExactly("a.com", "b.com", "c.com");
+              assertThat(
+                      threatMatches.stream()
+                          .map(Spec11ThreatMatch::getCheckDate)
+                          .collect(toImmutableSet()))
+                  .containsExactly(CURRENT_DATE, LocalDate.parse("2019-01-01"));
+            });
+  }
+
+  @Test
+  void testSuccess_empty() throws Exception {
+    runCommandForced();
+    assertInStdout("Backfill Spec11 results from 692 files?");
+    assertInStdout("Successfully parsed through 692 files with 0 threats.");
+  }
+
+  @Test
+  void testSuccess_sameDayTwice() throws Exception {
+    when(threatMatchesParser.getRegistrarThreatMatches(CURRENT_DATE))
+        .thenReturn(sampleThreatMatches());
+    runCommandForced();
+    runCommandForced();
+    verifyExactlyThreeEntriesInDbFromLastDay();
+  }
+
+  @Test
+  void testSuccess_threeDomainsForDomainName() throws Exception {
+    // We should use the repo ID from the proper DomainBase object at the scan's point in time.
+    // First, domain was created at START_OF_TIME and deleted one year ago
+    DateTime now = fakeClock.nowUtc();
+    domainA = persistResource(domainA.asBuilder().setDeletionTime(now.minusYears(1)).build());
+
+    // Next, domain was created six months ago and deleted two months ago
+    DomainBase secondSave =
+        persistResource(
+            newDomainBase("a.com")
+                .asBuilder()
+                .setCreationTimeForTest(now.minusMonths(6))
+                .setDeletionTime(now.minusMonths(2))
+                .build());
+
+    // Lastly, domain was created one month ago and is still valid
+    DomainBase thirdSave =
+        persistResource(
+            newDomainBase("a.com").asBuilder().setCreationTimeForTest(now.minusMonths(1)).build());
+
+    // If the scan result was from three months ago, we should use the second save
+    when(threatMatchesParser.getRegistrarThreatMatches(now.toLocalDate().minusMonths(3)))
+        .thenReturn(sampleThreatMatches());
+    runCommandForced();
+    String threatMatchRepoId =
+        jpaTm()
+            .transact(
+                () ->
+                    jpaTm().loadAll(Spec11ThreatMatch.class).stream()
+                        .filter((match) -> match.getDomainName().equals("a.com"))
+                        .findFirst()
+                        .get()
+                        .getDomainRepoId());
+    assertThat(threatMatchRepoId).isNotEqualTo(domainA.getRepoId());
+    assertThat(threatMatchRepoId).isEqualTo(secondSave.getRepoId());
+    assertThat(threatMatchRepoId).isNotEqualTo(thirdSave.getRepoId());
+  }
+
+  @Test
+  void testSuccess_skipsExistingDatesWithoutOverwrite() throws Exception {
+    when(threatMatchesParser.getRegistrarThreatMatches(CURRENT_DATE))
+        .thenReturn(sampleThreatMatches());
+    Spec11ThreatMatch previous =
+        new Spec11ThreatMatch.Builder()
+            .setCheckDate(CURRENT_DATE)
+            .setDomainName("previous.tld")
+            .setDomainRepoId("1-DOMAIN")
+            .setRegistrarId("TheRegistrar")
+            .setThreatTypes(ImmutableSet.of(ThreatType.MALWARE))
+            .build();
+    jpaTm().transact(() -> jpaTm().put(previous));
+
+    runCommandForced();
+    ImmutableList<Spec11ThreatMatch> threatMatches =
+        jpaTm().transact(() -> jpaTm().loadAll(Spec11ThreatMatch.class));
+    assertAboutImmutableObjects()
+        .that(Iterables.getOnlyElement(threatMatches))
+        .isEqualExceptFields(previous, "id");
+  }
+
+  @Test
+  void testSuccess_overwritesExistingDatesWhenSpecified() throws Exception {
+    when(threatMatchesParser.getRegistrarThreatMatches(CURRENT_DATE))
+        .thenReturn(sampleThreatMatches());
+    Spec11ThreatMatch previous =
+        new Spec11ThreatMatch.Builder()
+            .setCheckDate(CURRENT_DATE)
+            .setDomainName("previous.tld")
+            .setDomainRepoId("1-DOMAIN")
+            .setRegistrarId("TheRegistrar")
+            .setThreatTypes(ImmutableSet.of(ThreatType.MALWARE))
+            .build();
+    jpaTm().transact(() -> jpaTm().put(previous));
+
+    runCommandForced("--overwrite_existing_dates");
+    verifyExactlyThreeEntriesInDbFromLastDay();
+  }
+
+  @Test
+  void testFailure_oneFileFails() throws Exception {
+    // If there are any exceptions, we should fail loud and fast
+    when(threatMatchesParser.getRegistrarThreatMatches(CURRENT_DATE))
+        .thenReturn(sampleThreatMatches());
+    when(threatMatchesParser.getRegistrarThreatMatches(CURRENT_DATE.minusDays(1)))
+        .thenThrow(new IOException("hi"));
+    RuntimeException runtimeException =
+        assertThrows(RuntimeException.class, this::runCommandForced);
+    assertThat(runtimeException.getCause().getClass()).isEqualTo(IOException.class);
+    assertThat(runtimeException).hasCauseThat().hasMessageThat().isEqualTo("hi");
+    jpaTm().transact(() -> assertThat(jpaTm().loadAll(Spec11ThreatMatch.class)).isEmpty());
+  }
+
+  @Test
+  void testFailure_noDomainForDomainName() throws Exception {
+    deleteResource(domainA);
+    when(threatMatchesParser.getRegistrarThreatMatches(CURRENT_DATE))
+        .thenReturn(sampleThreatMatches());
+    assertThat(assertThrows(IllegalStateException.class, this::runCommandForced))
+        .hasMessageThat()
+        .isEqualTo("Domain name a.com had no associated DomainBase objects.");
+  }
+
+  @Test
+  void testFailure_noDomainAtTimeOfScan() throws Exception {
+    // If the domain existed at some point(s) in time but not the time of the scan, fail.
+    // First, domain was created at START_OF_TIME and deleted one year ago
+    DateTime now = fakeClock.nowUtc();
+    domainA = persistResource(domainA.asBuilder().setDeletionTime(now.minusYears(1)).build());
+
+    // Second, domain was created one month ago and is still valid
+    persistResource(
+        newDomainBase("a.com").asBuilder().setCreationTimeForTest(now.minusMonths(1)).build());
+
+    // If we have a result for this domain from 3 months ago when it didn't exist, fail.
+    when(threatMatchesParser.getRegistrarThreatMatches(now.toLocalDate().minusMonths(3)))
+        .thenReturn(sampleThreatMatches());
+    assertThat(assertThrows(IllegalStateException.class, this::runCommandForced))
+        .hasMessageThat()
+        .isEqualTo("Could not find a DomainBase valid for a.com on day 2020-08-22.");
+  }
+
+  private void verifyExactlyThreeEntriesInDbFromLastDay() {
+    jpaTm()
+        .transact(
+            () -> {
+              ImmutableList<Spec11ThreatMatch> threatMatches =
+                  jpaTm().loadAll(Spec11ThreatMatch.class);
+              assertThat(threatMatches)
+                  .comparingElementsUsing(immutableObjectCorrespondence("id", "domainRepoId"))
+                  .containsExactly(
+                      expectedThreatMatch("TheRegistrar", "a.com"),
+                      expectedThreatMatch("NewRegistrar", "b.com"),
+                      expectedThreatMatch("NewRegistrar", "c.com"));
+            });
+  }
+
+  private Spec11ThreatMatch expectedThreatMatch(String registrarId, String domainName) {
+    return new Spec11ThreatMatch.Builder()
+        .setDomainRepoId("ignored")
+        .setDomainName(domainName)
+        .setRegistrarId(registrarId)
+        .setCheckDate(CURRENT_DATE)
+        .setThreatTypes(ImmutableSet.of(ThreatType.MALWARE))
+        .build();
+  }
+}

core/src/test/resources/google/registry/model/schema.txt

@@ -278,8 +278,8 @@ class google.registry.model.domain.DomainHistory {
 }
 class google.registry.model.domain.GracePeriod {
   google.registry.model.domain.rgp.GracePeriodStatus type;
-  google.registry.persistence.VKey<google.registry.model.billing.BillingEvent$OneTime> billingEventOneTime;
-  google.registry.persistence.VKey<google.registry.model.billing.BillingEvent$Recurring> billingEventRecurring;
+  google.registry.persistence.BillingVKey$BillingEventVKey billingEventOneTime;
+  google.registry.persistence.BillingVKey$BillingRecurrenceVKey billingEventRecurring;
   java.lang.Long gracePeriodId;
   java.lang.String clientId;
   org.joda.time.DateTime expirationTime;
@@ -888,10 +888,17 @@ enum google.registry.model.transfer.TransferStatus {
   SERVER_APPROVED;
   SERVER_CANCELLED;
 }
+class google.registry.persistence.BillingVKey$BillingEventVKey {
+  java.lang.Long billingId;
+  java.lang.Long historyRevisionId;
+  java.lang.String repoId;
+}
+class google.registry.persistence.BillingVKey$BillingRecurrenceVKey {
+  java.lang.Long billingId;
+  java.lang.Long historyRevisionId;
+  java.lang.String repoId;
+}
 class google.registry.persistence.DomainHistoryVKey {
-  com.googlecode.objectify.Key<T> ofyKey;
-  java.lang.Class<? extends T> kind;
-  java.lang.Long domainHistoryId;
-  java.lang.Object sqlKey;
-  java.lang.String domainRepoId;
+  java.lang.Long historyRevisionId;
+  java.lang.String repoId;
 }

db/src/main/resources/sql/flyway.txt

@@ -78,3 +78,6 @@ V77__fixes_for_replay.sql
 V78__add_history_id_for_redemption_history_entry.sql
 V79__drop_foreign_keys_on_pollmessage.sql
 V80__defer_bill_event_key.sql
+V81__drop_spec11_fkeys.sql
+V82__add_columns_to_restore_symmetric_billing_vkey.sql
+V83__add_indexes_on_domainhost.sql

db/src/main/resources/sql/flyway/V81__drop_spec11_fkeys.sql

@@ -0,0 +1,18 @@
+-- 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.
+
+-- Some objects haven't been fully populated in SQL yet, so don't depend on them.
+ALTER TABLE "Spec11ThreatMatch" DROP CONSTRAINT "fk_spec11_threat_match_domain_repo_id";
+ALTER TABLE "Spec11ThreatMatch" DROP CONSTRAINT "fk_spec11_threat_match_registrar_id";
+ALTER TABLE "Spec11ThreatMatch" DROP CONSTRAINT "fk_spec11_threat_match_tld";

db/src/main/resources/sql/flyway/V82__add_columns_to_restore_symmetric_billing_vkey.sql

@@ -0,0 +1,22 @@
+-- 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.
+
+alter table "GracePeriod"
+    add column if not exists "billing_event_domain_repo_id" text;
+alter table "GracePeriod"
+    add column if not exists "billing_recurrence_domain_repo_id" text;
+alter table "GracePeriodHistory"
+    add column if not exists "billing_event_domain_repo_id" text;
+alter table "GracePeriodHistory"
+    add column if not exists "billing_recurrence_domain_repo_id" text;

db/src/main/resources/sql/flyway/V83__add_indexes_on_domainhost.sql

@@ -0,0 +1,23 @@
+-- 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.
+
+alter table if exists "DomainHistoryHost"
+   add constraint UKt2e7ae3t8gcsxd13wjx2ka7ij unique (
+       domain_history_history_revision_id,
+       domain_history_domain_repo_id,
+       host_repo_id);
+
+alter table if exists "DomainHost"
+   add constraint UKat9erbh52e4lg3jw6ai9wkjj9 unique (
+       domain_repo_id, host_repo_id);

db/src/main/resources/sql/schema/db-schema.sql.generated

@@ -398,8 +398,10 @@
        grace_period_id int8 not null,
         billing_event_id int8,
         billing_event_history_id int8,
+        billing_event_domain_repo_id text,
         billing_recurrence_id int8,
         billing_recurrence_history_id int8,
+        billing_recurrence_domain_repo_id text,
         registrar_id text not null,
         domain_repo_id text not null,
         expiration_time timestamptz not null,
@@ -411,8 +413,10 @@
        grace_period_history_revision_id int8 not null,
         billing_event_id int8,
         billing_event_history_id int8,
+        billing_event_domain_repo_id text,
         billing_recurrence_id int8,
         billing_recurrence_history_id int8,
+        billing_recurrence_domain_repo_id text,
         registrar_id text not null,
         domain_repo_id text not null,
         expiration_time timestamptz not null,
@@ -764,6 +768,12 @@ create index IDXrh4xmrot9bd63o382ow9ltfig on "DomainHistory" (creation_time);
 create index IDXaro1omfuaxjwmotk3vo00trwm on "DomainHistory" (history_registrar_id);
 create index IDXsu1nam10cjes9keobapn5jvxj on "DomainHistory" (history_type);
 create index IDX6w3qbtgce93cal2orjg1tw7b7 on "DomainHistory" (history_modification_time);
+
+    alter table if exists "DomainHistoryHost" 
+       add constraint UKt2e7ae3t8gcsxd13wjx2ka7ij unique (domain_history_history_revision_id, domain_history_domain_repo_id, host_repo_id);
+
+    alter table if exists "DomainHost" 
+       add constraint UKat9erbh52e4lg3jw6ai9wkjj9 unique (domain_repo_id, host_repo_id);
 create index IDXj1mtx98ndgbtb1bkekahms18w on "GracePeriod" (domain_repo_id);
 create index IDXd01j17vrpjxaerxdmn8bwxs7s on "GracePeriodHistory" (domain_repo_id);
 create index IDXfg2nnjlujxo6cb9fha971bq2n on "HostHistory" (creation_time);

db/src/main/resources/sql/schema/nomulus.golden.sql

@@ -533,7 +533,9 @@ CREATE TABLE public."GracePeriod" (
     expiration_time timestamp with time zone NOT NULL,
     type text NOT NULL,
     billing_event_history_id bigint,
-    billing_recurrence_history_id bigint
+    billing_recurrence_history_id bigint,
+    billing_event_domain_repo_id text,
+    billing_recurrence_domain_repo_id text
 );
 
 
@@ -552,7 +554,9 @@ CREATE TABLE public."GracePeriodHistory" (
     expiration_time timestamp with time zone NOT NULL,
     type text NOT NULL,
     domain_history_revision_id bigint,
-    grace_period_id bigint NOT NULL
+    grace_period_id bigint NOT NULL,
+    billing_event_domain_repo_id text,
+    billing_recurrence_domain_repo_id text
 );
 
 
@@ -1424,6 +1428,22 @@ ALTER TABLE ONLY public."RegistryLock"
     ADD CONSTRAINT idx_registry_lock_repo_id_revision_id UNIQUE (repo_id, revision_id);
 
 
+--
+-- Name: DomainHost ukat9erbh52e4lg3jw6ai9wkjj9; Type: CONSTRAINT; Schema: public; Owner: -
+--
+
+ALTER TABLE ONLY public."DomainHost"
+    ADD CONSTRAINT ukat9erbh52e4lg3jw6ai9wkjj9 UNIQUE (domain_repo_id, host_repo_id);
+
+
+--
+-- Name: DomainHistoryHost ukt2e7ae3t8gcsxd13wjx2ka7ij; Type: CONSTRAINT; Schema: public; Owner: -
+--
+
+ALTER TABLE ONLY public."DomainHistoryHost"
+    ADD CONSTRAINT ukt2e7ae3t8gcsxd13wjx2ka7ij UNIQUE (domain_history_history_revision_id, domain_history_domain_repo_id, host_repo_id);
+
+
 --
 -- Name: allocation_token_domain_name_idx; Type: INDEX; Schema: public; Owner: -
 --
@@ -2251,30 +2271,6 @@ ALTER TABLE ONLY public."RegistrarPoc"
     ADD CONSTRAINT fk_registrar_poc_registrar_id FOREIGN KEY (registrar_id) REFERENCES public."Registrar"(registrar_id);
 
 
---
--- Name: Spec11ThreatMatch fk_spec11_threat_match_domain_repo_id; Type: FK CONSTRAINT; Schema: public; Owner: -
---
-
-ALTER TABLE ONLY public."Spec11ThreatMatch"
-    ADD CONSTRAINT fk_spec11_threat_match_domain_repo_id FOREIGN KEY (domain_repo_id) REFERENCES public."Domain"(repo_id);
-
-
---
--- Name: Spec11ThreatMatch fk_spec11_threat_match_registrar_id; Type: FK CONSTRAINT; Schema: public; Owner: -
---
-
-ALTER TABLE ONLY public."Spec11ThreatMatch"
-    ADD CONSTRAINT fk_spec11_threat_match_registrar_id FOREIGN KEY (registrar_id) REFERENCES public."Registrar"(registrar_id);
-
-
---
--- Name: Spec11ThreatMatch fk_spec11_threat_match_tld; Type: FK CONSTRAINT; Schema: public; Owner: -
---
-
-ALTER TABLE ONLY public."Spec11ThreatMatch"
-    ADD CONSTRAINT fk_spec11_threat_match_tld FOREIGN KEY (tld) REFERENCES public."Tld"(tld_name);
-
-
 --
 -- Name: DomainHistoryHost fka9woh3hu8gx5x0vly6bai327n; Type: FK CONSTRAINT; Schema: public; Owner: -
 --

java_common.gradle

@@ -74,10 +74,12 @@ test {
   useJUnitPlatform()
 }
 
-// Sets up integration test with a registry environment. The target environment is
-// passed by the 'test.gcp_integration.env' property. Test runner must have been
-// authorized to access the corresponding GCP project, e.g., by running 'gcloud auth'
-// or placing a credential file at a well known place.
+// Sets up integration test with a registry environment. The target environment
+// is passed by the 'test.gcp_integration.env' property. Test runner must have
+// been authorized to access the corresponding GCP project, e.g., by running
+// 'gcloud auth application-default login' or by downloading a credential file
+// and assign the path to it to the GOOGLE_APPLICATION_CREDENTIALS environment
+// variable.
 //
 // A typical use case is to run tests from desktop that accesses Cloud resources. See
 // core/src/test/java/google/registry/beam/initsql/BeamJpaModuleTest.java for an example.

proxy/src/main/java/google/registry/proxy/handler/EppServiceHandler.java

@@ -34,7 +34,9 @@ import io.netty.handler.codec.http.HttpResponse;
 import io.netty.handler.ssl.SslHandshakeCompletionEvent;
 import io.netty.util.AttributeKey;
 import io.netty.util.concurrent.Promise;
+import java.security.cert.CertificateEncodingException;
 import java.security.cert.X509Certificate;
+import java.util.Base64;
 import java.util.function.Supplier;
 
 /** Handler that processes EPP protocol logic. */
@@ -52,18 +54,26 @@ public class EppServiceHandler extends HttpsRelayServiceHandler {
   /** Name of the HTTP header that stores the client certificate hash. */
   public static final String SSL_CLIENT_CERTIFICATE_HASH_FIELD = "X-SSL-Certificate";
 
+  /** Name of the HTTP header that stores the full client certificate. */
+  public static final String SSL_CLIENT_FULL_CERTIFICATE_FIELD = "X-SSL-Full-Certificate";
+
   /** Name of the HTTP header that stores the client IP address. */
   public static final String FORWARDED_FOR_FIELD = "X-Forwarded-For";
 
   /** Name of the HTTP header that indicates if the EPP session should be closed. */
   public static final String EPP_SESSION_FIELD = "Epp-Session";
 
+  /** Name of the HTTP header that indicates a successful login has occurred. */
+  public static final String EPP_LOGGED_IN_FIELD = "Logged-In";
+
   public static final String EPP_CONTENT_TYPE = "application/epp+xml";
 
   private final byte[] helloBytes;
 
   private String sslClientCertificateHash;
+  private X509Certificate sslClientCertificate;
   private String clientAddress;
+  private boolean isLoggedIn = false;
 
   public EppServiceHandler(
       String relayHost,
@@ -103,7 +113,8 @@ public class EppServiceHandler extends HttpsRelayServiceHandler {
             .addListener(
                 (Promise<X509Certificate> promise) -> {
                   if (promise.isSuccess()) {
-                    sslClientCertificateHash = getCertificateHash(promise.get());
+                    sslClientCertificate = promise.get();
+                    sslClientCertificateHash = getCertificateHash(sslClientCertificate);
                     // Set the client cert hash key attribute for both this channel,
                     // used for collecting metrics on specific clients.
                     ctx.channel().attr(CLIENT_CERTIFICATE_HASH_KEY).set(sslClientCertificateHash);
@@ -132,6 +143,17 @@ public class EppServiceHandler extends HttpsRelayServiceHandler {
         .set(FORWARDED_FOR_FIELD, clientAddress)
         .set(HttpHeaderNames.CONTENT_TYPE, EPP_CONTENT_TYPE)
         .set(HttpHeaderNames.ACCEPT, EPP_CONTENT_TYPE);
+    if (!isLoggedIn) {
+      try {
+        request
+            .headers()
+            .set(
+                SSL_CLIENT_FULL_CERTIFICATE_FIELD,
+                Base64.getEncoder().encodeToString(sslClientCertificate.getEncoded()));
+      } catch (CertificateEncodingException e) {
+        throw new RuntimeException("Cannot encode client certificate", e);
+      }
+    }
     return request;
   }
 
@@ -141,9 +163,13 @@ public class EppServiceHandler extends HttpsRelayServiceHandler {
     checkArgument(msg instanceof HttpResponse);
     HttpResponse response = (HttpResponse) msg;
     String sessionAliveValue = response.headers().get(EPP_SESSION_FIELD);
+    String loginValue = response.headers().get(EPP_LOGGED_IN_FIELD);
     if (sessionAliveValue != null && sessionAliveValue.equals("close")) {
       promise.addListener(ChannelFutureListener.CLOSE);
     }
+    if (loginValue != null && loginValue.equals("true")) {
+      isLoggedIn = true;
+    }
     super.write(ctx, msg, promise);
   }
 }

proxy/src/test/java/google/registry/proxy/EppProtocolModuleTest.java

@@ -16,6 +16,7 @@ package google.registry.proxy;
 
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.networking.handler.SslServerInitializer.CLIENT_CERTIFICATE_PROMISE_KEY;
+import static google.registry.proxy.TestUtils.SAMPLE_CERT;
 import static google.registry.proxy.handler.ProxyProtocolHandler.REMOTE_ADDRESS_KEY;
 import static google.registry.util.ResourceUtils.readResourceBytes;
 import static google.registry.util.X509Utils.getCertificateHash;
@@ -25,7 +26,6 @@ import static org.junit.jupiter.api.Assertions.assertThrows;
 import com.google.common.base.Throwables;
 import google.registry.proxy.handler.HttpsRelayServiceHandler.NonOkHttpResponseException;
 import google.registry.testing.FakeClock;
-import google.registry.util.SelfSignedCaCertificate;
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.Unpooled;
 import io.netty.channel.embedded.EmbeddedChannel;
@@ -36,6 +36,8 @@ import io.netty.handler.codec.http.HttpResponseStatus;
 import io.netty.handler.codec.http.cookie.Cookie;
 import io.netty.handler.codec.http.cookie.DefaultCookie;
 import io.netty.util.concurrent.Promise;
+import java.io.ByteArrayInputStream;
+import java.security.cert.CertificateFactory;
 import java.security.cert.X509Certificate;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -96,8 +98,8 @@ class EppProtocolModuleTest extends ProtocolModuleTest {
     return buffer;
   }
 
-  private FullHttpRequest makeEppHttpRequest(byte[] content, Cookie... cookies) {
-    return TestUtils.makeEppHttpRequest(
+  private FullHttpRequest makeEppHttpRequestWithCertificate(byte[] content, Cookie... cookies) {
+    return TestUtils.makeEppHttpRequestWithCertificate(
         new String(content, UTF_8),
         PROXY_CONFIG.epp.relayHost,
         PROXY_CONFIG.epp.relayPath,
@@ -120,7 +122,10 @@ class EppProtocolModuleTest extends ProtocolModuleTest {
   @Override
   void beforeEach() throws Exception {
     testComponent = makeTestComponent(new FakeClock());
-    certificate = SelfSignedCaCertificate.create().cert();
+    CertificateFactory cf = CertificateFactory.getInstance("X.509");
+    certificate =
+        (X509Certificate)
+            cf.generateCertificate(new ByteArrayInputStream(SAMPLE_CERT.getBytes(UTF_8)));
     initializeChannel(
         ch -> {
           ch.attr(REMOTE_ADDRESS_KEY).set(CLIENT_ADDRESS);
@@ -134,13 +139,15 @@ class EppProtocolModuleTest extends ProtocolModuleTest {
   @Test
   void testSuccess_singleFrameInboundMessage() throws Exception {
     // First inbound message is hello.
-    assertThat((FullHttpRequest) channel.readInbound()).isEqualTo(makeEppHttpRequest(HELLO_BYTES));
+    assertThat((FullHttpRequest) channel.readInbound())
+        .isEqualTo(makeEppHttpRequestWithCertificate(HELLO_BYTES));
 
     byte[] inputBytes = readResourceBytes(getClass(), "login.xml").read();
 
     // Verify inbound message is as expected.
     assertThat(channel.writeInbound(getByteBufFromContent(inputBytes))).isTrue();
-    assertThat((FullHttpRequest) channel.readInbound()).isEqualTo(makeEppHttpRequest(inputBytes));
+    assertThat((FullHttpRequest) channel.readInbound())
+        .isEqualTo(makeEppHttpRequestWithCertificate(inputBytes));
 
     // Nothing more to read.
     assertThat((Object) channel.readInbound()).isNull();
@@ -161,8 +168,10 @@ class EppProtocolModuleTest extends ProtocolModuleTest {
                 Unpooled.wrappedBuffer(
                     getByteBufFromContent(inputBytes1), getByteBufFromContent(inputBytes2))))
         .isTrue();
-    assertThat((FullHttpRequest) channel.readInbound()).isEqualTo(makeEppHttpRequest(inputBytes1));
-    assertThat((FullHttpRequest) channel.readInbound()).isEqualTo(makeEppHttpRequest(inputBytes2));
+    assertThat((FullHttpRequest) channel.readInbound())
+        .isEqualTo(makeEppHttpRequestWithCertificate(inputBytes1));
+    assertThat((FullHttpRequest) channel.readInbound())
+        .isEqualTo(makeEppHttpRequestWithCertificate(inputBytes2));
 
     // Nothing more to read.
     assertThat((Object) channel.readInbound()).isNull();
@@ -186,11 +195,13 @@ class EppProtocolModuleTest extends ProtocolModuleTest {
 
     // The second frame contains the first message, and part of the second message.
     assertThat(channel.writeInbound(inputBuffer.readBytes(inputBytes2.length))).isTrue();
-    assertThat((FullHttpRequest) channel.readInbound()).isEqualTo(makeEppHttpRequest(inputBytes1));
+    assertThat((FullHttpRequest) channel.readInbound())
+        .isEqualTo(makeEppHttpRequestWithCertificate(inputBytes1));
 
     // The third frame contains the rest of the second message.
     assertThat(channel.writeInbound(inputBuffer)).isTrue();
-    assertThat((FullHttpRequest) channel.readInbound()).isEqualTo(makeEppHttpRequest(inputBytes2));
+    assertThat((FullHttpRequest) channel.readInbound())
+        .isEqualTo(makeEppHttpRequestWithCertificate(inputBytes2));
 
     // Nothing more to read.
     assertThat((Object) channel.readInbound()).isNull();
@@ -252,7 +263,7 @@ class EppProtocolModuleTest extends ProtocolModuleTest {
     byte[] inputBytes1 = readResourceBytes(getClass(), "logout.xml").read();
     assertThat(channel.writeInbound(getByteBufFromContent(inputBytes1))).isTrue();
     assertThat((FullHttpRequest) channel.readInbound())
-        .isEqualTo(makeEppHttpRequest(inputBytes1, cookie1, cookie2));
+        .isEqualTo(makeEppHttpRequestWithCertificate(inputBytes1, cookie1, cookie2));
 
     // Second outbound message change cookies.
     byte[] outputBytes2 = readResourceBytes(getClass(), "logout_response.xml").read();
@@ -267,7 +278,7 @@ class EppProtocolModuleTest extends ProtocolModuleTest {
     byte[] inputBytes2 = readResourceBytes(getClass(), "login.xml").read();
     assertThat(channel.writeInbound(getByteBufFromContent(inputBytes2))).isTrue();
     assertThat((FullHttpRequest) channel.readInbound())
-        .isEqualTo(makeEppHttpRequest(inputBytes2, cookie1, cookie2, cookie3));
+        .isEqualTo(makeEppHttpRequestWithCertificate(inputBytes2, cookie1, cookie2, cookie3));
 
     // Nothing more to write or read.
     assertThat((Object) channel.readOutbound()).isNull();

proxy/src/test/java/google/registry/proxy/TestUtils.java

@@ -36,6 +36,54 @@ import io.netty.handler.codec.http.cookie.ServerCookieEncoder;
 /** Utility class for various helper methods used in testing. */
 public class TestUtils {
 
+  public static final String SAMPLE_CERT =
+      "-----BEGIN CERTIFICATE-----\n"
+          + "MIIDvTCCAqWgAwIBAgIJAK/PgPT0jTwRMA0GCSqGSIb3DQEBCwUAMHUxCzAJBgNV\n"
+          + "BAYTAlVTMREwDwYDVQQIDAhOZXcgWW9yazERMA8GA1UEBwwITmV3IFlvcmsxDzAN\n"
+          + "BgNVBAoMBkdvb2dsZTEdMBsGA1UECwwUZG9tYWluLXJlZ2lzdHJ5LXRlc3QxEDAO\n"
+          + "BgNVBAMMB2NsaWVudDEwHhcNMTUwODI2MTkxODA4WhcNNDMwMTExMTkxODA4WjB1\n"
+          + "MQswCQYDVQQGEwJVUzERMA8GA1UECAwITmV3IFlvcmsxETAPBgNVBAcMCE5ldyBZ\n"
+          + "b3JrMQ8wDQYDVQQKDAZHb29nbGUxHTAbBgNVBAsMFGRvbWFpbi1yZWdpc3RyeS10\n"
+          + "ZXN0MRAwDgYDVQQDDAdjbGllbnQxMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIB\n"
+          + "CgKCAQEAvoE/IoFJyzb0dU4NFhL8FYgy+B/GnUd5aA66CMx5xKRMbEAtIgxU8TTO\n"
+          + "W+9jdTsE00Grk3Ct4KdY73CYW+6IFXL4O0K/m5S+uajh+I2UMVZJV38RAIqNxue0\n"
+          + "Egv9M4haSsCVIPcX9b+6McywfYSF1bzPb2Gb2FAQO7Jb0BjlPhPMIROCrbG40qPg\n"
+          + "LWrl33dz+O52kO+DyZEzHqI55xH6au77sMITsJe+X23lzQcMFUUm8moiOw0EKrj/\n"
+          + "GaMTZLHP46BCRoJDAPTNx55seIwgAHbKA2VVtqrvmA2XYJQA6ipdhfKRoJFy8Z8H\n"
+          + "DYsorGtazQL2HhF/5uJD25z1m5eQHQIDAQABo1AwTjAdBgNVHQ4EFgQUParEmiSR\n"
+          + "U/Oqy8hr7k+MBKhZwVkwHwYDVR0jBBgwFoAUParEmiSRU/Oqy8hr7k+MBKhZwVkw\n"
+          + "DAYDVR0TBAUwAwEB/zANBgkqhkiG9w0BAQsFAAOCAQEAojsUhF6PtZrStnHBFWNR\n"
+          + "ryzvANB8krZlYeX9Hkqn8zIVfAkpbVmL8aZQ7yj17jSpw47PQh3x5gwA9yc/SS0G\n"
+          + "E1rGuxYH02UGbua8G0+vviSQfLtskPQzK7EIR63WNhHEo/Q9umLJkZ0LguWEBf3L\n"
+          + "q8CoXv2i/RNvqVPcTNp/zCKXJZAa8wAjNRJs834AZj4k5xwyYZ3F8D5PGz+YMOmV\n"
+          + "M9Qd+NdXSC/Qn7HQzFhE8p5elBV35P8oX5dXEfn0S7zOXDenp5JvvLoggOWOcKsq\n"
+          + "KiWDQrsT+TMKmHL94/h4t7FghtQLMzY5SGYJsYTv/LG8tewrz6KRb/Wj3JNojyEw\n"
+          + "Ug==\n"
+          + "-----END CERTIFICATE-----\n";
+
+  public static final String SAMPLE_CERT_ENCODED =
+      "MIIDvTCCAqWgAwIBAgIJAK/PgPT0jTwRMA0GCSqGSIb3DQEBCwUAMHUxCzAJBgNV"
+          + "BAYTAlVTMREwDwYDVQQIDAhOZXcgWW9yazERMA8GA1UEBwwITmV3IFlvcmsxDzAN"
+          + "BgNVBAoMBkdvb2dsZTEdMBsGA1UECwwUZG9tYWluLXJlZ2lzdHJ5LXRlc3QxEDAO"
+          + "BgNVBAMMB2NsaWVudDEwHhcNMTUwODI2MTkxODA4WhcNNDMwMTExMTkxODA4WjB1"
+          + "MQswCQYDVQQGEwJVUzERMA8GA1UECAwITmV3IFlvcmsxETAPBgNVBAcMCE5ldyBZ"
+          + "b3JrMQ8wDQYDVQQKDAZHb29nbGUxHTAbBgNVBAsMFGRvbWFpbi1yZWdpc3RyeS10"
+          + "ZXN0MRAwDgYDVQQDDAdjbGllbnQxMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIB"
+          + "CgKCAQEAvoE/IoFJyzb0dU4NFhL8FYgy+B/GnUd5aA66CMx5xKRMbEAtIgxU8TTO"
+          + "W+9jdTsE00Grk3Ct4KdY73CYW+6IFXL4O0K/m5S+uajh+I2UMVZJV38RAIqNxue0"
+          + "Egv9M4haSsCVIPcX9b+6McywfYSF1bzPb2Gb2FAQO7Jb0BjlPhPMIROCrbG40qPg"
+          + "LWrl33dz+O52kO+DyZEzHqI55xH6au77sMITsJe+X23lzQcMFUUm8moiOw0EKrj/"
+          + "GaMTZLHP46BCRoJDAPTNx55seIwgAHbKA2VVtqrvmA2XYJQA6ipdhfKRoJFy8Z8H"
+          + "DYsorGtazQL2HhF/5uJD25z1m5eQHQIDAQABo1AwTjAdBgNVHQ4EFgQUParEmiSR"
+          + "U/Oqy8hr7k+MBKhZwVkwHwYDVR0jBBgwFoAUParEmiSRU/Oqy8hr7k+MBKhZwVkw"
+          + "DAYDVR0TBAUwAwEB/zANBgkqhkiG9w0BAQsFAAOCAQEAojsUhF6PtZrStnHBFWNR"
+          + "ryzvANB8krZlYeX9Hkqn8zIVfAkpbVmL8aZQ7yj17jSpw47PQh3x5gwA9yc/SS0G"
+          + "E1rGuxYH02UGbua8G0+vviSQfLtskPQzK7EIR63WNhHEo/Q9umLJkZ0LguWEBf3L"
+          + "q8CoXv2i/RNvqVPcTNp/zCKXJZAa8wAjNRJs834AZj4k5xwyYZ3F8D5PGz+YMOmV"
+          + "M9Qd+NdXSC/Qn7HQzFhE8p5elBV35P8oX5dXEfn0S7zOXDenp5JvvLoggOWOcKsq"
+          + "KiWDQrsT+TMKmHL94/h4t7FghtQLMzY5SGYJsYTv/LG8tewrz6KRb/Wj3JNojyEw"
+          + "Ug==";
+
   public static FullHttpRequest makeHttpPostRequest(String content, String host, String path) {
     ByteBuf buf = Unpooled.wrappedBuffer(content.getBytes(US_ASCII));
     FullHttpRequest request =
@@ -101,6 +149,21 @@ public class TestUtils {
     return request;
   }
 
+  public static FullHttpRequest makeEppHttpRequestWithCertificate(
+      String content,
+      String host,
+      String path,
+      String accessToken,
+      String sslClientCertificateHash,
+      String clientAddress,
+      Cookie... cookies) {
+    FullHttpRequest request =
+        makeEppHttpRequest(
+            content, host, path, accessToken, sslClientCertificateHash, clientAddress, cookies);
+    request.headers().set("X-SSL-Full-Certificate", SAMPLE_CERT_ENCODED);
+    return request;
+  }
+
   public static FullHttpResponse makeWhoisHttpResponse(String content, HttpResponseStatus status) {
     FullHttpResponse response = makeHttpResponse(content, status);
     response.headers().set("content-type", "text/plain");

proxy/src/test/java/google/registry/proxy/handler/EppServiceHandlerTest.java

@@ -16,6 +16,7 @@ package google.registry.proxy.handler;
 
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.networking.handler.SslServerInitializer.CLIENT_CERTIFICATE_PROMISE_KEY;
+import static google.registry.proxy.TestUtils.SAMPLE_CERT;
 import static google.registry.proxy.TestUtils.assertHttpRequestEquivalent;
 import static google.registry.proxy.TestUtils.makeEppHttpResponse;
 import static google.registry.proxy.handler.ProxyProtocolHandler.REMOTE_ADDRESS_KEY;
@@ -43,6 +44,8 @@ import io.netty.handler.codec.http.HttpResponseStatus;
 import io.netty.handler.codec.http.cookie.Cookie;
 import io.netty.handler.codec.http.cookie.DefaultCookie;
 import io.netty.util.concurrent.Promise;
+import java.io.ByteArrayInputStream;
+import java.security.cert.CertificateFactory;
 import java.security.cert.X509Certificate;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -109,9 +112,23 @@ class EppServiceHandlerTest {
         cookies);
   }
 
+  private FullHttpRequest makeEppHttpRequestWithCertificate(String content, Cookie... cookies) {
+    return TestUtils.makeEppHttpRequestWithCertificate(
+        content,
+        RELAY_HOST,
+        RELAY_PATH,
+        ACCESS_TOKEN,
+        getCertificateHash(clientCertificate),
+        CLIENT_ADDRESS,
+        cookies);
+  }
+
   @BeforeEach
   void beforeEach() throws Exception {
-    clientCertificate = SelfSignedCaCertificate.create().cert();
+    CertificateFactory cf = CertificateFactory.getInstance("X.509");
+    clientCertificate =
+        (X509Certificate)
+            cf.generateCertificate(new ByteArrayInputStream(SAMPLE_CERT.getBytes(UTF_8)));
     channel = setUpNewChannel(eppServiceHandler);
   }
 
@@ -194,7 +211,7 @@ class EppServiceHandlerTest {
     setHandshakeSuccess();
     // hello bytes should be passed to the next handler.
     FullHttpRequest helloRequest = channel.readInbound();
-    assertThat(helloRequest).isEqualTo(makeEppHttpRequest(HELLO));
+    assertThat(helloRequest).isEqualTo(makeEppHttpRequestWithCertificate(HELLO));
     // Nothing further to pass to the next handler.
     assertThat((Object) channel.readInbound()).isNull();
     assertThat(channel.isActive()).isTrue();
@@ -216,12 +233,34 @@ class EppServiceHandlerTest {
     String content = "<epp>stuff</epp>";
     channel.writeInbound(Unpooled.wrappedBuffer(content.getBytes(UTF_8)));
     FullHttpRequest request = channel.readInbound();
-    assertThat(request).isEqualTo(makeEppHttpRequest(content));
+    assertThat(request).isEqualTo(makeEppHttpRequestWithCertificate(content));
     // Nothing further to pass to the next handler.
     assertThat((Object) channel.readInbound()).isNull();
     assertThat(channel.isActive()).isTrue();
   }
 
+  @Test
+  void testSuccess_sendCertificateOnlyBeforeLogin() throws Exception {
+    setHandshakeSuccess();
+    // First inbound message is hello.
+    channel.readInbound();
+    String content = "<epp>stuff</epp>";
+    channel.writeInbound(Unpooled.wrappedBuffer(content.getBytes(UTF_8)));
+    FullHttpRequest request = channel.readInbound();
+    assertThat(request).isEqualTo(makeEppHttpRequestWithCertificate(content));
+    // Receive response indicating session is logged in
+    HttpResponse response = makeEppHttpResponse(content, HttpResponseStatus.OK);
+    response.headers().set("Logged-In", "true");
+    // Send another inbound message after login
+    channel.writeOutbound(response);
+    channel.writeInbound(Unpooled.wrappedBuffer(content.getBytes(UTF_8)));
+    request = channel.readInbound();
+    // Second request should not have full certificate
+    assertThat(request).isEqualTo(makeEppHttpRequest(content));
+    assertThat((Object) channel.readInbound()).isNull();
+    assertThat(channel.isActive()).isTrue();
+  }
+
   @Test
   void testSuccess_sendResponseToNextHandler() throws Exception {
     setHandshakeSuccess();
@@ -294,7 +333,8 @@ class EppServiceHandlerTest {
     String requestContent = "<epp>request</epp>";
     channel.writeInbound(Unpooled.wrappedBuffer(requestContent.getBytes(UTF_8)));
     FullHttpRequest request = channel.readInbound();
-    assertHttpRequestEquivalent(request, makeEppHttpRequest(requestContent, cookie1, cookie2));
+    assertHttpRequestEquivalent(
+        request, makeEppHttpRequestWithCertificate(requestContent, cookie1, cookie2));
     // Nothing further to pass to the next handler.
     assertThat((Object) channel.readInbound()).isNull();
     assertThat((Object) channel.readOutbound()).isNull();
@@ -317,13 +357,16 @@ class EppServiceHandlerTest {
     // First request written.
     channel.writeInbound(Unpooled.wrappedBuffer(requestContent1.getBytes(UTF_8)));
     FullHttpRequest request1 = channel.readInbound();
-    assertHttpRequestEquivalent(request1, makeEppHttpRequest(requestContent1, cookie1, cookie2));
+    assertHttpRequestEquivalent(
+        request1, makeEppHttpRequestWithCertificate(requestContent1, cookie1, cookie2));
     String responseContent2 = "<epp>response2</epp>";
     Cookie cookie3 = new DefaultCookie("name3", "value3");
     Cookie newCookie2 = new DefaultCookie("name2", "newValue");
     // Second response written.
-    channel.writeOutbound(
-        makeEppHttpResponse(responseContent2, HttpResponseStatus.OK, cookie3, newCookie2));
+    HttpResponse response =
+        makeEppHttpResponse(responseContent2, HttpResponseStatus.OK, cookie3, newCookie2);
+    response.headers().set("Logged-In", "true");
+    channel.writeOutbound(response);
     channel.readOutbound();
     String requestContent2 = "<epp>request2</epp>";
     // Second request written.