Fix beam deployment script again. (#1369)

uberjar task and uberjar name are now different (beamPipelineCommon and
beam_pipeline_common, respectively). This is more idiomatic with regard
to naming conventions but we need to take two different variables now.

<!-- Reviewable:start -->
---
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/google/nomulus/1369)
<!-- Reviewable:end -->

config/nom_build.py

@@ -140,6 +140,8 @@ PROPERTIES = [
              'a BEAM pipeline to image. Setting this property to empty string '
              'will disable image generation.',
              '/usr/bin/dot'),
+    Property('pipeline',
+             'The name of the Beam pipeline being staged.')
 ]
 
 GRADLE_FLAGS = [

core/build.gradle

@@ -772,50 +772,57 @@ createUberJar('nomulus', 'nomulus', 'google.registry.tools.RegistryTool')
 // This packages more code and dependency than necessary. However, without
 // restructuring the source tree it is difficult to generate leaner jars.
 createUberJar(
-    'beam_pipeline_common',
+    'beamPipelineCommon',
     'beam_pipeline_common',
     '')
 
-// Create beam staging task if environment is alpha or crash.
-// All other environments use formally released pipelines through CloudBuild.
+// Create beam staging task if the environment is alpha. Production, sandbox and
+// qa use formally released pipelines through CloudBuild, whereas crash and
+// alpha use the pipelines staged on alpha deployment project.
 //
 // User should install gcloud and login to GCP before invoking this tasks.
-if (environment in ['alpha', 'crash']) {
+if (environment == 'alpha') {
   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'
-      ],
-      [
-          mainClass: 'google.registry.beam.rde.RdePipeline',
-          metaData: 'google/registry/beam/rde_pipeline_metadata.json'
-      ],
+      InitSql            :
+          [
+              mainClass: 'google.registry.beam.initsql.InitSqlPipeline',
+              metaData : 'google/registry/beam/init_sql_pipeline_metadata.json'
+          ],
+      BulkDeleteDatastore:
+          [
+              mainClass: 'google.registry.beam.datastore.BulkDeleteDatastorePipeline',
+              metaData : 'google/registry/beam/bulk_delete_datastore_pipeline_metadata.json'
+          ],
+      Spec11             :
+          [
+              mainClass: 'google.registry.beam.spec11.Spec11Pipeline',
+              metaData : 'google/registry/beam/spec11_pipeline_metadata.json'
+          ],
+      Invoicing          :
+          [
+              mainClass: 'google.registry.beam.invoicing.InvoicingPipeline',
+              metaData : 'google/registry/beam/invoicing_pipeline_metadata.json'
+          ],
+      Rde                :
+          [
+              mainClass: 'google.registry.beam.rde.RdePipeline',
+              metaData : 'google/registry/beam/rde_pipeline_metadata.json'
+          ],
   ]
-  project.tasks.create("stage_beam_pipelines") {
+  project.tasks.create("stageBeamPipelines") {
     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
+        if (rootProject.pipeline == ''|| rootProject.pipeline == it.key) {
+          def mainClass = it.value['mainClass']
+          def metaData = it.value['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.beamPipelineCommon.outputs.files.asPath
 
-        def command = "\
+          def command = "\
             gcloud dataflow flex-template build \
             gs://${gcpProject}-deploy/live/beam/${metaDataBaseName} \
             --image-gcr-path ${imageName}:live \
@@ -825,10 +832,11 @@ if (environment in ['alpha', 'crash']) {
             --jar ${uberJarName} \
             --env FLEX_TEMPLATE_JAVA_MAIN_CLASS=${mainClass} \
             --project ${gcpProject}".toString()
-        rootProject.ext.execInBash(command, '/tmp')
+          rootProject.ext.execInBash(command, '/tmp')
+        }
       }
     }
-  }.dependsOn(tasks.beam_pipeline_common)
+  }.dependsOn(tasks.beamPipelineCommon)
 }
 
 // A jar with classes and resources from main sourceSet, excluding internal
@@ -1103,6 +1111,10 @@ test {
   // TODO(weiminyu): Remove dependency on sqlIntegrationTest
 }.dependsOn(fragileTest, outcastTest, standardTest, registryToolIntegrationTest, sqlIntegrationTest)
 
+// When we override tests, we also break the cleanTest command.
+cleanTest.dependsOn(cleanFragileTest, cleanOutcastTest, cleanStandardTest,
+                    cleanRegistryToolIntegrationTest, cleanSqlIntegrationTest)
+
 project.build.dependsOn devtool
 project.build.dependsOn buildToolImage
 project.build.dependsOn ':stage'

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

@@ -57,8 +57,7 @@ public final class CommitLogImports {
    */
   public static ImmutableList<ImmutableList<VersionedEntity>> loadEntitiesByTransaction(
       InputStream inputStream) {
-    try (AppEngineEnvironment appEngineEnvironment = new AppEngineEnvironment();
-        InputStream input = new BufferedInputStream(inputStream)) {
+    try (InputStream input = new BufferedInputStream(inputStream)) {
       Iterator<ImmutableObject> commitLogs = createDeserializingIterator(input, false);
       checkState(commitLogs.hasNext());
       checkState(commitLogs.next() instanceof CommitLogCheckpoint);

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

@@ -0,0 +1,134 @@
+// 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 google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static org.apache.http.HttpStatus.SC_INTERNAL_SERVER_ERROR;
+import static org.apache.http.HttpStatus.SC_OK;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.flogger.FluentLogger;
+import com.google.common.net.MediaType;
+import google.registry.config.RegistryConfig.Config;
+import google.registry.model.contact.ContactHistory;
+import google.registry.request.Action;
+import google.registry.request.Action.Service;
+import google.registry.request.Response;
+import google.registry.request.auth.Auth;
+import google.registry.util.Clock;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Stream;
+import javax.inject.Inject;
+import org.joda.time.DateTime;
+
+/**
+ * An action that wipes out Personal Identifiable Information (PII) fields of {@link ContactHistory}
+ * entities.
+ *
+ * <p>ContactHistory entities should be retained in the database for only certain amount of time.
+ * This periodic wipe out action only applies to SQL.
+ */
+@Action(
+    service = Service.BACKEND,
+    path = WipeOutContactHistoryPiiAction.PATH,
+    auth = Auth.AUTH_INTERNAL_OR_ADMIN)
+public class WipeOutContactHistoryPiiAction implements Runnable {
+
+  public static final String PATH = "/_dr/task/wipeOutContactHistoryPii";
+  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
+  private final Clock clock;
+  private final Response response;
+  private final int minMonthsBeforeWipeOut;
+  private final int wipeOutQueryBatchSize;
+
+  @Inject
+  public WipeOutContactHistoryPiiAction(
+      Clock clock,
+      @Config("minMonthsBeforeWipeOut") int minMonthsBeforeWipeOut,
+      @Config("wipeOutQueryBatchSize") int wipeOutQueryBatchSize,
+      Response response) {
+    this.clock = clock;
+    this.response = response;
+    this.minMonthsBeforeWipeOut = minMonthsBeforeWipeOut;
+    this.wipeOutQueryBatchSize = wipeOutQueryBatchSize;
+  }
+
+  @Override
+  public void run() {
+    response.setContentType(MediaType.PLAIN_TEXT_UTF_8);
+    try {
+      int totalNumOfWipedEntities = 0;
+      DateTime wipeOutTime = clock.nowUtc().minusMonths(minMonthsBeforeWipeOut);
+      logger.atInfo().log(
+          "About to wipe out all PII of contact history entities prior to %s.", wipeOutTime);
+
+      int numOfWipedEntities = 0;
+      do {
+        numOfWipedEntities =
+            jpaTm()
+                .transact(
+                    () ->
+                        wipeOutContactHistoryData(
+                            getNextContactHistoryEntitiesWithPiiBatch(wipeOutTime)));
+        totalNumOfWipedEntities += numOfWipedEntities;
+      } while (numOfWipedEntities > 0);
+      logger.atInfo().log(
+          "Wiped out PII of %d ContactHistory entities in total.", totalNumOfWipedEntities);
+      response.setStatus(SC_OK);
+
+    } catch (Exception e) {
+      logger.atSevere().withCause(e).log(
+          "Exception thrown during the process of wiping out contact history PII.");
+      response.setStatus(SC_INTERNAL_SERVER_ERROR);
+      response.setPayload(
+          String.format(
+              "Exception thrown during the process of wiping out contact history PII with cause"
+                  + ": %s",
+              e));
+    }
+  }
+
+  /**
+   * Returns a stream of up to {@link #wipeOutQueryBatchSize} {@link ContactHistory} entities
+   * containing PII that are prior to @param wipeOutTime.
+   */
+  @VisibleForTesting
+  Stream<ContactHistory> getNextContactHistoryEntitiesWithPiiBatch(DateTime wipeOutTime) {
+    // email is one of the required fields in EPP, meaning it's initially not null.
+    // Therefore, checking if it's null is one way to avoid processing contact history entities
+    // that have been processed previously. Refer to RFC 5733 for more information.
+    return jpaTm()
+        .query(
+            "FROM ContactHistory WHERE modificationTime < :wipeOutTime " + "AND email IS NOT NULL",
+            ContactHistory.class)
+        .setParameter("wipeOutTime", wipeOutTime)
+        .setMaxResults(wipeOutQueryBatchSize)
+        .getResultStream();
+  }
+
+  /** Wipes out the PII of each of the {@link ContactHistory} entities in the stream. */
+  @VisibleForTesting
+  int wipeOutContactHistoryData(Stream<ContactHistory> contactHistoryEntities) {
+    AtomicInteger numOfEntities = new AtomicInteger(0);
+    contactHistoryEntities.forEach(
+        contactHistoryEntity -> {
+          jpaTm().update(contactHistoryEntity.asBuilder().wipeOutPii().build());
+          numOfEntities.incrementAndGet();
+        });
+    logger.atInfo().log(
+        "Wiped out all PII fields of %d ContactHistory entities.", numOfEntities.get());
+    return numOfEntities.get();
+  }
+}

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

@@ -17,7 +17,6 @@ package google.registry.beam.common;
 import static com.google.common.base.Verify.verify;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 
-import google.registry.backup.AppEngineEnvironment;
 import google.registry.model.contact.ContactResource;
 import google.registry.persistence.transaction.CriteriaQueryBuilder;
 import google.registry.persistence.transaction.JpaTransactionManager;
@@ -59,18 +58,16 @@ public class JpaDemoPipeline implements Serializable {
                   public void processElement() {
                     // AppEngineEnvironment is needed as long as JPA entity classes still depends
                     // on Objectify.
-                    try (AppEngineEnvironment allowOfyEntity = new AppEngineEnvironment()) {
-                      int result =
-                          (Integer)
-                              jpaTm()
-                                  .transact(
-                                      () ->
-                                          jpaTm()
-                                              .getEntityManager()
-                                              .createNativeQuery("select 1;")
-                                              .getSingleResult());
-                      verify(result == 1, "Expecting 1, got %s.", result);
-                    }
+                    int result =
+                        (Integer)
+                            jpaTm()
+                                .transact(
+                                    () ->
+                                        jpaTm()
+                                            .getEntityManager()
+                                            .createNativeQuery("select 1;")
+                                            .getSingleResult());
+                    verify(result == 1, "Expecting 1, got %s.", result);
                     counter.inc();
                   }
                 }));

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

@@ -20,11 +20,9 @@ import static org.apache.beam.sdk.values.TypeDescriptors.integers;
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Streams;
-import google.registry.backup.AppEngineEnvironment;
 import google.registry.beam.common.RegistryQuery.CriteriaQuerySupplier;
 import google.registry.model.UpdateAutoTimestamp;
 import google.registry.model.UpdateAutoTimestamp.DisableAutoUpdateResource;
-import google.registry.model.ofy.ObjectifyService;
 import google.registry.model.replay.SqlEntity;
 import google.registry.persistence.transaction.JpaTransactionManager;
 import google.registry.persistence.transaction.TransactionManagerFactory;
@@ -211,14 +209,9 @@ public final class RegistryJpaIO {
 
       @ProcessElement
       public void processElement(OutputReceiver<T> outputReceiver) {
-        // AppEngineEnvironment is need for handling VKeys, which involve Ofy keys. Unlike
-        // SqlBatchWriter, it is unnecessary to initialize ObjectifyService in this class.
-        try (AppEngineEnvironment env = new AppEngineEnvironment()) {
-          // TODO(b/187210388): JpaTransactionManager should support non-transactional query.
-          jpaTm()
-              .transactNoRetry(
-                  () -> query.stream().map(resultMapper::apply).forEach(outputReceiver::output));
-        }
+        jpaTm()
+            .transactNoRetry(
+                () -> query.stream().map(resultMapper::apply).forEach(outputReceiver::output));
       }
     }
   }
@@ -364,16 +357,6 @@ public final class RegistryJpaIO {
       this.withAutoTimestamp = withAutoTimestamp;
     }
 
-    @Setup
-    public void setup() {
-      // AppEngineEnvironment is needed as long as Objectify keys are still involved in the handling
-      // of SQL entities (e.g., in VKeys). ObjectifyService needs to be initialized when conversion
-      // between Ofy entity and Datastore entity is needed.
-      try (AppEngineEnvironment env = new AppEngineEnvironment()) {
-        ObjectifyService.initOfy();
-      }
-    }
-
     @ProcessElement
     public void processElement(@Element KV<ShardedKey<Integer>, Iterable<T>> kv) {
       if (withAutoTimestamp) {
@@ -386,19 +369,17 @@ public final class RegistryJpaIO {
     }
 
     private void actuallyProcessElement(@Element KV<ShardedKey<Integer>, Iterable<T>> kv) {
-      try (AppEngineEnvironment env = new AppEngineEnvironment()) {
-        ImmutableList<Object> entities =
-            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(entities));
-          counter.inc(entities.size());
-        } catch (RuntimeException e) {
-          processSingly(entities);
-        }
+      ImmutableList<Object> entities =
+          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(entities));
+        counter.inc(entities.size());
+      } catch (RuntimeException e) {
+        processSingly(entities);
       }
     }
 

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

@@ -20,6 +20,8 @@ import com.google.auto.service.AutoService;
 import com.google.common.flogger.FluentLogger;
 import dagger.Lazy;
 import google.registry.config.RegistryEnvironment;
+import google.registry.config.SystemPropertySetter;
+import google.registry.model.AppEngineEnvironment;
 import google.registry.persistence.transaction.JpaTransactionManager;
 import google.registry.persistence.transaction.TransactionManagerFactory;
 import org.apache.beam.sdk.harness.JvmInitializer;
@@ -35,18 +37,26 @@ import org.apache.beam.sdk.options.PipelineOptions;
 @AutoService(JvmInitializer.class)
 public class RegistryPipelineWorkerInitializer implements JvmInitializer {
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
+  public static final String PROPERTY = "google.registry.beam";
 
   @Override
   public void beforeProcessing(PipelineOptions options) {
     RegistryPipelineOptions registryOptions = options.as(RegistryPipelineOptions.class);
     RegistryEnvironment environment = registryOptions.getRegistryEnvironment();
     if (environment == null || environment.equals(RegistryEnvironment.UNITTEST)) {
-      return;
+      throw new RuntimeException(
+          "A registry environment must be specified in the pipeline options.");
     }
     logger.atInfo().log("Setting up RegistryEnvironment %s.", environment);
     environment.setup();
     Lazy<JpaTransactionManager> transactionManagerLazy =
         toRegistryPipelineComponent(registryOptions).getJpaTransactionManager();
     TransactionManagerFactory.setJpaTmOnBeamWorker(transactionManagerLazy::get);
+    // Masquarade all threads as App Engine threads so we can create Ofy keys in the pipeline. Also
+    // loads all ofy entities.
+    new AppEngineEnvironment("Beam").setEnvironmentForAllThreads();
+    // Set the system property so that we can call IdService.allocateId() without access to
+    // datastore.
+    SystemPropertySetter.PRODUCTION_IMPL.setProperty(PROPERTY, "true");
   }
 }

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

@@ -20,7 +20,6 @@ import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
-import google.registry.backup.AppEngineEnvironment;
 import google.registry.backup.VersionedEntity;
 import google.registry.beam.common.RegistryJpaIO;
 import google.registry.beam.initsql.Transforms.RemoveDomainBaseForeignKeys;
