Set up deployment of the Spec11 pipeline with JPA TM (#716)

* Set up deployment of the Spec11 pipeline with JPA TM

* Remove unnecessarily pipeline options setting

* Use enviroment name in BeamJpaModuleTest

* Fix checkstyle error

.gitignore

@@ -79,6 +79,10 @@ nomulus.iml
 nomulus.ipr
 nomulus.iws
 
+# Auto-generated java classes by Intellij
+*/src/main/generated/
+*/src/test/generated_tests/
+
 # VScode
 .vscode
 

common/src/testing/java/google/registry/testing/truth/TextDiffSubject.java

@@ -16,6 +16,7 @@ package google.registry.testing.truth;
 
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.truth.Truth.assertAbout;
+import static com.google.common.truth.Truth.assertWithMessage;
 import static java.nio.charset.StandardCharsets.UTF_8;
 
 import com.github.difflib.DiffUtils;
@@ -31,6 +32,7 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.io.Resources;
 import com.google.common.truth.Fact;
 import com.google.common.truth.FailureMetadata;
+import com.google.common.truth.SimpleSubjectBuilder;
 import com.google.common.truth.Subject;
 import java.io.IOException;
 import java.net.URL;
@@ -68,6 +70,15 @@ public class TextDiffSubject extends Subject {
     this.actual = ImmutableList.copyOf(actual);
   }
 
+  protected TextDiffSubject(FailureMetadata metadata, URL actual) {
+    super(metadata, actual);
+    try {
+      this.actual = ImmutableList.copyOf(Resources.asCharSource(actual, UTF_8).readLines());
+    } catch (IOException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
   public TextDiffSubject withDiffFormat(DiffFormat format) {
     this.diffFormat = format;
     return this;
@@ -100,6 +111,11 @@ public class TextDiffSubject extends Subject {
     return assertThat(Resources.asCharSource(resourceUrl, UTF_8).readLines());
   }
 
+  public static SimpleSubjectBuilder<TextDiffSubject, URL> assertWithMessageAboutUrlSource(
+      String format, Object... params) {
+    return assertWithMessage(format, params).about(urlFactory());
+  }
+
   private static final Subject.Factory<TextDiffSubject, ImmutableList<String>>
       TEXT_DIFF_SUBJECT_TEXT_FACTORY = TextDiffSubject::new;
 
@@ -107,6 +123,13 @@ public class TextDiffSubject extends Subject {
     return TEXT_DIFF_SUBJECT_TEXT_FACTORY;
   }
 
+  private static final Subject.Factory<TextDiffSubject, URL> TEXT_DIFF_SUBJECT_URL_FACTORY =
+      TextDiffSubject::new;
+
+  public static Subject.Factory<TextDiffSubject, URL> urlFactory() {
+    return TEXT_DIFF_SUBJECT_URL_FACTORY;
+  }
+
   static String generateUnifiedDiff(
       ImmutableList<String> expectedContent, ImmutableList<String> actualContent) {
     Patch<String> diff;

config/nom_build.py

@@ -19,6 +19,7 @@ import argparse
 import attr
 import io
 import os
+import shutil
 import subprocess
 import sys
 from typing import List, Union
@@ -49,15 +50,30 @@ PROPERTIES_HEADER = """\
 # This file defines properties used by the gradle build.  It must be kept in
 # sync with config/nom_build.py.
 #
-# To regenerate, run config/nom_build.py --generate-gradle-properties
+# To regenerate, run ./nom_build --generate-gradle-properties
 #
 # To view property descriptions (which are command line flags for
-# nom_build), run config/nom_build.py --help.
+# nom_build), run ./nom_build --help.
 #
 # DO NOT EDIT THIS FILE BY HAND
 org.gradle.jvmargs=-Xmx1024m
 """
 
+# Help text to be displayed (in addition to the synopsis and flag help, which
+# are displayed automatically).
+HELP_TEXT = """\
+A wrapper around the gradle build that provides the following features:
+
+-   Converts properties into flags to guard against property name spelling errors
+    and to provide help descriptions for all properties.
+-   Provides pseudo-commands (with the ":nom:" prefix) that encapsulate common
+    actions that are difficult to implement in gradle.
+
+Pseudo-commands:
+    :nom:generate_golden_file - regenerates the golden file from the current
+        set of flyway files.
+"""
+
 # Define all of our special gradle properties here.
 PROPERTIES = [
     Property('mavenUrl',
@@ -114,6 +130,11 @@ PROPERTIES = [
     Property('nomulus_version',
              'The version of nomulus to test against in a database '
              'integration test.'),
+    Property('dot_path',
+             'The path to "dot", part of the graphviz package that converts '
+             'a BEAM pipeline to image. Setting this property to empty string '
+             'will disable image generation.',
+             '/usr/bin/dot'),
 ]
 
 GRADLE_FLAGS = [
@@ -251,8 +272,42 @@ def get_root() -> str:
     return cur_dir
 
 
-def main(args):
-    parser = argparse.ArgumentParser('nom_build')
+class Abort(Exception):
+    """Raised to terminate the process with a non-zero error code.
+
+    Parameters are ignored.
+    """
+
+
+def do_pseudo_task(task: str) -> None:
+    root = get_root()
+    if task == ':nom:generate_golden_file':
+        if not subprocess.call([f'{root}/gradlew', ':db:test']):
+            print('\033[33mWARNING:\033[0m Golden schema appears to be '
+                  'up-to-date.  If you are making schema changes,  be sure to '
+                  'add a flyway file for them.')
+            return
+        print('\033[33mWARNING:\033[0m Ignore the above failure, it is '
+              'expected.')
+
+        # Copy the new schema into place.
+        shutil.copy(f'{root}/db/build/resources/test/testcontainer/'
+                    'mount/dump.txt',
+                    f'{root}/db/src/main/resources/sql/schema/'
+                    'nomulus.golden.sql')
+
+        if subprocess.call([f'{root}/gradlew', ':db:test']):
+            print('\033[31mERROR:\033[0m Golden file test failed after '
+                  'copying schema.  Please check your flyway files.')
+            raise Abort()
+    else:
+        print(f'\033[31mERROR:\033[0m Unknown task {task}')
+        raise Abort()
+
+
+def main(args) -> int:
+    parser = argparse.ArgumentParser('nom_build', description=HELP_TEXT,
+                                     formatter_class=argparse.RawTextHelpFormatter)
     for prop in PROPERTIES:
         parser.add_argument('--' + prop.name, default=prop.default,
                             help=prop.desc)
@@ -291,7 +346,7 @@ def main(args):
     if args.generate_gradle_properties:
         with open(f'{root}/gradle.properties', 'w') as dst:
             dst.write(gradle_properties)
-        return
+        return 0
 
     # Verify that the gradle properties file is what we expect it to be.
     with open(f'{root}/gradle.properties') as src:
@@ -316,12 +371,39 @@ def main(args):
             if flag.has_arg:
                 gradle_command.append(arg_val)
 
+    # See if there are any special ":nom:" pseudo-tasks specified.
+    got_non_pseudo_tasks = False
+    for arg in args.non_flag_args[1:]:
+        if arg.startswith(':nom:'):
+            if got_non_pseudo_tasks:
+                # We can't currently deal with the situation of gradle tasks
+                # before pseudo-tasks.  This could be implemented by invoking
+                # gradle for only the set of gradle tasks before the pseudo
+                # task, but that's overkill for now.
+                print(f'\033[31mERROR:\033[0m Pseudo task ({arg}) must be '
+                      'specified prior to all actual gradle tasks.  Aborting.')
+                return 1
+            do_pseudo_task(arg)
+        else:
+            got_non_pseudo_tasks = True
+    non_flag_args = [
+        arg for arg in args.non_flag_args[1:] if not arg.startswith(':nom:')]
+
+    if not non_flag_args:
+        if not got_non_pseudo_tasks:
+            print('\033[33mWARNING:\033[0m No tasks specified.  Not '
+                  'doing anything')
+        return 0
+
     # Add the non-flag args (we exclude the first, which is the command name
     # itself) and run.
-    gradle_command.extend(args.non_flag_args[1:])
-    subprocess.call(gradle_command)
+    gradle_command.extend(non_flag_args)
+    return subprocess.call(gradle_command)
 
 
 if __name__ == '__main__':
-    main(sys.argv)
+    try:
+        sys.exit(main(sys.argv))
+    except Abort as ex:
+        sys.exit(1)
 

config/nom_build_test.py

@@ -14,6 +14,7 @@
 
 import io
 import os
+import shutil
 import unittest
 from unittest import mock
 import nom_build
@@ -67,6 +68,7 @@ class MyTest(unittest.TestCase):
             mock.patch.object(nom_build, 'print', self.print_fake).start())
 
         self.call_mock = mock.patch.object(subprocess, 'call').start()
+        self.copy_mock = mock.patch.object(shutil, 'copy').start()
 
         self.file_contents = {
             # Prefil with the actual file contents.
@@ -92,17 +94,32 @@ class MyTest(unittest.TestCase):
 
     def test_no_args(self):
         nom_build.main(['nom_build'])
-        self.assertEqual(self.printed, [])
-        self.call_mock.assert_called_with([GRADLEW])
+        self.assertEqual(self.printed,
+                         ['\x1b[33mWARNING:\x1b[0m No tasks specified.  Not '
+                         'doing anything'])
 
     def test_property_calls(self):
-        nom_build.main(['nom_build', '--testFilter=foo'])
-        self.call_mock.assert_called_with([GRADLEW, '-P', 'testFilter=foo'])
+        nom_build.main(['nom_build', 'task-name', '--testFilter=foo'])
+        self.call_mock.assert_called_with([GRADLEW, '-P', 'testFilter=foo',
+                                          'task-name'])
 
     def test_gradle_flags(self):
-        nom_build.main(['nom_build', '-d', '-b', 'foo'])
+        nom_build.main(['nom_build', 'task-name', '-d', '-b', 'foo'])
         self.call_mock.assert_called_with([GRADLEW, '--build-file', 'foo',
-                                          '--debug'])
+                                          '--debug', 'task-name'])
+
+    def test_generate_golden_file(self):
+        self.call_mock.side_effect = [1, 0]
+        nom_build.main(['nom_build', ':nom:generate_golden_file'])
+        self.call_mock.assert_has_calls([
+            mock.call([GRADLEW, ':db:test']),
+            mock.call([GRADLEW, ':db:test'])
+        ])
+
+    def test_generate_golden_file_nofail(self):
+        self.call_mock.return_value = 0
+        nom_build.main(['nom_build', ':nom:generate_golden_file'])
+        self.call_mock.assert_has_calls([mock.call([GRADLEW, ':db:test'])])
 
 unittest.main()
 

config/presubmits.py

@@ -22,7 +22,7 @@ import sys
 import re
 
 # We should never analyze any generated files
-UNIVERSALLY_SKIPPED_PATTERNS = {"/build/", "cloudbuild-caches", "/out/"}
+UNIVERSALLY_SKIPPED_PATTERNS = {"/build/", "cloudbuild-caches", "/out/", ".git/"}
 # We can't rely on CI to have the Enum package installed so we do this instead.
 FORBIDDEN = 1
 REQUIRED = 2

core/build.gradle

@@ -238,6 +238,7 @@ dependencies {
   compile deps['jline:jline']
   compile deps['joda-time:joda-time']
   compile deps['org.apache.avro:avro']
+  testCompile deps['org.apache.beam:beam-runners-core-construction-java']
   testCompile deps['org.apache.beam:beam-runners-direct-java']
   compile deps['org.apache.beam:beam-runners-google-cloud-dataflow-java']
   compile deps['org.apache.beam:beam-sdks-java-core']
@@ -256,6 +257,7 @@ dependencies {
   compile deps['org.bouncycastle:bcpg-jdk15on']
   testCompile deps['org.bouncycastle:bcpkix-jdk15on']
   compile deps['org.bouncycastle:bcprov-jdk15on']
+  testCompile deps['com.fasterxml.jackson.core:jackson-databind']
   runtime deps['org.glassfish.jaxb:jaxb-runtime']
   compile deps['org.hibernate:hibernate-core']
   compile deps['org.joda:joda-money']
@@ -967,6 +969,49 @@ task buildToolImage(dependsOn: nomulus, type: Exec) {
   commandLine 'docker', 'build', '-t', 'nomulus-tool', '.'
 }
 
+task generateInitSqlPipelineGraph(type: Test) {
+  include "**/InitSqlPipelineGraphTest.*"
+  testNameIncludePatterns = ["**createPipeline_compareGraph"]
+  ignoreFailures = true
+}
+
+task updateInitSqlPipelineGraph(type: Copy) {
+  def graphRelativePath = 'google/registry/beam/initsql/'
+  from ("${projectDir}/build/resources/test/${graphRelativePath}") {
+    include 'pipeline_curr.dot'
+    rename 'curr', 'golden'
+  }
+  into "src/test/resources/${graphRelativePath}"
+
+  dependsOn generateInitSqlPipelineGraph
+
+  doLast {
+    if (com.google.common.base.Strings.isNullOrEmpty(project.dot_path)) {
+      getLogger().info('Property dot_path is null. Not creating image for pipeline graph.')
+    }
+    def dotPath = project.dot_path
+    if (!new File(dotPath).exists()) {
+      throw new RuntimeException(
+          """\
+              ${dotPath} not found. Make sure graphviz is installed
+              and the dot_path property is set correctly."""
+              .stripIndent())
+    }
+    def goldenGraph = "src/test/resources/${graphRelativePath}/pipeline_golden.dot"
+    def goldenImage = "src/test/resources/${graphRelativePath}/pipeline_golden.png"
+    def cmd = "${dotPath} -Tpng -o \"${goldenImage}\" \"${goldenGraph}\""
+    try {
+      rootProject.ext.execInBash(cmd, projectDir)
+    } catch (Throwable throwable) {
+      throw new RuntimeException(
+          """\
+              Failed to generate golden image with command ${cmd}
+              Error: ${throwable.getMessage()}
+          """)
+    }
+  }
+}
+
 // Build the devtool jar.
 createUberJar(
     'devtool',

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

@@ -22,21 +22,35 @@ import java.lang.reflect.Method;
 import java.lang.reflect.Proxy;
 
 /**
- * Sets up a placeholder {@link Environment} on a non-AppEngine platform so that Datastore Entities
- * can be converted from/to Objectify entities. See {@code DatastoreEntityExtension} in test source
- * for more information.
+ * Sets up a fake {@link Environment} so that the following operations can be performed without the
+ * Datastore service:
+ *
+ * <ul>
+ *   <li>Create Objectify {@code Keys}.
+ *   <li>Instantiate Objectify objects.
+ *   <li>Convert Datastore {@code Entities} to their corresponding Objectify objects.
+ * </ul>
+ *
+ * <p>User has the option to specify their desired {@code appId} string, which forms part of an
+ * Objectify {@code Key} and is included in the equality check. This feature makes it easy to
+ * compare a migrated object in SQL with the original in Objectify.
+ *
+ * <p>Note that conversion from Objectify objects to Datastore {@code Entities} still requires the
+ * Datastore service.
  */
 public class AppEngineEnvironment implements Closeable {
 
-  private static final Environment PLACEHOLDER_ENV = createAppEngineEnvironment();
-
   private boolean isPlaceHolderNeeded;
 
   public AppEngineEnvironment() {
+    this("PlaceholderAppId");
+  }
+
+  public AppEngineEnvironment(String appId) {
     isPlaceHolderNeeded = ApiProxy.getCurrentEnvironment() == null;
     // isPlaceHolderNeeded may be true when we are invoked in a test with AppEngineRule.
     if (isPlaceHolderNeeded) {
-      ApiProxy.setEnvironmentForCurrentThread(PLACEHOLDER_ENV);
+      ApiProxy.setEnvironmentForCurrentThread(createAppEngineEnvironment(appId));
     }
   }
 
@@ -48,7 +62,7 @@ public class AppEngineEnvironment implements Closeable {
   }
 
   /** Returns a placeholder {@link Environment} that can return hardcoded AppId and Attributes. */
-  private static Environment createAppEngineEnvironment() {
+  private static Environment createAppEngineEnvironment(String appId) {
     return (Environment)
         Proxy.newProxyInstance(
             Environment.class.getClassLoader(),
@@ -56,7 +70,7 @@ public class AppEngineEnvironment implements Closeable {
             (Object proxy, Method method, Object[] args) -> {
               switch (method.getName()) {
                 case "getAppId":
-                  return "PlaceholderAppId";
+                  return appId;
                 case "getAttributes":
                   return ImmutableMap.<String, Object>of();
                 default:

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

@@ -169,10 +169,12 @@ public final class AsyncTaskEnqueuer {
         lock.getRelockDuration().isPresent(),
         "Lock with ID %s not configured for relock",
         lock.getRevisionId());
+    String backendHostname = appEngineServiceUtils.getServiceHostname("backend");
     addTaskToQueueWithRetry(
         asyncActionsPushQueue,
         TaskOptions.Builder.withUrl(RelockDomainAction.PATH)
             .method(Method.POST)
+            .header("Host", backendHostname)
             .param(
                 RelockDomainAction.OLD_UNLOCK_REVISION_ID_PARAM,
                 String.valueOf(lock.getRevisionId()))

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

@@ -532,7 +532,7 @@ public class DeleteContactsAndHostsAction implements Runnable {
           resource.getClass().getSimpleName());
       return new AutoValue_DeleteContactsAndHostsAction_DeletionRequest.Builder()
           .setKey(resourceKey)
-          .setLastUpdateTime(resource.getUpdateAutoTimestamp().getTimestamp())
+          .setLastUpdateTime(resource.getUpdateTimestamp().getTimestamp())
           .setRequestingClientId(
               checkNotNull(
                   params.get(PARAM_REQUESTING_CLIENT_ID), "Requesting client id not specified"))

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

@@ -319,13 +319,13 @@ public class RefreshDnsOnHostRenameAction implements Runnable {
       HostResource host =
           checkNotNull(ofy().load().key(hostKey).now(), "Host to refresh doesn't exist");
       boolean isHostDeleted =
-          isDeleted(host, latestOf(now, host.getUpdateAutoTimestamp().getTimestamp()));
+          isDeleted(host, latestOf(now, host.getUpdateTimestamp().getTimestamp()));
       if (isHostDeleted) {
         logger.atInfo().log("Host %s is already deleted, not refreshing DNS.", hostKey);
       }
       return new AutoValue_RefreshDnsOnHostRenameAction_DnsRefreshRequest.Builder()
           .setHostKey(hostKey)
-          .setLastUpdateTime(host.getUpdateAutoTimestamp().getTimestamp())
+          .setLastUpdateTime(host.getUpdateTimestamp().getTimestamp())
           .setRequestedTime(
               DateTime.parse(
                   checkNotNull(params.get(PARAM_REQUESTED_TIME), "Requested time not specified")))

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

@@ -19,7 +19,6 @@ import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.base.Strings.isNullOrEmpty;
 
 import com.google.common.base.Splitter;
-import dagger.Binds;
 import dagger.Component;
 import dagger.Lazy;
 import dagger.Module;
@@ -32,10 +31,6 @@ import google.registry.persistence.PersistenceModule;
 import google.registry.persistence.PersistenceModule.JdbcJpaTm;
 import google.registry.persistence.PersistenceModule.SocketFactoryJpaTm;
 import google.registry.persistence.transaction.JpaTransactionManager;
-import google.registry.util.Clock;
-import google.registry.util.Sleeper;
-import google.registry.util.SystemClock;
-import google.registry.util.SystemSleeper;
 import google.registry.util.UtilsModule;
 import java.io.BufferedReader;
 import java.io.IOException;
@@ -59,37 +54,39 @@ public class BeamJpaModule {
 
   private static final String GCS_SCHEME = "gs://";
 
-  @Nullable private final String credentialFilePath;
+  @Nullable private final String sqlAccessInfoFile;
+  @Nullable private final String cloudKmsProjectId;
 
   /**
    * Constructs a new instance of {@link BeamJpaModule}.
    *
    * <p>Note: it is an unfortunately necessary antipattern to check for the validity of
-   * credentialFilePath in {@link #provideCloudSqlAccessInfo} rather than in the constructor.
+   * 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 credentialFilePath the path to a Cloud SQL credential file. This must refer to either a
+   * @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.
    */
-  public BeamJpaModule(@Nullable String credentialFilePath) {
-    this.credentialFilePath = credentialFilePath;
+  public BeamJpaModule(@Nullable String sqlAccessInfoFile, @Nullable String cloudKmsProjectId) {
+    this.sqlAccessInfoFile = sqlAccessInfoFile;
+    this.cloudKmsProjectId = cloudKmsProjectId;
   }
 
   /** Returns true if the credential file is on GCS (and therefore expected to be encrypted). */
   private boolean isCloudSqlCredential() {
-    return credentialFilePath.startsWith(GCS_SCHEME);
+    return sqlAccessInfoFile.startsWith(GCS_SCHEME);
   }
 
   @Provides
   @Singleton
   SqlAccessInfo provideCloudSqlAccessInfo(Lazy<CloudSqlCredentialDecryptor> lazyDecryptor) {
-    checkArgument(!isNullOrEmpty(credentialFilePath), "Null or empty credentialFilePath");
+    checkArgument(!isNullOrEmpty(sqlAccessInfoFile), "Null or empty credentialFilePath");
     String line = readOnlyLineFromCredentialFile();
     if (isCloudSqlCredential()) {
       line = lazyDecryptor.get().decrypt(line);
@@ -106,7 +103,7 @@ public class BeamJpaModule {
 
   String readOnlyLineFromCredentialFile() {
     try {
-      ResourceId resourceId = FileSystems.matchSingleFileSpec(credentialFilePath).resourceId();
+      ResourceId resourceId = FileSystems.matchSingleFileSpec(sqlAccessInfoFile).resourceId();
       try (BufferedReader reader =
           new BufferedReader(
               new InputStreamReader(
@@ -146,8 +143,8 @@ public class BeamJpaModule {
 
   @Provides
   @Config("beamCloudKmsProjectId")
-  static String kmsProjectId() {
-    return "domain-registry-dev";
+  String kmsProjectId() {
+    return cloudKmsProjectId;
   }
 
   @Provides
@@ -159,19 +156,10 @@ public class BeamJpaModule {
   @Provides
   @Config("beamHibernateHikariMaximumPoolSize")
   static int getBeamHibernateHikariMaximumPoolSize() {
+    // TODO(weiminyu): make this configurable. Should be equal to number of cores.
     return 4;
   }
 
-  @Module
-  interface BindModule {
-
-    @Binds
-    Sleeper sleeper(SystemSleeper sleeper);
-
-    @Binds
-    Clock clock(SystemClock clock);
-  }
-
   @Singleton
   @Component(
       modules = {

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

@@ -0,0 +1,75 @@
+// 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.checkNotNull;
+
+import com.google.appengine.api.datastore.Entity;
+import java.util.Objects;
+
+/** Helper for manipulating {@code DomainBase} when migrating from Datastore to SQL database */
+final class DomainBaseUtil {
+
+  private DomainBaseUtil() {}
+
+  /**
+   * Removes {@link google.registry.model.billing.BillingEvent.Recurring}, {@link
+   * google.registry.model.poll.PollMessage PollMessages} and {@link
+   * google.registry.model.host.HostResource name servers} from a Datastore {@link Entity} that
+   * represents an Ofy {@link google.registry.model.domain.DomainBase}. This breaks the cycle of
+   * foreign key constraints between these entity kinds, allowing {@code DomainBases} to be inserted
+   * into the SQL database. See {@link InitSqlPipeline} for a use case, where the full {@code
+   * DomainBases} are written again during the last stage of the pipeline.
+   *
+   * <p>The returned object may be in bad state. Specifically, {@link
+   * google.registry.model.eppcommon.StatusValue#INACTIVE} is not added after name servers are
+   * removed. This only impacts tests.
+   *
+   * <p>This operation is performed on an Datastore {@link Entity} instead of Ofy Java object
+   * because Objectify requires access to a Datastore service when converting an Ofy object to a
+   * Datastore {@code Entity}. If we insist on working with Objectify objects, we face a few
+   * unsatisfactory options:
+   *
+   * <ul>
+   *   <li>Connect to our production Datastore, which incurs unnecessary security and code health
+   *       risk.
+   *   <li>Connect to a separate real Datastore instance, which is a waster and overkill.
+   *   <li>Use an in-memory test Datastore, which is a project health risk in that the test
+   *       Datastore would be added to Nomulus' production binary unless we create a separate
+   *       project for this pipeline.
+   * </ul>
+   *
+   * <p>Given our use case, operating on Datastore entities is the best option.
+   *
+   * @throws IllegalArgumentException if input does not represent a DomainBase
+   */
+  static Entity removeBillingAndPollAndHosts(Entity domainBase) {
+    checkNotNull(domainBase, "domainBase");
+    checkArgument(
+        Objects.equals(domainBase.getKind(), "DomainBase"),
+        "Expecting DomainBase, got %s",
+        domainBase.getKind());
+    Entity clone = domainBase.clone();
+    clone.removeProperty("autorenewBillingEvent");
+    clone.removeProperty("autorenewPollMessage");
+    clone.removeProperty("deletePollMessage");
+    clone.removeProperty("nsHosts");
+    domainBase.getProperties().keySet().stream()
+        .filter(s -> s.startsWith("transferData."))
+        .forEach(s -> clone.removeProperty(s));
+    return clone;
+  }
+}

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

@@ -0,0 +1,237 @@
+// 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 com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+import com.googlecode.objectify.Key;
+import google.registry.backup.AppEngineEnvironment;
+import google.registry.backup.VersionedEntity;
+import google.registry.beam.initsql.BeamJpaModule.JpaTransactionManagerComponent;
+import google.registry.beam.initsql.Transforms.RemoveDomainBaseForeignKeys;
+import google.registry.model.billing.BillingEvent;
+import google.registry.model.contact.ContactResource;
+import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.token.AllocationToken;
+import google.registry.model.host.HostResource;
+import google.registry.model.poll.PollMessage;
+import google.registry.model.registrar.Registrar;
+import google.registry.model.registrar.RegistrarContact;
+import google.registry.model.registry.Registry;
+import google.registry.model.reporting.HistoryEntry;
+import google.registry.persistence.transaction.JpaTransactionManager;
+import java.io.Serializable;
+import java.util.Collection;
+import java.util.Optional;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.PipelineResult;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.Wait;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.TupleTag;
+import org.joda.time.DateTime;
+
+/**
+ * A BEAM pipeline that populates a SQL database with data from a Datastore backup.
+ *
+ * <p>This pipeline migrates EPP resources and related entities that cross-reference each other. To
+ * avoid violating foreign key constraints, writes to SQL are ordered by entity kinds. In addition,
+ * the {@link DomainBase} kind is written twice (see details below). The write order is presented
+ * below. Although some kinds can be written concurrently, e.g. {@code ContactResource} and {@code
+ * RegistrarContact}, we do not expect any performance benefit since the limiting resource is the
+ * number of JDBC connections. Google internal users may refer to <a
+ * href="http://go/registry-r3-init-sql">the design doc</a> for more information.
+ *
+ * <ol>
+ *   <li>{@link Registry}: Assumes that {@code PremiumList} and {@code ReservedList} have been set
+ *       up in the SQL database.
+ *   <li>{@link Registrar}: Logically depends on {@code Registry}, Foreign key not modeled yet.
+ *   <li>{@link ContactResource}: references {@code Registrar}
+ *   <li>{@link RegistrarContact}: references {@code Registrar}.
+ *   <li>Cleansed {@link DomainBase}: with references to {@code BillingEvent}, {@code Recurring},
+ *       {@code Cancellation} and {@code HostResource} removed, still references {@code Registrar}
+ *       and {@code ContactResource}. The removal breaks circular Foreign Key references.
+ *   <li>{@link HostResource}: references {@code DomainBase}.
+ *   <li>{@link HistoryEntry}: maps to one of three SQL entity types and may reference {@code
+ *       Registrar}, {@code ContactResource}, {@code HostResource}, and {@code DomainBase}.
+ *   <li>{@link AllocationToken}: references {@code HistoryEntry}.
+ *   <li>{@link BillingEvent.Recurring}: references {@code Registrar}, {@code DomainBase} and {@code
+ *       HistoryEntry}.
+ *   <li>{@link BillingEvent.OneTime}: references {@code Registrar}, {@code DomainBase}, {@code
+ *       BillingEvent.Recurring}, {@code HistoryEntry} and {@code AllocationToken}.
+ *   <li>{@link BillingEvent.Modification}: SQL model TBD. Will reference {@code Registrar}, {@code
+ *       DomainBase} and {@code BillingEvent.OneTime}.
+ *   <li>{@link BillingEvent.Cancellation}: references {@code Registrar}, {@code DomainBase}, {@code
+ *       BillingEvent.Recurring}, {@code BillingEvent.OneTime}, and {@code HistoryEntry}.
+ *   <li>{@link PollMessage}: references {@code Registrar}, {@code DomainBase}, {@code
+ *       ContactResource}, {@code HostResource}, and {@code HistoryEntry}.
+ *   <li>{@link DomainBase}, original copy from Datastore.
+ * </ol>
+ */
+public class InitSqlPipeline implements Serializable {
+
+  /**
+   * Datastore kinds to be written to the SQL database before the cleansed version of {@link
+   * DomainBase}.
+   */
+  // TODO(weiminyu): include Registry.class when it is modeled in JPA.
+  private static final ImmutableList<Class<?>> PHASE_ONE_ORDERED =
+      ImmutableList.of(Registrar.class, ContactResource.class);
+
+  /**
+   * Datastore kinds to be written to the SQL database after the cleansed version of {@link
+   * DomainBase}.
+   *
+   * <p>The following entities are missing from the list:
+   *
+   * <ul>
+   *   <li>Those not modeled in JPA yet, e.g., {@code BillingEvent.Modification}.
+   *   <li>Those waiting for sanitation, e.g., {@code HistoryEntry}, which would have duplicate keys
+   *       after converting to SQL model.
+   *   <li>Those that have foreign key constraints on the above.
+   * </ul>
+   */
+  // TODO(weiminyu): add more entities when available.
+  private static final ImmutableList<Class<?>> PHASE_TWO_ORDERED =
+      ImmutableList.of(HostResource.class);
+
+  private final InitSqlPipelineOptions options;
+
+  private final Pipeline pipeline;
+
+  private final SerializableFunction<JpaTransactionManagerComponent, JpaTransactionManager>
+      jpaGetter;
+
+  InitSqlPipeline(InitSqlPipelineOptions options) {
+    this.options = options;
+    pipeline = Pipeline.create(options);
+    jpaGetter = JpaTransactionManagerComponent::cloudSqlJpaTransactionManager;
+  }
+
+  @VisibleForTesting
+  InitSqlPipeline(InitSqlPipelineOptions options, Pipeline pipeline) {
+    this.options = options;
+    this.pipeline = pipeline;
+    jpaGetter = JpaTransactionManagerComponent::localDbJpaTransactionManager;
+  }
+
+  public PipelineResult run() {
+    setupPipeline();
+    return pipeline.run();
+  }
+
+  @VisibleForTesting
+  void setupPipeline() {
+    PCollectionTuple datastoreSnapshot =
+        pipeline.apply(
+            "Load Datastore snapshot",
+            Transforms.loadDatastoreSnapshot(
+                options.getDatastoreExportDir(),
+                options.getCommitLogDir(),
+                DateTime.parse(options.getCommitLogStartTimestamp()),
+                DateTime.parse(options.getCommitLogEndTimestamp()),
+                ImmutableSet.<String>builder()
+                    .add("DomainBase")
+                    .addAll(toKindStrings(PHASE_ONE_ORDERED))
+                    .addAll(toKindStrings(PHASE_TWO_ORDERED))
+                    .build()));
+
+    // Set up the pipeline to write entity kinds from PHASE_ONE_ORDERED to SQL. Return a object
+    // that signals the completion of the phase.
+    PCollection<Void> blocker =
+        scheduleOnePhaseWrites(datastoreSnapshot, PHASE_ONE_ORDERED, Optional.empty(), null);
+    blocker =
+        writeToSql(
+            "DomainBase without circular foreign keys",
+            removeDomainBaseForeignKeys(datastoreSnapshot)
+                .apply("Wait on phase one", Wait.on(blocker)));
+    // Set up the pipeline to write entity kinds from PHASE_TWO_ORDERED to SQL. This phase won't
+    // start until all cleansed DomainBases have been written (started by line above).
+    scheduleOnePhaseWrites(
+        datastoreSnapshot, PHASE_TWO_ORDERED, Optional.of(blocker), "DomainBaseNoFkeys");
+  }
+
+  private PCollection<VersionedEntity> removeDomainBaseForeignKeys(
+      PCollectionTuple datastoreSnapshot) {
+    PCollection<VersionedEntity> domainBases =
+        datastoreSnapshot.get(Transforms.createTagForKind("DomainBase"));
+    return domainBases.apply(
+        "Remove circular foreign keys from DomainBase",
+        ParDo.of(new RemoveDomainBaseForeignKeys()));
+  }
+
+  /**
+   * Sets up the pipeline to write entities in {@code entityClasses} to SQL. Entities are written
+   * one kind at a time based on each kind's position in {@code entityClasses}. Concurrency exists
+   * within each kind.
+   *
+   * @param datastoreSnapshot the Datastore snapshot of all data to be migrated to SQL
+   * @param entityClasses the entity types in write order
+   * @param blockingPCollection the pipeline stage that blocks this phase
+   * @param blockingTag description of the stage (if exists) that blocks this phase. Needed for
+   *     generating unique transform ids
+   * @return the output {@code PCollection} from the writing of the last entity kind. Other parts of
+   *     the pipeline can {@link Wait} on this object
+   */
+  private PCollection<Void> scheduleOnePhaseWrites(
+      PCollectionTuple datastoreSnapshot,
+      Collection<Class<?>> entityClasses,
+      Optional<PCollection<Void>> blockingPCollection,
+      String blockingTag) {
+    checkArgument(!entityClasses.isEmpty(), "Each phase must have at least one kind.");
+    ImmutableList<TupleTag<VersionedEntity>> tags =
+        toKindStrings(entityClasses).stream()
+            .map(Transforms::createTagForKind)
+            .collect(ImmutableList.toImmutableList());
+
+    PCollection<Void> prev = blockingPCollection.orElse(null);
+    String prevTag = blockingTag;
+    for (TupleTag<VersionedEntity> tag : tags) {
+      PCollection<VersionedEntity> curr = datastoreSnapshot.get(tag);
+      if (prev != null) {
+        curr = curr.apply("Wait on " + prevTag, Wait.on(prev));
+      }
+      prev = writeToSql(tag.getId(), curr);
+      prevTag = tag.getId();
+    }
+    return prev;
+  }
+
+  private PCollection<Void> writeToSql(String transformId, PCollection<VersionedEntity> data) {
+    String credentialFileUrl =
+        options.getSqlCredentialUrlOverride() != null
+            ? options.getSqlCredentialUrlOverride()
+            : BackupPaths.getCloudSQLCredentialFilePatterns(options.getEnvironment()).get(0);
+
+    return data.apply(
+        "Write to sql: " + transformId,
+        Transforms.writeToSql(
+            transformId,
+            options.getMaxConcurrentSqlWriters(),
+            options.getSqlWriteBatchSize(),
+            new JpaSupplierFactory(credentialFileUrl, options.getCloudKmsProjectId(), jpaGetter)));
+  }
+
+  private static ImmutableList<String> toKindStrings(Collection<Class<?>> entityClasses) {
+    try (AppEngineEnvironment env = new AppEngineEnvironment()) {
+      return entityClasses.stream().map(Key::getKind).collect(ImmutableList.toImmutableList());
+    }
+  }
+}

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

@@ -0,0 +1,91 @@
+// 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 javax.annotation.Nullable;
+import org.apache.beam.sdk.extensions.gcp.options.GcpOptions;
+import org.apache.beam.sdk.options.Default;
+import org.apache.beam.sdk.options.Description;
+import org.apache.beam.sdk.options.Validation;
+
+/** Pipeline options for {@link InitSqlPipeline} */
+public interface InitSqlPipelineOptions extends GcpOptions {
+
+  @Description(
+      "Overrides the URL to the SQL credential file. " + "Required if environment is not provided.")
+  @Nullable
+  String getSqlCredentialUrlOverride();
+
+  void setSqlCredentialUrlOverride(String credentialUrlOverride);
+
+  @Description("The root directory of the export to load.")
+  String getDatastoreExportDir();
+
+  void setDatastoreExportDir(String datastoreExportDir);
+
+  @Description("The directory that contains all CommitLog files.")
+  String getCommitLogDir();
+
+  void setCommitLogDir(String commitLogDir);
+
+  @Description("The earliest CommitLogs to load, in ISO8601 format.")
+  @Validation.Required
+  String getCommitLogStartTimestamp();
+
+  void setCommitLogStartTimestamp(String commitLogStartTimestamp);
+
+  @Description("The latest CommitLogs to load, in ISO8601 format.")
+  @Validation.Required
+  String getCommitLogEndTimestamp();
+
+  void setCommitLogEndTimestamp(String commitLogEndTimestamp);
+
+  @Description(
+      "The deployed environment, alpha, crash, sandbox, or production. "
+          + "Not required only if sqlCredentialUrlOverride is provided.")
+  @Nullable
+  String getEnvironment();
+
+  void setEnvironment(String environment);
+
+  @Description("The GCP project that contains the keyring used for decrypting the "
+      + "SQL credential file.")
+  @Nullable
+  String getCloudKmsProjectId();
+
+  void setCloudKmsProjectId(String cloudKmsProjectId);
+
+  @Description(
+      "The maximum JDBC connection pool size on a VM. "
+          + "This value should be equal to or greater than the number of cores on the VM.")
+  @Default.Integer(4)
+  int getJdbcMaxPoolSize();
+
+  void setJdbcMaxPoolSize(int jdbcMaxPoolSize);
+
+  @Description(
+      "A hint to the pipeline runner of the maximum number of concurrent SQL writers to create. "
+          + "Note that multiple writers may run on the same VM and share the connection pool.")
+  @Default.Integer(4)
+  int getMaxConcurrentSqlWriters();
+
+  void setMaxConcurrentSqlWriters(int maxConcurrentSqlWriters);
+
+  @Description("The number of entities to be written to the SQL database in one transaction.")
+  @Default.Integer(20)
+  int getSqlWriteBatchSize();
+
+  void setSqlWriteBatchSize(int sqlWriteBatchSize);
+}

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

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

core/src/main/java/google/registry/beam/initsql/README.md

@@ -1,3 +1,17 @@
 ## Summary
 
-This package contains a BEAM pipeline that populates a Cloud SQL database from a Datastore backup.
+This package contains a BEAM pipeline that populates a Cloud SQL database from a
+Datastore backup. The pipeline uses an unsynchronized Datastore export and
+overlapping CommitLogs generated by the Nomulus server to recreate a consistent
+Datastore snapshot, and writes the data to a Cloud SQL instance.
+
+## Pipeline Visualization
+
+The golden flow graph of the InitSqlPipeline is saved both as a text-base
+[DOT file](../../../../../../test/resources/google/registry/beam/initsql/pipeline_golden.dot)
+and a
+[.png file](../../../../../../test/resources/google/registry/beam/initsql/pipeline_golden.png).
+A test compares the flow graph of the current pipeline with the golden graph,
+and will fail if changes are detected. When this happens, run the Gradle task
+':core:updateInitSqlPipelineGraph' to update the golden files and review the
+changes.

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

@@ -17,26 +17,42 @@ package google.registry.beam.initsql;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Preconditions.checkState;
+import static com.google.common.base.Throwables.throwIfUnchecked;
 import static google.registry.beam.initsql.BackupPaths.getCommitLogTimestamp;
 import static google.registry.beam.initsql.BackupPaths.getExportFilePatterns;
+import static google.registry.persistence.JpaRetries.isFailedTxnRetriable;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.setJpaTm;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
 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;
 
 import avro.shaded.com.google.common.collect.Iterators;
+import com.google.appengine.api.datastore.Entity;
+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.Streams;
+import google.registry.backup.AppEngineEnvironment;
 import google.registry.backup.CommitLogImports;
 import google.registry.backup.VersionedEntity;
+import google.registry.model.domain.DomainBase;
+import google.registry.model.ofy.ObjectifyService;
+import google.registry.model.ofy.Ofy;
+import google.registry.persistence.transaction.JpaTransactionManager;
 import google.registry.tools.LevelDbLogReader;
+import google.registry.util.SystemSleeper;
+import java.io.Serializable;
 import java.util.Collection;
 import java.util.Iterator;
 import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.ThreadLocalRandom;
+import java.util.function.Supplier;
 import org.apache.beam.sdk.coders.StringUtf8Coder;
 import org.apache.beam.sdk.io.Compression;
 import org.apache.beam.sdk.io.FileIO;
@@ -47,6 +63,7 @@ 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;
@@ -60,6 +77,7 @@ import org.apache.beam.sdk.values.TupleTag;
 import org.apache.beam.sdk.values.TupleTagList;
 import org.apache.beam.sdk.values.TypeDescriptor;
 import org.joda.time.DateTime;
+import org.joda.time.Duration;
 
 /**
  * {@link PTransform Pipeline transforms} used in pipelines that load from both Datastore export
@@ -209,7 +227,7 @@ public final class Transforms {
     return new PTransform<PCollection<String>, PCollection<Metadata>>() {
       @Override
       public PCollection<Metadata> expand(PCollection<String> input) {
-        return input.apply(FileIO.matchAll().withEmptyMatchTreatment(EmptyMatchTreatment.DISALLOW));
+        return input.apply(FileIO.matchAll().withEmptyMatchTreatment(EmptyMatchTreatment.ALLOW));
       }
     };
   }
@@ -245,6 +263,42 @@ public final class Transforms {
                     .iterator()));
   }
 
+  /**
+   * 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>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 new PTransform<PCollection<VersionedEntity>, PCollection<Void>>() {
+      @Override
+      public PCollection<Void> expand(PCollection<VersionedEntity> input) {
+        return input
+            .apply(
+                "Shard data for " + transformId,
+                MapElements.into(kvs(integers(), TypeDescriptor.of(VersionedEntity.class)))
+                    .via(ve -> KV.of(ThreadLocalRandom.current().nextInt(maxWriters), ve)))
+            .apply("Batch output by shard " + transformId, GroupIntoBatches.ofSize(batchSize))
+            .apply("Write in batch for " + transformId, ParDo.of(new SqlBatchWriter(jpaSupplier)));
+      }
+    };
+  }
+
+  /** Interface for serializable {@link Supplier suppliers}. */
+  public interface SerializableSupplier<T> extends Supplier<T>, Serializable {}
+
   /**
    * Returns a {@link PTransform} that produces a {@link PCollection} containing all elements in the
    * given {@link Iterable}.
@@ -322,4 +376,116 @@ 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 extends DoFn<KV<Integer, Iterable<VersionedEntity>>, Void> {
+
+    private static int instanceCount = 0;
+    private static JpaTransactionManager originalJpa;
+
+    private final SerializableSupplier<JpaTransactionManager> jpaSupplier;
+
+    private transient Ofy ofy;
+    private transient SystemSleeper sleeper;
+
+    SqlBatchWriter(SerializableSupplier<JpaTransactionManager> jpaSupplier) {
+      this.jpaSupplier = jpaSupplier;
+    }
+
+    @Setup
+    public void setup() {
+      sleeper = new SystemSleeper();
+
+      try (AppEngineEnvironment env = new AppEngineEnvironment()) {
+        ObjectifyService.initOfy();
+        ofy = ObjectifyService.ofy();
+      }
+
+      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<VersionedEntity>> kv) {
+      try (AppEngineEnvironment env = new AppEngineEnvironment()) {
+        ImmutableList<Object> ofyEntities =
+            Streams.stream(kv.getValue())
+                .map(VersionedEntity::getEntity)
+                .map(Optional::get)
+                .map(ofy::toPojo)
+                .collect(ImmutableList.toImmutableList());
+        retry(() -> jpaTm().transact(() -> jpaTm().saveNewOrUpdateAll(ofyEntities)));
+      }
+    }
+
+    // TODO(b/160632289): Enhance Retrier and use it here.
+    private void retry(Runnable runnable) {
+      int maxAttempts = 5;
+      int initialDelayMillis = 100;
+      double jitterRatio = 0.2;
+
+      for (int attempt = 0; attempt < maxAttempts; attempt++) {
+        try {
+          runnable.run();
+          return;
+        } catch (Throwable throwable) {
+          if (!isFailedTxnRetriable(throwable)) {
+            throwIfUnchecked(throwable);
+            throw new RuntimeException(throwable);
+          }
+          int sleepMillis = (1 << attempt) * initialDelayMillis;
+          int jitter =
+              ThreadLocalRandom.current().nextInt((int) (sleepMillis * jitterRatio))
+                  - (int) (sleepMillis * jitterRatio / 2);
+          sleeper.sleepUninterruptibly(Duration.millis(sleepMillis + jitter));
+        }
+      }
+    }
+  }
+
+  /**
+   * Removes BillingEvents, {@link google.registry.model.poll.PollMessage PollMessages} and {@link
+   * google.registry.model.host.HostResource} from a {@link DomainBase}. These are circular foreign
+   * key constraints that prevent migration of {@code DomainBase} to SQL databases.
+   *
+   * <p>See {@link InitSqlPipeline} for more information.
+   */
+  static class RemoveDomainBaseForeignKeys extends DoFn<VersionedEntity, VersionedEntity> {
+
+    @ProcessElement
+    public void processElement(
+        @Element VersionedEntity domainBase, OutputReceiver<VersionedEntity> out) {
+      checkArgument(
+          domainBase.getEntity().isPresent(), "Unexpected delete entity %s", domainBase.key());
+      Entity outputEntity =
+          DomainBaseUtil.removeBillingAndPollAndHosts(domainBase.getEntity().get());
+      out.output(
+          VersionedEntity.from(
+              domainBase.commitTimeMills(),
+              EntityTranslator.convertToPb(outputEntity).toByteArray()));
+    }
+  }
 }

core/src/main/java/google/registry/beam/spec11/sql/subdomains.sql

@@ -21,7 +21,7 @@
 SELECT
   domain.fullyQualifiedDomainName AS domainName,
   domain.__key__.name AS domainRepoId,
-  registrar.clientId AS clientId,
+  registrar.clientId AS registrarId,
   COALESCE(registrar.emailAddress, '') AS registrarEmailAddress
 FROM ( (
     SELECT

core/src/main/java/google/registry/env/common/backend/WEB-INF/web.xml

@@ -367,6 +367,12 @@
     <url-pattern>/_dr/task/linkRdeHosts</url-pattern>
   </servlet-mapping>
 
+  <!-- Action to automatically re-lock a domain after unlocking it -->
+  <servlet-mapping>
+    <servlet-name>backend-servlet</servlet-name>
+    <url-pattern>/_dr/task/relockDomain</url-pattern>
+  </servlet-mapping>
+
   <!-- Security config -->
   <security-constraint>
     <web-resource-collection>

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

@@ -18,7 +18,14 @@
       and streams it to cloud storage. When this job has finished successfully, it'll
       launch a separate task that uploads the deposit file to Iron Mountain via SFTP.
     </description>
-    <schedule>every day 00:07</schedule>
+    <!--
+      This only needs to run once per day, but we launch additional jobs in case the
+      cursor is lagging behind, so it'll catch up to the current date eventually.
+
+      See <a href="../../../production/default/WEB-INF/cron.xml">production config</a> for an
+      explanation of job starting times.
+    -->
+    <schedule>every 12 hours from 00:07 to 12:07</schedule>
     <target>backend</target>
   </cron>
 

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

@@ -347,8 +347,8 @@ public class DomainCreateFlow implements TransactionalFlow {
             .setRepoId(repoId)
             .setIdnTableName(validateDomainNameWithIdnTables(domainName))
             .setRegistrationExpirationTime(registrationExpirationTime)
-            .setAutorenewBillingEvent(Key.create(autorenewBillingEvent))
-            .setAutorenewPollMessage(Key.create(autorenewPollMessage))
+            .setAutorenewBillingEvent(autorenewBillingEvent.createVKey())
+            .setAutorenewPollMessage(autorenewPollMessage.createVKey())
             .setLaunchNotice(hasClaimsNotice ? launchCreate.get().getNotice() : null)
             .setSmdId(signedMarkId)
             .setDsData(secDnsCreate.isPresent() ? secDnsCreate.get().getDsData() : null)

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

@@ -31,7 +31,6 @@ import static google.registry.model.ResourceTransferUtils.handlePendingTransferO
 import static google.registry.model.ResourceTransferUtils.updateForeignKeyIndexDeletionTime;
 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.ofy.ObjectifyService.ofy;
 import static google.registry.model.reporting.DomainTransactionRecord.TransactionReportField.ADD_FIELDS;
 import static google.registry.model.reporting.DomainTransactionRecord.TransactionReportField.RENEW_FIELDS;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
@@ -209,7 +208,7 @@ public final class DomainDeleteFlow implements TransactionalFlow {
       PollMessage.OneTime deletePollMessage =
           createDeletePollMessage(existingDomain, historyEntry, deletionTime);
       entitiesToSave.add(deletePollMessage);
-      builder.setDeletePollMessage(Key.create(deletePollMessage));
+      builder.setDeletePollMessage(deletePollMessage.createVKey());
     }
 
     // Cancel any grace periods that were still active, and set the expiration time accordingly.
@@ -222,8 +221,7 @@ public final class DomainDeleteFlow implements TransactionalFlow {
         if (gracePeriod.getOneTimeBillingEvent() != null) {
           // Take the amount of amount of registration time being refunded off the expiration time.
           // This can be either add grace periods or renew grace periods.
-          BillingEvent.OneTime oneTime =
-              ofy().load().key(gracePeriod.getOneTimeBillingEvent()).now();
+          BillingEvent.OneTime oneTime = tm().load(gracePeriod.getOneTimeBillingEvent());
           newExpirationTime = newExpirationTime.minusYears(oneTime.getPeriodYears());
         } else if (gracePeriod.getRecurringBillingEvent() != null) {
           // Take 1 year off the registration if in the autorenew grace period (no need to load the
@@ -370,12 +368,12 @@ public final class DomainDeleteFlow implements TransactionalFlow {
   private Money getGracePeriodCost(GracePeriod gracePeriod, DateTime now) {
     if (gracePeriod.getType() == GracePeriodStatus.AUTO_RENEW) {
       DateTime autoRenewTime =
-          ofy().load().key(checkNotNull(gracePeriod.getRecurringBillingEvent())).now()
+          tm().load(checkNotNull(gracePeriod.getRecurringBillingEvent()))
               .getRecurrenceTimeOfYear()
-                  .getLastInstanceBeforeOrAt(now);
+              .getLastInstanceBeforeOrAt(now);
       return getDomainRenewCost(targetId, autoRenewTime, 1);
     }
-    return ofy().load().key(checkNotNull(gracePeriod.getOneTimeBillingEvent())).now().getCost();
+    return tm().load(checkNotNull(gracePeriod.getOneTimeBillingEvent())).getCost();
   }
 
   @Nullable

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

@@ -517,14 +517,14 @@ public class DomainFlowUtils {
    */
   public static void updateAutorenewRecurrenceEndTime(DomainBase domain, DateTime newEndTime) {
     Optional<PollMessage.Autorenew> autorenewPollMessage =
-        Optional.ofNullable(ofy().load().key(domain.getAutorenewPollMessage()).now());
+        tm().maybeLoad(domain.getAutorenewPollMessage());
 
     // Construct an updated autorenew poll message. If the autorenew poll message no longer exists,
     // create a new one at the same id. This can happen if a transfer was requested on a domain
     // where all autorenew poll messages had already been delivered (this would cause the poll
     // message to be deleted), and then subsequently the transfer was canceled, rejected, or deleted
     // (which would cause the poll message to be recreated here).
-    Key<PollMessage.Autorenew> existingAutorenewKey = domain.getAutorenewPollMessage();
+    Key<PollMessage.Autorenew> existingAutorenewKey = domain.getAutorenewPollMessage().getOfyKey();
     PollMessage.Autorenew updatedAutorenewPollMessage =
         autorenewPollMessage.isPresent()
             ? autorenewPollMessage.get().asBuilder().setAutorenewEndTime(newEndTime).build()
@@ -542,7 +542,7 @@ public class DomainFlowUtils {
       ofy().save().entity(updatedAutorenewPollMessage);
     }
 
-    Recurring recurring = ofy().load().key(domain.getAutorenewBillingEvent()).now();
+    Recurring recurring = tm().load(domain.getAutorenewBillingEvent());
     ofy().save().entity(recurring.asBuilder().setRecurrenceEndTime(newEndTime).build());
   }
 

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

@@ -33,7 +33,6 @@ 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;
@@ -181,8 +180,8 @@ public final class DomainRenewFlow implements TransactionalFlow {
             .setLastEppUpdateTime(now)
             .setLastEppUpdateClientId(clientId)
             .setRegistrationExpirationTime(newExpirationTime)
-            .setAutorenewBillingEvent(Key.create(newAutorenewEvent))
-            .setAutorenewPollMessage(Key.create(newAutorenewPollMessage))
+            .setAutorenewBillingEvent(newAutorenewEvent.createVKey())
+            .setAutorenewPollMessage(newAutorenewPollMessage.createVKey())
             .addGracePeriod(
                 GracePeriod.forBillingEvent(GracePeriodStatus.RENEW, explicitRenewEvent))
             .build();

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

@@ -26,7 +26,6 @@ import static google.registry.flows.domain.DomainFlowUtils.verifyNotReserved;
 import static google.registry.flows.domain.DomainFlowUtils.verifyPremiumNameIsNotBlocked;
 import static google.registry.flows.domain.DomainFlowUtils.verifyRegistrarIsActive;
 import static google.registry.model.ResourceTransferUtils.updateForeignKeyIndexDeletionTime;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.DateTimeUtils.END_OF_TIME;
 
@@ -174,8 +173,8 @@ public final class DomainRestoreRequestFlow implements TransactionalFlow  {
             existingDomain, newExpirationTime, autorenewEvent, autorenewPollMessage, now, clientId);
     updateForeignKeyIndexDeletionTime(newDomain);
     entitiesToSave.add(newDomain, historyEntry, autorenewEvent, autorenewPollMessage);
-    ofy().save().entities(entitiesToSave.build());
-    ofy().delete().key(existingDomain.getDeletePollMessage());
+    tm().saveNewOrUpdateAll(entitiesToSave.build());
+    tm().delete(existingDomain.getDeletePollMessage());
     dnsQueue.addDomainRefreshTask(existingDomain.getDomainName());
     return responseBuilder
         .setExtensions(createResponseExtensions(feesAndCredits, feeUpdate, isExpired))
@@ -232,8 +231,8 @@ public final class DomainRestoreRequestFlow implements TransactionalFlow  {
         .setStatusValues(null)
         .setGracePeriods(null)
         .setDeletePollMessage(null)
-        .setAutorenewBillingEvent(Key.create(autorenewEvent))
-        .setAutorenewPollMessage(Key.create(autorenewPollMessage))
+        .setAutorenewBillingEvent(autorenewEvent.createVKey())
+        .setAutorenewPollMessage(autorenewPollMessage.createVKey())
         .setLastEppUpdateTime(now)
         .setLastEppUpdateClientId(clientId)
         .build();

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

@@ -186,8 +186,8 @@ public final class DomainTransferApproveFlow implements TransactionalFlow {
                     .setTransferredRegistrationExpirationTime(newExpirationTime)
                     .build())
             .setRegistrationExpirationTime(newExpirationTime)
-            .setAutorenewBillingEvent(Key.create(autorenewEvent))
-            .setAutorenewPollMessage(Key.create(gainingClientAutorenewPollMessage))
+            .setAutorenewBillingEvent(autorenewEvent.createVKey())
+            .setAutorenewPollMessage(gainingClientAutorenewPollMessage.createVKey())
             // Remove all the old grace periods and add a new one for the transfer.
             .setGracePeriods(
                 billingEvent.isPresent()

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

@@ -14,16 +14,21 @@
 
 package google.registry.model;
 
+import com.google.common.annotations.VisibleForTesting;
+import javax.persistence.Access;
+import javax.persistence.AccessType;
+import javax.persistence.MappedSuperclass;
 import javax.xml.bind.annotation.XmlTransient;
 
 /**
  * Base class for entities that are the root of a Registry 2.0 entity group that gets enrolled in
  * commit logs for backup purposes.
  *
- * <p>The commit log system needs to preserve the ordering of closely timed mutations to entities
- * in a single entity group. We require an {@link UpdateAutoTimestamp} field on the root of a group
- * so that we can enforce strictly increasing timestamps.
+ * <p>The commit log system needs to preserve the ordering of closely timed mutations to entities in
+ * a single entity group. We require an {@link UpdateAutoTimestamp} field on the root of a group so
+ * that we can enforce strictly increasing timestamps.
  */
+@MappedSuperclass
 public abstract class BackupGroupRoot extends ImmutableObject {
   /**
    * An automatically managed timestamp of when this object was last written to Datastore.
@@ -32,10 +37,14 @@ public abstract class BackupGroupRoot extends ImmutableObject {
    * that this is updated on every save, rather than only in response to an {@code <update>} command
    */
   @XmlTransient
+  // Prevents subclasses from unexpectedly accessing as property (e.g., HostResource), which would
+  // require an unnecessary non-private setter method.
+  @Access(AccessType.FIELD)
+  @VisibleForTesting
   UpdateAutoTimestamp updateTimestamp = UpdateAutoTimestamp.create(null);
 
   /** Get the {@link UpdateAutoTimestamp} for this entity. */
-  public final UpdateAutoTimestamp getUpdateAutoTimestamp() {
+  public UpdateAutoTimestamp getUpdateTimestamp() {
     return updateTimestamp;
   }
 }

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

@@ -155,7 +155,7 @@ public final class EppResourceUtils {
     // time for writes.
     return Optional.of(
         cloneProjectedAtTime(
-            resource, latestOf(now, resource.getUpdateAutoTimestamp().getTimestamp())));
+            resource, latestOf(now, resource.getUpdateTimestamp().getTimestamp())));
   }
 
   /**
@@ -298,7 +298,7 @@ public final class EppResourceUtils {
     // and returns it projected forward to exactly the desired timestamp, or null if the resource is
     // deleted at that timestamp.
     final Result<T> loadResult =
-        isAtOrAfter(timestamp, resource.getUpdateAutoTimestamp().getTimestamp())
+        isAtOrAfter(timestamp, resource.getUpdateTimestamp().getTimestamp())
             ? new ResultNow<>(resource)
             : loadMostRecentRevisionAtTime(resource, timestamp);
     return () -> {

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

@@ -630,9 +630,9 @@ public abstract class BillingEvent extends ImmutableObject
           .setParent(historyEntry);
       // Set the grace period's billing event using the appropriate Cancellation builder method.
       if (gracePeriod.getOneTimeBillingEvent() != null) {
-        builder.setOneTimeEventKey(VKey.from(gracePeriod.getOneTimeBillingEvent()));
+        builder.setOneTimeEventKey(gracePeriod.getOneTimeBillingEvent());
       } else if (gracePeriod.getRecurringBillingEvent() != null) {
-        builder.setRecurringEventKey(VKey.from(gracePeriod.getRecurringBillingEvent()));
+        builder.setRecurringEventKey(gracePeriod.getRecurringBillingEvent());
       }
       return builder.build();
     }

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

@@ -224,12 +224,12 @@ public class ContactBase extends EppResource implements ResourceWithTransferData
     return disclose;
   }
 
-  public final String getCurrentSponsorClientId() {
+  public String getCurrentSponsorClientId() {
     return getPersistedCurrentSponsorClientId();
   }
 
   @Override
-  public final ContactTransferData getTransferData() {
+  public ContactTransferData getTransferData() {
     return Optional.ofNullable(transferData).orElse(ContactTransferData.EMPTY);
   }
 

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

@@ -227,7 +227,8 @@ public class DomainBase extends EppResource
    * refer to a {@link PollMessage} timed to when the domain is fully deleted. If the domain is
    * restored, the message should be deleted.
    */
-  @Transient Key<PollMessage.OneTime> deletePollMessage;
+  @Column(name = "deletion_poll_message_id")
+  VKey<PollMessage.OneTime> deletePollMessage;
 
   /**
    * The recurring billing event associated with this domain's autorenewals.
@@ -237,7 +238,8 @@ public class DomainBase extends EppResource
    * {@link #registrationExpirationTime} is changed the recurrence should be closed, a new one
    * should be created, and this field should be updated to point to the new one.
    */
-  @Transient Key<BillingEvent.Recurring> autorenewBillingEvent;
+  @Column(name = "billing_recurrence_id")
+  VKey<BillingEvent.Recurring> autorenewBillingEvent;
 
   /**
    * The recurring poll message associated with this domain's autorenewals.
@@ -247,7 +249,8 @@ public class DomainBase extends EppResource
    * {@link #registrationExpirationTime} is changed the recurrence should be closed, a new one
    * should be created, and this field should be updated to point to the new one.
    */
-  @Transient Key<PollMessage.Autorenew> autorenewPollMessage;
+  @Column(name = "autorenew_poll_message_id")
+  VKey<PollMessage.Autorenew> autorenewPollMessage;
 
   /** The unexpired grace periods for this domain (some of which may not be active yet). */
   @Transient @ElementCollection Set<GracePeriod> gracePeriods;
@@ -316,15 +319,15 @@ public class DomainBase extends EppResource
     return registrationExpirationTime;
   }
 
-  public Key<PollMessage.OneTime> getDeletePollMessage() {
+  public VKey<PollMessage.OneTime> getDeletePollMessage() {
     return deletePollMessage;
   }
 
-  public Key<BillingEvent.Recurring> getAutorenewBillingEvent() {
+  public VKey<BillingEvent.Recurring> getAutorenewBillingEvent() {
     return autorenewBillingEvent;
   }
 
-  public Key<PollMessage.Autorenew> getAutorenewPollMessage() {
+  public VKey<PollMessage.Autorenew> getAutorenewPollMessage() {
     return autorenewPollMessage;
   }
 
@@ -453,14 +456,8 @@ public class DomainBase extends EppResource
               .setRegistrationExpirationTime(expirationDate)
               // Set the speculatively-written new autorenew events as the domain's autorenew
               // events.
-              .setAutorenewBillingEvent(
-                  transferData.getServerApproveAutorenewEvent() == null
-                      ? null
-                      : transferData.getServerApproveAutorenewEvent().getOfyKey())
-              .setAutorenewPollMessage(
-                  transferData.getServerApproveAutorenewPollMessage() == null
-                      ? null
-                      : transferData.getServerApproveAutorenewPollMessage().getOfyKey());
+              .setAutorenewBillingEvent(transferData.getServerApproveAutorenewEvent())
+              .setAutorenewPollMessage(transferData.getServerApproveAutorenewPollMessage());
       if (transferData.getTransferPeriod().getValue() == 1) {
         // Set the grace period using a key to the prescheduled transfer billing event.  Not using
         // GracePeriod.forBillingEvent() here in order to avoid the actual Datastore fetch.
@@ -471,9 +468,7 @@ public class DomainBase extends EppResource
                     transferExpirationTime.plus(
                         Registry.get(getTld()).getTransferGracePeriodLength()),
                     transferData.getGainingClientId(),
-                    transferData.getServerApproveBillingEvent() == null
-                        ? null
-                        : transferData.getServerApproveBillingEvent().getOfyKey())));
+                    transferData.getServerApproveBillingEvent())));
       } else {
         // There won't be a billing event, so we don't need a grace period
         builder.setGracePeriods(ImmutableSet.of());
@@ -801,19 +796,17 @@ public class DomainBase extends EppResource
       return this;
     }
 
-    public Builder setDeletePollMessage(Key<PollMessage.OneTime> deletePollMessage) {
+    public Builder setDeletePollMessage(VKey<PollMessage.OneTime> deletePollMessage) {
       getInstance().deletePollMessage = deletePollMessage;
       return this;
     }
 
-    public Builder setAutorenewBillingEvent(
-        Key<BillingEvent.Recurring> autorenewBillingEvent) {
+    public Builder setAutorenewBillingEvent(VKey<BillingEvent.Recurring> autorenewBillingEvent) {
       getInstance().autorenewBillingEvent = autorenewBillingEvent;
       return this;
     }
 
-    public Builder setAutorenewPollMessage(
-        Key<PollMessage.Autorenew> autorenewPollMessage) {
+    public Builder setAutorenewPollMessage(VKey<PollMessage.Autorenew> autorenewPollMessage) {
       getInstance().autorenewPollMessage = autorenewPollMessage;
       return this;
     }

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

@@ -17,12 +17,13 @@ package google.registry.model.domain;
 import static com.google.common.base.Preconditions.checkArgument;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 
-import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Embed;
 import com.googlecode.objectify.annotation.Ignore;
 import google.registry.model.ImmutableObject;
 import google.registry.model.billing.BillingEvent;
+import google.registry.model.billing.BillingEvent.Recurring;
 import google.registry.model.domain.rgp.GracePeriodStatus;
+import google.registry.persistence.VKey;
 import javax.annotation.Nullable;
 import javax.persistence.Column;
 import javax.persistence.GeneratedValue;
@@ -57,18 +58,18 @@ public class GracePeriod extends ImmutableObject {
 
   /**
    * The one-time billing event corresponding to the action that triggered this grace period, or
-   * null if not applicable.  Not set for autorenew grace periods (which instead use the field
-   * {@code billingEventRecurring}) or for redemption grace periods (since deletes have no cost).
+   * null if not applicable. Not set for autorenew grace periods (which instead use the field {@code
+   * billingEventRecurring}) or for redemption grace periods (since deletes have no cost).
    */
   // NB: Would @IgnoreSave(IfNull.class), but not allowed for @Embed collections.
-  Key<BillingEvent.OneTime> billingEventOneTime = null;
+  VKey<BillingEvent.OneTime> billingEventOneTime = null;
 
   /**
    * The recurring billing event corresponding to the action that triggered this grace period, if
    * applicable - i.e. if the action was an autorenew - or null in all other cases.
    */
   // NB: Would @IgnoreSave(IfNull.class), but not allowed for @Embed collections.
-  Key<BillingEvent.Recurring> billingEventRecurring = null;
+  VKey<BillingEvent.Recurring> billingEventRecurring = null;
 
   public GracePeriodStatus getType() {
     return type;
@@ -91,8 +92,7 @@ public class GracePeriod extends ImmutableObject {
    * Returns the one time billing event. The value will only be non-null if the type of this grace
    * period is not AUTO_RENEW.
    */
-
-  public Key<BillingEvent.OneTime> getOneTimeBillingEvent() {
+  public VKey<BillingEvent.OneTime> getOneTimeBillingEvent() {
     return billingEventOneTime;
   }
 
@@ -100,16 +100,16 @@ public class GracePeriod extends ImmutableObject {
    * Returns the recurring billing event. The value will only be non-null if the type of this grace
    * period is AUTO_RENEW.
    */
-  public Key<BillingEvent.Recurring> getRecurringBillingEvent() {
+  public VKey<BillingEvent.Recurring> getRecurringBillingEvent() {
     return billingEventRecurring;
   }
 
   private static GracePeriod createInternal(
-       GracePeriodStatus type,
-       DateTime expirationTime,
-       String clientId,
-       @Nullable Key<BillingEvent.OneTime> billingEventOneTime,
-       @Nullable Key<BillingEvent.Recurring> billingEventRecurring) {
+      GracePeriodStatus type,
+      DateTime expirationTime,
+      String clientId,
+      @Nullable VKey<BillingEvent.OneTime> billingEventOneTime,
+      @Nullable VKey<BillingEvent.Recurring> billingEventRecurring) {
     checkArgument((billingEventOneTime == null) || (billingEventRecurring == null),
         "A grace period can have at most one billing event");
     checkArgument(
@@ -127,15 +127,15 @@ public class GracePeriod extends ImmutableObject {
   /**
    * Creates a GracePeriod for an (optional) OneTime billing event.
    *
-   * <p>Normal callers should always use {@link #forBillingEvent} instead, assuming they do not
-   * need to avoid loading the BillingEvent from Datastore.  This method should typically be
-   * called only from test code to explicitly construct GracePeriods.
+   * <p>Normal callers should always use {@link #forBillingEvent} instead, assuming they do not need
+   * to avoid loading the BillingEvent from Datastore. This method should typically be called only
+   * from test code to explicitly construct GracePeriods.
    */
   public static GracePeriod create(
       GracePeriodStatus type,
       DateTime expirationTime,
       String clientId,
-      @Nullable Key<BillingEvent.OneTime> billingEventOneTime) {
+      @Nullable VKey<BillingEvent.OneTime> billingEventOneTime) {
     return createInternal(type, expirationTime, clientId, billingEventOneTime, null);
   }
 
@@ -144,7 +144,7 @@ public class GracePeriod extends ImmutableObject {
       GracePeriodStatus type,
       DateTime expirationTime,
       String clientId,
-      Key<BillingEvent.Recurring> billingEventRecurring) {
+      VKey<Recurring> billingEventRecurring) {
     checkArgumentNotNull(billingEventRecurring, "billingEventRecurring cannot be null");
     return createInternal(type, expirationTime, clientId, null, billingEventRecurring);
   }
@@ -159,6 +159,6 @@ public class GracePeriod extends ImmutableObject {
   public static GracePeriod forBillingEvent(
       GracePeriodStatus type, BillingEvent.OneTime billingEvent) {
     return create(
-        type, billingEvent.getBillingTime(), billingEvent.getClientId(), Key.create(billingEvent));
+        type, billingEvent.getBillingTime(), billingEvent.getClientId(), billingEvent.createVKey());
   }
 }

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

@@ -168,7 +168,7 @@ class CommitLoggedWork<R> implements Runnable {
       DateTime transactionTime, Set<Entry<Key<BackupGroupRoot>, BackupGroupRoot>> bgrEntries) {
     ImmutableMap.Builder<Key<BackupGroupRoot>, DateTime> builder = new ImmutableMap.Builder<>();
     for (Entry<Key<BackupGroupRoot>, BackupGroupRoot> entry : bgrEntries) {
-      DateTime updateTime = entry.getValue().getUpdateAutoTimestamp().getTimestamp();
+      DateTime updateTime = entry.getValue().getUpdateTimestamp().getTimestamp();
       if (!updateTime.isBefore(transactionTime)) {
         builder.put(entry.getKey(), updateTime);
       }

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

@@ -48,6 +48,10 @@ public final class RdeRevision extends ImmutableObject {
    */
   int revision;
 
+  public int getRevision() {
+    return revision;
+  }
+
   /**
    * Returns next revision ID to use when staging a new deposit file for the given triplet.
    *

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

@@ -15,6 +15,7 @@
 package google.registry.model.registry.label;
 
 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 com.google.common.collect.ImmutableSet.toImmutableSet;
 import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
@@ -30,6 +31,7 @@ import com.google.common.collect.Multiset;
 import com.google.common.util.concurrent.UncheckedExecutionException;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Id;
+import com.googlecode.objectify.annotation.Ignore;
 import com.googlecode.objectify.annotation.Parent;
 import google.registry.model.Buildable;
 import google.registry.model.ImmutableObject;
@@ -42,6 +44,11 @@ import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.ExecutionException;
 import javax.annotation.Nullable;
+import javax.persistence.Column;
+import javax.persistence.GeneratedValue;
+import javax.persistence.GenerationType;
+import javax.persistence.MappedSuperclass;
+import javax.persistence.Transient;
 import org.joda.time.DateTime;
 
 /**
@@ -49,25 +56,46 @@ import org.joda.time.DateTime;
  *
  * @param <T> The type of the root value being listed, e.g. {@link ReservationType}.
  * @param <R> The type of domain label entry being listed, e.g. {@link ReservedListEntry} (note,
- *            must subclass {@link DomainLabelEntry}.
+ *     must subclass {@link DomainLabelEntry}.
  */
+@MappedSuperclass
 public abstract class BaseDomainLabelList<T extends Comparable<?>, R extends DomainLabelEntry<T, ?>>
     extends ImmutableObject implements Buildable {
 
+  @Ignore
+  @javax.persistence.Id
+  @GeneratedValue(strategy = GenerationType.IDENTITY)
+  Long revisionId;
+
   @Id
+  @Column(nullable = false)
   String name;
 
-  @Parent
-  Key<EntityGroupRoot> parent = getCrossTldKey();
+  @Parent @Transient Key<EntityGroupRoot> parent = getCrossTldKey();
 
-  DateTime creationTime;
+  @Transient DateTime creationTime;
 
+  // The list in Cloud SQL is immutable, we only have a creation_timestamp field and it should be
+  // set to the timestamp when the list is created. In Datastore, we have two fields and the
+  // lastUpdateTime is set to the current timestamp when creating and updating a list. So, we use
+  // lastUpdateTime as the creation_timestamp column during the dual-write phase for compatibility.
+  @Column(name = "creation_timestamp", nullable = false)
   DateTime lastUpdateTime;
 
+  /** Returns the ID of this revision, or throws if null. */
+  public long getRevisionId() {
+    checkState(
+        revisionId != null,
+        "revisionId is null because this object has not been persisted to the database yet");
+    return revisionId;
+  }
+
+  /** Returns the name of the reserved list. */
   public String getName() {
     return name;
   }
 
+  /** Returns the creation time of this revision of the reserved list. */
   public DateTime getCreationTime() {
     return creationTime;
   }
@@ -183,6 +211,9 @@ public abstract class BaseDomainLabelList<T extends Comparable<?>, R extends Dom
     @Override
     public T build() {
       checkArgument(!isNullOrEmpty(getInstance().name), "List must have a name");
+      // The list is immutable in Cloud SQL, so make sure the revision id is not set when the
+      // builder object is created from a list object
+      getInstance().revisionId = null;
       return super.build();
     }
   }

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

@@ -23,16 +23,20 @@ import com.google.common.net.InternetDomainName;
 import com.googlecode.objectify.annotation.Id;
 import google.registry.model.Buildable.GenericBuilder;
 import google.registry.model.ImmutableObject;
+import javax.persistence.Column;
+import javax.persistence.MappedSuperclass;
 
 /**
  * Represents a label entry parsed from a line in a reserved/premium list txt file.
  *
  * @param <T> The type of the value stored for the domain label, e.g. {@link ReservationType}.
  */
+@MappedSuperclass
 public abstract class DomainLabelEntry<T extends Comparable<?>, D extends DomainLabelEntry<?, ?>>
     extends ImmutableObject implements Comparable<D> {
 
   @Id
+  @Column(name = "domain_label", insertable = false, updatable = false)
   String label;
 
   String comment;

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

@@ -14,7 +14,9 @@
 
 package google.registry.model.registry.label;
 
+import static com.google.common.base.Charsets.US_ASCII;
 import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.hash.Funnels.stringFunnel;
 import static com.google.common.hash.Funnels.unencodedCharsFunnel;
 import static google.registry.config.RegistryConfig.getDomainLabelListCacheDuration;
 import static google.registry.config.RegistryConfig.getSingletonCachePersistDuration;
@@ -32,43 +34,82 @@ import com.google.common.cache.CacheLoader;
 import com.google.common.cache.CacheLoader.InvalidCacheLoadException;
 import com.google.common.cache.LoadingCache;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.hash.BloomFilter;
 import com.google.common.util.concurrent.UncheckedExecutionException;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Entity;
 import com.googlecode.objectify.annotation.Id;
+import com.googlecode.objectify.annotation.Ignore;
 import com.googlecode.objectify.annotation.Parent;
 import google.registry.model.Buildable;
 import google.registry.model.ImmutableObject;
 import google.registry.model.annotations.ReportedOn;
 import google.registry.model.registry.Registry;
+import google.registry.schema.replay.DatastoreAndSqlEntity;
 import google.registry.schema.replay.DatastoreEntity;
 import google.registry.schema.replay.SqlEntity;
+import google.registry.schema.tld.PremiumListDao;
 import google.registry.util.NonFinalForTesting;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
+import java.math.BigDecimal;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
 import javax.annotation.Nullable;
+import javax.persistence.CollectionTable;
+import javax.persistence.Column;
+import javax.persistence.ElementCollection;
+import javax.persistence.Index;
+import javax.persistence.JoinColumn;
+import javax.persistence.MapKeyColumn;
+import javax.persistence.PostLoad;
+import javax.persistence.PrePersist;
+import javax.persistence.Table;
+import javax.persistence.Transient;
+import org.hibernate.LazyInitializationException;
+import org.joda.money.CurrencyUnit;
 import org.joda.money.Money;
 import org.joda.time.Duration;
 
-/** A premium list entity, persisted to Datastore, that is used to check domain label prices. */
+/**
+ * A premium list entity that is used to check domain label prices.
+ *
+ * <p>Note that the primary key of this entity is {@link #revisionId}, which is auto-generated by
+ * the database. So, if a retry of insertion happens after the previous attempt unexpectedly
+ * succeeds, we will end up with having two exact same premium lists that differ only by revisionId.
+ * This is fine though, because we only use the list with the highest revisionId.
+ */
 @ReportedOn
 @Entity
+@javax.persistence.Entity
+@Table(indexes = {@Index(columnList = "name", name = "premiumlist_name_idx")})
 public final class PremiumList extends BaseDomainLabelList<Money, PremiumList.PremiumListEntry>
-    implements DatastoreEntity {
+    implements DatastoreAndSqlEntity {
 
   /** Stores the revision key for the set of currently used premium list entry entities. */
-  Key<PremiumListRevision> revisionKey;
+  @Transient Key<PremiumListRevision> revisionKey;
 
-  @Override
-  public ImmutableList<SqlEntity> toSqlEntities() {
-    return ImmutableList.of(); // PremiumList is dual-written
-  }
+  @Ignore
+  @Column(nullable = false)
+  CurrencyUnit currency;
+
+  @Ignore
+  @ElementCollection
+  @CollectionTable(
+      name = "PremiumEntry",
+      joinColumns = @JoinColumn(name = "revisionId", referencedColumnName = "revisionId"))
+  @MapKeyColumn(name = "domainLabel")
+  @Column(name = "price", nullable = false)
+  Map<String, BigDecimal> labelsToPrices;
+
+  @Ignore
+  @Column(nullable = false)
+  BloomFilter<String> bloomFilter;
 
   /** Virtual parent entity for premium list entry entities associated with a single revision. */
   @ReportedOn
@@ -247,6 +288,35 @@ public final class PremiumList extends BaseDomainLabelList<Money, PremiumList.Pr
     return Optional.ofNullable(loadPremiumList(name));
   }
 
+  /** Returns the {@link CurrencyUnit} used for this list. */
+  public CurrencyUnit getCurrency() {
+    return currency;
+  }
+
+  /**
+   * Returns a {@link Map} of domain labels to prices.
+   *
+   * <p>Note that this is lazily loaded and thus will throw a {@link LazyInitializationException} if
+   * used outside the transaction in which the given entity was loaded. You generally should not be
+   * using this anyway as it's inefficient to load all of the PremiumEntry rows if you don't need
+   * them. To check prices, use {@link PremiumListDao#getPremiumPrice} instead.
+   */
+  @Nullable
+  public ImmutableMap<String, BigDecimal> getLabelsToPrices() {
+    return labelsToPrices == null ? null : ImmutableMap.copyOf(labelsToPrices);
+  }
+
+  /**
+   * Returns a Bloom filter to determine whether a label might be premium, or is definitely not.
+   *
+   * <p>If the domain label might be premium, then the next step is to check for the existence of a
+   * corresponding row in the PremiumListEntry table. Otherwise, we know for sure it's not premium,
+   * and no DB load is required.
+   */
+  public BloomFilter<String> getBloomFilter() {
+    return bloomFilter;
+  }
+
   /**
    * A premium list entry entity, persisted to Datastore. Each instance represents the price of a
    * single label on a given TLD.
@@ -339,9 +409,39 @@ public final class PremiumList extends BaseDomainLabelList<Money, PremiumList.Pr
       return this;
     }
 
+    public Builder setCurrency(CurrencyUnit currency) {
+      getInstance().currency = currency;
+      return this;
+    }
+
+    public Builder setLabelsToPrices(Map<String, BigDecimal> labelsToPrices) {
+      getInstance().labelsToPrices = ImmutableMap.copyOf(labelsToPrices);
+      return this;
+    }
+
     @Override
     public PremiumList build() {
+      if (getInstance().labelsToPrices != null) {
+        // ASCII is used for the charset because all premium list domain labels are stored
+        // punycoded.
+        getInstance().bloomFilter =
+            BloomFilter.create(stringFunnel(US_ASCII), getInstance().labelsToPrices.size());
+        getInstance()
+            .labelsToPrices
+            .keySet()
+            .forEach(label -> getInstance().bloomFilter.put(label));
+      }
       return super.build();
     }
   }
+
+  @PrePersist
+  void prePersist() {
+    lastUpdateTime = creationTime;
+  }
+
+  @PostLoad
+  void postLoad() {
+    creationTime = lastUpdateTime;
+  }
 }

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

@@ -16,11 +16,8 @@ package google.registry.model.registry.label;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
-import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static google.registry.config.RegistryConfig.getDomainLabelListCacheDuration;
-import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.registry.label.ReservationType.FULLY_BLOCKED;
 import static google.registry.util.CollectionUtils.nullToEmpty;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
@@ -30,12 +27,8 @@ import com.google.common.base.Splitter;
 import com.google.common.cache.CacheBuilder;
 import com.google.common.cache.CacheLoader;
 import com.google.common.cache.LoadingCache;
-import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.MapDifference;
-import com.google.common.collect.Maps;
-import com.google.common.flogger.FluentLogger;
 import com.google.common.util.concurrent.UncheckedExecutionException;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Embed;
@@ -45,45 +38,58 @@ import com.googlecode.objectify.mapper.Mapper;
 import google.registry.model.Buildable;
 import google.registry.model.registry.Registry;
 import google.registry.model.registry.label.DomainLabelMetrics.MetricsReservedListMatch;
-import google.registry.schema.replay.DatastoreEntity;
-import google.registry.schema.replay.SqlEntity;
-import google.registry.schema.tld.ReservedList.ReservedEntry;
-import google.registry.schema.tld.ReservedListDao;
+import google.registry.schema.replay.DatastoreAndSqlEntity;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.ExecutionException;
 import javax.annotation.Nullable;
+import javax.persistence.CollectionTable;
+import javax.persistence.Column;
+import javax.persistence.ElementCollection;
+import javax.persistence.Embeddable;
+import javax.persistence.Index;
+import javax.persistence.JoinColumn;
+import javax.persistence.MapKeyColumn;
+import javax.persistence.Table;
 import org.joda.time.DateTime;
 
 /**
- * A reserved list entity, persisted to Datastore, that is used to check domain label reservations.
+ * A list of reserved domain labels that are blocked from being registered for various reasons.
+ *
+ * <p>Note that the primary key of this entity is {@link #revisionId}, which is auto-generated by
+ * the database. So, if a retry of insertion happens after the previous attempt unexpectedly
+ * succeeds, we will end up with having two exact same reserved lists that differ only by
+ * revisionId. This is fine though, because we only use the list with the highest revisionId.
  */
 @Entity
+@javax.persistence.Entity
+@Table(indexes = {@Index(columnList = "name", name = "reservedlist_name_idx")})
 public final class ReservedList
     extends BaseDomainLabelList<ReservationType, ReservedList.ReservedListEntry>
-    implements DatastoreEntity {
-
-  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
+    implements DatastoreAndSqlEntity {
 
   @Mapify(ReservedListEntry.LabelMapper.class)
+  @ElementCollection
+  @CollectionTable(
+      name = "ReservedEntry",
+      joinColumns = @JoinColumn(name = "revisionId", referencedColumnName = "revisionId"))
+  @MapKeyColumn(name = "domain_label")
   Map<String, ReservedListEntry> reservedListMap;
 
+  @Column(nullable = false)
   boolean shouldPublish = true;
 
-  @Override
-  public ImmutableList<SqlEntity> toSqlEntities() {
-    return ImmutableList.of(); // ReservedList is dual-written
-  }
-
   /**
    * A reserved list entry entity, persisted to Datastore, that represents a single label and its
    * reservation type.
    */
   @Embed
-  public static class ReservedListEntry
-      extends DomainLabelEntry<ReservationType, ReservedListEntry> implements Buildable {
+  @Embeddable
+  public static class ReservedListEntry extends DomainLabelEntry<ReservationType, ReservedListEntry>
+      implements Buildable {
 
+    @Column(nullable = false)
     ReservationType reservationType;
 
     /** Mapper for use with @Mapify */
@@ -150,6 +156,7 @@ public final class ReservedList
     return shouldPublish;
   }
 
+  /** Returns a {@link Map} of domain labels to {@link ReservedListEntry}. */
   public ImmutableMap<String, ReservedListEntry> getReservedListEntries() {
     return ImmutableMap.copyOf(nullToEmpty(reservedListMap));
   }
@@ -239,65 +246,10 @@ public final class ReservedList
               new CacheLoader<String, ReservedList>() {
                 @Override
                 public ReservedList load(String listName) {
-                  ReservedList datastoreList =
-                      ofy()
-                          .load()
-                          .type(ReservedList.class)
-                          .parent(getCrossTldKey())
-                          .id(listName)
-                          .now();
-                  // Also load the list from Cloud SQL, compare the two lists, and log if different.
-                  try {
-                    loadAndCompareCloudSqlList(datastoreList);
-                  } catch (Throwable t) {
-                    logger.atSevere().withCause(t).log("Error comparing reserved lists.");
-                  }
-                  return datastoreList;
+                  return ReservedListDualWriteDao.getLatestRevision(listName).orElse(null);
                 }
               });
 
-  private static final void loadAndCompareCloudSqlList(ReservedList datastoreList) {
-    Optional<google.registry.schema.tld.ReservedList> maybeCloudSqlList =
-        ReservedListDao.getLatestRevision(datastoreList.getName());
-    if (maybeCloudSqlList.isPresent()) {
-      Map<String, ReservedEntry> datastoreLabelsToReservations =
-          datastoreList.reservedListMap.entrySet().parallelStream()
-              .collect(
-                  toImmutableMap(
-                      Map.Entry::getKey,
-                      entry ->
-                          ReservedEntry.create(
-                              entry.getValue().reservationType, entry.getValue().comment)));
-
-      google.registry.schema.tld.ReservedList cloudSqlList = maybeCloudSqlList.get();
-      MapDifference<String, ReservedEntry> diff =
-          Maps.difference(datastoreLabelsToReservations, cloudSqlList.getLabelsToReservations());
-      if (!diff.areEqual()) {
-        if (diff.entriesDiffering().size() > 10) {
-          logger.atWarning().log(
-              String.format(
-                  "Unequal reserved lists detected, Cloud SQL list with revision"
-                      + " id %d has %d different records than the current"
-                      + " Datastore list.",
-                  cloudSqlList.getRevisionId(), diff.entriesDiffering().size()));
-        } else {
-          StringBuilder diffMessage = new StringBuilder("Unequal reserved lists detected:\n");
-          diff.entriesDiffering()
-              .forEach(
-                  (label, valueDiff) ->
-                      diffMessage.append(
-                          String.format(
-                              "Domain label %s has entry %s in Datastore and entry"
-                                  + " %s in Cloud SQL.\n",
-                              label, valueDiff.leftValue(), valueDiff.rightValue())));
-          logger.atWarning().log(diffMessage.toString());
-        }
-      }
-    } else {
-      logger.atWarning().log("Reserved list in Cloud SQL is empty.");
-    }
-  }
-
   /**
    * Gets the {@link ReservationType} of a label in a single ReservedList, or returns an absent
    * Optional if none exists in the list.

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

@@ -0,0 +1,124 @@
+// Copyright 2020 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.model.registry.label;
+
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
+import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
+import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
+
+import com.google.common.collect.MapDifference;
+import com.google.common.collect.MapDifference.ValueDifference;
+import com.google.common.collect.Maps;
+import com.google.common.flogger.FluentLogger;
+import com.googlecode.objectify.Key;
+import google.registry.model.registry.label.ReservedList.ReservedListEntry;
+import google.registry.persistence.VKey;
+import java.util.Map;
+import java.util.Optional;
+
+/**
+ * A {@link ReservedList} DAO that does dual-write and dual-read against Datastore and Cloud SQL. It
+ * still uses Datastore as the primary storage and suppresses any exception thrown by Cloud SQL.
+ *
+ * <p>TODO(b/160993806): Delete this DAO and switch to use the SQL only DAO after migrating to Cloud
+ * SQL.
+ */
+public class ReservedListDualWriteDao {
+
+  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
+
+  private ReservedListDualWriteDao() {}
+
+  /** Persist a new reserved list to Cloud SQL. */
+  public static void save(ReservedList reservedList) {
+    ofyTm().transact(() -> ofyTm().saveNewOrUpdate(reservedList));
+    try {
+      logger.atInfo().log("Saving reserved list %s to Cloud SQL", reservedList.getName());
+      ReservedListSqlDao.save(reservedList);
+      logger.atInfo().log(
+          "Saved reserved list %s with %d entries to Cloud SQL",
+          reservedList.getName(), reservedList.getReservedListEntries().size());
+    } catch (Throwable t) {
+      logger.atSevere().withCause(t).log("Error saving the reserved list to Cloud SQL.");
+    }
+  }
+
+  /**
+   * Returns the most recent revision of the {@link ReservedList} with the specified name, if it
+   * exists.
+   */
+  public static Optional<ReservedList> getLatestRevision(String reservedListName) {
+    Optional<ReservedList> maybeDatastoreList =
+        ofyTm()
+            .maybeLoad(
+                VKey.createOfy(
+                    ReservedList.class,
+                    Key.create(getCrossTldKey(), ReservedList.class, reservedListName)));
+    try {
+      // Also load the list from Cloud SQL, compare the two lists, and log if different.
+      maybeDatastoreList.ifPresent(ReservedListDualWriteDao::loadAndCompareCloudSqlList);
+    } catch (Throwable t) {
+      logger.atSevere().withCause(t).log("Error comparing reserved lists.");
+    }
+    return maybeDatastoreList;
+  }
+
+  private static void loadAndCompareCloudSqlList(ReservedList datastoreList) {
+    Optional<ReservedList> maybeCloudSqlList =
+        ReservedListSqlDao.getLatestRevision(datastoreList.getName());
+    if (maybeCloudSqlList.isPresent()) {
+      Map<String, ReservedListEntry> datastoreLabelsToReservations =
+          datastoreList.reservedListMap.entrySet().parallelStream()
+              .collect(
+                  toImmutableMap(
+                      Map.Entry::getKey,
+                      entry ->
+                          ReservedListEntry.create(
+                              entry.getKey(),
+                              entry.getValue().reservationType,
+                              entry.getValue().comment)));
+
+      ReservedList cloudSqlList = maybeCloudSqlList.get();
+      MapDifference<String, ReservedListEntry> diff =
+          Maps.difference(datastoreLabelsToReservations, cloudSqlList.reservedListMap);
+      if (!diff.areEqual()) {
+        if (diff.entriesDiffering().size() > 10) {
+          logger.atWarning().log(
+              String.format(
+                  "Unequal reserved lists detected, Cloud SQL list with revision"
+                      + " id %d has %d different records than the current"
+                      + " Datastore list.",
+                  cloudSqlList.getRevisionId(), diff.entriesDiffering().size()));
+        } else {
+          StringBuilder diffMessage = new StringBuilder("Unequal reserved lists detected:\n");
+          diff.entriesDiffering().entrySet().stream()
+              .forEach(
+                  entry -> {
+                    String label = entry.getKey();
+                    ValueDifference<ReservedListEntry> valueDiff = entry.getValue();
+                    diffMessage.append(
+                        String.format(
+                            "Domain label %s has entry %s in Datastore and entry"
+                                + " %s in Cloud SQL.\n",
+                            label, valueDiff.leftValue(), valueDiff.rightValue()));
+                  });
+          logger.atWarning().log(diffMessage.toString());
+        }
+      }
+    } else {
+      logger.atWarning().log("Reserved list in Cloud SQL is empty.");
+    }
+  }
+}

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

@@ -12,20 +12,46 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package google.registry.schema.tld;
+package google.registry.model.registry.label;
 
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 
-import com.google.common.util.concurrent.UncheckedExecutionException;
 import java.util.Optional;
-import java.util.concurrent.ExecutionException;
 
-/** Data access object class for {@link ReservedList} */
-public class ReservedListDao {
+/**
+ * A {@link ReservedList} DAO for Cloud SQL.
+ *
+ * <p>TODO(b/160993806): Rename this class to ReservedListDao after migrating to Cloud SQL.
+ */
+public class ReservedListSqlDao {
+
+  private ReservedListSqlDao() {}
 
   /** Persist a new reserved list to Cloud SQL. */
   public static void save(ReservedList reservedList) {
-    jpaTm().transact(() -> jpaTm().getEntityManager().persist(reservedList));
+    checkArgumentNotNull(reservedList, "Must specify reservedList");
+    jpaTm().transact(() -> jpaTm().saveNew(reservedList));
+  }
+
+  /**
+   * Returns the most recent revision of the {@link ReservedList} with the specified name, if it
+   * exists.
+   */
+  public static Optional<ReservedList> getLatestRevision(String reservedListName) {
+    return jpaTm()
+        .transact(
+            () ->
+                jpaTm()
+                    .getEntityManager()
+                    .createQuery(
+                        "FROM ReservedList rl LEFT JOIN FETCH rl.reservedListMap WHERE"
+                            + " rl.revisionId IN (SELECT MAX(revisionId) FROM ReservedList subrl"
+                            + " WHERE subrl.name = :name)",
+                        ReservedList.class)
+                    .setParameter("name", reservedListName)
+                    .getResultStream()
+                    .findFirst());
   }
 
   /**
@@ -46,39 +72,4 @@ public class ReservedListDao {
                         .size()
                     > 0);
   }
-
-  /**
-   * Returns the most recent revision of the {@link ReservedList} with the specified name, if it
-   * exists. TODO(shicong): Change this method to package level access after dual-read phase.
-   */
-  public static Optional<ReservedList> getLatestRevision(String reservedListName) {
-    return jpaTm()
-        .transact(
-            () ->
-                jpaTm()
-                    .getEntityManager()
-                    .createQuery(
-                        "FROM ReservedList rl LEFT JOIN FETCH rl.labelsToReservations WHERE"
-                            + " rl.revisionId IN (SELECT MAX(revisionId) FROM ReservedList subrl"
-                            + " WHERE subrl.name = :name)",
-                        ReservedList.class)
-                    .setParameter("name", reservedListName)
-                    .getResultStream()
-                    .findFirst());
-  }
-
-  /**
-   * Returns the most recent revision of the {@link ReservedList} with the specified name, from
-   * cache.
-   */
-  public static Optional<ReservedList> getLatestRevisionCached(String reservedListName) {
-    try {
-      return ReservedListCache.cacheReservedLists.get(reservedListName);
-    } catch (ExecutionException e) {
-      throw new UncheckedExecutionException(
-          "Could not retrieve reserved list named " + reservedListName, e);
-    }
-  }
-
-  private ReservedListDao() {}
 }

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

@@ -0,0 +1,58 @@
+// Copyright 2019 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.persistence;
+
+import com.google.common.base.Predicates;
+import com.google.common.collect.ImmutableSet;
+import java.sql.SQLException;
+import java.util.function.Predicate;
+import javax.persistence.OptimisticLockException;
+
+/** Helpers for identifying retriable database operations. */
+public final class JpaRetries {
+
+  private JpaRetries() {}
+
+  private static final ImmutableSet<String> RETRIABLE_TXN_SQL_STATE =
+      ImmutableSet.of(
+          "40001", // serialization_failure
+          "40P01", // deadlock_detected, PSQL-specific
+          "55006", // object_in_use, PSQL and DB2
+          "55P03" // lock_not_available, PSQL-specific
+          );
+
+  private static final Predicate<Throwable> RETRIABLE_TXN_PREDICATE =
+      Predicates.or(
+          OptimisticLockException.class::isInstance,
+          e ->
+              e instanceof SQLException
+                  && RETRIABLE_TXN_SQL_STATE.contains(((SQLException) e).getSQLState()));
+
+  public static boolean isFailedTxnRetriable(Throwable throwable) {
+    Throwable t = throwable;
+    while (t != null) {
+      if (RETRIABLE_TXN_PREDICATE.test(t)) {
+        return true;
+      }
+      t = t.getCause();
+    }
+    return false;
+  }
+
+  public static boolean isFailedQueryRetriable(Throwable throwable) {
+    // TODO(weiminyu): check for more error codes.
+    return isFailedTxnRetriable(throwable);
+  }
+}

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

@@ -58,6 +58,31 @@ public class VKey<T> extends ImmutableObject implements Serializable {
     return new VKey(kind, null, sqlKey);
   }
 
+  /** Creates a {@link VKey} which only contains the ofy primary key. */
+  public static <T> VKey<T> createOfy(
+      Class<? extends T> kind, com.googlecode.objectify.Key<? extends T> ofyKey) {
+    checkArgumentNotNull(kind, "kind must not be null");
+    checkArgumentNotNull(ofyKey, "ofyKey must not be null");
+    return new VKey(kind, ofyKey, null);
+  }
+
+  /**
+   * Creates a {@link VKey} which only contains the ofy primary key by specifying the id of the
+   * {@link Key}.
+   */
+  public static <T> VKey<T> createOfy(Class<? extends T> kind, long id) {
+    return createOfy(kind, Key.create(kind, id));
+  }
+
+  /**
+   * Creates a {@link VKey} which only contains the ofy primary key by specifying the name of the
+   * {@link Key}.
+   */
+  public static <T> VKey<T> createOfy(Class<? extends T> kind, String name) {
+    checkArgumentNotNull(kind, "name must not be null");
+    return createOfy(kind, Key.create(kind, name));
+  }
+
   /** Creates a {@link VKey} which only contains both sql and ofy primary key. */
   public static <T> VKey<T> create(
       Class<? extends T> kind, Object sqlKey, com.googlecode.objectify.Key ofyKey) {

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

@@ -25,4 +25,11 @@ public interface JpaTransactionManager extends TransactionManager {
 
   /** Deletes the entity by its id, throws exception if the entity is not deleted. */
   public abstract <T> void assertDelete(VKey<T> key);
+
+  /**
+   * Releases all resources and shuts down.
+   *
+   * <p>The errorprone check forbids injection of {@link java.io.Closeable} resources.
+   */
+  void teardown();
 }

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

@@ -62,6 +62,11 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     this.clock = clock;
   }
 
+  @Override
+  public void teardown() {
+    emf.close();
+  }
+
   @Override
   public EntityManager getEntityManager() {
     if (transactionInfo.get().entityManager == null) {

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

@@ -75,7 +75,12 @@ public class TransactionManagerFactory {
     return tm;
   }
 
-  /** Returns {@link JpaTransactionManager} instance. */
+  /**
+   * Returns {@link JpaTransactionManager} instance.
+   *
+   * <p>Between invocations of {@link TransactionManagerFactory#setJpaTm} every call to this method
+   * returns the same instance.
+   */
   public static JpaTransactionManager jpaTm() {
     return jpaTm.get();
   }
@@ -93,7 +98,7 @@ public class TransactionManagerFactory {
         RegistryEnvironment.get().equals(RegistryEnvironment.UNITTEST)
             || RegistryToolEnvironment.get() != null,
         "setJpamTm() should only be called by tools and tests.");
-    jpaTm = jpaTmSupplier;
+    jpaTm = Suppliers.memoize(jpaTmSupplier::get);
   }
 
   /** Sets the return of {@link #tm()} to the given instance of {@link TransactionManager}. */

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

@@ -89,7 +89,7 @@ public abstract class RdeModule {
   @Provides
   @Parameter(PARAM_LENIENT)
   static boolean provideLenient(HttpServletRequest req) {
-    return extractBooleanParameter(req, PARAM_REVISION);
+    return extractBooleanParameter(req, PARAM_LENIENT);
   }
 
   @Provides

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

@@ -29,7 +29,7 @@ public enum RdeResourceType {
   DOMAIN("urn:ietf:params:xml:ns:rdeDomain-1.0", EnumSet.of(FULL, THIN)),
   HOST("urn:ietf:params:xml:ns:rdeHost-1.0", EnumSet.of(FULL)),
   REGISTRAR("urn:ietf:params:xml:ns:rdeRegistrar-1.0", EnumSet.of(FULL, THIN)),
-  IDN("urn:ietf:params:xml:ns:rdeIDN-1.0", EnumSet.of(FULL, THIN)),
+  IDN("urn:ietf:params:xml:ns:rdeIDN-1.0", EnumSet.of(FULL)),
   HEADER("urn:ietf:params:xml:ns:rdeHeader-1.0", EnumSet.of(FULL, THIN));
 
   private final String uri;

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

@@ -77,7 +77,7 @@ public final class RdeStagingReducer extends Reducer<PendingDeposit, DepositFrag
   private final byte[] stagingKeyBytes;
   private final RdeMarshaller marshaller;
 
-  private RdeStagingReducer(
+  RdeStagingReducer(
       TaskQueueUtils taskQueueUtils,
       LockHandler lockHandler,
       int gcsBufferSize,
@@ -125,7 +125,7 @@ public final class RdeStagingReducer extends Reducer<PendingDeposit, DepositFrag
     final DateTime watermark = key.watermark();
     final int revision =
         Optional.ofNullable(key.revision())
-            .orElse(RdeRevision.getNextRevision(tld, watermark, mode));
+            .orElseGet(() -> RdeRevision.getNextRevision(tld, watermark, mode));
     String id = RdeUtil.timestampToId(watermark);
     String prefix = RdeNamingUtils.makeRydeFilename(tld, watermark, mode, 1, revision);
     if (key.manual()) {
@@ -168,9 +168,13 @@ public final class RdeStagingReducer extends Reducer<PendingDeposit, DepositFrag
           logger.atSevere().log("Fragment error: %s", fragment.error());
         }
       }
-      for (IdnTableEnum idn : IdnTableEnum.values()) {
-        output.write(marshaller.marshalIdn(idn.getTable()));
-        counter.increment(RdeResourceType.IDN);
+
+      // Don't write the IDN elements for BRDA.
+      if (mode == RdeMode.FULL) {
+        for (IdnTableEnum idn : IdnTableEnum.values()) {
+          output.write(marshaller.marshalIdn(idn.getTable()));
+          counter.increment(RdeResourceType.IDN);
+        }
       }
 
       // Output XML that says how many resources were emitted.

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

@@ -16,6 +16,7 @@ package google.registry.schema.tld;
 
 import com.google.common.collect.ImmutableList;
 import google.registry.model.ImmutableObject;
+import google.registry.model.registry.label.PremiumList;
 import google.registry.schema.replay.DatastoreEntity;
 import google.registry.schema.replay.SqlEntity;
 import java.io.Serializable;

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

@@ -1,151 +0,0 @@
-// Copyright 2019 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.schema.tld;
-
-import static com.google.common.base.Charsets.US_ASCII;
-import static com.google.common.base.Preconditions.checkState;
-import static com.google.common.hash.Funnels.stringFunnel;
-
-import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
-import com.google.common.hash.BloomFilter;
-import google.registry.model.CreateAutoTimestamp;
-import google.registry.schema.replay.DatastoreEntity;
-import google.registry.schema.replay.SqlEntity;
-import java.math.BigDecimal;
-import java.util.Map;
-import javax.annotation.Nullable;
-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.Index;
-import javax.persistence.JoinColumn;
-import javax.persistence.MapKeyColumn;
-import javax.persistence.Table;
-import org.hibernate.LazyInitializationException;
-import org.joda.money.CurrencyUnit;
-import org.joda.time.DateTime;
-
-/**
- * A list of premium prices for domain names.
- *
- * <p>Note that the primary key of this entity is {@link #revisionId}, which is auto-generated by
- * the database. So, if a retry of insertion happens after the previous attempt unexpectedly
- * succeeds, we will end up with having two exact same premium lists that differ only by revisionId.
- * This is fine though, because we only use the list with the highest revisionId.
- */
-@Entity
-@Table(indexes = {@Index(columnList = "name", name = "premiumlist_name_idx")})
-public class PremiumList implements SqlEntity {
-
-  @Column(nullable = false)
-  private String name;
-
-  @Id
-  @GeneratedValue(strategy = GenerationType.IDENTITY)
-  @Column(nullable = false)
-  private Long revisionId;
-
-  @Column(nullable = false)
-  private CreateAutoTimestamp creationTimestamp = CreateAutoTimestamp.create(null);
-
-  @Column(nullable = false)
-  private CurrencyUnit currency;
-
-  @ElementCollection
-  @CollectionTable(
-      name = "PremiumEntry",
-      joinColumns = @JoinColumn(name = "revisionId", referencedColumnName = "revisionId"))
-  @MapKeyColumn(name = "domainLabel")
-  @Column(name = "price", nullable = false)
-  private Map<String, BigDecimal> labelsToPrices;
-
-  @Column(nullable = false)
-  private BloomFilter<String> bloomFilter;
-
-  private PremiumList(String name, CurrencyUnit currency, Map<String, BigDecimal> labelsToPrices) {
-    this.name = name;
-    this.currency = currency;
-    this.labelsToPrices = labelsToPrices;
-    // ASCII is used for the charset because all premium list domain labels are stored punycoded.
-    this.bloomFilter = BloomFilter.create(stringFunnel(US_ASCII), labelsToPrices.size());
-    labelsToPrices.keySet().forEach(this.bloomFilter::put);
-  }
-
-  // Hibernate requires this default constructor.
-  private PremiumList() {}
-
-  /** Constructs a {@link PremiumList} object. */
-  public static PremiumList create(
-      String name, CurrencyUnit currency, Map<String, BigDecimal> labelsToPrices) {
-    return new PremiumList(name, currency, labelsToPrices);
-  }
-
-  /** Returns the name of the premium list, which is usually also a TLD string. */
-  public String getName() {
-    return name;
-  }
-
-  /** Returns the {@link CurrencyUnit} used for this list. */
-  public CurrencyUnit getCurrency() {
-    return currency;
-  }
-
-  /** Returns the ID of this revision, or throws if null. */
-  public Long getRevisionId() {
-    checkState(
-        revisionId != null,
-        "revisionId is null because this object has not yet been persisted to the DB");
-    return revisionId;
-  }
-
-  /** Returns the creation time of this revision of the premium list. */
-  public DateTime getCreationTimestamp() {
-    return creationTimestamp.getTimestamp();
-  }
-
-  /**
-   * Returns a {@link Map} of domain labels to prices.
-   *
-   * <p>Note that this is lazily loaded and thus will throw a {@link LazyInitializationException} if
-   * used outside the transaction in which the given entity was loaded. You generally should not be
-   * using this anyway as it's inefficient to load all of the PremiumEntry rows if you don't need
-   * them. To check prices, use {@link PremiumListDao#getPremiumPrice} instead.
-   */
-  @Nullable
-  public ImmutableMap<String, BigDecimal> getLabelsToPrices() {
-    return labelsToPrices == null ? null : ImmutableMap.copyOf(labelsToPrices);
-  }
-
-  /**
-   * Returns a Bloom filter to determine whether a label might be premium, or is definitely not.
-   *
-   * <p>If the domain label might be premium, then the next step is to check for the existence of a
-   * corresponding row in the PremiumListEntry table. Otherwise, we know for sure it's not premium,
-   * and no DB load is required.
-   */
-  public BloomFilter<String> getBloomFilter() {
-    return bloomFilter;
-  }
-
-  @Override
-  public ImmutableList<DatastoreEntity> toDatastoreEntities() {
-    return ImmutableList.of(); // PremiumList is dual-written
-  }
-}

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

@@ -25,6 +25,7 @@ import com.google.common.annotations.VisibleForTesting;
 import com.google.common.cache.CacheBuilder;
 import com.google.common.cache.CacheLoader;
 import com.google.common.cache.LoadingCache;
+import google.registry.model.registry.label.PremiumList;
 import google.registry.util.NonFinalForTesting;
 import java.math.BigDecimal;
 import java.util.Optional;

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

@@ -20,6 +20,7 @@ import static google.registry.persistence.transaction.TransactionManagerFactory.
 import com.google.common.cache.CacheLoader.InvalidCacheLoadException;
 import com.google.common.util.concurrent.UncheckedExecutionException;
 import google.registry.model.registry.Registry;
+import google.registry.model.registry.label.PremiumList;
 import google.registry.schema.tld.PremiumListCache.RevisionIdAndLabel;
 import java.math.BigDecimal;
 import java.util.Optional;

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

@@ -16,6 +16,7 @@ package google.registry.schema.tld;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
+import static org.joda.time.DateTimeZone.UTC;
 
 import com.google.common.base.Splitter;
 import com.google.common.collect.ImmutableMap;
@@ -23,11 +24,13 @@ import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedSet;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Maps;
+import google.registry.model.registry.label.PremiumList;
 import google.registry.model.registry.label.PremiumList.PremiumListEntry;
 import java.math.BigDecimal;
 import java.util.List;
 import java.util.Map;
 import org.joda.money.CurrencyUnit;
+import org.joda.time.DateTime;
 
 /** Static utility methods for {@link PremiumList}. */
 public class PremiumListUtils {
@@ -37,10 +40,7 @@ public class PremiumListUtils {
         Splitter.on('\n').omitEmptyStrings().splitToList(inputData);
 
     ImmutableMap<String, PremiumListEntry> prices =
-        new google.registry.model.registry.label.PremiumList.Builder()
-            .setName(name)
-            .build()
-            .parse(inputDataPreProcessed);
+        new PremiumList.Builder().setName(name).build().parse(inputDataPreProcessed);
     ImmutableSet<CurrencyUnit> currencies =
         prices.values().stream()
             .map(e -> e.getValue().getCurrencyUnit())
@@ -54,7 +54,12 @@ public class PremiumListUtils {
 
     Map<String, BigDecimal> priceAmounts =
         Maps.transformValues(prices, ple -> ple.getValue().getAmount());
-    return PremiumList.create(name, currency, priceAmounts);
+    return new PremiumList.Builder()
+        .setName(name)
+        .setCurrency(currency)
+        .setLabelsToPrices(priceAmounts)
+        .setCreationTime(DateTime.now(UTC))
+        .build();
   }
 
   private PremiumListUtils() {}

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

@@ -1,177 +0,0 @@
-// Copyright 2019 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.schema.tld;
-
-import static com.google.common.base.Preconditions.checkArgument;
-import static com.google.common.base.Preconditions.checkState;
-import static com.google.common.collect.ImmutableList.sortedCopyOf;
-import static com.google.common.collect.ImmutableList.toImmutableList;
-import static google.registry.util.DomainNameUtils.canonicalizeDomainName;
-
-import com.google.common.base.Joiner;
-import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
-import google.registry.model.CreateAutoTimestamp;
-import google.registry.model.ImmutableObject;
-import google.registry.model.registry.label.ReservationType;
-import google.registry.schema.replay.DatastoreEntity;
-import google.registry.schema.replay.SqlEntity;
-import java.util.Map;
-import java.util.stream.Stream;
-import javax.annotation.Nullable;
-import javax.persistence.CollectionTable;
-import javax.persistence.Column;
-import javax.persistence.ElementCollection;
-import javax.persistence.Embeddable;
-import javax.persistence.Entity;
-import javax.persistence.GeneratedValue;
-import javax.persistence.GenerationType;
-import javax.persistence.Id;
-import javax.persistence.Index;
-import javax.persistence.JoinColumn;
-import javax.persistence.MapKeyColumn;
-import javax.persistence.Table;
-import org.joda.time.DateTime;
-
-/**
- * A list of reserved domain labels that are blocked from being registered for various reasons.
- *
- * <p>Note that the primary key of this entity is {@link #revisionId}, which is auto-generated by
- * the database. So, if a retry of insertion happens after the previous attempt unexpectedly
- * succeeds, we will end up with having two exact same reserved lists that differ only by
- * revisionId. This is fine though, because we only use the list with the highest revisionId.
- */
-@Entity
-@Table(indexes = {@Index(columnList = "name", name = "reservedlist_name_idx")})
-public class ReservedList extends ImmutableObject implements SqlEntity {
-
-  @Column(nullable = false)
-  private String name;
-
-  @Id
-  @GeneratedValue(strategy = GenerationType.IDENTITY)
-  @Column(nullable = false)
-  private Long revisionId;
-
-  @Column(nullable = false)
-  private CreateAutoTimestamp creationTimestamp = CreateAutoTimestamp.create(null);
-
-  @Column(nullable = false)
-  private Boolean shouldPublish;
-
-  @ElementCollection
-  @CollectionTable(
-      name = "ReservedEntry",
-      joinColumns = @JoinColumn(name = "revisionId", referencedColumnName = "revisionId"))
-  @MapKeyColumn(name = "domainLabel")
-  private Map<String, ReservedEntry> labelsToReservations;
-
-  @Override
-  public ImmutableList<DatastoreEntity> toDatastoreEntities() {
-    return ImmutableList.of(); // ReservedList is dual-written\
-  }
-
-  @Embeddable
-  public static class ReservedEntry extends ImmutableObject {
-    @Column(nullable = false)
-    private ReservationType reservationType;
-
-    @Column(nullable = true)
-    private String comment;
-
-    private ReservedEntry(ReservationType reservationType, @Nullable String comment) {
-      this.reservationType = reservationType;
-      this.comment = comment;
-    }
-
-    // Hibernate requires this default constructor.
-    private ReservedEntry() {}
-
-    /** Constructs a {@link ReservedEntry} object. */
-    public static ReservedEntry create(ReservationType reservationType, @Nullable String comment) {
-      return new ReservedEntry(reservationType, comment);
-    }
-
-    /** Returns the reservation type for this entry. */
-    public ReservationType getReservationType() {
-      return reservationType;
-    }
-
-    /** Returns the comment for this entry. Retruns null if there is no comment. */
-    public String getComment() {
-      return comment;
-    }
-  }
-
-  private ReservedList(
-      String name, Boolean shouldPublish, Map<String, ReservedEntry> labelsToReservations) {
-    this.name = name;
-    this.shouldPublish = shouldPublish;
-    this.labelsToReservations = labelsToReservations;
-  }
-
-  // Hibernate requires this default constructor.
-  private ReservedList() {}
-
-  /** Constructs a {@link ReservedList} object. */
-  public static ReservedList create(
-      String name, Boolean shouldPublish, Map<String, ReservedEntry> labelsToReservations) {
-    ImmutableList<String> invalidLabels =
-        labelsToReservations.entrySet().parallelStream()
-            .flatMap(
-                entry -> {
-                  String label = entry.getKey();
-                  if (label.equals(canonicalizeDomainName(label))) {
-                    return Stream.empty();
-                  } else {
-                    return Stream.of(label);
-                  }
-                })
-            .collect(toImmutableList());
-    checkArgument(
-        invalidLabels.isEmpty(),
-        "Label(s) [%s] must be in puny-coded, lower-case form",
-        Joiner.on(",").join(sortedCopyOf(invalidLabels)));
-    return new ReservedList(name, shouldPublish, labelsToReservations);
-  }
-
-  /** Returns the name of the reserved list. */
-  public String getName() {
-    return name;
-  }
-
-  /** Returns the ID of this revision, or throws if null. */
-  public Long getRevisionId() {
-    checkState(
-        revisionId != null,
-        "revisionId is null because this object has not been persisted to the database yet");
-    return revisionId;
-  }
-
-  /** Returns the creation time of this revision of the reserved list. */
-  public DateTime getCreationTimestamp() {
-    return creationTimestamp.getTimestamp();
-  }
-
-  /** Returns a {@link Map} of domain labels to {@link ReservedEntry}. */
-  public ImmutableMap<String, ReservedEntry> getLabelsToReservations() {
-    return ImmutableMap.copyOf(labelsToReservations);
-  }
-
-  /** Returns true if the reserved list should be published. */
-  public Boolean getShouldPublish() {
-    return shouldPublish;
-  }
-}

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

@@ -1,54 +0,0 @@
-// Copyright 2019 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.schema.tld;
-
-import static google.registry.config.RegistryConfig.getDomainLabelListCacheDuration;
-import static java.util.concurrent.TimeUnit.MILLISECONDS;
-
-import com.google.common.annotations.VisibleForTesting;
-import com.google.common.cache.CacheBuilder;
-import com.google.common.cache.CacheLoader;
-import com.google.common.cache.LoadingCache;
-import google.registry.util.NonFinalForTesting;
-import java.util.Optional;
-import org.joda.time.Duration;
-
-/** Caching utils for {@link ReservedList} */
-public class ReservedListCache {
-
-  /**
-   * In-memory cache for reserved lists.
-   *
-   * <p>This is cached for a shorter duration because we need to periodically reload from the DB to
-   * check if a new revision has been published, and if so, then use that.
-   */
-  @NonFinalForTesting
-  static LoadingCache<String, Optional<ReservedList>> cacheReservedLists =
-      createCacheReservedLists(getDomainLabelListCacheDuration());
-
-  @VisibleForTesting
-  static LoadingCache<String, Optional<ReservedList>> createCacheReservedLists(
-      Duration cachePersistDuration) {
-    return CacheBuilder.newBuilder()
-        .expireAfterWrite(cachePersistDuration.getMillis(), MILLISECONDS)
-        .build(
-            new CacheLoader<String, Optional<ReservedList>>() {
-              @Override
-              public Optional<ReservedList> load(String reservedListName) {
-                return ReservedListDao.getLatestRevision(reservedListName);
-              }
-            });
-  }
-}

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

@@ -14,23 +14,12 @@
 
 package google.registry.tools;
 
-import static com.google.common.base.Preconditions.checkArgument;
-import static google.registry.model.registry.label.BaseDomainLabelList.splitOnComment;
-import static google.registry.util.DomainNameUtils.canonicalizeDomainName;
-
 import com.beust.jcommander.Parameter;
-import com.google.common.base.Splitter;
-import com.google.common.collect.HashMultiset;
-import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.Maps;
-import com.google.common.collect.Multiset;
 import com.google.common.flogger.FluentLogger;
-import google.registry.model.registry.label.ReservationType;
-import google.registry.schema.tld.ReservedList.ReservedEntry;
+import google.registry.model.registry.label.ReservedList;
+import google.registry.model.registry.label.ReservedListDualWriteDao;
 import google.registry.tools.params.PathParameter;
 import java.nio.file.Path;
-import java.util.List;
-import java.util.Map;
 import javax.annotation.Nullable;
 
 /**
@@ -62,69 +51,22 @@ public abstract class CreateOrUpdateReservedListCommand extends MutatingCommand
       arity = 1)
   Boolean shouldPublish;
 
-  google.registry.schema.tld.ReservedList cloudSqlReservedList;
-
-  abstract void saveToCloudSql();
+  ReservedList reservedList;
 
   @Override
-  protected String execute() throws Exception {
-    // Save the list to Datastore and output its response.
-    String output = super.execute();
-    logger.atInfo().log(output);
-
-    String cloudSqlMessage =
+  protected String execute() {
+    String message =
         String.format(
             "Saved reserved list %s with %d entries",
-            name, cloudSqlReservedList.getLabelsToReservations().size());
+            name, reservedList.getReservedListEntries().size());
     try {
-      logger.atInfo().log("Saving reserved list to Cloud SQL for TLD %s", name);
-      saveToCloudSql();
-      logger.atInfo().log(cloudSqlMessage);
+      logger.atInfo().log("Saving reserved list for TLD %s", name);
+      ReservedListDualWriteDao.save(reservedList);
+      logger.atInfo().log(message);
     } catch (Throwable e) {
-      cloudSqlMessage =
-          "Unexpected error saving reserved list to Cloud SQL from nomulus tool command";
-      logger.atSevere().withCause(e).log(cloudSqlMessage);
+      message = "Unexpected error saving reserved list from nomulus tool command";
+      logger.atSevere().withCause(e).log(message);
     }
-    return cloudSqlMessage;
-  }
-
-  /** Turns the list CSV data into a map of labels to {@link ReservedEntry}. */
-  static ImmutableMap<String, ReservedEntry> parseToReservationsByLabels(Iterable<String> lines) {
-    Map<String, ReservedEntry> labelsToEntries = Maps.newHashMap();
-    Multiset<String> duplicateLabels = HashMultiset.create();
-    for (String originalLine : lines) {
-      List<String> lineAndComment = splitOnComment(originalLine);
-      if (lineAndComment.isEmpty()) {
-        continue;
-      }
-      String line = lineAndComment.get(0);
-      String comment = lineAndComment.get(1);
-      List<String> parts = Splitter.on(',').trimResults().splitToList(line);
-      checkArgument(
-          parts.size() == 2 || parts.size() == 3,
-          "Could not parse line in reserved list: %s",
-          originalLine);
-      String label = parts.get(0);
-      checkArgument(
-          label.equals(canonicalizeDomainName(label)),
-          "Label '%s' must be in puny-coded, lower-case form",
-          label);
-      ReservationType reservationType = ReservationType.valueOf(parts.get(1));
-      ReservedEntry reservedEntry = ReservedEntry.create(reservationType, comment);
-      // Check if the label was already processed for this list (which is an error), and if so,
-      // accumulate it so that a list of all duplicates can be thrown.
-      if (labelsToEntries.containsKey(label)) {
-        duplicateLabels.add(label, duplicateLabels.contains(label) ? 1 : 2);
-      } else {
-        labelsToEntries.put(label, reservedEntry);
-      }
-    }
-    if (!duplicateLabels.isEmpty()) {
-      throw new IllegalStateException(
-          String.format(
-              "Reserved list cannot contain duplicate labels. Dupes (with counts) were: %s",
-              duplicateLabels));
-    }
-    return ImmutableMap.copyOf(labelsToEntries);
+    return message;
   }
 }

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

@@ -16,7 +16,6 @@ package google.registry.tools;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static google.registry.model.registry.Registries.assertTldExists;
-import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.util.ListNamingUtils.convertFilePathToName;
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.joda.time.DateTimeZone.UTC;
@@ -27,7 +26,6 @@ import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Splitter;
 import com.google.common.base.Strings;
 import google.registry.model.registry.label.ReservedList;
-import google.registry.schema.tld.ReservedListDao;
 import java.nio.file.Files;
 import java.util.List;
 import org.joda.time.DateTime;
@@ -50,15 +48,14 @@ final class CreateReservedListCommand extends CreateOrUpdateReservedListCommand
   protected void init() throws Exception {
     name = Strings.isNullOrEmpty(name) ? convertFilePathToName(input) : name;
     checkArgument(
-        !ReservedList.get(name).isPresent(),
-        "A reserved list already exists by this name");
+        !ReservedList.get(name).isPresent(), "A reserved list already exists by this name");
     if (!override) {
       validateListName(name);
     }
     DateTime now = DateTime.now(UTC);
     List<String> allLines = Files.readAllLines(input, UTF_8);
     boolean shouldPublish = this.shouldPublish == null || this.shouldPublish;
-    ReservedList reservedList =
+    reservedList =
         new ReservedList.Builder()
             .setName(name)
             .setReservedListMapFromLines(allLines)
@@ -66,23 +63,6 @@ final class CreateReservedListCommand extends CreateOrUpdateReservedListCommand
             .setCreationTime(now)
             .setLastUpdateTime(now)
             .build();
-    stageEntityChange(null, reservedList);
-    cloudSqlReservedList =
-        google.registry.schema.tld.ReservedList.create(
-            name, shouldPublish, parseToReservationsByLabels(allLines));
-  }
-
-  @Override
-  void saveToCloudSql() {
-    jpaTm()
-        .transact(
-            () -> {
-              checkArgument(
-                  !ReservedListDao.checkExists(cloudSqlReservedList.getName()),
-                  "A reserved list of this name already exists: %s.",
-                  cloudSqlReservedList.getName());
-              ReservedListDao.save(cloudSqlReservedList);
-            });
   }
 
   private static void validateListName(String name) {

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

@@ -16,17 +16,29 @@ package google.registry.tools;
 
 import com.beust.jcommander.Parameters;
 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 Spec11Pipeline spec11Pipeline;
+  @Inject @Config("projectId") String projectId;
+  @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() {
-    spec11Pipeline.deploy();
+    Spec11Pipeline pipeline = new Spec11Pipeline(projectId, beamStagingUrl, spec11TemplateUrl,
+        reportingBucketUrl, googleCredentialsBundle, retrier);
+    pipeline.deploy();
   }
 }
-

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

@@ -117,7 +117,7 @@ public final class DomainLockUtils {
               RegistryLock newLock =
                   RegistryLockDao.save(lock.asBuilder().setLockCompletionTimestamp(now).build());
               setAsRelock(newLock);
-              tm().transact(() -> applyLockStatuses(newLock, now));
+              tm().transact(() -> applyLockStatuses(newLock, now, isAdmin));
               return newLock;
             });
   }
@@ -171,7 +171,7 @@ public final class DomainLockUtils {
                       createLockBuilder(domainName, registrarId, registrarPocId, isAdmin)
                           .setLockCompletionTimestamp(now)
                           .build());
-              tm().transact(() -> applyLockStatuses(newLock, now));
+              tm().transact(() -> applyLockStatuses(newLock, now, isAdmin));
               setAsRelock(newLock);
               return newLock;
             });
@@ -222,18 +222,18 @@ public final class DomainLockUtils {
       String domainName, String registrarId, @Nullable String registrarPocId, boolean isAdmin) {
     DateTime now = jpaTm().getTransactionTime();
     DomainBase domainBase = getDomain(domainName, registrarId, now);
-    verifyDomainNotLocked(domainBase);
+    verifyDomainNotLocked(domainBase, isAdmin);
 
-    // Multiple pending actions are not allowed
+    // Multiple pending actions are not allowed for non-admins
     RegistryLockDao.getMostRecentByRepoId(domainBase.getRepoId())
         .ifPresent(
             previousLock ->
                 checkArgument(
                     previousLock.isLockRequestExpired(now)
-                        || previousLock.getUnlockCompletionTimestamp().isPresent(),
+                        || previousLock.getUnlockCompletionTimestamp().isPresent()
+                        || isAdmin,
                     "A pending or completed lock action already exists for %s",
                     previousLock.getDomainName()));
-
     return new RegistryLock.Builder()
         .setVerificationCode(stringGenerator.createString(VERIFICATION_CODE_LENGTH))
         .setDomainName(domainName)
@@ -250,6 +250,8 @@ public final class DomainLockUtils {
     Optional<RegistryLock> lockOptional =
         RegistryLockDao.getMostRecentVerifiedLockByRepoId(domainBase.getRepoId());
 
+    verifyDomainLocked(domainBase, isAdmin);
+
     RegistryLock.Builder newLockBuilder;
     if (isAdmin) {
       // Admins should always be able to unlock domains in case we get in a bad state
@@ -265,7 +267,6 @@ public final class DomainLockUtils {
                       .setLockCompletionTimestamp(now)
                       .setRegistrarId(registrarId));
     } else {
-      verifyDomainLocked(domainBase);
       RegistryLock lock =
           lockOptional.orElseThrow(
               () ->
@@ -293,16 +294,17 @@ public final class DomainLockUtils {
         .setRegistrarId(registrarId);
   }
 
-  private static void verifyDomainNotLocked(DomainBase domainBase) {
+  private static void verifyDomainNotLocked(DomainBase domainBase, boolean isAdmin) {
     checkArgument(
-        !domainBase.getStatusValues().containsAll(REGISTRY_LOCK_STATUSES),
+        isAdmin || !domainBase.getStatusValues().containsAll(REGISTRY_LOCK_STATUSES),
         "Domain %s is already locked",
         domainBase.getDomainName());
   }
 
-  private static void verifyDomainLocked(DomainBase domainBase) {
+  private static void verifyDomainLocked(DomainBase domainBase, boolean isAdmin) {
     checkArgument(
-        !Sets.intersection(domainBase.getStatusValues(), REGISTRY_LOCK_STATUSES).isEmpty(),
+        isAdmin || !Sets.intersection(domainBase.getStatusValues(), REGISTRY_LOCK_STATUSES)
+            .isEmpty(),
         "Domain %s is already unlocked",
         domainBase.getDomainName());
   }
@@ -311,7 +313,7 @@ public final class DomainLockUtils {
     DomainBase domain =
         loadByForeignKeyCached(DomainBase.class, domainName, now)
             .orElseThrow(
-                () -> new IllegalArgumentException(String.format("Unknown domain %s", domainName)));
+                () -> new IllegalArgumentException("Domain doesn't exist"));
     // The user must have specified either the correct registrar ID or the admin registrar ID
     checkArgument(
         registryAdminRegistrarId.equals(registrarId)
@@ -330,9 +332,9 @@ public final class DomainLockUtils {
                     String.format("Invalid verification code %s", verificationCode)));
   }
 
-  private void applyLockStatuses(RegistryLock lock, DateTime lockTime) {
+  private void applyLockStatuses(RegistryLock lock, DateTime lockTime, boolean isAdmin) {
     DomainBase domain = getDomain(lock.getDomainName(), lock.getRegistrarId(), lockTime);
-    verifyDomainNotLocked(domain);
+    verifyDomainNotLocked(domain, isAdmin);
 
     DomainBase newDomain =
         domain
@@ -345,9 +347,7 @@ public final class DomainLockUtils {
 
   private void removeLockStatuses(RegistryLock lock, boolean isAdmin, DateTime unlockTime) {
     DomainBase domain = getDomain(lock.getDomainName(), lock.getRegistrarId(), unlockTime);
-    if (!isAdmin) {
-      verifyDomainLocked(domain);
-    }
+    verifyDomainLocked(domain, isAdmin);
 
     DomainBase newDomain =
         domain

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

@@ -14,15 +14,7 @@
 
 package google.registry.tools;
 
-import static google.registry.model.EppResourceUtils.loadByForeignKey;
-
 import com.beust.jcommander.Parameters;
-import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.Sets;
-import com.google.common.flogger.FluentLogger;
-import google.registry.model.domain.DomainBase;
-import google.registry.model.eppcommon.StatusValue;
-import org.joda.time.DateTime;
 
 /**
  * A command to registry lock domain names.
@@ -32,25 +24,6 @@ import org.joda.time.DateTime;
 @Parameters(separators = " =", commandDescription = "Registry lock a domain via EPP.")
 public class LockDomainCommand extends LockOrUnlockDomainCommand {
 
-  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
-
-  @Override
-  protected boolean shouldApplyToDomain(String domain, DateTime now) {
-    DomainBase domainBase =
-        loadByForeignKey(DomainBase.class, domain, now)
-            .orElseThrow(
-                () ->
-                    new IllegalArgumentException(
-                        String.format("Domain '%s' does not exist or is deleted", domain)));
-    ImmutableSet<StatusValue> statusesToAdd =
-        Sets.difference(REGISTRY_LOCK_STATUSES, domainBase.getStatusValues()).immutableCopy();
-    if (statusesToAdd.isEmpty()) {
-      logger.atInfo().log("Domain '%s' is already locked and needs no updates.", domain);
-      return false;
-    }
-    return true;
-  }
-
   @Override
   protected void createAndApplyRequest(String domain) {
     domainLockUtils.administrativelyApplyLock(domain, clientId, null, true);

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

@@ -15,24 +15,28 @@
 package google.registry.tools;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static com.google.common.collect.Iterables.partition;
 import static google.registry.model.eppcommon.StatusValue.SERVER_DELETE_PROHIBITED;
 import static google.registry.model.eppcommon.StatusValue.SERVER_TRANSFER_PROHIBITED;
 import static google.registry.model.eppcommon.StatusValue.SERVER_UPDATE_PROHIBITED;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.CollectionUtils.findDuplicates;
 
 import com.beust.jcommander.Parameter;
 import com.google.common.base.Joiner;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.flogger.FluentLogger;
 import google.registry.config.RegistryConfig.Config;
 import google.registry.model.eppcommon.StatusValue;
 import java.util.List;
 import javax.inject.Inject;
-import org.joda.time.DateTime;
 
-/** Shared base class for commands to registry lock or unlock a domain via EPP. */
+/**
+ * Shared base class for commands to registry lock or unlock a domain via EPP.
+ */
 public abstract class LockOrUnlockDomainCommand extends ConfirmingCommand
     implements CommandWithRemoteApi {
 
@@ -57,7 +61,8 @@ public abstract class LockOrUnlockDomainCommand extends ConfirmingCommand
   @Config("registryAdminClientId")
   String registryAdminClientId;
 
-  @Inject DomainLockUtils domainLockUtils;
+  @Inject
+  DomainLockUtils domainLockUtils;
 
   protected ImmutableSet<String> getDomains() {
     return ImmutableSet.copyOf(mainParameters);
@@ -78,37 +83,34 @@ public abstract class LockOrUnlockDomainCommand extends ConfirmingCommand
   @Override
   protected String execute() {
     ImmutableSet.Builder<String> successfulDomainsBuilder = new ImmutableSet.Builder<>();
-    ImmutableSet.Builder<String> skippedDomainsBuilder = new ImmutableSet.Builder<>();
-    ImmutableSet.Builder<String> failedDomainsBuilder = new ImmutableSet.Builder<>();
+    ImmutableMap.Builder<String, String> failedDomainsToReasons = new ImmutableMap.Builder<>();
     partition(getDomains(), BATCH_SIZE)
         .forEach(
             batch ->
-                tm().transact(
-                        () -> {
-                          for (String domain : batch) {
-                            if (shouldApplyToDomain(domain, tm().getTransactionTime())) {
-                              try {
-                                createAndApplyRequest(domain);
-                              } catch (Throwable t) {
-                                logger.atSevere().withCause(t).log(
-                                    "Error when (un)locking domain %s.", domain);
-                                failedDomainsBuilder.add(domain);
-                              }
-                              successfulDomainsBuilder.add(domain);
-                            } else {
-                              skippedDomainsBuilder.add(domain);
-                            }
-                          }
-                        }));
+                // we require that the jpaTm is the outer transaction in DomainLockUtils
+                jpaTm().transact(() -> tm().transact(
+                    () -> {
+                      for (String domain : batch) {
+                        try {
+                          createAndApplyRequest(domain);
+                        } catch (Throwable t) {
+                          logger.atSevere().withCause(t).log(
+                              "Error when (un)locking domain %s.", domain);
+                          failedDomainsToReasons.put(domain, t.getMessage());
+                          continue;
+                        }
+                        successfulDomainsBuilder.add(domain);
+                      }
+                    })));
     ImmutableSet<String> successfulDomains = successfulDomainsBuilder.build();
-    ImmutableSet<String> skippedDomains = skippedDomainsBuilder.build();
-    ImmutableSet<String> failedDomains = failedDomainsBuilder.build();
+    ImmutableSet<String> failedDomains = failedDomainsToReasons.build().entrySet()
+        .stream()
+        .map(entry -> String.format("%s (%s)", entry.getKey(), entry.getValue()))
+        .collect(toImmutableSet());
     return String.format(
-        "Successfully locked/unlocked domains:\n%s\nSkipped domains:\n%s\nFailed domains:\n%s",
-        successfulDomains, skippedDomains, failedDomains);
+        "Successfully locked/unlocked domains:\n%s\nFailed domains:\n%s",
+        successfulDomains, failedDomains);
   }
 
-  protected abstract boolean shouldApplyToDomain(String domain, DateTime now);
-
   protected abstract void createAndApplyRequest(String domain);
 }

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

@@ -29,7 +29,7 @@ import com.google.appengine.tools.remoteapi.RemoteApiOptions;
 import com.google.common.base.Throwables;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Iterables;
-import google.registry.beam.initsql.BeamJpaModule;
+import google.registry.backup.AppEngineEnvironment;
 import google.registry.config.RegistryConfig;
 import google.registry.model.ofy.ObjectifyService;
 import google.registry.persistence.transaction.TransactionManagerFactory;
@@ -66,6 +66,12 @@ final class RegistryCli implements AutoCloseable, CommandRunner {
               + "If not set, credentials saved by running `nomulus login' will be used.")
   private String credentialJson = null;
 
+  @Parameter(
+      names = {"--sql_access_info"},
+      description = "Name of a file containing space-separated SQL access info used when deploying "
+          + "Beam pipelines")
+  private String sqlAccessInfoFile = null;
+
   // Do not make this final - compile-time constant inlining may interfere with JCommander.
   @ParametersDelegate
   private LoggingParameters loggingParams = new LoggingParameters();
@@ -161,7 +167,7 @@ final class RegistryCli implements AutoCloseable, CommandRunner {
     component =
         DaggerRegistryToolComponent.builder()
             .credentialFilePath(credentialJson)
-            .beamJpaModule(new BeamJpaModule(credentialJson))
+            .sqlAccessInfoFile(sqlAccessInfoFile)
             .build();
 
     // JCommander stores sub-commands as nested JCommander objects containing a list of user objects
@@ -172,7 +178,7 @@ final class RegistryCli implements AutoCloseable, CommandRunner {
             Iterables.getOnlyElement(jcommander.getCommands().get(parsedCommand).getObjects());
     loggingParams.configureLogging();  // Must be called after parameters are parsed.
 
-    try {
+    try (AppEngineEnvironment env = new AppEngineEnvironment()) {
       runCommand(command);
     } catch (RuntimeException ex) {
       if (Throwables.getRootCause(ex) instanceof LoginRequiredException) {

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

@@ -134,6 +134,9 @@ interface RegistryToolComponent {
     @BindsInstance
     Builder credentialFilePath(@Nullable @Config("credentialFilePath") String credentialFilePath);
 
+    @BindsInstance
+    Builder sqlAccessInfoFile(@Nullable @Config("sqlAccessInfoFile") String sqlAccessInfoFile);
+
     Builder beamJpaModule(BeamJpaModule beamJpaModule);
 
     RegistryToolComponent build();

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

@@ -14,16 +14,8 @@
 
 package google.registry.tools;
 
-import static google.registry.model.EppResourceUtils.loadByForeignKey;
-
 import com.beust.jcommander.Parameters;
-import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.Sets;
-import com.google.common.flogger.FluentLogger;
-import google.registry.model.domain.DomainBase;
-import google.registry.model.eppcommon.StatusValue;
 import java.util.Optional;
-import org.joda.time.DateTime;
 
 /**
  * A command to registry unlock domain names.
@@ -33,25 +25,6 @@ import org.joda.time.DateTime;
 @Parameters(separators = " =", commandDescription = "Registry unlock a domain via EPP.")
 public class UnlockDomainCommand extends LockOrUnlockDomainCommand {
 
-  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
-
-  @Override
-  protected boolean shouldApplyToDomain(String domain, DateTime now) {
-    DomainBase domainBase =
-        loadByForeignKey(DomainBase.class, domain, now)
-            .orElseThrow(
-                () ->
-                    new IllegalArgumentException(
-                        String.format("Domain '%s' does not exist or is deleted", domain)));
-    ImmutableSet<StatusValue> statusesToRemove =
-        Sets.intersection(domainBase.getStatusValues(), REGISTRY_LOCK_STATUSES).immutableCopy();
-    if (statusesToRemove.isEmpty()) {
-      logger.atInfo().log("Domain '%s' is already unlocked and needs no updates.", domain);
-      return false;
-    }
-    return true;
-  }
-
   @Override
   protected void createAndApplyRequest(String domain) {
     domainLockUtils.administrativelyApplyUnlock(domain, clientId, true, Optional.empty());

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

@@ -31,7 +31,6 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableMultimap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Sets;
-import com.googlecode.objectify.Key;
 import google.registry.model.billing.BillingEvent;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.Period;
@@ -224,8 +223,8 @@ class UnrenewDomainCommand extends ConfirmingCommand implements CommandWithRemot
             .setRegistrationExpirationTime(newExpirationTime)
             .setLastEppUpdateTime(now)
             .setLastEppUpdateClientId(domain.getCurrentSponsorClientId())
-            .setAutorenewBillingEvent(Key.create(newAutorenewEvent))
-            .setAutorenewPollMessage(Key.create(newAutorenewPollMessage))
+            .setAutorenewBillingEvent(newAutorenewEvent.createVKey())
+            .setAutorenewPollMessage(newAutorenewPollMessage.createVKey())
             .build();
     // In order to do it'll need to write out a new HistoryEntry (likely of type SYNTHETIC), a new
     // autorenew billing event and poll message, and a new one time poll message at the present time

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

@@ -15,18 +15,17 @@
 package google.registry.tools;
 
 import static com.google.common.base.Preconditions.checkArgument;
-import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 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 google.registry.model.registry.label.ReservedList;
-import google.registry.schema.tld.ReservedListDao;
 import google.registry.util.SystemClock;
 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.")
@@ -35,42 +34,20 @@ final class UpdateReservedListCommand extends CreateOrUpdateReservedListCommand
   @Override
   protected void init() throws Exception {
     name = Strings.isNullOrEmpty(name) ? convertFilePathToName(input) : name;
-    // TODO(shicong): Read existing entry from Cloud SQL
     Optional<ReservedList> existing = ReservedList.get(name);
     checkArgument(
         existing.isPresent(), "Could not update reserved list %s because it doesn't exist.", name);
     boolean shouldPublish =
         this.shouldPublish == null ? existing.get().getShouldPublish() : this.shouldPublish;
     List<String> allLines = Files.readAllLines(input, UTF_8);
+    DateTime now = new SystemClock().nowUtc();
     ReservedList.Builder updated =
         existing
             .get()
             .asBuilder()
             .setReservedListMapFromLines(allLines)
-            .setLastUpdateTime(new SystemClock().nowUtc())
+            .setLastUpdateTime(now)
             .setShouldPublish(shouldPublish);
-    stageEntityChange(existing.get(), updated.build());
-    cloudSqlReservedList =
-        google.registry.schema.tld.ReservedList.create(
-            name, shouldPublish, parseToReservationsByLabels(allLines));
-  }
-
-  @Override
-  void saveToCloudSql() {
-    jpaTm()
-        .transact(
-            () -> {
-              // This check is currently disabled because, during the Cloud SQL migration, we need
-              // to be able to update reserved lists in Datastore while simultaneously creating
-              // their first revision in Cloud SQL (i.e. if they haven't been migrated over yet).
-              // TODO(shicong): Re-instate this once all reserved lists are migrated to Cloud SQL,
-              //                 and add a unit test to verity that an exception will be thrown if
-              //                 the reserved list doesn't exist.
-              // checkArgument(
-              //     ReservedListDao.checkExists(cloudSqlReservedList.getName()),
-              //     "A reserved list of this name doesn't exist: %s.",
-              //     cloudSqlReservedList.getName());
-              ReservedListDao.save(cloudSqlReservedList);
-            });
+    reservedList = updated.build();
   }
 }

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

@@ -82,7 +82,7 @@ public class CreatePremiumListAction extends CreateOrUpdatePremiumListAction {
     logger.atInfo().log("Saving premium list to Cloud SQL for TLD %s", name);
     // TODO(mcilwain): Call logInputData() here once Datastore persistence is removed.
 
-    google.registry.schema.tld.PremiumList premiumList = parseToPremiumList(name, inputData);
+    PremiumList premiumList = parseToPremiumList(name, inputData);
     PremiumListDao.saveNew(premiumList);
 
     String message =

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

@@ -74,7 +74,7 @@ public class UpdatePremiumListAction extends CreateOrUpdatePremiumListAction {
   protected void saveToCloudSql() {
     logger.atInfo().log("Updating premium list '%s' in Cloud SQL.", name);
     // TODO(mcilwain): Add logInputData() call here once DB migration is complete.
-    google.registry.schema.tld.PremiumList premiumList = parseToPremiumList(name, inputData);
+    PremiumList premiumList = parseToPremiumList(name, inputData);
     PremiumListDao.update(premiumList);
     String message =
         String.format(

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

@@ -30,6 +30,7 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.flogger.FluentLogger;
 import com.google.gson.Gson;
 import google.registry.config.RegistryConfig.Config;
+import google.registry.flows.domain.DomainFlowUtils;
 import google.registry.model.registrar.Registrar;
 import google.registry.model.registrar.RegistrarContact;
 import google.registry.request.Action;
@@ -118,6 +119,7 @@ public class RegistryLockPostAction implements Runnable, JsonActionRunner.JsonAc
       String registrarId = postInput.registrarId;
       checkArgument(!Strings.isNullOrEmpty(registrarId), "Missing key for registrarId");
       checkArgument(!Strings.isNullOrEmpty(postInput.domainName), "Missing key for domainName");
+      DomainFlowUtils.validateDomainName(postInput.domainName);
       checkNotNull(postInput.isLock, "Missing key for isLock");
       UserAuthInfo userAuthInfo =
           authResult

core/src/main/javascript/google/registry/ui/js/registrar/registry_lock.js

@@ -92,6 +92,7 @@ registry.registrar.RegistryLock.prototype.fillLocksPage_ = function(e) {
             lockEnabledForContact: locksDetails.lockEnabledForContact});
 
     if (locksDetails.lockEnabledForContact) {
+      this.registryLockEmailAddress = locksDetails.email;
       // Listen to the lock-domain 'submit' button click
       var lockButton = goog.dom.getRequiredElement('button-lock-domain');
       goog.events.listen(lockButton, goog.events.EventType.CLICK, this.onLockDomain_, false, this);
@@ -116,7 +117,10 @@ registry.registrar.RegistryLock.prototype.showModal_ = function(targetElement, d
   // attach the modal to the parent element so focus remains correct if the user closes the modal
   var modalElement = goog.soy.renderAsElement(
       registry.soy.registrar.registrylock.confirmModal,
-      {domain: domain, isLock: isLock, isAdmin: this.isAdmin});
+      {domain: domain,
+          isLock: isLock,
+          isAdmin: this.isAdmin,
+          emailAddress: this.registryLockEmailAddress});
   parentElement.prepend(modalElement);
   if (domain == null) {
     goog.dom.getRequiredElement('domain-lock-input-value').focus();

core/src/main/resources/META-INF/persistence.xml

@@ -29,19 +29,19 @@
     <class>google.registry.model.host.HostResource</class>
     <class>google.registry.model.registrar.Registrar</class>
     <class>google.registry.model.registrar.RegistrarContact</class>
+    <class>google.registry.model.registry.label.PremiumList</class>
     <class>google.registry.model.reporting.Spec11ThreatMatch</class>
     <class>google.registry.schema.domain.RegistryLock</class>
     <class>google.registry.schema.tmch.ClaimsList</class>
     <class>google.registry.schema.cursor.Cursor</class>
     <class>google.registry.schema.server.Lock</class>
-    <class>google.registry.schema.tld.PremiumList</class>
     <class>google.registry.schema.tld.PremiumEntry</class>
-    <class>google.registry.schema.tld.ReservedList</class>
     <class>google.registry.model.domain.secdns.DelegationSignerData</class>
     <class>google.registry.model.domain.GracePeriod</class>
     <class>google.registry.model.poll.PollMessage</class>
     <class>google.registry.model.poll.PollMessage$OneTime</class>
     <class>google.registry.model.poll.PollMessage$Autorenew</class>
+    <class>google.registry.model.registry.label.ReservedList</class>
 
     <!-- Customized type converters -->
     <class>google.registry.persistence.converter.BillingCostTransitionConverter</class>

core/src/main/resources/google/registry/ui/soy/registrar/ContactSettings.soy

@@ -107,7 +107,7 @@
   {@param readonly: bool}
   {@param? registryLockAllowedForRegistrar: bool}
   <form name="item" class="{css('item')} {css('registrar')}">
-    <h1>Contact Details</h1>
+    <h1>Contact details</h1>
     {call .contactInfo data="all"}
       {param namePrefix: $namePrefix /}
       {param item: $item /}
@@ -142,11 +142,21 @@
         {param name: 'name' /}
       {/call}
       {call registry.soy.forms.inputFieldRow data="all"}
-        {param label: 'Email' /}
+        {param label: 'Primary account email' /}
         {param namePrefix: $namePrefix /}
         {param name: 'emailAddress' /}
         {param disabled: not $readonly and $item['emailAddress'] != null /}
       {/call}
+      {if isNonnull($item['registryLockEmailAddress'])}
+        {call registry.soy.forms.inputFieldRow data="all"}
+          {param label: 'Registry lock email address' /}
+          {param namePrefix: $namePrefix /}
+          {param name: 'registryLockEmailAddress' /}
+          {param disabled: not $readonly /}
+          {param description: 'Address to which registry (un)lock confirmation emails will be ' +
+              'sent. This is not necessarily the account email address that is used for login.' /}
+        {/call}
+      {/if}
       {call registry.soy.forms.inputFieldRow data="all"}
         {param label: 'Phone' /}
         {param namePrefix: $namePrefix /}
@@ -176,10 +186,6 @@
     {if isNonnull($item['gaeUserId'])}
       <input type="hidden" name="{$namePrefix}gaeUserId" value="{$item['gaeUserId']}">
     {/if}
-    {if isNonnull($item['registryLockEmailAddress'])}
-      <input type="hidden" name="{$namePrefix}registryLockEmailAddress"
-             value="{$item['registryLockEmailAddress']}">
-    {/if}
   </div>
 {/template}
 
@@ -282,19 +288,19 @@
       <hr>
   </tr>
   {call .whoisVisibleRadios_}
-    {param description: 'Show in Registrar WHOIS record as Admin contact' /}
+    {param description: 'Show in Registrar WHOIS record as admin contact' /}
     {param fieldName: $namePrefix + 'visibleInWhoisAsAdmin' /}
     {param visible: $item['visibleInWhoisAsAdmin'] == true /}
   {/call}
   {call .whoisVisibleRadios_}
-    {param description: 'Show in Registrar WHOIS record as Technical contact' /}
+    {param description: 'Show in Registrar WHOIS record as technical contact' /}
     {param fieldName: $namePrefix + 'visibleInWhoisAsTech' /}
     {param visible: $item['visibleInWhoisAsTech'] == true /}
   {/call}
   {call .whoisVisibleRadios_}
     {param description:
-        'Show Phone and Email in Domain WHOIS Record as Registrar Abuse Contact' +
-        ' (Per CL&D Requirements)' /}
+        'Show Phone and Email in Domain WHOIS Record as registrar abuse contact' +
+        ' (per CL&D requirements)' /}
     {param note:
         '*Can only apply to one contact. Selecting Yes for this contact will' +
         ' force this setting for all other contacts to be No.' /}

core/src/main/resources/google/registry/ui/soy/registrar/RegistryLock.soy

@@ -115,12 +115,12 @@
 {template .confirmModal}
   {@param isLock: bool}
   {@param isAdmin: bool}
+  {@param emailAddress: string}
   {@param? domain: string|null}
   <div id="lock-confirm-modal" class="{css('lock-confirm-modal')}">
     <div class="modal-content">
       <p>Are you sure you want to {if $isLock}lock a domain{else}unlock the domain {$domain}{/if}?
-        We will send an email to the email address on file to confirm the {if not $isLock}un{/if}
-        lock.</p>
+        We will send an email to {$emailAddress} to confirm the {if not $isLock}un{/if}lock.</p>
       <label for="domain-to-lock">Domain: </label>
       <input id="domain-lock-input-value"
           {if isNonnull($domain)}

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

@@ -33,29 +33,26 @@ import google.registry.testing.TaskQueueHelper.TaskMatcher;
 import google.registry.util.Retrier;
 import google.registry.util.TaskQueueUtils;
 import org.joda.time.DateTime;
-import org.junit.Before;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Unit tests for {@link CommitLogCheckpointAction}. */
-@RunWith(JUnit4.class)
 public class CommitLogCheckpointActionTest {
 
   private static final String QUEUE_NAME = "export-commits";
 
-  @Rule
+  @RegisterExtension
   public final AppEngineRule appEngine =
       AppEngineRule.builder().withDatastoreAndCloudSql().withTaskQueue().build();
 
-  CommitLogCheckpointStrategy strategy = mock(CommitLogCheckpointStrategy.class);
+  private CommitLogCheckpointStrategy strategy = mock(CommitLogCheckpointStrategy.class);
 
-  DateTime now = DateTime.now(UTC);
-  CommitLogCheckpointAction task = new CommitLogCheckpointAction();
+  private DateTime now = DateTime.now(UTC);
+  private CommitLogCheckpointAction task = new CommitLogCheckpointAction();
 
-  @Before
-  public void before() {
+  @BeforeEach
+  void beforeEach() {
     task.clock = new FakeClock(now);
     task.strategy = strategy;
     task.taskQueueUtils = new TaskQueueUtils(new Retrier(null, 1));
@@ -66,7 +63,7 @@ public class CommitLogCheckpointActionTest {
   }
 
   @Test
-  public void testRun_noCheckpointEverWritten_writesCheckpointAndEnqueuesTask() {
+  void testRun_noCheckpointEverWritten_writesCheckpointAndEnqueuesTask() {
     task.run();
     assertTasksEnqueued(
         QUEUE_NAME,
@@ -78,7 +75,7 @@ public class CommitLogCheckpointActionTest {
   }
 
   @Test
-  public void testRun_checkpointWrittenBeforeNow_writesCheckpointAndEnqueuesTask() {
+  void testRun_checkpointWrittenBeforeNow_writesCheckpointAndEnqueuesTask() {
     DateTime oneMinuteAgo = now.minusMinutes(1);
     persistResource(CommitLogCheckpointRoot.create(oneMinuteAgo));
     task.run();
@@ -92,7 +89,7 @@ public class CommitLogCheckpointActionTest {
   }
 
   @Test
-  public void testRun_checkpointWrittenAfterNow_doesntOverwrite_orEnqueueTask() {
+  void testRun_checkpointWrittenAfterNow_doesntOverwrite_orEnqueueTask() {
     DateTime oneMinuteFromNow = now.plusMinutes(1);
     persistResource(CommitLogCheckpointRoot.create(oneMinuteFromNow));
     task.run();

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

@@ -36,27 +36,22 @@ import google.registry.testing.FakeClock;
 import google.registry.testing.InjectRule;
 import org.joda.time.DateTime;
 import org.joda.time.Duration;
-import org.junit.Before;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Unit tests for {@link CommitLogCheckpointStrategy}. */
-@RunWith(JUnit4.class)
 public class CommitLogCheckpointStrategyTest {
 
-  @Rule
+  @RegisterExtension
   public final AppEngineRule appEngine = AppEngineRule.builder().withDatastoreAndCloudSql().build();
 
-  @Rule
-  public final InjectRule inject = new InjectRule();
+  @RegisterExtension public final InjectRule inject = new InjectRule();
 
-
-  final FakeClock clock = new FakeClock(DateTime.parse("2000-01-01TZ"));
-  final Ofy ofy = new Ofy(clock);
-  final TransactionManager tm = new DatastoreTransactionManager(ofy);
-  final CommitLogCheckpointStrategy strategy = new CommitLogCheckpointStrategy();
+  private final FakeClock clock = new FakeClock(DateTime.parse("2000-01-01TZ"));
+  private final Ofy ofy = new Ofy(clock);
+  private final TransactionManager tm = new DatastoreTransactionManager(ofy);
+  private final CommitLogCheckpointStrategy strategy = new CommitLogCheckpointStrategy();
 
   /**
    * Supplier to inject into CommitLogBucket for doling out predictable bucket IDs.
@@ -64,7 +59,7 @@ public class CommitLogCheckpointStrategyTest {
    * <p>If not overridden, the supplier returns 1 so that other saves won't hit an NPE (since even
    * if they use saveWithoutBackup() the transaction still selects a bucket key early).
    */
-  final FakeSupplier<Integer> fakeBucketIdSupplier = new FakeSupplier<>(1);
+  private final FakeSupplier<Integer> fakeBucketIdSupplier = new FakeSupplier<>(1);
 
   /** Gross but necessary supplier that can be modified to return the desired value. */
   private static class FakeSupplier<T> implements Supplier<T> {
@@ -74,7 +69,7 @@ public class CommitLogCheckpointStrategyTest {
     /** Set this value field to make the supplier return this value. */
     T value = null;
 
-    public FakeSupplier(T defaultValue) {
+    FakeSupplier(T defaultValue) {
       this.defaultValue = defaultValue;
     }
 
@@ -84,8 +79,8 @@ public class CommitLogCheckpointStrategyTest {
     }
   }
 
-  @Before
-  public void before() {
+  @BeforeEach
+  void beforeEach() {
     strategy.clock = clock;
     strategy.ofy = ofy;
 
@@ -102,13 +97,13 @@ public class CommitLogCheckpointStrategyTest {
   }
 
   @Test
-  public void test_readBucketTimestamps_noCommitLogs() {
+  void test_readBucketTimestamps_noCommitLogs() {
     assertThat(strategy.readBucketTimestamps())
         .containsExactly(1, START_OF_TIME, 2, START_OF_TIME, 3, START_OF_TIME);
   }
 
   @Test
-  public void test_readBucketTimestamps_withSomeCommitLogs() {
+  void test_readBucketTimestamps_withSomeCommitLogs() {
     DateTime startTime = clock.nowUtc();
     writeCommitLogToBucket(1);
     clock.advanceOneMilli();
@@ -118,7 +113,7 @@ public class CommitLogCheckpointStrategyTest {
   }
 
   @Test
-  public void test_readBucketTimestamps_againAfterUpdate_reflectsUpdate() {
+  void test_readBucketTimestamps_againAfterUpdate_reflectsUpdate() {
     DateTime firstTime = clock.nowUtc();
     writeCommitLogToBucket(1);
     writeCommitLogToBucket(2);
@@ -133,14 +128,14 @@ public class CommitLogCheckpointStrategyTest {
   }
 
   @Test
-  public void test_readNewCommitLogsAndFindThreshold_noCommitsAtAll_returnsEndOfTime() {
+  void test_readNewCommitLogsAndFindThreshold_noCommitsAtAll_returnsEndOfTime() {
     ImmutableMap<Integer, DateTime> bucketTimes =
         ImmutableMap.of(1, START_OF_TIME, 2, START_OF_TIME, 3, START_OF_TIME);
     assertThat(strategy.readNewCommitLogsAndFindThreshold(bucketTimes)).isEqualTo(END_OF_TIME);
   }
 
   @Test
-  public void test_readNewCommitLogsAndFindThreshold_noNewCommits_returnsEndOfTime() {
+  void test_readNewCommitLogsAndFindThreshold_noNewCommits_returnsEndOfTime() {
     DateTime now = clock.nowUtc();
     writeCommitLogToBucket(1);
     clock.advanceOneMilli();
@@ -153,7 +148,7 @@ public class CommitLogCheckpointStrategyTest {
   }
 
   @Test
-  public void test_readNewCommitLogsAndFindThreshold_tiedNewCommits_returnsCommitTimeMinusOne() {
+  void test_readNewCommitLogsAndFindThreshold_tiedNewCommits_returnsCommitTimeMinusOne() {
     DateTime now = clock.nowUtc();
     writeCommitLogToBucket(1);
     writeCommitLogToBucket(2);
@@ -164,7 +159,7 @@ public class CommitLogCheckpointStrategyTest {
   }
 
   @Test
-  public void test_readNewCommitLogsAndFindThreshold_someNewCommits_returnsEarliestTimeMinusOne() {
+  void test_readNewCommitLogsAndFindThreshold_someNewCommits_returnsEarliestTimeMinusOne() {
     DateTime now = clock.nowUtc();
     writeCommitLogToBucket(1);  // 1A
     writeCommitLogToBucket(2);  // 2A
@@ -191,7 +186,7 @@ public class CommitLogCheckpointStrategyTest {
   }
 
   @Test
-  public void test_readNewCommitLogsAndFindThreshold_commitsAtBucketTimes() {
+  void test_readNewCommitLogsAndFindThreshold_commitsAtBucketTimes() {
     DateTime now = clock.nowUtc();
     ImmutableMap<Integer, DateTime> bucketTimes =
         ImmutableMap.of(1, now.minusMillis(1), 2, now, 3, now.plusMillis(1));
@@ -199,7 +194,7 @@ public class CommitLogCheckpointStrategyTest {
   }
 
   @Test
-  public void test_computeBucketCheckpointTimes_earlyThreshold_setsEverythingToThreshold() {
+  void test_computeBucketCheckpointTimes_earlyThreshold_setsEverythingToThreshold() {
     DateTime now = clock.nowUtc();
     ImmutableMap<Integer, DateTime> bucketTimes =
         ImmutableMap.of(1, now.minusMillis(1), 2, now, 3, now.plusMillis(1));
@@ -208,7 +203,7 @@ public class CommitLogCheckpointStrategyTest {
   }
 
   @Test
-  public void test_computeBucketCheckpointTimes_middleThreshold_clampsToThreshold() {
+  void test_computeBucketCheckpointTimes_middleThreshold_clampsToThreshold() {
     DateTime now = clock.nowUtc();
     ImmutableMap<Integer, DateTime> bucketTimes =
         ImmutableMap.of(1, now.minusMillis(1), 2, now, 3, now.plusMillis(1));
@@ -217,7 +212,7 @@ public class CommitLogCheckpointStrategyTest {
   }
 
   @Test
-  public void test_computeBucketCheckpointTimes_lateThreshold_leavesBucketTimesAsIs() {
+  void test_computeBucketCheckpointTimes_lateThreshold_leavesBucketTimesAsIs() {
     DateTime now = clock.nowUtc();
     ImmutableMap<Integer, DateTime> bucketTimes =
         ImmutableMap.of(1, now.minusMillis(1), 2, now, 3, now.plusMillis(1));
@@ -226,7 +221,7 @@ public class CommitLogCheckpointStrategyTest {
   }
 
   @Test
-  public void test_computeCheckpoint_noCommitsAtAll_bucketCheckpointTimesAreStartOfTime() {
+  void test_computeCheckpoint_noCommitsAtAll_bucketCheckpointTimesAreStartOfTime() {
     assertThat(strategy.computeCheckpoint())
         .isEqualTo(CommitLogCheckpoint.create(
             clock.nowUtc(),
@@ -234,7 +229,7 @@ public class CommitLogCheckpointStrategyTest {
   }
 
   @Test
-  public void test_computeCheckpoint_noNewCommitLogs_bucketCheckpointTimesAreBucketTimes() {
+  void test_computeCheckpoint_noNewCommitLogs_bucketCheckpointTimesAreBucketTimes() {
     DateTime now = clock.nowUtc();
     writeCommitLogToBucket(1);
     clock.advanceOneMilli();
@@ -250,7 +245,7 @@ public class CommitLogCheckpointStrategyTest {
   }
 
   @Test
-  public void test_computeCheckpoint_someNewCommits_bucketCheckpointTimesAreClampedToThreshold() {
+  void test_computeCheckpoint_someNewCommits_bucketCheckpointTimesAreClampedToThreshold() {
     DateTime now = clock.nowUtc();
     writeCommitLogToBucket(1);  // 1A
     writeCommitLogToBucket(2);  // 2A

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

@@ -29,14 +29,11 @@ import google.registry.testing.InjectRule;
 import google.registry.testing.mapreduce.MapreduceTestCase;
 import org.joda.time.DateTime;
 import org.joda.time.Duration;
-import org.junit.Before;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Unit tests for {@link DeleteOldCommitLogsAction}. */
-@RunWith(JUnit4.class)
 public class DeleteOldCommitLogsActionTest
     extends MapreduceTestCase<DeleteOldCommitLogsAction> {
 
@@ -44,11 +41,10 @@ public class DeleteOldCommitLogsActionTest
   private final FakeResponse response = new FakeResponse();
   private ContactResource contact;
 
-  @Rule
-  public final InjectRule inject = new InjectRule();
+  @RegisterExtension public final InjectRule inject = new InjectRule();
 
-  @Before
-  public void setup() {
+  @BeforeEach
+  void beforeEach() {
     inject.setStaticField(Ofy.class, "clock", clock);
     action = new DeleteOldCommitLogsAction();
     action.mrRunner = makeDefaultRunner();
@@ -107,11 +103,9 @@ public class DeleteOldCommitLogsActionTest
     return ImmutableList.copyOf(ofy().load().type(clazz).iterable());
   }
 
-  /**
-   * Check that with very short maxAge, only the referenced elements remain.
-   */
+  /** Check that with very short maxAge, only the referenced elements remain. */
   @Test
-  public void test_shortMaxAge() throws Exception {
+  void test_shortMaxAge() throws Exception {
     runMapreduce(Duration.millis(1));
 
     assertThat(ImmutableList.copyOf(ofy().load().type(CommitLogManifest.class).keys().iterable()))
@@ -121,11 +115,9 @@ public class DeleteOldCommitLogsActionTest
     assertThat(ofyLoadType(CommitLogMutation.class)).hasSize(contact.getRevisions().size() * 3);
   }
 
-  /**
-   * Check that with very long maxAge, all the elements remain.
-   */
+  /** Check that with very long maxAge, all the elements remain. */
   @Test
-  public void test_longMaxAge() throws Exception {
+  void test_longMaxAge() throws Exception {
 
     ImmutableList<CommitLogManifest> initialManifests = ofyLoadType(CommitLogManifest.class);
     ImmutableList<CommitLogMutation> initialMutations = ofyLoadType(CommitLogMutation.class);

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

@@ -39,17 +39,14 @@ import google.registry.testing.GcsTestingUtils;
 import google.registry.testing.TestObject;
 import java.util.List;
 import org.joda.time.DateTime;
-import org.junit.Before;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Unit tests for {@link ExportCommitLogDiffAction}. */
-@RunWith(JUnit4.class)
 public class ExportCommitLogDiffActionTest {
 
-  @Rule
+  @RegisterExtension
   public final AppEngineRule appEngine =
       AppEngineRule.builder()
           .withDatastoreAndCloudSql()
@@ -64,15 +61,15 @@ public class ExportCommitLogDiffActionTest {
 
   private final ExportCommitLogDiffAction task = new ExportCommitLogDiffAction();
 
-  @Before
-  public void before() {
+  @BeforeEach
+  void beforeEach() {
     task.gcsService = gcsService;
     task.gcsBucket = "gcs bucket";
     task.batchSize = 5;
   }
 
   @Test
-  public void testRun_noCommitHistory_onlyUpperCheckpointExported() throws Exception {
+  void testRun_noCommitHistory_onlyUpperCheckpointExported() throws Exception {
     task.lowerCheckpointTime = oneMinuteAgo;
     task.upperCheckpointTime = now;
 
@@ -104,7 +101,7 @@ public class ExportCommitLogDiffActionTest {
   }
 
   @Test
-  public void testRun_regularCommitHistory_exportsCorrectCheckpointDiff() throws Exception {
+  void testRun_regularCommitHistory_exportsCorrectCheckpointDiff() throws Exception {
     task.lowerCheckpointTime = oneMinuteAgo;
     task.upperCheckpointTime = now;
 
@@ -175,7 +172,7 @@ public class ExportCommitLogDiffActionTest {
   }
 
   @Test
-  public void testRun_simultaneousTransactions_bothExported() throws Exception {
+  void testRun_simultaneousTransactions_bothExported() throws Exception {
     task.lowerCheckpointTime = oneMinuteAgo;
     task.upperCheckpointTime = now;
 
@@ -227,7 +224,7 @@ public class ExportCommitLogDiffActionTest {
   }
 
   @Test
-  public void testRun_exportsAcrossMultipleBatches() throws Exception {
+  void testRun_exportsAcrossMultipleBatches() throws Exception {
     task.batchSize = 2;
     task.lowerCheckpointTime = oneMinuteAgo;
     task.upperCheckpointTime = now;
@@ -288,7 +285,7 @@ public class ExportCommitLogDiffActionTest {
   }
 
   @Test
-  public void testRun_checkpointDiffWithNeverTouchedBuckets_exportsCorrectly() throws Exception {
+  void testRun_checkpointDiffWithNeverTouchedBuckets_exportsCorrectly() throws Exception {
     task.lowerCheckpointTime = oneMinuteAgo;
     task.upperCheckpointTime = now;
 
@@ -322,8 +319,7 @@ public class ExportCommitLogDiffActionTest {
   }
 
   @Test
-  public void testRun_checkpointDiffWithNonExistentBucketTimestamps_exportsCorrectly()
-      throws Exception {
+  void testRun_checkpointDiffWithNonExistentBucketTimestamps_exportsCorrectly() throws Exception {
     // Non-existent bucket timestamps can exist when the commit log bucket count was increased
     // recently.
 
@@ -404,7 +400,7 @@ public class ExportCommitLogDiffActionTest {
   }
 
   @Test
-  public void testRun_exportingFromStartOfTime_exportsAllCommits() throws Exception {
+  void testRun_exportingFromStartOfTime_exportsAllCommits() throws Exception {
     task.lowerCheckpointTime = START_OF_TIME;
     task.upperCheckpointTime = now;
 

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

@@ -44,28 +44,25 @@ import java.util.List;
 import java.util.concurrent.Callable;
 import java.util.logging.LogRecord;
 import org.joda.time.DateTime;
-import org.junit.Before;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Unit tests for {@link GcsDiffFileLister}. */
-@RunWith(JUnit4.class)
 public class GcsDiffFileListerTest {
 
-  static final String GCS_BUCKET = "gcs bucket";
+  private static final String GCS_BUCKET = "gcs bucket";
 
-  final DateTime now = DateTime.now(UTC);
-  final GcsDiffFileLister diffLister = new GcsDiffFileLister();
-  final GcsService gcsService = GcsServiceFactory.createGcsService();
+  private final DateTime now = DateTime.now(UTC);
+  private final GcsDiffFileLister diffLister = new GcsDiffFileLister();
+  private final GcsService gcsService = GcsServiceFactory.createGcsService();
   private final TestLogHandler logHandler = new TestLogHandler();
 
-  @Rule
+  @RegisterExtension
   public final AppEngineRule appEngine = AppEngineRule.builder().withDatastoreAndCloudSql().build();
 
-  @Before
-  public void before() throws Exception {
+  @BeforeEach
+  void beforeEach() throws Exception {
     diffLister.gcsService = gcsService;
     diffLister.gcsBucket = GCS_BUCKET;
     diffLister.executor = newDirectExecutorService();
@@ -111,13 +108,13 @@ public class GcsDiffFileListerTest {
   }
 
   @Test
-  public void testList_noFilesFound() {
+  void testList_noFilesFound() {
     DateTime fromTime = now.plusMillis(1);
     assertThat(listDiffFiles(fromTime, null)).isEmpty();
   }
 
   @Test
-  public void testList_patchesHoles() {
+  void testList_patchesHoles() {
     // Fake out the GCS list() method to return only the first and last file.
     // We can't use Mockito.spy() because GcsService's impl is final.
     diffLister.gcsService = (GcsService) newProxyInstance(
@@ -162,7 +159,7 @@ public class GcsDiffFileListerTest {
   }
 
   @Test
-  public void testList_failsOnFork() throws Exception {
+  void testList_failsOnFork() throws Exception {
     // We currently have files for now-4m ... now, construct the following sequence:
     //  now-8m <- now-7m <- now-6m  now-5m <- now-4m ... now
     //    ^___________________________|
@@ -179,7 +176,7 @@ public class GcsDiffFileListerTest {
   }
 
   @Test
-  public void testList_boundaries() {
+  void testList_boundaries() {
     assertThat(listDiffFiles(now.minusMinutes(4), now))
         .containsExactly(
             now.minusMinutes(4),
@@ -192,7 +189,7 @@ public class GcsDiffFileListerTest {
   }
 
   @Test
-  public void testList_failsOnGaps() throws Exception {
+  void testList_failsOnGaps() throws Exception {
     // We currently have files for now-4m ... now, construct the following sequence:
     //  now-8m <- now-7m <- now-6m  {missing} <- now-4m ... now
     for (int i = 6; i < 9; ++i) {
@@ -228,7 +225,7 @@ public class GcsDiffFileListerTest {
   }
 
   @Test
-  public void testList_toTimeSpecified() {
+  void testList_toTimeSpecified() {
     assertThat(listDiffFiles(
             now.minusMinutes(4).minusSeconds(1), now.minusMinutes(2).plusSeconds(1)))
         .containsExactly(

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

@@ -54,31 +54,28 @@ import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
 import org.joda.time.DateTime;
-import org.junit.Before;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Unit tests for {@link RestoreCommitLogsAction}. */
-@RunWith(JUnit4.class)
 public class RestoreCommitLogsActionTest {
 
-  static final String GCS_BUCKET = "gcs bucket";
+  private static final String GCS_BUCKET = "gcs bucket";
 
-  final DateTime now = DateTime.now(UTC);
-  final RestoreCommitLogsAction action = new RestoreCommitLogsAction();
-  final GcsService gcsService = createGcsService();
+  private final DateTime now = DateTime.now(UTC);
+  private final RestoreCommitLogsAction action = new RestoreCommitLogsAction();
+  private final GcsService gcsService = createGcsService();
 
-  @Rule
+  @RegisterExtension
   public final AppEngineRule appEngine =
       AppEngineRule.builder()
           .withDatastoreAndCloudSql()
           .withOfyTestEntities(TestObject.class)
           .build();
 
-  @Before
-  public void init() {
+  @BeforeEach
+  void beforeEach() {
     action.gcsService = gcsService;
     action.dryRun = false;
     action.datastoreService = DatastoreServiceFactory.getDatastoreService();
@@ -91,7 +88,7 @@ public class RestoreCommitLogsActionTest {
   }
 
   @Test
-  public void testRestore_multipleDiffFiles() throws Exception {
+  void testRestore_multipleDiffFiles() throws Exception {
     ofy().saveWithoutBackup().entities(
         TestObject.create("previous to keep"),
         TestObject.create("previous to delete")).now();
@@ -141,7 +138,7 @@ public class RestoreCommitLogsActionTest {
   }
 
   @Test
-  public void testRestore_noManifests() throws Exception {
+  void testRestore_noManifests() throws Exception {
     ofy().saveWithoutBackup().entity(
         TestObject.create("previous to keep")).now();
     saveDiffFileNotToRestore(now.minusMinutes(1));
@@ -155,7 +152,7 @@ public class RestoreCommitLogsActionTest {
   }
 
   @Test
-  public void testRestore_manifestWithNoDeletions() throws Exception {
+  void testRestore_manifestWithNoDeletions() throws Exception {
     ofy().saveWithoutBackup().entity(TestObject.create("previous to keep")).now();
     Key<CommitLogBucket> bucketKey = getBucketKey(1);
     Key<CommitLogManifest> manifestKey = CommitLogManifest.createKey(bucketKey, now);
@@ -174,7 +171,7 @@ public class RestoreCommitLogsActionTest {
 }
 
   @Test
-  public void testRestore_manifestWithNoMutations() throws Exception {
+  void testRestore_manifestWithNoMutations() throws Exception {
     ofy().saveWithoutBackup().entities(
         TestObject.create("previous to keep"),
         TestObject.create("previous to delete")).now();
@@ -195,7 +192,7 @@ public class RestoreCommitLogsActionTest {
 
   // This is a pathological case that shouldn't be possible, but we should be robust to it.
   @Test
-  public void testRestore_manifestWithNoMutationsOrDeletions() throws Exception {
+  void testRestore_manifestWithNoMutationsOrDeletions() throws Exception {
     ofy().saveWithoutBackup().entities(
         TestObject.create("previous to keep")).now();
     saveDiffFileNotToRestore(now.minusMinutes(1));
@@ -211,7 +208,7 @@ public class RestoreCommitLogsActionTest {
   }
 
   @Test
-  public void testRestore_mutateExistingEntity() throws Exception {
+  void testRestore_mutateExistingEntity() throws Exception {
     ofy().saveWithoutBackup().entity(TestObject.create("existing", "a")).now();
     Key<CommitLogManifest> manifestKey = CommitLogManifest.createKey(getBucketKey(1), now);
     saveDiffFileNotToRestore(now.minusMinutes(1));
@@ -229,7 +226,7 @@ public class RestoreCommitLogsActionTest {
 
   // This should be harmless; deletes are idempotent.
   @Test
-  public void testRestore_deleteMissingEntity() throws Exception {
+  void testRestore_deleteMissingEntity() throws Exception {
     ofy().saveWithoutBackup().entity(TestObject.create("previous to keep", "a")).now();
     saveDiffFileNotToRestore(now.minusMinutes(1));
     Iterable<ImmutableObject> commitLogs = saveDiffFile(

core/src/test/java/google/registry/batch/AsyncTaskEnqueuerTest.java

@@ -50,26 +50,24 @@ import google.registry.util.Retrier;
 import java.util.logging.Level;
 import org.joda.time.DateTime;
 import org.joda.time.Duration;
-import org.junit.Before;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.junit.jupiter.api.extension.RegisterExtension;
 import org.mockito.Mock;
-import org.mockito.junit.MockitoJUnit;
-import org.mockito.junit.MockitoRule;
+import org.mockito.junit.jupiter.MockitoExtension;
+import org.mockito.junit.jupiter.MockitoSettings;
+import org.mockito.quality.Strictness;
 
 /** Unit tests for {@link AsyncTaskEnqueuer}. */
-@RunWith(JUnit4.class)
+@ExtendWith(MockitoExtension.class)
 public class AsyncTaskEnqueuerTest {
 
-  @Rule
+  @RegisterExtension
   public final AppEngineRule appEngine =
       AppEngineRule.builder().withDatastoreAndCloudSql().withTaskQueue().build();
 
-  @Rule public final InjectRule inject = new InjectRule();
-
-  @Rule public final MockitoRule mocks = MockitoJUnit.rule();
+  @RegisterExtension public final InjectRule inject = new InjectRule();
 
   @Mock private AppEngineServiceUtils appEngineServiceUtils;
 
@@ -77,8 +75,8 @@ public class AsyncTaskEnqueuerTest {
   private final CapturingLogHandler logHandler = new CapturingLogHandler();
   private final FakeClock clock = new FakeClock(DateTime.parse("2015-05-18T12:34:56Z"));
 
-  @Before
-  public void setUp() {
+  @BeforeEach
+  void beforeEach() {
     LoggerConfig.getConfig(AsyncTaskEnqueuer.class).addHandler(logHandler);
     when(appEngineServiceUtils.getServiceHostname("backend")).thenReturn("backend.hostname.fake");
     asyncTaskEnqueuer = createForTesting(appEngineServiceUtils, clock, standardSeconds(90));
@@ -96,7 +94,7 @@ public class AsyncTaskEnqueuerTest {
   }
 
   @Test
-  public void test_enqueueAsyncResave_success() {
+  void test_enqueueAsyncResave_success() {
     ContactResource contact = persistActiveContact("jd23456");
     asyncTaskEnqueuer.enqueueAsyncResave(contact, clock.nowUtc(), clock.nowUtc().plusDays(5));
     assertTasksEnqueued(
@@ -114,7 +112,7 @@ public class AsyncTaskEnqueuerTest {
   }
 
   @Test
-  public void test_enqueueAsyncResave_multipleResaves() {
+  void test_enqueueAsyncResave_multipleResaves() {
     ContactResource contact = persistActiveContact("jd23456");
     DateTime now = clock.nowUtc();
     asyncTaskEnqueuer.enqueueAsyncResave(
@@ -130,16 +128,15 @@ public class AsyncTaskEnqueuerTest {
             .header("content-type", "application/x-www-form-urlencoded")
             .param(PARAM_RESOURCE_KEY, Key.create(contact).getString())
             .param(PARAM_REQUESTED_TIME, now.toString())
-            .param(
-                PARAM_RESAVE_TIMES,
-                "2015-05-20T14:34:56.000Z,2015-05-21T15:34:56.000Z")
+            .param(PARAM_RESAVE_TIMES, "2015-05-20T14:34:56.000Z,2015-05-21T15:34:56.000Z")
             .etaDelta(
                 standardHours(24).minus(standardSeconds(30)),
                 standardHours(24).plus(standardSeconds(30))));
   }
 
+  @MockitoSettings(strictness = Strictness.LENIENT)
   @Test
-  public void test_enqueueAsyncResave_ignoresTasksTooFarIntoFuture() throws Exception {
+  void test_enqueueAsyncResave_ignoresTasksTooFarIntoFuture() throws Exception {
     ContactResource contact = persistActiveContact("jd23456");
     asyncTaskEnqueuer.enqueueAsyncResave(contact, clock.nowUtc(), clock.nowUtc().plusDays(31));
     assertNoTasksEnqueued(QUEUE_ASYNC_ACTIONS);
@@ -147,7 +144,7 @@ public class AsyncTaskEnqueuerTest {
   }
 
   @Test
-  public void testEnqueueRelock() {
+  void testEnqueueRelock() {
     RegistryLock lock =
         saveRegistryLock(
             new RegistryLock.Builder()
@@ -168,6 +165,7 @@ public class AsyncTaskEnqueuerTest {
         new TaskMatcher()
             .url(RelockDomainAction.PATH)
             .method("POST")
+            .header("Host", "backend.hostname.fake")
             .param(
                 RelockDomainAction.OLD_UNLOCK_REVISION_ID_PARAM,
                 String.valueOf(lock.getRevisionId()))
@@ -176,8 +174,9 @@ public class AsyncTaskEnqueuerTest {
                 standardHours(6).plus(standardSeconds(30))));
   }
 
+  @MockitoSettings(strictness = Strictness.LENIENT)
   @Test
-  public void testFailure_enqueueRelock_noDuration() {
+  void testFailure_enqueueRelock_noDuration() {
     RegistryLock lockWithoutDuration =
         saveRegistryLock(
             new RegistryLock.Builder()

core/src/test/java/google/registry/batch/AsyncTaskMetricsTest.java

@@ -21,19 +21,16 @@ import static google.registry.batch.AsyncTaskMetrics.OperationType.CONTACT_AND_H
 
 import com.google.common.collect.ImmutableSet;
 import google.registry.testing.FakeClock;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.junit.jupiter.api.Test;
 
 /** Unit tests for {@link AsyncTaskMetrics}. */
-@RunWith(JUnit4.class)
-public class AsyncTaskMetricsTest {
+class AsyncTaskMetricsTest {
 
   private final FakeClock clock = new FakeClock();
   private final AsyncTaskMetrics asyncTaskMetrics = new AsyncTaskMetrics(clock);
 
   @Test
-  public void testRecordAsyncFlowResult_calculatesDurationMillisCorrectly() {
+  void testRecordAsyncFlowResult_calculatesDurationMillisCorrectly() {
     asyncTaskMetrics.recordAsyncFlowResult(
         CONTACT_AND_HOST_DELETE,
         SUCCESS,

core/src/test/java/google/registry/batch/DeleteContactsAndHostsActionTest.java

@@ -105,19 +105,16 @@ import google.registry.util.SystemSleeper;
 import java.util.Optional;
 import org.joda.time.DateTime;
 import org.joda.time.Duration;
-import org.junit.Before;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
 import org.mockito.Mock;
 
 /** Unit tests for {@link DeleteContactsAndHostsAction}. */
-@RunWith(JUnit4.class)
 public class DeleteContactsAndHostsActionTest
     extends MapreduceTestCase<DeleteContactsAndHostsAction> {
 
-  @Rule public final InjectRule inject = new InjectRule();
+  @RegisterExtension public final InjectRule inject = new InjectRule();
 
   private AsyncTaskEnqueuer enqueuer;
   private final FakeClock clock = new FakeClock(DateTime.parse("2015-01-15T11:22:33Z"));
@@ -146,8 +143,8 @@ public class DeleteContactsAndHostsActionTest
     ofy().clearSessionCache();
   }
 
-  @Before
-  public void setup() {
+  @BeforeEach
+  void beforeEach() {
     inject.setStaticField(Ofy.class, "clock", clock);
     enqueuer =
         AsyncTaskEnqueuerTest.createForTesting(
@@ -171,7 +168,7 @@ public class DeleteContactsAndHostsActionTest
   }
 
   @Test
-  public void testSuccess_contact_referencedByActiveDomain_doesNotGetDeleted() throws Exception {
+  void testSuccess_contact_referencedByActiveDomain_doesNotGetDeleted() throws Exception {
     ContactResource contact = persistContactPendingDelete("blah8221");
     persistResource(newDomainBase("example.tld", contact));
     DateTime timeEnqueued = clock.nowUtc();
@@ -211,17 +208,17 @@ public class DeleteContactsAndHostsActionTest
   }
 
   @Test
-  public void testSuccess_contact_notReferenced_getsDeleted_andPiiWipedOut() throws Exception {
+  void testSuccess_contact_notReferenced_getsDeleted_andPiiWipedOut() throws Exception {
     runSuccessfulContactDeletionTest(Optional.of("fakeClientTrid"));
   }
 
   @Test
-  public void testSuccess_contact_andNoClientTrid_deletesSuccessfully() throws Exception {
+  void testSuccess_contact_andNoClientTrid_deletesSuccessfully() throws Exception {
     runSuccessfulContactDeletionTest(Optional.empty());
   }
 
   @Test
-  public void test_cannotAcquireLock() {
+  void test_cannotAcquireLock() {
     // Make lock acquisition fail.
     acquireLock();
     enqueueMapreduceOnly();
@@ -229,7 +226,7 @@ public class DeleteContactsAndHostsActionTest
   }
 
   @Test
-  public void test_mapreduceHasWorkToDo_lockIsAcquired() {
+  void test_mapreduceHasWorkToDo_lockIsAcquired() {
     ContactResource contact = persistContactPendingDelete("blah8221");
     persistResource(newDomainBase("example.tld", contact));
     DateTime timeEnqueued = clock.nowUtc();
@@ -244,7 +241,7 @@ public class DeleteContactsAndHostsActionTest
   }
 
   @Test
-  public void test_noTasksToLease_releasesLockImmediately() {
+  void test_noTasksToLease_releasesLockImmediately() {
     enqueueMapreduceOnly();
     // If the Lock was correctly released, then we can acquire it now.
     assertThat(acquireLock()).isPresent();
@@ -293,8 +290,7 @@ public class DeleteContactsAndHostsActionTest
   }
 
   @Test
-  public void testSuccess_contactWithoutPendingTransfer_isDeletedAndHasNoTransferData()
-      throws Exception {
+  void testSuccess_contactWithoutPendingTransfer_isDeletedAndHasNoTransferData() throws Exception {
     ContactResource contact = persistContactPendingDelete("blah8221");
     enqueuer.enqueueAsyncDelete(
         contact,
@@ -308,7 +304,7 @@ public class DeleteContactsAndHostsActionTest
   }
 
   @Test
-  public void testSuccess_contactWithPendingTransfer_getsDeleted() throws Exception {
+  void testSuccess_contactWithPendingTransfer_getsDeleted() throws Exception {
     DateTime transferRequestTime = clock.nowUtc().minusDays(3);
     ContactResource contact =
         persistContactWithPendingTransfer(
@@ -371,7 +367,7 @@ public class DeleteContactsAndHostsActionTest
   }
 
   @Test
-  public void testSuccess_contact_referencedByDeletedDomain_getsDeleted() throws Exception {
+  void testSuccess_contact_referencedByDeletedDomain_getsDeleted() throws Exception {
     ContactResource contactUsed = persistContactPendingDelete("blah1234");
     persistResource(
         newDomainBase("example.tld", contactUsed)
@@ -410,7 +406,7 @@ public class DeleteContactsAndHostsActionTest
   }
 
   @Test
-  public void testSuccess_contact_notRequestedByOwner_doesNotGetDeleted() throws Exception {
+  void testSuccess_contact_notRequestedByOwner_doesNotGetDeleted() throws Exception {
     ContactResource contact = persistContactPendingDelete("jane0991");
     enqueuer.enqueueAsyncDelete(
         contact,
@@ -438,7 +434,7 @@ public class DeleteContactsAndHostsActionTest
   }
 
   @Test
-  public void testSuccess_contact_notRequestedByOwner_isSuperuser_getsDeleted() throws Exception {
+  void testSuccess_contact_notRequestedByOwner_isSuperuser_getsDeleted() throws Exception {
     ContactResource contact = persistContactWithPii("nate007");
     enqueuer.enqueueAsyncDelete(
         contact,
@@ -480,7 +476,7 @@ public class DeleteContactsAndHostsActionTest
   }
 
   @Test
-  public void testSuccess_targetResourcesDontExist_areDelayedForADay() throws Exception {
+  void testSuccess_targetResourcesDontExist_areDelayedForADay() throws Exception {
     ContactResource contactNotSaved = newContactResource("somecontact");
     HostResource hostNotSaved = newHostResource("a11.blah.foo");
     DateTime timeBeforeRun = clock.nowUtc();
@@ -519,7 +515,7 @@ public class DeleteContactsAndHostsActionTest
   }
 
   @Test
-  public void testSuccess_unparseableTasks_areDelayedForADay() throws Exception {
+  void testSuccess_unparseableTasks_areDelayedForADay() throws Exception {
     TaskOptions task =
         TaskOptions.Builder.withMethod(Method.PULL).param("gobbledygook", "kljhadfgsd9f7gsdfh");
     getQueue(QUEUE_ASYNC_DELETE).add(task);
@@ -535,7 +531,7 @@ public class DeleteContactsAndHostsActionTest
   }
 
   @Test
-  public void testSuccess_resourcesNotInPendingDelete_areSkipped() throws Exception {
+  void testSuccess_resourcesNotInPendingDelete_areSkipped() throws Exception {
     ContactResource contact = persistActiveContact("blah2222");
     HostResource host = persistActiveHost("rustles.your.jimmies");
     DateTime timeEnqueued = clock.nowUtc();
@@ -567,7 +563,7 @@ public class DeleteContactsAndHostsActionTest
   }
 
   @Test
-  public void testSuccess_alreadyDeletedResources_areSkipped() throws Exception {
+  void testSuccess_alreadyDeletedResources_areSkipped() throws Exception {
     ContactResource contactDeleted = persistDeletedContact("blah1236", clock.nowUtc().minusDays(2));
     HostResource hostDeleted = persistDeletedHost("a.lim.lop", clock.nowUtc().minusDays(3));
     enqueuer.enqueueAsyncDelete(
@@ -590,7 +586,7 @@ public class DeleteContactsAndHostsActionTest
   }
 
   @Test
-  public void testSuccess_host_referencedByActiveDomain_doesNotGetDeleted() throws Exception {
+  void testSuccess_host_referencedByActiveDomain_doesNotGetDeleted() throws Exception {
     HostResource host = persistHostPendingDelete("ns1.example.tld");
     persistUsedDomain("example.tld", persistActiveContact("abc456"), host);
     DateTime timeEnqueued = clock.nowUtc();
@@ -627,12 +623,12 @@ public class DeleteContactsAndHostsActionTest
   }
 
   @Test
-  public void testSuccess_host_notReferenced_getsDeleted() throws Exception {
+  void testSuccess_host_notReferenced_getsDeleted() throws Exception {
     runSuccessfulHostDeletionTest(Optional.of("fakeClientTrid"));
   }
 
   @Test
-  public void testSuccess_host_andNoClientTrid_deletesSuccessfully() throws Exception {
+  void testSuccess_host_andNoClientTrid_deletesSuccessfully() throws Exception {
     runSuccessfulHostDeletionTest(Optional.empty());
   }
 
@@ -675,7 +671,7 @@ public class DeleteContactsAndHostsActionTest
   }
 
   @Test
-  public void testSuccess_host_referencedByDeletedDomain_getsDeleted() throws Exception {
+  void testSuccess_host_referencedByDeletedDomain_getsDeleted() throws Exception {
     HostResource host = persistHostPendingDelete("ns1.example.tld");
     persistResource(
         newDomainBase("example.tld")
@@ -715,7 +711,7 @@ public class DeleteContactsAndHostsActionTest
   }
 
   @Test
-  public void testSuccess_subordinateHost_getsDeleted() throws Exception {
+  void testSuccess_subordinateHost_getsDeleted() throws Exception {
     DomainBase domain =
         persistResource(
             newDomainBase("example.tld")
@@ -766,7 +762,7 @@ public class DeleteContactsAndHostsActionTest
   }
 
   @Test
-  public void testSuccess_host_notRequestedByOwner_doesNotGetDeleted() throws Exception {
+  void testSuccess_host_notRequestedByOwner_doesNotGetDeleted() throws Exception {
     HostResource host = persistHostPendingDelete("ns2.example.tld");
     enqueuer.enqueueAsyncDelete(
         host,
@@ -794,7 +790,7 @@ public class DeleteContactsAndHostsActionTest
   }
 
   @Test
-  public void testSuccess_host_notRequestedByOwner_isSuperuser_getsDeleted() throws Exception {
+  void testSuccess_host_notRequestedByOwner_isSuperuser_getsDeleted() throws Exception {
     HostResource host = persistHostPendingDelete("ns66.example.tld");
     enqueuer.enqueueAsyncDelete(
         host,
@@ -828,7 +824,7 @@ public class DeleteContactsAndHostsActionTest
   }
 
   @Test
-  public void testSuccess_deleteABunchOfContactsAndHosts_butNotSome() throws Exception {
+  void testSuccess_deleteABunchOfContactsAndHosts_butNotSome() throws Exception {
     ContactResource c1 = persistContactPendingDelete("nsaid54");
     ContactResource c2 = persistContactPendingDelete("nsaid55");
     ContactResource c3 = persistContactPendingDelete("nsaid57");

core/src/test/java/google/registry/batch/DeleteProberDataActionTest.java

@@ -46,29 +46,26 @@ import google.registry.model.registry.Registry;
 import google.registry.model.registry.Registry.TldType;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.testing.FakeResponse;
-import google.registry.testing.SystemPropertyRule;
+import google.registry.testing.SystemPropertyExtension;
 import google.registry.testing.mapreduce.MapreduceTestCase;
 import java.util.Optional;
 import java.util.Set;
 import org.joda.money.Money;
 import org.joda.time.DateTime;
-import org.junit.Before;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Unit tests for {@link DeleteProberDataAction}. */
-@RunWith(JUnit4.class)
-public class DeleteProberDataActionTest extends MapreduceTestCase<DeleteProberDataAction> {
+class DeleteProberDataActionTest extends MapreduceTestCase<DeleteProberDataAction> {
 
   private static final DateTime DELETION_TIME = DateTime.parse("2010-01-01T00:00:00.000Z");
 
-  @Rule
-  public final SystemPropertyRule systemPropertyRule = new SystemPropertyRule();
+  @RegisterExtension
+  final SystemPropertyExtension systemPropertyExtension = new SystemPropertyExtension();
 
-  @Before
-  public void init() {
+  @BeforeEach
+  void beforeEach() {
     // Entities in these two should not be touched.
     createTld("tld", "TLD");
     // Since "example" doesn't end with .test, its entities won't be deleted even though it is of
@@ -96,7 +93,7 @@ public class DeleteProberDataActionTest extends MapreduceTestCase<DeleteProberDa
     action.isDryRun = false;
     action.tlds = ImmutableSet.of();
     action.registryAdminClientId = "TheRegistrar";
-    RegistryEnvironment.SANDBOX.setup(systemPropertyRule);
+    RegistryEnvironment.SANDBOX.setup(systemPropertyExtension);
   }
 
   private void runMapreduce() throws Exception {
@@ -105,7 +102,7 @@ public class DeleteProberDataActionTest extends MapreduceTestCase<DeleteProberDa
   }
 
   @Test
-  public void test_deletesAllAndOnlyProberData() throws Exception {
+  void test_deletesAllAndOnlyProberData() throws Exception {
     Set<ImmutableObject> tldEntities = persistLotsOfDomains("tld");
     Set<ImmutableObject> exampleEntities = persistLotsOfDomains("example");
     Set<ImmutableObject> notTestEntities = persistLotsOfDomains("not-test.test");
@@ -120,7 +117,7 @@ public class DeleteProberDataActionTest extends MapreduceTestCase<DeleteProberDa
   }
 
   @Test
-  public void testSuccess_deletesAllAndOnlyGivenTlds() throws Exception {
+  void testSuccess_deletesAllAndOnlyGivenTlds() throws Exception {
     Set<ImmutableObject> tldEntities = persistLotsOfDomains("tld");
     Set<ImmutableObject> exampleEntities = persistLotsOfDomains("example");
     Set<ImmutableObject> notTestEntities = persistLotsOfDomains("not-test.test");
@@ -136,7 +133,7 @@ public class DeleteProberDataActionTest extends MapreduceTestCase<DeleteProberDa
   }
 
   @Test
-  public void testFail_givenNonTestTld() {
+  void testFail_givenNonTestTld() {
     action.tlds = ImmutableSet.of("not-test.test");
     IllegalArgumentException thrown =
         assertThrows(IllegalArgumentException.class, this::runMapreduce);
@@ -146,7 +143,7 @@ public class DeleteProberDataActionTest extends MapreduceTestCase<DeleteProberDa
   }
 
   @Test
-  public void testFail_givenNonExistentTld() {
+  void testFail_givenNonExistentTld() {
     action.tlds = ImmutableSet.of("non-existent.test");
     IllegalArgumentException thrown =
         assertThrows(IllegalArgumentException.class, this::runMapreduce);
@@ -156,9 +153,9 @@ public class DeleteProberDataActionTest extends MapreduceTestCase<DeleteProberDa
   }
 
   @Test
-  public void testFail_givenNonDotTestTldOnProd() {
+  void testFail_givenNonDotTestTldOnProd() {
     action.tlds = ImmutableSet.of("example");
-    RegistryEnvironment.PRODUCTION.setup(systemPropertyRule);
+    RegistryEnvironment.PRODUCTION.setup(systemPropertyExtension);
     IllegalArgumentException thrown =
         assertThrows(IllegalArgumentException.class, this::runMapreduce);
     assertThat(thrown)
@@ -167,7 +164,7 @@ public class DeleteProberDataActionTest extends MapreduceTestCase<DeleteProberDa
   }
 
   @Test
-  public void testSuccess_doesntDeleteNicDomainForProbers() throws Exception {
+  void testSuccess_doesntDeleteNicDomainForProbers() throws Exception {
     DomainBase nic = persistActiveDomain("nic.ib-any.test");
     ForeignKeyIndex<DomainBase> fkiNic =
         ForeignKeyIndex.load(DomainBase.class, "nic.ib-any.test", START_OF_TIME);
@@ -178,7 +175,7 @@ public class DeleteProberDataActionTest extends MapreduceTestCase<DeleteProberDa
   }
 
   @Test
-  public void testDryRun_doesntDeleteData() throws Exception {
+  void testDryRun_doesntDeleteData() throws Exception {
     Set<ImmutableObject> tldEntities = persistLotsOfDomains("tld");
     Set<ImmutableObject> oaEntities = persistLotsOfDomains("oa-canary.test");
     action.isDryRun = true;
@@ -188,7 +185,7 @@ public class DeleteProberDataActionTest extends MapreduceTestCase<DeleteProberDa
   }
 
   @Test
-  public void testSuccess_activeDomain_isSoftDeleted() throws Exception {
+  void testSuccess_activeDomain_isSoftDeleted() throws Exception {
     DomainBase domain = persistResource(
         newDomainBase("blah.ib-any.test")
             .asBuilder()
@@ -203,7 +200,7 @@ public class DeleteProberDataActionTest extends MapreduceTestCase<DeleteProberDa
   }
 
   @Test
-  public void testSuccess_activeDomain_doubleMapSoftDeletes() throws Exception {
+  void testSuccess_activeDomain_doubleMapSoftDeletes() throws Exception {
     DomainBase domain = persistResource(
         newDomainBase("blah.ib-any.test")
             .asBuilder()
@@ -220,7 +217,7 @@ public class DeleteProberDataActionTest extends MapreduceTestCase<DeleteProberDa
   }
 
   @Test
-  public void test_recentlyCreatedDomain_isntDeletedYet() throws Exception {
+  void test_recentlyCreatedDomain_isntDeletedYet() throws Exception {
     persistResource(
         newDomainBase("blah.ib-any.test")
             .asBuilder()
@@ -234,7 +231,7 @@ public class DeleteProberDataActionTest extends MapreduceTestCase<DeleteProberDa
   }
 
   @Test
-  public void testDryRun_doesntSoftDeleteData() throws Exception {
+  void testDryRun_doesntSoftDeleteData() throws Exception {
     DomainBase domain = persistResource(
         newDomainBase("blah.ib-any.test")
             .asBuilder()
@@ -246,7 +243,7 @@ public class DeleteProberDataActionTest extends MapreduceTestCase<DeleteProberDa
   }
 
   @Test
-  public void test_domainWithSubordinateHosts_isSkipped() throws Exception {
+  void test_domainWithSubordinateHosts_isSkipped() throws Exception {
     persistActiveHost("ns1.blah.ib-any.test");
     DomainBase nakedDomain =
         persistDeletedDomain("todelete.ib-any.test", DateTime.now(UTC).minusYears(1));
@@ -263,7 +260,7 @@ public class DeleteProberDataActionTest extends MapreduceTestCase<DeleteProberDa
   }
 
   @Test
-  public void testFailure_registryAdminClientId_isRequiredForSoftDeletion() {
+  void testFailure_registryAdminClientId_isRequiredForSoftDeletion() {
     persistResource(
         newDomainBase("blah.ib-any.test")
             .asBuilder()

core/src/test/java/google/registry/batch/ExpandRecurringBillingEventsActionTest.java

@@ -59,28 +59,25 @@ import java.util.List;
 import java.util.Optional;
 import org.joda.money.Money;
 import org.joda.time.DateTime;
-import org.junit.Before;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Unit tests for {@link ExpandRecurringBillingEventsAction}. */
-@RunWith(JUnit4.class)
 public class ExpandRecurringBillingEventsActionTest
     extends MapreduceTestCase<ExpandRecurringBillingEventsAction> {
-  @Rule
-  public final InjectRule inject = new InjectRule();
+
+  @RegisterExtension public final InjectRule inject = new InjectRule();
 
   private final DateTime beginningOfTest = DateTime.parse("2000-10-02T00:00:00Z");
   private final FakeClock clock = new FakeClock(beginningOfTest);
 
-  DomainBase domain;
-  HistoryEntry historyEntry;
-  BillingEvent.Recurring recurring;
+  private DomainBase domain;
+  private HistoryEntry historyEntry;
+  private BillingEvent.Recurring recurring;
 
-  @Before
-  public void init() {
+  @BeforeEach
+  void beforeEach() {
     inject.setStaticField(Ofy.class, "clock", clock);
     action = new ExpandRecurringBillingEventsAction();
     action.mrRunner = makeDefaultRunner();
@@ -161,7 +158,7 @@ public class ExpandRecurringBillingEventsActionTest
   }
 
   @Test
-  public void testSuccess_expandSingleEvent() throws Exception {
+  void testSuccess_expandSingleEvent() throws Exception {
     persistResource(recurring);
     action.cursorTimeParam = Optional.of(START_OF_TIME);
     runMapreduce();
@@ -176,7 +173,7 @@ public class ExpandRecurringBillingEventsActionTest
   }
 
   @Test
-  public void testSuccess_expandSingleEvent_deletedDomain() throws Exception {
+  void testSuccess_expandSingleEvent_deletedDomain() throws Exception {
     DateTime deletionTime = DateTime.parse("2000-08-01T00:00:00Z");
     DomainBase deletedDomain = persistDeletedDomain("deleted.tld", deletionTime);
     historyEntry = persistResource(new HistoryEntry.Builder().setParent(deletedDomain).build());
@@ -208,7 +205,7 @@ public class ExpandRecurringBillingEventsActionTest
   }
 
   @Test
-  public void testSuccess_expandSingleEvent_idempotentForDuplicateRuns() throws Exception {
+  void testSuccess_expandSingleEvent_idempotentForDuplicateRuns() throws Exception {
     persistResource(recurring);
     action.cursorTimeParam = Optional.of(START_OF_TIME);
     runMapreduce();
@@ -225,7 +222,7 @@ public class ExpandRecurringBillingEventsActionTest
   }
 
   @Test
-  public void testSuccess_expandSingleEvent_idempotentForExistingOneTime() throws Exception {
+  void testSuccess_expandSingleEvent_idempotentForExistingOneTime() throws Exception {
     persistResource(recurring);
     BillingEvent.OneTime persisted = persistResource(defaultOneTimeBuilder()
         .setParent(historyEntry)
@@ -240,8 +237,7 @@ public class ExpandRecurringBillingEventsActionTest
   }
 
   @Test
-  public void testSuccess_expandSingleEvent_notIdempotentForDifferentBillingTime()
-      throws Exception {
+  void testSuccess_expandSingleEvent_notIdempotentForDifferentBillingTime() throws Exception {
     persistResource(recurring);
     action.cursorTimeParam = Optional.of(START_OF_TIME);
     runMapreduce();
@@ -259,8 +255,7 @@ public class ExpandRecurringBillingEventsActionTest
   }
 
   @Test
-  public void testSuccess_expandSingleEvent_notIdempotentForDifferentRecurring()
-      throws Exception {
+  void testSuccess_expandSingleEvent_notIdempotentForDifferentRecurring() throws Exception {
     persistResource(recurring);
     BillingEvent.Recurring recurring2 = persistResource(recurring.asBuilder()
         .setId(3L)
@@ -289,7 +284,7 @@ public class ExpandRecurringBillingEventsActionTest
   }
 
   @Test
-  public void testSuccess_ignoreRecurringBeforeWindow() throws Exception {
+  void testSuccess_ignoreRecurringBeforeWindow() throws Exception {
     recurring = persistResource(recurring.asBuilder()
         .setEventTime(DateTime.parse("1997-01-05T00:00:00Z"))
         .setRecurrenceEndTime(DateTime.parse("1999-10-05T00:00:00Z"))
@@ -303,7 +298,7 @@ public class ExpandRecurringBillingEventsActionTest
   }
 
   @Test
-  public void testSuccess_ignoreRecurringAfterWindow() throws Exception {
+  void testSuccess_ignoreRecurringAfterWindow() throws Exception {
     recurring = persistResource(recurring.asBuilder()
         .setEventTime(clock.nowUtc().plusYears(2))
         .build());
@@ -315,7 +310,7 @@ public class ExpandRecurringBillingEventsActionTest
   }
 
   @Test
-  public void testSuccess_expandSingleEvent_billingTimeAtCursorTime() throws Exception {
+  void testSuccess_expandSingleEvent_billingTimeAtCursorTime() throws Exception {
     persistResource(recurring);
     action.cursorTimeParam = Optional.of(DateTime.parse("2000-02-19T00:00:00Z"));
     runMapreduce();
@@ -328,8 +323,7 @@ public class ExpandRecurringBillingEventsActionTest
   }
 
   @Test
-  public void testSuccess_expandSingleEvent_cursorTimeBetweenEventAndBillingTime()
-      throws Exception {
+  void testSuccess_expandSingleEvent_cursorTimeBetweenEventAndBillingTime() throws Exception {
     persistResource(recurring);
     action.cursorTimeParam = Optional.of(DateTime.parse("2000-01-12T00:00:00Z"));
     runMapreduce();
@@ -342,7 +336,7 @@ public class ExpandRecurringBillingEventsActionTest
   }
 
   @Test
-  public void testSuccess_expandSingleEvent_billingTimeAtExecutionTime() throws Exception {
+  void testSuccess_expandSingleEvent_billingTimeAtExecutionTime() throws Exception {
     DateTime testTime = DateTime.parse("2000-02-19T00:00:00Z").minusMillis(1);
     persistResource(recurring);
     action.cursorTimeParam = Optional.of(START_OF_TIME);
@@ -359,7 +353,7 @@ public class ExpandRecurringBillingEventsActionTest
   }
 
   @Test
-  public void testSuccess_expandSingleEvent_multipleYearCreate() throws Exception {
+  void testSuccess_expandSingleEvent_multipleYearCreate() throws Exception {
     DateTime testTime = beginningOfTest.plusYears(2);
     action.cursorTimeParam = Optional.of(recurring.getEventTime());
     recurring =
@@ -381,7 +375,7 @@ public class ExpandRecurringBillingEventsActionTest
   }
 
   @Test
-  public void testSuccess_expandSingleEvent_withCursor() throws Exception {
+  void testSuccess_expandSingleEvent_withCursor() throws Exception {
     persistResource(recurring);
     saveCursor(START_OF_TIME);
     runMapreduce();
@@ -394,7 +388,7 @@ public class ExpandRecurringBillingEventsActionTest
   }
 
   @Test
-  public void testSuccess_expandSingleEvent_withCursorPastExpected() throws Exception {
+  void testSuccess_expandSingleEvent_withCursorPastExpected() throws Exception {
     persistResource(recurring);
     // Simulate a quick second run of the mapreduce (this should be a no-op).
     saveCursor(clock.nowUtc().minusSeconds(1));
@@ -406,7 +400,7 @@ public class ExpandRecurringBillingEventsActionTest
   }
 
   @Test
-  public void testSuccess_expandSingleEvent_recurrenceEndBeforeEvent() throws Exception {
+  void testSuccess_expandSingleEvent_recurrenceEndBeforeEvent() throws Exception {
     // This can occur when a domain is transferred or deleted before a domain comes up for renewal.
     recurring = persistResource(recurring.asBuilder()
         .setRecurrenceEndTime(recurring.getEventTime().minusDays(5))
@@ -420,7 +414,7 @@ public class ExpandRecurringBillingEventsActionTest
   }
 
   @Test
-  public void testSuccess_expandSingleEvent_dryRun() throws Exception {
+  void testSuccess_expandSingleEvent_dryRun() throws Exception {
     persistResource(recurring);
     action.isDryRun = true;
     saveCursor(START_OF_TIME); // Need a saved cursor to verify that it didn't move.
@@ -432,7 +426,7 @@ public class ExpandRecurringBillingEventsActionTest
   }
 
   @Test
-  public void testSuccess_expandSingleEvent_multipleYears() throws Exception {
+  void testSuccess_expandSingleEvent_multipleYears() throws Exception {
     DateTime testTime = clock.nowUtc().plusYears(5);
     clock.setTo(testTime);
     List<BillingEvent> expectedEvents = new ArrayList<>();
@@ -463,7 +457,7 @@ public class ExpandRecurringBillingEventsActionTest
   }
 
   @Test
-  public void testSuccess_expandSingleEvent_multipleYears_cursorInBetweenYears() throws Exception {
+  void testSuccess_expandSingleEvent_multipleYears_cursorInBetweenYears() throws Exception {
     DateTime testTime = clock.nowUtc().plusYears(5);
     clock.setTo(testTime);
     List<BillingEvent> expectedEvents = new ArrayList<>();
@@ -492,7 +486,7 @@ public class ExpandRecurringBillingEventsActionTest
   }
 
   @Test
-  public void testSuccess_singleEvent_beforeRenewal() throws Exception {
+  void testSuccess_singleEvent_beforeRenewal() throws Exception {
     DateTime testTime = DateTime.parse("2000-01-04T00:00:00Z");
     clock.setTo(testTime);
     persistResource(recurring);
@@ -505,7 +499,7 @@ public class ExpandRecurringBillingEventsActionTest
   }
 
   @Test
-  public void testSuccess_singleEvent_afterRecurrenceEnd_inAutorenewGracePeriod() throws Exception {
+  void testSuccess_singleEvent_afterRecurrenceEnd_inAutorenewGracePeriod() throws Exception {
     // The domain creation date is 1999-01-05, and the first renewal date is thus 2000-01-05.
     DateTime testTime = DateTime.parse("2001-02-06T00:00:00Z");
     clock.setTo(testTime);
@@ -530,8 +524,7 @@ public class ExpandRecurringBillingEventsActionTest
   }
 
   @Test
-  public void testSuccess_singleEvent_afterRecurrenceEnd_outsideAutorenewGracePeriod()
-      throws Exception {
+  void testSuccess_singleEvent_afterRecurrenceEnd_outsideAutorenewGracePeriod() throws Exception {
     // The domain creation date is 1999-01-05, and the first renewal date is thus 2000-01-05.
     DateTime testTime = DateTime.parse("2001-02-06T00:00:00Z");
     clock.setTo(testTime);
@@ -556,7 +549,7 @@ public class ExpandRecurringBillingEventsActionTest
   }
 
   @Test
-  public void testSuccess_expandSingleEvent_billingTimeOnLeapYear() throws Exception {
+  void testSuccess_expandSingleEvent_billingTimeOnLeapYear() throws Exception {
     recurring =
         persistResource(
             recurring.asBuilder().setEventTime(DateTime.parse("2000-01-15T00:00:00Z")).build());
@@ -575,7 +568,7 @@ public class ExpandRecurringBillingEventsActionTest
   }
 
   @Test
-  public void testSuccess_expandSingleEvent_billingTimeNotOnLeapYear() throws Exception {
+  void testSuccess_expandSingleEvent_billingTimeNotOnLeapYear() throws Exception {
     DateTime testTime = DateTime.parse("2001-12-01T00:00:00Z");
     recurring =
         persistResource(
@@ -597,7 +590,7 @@ public class ExpandRecurringBillingEventsActionTest
   }
 
   @Test
-  public void testSuccess_expandMultipleEvents() throws Exception {
+  void testSuccess_expandMultipleEvents() throws Exception {
     persistResource(recurring);
     BillingEvent.Recurring recurring2 = persistResource(recurring.asBuilder()
         .setEventTime(recurring.getEventTime().plusMonths(3))
@@ -630,7 +623,7 @@ public class ExpandRecurringBillingEventsActionTest
   }
 
   @Test
-  public void testSuccess_premiumDomain() throws Exception {
+  void testSuccess_premiumDomain() throws Exception {
     persistResource(
         Registry.get("tld")
             .asBuilder()
@@ -651,7 +644,7 @@ public class ExpandRecurringBillingEventsActionTest
   }
 
   @Test
-  public void testSuccess_varyingRenewPrices() throws Exception {
+  void testSuccess_varyingRenewPrices() throws Exception {
     DateTime testTime = beginningOfTest.plusYears(1);
     persistResource(
         Registry.get("tld")
@@ -691,7 +684,7 @@ public class ExpandRecurringBillingEventsActionTest
   }
 
   @Test
-  public void testFailure_cursorAfterExecutionTime() {
+  void testFailure_cursorAfterExecutionTime() {
     action.cursorTimeParam = Optional.of(clock.nowUtc().plusYears(1));
     IllegalArgumentException thrown =
         assertThrows(IllegalArgumentException.class, this::runMapreduce);
@@ -701,7 +694,7 @@ public class ExpandRecurringBillingEventsActionTest
   }
 
   @Test
-  public void testFailure_cursorAtExecutionTime() {
+  void testFailure_cursorAtExecutionTime() {
     // The clock advances one milli on runMapreduce.
     action.cursorTimeParam = Optional.of(clock.nowUtc().plusMillis(1));
     IllegalArgumentException thrown =
@@ -712,7 +705,7 @@ public class ExpandRecurringBillingEventsActionTest
   }
 
   @Test
-  public void testFailure_mapperException_doesNotMoveCursor() throws Exception {
+  void testFailure_mapperException_doesNotMoveCursor() throws Exception {
     saveCursor(START_OF_TIME); // Need a saved cursor to verify that it didn't move.
     // Set target to a TLD that doesn't exist.
     recurring = persistResource(recurring.asBuilder().setTargetId("domain.junk").build());

core/src/test/java/google/registry/batch/RefreshDnsOnHostRenameActionTest.java

@@ -61,27 +61,24 @@ import google.registry.util.SystemSleeper;
 import java.util.Optional;
 import org.joda.time.DateTime;
 import org.joda.time.Duration;
-import org.junit.Before;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
 import org.mockito.Mock;
 
 /** Unit tests for {@link RefreshDnsOnHostRenameAction}. */
-@RunWith(JUnit4.class)
 public class RefreshDnsOnHostRenameActionTest
     extends MapreduceTestCase<RefreshDnsOnHostRenameAction> {
 
-  @Rule public final InjectRule inject = new InjectRule();
+  @RegisterExtension public final InjectRule inject = new InjectRule();
 
   private AsyncTaskEnqueuer enqueuer;
   private final FakeClock clock = new FakeClock(DateTime.parse("2015-01-15T11:22:33Z"));
   private final FakeResponse fakeResponse = new FakeResponse();
   @Mock private RequestStatusChecker requestStatusChecker;
 
-  @Before
-  public void setup() {
+  @BeforeEach
+  void beforeEach() {
     createTld("tld");
     enqueuer =
         AsyncTaskEnqueuerTest.createForTesting(
@@ -124,7 +121,7 @@ public class RefreshDnsOnHostRenameActionTest
   }
 
   @Test
-  public void testSuccess_dnsUpdateEnqueued() throws Exception {
+  void testSuccess_dnsUpdateEnqueued() throws Exception {
     HostResource host = persistActiveHost("ns1.example.tld");
     persistResource(newDomainBase("example.tld", host));
     persistResource(newDomainBase("otherexample.tld", host));
@@ -141,7 +138,7 @@ public class RefreshDnsOnHostRenameActionTest
   }
 
   @Test
-  public void testSuccess_multipleHostsProcessedInBatch() throws Exception {
+  void testSuccess_multipleHostsProcessedInBatch() throws Exception {
     HostResource host1 = persistActiveHost("ns1.example.tld");
     HostResource host2 = persistActiveHost("ns2.example.tld");
     HostResource host3 = persistActiveHost("ns3.example.tld");
@@ -165,7 +162,7 @@ public class RefreshDnsOnHostRenameActionTest
   }
 
   @Test
-  public void testSuccess_deletedHost_doesntTriggerDnsRefresh() throws Exception {
+  void testSuccess_deletedHost_doesntTriggerDnsRefresh() throws Exception {
     HostResource host = persistDeletedHost("ns11.fakesss.tld", clock.nowUtc().minusDays(4));
     persistResource(newDomainBase("example1.tld", host));
     DateTime timeEnqueued = clock.nowUtc();
@@ -180,7 +177,7 @@ public class RefreshDnsOnHostRenameActionTest
   }
 
   @Test
-  public void testSuccess_noDnsTasksForDeletedDomain() throws Exception {
+  void testSuccess_noDnsTasksForDeletedDomain() throws Exception {
     HostResource renamedHost = persistActiveHost("ns1.example.tld");
     persistResource(
         newDomainBase("example.tld", renamedHost)
@@ -194,7 +191,7 @@ public class RefreshDnsOnHostRenameActionTest
   }
 
   @Test
-  public void testRun_hostDoesntExist_delaysTask() throws Exception {
+  void testRun_hostDoesntExist_delaysTask() throws Exception {
     HostResource host = newHostResource("ns1.example.tld");
     enqueuer.enqueueAsyncDnsRefresh(host, clock.nowUtc());
     enqueueMapreduceOnly();
@@ -208,7 +205,7 @@ public class RefreshDnsOnHostRenameActionTest
   }
 
   @Test
-  public void test_cannotAcquireLock() {
+  void test_cannotAcquireLock() {
     // Make lock acquisition fail.
     acquireLock();
     enqueueMapreduceOnly();
@@ -217,7 +214,7 @@ public class RefreshDnsOnHostRenameActionTest
   }
 
   @Test
-  public void test_mapreduceHasWorkToDo_lockIsAcquired() {
+  void test_mapreduceHasWorkToDo_lockIsAcquired() {
     HostResource host = persistActiveHost("ns1.example.tld");
     enqueuer.enqueueAsyncDnsRefresh(host, clock.nowUtc());
     enqueueMapreduceOnly();
@@ -225,7 +222,7 @@ public class RefreshDnsOnHostRenameActionTest
   }
 
   @Test
-  public void test_noTasksToLease_releasesLockImmediately() throws Exception {
+  void test_noTasksToLease_releasesLockImmediately() throws Exception {
     enqueueMapreduceOnly();
     assertNoDnsTasksEnqueued();
     assertNoTasksEnqueued(QUEUE_ASYNC_HOST_RENAME);

core/src/test/java/google/registry/batch/RelockDomainActionTest.java

@@ -44,14 +44,11 @@ import google.registry.util.AppEngineServiceUtils;
 import google.registry.util.StringGenerator.Alphabets;
 import java.util.Optional;
 import org.joda.time.Duration;
-import org.junit.Before;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Unit tests for {@link RelockDomainAction}. */
-@RunWith(JUnit4.class)
 public class RelockDomainActionTest {
 
   private static final String DOMAIN_NAME = "example.tld";
@@ -67,7 +64,7 @@ public class RelockDomainActionTest {
           AsyncTaskEnqueuerTest.createForTesting(
               mock(AppEngineServiceUtils.class), clock, Duration.ZERO));
 
-  @Rule
+  @RegisterExtension
   public final AppEngineRule appEngineRule =
       AppEngineRule.builder()
           .withDatastoreAndCloudSql()
@@ -78,8 +75,8 @@ public class RelockDomainActionTest {
   private RegistryLock oldLock;
   private RelockDomainAction action;
 
-  @Before
-  public void setup() {
+  @BeforeEach
+  void beforeEach() {
     createTlds("tld", "net");
     HostResource host = persistActiveHost("ns1.example.net");
     domain = persistResource(newDomainBase(DOMAIN_NAME, host));
@@ -95,7 +92,7 @@ public class RelockDomainActionTest {
   }
 
   @Test
-  public void testLock() {
+  void testLock() {
     action.run();
     assertThat(reloadDomain(domain).getStatusValues())
         .containsAtLeastElementsIn(REGISTRY_LOCK_STATUSES);
@@ -107,7 +104,7 @@ public class RelockDomainActionTest {
   }
 
   @Test
-  public void testFailure_unknownCode() {
+  void testFailure_unknownCode() {
     action = createAction(12128675309L);
     action.run();
     assertThat(response.getStatus()).isEqualTo(SC_NO_CONTENT);
@@ -115,7 +112,7 @@ public class RelockDomainActionTest {
   }
 
   @Test
-  public void testFailure_pendingDelete() {
+  void testFailure_pendingDelete() {
     persistResource(domain.asBuilder().setStatusValues(ImmutableSet.of(PENDING_DELETE)).build());
     action.run();
     assertThat(response.getStatus()).isEqualTo(SC_NO_CONTENT);
@@ -124,7 +121,7 @@ public class RelockDomainActionTest {
   }
 
   @Test
-  public void testFailure_pendingTransfer() {
+  void testFailure_pendingTransfer() {
     persistResource(domain.asBuilder().setStatusValues(ImmutableSet.of(PENDING_TRANSFER)).build());
     action.run();
     assertThat(response.getStatus()).isEqualTo(SC_NO_CONTENT);
@@ -133,7 +130,7 @@ public class RelockDomainActionTest {
   }
 
   @Test
-  public void testFailure_domainAlreadyLocked() {
+  void testFailure_domainAlreadyLocked() {
     domainLockUtils.administrativelyApplyLock(DOMAIN_NAME, CLIENT_ID, null, true);
     action.run();
     assertThat(response.getStatus()).isEqualTo(SC_NO_CONTENT);
@@ -142,7 +139,7 @@ public class RelockDomainActionTest {
   }
 
   @Test
-  public void testFailure_domainDeleted() {
+  void testFailure_domainDeleted() {
     persistDomainAsDeleted(domain, clock.nowUtc());
     action.run();
     assertThat(response.getStatus()).isEqualTo(SC_NO_CONTENT);
@@ -151,7 +148,7 @@ public class RelockDomainActionTest {
   }
 
   @Test
-  public void testFailure_domainTransferred() {
+  void testFailure_domainTransferred() {
     persistResource(domain.asBuilder().setPersistedCurrentSponsorClientId("NewRegistrar").build());
     action.run();
     assertThat(response.getStatus()).isEqualTo(SC_NO_CONTENT);
@@ -164,7 +161,7 @@ public class RelockDomainActionTest {
   }
 
   @Test
-  public void testFailure_relockAlreadySet() {
+  void testFailure_relockAlreadySet() {
     RegistryLock newLock =
         domainLockUtils.administrativelyApplyLock(DOMAIN_NAME, CLIENT_ID, null, true);
     saveRegistryLock(oldLock.asBuilder().setRelock(newLock).build());

core/src/test/java/google/registry/batch/ResaveAllEppResourcesActionTest.java

@@ -25,18 +25,14 @@ import google.registry.model.transfer.TransferStatus;
 import google.registry.testing.FakeResponse;
 import google.registry.testing.mapreduce.MapreduceTestCase;
 import org.joda.time.DateTime;
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
 
 /** Unit tests for {@link ResaveAllEppResourcesAction}. */
-@RunWith(JUnit4.class)
-public class ResaveAllEppResourcesActionTest
-    extends MapreduceTestCase<ResaveAllEppResourcesAction> {
+class ResaveAllEppResourcesActionTest extends MapreduceTestCase<ResaveAllEppResourcesAction> {
 
-  @Before
-  public void init() {
+  @BeforeEach
+  void beforeEach() {
     action = new ResaveAllEppResourcesAction();
     action.mrRunner = makeDefaultRunner();
     action.response = new FakeResponse();
@@ -48,19 +44,19 @@ public class ResaveAllEppResourcesActionTest
   }
 
   @Test
-  public void test_mapreduceSuccessfullyResavesEntity() throws Exception {
+  void test_mapreduceSuccessfullyResavesEntity() throws Exception {
     ContactResource contact = persistActiveContact("test123");
-    DateTime creationTime = contact.getUpdateAutoTimestamp().getTimestamp();
-    assertThat(ofy().load().entity(contact).now().getUpdateAutoTimestamp().getTimestamp())
+    DateTime creationTime = contact.getUpdateTimestamp().getTimestamp();
+    assertThat(ofy().load().entity(contact).now().getUpdateTimestamp().getTimestamp())
         .isEqualTo(creationTime);
     ofy().clearSessionCache();
     runMapreduce();
-    assertThat(ofy().load().entity(contact).now().getUpdateAutoTimestamp().getTimestamp())
+    assertThat(ofy().load().entity(contact).now().getUpdateTimestamp().getTimestamp())
         .isGreaterThan(creationTime);
   }
 
   @Test
-  public void test_mapreduceResolvesPendingTransfer() throws Exception {
+  void test_mapreduceResolvesPendingTransfer() throws Exception {
     DateTime now = DateTime.now(UTC);
     // Set up a contact with a transfer that implicitly completed five days ago.
     ContactResource contact =

core/src/test/java/google/registry/batch/ResaveEntityActionTest.java

@@ -49,33 +49,32 @@ import google.registry.testing.TaskQueueHelper.TaskMatcher;
 import google.registry.util.AppEngineServiceUtils;
 import org.joda.time.DateTime;
 import org.joda.time.Duration;
-import org.junit.Before;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.junit.jupiter.api.extension.RegisterExtension;
 import org.mockito.Mock;
-import org.mockito.junit.MockitoJUnit;
-import org.mockito.junit.MockitoRule;
+import org.mockito.junit.jupiter.MockitoExtension;
+import org.mockito.junit.jupiter.MockitoSettings;
+import org.mockito.quality.Strictness;
 
 /** Unit tests for {@link ResaveEntityAction}. */
-@RunWith(JUnit4.class)
+@ExtendWith(MockitoExtension.class)
 public class ResaveEntityActionTest {
 
-  @Rule
+  @RegisterExtension
   public final AppEngineRule appEngine =
       AppEngineRule.builder().withDatastoreAndCloudSql().withTaskQueue().build();
 
-  @Rule public final InjectRule inject = new InjectRule();
-  @Rule public final MockitoRule mocks = MockitoJUnit.rule();
+  @RegisterExtension public final InjectRule inject = new InjectRule();
 
   @Mock private AppEngineServiceUtils appEngineServiceUtils;
   @Mock private Response response;
   private final FakeClock clock = new FakeClock(DateTime.parse("2016-02-11T10:00:00Z"));
   private AsyncTaskEnqueuer asyncTaskEnqueuer;
 
-  @Before
-  public void before() {
+  @BeforeEach
+  void beforeEach() {
     inject.setStaticField(Ofy.class, "clock", clock);
     when(appEngineServiceUtils.getServiceHostname("backend")).thenReturn("backend.hostname.fake");
     asyncTaskEnqueuer =
@@ -93,8 +92,9 @@ public class ResaveEntityActionTest {
     action.run();
   }
 
+  @MockitoSettings(strictness = Strictness.LENIENT)
   @Test
-  public void test_domainPendingTransfer_isResavedAndTransferCompleted() {
+  void test_domainPendingTransfer_isResavedAndTransferCompleted() {
     DomainBase domain =
         persistDomainWithPendingTransfer(
             persistDomainWithDependentResources(
@@ -116,7 +116,7 @@ public class ResaveEntityActionTest {
   }
 
   @Test
-  public void test_domainPendingDeletion_isResavedAndReenqueued() {
+  void test_domainPendingDeletion_isResavedAndReenqueued() {
     DomainBase domain =
         persistResource(
             newDomainBase("domain.tld")

core/src/test/java/google/registry/beam/BeamUtilsTest.java

@@ -22,14 +22,11 @@ import org.apache.avro.Schema;
 import org.apache.avro.generic.GenericData;
 import org.apache.avro.generic.GenericRecord;
 import org.apache.beam.sdk.io.gcp.bigquery.SchemaAndRecord;
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
 
 /** Unit tests for {@link BeamUtils} */
-@RunWith(JUnit4.class)
-public class BeamUtilsTest {
+class BeamUtilsTest {
 
   private static final String GENERIC_SCHEMA =
       "{\"name\": \"AnObject\", "
@@ -41,8 +38,8 @@ public class BeamUtilsTest {
 
   private SchemaAndRecord schemaAndRecord;
 
-  @Before
-  public void initializeRecord() {
+  @BeforeEach
+  void beforeEach() {
     // Create a record with a given JSON schema.
     GenericRecord record = new GenericData.Record(new Schema.Parser().parse(GENERIC_SCHEMA));
     record.put("aString", "hello world");
@@ -51,26 +48,26 @@ public class BeamUtilsTest {
   }
 
   @Test
-  public void testExtractField_fieldExists_returnsExpectedStringValues() {
+  void testExtractField_fieldExists_returnsExpectedStringValues() {
     assertThat(BeamUtils.extractField(schemaAndRecord.getRecord(), "aString"))
         .isEqualTo("hello world");
     assertThat(BeamUtils.extractField(schemaAndRecord.getRecord(), "aFloat")).isEqualTo("2.54");
   }
 
   @Test
-  public void testExtractField_fieldDoesntExist_returnsNull() {
+  void testExtractField_fieldDoesntExist_returnsNull() {
     schemaAndRecord.getRecord().put("aFloat", null);
     assertThat(BeamUtils.extractField(schemaAndRecord.getRecord(), "aFloat")).isEqualTo("null");
     assertThat(BeamUtils.extractField(schemaAndRecord.getRecord(), "missing")).isEqualTo("null");
   }
 
   @Test
-  public void testCheckFieldsNotNull_noExceptionIfAllPresent() {
+  void testCheckFieldsNotNull_noExceptionIfAllPresent() {
     BeamUtils.checkFieldsNotNull(ImmutableList.of("aString", "aFloat"), schemaAndRecord);
   }
 
   @Test
-  public void testCheckFieldsNotNull_fieldMissing_throwsException() {
+  void testCheckFieldsNotNull_fieldMissing_throwsException() {
     IllegalStateException expected =
         assertThrows(
             IllegalStateException.class,

core/src/test/java/google/registry/beam/TestPipelineExtension.java

@@ -0,0 +1,542 @@
+// Copyright 2020 The Nomulus Authors. All Rights Reserved.
+// This applies to our modifications; the base file's license header is:
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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;
+
+import static com.google.common.base.Preconditions.checkState;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.is;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.base.Strings;
+import com.google.common.collect.Maps;
+import java.io.IOException;
+import java.lang.annotation.Annotation;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.UUID;
+import java.util.function.Predicate;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.PipelineResult;
+import org.apache.beam.sdk.annotations.Internal;
+import org.apache.beam.sdk.io.FileSystems;
+import org.apache.beam.sdk.metrics.MetricNameFilter;
+import org.apache.beam.sdk.metrics.MetricResult;
+import org.apache.beam.sdk.metrics.MetricsEnvironment;
+import org.apache.beam.sdk.metrics.MetricsFilter;
+import org.apache.beam.sdk.options.ApplicationNameOptions;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.options.PipelineOptions.CheckEnabled;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.options.ValueProvider;
+import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
+import org.apache.beam.sdk.runners.TransformHierarchy;
+import org.apache.beam.sdk.runners.TransformHierarchy.Node;
+import org.apache.beam.sdk.testing.CrashingRunner;
+import org.apache.beam.sdk.testing.NeedsRunner;
+import org.apache.beam.sdk.testing.PAssert;
+import org.apache.beam.sdk.testing.TestPipelineOptions;
+import org.apache.beam.sdk.testing.ValidatesRunner;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.util.common.ReflectHelpers;
+import org.junit.experimental.categories.Category;
+import org.junit.jupiter.api.extension.AfterEachCallback;
+import org.junit.jupiter.api.extension.BeforeEachCallback;
+import org.junit.jupiter.api.extension.ExtensionContext;
+
+// NOTE: This file is copied from the Apache Beam distribution so that it can be locally modified to
+// support JUnit 5.
+
+/**
+ * A creator of test pipelines that can be used inside of tests that can be configured to run
+ * locally or against a remote pipeline runner.
+ *
+ * <p>In order to run tests on a pipeline runner, the following conditions must be met:
+ *
+ * <ul>
+ *   <li>System property "beamTestPipelineOptions" must contain a JSON delimited list of pipeline
+ *       options. For example:
+ *       <pre>{@code [
+ *     "--runner=TestDataflowRunner",
+ *     "--project=mygcpproject",
+ *     "--stagingLocation=gs://mygcsbucket/path"
+ * ]}</pre>
+ *       Note that the set of pipeline options required is pipeline runner specific.
+ *   <li>Jars containing the SDK and test classes must be available on the classpath.
+ * </ul>
+ *
+ * <p>Use {@link PAssert} for tests, as it integrates with this test harness in both direct and
+ * remote execution modes. For example:
+ *
+ * <pre><code>
+ * {@literal @Rule}
+ *  public final transient TestPipeline p = TestPipeline.create();
+ *
+ * {@literal @Test}
+ * {@literal @Category}(NeedsRunner.class)
+ *  public void myPipelineTest() throws Exception {
+ *    final PCollection&lt;String&gt; pCollection = pipeline.apply(...)
+ *    PAssert.that(pCollection).containsInAnyOrder(...);
+ *    pipeline.run();
+ *  }
+ * </code></pre>
+ *
+ * <p>For pipeline runners, it is required that they must throw an {@link AssertionError} containing
+ * the message from the {@link PAssert} that failed.
+ *
+ * <p>See also the <a href="https://beam.apache.org/contribute/testing/">Testing</a> documentation
+ * section.
+ */
+public class TestPipelineExtension extends Pipeline
+    implements BeforeEachCallback, AfterEachCallback {
+
+  private final PipelineOptions options;
+
+  private static class PipelineRunEnforcement {
+
+    @SuppressWarnings("WeakerAccess")
+    protected boolean enableAutoRunIfMissing;
+
+    protected final Pipeline pipeline;
+
+    protected boolean runAttempted;
+
+    private PipelineRunEnforcement(final Pipeline pipeline) {
+      this.pipeline = pipeline;
+    }
+
+    protected void enableAutoRunIfMissing(final boolean enable) {
+      enableAutoRunIfMissing = enable;
+    }
+
+    protected void beforePipelineExecution() {
+      runAttempted = true;
+    }
+
+    protected void afterPipelineExecution() {}
+
+    protected void afterUserCodeFinished() {
+      if (!runAttempted && enableAutoRunIfMissing) {
+        pipeline.run().waitUntilFinish();
+      }
+    }
+  }
+
+  private static class PipelineAbandonedNodeEnforcement extends PipelineRunEnforcement {
+
+    // Null until the pipeline has been run
+    @Nullable private List<TransformHierarchy.Node> runVisitedNodes;
+
+    private final Predicate<Node> isPAssertNode =
+        node ->
+            node.getTransform() instanceof PAssert.GroupThenAssert
+                || node.getTransform() instanceof PAssert.GroupThenAssertForSingleton
+                || node.getTransform() instanceof PAssert.OneSideInputAssert;
+
+    private static class NodeRecorder extends PipelineVisitor.Defaults {
+
+      private final List<TransformHierarchy.Node> visited = new ArrayList<>();
+
+      @Override
+      public void leaveCompositeTransform(final TransformHierarchy.Node node) {
+        visited.add(node);
+      }
+
+      @Override
+      public void visitPrimitiveTransform(final TransformHierarchy.Node node) {
+        visited.add(node);
+      }
+    }
+
+    private PipelineAbandonedNodeEnforcement(final TestPipelineExtension pipeline) {
+      super(pipeline);
+      runVisitedNodes = null;
+    }
+
+    private List<TransformHierarchy.Node> recordPipelineNodes(final Pipeline pipeline) {
+      final NodeRecorder nodeRecorder = new NodeRecorder();
+      pipeline.traverseTopologically(nodeRecorder);
+      return nodeRecorder.visited;
+    }
+
+    private boolean isEmptyPipeline(final Pipeline pipeline) {
+      final IsEmptyVisitor isEmptyVisitor = new IsEmptyVisitor();
+      pipeline.traverseTopologically(isEmptyVisitor);
+      return isEmptyVisitor.isEmpty();
+    }
+
+    private void verifyPipelineExecution() {
+      if (!isEmptyPipeline(pipeline)) {
+        if (!runAttempted && !enableAutoRunIfMissing) {
+          throw new PipelineRunMissingException("The pipeline has not been run.");
+
+        } else {
+          final List<TransformHierarchy.Node> pipelineNodes = recordPipelineNodes(pipeline);
+          if (pipelineRunSucceeded() && !visitedAll(pipelineNodes)) {
+            final boolean hasDanglingPAssert =
+                pipelineNodes.stream()
+                    .filter(pn -> !runVisitedNodes.contains(pn))
+                    .anyMatch(isPAssertNode);
+            if (hasDanglingPAssert) {
+              throw new AbandonedNodeException("The pipeline contains abandoned PAssert(s).");
+            } else {
+              throw new AbandonedNodeException("The pipeline contains abandoned PTransform(s).");
+            }
+          }
+        }
+      }
+    }
+
+    private boolean visitedAll(final List<TransformHierarchy.Node> pipelineNodes) {
+      return runVisitedNodes.equals(pipelineNodes);
+    }
+
+    private boolean pipelineRunSucceeded() {
+      return runVisitedNodes != null;
+    }
+
+    @Override
+    protected void afterPipelineExecution() {
+      runVisitedNodes = recordPipelineNodes(pipeline);
+      super.afterPipelineExecution();
+    }
+
+    @Override
+    protected void afterUserCodeFinished() {
+      super.afterUserCodeFinished();
+      verifyPipelineExecution();
+    }
+  }
+
+  /**
+   * An exception thrown in case an abandoned {@link org.apache.beam.sdk.transforms.PTransform} is
+   * detected, that is, a {@link org.apache.beam.sdk.transforms.PTransform} that has not been run.
+   */
+  public static class AbandonedNodeException extends RuntimeException {
+
+    AbandonedNodeException(final String msg) {
+      super(msg);
+    }
+  }
+
+  /** An exception thrown in case a test finishes without invoking {@link Pipeline#run()}. */
+  public static class PipelineRunMissingException extends RuntimeException {
+
+    PipelineRunMissingException(final String msg) {
+      super(msg);
+    }
+  }
+
+  /** System property used to set {@link TestPipelineOptions}. */
+  public static final String PROPERTY_BEAM_TEST_PIPELINE_OPTIONS = "beamTestPipelineOptions";
+
+  static final String PROPERTY_USE_DEFAULT_DUMMY_RUNNER = "beamUseDummyRunner";
+
+  private static final ObjectMapper MAPPER =
+      new ObjectMapper()
+          .registerModules(ObjectMapper.findModules(ReflectHelpers.findClassLoader()));
+
+  @SuppressWarnings("OptionalUsedAsFieldOrParameterType")
+  private Optional<? extends PipelineRunEnforcement> enforcement = Optional.empty();
+
+  /**
+   * Creates and returns a new test pipeline.
+   *
+   * <p>Use {@link PAssert} to add tests, then call {@link Pipeline#run} to execute the pipeline and
+   * check the tests.
+   */
+  public static TestPipelineExtension create() {
+    return fromOptions(testingPipelineOptions());
+  }
+
+  public static TestPipelineExtension fromOptions(PipelineOptions options) {
+    return new TestPipelineExtension(options);
+  }
+
+  private TestPipelineExtension(final PipelineOptions options) {
+    super(options);
+    this.options = options;
+  }
+
+  @Override
+  public PipelineOptions getOptions() {
+    return this.options;
+  }
+
+  @Override
+  public void beforeEach(ExtensionContext context) throws Exception {
+    options.as(ApplicationNameOptions.class).setAppName(getAppName(context));
+
+    // if the enforcement level has not been set by the user do auto-inference
+    if (!enforcement.isPresent()) {
+      final boolean isCrashingRunner = CrashingRunner.class.isAssignableFrom(options.getRunner());
+
+      checkState(
+          !isCrashingRunner,
+          "Cannot test using a [%s] runner. Please re-check your configuration.",
+          CrashingRunner.class.getSimpleName());
+
+      enableAbandonedNodeEnforcement(true);
+    }
+  }
+
+  @Override
+  public void afterEach(ExtensionContext context) throws Exception {
+    enforcement.get().afterUserCodeFinished();
+  }
+
+  /** Returns the class + method name of the test. */
+  private String getAppName(ExtensionContext context) {
+    String methodName = context.getRequiredTestMethod().getName();
+    Class<?> testClass = context.getRequiredTestClass();
+    if (testClass.isMemberClass()) {
+      return String.format(
+          "%s$%s-%s",
+          testClass.getEnclosingClass().getSimpleName(), testClass.getSimpleName(), methodName);
+    } else {
+      return String.format("%s-%s", testClass.getSimpleName(), methodName);
+    }
+  }
+
+  /**
+   * Runs this {@link TestPipelineExtension}, unwrapping any {@code AssertionError} that is raised
+   * during testing.
+   */
+  @Override
+  public PipelineResult run() {
+    return run(getOptions());
+  }
+
+  /** Like {@link #run} but with the given potentially modified options. */
+  @Override
+  public PipelineResult run(PipelineOptions options) {
+    checkState(
+        enforcement.isPresent(),
+        "Is your TestPipeline declaration missing a @Rule annotation? Usage: "
+            + "@Rule public final transient TestPipeline pipeline = TestPipeline.create();");
+
+    final PipelineResult pipelineResult;
+    try {
+      enforcement.get().beforePipelineExecution();
+      PipelineOptions updatedOptions =
+          MAPPER.convertValue(MAPPER.valueToTree(options), PipelineOptions.class);
+      updatedOptions
+          .as(TestValueProviderOptions.class)
+          .setProviderRuntimeValues(StaticValueProvider.of(providerRuntimeValues));
+      pipelineResult = super.run(updatedOptions);
+      verifyPAssertsSucceeded(this, pipelineResult);
+    } catch (RuntimeException exc) {
+      Throwable cause = exc.getCause();
+      if (cause instanceof AssertionError) {
+        throw (AssertionError) cause;
+      } else {
+        throw exc;
+      }
+    }
+
+    // If we reach this point, the pipeline has been run and no exceptions have been thrown during
+    // its execution.
+    enforcement.get().afterPipelineExecution();
+    return pipelineResult;
+  }
+
+  /** Implementation detail of {@link #newProvider}, do not use. */
+  @Internal
+  public interface TestValueProviderOptions extends PipelineOptions {
+    ValueProvider<Map<String, Object>> getProviderRuntimeValues();
+
+    void setProviderRuntimeValues(ValueProvider<Map<String, Object>> runtimeValues);
+  }
+
+  /**
+   * Returns a new {@link ValueProvider} that is inaccessible before {@link #run}, but will be
+   * accessible while the pipeline runs.
+   */
+  public <T> ValueProvider<T> newProvider(T runtimeValue) {
+    String uuid = UUID.randomUUID().toString();
+    providerRuntimeValues.put(uuid, runtimeValue);
+    return ValueProvider.NestedValueProvider.of(
+        options.as(TestValueProviderOptions.class).getProviderRuntimeValues(),
+        new GetFromRuntimeValues<T>(uuid));
+  }
+
+  private final Map<String, Object> providerRuntimeValues = Maps.newHashMap();
+
+  private static class GetFromRuntimeValues<T>
+      implements SerializableFunction<Map<String, Object>, T> {
+    private final String key;
+
+    private GetFromRuntimeValues(String key) {
+      this.key = key;
+    }
+
+    @Override
+    public T apply(Map<String, Object> input) {
+      return (T) input.get(key);
+    }
+  }
+
+  /**
+   * Enables the abandoned node detection. Abandoned nodes are <code>PTransforms</code>, <code>
+   * PAsserts</code> included, that were not executed by the pipeline runner. Abandoned nodes are
+   * most likely to occur due to the one of the following scenarios:
+   *
+   * <ul>
+   *   <li>Lack of a <code>pipeline.run()</code> statement at the end of a test.
+   *   <li>Addition of PTransforms after the pipeline has already run.
+   * </ul>
+   *
+   * Abandoned node detection is automatically enabled when a real pipeline runner (i.e. not a
+   * {@link CrashingRunner}) and/or a {@link NeedsRunner} or a {@link ValidatesRunner} annotation
+   * are detected.
+   */
+  public TestPipelineExtension enableAbandonedNodeEnforcement(final boolean enable) {
+    enforcement =
+        enable
+            ? Optional.of(new PipelineAbandonedNodeEnforcement(this))
+            : Optional.of(new PipelineRunEnforcement(this));
+
+    return this;
+  }
+
+  /**
+   * If enabled, a <code>pipeline.run()</code> statement will be added automatically in case it is
+   * missing in the test.
+   */
+  public TestPipelineExtension enableAutoRunIfMissing(final boolean enable) {
+    enforcement.get().enableAutoRunIfMissing(enable);
+    return this;
+  }
+
+  @Override
+  public String toString() {
+    return "TestPipeline#" + options.as(ApplicationNameOptions.class).getAppName();
+  }
+
+  /** Creates {@link PipelineOptions} for testing. */
+  public static PipelineOptions testingPipelineOptions() {
+    try {
+      @Nullable
+      String beamTestPipelineOptions = System.getProperty(PROPERTY_BEAM_TEST_PIPELINE_OPTIONS);
+
+      PipelineOptions options =
+          Strings.isNullOrEmpty(beamTestPipelineOptions)
+              ? PipelineOptionsFactory.create()
+              : PipelineOptionsFactory.fromArgs(
+                      MAPPER.readValue(beamTestPipelineOptions, String[].class))
+                  .as(TestPipelineOptions.class);
+
+      // If no options were specified, set some reasonable defaults
+      if (Strings.isNullOrEmpty(beamTestPipelineOptions)) {
+        // If there are no provided options, check to see if a dummy runner should be used.
+        String useDefaultDummy = System.getProperty(PROPERTY_USE_DEFAULT_DUMMY_RUNNER);
+        if (!Strings.isNullOrEmpty(useDefaultDummy) && Boolean.valueOf(useDefaultDummy)) {
+          options.setRunner(CrashingRunner.class);
+        }
+      }
+      options.setStableUniqueNames(CheckEnabled.ERROR);
+
+      FileSystems.setDefaultPipelineOptions(options);
+      return options;
+    } catch (IOException e) {
+      throw new RuntimeException(
+          "Unable to instantiate test options from system property "
+              + PROPERTY_BEAM_TEST_PIPELINE_OPTIONS
+              + ":"
+              + System.getProperty(PROPERTY_BEAM_TEST_PIPELINE_OPTIONS),
+          e);
+    }
+  }
+
+  /**
+   * Verifies all {{@link PAssert PAsserts}} in the pipeline have been executed and were successful.
+   *
+   * <p>Note this only runs for runners which support Metrics. Runners which do not should verify
+   * this in some other way. See: https://issues.apache.org/jira/browse/BEAM-2001
+   */
+  public static void verifyPAssertsSucceeded(Pipeline pipeline, PipelineResult pipelineResult) {
+    if (MetricsEnvironment.isMetricsSupported()) {
+      long expectedNumberOfAssertions = (long) PAssert.countAsserts(pipeline);
+
+      long successfulAssertions = 0;
+      Iterable<MetricResult<Long>> successCounterResults =
+          pipelineResult
+              .metrics()
+              .queryMetrics(
+                  MetricsFilter.builder()
+                      .addNameFilter(MetricNameFilter.named(PAssert.class, PAssert.SUCCESS_COUNTER))
+                      .build())
+              .getCounters();
+      for (MetricResult<Long> counter : successCounterResults) {
+        if (counter.getAttempted() > 0) {
+          successfulAssertions++;
+        }
+      }
+
+      assertThat(
+          String.format(
+              "Expected %d successful assertions, but found %d.",
+              expectedNumberOfAssertions, successfulAssertions),
+          successfulAssertions,
+          is(expectedNumberOfAssertions));
+    }
+  }
+
+  private static class IsEmptyVisitor extends PipelineVisitor.Defaults {
+    private boolean empty = true;
+
+    public boolean isEmpty() {
+      return empty;
+    }
+
+    @Override
+    public void visitPrimitiveTransform(TransformHierarchy.Node node) {
+      empty = false;
+    }
+  }
+
+  /**
+   * A utility class for querying annotations.
+   *
+   * <p>NOTE: This was copied from the Apache Beam project from a separate file only for visibility
+   * reasons (it's package-private there).
+   */
+  static class Annotations {
+
+    /** Annotation predicates. */
+    static class Predicates {
+
+      static Predicate<Annotation> isAnnotationOfType(final Class<? extends Annotation> clazz) {
+        return annotation ->
+            annotation.annotationType() != null && annotation.annotationType().equals(clazz);
+      }
+
+      static Predicate<Annotation> isCategoryOf(final Class<?> value, final boolean allowDerived) {
+        return category ->
+            Arrays.stream(((Category) category).value())
+                .anyMatch(
+                    aClass -> allowDerived ? value.isAssignableFrom(aClass) : value.equals(aClass));
+      }
+    }
+  }
+}

core/src/test/java/google/registry/beam/initsql/BeamJpaExtension.java

@@ -0,0 +1,72 @@
+// 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 java.io.File;
+import java.io.IOException;
+import java.io.PrintStream;
+import java.io.Serializable;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.function.Supplier;
+import org.junit.jupiter.api.extension.AfterEachCallback;
+import org.junit.jupiter.api.extension.BeforeEachCallback;
+import org.junit.jupiter.api.extension.ExtensionContext;
+import org.testcontainers.containers.JdbcDatabaseContainer;
+
+/**
+ * Helpers for setting up {@link BeamJpaModule} in tests.
+ *
+ * <p>This extension is often used with a Database container and/or temporary file folder. User must
+ * make sure that all dependent extensions are set up before this extension, e.g., by assigning
+ * {@link org.junit.jupiter.api.Order orders}.
+ */
+public final class BeamJpaExtension implements BeforeEachCallback, AfterEachCallback, Serializable {
+
+  private final transient JdbcDatabaseContainer<?> database;
+  private final transient Supplier<Path> credentialPathSupplier;
+  private transient BeamJpaModule beamJpaModule;
+
+  private File credentialFile;
+
+  public BeamJpaExtension(Supplier<Path> credentialPathSupplier, JdbcDatabaseContainer database) {
+    this.database = database;
+    this.credentialPathSupplier = credentialPathSupplier;
+  }
+
+  public File getCredentialFile() {
+    return credentialFile;
+  }
+
+  public BeamJpaModule getBeamJpaModule() {
+    if (beamJpaModule != null) {
+      return beamJpaModule;
+    }
+    return beamJpaModule = new BeamJpaModule(credentialFile.getAbsolutePath(), null);
+  }
+
+  @Override
+  public void beforeEach(ExtensionContext context) throws IOException {
+    credentialFile = Files.createFile(credentialPathSupplier.get()).toFile();
+    new PrintStream(credentialFile)
+        .printf("%s %s %s", database.getJdbcUrl(), database.getUsername(), database.getPassword())
+        .close();
+  }
+
+  @Override
+  public void afterEach(ExtensionContext context) {
+    credentialFile.delete();
+  }
+}

core/src/test/java/google/registry/beam/initsql/BeamJpaModuleTest.java

@@ -18,14 +18,14 @@ import static com.google.common.truth.Truth.assertThat;
 
 import google.registry.persistence.NomulusPostgreSql;
 import google.registry.persistence.transaction.JpaTransactionManager;
-import java.io.File;
-import java.io.IOException;
-import java.io.PrintStream;
+import google.registry.testing.DatastoreEntityExtension;
+import java.nio.file.Path;
 import org.apache.beam.sdk.io.FileSystems;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
-import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Order;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.condition.EnabledIfSystemProperty;
+import org.junit.jupiter.api.extension.RegisterExtension;
 import org.junit.jupiter.api.io.TempDir;
 import org.testcontainers.containers.PostgreSQLContainer;
 import org.testcontainers.junit.jupiter.Container;
@@ -33,28 +33,28 @@ import org.testcontainers.junit.jupiter.Testcontainers;
 
 /** Unit tests for {@link BeamJpaModule}. */
 @Testcontainers
-public class BeamJpaModuleTest {
+class BeamJpaModuleTest {
+
+  @RegisterExtension
+  final DatastoreEntityExtension datastoreEntityExtension = new DatastoreEntityExtension();
 
   @Container
-  public PostgreSQLContainer database = new PostgreSQLContainer(NomulusPostgreSql.getDockerTag());
+  final PostgreSQLContainer database = new PostgreSQLContainer(NomulusPostgreSql.getDockerTag());
 
-  @TempDir File tempFolder;
+  @SuppressWarnings("WeakerAccess")
+  @TempDir
+  Path tmpDir;
 
-  private File credentialFile;
-
-  @BeforeEach
-  public void beforeEach() throws IOException {
-    credentialFile = new File(tempFolder, "credential");
-    new PrintStream(credentialFile)
-        .printf("%s %s %s", database.getJdbcUrl(), database.getUsername(), database.getPassword())
-        .close();
-  }
+  @RegisterExtension
+  @Order(Order.DEFAULT + 1)
+  final BeamJpaExtension beamJpaExtension =
+      new BeamJpaExtension(() -> tmpDir.resolve("credential.dat"), database);
 
   @Test
-  public void getJpaTransactionManager_local() {
+  void getJpaTransactionManager_local() {
     JpaTransactionManager jpa =
         DaggerBeamJpaModule_JpaTransactionManagerComponent.builder()
-            .beamJpaModule(new BeamJpaModule(credentialFile.getAbsolutePath()))
+            .beamJpaModule(beamJpaExtension.getBeamJpaModule())
             .build()
             .localDbJpaTransactionManager();
     assertThat(
@@ -75,14 +75,15 @@ public class BeamJpaModuleTest {
    */
   @Test
   @EnabledIfSystemProperty(named = "test.gcp_integration.env", matches = "\\S+")
-  public void getJpaTransactionManager_cloudSql_authRequired() {
+  void getJpaTransactionManager_cloudSql_authRequired() {
     String environmentName = System.getProperty("test.gcp_integration.env");
     FileSystems.setDefaultPipelineOptions(PipelineOptionsFactory.create());
     JpaTransactionManager jpa =
         DaggerBeamJpaModule_JpaTransactionManagerComponent.builder()
             .beamJpaModule(
                 new BeamJpaModule(
-                    BackupPaths.getCloudSQLCredentialFilePatterns(environmentName).get(0)))
+                    BackupPaths.getCloudSQLCredentialFilePatterns(environmentName).get(0),
+                    String.format("domain-registry-%s", environmentName)))
             .build()
             .cloudSqlJpaTransactionManager();
     assertThat(

core/src/test/java/google/registry/beam/initsql/CommitLogTransformsTest.java

@@ -21,6 +21,7 @@ import static google.registry.testing.DatastoreHelper.newRegistry;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import google.registry.backup.VersionedEntity;
+import google.registry.beam.TestPipelineExtension;
 import google.registry.model.contact.ContactResource;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.ofy.Ofy;
@@ -30,40 +31,37 @@ import google.registry.testing.InjectRule;
 import java.io.File;
 import java.io.IOException;
 import java.io.Serializable;
+import java.nio.file.Files;
+import java.nio.file.Path;
 import org.apache.beam.sdk.coders.StringUtf8Coder;
 import org.apache.beam.sdk.io.fs.MatchResult.Metadata;
-import org.apache.beam.sdk.testing.NeedsRunner;
 import org.apache.beam.sdk.testing.PAssert;
-import org.apache.beam.sdk.testing.TestPipeline;
 import org.apache.beam.sdk.transforms.Create;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.ParDo;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.PCollection;
 import org.joda.time.DateTime;
-import org.junit.After;
-import org.junit.Before;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.experimental.categories.Category;
-import org.junit.rules.TemporaryFolder;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.api.io.TempDir;
 
 /** Unit tests for {@link Transforms} related to loading CommitLogs. */
-// TODO(weiminyu): Upgrade to JUnit5 when TestPipeline is upgraded. It is also easy to adapt with
-// a wrapper.
-@RunWith(JUnit4.class)
-public class CommitLogTransformsTest implements Serializable {
+class CommitLogTransformsTest implements Serializable {
+
   private static final DateTime START_TIME = DateTime.parse("2000-01-01T00:00:00.0Z");
 
-  @Rule public final transient TemporaryFolder temporaryFolder = new TemporaryFolder();
+  @SuppressWarnings("WeakerAccess")
+  @TempDir
+  transient Path tmpDir;
 
-  @Rule public final transient InjectRule injectRule = new InjectRule();
+  @RegisterExtension final transient InjectRule injectRule = new InjectRule();
 
-  @Rule
-  public final transient TestPipeline pipeline =
-      TestPipeline.create().enableAbandonedNodeEnforcement(true);
+  @RegisterExtension
+  final transient TestPipelineExtension testPipeline =
+      TestPipelineExtension.create().enableAbandonedNodeEnforcement(true);
 
   private FakeClock fakeClock;
   private transient BackupTestStore store;
@@ -75,8 +73,8 @@ public class CommitLogTransformsTest implements Serializable {
   private transient ContactResource contact;
   private transient DomainBase domain;
 
-  @Before
-  public void beforeEach() throws Exception {
+  @BeforeEach
+  void beforeEach() throws Exception {
     fakeClock = new FakeClock(START_TIME);
     store = new BackupTestStore(fakeClock);
     injectRule.setStaticField(Ofy.class, "clock", fakeClock);
@@ -92,12 +90,12 @@ public class CommitLogTransformsTest implements Serializable {
     contact = (ContactResource) store.loadAsOfyEntity(contact);
     domain = (DomainBase) store.loadAsOfyEntity(domain);
 
-    commitLogsDir = temporaryFolder.newFolder();
+    commitLogsDir = Files.createDirectory(tmpDir.resolve("commit_logs")).toFile();
     firstCommitLogFile = store.saveCommitLogs(commitLogsDir.getAbsolutePath());
   }
 
-  @After
-  public void afterEach() throws Exception {
+  @AfterEach
+  void afterEach() throws Exception {
     if (store != null) {
       store.close();
       store = null;
@@ -105,10 +103,9 @@ public class CommitLogTransformsTest implements Serializable {
   }
 
   @Test
-  @Category(NeedsRunner.class)
-  public void getCommitLogFilePatterns() {
+  void getCommitLogFilePatterns() {
     PCollection<String> patterns =
-        pipeline.apply(
+        testPipeline.apply(
             "Get CommitLog file patterns",
             Transforms.getCommitLogFilePatterns(commitLogsDir.getAbsolutePath()));
 
@@ -117,14 +114,13 @@ public class CommitLogTransformsTest implements Serializable {
 
     PAssert.that(patterns).containsInAnyOrder(expectedPatterns);
 
-    pipeline.run();
+    testPipeline.run();
   }
 
   @Test
-  @Category(NeedsRunner.class)
-  public void getFilesByPatterns() {
+  void getFilesByPatterns() {
     PCollection<Metadata> fileMetas =
-        pipeline
+        testPipeline
             .apply(
                 "File patterns to metadata",
                 Create.of(commitLogsDir.getAbsolutePath() + "/commit_diff_until_*")
@@ -149,12 +145,11 @@ public class CommitLogTransformsTest implements Serializable {
 
     PAssert.that(fileNames).containsInAnyOrder(expectedFilenames);
 
-    pipeline.run();
+    testPipeline.run();
   }
 
   @Test
-  @Category(NeedsRunner.class)
-  public void filterCommitLogsByTime() throws IOException {
+  void filterCommitLogsByTime() throws IOException {
     ImmutableList<String> commitLogFilenames =
         ImmutableList.of(
             "commit_diff_until_2000-01-01T00:00:00.000Z",
@@ -163,16 +158,15 @@ public class CommitLogTransformsTest implements Serializable {
             "commit_diff_until_2000-01-01T00:00:00.003Z",
             "commit_diff_until_2000-01-01T00:00:00.004Z");
 
-    File commitLogDir = temporaryFolder.newFolder();
     for (String name : commitLogFilenames) {
-      new File(commitLogDir, name).createNewFile();
+      new File(commitLogsDir, name).createNewFile();
     }
 
     PCollection<String> filteredFilenames =
-        pipeline
+        testPipeline
             .apply(
                 "Get commitlog file patterns",
-                Transforms.getCommitLogFilePatterns(commitLogDir.getAbsolutePath()))
+                Transforms.getCommitLogFilePatterns(commitLogsDir.getAbsolutePath()))
             .apply("Find commitlog files", Transforms.getFilesByPatterns())
             .apply(
                 "Filtered by Time",
@@ -194,14 +188,13 @@ public class CommitLogTransformsTest implements Serializable {
             "commit_diff_until_2000-01-01T00:00:00.001Z",
             "commit_diff_until_2000-01-01T00:00:00.002Z");
 
-    pipeline.run();
+    testPipeline.run();
   }
 
   @Test
-  @Category(NeedsRunner.class)
-  public void loadOneCommitLogFile() {
+  void loadOneCommitLogFile() {
     PCollection<VersionedEntity> entities =
-        pipeline
+        testPipeline
             .apply(
                 "Get CommitLog file patterns",
                 Transforms.getCommitLogFilePatterns(commitLogsDir.getAbsolutePath()))
@@ -216,14 +209,13 @@ public class CommitLogTransformsTest implements Serializable {
         KV.of(fakeClock.nowUtc().getMillis() - 1, store.loadAsDatastoreEntity(contact)),
         KV.of(fakeClock.nowUtc().getMillis() - 1, store.loadAsDatastoreEntity(domain)));
 
-    pipeline.run();
+    testPipeline.run();
   }
 
   @Test
-  @Category(NeedsRunner.class)
-  public void loadOneCommitLogFile_filterByKind() {
+  void loadOneCommitLogFile_filterByKind() {
     PCollection<VersionedEntity> entities =
-        pipeline
+        testPipeline
             .apply(
                 "Get CommitLog file patterns",
                 Transforms.getCommitLogFilePatterns(commitLogsDir.getAbsolutePath()))
@@ -236,6 +228,6 @@ public class CommitLogTransformsTest implements Serializable {
         KV.of(fakeClock.nowUtc().getMillis() - 2, store.loadAsDatastoreEntity(registry)),
         KV.of(fakeClock.nowUtc().getMillis() - 1, store.loadAsDatastoreEntity(contact)));
 
-    pipeline.run();
+    testPipeline.run();
   }
 }

core/src/test/java/google/registry/beam/initsql/DomainBaseUtilTest.java

@@ -0,0 +1,221 @@
+// 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 google.registry.model.ImmutableObjectSubject.assertAboutImmutableObjects;
+import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.testing.DatastoreHelper.cloneAndSetAutoTimestamps;
+import static google.registry.testing.DatastoreHelper.createTld;
+import static google.registry.testing.DatastoreHelper.persistResource;
+import static google.registry.util.DateTimeUtils.START_OF_TIME;
+import static org.junit.Assert.assertThrows;
+
+import com.google.appengine.api.datastore.Entity;
+import com.google.common.collect.ImmutableSet;
+import com.googlecode.objectify.Key;
+import google.registry.model.billing.BillingEvent;
+import google.registry.model.billing.BillingEvent.OneTime;
+import google.registry.model.contact.ContactResource;
+import google.registry.model.domain.DesignatedContact;
+import google.registry.model.domain.DomainAuthInfo;
+import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.GracePeriod;
+import google.registry.model.domain.launch.LaunchNotice;
+import google.registry.model.domain.rgp.GracePeriodStatus;
+import google.registry.model.domain.secdns.DelegationSignerData;
+import google.registry.model.eppcommon.AuthInfo.PasswordAuth;
+import google.registry.model.eppcommon.StatusValue;
+import google.registry.model.eppcommon.Trid;
+import google.registry.model.host.HostResource;
+import google.registry.model.ofy.Ofy;
+import google.registry.model.poll.PollMessage;
+import google.registry.model.reporting.HistoryEntry;
+import google.registry.model.transfer.DomainTransferData;
+import google.registry.model.transfer.TransferStatus;
+import google.registry.persistence.VKey;
+import google.registry.testing.AppEngineRule;
+import google.registry.testing.DatastoreHelper;
+import google.registry.testing.FakeClock;
+import google.registry.testing.InjectRule;
+import org.joda.time.Instant;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+/** Unit tests for {@link DomainBaseUtil}. */
+public class DomainBaseUtilTest {
+
+  private final FakeClock fakeClock = new FakeClock(Instant.now());
+
+  private DomainBase domain;
+  private Entity domainEntity;
+  private Key<OneTime> oneTimeBillKey;
+  private VKey<BillingEvent.Recurring> recurringBillKey;
+  private Key<DomainBase> domainKey;
+
+  @RegisterExtension
+  AppEngineRule appEngineRule =
+      AppEngineRule.builder().withDatastore().withClock(fakeClock).build();
+
+  @RegisterExtension InjectRule injectRule = new InjectRule();
+
+  @BeforeEach
+  void beforeEach() {
+    injectRule.setStaticField(Ofy.class, "clock", fakeClock);
+    createTld("com");
+    domainKey = Key.create(null, DomainBase.class, "4-COM");
+    VKey<HostResource> hostKey =
+        persistResource(
+                new HostResource.Builder()
+                    .setHostName("ns1.example.com")
+                    .setSuperordinateDomain(VKey.from(domainKey))
+                    .setRepoId("1-COM")
+                    .build())
+            .createVKey();
+    VKey<ContactResource> contact1Key =
+        persistResource(
+                new ContactResource.Builder()
+                    .setContactId("contact_id1")
+                    .setRepoId("2-COM")
+                    .build())
+            .createVKey();
+    VKey<ContactResource> contact2Key =
+        persistResource(
+                new ContactResource.Builder()
+                    .setContactId("contact_id2")
+                    .setRepoId("3-COM")
+                    .build())
+            .createVKey();
+    Key<HistoryEntry> historyEntryKey =
+        Key.create(persistResource(new HistoryEntry.Builder().setParent(domainKey).build()));
+    oneTimeBillKey = Key.create(historyEntryKey, BillingEvent.OneTime.class, 1);
+    recurringBillKey = VKey.from(Key.create(historyEntryKey, BillingEvent.Recurring.class, 2));
+    VKey<PollMessage.Autorenew> autorenewPollKey =
+        VKey.from(Key.create(historyEntryKey, PollMessage.Autorenew.class, 3));
+    VKey<PollMessage.OneTime> onetimePollKey =
+        VKey.from(Key.create(historyEntryKey, PollMessage.OneTime.class, 1));
+    // Set up a new persisted domain entity.
+    domain =
+        persistResource(
+            cloneAndSetAutoTimestamps(
+                new DomainBase.Builder()
+                    .setDomainName("example.com")
+                    .setRepoId("4-COM")
+                    .setCreationClientId("a registrar")
+                    .setLastEppUpdateTime(fakeClock.nowUtc())
+                    .setLastEppUpdateClientId("AnotherRegistrar")
+                    .setLastTransferTime(fakeClock.nowUtc())
+                    .setStatusValues(
+                        ImmutableSet.of(
+                            StatusValue.CLIENT_DELETE_PROHIBITED,
+                            StatusValue.SERVER_DELETE_PROHIBITED,
+                            StatusValue.SERVER_TRANSFER_PROHIBITED,
+                            StatusValue.SERVER_UPDATE_PROHIBITED,
+                            StatusValue.SERVER_RENEW_PROHIBITED,
+                            StatusValue.SERVER_HOLD))
+                    .setRegistrant(contact1Key)
+                    .setContacts(
+                        ImmutableSet.of(
+                            DesignatedContact.create(DesignatedContact.Type.ADMIN, contact2Key)))
+                    .setNameservers(ImmutableSet.of(hostKey))
+                    .setSubordinateHosts(ImmutableSet.of("ns1.example.com"))
+                    .setPersistedCurrentSponsorClientId("losing")
+                    .setRegistrationExpirationTime(fakeClock.nowUtc().plusYears(1))
+                    .setAuthInfo(DomainAuthInfo.create(PasswordAuth.create("password")))
+                    .setDsData(
+                        ImmutableSet.of(DelegationSignerData.create(1, 2, 3, new byte[] {0, 1, 2})))
+                    .setLaunchNotice(
+                        LaunchNotice.create("tcnid", "validatorId", START_OF_TIME, START_OF_TIME))
+                    .setTransferData(
+                        new DomainTransferData.Builder()
+                            .setGainingClientId("gaining")
+                            .setLosingClientId("losing")
+                            .setPendingTransferExpirationTime(fakeClock.nowUtc())
+                            .setServerApproveEntities(
+                                ImmutableSet.of(
+                                    VKey.from(oneTimeBillKey), recurringBillKey, autorenewPollKey))
+                            .setServerApproveBillingEvent(VKey.from(oneTimeBillKey))
+                            .setServerApproveAutorenewEvent(recurringBillKey)
+                            .setServerApproveAutorenewPollMessage(autorenewPollKey)
+                            .setTransferRequestTime(fakeClock.nowUtc().plusDays(1))
+                            .setTransferStatus(TransferStatus.SERVER_APPROVED)
+                            .setTransferRequestTrid(Trid.create("client-trid", "server-trid"))
+                            .build())
+                    .setDeletePollMessage(onetimePollKey)
+                    .setAutorenewBillingEvent(recurringBillKey)
+                    .setAutorenewPollMessage(autorenewPollKey)
+                    .setSmdId("smdid")
+                    .addGracePeriod(
+                        GracePeriod.create(
+                            GracePeriodStatus.ADD,
+                            fakeClock.nowUtc().plusDays(1),
+                            "registrar",
+                            null))
+                    .build()));
+    domainEntity = tm().transact(() -> ofy().toEntity(domain));
+  }
+
+  @Test
+  void removeBillingAndPollAndHosts_allFkeysPresent() {
+    DomainBase domainTransformedByOfy =
+        domain
+            .asBuilder()
+            .setAutorenewBillingEvent(null)
+            .setAutorenewPollMessage(null)
+            .setNameservers(ImmutableSet.of())
+            .setDeletePollMessage(null)
+            .setTransferData(null)
+            .build();
+    DomainBase domainTransformedByUtil =
+        (DomainBase) ofy().toPojo(DomainBaseUtil.removeBillingAndPollAndHosts(domainEntity));
+    // Compensates for the missing INACTIVE status.
+    domainTransformedByUtil = domainTransformedByUtil.asBuilder().build();
+    assertAboutImmutableObjects()
+        .that(domainTransformedByUtil)
+        .isEqualExceptFields(domainTransformedByOfy, "revisions");
+  }
+
+  @Test
+  void removeBillingAndPollAndHosts_noFkeysPresent() {
+    DomainBase domainWithoutFKeys =
+        domain
+            .asBuilder()
+            .setAutorenewBillingEvent(null)
+            .setAutorenewPollMessage(null)
+            .setNameservers(ImmutableSet.of())
+            .setDeletePollMessage(null)
+            .setTransferData(null)
+            .build();
+    Entity entityWithoutFkeys = tm().transact(() -> ofy().toEntity(domainWithoutFKeys));
+    DomainBase domainTransformedByUtil =
+        (DomainBase) ofy().toPojo(DomainBaseUtil.removeBillingAndPollAndHosts(entityWithoutFkeys));
+    // Compensates for the missing INACTIVE status.
+    domainTransformedByUtil = domainTransformedByUtil.asBuilder().build();
+    assertAboutImmutableObjects()
+        .that(domainTransformedByUtil)
+        .isEqualExceptFields(domainWithoutFKeys, "revisions");
+  }
+
+  @Test
+  void removeBillingAndPollAndHosts_notDomainBase() {
+    Entity contactEntity =
+        tm().transact(() -> ofy().toEntity(DatastoreHelper.newContactResource("contact")));
+
+    assertThrows(
+        IllegalArgumentException.class,
+        () -> DomainBaseUtil.removeBillingAndPollAndHosts(contactEntity));
+  }
+}

core/src/test/java/google/registry/beam/initsql/ExportloadingTransformsTest.java

@@ -22,6 +22,7 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
 import google.registry.backup.VersionedEntity;
+import google.registry.beam.TestPipelineExtension;
 import google.registry.model.contact.ContactResource;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.ofy.Ofy;
@@ -30,36 +31,30 @@ import google.registry.testing.FakeClock;
 import google.registry.testing.InjectRule;
 import java.io.File;
 import java.io.Serializable;
+import java.nio.file.Path;
 import java.util.Collections;
 import org.apache.beam.sdk.coders.StringUtf8Coder;
 import org.apache.beam.sdk.io.fs.MatchResult.Metadata;
-import org.apache.beam.sdk.testing.NeedsRunner;
 import org.apache.beam.sdk.testing.PAssert;
-import org.apache.beam.sdk.testing.TestPipeline;
 import org.apache.beam.sdk.transforms.Create;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.ParDo;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.PCollection;
 import org.joda.time.DateTime;
-import org.junit.After;
-import org.junit.Before;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.experimental.categories.Category;
-import org.junit.rules.TemporaryFolder;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.api.io.TempDir;
 
 /**
  * Unit tests for {@link Transforms} related to loading Datastore exports.
  *
  * <p>This class implements {@link Serializable} so that test {@link DoFn} classes may be inlined.
  */
-// TODO(weiminyu): Upgrade to JUnit5 when TestPipeline is upgraded. It is also easy to adapt with
-// a wrapper.
-@RunWith(JUnit4.class)
-public class ExportloadingTransformsTest implements Serializable {
+class ExportloadingTransformsTest implements Serializable {
+
   private static final DateTime START_TIME = DateTime.parse("2000-01-01T00:00:00.0Z");
 
   private static final ImmutableList<Class<?>> ALL_KINDS =
@@ -67,13 +62,15 @@ public class ExportloadingTransformsTest implements Serializable {
   private static final ImmutableSet<String> ALL_KIND_STRS =
       ALL_KINDS.stream().map(Key::getKind).collect(ImmutableSet.toImmutableSet());
 
-  @Rule public final transient TemporaryFolder exportRootDir = new TemporaryFolder();
+  @SuppressWarnings("WeakerAccess")
+  @TempDir
+  transient Path tmpDir;
 
-  @Rule public final transient InjectRule injectRule = new InjectRule();
+  @RegisterExtension final transient InjectRule injectRule = new InjectRule();
 
-  @Rule
-  public final transient TestPipeline pipeline =
-      TestPipeline.create().enableAbandonedNodeEnforcement(true);
+  @RegisterExtension
+  final transient TestPipelineExtension testPipeline =
+      TestPipelineExtension.create().enableAbandonedNodeEnforcement(true);
 
   private FakeClock fakeClock;
   private transient BackupTestStore store;
@@ -84,8 +81,8 @@ public class ExportloadingTransformsTest implements Serializable {
   private transient ContactResource contact;
   private transient DomainBase domain;
 
-  @Before
-  public void beforeEach() throws Exception {
+  @BeforeEach
+  void beforeEach() throws Exception {
     fakeClock = new FakeClock(START_TIME);
     store = new BackupTestStore(fakeClock);
     injectRule.setStaticField(Ofy.class, "clock", fakeClock);
@@ -102,12 +99,11 @@ public class ExportloadingTransformsTest implements Serializable {
     contact = (ContactResource) store.loadAsOfyEntity(contact);
     domain = (DomainBase) store.loadAsOfyEntity(domain);
 
-    exportDir =
-        store.export(exportRootDir.getRoot().getAbsolutePath(), ALL_KINDS, Collections.EMPTY_SET);
+    exportDir = store.export(tmpDir.toAbsolutePath().toString(), ALL_KINDS, Collections.EMPTY_SET);
   }
 
-  @After
-  public void afterEach() throws Exception {
+  @AfterEach
+  void afterEach() throws Exception {
     if (store != null) {
       store.close();
       store = null;
@@ -115,10 +111,9 @@ public class ExportloadingTransformsTest implements Serializable {
   }
 
   @Test
-  @Category(NeedsRunner.class)
-  public void getExportFilePatterns() {
+  void getExportFilePatterns() {
     PCollection<String> patterns =
-        pipeline.apply(
+        testPipeline.apply(
             "Get Datastore file patterns",
             Transforms.getDatastoreExportFilePatterns(exportDir.getAbsolutePath(), ALL_KIND_STRS));
 
@@ -130,14 +125,13 @@ public class ExportloadingTransformsTest implements Serializable {
 
     PAssert.that(patterns).containsInAnyOrder(expectedPatterns);
 
-    pipeline.run();
+    testPipeline.run();
   }
 
   @Test
-  @Category(NeedsRunner.class)
-  public void getFilesByPatterns() {
+  void getFilesByPatterns() {
     PCollection<Metadata> fileMetas =
-        pipeline
+        testPipeline
             .apply(
                 "File patterns to metadata",
                 Create.of(
@@ -169,13 +163,13 @@ public class ExportloadingTransformsTest implements Serializable {
 
     PAssert.that(fileNames).containsInAnyOrder(expectedFilenames);
 
-    pipeline.run();
+    testPipeline.run();
   }
 
   @Test
-  public void loadDataFromFiles() {
+  void loadDataFromFiles() {
     PCollection<VersionedEntity> entities =
-        pipeline
+        testPipeline
             .apply(
                 "Get Datastore file patterns",
                 Transforms.getDatastoreExportFilePatterns(
@@ -189,6 +183,6 @@ public class ExportloadingTransformsTest implements Serializable {
         KV.of(Transforms.EXPORT_ENTITY_TIME_STAMP, store.loadAsDatastoreEntity(contact)),
         KV.of(Transforms.EXPORT_ENTITY_TIME_STAMP, store.loadAsDatastoreEntity(domain)));
 
-    pipeline.run();
+    testPipeline.run();
   }
 }

core/src/test/java/google/registry/beam/initsql/InitSqlPipelineGraphTest.java

@@ -0,0 +1,67 @@
+// 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 google.registry.testing.truth.TextDiffSubject.assertWithMessageAboutUrlSource;
+
+import com.google.common.io.Resources;
+import google.registry.beam.TestPipelineExtension;
+import java.io.File;
+import java.io.IOException;
+import java.io.PrintStream;
+import java.net.URL;
+import org.apache.beam.runners.core.construction.renderer.PipelineDotRenderer;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+/** Manages visualization of {@link InitSqlPipeline}. */
+class InitSqlPipelineGraphTest {
+
+  private static final String GOLDEN_DOT_FILE = "pipeline_golden.dot";
+
+  private static final String[] OPTIONS_ARGS =
+      new String[] {
+        "--commitLogStartTimestamp=2000-01-01TZ",
+        "--commitLogEndTimestamp=2000-01-02TZ",
+        "--datastoreExportDir=/somedir",
+        "--commitLogDir=/someotherdir",
+        "--environment=alpha"
+      };
+
+  private static final transient InitSqlPipelineOptions options =
+      PipelineOptionsFactory.fromArgs(OPTIONS_ARGS)
+          .withValidation()
+          .as(InitSqlPipelineOptions.class);
+
+  @RegisterExtension
+  final transient TestPipelineExtension testPipeline =
+      TestPipelineExtension.create().enableAbandonedNodeEnforcement(false);
+
+  @Test
+  public void createPipeline_compareGraph() throws IOException {
+    new InitSqlPipeline(options, testPipeline).setupPipeline();
+    String dotString = PipelineDotRenderer.toDotString(testPipeline);
+    URL goldenDotUrl = Resources.getResource(InitSqlPipelineGraphTest.class, GOLDEN_DOT_FILE);
+    File outputFile = new File(new File(goldenDotUrl.getFile()).getParent(), "pipeline_curr.dot");
+    try (PrintStream ps = new PrintStream(outputFile)) {
+      ps.print(dotString);
+    }
+    assertWithMessageAboutUrlSource(
+            "InitSqlPipeline graph changed. Run :core:updateInitSqlPipelineGraph to update.")
+        .that(outputFile.toURI().toURL())
+        .hasSameContentAs(goldenDotUrl);
+  }
+}

core/src/test/java/google/registry/beam/initsql/InitSqlPipelineOptionsTest.java

@@ -1,4 +1,4 @@
-// Copyright 2017 The Nomulus Authors. All Rights Reserved.
+// 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.
@@ -12,20 +12,16 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package google.registry.bigquery;
+package google.registry.beam.initsql;
 
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.junit.jupiter.api.Test;
 
-/** Unit tests for {@link BigqueryConnection}. */
-@RunWith(JUnit4.class)
-public class BigqueryConnectionTest {
+/** Unit tests for {@link google.registry.beam.initsql.InitSqlPipelineOptions}. * */
+public class InitSqlPipelineOptionsTest {
 
   @Test
-  public void testNothing() {
-    // Placeholder test class for now.
-    // TODO(b/16569089): figure out a good way for testing our Bigquery usage overall - maybe unit
-    // tests here, maybe end-to-end testing.
+  void registerToValidate() {
+    PipelineOptionsFactory.register(InitSqlPipelineOptions.class);
   }
 }

core/src/test/java/google/registry/beam/initsql/InitSqlPipelineTest.java

@@ -0,0 +1,281 @@
+// 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.truth.Truth.assertThat;
+import static google.registry.model.ImmutableObjectSubject.assertAboutImmutableObjects;
+import static google.registry.model.ImmutableObjectSubject.immutableObjectCorrespondence;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.testing.DatastoreHelper.newRegistry;
+import static google.registry.testing.DatastoreHelper.persistResource;
+import static google.registry.util.DateTimeUtils.START_OF_TIME;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+import com.googlecode.objectify.Key;
+import google.registry.backup.AppEngineEnvironment;
+import google.registry.beam.TestPipelineExtension;
+import google.registry.model.billing.BillingEvent;
+import google.registry.model.contact.ContactResource;
+import google.registry.model.domain.DesignatedContact;
+import google.registry.model.domain.DomainAuthInfo;
+import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.GracePeriod;
+import google.registry.model.domain.launch.LaunchNotice;
+import google.registry.model.domain.rgp.GracePeriodStatus;
+import google.registry.model.domain.secdns.DelegationSignerData;
+import google.registry.model.eppcommon.AuthInfo.PasswordAuth;
+import google.registry.model.eppcommon.StatusValue;
+import google.registry.model.eppcommon.Trid;
+import google.registry.model.host.HostResource;
+import google.registry.model.ofy.Ofy;
+import google.registry.model.poll.PollMessage;
+import google.registry.model.registrar.Registrar;
+import google.registry.model.registry.Registry;
+import google.registry.model.reporting.HistoryEntry;
+import google.registry.model.transfer.DomainTransferData;
+import google.registry.model.transfer.TransferStatus;
+import google.registry.persistence.VKey;
+import google.registry.persistence.transaction.JpaTestRules;
+import google.registry.persistence.transaction.JpaTestRules.JpaIntegrationTestRule;
+import google.registry.testing.AppEngineRule;
+import google.registry.testing.DatastoreEntityExtension;
+import google.registry.testing.FakeClock;
+import google.registry.testing.InjectRule;
+import java.io.File;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.joda.time.DateTime;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Order;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.api.io.TempDir;
+
+/** Unit tests for {@link InitSqlPipeline}. */
+class InitSqlPipelineTest {
+  private static final DateTime START_TIME = DateTime.parse("2000-01-01T00:00:00.0Z");
+
+  private static final ImmutableList<Class<?>> ALL_KINDS =
+      ImmutableList.of(
+          Registry.class,
+          Registrar.class,
+          ContactResource.class,
+          HostResource.class,
+          DomainBase.class,
+          HistoryEntry.class);
+
+  private transient FakeClock fakeClock = new FakeClock(START_TIME);
+
+  @RegisterExtension
+  @Order(Order.DEFAULT - 1)
+  final transient DatastoreEntityExtension datastore = new DatastoreEntityExtension();
+
+  @RegisterExtension final transient InjectRule injectRule = new InjectRule();
+
+  @SuppressWarnings("WeakerAccess")
+  @TempDir
+  transient Path tmpDir;
+
+  @RegisterExtension
+  final transient TestPipelineExtension testPipeline =
+      TestPipelineExtension.create().enableAbandonedNodeEnforcement(true);
+
+  @RegisterExtension
+  final transient JpaIntegrationTestRule database =
+      new JpaTestRules.Builder().withClock(fakeClock).buildIntegrationTestRule();
+
+  // Must not be transient!
+  @RegisterExtension
+  @Order(Order.DEFAULT + 1)
+  final BeamJpaExtension beamJpaExtension =
+      new BeamJpaExtension(() -> tmpDir.resolve("credential.dat"), database.getDatabase());
+
+  private File exportRootDir;
+  private File exportDir;
+  private File commitLogDir;
+
+  private transient Registrar registrar1;
+  private transient Registrar registrar2;
+  private transient DomainBase domain;
+  private transient ContactResource contact1;
+  private transient ContactResource contact2;
+  private transient HostResource hostResource;
+
+  private transient HistoryEntry historyEntry;
+
+  @BeforeEach
+  public void beforeEach() throws Exception {
+    try (BackupTestStore store = new BackupTestStore(fakeClock)) {
+      injectRule.setStaticField(Ofy.class, "clock", fakeClock);
+      exportRootDir = Files.createDirectory(tmpDir.resolve("exports")).toFile();
+
+      persistResource(newRegistry("com", "COM"));
+      registrar1 = persistResource(AppEngineRule.makeRegistrar1());
+      registrar2 = persistResource(AppEngineRule.makeRegistrar2());
+      Key<DomainBase> domainKey = Key.create(null, DomainBase.class, "4-COM");
+      hostResource =
+          persistResource(
+              new HostResource.Builder()
+                  .setHostName("ns1.example.com")
+                  .setSuperordinateDomain(VKey.from(domainKey))
+                  .setRepoId("1-COM")
+                  .setCreationClientId(registrar1.getClientId())
+                  .setPersistedCurrentSponsorClientId(registrar2.getClientId())
+                  .build());
+      contact1 =
+          persistResource(
+              new ContactResource.Builder()
+                  .setContactId("contact_id1")
+                  .setRepoId("2-COM")
+                  .setCreationClientId(registrar1.getClientId())
+                  .setPersistedCurrentSponsorClientId(registrar2.getClientId())
+                  .build());
+      contact2 =
+          persistResource(
+              new ContactResource.Builder()
+                  .setContactId("contact_id2")
+                  .setRepoId("3-COM")
+                  .setCreationClientId(registrar1.getClientId())
+                  .setPersistedCurrentSponsorClientId(registrar1.getClientId())
+                  .build());
+      historyEntry = persistResource(new HistoryEntry.Builder().setParent(domainKey).build());
+      Key<HistoryEntry> historyEntryKey = Key.create(historyEntry);
+      Key<BillingEvent.OneTime> oneTimeBillKey =
+          Key.create(historyEntryKey, BillingEvent.OneTime.class, 1);
+      VKey<BillingEvent.Recurring> recurringBillKey =
+          VKey.from(Key.create(historyEntryKey, BillingEvent.Recurring.class, 2));
+      VKey<PollMessage.Autorenew> autorenewPollKey =
+          VKey.from(Key.create(historyEntryKey, PollMessage.Autorenew.class, 3));
+      VKey<PollMessage.OneTime> onetimePollKey =
+          VKey.from(Key.create(historyEntryKey, PollMessage.OneTime.class, 1));
+      domain =
+          persistResource(
+              new DomainBase.Builder()
+                  .setDomainName("example.com")
+                  .setRepoId("4-COM")
+                  .setCreationClientId(registrar1.getClientId())
+                  .setLastEppUpdateTime(fakeClock.nowUtc())
+                  .setLastEppUpdateClientId(registrar2.getClientId())
+                  .setLastTransferTime(fakeClock.nowUtc())
+                  .setStatusValues(
+                      ImmutableSet.of(
+                          StatusValue.CLIENT_DELETE_PROHIBITED,
+                          StatusValue.SERVER_DELETE_PROHIBITED,
+                          StatusValue.SERVER_TRANSFER_PROHIBITED,
+                          StatusValue.SERVER_UPDATE_PROHIBITED,
+                          StatusValue.SERVER_RENEW_PROHIBITED,
+                          StatusValue.SERVER_HOLD))
+                  .setRegistrant(contact1.createVKey())
+                  .setContacts(
+                      ImmutableSet.of(
+                          DesignatedContact.create(
+                              DesignatedContact.Type.ADMIN, contact2.createVKey())))
+                  .setNameservers(ImmutableSet.of(hostResource.createVKey()))
+                  .setSubordinateHosts(ImmutableSet.of("ns1.example.com"))
+                  .setPersistedCurrentSponsorClientId(registrar2.getClientId())
+                  .setRegistrationExpirationTime(fakeClock.nowUtc().plusYears(1))
+                  .setAuthInfo(DomainAuthInfo.create(PasswordAuth.create("password")))
+                  .setDsData(
+                      ImmutableSet.of(DelegationSignerData.create(1, 2, 3, new byte[] {0, 1, 2})))
+                  .setLaunchNotice(
+                      LaunchNotice.create("tcnid", "validatorId", START_OF_TIME, START_OF_TIME))
+                  .setTransferData(
+                      new DomainTransferData.Builder()
+                          .setGainingClientId(registrar1.getClientId())
+                          .setLosingClientId(registrar2.getClientId())
+                          .setPendingTransferExpirationTime(fakeClock.nowUtc())
+                          .setServerApproveEntities(
+                              ImmutableSet.of(
+                                  VKey.from(oneTimeBillKey), recurringBillKey, autorenewPollKey))
+                          .setServerApproveBillingEvent(VKey.from(oneTimeBillKey))
+                          .setServerApproveAutorenewEvent(recurringBillKey)
+                          .setServerApproveAutorenewPollMessage(autorenewPollKey)
+                          .setTransferRequestTime(fakeClock.nowUtc().plusDays(1))
+                          .setTransferStatus(TransferStatus.SERVER_APPROVED)
+                          .setTransferRequestTrid(Trid.create("client-trid", "server-trid"))
+                          .build())
+                  .setDeletePollMessage(onetimePollKey)
+                  .setAutorenewBillingEvent(recurringBillKey)
+                  .setAutorenewPollMessage(autorenewPollKey)
+                  .setSmdId("smdid")
+                  .addGracePeriod(
+                      GracePeriod.create(
+                          GracePeriodStatus.ADD, fakeClock.nowUtc().plusDays(1), "registrar", null))
+                  .build());
+      exportDir = store.export(exportRootDir.getAbsolutePath(), ALL_KINDS, ImmutableSet.of());
+      commitLogDir = Files.createDirectory(tmpDir.resolve("commits")).toFile();
+    }
+  }
+
+  @Test
+  public void runPipeline() {
+    InitSqlPipelineOptions options =
+        PipelineOptionsFactory.fromArgs(
+                "--sqlCredentialUrlOverride="
+                    + beamJpaExtension.getCredentialFile().getAbsolutePath(),
+                "--commitLogStartTimestamp=" + START_TIME,
+                "--commitLogEndTimestamp=" + fakeClock.nowUtc().plusMillis(1),
+                "--datastoreExportDir=" + exportDir.getAbsolutePath(),
+                "--commitLogDir=" + commitLogDir.getAbsolutePath())
+            .withValidation()
+            .as(InitSqlPipelineOptions.class);
+    InitSqlPipeline initSqlPipeline = new InitSqlPipeline(options, testPipeline);
+    initSqlPipeline.run().waitUntilFinish();
+    try (AppEngineEnvironment env = new AppEngineEnvironment("test")) {
+      assertHostResourceEquals(
+          jpaTm().transact(() -> jpaTm().load(hostResource.createVKey())), hostResource);
+      assertThat(jpaTm().transact(() -> jpaTm().loadAll(Registrar.class)))
+          .comparingElementsUsing(immutableObjectCorrespondence("lastUpdateTime"))
+          .containsExactly(registrar1, registrar2);
+      assertThat(jpaTm().transact(() -> jpaTm().loadAll(ContactResource.class)))
+          .comparingElementsUsing(immutableObjectCorrespondence("revisions", "updateTimestamp"))
+          .containsExactly(contact1, contact2);
+      assertCleansedDomainEquals(jpaTm().transact(() -> jpaTm().load(domain.createVKey())), domain);
+    }
+  }
+
+  private static void assertHostResourceEquals(HostResource actual, HostResource expected) {
+    assertAboutImmutableObjects()
+        .that(actual)
+        .isEqualExceptFields(expected, "superordinateDomain", "revisions", "updateTimestamp");
+    assertThat(actual.getSuperordinateDomain().getSqlKey())
+        .isEqualTo(expected.getSuperordinateDomain().getSqlKey());
+  }
+
+  private static void assertCleansedDomainEquals(DomainBase actual, DomainBase expected) {
+    assertAboutImmutableObjects()
+        .that(actual)
+        .isEqualExceptFields(
+            expected,
+            "adminContact",
+            "registrantContact",
+            "gracePeriods",
+            "dsData",
+            "allContacts",
+            "revisions",
+            "updateTimestamp",
+            "autorenewBillingEvent",
+            "autorenewPollMessage",
+            "deletePollMessage",
+            "nsHosts",
+            "transferData");
+    assertThat(actual.getAdminContact().getSqlKey())
+        .isEqualTo(expected.getAdminContact().getSqlKey());
+    assertThat(actual.getRegistrant().getSqlKey()).isEqualTo(expected.getRegistrant().getSqlKey());
+    // TODO(weiminyu): compare gracePeriods, allContacts and dsData, when SQL model supports them.
+  }
+}

core/src/test/java/google/registry/beam/initsql/LoadDatastoreSnapshotTest.java

@@ -22,6 +22,7 @@ import com.google.appengine.api.datastore.Entity;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
+import google.registry.beam.TestPipelineExtension;
 import google.registry.model.contact.ContactResource;
 import google.registry.model.domain.DomainAuthInfo;
 import google.registry.model.domain.DomainBase;
@@ -31,16 +32,15 @@ import google.registry.model.registry.Registry;
 import google.registry.testing.FakeClock;
 import google.registry.testing.InjectRule;
 import java.io.File;
-import org.apache.beam.sdk.testing.TestPipeline;
+import java.nio.file.Files;
+import java.nio.file.Path;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.PCollectionTuple;
 import org.joda.time.DateTime;
-import org.junit.Before;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.rules.TemporaryFolder;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.api.io.TempDir;
 
 /**
  * Unit test for {@link Transforms#loadDatastoreSnapshot}.
@@ -69,8 +69,8 @@ import org.junit.runners.JUnit4;
  *   <li>Deletes are properly handled.
  * </ul>
  */
-@RunWith(JUnit4.class)
-public class LoadDatastoreSnapshotTest {
+class LoadDatastoreSnapshotTest {
+
   private static final DateTime START_TIME = DateTime.parse("2000-01-01T00:00:00.0Z");
 
   private static final ImmutableList<Class<?>> ALL_KINDS =
@@ -78,13 +78,15 @@ public class LoadDatastoreSnapshotTest {
   private static final ImmutableSet<String> ALL_KIND_STRS =
       ALL_KINDS.stream().map(Key::getKind).collect(ImmutableSet.toImmutableSet());
 
-  @Rule public final transient TemporaryFolder temporaryFolder = new TemporaryFolder();
+  @SuppressWarnings("WeakerAccess")
+  @TempDir
+  transient Path tmpDir;
 
-  @Rule public final transient InjectRule injectRule = new InjectRule();
+  @RegisterExtension final transient InjectRule injectRule = new InjectRule();
 
-  @Rule
-  public final transient TestPipeline pipeline =
-      TestPipeline.create().enableAbandonedNodeEnforcement(true);
+  @RegisterExtension
+  final transient TestPipelineExtension testPipeline =
+      TestPipelineExtension.create().enableAbandonedNodeEnforcement(true);
 
   private FakeClock fakeClock;
   private File exportRootDir;
@@ -100,14 +102,14 @@ public class LoadDatastoreSnapshotTest {
   private transient DateTime contactLastUpdateTime;
   private transient DateTime domainLastUpdateTime;
 
-  @Before
-  public void beforeEach() throws Exception {
+  @BeforeEach
+  void beforeEach() throws Exception {
     fakeClock = new FakeClock(START_TIME);
     try (BackupTestStore store = new BackupTestStore(fakeClock)) {
       injectRule.setStaticField(Ofy.class, "clock", fakeClock);
 
-      exportRootDir = temporaryFolder.newFolder();
-      commitLogsDir = temporaryFolder.newFolder();
+      exportRootDir = Files.createDirectory(tmpDir.resolve("export_root")).toFile();
+      commitLogsDir = Files.createDirectory(tmpDir.resolve("commit_logs")).toFile();
 
       Registry registry = newRegistry("tld1", "TLD1");
       ContactResource fillerContact = newContactResource("contact_filler");
@@ -152,9 +154,9 @@ public class LoadDatastoreSnapshotTest {
   }
 
   @Test
-  public void loadDatastoreSnapshot() {
+  void loadDatastoreSnapshot() {
     PCollectionTuple snapshot =
-        pipeline.apply(
+        testPipeline.apply(
             Transforms.loadDatastoreSnapshot(
                 exportDir.getAbsolutePath(),
                 commitLogsDir.getAbsolutePath(),
@@ -170,6 +172,6 @@ public class LoadDatastoreSnapshotTest {
     InitSqlTestUtils.assertContainsExactlyElementsIn(
         snapshot.get(Transforms.createTagForKind("ContactResource")),
         KV.of(contactLastUpdateTime.getMillis(), dsContact));
-    pipeline.run();
+    testPipeline.run();
   }
 }

core/src/test/java/google/registry/beam/initsql/WriteToSqlTest.java

@@ -0,0 +1,131 @@
+// 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.truth.Truth.assertThat;
+import static google.registry.model.ImmutableObjectSubject.immutableObjectCorrespondence;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+
+import com.google.appengine.api.datastore.Entity;
+import com.google.common.collect.ImmutableList;
+import google.registry.backup.VersionedEntity;
+import google.registry.beam.TestPipelineExtension;
+import google.registry.model.ImmutableObject;
+import google.registry.model.contact.ContactResource;
+import google.registry.model.ofy.Ofy;
+import google.registry.model.registrar.Registrar;
+import google.registry.persistence.transaction.JpaTestRules;
+import google.registry.persistence.transaction.JpaTestRules.JpaIntegrationTestRule;
+import google.registry.testing.AppEngineRule;
+import google.registry.testing.DatastoreEntityExtension;
+import google.registry.testing.DatastoreHelper;
+import google.registry.testing.FakeClock;
+import google.registry.testing.InjectRule;
+import java.io.Serializable;
+import java.nio.file.Path;
+import java.util.stream.Collectors;
+import org.apache.beam.sdk.transforms.Create;
+import org.joda.time.DateTime;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Order;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.api.io.TempDir;
+
+/** Unit test for {@link Transforms#writeToSql}. */
+class WriteToSqlTest implements Serializable {
+
+  private static final DateTime START_TIME = DateTime.parse("2000-01-01T00:00:00.0Z");
+
+  private final FakeClock fakeClock = new FakeClock(START_TIME);
+
+  @RegisterExtension
+  @Order(Order.DEFAULT - 1)
+  final transient DatastoreEntityExtension datastore = new DatastoreEntityExtension();
+
+  @RegisterExtension final transient InjectRule injectRule = new InjectRule();
+
+  @RegisterExtension
+  final transient JpaIntegrationTestRule database =
+      new JpaTestRules.Builder().withClock(fakeClock).buildIntegrationTestRule();
+
+  @SuppressWarnings("WeakerAccess")
+  @TempDir
+  transient Path tmpDir;
+
+  @RegisterExtension
+  final transient TestPipelineExtension testPipeline =
+      TestPipelineExtension.create().enableAbandonedNodeEnforcement(true);
+
+  // Must not be transient!
+  @RegisterExtension
+  @Order(Order.DEFAULT + 1)
+  public final BeamJpaExtension beamJpaExtension =
+      new BeamJpaExtension(() -> tmpDir.resolve("credential.dat"), database.getDatabase());
+
+  private ImmutableList<Entity> contacts;
+
+  @BeforeEach
+  void beforeEach() throws Exception {
+    try (BackupTestStore store = new BackupTestStore(fakeClock)) {
+      injectRule.setStaticField(Ofy.class, "clock", fakeClock);
+
+      // Required for contacts created below.
+      Registrar ofyRegistrar = AppEngineRule.makeRegistrar2();
+      store.insertOrUpdate(ofyRegistrar);
+      jpaTm().transact(() -> jpaTm().saveNewOrUpdate(store.loadAsOfyEntity(ofyRegistrar)));
+
+      ImmutableList.Builder<Entity> builder = new ImmutableList.Builder<>();
+
+      for (int i = 0; i < 3; i++) {
+        ContactResource contact = DatastoreHelper.newContactResource("contact_" + i);
+        store.insertOrUpdate(contact);
+        builder.add(store.loadAsDatastoreEntity(contact));
+      }
+      contacts = builder.build();
+    }
+  }
+
+  @Test
+  void writeToSql_twoWriters() {
+    testPipeline
+        .apply(
+            Create.of(
+                contacts.stream()
+                    .map(InitSqlTestUtils::entityToBytes)
+                    .map(bytes -> VersionedEntity.from(0L, bytes))
+                    .collect(Collectors.toList())))
+        .apply(
+            Transforms.writeToSql(
+                "ContactResource",
+                2,
+                4,
+                () ->
+                    DaggerBeamJpaModule_JpaTransactionManagerComponent.builder()
+                        .beamJpaModule(beamJpaExtension.getBeamJpaModule())
+                        .build()
+                        .localDbJpaTransactionManager()));
+    testPipeline.run().waitUntilFinish();
+
+    ImmutableList<?> sqlContacts = jpaTm().transact(() -> jpaTm().loadAll(ContactResource.class));
+    assertThat(sqlContacts)
+        .comparingElementsUsing(immutableObjectCorrespondence("revisions", "updateTimestamp"))
+        .containsExactlyElementsIn(
+            contacts.stream()
+                .map(InitSqlTestUtils::datastoreToOfyEntity)
+                .map(ImmutableObject.class::cast)
+                .collect(ImmutableList.toImmutableList()));
+  }
+}