Fix flaky Spec11PipelineTest (#1133)

In SQL the contact of a domain is an indexed field and therefore we can
find linked domains synchronously, without the need for MapReduce.

The delete logic is mostly lifted from DeleteContactsAndHostsAction, but
because everything happens in a transaction we do not need to recheck a
lot of the preconditions that were necessary to ensure that the async
delete request still meets the conditions that when the request was
enqueued.

.gcloudignore

@@ -1,4 +1,5 @@
 python/
+node_modules/
 **/build/
 **/out/
 .*/

config/nom_build.py

@@ -256,6 +256,7 @@ GRADLE_FLAGS = [
                'Specify a task to be excluded from execution.',
                True),
 ]
+
 def generate_gradle_properties() -> str:
     """Returns the expected contents of gradle.properties."""
     out = io.StringIO()
@@ -270,7 +271,7 @@ def generate_gradle_properties() -> str:
 def get_root() -> str:
     """Returns the root of the nomulus build tree."""
     cur_dir = os.getcwd()
-    if not os.path.exists(os.path.join(cur_dir, '.git')) or \
+    if not os.path.exists(os.path.join(cur_dir, 'buildSrc')) or \
        not os.path.exists(os.path.join(cur_dir, 'core')) or \
        not os.path.exists(os.path.join(cur_dir, 'gradle.properties')):
         raise Exception('You must run this script from the root directory')

core/build.gradle

@@ -330,6 +330,7 @@ dependencies {
   testCompile deps['org.junit.platform:junit-platform-suite-api']
   testCompile deps['org.mockito:mockito-core']
   testCompile deps['org.mockito:mockito-junit-jupiter']
+  testCompile 'org.checkerframework:checker-qual:3.9.1'
   runtime deps['org.postgresql:postgresql']
 
   // Indirect dependency found by undeclared-dependency check. Such
@@ -633,8 +634,8 @@ compileProdJS.dependsOn soyToJS
 task karmaTest(type: Exec) {
   dependsOn ':npmInstall'
   workingDir rootProject.projectDir
-  executable 'node_modules/karma/bin/karma'
-  args('start', "${project.projectDir}/karma.conf.js")
+  executable '.gradle/nodejs/node-v14.15.5-linux-x64/bin/node'
+  args('node_modules/karma/bin/karma', 'start', "${project.projectDir}/karma.conf.js")
 }
 
 test.dependsOn karmaTest
@@ -805,6 +806,14 @@ if (environment in ['alpha', 'crash']) {
           mainClass: 'google.registry.beam.datastore.BulkDeleteDatastorePipeline',
           metaData: 'google/registry/beam/bulk_delete_datastore_pipeline_metadata.json'
       ],
+      [
+          mainClass: 'google.registry.beam.spec11.Spec11Pipeline',
+          metaData: 'google/registry/beam/spec11_pipeline_metadata.json'
+      ],
+      [
+          mainClass: 'google.registry.beam.invoicing.InvoicingPipeline',
+          metaData: 'google/registry/beam/invoicing_pipeline_metadata.json'
+      ],
   ]
   project.tasks.create("stage_beam_pipelines") {
     doLast {

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

@@ -109,6 +109,7 @@ import org.joda.time.Duration;
  * A mapreduce that processes batch asynchronous deletions of contact and host resources by mapping
  * over all domains and checking for any references to the contacts/hosts in pending deletion.
  */
+@Deprecated
 @Action(
     service = Action.Service.BACKEND,
     path = "/_dr/task/deleteContactsAndHosts",

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

@@ -19,6 +19,7 @@ import static javax.servlet.http.HttpServletResponse.SC_FORBIDDEN;
 import static javax.servlet.http.HttpServletResponse.SC_INTERNAL_SERVER_ERROR;
 import static javax.servlet.http.HttpServletResponse.SC_OK;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.flogger.FluentLogger;
 import google.registry.config.RegistryConfig.Config;
@@ -28,10 +29,11 @@ import google.registry.request.Response;
 import google.registry.request.auth.Auth;
 import google.registry.util.Retrier;
 import java.sql.Connection;
+import java.sql.ResultSet;
+import java.sql.SQLException;
 import java.sql.Statement;
 import java.util.function.Supplier;
 import javax.inject.Inject;
-import org.flywaydb.core.api.FlywayException;
 
 /**
  * Wipes out all Cloud SQL data in a Nomulus GCP environment.
@@ -80,13 +82,13 @@ public class WipeOutCloudSqlAction implements Runnable {
     try {
       retrier.callWithRetry(
           () -> {
-            try (Connection conn = connectionSupplier.get();
-                Statement statement = conn.createStatement()) {
-              statement.execute("drop owned by schema_deployer;");
+            try (Connection conn = connectionSupplier.get()) {
+              dropAllTables(conn, listTables(conn));
+              dropAllSequences(conn, listSequences(conn));
             }
             return null;
           },
-          e -> !(e instanceof FlywayException));
+          e -> !(e instanceof SQLException));
       response.setStatus(SC_OK);
       response.setPayload("Wiped out Cloud SQL in " + projectId);
     } catch (RuntimeException e) {
@@ -95,4 +97,69 @@ public class WipeOutCloudSqlAction implements Runnable {
       response.setPayload("Failed to wipe out Cloud SQL in " + projectId);
     }
   }
+
+  /** Returns a list of all tables in the public schema of a Postgresql database. */
+  static ImmutableList<String> listTables(Connection connection) throws SQLException {
+    try (ResultSet resultSet =
+        connection.getMetaData().getTables(null, null, null, new String[] {"TABLE"})) {
+      ImmutableList.Builder<String> tables = new ImmutableList.Builder<>();
+      while (resultSet.next()) {
+        String schema = resultSet.getString("TABLE_SCHEM");
+        if (schema == null || !schema.equalsIgnoreCase("public")) {
+          continue;
+        }
+        String tableName = resultSet.getString("TABLE_NAME");
+        tables.add("public.\"" + tableName + "\"");
+      }
+      return tables.build();
+    }
+  }
+
+  static void dropAllTables(Connection conn, ImmutableList<String> tables) throws SQLException {
+    if (tables.isEmpty()) {
+      return;
+    }
+
+    try (Statement statement = conn.createStatement()) {
+      for (String table : tables) {
+        statement.addBatch(String.format("DROP TABLE IF EXISTS %s CASCADE;", table));
+      }
+      for (int code : statement.executeBatch()) {
+        if (code == Statement.EXECUTE_FAILED) {
+          throw new RuntimeException("Failed to drop some tables. Please check.");
+        }
+      }
+    }
+  }
+
+  /** Returns a list of all sequences in a Postgresql database. */
+  static ImmutableList<String> listSequences(Connection conn) throws SQLException {
+    try (Statement statement = conn.createStatement();
+        ResultSet resultSet =
+            statement.executeQuery("SELECT c.relname FROM pg_class c WHERE c.relkind = 'S';")) {
+      ImmutableList.Builder<String> sequences = new ImmutableList.Builder<>();
+      while (resultSet.next()) {
+        sequences.add('\"' + resultSet.getString(1) + '\"');
+      }
+      return sequences.build();
+    }
+  }
+
+  static void dropAllSequences(Connection conn, ImmutableList<String> sequences)
+      throws SQLException {
+    if (sequences.isEmpty()) {
+      return;
+    }
+
+    try (Statement statement = conn.createStatement()) {
+      for (String sequence : sequences) {
+        statement.addBatch(String.format("DROP SEQUENCE IF EXISTS %s CASCADE;", sequence));
+      }
+      for (int code : statement.executeBatch()) {
+        if (code == Statement.EXECUTE_FAILED) {
+          throw new RuntimeException("Failed to drop some sequences. Please check.");
+        }
+      }
+    }
+  }
 }

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

@@ -15,6 +15,7 @@
 package google.registry.batch;
 
 import static com.google.common.net.MediaType.PLAIN_TEXT_UTF_8;
+import static google.registry.beam.BeamUtils.createJobName;
 import static javax.servlet.http.HttpServletResponse.SC_FORBIDDEN;
 import static javax.servlet.http.HttpServletResponse.SC_INTERNAL_SERVER_ERROR;
 import static javax.servlet.http.HttpServletResponse.SC_OK;
@@ -30,9 +31,8 @@ import google.registry.config.RegistryConfig.Config;
 import google.registry.request.Action;
 import google.registry.request.Response;
 import google.registry.request.auth.Auth;
+import google.registry.util.Clock;
 import javax.inject.Inject;
-import org.joda.time.DateTime;
-import org.joda.time.DateTimeZone;
 
 /**
  * Wipes out all Cloud Datastore data in a Nomulus GCP environment.
@@ -58,17 +58,20 @@ public class WipeoutDatastoreAction implements Runnable {
   private final Response response;
   private final Dataflow dataflow;
   private final String stagingBucketUrl;
+  private final Clock clock;
 
   @Inject
   WipeoutDatastoreAction(
       @Config("projectId") String projectId,
       @Config("defaultJobRegion") String jobRegion,
       @Config("beamStagingBucketUrl") String stagingBucketUrl,
+      Clock clock,
       Response response,
       Dataflow dataflow) {
     this.projectId = projectId;
     this.jobRegion = jobRegion;
     this.stagingBucketUrl = stagingBucketUrl;
+    this.clock = clock;
     this.response = response;
     this.dataflow = dataflow;
   }
@@ -86,10 +89,7 @@ public class WipeoutDatastoreAction implements Runnable {
     try {
       LaunchFlexTemplateParameter parameters =
           new LaunchFlexTemplateParameter()
-              // Job name must be unique and in [-a-z0-9].
-              .setJobName(
-                  "bulk-delete-datastore-"
-                      + DateTime.now(DateTimeZone.UTC).toString("yyyy-MM-dd'T'HH-mm-ss'Z'"))
+              .setJobName(createJobName("bulk-delete-datastore-", clock))
               .setContainerSpecGcsPath(
                   String.format("%s/%s_metadata.json", stagingBucketUrl, PIPELINE_NAME))
               .setParameters(ImmutableMap.of("kindsToDelete", "*"));

core/src/main/java/google/registry/beam/BeamUtils.java

@@ -14,10 +14,14 @@
 
 package google.registry.beam;
 
+import static com.google.common.base.Preconditions.checkArgument;
+
 import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableList;
 import com.google.common.io.Resources;
+import google.registry.util.Clock;
 import google.registry.util.ResourceUtils;
+import java.util.regex.Pattern;
 import org.apache.avro.generic.GenericRecord;
 import org.apache.beam.sdk.io.gcp.bigquery.SchemaAndRecord;
 
@@ -41,8 +45,7 @@ public class BeamUtils {
       ImmutableList<String> fieldNames, SchemaAndRecord schemaAndRecord) {
     GenericRecord record = schemaAndRecord.getRecord();
     ImmutableList<String> nullFields =
-        fieldNames
-            .stream()
+        fieldNames.stream()
             .filter(fieldName -> record.get(fieldName) == null)
             .collect(ImmutableList.toImmutableList());
     String missingFieldList = Joiner.on(", ").join(nullFields);
@@ -61,4 +64,19 @@ public class BeamUtils {
   public static String getQueryFromFile(Class<?> clazz, String filename) {
     return ResourceUtils.readResourceUtf8(Resources.getResource(clazz, "sql/" + filename));
   }
+
+  /** Creates a beam job name and validates that it conforms to the requirements. */
+  public static String createJobName(String prefix, Clock clock) {
+    // Flex template job name must be unique and consists of only characters [-a-z0-9], starting
+    // with a letter and ending with a letter or number. So we replace the "T" and "Z" in ISO 8601
+    // with lowercase letters.
+    String jobName =
+        String.format("%s-%s", prefix, clock.nowUtc().toString("yyyy-MM-dd't'HH-mm-ss'z'"));
+    checkArgument(
+        Pattern.compile("^[a-z][-a-z0-9]*[a-z0-9]*").matcher(jobName).matches(),
+        "The job name %s is illegal, it consists of only characters [-a-z0-9], "
+            + "starting with a letter and ending with a letter or number,",
+        jobName);
+    return jobName;
+  }
 }

core/src/main/java/google/registry/beam/common/RegistryJpaIO.java

@@ -23,11 +23,18 @@ import com.google.common.collect.Streams;
 import google.registry.backup.AppEngineEnvironment;
 import google.registry.model.ofy.ObjectifyService;
 import google.registry.persistence.transaction.JpaTransactionManager;
+import google.registry.persistence.transaction.QueryComposer;
 import google.registry.persistence.transaction.TransactionManagerFactory;
+import java.io.Serializable;
 import java.util.Objects;
 import java.util.concurrent.ThreadLocalRandom;
+import javax.persistence.criteria.CriteriaQuery;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.SerializableCoder;
 import org.apache.beam.sdk.metrics.Counter;
 import org.apache.beam.sdk.metrics.Metrics;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.transforms.Deduplicate;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.GroupIntoBatches;
 import org.apache.beam.sdk.transforms.PTransform;
@@ -36,6 +43,7 @@ import org.apache.beam.sdk.transforms.SerializableFunction;
 import org.apache.beam.sdk.transforms.WithKeys;
 import org.apache.beam.sdk.util.ShardedKey;
 import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PBegin;
 import org.apache.beam.sdk.values.PCollection;
 
 /**
@@ -51,10 +59,143 @@ public final class RegistryJpaIO {
 
   private RegistryJpaIO() {}
 
+  public static <R> Read<R, R> read(QueryComposerFactory<R> queryFactory) {
+    return Read.<R, R>builder().queryFactory(queryFactory).build();
+  }
+
+  public static <R, T> Read<R, T> read(
+      QueryComposerFactory<R> queryFactory, SerializableFunction<R, T> resultMapper) {
+    return Read.<R, T>builder().queryFactory(queryFactory).resultMapper(resultMapper).build();
+  }
+
   public static <T> Write<T> write() {
     return Write.<T>builder().build();
   }
 
+  // TODO(mmuller): Consider detached JpaQueryComposer that works with any JpaTransactionManager
+  // instance, i.e., change composer.buildQuery() to composer.buildQuery(JpaTransactionManager).
+  // This way QueryComposer becomes reusable and serializable (at least with Hibernate), and this
+  // interface would no longer be necessary.
+  public interface QueryComposerFactory<T>
+      extends SerializableFunction<JpaTransactionManager, QueryComposer<T>> {}
+
+  /**
+   * A {@link PTransform transform} that executes a JPA {@link CriteriaQuery} and adds the results
+   * to the BEAM pipeline. Users have the option to transform the results before sending them to the
+   * next stages.
+   *
+   * <p>The BEAM pipeline may execute this transform multiple times due to transient failures,
+   * loading duplicate results into the pipeline. Before we add dedepuplication support, the easiest
+   * workaround is to map results to {@link KV} pairs, and apply the {@link Deduplicate} transform
+   * to the output of this transform:
+   *
+   * <pre>{@code
+   * PCollection<String> contactIds =
+   *     pipeline
+   *         .apply(RegistryJpaIO.read(
+   *             (JpaTransactionManager tm) -> tm.createQueryComposer...,
+   *             contact -> KV.of(contact.getRepoId(), contact.getContactId()))
+   *         .withCoder(KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())))
+   *         .apply(Deduplicate.keyedValues())
+   *         .apply(Values.create());
+   * }</pre>
+   */
+  @AutoValue
+  public abstract static class Read<R, T> extends PTransform<PBegin, PCollection<T>> {
+
+    public static final String DEFAULT_NAME = "RegistryJpaIO.Read";
+
+    abstract String name();
+
+    abstract RegistryJpaIO.QueryComposerFactory<R> queryFactory();
+
+    abstract SerializableFunction<R, T> resultMapper();
+
+    abstract TransactionMode transactionMode();
+
+    abstract Coder<T> coder();
+
+    abstract Builder<R, T> toBuilder();
+
+    @Override
+    public PCollection<T> expand(PBegin input) {
+      return input
+          .apply("Starting " + name(), Create.of((Void) null))
+          .apply(
+              "Run query for " + name(),
+              ParDo.of(new QueryRunner<>(queryFactory(), resultMapper())))
+          .setCoder(coder());
+    }
+
+    public Read<R, T> withName(String name) {
+      return toBuilder().name(name).build();
+    }
+
+    public Read<R, T> withResultMapper(SerializableFunction<R, T> mapper) {
+      return toBuilder().resultMapper(mapper).build();
+    }
+
+    public Read<R, T> withTransactionMode(TransactionMode transactionMode) {
+      return toBuilder().transactionMode(transactionMode).build();
+    }
+
+    public Read<R, T> withCoder(Coder<T> coder) {
+      return toBuilder().coder(coder).build();
+    }
+
+    static <R, T> Builder<R, T> builder() {
+      return new AutoValue_RegistryJpaIO_Read.Builder()
+          .name(DEFAULT_NAME)
+          .resultMapper(x -> x)
+          .transactionMode(TransactionMode.TRANSACTIONAL)
+          .coder(SerializableCoder.of(Serializable.class));
+    }
+
+    @AutoValue.Builder
+    public abstract static class Builder<R, T> {
+
+      abstract Builder<R, T> name(String name);
+
+      abstract Builder<R, T> queryFactory(RegistryJpaIO.QueryComposerFactory<R> queryFactory);
+
+      abstract Builder<R, T> resultMapper(SerializableFunction<R, T> mapper);
+
+      abstract Builder<R, T> transactionMode(TransactionMode transactionMode);
+
+      abstract Builder<R, T> coder(Coder coder);
+
+      abstract Read<R, T> build();
+    }
+
+    static class QueryRunner<R, T> extends DoFn<Void, T> {
+      private final QueryComposerFactory<R> querySupplier;
+      private final SerializableFunction<R, T> resultMapper;
+
+      QueryRunner(QueryComposerFactory<R> querySupplier, SerializableFunction<R, T> resultMapper) {
+        this.querySupplier = querySupplier;
+        this.resultMapper = resultMapper;
+      }
+
+      @ProcessElement
+      public void processElement(OutputReceiver<T> outputReceiver) {
+        // TODO(b/187210388): JpaTransactionManager should support non-transactional query.
+        // TODO(weiminyu): add deduplication
+        jpaTm()
+            .transactNoRetry(
+                () ->
+                    querySupplier.apply(jpaTm()).stream()
+                        .map(resultMapper::apply)
+                        .forEach(outputReceiver::output));
+        // TODO(weiminyu): improve performance by reshuffle.
+      }
+    }
+
+    public enum TransactionMode {
+      NOT_TRANSACTIONAL,
+      TRANSACTIONAL;
+    }
+  }
+
   /**
    * A {@link PTransform transform} that writes a PCollection of entities to the SQL database using
    * the {@link JpaTransactionManager}.

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

@@ -87,20 +87,18 @@ public class BulkDeleteDatastorePipeline {
 
   private final BulkDeletePipelineOptions options;
 
-  private final Pipeline pipeline;
-
   BulkDeleteDatastorePipeline(BulkDeletePipelineOptions options) {
     this.options = options;
-    pipeline = Pipeline.create(options);
   }
 
   public void run() {
-    setupPipeline();
+    Pipeline pipeline = Pipeline.create(options);
+    setupPipeline(pipeline);
     pipeline.run();
   }
 
   @SuppressWarnings("deprecation") // org.apache.beam.sdk.transforms.Reshuffle
-  private void setupPipeline() {
+  private void setupPipeline(Pipeline pipeline) {
     checkState(
         !FORBIDDEN_PROJECTS.contains(options.getProject()),
         "Bulk delete is forbidden in %s",

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

@@ -1,201 +0,0 @@
-// Copyright 2020 The Nomulus Authors. All Rights Reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package google.registry.beam.initsql;
-
-import static com.google.common.base.Preconditions.checkArgument;
-import static com.google.common.base.Preconditions.checkState;
-import static com.google.common.base.Strings.isNullOrEmpty;
-
-import com.google.common.base.Splitter;
-import dagger.Component;
-import dagger.Lazy;
-import dagger.Module;
-import dagger.Provides;
-import google.registry.config.CredentialModule;
-import google.registry.config.RegistryConfig.Config;
-import google.registry.config.RegistryConfig.ConfigModule;
-import google.registry.keyring.kms.KmsModule;
-import google.registry.persistence.PersistenceModule;
-import google.registry.persistence.PersistenceModule.JdbcJpaTm;
-import google.registry.persistence.PersistenceModule.SocketFactoryJpaTm;
-import google.registry.persistence.PersistenceModule.TransactionIsolationLevel;
-import google.registry.persistence.transaction.JpaTransactionManager;
-import google.registry.privileges.secretmanager.SecretManagerModule;
-import google.registry.util.UtilsModule;
-import java.io.BufferedReader;
-import java.io.IOException;
-import java.io.InputStreamReader;
-import java.nio.channels.Channels;
-import java.nio.charset.StandardCharsets;
-import java.util.List;
-import javax.annotation.Nullable;
-import javax.inject.Singleton;
-import org.apache.beam.sdk.io.FileSystems;
-import org.apache.beam.sdk.io.fs.ResourceId;
-
-/**
- * Provides bindings for {@link JpaTransactionManager} to Cloud SQL.
- *
- * <p>This module is intended for use in BEAM pipelines, and uses a BEAM utility to access GCS like
- * a regular file system.
- */
-@Module
-public class BeamJpaModule {
-
-  private static final String GCS_SCHEME = "gs://";
-
-  @Nullable private final String sqlAccessInfoFile;
-  @Nullable private final String cloudKmsProjectId;
-  @Nullable private final TransactionIsolationLevel isolationOverride;
-
-  /**
-   * Constructs a new instance of {@link BeamJpaModule}.
-   *
-   * <p>Note: it is an unfortunately necessary antipattern to check for the validity of
-   * sqlAccessInfoFile in {@link #provideCloudSqlAccessInfo} rather than in the constructor.
-   * Unfortunately, this is a restriction imposed upon us by Dagger. Specifically, because we use
-   * this in at least one 1 {@link google.registry.tools.RegistryTool} command(s), it must be
-   * instantiated in {@code google.registry.tools.RegistryToolComponent} for all possible commands;
-   * Dagger doesn't permit it to ever be null. For the vast majority of commands, it will never be
-   * used (so a null credential file path is fine in those cases).
-   *
-   * @param sqlAccessInfoFile the path to a Cloud SQL credential file. This must refer to either a
-   *     real encrypted file on GCS as returned by {@link
-   *     BackupPaths#getCloudSQLCredentialFilePatterns} or an unencrypted file on local filesystem
-   *     with credentials to a test database.
-   * @param cloudKmsProjectId the GCP project where the credential decryption key can be found
-   * @param isolationOverride the desired Transaction Isolation level for all JDBC connections
-   */
-  public BeamJpaModule(
-      @Nullable String sqlAccessInfoFile,
-      @Nullable String cloudKmsProjectId,
-      @Nullable TransactionIsolationLevel isolationOverride) {
-    this.sqlAccessInfoFile = sqlAccessInfoFile;
-    this.cloudKmsProjectId = cloudKmsProjectId;
-    this.isolationOverride = isolationOverride;
-  }
-
-  public BeamJpaModule(@Nullable String sqlAccessInfoFile, @Nullable String cloudKmsProjectId) {
-    this(sqlAccessInfoFile, cloudKmsProjectId, null);
-  }
-
-  /** Returns true if the credential file is on GCS (and therefore expected to be encrypted). */
-  private boolean isCloudSqlCredential() {
-    return sqlAccessInfoFile.startsWith(GCS_SCHEME);
-  }
-
-  @Provides
-  @Singleton
-  SqlAccessInfo provideCloudSqlAccessInfo(Lazy<CloudSqlCredentialDecryptor> lazyDecryptor) {
-    checkArgument(!isNullOrEmpty(sqlAccessInfoFile), "Null or empty credentialFilePath");
-    String line = readOnlyLineFromCredentialFile();
-    if (isCloudSqlCredential()) {
-      line = lazyDecryptor.get().decrypt(line);
-    }
-    // See ./BackupPaths.java for explanation of the line format.
-    List<String> parts = Splitter.on(' ').splitToList(line.trim());
-    checkState(parts.size() == 3, "Expecting three phrases in %s", line);
-    if (isCloudSqlCredential()) {
-      return SqlAccessInfo.createCloudSqlAccessInfo(parts.get(0), parts.get(1), parts.get(2));
-    } else {
-      return SqlAccessInfo.createLocalSqlAccessInfo(parts.get(0), parts.get(1), parts.get(2));
-    }
-  }
-
-  String readOnlyLineFromCredentialFile() {
-    try {
-      ResourceId resourceId = FileSystems.matchSingleFileSpec(sqlAccessInfoFile).resourceId();
-      try (BufferedReader reader =
-          new BufferedReader(
-              new InputStreamReader(
-                  Channels.newInputStream(FileSystems.open(resourceId)), StandardCharsets.UTF_8))) {
-        return reader.readLine();
-      }
-    } catch (IOException e) {
-      throw new RuntimeException(e);
-    }
-  }
-
-  @Provides
-  @Config("beamCloudSqlJdbcUrl")
-  String provideJdbcUrl(SqlAccessInfo sqlAccessInfo) {
-    return sqlAccessInfo.jdbcUrl();
-  }
-
-  @Provides
-  @Config("beamCloudSqlInstanceConnectionName")
-  String provideSqlInstanceName(SqlAccessInfo sqlAccessInfo) {
-    return sqlAccessInfo
-        .cloudSqlInstanceName()
-        .orElseThrow(() -> new IllegalStateException("Cloud SQL not provisioned."));
-  }
-
-  @Provides
-  @Config("beamCloudSqlUsername")
-  String provideSqlUsername(SqlAccessInfo sqlAccessInfo) {
-    return sqlAccessInfo.user();
-  }
-
-  @Provides
-  @Config("beamCloudSqlPassword")
-  String provideSqlPassword(SqlAccessInfo sqlAccessInfo) {
-    return sqlAccessInfo.password();
-  }
-
-  @Provides
-  @Config("beamCloudKmsProjectId")
-  String kmsProjectId() {
-    return cloudKmsProjectId;
-  }
-
-  @Provides
-  @Config("beamCloudKmsKeyRing")
-  static String keyRingName() {
-    return "nomulus-tool-keyring";
-  }
-
-  @Provides
-  @Config("beamIsolationOverride")
-  @Nullable
-  TransactionIsolationLevel providesIsolationOverride() {
-    return isolationOverride;
-  }
-
-  @Provides
-  @Config("beamHibernateHikariMaximumPoolSize")
-  static int getBeamHibernateHikariMaximumPoolSize() {
-    // TODO(weiminyu): make this configurable. Should be equal to number of cores.
-    return 4;
-  }
-
-  @Singleton
-  @Component(
-      modules = {
-        ConfigModule.class,
-        CredentialModule.class,
-        BeamJpaModule.class,
-        KmsModule.class,
-        PersistenceModule.class,
-        SecretManagerModule.class,
-        UtilsModule.class
-      })
-  public interface JpaTransactionManagerComponent {
-    @SocketFactoryJpaTm
-    JpaTransactionManager cloudSqlJpaTransactionManager();
-
-    @JdbcJpaTm
-    JpaTransactionManager localDbJpaTransactionManager();
-  }
-}

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

@@ -120,26 +120,22 @@ public class InitSqlPipeline implements Serializable {
 
   private final InitSqlPipelineOptions options;
 
-  private final Pipeline pipeline;
-
   InitSqlPipeline(InitSqlPipelineOptions options) {
     this.options = options;
-    pipeline = Pipeline.create(options);
+  }
+
+  PipelineResult run() {
+    return run(Pipeline.create(options));
   }
 
   @VisibleForTesting
-  InitSqlPipeline(InitSqlPipelineOptions options, Pipeline pipeline) {
-    this.options = options;
-    this.pipeline = pipeline;
-  }
-
-  public PipelineResult run() {
-    setupPipeline();
+  PipelineResult run(Pipeline pipeline) {
+    setupPipeline(pipeline);
     return pipeline.run();
   }
 
   @VisibleForTesting
-  void setupPipeline() {
+  void setupPipeline(Pipeline pipeline) {
     options.setIsolationOverride(TransactionIsolationLevel.TRANSACTION_READ_UNCOMMITTED);
     PCollectionTuple datastoreSnapshot =
         pipeline.apply(

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

@@ -1,60 +0,0 @@
-// Copyright 2020 The Nomulus Authors. All Rights Reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package google.registry.beam.initsql;
-
-import google.registry.beam.initsql.BeamJpaModule.JpaTransactionManagerComponent;
-import google.registry.beam.initsql.Transforms.SerializableSupplier;
-import google.registry.persistence.PersistenceModule.TransactionIsolationLevel;
-import google.registry.persistence.transaction.JpaTransactionManager;
-import javax.annotation.Nullable;
-import org.apache.beam.sdk.transforms.SerializableFunction;
-
-public class JpaSupplierFactory implements SerializableSupplier<JpaTransactionManager> {
-
-  private static final long serialVersionUID = 1L;
-
-  private final String credentialFileUrl;
-  @Nullable private final String cloudKmsProjectId;
-  private final SerializableFunction<JpaTransactionManagerComponent, JpaTransactionManager>
-      jpaGetter;
-  @Nullable private final TransactionIsolationLevel isolationLevelOverride;
-
-  public JpaSupplierFactory(
-      String credentialFileUrl,
-      @Nullable String cloudKmsProjectId,
-      SerializableFunction<JpaTransactionManagerComponent, JpaTransactionManager> jpaGetter) {
-    this(credentialFileUrl, cloudKmsProjectId, jpaGetter, null);
-  }
-
-  public JpaSupplierFactory(
-      String credentialFileUrl,
-      @Nullable String cloudKmsProjectId,
-      SerializableFunction<JpaTransactionManagerComponent, JpaTransactionManager> jpaGetter,
-      @Nullable TransactionIsolationLevel isolationLevelOverride) {
-    this.credentialFileUrl = credentialFileUrl;
-    this.cloudKmsProjectId = cloudKmsProjectId;
-    this.jpaGetter = jpaGetter;
-    this.isolationLevelOverride = isolationLevelOverride;
-  }
-
-  @Override
-  public JpaTransactionManager get() {
-    return jpaGetter.apply(
-        DaggerBeamJpaModule_JpaTransactionManagerComponent.builder()
-            .beamJpaModule(
-                new BeamJpaModule(credentialFileUrl, cloudKmsProjectId, isolationLevelOverride))
-            .build());
-  }
-}

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

@@ -20,12 +20,9 @@ import static com.google.common.base.Preconditions.checkState;
 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.transaction.TransactionManagerFactory.jpaTm;
-import static google.registry.persistence.transaction.TransactionManagerFactory.setJpaTm;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
 import static google.registry.util.DateTimeUtils.isBeforeOrAt;
 import static java.util.Comparator.comparing;
-import static org.apache.beam.sdk.values.TypeDescriptors.integers;
 import static org.apache.beam.sdk.values.TypeDescriptors.kvs;
 import static org.apache.beam.sdk.values.TypeDescriptors.strings;
 
@@ -35,15 +32,14 @@ import com.google.appengine.api.datastore.EntityTranslator;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Streams;
-import com.googlecode.objectify.Key;
-import google.registry.backup.AppEngineEnvironment;
 import google.registry.backup.CommitLogImports;
 import google.registry.backup.VersionedEntity;
+import google.registry.model.billing.BillingEvent.Flag;
+import google.registry.model.billing.BillingEvent.Reason;
 import google.registry.model.domain.DomainBase;
-import google.registry.model.ofy.ObjectifyService;
 import google.registry.model.reporting.HistoryEntry;
-import google.registry.persistence.transaction.JpaTransactionManager;
 import google.registry.schema.replay.DatastoreAndSqlEntity;
 import google.registry.schema.replay.SqlEntity;
 import google.registry.tools.LevelDbLogReader;
@@ -53,7 +49,6 @@ import java.util.Iterator;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
-import java.util.concurrent.ThreadLocalRandom;
 import java.util.function.Supplier;
 import javax.annotation.Nullable;
 import org.apache.beam.sdk.coders.StringUtf8Coder;
@@ -62,18 +57,14 @@ import org.apache.beam.sdk.io.FileIO;
 import org.apache.beam.sdk.io.FileIO.ReadableFile;
 import org.apache.beam.sdk.io.fs.EmptyMatchTreatment;
 import org.apache.beam.sdk.io.fs.MatchResult.Metadata;
-import org.apache.beam.sdk.metrics.Counter;
-import org.apache.beam.sdk.metrics.Metrics;
 import org.apache.beam.sdk.transforms.Create;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.Flatten;
 import org.apache.beam.sdk.transforms.GroupByKey;
-import org.apache.beam.sdk.transforms.GroupIntoBatches;
 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;
@@ -268,81 +259,58 @@ public final class Transforms {
                     .iterator()));
   }
 
-  /**
-   * 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.
-   *
-   * @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}
-   */
-  public static PTransform<PCollection<VersionedEntity>, PCollection<Void>> writeToSql(
-      String transformId,
-      int maxWriters,
-      int batchSize,
-      SerializableSupplier<JpaTransactionManager> jpaSupplier) {
-    return writeToSql(
-        transformId,
-        maxWriters,
-        batchSize,
-        jpaSupplier,
-        Transforms::convertVersionedEntityToSqlEntity,
-        TypeDescriptor.of(VersionedEntity.class));
-  }
+  // Production data repair configs go below. See b/185954992.
 
-  /**
-   * 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<T> input) {
-        return input
-            .apply(
-                "Shard data for " + transformId,
-                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<T>(transformId, jpaSupplier, jpaConverter)));
-      }
-    };
-  }
+  // Prober domains in bad state, without associated contacts, hosts, billings, and history.
+  // They can be safely ignored.
+  private static final ImmutableSet<String> IGNORED_DOMAINS =
+      ImmutableSet.of("6AF6D2-IQCANT", "2-IQANYT");
 
-  private static Key toOfyKey(Object ofyEntity) {
-    return Key.create(ofyEntity);
-  }
+  // Prober hosts referencing phantom registrars. They and their associated history entries can be
+  // safely ignored.
+  private static final ImmutableSet<String> IGNORED_HOSTS =
+      ImmutableSet.of(
+          "4E21_WJ0TEST-GOOGLE",
+          "4E21_WJ1TEST-GOOGLE",
+          "4E21_WJ2TEST-GOOGLE",
+          "4E21_WJ3TEST-GOOGLE");
+
+  // Prober contacts referencing phantom registrars. They and their associated history entries can
+  // be safely ignored.
+  private static final ImmutableSet IGNORED_CONTACTS =
+      ImmutableSet.of(
+          "1_WJ0TEST-GOOGLE", "1_WJ1TEST-GOOGLE", "1_WJ2TEST-GOOGLE", "1_WJ3TEST-GOOGLE");
 
   private static boolean isMigratable(Entity entity) {
+    // Checks specific to production data. See b/185954992 for details.
+    // The names of these bad entities in production do not conflict with other environments. For
+    // simplicities sake we apply them regardless of the source of the data.
+    if (entity.getKind().equals("DomainBase")
+        && IGNORED_DOMAINS.contains(entity.getKey().getName())) {
+      return false;
+    }
+    if (entity.getKind().equals("ContactResource")) {
+      String roid = entity.getKey().getName();
+      return !IGNORED_CONTACTS.contains(roid);
+    }
+    if (entity.getKind().equals("HostResource")) {
+      String roid = entity.getKey().getName();
+      return !IGNORED_HOSTS.contains(roid);
+    }
+    if (entity.getKind().equals("HistoryEntry")) {
+      // Remove production bad data: History of the contacts to be ignored:
+      com.google.appengine.api.datastore.Key parentKey = entity.getKey().getParent();
+      if (parentKey.getKind().equals("ContactResource")) {
+        String contactRoid = parentKey.getName();
+        return !IGNORED_CONTACTS.contains(contactRoid);
+      }
+      if (parentKey.getKind().equals("HostResource")) {
+        String hostRoid = parentKey.getName();
+        return !IGNORED_HOSTS.contains(hostRoid);
+      }
+    }
+    // End of production-specific checks.
+
     if (entity.getKind().equals("HistoryEntry")) {
       // DOMAIN_APPLICATION_CREATE is deprecated type and should not be migrated.
       // The Enum name DOMAIN_APPLICATION_CREATE no longer exists in Java and cannot
@@ -352,6 +320,18 @@ public final class Transforms {
     return true;
   }
 
+  private static Entity repairBadData(Entity entity) {
+    if (entity.getKind().equals("Cancellation")
+        && Objects.equals(entity.getProperty("reason"), "AUTO_RENEW")) {
+      // AUTO_RENEW has been moved from 'reason' to flags. Change reason to RENEW and add the
+      // AUTO_RENEW flag. Note: all affected entities have empty flags so we can simply assign
+      // instead of append. See b/185954992.
+      entity.setUnindexedProperty("reason", Reason.RENEW.name());
+      entity.setUnindexedProperty("flags", ImmutableList.of(Flag.AUTO_RENEW.name()));
+    }
+    return entity;
+  }
+
   private static SqlEntity toSqlEntity(Object ofyEntity) {
     if (ofyEntity instanceof HistoryEntry) {
       HistoryEntry ofyHistory = (HistoryEntry) ofyEntity;
@@ -372,6 +352,7 @@ public final class Transforms {
     return dsEntity
         .getEntity()
         .filter(Transforms::isMigratable)
+        .map(Transforms::repairBadData)
         .map(e -> ofy().toPojo(e))
         .map(Transforms::toSqlEntity)
         .orElse(null);
@@ -458,93 +439,6 @@ public final class Transforms {
     }
   }
 
-  /**
-   * Writes a batch of entities to a SQL database.
-   *
-   * <p>Note that an arbitrary number of instances of this class may be created and freed in
-   * arbitrary order in a single JVM. Due to the tech debt that forced us to use a static variable
-   * 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<T> extends DoFn<KV<Integer, Iterable<T>>, Void> {
-
-    private static int instanceCount = 0;
-    private static JpaTransactionManager originalJpa;
-
-    private Counter counter;
-
-    private final SerializableSupplier<JpaTransactionManager> jpaSupplier;
-    private final SerializableFunction<T, Object> jpaConverter;
-
-    SqlBatchWriter(
-        String type,
-        SerializableSupplier<JpaTransactionManager> jpaSupplier,
-        SerializableFunction<T, Object> jpaConverter) {
-      counter = Metrics.counter("SQL_WRITE", type);
-      this.jpaSupplier = jpaSupplier;
-      this.jpaConverter = jpaConverter;
-    }
-
-    @Setup
-    public void setup() {
-      try (AppEngineEnvironment env = new AppEngineEnvironment()) {
-        ObjectifyService.initOfy();
-      }
-
-      synchronized (SqlBatchWriter.class) {
-        if (instanceCount == 0) {
-          originalJpa = jpaTm();
-          setJpaTm(jpaSupplier);
-        }
-        instanceCount++;
-      }
-    }
-
-    @Teardown
-    public void teardown() {
-      synchronized (SqlBatchWriter.class) {
-        instanceCount--;
-        if (instanceCount == 0) {
-          jpaTm().teardown();
-          setJpaTm(() -> originalJpa);
-        }
-      }
-    }
-
-    @ProcessElement
-    public void processElement(@Element KV<Integer, Iterable<T>> kv) {
-      try (AppEngineEnvironment env = new AppEngineEnvironment()) {
-        ImmutableList<Object> ofyEntities =
-            Streams.stream(kv.getValue())
-                .map(this.jpaConverter::apply)
-                // TODO(b/177340730): post migration delete the line below.
-                .filter(Objects::nonNull)
-                .collect(ImmutableList.toImmutableList());
-        try {
-          jpaTm().transact(() -> jpaTm().putAll(ofyEntities));
-          counter.inc(ofyEntities.size());
-        } catch (RuntimeException e) {
-          processSingly(ofyEntities);
-        }
-      }
-    }
-
-    /**
-     * Writes entities in a failed batch one by one to identify the first bad entity and throws a
-     * {@link RuntimeException} on it.
-     */
-    private void processSingly(ImmutableList<Object> ofyEntities) {
-      for (Object ofyEntity : ofyEntities) {
-        try {
-          jpaTm().transact(() -> jpaTm().put(ofyEntity));
-          counter.inc();
-        } catch (RuntimeException e) {
-          throw new RuntimeException(toOfyKey(ofyEntity).toString(), e);
-        }
-      }
-    }
-  }
-
   /**
    * Removes BillingEvents, {@link google.registry.model.poll.PollMessage PollMessages} and {@link
    * google.registry.model.host.HostResource} from a {@link DomainBase}. These are circular foreign

core/src/main/java/google/registry/beam/invoicing/BillingEvent.java

@@ -260,6 +260,11 @@ public abstract class BillingEvent implements Serializable {
         poNumber());
   }
 
+  /** Returns the grouping key for this {@code BillingEvent}, to generate the detailed report. */
+  String getDetailedReportGroupingKey() {
+    return String.format("%s_%s", registrarId(), tld());
+  }
+
   /** Key for each {@code BillingEvent}, when aggregating for the overall invoice. */
   @AutoValue
   abstract static class InvoiceGroupingKey implements Serializable {

core/src/main/java/google/registry/beam/invoicing/InvoicingPipeline.java

@@ -14,28 +14,27 @@
 
 package google.registry.beam.invoicing;
 
-import com.google.auth.oauth2.GoogleCredentials;
+import static google.registry.beam.BeamUtils.getQueryFromFile;
+import static org.apache.beam.sdk.values.TypeDescriptors.strings;
+
 import google.registry.beam.invoicing.BillingEvent.InvoiceGroupingKey;
 import google.registry.beam.invoicing.BillingEvent.InvoiceGroupingKey.InvoiceGroupingKeyCoder;
-import google.registry.config.CredentialModule.LocalCredential;
-import google.registry.config.RegistryConfig.Config;
 import google.registry.reporting.billing.BillingModule;
-import google.registry.reporting.billing.GenerateInvoicesAction;
-import google.registry.util.GoogleCredentialsBundle;
+import google.registry.util.SqlTemplate;
 import java.io.Serializable;
-import javax.inject.Inject;
-import org.apache.beam.runners.dataflow.DataflowRunner;
-import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions;
+import java.time.LocalDateTime;
+import java.time.LocalTime;
+import java.time.YearMonth;
+import java.time.format.DateTimeFormatter;
 import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.PipelineResult;
 import org.apache.beam.sdk.coders.SerializableCoder;
-import org.apache.beam.sdk.io.DefaultFilenamePolicy.Params;
-import org.apache.beam.sdk.io.FileBasedSink;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.io.FileIO;
 import org.apache.beam.sdk.io.TextIO;
 import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO;
-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.Contextful;
 import org.apache.beam.sdk.transforms.Count;
 import org.apache.beam.sdk.transforms.Filter;
 import org.apache.beam.sdk.transforms.MapElements;
@@ -43,107 +42,48 @@ import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.sdk.values.TypeDescriptor;
-import org.apache.beam.sdk.values.TypeDescriptors;
 
 /**
- * Definition of a Dataflow pipeline template, which generates a given month's invoices.
+ * Definition of a Dataflow Flex pipeline template, which generates a given month's invoices.
  *
- * <p>To stage this template on GCS, run the {@link
- * google.registry.tools.DeployInvoicingPipelineCommand} Nomulus command.
+ * <p>To stage this template locally, run the {@code stage_beam_pipeline.sh} shell script.
  *
  * <p>Then, you can run the staged template via the API client library, gCloud or a raw REST call.
- * For an example using the API client library, see {@link GenerateInvoicesAction}.
  *
- * @see <a href="https://cloud.google.com/dataflow/docs/templates/overview">Dataflow Templates</a>
+ * @see <a href="https://cloud.google.com/dataflow/docs/guides/templates/using-flex-templates">Using
+ *     Flex Templates</a>
  */
 public class InvoicingPipeline implements Serializable {
 
-  private final String projectId;
-  private final String beamJobRegion;
-  private final String beamBucketUrl;
-  private final String invoiceTemplateUrl;
-  private final String beamStagingUrl;
-  private final String billingBucketUrl;
-  private final String invoiceFilePrefix;
-  private final GoogleCredentials googleCredentials;
+  private static final DateTimeFormatter TIMESTAMP_FORMATTER =
+      DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss.SSSSSS");
 
-  @Inject
-  public InvoicingPipeline(
-      @Config("projectId") String projectId,
-      @Config("defaultJobRegion") String beamJobRegion,
-      @Config("apacheBeamBucketUrl") String beamBucketUrl,
-      @Config("invoiceTemplateUrl") String invoiceTemplateUrl,
-      @Config("beamStagingUrl") String beamStagingUrl,
-      @Config("billingBucketUrl") String billingBucketUrl,
-      @Config("invoiceFilePrefix") String invoiceFilePrefix,
-      @LocalCredential GoogleCredentialsBundle googleCredentialsBundle) {
-    this.projectId = projectId;
-    this.beamJobRegion = beamJobRegion;
-    this.beamBucketUrl = beamBucketUrl;
-    this.invoiceTemplateUrl = invoiceTemplateUrl;
-    this.beamStagingUrl = beamStagingUrl;
-    this.billingBucketUrl = billingBucketUrl;
-    this.invoiceFilePrefix = invoiceFilePrefix;
-    this.googleCredentials = googleCredentialsBundle.getGoogleCredentials();
+  private final InvoicingPipelineOptions options;
+
+  InvoicingPipeline(InvoicingPipelineOptions options) {
+    this.options = options;
   }
 
-  /** Custom options for running the invoicing pipeline. */
-  public interface InvoicingPipelineOptions extends DataflowPipelineOptions {
-    /** Returns the yearMonth we're generating invoices for, in yyyy-MM format. */
-    @Description("The yearMonth we generate invoices for, in yyyy-MM format.")
-    ValueProvider<String> getYearMonth();
-    /**
-     * Sets the yearMonth we generate invoices for.
-     *
-     * <p>This is implicitly set when executing the Dataflow template, by specifying the 'yearMonth
-     * parameter.
-     */
-    void setYearMonth(ValueProvider<String> value);
+  PipelineResult run() {
+    Pipeline pipeline = Pipeline.create(options);
+    setupPipeline(pipeline);
+    return pipeline.run();
   }
 
-  /** Deploys the invoicing pipeline as a template on GCS, for a given projectID and GCS bucket. */
-  public void deploy() {
-    // We can't store options as a member variable due to serialization concerns.
-    InvoicingPipelineOptions options = PipelineOptionsFactory.as(InvoicingPipelineOptions.class);
-    options.setProject(projectId);
-    options.setRegion(beamJobRegion);
-    options.setRunner(DataflowRunner.class);
-    // This causes p.run() to stage the pipeline as a template on GCS, as opposed to running it.
-    options.setTemplateLocation(invoiceTemplateUrl);
-    options.setStagingLocation(beamStagingUrl);
-    // This credential is used when Dataflow deploys the template to GCS in target GCP project.
-    // So, make sure the credential has write permission to GCS in that project.
-    options.setGcpCredential(googleCredentials);
-
-    Pipeline p = Pipeline.create(options);
-
+  void setupPipeline(Pipeline pipeline) {
     PCollection<BillingEvent> billingEvents =
-        p.apply(
+        pipeline.apply(
             "Read BillingEvents from Bigquery",
             BigQueryIO.read(BillingEvent::parseFromRecord)
-                .fromQuery(InvoicingUtils.makeQueryProvider(options.getYearMonth(), projectId))
+                .fromQuery(makeQuery(options.getYearMonth(), options.getProject()))
                 .withCoder(SerializableCoder.of(BillingEvent.class))
                 .usingStandardSql()
                 .withoutValidation()
                 .withTemplateCompatibility());
-    applyTerminalTransforms(billingEvents, options.getYearMonth());
-    p.run();
-  }
 
-  /**
-   * Applies output transforms to the {@code BillingEvent} source collection.
-   *
-   * <p>This is factored out purely to facilitate testing.
-   */
-  void applyTerminalTransforms(
-      PCollection<BillingEvent> billingEvents, ValueProvider<String> yearMonthProvider) {
-    billingEvents
-        .apply("Generate overall invoice rows", new GenerateInvoiceRows())
-        .apply("Write overall invoice to CSV", writeInvoice(yearMonthProvider));
+    saveInvoiceCsv(billingEvents, options);
 
-    billingEvents.apply(
-        "Write detail reports to separate CSVs keyed by registrarId_tld pair",
-        writeDetailReports(yearMonthProvider));
+    saveDetailedCsv(billingEvents, options);
   }
 
   /** Transform that converts a {@code BillingEvent} into an invoice CSV row. */
@@ -156,49 +96,85 @@ public class InvoicingPipeline implements Serializable {
               "Map to invoicing key",
               MapElements.into(TypeDescriptor.of(InvoiceGroupingKey.class))
                   .via(BillingEvent::getInvoiceGroupingKey))
-          .apply(Filter.by((InvoiceGroupingKey key) -> key.unitPrice() != 0))
+          .apply(
+              "Filter out free events", Filter.by((InvoiceGroupingKey key) -> key.unitPrice() != 0))
           .setCoder(new InvoiceGroupingKeyCoder())
           .apply("Count occurrences", Count.perElement())
           .apply(
               "Format as CSVs",
-              MapElements.into(TypeDescriptors.strings())
+              MapElements.into(strings())
                   .via((KV<InvoiceGroupingKey, Long> kv) -> kv.getKey().toCsv(kv.getValue())));
     }
   }
 
-  /** Returns an IO transform that writes the overall invoice to a single CSV file. */
-  private TextIO.Write writeInvoice(ValueProvider<String> yearMonthProvider) {
-    return TextIO.write()
-        .to(
-            NestedValueProvider.of(
-                yearMonthProvider,
-                yearMonth ->
+  /** Saves the billing events to a single overall invoice CSV file. */
+  static void saveInvoiceCsv(
+      PCollection<BillingEvent> billingEvents, InvoicingPipelineOptions options) {
+    billingEvents
+        .apply("Generate overall invoice rows", new GenerateInvoiceRows())
+        .apply(
+            "Write overall invoice to CSV",
+            TextIO.write()
+                .to(
                     String.format(
                         "%s/%s/%s/%s-%s",
-                        billingBucketUrl,
+                        options.getBillingBucketUrl(),
                         BillingModule.INVOICES_DIRECTORY,
-                        yearMonth,
-                        invoiceFilePrefix,
-                        yearMonth)))
-        .withHeader(InvoiceGroupingKey.invoiceHeader())
-        .withoutSharding()
-        .withSuffix(".csv");
+                        options.getYearMonth(),
+                        options.getInvoiceFilePrefix(),
+                        options.getYearMonth()))
+                .withHeader(InvoiceGroupingKey.invoiceHeader())
+                .withoutSharding()
+                .withSuffix(".csv"));
   }
 
-  /** Returns an IO transform that writes detail reports to registrar-tld keyed CSV files. */
-  private TextIO.TypedWrite<BillingEvent, Params> writeDetailReports(
-      ValueProvider<String> yearMonthProvider) {
-    return TextIO.<BillingEvent>writeCustomType()
-        .to(
-            InvoicingUtils.makeDestinationFunction(
-                String.format("%s/%s", billingBucketUrl, BillingModule.INVOICES_DIRECTORY),
-                yearMonthProvider),
-            InvoicingUtils.makeEmptyDestinationParams(billingBucketUrl + "/errors"))
-        .withFormatFunction(BillingEvent::toCsv)
-        .withoutSharding()
-        .withTempDirectory(
-            FileBasedSink.convertToFileResourceIfPossible(beamBucketUrl + "/temporary"))
-        .withHeader(BillingEvent.getHeader())
-        .withSuffix(".csv");
+  /** Saves the billing events to detailed report CSV files keyed by registrar-tld pairs. */
+  static void saveDetailedCsv(
+      PCollection<BillingEvent> billingEvents, InvoicingPipelineOptions options) {
+    String yearMonth = options.getYearMonth();
+    billingEvents.apply(
+        "Write detailed report for each registrar-tld pair",
+        FileIO.<String, BillingEvent>writeDynamic()
+            .to(
+                String.format(
+                    "%s/%s/%s",
+                    options.getBillingBucketUrl(), BillingModule.INVOICES_DIRECTORY, yearMonth))
+            .by(BillingEvent::getDetailedReportGroupingKey)
+            .withNumShards(1)
+            .withDestinationCoder(StringUtf8Coder.of())
+            .withNaming(
+                key ->
+                    (window, pane, numShards, shardIndex, compression) ->
+                        String.format(
+                            "%s_%s_%s.csv", BillingModule.DETAIL_REPORT_PREFIX, yearMonth, key))
+            .via(
+                Contextful.fn(BillingEvent::toCsv),
+                TextIO.sink().withHeader(BillingEvent.getHeader())));
+  }
+
+  /** Create the Bigquery query for a given project and yearMonth at runtime. */
+  static String makeQuery(String yearMonth, String projectId) {
+    // Get the timestamp endpoints capturing the entire month with microsecond precision
+    YearMonth reportingMonth = YearMonth.parse(yearMonth);
+    LocalDateTime firstMoment = reportingMonth.atDay(1).atTime(LocalTime.MIDNIGHT);
+    LocalDateTime lastMoment = reportingMonth.atEndOfMonth().atTime(LocalTime.MAX);
+    // Construct the month's query by filling in the billing_events.sql template
+    return SqlTemplate.create(getQueryFromFile(InvoicingPipeline.class, "billing_events.sql"))
+        .put("FIRST_TIMESTAMP_OF_MONTH", firstMoment.format(TIMESTAMP_FORMATTER))
+        .put("LAST_TIMESTAMP_OF_MONTH", lastMoment.format(TIMESTAMP_FORMATTER))
+        .put("PROJECT_ID", projectId)
+        .put("DATASTORE_EXPORT_DATA_SET", "latest_datastore_export")
+        .put("ONETIME_TABLE", "OneTime")
+        .put("REGISTRY_TABLE", "Registry")
+        .put("REGISTRAR_TABLE", "Registrar")
+        .put("CANCELLATION_TABLE", "Cancellation")
+        .build();
+  }
+
+  public static void main(String[] args) {
+    PipelineOptionsFactory.register(InvoicingPipelineOptions.class);
+    InvoicingPipelineOptions options =
+        PipelineOptionsFactory.fromArgs(args).withValidation().as(InvoicingPipelineOptions.class);
+    new InvoicingPipeline(options).run();
   }
 }

core/src/main/java/google/registry/beam/invoicing/InvoicingPipelineOptions.java

@@ -0,0 +1,37 @@
+// Copyright 2021 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.beam.invoicing;
+
+import google.registry.beam.common.RegistryPipelineOptions;
+import org.apache.beam.sdk.options.Description;
+
+/** Custom options for running the invoicing pipeline. */
+public interface InvoicingPipelineOptions extends RegistryPipelineOptions {
+
+  @Description("The year and month we generate invoices for, in yyyy-MM format.")
+  String getYearMonth();
+
+  void setYearMonth(String value);
+
+  @Description("Filename prefix for the invoice CSV file.")
+  String getInvoiceFilePrefix();
+
+  void setInvoiceFilePrefix(String value);
+
+  @Description("The GCS bucket URL for invoices and detailed  reports to be uploaded.")
+  String getBillingBucketUrl();
+
+  void setBillingBucketUrl(String value);
+}

core/src/main/java/google/registry/beam/invoicing/InvoicingUtils.java

@@ -1,106 +0,0 @@
-// Copyright 2018 The Nomulus Authors. All Rights Reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package google.registry.beam.invoicing;
-
-import static google.registry.beam.BeamUtils.getQueryFromFile;
-
-import google.registry.util.SqlTemplate;
-import java.time.LocalDateTime;
-import java.time.LocalTime;
-import java.time.YearMonth;
-import java.time.format.DateTimeFormatter;
-import org.apache.beam.sdk.io.DefaultFilenamePolicy.Params;
-import org.apache.beam.sdk.io.FileBasedSink;
-import org.apache.beam.sdk.options.ValueProvider;
-import org.apache.beam.sdk.options.ValueProvider.NestedValueProvider;
-import org.apache.beam.sdk.transforms.SerializableFunction;
-
-/** Pipeline helper functions used to generate invoices from instances of {@link BillingEvent}. */
-public class InvoicingUtils {
-
-  private InvoicingUtils() {}
-
-  private static final DateTimeFormatter TIMESTAMP_FORMATTER =
-      DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss.SSSSSS");
-
-  /**
-   * Returns a function mapping from {@code BillingEvent} to filename {@code Params}.
-   *
-   * <p>Beam uses this to determine which file a given {@code BillingEvent} should get placed into.
-   *
-   * @param outputBucket the GCS bucket we're outputting reports to
-   * @param yearMonthProvider a runtime provider for the yyyy-MM we're generating the invoice for
-   */
-  static SerializableFunction<BillingEvent, Params> makeDestinationFunction(
-      String outputBucket, ValueProvider<String> yearMonthProvider) {
-    return billingEvent ->
-        new Params()
-            .withShardTemplate("")
-            .withSuffix(".csv")
-            .withBaseFilename(
-                NestedValueProvider.of(
-                    yearMonthProvider,
-                    yearMonth ->
-                        FileBasedSink.convertToFileResourceIfPossible(
-                            String.format(
-                                "%s/%s/%s",
-                                outputBucket, yearMonth, billingEvent.toFilename(yearMonth)))));
-  }
-
-  /**
-   * Returns the default filename parameters for an unmappable {@code BillingEvent}.
-   *
-   * <p>The "failed" file should only be populated when an error occurs, which warrants further
-   * investigation.
-   */
-  static Params makeEmptyDestinationParams(String outputBucket) {
-    return new Params()
-        .withBaseFilename(
-            FileBasedSink.convertToFileResourceIfPossible(
-                String.format("%s/%s", outputBucket, "FAILURES")));
-  }
-
-  /**
-   * Returns a provider that creates a Bigquery query for a given project and yearMonth at runtime.
-   *
-   * <p>We only know yearMonth at runtime, so this provider fills in the {@code
-   * sql/billing_events.sql} template at runtime.
-   *
-   * @param yearMonthProvider a runtime provider that returns which month we're invoicing for.
-   * @param projectId the projectId we're generating invoicing for.
-   */
-  static ValueProvider<String> makeQueryProvider(
-      ValueProvider<String> yearMonthProvider, String projectId) {
-    return NestedValueProvider.of(
-        yearMonthProvider,
-        (yearMonth) -> {
-          // Get the timestamp endpoints capturing the entire month with microsecond precision
-          YearMonth reportingMonth = YearMonth.parse(yearMonth);
-          LocalDateTime firstMoment = reportingMonth.atDay(1).atTime(LocalTime.MIDNIGHT);
-          LocalDateTime lastMoment = reportingMonth.atEndOfMonth().atTime(LocalTime.MAX);
-          // Construct the month's query by filling in the billing_events.sql template
-          return SqlTemplate.create(getQueryFromFile(InvoicingPipeline.class, "billing_events.sql"))
-              .put("FIRST_TIMESTAMP_OF_MONTH", firstMoment.format(TIMESTAMP_FORMATTER))
-              .put("LAST_TIMESTAMP_OF_MONTH", lastMoment.format(TIMESTAMP_FORMATTER))
-              .put("PROJECT_ID", projectId)
-              .put("DATASTORE_EXPORT_DATA_SET", "latest_datastore_export")
-              .put("ONETIME_TABLE", "OneTime")
-              .put("REGISTRY_TABLE", "Registry")
-              .put("REGISTRAR_TABLE", "Registrar")
-              .put("CANCELLATION_TABLE", "Cancellation")
-              .build();
-        });
-  }
-}

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

@@ -14,7 +14,6 @@
 
 package google.registry.beam.spec11;
 
-
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.apache.http.HttpStatus.SC_OK;
 
@@ -30,7 +29,6 @@ import java.net.URISyntaxException;
 import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.function.Supplier;
-import org.apache.beam.sdk.options.ValueProvider;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
 import org.apache.beam.sdk.values.KV;
@@ -73,7 +71,7 @@ public class SafeBrowsingTransforms {
     private static final int BATCH_SIZE = 490;
 
     /** Provides the SafeBrowsing API key at runtime. */
-    private final ValueProvider<String> apiKeyProvider;
+    private final String apiKey;
 
     /**
      * Maps a subdomain's {@code fullyQualifiedDomainName} to its corresponding {@link Subdomain} to
@@ -93,20 +91,18 @@ public class SafeBrowsingTransforms {
     private final Retrier retrier;
 
     /**
-     * Constructs a {@link EvaluateSafeBrowsingFn} that gets its API key from the given provider.
+     * Constructs a {@link EvaluateSafeBrowsingFn} with a given API key.
      *
      * <p>We need to dual-cast the closeableHttpClientSupplier lambda because all {@code DoFn}
      * member variables need to be serializable. The (Supplier & Serializable) dual cast is safe
      * because class methods are generally serializable, especially a static function such as {@link
      * HttpClients#createDefault()}.
-     *
-     * @param apiKeyProvider provides the SafeBrowsing API key from {@code KMS} at runtime
      */
     @SuppressWarnings("unchecked")
-    EvaluateSafeBrowsingFn(ValueProvider<String> apiKeyProvider, Retrier retrier) {
-      this.apiKeyProvider = apiKeyProvider;
+    EvaluateSafeBrowsingFn(String apiKey, Retrier retrier) {
+      this.apiKey = apiKey;
       this.retrier = retrier;
-      this.closeableHttpClientSupplier = (Supplier & Serializable) HttpClients::createDefault;
+      closeableHttpClientSupplier = (Supplier & Serializable) HttpClients::createDefault;
     }
 
     /**
@@ -117,12 +113,10 @@ public class SafeBrowsingTransforms {
      */
     @VisibleForTesting
     EvaluateSafeBrowsingFn(
-        ValueProvider<String> apiKeyProvider,
-        Retrier retrier,
-        Supplier<CloseableHttpClient> clientSupplier) {
-      this.apiKeyProvider = apiKeyProvider;
+        String apiKey, Retrier retrier, Supplier<CloseableHttpClient> clientSupplier) {
+      this.apiKey = apiKey;
       this.retrier = retrier;
-      this.closeableHttpClientSupplier = clientSupplier;
+      closeableHttpClientSupplier = clientSupplier;
     }
 
     /** Evaluates any buffered {@link Subdomain} objects upon completing the bundle. */
@@ -159,7 +153,7 @@ public class SafeBrowsingTransforms {
       try {
         URIBuilder uriBuilder = new URIBuilder(SAFE_BROWSING_URL);
         // Add the API key param
-        uriBuilder.addParameter("key", apiKeyProvider.get());
+        uriBuilder.addParameter("key", apiKey);
 
         HttpPost httpPost = new HttpPost(uriBuilder.build());
         httpPost.addHeader(HTTP.CONTENT_TYPE, ContentType.APPLICATION_JSON.toString());
@@ -175,7 +169,7 @@ public class SafeBrowsingTransforms {
               }
             },
             IOException.class);
-      } catch (URISyntaxException | JSONException  e) {
+      } catch (URISyntaxException | JSONException e) {
         // Fail the pipeline on a parsing exception- this indicates the API likely changed.
         throw new RuntimeException("Caught parsing exception, failing pipeline.", e);
       } finally {
@@ -239,7 +233,9 @@ public class SafeBrowsingTransforms {
             String url = match.getJSONObject("threat").getString("url");
             Subdomain subdomain = subdomainBuffer.get(url);
             resultBuilder.add(
-                KV.of(subdomain, ThreatMatch.create(match, subdomain.domainName())));
+                KV.of(
+                    subdomain,
+                    ThreatMatch.create(match.getString("threatType"), subdomain.domainName())));
           }
         }
       }

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

@@ -17,32 +17,27 @@ package google.registry.beam.spec11;
 import static com.google.common.base.Preconditions.checkArgument;
 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.beam.initsql.Transforms;
-import google.registry.beam.initsql.Transforms.SerializableSupplier;
+import dagger.Component;
+import dagger.Module;
+import dagger.Provides;
+import google.registry.beam.common.RegistryJpaIO;
 import google.registry.beam.spec11.SafeBrowsingTransforms.EvaluateSafeBrowsingFn;
-import google.registry.config.CredentialModule.LocalCredential;
-import google.registry.config.RegistryConfig.Config;
+import google.registry.config.RegistryConfig.ConfigModule;
 import google.registry.model.reporting.Spec11ThreatMatch;
 import google.registry.model.reporting.Spec11ThreatMatch.ThreatType;
-import google.registry.persistence.transaction.JpaTransactionManager;
-import google.registry.util.GoogleCredentialsBundle;
 import google.registry.util.Retrier;
 import google.registry.util.SqlTemplate;
+import google.registry.util.UtilsModule;
 import java.io.Serializable;
-import javax.inject.Inject;
-import org.apache.beam.runners.dataflow.DataflowRunner;
-import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions;
+import javax.inject.Singleton;
 import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.PipelineResult;
 import org.apache.beam.sdk.coders.SerializableCoder;
 import org.apache.beam.sdk.io.TextIO;
 import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO;
-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.GroupByKey;
 import org.apache.beam.sdk.transforms.MapElements;
 import org.apache.beam.sdk.transforms.ParDo;
@@ -58,21 +53,20 @@ import org.json.JSONException;
 import org.json.JSONObject;
 
 /**
- * Definition of a Dataflow pipeline template, which generates a given month's spec11 report.
+ * Definition of a Dataflow Flex template, which generates a given month's spec11 report.
  *
- * <p>To stage this template on GCS, run the {@link
- * google.registry.tools.DeploySpec11PipelineCommand} Nomulus command.
+ * <p>To stage this template locally, run the {@code stage_beam_pipeline.sh} shell script.
  *
  * <p>Then, you can run the staged template via the API client library, gCloud or a raw REST call.
  *
- * @see <a href="https://cloud.google.com/dataflow/docs/templates/overview">Dataflow Templates</a>
+ * @see <a href="https://cloud.google.com/dataflow/docs/guides/templates/using-flex-templates">Using
+ *     Flex Templates</a>
  */
 public class Spec11Pipeline implements Serializable {
 
   /**
    * Returns the subdirectory spec11 reports reside in for a given local date in yyyy-MM-dd format.
    *
-   * @see google.registry.beam.spec11.Spec11Pipeline
    * @see google.registry.reporting.spec11.Spec11EmailUtils
    */
   public static String getSpec11ReportFilePath(LocalDate localDate) {
@@ -87,84 +81,28 @@ public class Spec11Pipeline implements Serializable {
   /** The JSON object field into which we put the threat match array for Spec11 reports. */
   public static final String THREAT_MATCHES_FIELD = "threatMatches";
 
-  private final String projectId;
-  private final String beamJobRegion;
-  private final String beamStagingUrl;
-  private final String spec11TemplateUrl;
-  private final String reportingBucketUrl;
-  private final GoogleCredentials googleCredentials;
-  private final Retrier retrier;
-  private final SerializableSupplier<JpaTransactionManager> jpaSupplierFactory;
+  private final Spec11PipelineOptions options;
+  private final EvaluateSafeBrowsingFn safeBrowsingFn;
 
-  @Inject
-  public Spec11Pipeline(
-      @Config("projectId") String projectId,
-      @Config("defaultJobRegion") String beamJobRegion,
-      @Config("beamStagingUrl") String beamStagingUrl,
-      @Config("spec11TemplateUrl") String spec11TemplateUrl,
-      @Config("reportingBucketUrl") String reportingBucketUrl,
-      SerializableSupplier<JpaTransactionManager> jpaSupplierFactory,
-      @LocalCredential GoogleCredentialsBundle googleCredentialsBundle,
-      Retrier retrier) {
-    this.projectId = projectId;
-    this.beamJobRegion = beamJobRegion;
-    this.beamStagingUrl = beamStagingUrl;
-    this.spec11TemplateUrl = spec11TemplateUrl;
-    this.reportingBucketUrl = reportingBucketUrl;
-    this.jpaSupplierFactory = jpaSupplierFactory;
-    this.googleCredentials = googleCredentialsBundle.getGoogleCredentials();
-    this.retrier = retrier;
+  Spec11Pipeline(Spec11PipelineOptions options, EvaluateSafeBrowsingFn safeBrowsingFn) {
+    this.options = options;
+    this.safeBrowsingFn = safeBrowsingFn;
   }
 
-  /** Custom options for running the spec11 pipeline. */
-  public interface Spec11PipelineOptions extends DataflowPipelineOptions {
-    /** Returns the local date we're generating the report for, in yyyy-MM-dd format. */
-    @Description("The local date we generate the report for, in yyyy-MM-dd format.")
-    ValueProvider<String> getDate();
-
-    /**
-     * Sets the local date we generate invoices for.
-     *
-     * <p>This is implicitly set when executing the Dataflow template, by specifying the "date"
-     * parameter.
-     */
-    void setDate(ValueProvider<String> value);
-
-    /** Returns the SafeBrowsing API key we use to evaluate subdomain health. */
-    @Description("The API key we use to access the SafeBrowsing API.")
-    ValueProvider<String> getSafeBrowsingApiKey();
-
-    /**
-     * Sets the SafeBrowsing API key we use.
-     *
-     * <p>This is implicitly set when executing the Dataflow template, by specifying the
-     * "safeBrowsingApiKey" parameter.
-     */
-    void setSafeBrowsingApiKey(ValueProvider<String> value);
+  PipelineResult run() {
+    Pipeline pipeline = Pipeline.create(options);
+    setupPipeline(pipeline);
+    return pipeline.run();
   }
 
-  /** Deploys the spec11 pipeline as a template on GCS. */
-  public void deploy() {
-    // We can't store options as a member variable due to serialization concerns.
-    Spec11PipelineOptions options = PipelineOptionsFactory.as(Spec11PipelineOptions.class);
-    options.setProject(projectId);
-    options.setRegion(beamJobRegion);
-    options.setRunner(DataflowRunner.class);
-    // This causes p.run() to stage the pipeline as a template on GCS, as opposed to running it.
-    options.setTemplateLocation(spec11TemplateUrl);
-    options.setStagingLocation(beamStagingUrl);
-    // This credential is used when Dataflow deploys the template to GCS in target GCP project.
-    // So, make sure the credential has write permission to GCS in that project.
-    options.setGcpCredential(googleCredentials);
-
-    Pipeline p = Pipeline.create(options);
+  void setupPipeline(Pipeline pipeline) {
     PCollection<Subdomain> domains =
-        p.apply(
+        pipeline.apply(
             "Read active domains from BigQuery",
             BigQueryIO.read(Subdomain::parseFromRecord)
                 .fromQuery(
                     SqlTemplate.create(getQueryFromFile(Spec11Pipeline.class, "subdomains.sql"))
-                        .put("PROJECT_ID", projectId)
+                        .put("PROJECT_ID", options.getProject())
                         .put("DATASTORE_EXPORT_DATASET", "latest_datastore_export")
                         .put("REGISTRAR_TABLE", "Registrar")
                         .put("DOMAIN_BASE_TABLE", "DomainBase")
@@ -174,48 +112,40 @@ public class Spec11Pipeline implements Serializable {
                 .withoutValidation()
                 .withTemplateCompatibility());
 
-    evaluateUrlHealth(
-        domains,
-        new EvaluateSafeBrowsingFn(options.getSafeBrowsingApiKey(), retrier),
-        options.getDate());
-    p.run();
+    PCollection<KV<Subdomain, ThreatMatch>> threatMatches =
+        domains.apply("Run through SafeBrowsing API", ParDo.of(safeBrowsingFn));
+
+    saveToSql(threatMatches, options);
+    saveToGcs(threatMatches, options);
   }
 
-  /**
-   * Evaluate each {@link Subdomain} URL via the SafeBrowsing API.
-   *
-   * <p>This is factored out to facilitate testing.
-   */
-  void evaluateUrlHealth(
-      PCollection<Subdomain> domains,
-      EvaluateSafeBrowsingFn evaluateSafeBrowsingFn,
-      ValueProvider<String> dateProvider) {
-    PCollection<KV<Subdomain, ThreatMatch>> subdomainsSql =
-        domains.apply("Run through SafeBrowsing API", ParDo.of(evaluateSafeBrowsingFn));
-    TypeDescriptor<KV<Subdomain, ThreatMatch>> descriptor =
-        new TypeDescriptor<KV<Subdomain, ThreatMatch>>() {};
-    subdomainsSql.apply(
-        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));
+  static void saveToSql(
+      PCollection<KV<Subdomain, ThreatMatch>> threatMatches, Spec11PipelineOptions options) {
+    String transformId = "Spec11 Threat Matches";
+    LocalDate date = LocalDate.parse(options.getDate(), ISODateTimeFormat.date());
+    threatMatches.apply(
+        "Write to Sql: " + transformId,
+        RegistryJpaIO.<KV<Subdomain, ThreatMatch>>write()
+            .withName(transformId)
+            .withBatchSize(options.getSqlWriteBatchSize())
+            .withShards(options.getSqlWriteShards())
+            .withJpaConverter(
+                (kv) -> {
+                  Subdomain subdomain = kv.getKey();
+                  return new Spec11ThreatMatch.Builder()
+                      .setThreatTypes(
+                          ImmutableSet.of(ThreatType.valueOf(kv.getValue().threatType())))
+                      .setCheckDate(date)
+                      .setDomainName(subdomain.domainName())
+                      .setDomainRepoId(subdomain.domainRepoId())
+                      .setRegistrarId(subdomain.registrarId())
+                      .build();
+                }));
+  }
 
-    /* Store ThreatMatch objects in JSON. */
-    PCollection<KV<Subdomain, ThreatMatch>> subdomainsJson =
-        domains.apply("Run through SafeBrowsingAPI", ParDo.of(evaluateSafeBrowsingFn));
-    subdomainsJson
+  static void saveToGcs(
+      PCollection<KV<Subdomain, ThreatMatch>> threatMatches, Spec11PipelineOptions options) {
+    threatMatches
         .apply(
             "Map registrar ID to email/ThreatMatch pair",
             MapElements.into(
@@ -260,17 +190,54 @@ public class Spec11Pipeline implements Serializable {
             "Output to text file",
             TextIO.write()
                 .to(
-                    NestedValueProvider.of(
-                        dateProvider,
-                        date ->
-                            String.format(
-                                "%s/%s",
-                                reportingBucketUrl,
-                                getSpec11ReportFilePath(LocalDate.parse(date)))))
+                    String.format(
+                        "%s/%s",
+                        options.getReportingBucketUrl(),
+                        getSpec11ReportFilePath(LocalDate.parse(options.getDate()))))
                 .withoutSharding()
                 .withHeader("Map from registrar email / name to detected subdomain threats:"));
   }
 
+  public static void main(String[] args) {
+    PipelineOptionsFactory.register(Spec11PipelineOptions.class);
+    DaggerSpec11Pipeline_Spec11PipelineComponent.builder()
+        .spec11PipelineModule(new Spec11PipelineModule(args))
+        .build()
+        .spec11Pipeline()
+        .run();
+  }
+
+  @Module
+  static class Spec11PipelineModule {
+    private final String[] args;
+
+    Spec11PipelineModule(String[] args) {
+      this.args = args;
+    }
+
+    @Provides
+    Spec11PipelineOptions provideOptions() {
+      return PipelineOptionsFactory.fromArgs(args).withValidation().as(Spec11PipelineOptions.class);
+    }
+
+    @Provides
+    EvaluateSafeBrowsingFn provideSafeBrowsingFn(Spec11PipelineOptions options, Retrier retrier) {
+      return new EvaluateSafeBrowsingFn(options.getSafeBrowsingApiKey(), retrier);
+    }
+
+    @Provides
+    Spec11Pipeline providePipeline(
+        Spec11PipelineOptions options, EvaluateSafeBrowsingFn safeBrowsingFn) {
+      return new Spec11Pipeline(options, safeBrowsingFn);
+    }
+  }
+
+  @Component(modules = {Spec11PipelineModule.class, UtilsModule.class, ConfigModule.class})
+  @Singleton
+  interface Spec11PipelineComponent {
+    Spec11Pipeline spec11Pipeline();
+  }
+
   @AutoValue
   abstract static class EmailAndThreatMatch implements Serializable {
 

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

@@ -0,0 +1,37 @@
+// Copyright 2021 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.beam.spec11;
+
+import google.registry.beam.common.RegistryPipelineOptions;
+import org.apache.beam.sdk.options.Description;
+
+/** Custom options for running the spec11 pipeline. */
+public interface Spec11PipelineOptions extends RegistryPipelineOptions {
+
+  @Description("The local date we generate the report for, in yyyy-MM-dd format.")
+  String getDate();
+
+  void setDate(String value);
+
+  @Description("The API key we use to access the SafeBrowsing API.")
+  String getSafeBrowsingApiKey();
+
+  void setSafeBrowsingApiKey(String value);
+
+  @Description("The GCS bucket URL for Spec11 reports to be uploaded.")
+  String getReportingBucketUrl();
+
+  void setReportingBucketUrl(String value);
+}

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

@@ -15,6 +15,7 @@
 package google.registry.beam.spec11;
 
 import com.google.auto.value.AutoValue;
+import com.google.common.annotations.VisibleForTesting;
 import java.io.Serializable;
 import org.json.JSONException;
 import org.json.JSONObject;
@@ -31,16 +32,9 @@ public abstract class ThreatMatch implements Serializable {
   /** Returns the fully qualified domain name [SLD].[TLD] of the matched threat. */
   public abstract String fullyQualifiedDomainName();
 
-  /**
-   * Constructs a {@link ThreatMatch} by parsing a {@code SafeBrowsing API} response {@link
-   * JSONObject}.
-   *
-   * @throws JSONException when encountering parse errors in the response format
-   */
-  static ThreatMatch create(JSONObject threatMatchJSON, String fullyQualifiedDomainName)
-      throws JSONException {
-    return new AutoValue_ThreatMatch(
-        threatMatchJSON.getString(THREAT_TYPE_FIELD), fullyQualifiedDomainName);
+  @VisibleForTesting
+  static ThreatMatch create(String threatType, String fullyQualifiedDomainName) {
+    return new AutoValue_ThreatMatch(threatType, fullyQualifiedDomainName);
   }
 
   /** Returns a {@link JSONObject} representing a subset of this object's data. */

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

@@ -384,19 +384,6 @@ public final class RegistryConfig {
       return Duration.standardHours(1);
     }
 
-    /**
-     * Number of sharded entity group roots used for performing strongly consistent scans.
-     *
-     * <p><b>Warning:</b> This number may increase but never decrease.
-     *
-     * @see google.registry.model.index.EppResourceIndex
-     */
-    @Provides
-    @Config("eppResourceIndexBucketCount")
-    public static int provideEppResourceIndexBucketCount(RegistryConfigSettings config) {
-      return config.datastore.eppResourceIndexBucketsNum;
-    }
-
     @Provides
     @Config("cloudSqlJdbcUrl")
     public static String providesCloudSqlJdbcUrl(RegistryConfigSettings config) {
@@ -564,53 +551,6 @@ public final class RegistryConfig {
       return config.gSuite.outgoingEmailDisplayName;
     }
 
-    /**
-     * Returns the name of the GCS bucket for storing Beam templates and results.
-     *
-     * @see google.registry.reporting.billing.GenerateInvoicesAction
-     */
-    @Provides
-    @Config("apacheBeamBucket")
-    public static String provideApacheBeamBucket(@Config("projectId") String projectId) {
-      return projectId + "-beam";
-    }
-
-    /**
-     * Returns the URL of the GCS location for storing Apache Beam related objects.
-     *
-     * @see google.registry.reporting.billing.GenerateInvoicesAction
-     */
-    @Provides
-    @Config("apacheBeamBucketUrl")
-    public static String provideApacheBeamBucketUrl(@Config("apacheBeamBucket") String beamBucket) {
-      return "gs://" + beamBucket;
-    }
-
-    /**
-     * Returns the URL of the GCS location for storing the monthly invoicing Beam template.
-     *
-     * @see google.registry.reporting.billing.GenerateInvoicesAction
-     * @see google.registry.beam.invoicing.InvoicingPipeline
-     */
-    @Provides
-    @Config("invoiceTemplateUrl")
-    public static String provideInvoiceTemplateUrl(
-        @Config("apacheBeamBucketUrl") String beamBucketUrl) {
-      return beamBucketUrl + "/templates/invoicing";
-    }
-
-    /**
-     * Returns the URL of the GCS location for storing the monthly spec11 Beam template.
-     *
-     * @see google.registry.beam.spec11.Spec11Pipeline
-     */
-    @Provides
-    @Config("spec11TemplateUrl")
-    public static String provideSpec11TemplateUrl(
-        @Config("apacheBeamBucketUrl") String beamBucketUrl) {
-      return beamBucketUrl + "/templates/spec11";
-    }
-
     /**
      * Returns whether an SSL certificate hash is required to log in via EPP and run flows.
      *
@@ -634,18 +574,6 @@ public final class RegistryConfig {
       return config.beam.defaultJobRegion;
     }
 
-    /**
-     * Returns the default job zone to run Apache Beam (Cloud Dataflow) jobs in.
-     *
-     * @see google.registry.reporting.billing.GenerateInvoicesAction
-     * @see google.registry.reporting.spec11.GenerateSpec11ReportAction
-     */
-    @Provides
-    @Config("defaultJobZone")
-    public static String provideDefaultJobZone(RegistryConfigSettings config) {
-      return config.beam.defaultJobZone;
-    }
-
     /** Returns the GCS bucket URL with all staged BEAM flex templates. */
     @Provides
     @Config("beamStagingBucketUrl")
@@ -653,19 +581,6 @@ public final class RegistryConfig {
       return config.beam.stagingBucketUrl;
     }
 
-    /**
-     * Returns the URL of the GCS location we store jar dependencies for beam pipelines.
-     *
-     * @see google.registry.beam.invoicing.InvoicingPipeline
-     * @see google.registry.beam.spec11.Spec11Pipeline
-     */
-    @Provides
-    @Config("beamStagingUrl")
-    public static String provideInvoiceStagingUrl(
-        @Config("apacheBeamBucketUrl") String beamBucketUrl) {
-      return beamBucketUrl + "/staging";
-    }
-
     /**
      * Returns the Google Cloud Storage bucket for Spec11 and ICANN transaction and activity reports
      * to be uploaded.
@@ -1227,14 +1142,6 @@ public final class RegistryConfig {
       return formatComments(config.registryPolicy.reservedTermsExportDisclaimer);
     }
 
-    /** Returns the clientId of the registrar used by the {@code CheckApiServlet}. */
-    // TODO(b/80417678): remove this once CheckApiAction no longer uses this id.
-    @Provides
-    @Config("checkApiServletRegistrarClientId")
-    public static String provideCheckApiServletRegistrarClientId(RegistryConfigSettings config) {
-      return config.registryPolicy.checkApiServletClientId;
-    }
-
     /**
      * Returns the clientId of the registrar that admins are automatically logged in as if they
      * aren't otherwise associated with one.

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

@@ -133,7 +133,6 @@ public class RegistryConfigSettings {
   /** Configuration for Apache Beam (Cloud Dataflow). */
   public static class Beam {
     public String defaultJobRegion;
-    public String defaultJobZone;
     public String stagingBucketUrl;
   }
 

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

@@ -420,9 +420,6 @@ misc:
 beam:
   # The default region to run Apache Beam (Cloud Dataflow) jobs in.
   defaultJobRegion: us-east1
-  # The default zone to run Apache Beam (Cloud Dataflow) jobs in.
-  # TODO(weiminyu): consider dropping zone config. No obvious needs for this.
-  defaultJobZone: us-east1-c
   stagingBucketUrl: gcs-bucket-with-staged-templates
 
 keyring:

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

@@ -20,6 +20,8 @@ import com.google.common.base.Strings;
 import dagger.Module;
 import dagger.Provides;
 import google.registry.flows.picker.FlowPicker;
+import google.registry.model.contact.ContactHistory;
+import google.registry.model.domain.DomainHistory;
 import google.registry.model.domain.metadata.MetadataExtension;
 import google.registry.model.eppcommon.AuthInfo;
 import google.registry.model.eppcommon.Trid;
@@ -239,6 +241,18 @@ public class FlowModule {
     return historyBuilder;
   }
 
+  @Provides
+  static ContactHistory.Builder provideContactHistoryBuilder(
+      HistoryEntry.Builder historyEntryBuilder) {
+    return new ContactHistory.Builder().copyFrom(historyEntryBuilder);
+  }
+
+  @Provides
+  static DomainHistory.Builder provideDomainHistoryBuilder(
+      HistoryEntry.Builder historyEntryBuilder) {
+    return new DomainHistory.Builder().copyFrom(historyEntryBuilder);
+  }
+
   /**
    * Provides a partially filled in {@link EppResponse} builder.
    *

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

@@ -22,15 +22,19 @@ import static java.nio.charset.StandardCharsets.UTF_8;
 
 import com.google.common.base.Throwables;
 import com.google.common.flogger.FluentLogger;
+import com.googlecode.objectify.Key;
 import google.registry.flows.EppException.CommandUseErrorException;
 import google.registry.flows.EppException.ParameterValueRangeErrorException;
 import google.registry.flows.EppException.SyntaxErrorException;
 import google.registry.flows.EppException.UnimplementedProtocolVersionException;
 import google.registry.flows.custom.EntityChanges;
+import google.registry.model.EppResource;
 import google.registry.model.eppcommon.EppXmlTransformer;
 import google.registry.model.eppinput.EppInput.WrongProtocolVersionException;
 import google.registry.model.eppoutput.EppOutput;
 import google.registry.model.host.InetAddressAdapter.IpVersionMismatchException;
+import google.registry.model.ofy.ObjectifyService;
+import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.translators.CurrencyUnitAdapter.UnknownCurrencyException;
 import google.registry.xml.XmlException;
 import java.util.List;
@@ -99,6 +103,11 @@ public final class FlowUtils {
     }
   }
 
+  public static <H extends HistoryEntry> Key<H> createHistoryKey(
+      EppResource parent, Class<H> clazz) {
+    return Key.create(Key.create(parent), clazz, ObjectifyService.allocateId());
+  }
+
   /** Registrar is not logged in. */
   public static class NotLoggedInException extends CommandUseErrorException {
     public NotLoggedInException() {

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

@@ -16,6 +16,7 @@ package google.registry.flows;
 
 import static com.google.common.collect.Sets.intersection;
 import static google.registry.model.EppResourceUtils.getLinkedDomainKeys;
+import static google.registry.model.EppResourceUtils.isLinked;
 import static google.registry.model.EppResourceUtils.loadByForeignKey;
 import static google.registry.model.index.ForeignKeyIndex.loadAndGetKey;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
@@ -62,7 +63,10 @@ public final class ResourceFlowUtils {
 
   private ResourceFlowUtils() {}
 
-  /** In {@link #failfastForAsyncDelete}, check this (arbitrary) number of query results. */
+  /**
+   * In {@link #checkLinkedDomains(String, DateTime, Class, Function)}, check this (arbitrary)
+   * number of query results.
+   */
   private static final int FAILFAST_CHECK_COUNT = 5;
 
   /** Check that the given clientId corresponds to the owner of given resource. */
@@ -73,36 +77,54 @@ public final class ResourceFlowUtils {
     }
   }
 
-  /** Check whether an asynchronous delete would obviously fail, and throw an exception if so. */
-  public static <R extends EppResource> void failfastForAsyncDelete(
+  /**
+   * Check whether if there are domains linked to the resource to be deleted. Throws an exception if
+   * so.
+   *
+   * <p>Note that in datastore this is a smoke test as the query for linked domains is eventually
+   * consistent, so we only check a few domains to fail fast.
+   */
+  public static <R extends EppResource> void checkLinkedDomains(
       final String targetId,
       final DateTime now,
       final Class<R> resourceClass,
-      final Function<DomainBase, ImmutableSet<?>> getPotentialReferences) throws EppException {
-    // Enter a transactionless context briefly.
+      final Function<DomainBase, ImmutableSet<?>> getPotentialReferences)
+      throws EppException {
     EppException failfastException =
-        tm().doTransactionless(
-                () -> {
-                  final ForeignKeyIndex<R> fki = ForeignKeyIndex.load(resourceClass, targetId, now);
-                  if (fki == null) {
-                    return new ResourceDoesNotExistException(resourceClass, targetId);
-                  }
-                  /* Query for the first few linked domains, and if found, actually load them. The
-                   * query is eventually consistent and so might be very stale, but the direct
-                   * load will not be stale, just non-transactional. If we find at least one
-                   * actual reference then we can reliably fail. If we don't find any, we can't
-                   * trust the query and need to do the full mapreduce.
-                   */
-                  Iterable<VKey<DomainBase>> keys =
-                      getLinkedDomainKeys(fki.getResourceKey(), now, FAILFAST_CHECK_COUNT);
+        tm().isOfy()
+            ? tm().doTransactionless(
+                    () -> {
+                      final ForeignKeyIndex<R> fki =
+                          ForeignKeyIndex.load(resourceClass, targetId, now);
+                      if (fki == null) {
+                        return new ResourceDoesNotExistException(resourceClass, targetId);
+                      }
+                      // Query for the first few linked domains, and if found, actually load them.
+                      // The query is eventually consistent and so might be very stale, but the
+                      // direct load will not be stale, just non-transactional. If we find at least
+                      // one  actual reference then we can reliably fail. If we don't find any,
+                      // we can't  trust the query and need to do the full mapreduce.
+                      Iterable<VKey<DomainBase>> keys =
+                          getLinkedDomainKeys(fki.getResourceKey(), now, FAILFAST_CHECK_COUNT);
 
-                  VKey<R> resourceVKey = fki.getResourceKey();
-                  Predicate<DomainBase> predicate =
-                      domain -> getPotentialReferences.apply(domain).contains(resourceVKey);
-                  return tm().loadByKeys(keys).values().stream().anyMatch(predicate)
-                      ? new ResourceToDeleteIsReferencedException()
-                      : null;
-                });
+                      VKey<R> resourceVKey = fki.getResourceKey();
+                      Predicate<DomainBase> predicate =
+                          domain -> getPotentialReferences.apply(domain).contains(resourceVKey);
+                      return tm().loadByKeys(keys).values().stream().anyMatch(predicate)
+                          ? new ResourceToDeleteIsReferencedException()
+                          : null;
+                    })
+            : tm().transact(
+                    () -> {
+                      final ForeignKeyIndex<R> fki =
+                          ForeignKeyIndex.load(resourceClass, targetId, now);
+                      if (fki == null) {
+                        return new ResourceDoesNotExistException(resourceClass, targetId);
+                      }
+                      return isLinked(fki.getResourceKey(), now)
+                          ? new ResourceToDeleteIsReferencedException()
+                          : null;
+                    });
     if (failfastException != null) {
       throw failfastException;
     }
@@ -123,8 +145,7 @@ public final class ResourceFlowUtils {
   }
 
   public static <R extends EppResource & ForeignKeyedEppResource> R loadAndVerifyExistence(
-      Class<R> clazz, String targetId, DateTime now)
-          throws ResourceDoesNotExistException {
+      Class<R> clazz, String targetId, DateTime now) throws ResourceDoesNotExistException {
     return verifyExistence(clazz, targetId, loadByForeignKey(clazz, targetId, now));
   }
 
@@ -156,16 +177,16 @@ public final class ResourceFlowUtils {
   }
 
   /** Check that the given AuthInfo is either missing or else is valid for the given resource. */
-  public static void verifyOptionalAuthInfo(
-      Optional<AuthInfo> authInfo, ContactResource contact) throws EppException {
+  public static void verifyOptionalAuthInfo(Optional<AuthInfo> authInfo, ContactResource contact)
+      throws EppException {
     if (authInfo.isPresent()) {
       verifyAuthInfo(authInfo.get(), contact);
     }
   }
 
   /** Check that the given AuthInfo is either missing or else is valid for the given resource. */
-  public static void verifyOptionalAuthInfo(
-      Optional<AuthInfo> authInfo, DomainBase domain) throws EppException {
+  public static void verifyOptionalAuthInfo(Optional<AuthInfo> authInfo, DomainBase domain)
+      throws EppException {
     if (authInfo.isPresent()) {
       verifyAuthInfo(authInfo.get(), domain);
     }
@@ -229,7 +250,7 @@ public final class ResourceFlowUtils {
   /** Check that the same values aren't being added and removed in an update command. */
   public static void checkSameValuesNotAddedAndRemoved(
       ImmutableSet<?> fieldsToAdd, ImmutableSet<?> fieldsToRemove)
-          throws AddRemoveSameValueException {
+      throws AddRemoveSameValueException {
     if (!intersection(fieldsToAdd, fieldsToRemove).isEmpty()) {
       throw new AddRemoveSameValueException();
     }

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

@@ -16,7 +16,6 @@ package google.registry.flows;
 
 import static com.google.common.base.MoreObjects.toStringHelper;
 import static google.registry.request.RequestParameters.extractOptionalHeader;
-import static google.registry.util.X509Utils.loadCertificate;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
@@ -26,24 +25,17 @@ import com.google.common.net.InetAddresses;
 import dagger.Module;
 import dagger.Provides;
 import google.registry.config.RegistryConfig.Config;
-import google.registry.config.RegistryEnvironment;
 import google.registry.flows.EppException.AuthenticationErrorException;
 import google.registry.flows.certs.CertificateChecker;
 import google.registry.flows.certs.CertificateChecker.InsecureCertificateException;
 import google.registry.model.registrar.Registrar;
 import google.registry.request.Header;
 import google.registry.util.CidrAddressBlock;
-import google.registry.util.Clock;
 import google.registry.util.ProxyHttpHeaders;
-import java.io.ByteArrayInputStream;
 import java.net.InetAddress;
-import java.security.cert.CertificateException;
-import java.security.cert.X509Certificate;
-import java.util.Base64;
 import java.util.Optional;
 import javax.inject.Inject;
 import javax.servlet.http.HttpServletRequest;
-import org.joda.time.DateTime;
 
 /**
  * Container and validation for TLS certificate and IP-allow-listing.
@@ -54,10 +46,6 @@ import org.joda.time.DateTime;
  *   <dt>X-SSL-Certificate
  *   <dd>This field should contain a base64 encoded digest of the client's TLS certificate. It is
  *       used only if the validation of the full certificate fails.
- *   <dt>X-SSL-Full-Certificate
- *   <dd>This field should contain a base64 encoding of the client's TLS certificate. It is
- *       validated during an EPP login command against a known good value that is transmitted out of
- *       band.
  *   <dt>X-Forwarded-For
  *   <dd>This field should contain the host and port of the connecting client. It is validated
  *       during an EPP login command against an IP allow list that is transmitted out of band.
@@ -66,30 +54,22 @@ import org.joda.time.DateTime;
 public class TlsCredentials implements TransportCredentials {
 
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
-  private static final DateTime CERT_ENFORCEMENT_START_TIME =
-      DateTime.parse("2021-03-01T16:00:00Z");
 
   private final boolean requireSslCertificates;
   private final Optional<String> clientCertificateHash;
-  private final Optional<String> clientCertificate;
   private final Optional<InetAddress> clientInetAddr;
   private final CertificateChecker certificateChecker;
-  private final Clock clock;
 
   @Inject
   public TlsCredentials(
       @Config("requireSslCertificates") boolean requireSslCertificates,
       @Header(ProxyHttpHeaders.CERTIFICATE_HASH) Optional<String> clientCertificateHash,
-      @Header(ProxyHttpHeaders.FULL_CERTIFICATE) Optional<String> clientCertificate,
       @Header(ProxyHttpHeaders.IP_ADDRESS) Optional<String> clientAddress,
-      CertificateChecker certificateChecker,
-      Clock clock) {
+      CertificateChecker certificateChecker) {
     this.requireSslCertificates = requireSslCertificates;
     this.clientCertificateHash = clientCertificateHash;
-    this.clientCertificate = clientCertificate;
     this.clientInetAddr = clientAddress.map(TlsCredentials::parseInetAddress);
     this.certificateChecker = certificateChecker;
-    this.clock = clock;
   }
 
   static InetAddress parseInetAddress(String asciiAddr) {
@@ -103,7 +83,7 @@ public class TlsCredentials implements TransportCredentials {
   @Override
   public void validate(Registrar registrar, String password) throws AuthenticationErrorException {
     validateIp(registrar);
-    validateCertificate(registrar);
+    validateCertificateHash(registrar);
     validatePassword(registrar, password);
   }
 
@@ -137,89 +117,8 @@ public class TlsCredentials implements TransportCredentials {
     throw new BadRegistrarIpAddressException();
   }
 
-  /**
-   * Verifies client SSL certificate is permitted to issue commands as {@code registrar}.
-   *
-   * @throws MissingRegistrarCertificateException if frontend didn't send certificate header
-   * @throws BadRegistrarCertificateException if registrar requires certificate and it didn't match
-   */
   @VisibleForTesting
-  void validateCertificate(Registrar registrar) throws AuthenticationErrorException {
-    // Check that certificate is present in registrar object
-    if (!registrar.getClientCertificate().isPresent()
-        && !registrar.getFailoverClientCertificate().isPresent()) {
-      // Log an error and validate using certificate hash instead
-      // TODO(sarahbot): throw a RegistrarCertificateNotConfiguredException once hash is no longer
-      // used as failover
-      logger.atWarning().log(
-          "There is no certificate configured for registrar %s.", registrar.getClientId());
-    } else if (!clientCertificate.isPresent()) {
-      // Check that the request included the full certificate
-      // Log an error and validate using certificate hash instead
-      // TODO(sarahbot): throw a MissingRegistrarCertificateException once hash is no longer used as
-      // failover
-      logger.atWarning().log(
-          "Request from registrar %s did not include X-SSL-Full-Certificate.",
-          registrar.getClientId());
-    } else {
-      X509Certificate passedCert;
-      Optional<X509Certificate> storedCert;
-      Optional<X509Certificate> storedFailoverCert;
-
-      try {
-        storedCert = deserializePemCert(registrar.getClientCertificate());
-        storedFailoverCert = deserializePemCert(registrar.getFailoverClientCertificate());
-        passedCert = decodeCertString(clientCertificate.get());
-      } catch (Exception e) {
-        // TODO(Sarahbot@): remove this catch once we know it's working
-        logger.atWarning().log(
-            "Error converting certificate string to certificate for %s: %s",
-            registrar.getClientId(), e);
-        validateCertificateHash(registrar);
-        return;
-      }
-
-      // Check if the certificate is equal to the one on file for the registrar.
-      if (passedCert.equals(storedCert.orElse(null))
-          || passedCert.equals(storedFailoverCert.orElse(null))) {
-        // Check certificate for any requirement violations
-        // TODO(Sarahbot@): Throw exceptions instead of just logging once requirement enforcement
-        // begins
-        try {
-          certificateChecker.validateCertificate(passedCert);
-        } catch (InsecureCertificateException e) {
-          // TODO(Sarahbot@): Remove this if statement after March 1. After March 1, exception
-          // should be thrown in all environments.
-          // throw exception in unit tests and Sandbox
-          if (RegistryEnvironment.get().equals(RegistryEnvironment.UNITTEST)
-              || RegistryEnvironment.get().equals(RegistryEnvironment.SANDBOX)
-              || clock.nowUtc().isAfter(CERT_ENFORCEMENT_START_TIME)) {
-            throw new CertificateContainsSecurityViolationsException(e);
-          }
-          logger.atWarning().log(
-              "Registrar certificate used for %s does not meet certificate requirements: %s",
-              registrar.getClientId(), e.getMessage());
-        } catch (Exception e) {
-          logger.atWarning().log(
-              "Error validating certificate for %s: %s", registrar.getClientId(), e);
-        }
-        // successfully validated, return here since hash validation is not necessary
-        return;
-      }
-      // Log an error and validate using certificate hash instead
-      // TODO(sarahbot): throw a BadRegistrarCertificateException once hash is no longer used as
-      // failover
-      logger.atWarning().log("Non-matching certificate for registrar %s.", registrar.getClientId());
-    }
-    validateCertificateHash(registrar);
-  }
-
-  private void validateCertificateHash(Registrar registrar) throws AuthenticationErrorException {
-    logger.atWarning().log(
-        "Error validating certificate for %s, attempting to validate using certificate hash.",
-        registrar.getClientId());
-    // Check the certificate hash as a failover
-    // TODO(sarahbot): Remove hash checks once certificate checks are working.
+  void validateCertificateHash(Registrar registrar) throws AuthenticationErrorException {
     if (!registrar.getClientCertificateHash().isPresent()
         && !registrar.getFailoverClientCertificateHash().isPresent()) {
       if (requireSslCertificates) {
@@ -247,6 +146,20 @@ public class TlsCredentials implements TransportCredentials {
           registrar.getFailoverClientCertificateHash());
       throw new BadRegistrarCertificateException();
     }
+    if (requireSslCertificates) {
+      String passedCert =
+          clientCertificateHash.equals(registrar.getClientCertificateHash())
+              ? registrar.getClientCertificate().get()
+              : registrar.getFailoverClientCertificate().get();
+      try {
+        certificateChecker.validateCertificate(passedCert);
+      } catch (InsecureCertificateException e) {
+        logger.atWarning().log(
+            "Registrar certificate used for %s does not meet certificate requirements: %s",
+            registrar.getClientId(), e.getMessage());
+        throw new CertificateContainsSecurityViolationsException(e);
+      }
+    }
   }
 
   private void validatePassword(Registrar registrar, String password)
@@ -256,26 +169,9 @@ public class TlsCredentials implements TransportCredentials {
     }
   }
 
-  // Converts a PEM formatted certificate string into an X509Certificate
-  private Optional<X509Certificate> deserializePemCert(Optional<String> certificateString)
-      throws CertificateException {
-    if (certificateString.isPresent()) {
-      return Optional.of(loadCertificate(certificateString.get()));
-    }
-    return Optional.empty();
-  }
-
-  // Decodes the string representation of an encoded certificate back into an X509Certificate
-  private X509Certificate decodeCertString(String encodedCertString) throws CertificateException {
-    byte decodedCert[] = Base64.getDecoder().decode(encodedCertString);
-    ByteArrayInputStream inputStream = new ByteArrayInputStream(decodedCert);
-    return loadCertificate(inputStream);
-  }
-
   @Override
   public String toString() {
     return toStringHelper(getClass())
-        .add("clientCertificate", clientCertificate.orElse(null))
         .add("clientCertificateHash", clientCertificateHash.orElse(null))
         .add("clientAddress", clientInetAddr.orElse(null))
         .toString();
@@ -336,14 +232,6 @@ public class TlsCredentials implements TransportCredentials {
       return extractOptionalHeader(req, ProxyHttpHeaders.CERTIFICATE_HASH);
     }
 
-    @Provides
-    @Header(ProxyHttpHeaders.FULL_CERTIFICATE)
-    static Optional<String> provideClientCertificate(HttpServletRequest req) {
-      // Note: This header is actually required, we just want to handle its absence explicitly
-      // by throwing an EPP exception rather than a generic Bad Request exception.
-      return extractOptionalHeader(req, ProxyHttpHeaders.FULL_CERTIFICATE);
-    }
-
     @Provides
     @Header(ProxyHttpHeaders.IP_ADDRESS)
     static Optional<String> provideIpAddress(HttpServletRequest req) {

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

@@ -33,6 +33,7 @@ import google.registry.flows.annotations.ReportingSpec;
 import google.registry.flows.exceptions.ResourceAlreadyExistsForThisClientException;
 import google.registry.flows.exceptions.ResourceCreateContentionException;
 import google.registry.model.contact.ContactCommand.Create;
+import google.registry.model.contact.ContactHistory;
 import google.registry.model.contact.ContactResource;
 import google.registry.model.domain.metadata.MetadataExtension;
 import google.registry.model.eppinput.ResourceCommand;
@@ -61,7 +62,7 @@ public final class ContactCreateFlow implements TransactionalFlow {
   @Inject ExtensionManager extensionManager;
   @Inject @ClientId String clientId;
   @Inject @TargetId String targetId;
-  @Inject HistoryEntry.Builder historyBuilder;
+  @Inject ContactHistory.Builder historyBuilder;
   @Inject EppResponse.Builder responseBuilder;
   @Inject @Config("contactAndHostRoidSuffix") String roidSuffix;
   @Inject ContactCreateFlow() {}
@@ -93,12 +94,12 @@ public final class ContactCreateFlow implements TransactionalFlow {
     historyBuilder
         .setType(HistoryEntry.Type.CONTACT_CREATE)
         .setModificationTime(now)
-        .setXmlBytes(null)  // We don't want to store contact details in the history entry.
-        .setParent(Key.create(newContact));
+        .setXmlBytes(null) // We don't want to store contact details in the history entry.
+        .setContactBase(newContact);
     tm().insertAll(
             ImmutableSet.of(
                 newContact,
-                historyBuilder.build().toChildHistoryEntity(),
+                historyBuilder.build(),
                 ForeignKeyIndex.create(newContact, newContact.getDeletionTime()),
                 EppResourceIndex.create(Key.create(newContact))));
     return responseBuilder

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

@@ -15,16 +15,19 @@
 package google.registry.flows.contact;
 
 import static google.registry.flows.FlowUtils.validateClientIsLoggedIn;
-import static google.registry.flows.ResourceFlowUtils.failfastForAsyncDelete;
+import static google.registry.flows.ResourceFlowUtils.checkLinkedDomains;
 import static google.registry.flows.ResourceFlowUtils.loadAndVerifyExistence;
 import static google.registry.flows.ResourceFlowUtils.verifyNoDisallowedStatuses;
 import static google.registry.flows.ResourceFlowUtils.verifyOptionalAuthInfo;
 import static google.registry.flows.ResourceFlowUtils.verifyResourceOwnership;
+import static google.registry.model.ResourceTransferUtils.denyPendingTransfer;
+import static google.registry.model.ResourceTransferUtils.handlePendingTransferOnDelete;
+import static google.registry.model.eppoutput.Result.Code.SUCCESS;
 import static google.registry.model.eppoutput.Result.Code.SUCCESS_WITH_ACTION_PENDING;
+import static google.registry.model.transfer.TransferStatus.SERVER_CANCELLED;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.common.collect.ImmutableSet;
-import com.googlecode.objectify.Key;
 import google.registry.batch.AsyncTaskEnqueuer;
 import google.registry.flows.EppException;
 import google.registry.flows.ExtensionManager;
@@ -33,6 +36,7 @@ import google.registry.flows.FlowModule.Superuser;
 import google.registry.flows.FlowModule.TargetId;
 import google.registry.flows.TransactionalFlow;
 import google.registry.flows.annotations.ReportingSpec;
+import google.registry.model.contact.ContactHistory;
 import google.registry.model.contact.ContactResource;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.metadata.MetadataExtension;
@@ -40,7 +44,8 @@ import google.registry.model.eppcommon.AuthInfo;
 import google.registry.model.eppcommon.StatusValue;
 import google.registry.model.eppcommon.Trid;
 import google.registry.model.eppoutput.EppResponse;
-import google.registry.model.reporting.HistoryEntry;
+import google.registry.model.eppoutput.Result.Code;
+import google.registry.model.reporting.HistoryEntry.Type;
 import google.registry.model.reporting.IcannReportingTypes.ActivityReportField;
 import java.util.Optional;
 import javax.inject.Inject;
@@ -63,10 +68,11 @@ import org.joda.time.DateTime;
 @ReportingSpec(ActivityReportField.CONTACT_DELETE)
 public final class ContactDeleteFlow implements TransactionalFlow {
 
-  private static final ImmutableSet<StatusValue> DISALLOWED_STATUSES = ImmutableSet.of(
-      StatusValue.CLIENT_DELETE_PROHIBITED,
-      StatusValue.PENDING_DELETE,
-      StatusValue.SERVER_DELETE_PROHIBITED);
+  private static final ImmutableSet<StatusValue> DISALLOWED_STATUSES =
+      ImmutableSet.of(
+          StatusValue.CLIENT_DELETE_PROHIBITED,
+          StatusValue.PENDING_DELETE,
+          StatusValue.SERVER_DELETE_PROHIBITED);
 
   @Inject ExtensionManager extensionManager;
   @Inject @ClientId String clientId;
@@ -74,10 +80,12 @@ public final class ContactDeleteFlow implements TransactionalFlow {
   @Inject Trid trid;
   @Inject @Superuser boolean isSuperuser;
   @Inject Optional<AuthInfo> authInfo;
-  @Inject HistoryEntry.Builder historyBuilder;
+  @Inject ContactHistory.Builder historyBuilder;
   @Inject AsyncTaskEnqueuer asyncTaskEnqueuer;
   @Inject EppResponse.Builder responseBuilder;
-  @Inject ContactDeleteFlow() {}
+
+  @Inject
+  ContactDeleteFlow() {}
 
   @Override
   public final EppResponse run() throws EppException {
@@ -85,23 +93,45 @@ public final class ContactDeleteFlow implements TransactionalFlow {
     extensionManager.validate();
     validateClientIsLoggedIn(clientId);
     DateTime now = tm().getTransactionTime();
-    failfastForAsyncDelete(targetId, now, ContactResource.class, DomainBase::getReferencedContacts);
+    checkLinkedDomains(targetId, now, ContactResource.class, DomainBase::getReferencedContacts);
     ContactResource existingContact = loadAndVerifyExistence(ContactResource.class, targetId, now);
     verifyNoDisallowedStatuses(existingContact, DISALLOWED_STATUSES);
     verifyOptionalAuthInfo(authInfo, existingContact);
     if (!isSuperuser) {
       verifyResourceOwnership(clientId, existingContact);
     }
-    asyncTaskEnqueuer.enqueueAsyncDelete(
-        existingContact, tm().getTransactionTime(), clientId, trid, isSuperuser);
-    ContactResource newContact =
-        existingContact.asBuilder().addStatusValue(StatusValue.PENDING_DELETE).build();
-    historyBuilder
-        .setType(HistoryEntry.Type.CONTACT_PENDING_DELETE)
-        .setModificationTime(now)
-        .setParent(Key.create(existingContact));
-    tm().insert(historyBuilder.build().toChildHistoryEntity());
+    Type historyEntryType;
+    Code resultCode;
+    ContactResource newContact;
+    if (tm().isOfy()) {
+      asyncTaskEnqueuer.enqueueAsyncDelete(
+          existingContact, tm().getTransactionTime(), clientId, trid, isSuperuser);
+      newContact = existingContact.asBuilder().addStatusValue(StatusValue.PENDING_DELETE).build();
+      historyEntryType = Type.CONTACT_PENDING_DELETE;
+      resultCode = SUCCESS_WITH_ACTION_PENDING;
+    } else {
+      // Handle pending transfers on contact deletion.
+      newContact =
+          existingContact.getStatusValues().contains(StatusValue.PENDING_TRANSFER)
+              ? denyPendingTransfer(existingContact, SERVER_CANCELLED, now, clientId)
+              : existingContact;
+      // Wipe out PII on contact deletion.
+      newContact =
+          newContact.asBuilder().wipeOut().setStatusValues(null).setDeletionTime(now).build();
+      historyEntryType = Type.CONTACT_DELETE;
+      resultCode = SUCCESS;
+    }
+    ContactHistory contactHistory =
+        historyBuilder
+            .setType(historyEntryType)
+            .setModificationTime(now)
+            .setContactBase(newContact)
+            .build();
+    if (!tm().isOfy()) {
+      handlePendingTransferOnDelete(existingContact, newContact, now, contactHistory);
+    }
+    tm().insert(contactHistory);
     tm().update(newContact);
-    return responseBuilder.setResultFromCode(SUCCESS_WITH_ACTION_PENDING).build();
+    return responseBuilder.setResultFromCode(resultCode).build();
   }
 }

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

@@ -20,19 +20,21 @@ import com.google.common.base.CharMatcher;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Sets;
+import com.googlecode.objectify.Key;
 import google.registry.flows.EppException;
 import google.registry.flows.EppException.ParameterValuePolicyErrorException;
 import google.registry.flows.EppException.ParameterValueSyntaxErrorException;
 import google.registry.model.contact.ContactAddress;
+import google.registry.model.contact.ContactHistory;
 import google.registry.model.contact.ContactResource;
 import google.registry.model.contact.PostalInfo;
 import google.registry.model.poll.PendingActionNotificationResponse.ContactPendingActionNotificationResponse;
 import google.registry.model.poll.PollMessage;
-import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.transfer.TransferData;
 import google.registry.model.transfer.TransferResponse.ContactTransferResponse;
 import java.util.Set;
 import javax.annotation.Nullable;
+import org.joda.time.DateTime;
 
 /** Static utility functions for contact flows. */
 public class ContactFlowUtils {
@@ -66,31 +68,35 @@ public class ContactFlowUtils {
 
   /** Create a poll message for the gaining client in a transfer. */
   static PollMessage createGainingTransferPollMessage(
-      String targetId, TransferData transferData, HistoryEntry historyEntry) {
+      String targetId,
+      TransferData transferData,
+      DateTime now,
+      Key<ContactHistory> contactHistoryKey) {
     return new PollMessage.OneTime.Builder()
         .setClientId(transferData.getGainingClientId())
         .setEventTime(transferData.getPendingTransferExpirationTime())
         .setMsg(transferData.getTransferStatus().getMessage())
-        .setResponseData(ImmutableList.of(
-            createTransferResponse(targetId, transferData),
-            ContactPendingActionNotificationResponse.create(
-                  targetId,
-                  transferData.getTransferStatus().isApproved(),
-                  transferData.getTransferRequestTrid(),
-                  historyEntry.getModificationTime())))
-        .setParent(historyEntry)
+        .setResponseData(
+            ImmutableList.of(
+                createTransferResponse(targetId, transferData),
+                ContactPendingActionNotificationResponse.create(
+                    targetId,
+                    transferData.getTransferStatus().isApproved(),
+                    transferData.getTransferRequestTrid(),
+                    now)))
+        .setParentKey(contactHistoryKey)
         .build();
   }
 
   /** Create a poll message for the losing client in a transfer. */
   static PollMessage createLosingTransferPollMessage(
-      String targetId, TransferData transferData, HistoryEntry historyEntry) {
+      String targetId, TransferData transferData, Key<ContactHistory> contactHistoryKey) {
     return new PollMessage.OneTime.Builder()
         .setClientId(transferData.getLosingClientId())
         .setEventTime(transferData.getPendingTransferExpirationTime())
         .setMsg(transferData.getTransferStatus().getMessage())
         .setResponseData(ImmutableList.of(createTransferResponse(targetId, transferData)))
-        .setParent(historyEntry)
+        .setParentKey(contactHistoryKey)
         .build();
   }
 

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

@@ -32,6 +32,7 @@ import google.registry.flows.FlowModule.ClientId;
 import google.registry.flows.FlowModule.TargetId;
 import google.registry.flows.TransactionalFlow;
 import google.registry.flows.annotations.ReportingSpec;
+import google.registry.model.contact.ContactHistory;
 import google.registry.model.contact.ContactResource;
 import google.registry.model.domain.metadata.MetadataExtension;
 import google.registry.model.eppcommon.AuthInfo;
@@ -66,7 +67,7 @@ public final class ContactTransferApproveFlow implements TransactionalFlow {
   @Inject @ClientId String clientId;
   @Inject @TargetId String targetId;
   @Inject Optional<AuthInfo> authInfo;
-  @Inject HistoryEntry.Builder historyBuilder;
+  @Inject ContactHistory.Builder historyBuilder;
   @Inject EppResponse.Builder responseBuilder;
   @Inject ContactTransferApproveFlow() {}
 
@@ -86,15 +87,17 @@ public final class ContactTransferApproveFlow implements TransactionalFlow {
     verifyResourceOwnership(clientId, existingContact);
     ContactResource newContact =
         approvePendingTransfer(existingContact, TransferStatus.CLIENT_APPROVED, now);
-    HistoryEntry historyEntry = historyBuilder
-        .setType(HistoryEntry.Type.CONTACT_TRANSFER_APPROVE)
-        .setModificationTime(now)
-        .setParent(Key.create(existingContact))
-        .build();
+    ContactHistory contactHistory =
+        historyBuilder
+            .setType(HistoryEntry.Type.CONTACT_TRANSFER_APPROVE)
+            .setModificationTime(now)
+            .setContactBase(newContact)
+            .build();
     // Create a poll message for the gaining client.
     PollMessage gainingPollMessage =
-        createGainingTransferPollMessage(targetId, newContact.getTransferData(), historyEntry);
-    tm().insertAll(ImmutableSet.of(historyEntry.toChildHistoryEntity(), gainingPollMessage));
+        createGainingTransferPollMessage(
+            targetId, newContact.getTransferData(), now, Key.create(contactHistory));
+    tm().insertAll(ImmutableSet.of(contactHistory, gainingPollMessage));
     tm().update(newContact);
     // Delete the billing event and poll messages that were written in case the transfer would have
     // been implicitly server approved.

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

@@ -32,6 +32,7 @@ import google.registry.flows.FlowModule.ClientId;
 import google.registry.flows.FlowModule.TargetId;
 import google.registry.flows.TransactionalFlow;
 import google.registry.flows.annotations.ReportingSpec;
+import google.registry.model.contact.ContactHistory;
 import google.registry.model.contact.ContactResource;
 import google.registry.model.domain.metadata.MetadataExtension;
 import google.registry.model.eppcommon.AuthInfo;
@@ -66,7 +67,7 @@ public final class ContactTransferCancelFlow implements TransactionalFlow {
   @Inject Optional<AuthInfo> authInfo;
   @Inject @ClientId String clientId;
   @Inject @TargetId String targetId;
-  @Inject HistoryEntry.Builder historyBuilder;
+  @Inject ContactHistory.Builder historyBuilder;
   @Inject EppResponse.Builder responseBuilder;
   @Inject ContactTransferCancelFlow() {}
 
@@ -82,15 +83,17 @@ public final class ContactTransferCancelFlow implements TransactionalFlow {
     verifyTransferInitiator(clientId, existingContact);
     ContactResource newContact =
         denyPendingTransfer(existingContact, TransferStatus.CLIENT_CANCELLED, now, clientId);
-    HistoryEntry historyEntry = historyBuilder
-        .setType(HistoryEntry.Type.CONTACT_TRANSFER_CANCEL)
-        .setModificationTime(now)
-        .setParent(Key.create(existingContact))
-        .build();
+    ContactHistory contactHistory =
+        historyBuilder
+            .setType(HistoryEntry.Type.CONTACT_TRANSFER_CANCEL)
+            .setModificationTime(now)
+            .setContactBase(newContact)
+            .build();
     // Create a poll message for the losing client.
     PollMessage losingPollMessage =
-        createLosingTransferPollMessage(targetId, newContact.getTransferData(), historyEntry);
-    tm().insertAll(ImmutableSet.of(historyEntry.toChildHistoryEntity(), losingPollMessage));
+        createLosingTransferPollMessage(
+            targetId, newContact.getTransferData(), Key.create(contactHistory));
+    tm().insertAll(ImmutableSet.of(contactHistory, losingPollMessage));
     tm().update(newContact);
     // Delete the billing event and poll messages that were written in case the transfer would have
     // been implicitly server approved.

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

@@ -32,6 +32,7 @@ import google.registry.flows.FlowModule.ClientId;
 import google.registry.flows.FlowModule.TargetId;
 import google.registry.flows.TransactionalFlow;
 import google.registry.flows.annotations.ReportingSpec;
+import google.registry.model.contact.ContactHistory;
 import google.registry.model.contact.ContactResource;
 import google.registry.model.domain.metadata.MetadataExtension;
 import google.registry.model.eppcommon.AuthInfo;
@@ -64,7 +65,7 @@ public final class ContactTransferRejectFlow implements TransactionalFlow {
   @Inject Optional<AuthInfo> authInfo;
   @Inject @ClientId String clientId;
   @Inject @TargetId String targetId;
-  @Inject HistoryEntry.Builder historyBuilder;
+  @Inject ContactHistory.Builder historyBuilder;
   @Inject EppResponse.Builder responseBuilder;
   @Inject ContactTransferRejectFlow() {}
 
@@ -80,14 +81,16 @@ public final class ContactTransferRejectFlow implements TransactionalFlow {
     verifyResourceOwnership(clientId, existingContact);
     ContactResource newContact =
         denyPendingTransfer(existingContact, TransferStatus.CLIENT_REJECTED, now, clientId);
-    HistoryEntry historyEntry = historyBuilder
-        .setType(HistoryEntry.Type.CONTACT_TRANSFER_REJECT)
-        .setModificationTime(now)
-        .setParent(Key.create(existingContact))
-        .build();
+    ContactHistory contactHistory =
+        historyBuilder
+            .setType(HistoryEntry.Type.CONTACT_TRANSFER_REJECT)
+            .setModificationTime(now)
+            .setContactBase(newContact)
+            .build();
     PollMessage gainingPollMessage =
-        createGainingTransferPollMessage(targetId, newContact.getTransferData(), historyEntry);
-    tm().insertAll(ImmutableSet.of(historyEntry.toChildHistoryEntity(), gainingPollMessage));
+        createGainingTransferPollMessage(
+            targetId, newContact.getTransferData(), now, Key.create(contactHistory));
+    tm().insertAll(ImmutableSet.of(contactHistory, gainingPollMessage));
     tm().update(newContact);
     // Delete the billing event and poll messages that were written in case the transfer would have
     // been implicitly server approved.

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

@@ -14,6 +14,7 @@
 
 package google.registry.flows.contact;
 
+import static google.registry.flows.FlowUtils.createHistoryKey;
 import static google.registry.flows.FlowUtils.validateClientIsLoggedIn;
 import static google.registry.flows.ResourceFlowUtils.loadAndVerifyExistence;
 import static google.registry.flows.ResourceFlowUtils.verifyAuthInfo;
@@ -36,6 +37,7 @@ import google.registry.flows.TransactionalFlow;
 import google.registry.flows.annotations.ReportingSpec;
 import google.registry.flows.exceptions.AlreadyPendingTransferException;
 import google.registry.flows.exceptions.ObjectAlreadySponsoredException;
+import google.registry.model.contact.ContactHistory;
 import google.registry.model.contact.ContactResource;
 import google.registry.model.domain.metadata.MetadataExtension;
 import google.registry.model.eppcommon.AuthInfo;
@@ -81,7 +83,8 @@ public final class ContactTransferRequestFlow implements TransactionalFlow {
   @Inject @ClientId String gainingClientId;
   @Inject @TargetId String targetId;
   @Inject @Config("contactAutomaticTransferLength") Duration automaticTransferLength;
-  @Inject HistoryEntry.Builder historyBuilder;
+
+  @Inject ContactHistory.Builder historyBuilder;
   @Inject Trid trid;
   @Inject EppResponse.Builder responseBuilder;
   @Inject ContactTransferRequestFlow() {}
@@ -105,11 +108,7 @@ public final class ContactTransferRequestFlow implements TransactionalFlow {
       throw new ObjectAlreadySponsoredException();
     }
     verifyNoDisallowedStatuses(existingContact, DISALLOWED_STATUSES);
-    HistoryEntry historyEntry = historyBuilder
-        .setType(HistoryEntry.Type.CONTACT_TRANSFER_REQUEST)
-        .setModificationTime(now)
-        .setParent(Key.create(existingContact))
-        .build();
+
     DateTime transferExpirationTime = now.plus(automaticTransferLength);
     ContactTransferData serverApproveTransferData =
         new ContactTransferData.Builder()
@@ -120,12 +119,18 @@ public final class ContactTransferRequestFlow implements TransactionalFlow {
             .setPendingTransferExpirationTime(transferExpirationTime)
             .setTransferStatus(TransferStatus.SERVER_APPROVED)
             .build();
+    Key<ContactHistory> contactHistoryKey = createHistoryKey(existingContact, ContactHistory.class);
+    historyBuilder
+        .setId(contactHistoryKey.getId())
+        .setType(HistoryEntry.Type.CONTACT_TRANSFER_REQUEST)
+        .setModificationTime(now);
     // If the transfer is server approved, this message will be sent to the losing registrar. */
     PollMessage serverApproveLosingPollMessage =
-        createLosingTransferPollMessage(targetId, serverApproveTransferData, historyEntry);
+        createLosingTransferPollMessage(targetId, serverApproveTransferData, contactHistoryKey);
     // If the transfer is server approved, this message will be sent to the gaining registrar. */
     PollMessage serverApproveGainingPollMessage =
-        createGainingTransferPollMessage(targetId, serverApproveTransferData, historyEntry);
+        createGainingTransferPollMessage(
+            targetId, serverApproveTransferData, now, contactHistoryKey);
     ContactTransferData pendingTransferData =
         serverApproveTransferData
             .asBuilder()
@@ -137,8 +142,9 @@ public final class ContactTransferRequestFlow implements TransactionalFlow {
             .build();
     // When a transfer is requested, a poll message is created to notify the losing registrar.
     PollMessage requestPollMessage =
-        createLosingTransferPollMessage(targetId, pendingTransferData, historyEntry).asBuilder()
-            .setEventTime(now)  // Unlike the serverApprove messages, this applies immediately.
+        createLosingTransferPollMessage(targetId, pendingTransferData, contactHistoryKey)
+            .asBuilder()
+            .setEventTime(now) // Unlike the serverApprove messages, this applies immediately.
             .build();
     ContactResource newContact = existingContact.asBuilder()
         .setTransferData(pendingTransferData)
@@ -147,7 +153,7 @@ public final class ContactTransferRequestFlow implements TransactionalFlow {
     tm().update(newContact);
     tm().insertAll(
             ImmutableSet.of(
-                historyEntry.toChildHistoryEntity(),
+                historyBuilder.setContactBase(newContact).build(),
                 requestPollMessage,
                 serverApproveGainingPollMessage,
                 serverApproveLosingPollMessage));

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

@@ -27,7 +27,6 @@ import static google.registry.flows.contact.ContactFlowUtils.validateContactAgai
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.common.collect.ImmutableSet;
-import com.googlecode.objectify.Key;
 import google.registry.flows.EppException;
 import google.registry.flows.ExtensionManager;
 import google.registry.flows.FlowModule.ClientId;
@@ -38,6 +37,7 @@ import google.registry.flows.annotations.ReportingSpec;
 import google.registry.flows.exceptions.ResourceHasClientUpdateProhibitedException;
 import google.registry.model.contact.ContactCommand.Update;
 import google.registry.model.contact.ContactCommand.Update.Change;
+import google.registry.model.contact.ContactHistory;
 import google.registry.model.contact.ContactResource;
 import google.registry.model.contact.PostalInfo;
 import google.registry.model.domain.metadata.MetadataExtension;
@@ -82,7 +82,7 @@ public final class ContactUpdateFlow implements TransactionalFlow {
   @Inject @ClientId String clientId;
   @Inject @TargetId String targetId;
   @Inject @Superuser boolean isSuperuser;
-  @Inject HistoryEntry.Builder historyBuilder;
+  @Inject ContactHistory.Builder historyBuilder;
   @Inject EppResponse.Builder responseBuilder;
   @Inject ContactUpdateFlow() {}
 
@@ -102,11 +102,6 @@ public final class ContactUpdateFlow implements TransactionalFlow {
       verifyAllStatusesAreClientSettable(union(statusesToAdd, statusToRemove));
     }
     verifyNoDisallowedStatuses(existingContact, DISALLOWED_STATUSES);
-    historyBuilder
-        .setType(HistoryEntry.Type.CONTACT_UPDATE)
-        .setModificationTime(now)
-        .setXmlBytes(null)  // We don't want to store contact details in the history entry.
-        .setParent(Key.create(existingContact));
     checkSameValuesNotAddedAndRemoved(statusesToAdd, statusToRemove);
     ContactResource.Builder builder = existingContact.asBuilder();
     Change change = command.getInnerChange();
@@ -150,7 +145,12 @@ public final class ContactUpdateFlow implements TransactionalFlow {
     }
     validateAsciiPostalInfo(newContact.getInternationalizedPostalInfo());
     validateContactAgainstPolicy(newContact);
-    tm().insert(historyBuilder.build().toChildHistoryEntity());
+    historyBuilder
+        .setType(HistoryEntry.Type.CONTACT_UPDATE)
+        .setModificationTime(now)
+        .setXmlBytes(null) // We don't want to store contact details in the history entry.
+        .setContactBase(newContact);
+    tm().insert(historyBuilder.build());
     tm().update(newContact);
     return responseBuilder.build();
   }

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

@@ -80,6 +80,7 @@ import google.registry.model.billing.BillingEvent.Recurring;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.DomainCommand;
 import google.registry.model.domain.DomainCommand.Create;
+import google.registry.model.domain.DomainHistory;
 import google.registry.model.domain.GracePeriod;
 import google.registry.model.domain.Period;
 import google.registry.model.domain.fee.FeeCreateCommandExtension;
@@ -96,7 +97,6 @@ import google.registry.model.eppinput.EppInput;
 import google.registry.model.eppinput.ResourceCommand;
 import google.registry.model.eppoutput.CreateData.DomainCreateData;
 import google.registry.model.eppoutput.EppResponse;
-import google.registry.model.host.HostResource;
 import google.registry.model.index.EppResourceIndex;
 import google.registry.model.index.ForeignKeyIndex;
 import google.registry.model.ofy.ObjectifyService;
@@ -106,11 +106,11 @@ import google.registry.model.poll.PollMessage.Autorenew;
 import google.registry.model.registry.Registry;
 import google.registry.model.registry.Registry.TldState;
 import google.registry.model.registry.Registry.TldType;
+import google.registry.model.registry.label.ReservationType;
 import google.registry.model.reporting.DomainTransactionRecord;
 import google.registry.model.reporting.DomainTransactionRecord.TransactionReportField;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.reporting.IcannReportingTypes.ActivityReportField;
-import google.registry.persistence.VKey;
 import google.registry.tmch.LordnTaskUtils;
 import java.util.Optional;
 import javax.inject.Inject;
@@ -206,7 +206,7 @@ public class DomainCreateFlow implements TransactionalFlow {
   @Inject @ClientId String clientId;
   @Inject @TargetId String targetId;
   @Inject @Superuser boolean isSuperuser;
-  @Inject HistoryEntry.Builder historyBuilder;
+  @Inject DomainHistory.Builder historyBuilder;
   @Inject EppResponse.Builder responseBuilder;
   @Inject AllocationTokenFlowUtils allocationTokenFlowUtils;
   @Inject DomainCreateFlowCustomLogic flowCustomLogic;
@@ -301,10 +301,14 @@ public class DomainCreateFlow implements TransactionalFlow {
     validateFeeChallenge(targetId, now, feeCreate, feesAndCredits);
     Optional<SecDnsCreateExtension> secDnsCreate =
         validateSecDnsExtension(eppInput.getSingleExtension(SecDnsCreateExtension.class));
-    String repoId = createDomainRepoId(ObjectifyService.allocateId(), registry.getTldStr());
     DateTime registrationExpirationTime = leapSafeAddYears(now, years);
-    HistoryEntry historyEntry = buildHistoryEntry(
-        repoId, registry, now, period, registry.getAddGracePeriodLength());
+    String repoId = createDomainRepoId(ObjectifyService.allocateId(), registry.getTldStr());
+    Key<DomainHistory> domainHistoryKey =
+        Key.create(
+            Key.create(DomainBase.class, repoId),
+            DomainHistory.class,
+            ObjectifyService.allocateId());
+    historyBuilder.setId(domainHistoryKey.getId());
     // Bill for the create.
     BillingEvent.OneTime createBillingEvent =
         createOneTimeBillingEvent(
@@ -314,33 +318,27 @@ public class DomainCreateFlow implements TransactionalFlow {
             isReserved(domainName, isSunriseCreate),
             years,
             feesAndCredits,
-            historyEntry,
+            domainHistoryKey,
             allocationToken,
             now);
     // Create a new autorenew billing event and poll message starting at the expiration time.
     BillingEvent.Recurring autorenewBillingEvent =
-        createAutorenewBillingEvent(historyEntry, registrationExpirationTime);
+        createAutorenewBillingEvent(domainHistoryKey, registrationExpirationTime);
     PollMessage.Autorenew autorenewPollMessage =
-        createAutorenewPollMessage(historyEntry, registrationExpirationTime);
+        createAutorenewPollMessage(domainHistoryKey, registrationExpirationTime);
     ImmutableSet.Builder<ImmutableObject> entitiesToSave = new ImmutableSet.Builder<>();
-    entitiesToSave.add(
-        historyEntry,
-        createBillingEvent,
-        autorenewBillingEvent,
-        autorenewPollMessage);
+    entitiesToSave.add(createBillingEvent, autorenewBillingEvent, autorenewPollMessage);
     // Bill for EAP cost, if any.
     if (!feesAndCredits.getEapCost().isZero()) {
       entitiesToSave.add(createEapBillingEvent(feesAndCredits, createBillingEvent));
     }
 
-    ImmutableSet.Builder<StatusValue> statuses = new ImmutableSet.Builder<>();
-    if (getReservationTypes(domainName).contains(NAME_COLLISION)) {
-      statuses.add(SERVER_HOLD);
-      entitiesToSave.add(
-          createNameCollisionOneTimePollMessage(targetId, historyEntry, clientId, now));
-    }
-
-    DomainBase newDomain =
+    ImmutableSet<ReservationType> reservationTypes = getReservationTypes(domainName);
+    ImmutableSet<StatusValue> statuses =
+        reservationTypes.contains(NAME_COLLISION)
+            ? ImmutableSet.of(SERVER_HOLD)
+            : ImmutableSet.of();
+    DomainBase domain =
         new DomainBase.Builder()
             .setCreationClientId(clientId)
             .setPersistedCurrentSponsorClientId(clientId)
@@ -351,35 +349,39 @@ public class DomainCreateFlow implements TransactionalFlow {
             .setAutorenewPollMessage(autorenewPollMessage.createVKey())
             .setLaunchNotice(hasClaimsNotice ? launchCreate.get().getNotice() : null)
             .setSmdId(signedMarkId)
-            .setDsData(secDnsCreate.isPresent() ? secDnsCreate.get().getDsData() : null)
+            .setDsData(secDnsCreate.map(SecDnsCreateExtension::getDsData).orElse(null))
             .setRegistrant(command.getRegistrant())
             .setAuthInfo(command.getAuthInfo())
             .setDomainName(targetId)
-            .setNameservers(
-                (ImmutableSet<VKey<HostResource>>)
-                    command.getNameservers().stream().collect(toImmutableSet()))
-            .setStatusValues(statuses.build())
+            .setNameservers(command.getNameservers().stream().collect(toImmutableSet()))
+            .setStatusValues(statuses)
             .setContacts(command.getContacts())
             .addGracePeriod(
                 GracePeriod.forBillingEvent(GracePeriodStatus.ADD, repoId, createBillingEvent))
             .build();
+    DomainHistory domainHistory =
+        buildDomainHistory(domain, registry, now, period, registry.getAddGracePeriodLength());
+    if (reservationTypes.contains(NAME_COLLISION)) {
+      entitiesToSave.add(
+          createNameCollisionOneTimePollMessage(targetId, domainHistory, clientId, now));
+    }
     entitiesToSave.add(
-        newDomain,
-        ForeignKeyIndex.create(newDomain, newDomain.getDeletionTime()),
-        EppResourceIndex.create(Key.create(newDomain)));
+        domain,
+        domainHistory,
+        ForeignKeyIndex.create(domain, domain.getDeletionTime()),
+        EppResourceIndex.create(Key.create(domain)));
     if (allocationToken.isPresent()
         && TokenType.SINGLE_USE.equals(allocationToken.get().getTokenType())) {
       entitiesToSave.add(
-          allocationTokenFlowUtils.redeemToken(
-              allocationToken.get(), HistoryEntry.createVKey(Key.create(historyEntry))));
+          allocationTokenFlowUtils.redeemToken(allocationToken.get(), domainHistory.createVKey()));
     }
-    enqueueTasks(newDomain, hasSignedMarks, hasClaimsNotice);
+    enqueueTasks(domain, hasSignedMarks, hasClaimsNotice);
 
     EntityChanges entityChanges =
         flowCustomLogic.beforeSave(
             DomainCreateFlowCustomLogic.BeforeSaveParameters.newBuilder()
-                .setNewDomain(newDomain)
-                .setHistoryEntry(historyEntry)
+                .setNewDomain(domain)
+                .setHistoryEntry(domainHistory)
                 .setEntityChanges(
                     EntityChanges.newBuilder().setSaves(entitiesToSave.build()).build())
                 .setYears(years)
@@ -483,8 +485,8 @@ public class DomainCreateFlow implements TransactionalFlow {
             : null);
   }
 
-  private HistoryEntry buildHistoryEntry(
-      String repoId, Registry registry, DateTime now, Period period, Duration addGracePeriod) {
+  private DomainHistory buildDomainHistory(
+      DomainBase domain, Registry registry, DateTime now, Period period, Duration addGracePeriod) {
     // We ignore prober transactions
     if (registry.getTldType() == TldType.REAL) {
       historyBuilder
@@ -500,7 +502,7 @@ public class DomainCreateFlow implements TransactionalFlow {
         .setType(HistoryEntry.Type.DOMAIN_CREATE)
         .setPeriod(period)
         .setModificationTime(now)
-        .setParent(Key.create(DomainBase.class, repoId))
+        .setDomainContent(domain)
         .build();
   }
 
@@ -511,7 +513,7 @@ public class DomainCreateFlow implements TransactionalFlow {
       boolean isReserved,
       int years,
       FeesAndCredits feesAndCredits,
-      HistoryEntry historyEntry,
+      Key<DomainHistory> domainHistoryKey,
       Optional<AllocationToken> allocationToken,
       DateTime now) {
     ImmutableSet.Builder<Flag> flagsBuilder = new ImmutableSet.Builder<>();
@@ -540,12 +542,12 @@ public class DomainCreateFlow implements TransactionalFlow {
                     ? registry.getAnchorTenantAddGracePeriodLength()
                     : registry.getAddGracePeriodLength()))
         .setFlags(flagsBuilder.build())
-        .setParent(historyEntry)
+        .setParent(domainHistoryKey)
         .build();
   }
 
   private Recurring createAutorenewBillingEvent(
-      HistoryEntry historyEntry, DateTime registrationExpirationTime) {
+      Key<DomainHistory> domainHistoryKey, DateTime registrationExpirationTime) {
     return new BillingEvent.Recurring.Builder()
         .setReason(Reason.RENEW)
         .setFlags(ImmutableSet.of(Flag.AUTO_RENEW))
@@ -553,18 +555,18 @@ public class DomainCreateFlow implements TransactionalFlow {
         .setClientId(clientId)
         .setEventTime(registrationExpirationTime)
         .setRecurrenceEndTime(END_OF_TIME)
-        .setParent(historyEntry)
+        .setParent(domainHistoryKey)
         .build();
   }
 
   private Autorenew createAutorenewPollMessage(
-      HistoryEntry historyEntry, DateTime registrationExpirationTime) {
+      Key<DomainHistory> domainHistoryKey, DateTime registrationExpirationTime) {
     return new PollMessage.Autorenew.Builder()
         .setTargetId(targetId)
         .setClientId(clientId)
         .setEventTime(registrationExpirationTime)
         .setMsg("Domain was auto-renewed.")
-        .setParent(historyEntry)
+        .setParentKey(domainHistoryKey)
         .build();
   }
 

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

@@ -16,6 +16,7 @@ package google.registry.flows.domain;
 
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Strings.isNullOrEmpty;
+import static google.registry.flows.FlowUtils.createHistoryKey;
 import static google.registry.flows.FlowUtils.persistEntityChanges;
 import static google.registry.flows.FlowUtils.validateClientIsLoggedIn;
 import static google.registry.flows.ResourceFlowUtils.loadAndVerifyExistence;
@@ -63,6 +64,7 @@ import google.registry.flows.custom.EntityChanges;
 import google.registry.model.ImmutableObject;
 import google.registry.model.billing.BillingEvent;
 import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.DomainHistory;
 import google.registry.model.domain.GracePeriod;
 import google.registry.model.domain.fee.BaseFee.FeeType;
 import google.registry.model.domain.fee.Credit;
@@ -125,7 +127,7 @@ public final class DomainDeleteFlow implements TransactionalFlow {
   @Inject @ClientId String clientId;
   @Inject @TargetId String targetId;
   @Inject @Superuser boolean isSuperuser;
-  @Inject HistoryEntry.Builder historyBuilder;
+  @Inject DomainHistory.Builder historyBuilder;
   @Inject DnsQueue dnsQueue;
   @Inject Trid trid;
   @Inject AsyncTaskEnqueuer asyncTaskEnqueuer;
@@ -177,8 +179,8 @@ public final class DomainDeleteFlow implements TransactionalFlow {
             ? Duration.ZERO
             // By default, this should be 30 days of grace, and 5 days of pending delete.
             : redemptionGracePeriodLength.plus(pendingDeleteLength);
-    HistoryEntry historyEntry =
-        buildHistoryEntry(existingDomain, registry, now, durationUntilDelete, inAddGracePeriod);
+    Key<DomainHistory> domainHistoryKey = createHistoryKey(existingDomain, DomainHistory.class);
+    historyBuilder.setId(domainHistoryKey.getId());
     DateTime deletionTime = now.plus(durationUntilDelete);
     if (durationUntilDelete.equals(Duration.ZERO)) {
       builder.setDeletionTime(now).setStatusValues(null);
@@ -208,20 +210,28 @@ public final class DomainDeleteFlow implements TransactionalFlow {
     // Enqueue the deletion poll message if the delete is asynchronous or if requested by a
     // superuser (i.e. the registrar didn't request this delete and thus should be notified even if
     // it is synchronous).
-    if (!durationUntilDelete.equals(Duration.ZERO) || isSuperuser) {
+    if (durationUntilDelete.isLongerThan(Duration.ZERO) || isSuperuser) {
       PollMessage.OneTime deletePollMessage =
-          createDeletePollMessage(existingDomain, historyEntry, deletionTime);
+          createDeletePollMessage(existingDomain, domainHistoryKey, deletionTime);
       entitiesToSave.add(deletePollMessage);
       builder.setDeletePollMessage(deletePollMessage.createVKey());
     }
 
+    // Send a second poll message immediately if the domain is being deleted asynchronously by a
+    // registrar other than the sponsoring registrar (which will necessarily be a superuser).
+    if (durationUntilDelete.isLongerThan(Duration.ZERO)
+        && !clientId.equals(existingDomain.getPersistedCurrentSponsorClientId())) {
+      entitiesToSave.add(
+          createImmediateDeletePollMessage(existingDomain, domainHistoryKey, now, deletionTime));
+    }
+
     // Cancel any grace periods that were still active, and set the expiration time accordingly.
     DateTime newExpirationTime = existingDomain.getRegistrationExpirationTime();
     for (GracePeriod gracePeriod : existingDomain.getGracePeriods()) {
       // No cancellation is written if the grace period was not for a billable event.
       if (gracePeriod.hasBillingEvent()) {
         entitiesToSave.add(
-            BillingEvent.Cancellation.forGracePeriod(gracePeriod, historyEntry, targetId));
+            BillingEvent.Cancellation.forGracePeriod(gracePeriod, now, domainHistoryKey, targetId));
         if (gracePeriod.getOneTimeBillingEvent() != null) {
           // Take the amount of amount of registration time being refunded off the expiration time.
           // This can be either add grace periods or renew grace periods.
@@ -237,8 +247,10 @@ public final class DomainDeleteFlow implements TransactionalFlow {
     builder.setRegistrationExpirationTime(newExpirationTime);
 
     DomainBase newDomain = builder.build();
+    DomainHistory domainHistory =
+        buildDomainHistory(newDomain, registry, now, durationUntilDelete, inAddGracePeriod);
     updateForeignKeyIndexDeletionTime(newDomain);
-    handlePendingTransferOnDelete(existingDomain, newDomain, now, historyEntry);
+    handlePendingTransferOnDelete(existingDomain, newDomain, now, domainHistory);
     // Close the autorenew billing event and poll message. This may delete the poll message.
     updateAutorenewRecurrenceEndTime(existingDomain, now);
     // If there's a pending transfer, the gaining client's autorenew billing
@@ -246,15 +258,16 @@ public final class DomainDeleteFlow implements TransactionalFlow {
     // ResourceDeleteFlow since it's listed in serverApproveEntities.
     dnsQueue.addDomainRefreshTask(existingDomain.getDomainName());
 
-    entitiesToSave.add(newDomain, historyEntry);
-    EntityChanges entityChanges = flowCustomLogic.beforeSave(
-        BeforeSaveParameters.newBuilder()
-            .setExistingDomain(existingDomain)
-            .setNewDomain(newDomain)
-            .setHistoryEntry(historyEntry)
-            .setEntityChanges(EntityChanges.newBuilder().setSaves(entitiesToSave.build()).build())
-            .build());
-    persistEntityChanges(entityChanges);
+    entitiesToSave.add(newDomain, domainHistory);
+    EntityChanges entityChanges =
+        flowCustomLogic.beforeSave(
+            BeforeSaveParameters.newBuilder()
+                .setExistingDomain(existingDomain)
+                .setNewDomain(newDomain)
+                .setHistoryEntry(domainHistory)
+                .setEntityChanges(
+                    EntityChanges.newBuilder().setSaves(entitiesToSave.build()).build())
+                .build());
     BeforeResponseReturnData responseData =
         flowCustomLogic.beforeResponse(
             BeforeResponseParameters.newBuilder()
@@ -264,6 +277,7 @@ public final class DomainDeleteFlow implements TransactionalFlow {
                         : SUCCESS)
                 .setResponseExtensions(getResponseExtensions(existingDomain, now))
                 .build());
+    persistEntityChanges(entityChanges);
     return responseBuilder
         .setResultFromCode(responseData.resultCode())
         .setExtensions(responseData.responseExtensions())
@@ -284,8 +298,8 @@ public final class DomainDeleteFlow implements TransactionalFlow {
     }
   }
 
-  private HistoryEntry buildHistoryEntry(
-      DomainBase existingResource,
+  private DomainHistory buildDomainHistory(
+      DomainBase domain,
       Registry registry,
       DateTime now,
       Duration durationUntilDelete,
@@ -299,31 +313,30 @@ public final class DomainDeleteFlow implements TransactionalFlow {
               registry.getRenewGracePeriodLength()));
       ImmutableSet<DomainTransactionRecord> cancelledRecords =
           createCancelingRecords(
-              existingResource,
+              domain,
               now,
               maxGracePeriod,
               Sets.immutableEnumSet(Sets.union(ADD_FIELDS, RENEW_FIELDS)));
-      historyBuilder
-          .setDomainTransactionRecords(
-              union(
-                  cancelledRecords,
-                  DomainTransactionRecord.create(
-                      existingResource.getTld(),
-                      now.plus(durationUntilDelete),
-                      inAddGracePeriod
-                          ? TransactionReportField.DELETED_DOMAINS_GRACE
-                          : TransactionReportField.DELETED_DOMAINS_NOGRACE,
-                      1)));
+      historyBuilder.setDomainTransactionRecords(
+          union(
+              cancelledRecords,
+              DomainTransactionRecord.create(
+                  domain.getTld(),
+                  now.plus(durationUntilDelete),
+                  inAddGracePeriod
+                      ? TransactionReportField.DELETED_DOMAINS_GRACE
+                      : TransactionReportField.DELETED_DOMAINS_NOGRACE,
+                  1)));
     }
     return historyBuilder
         .setType(HistoryEntry.Type.DOMAIN_DELETE)
         .setModificationTime(now)
-        .setParent(Key.create(existingResource))
+        .setDomainContent(domain)
         .build();
   }
 
   private PollMessage.OneTime createDeletePollMessage(
-      DomainBase existingDomain, HistoryEntry historyEntry, DateTime deletionTime) {
+      DomainBase existingDomain, Key<DomainHistory> domainHistoryKey, DateTime deletionTime) {
     Optional<MetadataExtension> metadataExtension =
         eppInput.getSingleExtension(MetadataExtension.class);
     boolean hasMetadataMessage =
@@ -342,7 +355,23 @@ public final class DomainDeleteFlow implements TransactionalFlow {
             ImmutableList.of(
                 DomainPendingActionNotificationResponse.create(
                     existingDomain.getDomainName(), true, trid, deletionTime)))
-        .setParent(historyEntry)
+        .setParentKey(domainHistoryKey)
+        .build();
+  }
+
+  private PollMessage.OneTime createImmediateDeletePollMessage(
+      DomainBase existingDomain,
+      Key<DomainHistory> domainHistoryKey,
+      DateTime now,
+      DateTime deletionTime) {
+    return new PollMessage.OneTime.Builder()
+        .setClientId(existingDomain.getPersistedCurrentSponsorClientId())
+        .setEventTime(now)
+        .setParentKey(domainHistoryKey)
+        .setMsg(
+            String.format(
+                "Domain %s was deleted by registry administrator with final deletion effective: %s",
+                existingDomain.getDomainName(), deletionTime))
         .build();
   }
 

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

@@ -14,6 +14,7 @@
 
 package google.registry.flows.domain;
 
+import static google.registry.flows.FlowUtils.createHistoryKey;
 import static google.registry.flows.FlowUtils.persistEntityChanges;
 import static google.registry.flows.FlowUtils.validateClientIsLoggedIn;
 import static google.registry.flows.ResourceFlowUtils.loadAndVerifyExistence;
@@ -33,6 +34,7 @@ import static google.registry.util.DateTimeUtils.leapSafeAddYears;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
+import com.googlecode.objectify.Key;
 import google.registry.flows.EppException;
 import google.registry.flows.EppException.ParameterValueRangeErrorException;
 import google.registry.flows.ExtensionManager;
@@ -52,6 +54,7 @@ import google.registry.model.billing.BillingEvent.OneTime;
 import google.registry.model.billing.BillingEvent.Reason;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.DomainCommand.Renew;
+import google.registry.model.domain.DomainHistory;
 import google.registry.model.domain.DomainRenewData;
 import google.registry.model.domain.GracePeriod;
 import google.registry.model.domain.Period;
@@ -123,7 +126,7 @@ public final class DomainRenewFlow implements TransactionalFlow {
   @Inject @ClientId String clientId;
   @Inject @TargetId String targetId;
   @Inject @Superuser boolean isSuperuser;
-  @Inject HistoryEntry.Builder historyBuilder;
+  @Inject DomainHistory.Builder historyBuilder;
   @Inject EppResponse.Builder responseBuilder;
   @Inject DomainRenewFlowCustomLogic flowCustomLogic;
   @Inject DomainPricingLogic pricingLogic;
@@ -156,22 +159,23 @@ public final class DomainRenewFlow implements TransactionalFlow {
             .setNow(now)
             .setYears(years)
             .build());
-    Registry registry = Registry.get(existingDomain.getTld());
-    HistoryEntry historyEntry = buildHistoryEntry(
-        existingDomain, now, command.getPeriod(), registry.getRenewGracePeriodLength());
+    Key<DomainHistory> domainHistoryKey = createHistoryKey(existingDomain, DomainHistory.class);
+    historyBuilder.setId(domainHistoryKey.getId());
     String tld = existingDomain.getTld();
     // Bill for this explicit renew itself.
     BillingEvent.OneTime explicitRenewEvent =
-        createRenewBillingEvent(tld, feesAndCredits.getTotalCost(), years, historyEntry, now);
+        createRenewBillingEvent(tld, feesAndCredits.getTotalCost(), years, domainHistoryKey, now);
     // Create a new autorenew billing event and poll message starting at the new expiration time.
-    BillingEvent.Recurring newAutorenewEvent = newAutorenewBillingEvent(existingDomain)
-        .setEventTime(newExpirationTime)
-        .setParent(historyEntry)
-        .build();
-    PollMessage.Autorenew newAutorenewPollMessage = newAutorenewPollMessage(existingDomain)
-        .setEventTime(newExpirationTime)
-        .setParent(historyEntry)
-        .build();
+    BillingEvent.Recurring newAutorenewEvent =
+        newAutorenewBillingEvent(existingDomain)
+            .setEventTime(newExpirationTime)
+            .setParent(domainHistoryKey)
+            .build();
+    PollMessage.Autorenew newAutorenewPollMessage =
+        newAutorenewPollMessage(existingDomain)
+            .setEventTime(newExpirationTime)
+            .setParentKey(domainHistoryKey)
+            .build();
     // End the old autorenew billing event and poll message now. This may delete the poll message.
     updateAutorenewRecurrenceEndTime(existingDomain, now);
     DomainBase newDomain =
@@ -186,6 +190,10 @@ public final class DomainRenewFlow implements TransactionalFlow {
                 GracePeriod.forBillingEvent(
                     GracePeriodStatus.RENEW, existingDomain.getRepoId(), explicitRenewEvent))
             .build();
+    Registry registry = Registry.get(existingDomain.getTld());
+    DomainHistory domainHistory =
+        buildDomainHistory(
+            newDomain, now, command.getPeriod(), registry.getRenewGracePeriodLength());
     EntityChanges entityChanges =
         flowCustomLogic.beforeSave(
             BeforeSaveParameters.newBuilder()
@@ -193,19 +201,18 @@ public final class DomainRenewFlow implements TransactionalFlow {
                 .setNewDomain(newDomain)
                 .setNow(now)
                 .setYears(years)
-                .setHistoryEntry(historyEntry)
+                .setHistoryEntry(domainHistory)
                 .setEntityChanges(
                     EntityChanges.newBuilder()
                         .setSaves(
                             ImmutableSet.of(
                                 newDomain,
-                                historyEntry,
+                                domainHistory,
                                 explicitRenewEvent,
                                 newAutorenewEvent,
                                 newAutorenewPollMessage))
                         .build())
                 .build());
-    persistEntityChanges(entityChanges);
     BeforeResponseReturnData responseData =
         flowCustomLogic.beforeResponse(
             BeforeResponseParameters.newBuilder()
@@ -213,23 +220,24 @@ public final class DomainRenewFlow implements TransactionalFlow {
                 .setResData(DomainRenewData.create(targetId, newExpirationTime))
                 .setResponseExtensions(createResponseExtensions(feesAndCredits, feeRenew))
                 .build());
+    persistEntityChanges(entityChanges);
     return responseBuilder
         .setResData(responseData.resData())
         .setExtensions(responseData.responseExtensions())
         .build();
   }
 
-  private HistoryEntry buildHistoryEntry(
-      DomainBase existingDomain, DateTime now, Period period, Duration renewGracePeriod) {
+  private DomainHistory buildDomainHistory(
+      DomainBase newDomain, DateTime now, Period period, Duration renewGracePeriod) {
     return historyBuilder
         .setType(HistoryEntry.Type.DOMAIN_RENEW)
         .setPeriod(period)
         .setModificationTime(now)
-        .setParent(existingDomain)
+        .setDomainContent(newDomain)
         .setDomainTransactionRecords(
             ImmutableSet.of(
                 DomainTransactionRecord.create(
-                    existingDomain.getTld(),
+                    newDomain.getTld(),
                     now.plus(renewGracePeriod),
                     TransactionReportField.netRenewsFieldFromYears(period.getValue()),
                     1)))
@@ -255,7 +263,7 @@ public final class DomainRenewFlow implements TransactionalFlow {
   }
 
   private OneTime createRenewBillingEvent(
-      String tld, Money renewCost, int years, HistoryEntry historyEntry, DateTime now) {
+      String tld, Money renewCost, int years, Key<DomainHistory> domainHistoryKey, DateTime now) {
     return new BillingEvent.OneTime.Builder()
         .setReason(Reason.RENEW)
         .setTargetId(targetId)
@@ -264,7 +272,7 @@ public final class DomainRenewFlow implements TransactionalFlow {
         .setCost(renewCost)
         .setEventTime(now)
         .setBillingTime(now.plus(Registry.get(tld).getRenewGracePeriodLength()))
-        .setParent(historyEntry)
+        .setParent(domainHistoryKey)
         .build();
   }
 

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

@@ -14,6 +14,7 @@
 
 package google.registry.flows.domain;
 
+import static google.registry.flows.FlowUtils.createHistoryKey;
 import static google.registry.flows.FlowUtils.validateClientIsLoggedIn;
 import static google.registry.flows.ResourceFlowUtils.loadAndVerifyExistence;
 import static google.registry.flows.ResourceFlowUtils.verifyOptionalAuthInfo;
@@ -49,6 +50,7 @@ import google.registry.model.billing.BillingEvent.OneTime;
 import google.registry.model.billing.BillingEvent.Reason;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.DomainCommand.Update;
+import google.registry.model.domain.DomainHistory;
 import google.registry.model.domain.fee.BaseFee.FeeType;
 import google.registry.model.domain.fee.Fee;
 import google.registry.model.domain.fee.FeeTransformResponseExtension;
@@ -117,7 +119,7 @@ public final class DomainRestoreRequestFlow implements TransactionalFlow  {
   @Inject @ClientId String clientId;
   @Inject @TargetId String targetId;
   @Inject @Superuser boolean isSuperuser;
-  @Inject HistoryEntry.Builder historyBuilder;
+  @Inject DomainHistory.Builder historyBuilder;
   @Inject DnsQueue dnsQueue;
   @Inject EppResponse.Builder responseBuilder;
   @Inject DomainPricingLogic pricingLogic;
@@ -142,7 +144,8 @@ public final class DomainRestoreRequestFlow implements TransactionalFlow  {
     Optional<FeeUpdateCommandExtension> feeUpdate =
         eppInput.getSingleExtension(FeeUpdateCommandExtension.class);
     verifyRestoreAllowed(command, existingDomain, feeUpdate, feesAndCredits, now);
-    HistoryEntry historyEntry = buildHistoryEntry(existingDomain, now);
+    Key<DomainHistory> domainHistoryKey = createHistoryKey(existingDomain, DomainHistory.class);
+    historyBuilder.setId(domainHistoryKey.getId());
     ImmutableSet.Builder<ImmutableObject> entitiesToSave = new ImmutableSet.Builder<>();
 
     DateTime newExpirationTime =
@@ -150,29 +153,31 @@ public final class DomainRestoreRequestFlow implements TransactionalFlow  {
     // Restore the expiration time on the deleted domain, except if that's already passed, then add
     // a year and bill for it immediately, with no grace period.
     if (isExpired) {
-      entitiesToSave.add(createRenewBillingEvent(historyEntry, feesAndCredits.getRenewCost(), now));
+      entitiesToSave.add(
+          createRenewBillingEvent(domainHistoryKey, feesAndCredits.getRenewCost(), now));
     }
     // Always bill for the restore itself.
     entitiesToSave.add(
-        createRestoreBillingEvent(historyEntry, feesAndCredits.getRestoreCost(), now));
+        createRestoreBillingEvent(domainHistoryKey, feesAndCredits.getRestoreCost(), now));
 
     BillingEvent.Recurring autorenewEvent =
         newAutorenewBillingEvent(existingDomain)
             .setEventTime(newExpirationTime)
             .setRecurrenceEndTime(END_OF_TIME)
-            .setParent(historyEntry)
+            .setParent(domainHistoryKey)
             .build();
     PollMessage.Autorenew autorenewPollMessage =
         newAutorenewPollMessage(existingDomain)
             .setEventTime(newExpirationTime)
             .setAutorenewEndTime(END_OF_TIME)
-            .setParent(historyEntry)
+            .setParentKey(domainHistoryKey)
             .build();
     DomainBase newDomain =
         performRestore(
             existingDomain, newExpirationTime, autorenewEvent, autorenewPollMessage, now, clientId);
     updateForeignKeyIndexDeletionTime(newDomain);
-    entitiesToSave.add(newDomain, historyEntry, autorenewEvent, autorenewPollMessage);
+    DomainHistory domainHistory = buildDomainHistory(newDomain, now);
+    entitiesToSave.add(newDomain, domainHistory, autorenewEvent, autorenewPollMessage);
     tm().putAll(entitiesToSave.build());
     tm().delete(existingDomain.getDeletePollMessage());
     dnsQueue.addDomainRefreshTask(existingDomain.getDomainName());
@@ -181,15 +186,15 @@ public final class DomainRestoreRequestFlow implements TransactionalFlow  {
         .build();
   }
 
-  private HistoryEntry buildHistoryEntry(DomainBase existingDomain, DateTime now) {
+  private DomainHistory buildDomainHistory(DomainBase newDomain, DateTime now) {
     return historyBuilder
         .setType(HistoryEntry.Type.DOMAIN_RESTORE)
         .setModificationTime(now)
-        .setParent(Key.create(existingDomain))
+        .setDomainContent(newDomain)
         .setDomainTransactionRecords(
             ImmutableSet.of(
                 DomainTransactionRecord.create(
-                    existingDomain.getTld(), now, TransactionReportField.RESTORED_DOMAINS, 1)))
+                    newDomain.getTld(), now, TransactionReportField.RESTORED_DOMAINS, 1)))
         .build();
   }
 
@@ -242,20 +247,19 @@ public final class DomainRestoreRequestFlow implements TransactionalFlow  {
   }
 
   private OneTime createRenewBillingEvent(
-      HistoryEntry historyEntry, Money renewCost, DateTime now) {
-    return prepareBillingEvent(historyEntry, renewCost, now)
-        .setReason(Reason.RENEW)
-        .build();
+      Key<DomainHistory> domainHistoryKey, Money renewCost, DateTime now) {
+    return prepareBillingEvent(domainHistoryKey, renewCost, now).setReason(Reason.RENEW).build();
   }
 
   private BillingEvent.OneTime createRestoreBillingEvent(
-      HistoryEntry historyEntry, Money restoreCost, DateTime now) {
-    return prepareBillingEvent(historyEntry, restoreCost, now)
+      Key<DomainHistory> domainHistoryKey, Money restoreCost, DateTime now) {
+    return prepareBillingEvent(domainHistoryKey, restoreCost, now)
         .setReason(Reason.RESTORE)
         .build();
   }
 
-  private OneTime.Builder prepareBillingEvent(HistoryEntry historyEntry, Money cost, DateTime now) {
+  private OneTime.Builder prepareBillingEvent(
+      Key<DomainHistory> domainHistoryKey, Money cost, DateTime now) {
     return new BillingEvent.OneTime.Builder()
         .setTargetId(targetId)
         .setClientId(clientId)
@@ -263,7 +267,7 @@ public final class DomainRestoreRequestFlow implements TransactionalFlow  {
         .setBillingTime(now)
         .setPeriodYears(1)
         .setCost(cost)
-        .setParent(historyEntry);
+        .setParent(domainHistoryKey);
   }
 
   private static ImmutableList<FeeTransformResponseExtension> createResponseExtensions(

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

@@ -15,6 +15,7 @@
 package google.registry.flows.domain;
 
 import static com.google.common.collect.Iterables.getOnlyElement;
+import static google.registry.flows.FlowUtils.createHistoryKey;
 import static google.registry.flows.FlowUtils.validateClientIsLoggedIn;
 import static google.registry.flows.ResourceFlowUtils.computeExDateForApprovalTime;
 import static google.registry.flows.ResourceFlowUtils.loadAndVerifyExistence;
@@ -48,6 +49,7 @@ import google.registry.model.billing.BillingEvent;
 import google.registry.model.billing.BillingEvent.Flag;
 import google.registry.model.billing.BillingEvent.Reason;
 import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.DomainHistory;
 import google.registry.model.domain.GracePeriod;
 import google.registry.model.domain.metadata.MetadataExtension;
 import google.registry.model.domain.rgp.GracePeriodStatus;
@@ -90,7 +92,7 @@ public final class DomainTransferApproveFlow implements TransactionalFlow {
   @Inject @ClientId String clientId;
   @Inject @TargetId String targetId;
   @Inject @Superuser boolean isSuperuser;
-  @Inject HistoryEntry.Builder historyBuilder;
+  @Inject DomainHistory.Builder historyBuilder;
   @Inject EppResponse.Builder responseBuilder;
   @Inject DomainTransferApproveFlow() {}
 
@@ -114,10 +116,10 @@ public final class DomainTransferApproveFlow implements TransactionalFlow {
     }
     DomainTransferData transferData = existingDomain.getTransferData();
     String gainingClientId = transferData.getGainingClientId();
-    Registry registry = Registry.get(existingDomain.getTld());
-    HistoryEntry historyEntry = buildHistoryEntry(existingDomain, registry, now, gainingClientId);
     // Create a transfer billing event for 1 year, unless the superuser extension was used to set
     // the transfer period to zero. There is not a transfer cost if the transfer period is zero.
+    Key<DomainHistory> domainHistoryKey = createHistoryKey(existingDomain, DomainHistory.class);
+    historyBuilder.setId(domainHistoryKey.getId());
     Optional<BillingEvent.OneTime> billingEvent =
         (transferData.getTransferPeriod().getValue() == 0)
             ? Optional.empty()
@@ -130,10 +132,9 @@ public final class DomainTransferApproveFlow implements TransactionalFlow {
                     .setCost(getDomainRenewCost(targetId, transferData.getTransferRequestTime(), 1))
                     .setEventTime(now)
                     .setBillingTime(now.plus(Registry.get(tld).getTransferGracePeriodLength()))
-                    .setParent(historyEntry)
+                    .setParent(domainHistoryKey)
                     .build());
     ImmutableList.Builder<ImmutableObject> entitiesToSave = new ImmutableList.Builder<>();
-    entitiesToSave.add(historyEntry);
     // If we are within an autorenew grace period, cancel the autorenew billing event and don't
     // increase the registration time, since the transfer subsumes the autorenew's extra year.
     GracePeriod autorenewGrace =
@@ -146,7 +147,8 @@ public final class DomainTransferApproveFlow implements TransactionalFlow {
       // still needs to be charged for the auto-renew.
       if (billingEvent.isPresent()) {
         entitiesToSave.add(
-            BillingEvent.Cancellation.forGracePeriod(autorenewGrace, historyEntry, targetId));
+            BillingEvent.Cancellation.forGracePeriod(
+                autorenewGrace, now, domainHistoryKey, targetId));
       }
     }
     // Close the old autorenew event and poll message at the transfer time (aka now). This may end
@@ -155,24 +157,26 @@ public final class DomainTransferApproveFlow implements TransactionalFlow {
     DateTime newExpirationTime =
         computeExDateForApprovalTime(existingDomain, now, transferData.getTransferPeriod());
     // Create a new autorenew event starting at the expiration time.
-    BillingEvent.Recurring autorenewEvent = new BillingEvent.Recurring.Builder()
-        .setReason(Reason.RENEW)
-        .setFlags(ImmutableSet.of(Flag.AUTO_RENEW))
-        .setTargetId(targetId)
-        .setClientId(gainingClientId)
-        .setEventTime(newExpirationTime)
-        .setRecurrenceEndTime(END_OF_TIME)
-        .setParent(historyEntry)
-        .build();
+    BillingEvent.Recurring autorenewEvent =
+        new BillingEvent.Recurring.Builder()
+            .setReason(Reason.RENEW)
+            .setFlags(ImmutableSet.of(Flag.AUTO_RENEW))
+            .setTargetId(targetId)
+            .setClientId(gainingClientId)
+            .setEventTime(newExpirationTime)
+            .setRecurrenceEndTime(END_OF_TIME)
+            .setParent(domainHistoryKey)
+            .build();
     // Create a new autorenew poll message.
-    PollMessage.Autorenew gainingClientAutorenewPollMessage = new PollMessage.Autorenew.Builder()
-        .setTargetId(targetId)
-        .setClientId(gainingClientId)
-        .setEventTime(newExpirationTime)
-        .setAutorenewEndTime(END_OF_TIME)
-        .setMsg("Domain was auto-renewed.")
-        .setParent(historyEntry)
-        .build();
+    PollMessage.Autorenew gainingClientAutorenewPollMessage =
+        new PollMessage.Autorenew.Builder()
+            .setTargetId(targetId)
+            .setClientId(gainingClientId)
+            .setEventTime(newExpirationTime)
+            .setAutorenewEndTime(END_OF_TIME)
+            .setMsg("Domain was auto-renewed.")
+            .setParentKey(domainHistoryKey)
+            .build();
     // Construct the post-transfer domain.
     DomainBase partiallyApprovedDomain =
         approvePendingTransfer(existingDomain, TransferStatus.CLIENT_APPROVED, now);
@@ -204,13 +208,19 @@ public final class DomainTransferApproveFlow implements TransactionalFlow {
             .setLastEppUpdateTime(now)
             .setLastEppUpdateClientId(clientId)
             .build();
+    Registry registry = Registry.get(existingDomain.getTld());
+    DomainHistory domainHistory = buildDomainHistory(newDomain, registry, now, gainingClientId);
     // Create a poll message for the gaining client.
     PollMessage gainingClientPollMessage =
         createGainingTransferPollMessage(
-            targetId, newDomain.getTransferData(), newExpirationTime, historyEntry);
+            targetId, newDomain.getTransferData(), newExpirationTime, now, domainHistoryKey);
     billingEvent.ifPresent(entitiesToSave::add);
     entitiesToSave.add(
-        autorenewEvent, gainingClientPollMessage, gainingClientAutorenewPollMessage, newDomain);
+        autorenewEvent,
+        gainingClientPollMessage,
+        gainingClientAutorenewPollMessage,
+        newDomain,
+        domainHistory);
     tm().putAll(entitiesToSave.build());
     // Delete the billing event and poll messages that were written in case the transfer would have
     // been implicitly server approved.
@@ -221,11 +231,11 @@ public final class DomainTransferApproveFlow implements TransactionalFlow {
         .build();
   }
 
-  private HistoryEntry buildHistoryEntry(
-      DomainBase existingDomain, Registry registry, DateTime now, String gainingClientId) {
+  private DomainHistory buildDomainHistory(
+      DomainBase newDomain, Registry registry, DateTime now, String gainingClientId) {
     ImmutableSet<DomainTransactionRecord> cancelingRecords =
         createCancelingRecords(
-            existingDomain,
+            newDomain,
             now,
             registry.getAutomaticTransferLength().plus(registry.getTransferGracePeriodLength()),
             ImmutableSet.of(TRANSFER_SUCCESSFUL));
@@ -233,15 +243,15 @@ public final class DomainTransferApproveFlow implements TransactionalFlow {
         .setType(HistoryEntry.Type.DOMAIN_TRANSFER_APPROVE)
         .setModificationTime(now)
         .setOtherClientId(gainingClientId)
-        .setParent(Key.create(existingDomain))
+        .setDomainContent(newDomain)
         .setDomainTransactionRecords(
             union(
                 cancelingRecords,
                 DomainTransactionRecord.create(
-                        existingDomain.getTld(),
-                        now.plus(registry.getTransferGracePeriodLength()),
-                        TRANSFER_SUCCESSFUL,
-                        1)))
+                    newDomain.getTld(),
+                    now.plus(registry.getTransferGracePeriodLength()),
+                    TRANSFER_SUCCESSFUL,
+                    1)))
         .build();
   }
 }

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

@@ -14,6 +14,7 @@
 
 package google.registry.flows.domain;
 
+import static google.registry.flows.FlowUtils.createHistoryKey;
 import static google.registry.flows.FlowUtils.validateClientIsLoggedIn;
 import static google.registry.flows.ResourceFlowUtils.loadAndVerifyExistence;
 import static google.registry.flows.ResourceFlowUtils.verifyHasPendingTransfer;
@@ -39,6 +40,7 @@ import google.registry.flows.FlowModule.TargetId;
 import google.registry.flows.TransactionalFlow;
 import google.registry.flows.annotations.ReportingSpec;
 import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.DomainHistory;
 import google.registry.model.domain.metadata.MetadataExtension;
 import google.registry.model.eppcommon.AuthInfo;
 import google.registry.model.eppoutput.EppResponse;
@@ -77,7 +79,7 @@ public final class DomainTransferCancelFlow implements TransactionalFlow {
   @Inject @ClientId String clientId;
   @Inject @TargetId String targetId;
   @Inject @Superuser boolean isSuperuser;
-  @Inject HistoryEntry.Builder historyBuilder;
+  @Inject DomainHistory.Builder historyBuilder;
   @Inject EppResponse.Builder responseBuilder;
   @Inject DomainTransferCancelFlow() {}
 
@@ -95,14 +97,20 @@ public final class DomainTransferCancelFlow implements TransactionalFlow {
       checkAllowedAccessToTld(clientId, existingDomain.getTld());
     }
     Registry registry = Registry.get(existingDomain.getTld());
-    HistoryEntry historyEntry = buildHistoryEntry(existingDomain, registry, now);
+
+    Key<DomainHistory> domainHistoryKey = createHistoryKey(existingDomain, DomainHistory.class);
+    historyBuilder
+        .setId(domainHistoryKey.getId())
+        .setOtherClientId(existingDomain.getTransferData().getLosingClientId());
+
     DomainBase newDomain =
         denyPendingTransfer(existingDomain, TransferStatus.CLIENT_CANCELLED, now, clientId);
+    DomainHistory domainHistory = buildDomainHistory(newDomain, registry, now);
     tm().putAll(
             newDomain,
-            historyEntry,
+            domainHistory,
             createLosingTransferPollMessage(
-                targetId, newDomain.getTransferData(), null, historyEntry));
+                targetId, newDomain.getTransferData(), null, domainHistoryKey));
     // Reopen the autorenew event and poll message that we closed for the implicit transfer. This
     // may recreate the autorenew poll message if it was deleted when the transfer request was made.
     updateAutorenewRecurrenceEndTime(existingDomain, END_OF_TIME);
@@ -114,19 +122,17 @@ public final class DomainTransferCancelFlow implements TransactionalFlow {
         .build();
   }
 
-  private HistoryEntry buildHistoryEntry(
-      DomainBase existingDomain, Registry registry, DateTime now) {
+  private DomainHistory buildDomainHistory(DomainBase newDomain, Registry registry, DateTime now) {
     ImmutableSet<DomainTransactionRecord> cancelingRecords =
         createCancelingRecords(
-            existingDomain,
+            newDomain,
             now,
             registry.getAutomaticTransferLength().plus(registry.getTransferGracePeriodLength()),
             ImmutableSet.of(TRANSFER_SUCCESSFUL));
     return historyBuilder
         .setType(HistoryEntry.Type.DOMAIN_TRANSFER_CANCEL)
-        .setOtherClientId(existingDomain.getTransferData().getLosingClientId())
         .setModificationTime(now)
-        .setParent(Key.create(existingDomain))
+        .setDomainContent(newDomain)
         .setDomainTransactionRecords(cancelingRecords)
         .build();
   }

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

@@ -14,6 +14,7 @@
 
 package google.registry.flows.domain;
 
+import static google.registry.flows.FlowUtils.createHistoryKey;
 import static google.registry.flows.FlowUtils.validateClientIsLoggedIn;
 import static google.registry.flows.ResourceFlowUtils.loadAndVerifyExistence;
 import static google.registry.flows.ResourceFlowUtils.verifyHasPendingTransfer;
@@ -41,6 +42,7 @@ import google.registry.flows.FlowModule.TargetId;
 import google.registry.flows.TransactionalFlow;
 import google.registry.flows.annotations.ReportingSpec;
 import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.DomainHistory;
 import google.registry.model.domain.metadata.MetadataExtension;
 import google.registry.model.eppcommon.AuthInfo;
 import google.registry.model.eppoutput.EppResponse;
@@ -79,7 +81,7 @@ public final class DomainTransferRejectFlow implements TransactionalFlow {
   @Inject @ClientId String clientId;
   @Inject @TargetId String targetId;
   @Inject @Superuser boolean isSuperuser;
-  @Inject HistoryEntry.Builder historyBuilder;
+  @Inject DomainHistory.Builder historyBuilder;
   @Inject EppResponse.Builder responseBuilder;
   @Inject DomainTransferRejectFlow() {}
 
@@ -91,7 +93,11 @@ public final class DomainTransferRejectFlow implements TransactionalFlow {
     DateTime now = tm().getTransactionTime();
     DomainBase existingDomain = loadAndVerifyExistence(DomainBase.class, targetId, now);
     Registry registry = Registry.get(existingDomain.getTld());
-    HistoryEntry historyEntry = buildHistoryEntry(existingDomain, registry, now);
+    Key<DomainHistory> domainHistoryKey = createHistoryKey(existingDomain, DomainHistory.class);
+    historyBuilder
+        .setId(domainHistoryKey.getId())
+        .setOtherClientId(existingDomain.getTransferData().getGainingClientId());
+
     verifyOptionalAuthInfo(authInfo, existingDomain);
     verifyHasPendingTransfer(existingDomain);
     verifyResourceOwnership(clientId, existingDomain);
@@ -100,11 +106,12 @@ public final class DomainTransferRejectFlow implements TransactionalFlow {
     }
     DomainBase newDomain =
         denyPendingTransfer(existingDomain, TransferStatus.CLIENT_REJECTED, now, clientId);
+    DomainHistory domainHistory = buildDomainHistory(newDomain, registry, now);
     tm().putAll(
             newDomain,
-            historyEntry,
+            domainHistory,
             createGainingTransferPollMessage(
-                targetId, newDomain.getTransferData(), null, historyEntry));
+                targetId, newDomain.getTransferData(), null, now, domainHistoryKey));
     // Reopen the autorenew event and poll message that we closed for the implicit transfer. This
     // may end up recreating the poll message if it was deleted upon the transfer request.
     updateAutorenewRecurrenceEndTime(existingDomain, END_OF_TIME);
@@ -116,24 +123,21 @@ public final class DomainTransferRejectFlow implements TransactionalFlow {
         .build();
   }
 
-  private HistoryEntry buildHistoryEntry(
-      DomainBase existingDomain, Registry registry, DateTime now) {
+  private DomainHistory buildDomainHistory(DomainBase newDomain, Registry registry, DateTime now) {
     ImmutableSet<DomainTransactionRecord> cancelingRecords =
         createCancelingRecords(
-            existingDomain,
+            newDomain,
             now,
             registry.getAutomaticTransferLength().plus(registry.getTransferGracePeriodLength()),
             ImmutableSet.of(TRANSFER_SUCCESSFUL));
     return historyBuilder
         .setType(HistoryEntry.Type.DOMAIN_TRANSFER_REJECT)
         .setModificationTime(now)
-        .setOtherClientId(existingDomain.getTransferData().getGainingClientId())
-        .setParent(Key.create(existingDomain))
         .setDomainTransactionRecords(
             union(
                 cancelingRecords,
-                DomainTransactionRecord.create(
-                    existingDomain.getTld(), now, TRANSFER_NACKED, 1)))
+                DomainTransactionRecord.create(newDomain.getTld(), now, TRANSFER_NACKED, 1)))
+        .setDomainContent(newDomain)
         .build();
   }
 }

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

@@ -14,6 +14,7 @@
 
 package google.registry.flows.domain;
 
+import static google.registry.flows.FlowUtils.createHistoryKey;
 import static google.registry.flows.FlowUtils.validateClientIsLoggedIn;
 import static google.registry.flows.ResourceFlowUtils.computeExDateForApprovalTime;
 import static google.registry.flows.ResourceFlowUtils.loadAndVerifyExistence;
@@ -51,6 +52,7 @@ import google.registry.flows.exceptions.TransferPeriodMustBeOneYearException;
 import google.registry.flows.exceptions.TransferPeriodZeroAndFeeTransferExtensionException;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.DomainCommand.Transfer;
+import google.registry.model.domain.DomainHistory;
 import google.registry.model.domain.Period;
 import google.registry.model.domain.fee.FeeTransferCommandExtension;
 import google.registry.model.domain.fee.FeeTransformResponseExtension;
@@ -126,7 +128,7 @@ public final class DomainTransferRequestFlow implements TransactionalFlow {
   @Inject @ClientId String gainingClientId;
   @Inject @TargetId String targetId;
   @Inject @Superuser boolean isSuperuser;
-  @Inject HistoryEntry.Builder historyBuilder;
+  @Inject DomainHistory.Builder historyBuilder;
   @Inject Trid trid;
   @Inject AsyncTaskEnqueuer asyncTaskEnqueuer;
   @Inject EppResponse.Builder responseBuilder;
@@ -169,7 +171,10 @@ public final class DomainTransferRequestFlow implements TransactionalFlow {
     if (feesAndCredits.isPresent()) {
       validateFeeChallenge(targetId, now, feeTransfer, feesAndCredits.get());
     }
-    HistoryEntry historyEntry = buildHistoryEntry(existingDomain, registry, now, period);
+    Key<DomainHistory> domainHistoryKey = createHistoryKey(existingDomain, DomainHistory.class);
+    historyBuilder
+        .setId(domainHistoryKey.getId())
+        .setOtherClientId(existingDomain.getCurrentSponsorClientId());
     DateTime automaticTransferTime =
         superuserExtension.isPresent()
             ? now.plusDays(superuserExtension.get().getAutomaticTransferLength())
@@ -190,7 +195,7 @@ public final class DomainTransferRequestFlow implements TransactionalFlow {
         createTransferServerApproveEntities(
             automaticTransferTime,
             serverApproveNewExpirationTime,
-            historyEntry,
+            domainHistoryKey,
             existingDomain,
             trid,
             gainingClientId,
@@ -209,9 +214,12 @@ public final class DomainTransferRequestFlow implements TransactionalFlow {
             serverApproveEntities,
             period);
     // Create a poll message to notify the losing registrar that a transfer was requested.
-    PollMessage requestPollMessage = createLosingTransferPollMessage(
-        targetId, pendingTransferData, serverApproveNewExpirationTime, historyEntry)
-            .asBuilder().setEventTime(now).build();
+    PollMessage requestPollMessage =
+        createLosingTransferPollMessage(
+                targetId, pendingTransferData, serverApproveNewExpirationTime, domainHistoryKey)
+            .asBuilder()
+            .setEventTime(now)
+            .build();
     // End the old autorenew event and poll message at the implicit transfer time. This may delete
     // the poll message if it has no events left. Note that if the automatic transfer succeeds, then
     // cloneProjectedAtTime() will replace these old autorenew entities with the server approve ones
@@ -225,10 +233,12 @@ public final class DomainTransferRequestFlow implements TransactionalFlow {
             .setLastEppUpdateTime(now)
             .setLastEppUpdateClientId(gainingClientId)
             .build();
+    DomainHistory domainHistory = buildDomainHistory(newDomain, registry, now, period);
+
     asyncTaskEnqueuer.enqueueAsyncResave(newDomain, now, automaticTransferTime);
     tm().putAll(
             new ImmutableSet.Builder<>()
-                .add(newDomain, historyEntry, requestPollMessage)
+                .add(newDomain, domainHistory, requestPollMessage)
                 .addAll(serverApproveEntities)
                 .build());
     return responseBuilder
@@ -302,14 +312,13 @@ public final class DomainTransferRequestFlow implements TransactionalFlow {
     }
   }
 
-  private HistoryEntry buildHistoryEntry(
-      DomainBase existingDomain, Registry registry, DateTime now, Period period) {
+  private DomainHistory buildDomainHistory(
+      DomainBase newDomain, Registry registry, DateTime now, Period period) {
     return historyBuilder
         .setType(HistoryEntry.Type.DOMAIN_TRANSFER_REQUEST)
-        .setOtherClientId(existingDomain.getCurrentSponsorClientId())
         .setPeriod(period)
         .setModificationTime(now)
-        .setParent(Key.create(existingDomain))
+        .setDomainContent(newDomain)
         .setDomainTransactionRecords(
             ImmutableSet.of(
                 DomainTransactionRecord.create(

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

@@ -20,10 +20,12 @@ import static google.registry.util.DateTimeUtils.END_OF_TIME;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
+import com.googlecode.objectify.Key;
 import google.registry.model.billing.BillingEvent;
 import google.registry.model.billing.BillingEvent.Flag;
 import google.registry.model.billing.BillingEvent.Reason;
 import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.DomainHistory;
 import google.registry.model.domain.GracePeriod;
 import google.registry.model.domain.Period;
 import google.registry.model.domain.rgp.GracePeriodStatus;
@@ -31,7 +33,6 @@ import google.registry.model.eppcommon.Trid;
 import google.registry.model.poll.PendingActionNotificationResponse.DomainPendingActionNotificationResponse;
 import google.registry.model.poll.PollMessage;
 import google.registry.model.registry.Registry;
-import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.transfer.DomainTransferData;
 import google.registry.model.transfer.TransferData;
 import google.registry.model.transfer.TransferData.TransferServerApproveEntity;
@@ -90,20 +91,21 @@ public final class DomainTransferUtils {
    * Returns a set of entities created speculatively in anticipation of a server approval.
    *
    * <p>This set consists of:
+   *
    * <ul>
-   *    <li>The one-time billing event charging the gaining registrar for the transfer
-   *    <li>A cancellation of an autorenew charge for the losing registrar, if the autorenew grace
-   *        period will apply at transfer time
-   *    <li>A new post-transfer autorenew billing event for the domain (and gaining registrar)
-   *    <li>A new post-transfer autorenew poll message for the domain (and gaining registrar)
-   *    <li>A poll message for the gaining registrar
-   *    <li>A poll message for the losing registrar
+   *   <li>The one-time billing event charging the gaining registrar for the transfer
+   *   <li>A cancellation of an autorenew charge for the losing registrar, if the autorenew grace
+   *       period will apply at transfer time
+   *   <li>A new post-transfer autorenew billing event for the domain (and gaining registrar)
+   *   <li>A new post-transfer autorenew poll message for the domain (and gaining registrar)
+   *   <li>A poll message for the gaining registrar
+   *   <li>A poll message for the losing registrar
    * </ul>
    */
   public static ImmutableSet<TransferServerApproveEntity> createTransferServerApproveEntities(
       DateTime automaticTransferTime,
       DateTime serverApproveNewExpirationTime,
-      HistoryEntry historyEntry,
+      Key<DomainHistory> domainHistoryKey,
       DomainBase existingDomain,
       Trid trid,
       String gainingClientId,
@@ -123,32 +125,39 @@ public final class DomainTransferUtils {
             .build();
     Registry registry = Registry.get(existingDomain.getTld());
     ImmutableSet.Builder<TransferServerApproveEntity> builder = new ImmutableSet.Builder<>();
-    if (transferCost.isPresent()) {
-      builder.add(
-          createTransferBillingEvent(
-              automaticTransferTime,
-              historyEntry,
-              targetId,
-              gainingClientId,
-              registry,
-              transferCost.get()));
-    }
+    transferCost.ifPresent(
+        cost ->
+            builder.add(
+                createTransferBillingEvent(
+                    automaticTransferTime,
+                    domainHistoryKey,
+                    targetId,
+                    gainingClientId,
+                    registry,
+                    cost)));
     createOptionalAutorenewCancellation(
-            automaticTransferTime, historyEntry, targetId, existingDomain, transferCost)
+            automaticTransferTime, now, domainHistoryKey, targetId, existingDomain, transferCost)
         .ifPresent(builder::add);
     return builder
         .add(
             createGainingClientAutorenewEvent(
-                serverApproveNewExpirationTime, historyEntry, targetId, gainingClientId))
+                serverApproveNewExpirationTime, domainHistoryKey, targetId, gainingClientId))
         .add(
             createGainingClientAutorenewPollMessage(
-                serverApproveNewExpirationTime, historyEntry, targetId, gainingClientId))
+                serverApproveNewExpirationTime, domainHistoryKey, targetId, gainingClientId))
         .add(
             createGainingTransferPollMessage(
-                targetId, serverApproveTransferData, serverApproveNewExpirationTime, historyEntry))
+                targetId,
+                serverApproveTransferData,
+                serverApproveNewExpirationTime,
+                now,
+                domainHistoryKey))
         .add(
             createLosingTransferPollMessage(
-                targetId, serverApproveTransferData, serverApproveNewExpirationTime, historyEntry))
+                targetId,
+                serverApproveTransferData,
+                serverApproveNewExpirationTime,
+                domainHistoryKey))
         .build();
   }
 
@@ -157,19 +166,21 @@ public final class DomainTransferUtils {
       String targetId,
       TransferData transferData,
       @Nullable DateTime extendedRegistrationExpirationTime,
-      HistoryEntry historyEntry) {
+      DateTime now,
+      Key<DomainHistory> domainHistoryKey) {
     return new PollMessage.OneTime.Builder()
         .setClientId(transferData.getGainingClientId())
         .setEventTime(transferData.getPendingTransferExpirationTime())
         .setMsg(transferData.getTransferStatus().getMessage())
-        .setResponseData(ImmutableList.of(
-            createTransferResponse(targetId, transferData, extendedRegistrationExpirationTime),
-            DomainPendingActionNotificationResponse.create(
-                  targetId,
-                  transferData.getTransferStatus().isApproved(),
-                  transferData.getTransferRequestTrid(),
-                  historyEntry.getModificationTime())))
-        .setParent(historyEntry)
+        .setResponseData(
+            ImmutableList.of(
+                createTransferResponse(targetId, transferData, extendedRegistrationExpirationTime),
+                DomainPendingActionNotificationResponse.create(
+                    targetId,
+                    transferData.getTransferStatus().isApproved(),
+                    transferData.getTransferRequestTrid(),
+                    now)))
+        .setParentKey(domainHistoryKey)
         .build();
   }
 
@@ -178,14 +189,15 @@ public final class DomainTransferUtils {
       String targetId,
       TransferData transferData,
       @Nullable DateTime extendedRegistrationExpirationTime,
-      HistoryEntry historyEntry) {
+      Key<DomainHistory> domainHistoryKey) {
     return new PollMessage.OneTime.Builder()
         .setClientId(transferData.getLosingClientId())
         .setEventTime(transferData.getPendingTransferExpirationTime())
         .setMsg(transferData.getTransferStatus().getMessage())
-        .setResponseData(ImmutableList.of(
-            createTransferResponse(targetId, transferData, extendedRegistrationExpirationTime)))
-        .setParent(historyEntry)
+        .setResponseData(
+            ImmutableList.of(
+                createTransferResponse(targetId, transferData, extendedRegistrationExpirationTime)))
+        .setParentKey(domainHistoryKey)
         .build();
   }
 
@@ -207,7 +219,7 @@ public final class DomainTransferUtils {
 
   private static PollMessage.Autorenew createGainingClientAutorenewPollMessage(
       DateTime serverApproveNewExpirationTime,
-      HistoryEntry historyEntry,
+      Key<DomainHistory> domainHistoryKey,
       String targetId,
       String gainingClientId) {
     return new PollMessage.Autorenew.Builder()
@@ -216,13 +228,13 @@ public final class DomainTransferUtils {
         .setEventTime(serverApproveNewExpirationTime)
         .setAutorenewEndTime(END_OF_TIME)
         .setMsg("Domain was auto-renewed.")
-        .setParent(historyEntry)
+        .setParentKey(domainHistoryKey)
         .build();
   }
 
   private static BillingEvent.Recurring createGainingClientAutorenewEvent(
       DateTime serverApproveNewExpirationTime,
-      HistoryEntry historyEntry,
+      Key<DomainHistory> domainHistoryKey,
       String targetId,
       String gainingClientId) {
     return new BillingEvent.Recurring.Builder()
@@ -232,7 +244,7 @@ public final class DomainTransferUtils {
         .setClientId(gainingClientId)
         .setEventTime(serverApproveNewExpirationTime)
         .setRecurrenceEndTime(END_OF_TIME)
-        .setParent(historyEntry)
+        .setParent(domainHistoryKey)
         .build();
   }
 
@@ -254,7 +266,8 @@ public final class DomainTransferUtils {
    */
   private static Optional<BillingEvent.Cancellation> createOptionalAutorenewCancellation(
       DateTime automaticTransferTime,
-      HistoryEntry historyEntry,
+      DateTime now,
+      Key<DomainHistory> domainHistoryKey,
       String targetId,
       DomainBase existingDomain,
       Optional<Money> transferCost) {
@@ -265,7 +278,8 @@ public final class DomainTransferUtils {
             domainAtTransferTime.getGracePeriodsOfType(GracePeriodStatus.AUTO_RENEW), null);
     if (autorenewGracePeriod != null && transferCost.isPresent()) {
       return Optional.of(
-          BillingEvent.Cancellation.forGracePeriod(autorenewGracePeriod, historyEntry, targetId)
+          BillingEvent.Cancellation.forGracePeriod(
+                  autorenewGracePeriod, now, domainHistoryKey, targetId)
               .asBuilder()
               .setEventTime(automaticTransferTime)
               .build());
@@ -275,7 +289,7 @@ public final class DomainTransferUtils {
 
   private static BillingEvent.OneTime createTransferBillingEvent(
       DateTime automaticTransferTime,
-      HistoryEntry historyEntry,
+      Key<DomainHistory> domainHistoryKey,
       String targetId,
       String gainingClientId,
       Registry registry,
@@ -288,7 +302,7 @@ public final class DomainTransferUtils {
         .setPeriodYears(1)
         .setEventTime(automaticTransferTime)
         .setBillingTime(automaticTransferTime.plus(registry.getTransferGracePeriodLength()))
-        .setParent(historyEntry)
+        .setParent(domainHistoryKey)
         .build();
   }
 

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

@@ -43,7 +43,6 @@ import static google.registry.persistence.transaction.TransactionManagerFactory.
 
 import com.google.common.collect.ImmutableSet;
 import com.google.common.net.InternetDomainName;
-import com.googlecode.objectify.Key;
 import google.registry.dns.DnsQueue;
 import google.registry.flows.EppException;
 import google.registry.flows.ExtensionManager;
@@ -65,6 +64,7 @@ import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.DomainCommand.Update;
 import google.registry.model.domain.DomainCommand.Update.AddRemove;
 import google.registry.model.domain.DomainCommand.Update.Change;
+import google.registry.model.domain.DomainHistory;
 import google.registry.model.domain.fee.FeeUpdateCommandExtension;
 import google.registry.model.domain.metadata.MetadataExtension;
 import google.registry.model.domain.secdns.DelegationSignerData;
@@ -142,7 +142,7 @@ public final class DomainUpdateFlow implements TransactionalFlow {
   @Inject @ClientId String clientId;
   @Inject @TargetId String targetId;
   @Inject @Superuser boolean isSuperuser;
-  @Inject HistoryEntry.Builder historyBuilder;
+  @Inject DomainHistory.Builder historyBuilder;
   @Inject DnsQueue dnsQueue;
   @Inject EppResponse.Builder responseBuilder;
   @Inject DomainUpdateFlowCustomLogic flowCustomLogic;
@@ -165,19 +165,19 @@ public final class DomainUpdateFlow implements TransactionalFlow {
     verifyUpdateAllowed(command, existingDomain, now);
     flowCustomLogic.afterValidation(
         AfterValidationParameters.newBuilder().setExistingDomain(existingDomain).build());
-    HistoryEntry historyEntry = buildHistoryEntry(existingDomain, now);
     DomainBase newDomain = performUpdate(command, existingDomain, now);
+    DomainHistory domainHistory = buildDomainHistory(newDomain, now);
     validateNewState(newDomain);
     dnsQueue.addDomainRefreshTask(targetId);
     ImmutableSet.Builder<ImmutableObject> entitiesToSave = new ImmutableSet.Builder<>();
-    entitiesToSave.add(newDomain, historyEntry);
+    entitiesToSave.add(newDomain, domainHistory);
     Optional<BillingEvent.OneTime> statusUpdateBillingEvent =
-        createBillingEventForStatusUpdates(existingDomain, newDomain, historyEntry, now);
+        createBillingEventForStatusUpdates(existingDomain, newDomain, domainHistory, now);
     statusUpdateBillingEvent.ifPresent(entitiesToSave::add);
     EntityChanges entityChanges =
         flowCustomLogic.beforeSave(
             BeforeSaveParameters.newBuilder()
-                .setHistoryEntry(historyEntry)
+                .setHistoryEntry(domainHistory)
                 .setNewDomain(newDomain)
                 .setExistingDomain(existingDomain)
                 .setEntityChanges(
@@ -217,11 +217,11 @@ public final class DomainUpdateFlow implements TransactionalFlow {
         tld, add.getNameserverFullyQualifiedHostNames());
   }
 
-  private HistoryEntry buildHistoryEntry(DomainBase existingDomain, DateTime now) {
+  private DomainHistory buildDomainHistory(DomainBase newDomain, DateTime now) {
     return historyBuilder
         .setType(HistoryEntry.Type.DOMAIN_UPDATE)
         .setModificationTime(now)
-        .setParent(Key.create(existingDomain))
+        .setDomainContent(newDomain)
         .build();
   }
 

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

@@ -16,6 +16,7 @@ package google.registry.flows.domain.token;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 
 import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
@@ -153,7 +154,7 @@ public class AllocationTokenFlowUtils {
       throw new InvalidAllocationTokenException();
     }
     Optional<AllocationToken> maybeTokenEntity =
-        tm().loadByKeyIfPresent(VKey.create(AllocationToken.class, token));
+        transactIfJpaTm(() -> tm().loadByKeyIfPresent(VKey.create(AllocationToken.class, token)));
     if (!maybeTokenEntity.isPresent()) {
       throw new InvalidAllocationTokenException();
     }

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

@@ -15,7 +15,7 @@
 package google.registry.flows.host;
 
 import static google.registry.flows.FlowUtils.validateClientIsLoggedIn;
-import static google.registry.flows.ResourceFlowUtils.failfastForAsyncDelete;
+import static google.registry.flows.ResourceFlowUtils.checkLinkedDomains;
 import static google.registry.flows.ResourceFlowUtils.loadAndVerifyExistence;
 import static google.registry.flows.ResourceFlowUtils.verifyNoDisallowedStatuses;
 import static google.registry.flows.ResourceFlowUtils.verifyResourceOwnership;
@@ -65,10 +65,11 @@ import org.joda.time.DateTime;
 @ReportingSpec(ActivityReportField.HOST_DELETE)
 public final class HostDeleteFlow implements TransactionalFlow {
 
-  private static final ImmutableSet<StatusValue> DISALLOWED_STATUSES = ImmutableSet.of(
-      StatusValue.CLIENT_DELETE_PROHIBITED,
-      StatusValue.PENDING_DELETE,
-      StatusValue.SERVER_DELETE_PROHIBITED);
+  private static final ImmutableSet<StatusValue> DISALLOWED_STATUSES =
+      ImmutableSet.of(
+          StatusValue.CLIENT_DELETE_PROHIBITED,
+          StatusValue.PENDING_DELETE,
+          StatusValue.SERVER_DELETE_PROHIBITED);
 
   @Inject ExtensionManager extensionManager;
   @Inject @ClientId String clientId;
@@ -78,7 +79,9 @@ public final class HostDeleteFlow implements TransactionalFlow {
   @Inject HistoryEntry.Builder historyBuilder;
   @Inject AsyncTaskEnqueuer asyncTaskEnqueuer;
   @Inject EppResponse.Builder responseBuilder;
-  @Inject HostDeleteFlow() {}
+
+  @Inject
+  HostDeleteFlow() {}
 
   @Override
   public final EppResponse run() throws EppException {
@@ -87,7 +90,7 @@ public final class HostDeleteFlow implements TransactionalFlow {
     validateClientIsLoggedIn(clientId);
     DateTime now = tm().getTransactionTime();
     validateHostName(targetId);
-    failfastForAsyncDelete(targetId, now, HostResource.class, DomainBase::getNameservers);
+    checkLinkedDomains(targetId, now, HostResource.class, DomainBase::getNameservers);
     HostResource existingHost = loadAndVerifyExistence(HostResource.class, targetId, now);
     verifyNoDisallowedStatuses(existingHost, DISALLOWED_STATUSES);
     if (!isSuperuser) {

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

@@ -19,6 +19,7 @@ import static google.registry.flows.ResourceFlowUtils.loadAndVerifyExistence;
 import static google.registry.flows.host.HostFlowUtils.validateHostName;
 import static google.registry.model.EppResourceUtils.isLinked;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 
 import com.google.common.collect.ImmutableSet;
 import google.registry.flows.EppException;
@@ -76,7 +77,8 @@ public final class HostInfoFlow implements Flow {
     // there is no superordinate domain, the host's own values for these fields will be correct.
     if (host.isSubordinate()) {
       DomainBase superordinateDomain =
-          tm().loadByKey(host.getSuperordinateDomain()).cloneProjectedAtTime(now);
+          transactIfJpaTm(
+              () -> tm().loadByKey(host.getSuperordinateDomain()).cloneProjectedAtTime(now));
       hostInfoDataBuilder
           .setCurrentSponsorClientId(superordinateDomain.getCurrentSponsorClientId())
           .setLastTransferTime(host.computeLastTransferTime(superordinateDomain));

core/src/main/java/google/registry/keyring/kms/KmsKeyring.java

@@ -266,7 +266,13 @@ public class KmsKeyring implements Keyring {
             .collect(ImmutableList.toImmutableList());
 
     for (String keyName : labels) {
-      byte[] dsData = getDecryptedDataFromDatastore(keyName);
+      byte[] dsData;
+      try {
+        dsData = getDecryptedDataFromDatastore(keyName);
+      } catch (IllegalStateException e) {
+        logger.atWarning().log("Cannot load %s from Datastore. Skipping...", keyName);
+        continue;
+      }
       byte[] secretStoreData = getDataFromSecretStore(keyName);
       if (Arrays.equals(dsData, secretStoreData)) {
         logger.atInfo().log("%s is already up to date.\n", keyName);

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

@@ -47,7 +47,6 @@ import google.registry.model.server.KmsSecret;
 import google.registry.model.server.KmsSecretRevision;
 import google.registry.model.server.Lock;
 import google.registry.model.server.ServerSecret;
-import google.registry.model.smd.SignedMarkRevocationList;
 import google.registry.model.tmch.ClaimsListShard;
 import google.registry.model.tmch.ClaimsListShard.ClaimsListRevision;
 import google.registry.model.tmch.ClaimsListShard.ClaimsListSingleton;
@@ -105,7 +104,6 @@ public final class EntityClasses {
           Registry.class,
           ReservedList.class,
           ServerSecret.class,
-          SignedMarkRevocationList.class,
           TmchCrl.class);
 
   private EntityClasses() {}

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

@@ -42,6 +42,7 @@ import google.registry.model.ImmutableObject;
 import google.registry.model.annotations.ReportedOn;
 import google.registry.model.common.TimeOfYear;
 import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.DomainHistory;
 import google.registry.model.domain.GracePeriod;
 import google.registry.model.domain.rgp.GracePeriodStatus;
 import google.registry.model.domain.token.AllocationToken;
@@ -114,7 +115,7 @@ public abstract class BillingEvent extends ImmutableObject
   /** Entity id. */
   @Id @javax.persistence.Id Long id;
 
-  @Parent @DoNotHydrate @Transient Key<HistoryEntry> parent;
+  @Parent @DoNotHydrate @Transient Key<? extends HistoryEntry> parent;
 
   /** The registrar to bill. */
   @Index
@@ -191,7 +192,7 @@ public abstract class BillingEvent extends ImmutableObject
     return targetId;
   }
 
-  public Key<HistoryEntry> getParentKey() {
+  public Key<? extends HistoryEntry> getParentKey() {
     return parent;
   }
 
@@ -258,7 +259,7 @@ public abstract class BillingEvent extends ImmutableObject
       return thisCastToDerived();
     }
 
-    public B setParent(Key<HistoryEntry> parentKey) {
+    public B setParent(Key<? extends HistoryEntry> parentKey) {
       getInstance().parent = parentKey;
       return thisCastToDerived();
     }
@@ -602,23 +603,27 @@ public abstract class BillingEvent extends ImmutableObject
             GracePeriodStatus.TRANSFER, Reason.TRANSFER);
 
     /**
-     * Creates a cancellation billing event (parented on the provided history entry, and with the
-     * history entry's event time) that will cancel out the provided grace period's billing event,
+     * Creates a cancellation billing event (parented on the provided history key, and with the
+     * corresponding event time) that will cancel out the provided grace period's billing event,
      * using the supplied targetId and deriving other metadata (clientId, billing time, and the
      * cancellation reason) from the grace period.
      */
     public static BillingEvent.Cancellation forGracePeriod(
-        GracePeriod gracePeriod, HistoryEntry historyEntry, String targetId) {
+        GracePeriod gracePeriod,
+        DateTime eventTime,
+        Key<DomainHistory> domainHistoryKey,
+        String targetId) {
       checkArgument(gracePeriod.hasBillingEvent(),
           "Cannot create cancellation for grace period without billing event");
-      BillingEvent.Cancellation.Builder builder = new BillingEvent.Cancellation.Builder()
-          .setReason(checkNotNull(GRACE_PERIOD_TO_REASON.get(gracePeriod.getType())))
-          .setTargetId(targetId)
-          .setClientId(gracePeriod.getClientId())
-          .setEventTime(historyEntry.getModificationTime())
-          // The charge being cancelled will take place at the grace period's expiration time.
-          .setBillingTime(gracePeriod.getExpirationTime())
-          .setParent(historyEntry);
+      BillingEvent.Cancellation.Builder builder =
+          new BillingEvent.Cancellation.Builder()
+              .setReason(checkNotNull(GRACE_PERIOD_TO_REASON.get(gracePeriod.getType())))
+              .setTargetId(targetId)
+              .setClientId(gracePeriod.getClientId())
+              .setEventTime(eventTime)
+              // The charge being cancelled will take place at the grace period's expiration time.
+              .setBillingTime(gracePeriod.getExpirationTime())
+              .setParent(domainHistoryKey);
       // Set the grace period's billing event using the appropriate Cancellation builder method.
       if (gracePeriod.getOneTimeBillingEvent() != null) {
         builder.setOneTimeEventKey(gracePeriod.getOneTimeBillingEvent());

core/src/main/java/google/registry/model/common/Cursor.java

@@ -36,6 +36,7 @@ import google.registry.persistence.VKey;
 import google.registry.schema.replay.DatastoreAndSqlEntity;
 import java.io.Serializable;
 import java.util.List;
+import java.util.Optional;
 import javax.persistence.Column;
 import javax.persistence.EnumType;
 import javax.persistence.Enumerated;
@@ -280,8 +281,8 @@ public class Cursor extends ImmutableObject implements DatastoreAndSqlEntity {
   /**
    * Returns the current time for a given cursor, or {@code START_OF_TIME} if the cursor is null.
    */
-  public static DateTime getCursorTimeOrStartOfTime(Cursor cursor) {
-    return cursor != null ? cursor.getCursorTime() : START_OF_TIME;
+  public static DateTime getCursorTimeOrStartOfTime(Optional<Cursor> cursor) {
+    return cursor.map(Cursor::getCursorTime).orElse(START_OF_TIME);
   }
 
   public DateTime getCursorTime() {

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

@@ -287,7 +287,7 @@ public class ContactBase extends EppResource implements ResourceWithTransferData
   }
 
   @Override
-  public Builder asBuilder() {
+  public Builder<? extends ContactBase, ?> asBuilder() {
     return new Builder<>(clone(this));
   }
 

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

@@ -56,6 +56,7 @@ import javax.persistence.PostLoad;
 public class ContactHistory extends HistoryEntry implements SqlEntity {
 
   // Store ContactBase instead of ContactResource so we don't pick up its @Id
+  // Nullable for the sake of pre-Registry-3.0 history objects
   @Nullable ContactBase contactBase;
 
   @Id
@@ -108,6 +109,9 @@ public class ContactHistory extends HistoryEntry implements SqlEntity {
     if (contactBase != null && contactBase.getContactId() == null) {
       contactBase = null;
     }
+    if (contactBase != null && contactBase.getRepoId() == null) {
+      contactBase = contactBase.asBuilder().setRepoId(parent.getName()).build();
+    }
   }
 
   // In Datastore, save as a HistoryEntry object regardless of this object's type
@@ -193,9 +197,13 @@ public class ContactHistory extends HistoryEntry implements SqlEntity {
       super(instance);
     }
 
-    public Builder setContactBase(ContactBase contactBase) {
+    public Builder setContactBase(@Nullable ContactBase contactBase) {
+      // Nullable for the sake of pre-Registry-3.0 history objects
+      if (contactBase == null) {
+        return this;
+      }
       getInstance().contactBase = contactBase;
-      return this;
+      return super.setParent(contactBase);
     }
 
     public Builder setContactRepoId(String contactRepoId) {

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

@@ -33,6 +33,7 @@ import google.registry.persistence.VKey;
 import google.registry.schema.replay.DatastoreEntity;
 import google.registry.schema.replay.SqlEntity;
 import java.io.Serializable;
+import java.util.HashSet;
 import java.util.Optional;
 import java.util.Set;
 import javax.annotation.Nullable;
@@ -76,6 +77,7 @@ import javax.persistence.Table;
 public class DomainHistory extends HistoryEntry implements SqlEntity {
 
   // Store DomainContent instead of DomainBase so we don't pick up its @Id
+  // Nullable for the sake of pre-Registry-3.0 history objects
   @Nullable DomainContent domainContent;
 
   @Id
@@ -126,7 +128,8 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
         insertable = false,
         updatable = false)
   })
-  Set<DomainDsDataHistory> dsDataHistories = ImmutableSet.of();
+  // HashSet rather than ImmutableSet so that Hibernate can fill them out lazily on request
+  Set<DomainDsDataHistory> dsDataHistories = new HashSet<>();
 
   @Ignore
   @OneToMany(
@@ -145,7 +148,8 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
         insertable = false,
         updatable = false)
   })
-  Set<GracePeriodHistory> gracePeriodHistories = ImmutableSet.of();
+  // HashSet rather than ImmutableSet so that Hibernate can fill them out lazily on request
+  Set<GracePeriodHistory> gracePeriodHistories = new HashSet<>();
 
   @Override
   @Nullable
@@ -341,9 +345,13 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
       super(instance);
     }
 
-    public Builder setDomainContent(DomainContent domainContent) {
+    public Builder setDomainContent(@Nullable DomainContent domainContent) {
+      // Nullable for the sake of pre-Registry-3.0 history objects
+      if (domainContent == null) {
+        return this;
+      }
       getInstance().domainContent = domainContent;
-      return this;
+      return super.setParent(domainContent);
     }
 
     public Builder setDomainRepoId(String domainRepoId) {

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

@@ -38,7 +38,7 @@ import org.joda.time.DateTime;
 public class GracePeriodBase extends ImmutableObject {
 
   /** Unique id required for hibernate representation. */
-  @Transient Long gracePeriodId;
+  @Transient long gracePeriodId;
 
   /** Repository id for the domain which this grace period belongs to. */
   @Ignore

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

@@ -57,6 +57,7 @@ import javax.persistence.PostLoad;
 public class HostHistory extends HistoryEntry implements SqlEntity {
 
   // Store HostBase instead of HostResource so we don't pick up its @Id
+  // Nullable for the sake of pre-Registry-3.0 history objects
   @Nullable HostBase hostBase;
 
   @Id
@@ -194,9 +195,13 @@ public class HostHistory extends HistoryEntry implements SqlEntity {
       super(instance);
     }
 
-    public Builder setHostBase(HostBase hostBase) {
+    public Builder setHostBase(@Nullable HostBase hostBase) {
+      // Nullable for the sake of pre-Registry-3.0 history objects
+      if (hostBase == null) {
+        return this;
+      }
       getInstance().hostBase = hostBase;
-      return this;
+      return super.setParent(hostBase);
     }
 
     public Builder setHostRepoId(String hostRepoId) {

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

@@ -37,12 +37,19 @@ import google.registry.model.domain.DomainHistory;
 import google.registry.model.host.HostHistory;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.persistence.VKey;
+import google.registry.persistence.transaction.QueryComposer;
 import google.registry.persistence.transaction.TransactionManager;
+import google.registry.schema.replay.DatastoreEntity;
+import google.registry.schema.replay.SqlEntity;
+import java.util.List;
 import java.util.NoSuchElementException;
 import java.util.Optional;
 import java.util.function.Supplier;
+import java.util.stream.Stream;
 import java.util.stream.StreamSupport;
 import javax.annotation.Nullable;
+import javax.persistence.NoResultException;
+import javax.persistence.NonUniqueResultException;
 import org.joda.time.DateTime;
 
 /** Datastore implementation of {@link TransactionManager}. */
@@ -136,22 +143,23 @@ public class DatastoreTransactionManager implements TransactionManager {
 
   @Override
   public void putAll(Object... entities) {
-    syncIfTransactionless(getOfy().save().entities(entities));
+    syncIfTransactionless(
+        getOfy().save().entities(toDatastoreEntities(ImmutableList.copyOf(entities))));
   }
 
   @Override
   public void putAll(ImmutableCollection<?> entities) {
-    syncIfTransactionless(getOfy().save().entities(entities));
+    syncIfTransactionless(getOfy().save().entities(toDatastoreEntities(entities)));
   }
 
   @Override
   public void putWithoutBackup(Object entity) {
-    syncIfTransactionless(getOfy().saveWithoutBackup().entities(entity));
+    syncIfTransactionless(getOfy().saveWithoutBackup().entities(toDatastoreEntity(entity)));
   }
 
   @Override
   public void putAllWithoutBackup(ImmutableCollection<?> entities) {
-    syncIfTransactionless(getOfy().saveWithoutBackup().entities(entities));
+    syncIfTransactionless(getOfy().saveWithoutBackup().entities(toDatastoreEntities(entities)));
   }
 
   @Override
@@ -176,7 +184,7 @@ public class DatastoreTransactionManager implements TransactionManager {
 
   @Override
   public boolean exists(Object entity) {
-    return getOfy().load().key(Key.create(entity)).now() != null;
+    return getOfy().load().key(Key.create(toDatastoreEntity(entity))).now() != null;
   }
 
   @Override
@@ -205,8 +213,7 @@ public class DatastoreTransactionManager implements TransactionManager {
     return getOfy().load().keys(keyMap.keySet()).entrySet().stream()
         .collect(
             toImmutableMap(
-                entry -> keyMap.get(entry.getKey()),
-                entry -> toChildHistoryEntryIfPossible(entry.getValue())));
+                entry -> keyMap.get(entry.getKey()), entry -> toSqlEntity(entry.getValue())));
   }
 
   @Override
@@ -239,7 +246,7 @@ public class DatastoreTransactionManager implements TransactionManager {
 
   @Override
   public <T> T loadByEntity(T entity) {
-    return ofy().load().entity(entity).now();
+    return (T) toSqlEntity(ofy().load().entity(toDatastoreEntity(entity)).now());
   }
 
   @Override
@@ -281,7 +288,7 @@ public class DatastoreTransactionManager implements TransactionManager {
 
   @Override
   public void delete(Object entity) {
-    syncIfTransactionless(getOfy().delete().entity(entity));
+    syncIfTransactionless(getOfy().delete().entity(toDatastoreEntity(entity)));
   }
 
   @Override
@@ -299,7 +306,12 @@ public class DatastoreTransactionManager implements TransactionManager {
 
   @Override
   public void deleteWithoutBackup(Object entity) {
-    syncIfTransactionless(getOfy().deleteWithoutBackup().entity(entity));
+    syncIfTransactionless(getOfy().deleteWithoutBackup().entity(toDatastoreEntity(entity)));
+  }
+
+  @Override
+  public <T> QueryComposer<T> createQueryComposer(Class<T> entity) {
+    return new DatastoreQueryComposerImpl(entity);
   }
 
   @Override
@@ -338,29 +350,104 @@ public class DatastoreTransactionManager implements TransactionManager {
    */
   private void saveEntity(Object entity) {
     checkArgumentNotNull(entity, "entity must be specified");
-    if (entity instanceof HistoryEntry) {
-      entity = ((HistoryEntry) entity).asHistoryEntry();
-    }
-    syncIfTransactionless(getOfy().save().entity(entity));
+    syncIfTransactionless(getOfy().save().entity(toDatastoreEntity(entity)));
   }
 
   @Nullable
   private <T> T loadNullable(VKey<T> key) {
-    return toChildHistoryEntryIfPossible(getOfy().load().key(key.getOfyKey()).now());
+    return toSqlEntity(getOfy().load().key(key.getOfyKey()).now());
   }
 
-  /** Converts a nonnull {@link HistoryEntry} to the child format, e.g. {@link DomainHistory} */
-  @SuppressWarnings("unchecked")
-  public static <T> T toChildHistoryEntryIfPossible(@Nullable T obj) {
-    // NB: The Key of the object in question may not necessarily be the resulting class that we
-    // wish to have. Because all *History classes are @EntitySubclasses, their Keys will have type
-    // HistoryEntry -- even if you create them based off the *History class.
-    if (obj instanceof HistoryEntry
-        && !(obj instanceof ContactHistory)
-        && !(obj instanceof DomainHistory)
-        && !(obj instanceof HostHistory)) {
-      return (T) ((HistoryEntry) obj).toChildHistoryEntity();
+  /**
+   * Converts a possible {@link SqlEntity} to a {@link DatastoreEntity}.
+   *
+   * <p>One example is that this would convert a {@link DomainHistory} to a {@link HistoryEntry}.
+   */
+  private static Object toDatastoreEntity(@Nullable Object obj) {
+    if (obj instanceof SqlEntity) {
+      Optional<DatastoreEntity> possibleDatastoreEntity = ((SqlEntity) obj).toDatastoreEntity();
+      if (possibleDatastoreEntity.isPresent()) {
+        return possibleDatastoreEntity.get();
+      }
     }
     return obj;
   }
+
+  /** Converts many possible {@link SqlEntity} objects to {@link DatastoreEntity} objects. */
+  private static ImmutableList<Object> toDatastoreEntities(ImmutableCollection<?> collection) {
+    return collection.stream()
+        .map(DatastoreTransactionManager::toDatastoreEntity)
+        .collect(toImmutableList());
+  }
+
+  /**
+   * Converts an object to the corresponding {@link SqlEntity} if necessary and possible.
+   *
+   * <p>This should be used when returning objects from Datastore to make sure they reflect the most
+   * recent type of the object in question.
+   */
+  @SuppressWarnings("unchecked")
+  public static <T> T toSqlEntity(@Nullable T obj) {
+    // NB: The Key of the object in question may not necessarily be the resulting class that we
+    // wish to have. For example, because all *History classes are @EntitySubclasses, their Keys
+    // will have type HistoryEntry -- even if you create them based off the *History class.
+    if (obj instanceof DatastoreEntity && !(obj instanceof SqlEntity)) {
+      Optional<SqlEntity> possibleSqlEntity = ((DatastoreEntity) obj).toSqlEntity();
+      if (possibleSqlEntity.isPresent()) {
+        return (T) possibleSqlEntity.get();
+      }
+    }
+    return obj;
+  }
+
+  private static class DatastoreQueryComposerImpl<T> extends QueryComposer<T> {
+    DatastoreQueryComposerImpl(Class<T> entityClass) {
+      super(entityClass);
+    }
+
+    Query<T> buildQuery() {
+      Query<T> result = ofy().load().type(entityClass);
+      for (WhereClause pred : predicates) {
+        result = result.filter(pred.fieldName + pred.comparator.getDatastoreString(), pred.value);
+      }
+
+      if (orderBy != null) {
+        result = result.order(orderBy);
+      }
+
+      return result;
+    }
+
+    @Override
+    public Optional<T> first() {
+      return Optional.ofNullable(buildQuery().first().now());
+    }
+
+    @Override
+    public T getSingleResult() {
+      List<T> results = buildQuery().limit(2).list();
+      if (results.size() == 0) {
+        // The exception text here is the same as what we get for JPA queries.
+        throw new NoResultException("No entity found for query");
+      } else if (results.size() > 1) {
+        throw new NonUniqueResultException("More than one result found for getSingleResult query");
+      }
+      return results.get(0);
+    }
+
+    @Override
+    public Stream<T> stream() {
+      return Streams.stream(buildQuery());
+    }
+
+    @Override
+    public long count() {
+      return buildQuery().count();
+    }
+
+    @Override
+    public List<T> list() {
+      return buildQuery().list();
+    }
+  }
 }

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

@@ -103,6 +103,15 @@ public final class RdeRevision extends BackupGroupRoot implements NonReplicatedE
     return revisionOptional.map(rdeRevision -> rdeRevision.revision + 1).orElse(0);
   }
 
+  /** Returns the latest revision of the report already generated for the given triplet. */
+  public static Optional<Integer> getCurrentRevision(String tld, DateTime date, RdeMode mode) {
+    int nextRevision = getNextRevision(tld, date, mode);
+    if (nextRevision == 0) {
+      return Optional.empty();
+    }
+    return Optional.of(nextRevision - 1);
+  }
+
   /**
    * Sets the revision ID for a given triplet.
    *

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

@@ -106,6 +106,17 @@ public final class Registries {
     return tld;
   }
 
+  /**
+   * Pass-through check that the TLD exists, otherwise throw using the given error format message.
+   *
+   * <p>The specified TLD will be passed to the format message string.
+   */
+  public static String assertTldExists(String tld, String fmtMessage) {
+    String message = String.format(fmtMessage, tld);
+    checkArgument(getTlds().contains(checkArgumentNotNull(emptyToNull(tld), message)), message);
+    return tld;
+  }
+
   /** Pass-through check that every TLD in the given iterable exists, otherwise throw an IAE. */
   public static Iterable<String> assertTldsExist(Iterable<String> tlds) {
     for (String tld : tlds) {

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

@@ -14,12 +14,11 @@
 
 package google.registry.model.registry.label;
 
+import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static google.registry.model.DatabaseMigrationUtils.suppressExceptionUnlessInTest;
 
 import com.google.common.collect.Streams;
-import google.registry.model.DatabaseMigrationUtils;
-import google.registry.model.common.DatabaseTransitionSchedule.TransitionId;
 import google.registry.model.registry.Registry;
 import google.registry.model.registry.label.PremiumList.PremiumListEntry;
 import google.registry.schema.tld.PremiumListSqlDao;
@@ -32,27 +31,20 @@ import org.joda.money.Money;
 /**
  * DAO for {@link PremiumList} objects that handles the branching paths for SQL and Datastore.
  *
- * <p>For write actions, this class will perform the action against the primary database then, after
- * that success or failure, against the secondary database. If the secondary database fails, an
- * error is logged (but not thrown).
+ * <p>For write actions, this class will perform the action against Cloud SQL then, after that
+ * success or failure, against Datastore. If Datastore fails, an error is logged (but not thrown).
  *
  * <p>For read actions, when retrieving a price, we will log if the primary and secondary databases
- * have different values (or if the retrieval from the second database fails).
- *
- * <p>TODO (gbrodman): Change the isOfy() calls to the runtime selection of DBs when available
+ * have different values (or if the retrieval from Datastore fails).
  */
 public class PremiumListDualDao {
 
   /**
-   * Retrieves from the appropriate DB and returns the most recent premium list with the given name,
-   * or absent if no such list exists.
+   * Retrieves from Cloud SQL and returns the most recent premium list with the given name, or
+   * absent if no such list exists.
    */
   public static Optional<PremiumList> getLatestRevision(String premiumListName) {
-    if (DatabaseMigrationUtils.isDatastore(TransitionId.DOMAIN_LABEL_LISTS)) {
-      return PremiumListDatastoreDao.getLatestRevision(premiumListName);
-    } else {
       return PremiumListSqlDao.getLatestRevision(premiumListName);
-    }
   }
 
   /**
@@ -61,7 +53,7 @@ public class PremiumListDualDao {
    * <p>Returns absent if the label is not premium or there is no premium list for this registry.
    *
    * <p>Retrieves the price from both primary and secondary databases, and logs in the event of a
-   * failure in the secondary (but does not throw an exception).
+   * failure in Datastore (but does not throw an exception).
    */
   public static Optional<Money> getPremiumPrice(String label, Registry registry) {
     if (registry.getPremiumList() == null) {
@@ -69,98 +61,56 @@ public class PremiumListDualDao {
     }
     String premiumListName = registry.getPremiumList().getName();
     Optional<Money> primaryResult;
-    if (DatabaseMigrationUtils.isDatastore(TransitionId.DOMAIN_LABEL_LISTS)) {
-      primaryResult =
-          PremiumListDatastoreDao.getPremiumPrice(premiumListName, label, registry.getTldStr());
-    } else {
       primaryResult = PremiumListSqlDao.getPremiumPrice(premiumListName, label);
-    }
-    // Also load the value from the secondary DB, compare the two results, and log if different.
-    if (DatabaseMigrationUtils.isDatastore(TransitionId.DOMAIN_LABEL_LISTS)) {
-      suppressExceptionUnlessInTest(
-          () -> {
-            Optional<Money> secondaryResult =
-                PremiumListSqlDao.getPremiumPrice(premiumListName, label);
-            if (!primaryResult.equals(secondaryResult)) {
-              throw new IllegalStateException(
-                  String.format(
-                      "Unequal prices for domain %s.%s from primary Datastore DB (%s) and "
-                          + "secondary SQL db (%s).",
-                      label, registry.getTldStr(), primaryResult, secondaryResult));
-            }
-          },
-          String.format(
-              "Error loading price of domain %s.%s from Cloud SQL.", label, registry.getTldStr()));
-    } else {
-      suppressExceptionUnlessInTest(
-          () -> {
-            Optional<Money> secondaryResult =
-                PremiumListDatastoreDao.getPremiumPrice(
-                    premiumListName, label, registry.getTldStr());
-            if (!primaryResult.equals(secondaryResult)) {
-              throw new IllegalStateException(
-                  String.format(
-                      "Unequal prices for domain %s.%s from primary SQL DB (%s) and secondary "
-                          + "Datastore db (%s).",
-                      label, registry.getTldStr(), primaryResult, secondaryResult));
-            }
-          },
-          String.format(
-              "Error loading price of domain %s.%s from Datastore.", label, registry.getTldStr()));
-    }
+    // Also load the value from Datastore, compare the two results, and log if different.
+    suppressExceptionUnlessInTest(
+        () -> {
+          Optional<Money> secondaryResult =
+              PremiumListDatastoreDao.getPremiumPrice(premiumListName, label, registry.getTldStr());
+          checkState(
+              primaryResult.equals(secondaryResult),
+              "Unequal prices for domain %s.%s from primary SQL DB (%s) and secondary Datastore db"
+                  + " (%s).",
+              label,
+              registry.getTldStr(),
+              primaryResult,
+              secondaryResult);
+        },
+        String.format(
+            "Error loading price of domain %s.%s from Datastore.", label, registry.getTldStr()));
     return primaryResult;
   }
 
   /**
    * Saves the given list data to both primary and secondary databases.
    *
-   * <p>Logs but doesn't throw an exception in the event of a failure when writing to the secondary
-   * database.
+   * <p>Logs but doesn't throw an exception in the event of a failure when writing to Datastore.
    */
   public static PremiumList save(String name, List<String> inputData) {
-    PremiumList result;
-    if (DatabaseMigrationUtils.isDatastore(TransitionId.DOMAIN_LABEL_LISTS)) {
-      result = PremiumListDatastoreDao.save(name, inputData);
-      suppressExceptionUnlessInTest(
-          () -> PremiumListSqlDao.save(name, inputData), "Error when saving premium list to SQL.");
-    } else {
-      result = PremiumListSqlDao.save(name, inputData);
+    PremiumList result = PremiumListSqlDao.save(name, inputData);
       suppressExceptionUnlessInTest(
           () -> PremiumListDatastoreDao.save(name, inputData),
           "Error when saving premium list to Datastore.");
-    }
     return result;
   }
 
   /**
    * Deletes the premium list.
    *
-   * <p>Logs but doesn't throw an exception in the event of a failure when deleting from the
-   * secondary database.
+   * <p>Logs but doesn't throw an exception in the event of a failure when deleting from Datastore.
    */
   public static void delete(PremiumList premiumList) {
-    if (DatabaseMigrationUtils.isDatastore(TransitionId.DOMAIN_LABEL_LISTS)) {
-      PremiumListDatastoreDao.delete(premiumList);
-      suppressExceptionUnlessInTest(
-          () -> PremiumListSqlDao.delete(premiumList),
-          "Error when deleting premium list from SQL.");
-    } else {
       PremiumListSqlDao.delete(premiumList);
       suppressExceptionUnlessInTest(
           () -> PremiumListDatastoreDao.delete(premiumList),
           "Error when deleting premium list from Datastore.");
-    }
   }
 
   /** Returns whether or not there exists a premium list with the given name. */
   public static boolean exists(String premiumListName) {
     // It may seem like overkill, but loading the list has ways been the way we check existence and
     // given that we usually load the list around the time we check existence, we'll hit the cache
-    if (DatabaseMigrationUtils.isDatastore(TransitionId.DOMAIN_LABEL_LISTS)) {
-      return PremiumListDatastoreDao.getLatestRevision(premiumListName).isPresent();
-    } else {
-      return PremiumListSqlDao.getLatestRevision(premiumListName).isPresent();
-    }
+    return PremiumListSqlDao.getLatestRevision(premiumListName).isPresent();
   }
 
   /**
@@ -175,9 +125,6 @@ public class PremiumListDualDao {
                 () ->
                     new IllegalArgumentException(
                         String.format("No premium list with name %s.", premiumListName)));
-    if (DatabaseMigrationUtils.isDatastore(TransitionId.DOMAIN_LABEL_LISTS)) {
-      return PremiumListDatastoreDao.loadPremiumListEntriesUncached(premiumList);
-    } else {
       CurrencyUnit currencyUnit = premiumList.getCurrency();
       return Streams.stream(PremiumListSqlDao.loadPremiumListEntriesUncached(premiumList))
           .map(
@@ -188,7 +135,6 @@ public class PremiumListDualDao {
                       .build())
           .collect(toImmutableList());
     }
-  }
 
   private PremiumListDualDao() {}
 }

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

@@ -15,13 +15,11 @@
 package google.registry.model.registry.label;
 
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
-import static google.registry.model.DatabaseMigrationUtils.isDatastore;
 
 import com.google.common.collect.MapDifference;
 import com.google.common.collect.MapDifference.ValueDifference;
 import com.google.common.collect.Maps;
 import google.registry.model.DatabaseMigrationUtils;
-import google.registry.model.common.DatabaseTransitionSchedule.TransitionId;
 import google.registry.model.registry.label.ReservedList.ReservedListEntry;
 import java.util.Map;
 import java.util.Optional;
@@ -38,32 +36,18 @@ public class ReservedListDualDatabaseDao {
 
   /** Persist a new reserved list to the database. */
   public static void save(ReservedList reservedList) {
-    if (isDatastore(TransitionId.DOMAIN_LABEL_LISTS)) {
-      ReservedListDatastoreDao.save(reservedList);
-      DatabaseMigrationUtils.suppressExceptionUnlessInTest(
-          () -> ReservedListSqlDao.save(reservedList),
-          "Error saving the reserved list to Cloud SQL.");
-    } else {
       ReservedListSqlDao.save(reservedList);
       DatabaseMigrationUtils.suppressExceptionUnlessInTest(
           () -> ReservedListDatastoreDao.save(reservedList),
           "Error saving the reserved list to Datastore.");
-    }
   }
 
   /** Delete a reserved list from both databases. */
   public static void delete(ReservedList reservedList) {
-    if (isDatastore(TransitionId.DOMAIN_LABEL_LISTS)) {
-      ReservedListDatastoreDao.delete(reservedList);
-      DatabaseMigrationUtils.suppressExceptionUnlessInTest(
-          () -> ReservedListSqlDao.delete(reservedList),
-          "Error deleting the reserved list from Cloud SQL.");
-    } else {
       ReservedListSqlDao.delete(reservedList);
       DatabaseMigrationUtils.suppressExceptionUnlessInTest(
           () -> ReservedListDatastoreDao.delete(reservedList),
           "Error deleting the reserved list from Datastore.");
-    }
   }
 
   /**
@@ -72,9 +56,7 @@ public class ReservedListDualDatabaseDao {
    */
   public static Optional<ReservedList> getLatestRevision(String reservedListName) {
     Optional<ReservedList> maybePrimaryList =
-        isDatastore(TransitionId.DOMAIN_LABEL_LISTS)
-            ? ReservedListDatastoreDao.getLatestRevision(reservedListName)
-            : ReservedListSqlDao.getLatestRevision(reservedListName);
+        ReservedListSqlDao.getLatestRevision(reservedListName);
     DatabaseMigrationUtils.suppressExceptionUnlessInTest(
         () -> maybePrimaryList.ifPresent(primaryList -> loadAndCompare(primaryList)),
         "Error comparing reserved lists.");
@@ -83,14 +65,9 @@ public class ReservedListDualDatabaseDao {
 
   private static void loadAndCompare(ReservedList primaryList) {
     Optional<ReservedList> maybeSecondaryList =
-        isDatastore(TransitionId.DOMAIN_LABEL_LISTS)
-            ? ReservedListSqlDao.getLatestRevision(primaryList.getName())
-            : ReservedListDatastoreDao.getLatestRevision(primaryList.getName());
+        ReservedListDatastoreDao.getLatestRevision(primaryList.getName());
     if (!maybeSecondaryList.isPresent()) {
-      throw new IllegalStateException(
-          String.format(
-              "Reserved list in the secondary database (%s) is empty.",
-              isDatastore(TransitionId.DOMAIN_LABEL_LISTS) ? "Cloud SQL" : "Datastore"));
+      throw new IllegalStateException("Reserved list in Datastore is empty.");
     }
     Map<String, ReservedListEntry> labelsToReservations =
         primaryList.reservedListMap.entrySet().parallelStream()
@@ -110,12 +87,10 @@ public class ReservedListDualDatabaseDao {
         if (diff.entriesDiffering().size() > 10) {
         throw new IllegalStateException(
             String.format(
-                "Unequal reserved lists detected, %s list with revision"
+                "Unequal reserved lists detected, Datastore list with revision"
                     + " id %d has %d different records than the current"
-                    + " primary database list.",
-                isDatastore(TransitionId.DOMAIN_LABEL_LISTS) ? "Cloud SQL" : "Datastore",
-                secondaryList.getRevisionId(),
-                diff.entriesDiffering().size()));
+                    + " Cloud SQL list.",
+                secondaryList.getRevisionId(), diff.entriesDiffering().size()));
       }
           StringBuilder diffMessage = new StringBuilder("Unequal reserved lists detected:\n");
       diff.entriesDiffering().entrySet().stream()
@@ -125,12 +100,9 @@ public class ReservedListDualDatabaseDao {
                 ValueDifference<ReservedListEntry> valueDiff = entry.getValue();
                 diffMessage.append(
                     String.format(
-                        "Domain label %s has entry %s in %s and entry"
-                            + " %s in the secondary database.\n",
-                        label,
-                        valueDiff.leftValue(),
-                        isDatastore(TransitionId.DOMAIN_LABEL_LISTS) ? "Datastore" : "Cloud SQL",
-                        valueDiff.rightValue()));
+                        "Domain label %s has entry %s in Cloud SQL and entry"
+                            + " %s in the Datastore.\n",
+                        label, valueDiff.leftValue(), valueDiff.rightValue()));
               });
       diff.entriesOnlyOnLeft().entrySet().stream()
           .forEach(
@@ -138,9 +110,7 @@ public class ReservedListDualDatabaseDao {
                 String label = entry.getKey();
                 diffMessage.append(
                     String.format(
-                        "Domain label %s has entry in %s, but not in the secondary database.\n",
-                        label,
-                        isDatastore(TransitionId.DOMAIN_LABEL_LISTS) ? "Datastore" : "Cloud SQL"));
+                        "Domain label %s has entry in Cloud SQL, but not in Datastore.\n", label));
               });
       diff.entriesOnlyOnRight().entrySet().stream()
           .forEach(
@@ -148,9 +118,7 @@ public class ReservedListDualDatabaseDao {
                 String label = entry.getKey();
                 diffMessage.append(
                     String.format(
-                        "Domain label %s has entry in %s, but not in the primary database.\n",
-                        label,
-                        isDatastore(TransitionId.DOMAIN_LABEL_LISTS) ? "Cloud SQL" : "Datastore"));
+                        "Domain label %s has entry in Datastore, but not in Cloud SQL.\n", label));
               });
           throw new IllegalStateException(diffMessage.toString());
       }

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

@@ -377,12 +377,16 @@ public class HistoryEntry extends ImmutableObject implements Buildable, Datastor
       return thisCastToDerived();
     }
 
+    public B copyFrom(HistoryEntry.Builder<? extends HistoryEntry, ?> builder) {
+      return copyFrom(builder.getInstance());
+    }
+
     @Override
     public T build() {
       return super.build();
     }
 
-    public B setId(long id) {
+    public B setId(Long id) {
       getInstance().id = id;
       return thisCastToDerived();
     }

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

@@ -22,7 +22,7 @@ import static google.registry.util.DateTimeUtils.END_OF_TIME;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
 
 import com.google.common.collect.ImmutableCollection;
-import com.google.common.collect.Iterables;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Streams;
 import google.registry.model.EppResource;
 import google.registry.model.contact.ContactHistory;
@@ -34,6 +34,7 @@ import google.registry.model.host.HostResource;
 import google.registry.persistence.VKey;
 import google.registry.persistence.transaction.CriteriaQueryBuilder;
 import java.util.Comparator;
+import java.util.List;
 import java.util.stream.Stream;
 import javax.persistence.criteria.CriteriaBuilder;
 import javax.persistence.criteria.CriteriaQuery;
@@ -48,43 +49,53 @@ import org.joda.time.DateTime;
 public class HistoryEntryDao {
 
   /** Loads all history objects in the times specified, including all types. */
-  public static Iterable<? extends HistoryEntry> loadAllHistoryObjects(
+  public static ImmutableList<? extends HistoryEntry> loadAllHistoryObjects(
       DateTime afterTime, DateTime beforeTime) {
     if (tm().isOfy()) {
-      return ofy()
-          .load()
-          .type(HistoryEntry.class)
-          .order("modificationTime")
-          .filter("modificationTime >=", afterTime)
-          .filter("modificationTime <=", beforeTime);
+      return Streams.stream(
+              ofy()
+                  .load()
+                  .type(HistoryEntry.class)
+                  .order("modificationTime")
+                  .filter("modificationTime >=", afterTime)
+                  .filter("modificationTime <=", beforeTime))
+          .map(HistoryEntry::toChildHistoryEntity)
+          .collect(toImmutableList());
     } else {
       return jpaTm()
           .transact(
               () ->
-                  Iterables.concat(
-                      loadAllHistoryObjectsFromSql(ContactHistory.class, afterTime, beforeTime),
-                      loadAllHistoryObjectsFromSql(DomainHistory.class, afterTime, beforeTime),
-                      loadAllHistoryObjectsFromSql(HostHistory.class, afterTime, beforeTime)));
+                  new ImmutableList.Builder<HistoryEntry>()
+                      .addAll(
+                          loadAllHistoryObjectsFromSql(ContactHistory.class, afterTime, beforeTime))
+                      .addAll(
+                          loadAllHistoryObjectsFromSql(DomainHistory.class, afterTime, beforeTime))
+                      .addAll(
+                          loadAllHistoryObjectsFromSql(HostHistory.class, afterTime, beforeTime))
+                      .build());
     }
   }
 
   /** Loads all history objects corresponding to the given {@link EppResource}. */
-  public static Iterable<? extends HistoryEntry> loadHistoryObjectsForResource(
+  public static ImmutableList<? extends HistoryEntry> loadHistoryObjectsForResource(
       VKey<? extends EppResource> parentKey) {
     return loadHistoryObjectsForResource(parentKey, START_OF_TIME, END_OF_TIME);
   }
 
   /** Loads all history objects in the time period specified for the given {@link EppResource}. */
-  public static Iterable<? extends HistoryEntry> loadHistoryObjectsForResource(
+  public static ImmutableList<? extends HistoryEntry> loadHistoryObjectsForResource(
       VKey<? extends EppResource> parentKey, DateTime afterTime, DateTime beforeTime) {
     if (tm().isOfy()) {
-      return ofy()
-          .load()
-          .type(HistoryEntry.class)
-          .ancestor(parentKey.getOfyKey())
-          .order("modificationTime")
-          .filter("modificationTime >=", afterTime)
-          .filter("modificationTime <=", beforeTime);
+      return Streams.stream(
+              ofy()
+                  .load()
+                  .type(HistoryEntry.class)
+                  .ancestor(parentKey.getOfyKey())
+                  .order("modificationTime")
+                  .filter("modificationTime >=", afterTime)
+                  .filter("modificationTime <=", beforeTime))
+          .map(HistoryEntry::toChildHistoryEntity)
+          .collect(toImmutableList());
     } else {
       return jpaTm()
           .transact(() -> loadHistoryObjectsForResourceFromSql(parentKey, afterTime, beforeTime));
@@ -124,7 +135,7 @@ public class HistoryEntryDao {
         .getResultStream();
   }
 
-  private static Iterable<? extends HistoryEntry> loadHistoryObjectsForResourceFromSql(
+  private static ImmutableList<? extends HistoryEntry> loadHistoryObjectsForResourceFromSql(
       VKey<? extends EppResource> parentKey, DateTime afterTime, DateTime beforeTime) {
     // The class we're searching from is based on which parent type (e.g. Domain) we have
     Class<? extends HistoryEntry> historyClass = getHistoryClassFromParent(parentKey.getKind());
@@ -138,12 +149,9 @@ public class HistoryEntryDao {
             .where(repoIdFieldName, criteriaBuilder::equal, parentKey.getSqlKey().toString())
             .build();
 
-    return jpaTm()
-        .getEntityManager()
-        .createQuery(criteriaQuery)
-        .getResultStream()
-        .sorted(Comparator.comparing(HistoryEntry::getModificationTime))
-        .collect(toImmutableList());
+    return ImmutableList.sortedCopyOf(
+        Comparator.comparing(HistoryEntry::getModificationTime),
+        jpaTm().getEntityManager().createQuery(criteriaQuery).getResultList());
   }
 
   private static Class<? extends HistoryEntry> getHistoryClassFromParent(
@@ -166,7 +174,7 @@ public class HistoryEntryDao {
         : historyClass.equals(DomainHistory.class) ? "domainRepoId" : "hostRepoId";
   }
 
-  private static Iterable<? extends HistoryEntry> loadAllHistoryObjectsFromSql(
+  private static List<? extends HistoryEntry> loadAllHistoryObjectsFromSql(
       Class<? extends HistoryEntry> historyClass, DateTime afterTime, DateTime beforeTime) {
     CriteriaBuilder criteriaBuilder = jpaTm().getEntityManager().getCriteriaBuilder();
     return jpaTm()

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

@@ -16,34 +16,24 @@ package google.registry.model.smd;
 
 import static com.google.common.base.Preconditions.checkNotNull;
 import static google.registry.model.CacheUtils.memoizeWithShortExpiration;
-import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
 import static google.registry.util.DateTimeUtils.isBeforeOrAt;
 
-import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Supplier;
 import com.google.common.collect.ImmutableMap;
-import com.googlecode.objectify.Key;
-import com.googlecode.objectify.annotation.EmbedMap;
-import com.googlecode.objectify.annotation.Entity;
-import com.googlecode.objectify.annotation.Id;
-import com.googlecode.objectify.annotation.Ignore;
-import com.googlecode.objectify.annotation.OnSave;
-import com.googlecode.objectify.annotation.Parent;
 import google.registry.model.ImmutableObject;
-import google.registry.model.annotations.InCrossTld;
-import google.registry.model.annotations.NotBackedUp;
-import google.registry.model.annotations.NotBackedUp.Reason;
-import google.registry.model.common.EntityGroupRoot;
-import google.registry.schema.replay.NonReplicatedEntity;
+import google.registry.schema.replay.DatastoreEntity;
+import google.registry.schema.replay.SqlEntity;
 import java.util.Map;
+import java.util.Optional;
 import javax.persistence.CollectionTable;
 import javax.persistence.Column;
 import javax.persistence.ElementCollection;
+import javax.persistence.Entity;
 import javax.persistence.GeneratedValue;
 import javax.persistence.GenerationType;
+import javax.persistence.Id;
 import javax.persistence.JoinColumn;
 import javax.persistence.MapKeyColumn;
-import javax.persistence.Transient;
 import org.joda.time.DateTime;
 
 /**
@@ -53,34 +43,14 @@ import org.joda.time.DateTime;
  * all the {@link SignedMark SignedMarks} that have been revoked. A new list is created for each new
  * file that's created, depending on the timestamp.
  *
- * <p>We'll be putting the entire table into a single entity for the sake of performance. But in
- * order to avoid exceeding the one megabyte max entity size limit, we'll also be sharding that
- * entity into multiple entities, each entity containing {@value #SHARD_SIZE} rows.
- *
- * <p>TODO: We can remove the sharding once we have converted entirely to Cloud SQL storage during
- * the Registry 3.0 migration. Then, the entire table will be stored conceptually as one entity (in
- * fact in SignedMarkRevocationList and SignedMarkRevocationEntry tables).
- *
  * @see google.registry.tmch.SmdrlCsvParser
  * @see <a href="http://tools.ietf.org/html/draft-lozano-tmch-func-spec-08#section-6.2">TMCH
  *     functional specifications - SMD Revocation List</a>
  */
 @Entity
-@javax.persistence.Entity
-@NotBackedUp(reason = Reason.EXTERNALLY_SOURCED)
-@InCrossTld
-public class SignedMarkRevocationList extends ImmutableObject implements NonReplicatedEntity {
+public class SignedMarkRevocationList extends ImmutableObject implements SqlEntity {
 
-  @VisibleForTesting static final int SHARD_SIZE = 10000;
-
-  /** Common ancestor for queries. */
-  @Parent @Transient Key<EntityGroupRoot> parent = getCrossTldKey();
-
-  /** ID for the sharded entity. */
-  @Id @Transient long id;
-
-  @Ignore
-  @javax.persistence.Id
+  @Id
   @GeneratedValue(strategy = GenerationType.IDENTITY)
   Long revisionId;
 
@@ -88,7 +58,6 @@ public class SignedMarkRevocationList extends ImmutableObject implements NonRepl
   DateTime creationTime;
 
   /** A map from SMD IDs to revocation time. */
-  @EmbedMap
   @ElementCollection
   @CollectionTable(
       name = "SignedMarkRevocationEntry",
@@ -97,17 +66,10 @@ public class SignedMarkRevocationList extends ImmutableObject implements NonRepl
   @Column(name = "revocationTime", nullable = false)
   Map</*@MatchesPattern("[0-9]+-[0-9]+")*/ String, DateTime> revokes;
 
-  /** Indicates that this is a shard rather than a "full" list. */
-  @Ignore @Transient boolean isShard;
-
-  /**
-   * A cached supplier that fetches the SMDRL shards from Datastore and recombines them into a
-   * single {@link SignedMarkRevocationList} object.
-   */
+  /** A cached supplier that fetches the {@link SignedMarkRevocationList} object. */
   private static final Supplier<SignedMarkRevocationList> CACHE =
       memoizeWithShortExpiration(SignedMarkRevocationListDao::load);
 
-  /** Return a single logical instance that combines all Datastore shards. */
   public static SignedMarkRevocationList get() {
     return CACHE.get();
   }
@@ -137,20 +99,14 @@ public class SignedMarkRevocationList extends ImmutableObject implements NonRepl
     return revokes.size();
   }
 
-  /** Save this list to Datastore in sharded form and to Cloud SQL. Returns {@code this}. */
+  /** Save this list to Cloud SQL. Returns {@code this}. */
   public SignedMarkRevocationList save() {
     SignedMarkRevocationListDao.save(this);
     return this;
   }
 
-  /** As a safety mechanism, fail if someone tries to save this class directly. */
-  @OnSave
-  void disallowUnshardedSaves() {
-    if (!isShard) {
-      throw new UnshardedSaveException();
-    }
+  @Override
+  public Optional<DatastoreEntity> toDatastoreEntity() {
+    return Optional.empty(); // Not persisted in Datastore
   }
-
-  /** Exception when trying to directly save a {@link SignedMarkRevocationList} without sharding. */
-  public static class UnshardedSaveException extends RuntimeException {}
 }

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

@@ -14,221 +14,44 @@
 
 package google.registry.model.smd;
 
-import static com.google.common.base.Preconditions.checkNotNull;
-import static com.google.common.base.Preconditions.checkState;
-import static com.google.common.collect.ImmutableList.toImmutableList;
-import static com.google.common.collect.Iterables.isEmpty;
-import static google.registry.model.DatabaseMigrationUtils.suppressExceptionUnlessInTest;
-import static google.registry.model.common.DatabaseTransitionSchedule.PrimaryDatabase.DATASTORE;
-import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
-import static google.registry.model.ofy.ObjectifyService.allocateId;
-import static google.registry.model.ofy.ObjectifyService.ofy;
-import static google.registry.model.smd.SignedMarkRevocationList.SHARD_SIZE;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
-import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
-import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
-import static google.registry.util.CollectionUtils.isNullOrEmpty;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
 
 import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.Iterables;
-import com.google.common.collect.MapDifference;
-import com.google.common.collect.Maps;
 import com.google.common.flogger.FluentLogger;
-import google.registry.model.DatabaseMigrationUtils;
-import google.registry.model.common.DatabaseTransitionSchedule.PrimaryDatabase;
-import google.registry.model.common.DatabaseTransitionSchedule.TransitionId;
-import google.registry.util.CollectionUtils;
-import java.util.Map;
 import java.util.Optional;
-import org.joda.time.DateTime;
 
 public class SignedMarkRevocationListDao {
 
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
 
-  /**
-   * Loads the {@link SignedMarkRevocationList}.
-   *
-   * <p>Loads the list from the specified primary database, and attempts to load from the secondary
-   * database. If the load the secondary database fails, or the list from the secondary database
-   * does not match the list from the primary database, the error will be logged but no exception
-   * will be thrown.
-   */
+  /** Loads the {@link SignedMarkRevocationList}. */
   static SignedMarkRevocationList load() {
-    PrimaryDatabase primaryDatabase =
-        tm().transactNew(
-                () ->
-                    DatabaseMigrationUtils.getPrimaryDatabase(
-                        TransitionId.SIGNED_MARK_REVOCATION_LIST));
-    Optional<SignedMarkRevocationList> primaryList =
-        primaryDatabase.equals(DATASTORE) ? loadFromDatastore() : loadFromCloudSql();
-    if (!primaryList.isPresent()) {
-      throw new IllegalStateException(
-          String.format(
-              "SignedMarkRevocationList not found in the primary database (%s).",
-              primaryDatabase.name()));
-    }
-    suppressExceptionUnlessInTest(
-        () -> loadAndCompare(primaryDatabase, primaryList.get()),
-        String.format(
-            "Error loading and comparing the SignedMarkRevocationList from the secondary database"
-                + " (%s).",
-            primaryDatabase.equals(DATASTORE) ? "Cloud SQL" : "Datastore"));
-    return primaryList.get();
+    Optional<SignedMarkRevocationList> smdrl =
+        jpaTm()
+            .transact(
+                () -> {
+                  Long revisionId =
+                      jpaTm()
+                          .query("SELECT MAX(revisionId) FROM SignedMarkRevocationList", Long.class)
+                          .getSingleResult();
+                  return jpaTm()
+                      .query(
+                          "FROM SignedMarkRevocationList smrl LEFT JOIN FETCH smrl.revokes "
+                              + "WHERE smrl.revisionId = :revisionId",
+                          SignedMarkRevocationList.class)
+                      .setParameter("revisionId", revisionId)
+                      .getResultStream()
+                      .findFirst();
+                });
+    return smdrl.orElseGet(() -> SignedMarkRevocationList.create(START_OF_TIME, ImmutableMap.of()));
   }
 
-  /**
-   * Loads the list from the secondary database and compares it to the list from the primary
-   * database.
-   */
-  private static void loadAndCompare(
-      PrimaryDatabase primaryDatabase, SignedMarkRevocationList primaryList) {
-    Optional<SignedMarkRevocationList> secondaryList =
-        primaryDatabase.equals(DATASTORE) ? loadFromCloudSql() : loadFromDatastore();
-    if (secondaryList.isPresent() && !isNullOrEmpty(secondaryList.get().revokes)) {
-      MapDifference<String, DateTime> diff =
-          Maps.difference(primaryList.revokes, secondaryList.get().revokes);
-      if (!diff.areEqual()) {
-        if (diff.entriesDiffering().size() > 10) {
-          String message =
-              String.format(
-                  "Unequal SignedMarkRevocationList detected, %s list with revision id"
-                      + " %d has %d different records than the current primary database list.",
-                  primaryDatabase.equals(DATASTORE) ? "Cloud SQL" : "Datastore",
-                  secondaryList.get().revisionId,
-                  diff.entriesDiffering().size());
-          throw new IllegalStateException(message);
-        } else {
-          StringBuilder diffMessage =
-              new StringBuilder("Unequal SignedMarkRevocationList detected:\n");
-          diff.entriesDiffering()
-              .forEach(
-                  (label, valueDiff) ->
-                      diffMessage.append(
-                          String.format(
-                              "SMD %s has key %s in %s and key %s in secondary database.\n",
-                              label,
-                              valueDiff.leftValue(),
-                              primaryDatabase.name(),
-                              valueDiff.rightValue())));
-          throw new IllegalStateException(diffMessage.toString());
-        }
-      }
-    } else {
-      if (primaryList.size() != 0) {
-        throw new IllegalStateException(
-            String.format(
-                "SignedMarkRevocationList in %s is empty while it is not empty in the primary"
-                    + " database.",
-                primaryDatabase.equals(DATASTORE) ? "Cloud SQL" : "Datastore"));
-      }
-    }
-  }
-
-  /** Loads the shards from Datastore and combines them into one list. */
-  private static Optional<SignedMarkRevocationList> loadFromDatastore() {
-    return ofyTm()
-        .transactNewReadOnly(
-            () -> {
-              Iterable<SignedMarkRevocationList> shards =
-                  ofy().load().type(SignedMarkRevocationList.class).ancestor(getCrossTldKey());
-              DateTime creationTime =
-                  isEmpty(shards)
-                      ? START_OF_TIME
-                      : checkNotNull(Iterables.get(shards, 0).creationTime, "creationTime");
-              ImmutableMap.Builder<String, DateTime> revokes = new ImmutableMap.Builder<>();
-              for (SignedMarkRevocationList shard : shards) {
-                revokes.putAll(shard.revokes);
-                checkState(
-                    creationTime.equals(shard.creationTime),
-                    "Inconsistent creation times in Datastore shard: %s vs. %s",
-                    creationTime,
-                    shard.creationTime);
-              }
-              return Optional.of(SignedMarkRevocationList.create(creationTime, revokes.build()));
-            });
-  }
-
-  private static Optional<SignedMarkRevocationList> loadFromCloudSql() {
-    return jpaTm()
-        .transact(
-            () -> {
-              Long revisionId =
-                  jpaTm()
-                      .query("SELECT MAX(revisionId) FROM SignedMarkRevocationList", Long.class)
-                      .getSingleResult();
-              return jpaTm()
-                  .query(
-                      "FROM SignedMarkRevocationList smrl LEFT JOIN FETCH smrl.revokes "
-                          + "WHERE smrl.revisionId = :revisionId",
-                      SignedMarkRevocationList.class)
-                  .setParameter("revisionId", revisionId)
-                  .getResultStream()
-                  .findFirst();
-            });
-  }
-
-  /**
-   * Save the given {@link SignedMarkRevocationList}
-   *
-   * <p>Saves the list to the specified primary database, and attempts to save to the secondary
-   * database. If the save to the secondary database fails, the error will be logged but no
-   * exception will be thrown.
-   */
+  /** Save the given {@link SignedMarkRevocationList} */
   static void save(SignedMarkRevocationList signedMarkRevocationList) {
-    PrimaryDatabase primaryDatabase =
-        tm().transactNew(
-                () ->
-                    DatabaseMigrationUtils.getPrimaryDatabase(
-                        TransitionId.SIGNED_MARK_REVOCATION_LIST));
-    if (primaryDatabase.equals(DATASTORE)) {
-      saveToDatastore(signedMarkRevocationList.revokes, signedMarkRevocationList.creationTime);
-      suppressExceptionUnlessInTest(
-          () -> SignedMarkRevocationListDao.saveToCloudSql(signedMarkRevocationList),
-          "Error inserting signed mark revocations into secondary database (Cloud SQL).");
-    } else {
-      SignedMarkRevocationListDao.saveToCloudSql(signedMarkRevocationList);
-      suppressExceptionUnlessInTest(
-          () ->
-              saveToDatastore(
-                  signedMarkRevocationList.revokes, signedMarkRevocationList.creationTime),
-          "Error inserting signed mark revocations into secondary database (Datastore).");
-    }
-  }
-
-  private static void saveToCloudSql(SignedMarkRevocationList signedMarkRevocationList) {
     jpaTm().transact(() -> jpaTm().insert(signedMarkRevocationList));
     logger.atInfo().log(
         "Inserted %,d signed mark revocations into Cloud SQL.",
         signedMarkRevocationList.revokes.size());
   }
-
-  private static void saveToDatastore(Map<String, DateTime> revokes, DateTime creationTime) {
-    tm().transact(
-            () -> {
-              ofy()
-                  .deleteWithoutBackup()
-                  .keys(
-                      ofy()
-                          .load()
-                          .type(SignedMarkRevocationList.class)
-                          .ancestor(getCrossTldKey())
-                          .keys());
-              ofy()
-                  .saveWithoutBackup()
-                  .entities(
-                      CollectionUtils.partitionMap(revokes, SHARD_SIZE).stream()
-                          .map(
-                              shardRevokes -> {
-                                SignedMarkRevocationList shard =
-                                    SignedMarkRevocationList.create(creationTime, shardRevokes);
-                                shard.id = allocateId();
-                                shard.isShard =
-                                    true; // Avoid the exception in disallowUnshardedSaves().
-                                return shard;
-                              })
-                          .collect(toImmutableList()));
-            });
-  }
 }

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

@@ -164,6 +164,7 @@ public class DomainTransferData extends TransferData<DomainTransferData.Builder>
       serverApproveEntities = null;
       postLoad();
     }
+    hashCode = null; // reset the hash code since we may have changed the entities
   }
 
   /**
@@ -271,6 +272,7 @@ public class DomainTransferData extends TransferData<DomainTransferData.Builder>
       serverApproveEntitiesBuilder.add(billingCancellationId);
     }
     serverApproveEntities = forceEmptyToNull(serverApproveEntitiesBuilder.build());
+    hashCode = null; // reset the hash code since we may have changed the entities
   }
 
   @Override

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

@@ -47,7 +47,6 @@ import java.sql.DriverManager;
 import java.sql.SQLException;
 import java.util.HashMap;
 import java.util.Map;
-import java.util.Objects;
 import java.util.Optional;
 import java.util.Properties;
 import java.util.function.Supplier;
@@ -260,55 +259,6 @@ public abstract class PersistenceModule {
     return new JpaTransactionManagerImpl(create(overrides), clock);
   }
 
-  @Provides
-  @Singleton
-  @SocketFactoryJpaTm
-  static JpaTransactionManager provideSocketFactoryJpaTm(
-      SqlCredentialStore credentialStore,
-      @Config("beamCloudSqlUsername") String username,
-      @Config("beamCloudSqlPassword") String password,
-      @Config("beamHibernateHikariMaximumPoolSize") int hikariMaximumPoolSize,
-      @BeamPipelineCloudSqlConfigs ImmutableMap<String, String> cloudSqlConfigs,
-      Clock clock) {
-    HashMap<String, String> overrides = Maps.newHashMap(cloudSqlConfigs);
-    overrides.put(HIKARI_MAXIMUM_POOL_SIZE, String.valueOf(hikariMaximumPoolSize));
-    overrides.put(Environment.USER, username);
-    overrides.put(Environment.PASS, password);
-    // TODO(b/175700623): consider assigning different logins to pipelines
-    // TODO(b/179839014): Make SqlCredentialStore injectable in BEAM
-    // Note: the logs below appear in the pipeline's Worker logs, not the Job log.
-    try {
-      SqlCredential credential = credentialStore.getCredential(new RobotUser(RobotId.NOMULUS));
-      if (!Objects.equals(username, credential.login())) {
-        logger.atWarning().log(
-            "Wrong username for nomulus. Expecting %s, found %s.", username, credential.login());
-      } else if (!Objects.equals(password, credential.password())) {
-        logger.atWarning().log("Wrong password for nomulus.");
-      } else {
-        logger.atWarning().log("Credentials in the kerying and the secret manager match.");
-      }
-    } catch (Exception e) {
-      logger.atWarning().withCause(e).log("Failed to get SQL credential from Secret Manager.");
-    }
-    return new JpaTransactionManagerImpl(create(overrides), clock);
-  }
-
-  @Provides
-  @Singleton
-  @JdbcJpaTm
-  static JpaTransactionManager provideLocalJpaTm(
-      @Config("beamCloudSqlJdbcUrl") String jdbcUrl,
-      @Config("beamCloudSqlUsername") String username,
-      @Config("beamCloudSqlPassword") String password,
-      @DefaultHibernateConfigs ImmutableMap<String, String> defaultConfigs,
-      Clock clock) {
-    HashMap<String, String> overrides = Maps.newHashMap(defaultConfigs);
-    overrides.put(Environment.URL, jdbcUrl);
-    overrides.put(Environment.USER, username);
-    overrides.put(Environment.PASS, password);
-    return new JpaTransactionManagerImpl(create(overrides), clock);
-  }
-
   /** Constructs the {@link EntityManagerFactory} instance. */
   @VisibleForTesting
   static EntityManagerFactory create(
@@ -399,23 +349,6 @@ public abstract class PersistenceModule {
   @Documented
   public @interface NomulusToolJpaTm {}
 
-  /**
-   * Dagger qualifier for {@link JpaTransactionManager} that accesses Cloud SQL using socket
-   * factory. This is meant for applications not running on AppEngine, therefore without access to a
-   * {@link google.registry.keyring.api.Keyring}.
-   */
-  @Qualifier
-  @Documented
-  public @interface SocketFactoryJpaTm {}
-
-  /**
-   * Dagger qualifier for {@link JpaTransactionManager} backed by plain JDBC connections. This is
-   * mainly used by tests.
-   */
-  @Qualifier
-  @Documented
-  public @interface JdbcJpaTm {}
-
   /** Dagger qualifier for the partial Cloud SQL configs. */
   @Qualifier
   @Documented

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

@@ -18,6 +18,7 @@ import static google.registry.persistence.transaction.TransactionManagerFactory.
 
 import com.google.common.collect.ImmutableList;
 import java.util.Collection;
+import javax.persistence.EntityManager;
 import javax.persistence.criteria.CriteriaBuilder;
 import javax.persistence.criteria.CriteriaQuery;
 import javax.persistence.criteria.Expression;
@@ -35,22 +36,23 @@ import javax.persistence.criteria.Root;
 public class CriteriaQueryBuilder<T> {
 
   /** Functional interface that defines the 'where' operator, e.g. {@link CriteriaBuilder#equal}. */
-  public interface WhereClause<U> {
+  public interface WhereOperator<U> {
     Predicate predicate(Expression<U> expression, U object);
   }
 
   private final CriteriaQuery<T> query;
-  private final Root<T> root;
+  private final Root<?> root;
   private final ImmutableList.Builder<Predicate> predicates = new ImmutableList.Builder<>();
   private final ImmutableList.Builder<Order> orders = new ImmutableList.Builder<>();
 
-  private CriteriaQueryBuilder(CriteriaQuery<T> query, Root<T> root) {
+  private CriteriaQueryBuilder(CriteriaQuery<T> query, Root<?> root) {
     this.query = query;
     this.root = root;
   }
 
   /** Adds a WHERE clause to the query, given the specified operation, field, and value. */
-  public <V> CriteriaQueryBuilder<T> where(String fieldName, WhereClause<V> whereClause, V value) {
+  public <V> CriteriaQueryBuilder<T> where(
+      String fieldName, WhereOperator<V> whereClause, V value) {
     Expression<V> expression = root.get(fieldName);
     return where(whereClause.predicate(expression, value));
   }
@@ -94,9 +96,23 @@ public class CriteriaQueryBuilder<T> {
 
   /** Creates a query builder that will SELECT from the given class. */
   public static <T> CriteriaQueryBuilder<T> create(Class<T> clazz) {
-    CriteriaQuery<T> query = jpaTm().getEntityManager().getCriteriaBuilder().createQuery(clazz);
+    return create(jpaTm().getEntityManager(), clazz);
+  }
+
+  /** Creates a query builder for the given entity manager. */
+  public static <T> CriteriaQueryBuilder<T> create(EntityManager em, Class<T> clazz) {
+    CriteriaQuery<T> query = em.getCriteriaBuilder().createQuery(clazz);
     Root<T> root = query.from(clazz);
     query = query.select(root);
     return new CriteriaQueryBuilder<>(query, root);
   }
+
+  /** Creates a "count" query for the table for the class. */
+  public static <T> CriteriaQueryBuilder<Long> createCount(EntityManager em, Class<T> clazz) {
+    CriteriaBuilder builder = em.getCriteriaBuilder();
+    CriteriaQuery<Long> query = builder.createQuery(Long.class);
+    Root<T> root = query.from(clazz);
+    query = query.select(builder.count(root));
+    return new CriteriaQueryBuilder<>(query, root);
+  }
 }

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

@@ -18,7 +18,7 @@ import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
-import static google.registry.model.ofy.DatastoreTransactionManager.toChildHistoryEntryIfPossible;
+import static google.registry.model.ofy.DatastoreTransactionManager.toSqlEntity;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 import static java.util.AbstractMap.SimpleEntry;
 import static java.util.stream.Collectors.joining;
@@ -37,16 +37,19 @@ import google.registry.model.index.ForeignKeyIndex.ForeignKeyDomainIndex;
 import google.registry.model.index.ForeignKeyIndex.ForeignKeyHostIndex;
 import google.registry.model.ofy.DatastoreTransactionManager;
 import google.registry.model.server.KmsSecret;
+import google.registry.model.tmch.ClaimsListShard.ClaimsListSingleton;
 import google.registry.persistence.JpaRetries;
 import google.registry.persistence.VKey;
 import google.registry.util.Clock;
 import google.registry.util.Retrier;
 import google.registry.util.SystemSleeper;
 import java.lang.reflect.Field;
+import java.util.List;
 import java.util.Map;
 import java.util.NoSuchElementException;
 import java.util.Optional;
 import java.util.function.Supplier;
+import java.util.stream.Stream;
 import java.util.stream.StreamSupport;
 import javax.persistence.EntityManager;
 import javax.persistence.EntityManagerFactory;
@@ -71,6 +74,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   // TODO(b/176108270): Remove this property after database migration.
   private static final ImmutableSet<Class<? extends ImmutableObject>> IGNORED_ENTITY_CLASSES =
       ImmutableSet.of(
+          ClaimsListSingleton.class,
           EppResourceIndex.class,
           ForeignKeyContactIndex.class,
           ForeignKeyDomainIndex.class,
@@ -264,7 +268,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     }
     assertInTransaction();
     // Necessary due to the changes in HistoryEntry representation during the migration to SQL
-    Object toPersist = toChildHistoryEntryIfPossible(entity);
+    Object toPersist = toSqlEntity(entity);
     getEntityManager().persist(toPersist);
     transactionInfo.get().addUpdate(toPersist);
   }
@@ -294,7 +298,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     }
     assertInTransaction();
     // Necessary due to the changes in HistoryEntry representation during the migration to SQL
-    Object toPersist = toChildHistoryEntryIfPossible(entity);
+    Object toPersist = toSqlEntity(entity);
     getEntityManager().merge(toPersist);
     transactionInfo.get().addUpdate(toPersist);
   }
@@ -334,7 +338,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     assertInTransaction();
     checkArgument(exists(entity), "Given entity does not exist");
     // Necessary due to the changes in HistoryEntry representation during the migration to SQL
-    Object toPersist = toChildHistoryEntryIfPossible(entity);
+    Object toPersist = toSqlEntity(entity);
     getEntityManager().merge(toPersist);
     transactionInfo.get().addUpdate(toPersist);
   }
@@ -367,7 +371,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   @Override
   public boolean exists(Object entity) {
     checkArgumentNotNull(entity, "entity must be specified");
-    entity = toChildHistoryEntryIfPossible(entity);
+    entity = toSqlEntity(entity);
     EntityType<?> entityType = getEntityType(entity.getClass());
     ImmutableSet<EntityId> entityIds = getEntityIdsFromEntity(entityType, entity);
     return exists(entityType.getName(), entityIds);
@@ -410,7 +414,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   @Override
   public <T> ImmutableList<T> loadByEntitiesIfPresent(Iterable<T> entities) {
     return Streams.stream(entities)
-        .map(DatastoreTransactionManager::toChildHistoryEntryIfPossible)
+        .map(DatastoreTransactionManager::toSqlEntity)
         .filter(this::exists)
         .map(this::loadByEntity)
         .collect(toImmutableList());
@@ -445,9 +449,8 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   public <T> T loadByEntity(T entity) {
     checkArgumentNotNull(entity, "entity must be specified");
     assertInTransaction();
-    entity = toChildHistoryEntryIfPossible(entity);
     // If the caller requested a HistoryEntry, load the corresponding *History class
-    T possibleChild = toChildHistoryEntryIfPossible(entity);
+    T possibleChild = toSqlEntity(entity);
     return (T)
         loadByKey(
             VKey.createSql(
@@ -507,7 +510,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
       return;
     }
     assertInTransaction();
-    entity = toChildHistoryEntryIfPossible(entity);
+    entity = toSqlEntity(entity);
     Object managedEntity = entity;
     if (!getEntityManager().contains(entity)) {
       managedEntity = getEntityManager().merge(entity);
@@ -530,6 +533,11 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     delete(entity);
   }
 
+  @Override
+  public <T> QueryComposer<T> createQueryComposer(Class<T> entity) {
+    return new JpaQueryComposerImpl<T>(entity, getEntityManager());
+  }
+
   @Override
   public void clearSessionCache() {
     // This is an intended no-op method as there is no session cache in Postgresql.
@@ -681,4 +689,58 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
       }
     }
   }
+
+  private static class JpaQueryComposerImpl<T> extends QueryComposer<T> {
+
+    EntityManager em;
+
+    JpaQueryComposerImpl(Class<T> entityClass, EntityManager em) {
+      super(entityClass);
+      this.em = em;
+    }
+
+    private TypedQuery<T> buildQuery() {
+      CriteriaQueryBuilder<T> queryBuilder = CriteriaQueryBuilder.create(em, entityClass);
+      return addCriteria(queryBuilder);
+    }
+
+    private <U> TypedQuery<U> addCriteria(CriteriaQueryBuilder<U> queryBuilder) {
+      for (WhereClause<?> pred : predicates) {
+        pred.addToCriteriaQueryBuilder(queryBuilder);
+      }
+
+      if (orderBy != null) {
+        queryBuilder.orderByAsc(orderBy);
+      }
+
+      return em.createQuery(queryBuilder.build());
+    }
+
+    @Override
+    public Optional<T> first() {
+      List<T> results = buildQuery().setMaxResults(1).getResultList();
+      return results.size() > 0 ? Optional.of(results.get(0)) : Optional.empty();
+    }
+
+    @Override
+    public T getSingleResult() {
+      return buildQuery().getSingleResult();
+    }
+
+    @Override
+    public Stream<T> stream() {
+      return buildQuery().getResultStream();
+    }
+
+    @Override
+    public long count() {
+      CriteriaQueryBuilder<Long> queryBuilder = CriteriaQueryBuilder.createCount(em, entityClass);
+      return addCriteria(queryBuilder).getSingleResult();
+    }
+
+    @Override
+    public List<T> list() {
+      return buildQuery().getResultList();
+    }
+  }
 }

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

@@ -0,0 +1,196 @@
+// Copyright 2021 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.persistence.transaction;
+
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+
+import com.google.common.base.Function;
+import google.registry.persistence.transaction.CriteriaQueryBuilder.WhereOperator;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Optional;
+import java.util.stream.Stream;
+import javax.annotation.Nullable;
+import javax.persistence.criteria.CriteriaBuilder;
+
+/**
+ * Creates queries that can be used both for objectify and JPA.
+ *
+ * <p>Example usage:
+ *
+ * <pre>
+ *    tm().createQueryComposer(EntityType.class)
+ *        .where("fieldName", Comparator.EQ, "value")
+ *        .orderBy("fieldName")
+ *        .stream()
+ * </pre>
+ */
+public abstract class QueryComposer<T> {
+
+  // The class whose entities we're querying.  Note that this limits us to single table queries in
+  // SQL.  In datastore, there's really no other kind of query.
+  protected Class<T> entityClass;
+
+  // Field to order by, if any.  Null if we don't care about order.
+  @Nullable protected String orderBy;
+
+  protected List<WhereClause<?>> predicates = new ArrayList<WhereClause<?>>();
+
+  protected QueryComposer(Class<T> entityClass) {
+    this.entityClass = entityClass;
+  }
+
+  /**
+   * Introduce a "where" clause to the query.
+   *
+   * <p>Causes the query to return only results where the field and value have the relationship
+   * specified by the comparator. For example, "field EQ value", "field GT value" etc.
+   */
+  public <U extends Comparable<? super U>> QueryComposer<T> where(
+      String fieldName, Comparator comparator, U value) {
+    predicates.add(new WhereClause(fieldName, comparator, value));
+    return this;
+  }
+
+  /**
+   * Order the query results by the value of the specified field.
+   *
+   * <p>TODO(mmuller): add the ability to do descending sort order.
+   */
+  public QueryComposer<T> orderBy(String fieldName) {
+    orderBy = fieldName;
+    return this;
+  }
+
+  /** Returns the first result of the query or an empty optional if there is none. */
+  public abstract Optional<T> first();
+
+  /**
+   * Returns the one and only result of a query.
+   *
+   * <p>Throws a {@link javax.persistence.NonUniqueResultException} if there is more than one
+   * result, throws {@link javax.persistence.NoResultException} if no results are found.
+   */
+  public abstract T getSingleResult();
+
+  /** Returns the results of the query as a stream. */
+  public abstract Stream<T> stream();
+
+  /** Returns the number of results of the query. */
+  public abstract long count();
+
+  /** Returns the results of the query as a list. */
+  public abstract List<T> list();
+
+  // We have to wrap the CriteriaQueryBuilder predicate factories in our own functions because at
+  // the point where we pass them to the Comparator constructor, the compiler can't determine which
+  // of the overloads to use since there is no "value" object for context.
+
+  public static <U extends Comparable<? super U>> WhereOperator<U> equal(
+      CriteriaBuilder criteriaBuilder) {
+    return criteriaBuilder::equal;
+  }
+
+  public static <U extends Comparable<? super U>> WhereOperator<U> lessThan(
+      CriteriaBuilder criteriaBuilder) {
+    return criteriaBuilder::lessThan;
+  }
+
+  public static <U extends Comparable<? super U>> WhereOperator<U> lessThanOrEqualTo(
+      CriteriaBuilder criteriaBuilder) {
+    return criteriaBuilder::lessThanOrEqualTo;
+  }
+
+  public static <U extends Comparable<? super U>> WhereOperator<U> greaterThanOrEqualTo(
+      CriteriaBuilder criteriaBuilder) {
+    return criteriaBuilder::greaterThanOrEqualTo;
+  }
+
+  public static <U extends Comparable<? super U>> WhereOperator<U> greaterThan(
+      CriteriaBuilder criteriaBuilder) {
+    return criteriaBuilder::greaterThan;
+  }
+
+  /**
+   * Enum used to specify comparison operations, e.g. {@code where("fieldName", Comparator.NE,
+   * "someval")'}.
+   *
+   * <p>These contain values that specify the comparison behavior for both objectify and criteria
+   * queries. For objectify, we provide a string to be appended to the field name in a {@code
+   * filter()} expression. For criteria queries we provide a function that knows how to obtain a
+   * {@link WhereOperator} from a {@link CriteriaBuilder}.
+   *
+   * <p>Note that the objectify strings for comparators other than equality are preceded by a space
+   * because {@code filter()} expects the fieldname to be separated from the operator by a space.
+   */
+  public enum Comparator {
+    /**
+     * Return only records whose field is equal to the value.
+     *
+     * <p>Note that the datastore string for this is empty, which is consistent with the way {@code
+     * filter()} works (it uses an unadorned field name to check for equality).
+     */
+    EQ("", QueryComposer::equal),
+
+    /** Return only records whose field is less than the value. */
+    LT(" <", QueryComposer::lessThan),
+
+    /** Return only records whose field is less than or equal to the value. */
+    LTE(" <=", QueryComposer::lessThanOrEqualTo),
+
+    /** Return only records whose field is greater than or equal to the value. */
+    GTE(" >=", QueryComposer::greaterThanOrEqualTo),
+
+    /** Return only records whose field is greater than the value. */
+    GT(" >", QueryComposer::greaterThan);
+
+    private final String datastoreString;
+
+    @SuppressWarnings("ImmutableEnumChecker") // Functions are immutable.
+    private final Function<CriteriaBuilder, WhereOperator<?>> operatorFactory;
+
+    Comparator(
+        String datastoreString, Function<CriteriaBuilder, WhereOperator<?>> operatorFactory) {
+      this.datastoreString = datastoreString;
+      this.operatorFactory = operatorFactory;
+    }
+
+    public String getDatastoreString() {
+      return datastoreString;
+    }
+
+    public Function<CriteriaBuilder, WhereOperator<?>> getComparisonFactory() {
+      return operatorFactory;
+    }
+  };
+
+  protected static class WhereClause<U extends Comparable<? super U>> {
+    public String fieldName;
+    public Comparator comparator;
+    public U value;
+
+    WhereClause(String fieldName, Comparator comparator, U value) {
+      this.fieldName = fieldName;
+      this.comparator = comparator;
+      this.value = value;
+    }
+
+    public void addToCriteriaQueryBuilder(CriteriaQueryBuilder queryBuilder) {
+      CriteriaBuilder criteriaBuilder = jpaTm().getEntityManager().getCriteriaBuilder();
+      queryBuilder.where(
+          fieldName, comparator.getComparisonFactory().apply(criteriaBuilder), value);
+    }
+  }
+}

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

@@ -273,6 +273,9 @@ public interface TransactionManager {
    */
   void deleteWithoutBackup(Object entity);
 
+  /** Returns a QueryComposer which can be used to perform queries against the current database. */
+  <T> QueryComposer<T> createQueryComposer(Class<T> entity);
+
   /** Clears the session cache if the underlying database is Datastore, otherwise it is a no-op. */
   void clearSessionCache();
 

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

@@ -24,6 +24,7 @@ import static google.registry.request.Action.Method.POST;
 import static google.registry.util.DateTimeUtils.isBeforeOrAt;
 
 import com.google.appengine.tools.cloudstorage.GcsFilename;
+import com.google.common.flogger.FluentLogger;
 import com.google.common.io.ByteStreams;
 import google.registry.config.RegistryConfig.Config;
 import google.registry.gcs.GcsUtils;
@@ -31,6 +32,7 @@ import google.registry.keyring.api.KeyModule.Key;
 import google.registry.model.common.Cursor;
 import google.registry.model.common.Cursor.CursorType;
 import google.registry.model.rde.RdeNamingUtils;
+import google.registry.model.rde.RdeRevision;
 import google.registry.model.registry.Registry;
 import google.registry.rde.EscrowTaskRunner.EscrowTask;
 import google.registry.request.Action;
@@ -41,6 +43,7 @@ import google.registry.request.Response;
 import google.registry.request.auth.Auth;
 import java.io.IOException;
 import java.io.InputStream;
+import java.util.Optional;
 import javax.inject.Inject;
 import org.bouncycastle.openpgp.PGPPrivateKey;
 import org.joda.time.DateTime;
@@ -56,6 +59,8 @@ import org.joda.time.Duration;
     auth = Auth.AUTH_INTERNAL_OR_ADMIN)
 public final class RdeReportAction implements Runnable, EscrowTask {
 
+  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
+
   static final String PATH = "/_dr/task/rdeReport";
 
   @Inject GcsUtils gcsUtils;
@@ -76,8 +81,9 @@ public final class RdeReportAction implements Runnable, EscrowTask {
 
   @Override
   public void runWithLock(DateTime watermark) throws Exception {
-    Cursor cursor =
-        transactIfJpaTm(() -> tm().loadByKey(Cursor.createVKey(CursorType.RDE_UPLOAD, tld)));
+    Optional<Cursor> cursor =
+        transactIfJpaTm(
+            () -> tm().loadByKeyIfPresent(Cursor.createVKey(CursorType.RDE_UPLOAD, tld)));
     DateTime cursorTime = getCursorTimeOrStartOfTime(cursor);
     if (isBeforeOrAt(cursorTime, watermark)) {
       throw new NoContentException(
@@ -86,12 +92,17 @@ public final class RdeReportAction implements Runnable, EscrowTask {
                   + "last upload completion was at %s",
               tld, watermark, cursorTime));
     }
-    String prefix = RdeNamingUtils.makeRydeFilename(tld, watermark, FULL, 1, 0);
+    int revision =
+        RdeRevision.getCurrentRevision(tld, watermark, FULL)
+            .orElseThrow(
+                () -> new IllegalStateException("RdeRevision was not set on generated deposit"));
+    String prefix = RdeNamingUtils.makeRydeFilename(tld, watermark, FULL, 1, revision);
     GcsFilename reportFilename = new GcsFilename(bucket, prefix + "-report.xml.ghostryde");
     verify(gcsUtils.existsAndNotEmpty(reportFilename), "Missing file: %s", reportFilename);
     reporter.send(readReportFromGcs(reportFilename));
     response.setContentType(PLAIN_TEXT_UTF_8);
     response.setPayload(String.format("OK %s %s\n", tld, watermark));
+    logger.atInfo().log("Successfully sent report %s.", reportFilename);
   }
 
   /** Reads and decrypts the XML file from cloud storage. */

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

@@ -20,8 +20,8 @@ import static com.google.appengine.tools.cloudstorage.GcsServiceFactory.createGc
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.base.Verify.verify;
 import static google.registry.model.common.Cursor.getCursorTimeOrStartOfTime;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 import static java.nio.charset.StandardCharsets.UTF_8;
 
 import com.google.appengine.tools.cloudstorage.GcsFilename;
@@ -210,7 +210,11 @@ public final class RdeStagingReducer extends Reducer<PendingDeposit, DepositFrag
     tm().transact(
             () -> {
               Registry registry = Registry.get(tld);
-              Cursor cursor = ofy().load().key(Cursor.createKey(key.cursor(), registry)).now();
+              Optional<Cursor> cursor =
+                  transactIfJpaTm(
+                      () ->
+                          tm().loadByKeyIfPresent(
+                                  Cursor.createVKey(key.cursor(), registry.getTldStr())));
               DateTime position = getCursorTimeOrStartOfTime(cursor);
               checkState(key.interval() != null, "Interval must be present");
               DateTime newPosition = key.watermark().plus(key.interval());

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

@@ -64,6 +64,7 @@ import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.net.URI;
+import java.util.Optional;
 import javax.inject.Inject;
 import javax.inject.Named;
 import org.bouncycastle.openpgp.PGPKeyPair;
@@ -133,7 +134,8 @@ public final class RdeUploadAction implements Runnable, EscrowTask {
   @Override
   public void runWithLock(final DateTime watermark) throws Exception {
     logger.atInfo().log("Verifying readiness to upload the RDE deposit.");
-    Cursor cursor = transactIfJpaTm(() -> tm().loadByKey(Cursor.createVKey(RDE_STAGING, tld)));
+    Optional<Cursor> cursor =
+        transactIfJpaTm(() -> tm().loadByKeyIfPresent(Cursor.createVKey(RDE_STAGING, tld)));
     DateTime stagingCursorTime = getCursorTimeOrStartOfTime(cursor);
     if (isBeforeOrAt(stagingCursorTime, watermark)) {
       throw new NoContentException(
@@ -158,8 +160,10 @@ public final class RdeUploadAction implements Runnable, EscrowTask {
               sftpCursorTime,
               timeSinceLastSftp.getStandardMinutes()));
     }
-    int revision = RdeRevision.getNextRevision(tld, watermark, FULL) - 1;
-    verify(revision >= 0, "RdeRevision was not set on generated deposit");
+    int revision =
+        RdeRevision.getCurrentRevision(tld, watermark, FULL)
+            .orElseThrow(
+                () -> new IllegalStateException("RdeRevision was not set on generated deposit"));
     final String name = RdeNamingUtils.makeRydeFilename(tld, watermark, FULL, 1, revision);
     final GcsFilename xmlFilename = new GcsFilename(bucket, name + ".xml.ghostryde");
     final GcsFilename xmlLengthFilename = new GcsFilename(bucket, name + ".xml.length");

core/src/main/java/google/registry/reporting/billing/GenerateInvoicesAction.java

@@ -14,6 +14,7 @@
 
 package google.registry.reporting.billing;
 
+import static google.registry.beam.BeamUtils.createJobName;
 import static google.registry.reporting.ReportingUtils.enqueueBeamReportingTask;
 import static google.registry.reporting.billing.BillingModule.PARAM_SHOULD_PUBLISH;
 import static google.registry.request.Action.Method.POST;
@@ -21,9 +22,9 @@ import static javax.servlet.http.HttpServletResponse.SC_INTERNAL_SERVER_ERROR;
 import static javax.servlet.http.HttpServletResponse.SC_OK;
 
 import com.google.api.services.dataflow.Dataflow;
-import com.google.api.services.dataflow.model.LaunchTemplateParameters;
-import com.google.api.services.dataflow.model.LaunchTemplateResponse;
-import com.google.api.services.dataflow.model.RuntimeEnvironment;
+import com.google.api.services.dataflow.model.LaunchFlexTemplateParameter;
+import com.google.api.services.dataflow.model.LaunchFlexTemplateRequest;
+import com.google.api.services.dataflow.model.LaunchFlexTemplateResponse;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.flogger.FluentLogger;
 import com.google.common.net.MediaType;
@@ -33,6 +34,7 @@ import google.registry.request.Action;
 import google.registry.request.Parameter;
 import google.registry.request.Response;
 import google.registry.request.auth.Auth;
+import google.registry.util.Clock;
 import java.io.IOException;
 import java.util.Map;
 import javax.inject.Inject;
@@ -42,9 +44,8 @@ import org.joda.time.YearMonth;
  * Invokes the {@code InvoicingPipeline} beam template via the REST api, and enqueues the {@link
  * PublishInvoicesAction} to publish the subsequent output.
  *
- * <p>This action runs the {@link google.registry.beam.invoicing.InvoicingPipeline} beam template,
- * staged at gs://&lt;projectId&gt;-beam/templates/invoicing. The pipeline then generates invoices
- * for the month and stores them on GCS.
+ * <p>This action runs the {@link google.registry.beam.invoicing.InvoicingPipeline} beam flex
+ * template. The pipeline then generates invoices for the month and stores them on GCS.
  */
 @Action(
     service = Action.Service.BACKEND,
@@ -56,57 +57,73 @@ public class GenerateInvoicesAction implements Runnable {
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
 
   static final String PATH = "/_dr/task/generateInvoices";
+  static final String PIPELINE_NAME = "invoicing_pipeline";
 
   private final String projectId;
-  private final String beamBucketUrl;
-  private final String invoiceTemplateUrl;
-  private final String jobZone;
+  private final String jobRegion;
+  private final String stagingBucketUrl;
+  private final String billingBucketUrl;
+  private final String invoiceFilePrefix;
   private final boolean shouldPublish;
   private final YearMonth yearMonth;
-  private final Dataflow dataflow;
-  private final Response response;
   private final BillingEmailUtils emailUtils;
+  private final Clock clock;
+  private final Response response;
+  private final Dataflow dataflow;
 
   @Inject
   GenerateInvoicesAction(
       @Config("projectId") String projectId,
-      @Config("apacheBeamBucketUrl") String beamBucketUrl,
-      @Config("invoiceTemplateUrl") String invoiceTemplateUrl,
-      @Config("defaultJobZone") String jobZone,
+      @Config("defaultJobRegion") String jobRegion,
+      @Config("beamStagingBucketUrl") String stagingBucketUrl,
+      @Config("billingBucketUrl") String billingBucketUrl,
+      @Config("invoiceFilePrefix") String invoiceFilePrefix,
       @Parameter(PARAM_SHOULD_PUBLISH) boolean shouldPublish,
       YearMonth yearMonth,
-      Dataflow dataflow,
+      BillingEmailUtils emailUtils,
+      Clock clock,
       Response response,
-      BillingEmailUtils emailUtils) {
+      Dataflow dataflow) {
     this.projectId = projectId;
-    this.beamBucketUrl = beamBucketUrl;
-    this.invoiceTemplateUrl = invoiceTemplateUrl;
-    this.jobZone = jobZone;
+    this.jobRegion = jobRegion;
+    this.stagingBucketUrl = stagingBucketUrl;
+    this.billingBucketUrl = billingBucketUrl;
+    this.invoiceFilePrefix = invoiceFilePrefix;
     this.shouldPublish = shouldPublish;
     this.yearMonth = yearMonth;
-    this.dataflow = dataflow;
-    this.response = response;
     this.emailUtils = emailUtils;
+    this.clock = clock;
+    this.response = response;
+    this.dataflow = dataflow;
   }
 
   @Override
   public void run() {
+    response.setContentType(MediaType.PLAIN_TEXT_UTF_8);
     logger.atInfo().log("Launching invoicing pipeline for %s", yearMonth);
     try {
-      LaunchTemplateParameters params =
-          new LaunchTemplateParameters()
-              .setJobName(String.format("invoicing-%s", yearMonth))
-              .setEnvironment(
-                  new RuntimeEnvironment()
-                      .setZone(jobZone)
-                      .setTempLocation(beamBucketUrl + "/temporary"))
-              .setParameters(ImmutableMap.of("yearMonth", yearMonth.toString("yyyy-MM")));
-      LaunchTemplateResponse launchResponse =
+      LaunchFlexTemplateParameter parameter =
+          new LaunchFlexTemplateParameter()
+              .setJobName(createJobName("invoicing", clock))
+              .setContainerSpecGcsPath(
+                  String.format("%s/%s_metadata.json", stagingBucketUrl, PIPELINE_NAME))
+              .setParameters(
+                  ImmutableMap.of(
+                      "yearMonth",
+                      yearMonth.toString("yyyy-MM"),
+                      "invoiceFilePrefix",
+                      invoiceFilePrefix,
+                      "billingBucketUrl",
+                      billingBucketUrl));
+      LaunchFlexTemplateResponse launchResponse =
           dataflow
               .projects()
-              .templates()
-              .launch(projectId, params)
-              .setGcsPath(invoiceTemplateUrl)
+              .locations()
+              .flexTemplates()
+              .launch(
+                  projectId,
+                  jobRegion,
+                  new LaunchFlexTemplateRequest().setLaunchParameter(parameter))
               .execute();
       logger.atInfo().log("Got response: %s", launchResponse.getJob().toPrettyString());
       String jobId = launchResponse.getJob().getId();
@@ -123,12 +140,10 @@ public class GenerateInvoicesAction implements Runnable {
       logger.atWarning().withCause(e).log("Template Launch failed");
       emailUtils.sendAlertEmail(String.format("Template Launch failed due to %s", e.getMessage()));
       response.setStatus(SC_INTERNAL_SERVER_ERROR);
-      response.setContentType(MediaType.PLAIN_TEXT_UTF_8);
       response.setPayload(String.format("Template launch failed: %s", e.getMessage()));
       return;
     }
     response.setStatus(SC_OK);
-    response.setContentType(MediaType.PLAIN_TEXT_UTF_8);
     response.setPayload("Launched dataflow template.");
   }
 }

core/src/main/java/google/registry/reporting/billing/PublishInvoicesAction.java

@@ -59,6 +59,7 @@ public class PublishInvoicesAction implements Runnable {
   private static final String JOB_FAILED = "JOB_STATE_FAILED";
 
   private final String projectId;
+  private final String jobRegion;
   private final String jobId;
   private final BillingEmailUtils emailUtils;
   private final Dataflow dataflow;
@@ -68,12 +69,14 @@ public class PublishInvoicesAction implements Runnable {
   @Inject
   PublishInvoicesAction(
       @Config("projectId") String projectId,
+      @Config("defaultJobRegion") String jobRegion,
       @Parameter(ReportingModule.PARAM_JOB_ID) String jobId,
       BillingEmailUtils emailUtils,
       Dataflow dataflow,
       Response response,
       YearMonth yearMonth) {
     this.projectId = projectId;
+    this.jobRegion = jobRegion;
     this.jobId = jobId;
     this.emailUtils = emailUtils;
     this.dataflow = dataflow;
@@ -87,7 +90,7 @@ public class PublishInvoicesAction implements Runnable {
   public void run() {
     try {
       logger.atInfo().log("Starting publish job.");
-      Job job = dataflow.projects().jobs().get(projectId, jobId).execute();
+      Job job = dataflow.projects().locations().jobs().get(projectId, jobRegion, jobId).execute();
       String state = job.getCurrentState();
       switch (state) {
         case JOB_DONE:

core/src/main/java/google/registry/reporting/icann/IcannReportingUploadAction.java

@@ -16,7 +16,6 @@ package google.registry.reporting.icann;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.net.MediaType.PLAIN_TEXT_UTF_8;
-import static google.registry.model.common.Cursor.getCursorTimeOrStartOfTime;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.request.Action.Method.POST;
 import static javax.servlet.http.HttpServletResponse.SC_OK;
@@ -107,10 +106,10 @@ public final class IcannReportingUploadAction implements Runnable {
 
           // If cursor time is before now, upload the corresponding report
           cursors.entrySet().stream()
-              .filter(entry -> getCursorTimeOrStartOfTime(entry.getKey()).isBefore(clock.nowUtc()))
+              .filter(entry -> entry.getKey().getCursorTime().isBefore(clock.nowUtc()))
               .forEach(
                   entry -> {
-                    DateTime cursorTime = getCursorTimeOrStartOfTime(entry.getKey());
+                    DateTime cursorTime = entry.getKey().getCursorTime();
                     uploadReport(
                         cursorTime,
                         entry.getKey().getType(),

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

@@ -14,6 +14,7 @@
 
 package google.registry.reporting.spec11;
 
+import static google.registry.beam.BeamUtils.createJobName;
 import static google.registry.reporting.ReportingModule.PARAM_DATE;
 import static google.registry.reporting.ReportingUtils.enqueueBeamReportingTask;
 import static google.registry.request.Action.Method.POST;
@@ -21,19 +22,21 @@ import static javax.servlet.http.HttpServletResponse.SC_INTERNAL_SERVER_ERROR;
 import static javax.servlet.http.HttpServletResponse.SC_OK;
 
 import com.google.api.services.dataflow.Dataflow;
-import com.google.api.services.dataflow.model.LaunchTemplateParameters;
-import com.google.api.services.dataflow.model.LaunchTemplateResponse;
-import com.google.api.services.dataflow.model.RuntimeEnvironment;
+import com.google.api.services.dataflow.model.LaunchFlexTemplateParameter;
+import com.google.api.services.dataflow.model.LaunchFlexTemplateRequest;
+import com.google.api.services.dataflow.model.LaunchFlexTemplateResponse;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.flogger.FluentLogger;
 import com.google.common.net.MediaType;
 import google.registry.config.RegistryConfig.Config;
+import google.registry.config.RegistryEnvironment;
 import google.registry.keyring.api.KeyModule.Key;
 import google.registry.reporting.ReportingModule;
 import google.registry.request.Action;
 import google.registry.request.Parameter;
 import google.registry.request.Response;
 import google.registry.request.auth.Auth;
+import google.registry.util.Clock;
 import java.io.IOException;
 import java.util.Map;
 import javax.inject.Inject;
@@ -55,55 +58,68 @@ public class GenerateSpec11ReportAction implements Runnable {
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
 
   static final String PATH = "/_dr/task/generateSpec11";
+  static final String PIPELINE_NAME = "spec11_pipeline";
 
   private final String projectId;
-  private final String beamBucketUrl;
-  private final String spec11TemplateUrl;
-  private final String jobZone;
+  private final String jobRegion;
+  private final String stagingBucketUrl;
+  private final String reportingBucketUrl;
   private final String apiKey;
   private final LocalDate date;
+  private final Clock clock;
   private final Response response;
   private final Dataflow dataflow;
 
   @Inject
   GenerateSpec11ReportAction(
       @Config("projectId") String projectId,
-      @Config("apacheBeamBucketUrl") String beamBucketUrl,
-      @Config("spec11TemplateUrl") String spec11TemplateUrl,
-      @Config("defaultJobZone") String jobZone,
+      @Config("defaultJobRegion") String jobRegion,
+      @Config("beamStagingBucketUrl") String stagingBucketUrl,
+      @Config("reportingBucketUrl") String reportingBucketUrl,
       @Key("safeBrowsingAPIKey") String apiKey,
       @Parameter(PARAM_DATE) LocalDate date,
+      Clock clock,
       Response response,
       Dataflow dataflow) {
     this.projectId = projectId;
-    this.beamBucketUrl = beamBucketUrl;
-    this.spec11TemplateUrl = spec11TemplateUrl;
-    this.jobZone = jobZone;
+    this.jobRegion = jobRegion;
+    this.stagingBucketUrl = stagingBucketUrl;
+    this.reportingBucketUrl = reportingBucketUrl;
     this.apiKey = apiKey;
     this.date = date;
+    this.clock = clock;
     this.response = response;
     this.dataflow = dataflow;
   }
 
   @Override
   public void run() {
+    response.setContentType(MediaType.PLAIN_TEXT_UTF_8);
     try {
-      LaunchTemplateParameters params =
-          new LaunchTemplateParameters()
-              .setJobName(String.format("spec11_%s", date))
-              .setEnvironment(
-                  new RuntimeEnvironment()
-                      .setZone(jobZone)
-                      .setTempLocation(beamBucketUrl + "/temporary"))
+      LaunchFlexTemplateParameter parameter =
+          new LaunchFlexTemplateParameter()
+              .setJobName(createJobName("spec11", clock))
+              .setContainerSpecGcsPath(
+                  String.format("%s/%s_metadata.json", stagingBucketUrl, PIPELINE_NAME))
               .setParameters(
                   ImmutableMap.of(
-                      "safeBrowsingApiKey", apiKey, ReportingModule.PARAM_DATE, date.toString()));
-      LaunchTemplateResponse launchResponse =
+                      "safeBrowsingApiKey",
+                      apiKey,
+                      ReportingModule.PARAM_DATE,
+                      date.toString(),
+                      "reportingBucketUrl",
+                      reportingBucketUrl,
+                      "registryEnvironment",
+                      RegistryEnvironment.get().name()));
+      LaunchFlexTemplateResponse launchResponse =
           dataflow
               .projects()
-              .templates()
-              .launch(projectId, params)
-              .setGcsPath(spec11TemplateUrl)
+              .locations()
+              .flexTemplates()
+              .launch(
+                  projectId,
+                  jobRegion,
+                  new LaunchFlexTemplateRequest().setLaunchParameter(parameter))
               .execute();
       Map<String, String> beamTaskParameters =
           ImmutableMap.of(
@@ -116,12 +132,10 @@ public class GenerateSpec11ReportAction implements Runnable {
     } catch (IOException e) {
       logger.atWarning().withCause(e).log("Template Launch failed");
       response.setStatus(SC_INTERNAL_SERVER_ERROR);
-      response.setContentType(MediaType.PLAIN_TEXT_UTF_8);
       response.setPayload(String.format("Template launch failed: %s", e.getMessage()));
       return;
     }
     response.setStatus(SC_OK);
-    response.setContentType(MediaType.PLAIN_TEXT_UTF_8);
     response.setPayload("Launched Spec11 dataflow template.");
   }
 }

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

@@ -69,6 +69,7 @@ public class PublishSpec11ReportAction implements Runnable {
   private static final String JOB_FAILED = "JOB_STATE_FAILED";
 
   private final String projectId;
+  private final String jobRegion;
   private final String registryName;
   private final String jobId;
   private final Spec11EmailUtils emailUtils;
@@ -80,6 +81,7 @@ public class PublishSpec11ReportAction implements Runnable {
   @Inject
   PublishSpec11ReportAction(
       @Config("projectId") String projectId,
+      @Config("defaultJobRegion") String jobRegion,
       @Config("registryName") String registryName,
       @Parameter(ReportingModule.PARAM_JOB_ID) String jobId,
       Spec11EmailUtils emailUtils,
@@ -88,6 +90,7 @@ public class PublishSpec11ReportAction implements Runnable {
       Response response,
       @Parameter(PARAM_DATE) LocalDate date) {
     this.projectId = projectId;
+    this.jobRegion = jobRegion;
     this.registryName = registryName;
     this.jobId = jobId;
     this.emailUtils = emailUtils;
@@ -101,7 +104,7 @@ public class PublishSpec11ReportAction implements Runnable {
   public void run() {
     try {
       logger.atInfo().log("Starting publish job.");
-      Job job = dataflow.projects().jobs().get(projectId, jobId).execute();
+      Job job = dataflow.projects().locations().jobs().get(projectId, jobRegion, jobId).execute();
       String state = job.getCurrentState();
       switch (state) {
         case JOB_DONE:

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

@@ -17,7 +17,9 @@ package google.registry.reporting.spec11;
 import static com.google.common.base.Throwables.getRootCause;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.io.Resources.getResource;
-import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.QueryComposer.Comparator;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -129,17 +131,21 @@ public class Spec11EmailUtils {
   private RegistrarThreatMatches filterOutNonPublishedMatches(
       RegistrarThreatMatches registrarThreatMatches) {
     ImmutableList<ThreatMatch> filteredMatches =
-        registrarThreatMatches.threatMatches().stream()
-            .filter(
-                threatMatch ->
-                    ofy()
-                        .load()
-                        .type(DomainBase.class)
-                        .filter("fullyQualifiedDomainName", threatMatch.fullyQualifiedDomainName())
-                        .first()
-                        .now()
-                        .shouldPublishToDns())
-            .collect(toImmutableList());
+        transactIfJpaTm(
+            () -> {
+              return registrarThreatMatches.threatMatches().stream()
+                  .filter(
+                      threatMatch ->
+                          tm()
+                              .createQueryComposer(DomainBase.class)
+                              .where(
+                                  "fullyQualifiedDomainName",
+                                  Comparator.EQ,
+                                  threatMatch.fullyQualifiedDomainName())
+                              .stream()
+                              .anyMatch(DomainBase::shouldPublishToDns))
+                  .collect(toImmutableList());
+            });
     return RegistrarThreatMatches.create(registrarThreatMatches.clientId(), filteredMatches);
   }
 

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

@@ -14,12 +14,13 @@
 
 package google.registry.tools;
 
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.registry.Registries.assertTldsExist;
+import static google.registry.persistence.transaction.QueryComposer.Comparator;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 
 import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
-import com.google.common.collect.Iterables;
 import google.registry.model.domain.DomainBase;
 import google.registry.util.Clock;
 import java.util.List;
@@ -45,14 +46,12 @@ final class CountDomainsCommand implements CommandWithRemoteApi {
         .forEach(tld -> System.out.printf("%s,%d\n", tld, getCountForTld(tld, now)));
   }
 
-  private int getCountForTld(String tld, DateTime now) {
-    return Iterables.size(
-        ofy()
-            .load()
-            .type(DomainBase.class)
-            .filter("tld", tld)
-            .filter("deletionTime >", now)
-            .chunkAll()
-            .keys());
+  private long getCountForTld(String tld, DateTime now) {
+    return transactIfJpaTm(
+        () ->
+            tm().createQueryComposer(DomainBase.class)
+                .where("tld", Comparator.EQ, tld)
+                .where("deletionTime", Comparator.GT, now)
+                .count());
   }
 }

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

@@ -27,9 +27,7 @@ import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
 
-/**
- * Command to create groups in Google Groups for all contact types for a registrar.
- */
+/** Command to create groups in Google Groups for all contact types for a registrar. */
 @Parameters(separators = " =", commandDescription = "Create groups for a registrar.")
 public class CreateRegistrarGroupsCommand extends ConfirmingCommand
     implements CommandWithConnection, CommandWithRemoteApi {

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

@@ -30,8 +30,8 @@ import java.nio.file.Files;
 import java.util.List;
 import org.joda.time.DateTime;
 
-/** Command to create a {@link ReservedList} on Datastore. */
-@Parameters(separators = " =", commandDescription = "Create a ReservedList in Datastore.")
+/** Command to create a {@link ReservedList}. */
+@Parameters(separators = " =", commandDescription = "Create a ReservedList.")
 final class CreateReservedListCommand extends CreateOrUpdateReservedListCommand {
 
   @VisibleForTesting
@@ -70,7 +70,8 @@ final class CreateReservedListCommand extends CreateOrUpdateReservedListCommand
     checkArgument(nameParts.size() == 2, INVALID_FORMAT_ERROR_MESSAGE);
     String tld = nameParts.get(0);
     if (!tld.equals("common")) {
-      assertTldExists(tld);
+      assertTldExists(
+          tld, "The name must be in the format {tld|common}_list-name, yet TLD %s does not exist");
     }
     checkArgument(nameParts.get(1).matches("[-a-zA-Z0-9]+"), INVALID_FORMAT_ERROR_MESSAGE);
   }

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

@@ -25,10 +25,10 @@ import google.registry.model.registry.label.PremiumListDualDao;
 import javax.annotation.Nullable;
 
 /**
- * Command to delete a {@link PremiumList} in Datastore. This command will fail if the premium list
- * is currently in use on a tld.
+ * Command to delete a {@link PremiumList}. This command will fail if the premium list is currently
+ * in use on a tld.
  */
-@Parameters(separators = " =", commandDescription = "Delete a PremiumList from Datastore.")
+@Parameters(separators = " =", commandDescription = "Delete a PremiumList.")
 final class DeletePremiumListCommand extends ConfirmingCommand implements CommandWithRemoteApi {
 
   @Nullable PremiumList premiumList;
@@ -55,7 +55,7 @@ final class DeletePremiumListCommand extends ConfirmingCommand implements Comman
 
   @Override
   protected String prompt() {
-    return "You are about to delete the premium list: \n" + premiumList;
+    return "You are about to delete the premium list: \n" + premiumList.getName();
   }
 
   @Override

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

@@ -1,31 +0,0 @@
-// Copyright 2017 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 google.registry.beam.invoicing.InvoicingPipeline;
-import javax.inject.Inject;
-
-/** Nomulus command that deploys the {@link InvoicingPipeline} template. */
-@Parameters(commandDescription = "Deploy the invoicing pipeline to GCS.")
-public class DeployInvoicingPipelineCommand implements Command {
-
-  @Inject InvoicingPipeline invoicingPipeline;
-
-  @Override
-  public void run() {
-    invoicingPipeline.deploy();
-  }
-}

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

@@ -1,87 +0,0 @@
-// Copyright 2018 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.Parameter;
-import com.beust.jcommander.Parameters;
-import google.registry.beam.initsql.BeamJpaModule.JpaTransactionManagerComponent;
-import google.registry.beam.initsql.JpaSupplierFactory;
-import google.registry.beam.spec11.Spec11Pipeline;
-import google.registry.config.CredentialModule.LocalCredential;
-import google.registry.config.RegistryConfig.Config;
-import google.registry.util.GoogleCredentialsBundle;
-import google.registry.util.Retrier;
-import javax.annotation.Nullable;
-import javax.inject.Inject;
-
-/** Nomulus command that deploys the {@link Spec11Pipeline} template. */
-@Parameters(commandDescription = "Deploy the Spec11 pipeline to GCS.")
-public class DeploySpec11PipelineCommand implements Command {
-
-  @Inject
-  @Config("projectId")
-  String projectId;
-
-  @Inject
-  @Config("defaultJobRegion")
-  String beamJobRegion;
-
-  @Parameter(
-      names = {"-p", "--project"},
-      description = "Cloud KMS project ID",
-      required = true)
-  String cloudKmsProjectId;
-
-  @Inject
-  @Config("beamStagingUrl")
-  String beamStagingUrl;
-
-  @Inject
-  @Config("spec11TemplateUrl")
-  String spec11TemplateUrl;
-
-  @Inject
-  @Config("reportingBucketUrl")
-  String reportingBucketUrl;
-
-  @Inject @LocalCredential GoogleCredentialsBundle googleCredentialsBundle;
-  @Inject Retrier retrier;
-
-  @Inject
-  @Nullable
-  @Config("sqlAccessInfoFile")
-  String sqlAccessInfoFile;
-
-  @Override
-  public void run() {
-    JpaSupplierFactory jpaSupplierFactory =
-        new JpaSupplierFactory(
-            sqlAccessInfoFile,
-            cloudKmsProjectId,
-            JpaTransactionManagerComponent::cloudSqlJpaTransactionManager);
-
-    Spec11Pipeline pipeline =
-        new Spec11Pipeline(
-            projectId,
-            beamJobRegion,
-            beamStagingUrl,
-            spec11TemplateUrl,
-            reportingBucketUrl,
-            jpaSupplierFactory,
-            googleCredentialsBundle,
-            retrier);
-    pipeline.deploy();
-  }
-}

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

@@ -14,9 +14,10 @@
 
 package google.registry.tools;
 
-import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.QueryComposer.Comparator;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 import static java.nio.charset.StandardCharsets.UTF_8;
-import static org.joda.time.DateTimeZone.UTC;
 
 import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
@@ -24,9 +25,11 @@ import com.google.common.collect.ImmutableList;
 import google.registry.model.domain.DomainBase;
 import google.registry.tmch.LordnTaskUtils;
 import google.registry.tools.params.PathParameter;
+import google.registry.util.Clock;
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import javax.inject.Inject;
 import org.joda.time.DateTime;
 
 /** Command to generate a LORDN CSV file for an entire TLD. */
@@ -53,22 +56,21 @@ final class GenerateLordnCommand implements CommandWithRemoteApi {
       required = true)
   private Path sunriseOutputPath;
 
+  @Inject Clock clock;
+
   @Override
   public void run() throws IOException {
-    DateTime now = DateTime.now(UTC);
+    DateTime now = clock.nowUtc();
     ImmutableList.Builder<String> claimsCsv = new ImmutableList.Builder<>();
     ImmutableList.Builder<String> sunriseCsv = new ImmutableList.Builder<>();
-    for (DomainBase domain : ofy().load().type(DomainBase.class).filter("tld", tld)) {
-      String status = " ";
-      if (domain.getLaunchNotice() == null && domain.getSmdId() != null) {
-        sunriseCsv.add(LordnTaskUtils.getCsvLineForSunriseDomain(domain, domain.getCreationTime()));
-        status = "S";
-      } else if (domain.getLaunchNotice() != null || domain.getSmdId() != null) {
-        claimsCsv.add(LordnTaskUtils.getCsvLineForClaimsDomain(domain, domain.getCreationTime()));
-        status = "C";
-      }
-      System.out.printf("%s[%s] ", domain.getDomainName(), status);
-    }
+    transactIfJpaTm(
+        () ->
+            tm()
+                .createQueryComposer(DomainBase.class)
+                .where("tld", Comparator.EQ, tld)
+                .orderBy("repoId")
+                .stream()
+                .forEach(domain -> processDomain(claimsCsv, sunriseCsv, domain)));
     ImmutableList<String> claimsRows = claimsCsv.build();
     ImmutableList<String> claimsAll =
         new ImmutableList.Builder<String>()
@@ -86,4 +88,19 @@ final class GenerateLordnCommand implements CommandWithRemoteApi {
     Files.write(claimsOutputPath, claimsAll, UTF_8);
     Files.write(sunriseOutputPath, sunriseAll, UTF_8);
   }
+
+  private static void processDomain(
+      ImmutableList.Builder<String> claimsCsv,
+      ImmutableList.Builder<String> sunriseCsv,
+      DomainBase domain) {
+    String status = " ";
+    if (domain.getLaunchNotice() == null && domain.getSmdId() != null) {
+      sunriseCsv.add(LordnTaskUtils.getCsvLineForSunriseDomain(domain, domain.getCreationTime()));
+      status = "S";
+    } else if (domain.getLaunchNotice() != null || domain.getSmdId() != null) {
+      claimsCsv.add(LordnTaskUtils.getCsvLineForClaimsDomain(domain, domain.getCreationTime()));
+      status = "C";
+    }
+    System.out.printf("%s[%s] ", domain.getDomainName(), status);
+  }
 }

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

@@ -55,7 +55,9 @@ public abstract class MutatingCommand extends ConfirmingCommand implements Comma
 
     /** The possible types of mutation that can be performed on an entity. */
     public enum ChangeType {
-      CREATE, DELETE, UPDATE;
+      CREATE,
+      DELETE,
+      UPDATE;
 
       /** Return the ChangeType corresponding to the given combination of version existences. */
       public static ChangeType get(boolean hasOldVersion, boolean hasNewVersion) {
@@ -78,7 +80,7 @@ public abstract class MutatingCommand extends ConfirmingCommand implements Comma
     /** The key that points to the entity being changed. */
     final VKey<?> key;
 
-    public EntityChange(ImmutableObject oldEntity, ImmutableObject newEntity) {
+    private EntityChange(ImmutableObject oldEntity, ImmutableObject newEntity) {
       type = ChangeType.get(oldEntity != null, newEntity != null);
       checkArgument(
           type != ChangeType.UPDATE || Key.create(oldEntity).equals(Key.create(newEntity)),
@@ -96,6 +98,34 @@ public abstract class MutatingCommand extends ConfirmingCommand implements Comma
               : VKey.createOfy(entity.getClass(), Key.create(entity));
     }
 
+    /**
+     * EntityChange constructor that supports Vkey override. A Vkey is a key of an entity. This is a
+     * workaround to handle cases when a SqlEntity instance does not have a primary key before being
+     * persisted.
+     */
+    private EntityChange(ImmutableObject oldEntity, ImmutableObject newEntity, VKey<?> vkey) {
+      type = ChangeType.get(oldEntity != null, newEntity != null);
+      Key<?> oldKey = Key.create(oldEntity), newKey = Key.create(newEntity);
+      if (type == ChangeType.UPDATE) {
+        checkArgument(
+            oldKey.equals(newKey), "Both entity versions in an update must have the same Key.");
+        checkArgument(
+            oldKey.equals(vkey.getOfyKey()),
+            "The Key of the entity must be the same as the OfyKey of the vkey");
+      } else if (type == ChangeType.CREATE) {
+        checkArgument(
+            newKey.equals(vkey.getOfyKey()),
+            "Both entity versions in an update must have the same Key.");
+      } else if (type == ChangeType.DELETE) {
+        checkArgument(
+            oldKey.equals(vkey.getOfyKey()),
+            "The Key of the entity must be the same as the OfyKey of the vkey");
+      }
+      this.oldEntity = oldEntity;
+      this.newEntity = newEntity;
+      key = vkey;
+    }
+
     /** Returns a human-readable ID string for the entity being changed. */
     public String getEntityId() {
       return String.format(
@@ -110,8 +140,9 @@ public abstract class MutatingCommand extends ConfirmingCommand implements Comma
     public String toString() {
       String changeText;
       if (type == ChangeType.UPDATE) {
-        String diffText = prettyPrintEntityDeepDiff(
-            oldEntity.toDiffableFieldMap(), newEntity.toDiffableFieldMap());
+        String diffText =
+            prettyPrintEntityDeepDiff(
+                oldEntity.toDiffableFieldMap(), newEntity.toDiffableFieldMap());
         changeText = Optional.ofNullable(emptyToNull(diffText)).orElse("[no changes]\n");
       } else {
         changeText = MoreObjects.firstNonNull(oldEntity, newEntity) + "\n";
@@ -205,8 +236,8 @@ public abstract class MutatingCommand extends ConfirmingCommand implements Comma
   }
 
   /**
-   * Subclasses can call this to stage a mutation to an entity that will be applied by execute().
-   * Note that both objects passed must correspond to versions of the same entity with the same key.
+   * Stages an entity change that will be applied by execute(). Both ImmutableObject instances must
+   * be some version of the same entity with the same key.
    *
    * @param oldEntity the existing version of the entity, or null to create a new entity
    * @param newEntity the new version of the entity to save, or null to delete the entity
@@ -222,6 +253,25 @@ public abstract class MutatingCommand extends ConfirmingCommand implements Comma
     lastAddedKey = change.key;
   }
 
+  /**
+   * Stages an entity change which will be applied by execute(), with the support of Vkey override.
+   * It supports cases of SqlEntity instances that do not have primary keys before being persisted.
+   *
+   * @param oldEntity the existing version of the entity, or null to create a new entity
+   * @param newEntity the new version of the entity to save, or null to delete the entity
+   * @param vkey the key of the entity
+   */
+  protected void stageEntityChange(
+      @Nullable ImmutableObject oldEntity, @Nullable ImmutableObject newEntity, VKey vkey) {
+    EntityChange change = new EntityChange(oldEntity, newEntity, vkey);
+    checkArgument(
+        !changedEntitiesMap.containsKey(change.key),
+        "Cannot apply multiple changes for the same entity: %s",
+        change.getEntityId());
+    changedEntitiesMap.put(change.key, change);
+    lastAddedKey = change.key;
+  }
+
   /**
    * Subclasses can call this to write out all previously requested entity changes since the last
    * transaction flush in a transaction.

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

@@ -62,8 +62,6 @@ public final class RegistryTool {
           .put("delete_premium_list", DeletePremiumListCommand.class)
           .put("delete_reserved_list", DeleteReservedListCommand.class)
           .put("delete_tld", DeleteTldCommand.class)
-          .put("deploy_invoicing_pipeline", DeployInvoicingPipelineCommand.class)
-          .put("deploy_spec11_pipeline", DeploySpec11PipelineCommand.class)
           .put("encrypt_escrow_deposit", EncryptEscrowDepositCommand.class)
           .put("execute_epp", ExecuteEppCommand.class)
           .put("generate_allocation_tokens", GenerateAllocationTokensCommand.class)

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

@@ -18,7 +18,6 @@ import dagger.BindsInstance;
 import dagger.Component;
 import dagger.Lazy;
 import google.registry.batch.BatchModule;
-import google.registry.beam.initsql.BeamJpaModule;
 import google.registry.bigquery.BigqueryModule;
 import google.registry.config.CredentialModule.LocalCredentialJson;
 import google.registry.config.RegistryConfig.Config;
@@ -60,7 +59,6 @@ import javax.inject.Singleton;
       AppEngineAdminApiModule.class,
       AuthModule.class,
       BatchModule.class,
-      BeamJpaModule.class,
       BigqueryModule.class,
       ConfigModule.class,
       CloudDnsWriterModule.class,
@@ -107,10 +105,6 @@ interface RegistryToolComponent {
 
   void inject(DeleteContactByRoidCommand command);
 
-  void inject(DeployInvoicingPipelineCommand command);
-
-  void inject(DeploySpec11PipelineCommand command);
-
   void inject(EncryptEscrowDepositCommand command);
 
   void inject(GenerateAllocationTokensCommand command);
@@ -193,8 +187,6 @@ interface RegistryToolComponent {
     @BindsInstance
     Builder sqlAccessInfoFile(@Nullable @Config("sqlAccessInfoFile") String sqlAccessInfoFile);
 
-    Builder beamJpaModule(BeamJpaModule beamJpaModule);
-
     RegistryToolComponent build();
   }
 }

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

@@ -16,6 +16,7 @@ package google.registry.tools;
 
 import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
+import dagger.Lazy;
 import google.registry.rde.RdeReporter;
 import google.registry.tools.params.PathParameter;
 import java.nio.file.Files;
@@ -33,13 +34,12 @@ final class SendEscrowReportToIcannCommand implements CommandWithRemoteApi {
       required = true)
   private List<Path> files;
 
-  @Inject
-  RdeReporter rdeReporter;
+  @Inject Lazy<RdeReporter> rdeReporter;
 
   @Override
   public void run() throws Exception {
     for (Path file : files) {
-      rdeReporter.send(Files.readAllBytes(file));
+      rdeReporter.get().send(Files.readAllBytes(file));
       System.out.printf("Uploaded: %s\n", file);
     }
   }

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

@@ -14,6 +14,8 @@
 
 package google.registry.tools;
 
+import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
+
 import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
 import com.google.common.collect.ImmutableSortedMap;
@@ -23,14 +25,14 @@ import google.registry.model.common.DatabaseTransitionSchedule.PrimaryDatabaseTr
 import google.registry.model.common.DatabaseTransitionSchedule.TransitionId;
 import google.registry.model.common.TimedTransitionProperty;
 import google.registry.tools.params.TransitionListParameter.PrimaryDatabaseTransitions;
-import java.util.Optional;
 import org.joda.time.DateTime;
 
 /** Command to update {@link DatabaseTransitionSchedule}. */
 @Parameters(
     separators = " =",
     commandDescription = "Set the database transition schedule for transition id.")
-public class SetDatabaseTransitionScheduleCommand extends MutatingCommand {
+public class SetDatabaseTransitionScheduleCommand extends ConfirmingCommand
+    implements CommandWithRemoteApi {
 
   @Parameter(
       names = "--transition_schedule",
@@ -43,20 +45,25 @@ public class SetDatabaseTransitionScheduleCommand extends MutatingCommand {
 
   @Parameter(
       names = "--transition_id",
+      required = true,
       description = "Transition id string for the schedule being updated")
   private TransitionId transitionId;
 
   @Override
-  protected void init() {
-    Optional<DatabaseTransitionSchedule> currentSchedule =
-        DatabaseTransitionSchedule.get(transitionId);
+  protected String prompt() {
+    return String.format(
+        "Insert new schedule %s for transition ID %s?", transitionSchedule, transitionId);
+  }
 
+  @Override
+  protected String execute() {
     DatabaseTransitionSchedule newSchedule =
         DatabaseTransitionSchedule.create(
             transitionId,
             TimedTransitionProperty.fromValueMap(
                 transitionSchedule, PrimaryDatabaseTransition.class));
-
-    stageEntityChange(currentSchedule.orElse(null), newSchedule);
+    ofyTm().transact(() -> ofyTm().put(newSchedule));
+    return String.format(
+        "Inserted new schedule %s for transition ID %s.", transitionSchedule, transitionId);
   }
 }

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

@@ -14,40 +14,54 @@
 
 package google.registry.tools;
 
-import static com.google.common.base.Preconditions.checkArgument;
 import static google.registry.util.ListNamingUtils.convertFilePathToName;
 import static java.nio.charset.StandardCharsets.UTF_8;
 
 import com.beust.jcommander.Parameters;
 import com.google.common.base.Strings;
+import com.googlecode.objectify.Key;
 import google.registry.model.registry.label.ReservedList;
-import google.registry.util.SystemClock;
+import google.registry.persistence.VKey;
 import java.nio.file.Files;
 import java.util.List;
-import java.util.Optional;
-import org.joda.time.DateTime;
 
-/** Command to safely update {@link ReservedList} on Datastore. */
-@Parameters(separators = " =", commandDescription = "Update a ReservedList in Datastore.")
+/** Command to safely update {@link ReservedList}. */
+@Parameters(separators = " =", commandDescription = "Update a ReservedList.")
 final class UpdateReservedListCommand extends CreateOrUpdateReservedListCommand {
 
   @Override
   protected void init() throws Exception {
     name = Strings.isNullOrEmpty(name) ? convertFilePathToName(input) : name;
-    Optional<ReservedList> existing = ReservedList.get(name);
-    checkArgument(
-        existing.isPresent(), "Could not update reserved list %s because it doesn't exist.", name);
+    ReservedList existingReservedList =
+        ReservedList.get(name)
+            .orElseThrow(
+                () ->
+                    new IllegalArgumentException(
+                        String.format(
+                            "Could not update reserved list %s because it doesn't exist.", name)));
     boolean shouldPublish =
-        this.shouldPublish == null ? existing.get().getShouldPublish() : this.shouldPublish;
+        this.shouldPublish == null ? existingReservedList.getShouldPublish() : this.shouldPublish;
     List<String> allLines = Files.readAllLines(input, UTF_8);
-    DateTime now = new SystemClock().nowUtc();
     ReservedList.Builder updated =
-        existing
-            .get()
+        existingReservedList
             .asBuilder()
             .setReservedListMapFromLines(allLines)
-            .setLastUpdateTime(now)
             .setShouldPublish(shouldPublish);
     reservedList = updated.build();
+    // only call stageEntityChange if there are changes in entries
+
+    if (!existingReservedList
+        .getReservedListEntries()
+        .equals(reservedList.getReservedListEntries())) {
+      // calls the stageEntityChange method that takes old entity, new entity and a new vkey;
+      // a vkey has to be created here explicitly for ReservedList instances.
+      // ReservedList is a sqlEntity; it triggers the static method Vkey.create(Key<?> ofyCall),
+      // which invokes a static ReservedList.createVkey(Key ofyKey) method that does not exist.
+      // the sql primary key field (revisionId) is only set when it's being persisted;
+      stageEntityChange(
+          existingReservedList,
+          reservedList,
+          VKey.createOfy(ReservedList.class, Key.create(existingReservedList)));
+    }
   }
 }

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

@@ -18,10 +18,8 @@ import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.base.Strings.isNullOrEmpty;
 import static google.registry.util.PreconditionsUtils.checkArgumentPresent;
-import static google.registry.util.X509Utils.encodeX509CertificateFromPemString;
 import static google.registry.util.X509Utils.getCertificateHash;
 import static google.registry.util.X509Utils.loadCertificate;
-import static java.nio.charset.StandardCharsets.US_ASCII;
 
 import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
@@ -30,7 +28,6 @@ import google.registry.flows.certs.CertificateChecker;
 import google.registry.model.registrar.Registrar;
 import google.registry.tools.params.PathParameter;
 import google.registry.util.Clock;
-import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.Optional;
 import javax.annotation.Nullable;
@@ -80,10 +77,7 @@ final class ValidateLoginCredentialsCommand implements CommandWithRemoteApi {
     checkArgument(
         clientCertificatePath == null || isNullOrEmpty(clientCertificateHash),
         "Can't specify both --cert_hash and --cert_file");
-    String encodedCertificate = "";
     if (clientCertificatePath != null) {
-      String certificateString = new String(Files.readAllBytes(clientCertificatePath), US_ASCII);
-      encodedCertificate = encodeX509CertificateFromPemString(certificateString);
       clientCertificateHash = getCertificateHash(loadCertificate(clientCertificatePath));
     }
     Registrar registrar =
@@ -92,10 +86,8 @@ final class ValidateLoginCredentialsCommand implements CommandWithRemoteApi {
     new TlsCredentials(
             true,
             Optional.ofNullable(clientCertificateHash),
-            Optional.ofNullable(encodedCertificate),
             Optional.ofNullable(clientIpAddress),
-            certificateChecker,
-            clock)
+            certificateChecker)
         .validate(registrar, password);
     checkState(
         registrar.isLive(), "Registrar %s has non-live state: %s", clientId, registrar.getState());

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

@@ -15,6 +15,7 @@
 package google.registry.tools.server;
 
 import static com.google.common.flogger.LazyArgs.lazy;
+import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 
 import com.google.common.collect.ImmutableMap;
 import com.google.common.flogger.FluentLogger;
@@ -45,27 +46,29 @@ public abstract class CreateOrUpdatePremiumListAction implements Runnable {
   @Override
   public void run() {
     try {
+      checkArgumentNotNull(inputData, "Input data must not be null");
       save();
     } catch (IllegalArgumentException e) {
       logger.atInfo().withCause(e).log(
           "Usage error in attempting to save premium list from nomulus tool command");
-      response.setPayload(ImmutableMap.of("error", e.toString(), "status", "error"));
+      response.setPayload(ImmutableMap.of("error", e.getMessage(), "status", "error"));
     } catch (Exception e) {
       logger.atSevere().withCause(e).log(
           "Unexpected error saving premium list to Datastore from nomulus tool command");
-      response.setPayload(ImmutableMap.of("error", e.toString(), "status", "error"));
+      response.setPayload(ImmutableMap.of("error", e.getMessage(), "status", "error"));
     }
   }
 
   /** Logs the premium list data at INFO, truncated if too long. */
   void logInputData() {
+    String logData = (inputData == null) ? "(null)" : inputData;
     logger.atInfo().log(
         "Received the following input data: %s",
         lazy(
             () ->
-                (inputData.length() < MAX_LOGGING_PREMIUM_LIST_LENGTH)
-                    ? inputData
-                    : (inputData.substring(0, MAX_LOGGING_PREMIUM_LIST_LENGTH) + "<truncated>")));
+                (logData.length() < MAX_LOGGING_PREMIUM_LIST_LENGTH)
+                    ? logData
+                    : (logData.substring(0, MAX_LOGGING_PREMIUM_LIST_LENGTH) + "<truncated>")));
   }
 
   /** Saves the premium list to both Datastore and Cloud SQL. */

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

@@ -51,9 +51,12 @@ public class CreatePremiumListAction extends CreateOrUpdatePremiumListAction {
   @Override
   protected void save() {
     checkArgument(
-        !PremiumListDualDao.exists(name), "A premium list of this name already exists: %s.", name);
+        !PremiumListDualDao.exists(name), "A premium list of this name already exists: %s", name);
     if (!override) {
-      assertTldExists(name);
+      assertTldExists(
+          name,
+          "Premium names must match the name of the TLD they are intended to be used on"
+              + " (unless --override is specified), yet TLD %s does not exist");
     }
     logger.atInfo().log("Saving premium list for TLD %s", name);
     logInputData();

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

@@ -15,8 +15,7 @@
 package google.registry.tools.server;
 
 import static com.google.common.collect.ImmutableSortedSet.toImmutableSortedSet;
-import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
-import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.request.Action.Method.GET;
 import static google.registry.request.Action.Method.POST;
 
@@ -51,14 +50,15 @@ public final class ListPremiumListsAction extends ListObjectsAction<PremiumList>
 
   @Override
   public ImmutableSet<PremiumList> loadObjects() {
-    return transactIfJpaTm(
-        () ->
-            tm().loadAllOf(PremiumList.class).stream()
-                .map(PremiumList::getName)
-                .map(PremiumListDualDao::getLatestRevision)
-                .filter(Optional::isPresent)
-                .map(Optional::get)
-                .peek(list -> Hibernate.initialize(list.getLabelsToPrices()))
-                .collect(toImmutableSortedSet(Comparator.comparing(PremiumList::getName))));
+    return jpaTm()
+        .transact(
+            () ->
+                jpaTm().loadAllOf(PremiumList.class).stream()
+                    .map(PremiumList::getName)
+                    .map(PremiumListDualDao::getLatestRevision)
+                    .filter(Optional::isPresent)
+                    .map(Optional::get)
+                    .peek(list -> Hibernate.initialize(list.getLabelsToPrices()))
+                    .collect(toImmutableSortedSet(Comparator.comparing(PremiumList::getName))));
   }
 }

core/src/main/resources/google/registry/beam/init_sql_pipeline_metadata.json

@@ -5,7 +5,7 @@
     {
       "name": "registryEnvironment",
       "label": "The Registry environment.",
-      "helpText": "The Registry environment, required if environment-specific initialization is needed on worker VMs.",
+      "helpText": "The Registry environment, required if environment-specific initialization (such as JPA) is needed on worker VMs.",
       "is_optional": true,
       "regexes": [
         "^[0-9A-Z_]+$"