@@ -230,9 +229,7 @@ public class InitSqlPipeline implements Serializable {
   }
 
   private static ImmutableList<String> toKindStrings(Collection<Class<?>> entityClasses) {
-    try (AppEngineEnvironment env = new AppEngineEnvironment()) {
-      return entityClasses.stream().map(Key::getKind).collect(ImmutableList.toImmutableList());
-    }
+    return entityClasses.stream().map(Key::getKind).collect(ImmutableList.toImmutableList());
   }
 
   public static void main(String[] args) {

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

@@ -16,6 +16,7 @@ package google.registry.beam.spec11;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static google.registry.beam.BeamUtils.getQueryFromFile;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableSet;
@@ -30,6 +31,7 @@ import google.registry.model.domain.DomainBase;
 import google.registry.model.reporting.Spec11ThreatMatch;
 import google.registry.model.reporting.Spec11ThreatMatch.ThreatType;
 import google.registry.persistence.PersistenceModule.TransactionIsolationLevel;
+import google.registry.persistence.VKey;
 import google.registry.util.Retrier;
 import google.registry.util.SqlTemplate;
 import google.registry.util.UtilsModule;
@@ -41,6 +43,7 @@ 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.PipelineOptionsFactory;
+import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.GroupByKey;
 import org.apache.beam.sdk.transforms.MapElements;
 import org.apache.beam.sdk.transforms.ParDo;
@@ -113,15 +116,43 @@ public class Spec11Pipeline implements Serializable {
   }
 
   static PCollection<DomainNameInfo> readFromCloudSql(Pipeline pipeline) {
-    Read<Object[], DomainNameInfo> read =
+    Read<Object[], KV<String, String>> read =
         RegistryJpaIO.read(
-            "select d, r.emailAddress from Domain d join Registrar r on"
-                + " d.currentSponsorClientId = r.clientIdentifier where r.type = 'REAL'"
-                + " and d.deletionTime > now()",
+            "select d.repoId, r.emailAddress from Domain d join Registrar r on"
+                + " d.currentSponsorClientId = r.clientIdentifier where r.type = 'REAL' and"
+                + " d.deletionTime > now()",
             false,
             Spec11Pipeline::parseRow);
 
-    return pipeline.apply("Read active domains from Cloud SQL", read);
+    return pipeline
+        .apply("Read active domains from Cloud SQL", read)
+        .apply(
+            "Build DomainNameInfo",
+            ParDo.of(
+                new DoFn<KV<String, String>, DomainNameInfo>() {
+                  @ProcessElement
+                  public void processElement(
+                      @Element KV<String, String> input, OutputReceiver<DomainNameInfo> output) {
+                    DomainBase domainBase =
+                        jpaTm()
+                            .transact(
+                                () ->
+                                    jpaTm()
+                                        .loadByKey(
+                                            VKey.createSql(DomainBase.class, input.getKey())));
+                    String emailAddress = input.getValue();
+                    if (emailAddress == null) {
+                      emailAddress = "";
+                    }
+                    DomainNameInfo domainNameInfo =
+                        DomainNameInfo.create(
+                            domainBase.getDomainName(),
+                            domainBase.getRepoId(),
+                            domainBase.getCurrentSponsorRegistrarId(),
+                            emailAddress);
+                    output.output(domainNameInfo);
+                  }
+                }));
   }
 
   static PCollection<DomainNameInfo> readFromBigQuery(
@@ -142,17 +173,8 @@ public class Spec11Pipeline implements Serializable {
             .withTemplateCompatibility());
   }
 
-  private static DomainNameInfo parseRow(Object[] row) {
-    DomainBase domainBase = (DomainBase) row[0];
-    String emailAddress = (String) row[1];
-    if (emailAddress == null) {
-      emailAddress = "";
-    }
-    return DomainNameInfo.create(
-        domainBase.getDomainName(),
-        domainBase.getRepoId(),
-        domainBase.getCurrentSponsorRegistrarId(),
-        emailAddress);
+  private static KV<String, String> parseRow(Object[] row) {
+    return KV.of((String) row[0], (String) row[1]);
   }
 
   static void saveToSql(

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

@@ -1306,6 +1306,18 @@ public final class RegistryConfig {
     public static ImmutableSet<String> provideAllowedEcdsaCurves(RegistryConfigSettings config) {
       return ImmutableSet.copyOf(config.sslCertificateValidation.allowedEcdsaCurves);
     }
+
+    @Provides
+    @Config("minMonthsBeforeWipeOut")
+    public static int provideMinMonthsBeforeWipeOut(RegistryConfigSettings config) {
+      return config.contactHistory.minMonthsBeforeWipeOut;
+    }
+
+    @Provides
+    @Config("wipeOutQueryBatchSize")
+    public static int provideWipeOutQueryBatchSize(RegistryConfigSettings config) {
+      return config.contactHistory.wipeOutQueryBatchSize;
+    }
   }
 
   /** Returns the App Engine project ID, which is based off the environment name. */

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

@@ -41,6 +41,7 @@ public class RegistryConfigSettings {
   public Keyring keyring;
   public RegistryTool registryTool;
   public SslCertificateValidation sslCertificateValidation;
+  public ContactHistory contactHistory;
 
   /** Configuration options that apply to the entire App Engine project. */
   public static class AppEngine {
@@ -234,4 +235,10 @@ public class RegistryConfigSettings {
     public String expirationWarningEmailBodyText;
     public String expirationWarningEmailSubjectText;
   }
+
+  /** Configuration for contact history. */
+  public static class ContactHistory {
+    public int minMonthsBeforeWipeOut;
+    public int wipeOutQueryBatchSize;
+  }
 }

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

@@ -442,6 +442,13 @@ registryTool:
   # OAuth client secret used by the tool.
   clientSecret: YOUR_CLIENT_SECRET
 
+# Configuration options for handling contact history.
+contactHistory:
+  # The number of months that a ContactHistory entity should be stored in the database.
+  minMonthsBeforeWipeOut: 18
+  # The batch size for querying ContactHistory table in the database.
+  wipeOutQueryBatchSize: 500
+
 # Configuration options for checking SSL certificates.
 sslCertificateValidation:
   # A map specifying the maximum amount of days the certificate can be valid.

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

@@ -391,6 +391,13 @@
     <url-pattern>/_dr/task/relockDomain</url-pattern>
   </servlet-mapping>
 
+  <!-- Background action to wipe out PII fields of ContactHistory entities that
+have been in the database for a certain period of time. -->
+  <servlet-mapping>
+    <servlet-name>backend-servlet</servlet-name>
+    <url-pattern>/_dr/task/wipeOutContactHistoryPii</url-pattern>
+  </servlet-mapping>
+
   <!-- Action to wipeout Cloud SQL data -->
   <servlet-mapping>
     <servlet-name>backend-servlet</servlet-name>

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

@@ -246,4 +246,14 @@
     <schedule>every 3 minutes</schedule>
     <target>backend</target>
   </cron>
+
+  <cron>
+    <url><![CDATA[/_dr/task/wipeOutContactHistoryPii]]></url>
+    <description>
+      This job runs weekly to wipe out PII fields of ContactHistory entities
+      that have been in the database for a certain period of time.
+    </description>
+    <schedule>every monday 15:00</schedule>
+    <target>backend</target>
+  </cron>
 </cronentries>

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

@@ -12,12 +12,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package google.registry.backup;
+package google.registry.model;
 
 import com.google.apphosting.api.ApiProxy;
 import com.google.apphosting.api.ApiProxy.Environment;
 import com.google.common.collect.ImmutableMap;
-import java.io.Closeable;
+import google.registry.model.ofy.ObjectifyService;
+import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.lang.reflect.Proxy;
 
@@ -38,26 +39,39 @@ import java.lang.reflect.Proxy;
  * <p>Note that conversion from Objectify objects to Datastore {@code Entities} still requires the
  * Datastore service.
  */
-public class AppEngineEnvironment implements Closeable {
+public class AppEngineEnvironment {
 
-  private boolean isPlaceHolderNeeded;
+  private Environment environment;
 
   public AppEngineEnvironment() {
     this("PlaceholderAppId");
   }
 
   public AppEngineEnvironment(String appId) {
-    isPlaceHolderNeeded = ApiProxy.getCurrentEnvironment() == null;
-    // isPlaceHolderNeeded may be true when we are invoked in a test with AppEngineExtension.
-    if (isPlaceHolderNeeded) {
-      ApiProxy.setEnvironmentForCurrentThread(createAppEngineEnvironment(appId));
-    }
+    environment = createAppEngineEnvironment(appId);
   }
 
-  @Override
-  public void close() {
-    if (isPlaceHolderNeeded) {
-      ApiProxy.setEnvironmentForCurrentThread(null);
+  public void setEnvironmentForCurrentThread() {
+    ApiProxy.setEnvironmentForCurrentThread(environment);
+    ObjectifyService.initOfy();
+  }
+
+  public void setEnvironmentForAllThreads() {
+    setEnvironmentForCurrentThread();
+    ApiProxy.setEnvironmentFactory(() -> environment);
+  }
+
+  public void unsetEnvironmentForCurrentThread() {
+    ApiProxy.clearEnvironmentForCurrentThread();
+  }
+
+  public void unsetEnvironmentForAllThreads() {
+    try {
+      Method method = ApiProxy.class.getDeclaredMethod("clearEnvironmentFactory");
+      method.setAccessible(true);
+      method.invoke(null);
+    } catch (NoSuchMethodException | InvocationTargetException | IllegalAccessException e) {
+      throw new RuntimeException(e);
     }
   }
 

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

@@ -18,13 +18,14 @@ import static com.google.common.base.Preconditions.checkState;
 
 import com.google.appengine.api.datastore.DatastoreServiceFactory;
 import com.google.common.annotations.VisibleForTesting;
+import google.registry.beam.common.RegistryPipelineWorkerInitializer;
 import google.registry.config.RegistryEnvironment;
 import java.util.concurrent.atomic.AtomicLong;
 
 /**
  * Allocates a globally unique {@link Long} number to use as an Ofy {@code @Id}.
  *
- * <p>In non-test environments the Id is generated by Datastore, whereas in tests it's from an
+ * <p>In non-test, non-beam environments the Id is generated by Datastore, otherwise it's from an
  * atomic long number that's incremented every time this method is called.
  */
 public final class IdService {
@@ -35,13 +36,25 @@ public final class IdService {
    */
   private static final String APP_WIDE_ALLOCATION_KIND = "common";
 
-  /** Counts of used ids for use in unit tests. Outside tests this is never used. */
-  private static final AtomicLong nextTestId = new AtomicLong(1); // ids cannot be zero
+  /**
+   * Counts of used ids for use in unit tests or Beam.
+   *
+   * <p>Note that one should only use self-allocate Ids in Beam for entities whose Ids are not
+   * important and are not persisted back to the database, i. e. nowhere the uniqueness of the ID is
+   * required.
+   */
+  private static final AtomicLong nextSelfAllocatedId = new AtomicLong(1); // ids cannot be zero
+
+  private static final boolean isSelfAllocated() {
+    return RegistryEnvironment.UNITTEST.equals(RegistryEnvironment.get())
+        || "true".equals(System.getProperty(RegistryPipelineWorkerInitializer.PROPERTY, "false"));
+  }
 
   /** Allocates an id. */
+  // TODO(b/201547855): Find a way to allocate a unique ID without datastore.
   public static long allocateId() {
-    return RegistryEnvironment.UNITTEST.equals(RegistryEnvironment.get())
-        ? nextTestId.getAndIncrement()
+    return isSelfAllocated()
+        ? nextSelfAllocatedId.getAndIncrement()
         : DatastoreServiceFactory.getDatastoreService()
             .allocateIds(APP_WIDE_ALLOCATION_KIND, 1)
             .iterator()
@@ -49,13 +62,11 @@ public final class IdService {
             .getId();
   }
 
-  /** Resets the global test id counter (i.e. sets the next id to 1). */
+  /** Resets the global self-allocated id counter (i.e. sets the next id to 1). */
   @VisibleForTesting
-  public static void resetNextTestId() {
+  public static void resetSelfAllocatedId() {
     checkState(
-        RegistryEnvironment.UNITTEST.equals(RegistryEnvironment.get()),
-        "Can't call resetTestIdCounts() from RegistryEnvironment.%s",
-        RegistryEnvironment.get());
-    nextTestId.set(1); // ids cannot be zero
+        isSelfAllocated(), "Can only call resetSelfAllocatedId() in unit tests or Beam pipelines");
+    nextSelfAllocatedId.set(1); // ids cannot be zero
   }
 }

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

@@ -227,5 +227,11 @@ public class ContactHistory extends HistoryEntry implements SqlEntity {
       getInstance().parent = Key.create(ContactResource.class, contactRepoId);
       return this;
     }
+
+    public Builder wipeOutPii() {
+      getInstance().contactBase =
+          getInstance().getContactBase().get().asBuilder().wipeOut().build();
+      return this;
+    }
   }
 }

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

@@ -134,7 +134,7 @@ public class DatastoreTransactionManager implements TransactionManager {
   }
 
   @Override
-  public void insertWithoutBackup(Object entity) {
+  public void insertWithoutBackup(ImmutableObject entity) {
     putWithoutBackup(entity);
   }
 
@@ -160,7 +160,7 @@ public class DatastoreTransactionManager implements TransactionManager {
   }
 
   @Override
-  public void putWithoutBackup(Object entity) {
+  public void putWithoutBackup(ImmutableObject entity) {
     syncIfTransactionless(getOfy().saveWithoutBackup().entities(toDatastoreEntity(entity)));
   }
 
@@ -180,12 +180,12 @@ public class DatastoreTransactionManager implements TransactionManager {
   }
 
   @Override
-  public void updateAll(Object... entities) {
-    updateAll(ImmutableList.of(entities));
+  public void updateAll(ImmutableObject... entities) {
+    updateAll(ImmutableList.copyOf(entities));
   }
 
   @Override
-  public void updateWithoutBackup(Object entity) {
+  public void updateWithoutBackup(ImmutableObject entity) {
     putWithoutBackup(entity);
   }
 

core/src/main/java/google/registry/module/backend/BackendRequestComponent.java

@@ -33,6 +33,7 @@ import google.registry.batch.ResaveAllEppResourcesAction;
 import google.registry.batch.ResaveEntityAction;
 import google.registry.batch.SendExpiringCertificateNotificationEmailAction;
 import google.registry.batch.WipeOutCloudSqlAction;
+import google.registry.batch.WipeOutContactHistoryPiiAction;
 import google.registry.batch.WipeoutDatastoreAction;
 import google.registry.cron.CommitLogFanoutAction;
 import google.registry.cron.CronModule;
@@ -219,6 +220,8 @@ interface BackendRequestComponent {
 
   WipeoutDatastoreAction wipeoutDatastoreAction();
 
+  WipeOutContactHistoryPiiAction wipeOutContactHistoryPiiAction();
+
   @Subcomponent.Builder
   abstract class Builder implements RequestComponentBuilder<BackendRequestComponent> {
 

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

@@ -318,7 +318,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   }
 
   @Override
-  public void insertWithoutBackup(Object entity) {
+  public void insertWithoutBackup(ImmutableObject entity) {
     insert(entity);
   }
 
@@ -356,7 +356,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   }
 
   @Override
-  public void putWithoutBackup(Object entity) {
+  public void putWithoutBackup(ImmutableObject entity) {
     put(entity);
   }
 
@@ -386,12 +386,12 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   }
 
   @Override
-  public void updateAll(Object... entities) {
-    updateAll(ImmutableList.of(entities));
+  public void updateAll(ImmutableObject... entities) {
+    updateAll(ImmutableList.copyOf(entities));
   }
 
   @Override
-  public void updateWithoutBackup(Object entity) {
+  public void updateWithoutBackup(ImmutableObject entity) {
     update(entity);
   }
 

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

@@ -111,7 +111,7 @@ public interface TransactionManager {
    * "WithoutBackup" in its method name because it is not necessary to have the backup mechanism in
    * SQL.
    */
-  void insertWithoutBackup(Object entity);
+  void insertWithoutBackup(ImmutableObject entity);
 
   /**
    * Persists all new entities in the database without writing any backup if the underlying database
@@ -144,7 +144,7 @@ public interface TransactionManager {
    * "WithoutBackup" in its method name because it is not necessary to have the backup mechanism in
    * SQL.
    */
-  void putWithoutBackup(Object entity);
+  void putWithoutBackup(ImmutableObject entity);
 
   /**
    * Persists all new entities or update the existing entities in the database without writing
@@ -165,7 +165,7 @@ public interface TransactionManager {
   void updateAll(ImmutableCollection<?> entities);
 
   /** Updates all entities in the database, throws exception if any entity does not exist. */
-  void updateAll(Object... entities);
+  void updateAll(ImmutableObject... entities);
 
   /**
    * Updates an entity in the database without writing commit logs if the underlying database is
@@ -177,7 +177,7 @@ public interface TransactionManager {
    * "WithoutBackup" in its method name because it is not necessary to have the backup mechanism in
    * SQL.
    */
-  void updateWithoutBackup(Object entity);
+  void updateWithoutBackup(ImmutableObject entity);
 
   /**
    * Updates all entities in the database without writing any backup if the underlying database is

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

@@ -132,7 +132,9 @@ public class TransactionManagerFactory {
   /**
    * Sets the return of {@link #tm()} to the given instance of {@link TransactionManager}.
    *
-   * <p>Used when overriding the per-test transaction manager for dual-database tests.
+   * <p>Used when overriding the per-test transaction manager for dual-database tests. Should be
+   * matched with a corresponding invocation of {@link #removeTmOverrideForTest()} either at the end
+   * of the test or in an <code>@AfterEach</code> handler.
    */
   @VisibleForTesting
   public static void setTmForTest(TransactionManager newTm) {

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

@@ -123,7 +123,7 @@ public final class RegistryTool {
           .put("update_allocation_tokens", UpdateAllocationTokensCommand.class)
           .put("update_cursors", UpdateCursorsCommand.class)
           .put("update_domain", UpdateDomainCommand.class)
-          .put("update_kms_keyring", UpdateKmsKeyringCommand.class)
+          .put("update_keyring_secret", UpdateKeyringSecretCommand.class)
           .put("update_premium_list", UpdatePremiumListCommand.class)
           .put("update_registrar", UpdateRegistrarCommand.class)
           .put("update_reserved_list", UpdateReservedListCommand.class)

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

@@ -160,7 +160,7 @@ interface RegistryToolComponent {
 
   void inject(UpdateDomainCommand command);
 
-  void inject(UpdateKmsKeyringCommand command);
+  void inject(UpdateKeyringSecretCommand command);
 
   void inject(UpdateRegistrarCommand command);
 

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

@@ -157,7 +157,12 @@ final class UniformRapidSuspensionCommand extends MutatingEppToolCommand {
             "newDsData",
             newDsData != null ? DsRecord.convertToSoy(newDsData) : new SoyListData(),
             "reason",
-            (undo ? "Undo " : "") + "Uniform Rapid Suspension"));
+            (undo ? "Undo " : "") + "Uniform Rapid Suspension",
+            // Domain auto-renewal is disabled as part of URS, and it's re-enabled if URS is undone.
+            // Therefore, autorenews is set to false by default and it's set to true only if the
+            // command is run in --undo mode.
+            "autorenews",
+            Boolean.toString(undo)));
   }
 
   private ImmutableSortedSet<String> getExistingNameservers(DomainBase domain) {

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

@@ -27,17 +27,16 @@ import java.nio.file.Files;
 import java.nio.file.Path;
 import javax.inject.Inject;
 
-/** Command to set and update {@code KmsKeyring} values. */
-@Parameters(
-  separators = " =",
-  commandDescription = "Update values of secrets in KmsKeyring."
-)
-final class UpdateKmsKeyringCommand implements CommandWithRemoteApi {
+/**
+ * Command to set and update ASCII-armored secret from the active {@code Keyring} implementation.
+ */
+@Parameters(separators = " =", commandDescription = "Update values of secret in the keyring.")
+final class UpdateKeyringSecretCommand implements CommandWithRemoteApi {
 
   @Inject KmsUpdater kmsUpdater;
 
   @Inject
-  UpdateKmsKeyringCommand() {}
+  UpdateKeyringSecretCommand() {}
 
   @Parameter(names  = "--keyname", description = "The secret to update", required = true)
   private KeyringKeyName keyringKeyName;

core/src/main/resources/google/registry/tools/soy/UniformRapidSuspension.soy

@@ -25,6 +25,7 @@
 {@param statusesToRemove: list<string>}
 {@param newDsData: list<[keyTag:int, alg:int, digestType:int, digest:string]>}
 {@param reason: string}
+{@param autorenews: string}
 <?xml version="1.0" encoding="UTF-8" standalone="yes"?>
 <epp xmlns="urn:ietf:params:xml:ns:epp-1.0">
   <command>
@@ -75,6 +76,9 @@
         </secDNS:add>
         {/if}
       </secDNS:update>
+      <superuser:domainUpdate xmlns:superuser="urn:google:params:xml:ns:superuser-1.0">
+        <superuser:autorenews>{$autorenews}</superuser:autorenews>
+      </superuser:domainUpdate>
       <metadata:metadata xmlns:metadata="urn:google:params:xml:ns:metadata-1.0">
         <metadata:reason>{$reason}</metadata:reason>
         <metadata:requestedByRegistrar>false</metadata:requestedByRegistrar>

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

@@ -27,6 +27,7 @@ import static google.registry.persistence.transaction.TransactionManagerFactory.
 import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatabaseHelper.createTld;
+import static google.registry.testing.DatabaseHelper.insertInDb;
 import static google.registry.testing.DatabaseHelper.newDomainBase;
 import static google.registry.testing.DatabaseHelper.persistActiveContact;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
@@ -163,12 +164,7 @@ public class ReplayCommitLogsToSqlActionTest {
 
   @Test
   void testReplay_multipleDiffFiles() throws Exception {
-    jpaTm()
-        .transact(
-            () -> {
-              jpaTm().insertWithoutBackup(TestObject.create("previous to keep"));
-              jpaTm().insertWithoutBackup(TestObject.create("previous to delete"));
-            });
+    insertInDb(TestObject.create("previous to keep"), TestObject.create("previous to delete"));
     DateTime now = fakeClock.nowUtc();
     // Create 3 transactions, across two diff files.
     // Before: {"previous to keep", "previous to delete"}
@@ -216,7 +212,7 @@ public class ReplayCommitLogsToSqlActionTest {
   @Test
   void testReplay_noManifests() throws Exception {
     DateTime now = fakeClock.nowUtc();
-    jpaTm().transact(() -> jpaTm().insertWithoutBackup(TestObject.create("previous to keep")));
+    insertInDb(TestObject.create("previous to keep"));
     saveDiffFileNotToRestore(gcsUtils, now.minusMinutes(1));
     saveDiffFile(gcsUtils, createCheckpoint(now.minusMillis(2)));
     jpaTm().transact(() -> SqlReplayCheckpoint.set(now.minusMillis(1)));
@@ -228,7 +224,7 @@ public class ReplayCommitLogsToSqlActionTest {
   void testReplay_dryRun() throws Exception {
     action.dryRun = true;
     DateTime now = fakeClock.nowUtc();
-    jpaTm().transact(() -> jpaTm().insertWithoutBackup(TestObject.create("previous to keep")));
+    insertInDb(TestObject.create("previous to keep"));
     Key<CommitLogBucket> bucketKey = getBucketKey(1);
     Key<CommitLogManifest> manifestKey = CommitLogManifest.createKey(bucketKey, now);
     saveDiffFileNotToRestore(gcsUtils, now.minusMinutes(2));
@@ -253,7 +249,7 @@ public class ReplayCommitLogsToSqlActionTest {
   @Test
   void testReplay_manifestWithNoDeletions() throws Exception {
     DateTime now = fakeClock.nowUtc();
-    jpaTm().transact(() -> jpaTm().insertWithoutBackup(TestObject.create("previous to keep")));
+    insertInDb(TestObject.create("previous to keep"));
     Key<CommitLogBucket> bucketKey = getBucketKey(1);
     Key<CommitLogManifest> manifestKey = CommitLogManifest.createKey(bucketKey, now);
     saveDiffFileNotToRestore(gcsUtils, now.minusMinutes(2));
@@ -271,12 +267,7 @@ public class ReplayCommitLogsToSqlActionTest {
   @Test
   void testReplay_manifestWithNoMutations() throws Exception {
     DateTime now = fakeClock.nowUtc();
-    jpaTm()
-        .transact(
-            () -> {
-              jpaTm().insertWithoutBackup(TestObject.create("previous to keep"));
-              jpaTm().insertWithoutBackup(TestObject.create("previous to delete"));
-            });
+    insertInDb(TestObject.create("previous to keep"), TestObject.create("previous to delete"));
     saveDiffFileNotToRestore(gcsUtils, now.minusMinutes(2));
     jpaTm().transact(() -> SqlReplayCheckpoint.set(now.minusMinutes(1).minusMillis(1)));
     saveDiffFile(
@@ -293,7 +284,7 @@ public class ReplayCommitLogsToSqlActionTest {
   @Test
   void testReplay_mutateExistingEntity() throws Exception {
     DateTime now = fakeClock.nowUtc();
-    jpaTm().transact(() -> jpaTm().put(TestObject.create("existing", "a")));
+    insertInDb(TestObject.create("existing", "a"));
     Key<CommitLogManifest> manifestKey = CommitLogManifest.createKey(getBucketKey(1), now);
     saveDiffFileNotToRestore(gcsUtils, now.minusMinutes(1).minusMillis(1));
     jpaTm().transact(() -> SqlReplayCheckpoint.set(now.minusMinutes(1)));
@@ -312,7 +303,7 @@ public class ReplayCommitLogsToSqlActionTest {
   @Test
   void testReplay_deleteMissingEntity() throws Exception {
     DateTime now = fakeClock.nowUtc();
-    jpaTm().transact(() -> jpaTm().put(TestObject.create("previous to keep", "a")));
+    insertInDb(TestObject.create("previous to keep", "a"));
     saveDiffFileNotToRestore(gcsUtils, now.minusMinutes(1).minusMillis(1));
     jpaTm().transact(() -> SqlReplayCheckpoint.set(now.minusMinutes(1)));
     saveDiffFile(
@@ -368,7 +359,7 @@ public class ReplayCommitLogsToSqlActionTest {
 
     // Create and save to SQL a registrar contact that we will delete
     RegistrarContact toDelete = AppEngineExtension.makeRegistrarContact1();
-    jpaTm().transact(() -> jpaTm().put(toDelete));
+    insertInDb(toDelete);
     jpaTm().transact(() -> SqlReplayCheckpoint.set(now.minusMinutes(1).minusMillis(1)));
 
     // spy the txn manager so we can see what order things were inserted/removed
@@ -569,8 +560,7 @@ public class ReplayCommitLogsToSqlActionTest {
             .asBuilder()
             .setDsData(ImmutableSet.of(DelegationSignerData.create(1, 2, 3, new byte[] {0, 1, 2})))
             .build();
-    jpaTm().transact(() -> jpaTm().put(domain));
-
+    insertInDb(domain);
     assertThat(jpaTm().transact(() -> jpaTm().loadAllOf(DelegationSignerData.class))).isNotEmpty();
 
     saveDiffFile(
@@ -580,12 +570,8 @@ public class ReplayCommitLogsToSqlActionTest {
             getBucketKey(1), now.minusMinutes(3), ImmutableSet.of(Key.create(domain))));
     runAndAssertSuccess(now.minusMinutes(1), 1, 1);
 
-    jpaTm()
-        .transact(
-            () -> {
-              assertThat(jpaTm().loadAllOf(DomainBase.class)).isEmpty();
-              assertThat(jpaTm().loadAllOf(DelegationSignerData.class)).isEmpty();
-            });
+    assertThat(jpaTm().transact(() -> jpaTm().loadAllOf(DomainBase.class))).isEmpty();
+    assertThat(jpaTm().transact(() -> jpaTm().loadAllOf(DelegationSignerData.class))).isEmpty();
   }
 
   @Test
@@ -627,7 +613,7 @@ public class ReplayCommitLogsToSqlActionTest {
         ImmutableSet.of(DelegationSignerData.create(1, 2, 3, new byte[] {4, 5, 6}));
     DomainBase domainWithDsData =
         newDomainBase("example.tld").asBuilder().setDsData(dsData).build();
-    jpaTm().transact(() -> jpaTm().put(domainWithDsData));
+    insertInDb(domainWithDsData);
 
     // Replay a version of that domain without the dsData
     Key<CommitLogManifest> manifestKeyOne =
@@ -639,7 +625,7 @@ public class ReplayCommitLogsToSqlActionTest {
 
     // Create an object (any object) to delete via replay to trigger Hibernate flush events
     TestObject testObject = TestObject.create("foo", "bar");
-    jpaTm().transact(() -> jpaTm().put(testObject));
+    insertInDb(testObject);
 
     // Replay the original domain, with the original dsData
     Key<CommitLogManifest> manifestKeyTwo =

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

@@ -0,0 +1,372 @@
+// 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.truth.Truth.assertThat;
+import static com.google.common.truth.Truth.assertWithMessage;
+import static com.google.common.truth.Truth8.assertThat;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.testing.DatabaseHelper.persistResource;
+import static org.apache.http.HttpStatus.SC_OK;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.truth.Truth8;
+import google.registry.model.contact.ContactAddress;
+import google.registry.model.contact.ContactAuthInfo;
+import google.registry.model.contact.ContactBase;
+import google.registry.model.contact.ContactHistory;
+import google.registry.model.contact.ContactPhoneNumber;
+import google.registry.model.contact.ContactResource;
+import google.registry.model.contact.Disclose;
+import google.registry.model.contact.PostalInfo;
+import google.registry.model.eppcommon.AuthInfo.PasswordAuth;
+import google.registry.model.eppcommon.PresenceMarker;
+import google.registry.model.eppcommon.StatusValue;
+import google.registry.testing.AppEngineExtension;
+import google.registry.testing.DatabaseHelper;
+import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.FakeClock;
+import google.registry.testing.FakeResponse;
+import google.registry.testing.InjectExtension;
+import google.registry.testing.TestSqlOnly;
+import org.joda.time.DateTime;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+/** Unit tests for {@link WipeOutContactHistoryPiiAction}. */
+@DualDatabaseTest
+class WipeOutContactHistoryPiiActionTest {
+
+  private static final int MIN_MONTHS_BEFORE_WIPE_OUT = 18;
+  private static final int BATCH_SIZE = 500;
+  private static final ContactResource defaultContactResource =
+      new ContactResource.Builder()
+          .setContactId("sh8013")
+          .setRepoId("2FF-ROID")
+          .setStatusValues(ImmutableSet.of(StatusValue.CLIENT_DELETE_PROHIBITED))
+          .setLocalizedPostalInfo(
+              new PostalInfo.Builder()
+                  .setType(PostalInfo.Type.LOCALIZED)
+                  .setAddress(
+                      new ContactAddress.Builder()
+                          .setStreet(ImmutableList.of("123 Grand Ave"))
+                          .build())
+                  .build())
+          .setInternationalizedPostalInfo(
+              new PostalInfo.Builder()
+                  .setType(PostalInfo.Type.INTERNATIONALIZED)
+                  .setName("John Doe")
+                  .setOrg("Example Inc.")
+                  .setAddress(
+                      new ContactAddress.Builder()
+                          .setStreet(ImmutableList.of("123 Example Dr.", "Suite 100"))
+                          .setCity("Dulles")
+                          .setState("VA")
+                          .setZip("20166-6503")
+                          .setCountryCode("US")
+                          .build())
+                  .build())
+          .setVoiceNumber(
+              new ContactPhoneNumber.Builder()
+                  .setPhoneNumber("+1.7035555555")
+                  .setExtension("1234")
+                  .build())
+          .setFaxNumber(new ContactPhoneNumber.Builder().setPhoneNumber("+1.7035555556").build())
+          .setEmailAddress("jdoe@example.com")
+          .setPersistedCurrentSponsorRegistrarId("TheRegistrar")
+          .setCreationRegistrarId("NewRegistrar")
+          .setLastEppUpdateRegistrarId("NewRegistrar")
+          .setCreationTimeForTest(DateTime.parse("1999-04-03T22:00:00.0Z"))
+          .setLastEppUpdateTime(DateTime.parse("1999-12-03T09:00:00.0Z"))
+          .setLastTransferTime(DateTime.parse("2000-04-08T09:00:00.0Z"))
+          .setAuthInfo(ContactAuthInfo.create(PasswordAuth.create("2fooBAR")))
+          .setDisclose(
+              new Disclose.Builder()
+                  .setFlag(true)
+                  .setVoice(new PresenceMarker())
+                  .setEmail(new PresenceMarker())
+                  .build())
+          .build();
+
+  @RegisterExtension
+  public final AppEngineExtension appEngine =
+      AppEngineExtension.builder().withDatastoreAndCloudSql().withTaskQueue().build();
+
+  @RegisterExtension public final InjectExtension inject = new InjectExtension();
+  private final FakeClock clock = new FakeClock(DateTime.parse("2021-08-26T20:21:22Z"));
+
+  private FakeResponse response;
+  private WipeOutContactHistoryPiiAction action;
+
+  @BeforeEach
+  void beforeEach() {
+    response = new FakeResponse();
+    action =
+        new WipeOutContactHistoryPiiAction(clock, MIN_MONTHS_BEFORE_WIPE_OUT, BATCH_SIZE, response);
+  }
+
+  @TestSqlOnly
+  void getAllHistoryEntitiesOlderThan_returnsAllPersistedEntities() {
+    ImmutableList<ContactHistory> expectedToBeWipedOut =
+        persistLotsOfContactHistoryEntities(
+            20, MIN_MONTHS_BEFORE_WIPE_OUT + 1, 0, defaultContactResource);
+    jpaTm()
+        .transact(
+            () ->
+                assertThat(
+                        action.getNextContactHistoryEntitiesWithPiiBatch(
+                            clock.nowUtc().minusMonths(MIN_MONTHS_BEFORE_WIPE_OUT)))
+                    .containsExactlyElementsIn(expectedToBeWipedOut));
+  }
+
+  @TestSqlOnly
+  void getAllHistoryEntitiesOlderThan_returnsOnlyPartOfThePersistedEntities() {
+    ImmutableList<ContactHistory> expectedToBeWipedOut =
+        persistLotsOfContactHistoryEntities(
+            40, MIN_MONTHS_BEFORE_WIPE_OUT + 2, 0, defaultContactResource);
+
+    // persisted entities that should not be part of the actual result
+    persistLotsOfContactHistoryEntities(
+        15, 17, MIN_MONTHS_BEFORE_WIPE_OUT - 1, defaultContactResource);
+
+    jpaTm()
+        .transact(
+            () ->
+                Truth8.assertThat(
+                        action.getNextContactHistoryEntitiesWithPiiBatch(
+                            clock.nowUtc().minusMonths(MIN_MONTHS_BEFORE_WIPE_OUT)))
+                    .containsExactlyElementsIn(expectedToBeWipedOut));
+  }
+
+  @TestSqlOnly
+  void run_withNoEntitiesToWipeOut_success() {
+    assertThat(
+            jpaTm()
+                .transact(
+                    () ->
+                        action
+                            .getNextContactHistoryEntitiesWithPiiBatch(
+                                clock.nowUtc().minusMonths(MIN_MONTHS_BEFORE_WIPE_OUT))
+                            .count()))
+        .isEqualTo(0);
+    action.run();
+
+    assertThat(
+            jpaTm()
+                .transact(
+                    () ->
+                        action
+                            .getNextContactHistoryEntitiesWithPiiBatch(
+                                clock.nowUtc().minusMonths(MIN_MONTHS_BEFORE_WIPE_OUT))
+                            .count()))
+        .isEqualTo(0);
+
+    assertThat(response.getStatus()).isEqualTo(SC_OK);
+  }
+
+  @TestSqlOnly
+  void run_withOneBatchOfEntities_success() {
+    int numOfMonthsFromNow = MIN_MONTHS_BEFORE_WIPE_OUT + 2;
+    ImmutableList<ContactHistory> expectedToBeWipedOut =
+        persistLotsOfContactHistoryEntities(20, numOfMonthsFromNow, 0, defaultContactResource);
+
+    // The query should return a stream of all persisted entities.
+    assertThat(
+            jpaTm()
+                .transact(
+                    () ->
+                        action
+                            .getNextContactHistoryEntitiesWithPiiBatch(
+                                clock.nowUtc().minusMonths(MIN_MONTHS_BEFORE_WIPE_OUT))
+                            .count()))
+        .isEqualTo(expectedToBeWipedOut.size());
+
+    assertAllEntitiesContainPii(DatabaseHelper.loadByEntitiesIfPresent(expectedToBeWipedOut));
+
+    action.run();
+
+    // The query should return an empty stream after the wipe out action.
+    assertThat(
+            jpaTm()
+                .transact(
+                    () ->
+                        action
+                            .getNextContactHistoryEntitiesWithPiiBatch(
+                                clock.nowUtc().minusMonths(MIN_MONTHS_BEFORE_WIPE_OUT))
+                            .count()))
+        .isEqualTo(0);
+
+    assertAllPiiFieldsAreWipedOut(DatabaseHelper.loadByEntitiesIfPresent(expectedToBeWipedOut));
+  }
+
+  @TestSqlOnly
+  void run_withMultipleBatches_numOfEntitiesAsNonMultipleOfBatchSize_success() {
+    int numOfMonthsFromNow = MIN_MONTHS_BEFORE_WIPE_OUT + 2;
+    ImmutableList<ContactHistory> expectedToBeWipedOut =
+        persistLotsOfContactHistoryEntities(1234, numOfMonthsFromNow, 0, defaultContactResource);
+
+    // The query should return a subset of all persisted data.
+    assertThat(
+            jpaTm()
+                .transact(
+                    () ->
+                        action
+                            .getNextContactHistoryEntitiesWithPiiBatch(
+                                clock.nowUtc().minusMonths(MIN_MONTHS_BEFORE_WIPE_OUT))
+                            .count()))
+        .isEqualTo(BATCH_SIZE);
+
+    assertAllEntitiesContainPii(DatabaseHelper.loadByEntitiesIfPresent(expectedToBeWipedOut));
+    action.run();
+
+    // The query should return an empty stream after the wipe out action.
+    assertThat(
+            jpaTm()
+                .transact(
+                    () ->
+                        action
+                            .getNextContactHistoryEntitiesWithPiiBatch(
+                                clock.nowUtc().minusMonths(MIN_MONTHS_BEFORE_WIPE_OUT))
+                            .count()))
+        .isEqualTo(0);
+
+    assertAllPiiFieldsAreWipedOut(DatabaseHelper.loadByEntitiesIfPresent(expectedToBeWipedOut));
+  }
+
+  @TestSqlOnly
+  void run_withMultipleBatches_numOfEntitiesAsMultiplesOfBatchSize_success() {
+    int numOfMonthsFromNow = MIN_MONTHS_BEFORE_WIPE_OUT + 2;
+    ImmutableList<ContactHistory> expectedToBeWipedOut =
+        persistLotsOfContactHistoryEntities(2000, numOfMonthsFromNow, 0, defaultContactResource);
+
+    // The query should return a subset of all persisted data.
+    assertThat(
+            jpaTm()
+                .transact(
+                    () ->
+                        action
+                            .getNextContactHistoryEntitiesWithPiiBatch(
+                                clock.nowUtc().minusMonths(MIN_MONTHS_BEFORE_WIPE_OUT))
+                            .count()))
+        .isEqualTo(BATCH_SIZE);
+
+    assertAllEntitiesContainPii(DatabaseHelper.loadByEntitiesIfPresent(expectedToBeWipedOut));
+    action.run();
+
+    // The query should return an empty stream after the wipe out action.
+    assertThat(
+            jpaTm()
+                .transact(
+                    () ->
+                        action
+                            .getNextContactHistoryEntitiesWithPiiBatch(
+                                clock.nowUtc().minusMonths(MIN_MONTHS_BEFORE_WIPE_OUT))
+                            .count()))
+        .isEqualTo(0);
+
+    assertAllPiiFieldsAreWipedOut(DatabaseHelper.loadByEntitiesIfPresent(expectedToBeWipedOut));
+  }
+
+  @TestSqlOnly
+  void wipeOutContactHistoryData_wipesOutNoEntity() {
+    jpaTm()
+        .transact(
+            () -> {
+              assertThat(
+                      action.wipeOutContactHistoryData(
+                          action.getNextContactHistoryEntitiesWithPiiBatch(
+                              clock.nowUtc().minusMonths(MIN_MONTHS_BEFORE_WIPE_OUT))))
+                  .isEqualTo(0);
+            });
+  }
+
+  @TestSqlOnly
+  void wipeOutContactHistoryData_wipesOutMultipleEntities() {
+    int numOfMonthsFromNow = MIN_MONTHS_BEFORE_WIPE_OUT + 3;
+    ImmutableList<ContactHistory> expectedToBeWipedOut =
+        persistLotsOfContactHistoryEntities(20, numOfMonthsFromNow, 0, defaultContactResource);
+
+    assertAllEntitiesContainPii(DatabaseHelper.loadByEntitiesIfPresent(expectedToBeWipedOut));
+
+    jpaTm()
+        .transact(
+            () -> {
+              action.wipeOutContactHistoryData(
+                  action.getNextContactHistoryEntitiesWithPiiBatch(
+                      clock.nowUtc().minusMonths(MIN_MONTHS_BEFORE_WIPE_OUT)));
+            });
+
+    assertAllPiiFieldsAreWipedOut(DatabaseHelper.loadByEntitiesIfPresent(expectedToBeWipedOut));
+  }
+
+  /** persists a number of ContactHistory entities for load and query testing. */
+  ImmutableList<ContactHistory> persistLotsOfContactHistoryEntities(
+      int numOfEntities, int minusMonths, int minusDays, ContactResource contact) {
+    ImmutableList.Builder<ContactHistory> expectedEntitesBuilder = new ImmutableList.Builder<>();
+    for (int i = 0; i < numOfEntities; i++) {
+      expectedEntitesBuilder.add(
+          persistResource(
+              new ContactHistory()
+                  .asBuilder()
+                  .setRegistrarId("NewRegistrar")
+                  .setModificationTime(clock.nowUtc().minusMonths(minusMonths).minusDays(minusDays))
+                  .setType(ContactHistory.Type.CONTACT_DELETE)
+                  .setContact(persistResource(contact))
+                  .build()));
+    }
+    return expectedEntitesBuilder.build();
+  }
+
+  boolean areAllPiiFieldsWiped(ContactBase contactBase) {
+    return contactBase.getEmailAddress() == null
+        && contactBase.getFaxNumber() == null
+        && contactBase.getInternationalizedPostalInfo() == null
+        && contactBase.getLocalizedPostalInfo() == null
+        && contactBase.getVoiceNumber() == null;
+  }
+
+  boolean containsPii(ContactBase contactBase) {
+    return contactBase.getEmailAddress() != null
+        || contactBase.getFaxNumber() != null
+        || contactBase.getInternationalizedPostalInfo() != null
+        || contactBase.getLocalizedPostalInfo() != null
+        || contactBase.getVoiceNumber() != null;
+  }
+
+  void assertAllPiiFieldsAreWipedOut(ImmutableList<ContactHistory> entities) {
+    ImmutableList.Builder<ContactHistory> notWipedEntities = new ImmutableList.Builder<>();
+    for (ContactHistory entity : entities) {
+      if (!areAllPiiFieldsWiped(entity.getContactBase().get())) {
+        notWipedEntities.add(entity);
+      }
+    }
+    assertWithMessage("Not all PII fields of the contact history entities were wiped.")
+        .that(notWipedEntities.build())
+        .isEmpty();
+  }
+
+  void assertAllEntitiesContainPii(ImmutableList<ContactHistory> entities) {
+    ImmutableList.Builder<ContactHistory> entitiesWithNoPii = new ImmutableList.Builder<>();
+    for (ContactHistory entity : entities) {
+      if (!containsPii(entity.getContactBase().get())) {
+        entitiesWithNoPii.add(entity);
+      }
+    }
+    assertWithMessage("Not all contact history entities contain PII.")
+        .that(entitiesWithNoPii.build())
+        .isEmpty();
+  }
+}

core/src/test/java/google/registry/beam/common/RegistryJpaReadTest.java

@@ -14,7 +14,6 @@
 
 package google.registry.beam.common;
 
-import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.testing.AppEngineExtension.makeRegistrar1;
 import static google.registry.testing.DatabaseHelper.insertInDb;
 import static google.registry.testing.DatabaseHelper.newRegistry;
@@ -63,7 +62,8 @@ public class RegistryJpaReadTest {
 
   @RegisterExtension
   @Order(Order.DEFAULT - 1)
-  final transient DatastoreEntityExtension datastore = new DatastoreEntityExtension();
+  final transient DatastoreEntityExtension datastore =
+      new DatastoreEntityExtension().allThreads(true);
 
   @RegisterExtension
   final transient JpaIntegrationTestExtension database =
@@ -78,7 +78,7 @@ public class RegistryJpaReadTest {
   @BeforeEach
   void beforeEach() {
     Registrar ofyRegistrar = AppEngineExtension.makeRegistrar2();
-    jpaTm().transact(() -> jpaTm().put(ofyRegistrar));
+    insertInDb(ofyRegistrar);
 
     ImmutableList.Builder<ContactResource> builder = new ImmutableList.Builder<>();
 
@@ -87,7 +87,7 @@ public class RegistryJpaReadTest {
       builder.add(contact);
     }
     contacts = builder.build();
-    jpaTm().transact(() -> jpaTm().putAll(contacts));
+    insertInDb(contacts);
   }
 
   @Test

core/src/test/java/google/registry/beam/common/RegistryJpaWriteTest.java

@@ -17,6 +17,8 @@ package google.registry.beam.common;
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.ImmutableObjectSubject.immutableObjectCorrespondence;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.testing.DatabaseHelper.newContactResource;
+import static google.registry.testing.DatabaseHelper.putInDb;
 
 import com.google.appengine.api.datastore.Entity;
 import com.google.common.collect.ImmutableList;
@@ -32,7 +34,6 @@ import google.registry.model.registrar.Registrar;
 import google.registry.persistence.transaction.JpaTestExtensions;
 import google.registry.persistence.transaction.JpaTestExtensions.JpaIntegrationTestExtension;
 import google.registry.testing.AppEngineExtension;
-import google.registry.testing.DatabaseHelper;
 import google.registry.testing.DatastoreEntityExtension;
 import google.registry.testing.FakeClock;
 import google.registry.testing.InjectExtension;
@@ -76,12 +77,12 @@ class RegistryJpaWriteTest implements Serializable {
       // Required for contacts created below.
       Registrar ofyRegistrar = AppEngineExtension.makeRegistrar2();
       store.insertOrUpdate(ofyRegistrar);
-      jpaTm().transact(() -> jpaTm().put(store.loadAsOfyEntity(ofyRegistrar)));
+      putInDb(store.loadAsOfyEntity(ofyRegistrar));
 
       ImmutableList.Builder<Entity> builder = new ImmutableList.Builder<>();
 
       for (int i = 0; i < 3; i++) {
-        ContactResource contact = DatabaseHelper.newContactResource("contact_" + i);
+        ContactResource contact = newContactResource("contact_" + i);
         store.insertOrUpdate(contact);
         builder.add(store.loadAsDatastoreEntity(contact));
       }
@@ -106,8 +107,7 @@ class RegistryJpaWriteTest implements Serializable {
                 .withJpaConverter(Transforms::convertVersionedEntityToSqlEntity));
     testPipeline.run().waitUntilFinish();
 
-    ImmutableList<?> sqlContacts = jpaTm().transact(() -> jpaTm().loadAllOf(ContactResource.class));
-    assertThat(sqlContacts)
+    assertThat(jpaTm().transact(() -> jpaTm().loadAllOf(ContactResource.class)))
         .comparingElementsUsing(immutableObjectCorrespondence("revisions", "updateTimestamp"))
         .containsExactlyElementsIn(
             contacts.stream()

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

@@ -25,6 +25,7 @@ import com.google.appengine.api.datastore.EntityNotFoundException;
 import com.googlecode.objectify.Key;
 import google.registry.backup.CommitLogExports;
 import google.registry.backup.VersionedEntity;
+import google.registry.model.ImmutableObject;
 import google.registry.model.ofy.CommitLogCheckpoint;
 import google.registry.testing.AppEngineExtension;
 import google.registry.testing.FakeClock;
@@ -124,7 +125,7 @@ public final class BackupTestStore implements AutoCloseable {
    *
    * <p>See {@link #loadAsDatastoreEntity} and {@link VersionedEntity} for more information.
    */
-  public Object loadAsOfyEntity(Object ofyEntity) {
+  public ImmutableObject loadAsOfyEntity(ImmutableObject ofyEntity) {
     try {
       return auditedOfy().load().fromEntity(datastoreService.get(Key.create(ofyEntity).getRaw()));
     } catch (EntityNotFoundException e) {

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

@@ -71,7 +71,7 @@ class CommitLogTransformsTest implements Serializable {
   @RegisterExtension
   @Order(value = 1)
   final transient DatastoreEntityExtension datastoreEntityExtension =
-      new DatastoreEntityExtension();
+      new DatastoreEntityExtension().allThreads(true);
 
   @RegisterExtension
   final transient TestPipelineExtension testPipeline =

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

@@ -30,7 +30,6 @@ import static google.registry.util.DateTimeUtils.START_OF_TIME;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
-import google.registry.backup.AppEngineEnvironment;
 import google.registry.beam.TestPipelineExtension;
 import google.registry.flows.domain.DomainFlowUtils;
 import google.registry.model.billing.BillingEvent;
@@ -106,7 +105,8 @@ class InitSqlPipelineTest {
 
   @RegisterExtension
   @Order(Order.DEFAULT - 1)
-  final transient DatastoreEntityExtension datastore = new DatastoreEntityExtension();
+  final transient DatastoreEntityExtension datastore =
+      new DatastoreEntityExtension("test").allThreads(true);
 
   @RegisterExtension final transient InjectExtension injectExtension = new InjectExtension();
 
@@ -331,20 +331,18 @@ class InitSqlPipelineTest {
             .as(InitSqlPipelineOptions.class);
     InitSqlPipeline initSqlPipeline = new InitSqlPipeline(options);
     initSqlPipeline.run(testPipeline).waitUntilFinish();
-    try (AppEngineEnvironment env = new AppEngineEnvironment("test")) {
-      assertHostResourceEquals(
-          jpaTm().transact(() -> jpaTm().loadByKey(hostResource.createVKey())), hostResource);
-      assertThat(jpaTm().transact(() -> jpaTm().loadAllOf(Registrar.class)))
-          .comparingElementsUsing(immutableObjectCorrespondence("lastUpdateTime"))
-          .containsExactly(registrar1, registrar2);
-      assertThat(jpaTm().transact(() -> jpaTm().loadAllOf(ContactResource.class)))
-          .comparingElementsUsing(immutableObjectCorrespondence("revisions", "updateTimestamp"))
-          .containsExactly(contact1, contact2);
-      assertDomainEquals(jpaTm().transact(() -> jpaTm().loadByKey(domain.createVKey())), domain);
-      assertThat(jpaTm().transact(() -> jpaTm().loadAllOf(Cursor.class)))
-          .comparingElementsUsing(immutableObjectCorrespondence())
-          .containsExactly(globalCursor, tldCursor);
-    }
+    assertHostResourceEquals(
+        jpaTm().transact(() -> jpaTm().loadByKey(hostResource.createVKey())), hostResource);
+    assertThat(jpaTm().transact(() -> jpaTm().loadAllOf(Registrar.class)))
+        .comparingElementsUsing(immutableObjectCorrespondence("lastUpdateTime"))
+        .containsExactly(registrar1, registrar2);
+    assertThat(jpaTm().transact(() -> jpaTm().loadAllOf(ContactResource.class)))
+        .comparingElementsUsing(immutableObjectCorrespondence("revisions", "updateTimestamp"))
+        .containsExactly(contact1, contact2);
+    assertDomainEquals(jpaTm().transact(() -> jpaTm().loadByKey(domain.createVKey())), domain);
+    assertThat(jpaTm().transact(() -> jpaTm().loadAllOf(Cursor.class)))
+        .comparingElementsUsing(immutableObjectCorrespondence())
+        .containsExactly(globalCursor, tldCursor);
   }
 
   private static void assertHostResourceEquals(HostResource actual, HostResource expected) {

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

@@ -26,7 +26,6 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Streams;
 import com.google.common.truth.Truth;
 import com.google.storage.onestore.v3.OnestoreEntity.EntityProto;
-import google.registry.backup.AppEngineEnvironment;
 import google.registry.backup.VersionedEntity;
 import java.io.Serializable;
 import java.util.Collection;
@@ -124,20 +123,18 @@ public final class InitSqlTestUtils {
                       public void processElement(
                           @Element KV<String, Iterable<VersionedEntity>> input,
                           OutputReceiver<String> out) {
-                        try (AppEngineEnvironment env = new AppEngineEnvironment()) {
-                          ImmutableList<KV<Long, Object>> actual =
-                              Streams.stream(input.getValue())
-                                  .map(InitSqlTestUtils::rawEntityToOfyWithTimestamp)
-                                  .collect(ImmutableList.toImmutableList());
-                          try {
-                            Truth.assertThat(actual)
-                                .containsExactlyElementsIn(
-                                    Stream.of(expected)
-                                        .map(InitSqlTestUtils::expectedToOfyWithTimestamp)
-                                        .collect(ImmutableList.toImmutableList()));
-                          } catch (AssertionError e) {
-                            out.output(e.toString());
-                          }
+                        ImmutableList<KV<Long, Object>> actual =
+                            Streams.stream(input.getValue())
+                                .map(InitSqlTestUtils::rawEntityToOfyWithTimestamp)
+                                .collect(ImmutableList.toImmutableList());
+                        try {
+                          Truth.assertThat(actual)
+                              .containsExactlyElementsIn(
+                                  Stream.of(expected)
+                                      .map(InitSqlTestUtils::expectedToOfyWithTimestamp)
+                                      .collect(ImmutableList.toImmutableList()));
+                        } catch (AssertionError e) {
+                          out.output(e.toString());
                         }
                       }
                     }));

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

@@ -95,7 +95,7 @@ class LoadDatastoreSnapshotTest {
   @RegisterExtension
   @Order(value = 1)
   final transient DatastoreEntityExtension datastoreEntityExtension =
-      new DatastoreEntityExtension();
+      new DatastoreEntityExtension().allThreads(true);
 
   @RegisterExtension
   final transient TestPipelineExtension testPipeline =

core/src/test/java/google/registry/beam/invoicing/InvoicingPipelineTest.java

@@ -227,7 +227,8 @@ class InvoicingPipelineTest {
 
   @RegisterExtension
   @Order(Order.DEFAULT - 1)
-  final transient DatastoreEntityExtension datastore = new DatastoreEntityExtension();
+  final transient DatastoreEntityExtension datastore =
+      new DatastoreEntityExtension().allThreads(true);
 
   @RegisterExtension
   final TestPipelineExtension pipeline =

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

@@ -116,7 +116,8 @@ class Spec11PipelineTest {
 
   @RegisterExtension
   @Order(Order.DEFAULT - 1)
-  final transient DatastoreEntityExtension datastore = new DatastoreEntityExtension();
+  final transient DatastoreEntityExtension datastore =
+      new DatastoreEntityExtension().allThreads(true);
 
   @TempDir Path tmpDir;
 

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

@@ -189,8 +189,8 @@ public abstract class EntityTestCase {
     return fields;
   }
 
-  /** Verify indexing for an entity. */
-  public void verifyIndexing(Object obj, String... indexed) throws Exception {
+  /** Verify Datastore indexing for an entity. */
+  public void verifyDatastoreIndexing(Object obj, String... indexed) throws Exception {
     Set<String> indexedSet = ImmutableSet.copyOf(indexed);
     Set<String> allSet = getAllPotentiallyIndexedFieldPaths(obj.getClass());
     // Sanity test that the indexed fields we listed were found.

core/src/test/java/google/registry/model/billing/BillingEventTest.java

@@ -236,24 +236,24 @@ public class BillingEventTest extends EntityTestCase {
 
   @TestOfyOnly
   void testIndexing() throws Exception {
-    verifyIndexing(
+    verifyDatastoreIndexing(
         oneTime,
         "clientId",
         "eventTime",
         "billingTime",
         "syntheticCreationTime",
         "allocationToken");
-    verifyIndexing(
+    verifyDatastoreIndexing(
         oneTimeSynthetic,
         "clientId",
         "eventTime",
         "billingTime",
         "syntheticCreationTime",
         "allocationToken");
-    verifyIndexing(
+    verifyDatastoreIndexing(
         recurring, "clientId", "eventTime", "recurrenceEndTime", "recurrenceTimeOfYear.timeString");
-    verifyIndexing(cancellationOneTime, "clientId", "eventTime", "billingTime");
-    verifyIndexing(modification, "clientId", "eventTime");
+    verifyDatastoreIndexing(cancellationOneTime, "clientId", "eventTime", "billingTime");
+    verifyDatastoreIndexing(modification, "clientId", "eventTime");
   }
 
   @TestOfyAndSql

core/src/test/java/google/registry/model/common/CursorTest.java

@@ -75,7 +75,7 @@ public class CursorTest extends EntityTestCase {
     final DateTime time = DateTime.parse("2012-07-12T03:30:00.000Z");
     tm().transact(() -> tm().put(Cursor.createGlobal(RECURRING_BILLING, time)));
     Cursor cursor = tm().transact(() -> tm().loadByKey(Cursor.createGlobalVKey(RECURRING_BILLING)));
-    verifyIndexing(cursor);
+    verifyDatastoreIndexing(cursor);
   }
 
   @TestOfyAndSql

core/src/test/java/google/registry/model/contact/ContactResourceTest.java

@@ -17,15 +17,15 @@ package google.registry.model.contact;
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth8.assertThat;
 import static google.registry.model.EppResourceUtils.loadByForeignKey;
-import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.model.ImmutableObjectSubject.assertAboutImmutableObjects;
 import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
 import static google.registry.testing.ContactResourceSubject.assertAboutContacts;
 import static google.registry.testing.DatabaseHelper.cloneAndSetAutoTimestamps;
 import static google.registry.testing.DatabaseHelper.createTld;
 import static google.registry.testing.DatabaseHelper.insertInDb;
+import static google.registry.testing.DatabaseHelper.loadByEntity;
 import static google.registry.testing.DatabaseHelper.persistResource;
 import static google.registry.testing.SqlHelper.assertThrowForeignKeyViolation;
-import static google.registry.testing.SqlHelper.saveRegistrar;
 import static google.registry.util.DateTimeUtils.END_OF_TIME;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
@@ -33,7 +33,6 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
 import google.registry.model.EntityTestCase;
-import google.registry.model.ImmutableObjectSubject;
 import google.registry.model.contact.Disclose.PostalInfoChoice;
 import google.registry.model.contact.PostalInfo.Type;
 import google.registry.model.eppcommon.AuthInfo.PasswordAuth;
@@ -45,11 +44,14 @@ import google.registry.model.index.ForeignKeyIndex;
 import google.registry.model.index.ForeignKeyIndex.ForeignKeyContactIndex;
 import google.registry.model.transfer.ContactTransferData;
 import google.registry.model.transfer.TransferStatus;
-import google.registry.persistence.VKey;
+import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.TestOfyAndSql;
+import google.registry.testing.TestOfyOnly;
+import google.registry.testing.TestSqlOnly;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
 
 /** Unit tests for {@link ContactResource}. */
+@DualDatabaseTest
 public class ContactResourceTest extends EntityTestCase {
 
   private ContactResource originalContact;
@@ -66,11 +68,11 @@ public class ContactResourceTest extends EntityTestCase {
         new ContactResource.Builder()
             .setContactId("contact_id")
             .setRepoId("1-FOOBAR")
-            .setCreationRegistrarId("registrar1")
+            .setCreationRegistrarId("TheRegistrar")
             .setLastEppUpdateTime(fakeClock.nowUtc())
-            .setLastEppUpdateRegistrarId("registrar2")
+            .setLastEppUpdateRegistrarId("NewRegistrar")
             .setLastTransferTime(fakeClock.nowUtc())
-            .setPersistedCurrentSponsorRegistrarId("registrar3")
+            .setPersistedCurrentSponsorRegistrarId("NewRegistrar")
             .setLocalizedPostalInfo(
                 new PostalInfo.Builder()
                     .setType(Type.LOCALIZED)
@@ -116,8 +118,8 @@ public class ContactResourceTest extends EntityTestCase {
             .setStatusValues(ImmutableSet.of(StatusValue.OK))
             .setTransferData(
                 new ContactTransferData.Builder()
-                    .setGainingRegistrarId("gaining")
-                    .setLosingRegistrarId("losing")
+                    .setGainingRegistrarId("TheRegistrar")
+                    .setLosingRegistrarId("NewRegistrar")
                     .setPendingTransferExpirationTime(fakeClock.nowUtc())
                     .setTransferRequestTime(fakeClock.nowUtc())
                     .setTransferStatus(TransferStatus.SERVER_APPROVED)
@@ -128,34 +130,28 @@ public class ContactResourceTest extends EntityTestCase {
     contactResource = persistResource(cloneAndSetAutoTimestamps(originalContact));
   }
 
-  @Test
+  @TestOfyAndSql
   void testContactBaseToContactResource() {
-    ImmutableObjectSubject.assertAboutImmutableObjects()
+    assertAboutImmutableObjects()
         .that(new ContactResource.Builder().copyFrom(contactResource).build())
         .isEqualExceptFields(contactResource, "updateTimestamp", "revisions");
   }
 
-  @Test
+  @TestSqlOnly
   void testCloudSqlPersistence_failWhenViolateForeignKeyConstraint() {
-    assertThrowForeignKeyViolation(() -> insertInDb(originalContact));
+    assertThrowForeignKeyViolation(
+        () ->
+            insertInDb(
+                originalContact
+                    .asBuilder()
+                    .setRepoId("2-FOOBAR")
+                    .setCreationRegistrarId("nonexistent-registrar")
+                    .build()));
   }
 
-  @Test
+  @TestSqlOnly
   void testCloudSqlPersistence_succeed() {
-    saveRegistrar("registrar1");
-    saveRegistrar("registrar2");
-    saveRegistrar("registrar3");
-    saveRegistrar("gaining");
-    saveRegistrar("losing");
-    insertInDb(originalContact);
-    ContactResource persisted =
-        jpaTm()
-            .transact(
-                () ->
-                    jpaTm()
-                        .loadByKey(
-                            VKey.createSql(ContactResource.class, originalContact.getRepoId())));
-
+    ContactResource persisted = loadByEntity(originalContact);
     ContactResource fixed =
         originalContact
             .asBuilder()
@@ -167,12 +163,10 @@ public class ContactResourceTest extends EntityTestCase {
                     .setServerApproveEntities(null)
                     .build())
             .build();
-    ImmutableObjectSubject.assertAboutImmutableObjects()
-        .that(persisted)
-        .isEqualExceptFields(fixed, "updateTimestamp");
+    assertAboutImmutableObjects().that(persisted).isEqualExceptFields(fixed, "updateTimestamp");
   }
 
-  @Test
+  @TestOfyAndSql
   void testPersistence() {
     assertThat(
             loadByForeignKey(
@@ -180,12 +174,13 @@ public class ContactResourceTest extends EntityTestCase {
         .hasValue(contactResource);
   }
 
-  @Test
+  @TestOfyOnly
   void testIndexing() throws Exception {
-    verifyIndexing(contactResource, "deletionTime", "currentSponsorClientId", "searchName");
+    verifyDatastoreIndexing(
+        contactResource, "deletionTime", "currentSponsorClientId", "searchName");
   }
 
-  @Test
+  @TestOfyAndSql
   void testEmptyStringsBecomeNull() {
     assertThat(new ContactResource.Builder().setContactId(null).build().getContactId()).isNull();
     assertThat(new ContactResource.Builder().setContactId("").build().getContactId()).isNull();
@@ -217,7 +212,7 @@ public class ContactResourceTest extends EntityTestCase {
         .isNotNull();
   }
 
-  @Test
+  @TestOfyAndSql
   void testEmptyTransferDataBecomesNull() {
     ContactResource withNull = new ContactResource.Builder().setTransferData(null).build();
     ContactResource withEmpty =
@@ -226,7 +221,7 @@ public class ContactResourceTest extends EntityTestCase {
     assertThat(withEmpty.transferData).isNull();
   }
 
-  @Test
+  @TestOfyAndSql
   void testImplicitStatusValues() {
     // OK is implicit if there's no other statuses.
     assertAboutContacts()
@@ -248,7 +243,7 @@ public class ContactResourceTest extends EntityTestCase {
         .hasExactlyStatusValues(StatusValue.CLIENT_HOLD);
   }
 
-  @Test
+  @TestOfyAndSql
   void testExpiredTransfer() {
     ContactResource afterTransfer =
         contactResource
@@ -269,7 +264,7 @@ public class ContactResourceTest extends EntityTestCase {
     assertThat(afterTransfer.getLastTransferTime()).isEqualTo(fakeClock.nowUtc().plusDays(1));
   }
 
-  @Test
+  @TestOfyAndSql
   void testSetCreationTime_cantBeCalledTwice() {
     IllegalStateException thrown =
         assertThrows(
@@ -278,13 +273,13 @@ public class ContactResourceTest extends EntityTestCase {
     assertThat(thrown).hasMessageThat().contains("creationTime can only be set once");
   }
 
-  @Test
+  @TestOfyAndSql
   void testToHydratedString_notCircular() {
     // If there are circular references, this will overflow the stack.
     contactResource.toHydratedString();
   }
 
-  @Test
+  @TestOfyAndSql
   void testBeforeDatastoreSaveOnReplay_indexes() {
     ImmutableList<ForeignKeyContactIndex> foreignKeyIndexes =
         ofyTm().loadAllOf(ForeignKeyContactIndex.class);

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

@@ -19,6 +19,10 @@ import static google.registry.flows.domain.DomainTransferUtils.createPendingTran
 import static google.registry.model.ImmutableObjectSubject.assertAboutImmutableObjects;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.testing.DatabaseHelper.createTld;
+import static google.registry.testing.DatabaseHelper.insertInDb;
+import static google.registry.testing.DatabaseHelper.loadByEntity;
+import static google.registry.testing.DatabaseHelper.loadByKey;
+import static google.registry.testing.DatabaseHelper.updateInDb;
 import static google.registry.testing.SqlHelper.assertThrowForeignKeyViolation;
 import static google.registry.testing.SqlHelper.saveRegistrar;
 import static google.registry.util.DateTimeUtils.END_OF_TIME;
@@ -139,40 +143,19 @@ public class DomainBaseSqlTest {
   @TestSqlOnly
   void testDomainBasePersistence() {
     persistDomain();
-
-    jpaTm()
-        .transact(
-            () -> {
-              DomainBase result = jpaTm().loadByKey(domain.createVKey());
-              assertEqualDomainExcept(result);
-            });
+    assertEqualDomainExcept(loadByKey(domain.createVKey()));
   }
 
   @TestSqlOnly
   void testHostForeignKeyConstraints() {
-    assertThrowForeignKeyViolation(
-        () ->
-            jpaTm()
-                .transact(
-                    () -> {
-                      // Persist the domain without the associated host object.
-                      jpaTm().insert(contact);
-                      jpaTm().insert(contact2);
-                      jpaTm().insert(domain);
-                    }));
+    // Persist the domain without the associated host object.
+    assertThrowForeignKeyViolation(() -> insertInDb(contact, contact2, domain));
   }
 
   @TestSqlOnly
   void testContactForeignKeyConstraints() {
-    assertThrowForeignKeyViolation(
-        () ->
-            jpaTm()
-                .transact(
-                    () -> {
-                      // Persist the domain without the associated contact objects.
-                      jpaTm().insert(domain);
-                      jpaTm().insert(host);
-                    }));
+    // Persist the domain without the associated contact objects.
+    assertThrowForeignKeyViolation(() -> insertInDb(domain, host));
   }
 
   @TestSqlOnly
@@ -184,13 +167,8 @@ public class DomainBaseSqlTest {
               DomainBase persisted = jpaTm().loadByKey(domain.createVKey());
               jpaTm().put(persisted.asBuilder().build());
             });
-    jpaTm()
-        .transact(
-            () -> {
-              // Load the domain in its entirety.
-              DomainBase result = jpaTm().loadByKey(domain.createVKey());
-              assertEqualDomainExcept(result);
-            });
+    // Load the domain in its entirety.
+    assertEqualDomainExcept(loadByKey(domain.createVKey()));
   }
 
   @TestSqlOnly
@@ -379,24 +357,12 @@ public class DomainBaseSqlTest {
   @TestSqlOnly
   void testUpdates() {
     createTld("com");
-    jpaTm()
-        .transact(
-            () -> {
-              jpaTm().insert(contact);
-              jpaTm().insert(contact2);
-              jpaTm().insert(domain);
-              jpaTm().insert(host);
-            });
+    insertInDb(contact, contact2, domain, host);
     domain = domain.asBuilder().setNameservers(ImmutableSet.of()).build();
-    jpaTm().transact(() -> jpaTm().put(domain));
-    jpaTm()
-        .transact(
-            () -> {
-              DomainBase result = jpaTm().loadByKey(domain.createVKey());
-              assertAboutImmutableObjects()
-                  .that(result)
-                  .isEqualExceptFields(domain, "updateTimestamp", "creationTime");
-            });
+    updateInDb(domain);
+    assertAboutImmutableObjects()
+        .that(loadByEntity(domain))
+        .isEqualExceptFields(domain, "updateTimestamp", "creationTime");
   }
 
   static ContactResource makeContact(String repoId) {
@@ -410,128 +376,109 @@ public class DomainBaseSqlTest {
 
   private void persistDomain() {
     createTld("com");
-    jpaTm()
-        .transact(
-            () -> {
-              // Persist the contacts.  Note that these need to be persisted before the domain
-              // otherwise we get a foreign key constraint error.  If we ever decide to defer the
-              // relevant foreign key checks to commit time, then the order would not matter.
-              jpaTm().insert(contact);
-              jpaTm().insert(contact2);
-
-              // Persist the domain.
-              jpaTm().insert(domain);
-
-              // Persist the host.  This does _not_ need to be persisted before the domain,
-              // because only the row in the join table (DomainHost) is subject to foreign key
-              // constraints, and Hibernate knows to insert it after domain and host.
-              jpaTm().insert(host);
-            });
+    insertInDb(contact, contact2, domain, host);
   }
 
   @TestSqlOnly
   void persistDomainWithCompositeVKeys() {
     createTld("com");
-    jpaTm()
-        .transact(
-            () -> {
-              historyEntry =
-                  new DomainHistory.Builder()
-                      .setId(100L)
-                      .setType(HistoryEntry.Type.DOMAIN_CREATE)
-                      .setPeriod(Period.create(1, Period.Unit.YEARS))
-                      .setModificationTime(DateTime.now(UTC))
-                      .setDomainRepoId("4-COM")
-                      .setRegistrarId("registrar1")
-                      // These are non-null, but I don't think some tests set them.
-                      .setReason("felt like it")
-                      .setRequestedByRegistrar(false)
-                      .setXmlBytes(new byte[0])
-                      .build();
-              BillingEvent.Recurring billEvent =
-                  new BillingEvent.Recurring.Builder()
-                      .setId(200L)
-                      .setReason(Reason.RENEW)
-                      .setFlags(ImmutableSet.of(Flag.AUTO_RENEW))
-                      .setTargetId("example.com")
-                      .setRegistrarId("registrar1")
-                      .setDomainRepoId("4-COM")
-                      .setDomainHistoryRevisionId(1L)
-                      .setEventTime(DateTime.now(UTC).plusYears(1))
-                      .setRecurrenceEndTime(END_OF_TIME)
-                      .setParent(historyEntry)
-                      .build();
-              PollMessage.Autorenew autorenewPollMessage =
-                  new PollMessage.Autorenew.Builder()
-                      .setId(300L)
-                      .setRegistrarId("registrar1")
-                      .setEventTime(DateTime.now(UTC).plusYears(1))
-                      .setParent(historyEntry)
-                      .build();
-              PollMessage.OneTime deletePollMessage =
-                  new PollMessage.OneTime.Builder()
-                      .setId(400L)
-                      .setRegistrarId("registrar1")
-                      .setEventTime(DateTime.now(UTC).plusYears(1))
-                      .setParent(historyEntry)
-                      .build();
-              BillingEvent.OneTime oneTimeBillingEvent =
-                  new BillingEvent.OneTime.Builder()
-                      .setId(500L)
-                      // Use SERVER_STATUS so we don't have to add a period.
-                      .setReason(Reason.SERVER_STATUS)
-                      .setTargetId("example.com")
-                      .setRegistrarId("registrar1")
-                      .setDomainRepoId("4-COM")
-                      .setBillingTime(DateTime.now(UTC))
-                      .setCost(Money.of(USD, 100))
-                      .setEventTime(DateTime.now(UTC).plusYears(1))
-                      .setParent(historyEntry)
-                      .build();
-              DomainTransferData transferData =
-                  new DomainTransferData.Builder()
-                      .setServerApproveBillingEvent(oneTimeBillingEvent.createVKey())
-                      .setServerApproveAutorenewEvent(billEvent.createVKey())
-                      .setServerApproveAutorenewPollMessage(autorenewPollMessage.createVKey())
-                      .build();
-              gracePeriods =
-                  ImmutableSet.of(
-                      GracePeriod.create(
-                          GracePeriodStatus.ADD,
-                          "4-COM",
-                          END_OF_TIME,
-                          "registrar1",
-                          oneTimeBillingEvent.createVKey()),
-                      GracePeriod.createForRecurring(
-                          GracePeriodStatus.AUTO_RENEW,
-                          "4-COM",
-                          END_OF_TIME,
-                          "registrar1",
-                          billEvent.createVKey()));
+    historyEntry =
+        new DomainHistory.Builder()
+            .setId(100L)
+            .setType(HistoryEntry.Type.DOMAIN_CREATE)
+            .setPeriod(Period.create(1, Period.Unit.YEARS))
+            .setModificationTime(DateTime.now(UTC))
+            .setDomainRepoId("4-COM")
+            .setRegistrarId("registrar1")
+            // These are non-null, but I don't think some tests set them.
+            .setReason("felt like it")
+            .setRequestedByRegistrar(false)
+            .setXmlBytes(new byte[0])
+            .build();
+    BillingEvent.Recurring billEvent =
+        new BillingEvent.Recurring.Builder()
+            .setId(200L)
+            .setReason(Reason.RENEW)
+            .setFlags(ImmutableSet.of(Flag.AUTO_RENEW))
+            .setTargetId("example.com")
+            .setRegistrarId("registrar1")
+            .setDomainRepoId("4-COM")
+            .setDomainHistoryRevisionId(1L)
+            .setEventTime(DateTime.now(UTC).plusYears(1))
+            .setRecurrenceEndTime(END_OF_TIME)
+            .setParent(historyEntry)
+            .build();
+    PollMessage.Autorenew autorenewPollMessage =
+        new PollMessage.Autorenew.Builder()
+            .setId(300L)
+            .setRegistrarId("registrar1")
+            .setEventTime(DateTime.now(UTC).plusYears(1))
+            .setParent(historyEntry)
+            .build();
+    PollMessage.OneTime deletePollMessage =
+        new PollMessage.OneTime.Builder()
+            .setId(400L)
+            .setRegistrarId("registrar1")
+            .setEventTime(DateTime.now(UTC).plusYears(1))
+            .setParent(historyEntry)
+            .build();
+    BillingEvent.OneTime oneTimeBillingEvent =
+        new BillingEvent.OneTime.Builder()
+            .setId(500L)
+            // Use SERVER_STATUS so we don't have to add a period.
+            .setReason(Reason.SERVER_STATUS)
+            .setTargetId("example.com")
+            .setRegistrarId("registrar1")
+            .setDomainRepoId("4-COM")
+            .setBillingTime(DateTime.now(UTC))
+            .setCost(Money.of(USD, 100))
+            .setEventTime(DateTime.now(UTC).plusYears(1))
+            .setParent(historyEntry)
+            .build();
+    DomainTransferData transferData =
+        new DomainTransferData.Builder()
+            .setServerApproveBillingEvent(oneTimeBillingEvent.createVKey())
+            .setServerApproveAutorenewEvent(billEvent.createVKey())
+            .setServerApproveAutorenewPollMessage(autorenewPollMessage.createVKey())
+            .build();
+    gracePeriods =
+        ImmutableSet.of(
+            GracePeriod.create(
+                GracePeriodStatus.ADD,
+                "4-COM",
+                END_OF_TIME,
+                "registrar1",
+                oneTimeBillingEvent.createVKey()),
+            GracePeriod.createForRecurring(
+                GracePeriodStatus.AUTO_RENEW,
+                "4-COM",
+                END_OF_TIME,
+                "registrar1",
+                billEvent.createVKey()));
 
-              jpaTm().insert(contact);
-              jpaTm().insert(contact2);
-              jpaTm().insert(host);
-              domain =
-                  domain
-                      .asBuilder()
-                      .setAutorenewBillingEvent(billEvent.createVKey())
-                      .setAutorenewPollMessage(autorenewPollMessage.createVKey())
-                      .setDeletePollMessage(deletePollMessage.createVKey())
-                      .setTransferData(transferData)
-                      .setGracePeriods(gracePeriods)
-                      .build();
-              historyEntry = historyEntry.asBuilder().setDomain(domain).build();
-              jpaTm().insert(historyEntry);
-              jpaTm().insert(autorenewPollMessage);
-              jpaTm().insert(billEvent);
-              jpaTm().insert(deletePollMessage);
-              jpaTm().insert(oneTimeBillingEvent);
-              jpaTm().insert(domain);
-            });
+    domain =
+        domain
+            .asBuilder()
+            .setAutorenewBillingEvent(billEvent.createVKey())
+            .setAutorenewPollMessage(autorenewPollMessage.createVKey())
+            .setDeletePollMessage(deletePollMessage.createVKey())
+            .setTransferData(transferData)
+            .setGracePeriods(gracePeriods)
+            .build();
+    historyEntry = historyEntry.asBuilder().setDomain(domain).build();
+    insertInDb(
+        contact,
+        contact2,
+        host,
+        historyEntry,
+        autorenewPollMessage,
+        billEvent,
+        deletePollMessage,
+        oneTimeBillingEvent,
+        domain);
 
     // Store the existing BillingRecurrence VKey.  This happens after the event has been persisted.
-    DomainBase persisted = jpaTm().transact(() -> jpaTm().loadByKey(domain.createVKey()));
+    DomainBase persisted = loadByKey(domain.createVKey());
 
     // Verify that the domain data has been persisted.
     // dsData still isn't persisted.  gracePeriods appears to have the same values but for some
@@ -539,8 +486,7 @@ public class DomainBaseSqlTest {
     assertEqualDomainExcept(persisted, "creationTime", "dsData", "gracePeriods");
 
     // Verify that the DomainContent object from the history record sets the fields correctly.
-    DomainHistory persistedHistoryEntry =
-        jpaTm().transact(() -> jpaTm().loadByKey(historyEntry.createVKey()));
+    DomainHistory persistedHistoryEntry = loadByKey(historyEntry.createVKey());
     assertThat(persistedHistoryEntry.getDomainContent().get().getAutorenewPollMessage())
         .isEqualTo(domain.getAutorenewPollMessage());
     assertThat(persistedHistoryEntry.getDomainContent().get().getAutorenewBillingEvent())
@@ -562,114 +508,110 @@ public class DomainBaseSqlTest {
   @TestSqlOnly
   void persistDomainWithLegacyVKeys() {
     createTld("com");
-    jpaTm()
-        .transact(
-            () -> {
-              historyEntry =
-                  new DomainHistory.Builder()
-                      .setId(100L)
-                      .setType(HistoryEntry.Type.DOMAIN_CREATE)
-                      .setPeriod(Period.create(1, Period.Unit.YEARS))
-                      .setModificationTime(DateTime.now(UTC))
-                      .setDomainRepoId("4-COM")
-                      .setRegistrarId("registrar1")
-                      // These are non-null, but I don't think some tests set them.
-                      .setReason("felt like it")
-                      .setRequestedByRegistrar(false)
-                      .setXmlBytes(new byte[0])
-                      .build();
-              BillingEvent.Recurring billEvent =
-                  new BillingEvent.Recurring.Builder()
-                      .setId(200L)
-                      .setReason(Reason.RENEW)
-                      .setFlags(ImmutableSet.of(Flag.AUTO_RENEW))
-                      .setTargetId("example.com")
-                      .setRegistrarId("registrar1")
-                      .setDomainRepoId("4-COM")
-                      .setDomainHistoryRevisionId(1L)
-                      .setEventTime(DateTime.now(UTC).plusYears(1))
-                      .setRecurrenceEndTime(END_OF_TIME)
-                      .setParent(historyEntry)
-                      .build();
-              PollMessage.Autorenew autorenewPollMessage =
-                  new PollMessage.Autorenew.Builder()
-                      .setId(300L)
-                      .setRegistrarId("registrar1")
-                      .setEventTime(DateTime.now(UTC).plusYears(1))
-                      .setParent(historyEntry)
-                      .build();
-              PollMessage.OneTime deletePollMessage =
-                  new PollMessage.OneTime.Builder()
-                      .setId(400L)
-                      .setRegistrarId("registrar1")
-                      .setEventTime(DateTime.now(UTC).plusYears(1))
-                      .setParent(historyEntry)
-                      .build();
-              BillingEvent.OneTime oneTimeBillingEvent =
-                  new BillingEvent.OneTime.Builder()
-                      .setId(500L)
-                      // Use SERVER_STATUS so we don't have to add a period.
-                      .setReason(Reason.SERVER_STATUS)
-                      .setTargetId("example.com")
-                      .setRegistrarId("registrar1")
-                      .setDomainRepoId("4-COM")
-                      .setBillingTime(DateTime.now(UTC))
-                      .setCost(Money.of(USD, 100))
-                      .setEventTime(DateTime.now(UTC).plusYears(1))
-                      .setParent(historyEntry)
-                      .build();
-              DomainTransferData transferData =
-                  createPendingTransferData(
-                      new DomainTransferData.Builder()
-                          .setTransferRequestTrid(Trid.create("foo", "bar"))
-                          .setTransferRequestTime(fakeClock.nowUtc())
-                          .setGainingRegistrarId("registrar2")
-                          .setLosingRegistrarId("registrar1")
-                          .setPendingTransferExpirationTime(fakeClock.nowUtc().plusDays(1)),
-                      ImmutableSet.of(oneTimeBillingEvent, billEvent, autorenewPollMessage),
-                      Period.create(0, Unit.YEARS));
-              gracePeriods =
-                  ImmutableSet.of(
-                      GracePeriod.create(
-                          GracePeriodStatus.ADD,
-                          "4-COM",
-                          END_OF_TIME,
-                          "registrar1",
-                          oneTimeBillingEvent.createVKey()),
-                      GracePeriod.createForRecurring(
-                          GracePeriodStatus.AUTO_RENEW,
-                          "4-COM",
-                          END_OF_TIME,
-                          "registrar1",
-                          billEvent.createVKey()));
+    historyEntry =
+        new DomainHistory.Builder()
+            .setId(100L)
+            .setType(HistoryEntry.Type.DOMAIN_CREATE)
+            .setPeriod(Period.create(1, Period.Unit.YEARS))
+            .setModificationTime(DateTime.now(UTC))
+            .setDomainRepoId("4-COM")
+            .setRegistrarId("registrar1")
+            // These are non-null, but I don't think some tests set them.
+            .setReason("felt like it")
+            .setRequestedByRegistrar(false)
+            .setXmlBytes(new byte[0])
+            .build();
+    BillingEvent.Recurring billEvent =
+        new BillingEvent.Recurring.Builder()
+            .setId(200L)
+            .setReason(Reason.RENEW)
+            .setFlags(ImmutableSet.of(Flag.AUTO_RENEW))
+            .setTargetId("example.com")
+            .setRegistrarId("registrar1")
+            .setDomainRepoId("4-COM")
+            .setDomainHistoryRevisionId(1L)
+            .setEventTime(DateTime.now(UTC).plusYears(1))
+            .setRecurrenceEndTime(END_OF_TIME)
+            .setParent(historyEntry)
+            .build();
+    PollMessage.Autorenew autorenewPollMessage =
+        new PollMessage.Autorenew.Builder()
+            .setId(300L)
+            .setRegistrarId("registrar1")
+            .setEventTime(DateTime.now(UTC).plusYears(1))
+            .setParent(historyEntry)
+            .build();
+    PollMessage.OneTime deletePollMessage =
+        new PollMessage.OneTime.Builder()
+            .setId(400L)
+            .setRegistrarId("registrar1")
+            .setEventTime(DateTime.now(UTC).plusYears(1))
+            .setParent(historyEntry)
+            .build();
+    BillingEvent.OneTime oneTimeBillingEvent =
+        new BillingEvent.OneTime.Builder()
+            .setId(500L)
+            // Use SERVER_STATUS so we don't have to add a period.
+            .setReason(Reason.SERVER_STATUS)
+            .setTargetId("example.com")
+            .setRegistrarId("registrar1")
+            .setDomainRepoId("4-COM")
+            .setBillingTime(DateTime.now(UTC))
+            .setCost(Money.of(USD, 100))
+            .setEventTime(DateTime.now(UTC).plusYears(1))
+            .setParent(historyEntry)
+            .build();
+    DomainTransferData transferData =
+        createPendingTransferData(
+            new DomainTransferData.Builder()
+                .setTransferRequestTrid(Trid.create("foo", "bar"))
+                .setTransferRequestTime(fakeClock.nowUtc())
+                .setGainingRegistrarId("registrar2")
+                .setLosingRegistrarId("registrar1")
+                .setPendingTransferExpirationTime(fakeClock.nowUtc().plusDays(1)),
+            ImmutableSet.of(oneTimeBillingEvent, billEvent, autorenewPollMessage),
+            Period.create(0, Unit.YEARS));
+    gracePeriods =
+        ImmutableSet.of(
+            GracePeriod.create(
+                GracePeriodStatus.ADD,
+                "4-COM",
+                END_OF_TIME,
+                "registrar1",
+                oneTimeBillingEvent.createVKey()),
+            GracePeriod.createForRecurring(
+                GracePeriodStatus.AUTO_RENEW,
+                "4-COM",
+                END_OF_TIME,
+                "registrar1",
+                billEvent.createVKey()));
 
-              jpaTm().insert(contact);
-              jpaTm().insert(contact2);
-              jpaTm().insert(host);
-              domain =
-                  domain
-                      .asBuilder()
-                      .setAutorenewBillingEvent(
-                          createLegacyVKey(BillingEvent.Recurring.class, billEvent.getId()))
-                      .setAutorenewPollMessage(
-                          createLegacyVKey(
-                              PollMessage.Autorenew.class, autorenewPollMessage.getId()))
-                      .setDeletePollMessage(
-                          createLegacyVKey(PollMessage.OneTime.class, deletePollMessage.getId()))
-                      .setTransferData(transferData)
-                      .setGracePeriods(gracePeriods)
-                      .build();
-              historyEntry = historyEntry.asBuilder().setDomain(domain).build();
-              jpaTm().insert(historyEntry);
-              jpaTm().insert(autorenewPollMessage);
-              jpaTm().insert(billEvent);
-              jpaTm().insert(deletePollMessage);
-              jpaTm().insert(oneTimeBillingEvent);
-              jpaTm().insert(domain);
-            });
+    domain =
+        domain
+            .asBuilder()
+            .setAutorenewBillingEvent(
+                createLegacyVKey(BillingEvent.Recurring.class, billEvent.getId()))
+            .setAutorenewPollMessage(
+                createLegacyVKey(PollMessage.Autorenew.class, autorenewPollMessage.getId()))
+            .setDeletePollMessage(
+                createLegacyVKey(PollMessage.OneTime.class, deletePollMessage.getId()))
+            .setTransferData(transferData)
+            .setGracePeriods(gracePeriods)
+            .build();
+    historyEntry = historyEntry.asBuilder().setDomain(domain).build();
+    insertInDb(
+        contact,
+        contact2,
+        host,
+        historyEntry,
+        autorenewPollMessage,
+        billEvent,
+        deletePollMessage,
+        oneTimeBillingEvent,
+        domain);
 
     // Store the existing BillingRecurrence VKey.  This happens after the event has been persisted.
-    DomainBase persisted = jpaTm().transact(() -> jpaTm().loadByKey(domain.createVKey()));
+    DomainBase persisted = loadByKey(domain.createVKey());
 
     // Verify that the domain data has been persisted.
     // dsData still isn't persisted.  gracePeriods appears to have the same values but for some
@@ -677,8 +619,7 @@ public class DomainBaseSqlTest {
     assertEqualDomainExcept(persisted, "creationTime", "dsData", "gracePeriods");
 
     // Verify that the DomainContent object from the history record sets the fields correctly.
-    DomainHistory persistedHistoryEntry =
-        jpaTm().transact(() -> jpaTm().loadByKey(historyEntry.createVKey()));
+    DomainHistory persistedHistoryEntry = loadByKey(historyEntry.createVKey());
     assertThat(persistedHistoryEntry.getDomainContent().get().getAutorenewPollMessage())
         .isEqualTo(domain.getAutorenewPollMessage());
     assertThat(persistedHistoryEntry.getDomainContent().get().getAutorenewBillingEvent())

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

@@ -208,7 +208,7 @@ public class DomainBaseTest extends EntityTestCase {
 
   @Test
   void testIndexing() throws Exception {
-    verifyIndexing(
+    verifyDatastoreIndexing(
         domain,
         "allContacts.contact",
         "fullyQualifiedDomainName",

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

@@ -97,7 +97,7 @@ public class AllocationTokenTest extends EntityTestCase {
   void testIndexing() throws Exception {
     DomainBase domain = persistActiveDomain("blahdomain.foo");
     Key<HistoryEntry> historyEntryKey = Key.create(Key.create(domain), HistoryEntry.class, 1);
-    verifyIndexing(
+    verifyDatastoreIndexing(
         persistResource(
             new AllocationToken.Builder()
                 .setToken("abc123")

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

@@ -19,36 +19,40 @@ import static google.registry.model.ImmutableObjectSubject.assertAboutImmutableO
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatabaseHelper.insertInDb;
+import static google.registry.testing.DatabaseHelper.loadByEntity;
 import static google.registry.testing.DatabaseHelper.newContactResource;
 import static google.registry.testing.DatabaseHelper.newContactResourceWithRoid;
-import static google.registry.testing.SqlHelper.saveRegistrar;
 import static java.nio.charset.StandardCharsets.UTF_8;
 
+import com.google.common.collect.ImmutableList;
 import com.googlecode.objectify.Key;
 import google.registry.model.EntityTestCase;
+import google.registry.model.contact.ContactAddress;
 import google.registry.model.contact.ContactBase;
 import google.registry.model.contact.ContactHistory;
+import google.registry.model.contact.ContactPhoneNumber;
 import google.registry.model.contact.ContactResource;
+import google.registry.model.contact.PostalInfo;
 import google.registry.model.eppcommon.Trid;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.persistence.VKey;
-import org.junit.jupiter.api.Test;
+import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.TestOfyOnly;
+import google.registry.testing.TestSqlOnly;
 
 /** Tests for {@link ContactHistory}. */
+@DualDatabaseTest
 public class ContactHistoryTest extends EntityTestCase {
 
   ContactHistoryTest() {
     super(JpaEntityCoverageCheck.ENABLED);
   }
 
-  @Test
+  @TestSqlOnly
   void testPersistence() {
-    saveRegistrar("TheRegistrar");
-
     ContactResource contact = newContactResourceWithRoid("contactId", "contact1");
     insertInDb(contact);
-    VKey<ContactResource> contactVKey = contact.createVKey();
-    ContactResource contactFromDb = jpaTm().transact(() -> jpaTm().loadByKey(contactVKey));
+    ContactResource contactFromDb = loadByEntity(contact);
     ContactHistory contactHistory = createContactHistory(contactFromDb);
     insertInDb(contactHistory);
     jpaTm()
@@ -60,14 +64,11 @@ public class ContactHistoryTest extends EntityTestCase {
             });
   }
 
-  @Test
+  @TestSqlOnly
   void testLegacyPersistence_nullContactBase() {
-    saveRegistrar("TheRegistrar");
-
     ContactResource contact = newContactResourceWithRoid("contactId", "contact1");
     insertInDb(contact);
-    VKey<ContactResource> contactVKey = contact.createVKey();
-    ContactResource contactFromDb = jpaTm().transact(() -> jpaTm().loadByKey(contactVKey));
+    ContactResource contactFromDb = loadByEntity(contact);
     ContactHistory contactHistory =
         createContactHistory(contactFromDb).asBuilder().setContact(null).build();
     insertInDb(contactHistory);
@@ -81,10 +82,8 @@ public class ContactHistoryTest extends EntityTestCase {
             });
   }
 
-  @Test
+  @TestOfyOnly
   void testOfyPersistence() {
-    saveRegistrar("TheRegistrar");
-
     ContactResource contact = newContactResourceWithRoid("contactId", "contact1");
     tm().transact(() -> tm().insert(contact));
     VKey<ContactResource> contactVKey = contact.createVKey();
@@ -105,9 +104,8 @@ public class ContactHistoryTest extends EntityTestCase {
     assertThat(hostHistoryFromDb).isEqualTo(historyEntryFromDb);
   }
 
-  @Test
+  @TestSqlOnly
   void testBeforeSqlSave_afterContactPersisted() {
-    saveRegistrar("TheRegistrar");
     ContactResource contactResource = newContactResource("contactId");
     ContactHistory contactHistory =
         new ContactHistory.Builder()
@@ -136,6 +134,60 @@ public class ContactHistoryTest extends EntityTestCase {
                     .hasFieldsEqualTo(jpaTm().loadByEntity(contactHistory).getContactBase().get()));
   }
 
+  @TestSqlOnly
+  void testWipeOutPii_assertsAllPiiFieldsAreNull() {
+    ContactHistory originalEntity =
+        createContactHistory(
+            new ContactResource.Builder()
+                .setRepoId("1-FOOBAR")
+                .setLocalizedPostalInfo(
+                    new PostalInfo.Builder()
+                        .setType(PostalInfo.Type.LOCALIZED)
+                        .setAddress(
+                            new ContactAddress.Builder()
+                                .setStreet(ImmutableList.of("111 8th Ave", "4th Floor"))
+                                .setCity("New York")
+                                .setState("NY")
+                                .setZip("10011")
+                                .setCountryCode("US")
+                                .build())
+                        .build())
+                .setInternationalizedPostalInfo(
+                    new PostalInfo.Builder()
+                        .setType(PostalInfo.Type.INTERNATIONALIZED)
+                        .setAddress(
+                            new ContactAddress.Builder()
+                                .setStreet(ImmutableList.of("111 8th Ave", "4th Floor"))
+                                .setCity("New York")
+                                .setState("NY")
+                                .setZip("10011")
+                                .setCountryCode("US")
+                                .build())
+                        .build())
+                .setVoiceNumber(new ContactPhoneNumber.Builder().setPhoneNumber("867-5309").build())
+                .setFaxNumber(
+                    new ContactPhoneNumber.Builder()
+                        .setPhoneNumber("867-5309")
+                        .setExtension("1000")
+                        .build())
+                .setEmailAddress("test@example.com")
+                .build());
+
+    assertThat(originalEntity.getContactBase().get().getEmailAddress()).isNotNull();
+    assertThat(originalEntity.getContactBase().get().getLocalizedPostalInfo()).isNotNull();
+    assertThat(originalEntity.getContactBase().get().getInternationalizedPostalInfo()).isNotNull();
+    assertThat(originalEntity.getContactBase().get().getVoiceNumber()).isNotNull();
+    assertThat(originalEntity.getContactBase().get().getFaxNumber()).isNotNull();
+
+    ContactHistory wipedEntity = originalEntity.asBuilder().wipeOutPii().build();
+
+    assertThat(wipedEntity.getContactBase().get().getEmailAddress()).isNull();
+    assertThat(wipedEntity.getContactBase().get().getLocalizedPostalInfo()).isNull();
+    assertThat(wipedEntity.getContactBase().get().getInternationalizedPostalInfo()).isNull();
+    assertThat(wipedEntity.getContactBase().get().getVoiceNumber()).isNull();
+    assertThat(wipedEntity.getContactBase().get().getFaxNumber()).isNull();
+  }
+
   private ContactHistory createContactHistory(ContactBase contact) {
     return new ContactHistory.Builder()
         .setType(HistoryEntry.Type.HOST_CREATE)

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

@@ -26,6 +26,7 @@ import static google.registry.testing.DatabaseHelper.insertInDb;
 import static google.registry.testing.DatabaseHelper.newContactResourceWithRoid;
 import static google.registry.testing.DatabaseHelper.newDomainBase;
 import static google.registry.testing.DatabaseHelper.newHostResourceWithRoid;
+import static google.registry.testing.DatabaseHelper.putInDb;
 import static google.registry.util.DateTimeUtils.END_OF_TIME;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
 import static java.nio.charset.StandardCharsets.UTF_8;
@@ -146,14 +147,8 @@ public class DomainHistoryTest extends EntityTestCase {
             "tld", "TLD", ImmutableSortedMap.of(START_OF_TIME, GENERAL_AVAILABILITY));
     tm().transact(() -> tm().insert(registry));
     Registries.resetCache();
-    jpaTm()
-        .transact(
-            () -> {
-              jpaTm().insert(registry);
-              jpaTm()
-                  .insert(
-                      makeRegistrar2().asBuilder().setAllowedTlds(ImmutableSet.of("tld")).build());
-            });
+    insertInDb(
+        registry, makeRegistrar2().asBuilder().setAllowedTlds(ImmutableSet.of("tld")).build());
 
     HostResource host = newHostResourceWithRoid("ns1.example.com", "host1");
     ContactResource contact = newContactResourceWithRoid("contactId", "contact1");
@@ -164,12 +159,7 @@ public class DomainHistoryTest extends EntityTestCase {
               tm().insert(host);
               tm().insert(contact);
             });
-    jpaTm()
-        .transact(
-            () -> {
-              jpaTm().insert(host);
-              jpaTm().insert(contact);
-            });
+    insertInDb(host, contact);
     DomainBase domain =
         newDomainBase("example.tld", "domainRepoId", contact)
             .asBuilder()
@@ -190,7 +180,7 @@ public class DomainHistoryTest extends EntityTestCase {
 
     // Reload and rewrite.
     DomainHistory domainHistoryFromDb2 = tm().transact(() -> tm().loadByKey(domainHistoryVKey));
-    jpaTm().transact(() -> jpaTm().put(domainHistoryFromDb2));
+    putInDb(domainHistoryFromDb2);
   }
 
   @TestSqlOnly

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

@@ -19,6 +19,7 @@ import static google.registry.model.ImmutableObjectSubject.assertAboutImmutableO
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatabaseHelper.insertInDb;
+import static google.registry.testing.DatabaseHelper.loadByEntity;
 import static google.registry.testing.DatabaseHelper.newHostResource;
 import static google.registry.testing.DatabaseHelper.newHostResourceWithRoid;
 import static google.registry.testing.SqlHelper.saveRegistrar;
@@ -32,24 +33,23 @@ import google.registry.model.host.HostHistory;
 import google.registry.model.host.HostResource;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.persistence.VKey;
-import org.junit.jupiter.api.Test;
+import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.TestOfyOnly;
+import google.registry.testing.TestSqlOnly;
 
 /** Tests for {@link HostHistory}. */
+@DualDatabaseTest
 public class HostHistoryTest extends EntityTestCase {
 
   HostHistoryTest() {
     super(JpaEntityCoverageCheck.ENABLED);
   }
 
-  @Test
+  @TestSqlOnly
   void testPersistence() {
-    saveRegistrar("TheRegistrar");
-
     HostResource host = newHostResourceWithRoid("ns1.example.com", "host1");
     insertInDb(host);
-    VKey<HostResource> hostVKey =
-        VKey.create(HostResource.class, "host1", Key.create(HostResource.class, "host1"));
-    HostResource hostFromDb = jpaTm().transact(() -> jpaTm().loadByKey(hostVKey));
+    HostResource hostFromDb = loadByEntity(host);
     HostHistory hostHistory = createHostHistory(hostFromDb);
     insertInDb(hostHistory);
     jpaTm()
@@ -61,13 +61,12 @@ public class HostHistoryTest extends EntityTestCase {
             });
   }
 
-  @Test
+  @TestSqlOnly
   void testLegacyPersistence_nullHostBase() {
-    saveRegistrar("TheRegistrar");
     HostResource host = newHostResourceWithRoid("ns1.example.com", "host1");
     insertInDb(host);
 
-    HostResource hostFromDb = jpaTm().transact(() -> jpaTm().loadByKey(host.createVKey()));
+    HostResource hostFromDb = loadByEntity(host);
     HostHistory hostHistory = createHostHistory(hostFromDb).asBuilder().setHost(null).build();
     insertInDb(hostHistory);
 
@@ -80,7 +79,7 @@ public class HostHistoryTest extends EntityTestCase {
             });
   }
 
-  @Test
+  @TestOfyOnly
   void testOfySave() {
     saveRegistrar("registrar1");
 
@@ -105,9 +104,8 @@ public class HostHistoryTest extends EntityTestCase {
     assertThat(hostHistoryFromDb).isEqualTo(historyEntryFromDb);
   }
 
-  @Test
+  @TestSqlOnly
   void testBeforeSqlSave_afterHostPersisted() {
-    saveRegistrar("TheRegistrar");
     HostResource hostResource = newHostResource("ns1.example.tld");
     HostHistory hostHistory =
         new HostHistory.Builder()

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

@@ -20,6 +20,7 @@ import static google.registry.persistence.transaction.TransactionManagerFactory.
 import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
 import static google.registry.testing.DatabaseHelper.createTld;
 import static google.registry.testing.DatabaseHelper.insertInDb;
+import static google.registry.testing.DatabaseHelper.loadByKey;
 import static google.registry.testing.DatabaseHelper.newContactResourceWithRoid;
 import static google.registry.testing.DatabaseHelper.newHostResourceWithRoid;
 import static google.registry.util.CollectionUtils.nullToEmpty;
@@ -79,14 +80,8 @@ public class LegacyHistoryObjectTest extends EntityTestCase {
     ContactHistory legacyContactHistory = (ContactHistory) fromObjectify;
 
     // Next, save that from-Datastore object in SQL and verify we can load it back in
-    jpaTm()
-        .transact(
-            () -> {
-              jpaTm().insert(contact);
-              jpaTm().insert(legacyContactHistory);
-            });
-    ContactHistory legacyHistoryFromSql =
-        jpaTm().transact(() -> jpaTm().loadByKey(legacyContactHistory.createVKey()));
+    insertInDb(contact, legacyContactHistory);
+    ContactHistory legacyHistoryFromSql = loadByKey(legacyContactHistory.createVKey());
     assertAboutImmutableObjects()
         .that(legacyContactHistory)
         .isEqualExceptFields(legacyHistoryFromSql);
@@ -171,14 +166,8 @@ public class LegacyHistoryObjectTest extends EntityTestCase {
     HostHistory legacyHostHistory = (HostHistory) fromObjectify;
 
     // Next, save that from-Datastore object in SQL and verify we can load it back in
-    jpaTm()
-        .transact(
-            () -> {
-              jpaTm().insert(host);
-              jpaTm().insert(legacyHostHistory);
-            });
-    HostHistory legacyHistoryFromSql =
-        jpaTm().transact(() -> jpaTm().loadByKey(legacyHostHistory.createVKey()));
+    insertInDb(host, legacyHostHistory);
+    HostHistory legacyHistoryFromSql = loadByKey(legacyHostHistory.createVKey());
     assertAboutImmutableObjects().that(legacyHostHistory).isEqualExceptFields(legacyHistoryFromSql);
     // can't compare hostRepoId directly since it doesn't save the ofy key in SQL
     assertThat(legacyHostHistory.getParentVKey().getSqlKey())

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

@@ -121,7 +121,7 @@ class HostResourceTest extends EntityTestCase {
   void testIndexing() throws Exception {
     // Clone it and save it before running the indexing test so that its transferData fields are
     // populated from the superordinate domain.
-    verifyIndexing(
+    verifyDatastoreIndexing(
         persistResource(host),
         "deletionTime",
         "fullyQualifiedHostName",

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

@@ -49,7 +49,7 @@ class EppResourceIndexTest extends EntityTestCase {
 
   @Test
   void testIndexing() throws Exception {
-    verifyIndexing(Iterables.getOnlyElement(getEppResourceIndexObjects()), "kind");
+    verifyDatastoreIndexing(Iterables.getOnlyElement(getEppResourceIndexObjects()), "kind");
   }
 
   @Test

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

@@ -79,7 +79,7 @@ class ForeignKeyIndexTest extends EntityTestCase {
   void testIndexing() throws Exception {
     // Persist a host and implicitly persist a ForeignKeyIndex for it.
     persistActiveHost("ns1.example.com");
-    verifyIndexing(
+    verifyDatastoreIndexing(
         ForeignKeyIndex.load(HostResource.class, "ns1.example.com", fakeClock.nowUtc()),
         "deletionTime");
   }

core/src/test/java/google/registry/model/poll/PollMessageTest.java

@@ -15,10 +15,10 @@
 package google.registry.model.poll;
 
 import static com.google.common.truth.Truth.assertThat;
-import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatabaseHelper.createTld;
 import static google.registry.testing.DatabaseHelper.insertInDb;
+import static google.registry.testing.DatabaseHelper.loadByKey;
 import static google.registry.testing.DatabaseHelper.newDomainBase;
 import static google.registry.testing.DatabaseHelper.persistActiveContact;
 import static google.registry.testing.DatabaseHelper.persistResource;
@@ -94,17 +94,13 @@ public class PollMessageTest extends EntityTestCase {
   @TestSqlOnly
   void testCloudSqlSupportForPolymorphicVKey() {
     insertInDb(oneTime);
-    PollMessage persistedOneTime =
-        jpaTm()
-            .transact(() -> jpaTm().loadByKey(VKey.createSql(PollMessage.class, oneTime.getId())));
+    PollMessage persistedOneTime = loadByKey(VKey.createSql(PollMessage.class, oneTime.getId()));
     assertThat(persistedOneTime).isInstanceOf(PollMessage.OneTime.class);
     assertThat(persistedOneTime).isEqualTo(oneTime);
 
     insertInDb(autoRenew);
     PollMessage persistedAutoRenew =
-        jpaTm()
-            .transact(
-                () -> jpaTm().loadByKey(VKey.createSql(PollMessage.class, autoRenew.getId())));
+        loadByKey(VKey.createSql(PollMessage.class, autoRenew.getId()));
     assertThat(persistedAutoRenew).isInstanceOf(PollMessage.Autorenew.class);
     assertThat(persistedAutoRenew).isEqualTo(autoRenew);
   }
@@ -149,6 +145,6 @@ public class PollMessageTest extends EntityTestCase {
                 .setAutorenewEndTime(fakeClock.nowUtc().plusDays(365))
                 .setTargetId("foobar.foo")
                 .build());
-    verifyIndexing(pollMessage);
+    verifyDatastoreIndexing(pollMessage);
   }
 }

core/src/test/java/google/registry/model/registrar/RegistrarTest.java

@@ -137,7 +137,7 @@ class RegistrarTest extends EntityTestCase {
 
   @TestOfyOnly
   void testIndexing() throws Exception {
-    verifyIndexing(registrar, "registrarName", "ianaIdentifier");
+    verifyDatastoreIndexing(registrar, "registrarName", "ianaIdentifier");
   }
 
   @TestOfyAndSql

core/src/test/java/google/registry/model/replay/ReplicateToDatastoreActionTest.java

@@ -100,12 +100,7 @@ public class ReplicateToDatastoreActionTest {
     TestObject bar = TestObject.create("bar");
     TestObject baz = TestObject.create("baz");
 
-    jpaTm()
-        .transact(
-            () -> {
-              jpaTm().insert(foo);
-              jpaTm().insert(bar);
-            });
+    insertInDb(foo, bar);
     runAndVerifySuccess();
 
     assertThat(ofyTm().transact(() -> ofyTm().loadByKey(foo.key()))).isEqualTo(foo);
@@ -219,7 +214,7 @@ public class ReplicateToDatastoreActionTest {
   @Test
   void testBeforeDatastoreSaveCallback() {
     TestObject testObject = TestObject.create("foo");
-    jpaTm().transact(() -> jpaTm().put(testObject));
+    insertInDb(testObject);
     action.run();
     assertThat(ofyTm().loadAllOf(TestObject.class)).containsExactly(testObject);
     assertThat(TestObject.beforeDatastoreSaveCallCount).isEqualTo(1);

core/src/test/java/google/registry/model/reporting/HistoryEntryTest.java

@@ -157,6 +157,6 @@ class HistoryEntryTest extends EntityTestCase {
 
   @TestOfyOnly
   void testIndexing() throws Exception {
-    verifyIndexing(domainHistory.asHistoryEntry(), "modificationTime", "clientId");
+    verifyDatastoreIndexing(domainHistory.asHistoryEntry(), "modificationTime", "clientId");
   }
 }

core/src/test/java/google/registry/model/reporting/Spec11ThreatMatchTest.java

@@ -14,11 +14,12 @@
 
 package google.registry.model.reporting;
 
-import static com.google.common.truth.Truth.assertThat;
+import static google.registry.model.ImmutableObjectSubject.assertAboutImmutableObjects;
 import static google.registry.model.reporting.Spec11ThreatMatch.ThreatType.MALWARE;
 import static google.registry.model.reporting.Spec11ThreatMatch.ThreatType.UNWANTED_SOFTWARE;
-import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.testing.DatabaseHelper.createTld;
+import static google.registry.testing.DatabaseHelper.insertInDb;
+import static google.registry.testing.DatabaseHelper.loadByEntity;
 import static google.registry.testing.SqlHelper.assertThrowForeignKeyViolation;
 import static google.registry.testing.SqlHelper.saveRegistrar;
 import static org.junit.jupiter.api.Assertions.assertThrows;
@@ -107,53 +108,20 @@ public final class Spec11ThreatMatchTest extends EntityTestCase {
   void testPersistence() {
     createTld("tld");
     saveRegistrar(REGISTRAR_ID);
-
-    jpaTm()
-        .transact(
-            () -> {
-              jpaTm().insert(registrantContact);
-              jpaTm().insert(domain);
-              jpaTm().insert(host);
-              jpaTm().insert(threat);
-            });
-
-    VKey<Spec11ThreatMatch> threatVKey = VKey.createSql(Spec11ThreatMatch.class, threat.getId());
-    Spec11ThreatMatch persistedThreat = jpaTm().transact(() -> jpaTm().loadByKey(threatVKey));
-
-    // Threat object saved for the first time doesn't have an ID; it is generated by SQL
-    threat.id = persistedThreat.id;
-    assertThat(threat).isEqualTo(persistedThreat);
+    insertInDb(registrantContact, domain, host, threat);
+    assertAboutImmutableObjects().that(loadByEntity(threat)).isEqualExceptFields(threat, "id");
   }
 
   @TestSqlOnly
   @Disabled("We can't rely on foreign keys until we've migrated to SQL")
   void testThreatForeignKeyConstraints() {
-    assertThrowForeignKeyViolation(
-        () -> {
-          jpaTm()
-              .transact(
-                  () -> {
-                    // Persist the threat without the associated registrar.
-                    jpaTm().insert(host);
-                    jpaTm().insert(registrantContact);
-                    jpaTm().insert(domain);
-                    jpaTm().insert(threat);
-                  });
-        });
+    // Persist the threat without the associated registrar.
+    assertThrowForeignKeyViolation(() -> insertInDb(host, registrantContact, domain, threat));
 
     saveRegistrar(REGISTRAR_ID);
 
-    assertThrowForeignKeyViolation(
-        () -> {
-          jpaTm()
-              .transact(
-                  () -> {
-                    // Persist the threat without the associated domain.
-                    jpaTm().insert(registrantContact);
-                    jpaTm().insert(host);
-                    jpaTm().insert(threat);
-                  });
-        });
+    // Persist the threat without the associated domain.
+    assertThrowForeignKeyViolation(() -> insertInDb(registrantContact, host, threat));
   }
 
   @TestOfyAndSql

core/src/test/java/google/registry/model/tld/RegistryTest.java

@@ -95,7 +95,7 @@ public final class RegistryTest extends EntityTestCase {
 
   @TestOfyOnly
   void testIndexing() throws Exception {
-    verifyIndexing(Registry.get("tld"));
+    verifyDatastoreIndexing(Registry.get("tld"));
   }
 
   @TestOfyAndSql

core/src/test/java/google/registry/persistence/transaction/CriteriaQueryBuilderTest.java

@@ -16,6 +16,7 @@ package google.registry.persistence.transaction;
 
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.testing.DatabaseHelper.insertInDb;
 
 import com.google.common.collect.ImmutableList;
 import google.registry.model.ImmutableObject;
@@ -49,7 +50,7 @@ class CriteriaQueryBuilderTest {
 
   @BeforeEach
   void beforeEach() {
-    jpaTm().transact(() -> jpaTm().putAll(ImmutableList.of(entity1, entity2, entity3)));
+    insertInDb(entity1, entity2, entity3);
   }
 
   @Test

core/src/test/java/google/registry/persistence/transaction/JpaTransactionManagerImplTest.java

@@ -18,7 +18,9 @@ import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.testing.DatabaseHelper.assertDetachedFromEntityManager;
+import static google.registry.testing.DatabaseHelper.existsInDb;
 import static google.registry.testing.DatabaseHelper.insertInDb;
+import static google.registry.testing.DatabaseHelper.loadByKey;
 import static google.registry.testing.TestDataHelper.fileClassPath;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
@@ -46,6 +48,8 @@ import javax.persistence.IdClass;
 import javax.persistence.OptimisticLockException;
 import javax.persistence.RollbackException;
 import org.hibernate.exception.JDBCConnectionException;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
@@ -80,6 +84,16 @@ class JpaTransactionManagerImplTest {
               TestEntity.class, TestCompoundIdEntity.class, TestNamedCompoundIdEntity.class)
           .buildUnitTestExtension();
 
+  @BeforeEach
+  void beforeEach() {
+    TransactionManagerFactory.setTmForTest(jpaTm());
+  }
+
+  @AfterEach
+  void afterEach() {
+    TransactionManagerFactory.removeTmOverrideForTest();
+  }
+
   @Test
   void transact_succeeds() {
     assertPersonEmpty();
@@ -139,10 +153,10 @@ class JpaTransactionManagerImplTest {
 
   @Test
   void insert_succeeds() {
-    assertThat(jpaTm().transact(() -> jpaTm().exists(theEntity))).isFalse();
-    insertInDb(theEntity);
-    assertThat(jpaTm().transact(() -> jpaTm().exists(theEntity))).isTrue();
-    assertThat(jpaTm().transact(() -> jpaTm().loadByKey(theEntityKey))).isEqualTo(theEntity);
+    assertThat(existsInDb(theEntity)).isFalse();
+    jpaTm().transact(() -> jpaTm().insert(theEntity));
+    assertThat(existsInDb(theEntity)).isTrue();
+    assertThat(loadByKey(theEntityKey)).isEqualTo(theEntity);
   }
 
   @Test
@@ -258,17 +272,17 @@ class JpaTransactionManagerImplTest {
 
   @Test
   void insert_throwsExceptionIfEntityExists() {
-    assertThat(jpaTm().transact(() -> jpaTm().exists(theEntity))).isFalse();
-    insertInDb(theEntity);
-    assertThat(jpaTm().transact(() -> jpaTm().exists(theEntity))).isTrue();
-    assertThat(jpaTm().transact(() -> jpaTm().loadByKey(theEntityKey))).isEqualTo(theEntity);
-    assertThrows(RollbackException.class, () -> insertInDb(theEntity));
+    assertThat(existsInDb(theEntity)).isFalse();
+    jpaTm().transact(() -> jpaTm().insert(theEntity));
+    assertThat(existsInDb(theEntity)).isTrue();
+    assertThat(loadByKey(theEntityKey)).isEqualTo(theEntity);
+    assertThrows(RollbackException.class, () -> jpaTm().transact(() -> jpaTm().insert(theEntity)));
   }
 
   @Test
   void createCompoundIdEntity_succeeds() {
     assertThat(jpaTm().transact(() -> jpaTm().exists(compoundIdEntity))).isFalse();
-    insertInDb(compoundIdEntity);
+    jpaTm().transact(() -> jpaTm().insert(compoundIdEntity));
     assertThat(jpaTm().transact(() -> jpaTm().exists(compoundIdEntity))).isTrue();
     assertThat(jpaTm().transact(() -> jpaTm().loadByKey(compoundIdEntityKey)))
         .isEqualTo(compoundIdEntity);
@@ -278,18 +292,12 @@ class JpaTransactionManagerImplTest {
   void createNamedCompoundIdEntity_succeeds() {
     // Compound IDs should also work even if the field names don't match up exactly
     TestNamedCompoundIdEntity entity = new TestNamedCompoundIdEntity("foo", 1);
-    insertInDb(entity);
-    jpaTm()
-        .transact(
-            () -> {
-              assertThat(jpaTm().exists(entity)).isTrue();
-              assertThat(
-                      jpaTm()
-                          .loadByKey(
-                              VKey.createSql(
-                                  TestNamedCompoundIdEntity.class, new NamedCompoundId("foo", 1))))
-                  .isEqualTo(entity);
-            });
+    jpaTm().transact(() -> jpaTm().insert(entity));
+    assertThat(existsInDb(entity)).isTrue();
+    assertThat(
+            loadByKey(
+                VKey.createSql(TestNamedCompoundIdEntity.class, new NamedCompoundId("foo", 1))))
+        .isEqualTo(entity);
   }
 
   @Test

core/src/test/java/google/registry/schema/registrar/RegistrarContactTest.java

@@ -17,7 +17,9 @@ package google.registry.schema.registrar;
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.registrar.RegistrarContact.Type.WHOIS;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.setTmForTest;
 import static google.registry.testing.DatabaseHelper.insertInDb;
+import static google.registry.testing.DatabaseHelper.loadByEntity;
 import static google.registry.testing.SqlHelper.saveRegistrar;
 
 import com.google.common.collect.ImmutableSet;
@@ -48,6 +50,7 @@ class RegistrarContactTest {
 
   @BeforeEach
   public void beforeEach() {
+    setTmForTest(jpaTm());
     testRegistrar = saveRegistrar("registrarId");
     testRegistrarPoc =
         new RegistrarContact.Builder()
@@ -67,8 +70,6 @@ class RegistrarContactTest {
   @Test
   void testPersistence_succeeds() {
     insertInDb(testRegistrarPoc);
-    RegistrarContact persisted =
-        jpaTm().transact(() -> jpaTm().loadByKey(testRegistrarPoc.createVKey()));
-    assertThat(persisted).isEqualTo(testRegistrarPoc);
+    assertThat(loadByEntity(testRegistrarPoc)).isEqualTo(testRegistrarPoc);
   }
 }

core/src/test/java/google/registry/schema/registrar/RegistrarDaoTest.java

@@ -16,7 +16,10 @@ package google.registry.schema.registrar;
 
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.testing.DatabaseHelper.existsInDb;
 import static google.registry.testing.DatabaseHelper.insertInDb;
+import static google.registry.testing.DatabaseHelper.loadByKey;
+import static google.registry.testing.DatabaseHelper.updateInDb;
 import static org.joda.time.DateTimeZone.UTC;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
@@ -26,9 +29,11 @@ import google.registry.model.registrar.RegistrarAddress;
 import google.registry.persistence.VKey;
 import google.registry.persistence.transaction.JpaTestExtensions;
 import google.registry.persistence.transaction.JpaTestExtensions.JpaIntegrationWithCoverageExtension;
+import google.registry.persistence.transaction.TransactionManagerFactory;
 import google.registry.testing.DatastoreEntityExtension;
 import google.registry.testing.FakeClock;
 import org.joda.time.DateTime;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Order;
 import org.junit.jupiter.api.Test;
@@ -52,7 +57,8 @@ public class RegistrarDaoTest {
   private Registrar testRegistrar;
 
   @BeforeEach
-  void setUp() {
+  void beforeEach() {
+    TransactionManagerFactory.setTmForTest(jpaTm());
     testRegistrar =
         new Registrar.Builder()
             .setType(Registrar.Type.TEST)
@@ -69,41 +75,39 @@ public class RegistrarDaoTest {
             .build();
   }
 
+  @AfterEach
+  void afterEach() {
+    TransactionManagerFactory.removeTmOverrideForTest();
+  }
+
   @Test
   void saveNew_worksSuccessfully() {
-    assertThat(jpaTm().transact(() -> jpaTm().exists(testRegistrar))).isFalse();
+    assertThat(existsInDb(testRegistrar)).isFalse();
     insertInDb(testRegistrar);
-    assertThat(jpaTm().transact(() -> jpaTm().exists(testRegistrar))).isTrue();
+    assertThat(existsInDb(testRegistrar)).isTrue();
   }
 
   @Test
   void update_worksSuccessfully() {
     insertInDb(testRegistrar);
-    Registrar persisted = jpaTm().transact(() -> jpaTm().loadByKey(registrarKey));
+    Registrar persisted = loadByKey(registrarKey);
     assertThat(persisted.getRegistrarName()).isEqualTo("registrarName");
-    jpaTm()
-        .transact(
-            () ->
-                jpaTm()
-                    .update(
-                        persisted.asBuilder().setRegistrarName("changedRegistrarName").build()));
-    Registrar updated = jpaTm().transact(() -> jpaTm().loadByKey(registrarKey));
+    updateInDb(persisted.asBuilder().setRegistrarName("changedRegistrarName").build());
+    Registrar updated = loadByKey(registrarKey);
     assertThat(updated.getRegistrarName()).isEqualTo("changedRegistrarName");
   }
 
   @Test
   void update_throwsExceptionWhenEntityDoesNotExist() {
-    assertThat(jpaTm().transact(() -> jpaTm().exists(testRegistrar))).isFalse();
-    assertThrows(
-        IllegalArgumentException.class,
-        () -> jpaTm().transact(() -> jpaTm().update(testRegistrar)));
+    assertThat(existsInDb(testRegistrar)).isFalse();
+    assertThrows(IllegalArgumentException.class, () -> updateInDb(testRegistrar));
   }
 
   @Test
   void load_worksSuccessfully() {
-    assertThat(jpaTm().transact(() -> jpaTm().exists(testRegistrar))).isFalse();
+    assertThat(existsInDb(testRegistrar)).isFalse();
     insertInDb(testRegistrar);
-    Registrar persisted = jpaTm().transact(() -> jpaTm().loadByKey(registrarKey));
+    Registrar persisted = loadByKey(registrarKey);
 
     assertThat(persisted.getRegistrarId()).isEqualTo("registrarId");
     assertThat(persisted.getRegistrarName()).isEqualTo("registrarName");

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

@@ -460,7 +460,7 @@ public final class AppEngineExtension implements BeforeEachCallback, AfterEachCa
     if (withDatastore) {
       ObjectifyService.initOfy();
       // Reset id allocation in ObjectifyService so that ids are deterministic in tests.
-      IdService.resetNextTestId();
+      IdService.resetSelfAllocatedId();
       this.ofyTestEntities.forEach(AppEngineExtension::register);
     }
   }

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

@@ -986,18 +986,18 @@ public class DatabaseHelper {
    * <p><b>Note:</b> Your resource will not be enrolled in a commit log. If you want backups, use
    * {@link #persistResourceWithCommitLog(Object)}.
    */
-  public static <R> R persistResource(final R resource) {
+  public static <R extends ImmutableObject> R persistResource(final R resource) {
     return persistResource(resource, false);
   }
 
   /** Same as {@link #persistResource(Object)} with backups enabled. */
-  public static <R> R persistResourceWithCommitLog(final R resource) {
+  public static <R extends ImmutableObject> R persistResourceWithCommitLog(final R resource) {
     return persistResource(resource, true);
   }
 
-  private static <R> void saveResource(R resource, boolean wantBackup) {
+  private static <R extends ImmutableObject> void saveResource(R resource, boolean wantBackup) {
     if (tm().isOfy()) {
-      Consumer<Object> saver =
+      Consumer<ImmutableObject> saver =
           wantBackup || alwaysSaveWithBackup ? tm()::put : tm()::putWithoutBackup;
       saver.accept(resource);
       if (resource instanceof EppResource) {
@@ -1011,7 +1011,7 @@ public class DatabaseHelper {
   }
 
   private static <R extends EppResource> void persistEppResourceExtras(
-      R resource, EppResourceIndex index, Consumer<Object> saver) {
+      R resource, EppResourceIndex index, Consumer<ImmutableObject> saver) {
     assertWithMessage("Cannot persist an EppResource with a missing repoId in tests")
         .that(resource.getRepoId())
         .isNotEmpty();
@@ -1021,7 +1021,8 @@ public class DatabaseHelper {
     }
   }
 
-  private static <R> R persistResource(final R resource, final boolean wantBackup) {
+  private static <R extends ImmutableObject> R persistResource(
+      final R resource, final boolean wantBackup) {
     assertWithMessage("Attempting to persist a Builder is almost certainly an error in test code")
         .that(resource)
         .isNotInstanceOf(Buildable.Builder.class);
@@ -1042,7 +1043,7 @@ public class DatabaseHelper {
     tm().transact(
             () -> {
               if (tm().isOfy()) {
-                Consumer<Object> saver = tm()::put;
+                Consumer<ImmutableObject> saver = tm()::put;
                 saver.accept(resource);
                 persistEppResourceExtras(resource, eppResourceIndex, saver);
               } else {
@@ -1054,11 +1055,12 @@ public class DatabaseHelper {
     return transactIfJpaTm(() -> tm().loadByEntity(resource));
   }
 
-  public static <R> void persistResources(final Iterable<R> resources) {
+  public static <R extends ImmutableObject> void persistResources(final Iterable<R> resources) {
     persistResources(resources, false);
   }
 
-  private static <R> void persistResources(final Iterable<R> resources, final boolean wantBackup) {
+  private static <R extends ImmutableObject> void persistResources(
+      final Iterable<R> resources, final boolean wantBackup) {
     for (R resource : resources) {
       assertWithMessage("Attempting to persist a Builder is almost certainly an error in test code")
           .that(resource)
@@ -1379,9 +1381,9 @@ public class DatabaseHelper {
     return transactIfJpaTm(() -> tm().loadByEntitiesIfPresent(entities));
   }
 
-  /** Returns whether or not the given entity exists in the database. */
-  public static boolean existsInDatabase(Object object) {
-    return transactIfJpaTm(() -> tm().exists(object));
+  /** Returns whether or not the given entity exists in Cloud SQL. */
+  public static boolean existsInDb(ImmutableObject object) {
+    return jpaTm().transact(() -> jpaTm().exists(object));
   }
 
   /** Inserts the given entity/entities into Cloud SQL in a single transaction. */
@@ -1394,6 +1396,26 @@ public class DatabaseHelper {
     jpaTm().transact(() -> jpaTm().insertAll(entities));
   }
 
+  /** Puts the given entity/entities into Cloud SQL in a single transaction. */
+  public static <T extends ImmutableObject> void putInDb(T... entities) {
+    jpaTm().transact(() -> jpaTm().putAll(entities));
+  }
+
+  /** Puts the given entities into Cloud SQL in a single transaction. */
+  public static <T extends ImmutableObject> void putInDb(ImmutableCollection<T> entities) {
+    jpaTm().transact(() -> jpaTm().putAll(entities));
+  }
+
+  /** Updates the given entities in Cloud SQL in a single transaction. */
+  public static <T extends ImmutableObject> void updateInDb(T... entities) {
+    jpaTm().transact(() -> jpaTm().updateAll(entities));
+  }
+
+  /** Updates the given entities in Cloud SQL in a single transaction. */
+  public static <T extends ImmutableObject> void updateInDb(ImmutableCollection<T> entities) {
+    jpaTm().transact(() -> jpaTm().updateAll(entities));
+  }
+
   /**
    * In JPA mode, asserts that the given entity is detached from the current entity manager.
    *

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

@@ -14,17 +14,10 @@
 
 package google.registry.testing;
 
-import static google.registry.model.ofy.ObjectifyService.auditedOfy;
-
-import com.google.apphosting.api.ApiProxy;
-import com.google.apphosting.api.ApiProxy.Environment;
-import java.lang.reflect.InvocationTargetException;
-import java.lang.reflect.Method;
-import java.util.Map;
+import google.registry.model.AppEngineEnvironment;
 import org.junit.jupiter.api.extension.AfterEachCallback;
 import org.junit.jupiter.api.extension.BeforeEachCallback;
 import org.junit.jupiter.api.extension.ExtensionContext;
-import org.testcontainers.shaded.com.google.common.collect.ImmutableMap;
 
 /**
  * Allows instantiation of Datastore {@code Entity}s without the heavyweight {@link
@@ -41,13 +34,23 @@ import org.testcontainers.shaded.com.google.common.collect.ImmutableMap;
  * google.registry.model.domain.DomainBaseSqlTest} for example, and to <a
  * href="https://junit.org/junit5/docs/current/user-guide/#extensions-registration-programmatic">
  * JUnit 5 User Guide</a> for details of extension ordering.
+ *
+ * @see AppEngineEnvironment
  */
 public class DatastoreEntityExtension implements BeforeEachCallback, AfterEachCallback {
 
-  private static final Environment PLACEHOLDER_ENV = new PlaceholderEnvironment();
+  private final AppEngineEnvironment environment;
 
   private boolean allThreads = false;
 
+  public DatastoreEntityExtension(String appId) {
+    environment = new AppEngineEnvironment(appId);
+  }
+
+  public DatastoreEntityExtension() {
+    environment = new AppEngineEnvironment();
+  }
+
   /**
    * Whether all threads should be masqueraded as GAE threads.
    *
@@ -69,79 +72,19 @@ public class DatastoreEntityExtension implements BeforeEachCallback, AfterEachCa
 
   @Override
   public void beforeEach(ExtensionContext context) {
-    ApiProxy.setEnvironmentForCurrentThread(PLACEHOLDER_ENV);
-    // In order to create keys for entities they must be registered with Ofy. Calling this method
-    // will load the ObjectifyService class, whose static initialization block registers all Ofy
-    // entities.
-    auditedOfy();
     if (allThreads) {
-      ApiProxy.setEnvironmentFactory(() -> PLACEHOLDER_ENV);
+      environment.setEnvironmentForAllThreads();
+    } else {
+      environment.setEnvironmentForCurrentThread();
     }
   }
 
   @Override
-  public void afterEach(ExtensionContext context)
-      throws NoSuchMethodException, InvocationTargetException, IllegalAccessException {
-    // Clear the cached instance.
-    ApiProxy.clearEnvironmentForCurrentThread();
+  public void afterEach(ExtensionContext context) {
     if (allThreads) {
-      Method method = ApiProxy.class.getDeclaredMethod("clearEnvironmentFactory");
-      method.setAccessible(true);
-      method.invoke(null);
-    }
-  }
-
-  private static final class PlaceholderEnvironment implements Environment {
-
-    @Override
-    public String getAppId() {
-      return "PlaceholderAppId";
-    }
-
-    @Override
-    public Map<String, Object> getAttributes() {
-      return ImmutableMap.of();
-    }
-
-    @Override
-    public String getModuleId() {
-      throw new UnsupportedOperationException();
-    }
-
-    @Override
-    public String getVersionId() {
-      throw new UnsupportedOperationException();
-    }
-
-    @Override
-    public String getEmail() {
-      throw new UnsupportedOperationException();
-    }
-
-    @Override
-    public boolean isLoggedIn() {
-      throw new UnsupportedOperationException();
-    }
-
-    @Override
-    public boolean isAdmin() {
-      throw new UnsupportedOperationException();
-    }
-
-    @Override
-    public String getAuthDomain() {
-      throw new UnsupportedOperationException();
-    }
-
-    @SuppressWarnings("deprecation")
-    @Override
-    public String getRequestNamespace() {
-      throw new UnsupportedOperationException();
-    }
-
-    @Override
-    public long getRemainingMillis() {
-      throw new UnsupportedOperationException();
+      environment.unsetEnvironmentForAllThreads();
+    } else {
+      environment.unsetEnvironmentForCurrentThread();
     }
   }
 }

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

@@ -186,6 +186,34 @@ class UniformRapidSuspensionCommandTest
     assertNotInStdout("--undo"); // Undo shouldn't print a new undo command.
   }
 
+  @Test
+  void testAutorenews_setToFalsebyDefault() throws Exception {
+    persistResource(
+        newDomainBase("evil.tld")
+            .asBuilder()
+            .addStatusValue(StatusValue.SERVER_DELETE_PROHIBITED)
+            .build());
+    runCommandForced("--domain_name=evil.tld");
+    eppVerifier.verifySentAny();
+    assertInStdout("<superuser:autorenews>false</superuser:autorenews>");
+  }
+
+  @Test
+  void testAutorenews_setToTrueWhenUndo() throws Exception {
+    persistResource(
+        newDomainBase("evil.tld")
+            .asBuilder()
+            .addStatusValue(StatusValue.SERVER_DELETE_PROHIBITED)
+            .build());
+    runCommandForced(
+        "--domain_name=evil.tld",
+        "--undo",
+        "--hosts=ns1.example.com,ns2.example.com",
+        "--restore_client_hold");
+    eppVerifier.verifySentAny();
+    assertInStdout("<superuser:autorenews>true</superuser:autorenews>");
+  }
+
   @Test
   void testFailure_locksToPreserveWithoutUndo() {
     persistActiveDomain("evil.tld");

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

@@ -22,6 +22,7 @@ import static google.registry.persistence.transaction.TransactionManagerFactory.
 import static google.registry.reporting.spec11.Spec11RegistrarThreatMatchesParserTest.sampleThreatMatches;
 import static google.registry.testing.DatabaseHelper.createTld;
 import static google.registry.testing.DatabaseHelper.deleteResource;
+import static google.registry.testing.DatabaseHelper.insertInDb;
 import static google.registry.testing.DatabaseHelper.newDomainBase;
 import static google.registry.testing.DatabaseHelper.persistActiveDomain;
 import static google.registry.testing.DatabaseHelper.persistResource;
@@ -175,7 +176,7 @@ public class BackfillSpec11ThreatMatchesCommandTest
             .setRegistrarId("TheRegistrar")
             .setThreatTypes(ImmutableSet.of(ThreatType.MALWARE))
             .build();
-    jpaTm().transact(() -> jpaTm().put(previous));
+    insertInDb(previous);
 
     runCommandForced();
     ImmutableList<Spec11ThreatMatch> threatMatches =
@@ -197,7 +198,7 @@ public class BackfillSpec11ThreatMatchesCommandTest
             .setRegistrarId("TheRegistrar")
             .setThreatTypes(ImmutableSet.of(ThreatType.MALWARE))
             .build();
-    jpaTm().transact(() -> jpaTm().put(previous));
+    insertInDb(previous);
 
     runCommandForced("--overwrite_existing_dates");
     verifyExactlyThreeEntriesInDbFromLastDay();
@@ -214,7 +215,7 @@ public class BackfillSpec11ThreatMatchesCommandTest
         assertThrows(RuntimeException.class, this::runCommandForced);
     assertThat(runtimeException.getCause().getClass()).isEqualTo(IOException.class);
     assertThat(runtimeException).hasCauseThat().hasMessageThat().isEqualTo("hi");
-    jpaTm().transact(() -> assertThat(jpaTm().loadAllOf(Spec11ThreatMatch.class)).isEmpty());
+    assertThat(jpaTm().transact(() -> jpaTm().loadAllOf(Spec11ThreatMatch.class))).isEmpty();
   }
 
   @TestOfyAndSql

core/src/test/resources/google/registry/module/backend/backend_routing.txt

@@ -47,4 +47,5 @@ PATH                                               CLASS
 /_dr/task/updateSnapshotView                       UpdateSnapshotViewAction                       POST     n  INTERNAL,API APP ADMIN
 /_dr/task/uploadDatastoreBackup                    UploadDatastoreBackupAction                    POST     n  INTERNAL,API APP ADMIN
 /_dr/task/wipeOutCloudSql                          WipeOutCloudSqlAction                          GET      n  INTERNAL,API APP ADMIN
+/_dr/task/wipeOutContactHistoryPii                 WipeOutContactHistoryPiiAction                 GET      n  INTERNAL,API APP ADMIN
 /_dr/task/wipeOutDatastore                         WipeoutDatastoreAction                         GET      n  INTERNAL,API APP ADMIN

core/src/test/resources/google/registry/tools/server/uniform_rapid_suspension.xml

@@ -35,6 +35,9 @@
           </secDNS:dsData>
         </secDNS:add>
       </secDNS:update>
+      <superuser:domainUpdate xmlns:superuser="urn:google:params:xml:ns:superuser-1.0">
+        <superuser:autorenews>false</superuser:autorenews>
+      </superuser:domainUpdate>
       <metadata:metadata xmlns:metadata="urn:google:params:xml:ns:metadata-1.0">
         <metadata:reason>Uniform Rapid Suspension</metadata:reason>
         <metadata:requestedByRegistrar>false</metadata:requestedByRegistrar>

core/src/test/resources/google/registry/tools/server/uniform_rapid_suspension_existing_host.xml

@@ -25,6 +25,9 @@
           <secDNS:all>true</secDNS:all>
         </secDNS:rem>
       </secDNS:update>
+      <superuser:domainUpdate xmlns:superuser="urn:google:params:xml:ns:superuser-1.0">
+        <superuser:autorenews>false</superuser:autorenews>
+      </superuser:domainUpdate>
       <metadata:metadata xmlns:metadata="urn:google:params:xml:ns:metadata-1.0">
         <metadata:reason>Uniform Rapid Suspension</metadata:reason>
         <metadata:requestedByRegistrar>false</metadata:requestedByRegistrar>

core/src/test/resources/google/registry/tools/server/uniform_rapid_suspension_undo.xml

@@ -27,6 +27,9 @@
           <secDNS:all>true</secDNS:all>
         </secDNS:rem>
       </secDNS:update>
+      <superuser:domainUpdate xmlns:superuser="urn:google:params:xml:ns:superuser-1.0">
+        <superuser:autorenews>true</superuser:autorenews>
+      </superuser:domainUpdate>
       <metadata:metadata xmlns:metadata="urn:google:params:xml:ns:metadata-1.0">
         <metadata:reason>Undo Uniform Rapid Suspension</metadata:reason>
         <metadata:requestedByRegistrar>false</metadata:requestedByRegistrar>

core/src/test/resources/google/registry/tools/server/uniform_rapid_suspension_undo_client_hold.xml

@@ -28,6 +28,9 @@
           <secDNS:all>true</secDNS:all>
         </secDNS:rem>
       </secDNS:update>
+      <superuser:domainUpdate xmlns:superuser="urn:google:params:xml:ns:superuser-1.0">
+        <superuser:autorenews>true</superuser:autorenews>
+      </superuser:domainUpdate>
       <metadata:metadata xmlns:metadata="urn:google:params:xml:ns:metadata-1.0">
         <metadata:reason>Undo Uniform Rapid Suspension</metadata:reason>
         <metadata:requestedByRegistrar>false</metadata:requestedByRegistrar>

core/src/test/resources/google/registry/tools/server/uniform_rapid_suspension_undo_preserve.xml

@@ -26,6 +26,9 @@
           <secDNS:all>true</secDNS:all>
         </secDNS:rem>
       </secDNS:update>
+      <superuser:domainUpdate xmlns:superuser="urn:google:params:xml:ns:superuser-1.0">
+        <superuser:autorenews>true</superuser:autorenews>
+      </superuser:domainUpdate>
       <metadata:metadata xmlns:metadata="urn:google:params:xml:ns:metadata-1.0">
         <metadata:reason>Undo Uniform Rapid Suspension</metadata:reason>
         <metadata:requestedByRegistrar>false</metadata:requestedByRegistrar>

core/src/test/resources/google/registry/tools/server/uniform_rapid_suspension_with_client_hold.xml

@@ -36,6 +36,9 @@
           </secDNS:dsData>
         </secDNS:add>
       </secDNS:update>
+      <superuser:domainUpdate xmlns:superuser="urn:google:params:xml:ns:superuser-1.0">
+        <superuser:autorenews>false</superuser:autorenews>
+      </superuser:domainUpdate>
       <metadata:metadata xmlns:metadata="urn:google:params:xml:ns:metadata-1.0">
         <metadata:reason>Uniform Rapid Suspension</metadata:reason>
         <metadata:requestedByRegistrar>false</metadata:requestedByRegistrar>

gradle.properties

@@ -28,3 +28,4 @@ baseSchemaTag=
 schema_version=
 nomulus_version=
 dot_path=/usr/bin/dot
+pipeline=

release/cloudbuild-nomulus.yaml

@@ -83,6 +83,7 @@ steps:
   - -c
   - |
     ./release/stage_beam_pipeline.sh \
+      beamPipelineCommon \
       beam_pipeline_common \
       ${TAG_NAME} \
       ${PROJECT_ID} \

release/stage_beam_pipeline.sh

@@ -40,19 +40,20 @@ set -e
 
 if (( "$#" < 5 ||  $(("$#" % 2)) == 0 ));
 then
-  echo "Usage: $0 uberjar_name  release_tag dev_project " \
+  echo "Usage: $0 uberjar_task uberjar_name  release_tag dev_project " \
        "main_class metadata_pathname [ main_class metadata_pathname ] ..."
   exit 1
 fi
 
-uberjar_name="$1"
-release_tag="$2"
-dev_project="$3"
-shift 3
+uberjar_task="$1"
+uberjar_name="$2"
+release_tag="$3"
+dev_project="$4"
+shift 4
 
 maven_gcs_prefix="gcs://domain-registry-maven-repository"
 nom_build_dir="$(dirname $0)/.."
-${nom_build_dir}/nom_build clean :core:"${uberjar_name}" \
+${nom_build_dir}/nom_build clean :core:"${uberjar_task}" \
     --mavenUrl="${maven_gcs_prefix}"/maven \
     --pluginsUrl="${maven_gcs_prefix}"/plugins