Handle bad production data when migrating to SQL (#1120)

* Handle bad production data when migrating to SQL

Ignore or fix bad entites when populating SQL with production data in
Datastore. These are mostly inconsistent foreign keys.

See b/185954992 for details.

build.gradle

@@ -24,7 +24,7 @@ buildscript {
   }
 
   dependencies {
-    classpath 'com.google.cloud.tools:appengine-gradle-plugin:2.0.1'
+    classpath 'com.google.cloud.tools:appengine-gradle-plugin:2.4.1'
     classpath 'net.ltgt.gradle:gradle-errorprone-plugin:0.6.1'
     classpath 'org.sonatype.aether:aether-api:1.13.1'
     classpath 'org.sonatype.aether:aether-impl:1.13.1'

buildSrc/src/main/java/google/registry/gradle/plugin/ProjectData.java

@@ -91,7 +91,7 @@ abstract class ProjectData {
       /** The task was actually run and has finished successfully. */
       SUCCESS,
       /** The task was up-to-date and successful, and hence didn't need to run again. */
-      UP_TO_DATE;
+      UP_TO_DATE
     }
 
     abstract String uniqueName();

common/src/testing/java/google/registry/testing/FakeClock.java

@@ -35,6 +35,8 @@ public final class FakeClock implements Clock {
   // threads should see a consistent flow.
   private final AtomicLong currentTimeMillis = new AtomicLong();
 
+  private volatile long autoIncrementStepMs;
+
   /** Creates a FakeClock that starts at START_OF_TIME. */
   public FakeClock() {
     this(START_OF_TIME);
@@ -48,7 +50,21 @@ public final class FakeClock implements Clock {
   /** Returns the current time. */
   @Override
   public DateTime nowUtc() {
-    return new DateTime(currentTimeMillis.get(), UTC);
+    return new DateTime(currentTimeMillis.addAndGet(autoIncrementStepMs), UTC);
+  }
+
+  /**
+   * Sets the increment applied to the clock whenever it is queried. The increment is zero by
+   * default: the clock is left unchanged when queried.
+   *
+   * <p>Passing a duration of zero to this method effectively unsets the auto increment mode.
+   *
+   * @param autoIncrementStep the new auto increment duration
+   * @return this
+   */
+  public FakeClock setAutoIncrementStep(ReadableDuration autoIncrementStep) {
+    this.autoIncrementStepMs = autoIncrementStep.getMillis();
+    return this;
   }
 
   /** Advances clock by one millisecond. */

config/nom_build.py

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

config/presubmits.py

@@ -196,7 +196,7 @@ PRESUBMITS = {
         #    - concatenation of literals: (\s*\+\s*"([^"]|\\")*")*
         # Line 3: , or the closing parenthesis, marking the end of the first
         #    parameter
-        r'.*\.create(Native)?Query\('
+        r'.*\.(query|createQuery|createNativeQuery)\('
         r'(?!(\s*([A-Z_]+|"([^"]|\\")*"(\s*\+\s*"([^"]|\\")*")*)'
         r'(,|\s*\))))',
         "java",
@@ -206,10 +206,12 @@ PRESUBMITS = {
          # using Criteria
          "ForeignKeyIndex.java",
          "HistoryEntryDao.java",
+         "JpaTransactionManager.java",
          "JpaTransactionManagerImpl.java",
          # CriteriaQueryBuilder is a false positive
          "CriteriaQueryBuilder.java",
          "RdapDomainSearchAction.java",
+         "RdapNameserverSearchAction.java",
          "RdapSearchActionBase.java",
          },
     ):

core/build.gradle

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import java.lang.reflect.Constructor
+import com.google.common.base.CaseFormat
 import java.util.Optional
 
 plugins {
@@ -330,6 +330,7 @@ dependencies {
   testCompile deps['org.junit.platform:junit-platform-suite-api']
   testCompile deps['org.mockito:mockito-core']
   testCompile deps['org.mockito:mockito-junit-jupiter']
+  testCompile 'org.checkerframework:checker-qual:3.9.1'
   runtime deps['org.postgresql:postgresql']
 
   // Indirect dependency found by undeclared-dependency check. Such
@@ -423,7 +424,7 @@ task jaxbToJava {
       }
     }
     execInBash(
-        'find . -name *.java -exec sed -i /\\*\\ \\<p\\>\\$/d {} +',
+        "find . -name *.java -exec sed -i -e '/" + /\* <p>$/ + "/d' {} +",
         generatedDir)
   }
 }
@@ -791,6 +792,57 @@ createUberJar(
     'beam_pipeline_common',
     '')
 
+// Create beam staging task if environment is alpha or crash.
+// All other environments use formally released pipelines through CloudBuild.
+//
+// User should install gcloud and login to GCP before invoking this tasks.
+if (environment in ['alpha', 'crash']) {
+  def pipelines = [
+      [
+          mainClass: 'google.registry.beam.initsql.InitSqlPipeline',
+          metaData: 'google/registry/beam/init_sql_pipeline_metadata.json'
+      ],
+      [
+          mainClass: 'google.registry.beam.datastore.BulkDeleteDatastorePipeline',
+          metaData: 'google/registry/beam/bulk_delete_datastore_pipeline_metadata.json'
+      ],
+      [
+          mainClass: 'google.registry.beam.spec11.Spec11Pipeline',
+          metaData: 'google/registry/beam/spec11_pipeline_metadata.json'
+      ],
+      [
+          mainClass: 'google.registry.beam.invoicing.InvoicingPipeline',
+          metaData: 'google/registry/beam/invoicing_pipeline_metadata.json'
+      ],
+  ]
+  project.tasks.create("stage_beam_pipelines") {
+    doLast {
+      pipelines.each {
+        def mainClass = it['mainClass']
+        def metaData = it['metaData']
+        def pipelineName = CaseFormat.UPPER_CAMEL.to(
+                CaseFormat.LOWER_UNDERSCORE,
+                mainClass.substring(mainClass.lastIndexOf('.') + 1))
+        def imageName = "gcr.io/${gcpProject}/beam/${pipelineName}"
+        def metaDataBaseName = metaData.substring(metaData.lastIndexOf('/') + 1)
+        def uberJarName = tasks.beam_pipeline_common.outputs.files.asPath
+
+        def command = "\
+            gcloud dataflow flex-template build \
+            gs://${gcpProject}-deploy/live/beam/${metaDataBaseName} \
+            --image-gcr-path ${imageName}:live \
+            --sdk-language JAVA \
+            --flex-template-base-image JAVA11 \
+            --metadata-file ${projectDir}/src/main/resources/${metaData} \
+            --jar ${uberJarName} \
+            --env FLEX_TEMPLATE_JAVA_MAIN_CLASS=${mainClass} \
+            --project ${gcpProject}".toString()
+        rootProject.ext.execInBash(command, '/tmp')
+      }
+    }
+  }.dependsOn(tasks.beam_pipeline_common)
+}
+
 // A jar with classes and resources from main sourceSet, excluding internal
 // data. See comments on configurations.nomulus_test above for details.
 // TODO(weiminyu): release process should build this using the public repo to eliminate the need

core/gradle/dependency-locks/compile.lockfile

@@ -253,10 +253,10 @@ org.postgresql:postgresql:42.2.18
 org.rnorth.duct-tape:duct-tape:1.0.8
 org.rnorth.visible-assertions:visible-assertions:2.1.2
 org.slf4j:slf4j-api:1.7.30
-org.testcontainers:database-commons:1.15.1
-org.testcontainers:jdbc:1.15.1
-org.testcontainers:postgresql:1.15.1
-org.testcontainers:testcontainers:1.15.1
+org.testcontainers:database-commons:1.15.2
+org.testcontainers:jdbc:1.15.2
+org.testcontainers:postgresql:1.15.2
+org.testcontainers:testcontainers:1.15.2
 org.threeten:threetenbp:1.5.0
 org.tukaani:xz:1.5
 org.w3c.css:sac:1.3

core/gradle/dependency-locks/compileClasspath.lockfile

@@ -246,10 +246,10 @@ org.postgresql:postgresql:42.2.18
 org.rnorth.duct-tape:duct-tape:1.0.8
 org.rnorth.visible-assertions:visible-assertions:2.1.2
 org.slf4j:slf4j-api:1.7.30
-org.testcontainers:database-commons:1.15.1
-org.testcontainers:jdbc:1.15.1
-org.testcontainers:postgresql:1.15.1
-org.testcontainers:testcontainers:1.15.1
+org.testcontainers:database-commons:1.15.2
+org.testcontainers:jdbc:1.15.2
+org.testcontainers:postgresql:1.15.2
+org.testcontainers:testcontainers:1.15.2
 org.threeten:threetenbp:1.5.0
 org.tukaani:xz:1.5
 org.w3c.css:sac:1.3

core/gradle/dependency-locks/default.lockfile

@@ -267,10 +267,10 @@ org.slf4j:jcl-over-slf4j:1.7.30
 org.slf4j:jul-to-slf4j:1.7.30
 org.slf4j:slf4j-api:1.7.30
 org.slf4j:slf4j-jdk14:1.7.28
-org.testcontainers:database-commons:1.15.1
-org.testcontainers:jdbc:1.15.1
-org.testcontainers:postgresql:1.15.1
-org.testcontainers:testcontainers:1.15.1
+org.testcontainers:database-commons:1.15.2
+org.testcontainers:jdbc:1.15.2
+org.testcontainers:postgresql:1.15.2
+org.testcontainers:testcontainers:1.15.2
 org.threeten:threetenbp:1.5.0
 org.tukaani:xz:1.5
 org.w3c.css:sac:1.3

core/gradle/dependency-locks/deploy_jar.lockfile

@@ -266,10 +266,10 @@ org.slf4j:jcl-over-slf4j:1.7.30
 org.slf4j:jul-to-slf4j:1.7.30
 org.slf4j:slf4j-api:1.7.30
 org.slf4j:slf4j-jdk14:1.7.28
-org.testcontainers:database-commons:1.15.1
-org.testcontainers:jdbc:1.15.1
-org.testcontainers:postgresql:1.15.1
-org.testcontainers:testcontainers:1.15.1
+org.testcontainers:database-commons:1.15.2
+org.testcontainers:jdbc:1.15.2
+org.testcontainers:postgresql:1.15.2
+org.testcontainers:testcontainers:1.15.2
 org.threeten:threetenbp:1.5.0
 org.tukaani:xz:1.5
 org.w3c.css:sac:1.3

core/gradle/dependency-locks/nonprodCompile.lockfile

@@ -253,10 +253,10 @@ org.postgresql:postgresql:42.2.18
 org.rnorth.duct-tape:duct-tape:1.0.8
 org.rnorth.visible-assertions:visible-assertions:2.1.2
 org.slf4j:slf4j-api:1.7.30
-org.testcontainers:database-commons:1.15.1
-org.testcontainers:jdbc:1.15.1
-org.testcontainers:postgresql:1.15.1
-org.testcontainers:testcontainers:1.15.1
+org.testcontainers:database-commons:1.15.2
+org.testcontainers:jdbc:1.15.2
+org.testcontainers:postgresql:1.15.2
+org.testcontainers:testcontainers:1.15.2
 org.threeten:threetenbp:1.5.0
 org.tukaani:xz:1.5
 org.w3c.css:sac:1.3

core/gradle/dependency-locks/nonprodCompileClasspath.lockfile

@@ -247,10 +247,10 @@ org.postgresql:postgresql:42.2.18
 org.rnorth.duct-tape:duct-tape:1.0.8
 org.rnorth.visible-assertions:visible-assertions:2.1.2
 org.slf4j:slf4j-api:1.7.30
-org.testcontainers:database-commons:1.15.1
-org.testcontainers:jdbc:1.15.1
-org.testcontainers:postgresql:1.15.1
-org.testcontainers:testcontainers:1.15.1
+org.testcontainers:database-commons:1.15.2
+org.testcontainers:jdbc:1.15.2
+org.testcontainers:postgresql:1.15.2
+org.testcontainers:testcontainers:1.15.2
 org.threeten:threetenbp:1.5.0
 org.tukaani:xz:1.5
 org.w3c.css:sac:1.3

core/gradle/dependency-locks/nonprodRuntime.lockfile

@@ -265,10 +265,10 @@ org.rnorth.visible-assertions:visible-assertions:2.1.2
 org.slf4j:jcl-over-slf4j:1.7.30
 org.slf4j:jul-to-slf4j:1.7.30
 org.slf4j:slf4j-api:1.7.30
-org.testcontainers:database-commons:1.15.1
-org.testcontainers:jdbc:1.15.1
-org.testcontainers:postgresql:1.15.1
-org.testcontainers:testcontainers:1.15.1
+org.testcontainers:database-commons:1.15.2
+org.testcontainers:jdbc:1.15.2
+org.testcontainers:postgresql:1.15.2
+org.testcontainers:testcontainers:1.15.2
 org.threeten:threetenbp:1.5.0
 org.tukaani:xz:1.5
 org.w3c.css:sac:1.3

core/gradle/dependency-locks/nonprodRuntimeClasspath.lockfile

@@ -265,10 +265,10 @@ org.rnorth.visible-assertions:visible-assertions:2.1.2
 org.slf4j:jcl-over-slf4j:1.7.30
 org.slf4j:jul-to-slf4j:1.7.30
 org.slf4j:slf4j-api:1.7.30
-org.testcontainers:database-commons:1.15.1
-org.testcontainers:jdbc:1.15.1
-org.testcontainers:postgresql:1.15.1
-org.testcontainers:testcontainers:1.15.1
+org.testcontainers:database-commons:1.15.2
+org.testcontainers:jdbc:1.15.2
+org.testcontainers:postgresql:1.15.2
+org.testcontainers:testcontainers:1.15.2
 org.threeten:threetenbp:1.5.0
 org.tukaani:xz:1.5
 org.w3c.css:sac:1.3

core/gradle/dependency-locks/runtime.lockfile

@@ -265,10 +265,10 @@ org.rnorth.visible-assertions:visible-assertions:2.1.2
 org.slf4j:jcl-over-slf4j:1.7.30
 org.slf4j:jul-to-slf4j:1.7.30
 org.slf4j:slf4j-api:1.7.30
-org.testcontainers:database-commons:1.15.1
-org.testcontainers:jdbc:1.15.1
-org.testcontainers:postgresql:1.15.1
-org.testcontainers:testcontainers:1.15.1
+org.testcontainers:database-commons:1.15.2
+org.testcontainers:jdbc:1.15.2
+org.testcontainers:postgresql:1.15.2
+org.testcontainers:testcontainers:1.15.2
 org.threeten:threetenbp:1.5.0
 org.tukaani:xz:1.5
 org.w3c.css:sac:1.3

core/gradle/dependency-locks/runtimeClasspath.lockfile

@@ -266,10 +266,10 @@ org.slf4j:jcl-over-slf4j:1.7.30
 org.slf4j:jul-to-slf4j:1.7.30
 org.slf4j:slf4j-api:1.7.30
 org.slf4j:slf4j-jdk14:1.7.28
-org.testcontainers:database-commons:1.15.1
-org.testcontainers:jdbc:1.15.1
-org.testcontainers:postgresql:1.15.1
-org.testcontainers:testcontainers:1.15.1
+org.testcontainers:database-commons:1.15.2
+org.testcontainers:jdbc:1.15.2
+org.testcontainers:postgresql:1.15.2
+org.testcontainers:testcontainers:1.15.2
 org.threeten:threetenbp:1.5.0
 org.tukaani:xz:1.5
 org.w3c.css:sac:1.3

core/gradle/dependency-locks/testCompile.lockfile

@@ -300,12 +300,12 @@ org.seleniumhq.selenium:selenium-remote-driver:3.141.59
 org.seleniumhq.selenium:selenium-safari-driver:3.141.59
 org.seleniumhq.selenium:selenium-support:3.141.59
 org.slf4j:slf4j-api:1.7.30
-org.testcontainers:database-commons:1.15.1
-org.testcontainers:jdbc:1.15.1
-org.testcontainers:junit-jupiter:1.15.1
-org.testcontainers:postgresql:1.15.1
-org.testcontainers:selenium:1.15.1
-org.testcontainers:testcontainers:1.15.1
+org.testcontainers:database-commons:1.15.2
+org.testcontainers:jdbc:1.15.2
+org.testcontainers:junit-jupiter:1.15.2
+org.testcontainers:postgresql:1.15.2
+org.testcontainers:selenium:1.15.2
+org.testcontainers:testcontainers:1.15.2
 org.threeten:threetenbp:1.5.0
 org.tukaani:xz:1.5
 org.w3c.css:sac:1.3

core/gradle/dependency-locks/testCompileClasspath.lockfile

@@ -294,12 +294,12 @@ org.seleniumhq.selenium:selenium-remote-driver:3.141.59
 org.seleniumhq.selenium:selenium-safari-driver:3.141.59
 org.seleniumhq.selenium:selenium-support:3.141.59
 org.slf4j:slf4j-api:1.7.30
-org.testcontainers:database-commons:1.15.1
-org.testcontainers:jdbc:1.15.1
-org.testcontainers:junit-jupiter:1.15.1
-org.testcontainers:postgresql:1.15.1
-org.testcontainers:selenium:1.15.1
-org.testcontainers:testcontainers:1.15.1
+org.testcontainers:database-commons:1.15.2
+org.testcontainers:jdbc:1.15.2
+org.testcontainers:junit-jupiter:1.15.2
+org.testcontainers:postgresql:1.15.2
+org.testcontainers:selenium:1.15.2
+org.testcontainers:testcontainers:1.15.2
 org.threeten:threetenbp:1.5.0
 org.tukaani:xz:1.5
 org.w3c.css:sac:1.3

core/gradle/dependency-locks/testRuntime.lockfile

@@ -313,12 +313,12 @@ org.seleniumhq.selenium:selenium-support:3.141.59
 org.slf4j:jcl-over-slf4j:1.7.30
 org.slf4j:jul-to-slf4j:1.7.30
 org.slf4j:slf4j-api:1.7.30
-org.testcontainers:database-commons:1.15.1
-org.testcontainers:jdbc:1.15.1
-org.testcontainers:junit-jupiter:1.15.1
-org.testcontainers:postgresql:1.15.1
-org.testcontainers:selenium:1.15.1
-org.testcontainers:testcontainers:1.15.1
+org.testcontainers:database-commons:1.15.2
+org.testcontainers:jdbc:1.15.2
+org.testcontainers:junit-jupiter:1.15.2
+org.testcontainers:postgresql:1.15.2
+org.testcontainers:selenium:1.15.2
+org.testcontainers:testcontainers:1.15.2
 org.threeten:threetenbp:1.5.0
 org.tukaani:xz:1.5
 org.w3c.css:sac:1.3

core/gradle/dependency-locks/testRuntimeClasspath.lockfile

@@ -314,12 +314,12 @@ org.slf4j:jcl-over-slf4j:1.7.30
 org.slf4j:jul-to-slf4j:1.7.30
 org.slf4j:slf4j-api:1.7.30
 org.slf4j:slf4j-jdk14:1.7.28
-org.testcontainers:database-commons:1.15.1
-org.testcontainers:jdbc:1.15.1
-org.testcontainers:junit-jupiter:1.15.1
-org.testcontainers:postgresql:1.15.1
-org.testcontainers:selenium:1.15.1
-org.testcontainers:testcontainers:1.15.1
+org.testcontainers:database-commons:1.15.2
+org.testcontainers:jdbc:1.15.2
+org.testcontainers:junit-jupiter:1.15.2
+org.testcontainers:postgresql:1.15.2
+org.testcontainers:selenium:1.15.2
+org.testcontainers:testcontainers:1.15.2
 org.threeten:threetenbp:1.5.0
 org.tukaani:xz:1.5
 org.w3c.css:sac:1.3

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

@@ -154,7 +154,13 @@ public class ReplayCommitLogsToSqlAction implements Runnable {
     Object ofyPojo = ofy().toPojo(entity);
     if (ofyPojo instanceof DatastoreEntity) {
       DatastoreEntity datastoreEntity = (DatastoreEntity) ofyPojo;
-      datastoreEntity.toSqlEntity().ifPresent(jpaTm()::put);
+      datastoreEntity
+          .toSqlEntity()
+          .ifPresent(
+              sqlEntity -> {
+                ReplaySpecializer.beforeSqlSave(sqlEntity);
+                jpaTm().put(sqlEntity);
+              });
     } else {
       // this should never happen, but we shouldn't fail on it
       logger.atSevere().log(

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

@@ -57,8 +57,6 @@ public final class AsyncTaskEnqueuer {
   public static final String QUEUE_ASYNC_DELETE = "async-delete-pull";
   public static final String QUEUE_ASYNC_HOST_RENAME = "async-host-rename-pull";
 
-  public static final String PATH_RESAVE_ENTITY = "/_dr/task/resaveEntity";
-
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
   private static final Duration MAX_ASYNC_ETA = Duration.standardDays(30);
 
@@ -112,7 +110,7 @@ public final class AsyncTaskEnqueuer {
     logger.atInfo().log("Enqueuing async re-save of %s to run at %s.", entityKey, whenToResave);
     String backendHostname = appEngineServiceUtils.getServiceHostname("backend");
     TaskOptions task =
-        TaskOptions.Builder.withUrl(PATH_RESAVE_ENTITY)
+        TaskOptions.Builder.withUrl(ResaveEntityAction.PATH)
             .method(Method.POST)
             .header("Host", backendHostname)
             .countdownMillis(etaDuration.getMillis())

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

@@ -25,8 +25,6 @@ import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_AUTORENEW;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.pricing.PricingEngineProxy.getDomainRenewCost;
-import static google.registry.schema.cursor.Cursor.GLOBAL;
-import static google.registry.schema.cursor.CursorDao.loadAndCompare;
 import static google.registry.util.CollectionUtils.union;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
 import static google.registry.util.DateTimeUtils.earliestOf;
@@ -60,7 +58,6 @@ import google.registry.request.Action;
 import google.registry.request.Parameter;
 import google.registry.request.Response;
 import google.registry.request.auth.Auth;
-import google.registry.schema.cursor.CursorDao;
 import google.registry.util.Clock;
 import java.util.Optional;
 import java.util.Set;
@@ -95,7 +92,6 @@ public class ExpandRecurringBillingEventsAction implements Runnable {
   @Override
   public void run() {
     Cursor cursor = ofy().load().key(Cursor.createGlobalKey(RECURRING_BILLING)).now();
-    loadAndCompare(cursor, GLOBAL);
     DateTime executeTime = clock.nowUtc();
     DateTime persistedCursorTime = (cursor == null ? START_OF_TIME : cursor.getCursorTime());
     DateTime cursorTime = cursorTimeParam.orElse(persistedCursorTime);
@@ -332,7 +328,6 @@ public class ExpandRecurringBillingEventsAction implements Runnable {
       tm().transact(
               () -> {
                 Cursor cursor = ofy().load().key(Cursor.createGlobalKey(RECURRING_BILLING)).now();
-                loadAndCompare(cursor, GLOBAL);
                 DateTime currentCursorTime =
                     (cursor == null ? START_OF_TIME : cursor.getCursorTime());
                 if (!currentCursorTime.equals(expectedPersistedCursorTime)) {
@@ -342,8 +337,7 @@ public class ExpandRecurringBillingEventsAction implements Runnable {
                   return;
                 }
                 if (!isDryRun) {
-                  CursorDao.saveCursor(
-                      Cursor.createGlobal(RECURRING_BILLING, executionTime), GLOBAL);
+                  tm().put(Cursor.createGlobal(RECURRING_BILLING, executionTime));
                 }
               });
     }

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

@@ -25,7 +25,7 @@ import static google.registry.batch.AsyncTaskMetrics.OperationType.DNS_REFRESH;
 import static google.registry.mapreduce.inputs.EppResourceInputs.createEntityInput;
 import static google.registry.model.EppResourceUtils.isActive;
 import static google.registry.model.EppResourceUtils.isDeleted;
-import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.DateTimeUtils.latestOf;
 import static java.util.concurrent.TimeUnit.DAYS;
 import static java.util.concurrent.TimeUnit.SECONDS;
@@ -44,7 +44,6 @@ import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.flogger.FluentLogger;
-import com.googlecode.objectify.Key;
 import google.registry.batch.AsyncTaskMetrics.OperationResult;
 import google.registry.dns.DnsQueue;
 import google.registry.mapreduce.MapreduceRunner;
@@ -64,6 +63,7 @@ import google.registry.util.SystemClock;
 import java.io.Serializable;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.NoSuchElementException;
 import java.util.Optional;
 import java.util.logging.Level;
 import javax.annotation.Nullable;
@@ -123,7 +123,7 @@ public class RefreshDnsOnHostRenameAction implements Runnable {
     }
 
     ImmutableList.Builder<DnsRefreshRequest> requestsBuilder = new ImmutableList.Builder<>();
-    ImmutableList.Builder<Key<HostResource>> hostKeys = new ImmutableList.Builder<>();
+    ImmutableList.Builder<VKey<HostResource>> hostKeys = new ImmutableList.Builder<>();
     final List<DnsRefreshRequest> requestsToDelete = new ArrayList<>();
 
     for (TaskHandle task : tasks) {
@@ -204,10 +204,10 @@ public class RefreshDnsOnHostRenameAction implements Runnable {
         emit(true, true);
         return;
       }
-      Key<HostResource> referencingHostKey = null;
+      VKey<HostResource> referencingHostKey = null;
       for (DnsRefreshRequest request : refreshRequests) {
         if (isActive(domain, request.lastUpdateTime())
-            && domain.getNameservers().contains(VKey.from(request.hostKey()))) {
+            && domain.getNameservers().contains(request.hostKey())) {
           referencingHostKey = request.hostKey();
           break;
         }
@@ -293,7 +293,8 @@ public class RefreshDnsOnHostRenameAction implements Runnable {
 
     private static final long serialVersionUID = 1772812852271288622L;
 
-    abstract Key<HostResource> hostKey();
+    abstract VKey<HostResource> hostKey();
+
     abstract DateTime lastUpdateTime();
     abstract DateTime requestedTime();
     abstract boolean isRefreshNeeded();
@@ -301,7 +302,8 @@ public class RefreshDnsOnHostRenameAction implements Runnable {
 
     @AutoValue.Builder
     abstract static class Builder {
-      abstract Builder setHostKey(Key<HostResource> hostKey);
+      abstract Builder setHostKey(VKey<HostResource> hostKey);
+
       abstract Builder setLastUpdateTime(DateTime lastUpdateTime);
       abstract Builder setRequestedTime(DateTime requestedTime);
       abstract Builder setIsRefreshNeeded(boolean isRefreshNeeded);
@@ -314,10 +316,12 @@ public class RefreshDnsOnHostRenameAction implements Runnable {
      */
     static DnsRefreshRequest createFromTask(TaskHandle task, DateTime now) throws Exception {
       ImmutableMap<String, String> params = ImmutableMap.copyOf(task.extractParams());
-      Key<HostResource> hostKey =
-          Key.create(checkNotNull(params.get(PARAM_HOST_KEY), "Host to refresh not specified"));
+      VKey<HostResource> hostKey =
+          VKey.fromWebsafeKey(
+              checkNotNull(params.get(PARAM_HOST_KEY), "Host to refresh not specified"));
       HostResource host =
-          checkNotNull(ofy().load().key(hostKey).now(), "Host to refresh doesn't exist");
+          tm().transact(() -> tm().loadByKeyIfPresent(hostKey))
+              .orElseThrow(() -> new NoSuchElementException("Host to refresh doesn't exist"));
       boolean isHostDeleted =
           isDeleted(host, latestOf(now, host.getUpdateTimestamp().getTimestamp()));
       if (isHostDeleted) {

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

@@ -16,7 +16,6 @@ package google.registry.batch;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.request.Action.Method.POST;
@@ -33,6 +32,7 @@ import google.registry.model.eppcommon.StatusValue;
 import google.registry.model.registrar.Registrar;
 import google.registry.model.registrar.RegistrarContact;
 import google.registry.model.registry.RegistryLockDao;
+import google.registry.persistence.VKey;
 import google.registry.request.Action;
 import google.registry.request.Parameter;
 import google.registry.request.Response;
@@ -125,6 +125,7 @@ public class RelockDomainAction implements Runnable {
     response.setContentType(MediaType.PLAIN_TEXT_UTF_8);
 
     // nb: DomainLockUtils relies on the JPA transaction being the outermost transaction
+    // if we have Datastore as the primary DB (if SQL is the primary DB, it's irrelevant)
     jpaTm().transact(() -> tm().transact(this::relockDomain));
   }
 
@@ -139,12 +140,8 @@ public class RelockDomainAction implements Runnable {
                       new IllegalArgumentException(
                           String.format("Unknown revision ID %d", oldUnlockRevisionId)));
       domain =
-          ofy()
-              .load()
-              .type(DomainBase.class)
-              .id(oldLock.getRepoId())
-              .now()
-              .cloneProjectedAtTime(jpaTm().getTransactionTime());
+          tm().loadByKey(VKey.create(DomainBase.class, oldLock.getRepoId()))
+              .cloneProjectedAtTime(tm().getTransactionTime());
     } catch (Throwable t) {
       handleTransientFailure(Optional.ofNullable(oldLock), t);
       return;

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

@@ -17,7 +17,6 @@ package google.registry.batch;
 import static google.registry.batch.AsyncTaskEnqueuer.PARAM_REQUESTED_TIME;
 import static google.registry.batch.AsyncTaskEnqueuer.PARAM_RESAVE_TIMES;
 import static google.registry.batch.AsyncTaskEnqueuer.PARAM_RESOURCE_KEY;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.common.collect.ImmutableSet;
@@ -26,6 +25,7 @@ import com.google.common.flogger.FluentLogger;
 import com.googlecode.objectify.Key;
 import google.registry.model.EppResource;
 import google.registry.model.ImmutableObject;
+import google.registry.persistence.VKey;
 import google.registry.request.Action;
 import google.registry.request.Action.Method;
 import google.registry.request.Parameter;
@@ -74,16 +74,17 @@ public class ResaveEntityAction implements Runnable {
   public void run() {
     logger.atInfo().log(
         "Re-saving entity %s which was enqueued at %s.", resourceKey, requestedTime);
-    tm().transact(() -> {
-      ImmutableObject entity = ofy().load().key(resourceKey).now();
-      ofy().save().entity(
-          (entity instanceof EppResource)
-              ? ((EppResource) entity).cloneProjectedAtTime(tm().getTransactionTime()) : entity
-      );
-      if (!resaveTimes.isEmpty()) {
-        asyncTaskEnqueuer.enqueueAsyncResave(entity, requestedTime, resaveTimes);
-      }
-    });
+    tm().transact(
+            () -> {
+              ImmutableObject entity = tm().loadByKey(VKey.from(resourceKey));
+              tm().put(
+                      (entity instanceof EppResource)
+                          ? ((EppResource) entity).cloneProjectedAtTime(tm().getTransactionTime())
+                          : entity);
+              if (!resaveTimes.isEmpty()) {
+                asyncTaskEnqueuer.enqueueAsyncResave(entity, requestedTime, resaveTimes);
+              }
+            });
     response.setPayload("Entity re-saved.");
   }
 }

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

@@ -0,0 +1,165 @@
+// Copyright 2021 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.batch;
+
+import static com.google.common.net.MediaType.PLAIN_TEXT_UTF_8;
+import static javax.servlet.http.HttpServletResponse.SC_FORBIDDEN;
+import static javax.servlet.http.HttpServletResponse.SC_INTERNAL_SERVER_ERROR;
+import static javax.servlet.http.HttpServletResponse.SC_OK;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.flogger.FluentLogger;
+import google.registry.config.RegistryConfig.Config;
+import google.registry.persistence.PersistenceModule.SchemaManagerConnection;
+import google.registry.request.Action;
+import google.registry.request.Response;
+import google.registry.request.auth.Auth;
+import google.registry.util.Retrier;
+import java.sql.Connection;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.function.Supplier;
+import javax.inject.Inject;
+
+/**
+ * Wipes out all Cloud SQL data in a Nomulus GCP environment.
+ *
+ * <p>This class is created for the QA environment, where migration testing with production data
+ * will happen. A regularly scheduled wipeout is a prerequisite to using production data there.
+ */
+@Action(
+    service = Action.Service.BACKEND,
+    path = "/_dr/task/wipeOutCloudSql",
+    auth = Auth.AUTH_INTERNAL_OR_ADMIN)
+public class WipeOutCloudSqlAction implements Runnable {
+  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
+
+  // As a short-lived class, hardcode allowed projects here instead of using config files.
+  private static final ImmutableSet<String> ALLOWED_PROJECTS =
+      ImmutableSet.of("domain-registry-qa");
+
+  private final String projectId;
+  private final Supplier<Connection> connectionSupplier;
+  private final Response response;
+  private final Retrier retrier;
+
+  @Inject
+  WipeOutCloudSqlAction(
+      @Config("projectId") String projectId,
+      @SchemaManagerConnection Supplier<Connection> connectionSupplier,
+      Response response,
+      Retrier retrier) {
+    this.projectId = projectId;
+    this.connectionSupplier = connectionSupplier;
+    this.response = response;
+    this.retrier = retrier;
+  }
+
+  @Override
+  public void run() {
+    response.setContentType(PLAIN_TEXT_UTF_8);
+
+    if (!ALLOWED_PROJECTS.contains(projectId)) {
+      response.setStatus(SC_FORBIDDEN);
+      response.setPayload("Wipeout is not allowed in " + projectId);
+      return;
+    }
+
+    try {
+      retrier.callWithRetry(
+          () -> {
+            try (Connection conn = connectionSupplier.get()) {
+              dropAllTables(conn, listTables(conn));
+              dropAllSequences(conn, listSequences(conn));
+            }
+            return null;
+          },
+          e -> !(e instanceof SQLException));
+      response.setStatus(SC_OK);
+      response.setPayload("Wiped out Cloud SQL in " + projectId);
+    } catch (RuntimeException e) {
+      logger.atSevere().withCause(e).log("Failed to wipe out Cloud SQL data.");
+      response.setStatus(SC_INTERNAL_SERVER_ERROR);
+      response.setPayload("Failed to wipe out Cloud SQL in " + projectId);
+    }
+  }
+
+  /** Returns a list of all tables in the public schema of a Postgresql database. */
+  static ImmutableList<String> listTables(Connection connection) throws SQLException {
+    try (ResultSet resultSet =
+        connection.getMetaData().getTables(null, null, null, new String[] {"TABLE"})) {
+      ImmutableList.Builder<String> tables = new ImmutableList.Builder<>();
+      while (resultSet.next()) {
+        String schema = resultSet.getString("TABLE_SCHEM");
+        if (schema == null || !schema.equalsIgnoreCase("public")) {
+          continue;
+        }
+        String tableName = resultSet.getString("TABLE_NAME");
+        tables.add("public.\"" + tableName + "\"");
+      }
+      return tables.build();
+    }
+  }
+
+  static void dropAllTables(Connection conn, ImmutableList<String> tables) throws SQLException {
+    if (tables.isEmpty()) {
+      return;
+    }
+
+    try (Statement statement = conn.createStatement()) {
+      for (String table : tables) {
+        statement.addBatch(String.format("DROP TABLE IF EXISTS %s CASCADE;", table));
+      }
+      for (int code : statement.executeBatch()) {
+        if (code == Statement.EXECUTE_FAILED) {
+          throw new RuntimeException("Failed to drop some tables. Please check.");
+        }
+      }
+    }
+  }
+
+  /** Returns a list of all sequences in a Postgresql database. */
+  static ImmutableList<String> listSequences(Connection conn) throws SQLException {
+    try (Statement statement = conn.createStatement();
+        ResultSet resultSet =
+            statement.executeQuery("SELECT c.relname FROM pg_class c WHERE c.relkind = 'S';")) {
+      ImmutableList.Builder<String> sequences = new ImmutableList.Builder<>();
+      while (resultSet.next()) {
+        sequences.add('\"' + resultSet.getString(1) + '\"');
+      }
+      return sequences.build();
+    }
+  }
+
+  static void dropAllSequences(Connection conn, ImmutableList<String> sequences)
+      throws SQLException {
+    if (sequences.isEmpty()) {
+      return;
+    }
+
+    try (Statement statement = conn.createStatement()) {
+      for (String sequence : sequences) {
+        statement.addBatch(String.format("DROP SEQUENCE IF EXISTS %s CASCADE;", sequence));
+      }
+      for (int code : statement.executeBatch()) {
+        if (code == Statement.EXECUTE_FAILED) {
+          throw new RuntimeException("Failed to drop some sequences. Please check.");
+        }
+      }
+    }
+  }
+}

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

@@ -0,0 +1,115 @@
+// Copyright 2021 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.batch;
+
+import static com.google.common.net.MediaType.PLAIN_TEXT_UTF_8;
+import static google.registry.beam.BeamUtils.createJobName;
+import static javax.servlet.http.HttpServletResponse.SC_FORBIDDEN;
+import static javax.servlet.http.HttpServletResponse.SC_INTERNAL_SERVER_ERROR;
+import static javax.servlet.http.HttpServletResponse.SC_OK;
+
+import com.google.api.services.dataflow.Dataflow;
+import com.google.api.services.dataflow.model.LaunchFlexTemplateParameter;
+import com.google.api.services.dataflow.model.LaunchFlexTemplateRequest;
+import com.google.api.services.dataflow.model.LaunchFlexTemplateResponse;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.flogger.FluentLogger;
+import google.registry.config.RegistryConfig.Config;
+import google.registry.request.Action;
+import google.registry.request.Response;
+import google.registry.request.auth.Auth;
+import google.registry.util.Clock;
+import javax.inject.Inject;
+
+/**
+ * Wipes out all Cloud Datastore data in a Nomulus GCP environment.
+ *
+ * <p>This class is created for the QA environment, where migration testing with production data
+ * will happen. A regularly scheduled wipeout is a prerequisite to using production data there.
+ */
+@Action(
+    service = Action.Service.BACKEND,
+    path = "/_dr/task/wipeOutDatastore",
+    auth = Auth.AUTH_INTERNAL_OR_ADMIN)
+public class WipeoutDatastoreAction implements Runnable {
+  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
+
+  private static final String PIPELINE_NAME = "bulk_delete_datastore_pipeline";
+
+  // As a short-lived class, hardcode allowed projects here instead of using config files.
+  private static final ImmutableSet<String> ALLOWED_PROJECTS =
+      ImmutableSet.of("domain-registry-qa");
+
+  private final String projectId;
+  private final String jobRegion;
+  private final Response response;
+  private final Dataflow dataflow;
+  private final String stagingBucketUrl;
+  private final Clock clock;
+
+  @Inject
+  WipeoutDatastoreAction(
+      @Config("projectId") String projectId,
+      @Config("defaultJobRegion") String jobRegion,
+      @Config("beamStagingBucketUrl") String stagingBucketUrl,
+      Clock clock,
+      Response response,
+      Dataflow dataflow) {
+    this.projectId = projectId;
+    this.jobRegion = jobRegion;
+    this.stagingBucketUrl = stagingBucketUrl;
+    this.clock = clock;
+    this.response = response;
+    this.dataflow = dataflow;
+  }
+
+  @Override
+  public void run() {
+    response.setContentType(PLAIN_TEXT_UTF_8);
+
+    if (!ALLOWED_PROJECTS.contains(projectId)) {
+      response.setStatus(SC_FORBIDDEN);
+      response.setPayload("Wipeout is not allowed in " + projectId);
+      return;
+    }
+
+    try {
+      LaunchFlexTemplateParameter parameters =
+          new LaunchFlexTemplateParameter()
+              .setJobName(createJobName("bulk-delete-datastore-", clock))
+              .setContainerSpecGcsPath(
+                  String.format("%s/%s_metadata.json", stagingBucketUrl, PIPELINE_NAME))
+              .setParameters(ImmutableMap.of("kindsToDelete", "*"));
+      LaunchFlexTemplateResponse launchResponse =
+          dataflow
+              .projects()
+              .locations()
+              .flexTemplates()
+              .launch(
+                  projectId,
+                  jobRegion,
+                  new LaunchFlexTemplateRequest().setLaunchParameter(parameters))
+              .execute();
+      response.setStatus(SC_OK);
+      response.setPayload("Launched " + launchResponse.getJob().getName());
+    } catch (Exception e) {
+      String msg = String.format("Failed to launch %s.", PIPELINE_NAME);
+      logger.atSevere().withCause(e).log(msg);
+      response.setStatus(SC_INTERNAL_SERVER_ERROR);
+      response.setPayload(msg);
+    }
+  }
+}

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

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

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

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

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

@@ -505,7 +505,7 @@ public class DatastoreV1 {
       }
 
       @StartBundle
-      public void startBundle(StartBundleContext c) throws Exception {
+      public void startBundle(StartBundleContext c) {
         datastore =
             datastoreFactory.getDatastore(
                 c.getPipelineOptions(), v1Options.getProjectId(), v1Options.getLocalhost());
@@ -548,7 +548,7 @@ public class DatastoreV1 {
       }
 
       @StartBundle
-      public void startBundle(StartBundleContext c) throws Exception {
+      public void startBundle(StartBundleContext c) {
         datastore =
             datastoreFactory.getDatastore(
                 c.getPipelineOptions(), options.getProjectId(), options.getLocalhost());
@@ -556,7 +556,7 @@ public class DatastoreV1 {
       }
 
       @ProcessElement
-      public void processElement(ProcessContext c) throws Exception {
+      public void processElement(ProcessContext c) {
         Query query = c.element();
 
         // If query has a user set limit, then do not split.
@@ -626,7 +626,7 @@ public class DatastoreV1 {
       }
 
       @StartBundle
-      public void startBundle(StartBundleContext c) throws Exception {
+      public void startBundle(StartBundleContext c) {
         datastore =
             datastoreFactory.getDatastore(
                 c.getPipelineOptions(), options.getProjectId(), options.getLocalhost());

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

@@ -93,7 +93,7 @@ public final class BackupPaths {
     checkArgument(!isNullOrEmpty(exportDir), "Null or empty exportDir.");
     checkArgument(!isNullOrEmpty(kind), "Null or empty kind.");
     checkArgument(shard >= 0, "Negative shard %s not allowed.", shard);
-    return String.format(EXPORT_PATTERN_TEMPLATE, exportDir, kind, Integer.toString(shard));
+    return String.format(EXPORT_PATTERN_TEMPLATE, exportDir, kind, shard);
   }
 
   /** Returns an {@link ImmutableList} of regex patterns that match all CommitLog files. */

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

core/src/main/java/google/registry/bigquery/BigqueryConnection.java

@@ -300,10 +300,9 @@ public class BigqueryConnection implements AutoCloseable {
    * Initializes the BigqueryConnection object by setting up the API client and creating the default
    * dataset if it doesn't exist.
    */
-  private BigqueryConnection initialize() throws Exception {
+  private void initialize() throws Exception {
     createDatasetIfNeeded(datasetId);
     createDatasetIfNeeded(TEMP_DATASET_NAME);
-    return this;
   }
 
   /**
@@ -378,13 +377,11 @@ public class BigqueryConnection implements AutoCloseable {
 
   /**
    * Starts an asynchronous load job to populate the specified destination table with the given
-   * source URIs and source format.  Returns a ListenableFuture that holds the same destination
-   * table object on success.
+   * source URIs and source format. Returns a ListenableFuture that holds the same destination table
+   * object on success.
    */
-  public ListenableFuture<DestinationTable> load(
-      DestinationTable dest,
-      SourceFormat sourceFormat,
-      Iterable<String> sourceUris) {
+  public ListenableFuture<DestinationTable> startLoad(
+      DestinationTable dest, SourceFormat sourceFormat, Iterable<String> sourceUris) {
     Job job = new Job()
         .setConfiguration(new JobConfiguration()
             .setLoad(new JobConfigurationLoad()
@@ -400,9 +397,7 @@ public class BigqueryConnection implements AutoCloseable {
    * of the specified query, or if the table is a view, to update the view to reflect that query.
    * Returns a ListenableFuture that holds the same destination table object on success.
    */
-  public ListenableFuture<DestinationTable> query(
-      String querySql,
-      DestinationTable dest) {
+  public ListenableFuture<DestinationTable> startQuery(String querySql, DestinationTable dest) {
     if (dest.type == TableType.VIEW) {
       // Use Futures.transform() rather than calling apply() directly so that any exceptions thrown
       // by calling updateTable will be propagated on the get() call, not from here.
@@ -562,20 +557,18 @@ public class BigqueryConnection implements AutoCloseable {
     // Tracking bug for query-to-GCS support is b/13777340.
     DestinationTable tempTable = buildTemporaryTable().build();
     return transformAsync(
-        query(querySql, tempTable),
+        startQuery(querySql, tempTable),
         tempTable1 -> extractTable(tempTable1, destinationUri, destinationFormat, printHeader),
         directExecutor());
   }
 
   /** @see #runJob(Job, AbstractInputStreamContent) */
-  public Job runJob(Job job) {
+  private Job runJob(Job job) {
     return runJob(job, null);
   }
 
-  /**
-   * Launch a job, wait for it to complete, but <i>do not</i> check for errors.
-   */
-  public Job runJob(Job job, @Nullable AbstractInputStreamContent data) {
+  /** Launch a job, wait for it to complete, but <i>do not</i> check for errors. */
+  private Job runJob(Job job, @Nullable AbstractInputStreamContent data) {
     return checkJob(waitForJob(launchJob(job, data)));
   }
 

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

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

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

@@ -123,6 +123,7 @@ public class RegistryConfigSettings {
   /** Configuration for Cloud SQL. */
   public static class CloudSql {
     public String jdbcUrl;
+    // TODO(05012021): remove username field after it is removed from all yaml files.
     public String username;
     public String instanceConnectionName;
     public boolean replicateTransactions;
@@ -132,7 +133,7 @@ public class RegistryConfigSettings {
   /** Configuration for Apache Beam (Cloud Dataflow). */
   public static class Beam {
     public String defaultJobRegion;
-    public String defaultJobZone;
+    public String stagingBucketUrl;
   }
 
   /** Configuration for Cloud DNS. */
@@ -220,6 +221,7 @@ public class RegistryConfigSettings {
   public static class RegistryTool {
     public String clientId;
     public String clientSecret;
+    // TODO(05012021): remove username field after it is removed from all yaml files.
     public String username;
   }
 

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

@@ -225,8 +225,6 @@ cloudSql:
   # If jdbcUrl in this file is moved elsewhere, be sure to move this notice
   # with it until the change is applied.
   jdbcUrl: jdbc:postgresql://localhost
-  # Username for the database user.
-  username: username
   # This name is used by Cloud SQL when connecting to the database.
   instanceConnectionName: project-id:region:instance-id
   # Set this to true to replicate cloud SQL transactions to datastore in the
@@ -422,9 +420,7 @@ misc:
 beam:
   # The default region to run Apache Beam (Cloud Dataflow) jobs in.
   defaultJobRegion: us-east1
-  # The default zone to run Apache Beam (Cloud Dataflow) jobs in.
-  # TODO(weiminyu): consider dropping zone config. No obvious needs for this.
-  defaultJobZone: us-east1-c
+  stagingBucketUrl: gcs-bucket-with-staged-templates
 
 keyring:
   # The name of the active keyring, either "KMS" or "Dummy".
@@ -446,7 +442,6 @@ registryTool:
   clientId: YOUR_CLIENT_ID
   # OAuth client secret used by the tool.
   clientSecret: YOUR_CLIENT_SECRET
-  username: toolusername
 
 # Configuration options for checking SSL certificates.
 sslCertificateValidation:

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

@@ -379,6 +379,18 @@
     <url-pattern>/_dr/task/relockDomain</url-pattern>
   </servlet-mapping>
 
+  <!-- Action to wipeout Cloud SQL data -->
+  <servlet-mapping>
+    <servlet-name>backend-servlet</servlet-name>
+    <url-pattern>/_dr/task/wipeOutCloudSql</url-pattern>
+  </servlet-mapping>
+
+  <!-- Action to wipeout Cloud Datastore data -->
+  <servlet-mapping>
+    <servlet-name>backend-servlet</servlet-name>
+    <url-pattern>/_dr/task/wipeOutDatastore</url-pattern>
+  </servlet-mapping>
+
   <!-- Security config -->
   <security-constraint>
     <web-resource-collection>

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

@@ -82,4 +82,22 @@
     <target>backend</target>
   </cron>
 
+  <cron>
+    <url><![CDATA[/_dr/task/wipeOutCloudSql]]></url>
+    <description>
+      This job runs an action that deletes all data in Cloud SQL.
+    </description>
+    <schedule>every saturday 03:07</schedule>
+    <target>backend</target>
+  </cron>
+
+  <cron>
+    <url><![CDATA[/_dr/task/wipeOutDatastore]]></url>
+    <description>
+      This job runs an action that deletes all data in Cloud Datastore.
+    </description>
+    <schedule>every saturday 03:07</schedule>
+    <target>backend</target>
+  </cron>
+
 </cronentries>

core/src/main/java/google/registry/export/AnnotatedEntities.java

@@ -21,12 +21,13 @@ import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Ordering;
 import com.googlecode.objectify.Key;
 import google.registry.model.EntityClasses;
+import google.registry.model.annotations.InCrossTld;
 import google.registry.model.annotations.NotBackedUp;
 import google.registry.model.annotations.ReportedOn;
 import google.registry.model.annotations.VirtualEntity;
 
 /** Constants related to export code. */
-public final class ExportConstants {
+public final class AnnotatedEntities {
 
   /** Returns the names of kinds to include in Datastore backups. */
   public static ImmutableSet<String> getBackupKinds() {
@@ -49,4 +50,13 @@ public final class ExportConstants {
         .map(Key::getKind)
         .collect(toImmutableSortedSet(Ordering.natural()));
   }
+
+  /** Returns the names of kinds that are in the cross-TLD entity group. */
+  public static ImmutableSet<String> getCrossTldKinds() {
+    return EntityClasses.ALL_CLASSES.stream()
+        .filter(hasAnnotation(InCrossTld.class))
+        .filter(hasAnnotation(VirtualEntity.class).negate())
+        .map(Key::getKind)
+        .collect(toImmutableSortedSet(Ordering.natural()));
+  }
 }

core/src/main/java/google/registry/export/BackupDatastoreAction.java

@@ -66,12 +66,13 @@ public class BackupDatastoreAction implements Runnable {
     try {
       Operation backup =
           datastoreAdmin
-              .export(RegistryConfig.getDatastoreBackupsBucket(), ExportConstants.getBackupKinds())
+              .export(
+                  RegistryConfig.getDatastoreBackupsBucket(), AnnotatedEntities.getBackupKinds())
               .execute();
 
       String backupName = backup.getName();
       // Enqueue a poll task to monitor the backup and load REPORTING-related kinds into bigquery.
-      enqueuePollTask(backupName, ExportConstants.getReportingKinds());
+      enqueuePollTask(backupName, AnnotatedEntities.getReportingKinds());
       String message =
           String.format(
               "Datastore backup started with name: %s\nSaving to %s",

core/src/main/java/google/registry/export/SyncGroupMembersAction.java

@@ -16,7 +16,6 @@ package google.registry.export;
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.request.Action.Method.POST;
 import static google.registry.util.CollectionUtils.nullToEmpty;
@@ -164,7 +163,7 @@ public final class SyncGroupMembersAction implements Runnable {
         registrarsToSave.add(result.getKey().asBuilder().setContactsRequireSyncing(false).build());
       }
     }
-    tm().transactNew(() -> ofy().save().entities(registrarsToSave.build()));
+    tm().transactNew(() -> tm().updateAll(registrarsToSave.build()));
     return errors;
   }
 

core/src/main/java/google/registry/export/sheet/SyncRegistrarsSheet.java

@@ -17,7 +17,6 @@ package google.registry.export.sheet;
 import static com.google.common.base.MoreObjects.firstNonNull;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static google.registry.model.common.Cursor.CursorType.SYNC_REGISTRAR_SHEET;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.registrar.RegistrarContact.Type.ABUSE;
 import static google.registry.model.registrar.RegistrarContact.Type.ADMIN;
 import static google.registry.model.registrar.RegistrarContact.Type.BILLING;
@@ -25,8 +24,8 @@ import static google.registry.model.registrar.RegistrarContact.Type.LEGAL;
 import static google.registry.model.registrar.RegistrarContact.Type.MARKETING;
 import static google.registry.model.registrar.RegistrarContact.Type.TECH;
 import static google.registry.model.registrar.RegistrarContact.Type.WHOIS;
-import static google.registry.schema.cursor.Cursor.GLOBAL;
-import static google.registry.schema.cursor.CursorDao.loadAndCompare;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
 
 import com.google.common.base.Joiner;
@@ -38,10 +37,10 @@ import google.registry.model.common.Cursor;
 import google.registry.model.registrar.Registrar;
 import google.registry.model.registrar.RegistrarAddress;
 import google.registry.model.registrar.RegistrarContact;
-import google.registry.schema.cursor.CursorDao;
 import google.registry.util.Clock;
 import google.registry.util.DateTimeUtils;
 import java.io.IOException;
+import java.util.Optional;
 import java.util.function.Predicate;
 import javax.annotation.Nullable;
 import javax.inject.Inject;
@@ -63,9 +62,10 @@ class SyncRegistrarsSheet {
    * successfully completed, as measured by a cursor.
    */
   boolean wereRegistrarsModified() {
-    Cursor cursor = ofy().load().key(Cursor.createGlobalKey(SYNC_REGISTRAR_SHEET)).now();
-    loadAndCompare(cursor, GLOBAL);
-    DateTime lastUpdateTime = (cursor == null) ? START_OF_TIME : cursor.getCursorTime();
+    Optional<Cursor> cursor =
+        transactIfJpaTm(
+            () -> tm().loadByKeyIfPresent(Cursor.createGlobalVKey(SYNC_REGISTRAR_SHEET)));
+    DateTime lastUpdateTime = !cursor.isPresent() ? START_OF_TIME : cursor.get().getCursorTime();
     for (Registrar registrar : Registrar.loadAllCached()) {
       if (DateTimeUtils.isAtOrAfter(registrar.getLastUpdateTime(), lastUpdateTime)) {
         return true;
@@ -155,9 +155,7 @@ class SyncRegistrarsSheet {
                   return builder.build();
                 })
             .collect(toImmutableList()));
-    CursorDao.saveCursor(
-        Cursor.createGlobal(SYNC_REGISTRAR_SHEET, executionTime),
-        google.registry.schema.cursor.Cursor.GLOBAL);
+    tm().transact(() -> tm().put(Cursor.createGlobal(SYNC_REGISTRAR_SHEET, executionTime)));
   }
 
   private static String convertContacts(

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

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

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

@@ -15,7 +15,7 @@
 package google.registry.flows;
 
 import static com.google.common.base.Preconditions.checkState;
-import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.xml.ValidationMode.LENIENT;
 import static google.registry.xml.ValidationMode.STRICT;
 import static java.nio.charset.StandardCharsets.UTF_8;
@@ -51,8 +51,8 @@ public final class FlowUtils {
 
   /** Persists the saves and deletes in an {@link EntityChanges} to Datastore. */
   public static void persistEntityChanges(EntityChanges entityChanges) {
-    ofy().save().entities(entityChanges.getSaves());
-    ofy().delete().keys(entityChanges.getDeletes());
+    tm().putAll(entityChanges.getSaves());
+    tm().delete(entityChanges.getDeletes());
   }
 
   /**

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

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

core/src/main/java/google/registry/flows/custom/EntityChanges.java

@@ -16,8 +16,8 @@ package google.registry.flows.custom;
 
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableSet;
-import com.googlecode.objectify.Key;
 import google.registry.model.ImmutableObject;
+import google.registry.persistence.VKey;
 
 /** A wrapper class that encapsulates Datastore entities to both save and delete. */
 @AutoValue
@@ -25,7 +25,7 @@ public abstract class EntityChanges {
 
   public abstract ImmutableSet<ImmutableObject> getSaves();
 
-  public abstract ImmutableSet<Key<ImmutableObject>> getDeletes();
+  public abstract ImmutableSet<VKey<ImmutableObject>> getDeletes();
 
   public static Builder newBuilder() {
     // Default both entities to save and entities to delete to empty sets, so that the build()
@@ -48,11 +48,11 @@ public abstract class EntityChanges {
       return this;
     }
 
-    public abstract Builder setDeletes(ImmutableSet<Key<ImmutableObject>> entitiesToDelete);
+    public abstract Builder setDeletes(ImmutableSet<VKey<ImmutableObject>> entitiesToDelete);
 
-    public abstract ImmutableSet.Builder<Key<ImmutableObject>> deletesBuilder();
+    public abstract ImmutableSet.Builder<VKey<ImmutableObject>> deletesBuilder();
 
-    public Builder addDelete(Key<ImmutableObject> entityToDelete) {
+    public Builder addDelete(VKey<ImmutableObject> entityToDelete) {
       deletesBuilder().add(entityToDelete);
       return this;
     }

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

@@ -45,7 +45,7 @@ import google.registry.model.eppinput.ResourceCommand;
 import google.registry.model.eppoutput.EppResponse;
 import google.registry.model.registry.Registry;
 import google.registry.model.reporting.IcannReportingTypes.ActivityReportField;
-import google.registry.model.tmch.ClaimsListShard;
+import google.registry.model.tmch.ClaimsListDualDatabaseDao;
 import google.registry.util.Clock;
 import java.util.HashSet;
 import java.util.Optional;
@@ -104,7 +104,8 @@ public final class DomainClaimsCheckFlow implements Flow {
           verifyClaimsPeriodNotEnded(registry, now);
         }
       }
-      Optional<String> claimKey = ClaimsListShard.get().getClaimKey(parsedDomain.parts().get(0));
+      Optional<String> claimKey =
+          ClaimsListDualDatabaseDao.get().getClaimKey(parsedDomain.parts().get(0));
       launchChecksBuilder.add(
           LaunchCheck.create(
               LaunchCheckName.create(claimKey.isPresent(), domainName), claimKey.orElse(null)));

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

@@ -80,6 +80,7 @@ import google.registry.model.billing.BillingEvent.Recurring;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.DomainCommand;
 import google.registry.model.domain.DomainCommand.Create;
+import google.registry.model.domain.DomainHistory;
 import google.registry.model.domain.GracePeriod;
 import google.registry.model.domain.Period;
 import google.registry.model.domain.fee.FeeCreateCommandExtension;
@@ -206,7 +207,7 @@ public class DomainCreateFlow implements TransactionalFlow {
   @Inject @ClientId String clientId;
   @Inject @TargetId String targetId;
   @Inject @Superuser boolean isSuperuser;
-  @Inject HistoryEntry.Builder historyBuilder;
+  @Inject DomainHistory.Builder historyBuilder;
   @Inject EppResponse.Builder responseBuilder;
   @Inject AllocationTokenFlowUtils allocationTokenFlowUtils;
   @Inject DomainCreateFlowCustomLogic flowCustomLogic;
@@ -303,8 +304,8 @@ public class DomainCreateFlow implements TransactionalFlow {
         validateSecDnsExtension(eppInput.getSingleExtension(SecDnsCreateExtension.class));
     String repoId = createDomainRepoId(ObjectifyService.allocateId(), registry.getTldStr());
     DateTime registrationExpirationTime = leapSafeAddYears(now, years);
-    HistoryEntry historyEntry = buildHistoryEntry(
-        repoId, registry, now, period, registry.getAddGracePeriodLength());
+    DomainHistory domainHistory =
+        buildHistoryEntry(repoId, registry, now, period, registry.getAddGracePeriodLength());
     // Bill for the create.
     BillingEvent.OneTime createBillingEvent =
         createOneTimeBillingEvent(
@@ -314,20 +315,16 @@ public class DomainCreateFlow implements TransactionalFlow {
             isReserved(domainName, isSunriseCreate),
             years,
             feesAndCredits,
-            historyEntry,
+            domainHistory,
             allocationToken,
             now);
     // Create a new autorenew billing event and poll message starting at the expiration time.
     BillingEvent.Recurring autorenewBillingEvent =
-        createAutorenewBillingEvent(historyEntry, registrationExpirationTime);
+        createAutorenewBillingEvent(domainHistory, registrationExpirationTime);
     PollMessage.Autorenew autorenewPollMessage =
-        createAutorenewPollMessage(historyEntry, registrationExpirationTime);
+        createAutorenewPollMessage(domainHistory, registrationExpirationTime);
     ImmutableSet.Builder<ImmutableObject> entitiesToSave = new ImmutableSet.Builder<>();
-    entitiesToSave.add(
-        historyEntry,
-        createBillingEvent,
-        autorenewBillingEvent,
-        autorenewPollMessage);
+    entitiesToSave.add(createBillingEvent, autorenewBillingEvent, autorenewPollMessage);
     // Bill for EAP cost, if any.
     if (!feesAndCredits.getEapCost().isZero()) {
       entitiesToSave.add(createEapBillingEvent(feesAndCredits, createBillingEvent));
@@ -337,7 +334,7 @@ public class DomainCreateFlow implements TransactionalFlow {
     if (getReservationTypes(domainName).contains(NAME_COLLISION)) {
       statuses.add(SERVER_HOLD);
       entitiesToSave.add(
-          createNameCollisionOneTimePollMessage(targetId, historyEntry, clientId, now));
+          createNameCollisionOneTimePollMessage(targetId, domainHistory, clientId, now));
     }
 
     DomainBase newDomain =
@@ -365,13 +362,13 @@ public class DomainCreateFlow implements TransactionalFlow {
             .build();
     entitiesToSave.add(
         newDomain,
+        domainHistory.asBuilder().setDomainContent(newDomain).build(),
         ForeignKeyIndex.create(newDomain, newDomain.getDeletionTime()),
         EppResourceIndex.create(Key.create(newDomain)));
     if (allocationToken.isPresent()
         && TokenType.SINGLE_USE.equals(allocationToken.get().getTokenType())) {
       entitiesToSave.add(
-          allocationTokenFlowUtils.redeemToken(
-              allocationToken.get(), HistoryEntry.createVKey(Key.create(historyEntry))));
+          allocationTokenFlowUtils.redeemToken(allocationToken.get(), domainHistory.createVKey()));
     }
     enqueueTasks(newDomain, hasSignedMarks, hasClaimsNotice);
 
@@ -379,7 +376,7 @@ public class DomainCreateFlow implements TransactionalFlow {
         flowCustomLogic.beforeSave(
             DomainCreateFlowCustomLogic.BeforeSaveParameters.newBuilder()
                 .setNewDomain(newDomain)
-                .setHistoryEntry(historyEntry)
+                .setHistoryEntry(domainHistory)
                 .setEntityChanges(
                     EntityChanges.newBuilder().setSaves(entitiesToSave.build()).build())
                 .setYears(years)
@@ -483,7 +480,7 @@ public class DomainCreateFlow implements TransactionalFlow {
             : null);
   }
 
-  private HistoryEntry buildHistoryEntry(
+  private DomainHistory buildHistoryEntry(
       String repoId, Registry registry, DateTime now, Period period, Duration addGracePeriod) {
     // We ignore prober transactions
     if (registry.getTldType() == TldType.REAL) {
@@ -511,7 +508,7 @@ public class DomainCreateFlow implements TransactionalFlow {
       boolean isReserved,
       int years,
       FeesAndCredits feesAndCredits,
-      HistoryEntry historyEntry,
+      DomainHistory domainHistory,
       Optional<AllocationToken> allocationToken,
       DateTime now) {
     ImmutableSet.Builder<Flag> flagsBuilder = new ImmutableSet.Builder<>();
@@ -540,12 +537,12 @@ public class DomainCreateFlow implements TransactionalFlow {
                     ? registry.getAnchorTenantAddGracePeriodLength()
                     : registry.getAddGracePeriodLength()))
         .setFlags(flagsBuilder.build())
-        .setParent(historyEntry)
+        .setParent(domainHistory)
         .build();
   }
 
   private Recurring createAutorenewBillingEvent(
-      HistoryEntry historyEntry, DateTime registrationExpirationTime) {
+      DomainHistory domainHistory, DateTime registrationExpirationTime) {
     return new BillingEvent.Recurring.Builder()
         .setReason(Reason.RENEW)
         .setFlags(ImmutableSet.of(Flag.AUTO_RENEW))
@@ -553,7 +550,7 @@ public class DomainCreateFlow implements TransactionalFlow {
         .setClientId(clientId)
         .setEventTime(registrationExpirationTime)
         .setRecurrenceEndTime(END_OF_TIME)
-        .setParent(historyEntry)
+        .setParent(domainHistory)
         .build();
   }
 

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

@@ -208,13 +208,21 @@ public final class DomainDeleteFlow implements TransactionalFlow {
     // Enqueue the deletion poll message if the delete is asynchronous or if requested by a
     // superuser (i.e. the registrar didn't request this delete and thus should be notified even if
     // it is synchronous).
-    if (!durationUntilDelete.equals(Duration.ZERO) || isSuperuser) {
+    if (durationUntilDelete.isLongerThan(Duration.ZERO) || isSuperuser) {
       PollMessage.OneTime deletePollMessage =
           createDeletePollMessage(existingDomain, historyEntry, deletionTime);
       entitiesToSave.add(deletePollMessage);
       builder.setDeletePollMessage(deletePollMessage.createVKey());
     }
 
+    // Send a second poll message immediately if the domain is being deleted asynchronously by a
+    // registrar other than the sponsoring registrar (which will necessarily be a superuser).
+    if (durationUntilDelete.isLongerThan(Duration.ZERO)
+        && !clientId.equals(existingDomain.getPersistedCurrentSponsorClientId())) {
+      entitiesToSave.add(
+          createImmediateDeletePollMessage(existingDomain, historyEntry, now, deletionTime));
+    }
+
     // Cancel any grace periods that were still active, and set the expiration time accordingly.
     DateTime newExpirationTime = existingDomain.getRegistrationExpirationTime();
     for (GracePeriod gracePeriod : existingDomain.getGracePeriods()) {
@@ -254,7 +262,6 @@ public final class DomainDeleteFlow implements TransactionalFlow {
             .setHistoryEntry(historyEntry)
             .setEntityChanges(EntityChanges.newBuilder().setSaves(entitiesToSave.build()).build())
             .build());
-    persistEntityChanges(entityChanges);
     BeforeResponseReturnData responseData =
         flowCustomLogic.beforeResponse(
             BeforeResponseParameters.newBuilder()
@@ -264,6 +271,7 @@ public final class DomainDeleteFlow implements TransactionalFlow {
                         : SUCCESS)
                 .setResponseExtensions(getResponseExtensions(existingDomain, now))
                 .build());
+    persistEntityChanges(entityChanges);
     return responseBuilder
         .setResultFromCode(responseData.resultCode())
         .setExtensions(responseData.responseExtensions())
@@ -346,6 +354,19 @@ public final class DomainDeleteFlow implements TransactionalFlow {
         .build();
   }
 
+  private PollMessage.OneTime createImmediateDeletePollMessage(
+      DomainBase existingDomain, HistoryEntry historyEntry, DateTime now, DateTime deletionTime) {
+    return new PollMessage.OneTime.Builder()
+        .setClientId(existingDomain.getPersistedCurrentSponsorClientId())
+        .setEventTime(now)
+        .setParent(historyEntry)
+        .setMsg(
+            String.format(
+                "Domain %s was deleted by registry administrator with final deletion effective: %s",
+                existingDomain.getDomainName(), deletionTime))
+        .build();
+  }
+
   @Nullable
   private ImmutableList<FeeTransformResponseExtension> getResponseExtensions(
       DomainBase existingDomain, DateTime now) {

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

@@ -129,7 +129,7 @@ import google.registry.model.registry.label.ReservedList;
 import google.registry.model.reporting.DomainTransactionRecord;
 import google.registry.model.reporting.DomainTransactionRecord.TransactionReportField;
 import google.registry.model.reporting.HistoryEntry;
-import google.registry.model.tmch.ClaimsListShard;
+import google.registry.model.tmch.ClaimsListDualDatabaseDao;
 import google.registry.persistence.VKey;
 import google.registry.tldconfig.idn.IdnLabelValidator;
 import google.registry.util.Idn;
@@ -994,7 +994,7 @@ public class DomainFlowUtils {
       InternetDomainName domainName, boolean hasSignedMarks, boolean hasClaimsNotice)
       throws EppException {
     boolean isInClaimsList =
-        ClaimsListShard.get().getClaimKey(domainName.parts().get(0)).isPresent();
+        ClaimsListDualDatabaseDao.get().getClaimKey(domainName.parts().get(0)).isPresent();
     if (hasClaimsNotice && !isInClaimsList) {
       throw new UnexpectedClaimsNoticeException(domainName.toString());
     }
@@ -1093,8 +1093,7 @@ public class DomainFlowUtils {
           .list();
     } else {
       return jpaTm()
-          .getEntityManager()
-          .createQuery(
+          .query(
               "FROM DomainHistory WHERE modificationTime >= :beginning "
                   + "ORDER BY modificationTime ASC",
               DomainHistory.class)

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

@@ -22,6 +22,7 @@ import static google.registry.flows.domain.DomainFlowUtils.handleFeeRequest;
 import static google.registry.flows.domain.DomainFlowUtils.loadForeignKeyedDesignatedContacts;
 import static google.registry.model.EppResourceUtils.loadByForeignKey;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
@@ -100,9 +101,11 @@ public final class DomainInfoFlow implements Flow {
     verifyOptionalAuthInfo(authInfo, domain);
     flowCustomLogic.afterValidation(
         AfterValidationParameters.newBuilder().setDomain(domain).build());
-    // Prefetch all referenced resources. Calling values() blocks until loading is done.
-    tm().loadByKeys(domain.getNameservers());
-    tm().loadByKeys(domain.getReferencedContacts());
+    // In ofy, refetch all referenced resources.
+    if (tm().isOfy()) {
+      tm().loadByKeys(domain.getNameservers());
+      tm().loadByKeys(domain.getReferencedContacts());
+    }
     // Registrars can only see a few fields on unauthorized domains.
     // This is a policy decision that is left up to us by the rfcs.
     DomainInfoData.Builder infoBuilder =
@@ -110,14 +113,16 @@ public final class DomainInfoFlow implements Flow {
             .setFullyQualifiedDomainName(domain.getDomainName())
             .setRepoId(domain.getRepoId())
             .setCurrentSponsorClientId(domain.getCurrentSponsorClientId())
-            .setRegistrant(tm().loadByKey(domain.getRegistrant()).getContactId());
+            .setRegistrant(
+                transactIfJpaTm(() -> tm().loadByKey(domain.getRegistrant())).getContactId());
     // If authInfo is non-null, then the caller is authorized to see the full information since we
     // will have already verified the authInfo is valid.
     if (clientId.equals(domain.getCurrentSponsorClientId()) || authInfo.isPresent()) {
       HostsRequest hostsRequest = ((Info) resourceCommand).getHostsRequest();
       infoBuilder
           .setStatusValues(domain.getStatusValues())
-          .setContacts(loadForeignKeyedDesignatedContacts(domain.getContacts()))
+          .setContacts(
+              transactIfJpaTm(() -> loadForeignKeyedDesignatedContacts(domain.getContacts())))
           .setNameservers(hostsRequest.requestDelegated() ? domain.loadNameserverHostNames() : null)
           .setSubordinateHosts(
               hostsRequest.requestSubordinate() ? domain.getSubordinateHosts() : null)

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

@@ -233,6 +233,9 @@ public final class DomainRestoreRequestFlow implements TransactionalFlow  {
         .setDeletePollMessage(null)
         .setAutorenewBillingEvent(autorenewEvent.createVKey())
         .setAutorenewPollMessage(autorenewPollMessage.createVKey())
+        // Clear the autorenew end time so if it had expired but is now explicitly being restored,
+        // it won't immediately be deleted again.
+        .setAutorenewEndTime(Optional.empty())
         .setLastEppUpdateTime(now)
         .setLastEppUpdateClientId(clientId)
         .build();

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

@@ -27,13 +27,13 @@ import static google.registry.flows.domain.DomainFlowUtils.updateAutorenewRecurr
 import static google.registry.flows.domain.DomainTransferUtils.createGainingTransferPollMessage;
 import static google.registry.flows.domain.DomainTransferUtils.createTransferResponse;
 import static google.registry.model.ResourceTransferUtils.approvePendingTransfer;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.reporting.DomainTransactionRecord.TransactionReportField.TRANSFER_SUCCESSFUL;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.pricing.PricingEngineProxy.getDomainRenewCost;
 import static google.registry.util.CollectionUtils.union;
 import static google.registry.util.DateTimeUtils.END_OF_TIME;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
 import google.registry.flows.EppException;
@@ -132,6 +132,8 @@ public final class DomainTransferApproveFlow implements TransactionalFlow {
                     .setBillingTime(now.plus(Registry.get(tld).getTransferGracePeriodLength()))
                     .setParent(historyEntry)
                     .build());
+    ImmutableList.Builder<ImmutableObject> entitiesToSave = new ImmutableList.Builder<>();
+    entitiesToSave.add(historyEntry);
     // If we are within an autorenew grace period, cancel the autorenew billing event and don't
     // increase the registration time, since the transfer subsumes the autorenew's extra year.
     GracePeriod autorenewGrace =
@@ -143,7 +145,7 @@ public final class DomainTransferApproveFlow implements TransactionalFlow {
       // then the gaining registrar is not charged for the one year renewal and the losing registrar
       // still needs to be charged for the auto-renew.
       if (billingEvent.isPresent()) {
-        ofy().save().entity(
+        entitiesToSave.add(
             BillingEvent.Cancellation.forGracePeriod(autorenewGrace, historyEntry, targetId));
       }
     }
@@ -190,31 +192,26 @@ public final class DomainTransferApproveFlow implements TransactionalFlow {
             .setAutorenewPollMessage(gainingClientAutorenewPollMessage.createVKey())
             // Remove all the old grace periods and add a new one for the transfer.
             .setGracePeriods(
-                billingEvent.isPresent()
-                    ? ImmutableSet.of(
-                        GracePeriod.forBillingEvent(
-                            GracePeriodStatus.TRANSFER,
-                            existingDomain.getRepoId(),
-                            billingEvent.get()))
-                    : ImmutableSet.of())
+                billingEvent
+                    .map(
+                        oneTime ->
+                            ImmutableSet.of(
+                                GracePeriod.forBillingEvent(
+                                    GracePeriodStatus.TRANSFER,
+                                    existingDomain.getRepoId(),
+                                    oneTime)))
+                    .orElseGet(ImmutableSet::of))
             .setLastEppUpdateTime(now)
             .setLastEppUpdateClientId(clientId)
             .build();
     // Create a poll message for the gaining client.
-    PollMessage gainingClientPollMessage = createGainingTransferPollMessage(
-        targetId,
-        newDomain.getTransferData(),
-        newExpirationTime,
-        historyEntry);
-    ImmutableSet.Builder<ImmutableObject> entitiesToSave = new ImmutableSet.Builder<>();
-    entitiesToSave.add(
-        newDomain,
-        historyEntry,
-        autorenewEvent,
-        gainingClientPollMessage,
-        gainingClientAutorenewPollMessage);
+    PollMessage gainingClientPollMessage =
+        createGainingTransferPollMessage(
+            targetId, newDomain.getTransferData(), newExpirationTime, historyEntry);
     billingEvent.ifPresent(entitiesToSave::add);
-    ofy().save().entities(entitiesToSave.build());
+    entitiesToSave.add(
+        autorenewEvent, gainingClientPollMessage, gainingClientAutorenewPollMessage, newDomain);
+    tm().putAll(entitiesToSave.build());
     // Delete the billing event and poll messages that were written in case the transfer would have
     // been implicitly server approved.
     tm().delete(existingDomain.getTransferData().getServerApproveEntities());

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

@@ -25,7 +25,6 @@ import static google.registry.flows.domain.DomainFlowUtils.updateAutorenewRecurr
 import static google.registry.flows.domain.DomainTransferUtils.createLosingTransferPollMessage;
 import static google.registry.flows.domain.DomainTransferUtils.createTransferResponse;
 import static google.registry.model.ResourceTransferUtils.denyPendingTransfer;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.reporting.DomainTransactionRecord.TransactionReportField.TRANSFER_SUCCESSFUL;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.DateTimeUtils.END_OF_TIME;
@@ -39,7 +38,6 @@ import google.registry.flows.FlowModule.Superuser;
 import google.registry.flows.FlowModule.TargetId;
 import google.registry.flows.TransactionalFlow;
 import google.registry.flows.annotations.ReportingSpec;
-import google.registry.model.ImmutableObject;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.metadata.MetadataExtension;
 import google.registry.model.eppcommon.AuthInfo;
@@ -100,11 +98,11 @@ public final class DomainTransferCancelFlow implements TransactionalFlow {
     HistoryEntry historyEntry = buildHistoryEntry(existingDomain, registry, now);
     DomainBase newDomain =
         denyPendingTransfer(existingDomain, TransferStatus.CLIENT_CANCELLED, now, clientId);
-    ofy().save().<ImmutableObject>entities(
-        newDomain,
-        historyEntry,
-        createLosingTransferPollMessage(
-            targetId, newDomain.getTransferData(), null, historyEntry));
+    tm().putAll(
+            newDomain,
+            historyEntry,
+            createLosingTransferPollMessage(
+                targetId, newDomain.getTransferData(), null, historyEntry));
     // Reopen the autorenew event and poll message that we closed for the implicit transfer. This
     // may recreate the autorenew poll message if it was deleted when the transfer request was made.
     updateAutorenewRecurrenceEndTime(existingDomain, END_OF_TIME);

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

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

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

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

core/src/main/java/google/registry/flows/poll/PollAckFlow.java

@@ -16,15 +16,13 @@ package google.registry.flows.poll;
 
 import static google.registry.flows.FlowUtils.validateClientIsLoggedIn;
 import static google.registry.flows.poll.PollFlowUtils.ackPollMessage;
-import static google.registry.flows.poll.PollFlowUtils.getPollMessagesQuery;
+import static google.registry.flows.poll.PollFlowUtils.getPollMessageCount;
 import static google.registry.model.eppoutput.Result.Code.SUCCESS_WITH_NO_MESSAGES;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.poll.PollMessageExternalKeyConverter.makePollMessageExternalId;
 import static google.registry.model.poll.PollMessageExternalKeyConverter.parsePollMessageExternalId;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.DateTimeUtils.isBeforeOrAt;
 
-import com.googlecode.objectify.Key;
 import google.registry.flows.EppException;
 import google.registry.flows.EppException.AuthorizationErrorException;
 import google.registry.flows.EppException.ObjectDoesNotExistException;
@@ -39,6 +37,8 @@ import google.registry.model.poll.MessageQueueInfo;
 import google.registry.model.poll.PollMessage;
 import google.registry.model.poll.PollMessageExternalKeyConverter;
 import google.registry.model.poll.PollMessageExternalKeyConverter.PollMessageExternalKeyParseException;
+import google.registry.persistence.VKey;
+import java.util.Optional;
 import javax.inject.Inject;
 import org.joda.time.DateTime;
 
@@ -71,7 +71,7 @@ public class PollAckFlow implements TransactionalFlow {
       throw new MissingMessageIdException();
     }
 
-    Key<PollMessage> pollMessageKey;
+    VKey<PollMessage> pollMessageKey;
     // Try parsing the messageId, and throw an exception if it's invalid.
     try {
       pollMessageKey = parsePollMessageExternalId(messageId);
@@ -84,12 +84,13 @@ public class PollAckFlow implements TransactionalFlow {
     // Load the message to be acked. If a message is queued to be delivered in the future, we treat
     // it as if it doesn't exist yet. Same for if the message ID year isn't the same as the actual
     // poll message's event time (that means they're passing in an old already-acked ID).
-    PollMessage pollMessage = ofy().load().key(pollMessageKey).now();
-    if (pollMessage == null
-        || !isBeforeOrAt(pollMessage.getEventTime(), now)
-        || !makePollMessageExternalId(pollMessage).equals(messageId)) {
+    Optional<PollMessage> maybePollMessage = tm().loadByKeyIfPresent(pollMessageKey);
+    if (!maybePollMessage.isPresent()
+        || !isBeforeOrAt(maybePollMessage.get().getEventTime(), now)
+        || !makePollMessageExternalId(maybePollMessage.get()).equals(messageId)) {
       throw new MessageDoesNotExistException(messageId);
     }
+    PollMessage pollMessage = maybePollMessage.get();
 
     // Make sure this client is authorized to ack this message. It could be that the message is
     // supposed to go to a different registrar.
@@ -106,8 +107,11 @@ public class PollAckFlow implements TransactionalFlow {
     // acked, then we return a special status code indicating that. Note that the query will
     // include the message being acked.
 
-    int messageCount = tm().doTransactionless(() -> getPollMessagesQuery(clientId, now).count());
-    if (!includeAckedMessageInCount) {
+    int messageCount = tm().doTransactionless(() -> getPollMessageCount(clientId, now));
+    // Within the same transaction, Datastore will not reflect the updated count (potentially
+    // reduced by one thanks to the acked poll message). SQL will, however, so we shouldn't reduce
+    // the count in the SQL case.
+    if (!includeAckedMessageInCount && tm().isOfy()) {
       messageCount--;
     }
     if (messageCount <= 0) {

core/src/main/java/google/registry/flows/poll/PollFlowUtils.java

@@ -16,25 +16,56 @@ package google.registry.flows.poll;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.DateTimeUtils.isBeforeOrAt;
 
 import com.googlecode.objectify.cmd.Query;
 import google.registry.model.poll.PollMessage;
+import java.util.Optional;
 import org.joda.time.DateTime;
 
 /** Static utility functions for poll flows. */
 public final class PollFlowUtils {
 
-  private PollFlowUtils() {}
+  public static final String SQL_POLL_MESSAGE_QUERY =
+      "FROM PollMessage WHERE clientId = :registrarId AND eventTime <= :now ORDER BY eventTime ASC";
+  private static final String SQL_POLL_MESSAGE_COUNT_QUERY =
+      "SELECT COUNT(*) FROM PollMessage WHERE clientId = :registrarId AND eventTime <= :now";
 
-  /** Returns a query for poll messages for the logged in registrar which are not in the future. */
-  public static Query<PollMessage> getPollMessagesQuery(String clientId, DateTime now) {
-    return ofy().load()
-        .type(PollMessage.class)
-        .filter("clientId", clientId)
-        .filter("eventTime <=", now.toDate())
-        .order("eventTime");
+  /** Returns the number of poll messages for the given registrar that are not in the future. */
+  public static int getPollMessageCount(String registrarId, DateTime now) {
+    if (tm().isOfy()) {
+      return datastorePollMessageQuery(registrarId, now).count();
+    } else {
+      return jpaTm()
+          .transact(
+              () ->
+                  jpaTm()
+                      .query(SQL_POLL_MESSAGE_COUNT_QUERY, Long.class)
+                      .setParameter("registrarId", registrarId)
+                      .setParameter("now", now)
+                      .getSingleResult()
+                      .intValue());
+    }
+  }
+
+  /** Returns the first (by event time) poll message not in the future for this registrar. */
+  public static Optional<PollMessage> getFirstPollMessage(String registrarId, DateTime now) {
+    if (tm().isOfy()) {
+      return Optional.ofNullable(datastorePollMessageQuery(registrarId, now).first().now());
+    } else {
+      return jpaTm()
+          .transact(
+              () ->
+                  jpaTm()
+                      .query(SQL_POLL_MESSAGE_QUERY, PollMessage.class)
+                      .setParameter("registrarId", registrarId)
+                      .setParameter("now", now)
+                      .setMaxResults(1)
+                      .getResultStream()
+                      .findFirst());
+    }
   }
 
   /**
@@ -74,4 +105,16 @@ public final class PollFlowUtils {
     }
     return includeAckedMessageInCount;
   }
+
+  /** A Datastore query for poll messages from the given registrar that are not in the future. */
+  public static Query<PollMessage> datastorePollMessageQuery(String registrarId, DateTime now) {
+    return ofy()
+        .load()
+        .type(PollMessage.class)
+        .filter("clientId", registrarId)
+        .filter("eventTime <=", now.toDate())
+        .order("eventTime");
+  }
+
+  private PollFlowUtils() {}
 }

core/src/main/java/google/registry/flows/poll/PollRequestFlow.java

@@ -15,7 +15,8 @@
 package google.registry.flows.poll;
 
 import static google.registry.flows.FlowUtils.validateClientIsLoggedIn;
-import static google.registry.flows.poll.PollFlowUtils.getPollMessagesQuery;
+import static google.registry.flows.poll.PollFlowUtils.getFirstPollMessage;
+import static google.registry.flows.poll.PollFlowUtils.getPollMessageCount;
 import static google.registry.model.eppoutput.Result.Code.SUCCESS_WITH_ACK_MESSAGE;
 import static google.registry.model.eppoutput.Result.Code.SUCCESS_WITH_NO_MESSAGES;
 import static google.registry.model.poll.PollMessageExternalKeyConverter.makePollMessageExternalId;
@@ -31,6 +32,7 @@ import google.registry.model.poll.MessageQueueInfo;
 import google.registry.model.poll.PollMessage;
 import google.registry.model.poll.PollMessageExternalKeyConverter;
 import google.registry.util.Clock;
+import java.util.Optional;
 import javax.inject.Inject;
 import org.joda.time.DateTime;
 
@@ -63,18 +65,20 @@ public class PollRequestFlow implements Flow {
     }
     // Return the oldest message from the queue.
     DateTime now = clock.nowUtc();
-    PollMessage pollMessage = getPollMessagesQuery(clientId, now).first().now();
-    if (pollMessage == null) {
+    Optional<PollMessage> maybePollMessage = getFirstPollMessage(clientId, now);
+    if (!maybePollMessage.isPresent()) {
       return responseBuilder.setResultFromCode(SUCCESS_WITH_NO_MESSAGES).build();
     }
+    PollMessage pollMessage = maybePollMessage.get();
     return responseBuilder
         .setResultFromCode(SUCCESS_WITH_ACK_MESSAGE)
-        .setMessageQueueInfo(new MessageQueueInfo.Builder()
-            .setQueueDate(pollMessage.getEventTime())
-            .setMsg(pollMessage.getMsg())
-            .setQueueLength(getPollMessagesQuery(clientId, now).count())
-            .setMessageId(makePollMessageExternalId(pollMessage))
-            .build())
+        .setMessageQueueInfo(
+            new MessageQueueInfo.Builder()
+                .setQueueDate(pollMessage.getEventTime())
+                .setMsg(pollMessage.getMsg())
+                .setQueueLength(getPollMessageCount(clientId, now))
+                .setMessageId(makePollMessageExternalId(pollMessage))
+                .build())
         .setMultipleResData(pollMessage.getResponseData())
         .build();
   }

core/src/main/java/google/registry/keyring/api/InMemoryKeyring.java

@@ -39,8 +39,6 @@ public final class InMemoryKeyring implements Keyring {
   private final String marksdbLordnPassword;
   private final String marksdbSmdrlLoginAndPassword;
   private final String jsonCredential;
-  private final String cloudSqlPassword;
-  private final String toolsCloudSqlPassword;
 
   public InMemoryKeyring(
       PGPKeyPair rdeStagingKey,
@@ -83,8 +81,6 @@ public final class InMemoryKeyring implements Keyring {
     this.marksdbSmdrlLoginAndPassword =
         checkNotNull(marksdbSmdrlLoginAndPassword, "marksdbSmdrlLoginAndPassword");
     this.jsonCredential = checkNotNull(jsonCredential, "jsonCredential");
-    this.cloudSqlPassword = checkNotNull(cloudSqlPassword, "cloudSqlPassword");
-    this.toolsCloudSqlPassword = checkNotNull(toolsCloudSqlPassword, "toolsCloudSqlPassword");
   }
 
   @Override
@@ -157,16 +153,6 @@ public final class InMemoryKeyring implements Keyring {
     return jsonCredential;
   }
 
-  @Override
-  public String getCloudSqlPassword() {
-    return cloudSqlPassword;
-  }
-
-  @Override
-  public String getToolsCloudSqlPassword() {
-    return toolsCloudSqlPassword;
-  }
-
   /** Does nothing. */
   @Override
   public void close() {}

core/src/main/java/google/registry/keyring/api/KeyModule.java

@@ -36,18 +36,6 @@ public final class KeyModule {
     String value();
   }
 
-  @Provides
-  @Key("cloudSqlPassword")
-  static String providesCloudSqlPassword(Keyring keyring) {
-    return keyring.getCloudSqlPassword();
-  }
-
-  @Provides
-  @Key("toolsCloudSqlPassword")
-  static String providesToolsCloudSqlPassword(Keyring keyring) {
-    return keyring.getToolsCloudSqlPassword();
-  }
-
   @Provides
   @Key("brdaReceiverKey")
   static PGPPublicKey provideBrdaReceiverKey(Keyring keyring) {

core/src/main/java/google/registry/keyring/api/Keyring.java

@@ -28,12 +28,6 @@ import org.bouncycastle.openpgp.PGPPublicKey;
 @ThreadSafe
 public interface Keyring extends AutoCloseable {
 
-  /** Returns the password which is used by App Engine to connect to the Cloud SQL database. */
-  String getCloudSqlPassword();
-
-  /** Returns the password which is used by nomulus tool to connect to the Cloud SQL database. */
-  String getToolsCloudSqlPassword();
-
   /**
    * Returns the key which should be used to sign RDE deposits being uploaded to a third-party.
    *

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

@@ -19,14 +19,25 @@ import static com.google.common.base.CaseFormat.UPPER_UNDERSCORE;
 import static com.google.common.base.Preconditions.checkState;
 import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
 import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Streams;
+import com.google.common.flogger.FluentLogger;
 import com.googlecode.objectify.Key;
 import google.registry.config.RegistryConfig.Config;
 import google.registry.keyring.api.KeySerializer;
 import google.registry.keyring.api.Keyring;
 import google.registry.keyring.api.KeyringException;
 import google.registry.model.server.KmsSecret;
+import google.registry.model.server.KmsSecretRevision;
+import google.registry.model.server.KmsSecretRevisionSqlDao;
+import google.registry.privileges.secretmanager.KeyringSecretStore;
 import java.io.IOException;
+import java.util.Arrays;
+import java.util.Optional;
+import java.util.stream.Stream;
 import javax.inject.Inject;
 import org.bouncycastle.openpgp.PGPException;
 import org.bouncycastle.openpgp.PGPKeyPair;
@@ -42,6 +53,8 @@ import org.bouncycastle.openpgp.PGPPublicKey;
  */
 public class KmsKeyring implements Keyring {
 
+  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
+
   /** Key labels for private key secrets. */
   enum PrivateKeyLabel {
     BRDA_SIGNING_PRIVATE,
@@ -68,7 +81,6 @@ public class KmsKeyring implements Keyring {
 
   /** Key labels for string secrets. */
   enum StringKeyLabel {
-    CLOUD_SQL_PASSWORD_STRING,
     SAFE_BROWSING_API_KEY,
     ICANN_REPORTING_PASSWORD_STRING,
     JSON_CREDENTIAL_STRING,
@@ -76,8 +88,7 @@ public class KmsKeyring implements Keyring {
     MARKSDB_LORDN_PASSWORD_STRING,
     MARKSDB_SMDRL_LOGIN_STRING,
     RDE_SSH_CLIENT_PRIVATE_STRING,
-    RDE_SSH_CLIENT_PUBLIC_STRING,
-    TOOLS_CLOUD_SQL_PASSWORD_STRING;
+    RDE_SSH_CLIENT_PUBLIC_STRING;
 
     String getLabel() {
       return UPPER_UNDERSCORE.to(LOWER_HYPHEN, name());
@@ -85,20 +96,13 @@ public class KmsKeyring implements Keyring {
   }
 
   private final KmsConnection kmsConnection;
+  private final KeyringSecretStore secretStore;
 
   @Inject
-  KmsKeyring(@Config("defaultKmsConnection") KmsConnection kmsConnection) {
+  KmsKeyring(
+      @Config("defaultKmsConnection") KmsConnection kmsConnection, KeyringSecretStore secretStore) {
     this.kmsConnection = kmsConnection;
-  }
-
-  @Override
-  public String getCloudSqlPassword() {
-    return getString(StringKeyLabel.CLOUD_SQL_PASSWORD_STRING);
-  }
-
-  @Override
-  public String getToolsCloudSqlPassword() {
-    return getString(StringKeyLabel.TOOLS_CLOUD_SQL_PASSWORD_STRING);
+    this.secretStore = secretStore;
   }
 
   @Override
@@ -200,17 +204,83 @@ public class KmsKeyring implements Keyring {
     return getKeyPair(keyLabel).getPrivateKey();
   }
 
-  private byte[] getDecryptedData(String keyName) {
-    KmsSecret secret =
-        ofy().load().key(Key.create(getCrossTldKey(), KmsSecret.class, keyName)).now();
-    checkState(secret != null, "Requested secret '%s' does not exist.", keyName);
-    String encryptedData = ofy().load().key(secret.getLatestRevision()).now().getEncryptedValue();
+  private byte[] getDecryptedDataFromDatastore(String keyName) {
+    String encryptedData;
+    if (tm().isOfy()) {
+      KmsSecret secret =
+          ofy().load().key(Key.create(getCrossTldKey(), KmsSecret.class, keyName)).now();
+      checkState(secret != null, "Requested secret '%s' does not exist.", keyName);
+      encryptedData = ofy().load().key(secret.getLatestRevision()).now().getEncryptedValue();
+    } else {
+      Optional<KmsSecretRevision> revision =
+          tm().transact(() -> KmsSecretRevisionSqlDao.getLatestRevision(keyName));
+      checkState(revision.isPresent(), "Requested secret '%s' does not exist.", keyName);
+      encryptedData = revision.get().getEncryptedValue();
+    }
 
     try {
-      return kmsConnection.decrypt(secret.getName(), encryptedData);
+      return kmsConnection.decrypt(keyName, encryptedData);
     } catch (Exception e) {
       throw new KeyringException(
           String.format("CloudKMS decrypt operation failed for secret %s", keyName), e);
     }
   }
+
+  private byte[] getDataFromSecretStore(String keyName) {
+    try {
+      return secretStore.getSecret(keyName);
+    } catch (Exception e) {
+      return new byte[0];
+    }
+  }
+
+  private byte[] getDecryptedData(String keyName) {
+    byte[] dsData = getDecryptedDataFromDatastore(keyName);
+    byte[] secretStoreData = getDataFromSecretStore(keyName);
+
+    if (Arrays.equals(dsData, secretStoreData)) {
+      logger.atInfo().log("Values for %s in Datastore and Secret Manager match.", keyName);
+      return secretStoreData;
+    }
+    logger.atWarning().log("Values for %s in Datastore and Secret Manager do not match.", keyName);
+    return dsData;
+  }
+
+  /**
+   * Generates the tasks to migrate secrets from Datastore to Secret Manager.
+   *
+   * <p>The keys in the returned {@link ImmutableMap} are the names of the secrets that need
+   * migration. The values in the map are {@link Runnable Runnables} that copy secret data from
+   * Datastore to Secret Manager for their corresponding keys. Only secrets that are absent in
+   * Secret Manager or have inconsistent values are included in the returned map.
+   */
+  public ImmutableMap<String, Runnable> migrationPlan() {
+
+    ImmutableMap.Builder<String, Runnable> tasks = new ImmutableMap.Builder<>();
+
+    ImmutableList<String> labels =
+        Streams.concat(
+                Stream.of(PrivateKeyLabel.values()).map(PrivateKeyLabel::getLabel),
+                Stream.of(PublicKeyLabel.values()).map(PublicKeyLabel::getLabel),
+                Stream.of(StringKeyLabel.values()).map(StringKeyLabel::getLabel))
+            .collect(ImmutableList.toImmutableList());
+
+    for (String keyName : labels) {
+      byte[] dsData;
+      try {
+        dsData = getDecryptedDataFromDatastore(keyName);
+      } catch (IllegalStateException e) {
+        logger.atWarning().log("Cannot load %s from Datastore. Skipping...", keyName);
+        continue;
+      }
+      byte[] secretStoreData = getDataFromSecretStore(keyName);
+      if (Arrays.equals(dsData, secretStoreData)) {
+        logger.atInfo().log("%s is already up to date.\n", keyName);
+        continue;
+      }
+      logger.atInfo().log("%s needs to be migrated.\n", keyName);
+      tasks.put(keyName, () -> secretStore.createOrUpdateSecret(keyName, dsData));
+    }
+    return tasks.build();
+  }
 }

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

@@ -24,7 +24,6 @@ import static google.registry.keyring.kms.KmsKeyring.PublicKeyLabel.BRDA_SIGNING
 import static google.registry.keyring.kms.KmsKeyring.PublicKeyLabel.RDE_RECEIVER_PUBLIC;
 import static google.registry.keyring.kms.KmsKeyring.PublicKeyLabel.RDE_SIGNING_PUBLIC;
 import static google.registry.keyring.kms.KmsKeyring.PublicKeyLabel.RDE_STAGING_PUBLIC;
-import static google.registry.keyring.kms.KmsKeyring.StringKeyLabel.CLOUD_SQL_PASSWORD_STRING;
 import static google.registry.keyring.kms.KmsKeyring.StringKeyLabel.ICANN_REPORTING_PASSWORD_STRING;
 import static google.registry.keyring.kms.KmsKeyring.StringKeyLabel.JSON_CREDENTIAL_STRING;
 import static google.registry.keyring.kms.KmsKeyring.StringKeyLabel.MARKSDB_DNL_LOGIN_STRING;
@@ -33,8 +32,6 @@ import static google.registry.keyring.kms.KmsKeyring.StringKeyLabel.MARKSDB_SMDR
 import static google.registry.keyring.kms.KmsKeyring.StringKeyLabel.RDE_SSH_CLIENT_PRIVATE_STRING;
 import static google.registry.keyring.kms.KmsKeyring.StringKeyLabel.RDE_SSH_CLIENT_PUBLIC_STRING;
 import static google.registry.keyring.kms.KmsKeyring.StringKeyLabel.SAFE_BROWSING_API_KEY;
-import static google.registry.keyring.kms.KmsKeyring.StringKeyLabel.TOOLS_CLOUD_SQL_PASSWORD_STRING;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 
@@ -46,6 +43,7 @@ import google.registry.keyring.kms.KmsKeyring.PublicKeyLabel;
 import google.registry.keyring.kms.KmsKeyring.StringKeyLabel;
 import google.registry.model.server.KmsSecret;
 import google.registry.model.server.KmsSecretRevision;
+import google.registry.privileges.secretmanager.KeyringSecretStore;
 import java.io.IOException;
 import java.util.HashMap;
 import java.util.LinkedHashMap;
@@ -62,20 +60,19 @@ import org.bouncycastle.openpgp.PGPPublicKey;
 public final class KmsUpdater {
 
   private final KmsConnection kmsConnection;
+  private final KeyringSecretStore secretStore;
   private final HashMap<String, byte[]> secretValues;
 
   @Inject
-  public KmsUpdater(@Config("defaultKmsConnection") KmsConnection kmsConnection) {
+  public KmsUpdater(
+      @Config("defaultKmsConnection") KmsConnection kmsConnection, KeyringSecretStore secretStore) {
     this.kmsConnection = kmsConnection;
+    this.secretStore = secretStore;
 
     // Use LinkedHashMap to preserve insertion order on update() to simplify testing and debugging
     this.secretValues = new LinkedHashMap<>();
   }
 
-  public KmsUpdater setCloudSqlPassword(String password) {
-    return setString(password, CLOUD_SQL_PASSWORD_STRING);
-  }
-
   public KmsUpdater setRdeSigningKey(PGPKeyPair keyPair) throws IOException, PGPException {
     return setKeyPair(keyPair, RDE_SIGNING_PRIVATE, RDE_SIGNING_PUBLIC);
   }
@@ -108,10 +105,6 @@ public final class KmsUpdater {
     return setString(apiKey, SAFE_BROWSING_API_KEY);
   }
 
-  public KmsUpdater setToolsCloudSqlPassword(String password) {
-    return setString(password, TOOLS_CLOUD_SQL_PASSWORD_STRING);
-  }
-
   public KmsUpdater setIcannReportingPassword(String password) {
     return setString(password, ICANN_REPORTING_PASSWORD_STRING);
   }
@@ -143,6 +136,19 @@ public final class KmsUpdater {
     checkState(!secretValues.isEmpty(), "At least one Keyring value must be persisted");
 
     persistEncryptedValues(encryptValues(secretValues));
+
+    // Errors when writing to secret store can be thrown to the top, since writes are always
+    // executed by a human user using the UpdateKmsKeyringCommand.
+    try {
+      secretValues
+          .entrySet()
+          .forEach(e -> secretStore.createOrUpdateSecret(e.getKey(), e.getValue()));
+    } catch (RuntimeException e) {
+      throw new RuntimeException(
+          "Failed to persist secrets to Secret Manager. "
+              + "Please check the status of Secret Manager and re-run the command.",
+          e);
+    }
   }
 
   /**
@@ -194,8 +200,7 @@ public final class KmsUpdater {
    */
   private static void persistEncryptedValues(
       final ImmutableMap<String, EncryptResponse> encryptedValues) {
-    tm()
-        .transact(
+    tm().transact(
             () -> {
               for (Map.Entry<String, EncryptResponse> entry : encryptedValues.entrySet()) {
                 String secretName = entry.getKey();
@@ -207,7 +212,7 @@ public final class KmsUpdater {
                         .setKmsCryptoKeyVersionName(revisionData.cryptoKeyVersionName())
                         .setParent(secretName)
                         .build();
-                ofy().save().entities(secretRevision, KmsSecret.create(secretName, secretRevision));
+                tm().putAll(secretRevision, KmsSecret.create(secretName, secretRevision));
               }
             });
   }

core/src/main/java/google/registry/mapreduce/inputs/CommitLogManifestReader.java

@@ -58,16 +58,16 @@ class CommitLogManifestReader
 
   @Override
   public QueryResultIterator<Key<CommitLogManifest>> getQueryIterator(@Nullable Cursor cursor) {
-    return startQueryAt(query(), cursor).keys().iterator();
+    return startQueryAt(createBucketQuery(), cursor).keys().iterator();
   }
 
   @Override
   public int getTotal() {
-    return query().count();
+    return createBucketQuery().count();
   }
 
   /** Query for children of this bucket. */
-  Query<CommitLogManifest> query() {
+  Query<CommitLogManifest> createBucketQuery() {
     Query<CommitLogManifest> query = ofy().load().type(CommitLogManifest.class).ancestor(bucketKey);
     if (olderThan != null) {
       query = query.filterKey(

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

@@ -16,6 +16,8 @@ package google.registry.model;
 
 import javax.persistence.Access;
 import javax.persistence.AccessType;
+import javax.persistence.AttributeOverride;
+import javax.persistence.Column;
 import javax.persistence.MappedSuperclass;
 import javax.xml.bind.annotation.XmlTransient;
 
@@ -40,6 +42,7 @@ public abstract class BackupGroupRoot extends ImmutableObject {
   // Prevents subclasses from unexpectedly accessing as property (e.g., HostResource), which would
   // require an unnecessary non-private setter method.
   @Access(AccessType.FIELD)
+  @AttributeOverride(name = "lastUpdateTime", column = @Column(name = "updateTimestamp"))
   UpdateAutoTimestamp updateTimestamp = UpdateAutoTimestamp.create(null);
 
   /** Get the {@link UpdateAutoTimestamp} for this entity. */

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

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

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

@@ -194,30 +194,6 @@ public final class EppResourceUtils {
     return ForeignKeyIndex.load(clazz, uniqueIds, now).keySet();
   }
 
-  /**
-   * Loads resources that match some filter and that have {@link EppResource#deletionTime} that is
-   * not before "now".
-   *
-   * <p>This is an eventually consistent query.
-   *
-   * @param clazz the resource type to load
-   * @param now the logical time of the check
-   * @param filterDefinition the filter to apply when loading resources
-   * @param filterValue the acceptable value for the filter
-   */
-  public static <T extends EppResource> Iterable<T> queryNotDeleted(
-      Class<T> clazz, DateTime now, String filterDefinition, Object filterValue) {
-    return ofy()
-        .load()
-        .type(clazz)
-        .filter(filterDefinition, filterValue)
-        .filter("deletionTime >", now.toDate())
-        .list()
-        .stream()
-        .map(EppResourceUtils.transformAtTime(now))
-        .collect(toImmutableSet());
-  }
-
   /**
    * Returns a Function that transforms an EppResource to the given DateTime, suitable for use with
    * Iterables.transform() over a collection of EppResources.
@@ -418,8 +394,7 @@ public final class EppResourceUtils {
                 if (isContactKey) {
                   query =
                       jpaTm()
-                          .getEntityManager()
-                          .createQuery(CONTACT_LINKED_DOMAIN_QUERY, String.class)
+                          .query(CONTACT_LINKED_DOMAIN_QUERY, String.class)
                           .setParameter("fkRepoId", key)
                           .setParameter("now", now);
                 } else {

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

@@ -17,6 +17,7 @@ package google.registry.model;
 import static com.google.common.collect.Iterables.transform;
 import static com.google.common.collect.Maps.transformValues;
 import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static java.lang.annotation.ElementType.FIELD;
 import static java.lang.annotation.RetentionPolicy.RUNTIME;
 import static java.util.stream.Collectors.toCollection;
@@ -187,7 +188,10 @@ public abstract class ImmutableObject implements Cloneable {
   /** Helper function to recursively hydrate an ImmutableObject. */
   private static Object hydrate(Object value) {
     if (value instanceof Key) {
-      return hydrate(ofy().load().key((Key<?>) value).now());
+      if (tm().isOfy()) {
+        return hydrate(ofy().load().key((Key<?>) value).now());
+      }
+      return value;
     } else if (value instanceof Map) {
       return transformValues((Map<?, ?>) value, ImmutableObject::hydrate);
     } else if (value instanceof Collection) {

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

@@ -18,7 +18,6 @@ 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.collect.ImmutableList.toImmutableList;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.registry.Registry.TldState.GENERAL_AVAILABILITY;
 import static google.registry.model.registry.Registry.TldState.START_DATE_SUNRISE;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
@@ -41,10 +40,10 @@ import google.registry.model.registry.Registry;
 import google.registry.model.registry.Registry.TldState;
 import google.registry.model.registry.label.PremiumList;
 import google.registry.model.registry.label.PremiumListDualDao;
+import google.registry.persistence.VKey;
 import google.registry.util.CidrAddressBlock;
 import java.util.Collection;
 import java.util.Optional;
-import java.util.Set;
 import java.util.function.Function;
 import java.util.regex.Pattern;
 import org.joda.money.CurrencyUnit;
@@ -225,15 +224,14 @@ public final class OteAccountBuilder {
   }
 
   /**
-   * Persists all the OT&E entities to datastore.
+   * Persists all the OT&E entities to the database.
    *
    * @return map from the new clientIds created to the new TLDs they have access to. Can be used to
    *     go over all the newly created Registrars / Registries / RegistrarContacts if any
    *     post-creation work is needed.
    */
   public ImmutableMap<String, String> buildAndPersist() {
-    // save all the entitiesl in a single transaction
-    tm().transact(this::saveAllEntities);
+    saveAllEntities();
     return clientIdToTld;
   }
 
@@ -246,30 +244,38 @@ public final class OteAccountBuilder {
 
   /** Saves all the OT&amp;E entities we created. */
   private void saveAllEntities() {
-    tm().assertInTransaction();
-
     // use ImmutableObject instead of Registry so that the Key generation doesn't break
-    ImmutableList<ImmutableObject> registries = ImmutableList.of(sunriseTld, gaTld, eapTld);
+    ImmutableList<Registry> registries = ImmutableList.of(sunriseTld, gaTld, eapTld);
     ImmutableList<RegistrarContact> contacts = contactsBuilder.build();
 
-    if (!replaceExisting) {
-      ImmutableList<Key<ImmutableObject>> keys =
-          Streams.concat(registries.stream(), registrars.stream(), contacts.stream())
-              .map(Key::create)
-              .collect(toImmutableList());
-      Set<Key<ImmutableObject>> existingKeys = ofy().load().keys(keys).keySet();
-      checkState(
-          existingKeys.isEmpty(),
-          "Found existing object(s) conflicting with OT&E objects: %s",
-          existingKeys);
-    }
-    // Save the Registries (TLDs) first
-    ofy().save().entities(registries).now();
-    // Now we can set the allowedTlds for the registrars
-    registrars = registrars.stream().map(this::addAllowedTld).collect(toImmutableList());
-    // and we can save the registrars and contacts!
-    ofy().save().entities(registrars);
-    ofy().save().entities(contacts);
+    tm().transact(
+            () -> {
+              if (!replaceExisting) {
+                ImmutableList<VKey<? extends ImmutableObject>> keys =
+                    Streams.concat(
+                            registries.stream()
+                                .map(registry -> Registry.createVKey(registry.getTldStr())),
+                            registrars.stream().map(Registrar::createVKey),
+                            contacts.stream().map(RegistrarContact::createVKey))
+                        .collect(toImmutableList());
+                ImmutableMap<VKey<? extends ImmutableObject>, ImmutableObject> existingObjects =
+                    tm().loadByKeysIfPresent(keys);
+                checkState(
+                    existingObjects.isEmpty(),
+                    "Found existing object(s) conflicting with OT&E objects: %s",
+                    existingObjects.keySet());
+              }
+              // Save the Registries (TLDs) first
+              tm().putAll(registries);
+            });
+    // Now we can set the allowedTlds for the registrars in a new transaction
+    tm().transact(
+            () -> {
+              registrars = registrars.stream().map(this::addAllowedTld).collect(toImmutableList());
+              // and we can save the registrars and contacts!
+              tm().putAll(registrars);
+              tm().putAll(contacts);
+            });
   }
 
   private Registrar addAllowedTld(Registrar registrar) {

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

@@ -17,7 +17,6 @@ package google.registry.model;
 import static com.google.common.base.Predicates.equalTo;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static google.registry.model.eppcommon.EppXmlTransformer.unmarshal;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.util.CollectionUtils.isNullOrEmpty;
 import static google.registry.util.DomainNameUtils.ACE_PREFIX;
 
@@ -28,7 +27,6 @@ import com.google.common.collect.ImmutableCollection;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Multiset;
 import com.googlecode.objectify.Key;
-import com.googlecode.objectify.cmd.Query;
 import google.registry.model.domain.DomainCommand;
 import google.registry.model.domain.fee.FeeCreateCommandExtension;
 import google.registry.model.domain.launch.LaunchCreateExtension;
@@ -39,6 +37,7 @@ import google.registry.model.eppinput.EppInput.ResourceCommandWrapper;
 import google.registry.model.host.HostCommand;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.reporting.HistoryEntry.Type;
+import google.registry.model.reporting.HistoryEntryDao;
 import google.registry.xml.XmlException;
 import java.util.Arrays;
 import java.util.EnumSet;
@@ -196,16 +195,10 @@ public class OteStats {
    * <p>Stops when it notices that all tests have passed.
    */
   private OteStats recordRegistrarHistory(String registrarName) {
-    ImmutableCollection<String> clientIds =
+    ImmutableCollection<String> registrarIds =
         OteAccountBuilder.createClientIdToTldMap(registrarName).keySet();
 
-    Query<HistoryEntry> query =
-        ofy()
-            .load()
-            .type(HistoryEntry.class)
-            .filter("clientId in", clientIds)
-            .order("modificationTime");
-    for (HistoryEntry historyEntry : query) {
+    for (HistoryEntry historyEntry : HistoryEntryDao.loadHistoryObjectsByRegistrars(registrarIds)) {
       try {
         record(historyEntry);
       } catch (XmlException e) {

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

@@ -17,7 +17,6 @@ package google.registry.model;
 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 google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.common.collect.ImmutableList;
@@ -107,14 +106,14 @@ public final class ResourceTransferUtils {
   /** Update the relevant {@link ForeignKeyIndex} to cache the new deletion time. */
   public static <R extends EppResource> void updateForeignKeyIndexDeletionTime(R resource) {
     if (resource instanceof ForeignKeyedEppResource) {
-      ofy().save().entity(ForeignKeyIndex.create(resource, resource.getDeletionTime()));
+      tm().insert(ForeignKeyIndex.create(resource, resource.getDeletionTime()));
     }
   }
 
   /** If there is a transfer out, delete the server-approve entities and enqueue a poll message. */
   public static <R extends EppResource & ResourceWithTransferData>
       void handlePendingTransferOnDelete(
-            R resource, R newResource, DateTime now, HistoryEntry historyEntry) {
+          R resource, R newResource, DateTime now, HistoryEntry historyEntry) {
     if (resource.getStatusValues().contains(StatusValue.PENDING_TRANSFER)) {
       TransferData oldTransferData = resource.getTransferData();
       tm().delete(oldTransferData.getServerApproveEntities());

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

@@ -14,11 +14,22 @@
 
 package google.registry.model;
 
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
 
+import com.googlecode.objectify.annotation.Ignore;
+import com.googlecode.objectify.annotation.OnLoad;
 import google.registry.model.translators.UpdateAutoTimestampTranslatorFactory;
+import google.registry.util.DateTimeUtils;
+import java.time.ZonedDateTime;
 import java.util.Optional;
 import javax.annotation.Nullable;
+import javax.persistence.Column;
+import javax.persistence.Embeddable;
+import javax.persistence.PostLoad;
+import javax.persistence.PrePersist;
+import javax.persistence.PreUpdate;
+import javax.persistence.Transient;
 import org.joda.time.DateTime;
 
 /**
@@ -26,14 +37,45 @@ import org.joda.time.DateTime;
  *
  * @see UpdateAutoTimestampTranslatorFactory
  */
+@Embeddable
 public class UpdateAutoTimestamp extends ImmutableObject {
 
-  // When set to true, database converters/translators should do tha auto update.  When set to
+  // When set to true, database converters/translators should do the auto update.  When set to
   // false, auto update should be suspended (this exists to allow us to preserve the original value
   // during a replay).
   private static ThreadLocal<Boolean> autoUpdateEnabled = ThreadLocal.withInitial(() -> true);
 
-  DateTime timestamp;
+  @Transient DateTime timestamp;
+
+  @Ignore
+  @Column(nullable = false)
+  ZonedDateTime lastUpdateTime;
+
+  // Unfortunately, we cannot use the @UpdateTimestamp annotation on "lastUpdateTime" in this class
+  // because Hibernate does not allow it to be used on @Embeddable classes, see
+  // https://hibernate.atlassian.net/browse/HHH-13235. This is a workaround.
+  @PrePersist
+  @PreUpdate
+  void setTimestamp() {
+    if (autoUpdateEnabled() || lastUpdateTime == null) {
+      timestamp = jpaTm().getTransactionTime();
+      lastUpdateTime = DateTimeUtils.toZonedDateTime(timestamp);
+    }
+  }
+
+  @OnLoad
+  void onLoad() {
+    if (timestamp != null) {
+      lastUpdateTime = DateTimeUtils.toZonedDateTime(timestamp);
+    }
+  }
+
+  @PostLoad
+  void postLoad() {
+    if (lastUpdateTime != null) {
+      timestamp = DateTimeUtils.toJodaDateTime(lastUpdateTime);
+    }
+  }
 
   /** Returns the timestamp, or {@code START_OF_TIME} if it's null. */
   public DateTime getTimestamp() {
@@ -43,6 +85,7 @@ public class UpdateAutoTimestamp extends ImmutableObject {
   public static UpdateAutoTimestamp create(@Nullable DateTime timestamp) {
     UpdateAutoTimestamp instance = new UpdateAutoTimestamp();
     instance.timestamp = timestamp;
+    instance.lastUpdateTime = timestamp == null ? null : DateTimeUtils.toZonedDateTime(timestamp);
     return instance;
   }
 

core/src/main/java/google/registry/model/annotations/InCrossTld.java

@@ -0,0 +1,34 @@
+// Copyright 2021 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.model.annotations;
+
+import com.googlecode.objectify.annotation.Entity;
+import google.registry.model.common.EntityGroupRoot;
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Inherited;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * Annotation for an Objectify {@link Entity} to indicate that it is in the cross-TLD entity group.
+ *
+ * <p>This means that the entity's <code>@Parent</code> field has to have the value of {@link
+ * EntityGroupRoot#getCrossTldKey}.
+ */
+@Retention(RetentionPolicy.RUNTIME)
+@Target({ElementType.TYPE})
+@Inherited
+public @interface InCrossTld {}

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

@@ -18,7 +18,8 @@ import static com.google.common.base.MoreObjects.firstNonNull;
 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 google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 import static google.registry.util.CollectionUtils.forceEmptyToNull;
 import static google.registry.util.CollectionUtils.nullToEmptyImmutableCopy;
 import static google.registry.util.DateTimeUtils.END_OF_TIME;
@@ -772,10 +773,10 @@ public abstract class BillingEvent extends ImmutableObject
         Modification instance = getInstance();
         checkNotNull(instance.reason);
         checkNotNull(instance.eventRef);
-        BillingEvent.OneTime billingEvent = ofy().load().key(instance.eventRef).now();
-        checkArgument(Objects.equals(
-            instance.cost.getCurrencyUnit(),
-            billingEvent.cost.getCurrencyUnit()),
+        BillingEvent.OneTime billingEvent =
+            transactIfJpaTm(() -> tm().loadByKey(VKey.from(instance.eventRef)));
+        checkArgument(
+            Objects.equals(instance.cost.getCurrencyUnit(), billingEvent.cost.getCurrencyUnit()),
             "Referenced billing event is in a different currency");
         return super.build();
       }

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

@@ -20,16 +20,18 @@ import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Id;
 import com.googlecode.objectify.annotation.Parent;
 import google.registry.model.ImmutableObject;
+import google.registry.model.annotations.InCrossTld;
 import javax.persistence.MappedSuperclass;
 import javax.persistence.Transient;
 
 /** A singleton entity in Datastore. */
 @MappedSuperclass
+@InCrossTld
 public abstract class CrossTldSingleton extends ImmutableObject {
 
   public static final long SINGLETON_ID = 1; // There is always exactly one of these.
 
-  @Id @Transient long id = SINGLETON_ID;
+  @Id @javax.persistence.Id long id = SINGLETON_ID;
 
   @Transient @Parent Key<EntityGroupRoot> parent = getCrossTldKey();
 }

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

@@ -24,12 +24,25 @@ import com.google.common.base.Splitter;
 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.OnLoad;
 import com.googlecode.objectify.annotation.Parent;
 import google.registry.model.ImmutableObject;
 import google.registry.model.UpdateAutoTimestamp;
+import google.registry.model.annotations.InCrossTld;
+import google.registry.model.common.Cursor.CursorId;
 import google.registry.model.registry.Registry;
-import google.registry.schema.replay.DatastoreOnlyEntity;
+import google.registry.persistence.VKey;
+import google.registry.schema.replay.DatastoreAndSqlEntity;
+import java.io.Serializable;
 import java.util.List;
+import java.util.Optional;
+import javax.persistence.Column;
+import javax.persistence.EnumType;
+import javax.persistence.Enumerated;
+import javax.persistence.IdClass;
+import javax.persistence.PostLoad;
+import javax.persistence.Transient;
 import org.joda.time.DateTime;
 
 /**
@@ -38,7 +51,13 @@ import org.joda.time.DateTime;
  * scoped on {@link EntityGroupRoot}.
  */
 @Entity
-public class Cursor extends ImmutableObject implements DatastoreOnlyEntity {
+@javax.persistence.Entity
+@IdClass(CursorId.class)
+@InCrossTld
+public class Cursor extends ImmutableObject implements DatastoreAndSqlEntity {
+
+  /** The scope of a global cursor. A global cursor is a cursor that is not specific to one tld. */
+  public static final String GLOBAL = "GLOBAL";
 
   /** The types of cursors, used as the string id field for each cursor in Datastore. */
   public enum CursorType {
@@ -104,9 +123,9 @@ public class Cursor extends ImmutableObject implements DatastoreOnlyEntity {
 
     /**
      * If there are multiple cursors for a given cursor type, a cursor must also have a scope
-     * defined (distinct from a parent, which is always the EntityGroupRoot key). For instance,
-     * for a cursor that is defined at the registry level, the scope type will be Registry.class.
-     * For a cursor (theoretically) defined for each EPP resource, the scope type will be
+     * defined (distinct from a parent, which is always the EntityGroupRoot key). For instance, for
+     * a cursor that is defined at the registry level, the scope type will be Registry.class. For a
+     * cursor (theoretically) defined for each EPP resource, the scope type will be
      * EppResource.class. For a global cursor, i.e. one that applies per environment, this will be
      * {@link EntityGroupRoot}.
      */
@@ -115,24 +134,73 @@ public class Cursor extends ImmutableObject implements DatastoreOnlyEntity {
     }
   }
 
-  @Parent
-  Key<EntityGroupRoot> parent = getCrossTldKey();
+  @Transient @Parent Key<EntityGroupRoot> parent = getCrossTldKey();
 
-  @Id
-  String id;
+  @Transient @Id String id;
 
+  @Ignore
+  @Enumerated(EnumType.STRING)
+  @Column(nullable = false)
+  @javax.persistence.Id
+  CursorType type;
+
+  @Ignore
+  @Column(nullable = false)
+  @javax.persistence.Id
+  String scope;
+
+  @Column(nullable = false)
   DateTime cursorTime = START_OF_TIME;
 
   /** An automatically managed timestamp of when this object was last written to Datastore. */
+  @Column(nullable = false)
   UpdateAutoTimestamp lastUpdateTime = UpdateAutoTimestamp.create(null);
 
+  @OnLoad
+  void onLoad() {
+    scope = getScopeFromId(id);
+    type = getTypeFromId(id);
+  }
+
+  @PostLoad
+  void postLoad() {
+    // "Generate" the ID based on the scope and type
+    Key<? extends ImmutableObject> scopeKey =
+        scope.equals(GLOBAL)
+            ? getCrossTldKey()
+            : Key.create(getCrossTldKey(), Registry.class, scope);
+    id = generateId(type, scopeKey);
+  }
+
+  public static VKey<Cursor> createVKey(Key<Cursor> key) {
+    String id = key.getName();
+    return VKey.create(Cursor.class, new CursorId(getTypeFromId(id), getScopeFromId(id)), key);
+  }
+
+  public VKey<Cursor> createVKey() {
+    return createVKey(type, scope);
+  }
+
+  public static VKey<Cursor> createGlobalVKey(CursorType type) {
+    return createVKey(type, GLOBAL);
+  }
+
+  public static VKey<Cursor> createVKey(CursorType type, String scope) {
+    Key<Cursor> key =
+        scope.equals(GLOBAL) ? createGlobalKey(type) : createKey(type, Registry.get(scope));
+    return VKey.create(Cursor.class, new CursorId(type, scope), key);
+  }
+
   public DateTime getLastUpdateTime() {
     return lastUpdateTime.getTimestamp();
   }
 
+  public String getScope() {
+    return scope;
+  }
+
   public CursorType getType() {
-    List<String> id = Splitter.on('_').splitToList(this.id);
-    return CursorType.valueOf(String.join("_", id.subList(1, id.size())));
+    return type;
   }
 
   /**
@@ -142,10 +210,12 @@ public class Cursor extends ImmutableObject implements DatastoreOnlyEntity {
   private static void checkValidCursorTypeForScope(
       CursorType cursorType, Key<? extends ImmutableObject> scope) {
     checkArgument(
-        cursorType.getScopeClass().equals(
-            scope.equals(EntityGroupRoot.getCrossTldKey())
-                ? EntityGroupRoot.class
-                : ofy().factory().getMetadata(scope).getEntityClass()),
+        cursorType
+            .getScopeClass()
+            .equals(
+                scope.equals(getCrossTldKey())
+                    ? EntityGroupRoot.class
+                    : ofy().factory().getMetadata(scope).getEntityClass()),
         "Class required for cursor does not match scope class");
   }
 
@@ -154,6 +224,20 @@ public class Cursor extends ImmutableObject implements DatastoreOnlyEntity {
     return String.format("%s_%s", scope.getString(), cursorType.name());
   }
 
+  private static String getScopeFromId(String id) {
+    List<String> idSplit = Splitter.on('_').splitToList(id);
+    // The "parent" is always the crossTldKey; in order to find the scope (either Registry or
+    // cross-tld-key) we have to parse the part of the ID
+    Key<?> scopeKey = Key.valueOf(idSplit.get(0));
+    return scopeKey.equals(getCrossTldKey()) ? GLOBAL : scopeKey.getName();
+  }
+
+  private static CursorType getTypeFromId(String id) {
+    List<String> idSplit = Splitter.on('_').splitToList(id);
+    // The cursor type is the second part of the ID string
+    return CursorType.valueOf(String.join("_", idSplit.subList(1, idSplit.size())));
+  }
+
   /** Creates a unique key for a given scope and cursor type. */
   public static Key<Cursor> createKey(CursorType cursorType, ImmutableObject scope) {
     Key<? extends ImmutableObject> scopeKey = Key.create(scope);
@@ -166,13 +250,12 @@ public class Cursor extends ImmutableObject implements DatastoreOnlyEntity {
     checkArgument(
         cursorType.getScopeClass().equals(EntityGroupRoot.class),
         "Cursor type is not a global cursor.");
-    return Key.create(
-        getCrossTldKey(), Cursor.class, generateId(cursorType, EntityGroupRoot.getCrossTldKey()));
+    return Key.create(getCrossTldKey(), Cursor.class, generateId(cursorType, getCrossTldKey()));
   }
 
   /** Creates a new global cursor instance. */
   public static Cursor createGlobal(CursorType cursorType, DateTime cursorTime) {
-    return create(cursorType, cursorTime, EntityGroupRoot.getCrossTldKey());
+    return create(cursorType, cursorTime, getCrossTldKey());
   }
 
   /** Creates a new cursor instance with a given {@link Key} scope. */
@@ -184,8 +267,10 @@ public class Cursor extends ImmutableObject implements DatastoreOnlyEntity {
     checkNotNull(cursorType, "Cursor type cannot be null");
     checkValidCursorTypeForScope(cursorType, scope);
     instance.id = generateId(cursorType, scope);
+    instance.type = cursorType;
+    instance.scope = scope.equals(getCrossTldKey()) ? GLOBAL : scope.getName();
     return instance;
-   }
+  }
 
   /** Creates a new cursor instance with a given {@link ImmutableObject} scope. */
   public static Cursor create(CursorType cursorType, DateTime cursorTime, ImmutableObject scope) {
@@ -196,11 +281,24 @@ public class Cursor extends ImmutableObject implements DatastoreOnlyEntity {
   /**
    * Returns the current time for a given cursor, or {@code START_OF_TIME} if the cursor is null.
    */
-  public static DateTime getCursorTimeOrStartOfTime(Cursor cursor) {
-    return cursor != null ? cursor.getCursorTime() : START_OF_TIME;
+  public static DateTime getCursorTimeOrStartOfTime(Optional<Cursor> cursor) {
+    return cursor.map(Cursor::getCursorTime).orElse(START_OF_TIME);
   }
 
   public DateTime getCursorTime() {
     return cursorTime;
   }
+
+  static class CursorId extends ImmutableObject implements Serializable {
+
+    public CursorType type;
+    public String scope;
+
+    private CursorId() {}
+
+    public CursorId(CursorType type, String scope) {
+      this.type = type;
+      this.scope = scope;
+    }
+  }
 }

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

@@ -32,11 +32,13 @@ import com.googlecode.objectify.annotation.Mapify;
 import com.googlecode.objectify.annotation.Parent;
 import google.registry.model.ImmutableObject;
 import google.registry.model.UpdateAutoTimestamp;
+import google.registry.model.annotations.InCrossTld;
 import google.registry.model.common.TimedTransitionProperty.TimeMapper;
 import google.registry.model.common.TimedTransitionProperty.TimedTransition;
 import google.registry.model.registry.label.PremiumList;
 import google.registry.model.registry.label.ReservedList;
 import google.registry.model.smd.SignedMarkRevocationList;
+import google.registry.model.tmch.ClaimsListShard;
 import google.registry.persistence.VKey;
 import google.registry.schema.replay.DatastoreOnlyEntity;
 import java.util.Optional;
@@ -45,6 +47,7 @@ import org.joda.time.DateTime;
 
 @Entity
 @Immutable
+@InCrossTld
 public class DatabaseTransitionSchedule extends ImmutableObject implements DatastoreOnlyEntity {
 
   /**
@@ -58,6 +61,8 @@ public class DatabaseTransitionSchedule extends ImmutableObject implements Datas
 
   /** The id of the transition schedule. */
   public enum TransitionId {
+    /** The schedule for migration of {@link ClaimsListShard} entities. */
+    CLAIMS_LIST,
     /** The schedule for the migration of {@link PremiumList} and {@link ReservedList}. */
     DOMAIN_LABEL_LISTS,
     /** The schedule for the migration of the {@link SignedMarkRevocationList} entity. */

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

@@ -16,8 +16,7 @@ package google.registry.model.common;
 
 import static com.google.common.base.Preconditions.checkState;
 import static google.registry.model.ofy.ObjectifyService.allocateId;
-import static google.registry.model.ofy.ObjectifyService.ofy;
-import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
 
 import com.google.appengine.api.users.User;
 import com.google.common.base.Splitter;
@@ -56,13 +55,22 @@ public class GaeUserIdConverter extends ImmutableObject {
     try {
       // Perform these operations in a transactionless context to avoid enlisting in some outer
       // transaction (if any).
-      tm().doTransactionless(() -> ofy().saveWithoutBackup().entity(gaeUserIdConverter).now());
+      ofyTm()
+          .doTransactionless(
+              () -> {
+                ofyTm().putWithoutBackup(gaeUserIdConverter);
+                return null;
+              });
 
       // The read must be done in its own transaction to avoid reading from the session cache.
-      return tm()
-          .transactNew(() -> ofy().load().entity(gaeUserIdConverter).safe().user.getUserId());
+      return ofyTm().transactNew(() -> ofyTm().loadByEntity(gaeUserIdConverter).user.getUserId());
     } finally {
-      tm().doTransactionless(() -> ofy().deleteWithoutBackup().entity(gaeUserIdConverter).now());
+      ofyTm()
+          .doTransactionless(
+              () -> {
+                ofyTm().deleteWithoutBackup(gaeUserIdConverter);
+                return null;
+              });
     }
   }
 }

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

@@ -21,6 +21,7 @@ import static google.registry.model.EppResourceUtils.projectResourceOntoBuilderA
 import com.google.common.collect.ImmutableList;
 import com.googlecode.objectify.annotation.IgnoreSave;
 import com.googlecode.objectify.annotation.Index;
+import com.googlecode.objectify.annotation.OnLoad;
 import com.googlecode.objectify.condition.IfNull;
 import google.registry.model.EppResource;
 import google.registry.model.EppResource.ResourceWithTransferData;
@@ -189,6 +190,17 @@ public class ContactBase extends EppResource implements ResourceWithTransferData
             + " use ContactResource instead");
   }
 
+  @OnLoad
+  void onLoad() {
+    if (voice != null && voice.hasNullFields()) {
+      voice = null;
+    }
+
+    if (fax != null && fax.hasNullFields()) {
+      fax = null;
+    }
+  }
+
   public String getContactId() {
     return contactId;
   }
@@ -325,11 +337,17 @@ public class ContactBase extends EppResource implements ResourceWithTransferData
     }
 
     public B setVoiceNumber(ContactPhoneNumber voiceNumber) {
+      if (voiceNumber != null && voiceNumber.hasNullFields()) {
+        voiceNumber = null;
+      }
       getInstance().voice = voiceNumber;
       return thisCastToDerived();
     }
 
     public B setFaxNumber(ContactPhoneNumber faxNumber) {
+      if (faxNumber != null && faxNumber.hasNullFields()) {
+        faxNumber = null;
+      }
       getInstance().fax = faxNumber;
       return thisCastToDerived();
     }

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

@@ -14,6 +14,8 @@
 
 package google.registry.model.contact;
 
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.EntitySubclass;
 import google.registry.model.ImmutableObject;
@@ -54,6 +56,7 @@ import javax.persistence.PostLoad;
 public class ContactHistory extends HistoryEntry implements SqlEntity {
 
   // Store ContactBase instead of ContactResource so we don't pick up its @Id
+  // Nullable for the sake of pre-Registry-3.0 history objects
   @Nullable ContactBase contactBase;
 
   @Id
@@ -114,6 +117,12 @@ public class ContactHistory extends HistoryEntry implements SqlEntity {
     return Optional.of(asHistoryEntry());
   }
 
+  // Used to fill out the contactBase field during asynchronous replay
+  public static void beforeSqlSave(ContactHistory contactHistory) {
+    contactHistory.contactBase =
+        jpaTm().loadByKey(VKey.createSql(ContactResource.class, contactHistory.getContactRepoId()));
+  }
+
   /** Class to represent the composite primary key of {@link ContactHistory} entity. */
   public static class ContactHistoryId extends ImmutableObject implements Serializable {
 
@@ -185,9 +194,13 @@ public class ContactHistory extends HistoryEntry implements SqlEntity {
       super(instance);
     }
 
-    public Builder setContactBase(ContactBase contactBase) {
+    public Builder setContactBase(@Nullable ContactBase contactBase) {
+      // Nullable for the sake of pre-Registry-3.0 history objects
+      if (contactBase == null) {
+        return this;
+      }
       getInstance().contactBase = contactBase;
-      return this;
+      return super.setParent(contactBase);
     }
 
     public Builder setContactRepoId(String contactRepoId) {

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

@@ -23,8 +23,9 @@ import static com.google.common.collect.Sets.difference;
 import static com.google.common.collect.Sets.intersection;
 import static google.registry.model.EppResourceUtils.projectResourceOntoBuilderAtTime;
 import static google.registry.model.EppResourceUtils.setAutomaticTransferSuccessProperties;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 import static google.registry.util.CollectionUtils.forceEmptyToNull;
 import static google.registry.util.CollectionUtils.nullToEmpty;
 import static google.registry.util.CollectionUtils.nullToEmptyImmutableCopy;
@@ -302,12 +303,6 @@ public class DomainContent extends EppResource
 
   @OnLoad
   void load() {
-    // Back fill with correct END_OF_TIME sentinel value.
-    // TODO(mcilwain): Remove this once back-filling is complete.
-    if (autorenewEndTime == null) {
-      autorenewEndTime = END_OF_TIME;
-    }
-
     // Reconstitute all of the contacts so that they have VKeys.
     allContacts =
         allContacts.stream().map(DesignatedContact::reconstitute).collect(toImmutableSet());
@@ -321,9 +316,6 @@ public class DomainContent extends EppResource
         nullToEmptyImmutableCopy(gracePeriods).stream()
             .map(gracePeriod -> gracePeriod.cloneAfterOfyLoad(getRepoId()))
             .collect(toImmutableSet());
-    // TODO(b/169873747): Remove this method after explicitly re-saving all domain entities.
-    // See also: GradePeriod.onLoad.
-    gracePeriods.forEach(GracePeriod::onLoad);
 
     // Restore history record ids.
     autorenewPollMessageHistoryId = getHistoryId(autorenewPollMessage);
@@ -346,7 +338,7 @@ public class DomainContent extends EppResource
 
   @PostLoad
   @SuppressWarnings("UnusedMethod")
-  private final void postLoad() {
+  private void postLoad() {
     // Reconstitute the contact list.
     ImmutableSet.Builder<DesignatedContact> contactsBuilder = new ImmutableSet.Builder<>();
 
@@ -387,13 +379,11 @@ public class DomainContent extends EppResource
   public static void beforeSqlDelete(VKey<DomainBase> key) {
     // Delete all grace periods associated with the domain.
     jpaTm()
-        .getEntityManager()
-        .createQuery("DELETE FROM GracePeriod WHERE domain_repo_id = :repo_id")
+        .query("DELETE FROM GracePeriod WHERE domain_repo_id = :repo_id")
         .setParameter("repo_id", key.getSqlKey())
         .executeUpdate();
     jpaTm()
-        .getEntityManager()
-        .createQuery("DELETE FROM DelegationSignerData WHERE domain_repo_id = :repo_id")
+        .query("DELETE FROM DelegationSignerData WHERE domain_repo_id = :repo_id")
         .setParameter("repo_id", key.getSqlKey())
         .executeUpdate();
   }
@@ -443,12 +433,7 @@ public class DomainContent extends EppResource
    * purposes of more legible business logic.
    */
   public Optional<DateTime> getAutorenewEndTime() {
-    // TODO(mcilwain): Remove null handling for autorenewEndTime once data migration away from null
-    //                 is complete.
-    return Optional.ofNullable(
-        (autorenewEndTime == null || autorenewEndTime.equals(END_OF_TIME))
-            ? null
-            : autorenewEndTime);
+    return Optional.ofNullable(autorenewEndTime.equals(END_OF_TIME) ? null : autorenewEndTime);
   }
 
   @Override
@@ -686,13 +671,11 @@ public class DomainContent extends EppResource
 
   /** Loads and returns the fully qualified host names of all linked nameservers. */
   public ImmutableSortedSet<String> loadNameserverHostNames() {
-    return ofy()
-        .load()
-        .keys(getNameservers().stream().map(VKey::getOfyKey).collect(toImmutableSet()))
-        .values()
-        .stream()
-        .map(HostResource::getHostName)
-        .collect(toImmutableSortedSet(Ordering.natural()));
+    return transactIfJpaTm(
+        () ->
+            tm().loadByKeys(getNameservers()).values().stream()
+                .map(HostResource::getHostName)
+                .collect(toImmutableSortedSet(Ordering.natural())));
   }
 
   /** A key to the registrant who registered this domain. */

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

@@ -15,6 +15,7 @@
 package google.registry.model.domain;
 
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.util.CollectionUtils.nullToEmptyImmutableCopy;
 
 import com.google.common.collect.ImmutableSet;
@@ -32,6 +33,7 @@ import google.registry.persistence.VKey;
 import google.registry.schema.replay.DatastoreEntity;
 import google.registry.schema.replay.SqlEntity;
 import java.io.Serializable;
+import java.util.HashSet;
 import java.util.Optional;
 import java.util.Set;
 import javax.annotation.Nullable;
@@ -75,6 +77,7 @@ import javax.persistence.Table;
 public class DomainHistory extends HistoryEntry implements SqlEntity {
 
   // Store DomainContent instead of DomainBase so we don't pick up its @Id
+  // Nullable for the sake of pre-Registry-3.0 history objects
   @Nullable DomainContent domainContent;
 
   @Id
@@ -125,7 +128,8 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
         insertable = false,
         updatable = false)
   })
-  Set<DomainDsDataHistory> dsDataHistories = ImmutableSet.of();
+  // HashSet rather than ImmutableSet so that Hibernate can fill them out lazily on request
+  Set<DomainDsDataHistory> dsDataHistories = new HashSet<>();
 
   @Ignore
   @OneToMany(
@@ -144,7 +148,8 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
         insertable = false,
         updatable = false)
   })
-  Set<GracePeriodHistory> gracePeriodHistories = ImmutableSet.of();
+  // HashSet rather than ImmutableSet so that Hibernate can fill them out lazily on request
+  Set<GracePeriodHistory> gracePeriodHistories = new HashSet<>();
 
   @Override
   @Nullable
@@ -263,6 +268,12 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
     return Optional.of(asHistoryEntry());
   }
 
+  // Used to fill out the domainContent field during asynchronous replay
+  public static void beforeSqlSave(DomainHistory domainHistory) {
+    domainHistory.domainContent =
+        jpaTm().loadByKey(VKey.createSql(DomainBase.class, domainHistory.getDomainRepoId()));
+  }
+
   /** Class to represent the composite primary key of {@link DomainHistory} entity. */
   public static class DomainHistoryId extends ImmutableObject implements Serializable {
 
@@ -334,9 +345,13 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
       super(instance);
     }
 
-    public Builder setDomainContent(DomainContent domainContent) {
+    public Builder setDomainContent(@Nullable DomainContent domainContent) {
+      // Nullable for the sake of pre-Registry-3.0 history objects
+      if (domainContent == null) {
+        return this;
+      }
       getInstance().domainContent = domainContent;
-      return this;
+      return super.setParent(domainContent);
     }
 
     public Builder setDomainRepoId(String domainRepoId) {

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

@@ -54,17 +54,6 @@ public class GracePeriod extends GracePeriodBase implements DatastoreAndSqlEntit
     return super.getGracePeriodId();
   }
 
-  // TODO(b/169873747): Remove this method after explicitly re-saving all domain entities.
-  // This method is invoked from DomainContent.load(): Objectify's @OnLoad annotation
-  // apparently does not work on embedded objects inside an entity.
-  // Changing signature to void onLoad(@AlsoLoad("gracePeriodId") Long gracePeriodId)
-  // would not work. Method is not called if gracePeriodId is null.
-  void onLoad() {
-    if (gracePeriodId == null) {
-      gracePeriodId = ObjectifyService.allocateId();
-    }
-  }
-
   private static GracePeriod createInternal(
       GracePeriodStatus type,
       String domainRepoId,

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

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

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

@@ -71,6 +71,11 @@ public class PhoneNumber extends ImmutableObject {
     return phoneNumber + (extension != null ? " x" + extension : "");
   }
 
+  /** Returns true if both fields of the phone number are null. */
+  public boolean hasNullFields() {
+    return phoneNumber == null && extension == null;
+  }
+
   /** A builder for constructing {@link PhoneNumber}. */
   public static class Builder<T extends PhoneNumber> extends Buildable.Builder<T> {
     @Override

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

@@ -14,6 +14,8 @@
 
 package google.registry.model.host;
 
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.EntitySubclass;
 import google.registry.model.ImmutableObject;
@@ -55,6 +57,7 @@ import javax.persistence.PostLoad;
 public class HostHistory extends HistoryEntry implements SqlEntity {
 
   // Store HostBase instead of HostResource so we don't pick up its @Id
+  // Nullable for the sake of pre-Registry-3.0 history objects
   @Nullable HostBase hostBase;
 
   @Id
@@ -115,6 +118,12 @@ public class HostHistory extends HistoryEntry implements SqlEntity {
     return Optional.of(asHistoryEntry());
   }
 
+  // Used to fill out the hostBase field during asynchronous replay
+  public static void beforeSqlSave(HostHistory hostHistory) {
+    hostHistory.hostBase =
+        jpaTm().loadByKey(VKey.createSql(HostResource.class, hostHistory.getHostRepoId()));
+  }
+
   /** Class to represent the composite primary key of {@link HostHistory} entity. */
   public static class HostHistoryId extends ImmutableObject implements Serializable {
 
@@ -186,9 +195,13 @@ public class HostHistory extends HistoryEntry implements SqlEntity {
       super(instance);
     }
 
-    public Builder setHostBase(HostBase hostBase) {
+    public Builder setHostBase(@Nullable HostBase hostBase) {
+      // Nullable for the sake of pre-Registry-3.0 history objects
+      if (hostBase == null) {
+        return this;
+      }
       getInstance().hostBase = hostBase;
-      return this;
+      return super.setParent(hostBase);
     }
 
     public Builder setHostRepoId(String hostRepoId) {

core/src/main/java/google/registry/model/index/ForeignKeyIndex.java

@@ -16,17 +16,20 @@ package google.registry.model.index;
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static google.registry.config.RegistryConfig.getEppResourceCachingDuration;
 import static google.registry.config.RegistryConfig.getEppResourceMaxCachedEntries;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.util.CollectionUtils.entriesToImmutableMap;
 import static google.registry.util.TypeUtils.instantiate;
 
 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 com.google.common.collect.ImmutableBiMap;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
@@ -45,10 +48,11 @@ import google.registry.model.contact.ContactResource;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.host.HostResource;
 import google.registry.persistence.VKey;
+import google.registry.persistence.transaction.CriteriaQueryBuilder;
 import google.registry.schema.replay.DatastoreOnlyEntity;
 import google.registry.util.NonFinalForTesting;
+import java.util.Collection;
 import java.util.Comparator;
-import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.ExecutionException;
@@ -81,10 +85,10 @@ public abstract class ForeignKeyIndex<E extends EppResource> extends BackupGroup
   public static class ForeignKeyHostIndex extends ForeignKeyIndex<HostResource>
       implements DatastoreOnlyEntity {}
 
-  private static final ImmutableMap<
+  private static final ImmutableBiMap<
           Class<? extends EppResource>, Class<? extends ForeignKeyIndex<?>>>
       RESOURCE_CLASS_TO_FKI_CLASS =
-          ImmutableMap.of(
+          ImmutableBiMap.of(
               ContactResource.class, ForeignKeyContactIndex.class,
               DomainBase.class, ForeignKeyDomainIndex.class,
               HostResource.class, ForeignKeyHostIndex.class);
@@ -184,72 +188,100 @@ public abstract class ForeignKeyIndex<E extends EppResource> extends BackupGroup
   }
 
   /**
-   * Load a list of {@link ForeignKeyIndex} instances by class and id strings that are active at or
+   * Load a map of {@link ForeignKeyIndex} instances by class and id strings that are active at or
    * after the specified moment in time.
    *
    * <p>The returned map will omit any keys for which the {@link ForeignKeyIndex} doesn't exist or
    * has been soft deleted.
    */
   public static <E extends EppResource> ImmutableMap<String, ForeignKeyIndex<E>> load(
-      Class<E> clazz, Iterable<String> foreignKeys, final DateTime now) {
+      Class<E> clazz, Collection<String> foreignKeys, final DateTime now) {
+    return loadIndexesFromStore(clazz, foreignKeys, true).entrySet().stream()
+        .filter(e -> now.isBefore(e.getValue().getDeletionTime()))
+        .collect(entriesToImmutableMap());
+  }
+
+  /**
+   * Helper method to load all of the most recent {@link ForeignKeyIndex}es for the given foreign
+   * keys, regardless of whether or not they have been soft-deleted.
+   *
+   * <p>Used by both the cached (w/o deletion check) and the non-cached (with deletion check) calls.
+   *
+   * <p>Note that in the cached case, we wish to run this outside of any transaction because we may
+   * be loading many entities, going over the Datastore limit on the number of enrolled entity
+   * groups per transaction (25). If we require consistency, however, we must use a transaction.
+   *
+   * @param inTransaction whether or not to use an Objectify transaction
+   */
+  private static <E extends EppResource>
+      ImmutableMap<String, ForeignKeyIndex<E>> loadIndexesFromStore(
+          Class<E> clazz, Collection<String> foreignKeys, boolean inTransaction) {
     if (tm().isOfy()) {
-      return ofy().load().type(mapToFkiClass(clazz)).ids(foreignKeys).entrySet().stream()
-          .filter(e -> now.isBefore(e.getValue().deletionTime))
-          .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue));
+      Class<ForeignKeyIndex<E>> fkiClass = mapToFkiClass(clazz);
+      return ImmutableMap.copyOf(
+          inTransaction
+              ? ofy().load().type(fkiClass).ids(foreignKeys)
+              : tm().doTransactionless(() -> ofy().load().type(fkiClass).ids(foreignKeys)));
     } else {
       String property = RESOURCE_CLASS_TO_FKI_PROPERTY.get(clazz);
-      List<E> entities =
+      ImmutableList<ForeignKeyIndex<E>> indexes =
           tm().transact(
-                  () -> {
-                    String entityName =
-                        jpaTm().getEntityManager().getMetamodel().entity(clazz).getName();
-                    return jpaTm()
-                        .getEntityManager()
-                        .createQuery(
-                            String.format(
-                                "FROM %s WHERE %s IN :propertyValue and deletionTime > :now ",
-                                entityName, property),
-                            clazz)
-                        .setParameter("propertyValue", foreignKeys)
-                        .setParameter("now", now)
-                        .getResultList();
-                  });
+                  () ->
+                      jpaTm()
+                          .getEntityManager()
+                          .createQuery(
+                              CriteriaQueryBuilder.create(clazz)
+                                  .whereFieldIsIn(property, foreignKeys)
+                                  .build())
+                          .getResultStream()
+                          .map(e -> ForeignKeyIndex.create(e, e.getDeletionTime()))
+                          .collect(toImmutableList()));
       // We need to find and return the entities with the maximum deletionTime for each foreign key.
-      return Multimaps.index(entities, EppResource::getForeignKey).asMap().entrySet().stream()
+      return Multimaps.index(indexes, ForeignKeyIndex::getForeignKey).asMap().entrySet().stream()
           .map(
               entry ->
                   Maps.immutableEntry(
                       entry.getKey(),
                       entry.getValue().stream()
-                          .max(Comparator.comparing(EppResource::getDeletionTime))
+                          .max(Comparator.comparing(ForeignKeyIndex::getDeletionTime))
                           .get()))
-          .collect(
-              toImmutableMap(
-                  Map.Entry::getKey,
-                  entry -> create(entry.getValue(), entry.getValue().getDeletionTime())));
+          .collect(entriesToImmutableMap());
     }
   }
 
-  static final CacheLoader<Key<ForeignKeyIndex<?>>, Optional<ForeignKeyIndex<?>>> CACHE_LOADER =
-      new CacheLoader<Key<ForeignKeyIndex<?>>, Optional<ForeignKeyIndex<?>>>() {
+  static final CacheLoader<VKey<ForeignKeyIndex<?>>, Optional<ForeignKeyIndex<?>>> CACHE_LOADER =
+      new CacheLoader<VKey<ForeignKeyIndex<?>>, Optional<ForeignKeyIndex<?>>>() {
 
         @Override
-        public Optional<ForeignKeyIndex<?>> load(Key<ForeignKeyIndex<?>> key) {
-          return Optional.ofNullable(tm().doTransactionless(() -> ofy().load().key(key).now()));
+        public Optional<ForeignKeyIndex<?>> load(VKey<ForeignKeyIndex<?>> key) {
+          String foreignKey = key.getSqlKey().toString();
+          return Optional.ofNullable(
+              loadIndexesFromStore(
+                      RESOURCE_CLASS_TO_FKI_CLASS.inverse().get(key.getKind()),
+                      ImmutableSet.of(foreignKey),
+                      false)
+                  .get(foreignKey));
         }
 
         @Override
-        public Map<Key<ForeignKeyIndex<?>>, Optional<ForeignKeyIndex<?>>> loadAll(
-            Iterable<? extends Key<ForeignKeyIndex<?>>> keys) {
-          ImmutableSet<Key<ForeignKeyIndex<?>>> typedKeys = ImmutableSet.copyOf(keys);
-          Map<Key<ForeignKeyIndex<?>>, ForeignKeyIndex<?>> existingFkis =
-              tm().doTransactionless(() -> ofy().load().keys(typedKeys));
+        public Map<VKey<ForeignKeyIndex<?>>, Optional<ForeignKeyIndex<?>>> loadAll(
+            Iterable<? extends VKey<ForeignKeyIndex<?>>> keys) {
+          if (!keys.iterator().hasNext()) {
+            return ImmutableMap.of();
+          }
+          Class<? extends EppResource> resourceClass =
+              RESOURCE_CLASS_TO_FKI_CLASS.inverse().get(keys.iterator().next().getKind());
+          ImmutableSet<String> foreignKeys =
+              Streams.stream(keys).map(v -> v.getSqlKey().toString()).collect(toImmutableSet());
+          ImmutableSet<VKey<ForeignKeyIndex<?>>> typedKeys = ImmutableSet.copyOf(keys);
+          ImmutableMap<String, ? extends ForeignKeyIndex<? extends EppResource>> existingFkis =
+              loadIndexesFromStore(resourceClass, foreignKeys, false);
           // ofy() omits keys that don't have values in Datastore, so re-add them in
           // here with Optional.empty() values.
           return Maps.asMap(
               typedKeys,
-              (Key<ForeignKeyIndex<?>> key) ->
-                  Optional.ofNullable(existingFkis.getOrDefault(key, null)));
+              (VKey<ForeignKeyIndex<?>> key) ->
+                  Optional.ofNullable(existingFkis.getOrDefault(key.getSqlKey().toString(), null)));
         }
       };
 
@@ -267,10 +299,10 @@ public abstract class ForeignKeyIndex<E extends EppResource> extends BackupGroup
    * given IDs (blah) don't exist."
    */
   @NonFinalForTesting
-  private static LoadingCache<Key<ForeignKeyIndex<?>>, Optional<ForeignKeyIndex<?>>>
+  private static LoadingCache<VKey<ForeignKeyIndex<?>>, Optional<ForeignKeyIndex<?>>>
       cacheForeignKeyIndexes = createForeignKeyIndexesCache(getEppResourceCachingDuration());
 
-  private static LoadingCache<Key<ForeignKeyIndex<?>>, Optional<ForeignKeyIndex<?>>>
+  private static LoadingCache<VKey<ForeignKeyIndex<?>>, Optional<ForeignKeyIndex<?>>>
       createForeignKeyIndexesCache(Duration expiry) {
     return CacheBuilder.newBuilder()
         .expireAfterWrite(java.time.Duration.ofMillis(expiry.getMillis()))
@@ -295,25 +327,28 @@ public abstract class ForeignKeyIndex<E extends EppResource> extends BackupGroup
    * reasons, and are OK with the trade-offs in loss of transactional consistency.
    */
   public static <E extends EppResource> ImmutableMap<String, ForeignKeyIndex<E>> loadCached(
-      Class<E> clazz, Iterable<String> foreignKeys, final DateTime now) {
+      Class<E> clazz, Collection<String> foreignKeys, final DateTime now) {
     if (!RegistryConfig.isEppResourceCachingEnabled()) {
       return tm().doTransactionless(() -> load(clazz, foreignKeys, now));
     }
-    ImmutableList<Key<ForeignKeyIndex<?>>> fkiKeys =
+    Class<? extends ForeignKeyIndex<?>> fkiClass = mapToFkiClass(clazz);
+    // Safe to cast VKey<FKI<E>> to VKey<FKI<?>>
+    @SuppressWarnings("unchecked")
+    ImmutableList<VKey<ForeignKeyIndex<?>>> fkiVKeys =
         Streams.stream(foreignKeys)
-            .map(fk -> Key.<ForeignKeyIndex<?>>create(mapToFkiClass(clazz), fk))
+            .map(fk -> (VKey<ForeignKeyIndex<?>>) VKey.create(fkiClass, fk))
             .collect(toImmutableList());
     try {
       // This cast is safe because when we loaded ForeignKeyIndexes above we used type clazz, which
       // is scoped to E.
       @SuppressWarnings("unchecked")
       ImmutableMap<String, ForeignKeyIndex<E>> fkisFromCache =
-          cacheForeignKeyIndexes.getAll(fkiKeys).entrySet().stream()
+          cacheForeignKeyIndexes.getAll(fkiVKeys).entrySet().stream()
               .filter(entry -> entry.getValue().isPresent())
               .filter(entry -> now.isBefore(entry.getValue().get().getDeletionTime()))
               .collect(
                   toImmutableMap(
-                      entry -> entry.getKey().getName(),
+                      entry -> entry.getKey().getSqlKey().toString(),
                       entry -> (ForeignKeyIndex<E>) entry.getValue().get()));
       return fkisFromCache;
     } catch (ExecutionException e) {