Parameterize the serialization of objects being written to SQL (#892)

* Parameterize the serialization of objects being written to SQL

We shouldn't require that objects written to SQL during a Beam pipeline
be VersionedEntity objects -- they may be non-Objectify entities. As a
result, we should allow the user to specify what the objects are that
should be written to SQL.

Note: we will need to clean up the Spec11PipelineTest more but that can
be out of the scope of this PR.

* Overload the method and add a bit of javadoc

* Actually use the overloaded function

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

@@ -15,10 +15,10 @@
 package google.registry.backup;
 
 import static com.google.common.base.Preconditions.checkState;
+import static com.google.common.collect.ImmutableList.toImmutableList;
 import static google.registry.backup.BackupUtils.createDeserializingIterator;
 
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.Streams;
 import google.registry.model.ImmutableObject;
 import google.registry.model.ofy.CommitLogCheckpoint;
 import google.registry.model.ofy.CommitLogManifest;
@@ -31,7 +31,6 @@ import java.io.InputStream;
 import java.nio.channels.Channels;
 import java.nio.channels.ReadableByteChannel;
 import java.util.Iterator;
-import java.util.stream.Stream;
 
 /**
  * Helpers for reading CommitLog records from a file.
@@ -43,6 +42,56 @@ public final class CommitLogImports {
 
   private CommitLogImports() {}
 
+  /**
+   * Returns entities in an {@code inputStream} (from a single CommitLog file) as an {@link
+   * ImmutableList} of {@link ImmutableList}s of {@link VersionedEntity} records where the inner
+   * lists each consist of one transaction. Upon completion the {@code inputStream} is closed.
+   *
+   * <p>The returned list may be empty, since CommitLogs are written at fixed intervals regardless
+   * if actual changes exist. Each sublist, however, will not be empty.
+   *
+   * <p>A CommitLog file starts with a {@link CommitLogCheckpoint}, followed by (repeated)
+   * subsequences of [{@link CommitLogManifest}, [{@link CommitLogMutation}] ...]. Each subsequence
+   * represents the changes in one transaction. The {@code CommitLogManifest} contains deleted
+   * entity keys, whereas each {@code CommitLogMutation} contains one whole entity.
+   */
+  public static ImmutableList<ImmutableList<VersionedEntity>> loadEntitiesByTransaction(
+      InputStream inputStream) {
+    try (AppEngineEnvironment appEngineEnvironment = new AppEngineEnvironment();
+        InputStream input = new BufferedInputStream(inputStream)) {
+      Iterator<ImmutableObject> commitLogs = createDeserializingIterator(input);
+      checkState(commitLogs.hasNext());
+      checkState(commitLogs.next() instanceof CommitLogCheckpoint);
+
+      ImmutableList.Builder<ImmutableList<VersionedEntity>> resultBuilder =
+          new ImmutableList.Builder<>();
+      ImmutableList.Builder<VersionedEntity> currentTransactionBuilder =
+          new ImmutableList.Builder<>();
+
+      while (commitLogs.hasNext()) {
+        ImmutableObject currentObject = commitLogs.next();
+        if (currentObject instanceof CommitLogManifest) {
+          // CommitLogManifest means we are starting a new transaction
+          addIfNonempty(resultBuilder, currentTransactionBuilder);
+          currentTransactionBuilder = new ImmutableList.Builder<>();
+          VersionedEntity.fromManifest((CommitLogManifest) currentObject)
+              .forEach(currentTransactionBuilder::add);
+        } else if (currentObject instanceof CommitLogMutation) {
+          currentTransactionBuilder.add(
+              VersionedEntity.fromMutation((CommitLogMutation) currentObject));
+        } else {
+          throw new IllegalStateException(
+              String.format("Unknown entity type %s in commit logs", currentObject.getClass()));
+        }
+      }
+      // Add the last transaction in (if it's not empty)
+      addIfNonempty(resultBuilder, currentTransactionBuilder);
+      return resultBuilder.build();
+    } catch (IOException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
   /**
    * Returns entities in an {@code inputStream} (from a single CommitLog file) as an {@link
    * ImmutableList} of {@link VersionedEntity} records. Upon completion the {@code inputStream} is
@@ -57,23 +106,9 @@ public final class CommitLogImports {
    * entity keys, whereas each {@code CommitLogMutation} contains one whole entity.
    */
   public static ImmutableList<VersionedEntity> loadEntities(InputStream inputStream) {
-    try (AppEngineEnvironment appEngineEnvironment = new AppEngineEnvironment();
-        InputStream input = new BufferedInputStream(inputStream)) {
-      Iterator<ImmutableObject> commitLogs = createDeserializingIterator(input);
-      checkState(commitLogs.hasNext());
-      checkState(commitLogs.next() instanceof CommitLogCheckpoint);
-
-      return Streams.stream(commitLogs)
-          .map(
-              e ->
-                  e instanceof CommitLogManifest
-                      ? VersionedEntity.fromManifest((CommitLogManifest) e)
-                      : Stream.of(VersionedEntity.fromMutation((CommitLogMutation) e)))
-          .flatMap(s -> s)
-          .collect(ImmutableList.toImmutableList());
-    } catch (IOException e) {
-      throw new RuntimeException(e);
-    }
+    return loadEntitiesByTransaction(inputStream).stream()
+        .flatMap(ImmutableList::stream)
+        .collect(toImmutableList());
   }
 
   /** Covenience method that adapts {@link #loadEntities(InputStream)} to a {@link File}. */
@@ -92,4 +127,13 @@ public final class CommitLogImports {
   public static ImmutableList<VersionedEntity> loadEntities(ReadableByteChannel channel) {
     return loadEntities(Channels.newInputStream(channel));
   }
+
+  private static void addIfNonempty(
+      ImmutableList.Builder<ImmutableList<VersionedEntity>> resultBuilder,
+      ImmutableList.Builder<VersionedEntity> currentTransactionBuilder) {
+    ImmutableList<VersionedEntity> currentTransaction = currentTransactionBuilder.build();
+    if (!currentTransaction.isEmpty()) {
+      resultBuilder.add(currentTransaction);
+    }
+  }
 }

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

@@ -122,7 +122,7 @@ class GcsDiffFileLister {
 
     // Reconstruct the sequence of files by traversing backwards from "lastUpperBoundTime" (i.e. the
     // last file that we found) and finding its previous file until we either run out of files or
-    // get to one that preceeds "fromTime".
+    // get to one that precedes "fromTime".
     //
     // GCS file listing is eventually consistent, so it's possible that we are missing a file. The
     // metadata of a file is sufficient to identify the preceding file, so if we start from the

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

@@ -0,0 +1,188 @@
+// 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.backup;
+
+import static google.registry.backup.ExportCommitLogDiffAction.DIFF_FILE_PREFIX;
+import static google.registry.model.ofy.EntityWritePriorities.getEntityPriority;
+import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static javax.servlet.http.HttpServletResponse.SC_NO_CONTENT;
+import static org.joda.time.Duration.standardHours;
+
+import com.google.appengine.api.datastore.Entity;
+import com.google.appengine.api.datastore.Key;
+import com.google.appengine.tools.cloudstorage.GcsFileMetadata;
+import com.google.appengine.tools.cloudstorage.GcsService;
+import com.google.common.collect.ImmutableList;
+import com.google.common.flogger.FluentLogger;
+import google.registry.config.RegistryConfig;
+import google.registry.model.server.Lock;
+import google.registry.model.translators.VKeyTranslatorFactory;
+import google.registry.persistence.VKey;
+import google.registry.request.Action;
+import google.registry.request.Response;
+import google.registry.request.auth.Auth;
+import google.registry.schema.replay.DatastoreEntity;
+import google.registry.schema.replay.DatastoreOnlyEntity;
+import google.registry.schema.replay.NonReplicatedEntity;
+import google.registry.schema.replay.SqlReplayCheckpoint;
+import google.registry.util.RequestStatusChecker;
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.channels.Channels;
+import java.util.Optional;
+import javax.inject.Inject;
+import javax.servlet.http.HttpServletResponse;
+import org.joda.time.DateTime;
+import org.joda.time.Duration;
+
+/** Action that replays commit logs to Cloud SQL to keep it up to date. */
+@Action(
+    service = Action.Service.BACKEND,
+    path = ReplayCommitLogsToSqlAction.PATH,
+    method = Action.Method.POST,
+    automaticallyPrintOk = true,
+    auth = Auth.AUTH_INTERNAL_OR_ADMIN)
+public class ReplayCommitLogsToSqlAction implements Runnable {
+
+  static final String PATH = "/_dr/task/replayCommitLogsToSql";
+
+  private static final int BLOCK_SIZE =
+      1024 * 1024; // Buffer 1mb at a time, for no particular reason.
+  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
+  private static final Duration LEASE_LENGTH = standardHours(1);
+
+  @Inject GcsService gcsService;
+  @Inject Response response;
+  @Inject RequestStatusChecker requestStatusChecker;
+  @Inject GcsDiffFileLister diffLister;
+
+  @Inject
+  ReplayCommitLogsToSqlAction() {}
+
+  @Override
+  public void run() {
+    if (!RegistryConfig.getCloudSqlReplayCommitLogs()) {
+      String message = "ReplayCommitLogsToSqlAction was called but disabled in the config.";
+      logger.atWarning().log(message);
+      // App Engine will retry on any non-2xx status code, which we don't want in this case.
+      response.setStatus(SC_NO_CONTENT);
+      response.setPayload(message);
+      return;
+    }
+    Optional<Lock> lock =
+        Lock.acquire(
+            this.getClass().getSimpleName(), null, LEASE_LENGTH, requestStatusChecker, false);
+    if (lock.isEmpty()) {
+      String message = "Can't acquire SQL commit log replay lock, aborting.";
+      logger.atSevere().log(message);
+      // App Engine will retry on any non-2xx status code, which we don't want in this case.
+      // Let the next run after the next export happen naturally.
+      response.setStatus(SC_NO_CONTENT);
+      response.setPayload(message);
+      return;
+    }
+    try {
+      replayFiles();
+      response.setStatus(HttpServletResponse.SC_OK);
+      logger.atInfo().log("ReplayCommitLogsToSqlAction completed successfully.");
+    } finally {
+      lock.ifPresent(Lock::release);
+    }
+  }
+
+  private void replayFiles() {
+    // Start at the first millisecond we haven't seen yet
+    DateTime fromTime = jpaTm().transact(() -> SqlReplayCheckpoint.get().plusMillis(1));
+    // If there's an inconsistent file set, this will throw IllegalStateException and the job
+    // will try later -- this is likely because an export hasn't finished yet.
+    ImmutableList<GcsFileMetadata> commitLogFiles =
+        diffLister.listDiffFiles(fromTime, /* current time */ null);
+    for (GcsFileMetadata metadata : commitLogFiles) {
+      // One transaction per GCS file
+      jpaTm().transact(() -> processFile(metadata));
+    }
+    logger.atInfo().log("Replayed %d commit log files to SQL successfully.", commitLogFiles.size());
+  }
+
+  private void processFile(GcsFileMetadata metadata) {
+    try (InputStream input =
+        Channels.newInputStream(
+            gcsService.openPrefetchingReadChannel(metadata.getFilename(), 0, BLOCK_SIZE))) {
+      // Load and process the Datastore transactions one at a time
+      ImmutableList<ImmutableList<VersionedEntity>> allTransactions =
+          CommitLogImports.loadEntitiesByTransaction(input);
+      allTransactions.forEach(this::replayTransaction);
+      // if we succeeded, set the last-seen time
+      DateTime checkpoint =
+          DateTime.parse(
+              metadata.getFilename().getObjectName().substring(DIFF_FILE_PREFIX.length()));
+      SqlReplayCheckpoint.set(checkpoint);
+      logger.atInfo().log("Replayed %d transactions from commit log file.", allTransactions.size());
+    } catch (IOException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  private void replayTransaction(ImmutableList<VersionedEntity> transaction) {
+    transaction.stream()
+        .sorted(ReplayCommitLogsToSqlAction::compareByWeight)
+        .forEach(
+            versionedEntity ->
+                versionedEntity
+                    .getEntity()
+                    .ifPresentOrElse(
+                        this::handleEntityPut, () -> handleEntityDelete(versionedEntity)));
+  }
+
+  private void handleEntityPut(Entity entity) {
+    Object ofyPojo = ofy().toPojo(entity);
+    if (ofyPojo instanceof DatastoreEntity) {
+      DatastoreEntity datastoreEntity = (DatastoreEntity) ofyPojo;
+      datastoreEntity.toSqlEntity().ifPresent(jpaTm()::put);
+    } else {
+      // this should never happen, but we shouldn't fail on it
+      logger.atSevere().log(
+          "%s does not implement DatastoreEntity, which is necessary for SQL replay.",
+          ofyPojo.getClass());
+    }
+  }
+
+  private void handleEntityDelete(VersionedEntity entityToDelete) {
+    Key key = entityToDelete.key();
+    VKey<?> entityVKey;
+    try {
+      entityVKey = VKeyTranslatorFactory.createVKey(key);
+    } catch (RuntimeException e) {
+      // This means that the key wasn't convertible to VKey through the standard methods or via
+      // a createVKey method. This means that the object isn't persisted in SQL so we ignore it.
+      logger.atInfo().log(
+          "Skipping SQL delete for kind %s since it is not convertible.", key.getKind());
+      return;
+    }
+    Class<?> entityClass = entityVKey.getKind();
+    // Delete the key iff the class represents a JPA entity that is replicated
+    if (!NonReplicatedEntity.class.isAssignableFrom(entityClass)
+        && !DatastoreOnlyEntity.class.isAssignableFrom(entityClass)
+        && entityClass.getAnnotation(javax.persistence.Entity.class) != null) {
+      jpaTm().delete(entityVKey);
+    }
+  }
+
+  private static int compareByWeight(VersionedEntity a, VersionedEntity b) {
+    return getEntityPriority(a.key().getKind(), a.getEntity().isEmpty())
+        - getEntityPriority(b.key().getKind(), b.getEntity().isEmpty());
+  }
+}

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

@@ -20,6 +20,7 @@ import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.base.Throwables.throwIfUnchecked;
 import static google.registry.beam.initsql.BackupPaths.getCommitLogTimestamp;
 import static google.registry.beam.initsql.BackupPaths.getExportFilePatterns;
+import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.JpaRetries.isFailedTxnRetriable;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.setJpaTm;
@@ -68,6 +69,7 @@ import org.apache.beam.sdk.transforms.MapElements;
 import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.transforms.ParDo;
 import org.apache.beam.sdk.transforms.ProcessFunction;
+import org.apache.beam.sdk.transforms.SerializableFunction;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.PBegin;
 import org.apache.beam.sdk.values.PCollection;
@@ -264,9 +266,9 @@ public final class Transforms {
   }
 
   /**
-   * Returns a {@link PTransform} that writes a {@link PCollection} of entities to a SQL database.
-   * and outputs an empty {@code PCollection<Void>}. This allows other operations to {@link
-   * org.apache.beam.sdk.transforms.Wait wait} for the completion of this transform.
+   * Returns a {@link PTransform} that writes a {@link PCollection} of {@link VersionedEntity}s to a
+   * SQL database. and outputs an empty {@code PCollection<Void>}. This allows other operations to
+   * {@link org.apache.beam.sdk.transforms.Wait wait} for the completion of this transform.
    *
    * <p>Errors are handled according to the pipeline runner's default policy. As part of a one-time
    * job, we will not add features unless proven necessary.
@@ -282,16 +284,53 @@ public final class Transforms {
       int maxWriters,
       int batchSize,
       SerializableSupplier<JpaTransactionManager> jpaSupplier) {
-    return new PTransform<PCollection<VersionedEntity>, PCollection<Void>>() {
+    return writeToSql(
+        transformId,
+        maxWriters,
+        batchSize,
+        jpaSupplier,
+        (e) -> ofy().toPojo(e.getEntity().get()),
+        TypeDescriptor.of(VersionedEntity.class));
+  }
+
+  /**
+   * Returns a {@link PTransform} that writes a {@link PCollection} of entities to a SQL database.
+   * and outputs an empty {@code PCollection<Void>}. This allows other operations to {@link
+   * org.apache.beam.sdk.transforms.Wait wait} for the completion of this transform.
+   *
+   * <p>The converter and type descriptor are generics so that we can convert any type of entity to
+   * an object to be placed in SQL.
+   *
+   * <p>Errors are handled according to the pipeline runner's default policy. As part of a one-time
+   * job, we will not add features unless proven necessary.
+   *
+   * @param transformId a unique ID for an instance of the returned transform
+   * @param maxWriters the max number of concurrent writes to SQL, which also determines the max
+   *     number of connection pools created
+   * @param batchSize the number of entities to write in each operation
+   * @param jpaSupplier supplier of a {@link JpaTransactionManager}
+   * @param jpaConverter the function that converts the input object to a JPA entity
+   * @param objectDescriptor the type descriptor of the input object
+   */
+  public static <T> PTransform<PCollection<T>, PCollection<Void>> writeToSql(
+      String transformId,
+      int maxWriters,
+      int batchSize,
+      SerializableSupplier<JpaTransactionManager> jpaSupplier,
+      SerializableFunction<T, Object> jpaConverter,
+      TypeDescriptor<T> objectDescriptor) {
+    return new PTransform<PCollection<T>, PCollection<Void>>() {
       @Override
-      public PCollection<Void> expand(PCollection<VersionedEntity> input) {
+      public PCollection<Void> expand(PCollection<T> input) {
         return input
             .apply(
                 "Shard data for " + transformId,
-                MapElements.into(kvs(integers(), TypeDescriptor.of(VersionedEntity.class)))
+                MapElements.into(kvs(integers(), objectDescriptor))
                     .via(ve -> KV.of(ThreadLocalRandom.current().nextInt(maxWriters), ve)))
             .apply("Batch output by shard " + transformId, GroupIntoBatches.ofSize(batchSize))
-            .apply("Write in batch for " + transformId, ParDo.of(new SqlBatchWriter(jpaSupplier)));
+            .apply(
+                "Write in batch for " + transformId,
+                ParDo.of(new SqlBatchWriter<T>(jpaSupplier, jpaConverter)));
       }
     };
   }
@@ -385,18 +424,22 @@ public final class Transforms {
    * to hold the {@code JpaTransactionManager} instance, we must ensure that JpaTransactionManager
    * is not changed or torn down while being used by some instance.
    */
-  private static class SqlBatchWriter extends DoFn<KV<Integer, Iterable<VersionedEntity>>, Void> {
+  private static class SqlBatchWriter<T> extends DoFn<KV<Integer, Iterable<T>>, Void> {
 
     private static int instanceCount = 0;
     private static JpaTransactionManager originalJpa;
 
     private final SerializableSupplier<JpaTransactionManager> jpaSupplier;
+    private final SerializableFunction<T, Object> jpaConverter;
 
     private transient Ofy ofy;
     private transient SystemSleeper sleeper;
 
-    SqlBatchWriter(SerializableSupplier<JpaTransactionManager> jpaSupplier) {
+    SqlBatchWriter(
+        SerializableSupplier<JpaTransactionManager> jpaSupplier,
+        SerializableFunction<T, Object> jpaConverter) {
       this.jpaSupplier = jpaSupplier;
+      this.jpaConverter = jpaConverter;
     }
 
     @Setup
@@ -429,13 +472,11 @@ public final class Transforms {
     }
 
     @ProcessElement
-    public void processElement(@Element KV<Integer, Iterable<VersionedEntity>> kv) {
+    public void processElement(@Element KV<Integer, Iterable<T>> kv) {
       try (AppEngineEnvironment env = new AppEngineEnvironment()) {
         ImmutableList<Object> ofyEntities =
             Streams.stream(kv.getValue())
-                .map(VersionedEntity::getEntity)
-                .map(Optional::get)
-                .map(ofy::toPojo)
+                .map(this.jpaConverter::apply)
                 .collect(ImmutableList.toImmutableList());
         retry(() -> jpaTm().transact(() -> jpaTm().putAll(ofyEntities)));
       }

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

@@ -20,7 +20,7 @@ import static google.registry.beam.BeamUtils.getQueryFromFile;
 import com.google.auth.oauth2.GoogleCredentials;
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableSet;
-import google.registry.backup.AppEngineEnvironment;
+import google.registry.beam.initsql.Transforms;
 import google.registry.beam.initsql.Transforms.SerializableSupplier;
 import google.registry.beam.spec11.SafeBrowsingTransforms.EvaluateSafeBrowsingFn;
 import google.registry.config.CredentialModule.LocalCredential;
@@ -43,7 +43,6 @@ import org.apache.beam.sdk.options.Description;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
 import org.apache.beam.sdk.options.ValueProvider;
 import org.apache.beam.sdk.options.ValueProvider.NestedValueProvider;
-import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.GroupByKey;
 import org.apache.beam.sdk.transforms.MapElements;
 import org.apache.beam.sdk.transforms.ParDo;
@@ -191,34 +190,27 @@ public class Spec11Pipeline implements Serializable {
       PCollection<Subdomain> domains,
       EvaluateSafeBrowsingFn evaluateSafeBrowsingFn,
       ValueProvider<String> dateProvider) {
-
     PCollection<KV<Subdomain, ThreatMatch>> subdomainsSql =
         domains.apply("Run through SafeBrowsing API", ParDo.of(evaluateSafeBrowsingFn));
-    /* Store ThreatMatch objects in SQL. */
+    TypeDescriptor<KV<Subdomain, ThreatMatch>> descriptor =
+        new TypeDescriptor<KV<Subdomain, ThreatMatch>>() {};
     subdomainsSql.apply(
-        ParDo.of(
-            new DoFn<KV<Subdomain, ThreatMatch>, Void>() {
-              @ProcessElement
-              public void processElement(ProcessContext context) {
-                // create the Spec11ThreatMatch from Subdomain and ThreatMatch
-                try (AppEngineEnvironment env = new AppEngineEnvironment()) {
-                  Subdomain subdomain = context.element().getKey();
-                  Spec11ThreatMatch threatMatch =
-                      new Spec11ThreatMatch.Builder()
-                          .setThreatTypes(
-                              ImmutableSet.of(
-                                  ThreatType.valueOf(context.element().getValue().threatType())))
-                          .setCheckDate(
-                              LocalDate.parse(dateProvider.get(), ISODateTimeFormat.date()))
-                          .setDomainName(subdomain.domainName())
-                          .setDomainRepoId(subdomain.domainRepoId())
-                          .setRegistrarId(subdomain.registrarId())
-                          .build();
-                  JpaTransactionManager jpaTransactionManager = jpaSupplierFactory.get();
-                  jpaTransactionManager.transact(() -> jpaTransactionManager.insert(threatMatch));
-                }
-              }
-            }));
+        Transforms.writeToSql(
+            "Spec11ThreatMatch",
+            4,
+            4,
+            jpaSupplierFactory,
+            (kv) -> {
+              Subdomain subdomain = kv.getKey();
+              return new Spec11ThreatMatch.Builder()
+                  .setThreatTypes(ImmutableSet.of(ThreatType.valueOf(kv.getValue().threatType())))
+                  .setCheckDate(LocalDate.parse(dateProvider.get(), ISODateTimeFormat.date()))
+                  .setDomainName(subdomain.domainName())
+                  .setDomainRepoId(subdomain.domainRepoId())
+                  .setRegistrarId(subdomain.registrarId())
+                  .build();
+            },
+            descriptor));
 
     /* Store ThreatMatch objects in JSON. */
     PCollection<KV<Subdomain, ThreatMatch>> subdomainsJson =

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

@@ -1592,6 +1592,22 @@ public final class RegistryConfig {
     CONFIG_SETTINGS.get().cloudSql.replicateTransactions = replicateTransactions;
   }
 
+  /**
+   * Returns whether or not to replay commit logs to the SQL database after export to GCS.
+   *
+   * <p>If true, we will trigger the {@link google.registry.backup.ReplayCommitLogsToSqlAction}
+   * after the {@link google.registry.backup.ExportCommitLogDiffAction} to load the commit logs and
+   * replay them to SQL.
+   */
+  public static boolean getCloudSqlReplayCommitLogs() {
+    return CONFIG_SETTINGS.get().cloudSql.replayCommitLogs;
+  }
+
+  @VisibleForTesting
+  public static void overrideCloudSqlReplayCommitLogs(boolean replayCommitLogs) {
+    CONFIG_SETTINGS.get().cloudSql.replayCommitLogs = replayCommitLogs;
+  }
+
   /** Returns the roid suffix to be used for the roids of all contacts and hosts. */
   public static String getContactAndHostRoidSuffix() {
     return CONFIG_SETTINGS.get().registryPolicy.contactAndHostRoidSuffix;

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

@@ -126,6 +126,7 @@ public class RegistryConfigSettings {
     public String username;
     public String instanceConnectionName;
     public boolean replicateTransactions;
+    public boolean replayCommitLogs;
   }
 
   /** Configuration for Apache Beam (Cloud Dataflow). */

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

@@ -233,6 +233,8 @@ cloudSql:
   # Set this to true to replicate cloud SQL transactions to datastore in the
   # background.
   replicateTransactions: false
+  # Set this to true to enable replay of commit logs to SQL
+  replayCommitLogs: false
 
 cloudDns:
   # Set both properties to null in Production.

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

@@ -208,6 +208,12 @@
     <max-concurrent-requests>5</max-concurrent-requests>
   </queue>
 
+  <!-- Queue for replaying commit logs to SQL during the transition from Datastore -> SQL. -->
+  <queue>
+    <name>replay-commit-logs-to-sql</name>
+    <rate>1/s</rate>
+  </queue>
+
   <!-- The load[0-9] queues are used for load-testing, and can be safely deleted
        in any environment that doesn't require load-testing. -->
   <queue>

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

@@ -74,6 +74,14 @@ public class EppRequestHandler {
           && eppOutput.getResponse().getResult().getCode() == SUCCESS_AND_CLOSE) {
         response.setHeader("Epp-Session", "close");
       }
+      // If a login request returns a success, a logged-in header is added to the response to inform
+      // the proxy that it is no longer necessary to send the full client certificate to the backend
+      // for this connection.
+      if (eppOutput.isResponse()
+          && eppOutput.getResponse().isLoginResponse()
+          && eppOutput.isSuccess()) {
+        response.setHeader("Logged-In", "true");
+      }
     } catch (Exception e) {
       logger.atWarning().withCause(e).log("handleEppCommand general exception");
       response.setStatus(SC_BAD_REQUEST);

core/src/main/java/google/registry/flows/session/LoginFlow.java

@@ -141,7 +141,7 @@ public class LoginFlow implements Flow {
     sessionMetadata.resetFailedLoginAttempts();
     sessionMetadata.setClientId(login.getClientId());
     sessionMetadata.setServiceExtensionUris(serviceExtensionUrisBuilder.build());
-    return responseBuilder.build();
+    return responseBuilder.setIsLoginResponse().build();
   }
 
   /** Registrar with this client ID could not be found. */

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

@@ -211,11 +211,6 @@ public final class OteAccountBuilder {
     return transformRegistrars(builder -> builder.setPassword(password));
   }
 
-  /** Sets the client certificate hash to all the OT&E Registrars. */
-  public OteAccountBuilder setCertificateHash(String certHash) {
-    return transformRegistrars(builder -> builder.setClientCertificateHash(certHash));
-  }
-
   /** Sets the client certificate to all the OT&E Registrars. */
   public OteAccountBuilder setCertificate(String asciiCert, DateTime now) {
     return transformRegistrars(builder -> builder.setClientCertificate(asciiCert, now));

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

@@ -290,7 +290,7 @@ public abstract class BillingEvent extends ImmutableObject
         @javax.persistence.Index(columnList = "allocationToken")
       })
   @AttributeOverride(name = "id", column = @Column(name = "billing_event_id"))
-  @WithLongVKey
+  @WithLongVKey(compositeKey = true)
   public static class OneTime extends BillingEvent implements DatastoreAndSqlEntity {
 
     /** The billable value. */
@@ -464,7 +464,7 @@ public abstract class BillingEvent extends ImmutableObject
         @javax.persistence.Index(columnList = "recurrence_time_of_year")
       })
   @AttributeOverride(name = "id", column = @Column(name = "billing_recurrence_id"))
-  @WithLongVKey
+  @WithLongVKey(compositeKey = true)
   public static class Recurring extends BillingEvent implements DatastoreAndSqlEntity {
 
     /**
@@ -559,7 +559,7 @@ public abstract class BillingEvent extends ImmutableObject
         @javax.persistence.Index(columnList = "billingTime")
       })
   @AttributeOverride(name = "id", column = @Column(name = "billing_cancellation_id"))
-  @WithLongVKey
+  @WithLongVKey(compositeKey = true)
   public static class Cancellation extends BillingEvent implements DatastoreAndSqlEntity {
 
     /** The billing time of the charge that is being cancelled. */
@@ -680,7 +680,7 @@ public abstract class BillingEvent extends ImmutableObject
   /** An event representing a modification of an existing one-time billing event. */
   @ReportedOn
   @Entity
-  @WithLongVKey
+  @WithLongVKey(compositeKey = true)
   public static class Modification extends BillingEvent implements DatastoreOnlyEntity {
 
     /** The change in cost that should be applied to the original billing event. */

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

@@ -14,7 +14,6 @@
 
 package google.registry.model.contact;
 
-import com.google.common.collect.ImmutableList;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.EntitySubclass;
 import google.registry.model.ImmutableObject;
@@ -111,8 +110,8 @@ public class ContactHistory extends HistoryEntry implements SqlEntity {
 
   // In Datastore, save as a HistoryEntry object regardless of this object's type
   @Override
-  public ImmutableList<DatastoreEntity> toDatastoreEntities() {
-    return ImmutableList.of(asHistoryEntry());
+  public Optional<DatastoreEntity> toDatastoreEntity() {
+    return Optional.of(asHistoryEntry());
   }
 
   /** Class to represent the composite primary key of {@link ContactHistory} entity. */

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

@@ -17,7 +17,6 @@ package google.registry.model.domain;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static google.registry.util.CollectionUtils.nullToEmptyImmutableCopy;
 
-import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.EntitySubclass;
@@ -250,8 +249,8 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
 
   // In Datastore, save as a HistoryEntry object regardless of this object's type
   @Override
-  public ImmutableList<DatastoreEntity> toDatastoreEntities() {
-    return ImmutableList.of(asHistoryEntry());
+  public Optional<DatastoreEntity> toDatastoreEntity() {
+    return Optional.of(asHistoryEntry());
   }
 
   /** Class to represent the composite primary key of {@link DomainHistory} entity. */

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

@@ -14,11 +14,11 @@
 
 package google.registry.model.domain.secdns;
 
-import com.google.common.collect.ImmutableList;
 import google.registry.model.domain.DomainHistory;
 import google.registry.model.ofy.ObjectifyService;
 import google.registry.schema.replay.DatastoreEntity;
 import google.registry.schema.replay.SqlEntity;
+import java.util.Optional;
 import javax.persistence.Access;
 import javax.persistence.AccessType;
 import javax.persistence.Column;
@@ -86,7 +86,7 @@ public class DomainDsDataHistory extends DomainDsDataBase implements SqlEntity {
   }
 
   @Override
-  public ImmutableList<DatastoreEntity> toDatastoreEntities() {
-    return ImmutableList.of(); // not persisted in Datastore
+  public Optional<DatastoreEntity> toDatastoreEntity() {
+    return Optional.empty(); // Not persisted in Datastore
   }
 }

core/src/main/java/google/registry/model/eppoutput/EppResponse.java

@@ -65,6 +65,7 @@ import javax.xml.bind.annotation.XmlElement;
 import javax.xml.bind.annotation.XmlElementRef;
 import javax.xml.bind.annotation.XmlElementRefs;
 import javax.xml.bind.annotation.XmlElementWrapper;
+import javax.xml.bind.annotation.XmlTransient;
 import javax.xml.bind.annotation.XmlType;
 
 /**
@@ -87,6 +88,9 @@ public class EppResponse extends ImmutableObject implements ResponseOrGreeting {
   /** The command result. The RFC allows multiple failure results, but we always return one. */
   Result result;
 
+  /** Indicates if this response is for a login request. */
+  @XmlTransient boolean isLoginResponse = false;
+
   /**
    * Information about messages queued for retrieval. This may appear in response to any EPP message
    * (if messages are queued), but in practice this will only be set in response to a poll request.
@@ -178,6 +182,10 @@ public class EppResponse extends ImmutableObject implements ResponseOrGreeting {
     return result;
   }
 
+  public boolean isLoginResponse() {
+    return isLoginResponse;
+  }
+
   /** Marker interface for types that can go in the {@link #resData} field. */
   public interface ResponseData {}
 
@@ -222,5 +230,10 @@ public class EppResponse extends ImmutableObject implements ResponseOrGreeting {
       getInstance().extensions = forceEmptyToNull(extensions);
       return this;
     }
+
+    public Builder setIsLoginResponse() {
+      getInstance().isLoginResponse = true;
+      return this;
+    }
   }
 }

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

@@ -14,7 +14,6 @@
 
 package google.registry.model.host;
 
-import com.google.common.collect.ImmutableList;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.EntitySubclass;
 import google.registry.model.ImmutableObject;
@@ -112,8 +111,8 @@ public class HostHistory extends HistoryEntry implements SqlEntity {
 
   // In Datastore, save as a HistoryEntry object regardless of this object's type
   @Override
-  public ImmutableList<DatastoreEntity> toDatastoreEntities() {
-    return ImmutableList.of(asHistoryEntry());
+  public Optional<DatastoreEntity> toDatastoreEntity() {
+    return Optional.of(asHistoryEntry());
   }
 
   /** Class to represent the composite primary key of {@link HostHistory} entity. */

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

@@ -0,0 +1,60 @@
+// 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.ofy;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.ImmutableMap;
+
+/**
+ * Contains the mapping from class names to SQL-replay-write priorities.
+ *
+ * <p>When replaying Datastore commit logs to SQL (asynchronous replication), in order to avoid
+ * issues with foreign keys, we should replay entity writes so that foreign key references are
+ * always written after the entity that they reference. This class represents that DAG, where lower
+ * values represent an earlier write (and later delete). Higher-valued classes can have foreign keys
+ * on lower-valued classes, but not vice versa.
+ */
+public class EntityWritePriorities {
+
+  /**
+   * Mapping from class name to "priority".
+   *
+   * <p>Here, "priority" means the order in which the class should be inserted / updated in a
+   * transaction with respect to instances of other classes. By default, all classes have a priority
+   * number of zero.
+   *
+   * <p>For each transaction, classes should be written in priority order from the lowest number to
+   * the highest, in order to maintain foreign-key write consistency. For the same reason, deletes
+   * should happen after all writes.
+   */
+  static final ImmutableMap<String, Integer> CLASS_PRIORITIES =
+      ImmutableMap.of(
+          "HistoryEntry", -10,
+          "AllocationToken", -9,
+          "ContactResource", 5,
+          "DomainBase", 10);
+
+  // The beginning of the range of priority numbers reserved for delete.  This must be greater than
+  // any of the values in CLASS_PRIORITIES by enough overhead to accommodate  any negative values in
+  // it. Note: by design, deletions will happen in the opposite order of insertions, which is
+  // necessary to make sure foreign keys aren't violated during deletion.
+  @VisibleForTesting static final int DELETE_RANGE = Integer.MAX_VALUE / 2;
+
+  /** Returns the priority of the entity type in the map entry. */
+  public static int getEntityPriority(String kind, boolean isDelete) {
+    int priority = CLASS_PRIORITIES.getOrDefault(kind, 0);
+    return isDelete ? DELETE_RANGE - priority : priority;
+  }
+}

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

@@ -20,6 +20,7 @@ import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static com.google.common.collect.Maps.filterValues;
 import static com.google.common.collect.Maps.toMap;
 import static google.registry.model.ofy.CommitLogBucket.getArbitraryBucketId;
+import static google.registry.model.ofy.EntityWritePriorities.getEntityPriority;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 
@@ -29,7 +30,6 @@ import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
 import google.registry.persistence.VKey;
 import google.registry.schema.replay.DatastoreEntity;
-import google.registry.schema.replay.SqlEntity;
 import java.util.Map;
 import org.joda.time.DateTime;
 
@@ -61,7 +61,7 @@ class TransactionInfo {
 
   TransactionInfo(DateTime now) {
     this.transactionTime = now;
-    ofy().load().key(bucketKey);  // Asynchronously load value into session cache.
+    ofy().load().key(bucketKey); // Asynchronously load value into session cache.
   }
 
   TransactionInfo setReadOnly() {
@@ -101,23 +101,10 @@ class TransactionInfo {
         .collect(toImmutableSet());
   }
 
-  // Mapping from class name to "weight" (which in this case is the order in which the class must
-  // be "put" in a transaction with respect to instances of other classes).  Lower weight classes
-  // are put first, by default all classes have a weight of zero.
-  static final ImmutableMap<String, Integer> CLASS_WEIGHTS =
-      ImmutableMap.of(
-          "HistoryEntry", -1,
-          "DomainBase", 1);
-
-  // The beginning of the range of weights reserved for delete.  This must be greater than any of
-  // the values in CLASS_WEIGHTS by enough overhead to accomodate any negative values in it.
-  @VisibleForTesting static final int DELETE_RANGE = Integer.MAX_VALUE / 2;
-
   /** Returns the weight of the entity type in the map entry. */
   @VisibleForTesting
   static int getWeight(ImmutableMap.Entry<Key<?>, Object> entry) {
-    int weight = CLASS_WEIGHTS.getOrDefault(entry.getKey().getKind(), 0);
-    return entry.getValue().equals(Delete.SENTINEL) ? DELETE_RANGE - weight : weight;
+    return getEntityPriority(entry.getKey().getKind(), entry.getValue().equals(Delete.SENTINEL));
   }
 
   private static int compareByWeight(
@@ -137,10 +124,9 @@ class TransactionInfo {
                         if (entry.getValue().equals(Delete.SENTINEL)) {
                           jpaTm().delete(VKey.from(entry.getKey()));
                         } else {
-                          for (SqlEntity entity :
-                              ((DatastoreEntity) entry.getValue()).toSqlEntities()) {
-                            jpaTm().put(entity);
-                          }
+                          ((DatastoreEntity) entry.getValue())
+                              .toSqlEntity()
+                              .ifPresent(jpaTm()::put);
                         }
                       });
             });

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

@@ -279,7 +279,7 @@ public abstract class PollMessage extends ImmutableObject
   @EntitySubclass(index = false)
   @javax.persistence.Entity
   @DiscriminatorValue("ONE_TIME")
-  @WithLongVKey
+  @WithLongVKey(compositeKey = true)
   public static class OneTime extends PollMessage {
 
     // Response data. Objectify cannot persist a base class type, so we must have a separate field
@@ -432,7 +432,7 @@ public abstract class PollMessage extends ImmutableObject
   @EntitySubclass(index = false)
   @javax.persistence.Entity
   @DiscriminatorValue("AUTORENEW")
-  @WithLongVKey
+  @WithLongVKey(compositeKey = true)
   public static class Autorenew extends PollMessage {
 
     /** The target id of the autorenew event. */

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

@@ -856,26 +856,6 @@ public class Registrar extends ImmutableObject
       }
     }
 
-    /**
-     * Sets client certificate hash, but not the certificate.
-     *
-     * <p><b>Warning:</b> {@link #setClientCertificate(String, DateTime)} sets the hash for you and
-     * is preferred. Calling this method will nullify the {@code clientCertificate} field.
-     */
-    public Builder setClientCertificateHash(String clientCertificateHash) {
-      if (clientCertificateHash != null) {
-        checkArgument(
-            Pattern.matches("[A-Za-z0-9+/]+", clientCertificateHash),
-            "--cert_hash not a valid base64 (no padding) value");
-        checkArgument(
-            base64().decode(clientCertificateHash).length == 256 / 8,
-            "--cert_hash base64 does not decode to 256 bits");
-      }
-      getInstance().clientCertificate = null;
-      getInstance().clientCertificateHash = clientCertificateHash;
-      return this;
-    }
-
     public Builder setContactsRequireSyncing(boolean contactsRequireSyncing) {
       getInstance().contactsRequireSyncing = contactsRequireSyncing;
       return this;

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

@@ -18,7 +18,6 @@ import static com.googlecode.objectify.Key.getKind;
 import static google.registry.util.CollectionUtils.nullToEmptyImmutableCopy;
 
 import com.google.common.annotations.VisibleForTesting;
-import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Entity;
@@ -45,6 +44,7 @@ import google.registry.model.host.HostResource;
 import google.registry.persistence.VKey;
 import google.registry.schema.replay.DatastoreEntity;
 import google.registry.schema.replay.SqlEntity;
+import java.util.Optional;
 import java.util.Set;
 import javax.annotation.Nullable;
 import javax.persistence.Access;
@@ -314,8 +314,8 @@ public class HistoryEntry extends ImmutableObject implements Buildable, Datastor
 
   // In SQL, save the child type
   @Override
-  public ImmutableList<SqlEntity> toSqlEntities() {
-    return ImmutableList.of((SqlEntity) toChildHistoryEntity());
+  public Optional<SqlEntity> toSqlEntity() {
+    return Optional.of((SqlEntity) toChildHistoryEntity());
   }
 
   /** Creates a {@link VKey} instance from a {@link Key} instance. */

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

@@ -18,13 +18,13 @@ import static com.google.common.base.Preconditions.checkArgument;
 import static google.registry.util.CollectionUtils.isNullOrEmpty;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 
-import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import google.registry.model.Buildable;
 import google.registry.model.ImmutableObject;
 import google.registry.schema.replay.DatastoreEntity;
 import google.registry.schema.replay.SqlEntity;
 import google.registry.util.DomainNameUtils;
+import java.util.Optional;
 import java.util.Set;
 import javax.persistence.Column;
 import javax.persistence.Entity;
@@ -111,8 +111,8 @@ public class Spec11ThreatMatch extends ImmutableObject implements Buildable, Sql
   }
 
   @Override
-  public ImmutableList<DatastoreEntity> toDatastoreEntities() {
-    return ImmutableList.of(); // not stored in Datastore
+  public Optional<DatastoreEntity> toDatastoreEntity() {
+    return Optional.empty(); // Not persisted in Datastore
   }
 
   @Override

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

@@ -20,7 +20,6 @@ import static google.registry.persistence.transaction.TransactionManagerFactory.
 import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
-import com.google.common.collect.ImmutableList;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Entity;
 import google.registry.model.annotations.NotBackedUp;
@@ -28,9 +27,7 @@ import google.registry.model.annotations.NotBackedUp.Reason;
 import google.registry.model.common.CrossTldSingleton;
 import google.registry.model.tmch.TmchCrl.TmchCrlId;
 import google.registry.persistence.VKey;
-import google.registry.schema.replay.DatastoreEntity;
 import google.registry.schema.replay.NonReplicatedEntity;
-import google.registry.schema.replay.SqlEntity;
 import java.io.Serializable;
 import java.util.Optional;
 import javax.annotation.concurrent.Immutable;
@@ -106,16 +103,6 @@ public final class TmchCrl extends CrossTldSingleton implements NonReplicatedEnt
     return updated;
   }
 
-  @Override
-  public ImmutableList<SqlEntity> toSqlEntities() {
-    return ImmutableList.of(); // dually-written
-  }
-
-  @Override
-  public ImmutableList<DatastoreEntity> toDatastoreEntities() {
-    return ImmutableList.of(); // dually-written
-  }
-
   static class TmchCrlId implements Serializable {
 
     @Column(name = "certificateRevocations")

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

@@ -31,4 +31,12 @@ public @interface WithLongVKey {
    * class name will be "VKeyConverter_" concatenated with the suffix.
    */
   String classNameSuffix() default "";
+
+  /**
+   * Set to true if this is a composite vkey.
+   *
+   * <p>For composite VKeys, we don't attempt to define an objectify key when loading from SQL: the
+   * enclosing class has to take care of that.
+   */
+  boolean compositeKey() default false;
 }

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

@@ -31,4 +31,12 @@ public @interface WithStringVKey {
    * class name will be "VKeyConverter_" concatenated with the suffix.
    */
   String classNameSuffix() default "";
+
+  /**
+   * Set to true if this is a composite vkey.
+   *
+   * <p>For composite VKeys, we don't attempt to define an objectify key when loading from SQL: the
+   * enclosing class has to take care of that.
+   */
+  boolean compositeKey() default false;
 }

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

@@ -14,6 +14,7 @@
 
 package google.registry.persistence.converter;
 
+import com.googlecode.objectify.Key;
 import google.registry.persistence.VKey;
 import javax.annotation.Nullable;
 import javax.persistence.AttributeConverter;
@@ -29,7 +30,28 @@ public abstract class VKeyConverter<T, C> implements AttributeConverter<VKey<? e
   @Override
   @Nullable
   public VKey<? extends T> convertToEntityAttribute(@Nullable C dbData) {
-    return dbData == null ? null : VKey.createSql(getAttributeClass(), dbData);
+    if (dbData == null) {
+      return null;
+    }
+    Class<? extends T> clazz = getAttributeClass();
+    Key ofyKey = null;
+    if (!hasCompositeOfyKey()) {
+      // If this isn't a composite key, we can create the Ofy key from the SQL key.
+      ofyKey =
+          dbData instanceof String
+              ? Key.create(clazz, (String) dbData)
+              : Key.create(clazz, (Long) dbData);
+      return VKey.create(clazz, dbData, ofyKey);
+    } else {
+      // We don't know how to create the Ofy key and probably don't have everything necessary to do
+      // it anyway, so just create an asymmetric key - the containing object will have to convert it
+      // into a symmetric key.
+      return VKey.createSql(clazz, dbData);
+    }
+  }
+
+  protected boolean hasCompositeOfyKey() {
+    return false;
   }
 
   /** Returns the class of the attribute. */

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

@@ -14,9 +14,9 @@
 
 package google.registry.persistence.transaction;
 
-import com.google.common.collect.ImmutableList;
 import google.registry.schema.replay.DatastoreEntity;
 import google.registry.schema.replay.SqlEntity;
+import java.util.Optional;
 import javax.persistence.Entity;
 import javax.persistence.GeneratedValue;
 import javax.persistence.GenerationType;
@@ -45,7 +45,7 @@ public class TransactionEntity implements SqlEntity {
   }
 
   @Override
-  public ImmutableList<DatastoreEntity> toDatastoreEntities() {
-    return ImmutableList.of(); // not stored in Datastore per se
+  public Optional<DatastoreEntity> toDatastoreEntity() {
+    return Optional.empty(); // Not persisted in Datastore per se
   }
 }

core/src/main/java/google/registry/schema/cursor/Cursor.java

@@ -16,7 +16,6 @@ package google.registry.schema.cursor;
 
 import static com.google.appengine.api.search.checkers.Preconditions.checkNotNull;
 
-import com.google.common.collect.ImmutableList;
 import google.registry.model.ImmutableObject;
 import google.registry.model.UpdateAutoTimestamp;
 import google.registry.model.common.Cursor.CursorType;
@@ -26,6 +25,7 @@ import google.registry.schema.replay.SqlEntity;
 import google.registry.util.DateTimeUtils;
 import java.io.Serializable;
 import java.time.ZonedDateTime;
+import java.util.Optional;
 import javax.persistence.Column;
 import javax.persistence.Entity;
 import javax.persistence.EnumType;
@@ -104,8 +104,8 @@ public class Cursor implements SqlEntity {
   }
 
   @Override
-  public ImmutableList<DatastoreEntity> toDatastoreEntities() {
-    return ImmutableList.of(); // Cursors are not converted since they are ephemeral
+  public Optional<DatastoreEntity> toDatastoreEntity() {
+    return Optional.empty(); // Cursors are not converted since they are ephemeral
   }
 
   static class CursorId extends ImmutableObject implements Serializable {

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

@@ -19,7 +19,6 @@ import static google.registry.util.DateTimeUtils.isBeforeOrAt;
 import static google.registry.util.DateTimeUtils.toZonedDateTime;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 
-import com.google.common.collect.ImmutableList;
 import google.registry.model.Buildable;
 import google.registry.model.CreateAutoTimestamp;
 import google.registry.model.ImmutableObject;
@@ -234,8 +233,8 @@ public final class RegistryLock extends ImmutableObject implements Buildable, Sq
   }
 
   @Override
-  public ImmutableList<DatastoreEntity> toDatastoreEntities() {
-    return ImmutableList.of(); // not stored in Datastore
+  public Optional<DatastoreEntity> toDatastoreEntity() {
+    return Optional.empty(); // Not persisted in Datastore
   }
 
   /** Builder for {@link google.registry.schema.domain.RegistryLock}. */

core/src/main/java/google/registry/schema/replay/DatastoreAndSqlEntity.java

@@ -14,18 +14,18 @@
 
 package google.registry.schema.replay;
 
-import com.google.common.collect.ImmutableList;
+import java.util.Optional;
 
 /** An entity that has the same Java object representation in SQL and Datastore. */
 public interface DatastoreAndSqlEntity extends DatastoreEntity, SqlEntity {
 
   @Override
-  default ImmutableList<DatastoreEntity> toDatastoreEntities() {
-    return ImmutableList.of(this);
+  default Optional<DatastoreEntity> toDatastoreEntity() {
+    return Optional.of(this);
   }
 
   @Override
-  default ImmutableList<SqlEntity> toSqlEntities() {
-    return ImmutableList.of(this);
+  default Optional<SqlEntity> toSqlEntity() {
+    return Optional.of(this);
   }
 }

core/src/main/java/google/registry/schema/replay/DatastoreEntity.java

@@ -14,7 +14,7 @@
 
 package google.registry.schema.replay;
 
-import com.google.common.collect.ImmutableList;
+import java.util.Optional;
 
 /**
  * An object that can be stored in Datastore and serialized using Objectify's {@link
@@ -26,5 +26,5 @@ import com.google.common.collect.ImmutableList;
  */
 public interface DatastoreEntity {
 
-  ImmutableList<SqlEntity> toSqlEntities();
+  Optional<SqlEntity> toSqlEntity();
 }

core/src/main/java/google/registry/schema/replay/DatastoreOnlyEntity.java

@@ -14,11 +14,12 @@
 
 package google.registry.schema.replay;
 
-import com.google.common.collect.ImmutableList;
+import java.util.Optional;
 
+/** An entity that is only stored in Datastore, that should not be replayed to SQL. */
 public interface DatastoreOnlyEntity extends DatastoreEntity {
   @Override
-  default ImmutableList<SqlEntity> toSqlEntities() {
-    return ImmutableList.of();
+  default Optional<SqlEntity> toSqlEntity() {
+    return Optional.empty();
   }
 }

core/src/main/java/google/registry/schema/replay/NonReplicatedEntity.java

@@ -14,7 +14,7 @@
 
 package google.registry.schema.replay;
 
-import com.google.common.collect.ImmutableList;
+import java.util.Optional;
 
 /**
  * Represents an entity that should not participate in asynchronous replication.
@@ -24,12 +24,12 @@ import com.google.common.collect.ImmutableList;
 public interface NonReplicatedEntity extends DatastoreEntity, SqlEntity {
 
   @Override
-  default ImmutableList<DatastoreEntity> toDatastoreEntities() {
-    return ImmutableList.of();
+  default Optional<DatastoreEntity> toDatastoreEntity() {
+    return Optional.empty();
   }
 
   @Override
-  default ImmutableList<SqlEntity> toSqlEntities() {
-    return ImmutableList.of();
+  default Optional<SqlEntity> toSqlEntity() {
+    return Optional.empty();
   }
 }

core/src/main/java/google/registry/schema/replay/SqlEntity.java

@@ -14,7 +14,7 @@
 
 package google.registry.schema.replay;
 
-import com.google.common.collect.ImmutableList;
+import java.util.Optional;
 
 /**
  * An object that can be stored in Cloud SQL using {@link
@@ -25,5 +25,5 @@ import com.google.common.collect.ImmutableList;
  */
 public interface SqlEntity {
 
-  ImmutableList<DatastoreEntity> toDatastoreEntities();
+  Optional<DatastoreEntity> toDatastoreEntity();
 }

core/src/main/java/google/registry/schema/replay/SqlReplayCheckpoint.java

@@ -18,7 +18,7 @@ import static google.registry.model.common.CrossTldSingleton.SINGLETON_ID;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
 
-import com.google.common.collect.ImmutableList;
+import java.util.Optional;
 import javax.persistence.Entity;
 import javax.persistence.Id;
 import org.joda.time.DateTime;
@@ -32,8 +32,8 @@ public class SqlReplayCheckpoint implements SqlEntity {
   private DateTime lastReplayTime;
 
   @Override
-  public ImmutableList<DatastoreEntity> toDatastoreEntities() {
-    return ImmutableList.of(); // not necessary to persist in Datastore
+  public Optional<DatastoreEntity> toDatastoreEntity() {
+    return Optional.empty(); // Not necessary to persist in Datastore
   }
 
   public static DateTime get() {

core/src/main/java/google/registry/schema/server/Lock.java

@@ -16,7 +16,6 @@ package google.registry.schema.server;
 
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 
-import com.google.common.collect.ImmutableList;
 import google.registry.model.ImmutableObject;
 import google.registry.schema.replay.DatastoreEntity;
 import google.registry.schema.replay.SqlEntity;
@@ -24,6 +23,7 @@ import google.registry.schema.server.Lock.LockId;
 import google.registry.util.DateTimeUtils;
 import java.io.Serializable;
 import java.time.ZonedDateTime;
+import java.util.Optional;
 import javax.persistence.Column;
 import javax.persistence.Entity;
 import javax.persistence.Id;
@@ -120,8 +120,8 @@ public class Lock implements SqlEntity {
   }
 
   @Override
-  public ImmutableList<DatastoreEntity> toDatastoreEntities() {
-    return ImmutableList.of(); // Locks are not converted since they are ephemeral
+  public Optional<DatastoreEntity> toDatastoreEntity() {
+    return Optional.empty(); // Locks are not converted since they are ephemeral
   }
 
   static class LockId extends ImmutableObject implements Serializable {

core/src/main/java/google/registry/schema/tld/PremiumEntry.java

@@ -14,13 +14,13 @@
 
 package google.registry.schema.tld;
 
-import com.google.common.collect.ImmutableList;
 import google.registry.model.ImmutableObject;
 import google.registry.model.registry.label.PremiumList;
 import google.registry.schema.replay.DatastoreEntity;
 import google.registry.schema.replay.SqlEntity;
 import java.io.Serializable;
 import java.math.BigDecimal;
+import java.util.Optional;
 import javax.persistence.Column;
 import javax.persistence.Entity;
 import javax.persistence.Id;
@@ -47,7 +47,7 @@ public class PremiumEntry extends ImmutableObject implements Serializable, SqlEn
   private PremiumEntry() {}
 
   @Override
-  public ImmutableList<DatastoreEntity> toDatastoreEntities() {
-    return ImmutableList.of(); // PremiumList is dually-written
+  public Optional<DatastoreEntity> toDatastoreEntity() {
+    return Optional.empty(); // PremiumList is dually-written
   }
 }

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

@@ -138,15 +138,6 @@ abstract class CreateOrUpdateRegistrarCommand extends MutatingCommand {
       validateWith = PathParameter.InputFile.class)
   Path clientCertificateFilename;
 
-  @Nullable
-  @Parameter(
-    names = "--cert_hash",
-    description =
-        "Hash of client certificate (SHA256 base64 no padding). Do not use this unless "
-            + "you want to store ONLY the hash and not the full certificate"
-  )
-  String clientCertificateHash;
-
   @Nullable
   @Parameter(
       names = "--failover_cert_file",
@@ -375,14 +366,6 @@ abstract class CreateOrUpdateRegistrarCommand extends MutatingCommand {
         }
         builder.setFailoverClientCertificate(asciiCert, now);
       }
-      if (!isNullOrEmpty(clientCertificateHash)) {
-        checkArgument(clientCertificateFilename == null,
-            "Can't specify both --cert_hash and --cert_file");
-        if ("null".equals(clientCertificateHash)) {
-          clientCertificateHash = null;
-        }
-        builder.setClientCertificateHash(clientCertificateHash);
-      }
       if (ianaId != null) {
         builder.setIanaIdentifier(ianaId.orElse(null));
       }

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

@@ -45,10 +45,10 @@ import google.registry.model.domain.DomainBase;
 @Parameters(
     separators = " =",
     commandDescription = "Dedupe BillingEvent.OneTime entities with duplicate IDs.")
-public class DedupeOneTimeBillingEventIdsCommand extends DedupeEntityIdsCommand<OneTime> {
+public class DedupeOneTimeBillingEventIdsCommand extends ReadEntityFromKeyPathCommand<OneTime> {
 
   @Override
-  void dedupe(OneTime entity) {
+  void process(OneTime entity) {
     Key<BillingEvent> key = Key.create(entity);
     Key<DomainBase> domainKey = getGrandParentAsDomain(key);
     DomainBase domain = ofy().load().key(domainKey).now();

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

@@ -54,11 +54,10 @@ import java.util.Set;
 @Parameters(
     separators = " =",
     commandDescription = "Dedupe BillingEvent.Recurring entities with duplicate IDs.")
-public class DedupeRecurringBillingEventIdsCommand
-    extends DedupeEntityIdsCommand<BillingEvent.Recurring> {
+public class DedupeRecurringBillingEventIdsCommand extends ReadEntityFromKeyPathCommand<Recurring> {
 
   @Override
-  void dedupe(Recurring recurring) {
+  void process(Recurring recurring) {
     // Loads the associated DomainBase and BillingEvent.OneTime entities that
     // may have reference to this BillingEvent.Recurring entity.
     Key<DomainBase> domainKey = getGrandParentAsDomain(Key.create(recurring));

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

@@ -33,8 +33,13 @@ import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.util.List;
 
-/** Base Command to dedupe entities with duplicate IDs. */
-abstract class DedupeEntityIdsCommand<T> extends MutatingCommand {
+/**
+ * Base Command to read entities from Datastore by their key paths retrieved from BigQuery.
+ *
+ * <p>The key path is the value of column __key__.path of the entity's BigQuery table. Its value is
+ * converted from the entity's key.
+ */
+abstract class ReadEntityFromKeyPathCommand<T> extends MutatingCommand {
 
   @Parameter(
       names = "--key_paths_file",
@@ -48,7 +53,7 @@ abstract class DedupeEntityIdsCommand<T> extends MutatingCommand {
 
   private StringBuilder changeMessage = new StringBuilder();
 
-  abstract void dedupe(T entity);
+  abstract void process(T entity);
 
   @Override
   protected void init() throws Exception {
@@ -69,7 +74,7 @@ abstract class DedupeEntityIdsCommand<T> extends MutatingCommand {
       }
       Class<T> clazz = new TypeInstantiator<T>(getClass()) {}.getExactType();
       if (clazz.isInstance(entity)) {
-        dedupe((T) entity);
+        process((T) entity);
       } else {
         throw new IllegalArgumentException("Unsupported entity key: " + untypedKey);
       }

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

@@ -99,6 +99,7 @@ public final class RegistryTool {
           .put("populate_null_registrar_fields", PopulateNullRegistrarFieldsCommand.class)
           .put("registrar_contact", RegistrarContactCommand.class)
           .put("remove_ip_address", RemoveIpAddressCommand.class)
+          .put("remove_registry_one_key", RemoveRegistryOneKeyCommand.class)
           .put("renew_domain", RenewDomainCommand.class)
           .put("resave_entities", ResaveEntitiesCommand.class)
           .put("resave_environment_entities", ResaveEnvironmentEntitiesCommand.class)

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

@@ -0,0 +1,69 @@
+// 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 com.beust.jcommander.Parameters;
+import com.googlecode.objectify.Key;
+import google.registry.model.domain.DomainBase;
+import google.registry.persistence.VKey;
+import javax.annotation.Nullable;
+import org.joda.time.DateTime;
+
+/** Command to remove the Registry 1.0 key in {@link DomainBase} entity. */
+@Parameters(separators = " =", commandDescription = "Remove .")
+public class RemoveRegistryOneKeyCommand extends ReadEntityFromKeyPathCommand<DomainBase> {
+
+  @Override
+  void process(DomainBase entity) {
+    // Assert that the DomainBase entity must be deleted before 2017-08-01(most of the problematic
+    // entities were deleted before 2017, though there are still a few entities deleted in 2017-07).
+    // This is because we finished the Registry 2.0 migration in 2017 and should not generate any
+    // Registry 1.0 key after it.
+    if (!entity.getDeletionTime().isBefore(DateTime.parse("2017-08-01T00:00:00Z"))) {
+      throw new IllegalStateException(
+          String.format(
+              "Entity's deletion time %s is not before 2017-08-01T00:00:00Z",
+              entity.getDeletionTime()));
+    }
+    boolean hasChange = false;
+    DomainBase.Builder domainBuilder = entity.asBuilder();
+    // We only found the registry 1.0 key existed in fields autorenewBillingEvent,
+    // autorenewPollMessage and deletePollMessage so we just need to check these fields for each
+    // entity.
+    if (isRegistryOneKey(entity.getAutorenewBillingEvent())) {
+      domainBuilder.setAutorenewBillingEvent(null);
+      hasChange = true;
+    }
+    if (isRegistryOneKey(entity.getAutorenewPollMessage())) {
+      domainBuilder.setAutorenewPollMessage(null);
+      hasChange = true;
+    }
+    if (isRegistryOneKey(entity.getDeletePollMessage())) {
+      domainBuilder.setDeletePollMessage(null);
+      hasChange = true;
+    }
+    if (hasChange) {
+      stageEntityChange(entity, domainBuilder.build());
+    }
+  }
+
+  private static boolean isRegistryOneKey(@Nullable VKey<?> vKey) {
+    if (vKey == null || vKey.getOfyKey() == null || vKey.getOfyKey().getParent() == null) {
+      return false;
+    }
+    Key<?> parentKey = vKey.getOfyKey().getParent();
+    return parentKey.getKind().equals("EntityGroupRoot") && parentKey.getName().equals("per-tld");
+  }
+}

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

@@ -65,13 +65,6 @@ final class SetupOteCommand extends ConfirmingCommand implements CommandWithRemo
       validateWith = PathParameter.InputFile.class)
   private Path certFile;
 
-  @Parameter(
-      names = {"-h", "--certhash"},
-      description =
-          "Hash of client certificate (SHA256 base64 no padding). Do not use this unless "
-              + "you want to store ONLY the hash and not the full certificate.")
-  private String certHash;
-
   @Parameter(
       names = {"--overwrite"},
       description = "Whether to replace existing entities if we encounter any, instead of failing.")
@@ -89,9 +82,7 @@ final class SetupOteCommand extends ConfirmingCommand implements CommandWithRemo
   /** Run any pre-execute command checks */
   @Override
   protected void init() throws Exception {
-    checkArgument(
-        certFile == null ^ certHash == null,
-        "Must specify exactly one of client certificate file or client certificate hash.");
+    checkArgument(certFile != null, "Must specify exactly one client certificate file.");
 
     password = passwordGenerator.createString(PASSWORD_LENGTH);
     oteAccountBuilder =
@@ -101,16 +92,10 @@ final class SetupOteCommand extends ConfirmingCommand implements CommandWithRemo
             .setIpAllowList(ipAllowList)
             .setReplaceExisting(overwrite);
 
-    if (certFile != null) {
       String asciiCert = MoreFiles.asCharSource(certFile, US_ASCII).read();
       // Don't wait for create_registrar to fail if it's a bad certificate file.
       loadCertificate(asciiCert);
       oteAccountBuilder.setCertificate(asciiCert, clock.nowUtc());
-    }
-
-    if (certHash != null) {
-      oteAccountBuilder.setCertificateHash(certHash);
-    }
   }
 
   @Override

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

@@ -0,0 +1,458 @@
+// 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.backup;
+
+import static com.google.common.collect.ImmutableList.toImmutableList;
+import static com.google.common.truth.Truth.assertThat;
+import static com.google.common.util.concurrent.MoreExecutors.newDirectExecutorService;
+import static google.registry.backup.RestoreCommitLogsActionTest.GCS_BUCKET;
+import static google.registry.backup.RestoreCommitLogsActionTest.createCheckpoint;
+import static google.registry.backup.RestoreCommitLogsActionTest.saveDiffFile;
+import static google.registry.backup.RestoreCommitLogsActionTest.saveDiffFileNotToRestore;
+import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
+import static google.registry.model.ofy.CommitLogBucket.getBucketKey;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.testing.DatabaseHelper.createTld;
+import static google.registry.testing.DatabaseHelper.newDomainBase;
+import static google.registry.testing.DatabaseHelper.persistActiveContact;
+import static javax.servlet.http.HttpServletResponse.SC_NO_CONTENT;
+import static javax.servlet.http.HttpServletResponse.SC_OK;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+import com.google.appengine.tools.cloudstorage.GcsService;
+import com.google.appengine.tools.cloudstorage.GcsServiceFactory;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.truth.Truth8;
+import com.googlecode.objectify.Key;
+import google.registry.config.RegistryConfig;
+import google.registry.model.common.Cursor;
+import google.registry.model.common.Cursor.CursorType;
+import google.registry.model.contact.ContactResource;
+import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.GracePeriod;
+import google.registry.model.domain.secdns.DelegationSignerData;
+import google.registry.model.ofy.CommitLogBucket;
+import google.registry.model.ofy.CommitLogManifest;
+import google.registry.model.ofy.CommitLogMutation;
+import google.registry.model.registrar.RegistrarContact;
+import google.registry.model.registry.label.ReservedList;
+import google.registry.model.server.Lock;
+import google.registry.model.tmch.ClaimsListShard;
+import google.registry.model.translators.VKeyTranslatorFactory;
+import google.registry.persistence.VKey;
+import google.registry.persistence.transaction.JpaTransactionManager;
+import google.registry.persistence.transaction.TransactionManagerFactory;
+import google.registry.schema.replay.SqlReplayCheckpoint;
+import google.registry.testing.AppEngineExtension;
+import google.registry.testing.FakeClock;
+import google.registry.testing.FakeResponse;
+import google.registry.testing.TestObject;
+import google.registry.util.RequestStatusChecker;
+import java.io.IOException;
+import org.joda.time.DateTime;
+import org.joda.time.Duration;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.mockito.ArgumentCaptor;
+import org.mockito.InOrder;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.junit.jupiter.MockitoExtension;
+
+/** Tests for {@link ReplayCommitLogsToSqlAction}. */
+@ExtendWith(MockitoExtension.class)
+public class ReplayCommitLogsToSqlActionTest {
+
+  private final FakeClock fakeClock = new FakeClock(DateTime.parse("2000-01-01TZ"));
+
+  @RegisterExtension
+  public final AppEngineExtension appEngine =
+      AppEngineExtension.builder()
+          .withDatastoreAndCloudSql()
+          .withClock(fakeClock)
+          .withOfyTestEntities(TestObject.class)
+          .withJpaUnitTestEntities(
+              RegistrarContact.class,
+              TestObject.class,
+              SqlReplayCheckpoint.class,
+              ContactResource.class,
+              DomainBase.class,
+              GracePeriod.class,
+              DelegationSignerData.class)
+          .build();
+
+  /** Local GCS service. */
+  private final GcsService gcsService = GcsServiceFactory.createGcsService();
+
+  private final ReplayCommitLogsToSqlAction action = new ReplayCommitLogsToSqlAction();
+  private final FakeResponse response = new FakeResponse();
+  @Mock private RequestStatusChecker requestStatusChecker;
+
+  @BeforeAll
+  static void beforeAll() {
+    VKeyTranslatorFactory.addTestEntityClass(TestObject.class);
+  }
+
+  @BeforeEach
+  void beforeEach() {
+    action.gcsService = gcsService;
+    action.response = response;
+    action.requestStatusChecker = requestStatusChecker;
+    action.diffLister = new GcsDiffFileLister();
+    action.diffLister.gcsService = gcsService;
+    action.diffLister.gcsBucket = GCS_BUCKET;
+    action.diffLister.executor = newDirectExecutorService();
+    RegistryConfig.overrideCloudSqlReplayCommitLogs(true);
+  }
+
+  @Test
+  void testReplay_multipleDiffFiles() throws Exception {
+    jpaTm()
+        .transact(
+            () -> {
+              jpaTm().insertWithoutBackup(TestObject.create("previous to keep"));
+              jpaTm().insertWithoutBackup(TestObject.create("previous to delete"));
+            });
+    DateTime now = fakeClock.nowUtc();
+    // Create 3 transactions, across two diff files.
+    // Before: {"previous to keep", "previous to delete"}
+    // 1a: Add {"a", "b"}, Delete {"previous to delete"}
+    // 1b: Add {"c", "d"}, Delete {"a"}
+    // 2:  Add {"e", "f"}, Delete {"c"}
+    // After:  {"previous to keep", "b", "d", "e", "f"}
+    Key<CommitLogManifest> manifest1aKey =
+        CommitLogManifest.createKey(getBucketKey(1), now.minusMinutes(3));
+    Key<CommitLogManifest> manifest1bKey =
+        CommitLogManifest.createKey(getBucketKey(2), now.minusMinutes(2));
+    Key<CommitLogManifest> manifest2Key =
+        CommitLogManifest.createKey(getBucketKey(1), now.minusMinutes(1));
+    saveDiffFileNotToRestore(gcsService, now.minusMinutes(2));
+    saveDiffFile(
+        gcsService,
+        createCheckpoint(now.minusMinutes(1)),
+        CommitLogManifest.create(
+            getBucketKey(1),
+            now.minusMinutes(3),
+            ImmutableSet.of(Key.create(TestObject.create("previous to delete")))),
+        CommitLogMutation.create(manifest1aKey, TestObject.create("a")),
+        CommitLogMutation.create(manifest1aKey, TestObject.create("b")),
+        CommitLogManifest.create(
+            getBucketKey(2),
+            now.minusMinutes(2),
+            ImmutableSet.of(Key.create(TestObject.create("a")))),
+        CommitLogMutation.create(manifest1bKey, TestObject.create("c")),
+        CommitLogMutation.create(manifest1bKey, TestObject.create("d")));
+    saveDiffFile(
+        gcsService,
+        createCheckpoint(now),
+        CommitLogManifest.create(
+            getBucketKey(1),
+            now.minusMinutes(1),
+            ImmutableSet.of(Key.create(TestObject.create("c")))),
+        CommitLogMutation.create(manifest2Key, TestObject.create("e")),
+        CommitLogMutation.create(manifest2Key, TestObject.create("f")));
+    jpaTm().transact(() -> SqlReplayCheckpoint.set(now.minusMinutes(1).minusMillis(1)));
+    fakeClock.advanceOneMilli();
+    runAndAssertSuccess(now);
+    assertExpectedIds("previous to keep", "b", "d", "e", "f");
+  }
+
+  @Test
+  void testReplay_noManifests() throws Exception {
+    DateTime now = fakeClock.nowUtc();
+    jpaTm().transact(() -> jpaTm().insertWithoutBackup(TestObject.create("previous to keep")));
+    saveDiffFileNotToRestore(gcsService, now.minusMinutes(1));
+    saveDiffFile(gcsService, createCheckpoint(now.minusMillis(2)));
+    jpaTm().transact(() -> SqlReplayCheckpoint.set(now.minusMillis(1)));
+    runAndAssertSuccess(now.minusMillis(1));
+    assertExpectedIds("previous to keep");
+  }
+
+  @Test
+  void testReplay_manifestWithNoDeletions() throws Exception {
+    DateTime now = fakeClock.nowUtc();
+    jpaTm().transact(() -> jpaTm().insertWithoutBackup(TestObject.create("previous to keep")));
+    Key<CommitLogBucket> bucketKey = getBucketKey(1);
+    Key<CommitLogManifest> manifestKey = CommitLogManifest.createKey(bucketKey, now);
+    saveDiffFileNotToRestore(gcsService, now.minusMinutes(2));
+    jpaTm().transact(() -> SqlReplayCheckpoint.set(now.minusMinutes(1).minusMillis(1)));
+    saveDiffFile(
+        gcsService,
+        createCheckpoint(now.minusMinutes(1)),
+        CommitLogManifest.create(bucketKey, now, null),
+        CommitLogMutation.create(manifestKey, TestObject.create("a")),
+        CommitLogMutation.create(manifestKey, TestObject.create("b")));
+    runAndAssertSuccess(now.minusMinutes(1));
+    assertExpectedIds("previous to keep", "a", "b");
+  }
+
+  @Test
+  void testReplay_manifestWithNoMutations() throws Exception {
+    DateTime now = fakeClock.nowUtc();
+    jpaTm()
+        .transact(
+            () -> {
+              jpaTm().insertWithoutBackup(TestObject.create("previous to keep"));
+              jpaTm().insertWithoutBackup(TestObject.create("previous to delete"));
+            });
+    saveDiffFileNotToRestore(gcsService, now.minusMinutes(2));
+    jpaTm().transact(() -> SqlReplayCheckpoint.set(now.minusMinutes(1).minusMillis(1)));
+    saveDiffFile(
+        gcsService,
+        createCheckpoint(now.minusMinutes(1)),
+        CommitLogManifest.create(
+            getBucketKey(1),
+            now,
+            ImmutableSet.of(Key.create(TestObject.create("previous to delete")))));
+    runAndAssertSuccess(now.minusMinutes(1));
+    assertExpectedIds("previous to keep");
+  }
+
+  @Test
+  void wtestReplay_mutateExistingEntity() throws Exception {
+    DateTime now = fakeClock.nowUtc();
+    jpaTm().transact(() -> jpaTm().put(TestObject.create("existing", "a")));
+    Key<CommitLogManifest> manifestKey = CommitLogManifest.createKey(getBucketKey(1), now);
+    saveDiffFileNotToRestore(gcsService, now.minusMinutes(1).minusMillis(1));
+    jpaTm().transact(() -> SqlReplayCheckpoint.set(now.minusMinutes(1)));
+    saveDiffFile(
+        gcsService,
+        createCheckpoint(now.minusMillis(1)),
+        CommitLogManifest.create(getBucketKey(1), now, null),
+        CommitLogMutation.create(manifestKey, TestObject.create("existing", "b")));
+    action.run();
+    TestObject fromDatabase =
+        jpaTm().transact(() -> jpaTm().load(VKey.createSql(TestObject.class, "existing")));
+    assertThat(fromDatabase.getField()).isEqualTo("b");
+  }
+
+  // This should be harmless
+  @Test
+  void testReplay_deleteMissingEntity() throws Exception {
+    DateTime now = fakeClock.nowUtc();
+    jpaTm().transact(() -> jpaTm().put(TestObject.create("previous to keep", "a")));
+    saveDiffFileNotToRestore(gcsService, now.minusMinutes(1).minusMillis(1));
+    jpaTm().transact(() -> SqlReplayCheckpoint.set(now.minusMinutes(1)));
+    saveDiffFile(
+        gcsService,
+        createCheckpoint(now.minusMillis(1)),
+        CommitLogManifest.create(
+            getBucketKey(1),
+            now,
+            ImmutableSet.of(Key.create(TestObject.create("previous to delete")))));
+    action.run();
+    assertExpectedIds("previous to keep");
+  }
+
+  @Test
+  @SuppressWarnings({"unchecked", "rawtypes"})
+  void testReplay_properlyWeighted() throws Exception {
+    DateTime now = fakeClock.nowUtc();
+    Key<CommitLogManifest> manifestKey =
+        CommitLogManifest.createKey(getBucketKey(1), now.minusMinutes(1));
+    // Create (but don't save to SQL) a domain + contact
+    createTld("tld");
+    DomainBase domain = newDomainBase("example.tld");
+    CommitLogMutation domainMutation =
+        tm().transact(() -> CommitLogMutation.create(manifestKey, domain));
+    ContactResource contact = tm().transact(() -> tm().load(domain.getRegistrant()));
+    CommitLogMutation contactMutation =
+        tm().transact(() -> CommitLogMutation.create(manifestKey, contact));
+
+    // Create and save to SQL a registrar contact that we will delete
+    RegistrarContact toDelete = AppEngineExtension.makeRegistrarContact1();
+    jpaTm().transact(() -> jpaTm().put(toDelete));
+    jpaTm().transact(() -> SqlReplayCheckpoint.set(now.minusMinutes(1).minusMillis(1)));
+
+    // spy the txn manager so we can see what order things were inserted/removed
+    JpaTransactionManager spy = spy(jpaTm());
+    TransactionManagerFactory.setJpaTm(() -> spy);
+    // Save in the commit logs the domain and contact (in that order) and the token deletion
+    saveDiffFile(
+        gcsService,
+        createCheckpoint(now.minusMinutes(1)),
+        CommitLogManifest.create(
+            getBucketKey(1), now.minusMinutes(1), ImmutableSet.of(Key.create(toDelete))),
+        domainMutation,
+        contactMutation);
+
+    runAndAssertSuccess(now.minusMinutes(1));
+    // Verify two things:
+    // 1. that the contact insert occurred before the domain insert (necessary for FK ordering)
+    //    even though the domain came first in the file
+    // 2. that the allocation token delete occurred after the insertions
+    InOrder inOrder = Mockito.inOrder(spy);
+    inOrder.verify(spy).put(any(ContactResource.class));
+    inOrder.verify(spy).put(any(DomainBase.class));
+    inOrder.verify(spy).delete(toDelete.createVKey());
+    inOrder.verify(spy).put(any(SqlReplayCheckpoint.class));
+  }
+
+  @Test
+  @SuppressWarnings({"unchecked", "rawtypes"})
+  void testReplay_properlyWeighted_doesNotApplyCrossTransactions() throws Exception {
+    DateTime now = fakeClock.nowUtc();
+    Key<CommitLogManifest> manifestKey =
+        CommitLogManifest.createKey(getBucketKey(1), now.minusMinutes(1));
+
+    // Create and save the standard contact
+    ContactResource contact = persistActiveContact("contact1234");
+    jpaTm().transact(() -> jpaTm().put(contact));
+
+    // Simulate a Datastore transaction with a new version of the contact
+    ContactResource contactWithEdit =
+        contact.asBuilder().setEmailAddress("replay@example.tld").build();
+    CommitLogMutation contactMutation =
+        tm().transact(() -> CommitLogMutation.create(manifestKey, contactWithEdit));
+
+    jpaTm().transact(() -> SqlReplayCheckpoint.set(now.minusMinutes(1).minusMillis(1)));
+
+    // spy the txn manager so we can see what order things were inserted
+    JpaTransactionManager spy = spy(jpaTm());
+    TransactionManagerFactory.setJpaTm(() -> spy);
+    // Save two commits -- the deletion, then the new version of the contact
+    saveDiffFile(
+        gcsService,
+        createCheckpoint(now.minusMinutes(1).plusMillis(1)),
+        CommitLogManifest.create(
+            getBucketKey(1), now.minusMinutes(1), ImmutableSet.of(Key.create(contact))),
+        CommitLogManifest.create(
+            getBucketKey(1), now.minusMinutes(1).plusMillis(1), ImmutableSet.of()),
+        contactMutation);
+    runAndAssertSuccess(now.minusMinutes(1).plusMillis(1));
+    // Verify that the delete occurred first (because it was in the first transaction) even though
+    // deletes have higher weight
+    ArgumentCaptor<Object> putCaptor = ArgumentCaptor.forClass(Object.class);
+    InOrder inOrder = Mockito.inOrder(spy);
+    inOrder.verify(spy).delete(contact.createVKey());
+    inOrder.verify(spy).put(putCaptor.capture());
+    assertThat(putCaptor.getValue().getClass()).isEqualTo(ContactResource.class);
+    assertThat(jpaTm().transact(() -> jpaTm().load(contact.createVKey()).getEmailAddress()))
+        .isEqualTo("replay@example.tld");
+  }
+
+  @Test
+  void testSuccess_nonReplicatedEntity_isNotReplayed() {
+    DateTime now = fakeClock.nowUtc();
+
+    // spy the txn manager so we can verify it's never called
+    JpaTransactionManager spy = spy(jpaTm());
+    TransactionManagerFactory.setJpaTm(() -> spy);
+
+    jpaTm().transact(() -> SqlReplayCheckpoint.set(now.minusMinutes(1).minusMillis(1)));
+    Key<CommitLogManifest> manifestKey =
+        CommitLogManifest.createKey(getBucketKey(1), now.minusMinutes(1));
+
+    // Have a commit log with a couple objects that shouldn't be replayed
+    ReservedList reservedList =
+        new ReservedList.Builder().setReservedListMap(ImmutableMap.of()).setName("name").build();
+    Cursor cursor = Cursor.createGlobal(CursorType.RECURRING_BILLING, now.minusHours(1));
+    tm().transact(
+            () -> {
+              try {
+                saveDiffFile(
+                    gcsService,
+                    createCheckpoint(now.minusMinutes(1)),
+                    CommitLogManifest.create(
+                        getBucketKey(1), now.minusMinutes(1), ImmutableSet.of()),
+                    // Reserved list is dually-written non-replicated
+                    CommitLogMutation.create(manifestKey, reservedList),
+                    // Cursors aren't replayed to SQL at all
+                    CommitLogMutation.create(manifestKey, cursor));
+              } catch (IOException e) {
+                throw new RuntimeException(e);
+              }
+            });
+    runAndAssertSuccess(now.minusMinutes(1));
+    // jpaTm()::put should only have been called with the checkpoint
+    verify(spy, times(2)).put(any(SqlReplayCheckpoint.class));
+    verify(spy, times(2)).put(any());
+  }
+
+  @Test
+  void testSuccess_nonReplicatedEntity_isNotDeleted() throws Exception {
+    DateTime now = fakeClock.nowUtc();
+    // spy the txn manager so we can verify it's never called
+    JpaTransactionManager spy = spy(jpaTm());
+    TransactionManagerFactory.setJpaTm(() -> spy);
+
+    jpaTm().transact(() -> SqlReplayCheckpoint.set(now.minusMinutes(1).minusMillis(1)));
+    // Save a couple deletes that aren't propagated to SQL (the objects deleted are irrelevant)
+    Key<ClaimsListShard> claimsListKey = Key.create(ClaimsListShard.class, 1L);
+    saveDiffFile(
+        gcsService,
+        createCheckpoint(now.minusMinutes(1)),
+        CommitLogManifest.create(
+            getBucketKey(1),
+            now.minusMinutes(1),
+            // one object only exists in Datastore, one is dually-written (so isn't replicated)
+            ImmutableSet.of(getCrossTldKey(), claimsListKey)));
+
+    runAndAssertSuccess(now.minusMinutes(1));
+    verify(spy, times(0)).delete(any(VKey.class));
+  }
+
+  @Test
+  void testFailure_notEnabled() {
+    RegistryConfig.overrideCloudSqlReplayCommitLogs(false);
+    action.run();
+    assertThat(response.getStatus()).isEqualTo(SC_NO_CONTENT);
+    assertThat(response.getPayload())
+        .isEqualTo("ReplayCommitLogsToSqlAction was called but disabled in the config.");
+  }
+
+  @Test
+  void testFailure_cannotAcquireLock() {
+    Truth8.assertThat(
+            Lock.acquire(
+                ReplayCommitLogsToSqlAction.class.getSimpleName(),
+                null,
+                Duration.standardHours(1),
+                requestStatusChecker,
+                false))
+        .isPresent();
+    fakeClock.advanceOneMilli();
+    action.run();
+    assertThat(response.getStatus()).isEqualTo(SC_NO_CONTENT);
+    assertThat(response.getPayload())
+        .isEqualTo("Can't acquire SQL commit log replay lock, aborting.");
+  }
+
+  private void runAndAssertSuccess(DateTime expectedCheckpointTime) {
+    action.run();
+    assertThat(response.getStatus()).isEqualTo(SC_OK);
+    assertThat(jpaTm().transact(SqlReplayCheckpoint::get)).isEqualTo(expectedCheckpointTime);
+  }
+
+  private void assertExpectedIds(String... expectedIds) {
+    ImmutableList<String> actualIds =
+        jpaTm()
+            .transact(
+                () ->
+                    jpaTm().loadAll(TestObject.class).stream()
+                        .map(TestObject::getId)
+                        .collect(toImmutableList()));
+    assertThat(actualIds).containsExactlyElementsIn(expectedIds);
+  }
+}

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

@@ -61,7 +61,7 @@ import org.junit.jupiter.api.extension.RegisterExtension;
 /** Unit tests for {@link RestoreCommitLogsAction}. */
 public class RestoreCommitLogsActionTest {
 
-  private static final String GCS_BUCKET = "gcs bucket";
+  static final String GCS_BUCKET = "gcs bucket";
 
   private final DateTime now = DateTime.now(UTC);
   private final RestoreCommitLogsAction action = new RestoreCommitLogsAction();
@@ -89,9 +89,10 @@ public class RestoreCommitLogsActionTest {
 
   @Test
   void testRestore_multipleDiffFiles() throws Exception {
-    ofy().saveWithoutBackup().entities(
-        TestObject.create("previous to keep"),
-        TestObject.create("previous to delete")).now();
+    ofy()
+        .saveWithoutBackup()
+        .entities(TestObject.create("previous to keep"), TestObject.create("previous to delete"))
+        .now();
     // Create 3 transactions, across two diff files.
     // Before: {"previous to keep", "previous to delete"}
     // 1a: Add {"a", "b"}, Delete {"previous to delete"}
@@ -104,29 +105,33 @@ public class RestoreCommitLogsActionTest {
         CommitLogManifest.createKey(getBucketKey(2), now.minusMinutes(2));
     Key<CommitLogManifest> manifest2Key =
         CommitLogManifest.createKey(getBucketKey(1), now.minusMinutes(1));
-    saveDiffFileNotToRestore(now.minusMinutes(2));
-    Iterable<ImmutableObject> file1CommitLogs = saveDiffFile(
-        createCheckpoint(now.minusMinutes(1)),
-        CommitLogManifest.create(
-            getBucketKey(1),
-            now.minusMinutes(3),
-            ImmutableSet.of(Key.create(TestObject.create("previous to delete")))),
-        CommitLogMutation.create(manifest1aKey, TestObject.create("a")),
-        CommitLogMutation.create(manifest1aKey, TestObject.create("b")),
-        CommitLogManifest.create(
-            getBucketKey(2),
-            now.minusMinutes(2),
-            ImmutableSet.of(Key.create(TestObject.create("a")))),
-        CommitLogMutation.create(manifest1bKey, TestObject.create("c")),
-        CommitLogMutation.create(manifest1bKey, TestObject.create("d")));
-    Iterable<ImmutableObject> file2CommitLogs = saveDiffFile(
-        createCheckpoint(now),
-        CommitLogManifest.create(
-            getBucketKey(1),
-            now.minusMinutes(1),
-            ImmutableSet.of(Key.create(TestObject.create("c")))),
-        CommitLogMutation.create(manifest2Key, TestObject.create("e")),
-        CommitLogMutation.create(manifest2Key, TestObject.create("f")));
+    saveDiffFileNotToRestore(gcsService, now.minusMinutes(2));
+    Iterable<ImmutableObject> file1CommitLogs =
+        saveDiffFile(
+            gcsService,
+            createCheckpoint(now.minusMinutes(1)),
+            CommitLogManifest.create(
+                getBucketKey(1),
+                now.minusMinutes(3),
+                ImmutableSet.of(Key.create(TestObject.create("previous to delete")))),
+            CommitLogMutation.create(manifest1aKey, TestObject.create("a")),
+            CommitLogMutation.create(manifest1aKey, TestObject.create("b")),
+            CommitLogManifest.create(
+                getBucketKey(2),
+                now.minusMinutes(2),
+                ImmutableSet.of(Key.create(TestObject.create("a")))),
+            CommitLogMutation.create(manifest1bKey, TestObject.create("c")),
+            CommitLogMutation.create(manifest1bKey, TestObject.create("d")));
+    Iterable<ImmutableObject> file2CommitLogs =
+        saveDiffFile(
+            gcsService,
+            createCheckpoint(now),
+            CommitLogManifest.create(
+                getBucketKey(1),
+                now.minusMinutes(1),
+                ImmutableSet.of(Key.create(TestObject.create("c")))),
+            CommitLogMutation.create(manifest2Key, TestObject.create("e")),
+            CommitLogMutation.create(manifest2Key, TestObject.create("f")));
     action.fromTime = now.minusMinutes(1).minusMillis(1);
     action.run();
     ofy().clearSessionCache();
@@ -139,10 +144,9 @@ public class RestoreCommitLogsActionTest {
 
   @Test
   void testRestore_noManifests() throws Exception {
-    ofy().saveWithoutBackup().entity(
-        TestObject.create("previous to keep")).now();
-    saveDiffFileNotToRestore(now.minusMinutes(1));
-    Iterable<ImmutableObject> commitLogs = saveDiffFile(createCheckpoint(now));
+    ofy().saveWithoutBackup().entity(TestObject.create("previous to keep")).now();
+    saveDiffFileNotToRestore(gcsService, now.minusMinutes(1));
+    Iterable<ImmutableObject> commitLogs = saveDiffFile(gcsService, createCheckpoint(now));
     action.run();
     ofy().clearSessionCache();
     assertExpectedIds("previous to keep");
@@ -156,32 +160,37 @@ public class RestoreCommitLogsActionTest {
     ofy().saveWithoutBackup().entity(TestObject.create("previous to keep")).now();
     Key<CommitLogBucket> bucketKey = getBucketKey(1);
     Key<CommitLogManifest> manifestKey = CommitLogManifest.createKey(bucketKey, now);
-    saveDiffFileNotToRestore(now.minusMinutes(1));
-    Iterable<ImmutableObject> commitLogs = saveDiffFile(
-        createCheckpoint(now),
-        CommitLogManifest.create(bucketKey, now, null),
-        CommitLogMutation.create(manifestKey, TestObject.create("a")),
-        CommitLogMutation.create(manifestKey, TestObject.create("b")));
+    saveDiffFileNotToRestore(gcsService, now.minusMinutes(1));
+    Iterable<ImmutableObject> commitLogs =
+        saveDiffFile(
+            gcsService,
+            createCheckpoint(now),
+            CommitLogManifest.create(bucketKey, now, null),
+            CommitLogMutation.create(manifestKey, TestObject.create("a")),
+            CommitLogMutation.create(manifestKey, TestObject.create("b")));
     action.run();
     ofy().clearSessionCache();
     assertExpectedIds("previous to keep", "a", "b");
     assertInDatastore(commitLogs);
     assertInDatastore(CommitLogCheckpointRoot.create(now));
     assertCommitLogBuckets(ImmutableMap.of(1, now));
-}
+  }
 
   @Test
   void testRestore_manifestWithNoMutations() throws Exception {
-    ofy().saveWithoutBackup().entities(
-        TestObject.create("previous to keep"),
-        TestObject.create("previous to delete")).now();
-    saveDiffFileNotToRestore(now.minusMinutes(1));
-    Iterable<ImmutableObject> commitLogs = saveDiffFile(
-        createCheckpoint(now),
-        CommitLogManifest.create(
-            getBucketKey(1),
-            now,
-            ImmutableSet.of(Key.create(TestObject.create("previous to delete")))));
+    ofy()
+        .saveWithoutBackup()
+        .entities(TestObject.create("previous to keep"), TestObject.create("previous to delete"))
+        .now();
+    saveDiffFileNotToRestore(gcsService, now.minusMinutes(1));
+    Iterable<ImmutableObject> commitLogs =
+        saveDiffFile(
+            gcsService,
+            createCheckpoint(now),
+            CommitLogManifest.create(
+                getBucketKey(1),
+                now,
+                ImmutableSet.of(Key.create(TestObject.create("previous to delete")))));
     action.run();
     ofy().clearSessionCache();
     assertExpectedIds("previous to keep");
@@ -193,12 +202,13 @@ public class RestoreCommitLogsActionTest {
   // This is a pathological case that shouldn't be possible, but we should be robust to it.
   @Test
   void testRestore_manifestWithNoMutationsOrDeletions() throws Exception {
-    ofy().saveWithoutBackup().entities(
-        TestObject.create("previous to keep")).now();
-    saveDiffFileNotToRestore(now.minusMinutes(1));
-    Iterable<ImmutableObject> commitLogs = saveDiffFile(
-        createCheckpoint(now),
-        CommitLogManifest.create(getBucketKey(1), now, null));
+    ofy().saveWithoutBackup().entities(TestObject.create("previous to keep")).now();
+    saveDiffFileNotToRestore(gcsService, now.minusMinutes(1));
+    Iterable<ImmutableObject> commitLogs =
+        saveDiffFile(
+            gcsService,
+            createCheckpoint(now),
+            CommitLogManifest.create(getBucketKey(1), now, null));
     action.run();
     ofy().clearSessionCache();
     assertExpectedIds("previous to keep");
@@ -211,11 +221,13 @@ public class RestoreCommitLogsActionTest {
   void testRestore_mutateExistingEntity() throws Exception {
     ofy().saveWithoutBackup().entity(TestObject.create("existing", "a")).now();
     Key<CommitLogManifest> manifestKey = CommitLogManifest.createKey(getBucketKey(1), now);
-    saveDiffFileNotToRestore(now.minusMinutes(1));
-    Iterable<ImmutableObject> commitLogs = saveDiffFile(
-        createCheckpoint(now),
-        CommitLogManifest.create(getBucketKey(1), now, null),
-        CommitLogMutation.create(manifestKey, TestObject.create("existing", "b")));
+    saveDiffFileNotToRestore(gcsService, now.minusMinutes(1));
+    Iterable<ImmutableObject> commitLogs =
+        saveDiffFile(
+            gcsService,
+            createCheckpoint(now),
+            CommitLogManifest.create(getBucketKey(1), now, null),
+            CommitLogMutation.create(manifestKey, TestObject.create("existing", "b")));
     action.run();
     ofy().clearSessionCache();
     assertThat(ofy().load().entity(TestObject.create("existing")).now().getField()).isEqualTo("b");
@@ -228,13 +240,15 @@ public class RestoreCommitLogsActionTest {
   @Test
   void testRestore_deleteMissingEntity() throws Exception {
     ofy().saveWithoutBackup().entity(TestObject.create("previous to keep", "a")).now();
-    saveDiffFileNotToRestore(now.minusMinutes(1));
-    Iterable<ImmutableObject> commitLogs = saveDiffFile(
-        createCheckpoint(now),
-        CommitLogManifest.create(
-            getBucketKey(1),
-            now,
-            ImmutableSet.of(Key.create(TestObject.create("previous to delete")))));
+    saveDiffFileNotToRestore(gcsService, now.minusMinutes(1));
+    Iterable<ImmutableObject> commitLogs =
+        saveDiffFile(
+            gcsService,
+            createCheckpoint(now),
+            CommitLogManifest.create(
+                getBucketKey(1),
+                now,
+                ImmutableSet.of(Key.create(TestObject.create("previous to delete")))));
     action.run();
     ofy().clearSessionCache();
     assertExpectedIds("previous to keep");
@@ -243,12 +257,13 @@ public class RestoreCommitLogsActionTest {
     assertInDatastore(CommitLogCheckpointRoot.create(now));
   }
 
-  private CommitLogCheckpoint createCheckpoint(DateTime now) {
+  static CommitLogCheckpoint createCheckpoint(DateTime now) {
     return CommitLogCheckpoint.create(now, toMap(getBucketIds(), x -> now));
   }
 
-  private Iterable<ImmutableObject> saveDiffFile(
-      CommitLogCheckpoint checkpoint, ImmutableObject... entities) throws IOException {
+  static Iterable<ImmutableObject> saveDiffFile(
+      GcsService gcsService, CommitLogCheckpoint checkpoint, ImmutableObject... entities)
+      throws IOException {
     DateTime now = checkpoint.getCheckpointTime();
     List<ImmutableObject> allEntities = Lists.asList(checkpoint, entities);
     ByteArrayOutputStream output = new ByteArrayOutputStream();
@@ -264,8 +279,9 @@ public class RestoreCommitLogsActionTest {
     return allEntities;
   }
 
-  private void saveDiffFileNotToRestore(DateTime now) throws Exception {
+  static void saveDiffFileNotToRestore(GcsService gcsService, DateTime now) throws Exception {
     saveDiffFile(
+        gcsService,
         createCheckpoint(now),
         CommitLogManifest.create(getBucketKey(1), now, null),
         CommitLogMutation.create(

core/src/test/java/google/registry/beam/spec11/Spec11PipelineTest.java

@@ -20,7 +20,6 @@ import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.verifyNoMoreInteractions;
 import static org.mockito.Mockito.when;
 import static org.mockito.Mockito.withSettings;
 
@@ -285,8 +284,7 @@ class Spec11PipelineTest {
             .build();
 
     verify(mockJpaTm).transact(any(Runnable.class));
-    verify(mockJpaTm).insert(expected);
-    verifyNoMoreInteractions(mockJpaTm);
+    verify(mockJpaTm).putAll(ImmutableList.of(expected));
   }
 
   /**

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

@@ -501,7 +501,7 @@ class EppLifecycleDomainTest extends EppTestCase {
 
   @Test
   void testEapDomainDeletion_withinAddGracePeriod_eapFeeIsNotRefunded() throws Exception {
-    assertThatCommand("login_valid_fee_extension.xml").hasResponse("generic_success_response.xml");
+    assertThatCommand("login_valid_fee_extension.xml").hasSuccessfulLogin();
     createContacts(DateTime.parse("2000-06-01T00:00:00Z"));
 
     // Set the EAP schedule.
@@ -718,7 +718,7 @@ class EppLifecycleDomainTest extends EppTestCase {
             START_OF_TIME, PREDELEGATION,
             gaDate, GENERAL_AVAILABILITY));
 
-    assertThatCommand("login_valid_fee_extension.xml").hasResponse("generic_success_response.xml");
+    assertThatCommand("login_valid_fee_extension.xml").hasSuccessfulLogin();
 
     assertThatCommand("domain_check_fee_premium.xml")
         .atTime(gaDate.plusDays(1))
@@ -1196,7 +1196,7 @@ class EppLifecycleDomainTest extends EppTestCase {
 
     assertThatLogin("NewRegistrar", "foo-BAR2")
         .atTime(sunriseDate.minusDays(3))
-        .hasResponse("generic_success_response.xml");
+        .hasSuccessfulLogin();
 
     createContactsAndHosts();
 
@@ -1292,7 +1292,7 @@ class EppLifecycleDomainTest extends EppTestCase {
 
     assertThatLogin("NewRegistrar", "foo-BAR2")
         .atTime(sunriseDate.minusDays(3))
-        .hasResponse("generic_success_response.xml");
+        .hasSuccessfulLogin();
 
     createContactsAndHosts();
 

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

@@ -44,13 +44,13 @@ class EppLoginTlsTest extends EppTestCase {
     persistResource(
         loadRegistrar("NewRegistrar")
             .asBuilder()
-            .setClientCertificateHash(CertificateSamples.SAMPLE_CERT_HASH)
+            .setClientCertificate(CertificateSamples.SAMPLE_CERT, DateTime.now(UTC))
             .build());
     // Set a cert for the second registrar, or else any cert will be allowed for login.
     persistResource(
         loadRegistrar("TheRegistrar")
             .asBuilder()
-            .setClientCertificateHash(CertificateSamples.SAMPLE_CERT2_HASH)
+            .setClientCertificate(CertificateSamples.SAMPLE_CERT2, DateTime.now(UTC))
             .build());
   }
 

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

@@ -121,6 +121,10 @@ public class EppTestCase {
       return assertCommandAndResponse(
           inputFilename, inputSubstitutions, outputFilename, outputSubstitutions, now);
     }
+
+    public String hasSuccessfulLogin() throws Exception {
+      return assertLoginCommandAndResponse(inputFilename, inputSubstitutions, null, now);
+    }
   }
 
   protected CommandAsserter assertThatCommand(String inputFilename) {
@@ -137,13 +141,33 @@ public class EppTestCase {
   }
 
   protected void assertThatLoginSucceeds(String clientId, String password) throws Exception {
-    assertThatLogin(clientId, password).hasResponse("generic_success_response.xml");
+    assertThatLogin(clientId, password).hasSuccessfulLogin();
   }
 
   protected void assertThatLogoutSucceeds() throws Exception {
     assertThatCommand("logout.xml").hasResponse("logout_response.xml");
   }
 
+  private String assertLoginCommandAndResponse(
+      String inputFilename,
+      @Nullable Map<String, String> inputSubstitutions,
+      @Nullable Map<String, String> outputSubstitutions,
+      DateTime now)
+      throws Exception {
+    String outputFilename = "generic_success_response.xml";
+    clock.setTo(now);
+    String input = loadFile(EppTestCase.class, inputFilename, inputSubstitutions);
+    String expectedOutput = loadFile(EppTestCase.class, outputFilename, outputSubstitutions);
+    setUpSession();
+    FakeResponse response = executeXmlCommand(input);
+
+    // Check that the logged-in header was added to the response
+    assertThat(response.getHeaders()).isEqualTo(ImmutableMap.of("Logged-In", "true"));
+
+    return verifyAndReturnOutput(
+        response.getPayload(), expectedOutput, inputFilename, outputFilename);
+  }
+
   private String assertCommandAndResponse(
       String inputFilename,
       @Nullable Map<String, String> inputSubstitutions,
@@ -154,6 +178,18 @@ public class EppTestCase {
     clock.setTo(now);
     String input = loadFile(EppTestCase.class, inputFilename, inputSubstitutions);
     String expectedOutput = loadFile(EppTestCase.class, outputFilename, outputSubstitutions);
+    setUpSession();
+    FakeResponse response = executeXmlCommand(input);
+
+    // Checks that the Logged-In header is not in the response. If testing the login command, use
+    // assertLoginCommandAndResponse instead of this method.
+    assertThat(response.getHeaders()).doesNotContainEntry("Logged-In", "true");
+
+    return verifyAndReturnOutput(
+        response.getPayload(), expectedOutput, inputFilename, outputFilename);
+  }
+
+  private void setUpSession() {
     if (sessionMetadata == null) {
       sessionMetadata =
           new HttpSessionMetadata(new FakeHttpSession()) {
@@ -165,7 +201,13 @@ public class EppTestCase {
             }
           };
     }
-    String actualOutput = executeXmlCommand(input);
+  }
+
+  private String verifyAndReturnOutput(
+      String actualOutput, String expectedOutput, String inputFilename, String outputFilename)
+      throws Exception {
+    // Run the resulting xml through the unmarshaller to verify that it was valid.
+    EppXmlTransformer.validateOutput(actualOutput);
     assertXmlEqualsWithMessage(
         expectedOutput,
         actualOutput,
@@ -176,7 +218,7 @@ public class EppTestCase {
     return actualOutput;
   }
 
-  private String executeXmlCommand(String inputXml) throws Exception {
+  private FakeResponse executeXmlCommand(String inputXml) throws Exception {
     EppRequestHandler handler = new EppRequestHandler();
     FakeResponse response = new FakeResponse();
     handler.response = response;
@@ -195,10 +237,7 @@ public class EppTestCase {
         inputXml.getBytes(UTF_8));
     assertThat(response.getStatus()).isEqualTo(SC_OK);
     assertThat(response.getContentType()).isEqualTo(APPLICATION_EPP_XML_UTF8);
-    String result = response.getPayload();
-    // Run the resulting xml through the unmarshaller to verify that it was valid.
-    EppXmlTransformer.validateOutput(result);
-    return result;
+    return response;
   }
 
   EppMetric getRecordedEppMetric() {

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

@@ -96,17 +96,24 @@ import google.registry.model.transfer.DomainTransferData;
 import google.registry.model.transfer.TransferResponse;
 import google.registry.model.transfer.TransferStatus;
 import google.registry.persistence.VKey;
+import google.registry.testing.ReplayExtension;
 import google.registry.testing.TaskQueueHelper.TaskMatcher;
 import java.util.Map;
 import org.joda.money.Money;
 import org.joda.time.DateTime;
 import org.joda.time.Duration;
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Order;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Unit tests for {@link DomainDeleteFlow}. */
 class DomainDeleteFlowTest extends ResourceFlowTestCase<DomainDeleteFlow, DomainBase> {
 
+  @Order(value = Order.DEFAULT - 2)
+  @RegisterExtension
+  final ReplayExtension replayExtension = new ReplayExtension(clock);
+
   private DomainBase domain;
   private HistoryEntry earlierHistoryEntry;
 
@@ -137,6 +144,7 @@ class DomainDeleteFlowTest extends ResourceFlowTestCase<DomainDeleteFlow, Domain
         persistResource(createAutorenewBillingEvent("TheRegistrar").build());
     PollMessage.Autorenew autorenewPollMessage =
         persistResource(createAutorenewPollMessage("TheRegistrar").build());
+
     domain =
         persistResource(
             domain
@@ -144,6 +152,7 @@ class DomainDeleteFlowTest extends ResourceFlowTestCase<DomainDeleteFlow, Domain
                 .setAutorenewBillingEvent(autorenewBillingEvent.createVKey())
                 .setAutorenewPollMessage(autorenewPollMessage.createVKey())
                 .build());
+
     assertTransactionalFlow(true);
   }
 
@@ -151,15 +160,20 @@ class DomainDeleteFlowTest extends ResourceFlowTestCase<DomainDeleteFlow, Domain
     // Persist a linked contact.
     ContactResource contact = persistActiveContact("sh8013");
     domain =
-        newDomainBase(getUniqueIdFromCommand())
-            .asBuilder()
-            .setCreationTimeForTest(TIME_BEFORE_FLOW)
-            .setRegistrant(contact.createVKey())
-            .setRegistrationExpirationTime(expirationTime)
-            .build();
+        persistResource(
+            newDomainBase(getUniqueIdFromCommand())
+                .asBuilder()
+                .setCreationTimeForTest(TIME_BEFORE_FLOW)
+                .setRegistrant(contact.createVKey())
+                .setRegistrationExpirationTime(expirationTime)
+                .build());
     earlierHistoryEntry =
         persistResource(
-            new HistoryEntry.Builder().setType(DOMAIN_CREATE).setParent(domain).build());
+            new HistoryEntry.Builder()
+                .setType(DOMAIN_CREATE)
+                .setParent(domain)
+                .setModificationTime(clock.nowUtc())
+                .build());
   }
 
   private void setUpGracePeriods(GracePeriod... gracePeriods) {
@@ -208,7 +222,14 @@ class DomainDeleteFlowTest extends ResourceFlowTestCase<DomainDeleteFlow, Domain
 
   private void assertAutorenewClosedAndCancellationCreatedFor(
       BillingEvent.OneTime graceBillingEvent, HistoryEntry historyEntryDomainDelete) {
-    DateTime eventTime = clock.nowUtc();
+    assertAutorenewClosedAndCancellationCreatedFor(
+        graceBillingEvent, historyEntryDomainDelete, clock.nowUtc());
+  }
+
+  private void assertAutorenewClosedAndCancellationCreatedFor(
+      BillingEvent.OneTime graceBillingEvent,
+      HistoryEntry historyEntryDomainDelete,
+      DateTime eventTime) {
     assertBillingEvents(
         createAutorenewBillingEvent("TheRegistrar").setRecurrenceEndTime(eventTime).build(),
         graceBillingEvent,
@@ -293,7 +314,11 @@ class DomainDeleteFlowTest extends ResourceFlowTestCase<DomainDeleteFlow, Domain
     setUpSuccessfulTest();
     setUpGracePeriods(
         GracePeriod.create(
-            GracePeriodStatus.ADD, domain.getRepoId(), TIME_BEFORE_FLOW.plusDays(1), "foo", null));
+            GracePeriodStatus.ADD,
+            domain.getRepoId(),
+            TIME_BEFORE_FLOW.plusDays(1),
+            "TheRegistrar",
+            null));
     dryRunFlowAssertResponse(loadFile("generic_success_response.xml"));
   }
 
@@ -398,7 +423,7 @@ class DomainDeleteFlowTest extends ResourceFlowTestCase<DomainDeleteFlow, Domain
             GracePeriodStatus.TRANSFER,
             domain.getRepoId(),
             TIME_BEFORE_FLOW.plusDays(1),
-            "foo",
+            "NewRegistrar",
             null));
     // We should see exactly one poll message, which is for the autorenew 1 month in the future.
     assertPollMessages(createAutorenewPollMessage("TheRegistrar").build());
@@ -730,11 +755,11 @@ class DomainDeleteFlowTest extends ResourceFlowTestCase<DomainDeleteFlow, Domain
             .setNameservers(ImmutableSet.of(host.createVKey()))
             .setDeletionTime(START_OF_TIME)
             .build());
-    clock.advanceOneMilli();
+    DateTime eventTime = clock.nowUtc();
     runFlowAssertResponse(loadFile("generic_success_response.xml"));
     assertDnsTasksEnqueued("example.tld");
     assertAutorenewClosedAndCancellationCreatedFor(
-        graceBillingEvent, getOnlyHistoryEntryOfType(domain, DOMAIN_DELETE));
+        graceBillingEvent, getOnlyHistoryEntryOfType(domain, DOMAIN_DELETE), eventTime);
   }
 
   @Test

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

@@ -26,6 +26,7 @@ import static google.registry.testing.DatabaseHelper.newDomainBase;
 import static google.registry.testing.DatabaseHelper.persistActiveDomain;
 import static google.registry.testing.DatabaseHelper.persistDeletedDomain;
 import static google.registry.testing.DatabaseHelper.persistResource;
+import static google.registry.testing.DatabaseHelper.persistResources;
 import static google.registry.testing.DomainBaseSubject.assertAboutDomains;
 import static google.registry.testing.EppExceptionSubject.assertAboutEppExceptions;
 import static google.registry.testing.HistoryEntrySubject.assertAboutHistoryEntries;
@@ -69,12 +70,17 @@ import google.registry.model.registry.Registry;
 import google.registry.model.reporting.DomainTransactionRecord;
 import google.registry.model.reporting.DomainTransactionRecord.TransactionReportField;
 import google.registry.model.reporting.HistoryEntry;
+import google.registry.testing.ReplayExtension;
 import java.util.Map;
 import org.joda.money.Money;
 import org.joda.time.DateTime;
 import org.joda.time.Duration;
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Order;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.BeforeEachCallback;
+import org.junit.jupiter.api.extension.ExtensionContext;
+import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Unit tests for {@link DomainRenewFlow}. */
 class DomainRenewFlowTest extends ResourceFlowTestCase<DomainRenewFlow, DomainBase> {
@@ -96,50 +102,66 @@ class DomainRenewFlowTest extends ResourceFlowTestCase<DomainRenewFlow, DomainBa
 
   private final DateTime expirationTime = DateTime.parse("2000-04-03T22:00:00.0Z");
 
+  @Order(value = Order.DEFAULT - 3)
+  @RegisterExtension
+  final SetClockExtension setClockExtension = new SetClockExtension();
+
+  @Order(value = Order.DEFAULT - 2)
+  @RegisterExtension
+  final ReplayExtension replayExtension = new ReplayExtension(clock);
+
   @BeforeEach
   void initDomainTest() {
     createTld("tld");
-    clock.setTo(expirationTime.minusMillis(2));
     setEppInput("domain_renew.xml", ImmutableMap.of("DOMAIN", "example.tld", "YEARS", "5"));
   }
 
   private void persistDomain(StatusValue... statusValues) throws Exception {
     DomainBase domain = newDomainBase(getUniqueIdFromCommand());
-    HistoryEntry historyEntryDomainCreate =
-        persistResource(
-            new HistoryEntry.Builder()
-                .setParent(domain)
-                .setType(HistoryEntry.Type.DOMAIN_CREATE)
-                .build());
-    BillingEvent.Recurring autorenewEvent =
-        persistResource(
-            new BillingEvent.Recurring.Builder()
-                .setReason(Reason.RENEW)
-                .setFlags(ImmutableSet.of(Flag.AUTO_RENEW))
-                .setTargetId(getUniqueIdFromCommand())
-                .setClientId("TheRegistrar")
-                .setEventTime(expirationTime)
-                .setRecurrenceEndTime(END_OF_TIME)
-                .setParent(historyEntryDomainCreate)
-                .build());
-    PollMessage.Autorenew autorenewPollMessage =
-        persistResource(
-            new PollMessage.Autorenew.Builder()
-                .setTargetId(getUniqueIdFromCommand())
-                .setClientId("TheRegistrar")
-                .setEventTime(expirationTime)
-                .setAutorenewEndTime(END_OF_TIME)
-                .setMsg("Domain was auto-renewed.")
-                .setParent(historyEntryDomainCreate)
-                .build());
-    persistResource(
-        domain
-            .asBuilder()
-            .setRegistrationExpirationTime(expirationTime)
-            .setStatusValues(ImmutableSet.copyOf(statusValues))
-            .setAutorenewBillingEvent(autorenewEvent.createVKey())
-            .setAutorenewPollMessage(autorenewPollMessage.createVKey())
-            .build());
+    tm().transact(
+            () -> {
+              try {
+                HistoryEntry historyEntryDomainCreate =
+                    new HistoryEntry.Builder()
+                        .setParent(domain)
+                        .setType(HistoryEntry.Type.DOMAIN_CREATE)
+                        .setModificationTime(clock.nowUtc())
+                        .build();
+                BillingEvent.Recurring autorenewEvent =
+                    new BillingEvent.Recurring.Builder()
+                        .setReason(Reason.RENEW)
+                        .setFlags(ImmutableSet.of(Flag.AUTO_RENEW))
+                        .setTargetId(getUniqueIdFromCommand())
+                        .setClientId("TheRegistrar")
+                        .setEventTime(expirationTime)
+                        .setRecurrenceEndTime(END_OF_TIME)
+                        .setParent(historyEntryDomainCreate)
+                        .build();
+                PollMessage.Autorenew autorenewPollMessage =
+                    new PollMessage.Autorenew.Builder()
+                        .setTargetId(getUniqueIdFromCommand())
+                        .setClientId("TheRegistrar")
+                        .setEventTime(expirationTime)
+                        .setAutorenewEndTime(END_OF_TIME)
+                        .setMsg("Domain was auto-renewed.")
+                        .setParent(historyEntryDomainCreate)
+                        .build();
+                DomainBase newDomain =
+                    domain
+                        .asBuilder()
+                        .setRegistrationExpirationTime(expirationTime)
+                        .setStatusValues(ImmutableSet.copyOf(statusValues))
+                        .setAutorenewBillingEvent(autorenewEvent.createVKey())
+                        .setAutorenewPollMessage(autorenewPollMessage.createVKey())
+                        .build();
+                persistResources(
+                    ImmutableSet.of(
+                        historyEntryDomainCreate, autorenewEvent,
+                        autorenewPollMessage, newDomain));
+              } catch (Exception e) {
+                throw new RuntimeException("Error persisting domain", e);
+              }
+            });
     clock.advanceOneMilli();
   }
 
@@ -255,6 +277,7 @@ class DomainRenewFlowTest extends ResourceFlowTestCase<DomainRenewFlow, DomainBa
 
   @Test
   void testSuccess() throws Exception {
+    clock.advanceOneMilli();
     persistDomain();
     doSuccessfulTest(
         "domain_renew_response.xml",
@@ -789,4 +812,21 @@ class DomainRenewFlowTest extends ResourceFlowTestCase<DomainRenewFlow, DomainBa
                 TransactionReportField.netRenewsFieldFromYears(5),
                 1));
   }
+
+  /**
+   * Local extension so we can set the clock correctly prior to using it to define basic entities in
+   * AppEngineExtension.
+   *
+   * <p>This has to happen first, we'll get timestamp inversions if we try to set the clock after
+   * these objects are created.
+   */
+  class SetClockExtension implements BeforeEachCallback {
+    @Override
+    public void beforeEach(ExtensionContext context) {
+      // we have to go far enough back before the expiration time so that the clock advances we do
+      // after each persist don't move us into a grace period.  The current value is likely beyond
+      // what is necessary, but this doesn't hurt.
+      clock.setTo(expirationTime.minusMillis(20));
+    }
+  }
 }

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

@@ -90,11 +90,14 @@ import google.registry.model.host.HostResource;
 import google.registry.model.registry.Registry;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.persistence.VKey;
+import google.registry.testing.ReplayExtension;
 import java.util.Optional;
 import org.joda.money.Money;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Order;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Unit tests for {@link DomainUpdateFlow}. */
 class DomainUpdateFlowTest extends ResourceFlowTestCase<DomainUpdateFlow, DomainBase> {
@@ -112,6 +115,10 @@ class DomainUpdateFlowTest extends ResourceFlowTestCase<DomainUpdateFlow, Domain
   private ContactResource mak21Contact;
   private ContactResource unusedContact;
 
+  @Order(value = Order.DEFAULT - 2)
+  @RegisterExtension
+  final ReplayExtension replayExtension = new ReplayExtension(clock);
+
   @BeforeEach
   void initDomainTest() {
     createTld("tld");
@@ -146,6 +153,7 @@ class DomainUpdateFlowTest extends ResourceFlowTestCase<DomainUpdateFlow, Domain
     persistResource(
         new HistoryEntry.Builder()
             .setType(HistoryEntry.Type.DOMAIN_CREATE)
+            .setModificationTime(clock.nowUtc())
             .setParent(domain)
             .build());
     clock.advanceOneMilli();
@@ -168,6 +176,7 @@ class DomainUpdateFlowTest extends ResourceFlowTestCase<DomainUpdateFlow, Domain
     persistResource(
         new HistoryEntry.Builder()
             .setType(HistoryEntry.Type.DOMAIN_CREATE)
+            .setModificationTime(clock.nowUtc())
             .setParent(domain)
             .build());
     clock.advanceOneMilli();

core/src/test/java/google/registry/flows/session/LoginFlowTestCase.java

@@ -14,6 +14,7 @@
 
 package google.registry.flows.session;
 
+import static com.google.common.truth.Truth.assertThat;
 import static google.registry.testing.DatabaseHelper.deleteResource;
 import static google.registry.testing.DatabaseHelper.loadRegistrar;
 import static google.registry.testing.DatabaseHelper.persistResource;
@@ -32,6 +33,7 @@ import google.registry.flows.session.LoginFlow.PasswordChangesNotSupportedExcept
 import google.registry.flows.session.LoginFlow.RegistrarAccountNotActiveException;
 import google.registry.flows.session.LoginFlow.TooManyFailedLoginsException;
 import google.registry.flows.session.LoginFlow.UnsupportedLanguageException;
+import google.registry.model.eppoutput.EppOutput;
 import google.registry.model.registrar.Registrar;
 import google.registry.model.registrar.Registrar.State;
 import org.junit.jupiter.api.BeforeEach;
@@ -74,6 +76,14 @@ public abstract class LoginFlowTestCase extends FlowTestCase<LoginFlow> {
     doSuccessfulTest("login_valid.xml");
   }
 
+  @Test
+  void testSuccess_setsIsLoginResponse() throws Exception {
+    setEppInput("login_valid.xml");
+    assertTransactionalFlow(false);
+    EppOutput output = runFlow();
+    assertThat(output.getResponse().isLoginResponse()).isTrue();
+  }
+
   @Test
   void testSuccess_suspendedRegistrar() throws Exception {
     persistResource(getRegistrarBuilder().setState(State.SUSPENDED).build());

core/src/test/java/google/registry/flows/session/LoginFlowViaTlsTest.java

@@ -15,6 +15,7 @@
 package google.registry.flows.session;
 
 import static google.registry.testing.DatabaseHelper.persistResource;
+import static org.joda.time.DateTimeZone.UTC;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.net.InetAddresses;
@@ -26,6 +27,7 @@ import google.registry.model.registrar.Registrar;
 import google.registry.testing.CertificateSamples;
 import google.registry.util.CidrAddressBlock;
 import java.util.Optional;
+import org.joda.time.DateTime;
 import org.junit.jupiter.api.Test;
 
 /** Unit tests for {@link LoginFlow} when accessed via a TLS transport. */
@@ -41,7 +43,7 @@ public class LoginFlowViaTlsTest extends LoginFlowTestCase {
   @Override
   protected Registrar.Builder getRegistrarBuilder() {
     return super.getRegistrarBuilder()
-        .setClientCertificateHash(GOOD_CERT)
+        .setClientCertificate(CertificateSamples.SAMPLE_CERT, DateTime.now(UTC))
         .setIpAddressAllowList(
             ImmutableList.of(CidrAddressBlock.create(InetAddresses.forString(GOOD_IP.get()), 32)));
   }

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

@@ -156,16 +156,6 @@ public final class OteAccountBuilderTest {
         .isTrue();
   }
 
-  @Test
-  void testCreateOteEntities_setCertificateHash() {
-    OteAccountBuilder.forClientId("myclientid")
-        .setCertificateHash(SAMPLE_CERT_HASH)
-        .buildAndPersist();
-
-    assertThat(Registrar.loadByClientId("myclientid-3").get().getClientCertificateHash())
-        .isEqualTo(SAMPLE_CERT_HASH);
-  }
-
   @Test
   void testCreateOteEntities_setCertificate() {
     OteAccountBuilder.forClientId("myclientid")

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

@@ -84,10 +84,10 @@ public class DomainBaseSqlTest {
     saveRegistrar("registrar1");
     saveRegistrar("registrar2");
     saveRegistrar("registrar3");
-    contactKey = VKey.createSql(ContactResource.class, "contact_id1");
-    contact2Key = VKey.createSql(ContactResource.class, "contact_id2");
+    contactKey = createKey(ContactResource.class, "contact_id1");
+    contact2Key = createKey(ContactResource.class, "contact_id2");
 
-    host1VKey = VKey.createSql(HostResource.class, "host1");
+    host1VKey = createKey(HostResource.class, "host1");
 
     domain =
         new DomainBase.Builder()
@@ -696,6 +696,10 @@ public class DomainBaseSqlTest {
     assertThat(domain.getGracePeriods()).isEqualTo(gracePeriods);
   }
 
+  private <T> VKey<T> createKey(Class<T> clazz, String name) {
+    return VKey.create(clazz, name, Key.create(clazz, name));
+  }
+
   private <T> VKey<T> createLegacyVKey(Class<T> clazz, long id) {
     return VKey.create(
         clazz, id, Key.create(Key.create(EntityGroupRoot.class, "per-tld"), clazz, id));

core/src/test/java/google/registry/model/history/DomainHistoryTest.java

@@ -14,7 +14,6 @@
 
 package google.registry.model.history;
 
-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;
@@ -92,10 +91,7 @@ public class DomainHistoryTest extends EntityTestCase {
               assertDomainHistoriesEqual(fromDatabase, domainHistory);
               assertThat(fromDatabase.getParentVKey()).isEqualTo(domainHistory.getParentVKey());
               assertThat(fromDatabase.getNsHosts())
-                  .containsExactlyElementsIn(
-                      domainHistory.getNsHosts().stream()
-                          .map(key -> VKey.createSql(HostResource.class, key.getSqlKey()))
-                          .collect(toImmutableSet()));
+                  .containsExactlyElementsIn(domainHistory.getNsHosts());
             });
   }
 

core/src/test/java/google/registry/model/ofy/EntityWritePrioritiesTest.java

@@ -18,19 +18,20 @@ import static com.google.common.truth.Truth.assertThat;
 
 import com.google.common.collect.ImmutableMap;
 import com.googlecode.objectify.Key;
+import google.registry.model.ofy.TransactionInfo.Delete;
 import google.registry.model.registrar.Registrar;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.testing.AppEngineExtension;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
-class TransactionInfoTest {
+class EntityWritePrioritiesTest {
 
   @RegisterExtension
   AppEngineExtension appEngine = new AppEngineExtension.Builder().withDatastore().build();
 
   @Test
-  void testGetWeight() {
+  void testGetPriority() {
     // just verify that the lowest is what we expect for both save and delete and verify that the
     // Registrar class is zero.
     ImmutableMap<Key<?>, Object> actions =
@@ -39,10 +40,12 @@ class TransactionInfoTest {
             Key.create(HistoryEntry.class, 200), "fake history entry",
             Key.create(Registrar.class, 300), "fake registrar");
     ImmutableMap<Long, Integer> expectedValues =
-        ImmutableMap.of(100L, TransactionInfo.DELETE_RANGE + 1, 200L, -1, 300L, 0);
+        ImmutableMap.of(100L, EntityWritePriorities.DELETE_RANGE + 10, 200L, -10, 300L, 0);
 
     for (ImmutableMap.Entry<Key<?>, Object> entry : actions.entrySet()) {
-      assertThat(TransactionInfo.getWeight(entry))
+      assertThat(
+              EntityWritePriorities.getEntityPriority(
+                  entry.getKey().getKind(), Delete.SENTINEL.equals(entry.getValue())))
           .isEqualTo(expectedValues.get(entry.getKey().getId()));
     }
   }

core/src/test/java/google/registry/persistence/converter/LongVKeyConverterTest.java

@@ -19,8 +19,7 @@ import static google.registry.persistence.transaction.TransactionManagerFactory.
 
 import google.registry.persistence.VKey;
 import google.registry.persistence.WithLongVKey;
-import google.registry.persistence.transaction.JpaTestRules;
-import google.registry.persistence.transaction.JpaTestRules.JpaUnitTestExtension;
+import google.registry.testing.AppEngineExtension;
 import javax.persistence.Entity;
 import javax.persistence.Id;
 import org.junit.jupiter.api.Test;
@@ -30,33 +29,56 @@ import org.junit.jupiter.api.extension.RegisterExtension;
 public class LongVKeyConverterTest {
 
   @RegisterExtension
-  public final JpaUnitTestExtension jpaExtension =
-      new JpaTestRules.Builder()
-          .withEntityClass(TestEntity.class, VKeyConverter_LongType.class)
-          .buildUnitTestRule();
+  public final AppEngineExtension appEngineExtension =
+      new AppEngineExtension.Builder()
+          .withDatastoreAndCloudSql()
+          .withoutCannedData()
+          .withJpaUnitTestEntities(
+              TestLongEntity.class,
+              VKeyConverter_LongType.class,
+              VKeyConverter_CompositeLongType.class)
+          .withOfyTestEntities(TestLongEntity.class, CompositeKeyTestLongEntity.class)
+          .build();
 
   @Test
   void testRoundTrip() {
-    TestEntity original = new TestEntity(VKey.createSql(TestEntity.class, 10L));
+    TestLongEntity original =
+        new TestLongEntity(
+            VKey.createSql(TestLongEntity.class, 10L),
+            VKey.createSql(CompositeKeyTestLongEntity.class, 20L));
     jpaTm().transact(() -> jpaTm().getEntityManager().persist(original));
 
-    TestEntity retrieved =
-        jpaTm().transact(() -> jpaTm().getEntityManager().find(TestEntity.class, "id"));
+    TestLongEntity retrieved =
+        jpaTm().transact(() -> jpaTm().getEntityManager().find(TestLongEntity.class, "id"));
     assertThat(retrieved.number.getSqlKey()).isEqualTo(10L);
+    assertThat(retrieved.number.getOfyKey().getId()).isEqualTo(10L);
+
+    assertThat(retrieved.composite.getSqlKey()).isEqualTo(20L);
+    assertThat(retrieved.composite.maybeGetOfyKey().isPresent()).isFalse();
   }
 
-  @Entity(name = "TestEntity")
+  @Entity(name = "TestLongEntity")
+  @com.googlecode.objectify.annotation.Entity
   @WithLongVKey(classNameSuffix = "LongType")
-  static class TestEntity {
-    @Id String id = "id";
+  static class TestLongEntity {
+    @com.googlecode.objectify.annotation.Id @Id String id = "id";
 
-    VKey<TestEntity> number;
+    VKey<TestLongEntity> number;
+    VKey<CompositeKeyTestLongEntity> composite;
 
-    TestEntity(VKey<TestEntity> number) {
+    TestLongEntity(VKey<TestLongEntity> number, VKey<CompositeKeyTestLongEntity> composite) {
       this.number = number;
+      this.composite = composite;
     }
 
     /** Default constructor, needed for hibernate. */
-    public TestEntity() {}
+    public TestLongEntity() {}
+  }
+
+  @Entity(name = "CompositeKeyTestLongEntity")
+  @com.googlecode.objectify.annotation.Entity
+  @WithLongVKey(classNameSuffix = "CompositeLongType", compositeKey = true)
+  static class CompositeKeyTestLongEntity {
+    @com.googlecode.objectify.annotation.Id @Id String id = "id";
   }
 }

core/src/test/java/google/registry/persistence/converter/StringVKeyConverterTest.java

@@ -19,8 +19,7 @@ import static google.registry.persistence.transaction.TransactionManagerFactory.
 
 import google.registry.persistence.VKey;
 import google.registry.persistence.WithStringVKey;
-import google.registry.persistence.transaction.JpaTestRules;
-import google.registry.persistence.transaction.JpaTestRules.JpaUnitTestExtension;
+import google.registry.testing.AppEngineExtension;
 import javax.persistence.Entity;
 import javax.persistence.Id;
 import org.junit.jupiter.api.Test;
@@ -30,36 +29,61 @@ import org.junit.jupiter.api.extension.RegisterExtension;
 public class StringVKeyConverterTest {
 
   @RegisterExtension
-  public final JpaUnitTestExtension jpaExtension =
-      new JpaTestRules.Builder()
-          .withEntityClass(TestEntity.class, VKeyConverter_StringType.class)
-          .buildUnitTestRule();
+  public final AppEngineExtension appEngineExtension =
+      new AppEngineExtension.Builder()
+          .withDatastoreAndCloudSql()
+          .withoutCannedData()
+          .withJpaUnitTestEntities(
+              TestStringEntity.class,
+              VKeyConverter_StringType.class,
+              VKeyConverter_CompositeStringType.class)
+          .withOfyTestEntities(TestStringEntity.class, CompositeKeyTestStringEntity.class)
+          .build();
 
   @Test
   void testRoundTrip() {
-    TestEntity original =
-        new TestEntity("TheRealSpartacus", VKey.createSql(TestEntity.class, "ImSpartacus!"));
+    TestStringEntity original =
+        new TestStringEntity(
+            "TheRealSpartacus",
+            VKey.createSql(TestStringEntity.class, "ImSpartacus!"),
+            VKey.createSql(CompositeKeyTestStringEntity.class, "NoImSpartacus!"));
     jpaTm().transact(() -> jpaTm().getEntityManager().persist(original));
 
-    TestEntity retrieved =
+    TestStringEntity retrieved =
         jpaTm()
-            .transact(() -> jpaTm().getEntityManager().find(TestEntity.class, "TheRealSpartacus"));
+            .transact(
+                () -> jpaTm().getEntityManager().find(TestStringEntity.class, "TheRealSpartacus"));
     assertThat(retrieved.other.getSqlKey()).isEqualTo("ImSpartacus!");
+    assertThat(retrieved.other.getOfyKey().getName()).isEqualTo("ImSpartacus!");
+
+    assertThat(retrieved.composite.getSqlKey()).isEqualTo("NoImSpartacus!");
+    assertThat(retrieved.composite.maybeGetOfyKey().isPresent()).isFalse();
   }
 
-  @Entity(name = "TestEntity")
+  @Entity(name = "TestStringEntity")
+  @com.googlecode.objectify.annotation.Entity
   @WithStringVKey(classNameSuffix = "StringType")
-  static class TestEntity {
-    @Id String id;
+  static class TestStringEntity {
+    @com.googlecode.objectify.annotation.Id @Id String id;
 
-    VKey<TestEntity> other;
+    VKey<TestStringEntity> other;
+    VKey<CompositeKeyTestStringEntity> composite;
 
-    TestEntity(String id, VKey<TestEntity> other) {
+    TestStringEntity(
+        String id, VKey<TestStringEntity> other, VKey<CompositeKeyTestStringEntity> composite) {
       this.id = id;
       this.other = other;
+      this.composite = composite;
     }
 
     /** Default constructor, needed for hibernate. */
-    public TestEntity() {}
+    public TestStringEntity() {}
+  }
+
+  @Entity(name = "CompositeKeyTestStringEntity")
+  @com.googlecode.objectify.annotation.Entity
+  @WithStringVKey(classNameSuffix = "CompositeStringType", compositeKey = true)
+  static class CompositeKeyTestStringEntity {
+    @com.googlecode.objectify.annotation.Id @Id String id = "id";
   }
 }

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

@@ -638,11 +638,13 @@ public class DatabaseHelper {
       DateTime requestTime,
       DateTime expirationTime,
       DateTime extendedRegistrationExpirationTime) {
-    HistoryEntry historyEntryDomainTransfer = persistResource(
-        new HistoryEntry.Builder()
-            .setType(HistoryEntry.Type.DOMAIN_TRANSFER_REQUEST)
-            .setParent(domain)
-            .build());
+    HistoryEntry historyEntryDomainTransfer =
+        persistResource(
+            new HistoryEntry.Builder()
+                .setType(HistoryEntry.Type.DOMAIN_TRANSFER_REQUEST)
+                .setModificationTime(tm().transact(() -> tm().getTransactionTime()))
+                .setParent(domain)
+                .build());
     BillingEvent.OneTime transferBillingEvent = persistResource(createBillingEventForTransfer(
             domain,
             historyEntryDomainTransfer,
@@ -1184,6 +1186,7 @@ public class DatabaseHelper {
   public static void deleteResource(final Object resource) {
     if (alwaysSaveWithBackup) {
       tm().transact(() -> tm().delete(resource));
+      maybeAdvanceClock();
     } else {
       transactIfJpaTm(() -> tm().deleteWithoutBackup(resource));
     }

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

@@ -24,16 +24,19 @@ import google.registry.model.ImmutableObject;
 import google.registry.model.annotations.VirtualEntity;
 import google.registry.model.common.EntityGroupRoot;
 import google.registry.persistence.VKey;
+import google.registry.schema.replay.DatastoreAndSqlEntity;
 import google.registry.schema.replay.EntityTest.EntityForTesting;
+import javax.persistence.Transient;
 
 /** A test model object that can be persisted in any entity group. */
 @Entity
+@javax.persistence.Entity
 @EntityForTesting
-public class TestObject extends ImmutableObject {
+public class TestObject extends ImmutableObject implements DatastoreAndSqlEntity {
 
-  @Parent Key<EntityGroupRoot> parent;
+  @Parent @Transient Key<EntityGroupRoot> parent;
 
-  @Id String id;
+  @Id @javax.persistence.Id String id;
 
   String field;
 

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

@@ -21,7 +21,6 @@ import static google.registry.persistence.transaction.TransactionManagerFactory.
 import static google.registry.testing.CertificateSamples.SAMPLE_CERT;
 import static google.registry.testing.CertificateSamples.SAMPLE_CERT3;
 import static google.registry.testing.CertificateSamples.SAMPLE_CERT3_HASH;
-import static google.registry.testing.CertificateSamples.SAMPLE_CERT_HASH;
 import static google.registry.testing.DatabaseHelper.createTlds;
 import static google.registry.testing.DatabaseHelper.persistNewRegistrar;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
@@ -447,29 +446,6 @@ class CreateRegistrarCommandTest extends CommandTestCase<CreateRegistrarCommand>
     assertThat(registrar).isEmpty();
   }
 
-  @Test
-  void testSuccess_clientCertHashFlag() throws Exception {
-    runCommandForced(
-        "--name=blobio",
-        "--password=some_password",
-        "--registrar_type=REAL",
-        "--iana_id=8",
-        "--cert_hash=" + SAMPLE_CERT_HASH,
-        "--passcode=01234",
-        "--icann_referral_email=foo@bar.test",
-        "--street=\"123 Fake St\"",
-        "--city Fakington",
-        "--state MA",
-        "--zip 00351",
-        "--cc US",
-        "clientz");
-
-    Optional<Registrar> registrar = Registrar.loadByClientId("clientz");
-    assertThat(registrar).isPresent();
-    assertThat(registrar.get().getClientCertificate()).isNull();
-    assertThat(registrar.get().getClientCertificateHash()).isEqualTo(SAMPLE_CERT_HASH);
-  }
-
   @Test
   void testSuccess_failoverClientCertFileFlag() throws Exception {
     fakeClock.setTo(DateTime.parse("2020-11-01T00:00:00Z"));
@@ -1182,74 +1158,6 @@ class CreateRegistrarCommandTest extends CommandTestCase<CreateRegistrarCommand>
                 "clientz"));
   }
 
-  @Test
-  void testFailure_certHashAndCertFile() {
-    assertThrows(
-        IllegalArgumentException.class,
-        () ->
-            runCommandForced(
-                "--name=blobio",
-                "--password=some_password",
-                "--registrar_type=REAL",
-                "--iana_id=8",
-                "--cert_file=" + getCertFilename(),
-                "--cert_hash=ABCDEF",
-                "--passcode=01234",
-                "--icann_referral_email=foo@bar.test",
-                "--street=\"123 Fake St\"",
-                "--city Fakington",
-                "--state MA",
-                "--zip 00351",
-                "--cc US",
-                "clientz"));
-  }
-
-  @Test
-  void testFailure_certHashNotBase64() {
-    IllegalArgumentException thrown =
-        assertThrows(
-            IllegalArgumentException.class,
-            () ->
-                runCommandForced(
-                    "--name=blobio",
-                    "--password=some_password",
-                    "--registrar_type=REAL",
-                    "--iana_id=8",
-                    "--cert_hash=!",
-                    "--passcode=01234",
-                    "--icann_referral_email=foo@bar.test",
-                    "--street=\"123 Fake St\"",
-                    "--city Fakington",
-                    "--state MA",
-                    "--zip 00351",
-                    "--cc US",
-                    "clientz"));
-    assertThat(thrown).hasMessageThat().contains("base64");
-  }
-
-  @Test
-  void testFailure_certHashNotA256BitValue() {
-    IllegalArgumentException thrown =
-        assertThrows(
-            IllegalArgumentException.class,
-            () ->
-                runCommandForced(
-                    "--name=blobio",
-                    "--password=some_password",
-                    "--registrar_type=REAL",
-                    "--iana_id=8",
-                    "--cert_hash=abc",
-                    "--passcode=01234",
-                    "--icann_referral_email=foo@bar.test",
-                    "--street=\"123 Fake St\"",
-                    "--city Fakington",
-                    "--state MA",
-                    "--zip 00351",
-                    "--cc US",
-                    "clientz"));
-    assertThat(thrown).hasMessageThat().contains("256");
-  }
-
   @Test
   void testFailure_missingName() {
     IllegalArgumentException thrown =

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

@@ -0,0 +1,122 @@
+// 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.model.ImmutableObjectSubject.immutableObjectCorrespondence;
+import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.testing.DatabaseHelper.createTld;
+import static google.registry.testing.DatabaseHelper.newDomainBase;
+import static google.registry.testing.DatabaseHelper.persistResource;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.truth.Correspondence;
+import com.googlecode.objectify.Key;
+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.poll.PollMessage;
+import google.registry.model.reporting.HistoryEntry;
+import google.registry.persistence.VKey;
+import org.joda.time.DateTime;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+/** Unit test for {@link RemoveRegistryOneKeyCommand}. */
+public class RemoveRegistryOneKeyCommandTest extends CommandTestCase<RemoveRegistryOneKeyCommand> {
+  DomainBase domain;
+  HistoryEntry historyEntry;
+
+  @BeforeEach
+  void beforeEach() {
+    createTld("foobar");
+    domain =
+        newDomainBase("foo.foobar")
+            .asBuilder()
+            .setDeletionTime(DateTime.parse("2016-01-01T00:00:00Z"))
+            .setAutorenewBillingEvent(createRegistryOneVKey(BillingEvent.Recurring.class, 100L))
+            .setAutorenewPollMessage(createRegistryOneVKey(PollMessage.Autorenew.class, 200L))
+            .setDeletePollMessage(createRegistryOneVKey(PollMessage.OneTime.class, 300L))
+            .build();
+  }
+
+  @Test
+  void removeRegistryOneKeyInDomainBase_succeeds() throws Exception {
+    DomainBase origin = persistResource(domain);
+
+    runCommand(
+        "--force",
+        "--key_paths_file",
+        writeToNamedTmpFile("keypath.txt", getKeyPathLiteral(domain)));
+
+    DomainBase persisted = ofy().load().key(domain.createVKey().getOfyKey()).now();
+    assertThat(ImmutableList.of(persisted))
+        .comparingElementsUsing(getDomainBaseCorrespondence())
+        .containsExactly(origin);
+    assertThat(persisted.getAutorenewBillingEvent()).isNull();
+    assertThat(persisted.getAutorenewPollMessage()).isNull();
+    assertThat(persisted.getDeletePollMessage()).isNull();
+  }
+
+  @Test
+  void removeRegistryOneKeyInDomainBase_notModifyRegistryTwoKey() throws Exception {
+    DomainBase origin =
+        persistResource(
+            domain
+                .asBuilder()
+                .setAutorenewBillingEvent(
+                    createRegistryTwoVKey(BillingEvent.Recurring.class, domain, 300L))
+                .build());
+
+    runCommand(
+        "--force",
+        "--key_paths_file",
+        writeToNamedTmpFile("keypath.txt", getKeyPathLiteral(domain)));
+
+    DomainBase persisted = ofy().load().key(domain.createVKey().getOfyKey()).now();
+    assertThat(ImmutableList.of(persisted))
+        .comparingElementsUsing(getDomainBaseCorrespondence())
+        .containsExactly(origin);
+    assertThat(persisted.getAutorenewBillingEvent())
+        .isEqualTo(createRegistryTwoVKey(BillingEvent.Recurring.class, domain, 300L));
+    assertThat(persisted.getAutorenewPollMessage()).isNull();
+    assertThat(persisted.getDeletePollMessage()).isNull();
+  }
+
+  private static String getKeyPathLiteral(Object entity) {
+    Key<?> key = Key.create(entity);
+    return String.format("\"DomainBase\", \"%s\"", key.getName());
+  }
+
+  private static <T> VKey<T> createRegistryOneVKey(Class<T> clazz, long id) {
+    Key<?> parent = Key.create(EntityGroupRoot.class, "per-tld");
+    return VKey.create(clazz, id, Key.create(parent, clazz, id));
+  }
+
+  private static <T> VKey<T> createRegistryTwoVKey(Class<T> clazz, DomainBase domain, long id) {
+    Key<?> parent = Key.create(domain.createVKey().getOfyKey(), HistoryEntry.class, 1000L);
+    return VKey.create(clazz, id, Key.create(parent, clazz, id));
+  }
+
+  private static Correspondence<ImmutableObject, ImmutableObject> getDomainBaseCorrespondence() {
+    return immutableObjectCorrespondence(
+        "revisions",
+        "updateTimestamp",
+        "autorenewBillingEvent",
+        "autorenewPollMessage",
+        "deletePollMessage");
+  }
+}

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

@@ -18,7 +18,6 @@ import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.registrar.Registrar.State.ACTIVE;
 import static google.registry.model.registry.Registry.TldState.GENERAL_AVAILABILITY;
 import static google.registry.model.registry.Registry.TldState.START_DATE_SUNRISE;
-import static google.registry.testing.CertificateSamples.SAMPLE_CERT;
 import static google.registry.testing.CertificateSamples.SAMPLE_CERT_HASH;
 import static google.registry.testing.DatabaseHelper.createTld;
 import static google.registry.testing.DatabaseHelper.loadRegistrar;
@@ -98,8 +97,7 @@ class SetupOteCommandTest extends CommandTestCase<SetupOteCommand> {
       String registrarName,
       String allowedTld,
       String password,
-      ImmutableList<CidrAddressBlock> ipAllowList,
-      boolean hashOnly) {
+      ImmutableList<CidrAddressBlock> ipAllowList) {
     Registrar registrar = loadRegistrar(registrarName);
     assertThat(registrar).isNotNull();
     assertThat(registrar.getAllowedTlds()).containsExactlyElementsIn(ImmutableSet.of(allowedTld));
@@ -108,18 +106,6 @@ class SetupOteCommandTest extends CommandTestCase<SetupOteCommand> {
     assertThat(registrar.verifyPassword(password)).isTrue();
     assertThat(registrar.getIpAddressAllowList()).isEqualTo(ipAllowList);
     assertThat(registrar.getClientCertificateHash()).isEqualTo(SAMPLE_CERT_HASH);
-    // If certificate hash is provided, there's no certificate file stored with the registrar.
-    if (!hashOnly) {
-      assertThat(registrar.getClientCertificate()).isEqualTo(SAMPLE_CERT);
-    }
-  }
-
-  private void verifyRegistrarCreation(
-      String registrarName,
-      String allowedTld,
-      String password,
-      ImmutableList<CidrAddressBlock> ipAllowList) {
-    verifyRegistrarCreation(registrarName, allowedTld, password, ipAllowList, false);
   }
 
   private void verifyRegistrarContactCreation(String registrarName, String email) {
@@ -184,24 +170,6 @@ class SetupOteCommandTest extends CommandTestCase<SetupOteCommand> {
     verifyRegistrarContactCreation("abc-5", "abc@email.com");
   }
 
-  @Test
-  void testSuccess_certificateHash() throws Exception {
-    runCommandForced(
-        "--ip_allow_list=1.1.1.1",
-        "--registrar=blobio",
-        "--email=contact@email.com",
-        "--certhash=" + SAMPLE_CERT_HASH);
-
-    verifyTldCreation("blobio-eap", "BLOBIOE3", GENERAL_AVAILABILITY, true);
-
-    ImmutableList<CidrAddressBlock> ipAddress =
-        ImmutableList.of(CidrAddressBlock.create("1.1.1.1"));
-
-    verifyRegistrarCreation("blobio-5", "blobio-eap", PASSWORD, ipAddress, true);
-
-    verifyRegistrarContactCreation("blobio-5", "contact@email.com");
-  }
-
   @Test
   void testSuccess_multipleIps() throws Exception {
     runCommandForced(
@@ -256,7 +224,7 @@ class SetupOteCommandTest extends CommandTestCase<SetupOteCommand> {
   }
 
   @Test
-  void testFailure_missingCertificateFileAndCertificateHash() {
+  void testFailure_missingCertificateFile() {
     IllegalArgumentException thrown =
         assertThrows(
             IllegalArgumentException.class,
@@ -265,26 +233,7 @@ class SetupOteCommandTest extends CommandTestCase<SetupOteCommand> {
                     "--ip_allow_list=1.1.1.1", "--email=contact@email.com", "--registrar=blobio"));
     assertThat(thrown)
         .hasMessageThat()
-        .contains(
-            "Must specify exactly one of client certificate file or client certificate hash.");
-  }
-
-  @Test
-  void testFailure_suppliedCertificateFileAndCertificateHash() {
-    IllegalArgumentException thrown =
-        assertThrows(
-            IllegalArgumentException.class,
-            () ->
-                runCommandForced(
-                    "--ip_allow_list=1.1.1.1",
-                    "--email=contact@email.com",
-                    "--registrar=blobio",
-                    "--certfile=" + getCertFilename(),
-                    "--certhash=" + SAMPLE_CERT_HASH));
-    assertThat(thrown)
-        .hasMessageThat()
-        .contains(
-            "Must specify exactly one of client certificate file or client certificate hash.");
+        .contains("Must specify exactly one client certificate file.");
   }
 
   @Test

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

@@ -21,7 +21,6 @@ import static google.registry.persistence.transaction.TransactionManagerFactory.
 import static google.registry.testing.CertificateSamples.SAMPLE_CERT;
 import static google.registry.testing.CertificateSamples.SAMPLE_CERT3;
 import static google.registry.testing.CertificateSamples.SAMPLE_CERT3_HASH;
-import static google.registry.testing.CertificateSamples.SAMPLE_CERT_HASH;
 import static google.registry.testing.DatabaseHelper.createTlds;
 import static google.registry.testing.DatabaseHelper.loadRegistrar;
 import static google.registry.testing.DatabaseHelper.persistResource;
@@ -339,14 +338,6 @@ class UpdateRegistrarCommandTest extends CommandTestCase<UpdateRegistrarCommand>
     assertThat(registrar.getFailoverClientCertificate()).isEqualTo(SAMPLE_CERT3);
   }
 
-  @Test
-  void testSuccess_certHash() throws Exception {
-    assertThat(loadRegistrar("NewRegistrar").getClientCertificateHash()).isNull();
-    runCommand("--cert_hash=" + SAMPLE_CERT_HASH, "--force", "NewRegistrar");
-    assertThat(loadRegistrar("NewRegistrar").getClientCertificateHash())
-        .isEqualTo(SAMPLE_CERT_HASH);
-  }
-
   @Test
   void testSuccess_clearCert() throws Exception {
     persistResource(
@@ -359,18 +350,6 @@ class UpdateRegistrarCommandTest extends CommandTestCase<UpdateRegistrarCommand>
     assertThat(loadRegistrar("NewRegistrar").getClientCertificate()).isNull();
   }
 
-  @Test
-  void testSuccess_clearCertHash() throws Exception {
-    persistResource(
-        loadRegistrar("NewRegistrar")
-            .asBuilder()
-            .setClientCertificateHash(SAMPLE_CERT_HASH)
-            .build());
-    assertThat(isNullOrEmpty(loadRegistrar("NewRegistrar").getClientCertificateHash())).isFalse();
-    runCommand("--cert_hash=null", "--force", "NewRegistrar");
-    assertThat(loadRegistrar("NewRegistrar").getClientCertificateHash()).isNull();
-  }
-
   @Test
   void testSuccess_ianaId() throws Exception {
     assertThat(loadRegistrar("NewRegistrar").getIanaIdentifier()).isEqualTo(8);
@@ -762,18 +741,6 @@ class UpdateRegistrarCommandTest extends CommandTestCase<UpdateRegistrarCommand>
         () -> runCommand("--cert_file=" + writeToTmpFile("ABCDEF"), "--force", "NewRegistrar"));
   }
 
-  @Test
-  void testFailure_certHashAndCertFile() {
-    assertThrows(
-        IllegalArgumentException.class,
-        () ->
-            runCommand(
-                "--cert_file=" + getCertFilename(SAMPLE_CERT3),
-                "--cert_hash=ABCDEF",
-                "--force",
-                "NewRegistrar"));
-  }
-
   @Test
   void testFailure_missingClientId() {
     assertThrows(ParameterException.class, () -> runCommand("--force"));

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

@@ -20,6 +20,7 @@ import static google.registry.testing.DatabaseHelper.createTld;
 import static google.registry.testing.DatabaseHelper.loadRegistrar;
 import static google.registry.testing.DatabaseHelper.persistResource;
 import static google.registry.testing.EppExceptionSubject.assertAboutEppExceptions;
+import static org.joda.time.DateTimeZone.UTC;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
 import com.beust.jcommander.ParameterException;
@@ -33,6 +34,7 @@ import google.registry.testing.CertificateSamples;
 import google.registry.util.CidrAddressBlock;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
@@ -50,7 +52,7 @@ class ValidateLoginCredentialsCommandTest extends CommandTestCase<ValidateLoginC
         loadRegistrar("NewRegistrar")
             .asBuilder()
             .setPassword(PASSWORD)
-            .setClientCertificateHash(CERT_HASH)
+            .setClientCertificate(CertificateSamples.SAMPLE_CERT, DateTime.now(UTC))
             .setIpAddressAllowList(ImmutableList.of(new CidrAddressBlock(CLIENT_IP)))
             .setState(ACTIVE)
             .setAllowedTlds(ImmutableSet.of("tld"))

core/src/test/java/google/registry/webdriver/RegistrarConsoleScreenshotTest.java

@@ -342,26 +342,6 @@ class RegistrarConsoleScreenshotTest extends WebDriverTestCase {
     driver.diffPage("edit");
   }
 
-  @RetryingTest(3)
-  void settingsSecurityWithHashOnly() throws Throwable {
-    server.runInAppEngineEnvironment(
-        () -> {
-          persistResource(
-              loadRegistrar("TheRegistrar")
-                  .asBuilder()
-                  .setClientCertificateHash(CertificateSamples.SAMPLE_CERT_HASH)
-                  .build());
-          return null;
-        });
-    driver.manage().window().setSize(new Dimension(1050, 2000));
-    driver.get(server.getUrl("/registrar#security-settings"));
-    driver.waitForDisplayedElement(By.tagName("h1"));
-    driver.diffPage("view");
-    driver.waitForDisplayedElement(By.id("reg-app-btn-edit")).click();
-    driver.waitForDisplayedElement(By.tagName("h1"));
-    driver.diffPage("edit");
-  }
-
   @RetryingTest(3)
   void index_registrarDisabled() throws Throwable {
     server.runInAppEngineEnvironment(

db/src/main/resources/sql/er_diagram/brief_er_diagram.html

@@ -261,11 +261,11 @@ td.section {
     </tr> 
     <tr> 
      <td class="property_name">generated on</td> 
-     <td class="property_value">2020-11-18 21:42:47.876348</td> 
+     <td class="property_value">2020-11-30 21:59:21.319302</td> 
     </tr> 
     <tr>
      <td class="property_name">last flyway file</td>
-     <td id="lastFlywayFile" class="property_value">V79__drop_foreign_keys_on_pollmessage.sql</td>
+     <td id="lastFlywayFile" class="property_value">V80__defer_bill_event_key.sql</td>
     </tr>
    </tbody>
   </table> 
@@ -284,7 +284,7 @@ td.section {
      generated on
     </text> 
     <text text-anchor="start" x="4027.94" y="-10.8" font-family="Helvetica,sans-Serif" font-size="14.00">
-     2020-11-18 21:42:47.876348
+     2020-11-30 21:59:21.319302
     </text> 
     <polygon fill="none" stroke="#888888" points="3940.44,-4 3940.44,-44 4205.44,-44 4205.44,-4 3940.44,-4" /> <!-- allocationtoken_a08ccbef --> 
     <g id="node1" class="node"> 

db/src/main/resources/sql/er_diagram/full_er_diagram.html

@@ -261,11 +261,11 @@ td.section {
     </tr> 
     <tr> 
      <td class="property_name">generated on</td> 
-     <td class="property_value">2020-11-18 21:42:46.037281</td> 
+     <td class="property_value">2020-11-30 21:59:18.813825</td> 
     </tr> 
     <tr>
      <td class="property_name">last flyway file</td>
-     <td id="lastFlywayFile" class="property_value">V79__drop_foreign_keys_on_pollmessage.sql</td>
+     <td id="lastFlywayFile" class="property_value">V80__defer_bill_event_key.sql</td>
     </tr>
    </tbody>
   </table> 
@@ -284,7 +284,7 @@ td.section {
      generated on
     </text> 
     <text text-anchor="start" x="4631.68" y="-10.8" font-family="Helvetica,sans-Serif" font-size="14.00">
-     2020-11-18 21:42:46.037281
+     2020-11-30 21:59:18.813825
     </text> 
     <polygon fill="none" stroke="#888888" points="4544.18,-4 4544.18,-44 4809.18,-44 4809.18,-4 4544.18,-4" /> <!-- allocationtoken_a08ccbef --> 
     <g id="node1" class="node"> 

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

@@ -77,3 +77,4 @@ V76__change_history_nullability.sql
 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

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

@@ -0,0 +1,27 @@
+-- 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.
+
+-- We have to make this foreign key initially deferred even though
+-- BillingEvent is saved before Domain (and deleted after).  If we don't, it
+-- appears that Hibernate can sporadically introduce FK constraint failures
+-- when updating a Domain to reference a new BillingEvent and then deleting
+-- the old BillingEvent.  This may be due to the fact that this FK
+-- relationships is not known to hibernate.
+ALTER TABLE "Domain" DROP CONSTRAINT fk_domain_transfer_billing_event_id;
+ALTER TABLE if exists "Domain"
+   ADD CONSTRAINT fk_domain_transfer_billing_event_id
+   FOREIGN KEY (transfer_billing_event_id)
+   REFERENCES "BillingEvent"
+   DEFERRABLE INITIALLY DEFERRED;
+

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

@@ -2064,7 +2064,7 @@ ALTER TABLE ONLY public."Domain"
 --
 
 ALTER TABLE ONLY public."Domain"
-    ADD CONSTRAINT fk_domain_transfer_billing_event_id FOREIGN KEY (transfer_billing_event_id) REFERENCES public."BillingEvent"(billing_event_id);
+    ADD CONSTRAINT fk_domain_transfer_billing_event_id FOREIGN KEY (transfer_billing_event_id) REFERENCES public."BillingEvent"(billing_event_id) DEFERRABLE INITIALLY DEFERRED;
 
 
 --

processor/build.gradle

@@ -18,6 +18,12 @@ plugins {
 
 dependencies {
     def deps = rootProject.dependencyMap
+
+    // Custom-built objectify jar at commit ecd5165, included in Nomulus
+    // release.
+    compile files(
+        "${rootDir}/third_party/objectify/v4_1/objectify-4.1.3.jar")
+
     compile deps['com.google.code.findbugs:jsr305']
     compile deps['com.google.guava:guava']
     compile deps['com.squareup:javapoet']

processor/src/main/java/google/registry/processors/AbstractVKeyProcessor.java

@@ -27,11 +27,14 @@ import com.squareup.javapoet.TypeSpec;
 import java.io.IOException;
 import java.io.UncheckedIOException;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import javax.annotation.processing.AbstractProcessor;
 import javax.annotation.processing.RoundEnvironment;
 import javax.lang.model.element.AnnotationMirror;
+import javax.lang.model.element.AnnotationValue;
 import javax.lang.model.element.Element;
+import javax.lang.model.element.ExecutableElement;
 import javax.lang.model.element.Modifier;
 import javax.lang.model.element.TypeElement;
 import javax.lang.model.type.DeclaredType;
@@ -47,9 +50,12 @@ import javax.persistence.Converter;
 public abstract class AbstractVKeyProcessor extends AbstractProcessor {
 
   private static final String CONVERTER_CLASS_NAME_TEMP = "VKeyConverter_%s";
-  // The method with same name should be defined in StringVKey and LongVKey
+  // The method with same name should be defined in WithStringVKey and WithLongVKey
   private static final String CLASS_NAME_SUFFIX_KEY = "classNameSuffix";
 
+  // Method in WithStringVKey and WithLongVKey to indicate that this is a composite key.
+  private static final String COMPOSITE_KEY_KEY = "compositeKey";
+
   abstract Class<?> getSqlColumnType();
 
   abstract String getAnnotationSimpleName();
@@ -76,18 +82,18 @@ public abstract class AbstractVKeyProcessor extends AbstractProcessor {
                         actualAnnotation.size() == 1,
                         String.format(
                             "type can have only 1 %s annotation", getAnnotationSimpleName()));
-                    String converterClassNameSuffix =
-                        actualAnnotation.get(0).getElementValues().entrySet().stream()
-                            .filter(
-                                entry ->
-                                    entry
-                                        .getKey()
-                                        .getSimpleName()
-                                        .toString()
-                                        .equals(CLASS_NAME_SUFFIX_KEY))
-                            .map(entry -> ((String) entry.getValue().getValue()).trim())
-                            .findFirst()
-                            .orElse("");
+                    String converterClassNameSuffix = "";
+                    boolean hasCompositeOfyKey = false;
+                    for (Map.Entry<? extends ExecutableElement, ? extends AnnotationValue> entry :
+                        actualAnnotation.get(0).getElementValues().entrySet()) {
+                      String keyName = entry.getKey().getSimpleName().toString();
+                      Object value = entry.getValue().getValue();
+                      if (keyName.equals(CLASS_NAME_SUFFIX_KEY)) {
+                        converterClassNameSuffix = ((String) value).trim();
+                      } else if (keyName.equals(COMPOSITE_KEY_KEY)) {
+                        hasCompositeOfyKey = (Boolean) value;
+                      }
+                    }
                     if (converterClassNameSuffix.isEmpty()) {
                       converterClassNameSuffix =
                           getTypeUtils().asElement(entityType).getSimpleName().toString();
@@ -97,7 +103,8 @@ public abstract class AbstractVKeyProcessor extends AbstractProcessor {
                       createJavaFile(
                               getPackageName(annotatedTypeElement),
                               String.format(CONVERTER_CLASS_NAME_TEMP, converterClassNameSuffix),
-                              entityType)
+                              entityType,
+                              hasCompositeOfyKey)
                           .writeTo(processingEnv.getFiler());
                     } catch (IOException e) {
                       throw new UncheckedIOException(e);
@@ -108,7 +115,10 @@ public abstract class AbstractVKeyProcessor extends AbstractProcessor {
   }
 
   private JavaFile createJavaFile(
-      String packageName, String converterClassName, TypeMirror entityTypeMirror) {
+      String packageName,
+      String converterClassName,
+      TypeMirror entityTypeMirror,
+      boolean hasCompositeOfyKey) {
     TypeName entityType = ClassName.get(entityTypeMirror);
 
     ParameterizedTypeName attributeConverter =
@@ -127,7 +137,7 @@ public abstract class AbstractVKeyProcessor extends AbstractProcessor {
             .addStatement("return $T.class", entityType)
             .build();
 
-    TypeSpec vKeyConverter =
+    TypeSpec.Builder classBuilder =
         TypeSpec.classBuilder(converterClassName)
             .addAnnotation(
                 AnnotationSpec.builder(ClassName.get(Converter.class))
@@ -135,9 +145,24 @@ public abstract class AbstractVKeyProcessor extends AbstractProcessor {
                     .build())
             .addModifiers(Modifier.FINAL)
             .superclass(attributeConverter)
-            .addMethod(getAttributeClass)
-            .build();
+            .addMethod(getAttributeClass);
 
+    // If this is a converter for a composite vkey type, generate an override for the default
+    // {@link google.registry.persistence.VKeyConverter.hasCompositeOfyKey()} method, which returns
+    // false.
+    if (hasCompositeOfyKey) {
+      MethodSpec hasCompositeOfyKeyMethod =
+          MethodSpec.methodBuilder("hasCompositeOfyKey")
+              .addAnnotation(Override.class)
+              .addModifiers(Modifier.PROTECTED)
+              .returns(boolean.class)
+              .addStatement("return true", entityType)
+              .build();
+
+      classBuilder.addMethod(hasCompositeOfyKeyMethod);
+    }
+
+    TypeSpec vKeyConverter = classBuilder.build();
     return JavaFile.builder(packageName, vKeyConverter).build();
   }
 

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

@@ -249,6 +249,20 @@ class EppServiceHandlerTest {
     assertThat(channel.isActive()).isFalse();
   }
 
+  @Test
+  void sendResponseToNextHandler_unrecognizedHeader() throws Exception {
+    setHandshakeSuccess();
+    String content = "<epp>stuff</epp>";
+    HttpResponse response = makeEppHttpResponse(content, HttpResponseStatus.OK);
+    response.headers().set("unrecognized-header", "test");
+    channel.writeOutbound(response);
+    ByteBuf expectedResponse = channel.readOutbound();
+    assertThat(Unpooled.wrappedBuffer(content.getBytes(UTF_8))).isEqualTo(expectedResponse);
+    // Nothing further to pass to the next handler.
+    assertThat((Object) channel.readOutbound()).isNull();
+    assertThat(channel.isActive()).isTrue();
+  }
+
   @Test
   void testFailure_disconnectOnNonOKResponseStatus() throws Exception {
     setHandshakeSuccess();

release/rollback/appengine.py

@@ -14,39 +14,20 @@
 """Helper for using the AppEngine Admin REST API."""
 
 import time
-from typing import Any, Dict, FrozenSet, Set
+from typing import FrozenSet, Optional, Set, Tuple
 
 from googleapiclient import discovery
-from googleapiclient import http
 
 import common
 
 # AppEngine services under management.
 SERVICES = frozenset(['backend', 'default', 'pubapi', 'tools'])
-# Forces 'list' calls (for services and versions) to return all
-# results in one shot, to avoid having to handle pagination. This values
-# should be greater than the maximum allowed services and versions in any
-# project (
-# https://cloud.google.com/appengine/docs/standard/python/an-overview-of-app-engine#limits).
-_PAGE_SIZE = 250
 # Number of times to check the status of an operation before timing out.
 _STATUS_CHECK_TIMES = 5
 # Delay between status checks of a long-running operation, in seconds
 _STATUS_CHECK_INTERVAL = 5
 
 
-class PagingError(Exception):
-    """Error for unexpected partial results.
-
-    List calls in this module do not handle pagination. This error is raised
-    when a partial result is received.
-    """
-    def __init__(self, uri: str):
-        super().__init__(
-            self, f'Received paged response unexpectedly when calling {uri}. '
-            'Consider increasing _PAGE_SIZE.')
-
-
 class AppEngineAdmin:
     """Wrapper around the AppEngine Admin REST API client.
 
@@ -55,9 +36,16 @@ class AppEngineAdmin:
     """
     def __init__(self,
                  project: str,
-                 service_lookup: discovery.Resource = None,
+                 service_lookup: Optional[discovery.Resource] = None,
                  status_check_interval: int = _STATUS_CHECK_INTERVAL) -> None:
-        """Initialize this instance for an AppEngine(GCP) project."""
+        """Initialize this instance for an AppEngine(GCP) project.
+
+        Args:
+            project: The GCP project name of this AppEngine instance.
+            service_lookup: The GCP discovery handle for service API lookup.
+            status_check_interval: The delay in seconds between status queries
+                when executing long running operations.
+        """
         self._project = project
 
         if service_lookup is not None:
@@ -66,6 +54,8 @@ class AppEngineAdmin:
             apps = discovery.build('appengine', 'v1beta').apps()
 
         self._services = apps.services()
+        self._versions = self._services.versions()
+        self._instances = self._versions.instances()
         self._operations = apps.operations()
         self._status_check_interval = status_check_interval
 
@@ -73,14 +63,6 @@ class AppEngineAdmin:
     def project(self):
         return self._project
 
-    def _checked_request(self, request: http.HttpRequest) -> Dict[str, Any]:
-        """Verifies that all results are returned for a request."""
-        response = request.execute()
-        if 'nextPageToken' in response:
-            raise PagingError(request.uri)
-
-        return response
-
     def get_serving_versions(self) -> FrozenSet[common.VersionKey]:
         """Returns the serving versions of every Nomulus service.
 
@@ -92,14 +74,15 @@ class AppEngineAdmin:
         Returns: An immutable collection of the serving versions grouped by
             service.
         """
-        response = self._checked_request(
-            self._services.list(appsId=self._project, pageSize=_PAGE_SIZE))
+        services = common.list_all_pages(self._services.list,
+                                         'services',
+                                         appsId=self._project)
 
         # Response format is specified at
-        # http://googleapis.github.io/google-api-python-client/docs/dyn/appengine_v1beta5.apps.services.html#list.
+        # http://googleapis.github.io/google-api-python-client/docs/dyn/appengine_v1beta.apps.services.html#list.
 
         versions = []
-        for service in response.get('services', []):
+        for service in services:
             if service['id'] in SERVICES:
                 # yapf: disable
                 versions_with_traffic = (
@@ -134,15 +117,15 @@ class AppEngineAdmin:
         # Sort the requested services for ease of testing. For now the mocked
         # AppEngine admin in appengine_test can only respond in a fixed order.
         for service_id in sorted(requested_services):
-            response = self._checked_request(self._services.versions().list(
-                appsId=self._project,
-                servicesId=service_id,
-                pageSize=_PAGE_SIZE))
+            response = common.list_all_pages(self._versions.list,
+                                             'versions',
+                                             appsId=self._project,
+                                             servicesId=service_id)
 
             # Format of version_list is defined at
-            # https://googleapis.github.io/google-api-python-client/docs/dyn/appengine_v1beta5.apps.services.versions.html#list.
+            # https://googleapis.github.io/google-api-python-client/docs/dyn/appengine_v1beta.apps.services.versions.html#list.
 
-            for version in response.get('versions', []):
+            for version in response:
                 if common.VersionKey(service_id, version['id']) in versions:
                     scalings = [
                         s for s in list(common.AppEngineScaling)
@@ -165,16 +148,34 @@ class AppEngineAdmin:
 
         return frozenset(version_configs)
 
+    def list_instances(
+            self,
+            version: common.VersionKey) -> Tuple[common.VmInstanceInfo, ...]:
+        instances = common.list_all_pages(self._versions.instances().list,
+                                          'instances',
+                                          appsId=self._project,
+                                          servicesId=version.service_id,
+                                          versionsId=version.version_id)
+
+        # Format of version_list is defined at
+        # https://googleapis.github.io/google-api-python-client/docs/dyn/appengine_v1beta.apps.services.versions.instances.html#list
+
+        return tuple([
+            common.VmInstanceInfo(
+                inst['id'], common.parse_gcp_timestamp(inst['startTime']))
+            for inst in instances
+        ])
+
     def set_manual_scaling_num_instance(self, service_id: str, version_id: str,
                                         manual_instances: int) -> None:
         """Creates an request to change an AppEngine version's status."""
         update_mask = 'manualScaling.instances'
         body = {'manualScaling': {'instances': manual_instances}}
-        response = self._services.versions().patch(appsId=self._project,
-                                                   servicesId=service_id,
-                                                   versionsId=version_id,
-                                                   updateMask=update_mask,
-                                                   body=body).execute()
+        response = self._versions.patch(appsId=self._project,
+                                        servicesId=service_id,
+                                        versionsId=version_id,
+                                        updateMask=update_mask,
+                                        body=body).execute()
 
         operation_id = response.get('name').split('operations/')[1]
         for _ in range(_STATUS_CHECK_TIMES):

release/rollback/appengine_test.py

@@ -17,11 +17,14 @@ import unittest
 from unittest import mock
 from unittest.mock import patch
 
+from googleapiclient import http
+
 import appengine
 import common
 
 
-def setup_appengine_admin() -> Tuple[object, object]:
+def setup_appengine_admin(
+) -> Tuple[appengine.AppEngineAdmin, http.HttpRequest]:
     """Helper for setting up a mocked AppEngineAdmin instance.
 
     Returns:
@@ -32,7 +35,7 @@ def setup_appengine_admin() -> Tuple[object, object]:
     # Assign mocked API response to mock_request.execute.
     mock_request = mock.MagicMock()
     mock_request.uri.return_value = 'myuri'
-    # Mocked resource shared by services, versions, and operations.
+    # Mocked resource shared by services, versions, instances, and operations.
     resource = mock.MagicMock()
     resource.list.return_value = mock_request
     resource.get.return_value = mock_request
@@ -41,6 +44,7 @@ def setup_appengine_admin() -> Tuple[object, object]:
     apps = mock.MagicMock()
     apps.services.return_value = resource
     resource.versions.return_value = resource
+    resource.instances.return_value = resource
     apps.operations.return_value = resource
     service_lookup = mock.MagicMock()
     service_lookup.apps.return_value = apps
@@ -66,11 +70,6 @@ class AppEngineTestCase(unittest.TestCase):
         else:
             self._mock_request.execute.return_value = responses
 
-    def test_checked_request_multipage_raises(self) -> None:
-        self._set_mocked_response({'nextPageToken': ''})
-        self.assertRaises(appengine.PagingError,
-                          self._client.get_serving_versions)
-
     def test_get_serving_versions(self) -> None:
         self._set_mocked_response({
             'services': [{

release/rollback/common.py

@@ -11,13 +11,16 @@
 # 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.
-"""Declares data types that describe AppEngine services and versions."""
+"""Data types and utilities common to the other modules in this package."""
 
 import dataclasses
+import datetime
 import enum
 import pathlib
 import re
-from typing import Optional
+from typing import Any, Optional, Tuple
+
+from google.protobuf import timestamp_pb2
 
 
 class CannotRollbackError(Exception):
@@ -91,6 +94,13 @@ class VersionConfig(VersionKey):
     manual_scaling_instances: Optional[int] = None
 
 
+@dataclasses.dataclass(frozen=True)
+class VmInstanceInfo:
+    """Information about an AppEngine VM instance."""
+    instance_name: str
+    start_time: datetime.datetime
+
+
 def get_nomulus_root() -> str:
     """Finds the current Nomulus root directory.
 
@@ -109,3 +119,63 @@ def get_nomulus_root() -> str:
 
     raise RuntimeError(
         'Do not move this file out of the Nomulus directory tree.')
+
+
+def list_all_pages(func, data_field: str, *args, **kwargs) -> Tuple[Any, ...]:
+    """Collects all data items from a paginator-based 'List' API.
+
+    Args:
+        func: The GCP API method that supports paged responses.
+        data_field: The field in a response object containing the data
+            items to be returned. This is guaranteed to be an Iterable
+            type.
+        *args: Positional arguments passed to func.
+        *kwargs: Keyword arguments passed to func.
+
+    Returns: An immutable collection of data items assembled from the
+        paged responses.
+    """
+    result_collector = []
+    page_token = None
+    while True:
+        request = func(*args, pageToken=page_token, **kwargs)
+        response = request.execute()
+        result_collector.extend(response.get(data_field, []))
+        page_token = response.get('nextPageToken')
+        if not page_token:
+            return tuple(result_collector)
+
+
+def parse_gcp_timestamp(timestamp: str) -> datetime.datetime:
+    """Parses a timestamp string in GCP API to datetime.
+
+    This method uses protobuf's Timestamp class to parse timestamp strings.
+    This class is used by GCP APIs to parse timestamp strings, and is tolerant
+    to certain cases which can break datetime as of Python 3.8, e.g., the
+    trailing 'Z' as timezone, and fractional seconds with number of digits
+    other than 3 or 6.
+
+    Args:
+        timestamp: A string in RFC 3339 format.
+
+    Returns: A datetime instance.
+    """
+    ts = timestamp_pb2.Timestamp()
+    ts.FromJsonString(timestamp)
+    return ts.ToDatetime()
+
+
+def to_gcp_timestamp(timestamp: datetime.datetime) -> str:
+    """Converts a datetime to string.
+
+    This method uses protobuf's Timestamp class to parse timestamp strings.
+    This class is used by GCP APIs to parse timestamp strings.
+
+    Args:
+        timestamp: The datetime instance to be converted.
+
+    Returns: A string in RFC 3339 format.
+    """
+    ts = timestamp_pb2.Timestamp()
+    ts.FromDatetime(timestamp)
+    return ts.ToJsonString()

release/rollback/common_test.py

@@ -0,0 +1,70 @@
+# 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.
+"""Unit tests for the common module."""
+import datetime
+import unittest
+from unittest import mock
+from unittest.mock import call, patch
+
+import common
+
+
+class CommonTestCase(unittest.TestCase):
+    """Unit tests for the common module."""
+    def setUp(self) -> None:
+        self._mock_request = mock.MagicMock()
+        self._mock_api = mock.MagicMock()
+        self._mock_api.list.return_value = self._mock_request
+        self.addCleanup(patch.stopall)
+
+    def test_list_all_pages_single_page(self):
+        self._mock_request.execute.return_value = {'data': [1]}
+        response = common.list_all_pages(self._mock_api.list,
+                                         'data',
+                                         appsId='project')
+        self.assertSequenceEqual(response, [1])
+        self._mock_api.list.assert_called_once_with(pageToken=None,
+                                                    appsId='project')
+
+    def test_list_all_pages_multi_page(self):
+        self._mock_request.execute.side_effect = [{
+            'data': [1],
+            'nextPageToken': 'token'
+        }, {
+            'data': [2]
+        }]
+        response = common.list_all_pages(self._mock_api.list,
+                                         'data',
+                                         appsId='project')
+        self.assertSequenceEqual(response, [1, 2])
+        self.assertSequenceEqual(self._mock_api.list.call_args_list, [
+            call(pageToken=None, appsId='project'),
+            call(pageToken='token', appsId='project')
+        ])
+
+    def test_parse_timestamp(self):
+        self.assertEqual(common.parse_gcp_timestamp('2020-01-01T00:00:00Z'),
+                         datetime.datetime(2020, 1, 1))
+
+    def test_parse_timestamp_irregular_nano_digits(self):
+        # datetime only accepts 3 or 6 digits in fractional second.
+        self.assertRaises(
+            ValueError,
+            lambda: datetime.datetime.fromisoformat('2020-01-01T00:00:00.9'))
+        self.assertEqual(common.parse_gcp_timestamp('2020-01-01T00:00:00.9Z'),
+                         datetime.datetime(2020, 1, 1, microsecond=900000))
+
+
+if __name__ == '__main__':
+    unittest.main()

release/rollback/plan.py

@@ -51,7 +51,7 @@ class ServiceRollback:
 def _get_service_rollback_plan(
         target_configs: FrozenSet[common.VersionConfig],
         serving_configs: FrozenSet[common.VersionConfig]
-) -> Tuple[ServiceRollback]:
+) -> Tuple[ServiceRollback, ...]:
     # yapf: enable
     """Determines the versions to bring up/down in each service.
 
@@ -111,7 +111,7 @@ def _generate_steps(
         appengine_admin: appengine.AppEngineAdmin,
         env: str,
         target_release: str,
-        rollback_plan: Tuple[ServiceRollback]
+        rollback_plan: Tuple[ServiceRollback, ...]
 ) -> Tuple[steps.RollbackStep, ...]:
     # yapf: enable
     """Generates the sequence of operations for execution.
@@ -158,11 +158,11 @@ def _generate_steps(
 
     for plan in rollback_plan:
         for version in plan.serving_versions:
-            if plan.target_version.scaling != common.AppEngineScaling.AUTOMATIC:
+            if version.scaling != common.AppEngineScaling.AUTOMATIC:
                 rollback_steps.append(
                     steps.start_or_stop_version(appengine_admin.project,
                                                 'stop', version))
-            if plan.target_version.scaling == common.AppEngineScaling.MANUAL:
+            if version.scaling == common.AppEngineScaling.MANUAL:
                 # Release all but one instances. Cannot set num_instances to 0
                 # with this api.
                 rollback_steps.append(
@@ -180,7 +180,7 @@ def _generate_steps(
 
 def get_rollback_plan(gcs_client: gcs.GcsClient,
                       appengine_admin: appengine.AppEngineAdmin, env: str,
-                      target_release: str) -> Tuple[steps.RollbackStep]:
+                      target_release: str) -> Tuple[steps.RollbackStep, ...]:
     """Generates the sequence of rollback operations for execution."""
     target_versions = gcs_client.get_versions_by_release(env, target_release)
     serving_versions = appengine_admin.get_serving_versions()

release/rollback/rolling_restart.py

@@ -0,0 +1,186 @@
+# 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.
+"""Script to rolling-restart the Nomulus server on AppEngine.
+
+This script effects a rolling restart of the Nomulus server by deleting VM
+instances at a controlled pace and leave it to the AppEngine scaling policy
+to bring up new VM instances.
+
+For each service, this script gets a list of VM instances and sequentially
+handles each instance as follows:
+1. Issue a gcloud delete command for this instance.
+2. Poll the AppEngine at fixed intervals until this instance no longer exists.
+
+Instance deletion is not instantaneous. An instance actively processing
+requests takes time to shutdown, and its replacement almost always comes
+up immediately after the shutdown. For this reason, we believe that our current
+implementation is sufficient safe, and will not pursue more sophisticated
+algorithms.
+
+Note that for backend instances that may handle large queries, it may take tens
+of seconds, even minutes, to shut down one of them.
+
+This script also accepts an optional start_time parameter that serves as a
+filter of instances to delete: only those instances that started before this
+time will be deleted. This parameter makes error handling easy. When this
+script fails, simply rerun with the same start_time until it succeeds.
+"""
+import argparse
+import datetime
+import sys
+import time
+from typing import Iterable, Optional, Tuple
+
+import appengine
+import common
+import steps
+
+HELP_MAIN = 'Script to rolling-restart the Nomulus server on AppEngine'
+HELP_MIN_DELAY = 'Minimum delay in seconds between instance deletions.'
+HELP_MIN_LIVE_INSTANCE_PERCENT = (
+    'Minimum number of instances to keep, as a percentage '
+    'of the total at the beginning of the restart process.')
+
+
+# yapf: disable
+def generate_steps(
+        appengine_admin: appengine.AppEngineAdmin,
+        version: common.VersionKey,
+        started_before: datetime.datetime
+) -> Tuple[steps.KillNomulusInstance, ...]:
+    # yapf: enable
+    instances = appengine_admin.list_instances(version)
+    return tuple([
+        steps.kill_nomulus_instance(appengine_admin.project, version,
+                                    inst.instance_name) for inst in instances
+        if inst.start_time <= started_before
+    ])
+
+
+def execute_steps(appengine_admin: appengine.AppEngineAdmin,
+                  version: common.VersionKey,
+                  cmds: Tuple[steps.KillNomulusInstance, ...], min_delay: int,
+                  configured_num_instances: Optional[int]) -> None:
+    print(f'Restarting {len(cmds)} instances in {version.service_id}')
+    for cmd in cmds:
+        print(cmd.info())
+        cmd.execute()
+
+        while True:
+            time.sleep(min_delay)
+            running_instances = [
+                inst.instance_name
+                for inst in appengine_admin.list_instances(version)
+            ]
+            if cmd.instance_name in running_instances:
+                print('Waiting for VM to shut down...')
+                continue
+            if (configured_num_instances is not None
+                    and len(running_instances) < configured_num_instances):
+                print('Waiting for new VM to come up...')
+                continue
+            break
+        print('VM instance has shut down.\n')
+
+    print(f'Done: {len(cmds)} instances in {version.service_id}\n')
+
+
+# yapf: disable
+def restart_one_service(appengine_admin: appengine.AppEngineAdmin,
+                        version: common.VersionKey,
+                        min_delay: int,
+                        started_before: datetime.datetime,
+                        configured_num_instances: Optional[int]) -> None:
+    # yapf: enable
+    """Restart VM instances in one service according to their start time.
+
+    Args:
+        appengine_admin: The client of AppEngine Admin API.
+        version: The Nomulus version to restart. This must be the currently
+            serving version.
+        min_delay: The minimum delay between successive deletions.
+        started_before: Only VM instances started before this time are to be
+            deleted.
+        configured_num_instances: When present, the constant number of instances
+            this version is configured with.
+    """
+    cmds = generate_steps(appengine_admin, version, started_before)
+    # yapf: disable
+    execute_steps(
+        appengine_admin, version, cmds, min_delay, configured_num_instances)
+    # yapf: enable
+
+
+# yapf: disable
+def rolling_restart(project: str,
+                    services: Iterable[str],
+                    min_delay: int,
+                    started_before: datetime.datetime):
+    # yapf: enable
+    print(f'Rolling restart {project} at '
+          f'{common.to_gcp_timestamp(started_before)}\n')
+    appengine_admin = appengine.AppEngineAdmin(project)
+    version_configs = appengine_admin.get_version_configs(
+        set(appengine_admin.get_serving_versions()))
+    restart_versions = [
+        version for version in version_configs
+        if version.service_id in services
+    ]
+    # yapf: disable
+    for version in restart_versions:
+        restart_one_service(appengine_admin,
+                            version,
+                            min_delay,
+                            started_before,
+                            version.manual_scaling_instances)
+    # yapf: enable
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(prog='rolling_restart',
+                                     description=HELP_MAIN)
+    parser.add_argument('--project',
+                        '-p',
+                        required=True,
+                        help='The GCP project of the Nomulus server.')
+    parser.add_argument('--services',
+                        '-s',
+                        nargs='+',
+                        choices=appengine.SERVICES,
+                        default=appengine.SERVICES,
+                        help='The services to rollback.')
+    parser.add_argument('--min_delay',
+                        '-d',
+                        type=int,
+                        default=5,
+                        choices=range(1, 100),
+                        help=HELP_MIN_DELAY)
+    parser.add_argument(
+        '--started_before',
+        '-b',
+        type=common.parse_gcp_timestamp,
+        default=datetime.datetime.utcnow(),
+        help='Only kill VM instances started before this time.')
+
+    args = parser.parse_args()
+    rolling_restart(**vars(args))
+    return 0
+
+
+if __name__ == '__main__':
+    try:
+        sys.exit(main())
+    except Exception as ex:  # pylint: disable=broad-except
+        print(ex)
+        sys.exit(1)

release/rollback/rolling_restart_test.py

@@ -0,0 +1,149 @@
+# 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.
+"""Unit tests of rolling_restart."""
+
+import datetime
+import unittest
+from unittest import mock
+
+import common
+import rolling_restart
+import steps
+
+import appengine_test
+
+
+class RollingRestartTestCase(unittest.TestCase):
+    """Tests for rolling_restart."""
+    def setUp(self) -> None:
+        self._appengine_admin, self._appengine_request = (
+            appengine_test.setup_appengine_admin())
+        self._version = common.VersionKey('my_service', 'my_version')
+        self.addCleanup(mock.patch.stopall)
+
+    def _setup_execute_steps_tests(self):
+        self._appengine_request.execute.side_effect = [
+            # First list_instance response.
+            {
+                'instances': [{
+                    'id': 'vm_to_delete',
+                    'startTime': '2019-01-01T00:00:00Z'
+                }, {
+                    'id': 'vm_to_stay',
+                    'startTime': '2019-01-01T00:00:00Z'
+                }]
+            },
+            # Second list_instance response
+            {
+                'instances': [{
+                    'id': 'vm_to_stay',
+                    'startTime': '2019-01-01T00:00:00Z'
+                }]
+            },
+            # Third list_instance response
+            {
+                'instances': [{
+                    'id': 'vm_to_stay',
+                    'startTime': '2019-01-01T00:00:00Z'
+                }, {
+                    'id': 'vm_new',
+                    'startTime': '2019-01-01T00:00:00Z'
+                }]
+            }
+        ]
+
+    def _setup_generate_steps_tests(self):
+        self._appengine_request.execute.side_effect = [
+            # First page of list_instance response.
+            {
+                'instances': [{
+                    'id': 'vm_2019',
+                    'startTime': '2019-01-01T00:00:00Z'
+                }],
+                'nextPageToken':
+                'token'
+            },
+            # Second and final page of list_instance response
+            {
+                'instances': [{
+                    'id': 'vm_2020',
+                    'startTime': '2020-01-01T00:00:00Z'
+                }]
+            }
+        ]
+
+    def test_kill_vm_command(self) -> None:
+        cmd = steps.kill_nomulus_instance(
+            'my_project', common.VersionKey('my_service', 'my_version'),
+            'my_inst')
+        self.assertEqual(cmd.instance_name, 'my_inst')
+        self.assertIn(('gcloud app instances delete my_inst --quiet '
+                       '--user-output-enabled=false --service my_service '
+                       '--version my_version --project my_project'),
+                      cmd.info())
+
+    def _generate_kill_vm_command(self, version: common.VersionKey,
+                                  instance_name: str):
+        return steps.kill_nomulus_instance(self._appengine_admin.project,
+                                           version, instance_name)
+
+    def test_generate_commands(self):
+        self._setup_generate_steps_tests()
+        commands = rolling_restart.generate_steps(self._appengine_admin,
+                                                  self._version,
+                                                  datetime.datetime.utcnow())
+        self.assertSequenceEqual(commands, [
+            self._generate_kill_vm_command(self._version, 'vm_2019'),
+            self._generate_kill_vm_command(self._version, 'vm_2020')
+        ])
+
+    def test_generate_commands_older_vm(self):
+        self._setup_generate_steps_tests()
+        version = common.VersionKey('my_service', 'my_version')
+        # yapf: disable
+        commands = rolling_restart.generate_steps(
+            self._appengine_admin,
+            version,
+            common.parse_gcp_timestamp('2019-12-01T00:00:00Z'))
+        # yapf: enable
+        self.assertSequenceEqual(
+            commands, [self._generate_kill_vm_command(version, 'vm_2019')])
+
+    def test_execute_steps_variable_instances(self):
+        self._setup_execute_steps_tests()
+        cmd = mock.MagicMock()
+        cmd.instance_name = 'vm_to_delete'
+        cmds = tuple([cmd])  # yapf does not format (cmd,) correctly.
+        rolling_restart.execute_steps(appengine_admin=self._appengine_admin,
+                                      version=self._version,
+                                      cmds=cmds,
+                                      min_delay=0,
+                                      configured_num_instances=None)
+        self.assertEqual(self._appengine_request.execute.call_count, 2)
+
+    def test_execute_steps_fixed_instances(self):
+        self._setup_execute_steps_tests()
+        cmd = mock.MagicMock()
+        cmd.instance_name = 'vm_to_delete'
+        cmds = tuple([cmd])  # yapf does not format (cmd,) correctly.
+        rolling_restart.execute_steps(appengine_admin=self._appengine_admin,
+                                      version=self._version,
+                                      cmds=cmds,
+                                      min_delay=0,
+                                      configured_num_instances=2)
+        self.assertEqual(self._appengine_request.execute.call_count, 3)
+
+
+if __name__ == '__main__':
+    unittest.main()

release/rollback/steps.py

@@ -122,6 +122,24 @@ def direct_service_traffic_to_version(
          '--quiet', f'--splits={version.version_id}=1', '--project', project))
 
 
+@dataclasses.dataclass(frozen=True)
+class KillNomulusInstance(RollbackStep):
+    """Step that kills a Nomulus VM instance."""
+    instance_name: str
+
+
+# yapf: disable
+def kill_nomulus_instance(project: str,
+                          version: common.VersionKey,
+                          instance_name: str) -> KillNomulusInstance:
+    # yapf: enable
+    return KillNomulusInstance(
+        'Delete one VM instance.',
+        ('gcloud', 'app', 'instances', 'delete', instance_name, '--quiet',
+         '--user-output-enabled=false', '--service', version.service_id,
+         '--version', version.version_id, '--project', project), instance_name)
+
+
 @dataclasses.dataclass(frozen=True)
 class _UpdateDeployTag(RollbackStep):
     """Updates the deployment tag on GCS."""