Revert update auto timestamp non-transactional fallback (#1380)

This was added recently in PR #1341 as an attempted fix for our test flakiness,
but it turns out that it didn't address the root issue (whereas PR #1361
did). So this removes the fallback, as there's no reason this should ever be
called outside of a transactional context.

build.gradle

@@ -209,19 +209,22 @@ rootProject.ext {
 
 task runPresubmits(type: Exec) {
 
-  // Find a python version greater than 3.7.3 (this is somewhat arbitrary, we
-  // know we'd like at least 3.6, but 3.7.3 is the latest that ships with
-  // Debian so it seems like that should be available anywhere).
-  def MIN_PY_VER = 0x3070300
-  if (pyver('python') >= MIN_PY_VER) {
-    executable 'python'
-  } else if (pyver('/usr/bin/python3') >= MIN_PY_VER) {
-    executable '/usr/bin/python3'
-  } else {
-    throw new GradleException("No usable Python version found (build " +
-                              "requires at least python 3.7.3)");
-  }
   args('config/presubmits.py')
+
+  doFirst {
+    // Find a python version greater than 3.7.3 (this is somewhat arbitrary, we
+    // know we'd like at least 3.6, but 3.7.3 is the latest that ships with
+    // Debian so it seems like that should be available anywhere).
+    def MIN_PY_VER = 0x3070300
+    if (pyver('python') >= MIN_PY_VER) {
+      executable 'python'
+    } else if (pyver('/usr/bin/python3') >= MIN_PY_VER) {
+      executable '/usr/bin/python3'
+    } else {
+      throw new GradleException("No usable Python version found (build " +
+                                "requires at least python 3.7.3)");
+    }
+  }
 }
 
 def javadocSource = []

config/presubmits.py

@@ -213,6 +213,7 @@ PRESUBMITS = {
         {"src/test", "ActivityReportingQueryBuilder.java",
          # This class contains helper method to make queries in Beam.
          "RegistryJpaIO.java",
+         "CreateSyntheticHistoryEntriesAction.java",
          # TODO(b/179158393): Remove everything below, which should be done
          # using Criteria
          "JpaTransactionManager.java",

core/build.gradle

@@ -783,27 +783,27 @@ createUberJar(
 // User should install gcloud and login to GCP before invoking this tasks.
 if (environment == 'alpha') {
   def pipelines = [
-      InitSql            :
+      initSql            :
           [
               mainClass: 'google.registry.beam.initsql.InitSqlPipeline',
               metaData : 'google/registry/beam/init_sql_pipeline_metadata.json'
           ],
-      BulkDeleteDatastore:
+      bulkDeleteDatastore:
           [
               mainClass: 'google.registry.beam.datastore.BulkDeleteDatastorePipeline',
               metaData : 'google/registry/beam/bulk_delete_datastore_pipeline_metadata.json'
           ],
-      Spec11             :
+      spec11             :
           [
               mainClass: 'google.registry.beam.spec11.Spec11Pipeline',
               metaData : 'google/registry/beam/spec11_pipeline_metadata.json'
           ],
-      Invoicing          :
+      invoicing          :
           [
               mainClass: 'google.registry.beam.invoicing.InvoicingPipeline',
               metaData : 'google/registry/beam/invoicing_pipeline_metadata.json'
           ],
-      Rde                :
+      rde                :
           [
               mainClass: 'google.registry.beam.rde.RdePipeline',
               metaData : 'google/registry/beam/rde_pipeline_metadata.json'

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

@@ -55,7 +55,7 @@ public final class CommitLogImports {
    * represents the changes in one transaction. The {@code CommitLogManifest} contains deleted
    * entity keys, whereas each {@code CommitLogMutation} contains one whole entity.
    */
-  public static ImmutableList<ImmutableList<VersionedEntity>> loadEntitiesByTransaction(
+  static ImmutableList<ImmutableList<VersionedEntity>> loadEntitiesByTransaction(
       InputStream inputStream) {
     try (InputStream input = new BufferedInputStream(inputStream)) {
       Iterator<ImmutableObject> commitLogs = createDeserializingIterator(input, false);
@@ -104,7 +104,7 @@ public final class CommitLogImports {
    * represents the changes in one transaction. The {@code CommitLogManifest} contains deleted
    * entity keys, whereas each {@code CommitLogMutation} contains one whole entity.
    */
-  public static ImmutableList<VersionedEntity> loadEntities(InputStream inputStream) {
+  static ImmutableList<VersionedEntity> loadEntities(InputStream inputStream) {
     return loadEntitiesByTransaction(inputStream).stream()
         .flatMap(ImmutableList::stream)
         .collect(toImmutableList());

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

@@ -105,29 +105,29 @@ public abstract class VersionedEntity implements Serializable {
    * VersionedEntity VersionedEntities}. See {@link CommitLogImports#loadEntities} for more
    * information.
    */
-  public static Stream<VersionedEntity> fromManifest(CommitLogManifest manifest) {
+  static Stream<VersionedEntity> fromManifest(CommitLogManifest manifest) {
     long commitTimeMillis = manifest.getCommitTime().getMillis();
     return manifest.getDeletions().stream()
         .map(com.googlecode.objectify.Key::getRaw)
-        .map(key -> builder().commitTimeMills(commitTimeMillis).key(key).build());
+        .map(key -> newBuilder().commitTimeMills(commitTimeMillis).key(key).build());
   }
 
   /* Converts a {@link CommitLogMutation} to a {@link VersionedEntity}. */
-  public static VersionedEntity fromMutation(CommitLogMutation mutation) {
+  static VersionedEntity fromMutation(CommitLogMutation mutation) {
     return from(
         com.googlecode.objectify.Key.create(mutation).getParent().getId(),
         mutation.getEntityProtoBytes());
   }
 
   public static VersionedEntity from(long commitTimeMillis, byte[] entityProtoBytes) {
-    return builder()
+    return newBuilder()
         .entityProtoBytes(entityProtoBytes)
         .key(EntityTranslator.createFromPbBytes(entityProtoBytes).getKey())
         .commitTimeMills(commitTimeMillis)
         .build();
   }
 
-  static Builder builder() {
+  private static Builder newBuilder() {
     return new AutoValue_VersionedEntity.Builder();
   }
 
@@ -142,7 +142,7 @@ public abstract class VersionedEntity implements Serializable {
 
     public abstract VersionedEntity build();
 
-    public Builder entityProtoBytes(byte[] bytes) {
+    Builder entityProtoBytes(byte[] bytes) {
       return entityProtoBytes(new ImmutableBytes(bytes));
     }
   }

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

@@ -84,8 +84,12 @@ public class WipeOutContactHistoryPiiAction implements Runnable {
                             getNextContactHistoryEntitiesWithPiiBatch(wipeOutTime)));
         totalNumOfWipedEntities += numOfWipedEntities;
       } while (numOfWipedEntities > 0);
-      logger.atInfo().log(
-          "Wiped out PII of %d ContactHistory entities in total.", totalNumOfWipedEntities);
+      String msg =
+          String.format(
+              "Done. Wiped out PII of %d ContactHistory entities in total.",
+              totalNumOfWipedEntities);
+      logger.atInfo().log(msg);
+      response.setPayload(msg);
       response.setStatus(SC_OK);
 
     } catch (Exception e) {

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

@@ -22,6 +22,7 @@ import static google.registry.beam.initsql.BackupPaths.getExportFilePatterns;
 import static google.registry.model.ofy.ObjectifyService.auditedOfy;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
 import static google.registry.util.DateTimeUtils.isBeforeOrAt;
+import static google.registry.util.DomainNameUtils.canonicalizeDomainName;
 import static java.util.Comparator.comparing;
 import static org.apache.beam.sdk.values.TypeDescriptors.kvs;
 import static org.apache.beam.sdk.values.TypeDescriptors.strings;
@@ -277,7 +278,7 @@ public final class Transforms {
 
   // Prober contacts referencing phantom registrars. They and their associated history entries can
   // be safely ignored.
-  private static final ImmutableSet IGNORED_CONTACTS =
+  private static final ImmutableSet<String> IGNORED_CONTACTS =
       ImmutableSet.of(
           "1_WJ0TEST-GOOGLE", "1_WJ1TEST-GOOGLE", "1_WJ2TEST-GOOGLE", "1_WJ3TEST-GOOGLE");
 
@@ -320,7 +321,8 @@ public final class Transforms {
     return true;
   }
 
-  private static Entity repairBadData(Entity entity) {
+  @VisibleForTesting
+  static Entity repairBadData(Entity entity) {
     if (entity.getKind().equals("Cancellation")
         && Objects.equals(entity.getProperty("reason"), "AUTO_RENEW")) {
       // AUTO_RENEW has been moved from 'reason' to flags. Change reason to RENEW and add the
@@ -328,6 +330,15 @@ public final class Transforms {
       // instead of append. See b/185954992.
       entity.setUnindexedProperty("reason", Reason.RENEW.name());
       entity.setUnindexedProperty("flags", ImmutableList.of(Flag.AUTO_RENEW.name()));
+    } else if (entity.getKind().equals("DomainBase")) {
+      // Canonicalize old domain/host names from 2016 and earlier before we were enforcing this.
+      entity.setIndexedProperty(
+          "fullyQualifiedDomainName",
+          canonicalizeDomainName((String) entity.getProperty("fullyQualifiedDomainName")));
+    } else if (entity.getKind().equals("HostResource")) {
+      entity.setIndexedProperty(
+          "fullyQualifiedHostName",
+          canonicalizeDomainName((String) entity.getProperty("fullyQualifiedHostName")));
     }
     return entity;
   }
@@ -365,7 +376,8 @@ public final class Transforms {
    * Returns a {@link PTransform} that produces a {@link PCollection} containing all elements in the
    * given {@link Iterable}.
    */
-  static PTransform<PBegin, PCollection<String>> toStringPCollection(Iterable<String> strings) {
+  private static PTransform<PBegin, PCollection<String>> toStringPCollection(
+      Iterable<String> strings) {
     return Create.of(strings).withCoder(StringUtf8Coder.of());
   }
 
@@ -373,7 +385,7 @@ public final class Transforms {
    * Returns a {@link PTransform} from file {@link Metadata} to {@link VersionedEntity} using
    * caller-provided {@code transformer}.
    */
-  static PTransform<PCollection<Metadata>, PCollection<VersionedEntity>> processFiles(
+  private static PTransform<PCollection<Metadata>, PCollection<VersionedEntity>> processFiles(
       DoFn<ReadableFile, VersionedEntity> transformer) {
     return new PTransform<PCollection<Metadata>, PCollection<VersionedEntity>>() {
       @Override
@@ -389,7 +401,7 @@ public final class Transforms {
     private final DateTime fromTime;
     private final DateTime toTime;
 
-    public FilterCommitLogFileByTime(DateTime fromTime, DateTime toTime) {
+    FilterCommitLogFileByTime(DateTime fromTime, DateTime toTime) {
       checkNotNull(fromTime, "fromTime");
       checkNotNull(toTime, "toTime");
       checkArgument(

core/src/main/java/google/registry/beam/rde/RdeIO.java

@@ -19,10 +19,13 @@ import static com.google.common.base.Verify.verify;
 import static google.registry.model.common.Cursor.getCursorTimeOrStartOfTime;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
+import static google.registry.rde.RdeModule.BRDA_QUEUE;
+import static google.registry.rde.RdeModule.RDE_UPLOAD_QUEUE;
 import static java.nio.charset.StandardCharsets.UTF_8;
 
 import com.google.auto.value.AutoValue;
 import com.google.cloud.storage.BlobId;
+import com.google.common.collect.ImmutableMultimap;
 import com.google.common.flogger.FluentLogger;
 import google.registry.gcs.GcsUtils;
 import google.registry.keyring.api.PgpHelper;
@@ -31,14 +34,20 @@ import google.registry.model.rde.RdeMode;
 import google.registry.model.rde.RdeNamingUtils;
 import google.registry.model.rde.RdeRevision;
 import google.registry.model.tld.Registry;
+import google.registry.rde.BrdaCopyAction;
 import google.registry.rde.DepositFragment;
 import google.registry.rde.Ghostryde;
 import google.registry.rde.PendingDeposit;
 import google.registry.rde.RdeCounter;
 import google.registry.rde.RdeMarshaller;
+import google.registry.rde.RdeModule;
 import google.registry.rde.RdeResourceType;
+import google.registry.rde.RdeUploadAction;
 import google.registry.rde.RdeUtil;
+import google.registry.request.Action.Service;
+import google.registry.request.RequestParameters;
 import google.registry.tldconfig.idn.IdnTableEnum;
+import google.registry.util.CloudTasksUtils;
 import google.registry.xjc.rdeheader.XjcRdeHeader;
 import google.registry.xjc.rdeheader.XjcRdeHeaderElement;
 import google.registry.xml.ValidationMode;
@@ -68,6 +77,8 @@ public class RdeIO {
 
     abstract GcsUtils gcsUtils();
 
+    abstract CloudTasksUtils cloudTasksUtils();
+
     abstract String rdeBucket();
 
     // It's OK to return a primitive array because we are only using it to construct the
@@ -83,7 +94,9 @@ public class RdeIO {
 
     @AutoValue.Builder
     abstract static class Builder {
-      abstract Builder setGcsUtils(GcsUtils gcsUtils);
+      abstract Builder setGcsUtils(GcsUtils value);
+
+      abstract Builder setCloudTasksUtils(CloudTasksUtils value);
 
       abstract Builder setRdeBucket(String value);
 
@@ -100,7 +113,8 @@ public class RdeIO {
           .apply(
               "Write to GCS",
               ParDo.of(new RdeWriter(gcsUtils(), rdeBucket(), stagingKeyBytes(), validationMode())))
-          .apply("Update cursors", ParDo.of(new CursorUpdater()));
+          .apply("Update cursors", ParDo.of(new CursorUpdater()))
+          .apply("Enqueue upload action", ParDo.of(new UploadEnqueuer(cloudTasksUtils())));
       return PDone.in(input.getPipeline());
     }
   }
@@ -236,11 +250,12 @@ public class RdeIO {
     }
   }
 
-  private static class CursorUpdater extends DoFn<KV<PendingDeposit, Integer>, Void> {
+  private static class CursorUpdater extends DoFn<KV<PendingDeposit, Integer>, PendingDeposit> {
     private static final FluentLogger logger = FluentLogger.forEnclosingClass();
 
     @ProcessElement
-    public void processElement(@Element KV<PendingDeposit, Integer> input) {
+    public void processElement(
+        @Element KV<PendingDeposit, Integer> input, OutputReceiver<PendingDeposit> outputReceiver) {
       tm().transact(
               () -> {
                 PendingDeposit key = input.getKey();
@@ -268,6 +283,45 @@ public class RdeIO {
                     "Rolled forward %s on %s cursor to %s.", key.cursor(), key.tld(), newPosition);
                 RdeRevision.saveRevision(key.tld(), key.watermark(), key.mode(), revision);
               });
+      outputReceiver.output(input.getKey());
+    }
+  }
+
+  private static class UploadEnqueuer extends DoFn<PendingDeposit, Void> {
+
+    private final CloudTasksUtils cloudTasksUtils;
+
+    private UploadEnqueuer(CloudTasksUtils cloudTasksUtils) {
+      this.cloudTasksUtils = cloudTasksUtils;
+    }
+
+    @ProcessElement
+    public void processElement(@Element PendingDeposit input, PipelineOptions options) {
+      if (input.mode() == RdeMode.FULL) {
+        cloudTasksUtils.enqueue(
+            RDE_UPLOAD_QUEUE,
+            CloudTasksUtils.createPostTask(
+                RdeUploadAction.PATH,
+                Service.BACKEND.getServiceId(),
+                ImmutableMultimap.of(
+                    RequestParameters.PARAM_TLD,
+                    input.tld(),
+                    RdeModule.PARAM_PREFIX,
+                    options.getJobName() + '/')));
+      } else {
+        cloudTasksUtils.enqueue(
+            BRDA_QUEUE,
+            CloudTasksUtils.createPostTask(
+                BrdaCopyAction.PATH,
+                Service.BACKEND.getServiceId(),
+                ImmutableMultimap.of(
+                    RequestParameters.PARAM_TLD,
+                    input.tld(),
+                    RdeModule.PARAM_WATERMARK,
+                    input.watermark().toString(),
+                    RdeModule.PARAM_PREFIX,
+                    options.getJobName() + '/')));
+      }
     }
   }
 }

core/src/main/java/google/registry/beam/rde/RdePipeline.java

@@ -27,6 +27,7 @@ import com.google.common.io.BaseEncoding;
 import dagger.BindsInstance;
 import dagger.Component;
 import google.registry.beam.common.RegistryJpaIO;
+import google.registry.config.CloudTasksUtilsModule;
 import google.registry.config.CredentialModule;
 import google.registry.config.RegistryConfig.ConfigModule;
 import google.registry.gcs.GcsUtils;
@@ -44,6 +45,8 @@ import google.registry.rde.PendingDeposit;
 import google.registry.rde.PendingDeposit.PendingDepositCoder;
 import google.registry.rde.RdeFragmenter;
 import google.registry.rde.RdeMarshaller;
+import google.registry.util.CloudTasksUtils;
+import google.registry.util.UtilsModule;
 import google.registry.xml.ValidationMode;
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
@@ -66,7 +69,6 @@ import org.apache.beam.sdk.options.PipelineOptionsFactory;
 import org.apache.beam.sdk.transforms.FlatMapElements;
 import org.apache.beam.sdk.transforms.Flatten;
 import org.apache.beam.sdk.transforms.GroupByKey;
-import org.apache.beam.sdk.transforms.Reshuffle;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.sdk.values.PCollectionList;
@@ -93,6 +95,7 @@ public class RdePipeline implements Serializable {
   private final String rdeBucket;
   private final byte[] stagingKeyBytes;
   private final GcsUtils gcsUtils;
+  private final CloudTasksUtils cloudTasksUtils;
 
   // Registrars to be excluded from data escrow. Not including the sandbox-only OTE type so that
   // if sneaks into production we would get an extra signal.
@@ -111,13 +114,14 @@ public class RdePipeline implements Serializable {
   }
 
   @Inject
-  RdePipeline(RdePipelineOptions options, GcsUtils gcsUtils) {
+  RdePipeline(RdePipelineOptions options, GcsUtils gcsUtils, CloudTasksUtils cloudTasksUtils) {
     this.options = options;
     this.mode = ValidationMode.valueOf(options.getValidationMode());
     this.pendings = decodePendings(options.getPendings());
-    this.rdeBucket = options.getGcsBucket();
+    this.rdeBucket = options.getRdeStagingBucket();
     this.stagingKeyBytes = BaseEncoding.base64Url().decode(options.getStagingKey());
     this.gcsUtils = gcsUtils;
+    this.cloudTasksUtils = cloudTasksUtils;
   }
 
   PipelineResult run() {
@@ -140,10 +144,11 @@ public class RdePipeline implements Serializable {
 
   void persistData(PCollection<KV<PendingDeposit, Iterable<DepositFragment>>> input) {
     input.apply(
-        "Write to GCS and update cursors",
+        "Write to GCS, update cursors, and enqueue upload tasks",
         RdeIO.Write.builder()
             .setRdeBucket(rdeBucket)
             .setGcsUtils(gcsUtils)
+            .setCloudTasksUtils(cloudTasksUtils)
             .setValidationMode(mode)
             .setStagingKeyBytes(stagingKeyBytes)
             .build());
@@ -177,18 +182,13 @@ public class RdePipeline implements Serializable {
                     }));
   }
 
-  @SuppressWarnings("deprecation") // Reshuffle is still recommended by Dataflow.
   <T extends EppResource>
       PCollection<KV<PendingDeposit, DepositFragment>> processNonRegistrarEntities(
           Pipeline pipeline, Class<T> clazz) {
     return createInputs(pipeline, clazz)
         .apply("Marshal " + clazz.getSimpleName() + " into DepositFragment", mapToFragments(clazz))
-        .setCoder(KvCoder.of(PendingDepositCoder.of(), SerializableCoder.of(DepositFragment.class)))
-        .apply(
-            "Reshuffle KV<PendingDeposit, DepositFragment> of "
-                + clazz.getSimpleName()
-                + " to prevent fusion",
-            Reshuffle.of());
+        .setCoder(
+            KvCoder.of(PendingDepositCoder.of(), SerializableCoder.of(DepositFragment.class)));
   }
 
   <T extends EppResource> PCollection<VKey<T>> createInputs(Pipeline pipeline, Class<T> clazz) {
@@ -202,7 +202,7 @@ public class RdePipeline implements Serializable {
             String.class,
             // TODO: consider adding coders for entities and pass them directly instead of using
             // VKeys.
-            x -> VKey.create(clazz, x)));
+            x -> VKey.createSql(clazz, x)));
   }
 
   <T extends EppResource>
@@ -270,7 +270,7 @@ public class RdePipeline implements Serializable {
    * Encodes the TLD to pending deposit map in an URL safe string that is sent to the pipeline
    * worker by the pipeline launcher as a pipeline option.
    */
-  static String encodePendings(ImmutableSetMultimap<String, PendingDeposit> pendings)
+  public static String encodePendings(ImmutableSetMultimap<String, PendingDeposit> pendings)
       throws IOException {
     try (ByteArrayOutputStream baos = new ByteArrayOutputStream()) {
       ObjectOutputStream oos = new ObjectOutputStream(baos);
@@ -282,13 +282,24 @@ public class RdePipeline implements Serializable {
 
   public static void main(String[] args) throws IOException, ClassNotFoundException {
     PipelineOptionsFactory.register(RdePipelineOptions.class);
-    RdePipelineOptions options = PipelineOptionsFactory.fromArgs(args).as(RdePipelineOptions.class);
+    RdePipelineOptions options =
+        PipelineOptionsFactory.fromArgs(args).withValidation().as(RdePipelineOptions.class);
+    // RegistryPipelineWorkerInitializer only initializes before pipeline executions, after the
+    // main() function constructed the graph. We need the registry environment set up so that we
+    // can create a CloudTasksUtils which uses the environment-dependent config file.
+    options.getRegistryEnvironment().setup();
     options.setIsolationOverride(TransactionIsolationLevel.TRANSACTION_READ_COMMITTED);
     DaggerRdePipeline_RdePipelineComponent.builder().options(options).build().rdePipeline().run();
   }
 
   @Singleton
-  @Component(modules = {CredentialModule.class, ConfigModule.class})
+  @Component(
+      modules = {
+        CredentialModule.class,
+        ConfigModule.class,
+        CloudTasksUtilsModule.class,
+        UtilsModule.class
+      })
   interface RdePipelineComponent {
     RdePipeline rdePipeline();
 

core/src/main/java/google/registry/beam/rde/RdePipelineOptions.java

@@ -31,9 +31,9 @@ public interface RdePipelineOptions extends RegistryPipelineOptions {
   void setValidationMode(String value);
 
   @Description("The GCS bucket where the encrypted RDE deposits will be uploaded to.")
-  String getGcsBucket();
+  String getRdeStagingBucket();
 
-  void setGcsBucket(String value);
+  void setRdeStagingBucket(String value);
 
   @Description("The Base64-encoded PGP public key to encrypt the deposits.")
   String getStagingKey();

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

@@ -22,10 +22,13 @@ import dagger.Provides;
 import google.registry.config.CredentialModule.DefaultCredential;
 import google.registry.config.RegistryConfig.Config;
 import google.registry.util.CloudTasksUtils;
+import google.registry.util.CloudTasksUtils.GcpCloudTasksClient;
+import google.registry.util.CloudTasksUtils.SerializableCloudTasksClient;
 import google.registry.util.GoogleCredentialsBundle;
 import google.registry.util.Retrier;
 import java.io.IOException;
-import javax.inject.Provider;
+import java.io.Serializable;
+import java.util.function.Supplier;
 import javax.inject.Singleton;
 
 /**
@@ -42,24 +45,35 @@ public abstract class CloudTasksUtilsModule {
   public static CloudTasksUtils provideCloudTasksUtils(
       @Config("projectId") String projectId,
       @Config("locationId") String locationId,
-      // Use a provider so that we can use try-with-resources with the client, which implements
-      // Autocloseable.
-      Provider<CloudTasksClient> clientProvider,
+      SerializableCloudTasksClient client,
       Retrier retrier) {
-    return new CloudTasksUtils(retrier, projectId, locationId, clientProvider);
+    return new CloudTasksUtils(retrier, projectId, locationId, client);
+  }
+
+  // Provides a supplier instead of using a Dagger @Provider because the latter is not serializable.
+  @Provides
+  public static Supplier<CloudTasksClient> provideCloudTasksClientSupplier(
+      @DefaultCredential GoogleCredentialsBundle credentials) {
+    return (Supplier<CloudTasksClient> & Serializable)
+        () -> {
+          CloudTasksClient client;
+          try {
+            client =
+                CloudTasksClient.create(
+                    CloudTasksSettings.newBuilder()
+                        .setCredentialsProvider(
+                            FixedCredentialsProvider.create(credentials.getGoogleCredentials()))
+                        .build());
+          } catch (IOException e) {
+            throw new RuntimeException(e);
+          }
+          return client;
+        };
   }
 
   @Provides
-  public static CloudTasksClient provideCloudTasksClient(
-      @DefaultCredential GoogleCredentialsBundle credentials) {
-    try {
-      return CloudTasksClient.create(
-          CloudTasksSettings.newBuilder()
-              .setCredentialsProvider(
-                  FixedCredentialsProvider.create(credentials.getGoogleCredentials()))
-              .build());
-    } catch (IOException e) {
-      throw new RuntimeException(e);
-    }
+  public static SerializableCloudTasksClient provideSerializableCloudTasksClient(
+      final Supplier<CloudTasksClient> clientSupplier) {
+    return new GcpCloudTasksClient(clientSupplier);
   }
 }

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

@@ -348,4 +348,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/UpdateAutoTimestamp.java

@@ -15,12 +15,8 @@
 package google.registry.model;
 
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
-import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
-import static org.joda.time.DateTimeZone.UTC;
 
-import com.google.common.flogger.FluentLogger;
-import com.google.common.flogger.StackSize;
 import com.googlecode.objectify.annotation.Ignore;
 import com.googlecode.objectify.annotation.OnLoad;
 import google.registry.model.translators.UpdateAutoTimestampTranslatorFactory;
@@ -44,12 +40,10 @@ import org.joda.time.DateTime;
 @Embeddable
 public class UpdateAutoTimestamp extends ImmutableObject {
 
-  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
-
   // When set to true, database converters/translators should do the auto update.  When set to
   // false, auto update should be suspended (this exists to allow us to preserve the original value
   // during a replay).
-  private static ThreadLocal<Boolean> autoUpdateEnabled = ThreadLocal.withInitial(() -> true);
+  private static final ThreadLocal<Boolean> autoUpdateEnabled = ThreadLocal.withInitial(() -> true);
 
   @Transient DateTime timestamp;
 
@@ -63,16 +57,7 @@ public class UpdateAutoTimestamp extends ImmutableObject {
   @PrePersist
   @PreUpdate
   void setTimestamp() {
-    // On the off chance that this is called outside of a transaction, log it instead of failing
-    // with an exception from attempting to call jpaTm().getTransactionTime(), and then fall back
-    // to DateTime.now(UTC).
-    if (!jpaTm().inTransaction()) {
-      logger.atSevere().withStackTrace(StackSize.MEDIUM).log(
-          "Failed to update automatic timestamp because this wasn't called in a JPA transaction%s.",
-          ofyTm().inTransaction() ? " (but there is an open Ofy transaction)" : "");
-      timestamp = DateTime.now(UTC);
-      lastUpdateTime = DateTimeUtils.toZonedDateTime(timestamp);
-    } else if (autoUpdateEnabled() || lastUpdateTime == null) {
+    if (autoUpdateEnabled() || lastUpdateTime == null) {
       timestamp = jpaTm().getTransactionTime();
       lastUpdateTime = DateTimeUtils.toZonedDateTime(timestamp);
     }

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

@@ -865,7 +865,8 @@ public class DomainContent extends EppResource
     public B setDomainName(String domainName) {
       checkArgument(
           domainName.equals(canonicalizeDomainName(domainName)),
-          "Domain name must be in puny-coded, lower-case form");
+          "Domain name %s not in puny-coded, lower-case form",
+          domainName);
       getInstance().fullyQualifiedDomainName = domainName;
       return thisCastToDerived();
     }

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

@@ -196,7 +196,8 @@ public class HostBase extends EppResource {
     public B setHostName(String hostName) {
       checkArgument(
           hostName.equals(canonicalizeDomainName(hostName)),
-          "Host name must be in puny-coded, lower-case form");
+          "Host name %s not in puny-coded, lower-case form",
+          hostName);
       getInstance().fullyQualifiedHostName = hostName;
       return thisCastToDerived();
     }

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

@@ -21,6 +21,7 @@ import static google.registry.config.RegistryConfig.getHibernateHikariIdleTimeou
 import static google.registry.config.RegistryConfig.getHibernateHikariMaximumPoolSize;
 import static google.registry.config.RegistryConfig.getHibernateHikariMinimumIdle;
 import static google.registry.config.RegistryConfig.getHibernateLogSqlQueries;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.api.client.auth.oauth2.Credential;
 import com.google.common.annotations.VisibleForTesting;
@@ -34,6 +35,7 @@ import google.registry.config.RegistryConfig.Config;
 import google.registry.persistence.transaction.CloudSqlCredentialSupplier;
 import google.registry.persistence.transaction.JpaTransactionManager;
 import google.registry.persistence.transaction.JpaTransactionManagerImpl;
+import google.registry.persistence.transaction.TransactionManager;
 import google.registry.privileges.secretmanager.SqlCredential;
 import google.registry.privileges.secretmanager.SqlCredentialStore;
 import google.registry.privileges.secretmanager.SqlUser;
@@ -206,6 +208,12 @@ public abstract class PersistenceModule {
     }
   }
 
+  @Provides
+  @Singleton
+  static TransactionManager provideTransactionManager() {
+    return tm();
+  }
+
   @Provides
   @Singleton
   @AppEngineJpaTm

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

@@ -98,14 +98,16 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   // EntityManagerFactory is thread safe.
   private final EntityManagerFactory emf;
   private final Clock clock;
+
   // TODO(b/177588434): Investigate alternatives for managing transaction information. ThreadLocal
   // adds an unnecessary restriction that each request has to be processed by one thread
   // synchronously.
-  private final ThreadLocal<TransactionInfo> transactionInfo =
+  private static final ThreadLocal<TransactionInfo> transactionInfo =
       ThreadLocal.withInitial(TransactionInfo::new);
+
   // If this value is present, use it to determine whether or not to replay SQL transactions to
   // Datastore, rather than using the schedule stored in Datastore.
-  private static ThreadLocal<Optional<Boolean>> replaySqlToDatastoreOverrideForTest =
+  private static final ThreadLocal<Optional<Boolean>> replaySqlToDatastoreOverrideForTest =
       ThreadLocal.withInitial(Optional::empty);
 
   public JpaTransactionManagerImpl(EntityManagerFactory emf, Clock clock) {

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

@@ -132,13 +132,16 @@ public class TransactionManagerFactory {
   /**
    * Sets the return of {@link #tm()} to the given instance of {@link TransactionManager}.
    *
+   * <p>DO NOT CALL THIS DIRECTLY IF POSSIBLE. Strongly prefer the use of <code>TmOverrideExtension
+   * </code> in test code instead.
+   *
    * <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) {
-    tmForTest = Optional.of(newTm);
+  public static void setTmOverrideForTest(TransactionManager newTmOverride) {
+    tmForTest = Optional.of(newTmOverride);
   }
 
   /** Resets the overridden transaction manager post-test. */
@@ -152,10 +155,10 @@ public class TransactionManagerFactory {
     }
   }
 
-  /** Thrown when a write is attempted when the DB is in read-only mode. */
+  /** Registry is currently undergoing maintenance and is in read-only mode. */
   public static class ReadOnlyModeException extends IllegalStateException {
-    public ReadOnlyModeException() {
-      super("Registry is currently in read-only mode");
+    ReadOnlyModeException() {
+      super("Registry is currently undergoing maintenance and is in read-only mode");
     }
   }
 }

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

@@ -31,6 +31,7 @@ import google.registry.request.auth.Auth;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
+import java.util.Optional;
 import javax.inject.Inject;
 import org.bouncycastle.openpgp.PGPKeyPair;
 import org.bouncycastle.openpgp.PGPPrivateKey;
@@ -60,7 +61,7 @@ import org.joda.time.DateTime;
     auth = Auth.AUTH_INTERNAL_OR_ADMIN)
 public final class BrdaCopyAction implements Runnable {
 
-  static final String PATH = "/_dr/task/brdaCopy";
+  public static final String PATH = "/_dr/task/brdaCopy";
 
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
 
@@ -69,6 +70,7 @@ public final class BrdaCopyAction implements Runnable {
   @Inject @Config("rdeBucket") String stagingBucket;
   @Inject @Parameter(RequestParameters.PARAM_TLD) String tld;
   @Inject @Parameter(RdeModule.PARAM_WATERMARK) DateTime watermark;
+  @Inject @Parameter(RdeModule.PARAM_PREFIX) Optional<String> prefix;
   @Inject @Key("brdaReceiverKey") PGPPublicKey receiverKey;
   @Inject @Key("brdaSigningKey") PGPKeyPair signingKey;
   @Inject @Key("rdeStagingDecryptionKey") PGPPrivateKey stagingDecryptionKey;
@@ -84,11 +86,12 @@ public final class BrdaCopyAction implements Runnable {
   }
 
   private void copyAsRyde() throws IOException {
-    String prefix = RdeNamingUtils.makeRydeFilename(tld, watermark, THIN, 1, 0);
-    BlobId xmlFilename = BlobId.of(stagingBucket, prefix + ".xml.ghostryde");
-    BlobId xmlLengthFilename = BlobId.of(stagingBucket, prefix + ".xml.length");
-    BlobId rydeFile = BlobId.of(brdaBucket, prefix + ".ryde");
-    BlobId sigFile = BlobId.of(brdaBucket, prefix + ".sig");
+    String nameWithoutPrefix = RdeNamingUtils.makeRydeFilename(tld, watermark, THIN, 1, 0);
+    String name = prefix.orElse("") + nameWithoutPrefix;
+    BlobId xmlFilename = BlobId.of(stagingBucket, name + ".xml.ghostryde");
+    BlobId xmlLengthFilename = BlobId.of(stagingBucket, name + ".xml.length");
+    BlobId rydeFile = BlobId.of(brdaBucket, nameWithoutPrefix + ".ryde");
+    BlobId sigFile = BlobId.of(brdaBucket, nameWithoutPrefix + ".sig");
 
     long xmlLength = readXmlLength(xmlLengthFilename);
 
@@ -97,11 +100,12 @@ public final class BrdaCopyAction implements Runnable {
         InputStream ghostrydeDecoder = Ghostryde.decoder(gcsInput, stagingDecryptionKey);
         OutputStream rydeOut = gcsUtils.openOutputStream(rydeFile);
         OutputStream sigOut = gcsUtils.openOutputStream(sigFile);
-        RydeEncoder rydeEncoder = new RydeEncoder.Builder()
-            .setRydeOutput(rydeOut, receiverKey)
-            .setSignatureOutput(sigOut, signingKey)
-            .setFileMetadata(prefix, xmlLength, watermark)
-            .build()) {
+        RydeEncoder rydeEncoder =
+            new RydeEncoder.Builder()
+                .setRydeOutput(rydeOut, receiverKey)
+                .setSignatureOutput(sigOut, signingKey)
+                .setFileMetadata(nameWithoutPrefix, xmlLength, watermark)
+                .build()) {
       ByteStreams.copy(ghostrydeDecoder, rydeEncoder);
     }
   }

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

@@ -49,6 +49,10 @@ public abstract class RdeModule {
   public static final String PARAM_MODE = "mode";
   public static final String PARAM_REVISION = "revision";
   public static final String PARAM_LENIENT = "lenient";
+  public static final String PARAM_PREFIX = "prefix";
+  public static final String RDE_UPLOAD_QUEUE = "rde-upload";
+  public static final String RDE_REPORT_QUEUE = "rde-report";
+  public static final String BRDA_QUEUE = "brda";
 
   @Provides
   @Parameter(PARAM_WATERMARK)
@@ -92,10 +96,11 @@ public abstract class RdeModule {
     return extractBooleanParameter(req, PARAM_LENIENT);
   }
 
+  // TODO (jianglai): Make it a required parameter once we migrate to Cloud SQL.
   @Provides
-  @Named("brda")
-  static Queue provideQueueBrda() {
-    return getQueue("brda");
+  @Parameter(PARAM_PREFIX)
+  static Optional<String> providePrefix(HttpServletRequest req) {
+    return extractOptionalParameter(req, PARAM_PREFIX);
   }
 
   @Provides

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

@@ -68,6 +68,7 @@ public final class RdeReportAction implements Runnable, EscrowTask {
   @Inject Response response;
   @Inject RdeReporter reporter;
   @Inject @Parameter(RequestParameters.PARAM_TLD) String tld;
+  @Inject @Parameter(RdeModule.PARAM_PREFIX) Optional<String> prefix;
   @Inject @Config("rdeBucket") String bucket;
   @Inject @Config("rdeInterval") Duration interval;
   @Inject @Config("rdeReportLockTimeout") Duration timeout;
@@ -96,8 +97,9 @@ public final class RdeReportAction implements Runnable, EscrowTask {
         RdeRevision.getCurrentRevision(tld, watermark, FULL)
             .orElseThrow(
                 () -> new IllegalStateException("RdeRevision was not set on generated deposit"));
-    String prefix = RdeNamingUtils.makeRydeFilename(tld, watermark, FULL, 1, revision);
-    BlobId reportFilename = BlobId.of(bucket, prefix + "-report.xml.ghostryde");
+    String name =
+        prefix.orElse("") + RdeNamingUtils.makeRydeFilename(tld, watermark, FULL, 1, revision);
+    BlobId reportFilename = BlobId.of(bucket, name + "-report.xml.ghostryde");
     verify(gcsUtils.existsAndNotEmpty(reportFilename), "Missing file: %s", reportFilename);
     reporter.send(readReportFromGcs(reportFilename));
     response.setContentType(PLAIN_TEXT_UTF_8);

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

@@ -14,20 +14,33 @@
 
 package google.registry.rde;
 
+import static google.registry.beam.BeamUtils.createJobName;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.request.Action.Method.GET;
 import static google.registry.request.Action.Method.POST;
 import static google.registry.xml.ValidationMode.LENIENT;
 import static google.registry.xml.ValidationMode.STRICT;
+import static javax.servlet.http.HttpServletResponse.SC_INTERNAL_SERVER_ERROR;
 import static javax.servlet.http.HttpServletResponse.SC_NO_CONTENT;
+import static javax.servlet.http.HttpServletResponse.SC_OK;
 
+import com.google.api.services.dataflow.Dataflow;
+import com.google.api.services.dataflow.model.LaunchFlexTemplateParameter;
+import com.google.api.services.dataflow.model.LaunchFlexTemplateRequest;
+import com.google.api.services.dataflow.model.LaunchFlexTemplateResponse;
 import com.google.common.base.Ascii;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSetMultimap;
 import com.google.common.collect.Multimaps;
 import com.google.common.flogger.FluentLogger;
+import com.google.common.io.BaseEncoding;
+import google.registry.beam.rde.RdePipeline;
 import google.registry.config.RegistryConfig.Config;
+import google.registry.config.RegistryEnvironment;
 import google.registry.gcs.GcsUtils;
+import google.registry.keyring.api.KeyModule.Key;
 import google.registry.mapreduce.MapreduceRunner;
 import google.registry.mapreduce.inputs.EppResourceInputs;
 import google.registry.mapreduce.inputs.NullInput;
@@ -47,6 +60,7 @@ import google.registry.request.Response;
 import google.registry.request.auth.Auth;
 import google.registry.util.Clock;
 import google.registry.xml.ValidationMode;
+import java.io.IOException;
 import java.util.Optional;
 import javax.inject.Inject;
 import org.joda.time.DateTime;
@@ -201,6 +215,7 @@ public final class RdeStagingAction implements Runnable {
 
   public static final String PATH = "/_dr/task/rdeStaging";
 
+  private static final String PIPELINE_NAME = "rde_pipeline";
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
 
   @Inject Clock clock;
@@ -209,7 +224,11 @@ public final class RdeStagingAction implements Runnable {
   @Inject Response response;
   @Inject GcsUtils gcsUtils;
   @Inject MapreduceRunner mrRunner;
+  @Inject @Config("projectId") String projectId;
+  @Inject @Config("defaultJobRegion") String jobRegion;
   @Inject @Config("transactionCooldown") Duration transactionCooldown;
+  @Inject @Config("beamStagingBucketUrl") String stagingBucketUrl;
+  @Inject @Config("rdeBucket") String rdeBucket;
   @Inject @Parameter(RdeModule.PARAM_MANUAL) boolean manual;
   @Inject @Parameter(RdeModule.PARAM_DIRECTORY) Optional<String> directory;
   @Inject @Parameter(RdeModule.PARAM_MODE) ImmutableSet<String> modeStrings;
@@ -217,7 +236,8 @@ public final class RdeStagingAction implements Runnable {
   @Inject @Parameter(RdeModule.PARAM_WATERMARKS) ImmutableSet<DateTime> watermarks;
   @Inject @Parameter(RdeModule.PARAM_REVISION) Optional<Integer> revision;
   @Inject @Parameter(RdeModule.PARAM_LENIENT) boolean lenient;
-
+  @Inject @Key("rdeStagingEncryptionKey") byte[] stagingKeyBytes;
+  @Inject Dataflow dataflow;
   @Inject RdeStagingAction() {}
 
   @Override
@@ -228,27 +248,66 @@ public final class RdeStagingAction implements Runnable {
       String message = "Nothing needs to be deposited.";
       logger.atInfo().log(message);
       response.setStatus(SC_NO_CONTENT);
-      response.setPayload(message);
+      // No need to set payload as HTTP 204 response status code does not allow a payload.
       return;
     }
     for (PendingDeposit pending : pendings.values()) {
       logger.atInfo().log("Pending deposit: %s", pending);
     }
     ValidationMode validationMode = lenient ? LENIENT : STRICT;
-    RdeStagingMapper mapper = new RdeStagingMapper(validationMode, pendings);
-    RdeStagingReducer reducer = reducerFactory.create(validationMode, gcsUtils);
-
-    mrRunner
-        .setJobName("Stage escrow deposits for all TLDs")
-        .setModuleName("backend")
-        .setDefaultReduceShards(pendings.size())
-        .runMapreduce(
-            mapper,
-            reducer,
-            ImmutableList.of(
-                // Add an extra shard that maps over a null resource. See the mapper code for why.
-                new NullInput<>(), EppResourceInputs.createEntityInput(EppResource.class)))
-        .sendLinkToMapreduceConsole(response);
+    if (tm().isOfy()) {
+      RdeStagingMapper mapper = new RdeStagingMapper(validationMode, pendings);
+      RdeStagingReducer reducer = reducerFactory.create(validationMode, gcsUtils);
+      mrRunner
+          .setJobName("Stage escrow deposits for all TLDs")
+          .setModuleName("backend")
+          .setDefaultReduceShards(pendings.size())
+          .runMapreduce(
+              mapper,
+              reducer,
+              ImmutableList.of(
+                  // Add an extra shard that maps over a null resource. See the mapper code for why.
+                  new NullInput<>(), EppResourceInputs.createEntityInput(EppResource.class)))
+          .sendLinkToMapreduceConsole(response);
+    } else {
+      try {
+        LaunchFlexTemplateParameter parameter =
+            new LaunchFlexTemplateParameter()
+                .setJobName(createJobName("rde", clock))
+                .setContainerSpecGcsPath(
+                    String.format("%s/%s_metadata.json", stagingBucketUrl, PIPELINE_NAME))
+                .setParameters(
+                    ImmutableMap.of(
+                        "pendings",
+                        RdePipeline.encodePendings(pendings),
+                        "validationMode",
+                        validationMode.name(),
+                        "rdeStagingBucket",
+                        rdeBucket,
+                        "stagingKey",
+                        BaseEncoding.base64Url().omitPadding().encode(stagingKeyBytes),
+                        "registryEnvironment",
+                        RegistryEnvironment.get().name()));
+        LaunchFlexTemplateResponse launchResponse =
+            dataflow
+                .projects()
+                .locations()
+                .flexTemplates()
+                .launch(
+                    projectId,
+                    jobRegion,
+                    new LaunchFlexTemplateRequest().setLaunchParameter(parameter))
+                .execute();
+        logger.atInfo().log("Got response: %s", launchResponse.getJob().toPrettyString());
+        response.setStatus(SC_OK);
+        response.setPayload(
+            String.format("Launched RDE pipeline: %s", launchResponse.getJob().getId()));
+      } catch (IOException e) {
+        logger.atWarning().withCause(e).log("Pipeline Launch failed");
+        response.setStatus(SC_INTERNAL_SERVER_ERROR);
+        response.setPayload(String.format("Pipeline launch failed: %s", e.getMessage()));
+      }
+    }
   }
 
   private ImmutableSetMultimap<String, PendingDeposit> getStandardPendingDeposits() {

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

@@ -14,7 +14,6 @@
 
 package google.registry.rde;
 
-import static com.google.appengine.api.taskqueue.TaskOptions.Builder.withUrl;
 import static com.google.common.base.Verify.verify;
 import static com.google.common.net.MediaType.PLAIN_TEXT_UTF_8;
 import static com.jcraft.jsch.ChannelSftp.OVERWRITE;
@@ -24,14 +23,15 @@ import static google.registry.model.common.Cursor.getCursorTimeOrStartOfTime;
 import static google.registry.model.rde.RdeMode.FULL;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
+import static google.registry.rde.RdeModule.RDE_REPORT_QUEUE;
 import static google.registry.request.Action.Method.POST;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
 import static google.registry.util.DateTimeUtils.isBeforeOrAt;
 import static java.util.Arrays.asList;
 
-import com.google.appengine.api.taskqueue.Queue;
 import com.google.cloud.storage.BlobId;
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.HashMultimap;
 import com.google.common.flogger.FluentLogger;
 import com.google.common.io.ByteStreams;
 import com.jcraft.jsch.JSch;
@@ -49,14 +49,15 @@ import google.registry.model.tld.Registry;
 import google.registry.rde.EscrowTaskRunner.EscrowTask;
 import google.registry.rde.JSchSshSession.JSchSshSessionFactory;
 import google.registry.request.Action;
+import google.registry.request.Action.Service;
 import google.registry.request.HttpException.NoContentException;
 import google.registry.request.Parameter;
 import google.registry.request.RequestParameters;
 import google.registry.request.Response;
 import google.registry.request.auth.Auth;
 import google.registry.util.Clock;
+import google.registry.util.CloudTasksUtils;
 import google.registry.util.Retrier;
-import google.registry.util.TaskQueueUtils;
 import google.registry.util.TeeOutputStream;
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
@@ -66,7 +67,6 @@ import java.io.OutputStream;
 import java.net.URI;
 import java.util.Optional;
 import javax.inject.Inject;
-import javax.inject.Named;
 import org.bouncycastle.openpgp.PGPKeyPair;
 import org.bouncycastle.openpgp.PGPPrivateKey;
 import org.bouncycastle.openpgp.PGPPublicKey;
@@ -90,7 +90,7 @@ import org.joda.time.Duration;
     auth = Auth.AUTH_INTERNAL_OR_ADMIN)
 public final class RdeUploadAction implements Runnable, EscrowTask {
 
-  static final String PATH = "/_dr/task/rdeUpload";
+  public static final String PATH = "/_dr/task/rdeUpload";
 
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
 
@@ -109,9 +109,10 @@ public final class RdeUploadAction implements Runnable, EscrowTask {
   @Inject JSchSshSessionFactory jschSshSessionFactory;
   @Inject Response response;
   @Inject SftpProgressMonitor sftpProgressMonitor;
-  @Inject TaskQueueUtils taskQueueUtils;
+  @Inject CloudTasksUtils cloudTasksUtils;
   @Inject Retrier retrier;
   @Inject @Parameter(RequestParameters.PARAM_TLD) String tld;
+  @Inject @Parameter(RdeModule.PARAM_PREFIX) Optional<String> prefix;
   @Inject @Config("rdeBucket") String bucket;
   @Inject @Config("rdeInterval") Duration interval;
   @Inject @Config("rdeUploadLockTimeout") Duration timeout;
@@ -120,15 +121,21 @@ public final class RdeUploadAction implements Runnable, EscrowTask {
   @Inject @Key("rdeReceiverKey") PGPPublicKey receiverKey;
   @Inject @Key("rdeSigningKey") PGPKeyPair signingKey;
   @Inject @Key("rdeStagingDecryptionKey") PGPPrivateKey stagingDecryptionKey;
-  @Inject @Named("rde-report") Queue reportQueue;
   @Inject RdeUploadAction() {}
 
   @Override
   public void run() {
     logger.atInfo().log("Attempting to acquire RDE upload lock for TLD '%s'.", tld);
     runner.lockRunAndRollForward(this, Registry.get(tld), timeout, CursorType.RDE_UPLOAD, interval);
-    taskQueueUtils.enqueue(
-        reportQueue, withUrl(RdeReportAction.PATH).param(RequestParameters.PARAM_TLD, tld));
+    HashMultimap<String, String> params = HashMultimap.create();
+    params.put(RequestParameters.PARAM_TLD, tld);
+    if (prefix.isPresent()) {
+      params.put(RdeModule.PARAM_PREFIX, prefix.get());
+    }
+    cloudTasksUtils.enqueue(
+        RDE_REPORT_QUEUE,
+        CloudTasksUtils.createPostTask(
+            RdeReportAction.PATH, Service.BACKEND.getServiceId(), params));
   }
 
   @Override
@@ -164,7 +171,9 @@ public final class RdeUploadAction implements Runnable, EscrowTask {
         RdeRevision.getCurrentRevision(tld, watermark, FULL)
             .orElseThrow(
                 () -> new IllegalStateException("RdeRevision was not set on generated deposit"));
-    final String name = RdeNamingUtils.makeRydeFilename(tld, watermark, FULL, 1, revision);
+    final String nameWithoutPrefix =
+        RdeNamingUtils.makeRydeFilename(tld, watermark, FULL, 1, revision);
+    final String name = prefix.orElse("") + nameWithoutPrefix;
     final BlobId xmlFilename = BlobId.of(bucket, name + ".xml.ghostryde");
     final BlobId xmlLengthFilename = BlobId.of(bucket, name + ".xml.length");
     BlobId reportFilename = BlobId.of(bucket, name + "-report.xml.ghostryde");
@@ -174,7 +183,8 @@ public final class RdeUploadAction implements Runnable, EscrowTask {
     logger.atInfo().log("Commencing RDE upload for TLD '%s' to '%s'.", tld, uploadUrl);
     final long xmlLength = readXmlLength(xmlLengthFilename);
     retrier.callWithRetry(
-        () -> upload(xmlFilename, xmlLength, watermark, name), JSchException.class);
+        () -> upload(xmlFilename, xmlLength, watermark, name, nameWithoutPrefix),
+        JSchException.class);
     logger.atInfo().log(
         "Updating RDE cursor '%s' for TLD '%s' following successful upload.", RDE_UPLOAD_SFTP, tld);
     tm().transact(
@@ -210,7 +220,8 @@ public final class RdeUploadAction implements Runnable, EscrowTask {
    * }</pre>
    */
   @VisibleForTesting
-  protected void upload(BlobId xmlFile, long xmlLength, DateTime watermark, String name)
+  protected void upload(
+      BlobId xmlFile, long xmlLength, DateTime watermark, String name, String nameWithoutPrefix)
       throws Exception {
     logger.atInfo().log("Uploading XML file '%s' to remote path '%s'.", xmlFile, uploadUrl);
     try (InputStream gcsInput = gcsUtils.openInputStream(xmlFile);
@@ -218,8 +229,8 @@ public final class RdeUploadAction implements Runnable, EscrowTask {
       try (JSchSshSession session = jschSshSessionFactory.create(lazyJsch.get(), uploadUrl);
           JSchSftpChannel ftpChan = session.openSftpChannel()) {
         ByteArrayOutputStream sigOut = new ByteArrayOutputStream();
-        String rydeFilename = name + ".ryde";
-        BlobId rydeGcsFilename = BlobId.of(bucket, rydeFilename);
+        String rydeFilename = nameWithoutPrefix + ".ryde";
+        BlobId rydeGcsFilename = BlobId.of(bucket, name + ".ryde");
         try (OutputStream ftpOutput =
                 ftpChan.get().put(rydeFilename, sftpProgressMonitor, OVERWRITE);
             OutputStream gcsOutput = gcsUtils.openOutputStream(rydeGcsFilename);
@@ -228,14 +239,15 @@ public final class RdeUploadAction implements Runnable, EscrowTask {
                 new RydeEncoder.Builder()
                     .setRydeOutput(teeOutput, receiverKey)
                     .setSignatureOutput(sigOut, signingKey)
-                    .setFileMetadata(name, xmlLength, watermark)
+                    .setFileMetadata(nameWithoutPrefix, xmlLength, watermark)
                     .build()) {
             long bytesCopied = ByteStreams.copy(ghostrydeDecoder, rydeEncoder);
           logger.atInfo().log("Uploaded %,d bytes to path '%s'.", bytesCopied, rydeFilename);
           }
-        String sigFilename = name + ".sig";
+        String sigFilename = nameWithoutPrefix + ".sig";
+        BlobId sigGcsFilename = BlobId.of(bucket, name + ".sig");
         byte[] signature = sigOut.toByteArray();
-        gcsUtils.createFromBytes(BlobId.of(bucket, sigFilename), signature);
+        gcsUtils.createFromBytes(sigGcsFilename, signature);
         ftpChan.get().put(new ByteArrayInputStream(signature), sigFilename);
         logger.atInfo().log("Uploaded %,d bytes to path '%s'.", signature.length, sigFilename);
       }

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

@@ -30,6 +30,7 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.flogger.FluentLogger;
 import com.google.common.net.MediaType;
 import google.registry.config.RegistryConfig.Config;
+import google.registry.config.RegistryEnvironment;
 import google.registry.model.common.DatabaseMigrationStateSchedule.PrimaryDatabase;
 import google.registry.reporting.ReportingModule;
 import google.registry.request.Action;
@@ -127,7 +128,9 @@ public class GenerateInvoicesAction implements Runnable {
                       "database",
                       database.name(),
                       "billingBucketUrl",
-                      billingBucketUrl));
+                      billingBucketUrl,
+                      "registryEnvironment",
+                      RegistryEnvironment.get().name()));
       LaunchFlexTemplateResponse launchResponse =
           dataflow
               .projects()

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

@@ -16,6 +16,7 @@ package google.registry.reporting.icann;
 
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.reporting.icann.IcannReportingModule.DATASTORE_EXPORT_DATA_SET;
+import static google.registry.reporting.icann.IcannReportingModule.ICANN_REPORTING_DATA_SET;
 import static google.registry.reporting.icann.QueryBuilderUtils.getQueryFromFile;
 import static google.registry.reporting.icann.QueryBuilderUtils.getTableName;
 
@@ -23,6 +24,7 @@ import com.google.common.collect.ImmutableMap;
 import google.registry.config.RegistryConfig.Config;
 import google.registry.util.SqlTemplate;
 import javax.inject.Inject;
+import javax.inject.Named;
 import org.joda.time.LocalDate;
 import org.joda.time.YearMonth;
 import org.joda.time.format.DateTimeFormat;
@@ -38,29 +40,32 @@ public final class ActivityReportingQueryBuilder implements QueryBuilder {
   static final String EPP_METRICS = "epp_metrics";
   static final String WHOIS_COUNTS = "whois_counts";
   static final String ACTIVITY_REPORT_AGGREGATION = "activity_report_aggregation";
-  final String BIGQUERY_DATA_SET = tm().isOfy() ? "icann_reporting" : "cloud_sql_icann_reporting";
+
+  private final String projectId;
+  private final DnsCountQueryCoordinator dnsCountQueryCoordinator;
+  private final String icannReportingDataSet;
 
   @Inject
-  @Config("projectId")
-  String projectId;
-
-  @Inject DnsCountQueryCoordinator dnsCountQueryCoordinator;
-
-  @Inject
-  ActivityReportingQueryBuilder() {}
+  ActivityReportingQueryBuilder(
+      @Config("projectId") String projectId,
+      @Named(ICANN_REPORTING_DATA_SET) String icannReportingDataSet,
+      DnsCountQueryCoordinator dnsCountQueryCoordinator) {
+    this.projectId = projectId;
+    this.dnsCountQueryCoordinator = dnsCountQueryCoordinator;
+    this.icannReportingDataSet = icannReportingDataSet;
+  }
 
   /** Returns the aggregate query which generates the activity report from the saved view. */
   @Override
   public String getReportQuery(YearMonth yearMonth) {
     return String.format(
         "#standardSQL\nSELECT * FROM `%s.%s.%s`",
-        projectId, BIGQUERY_DATA_SET, getTableName(ACTIVITY_REPORT_AGGREGATION, yearMonth));
+        projectId, icannReportingDataSet, getTableName(ACTIVITY_REPORT_AGGREGATION, yearMonth));
   }
 
   /** Sets the month we're doing activity reporting for, and returns the view query map. */
   @Override
   public ImmutableMap<String, String> getViewQueryMap(YearMonth yearMonth) {
-
     LocalDate firstDayOfMonth = yearMonth.toLocalDate(1);
     // The pattern-matching is inclusive, so we subtract 1 day to only report that month's data.
     LocalDate lastDayOfMonth = yearMonth.toLocalDate(1).plusMonths(1).minusDays(1);
@@ -102,7 +107,7 @@ public final class ActivityReportingQueryBuilder implements QueryBuilder {
     String eppQuery =
         SqlTemplate.create(getQueryFromFile("epp_metrics.sql"))
             .put("PROJECT_ID", projectId)
-            .put("ICANN_REPORTING_DATA_SET", BIGQUERY_DATA_SET)
+            .put("ICANN_REPORTING_DATA_SET", icannReportingDataSet)
             .put("MONTHLY_LOGS_TABLE", getTableName(MONTHLY_LOGS, yearMonth))
             // All metadata logs for reporting come from google.registry.flows.FlowReporter.
             .put(
@@ -114,7 +119,7 @@ public final class ActivityReportingQueryBuilder implements QueryBuilder {
     String whoisQuery =
         SqlTemplate.create(getQueryFromFile("whois_counts.sql"))
             .put("PROJECT_ID", projectId)
-            .put("ICANN_REPORTING_DATA_SET", BIGQUERY_DATA_SET)
+            .put("ICANN_REPORTING_DATA_SET", icannReportingDataSet)
             .put("MONTHLY_LOGS_TABLE", getTableName(MONTHLY_LOGS, yearMonth))
             .build();
     queriesBuilder.put(getTableName(WHOIS_COUNTS, yearMonth), whoisQuery);
@@ -129,7 +134,7 @@ public final class ActivityReportingQueryBuilder implements QueryBuilder {
             .put(
                 "REGISTRAR_OPERATING_STATUS_TABLE",
                 getTableName(REGISTRAR_OPERATING_STATUS, yearMonth))
-            .put("ICANN_REPORTING_DATA_SET", BIGQUERY_DATA_SET)
+            .put("ICANN_REPORTING_DATA_SET", icannReportingDataSet)
             .put("DNS_COUNTS_TABLE", getTableName(DNS_COUNTS, yearMonth))
             .put("EPP_METRICS_TABLE", getTableName(EPP_METRICS, yearMonth))
             .put("WHOIS_COUNTS_TABLE", getTableName(WHOIS_COUNTS, yearMonth));
@@ -147,7 +152,7 @@ public final class ActivityReportingQueryBuilder implements QueryBuilder {
     return queriesBuilder.build();
   }
 
-  public void prepareForQuery(YearMonth yearMonth) throws InterruptedException {
+  void prepareForQuery(YearMonth yearMonth) throws InterruptedException {
     dnsCountQueryCoordinator.prepareForQuery(yearMonth);
   }
 }

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

@@ -14,7 +14,6 @@
 
 package google.registry.reporting.icann;
 
-import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.request.RequestParameters.extractOptionalParameter;
 import static google.registry.request.RequestParameters.extractRequiredParameter;
 import static google.registry.request.RequestParameters.extractSetOfEnumParameters;
@@ -24,9 +23,11 @@ import com.google.common.util.concurrent.MoreExecutors;
 import dagger.Module;
 import dagger.Provides;
 import google.registry.bigquery.BigqueryConnection;
+import google.registry.persistence.transaction.TransactionManager;
 import google.registry.request.HttpException.BadRequestException;
 import google.registry.request.Parameter;
 import java.util.Optional;
+import javax.inject.Named;
 import javax.servlet.http.HttpServletRequest;
 import org.joda.time.Duration;
 
@@ -42,8 +43,7 @@ public final class IcannReportingModule {
 
   static final String PARAM_SUBDIR = "subdir";
   static final String PARAM_REPORT_TYPES = "reportTypes";
-  static final String ICANN_REPORTING_DATA_SET =
-      tm().isOfy() ? "icann_reporting" : "cloud_sql_icann_reporting";
+  static final String ICANN_REPORTING_DATA_SET = "icannReportingDataSet";
   static final String DATASTORE_EXPORT_DATA_SET = "latest_datastore_export";
   static final String MANIFEST_FILE_NAME = "MANIFEST.txt";
 
@@ -88,11 +88,12 @@ public final class IcannReportingModule {
    */
   @Provides
   static BigqueryConnection provideBigqueryConnection(
-      BigqueryConnection.Builder bigQueryConnectionBuilder) {
+      BigqueryConnection.Builder bigQueryConnectionBuilder,
+      @Named(ICANN_REPORTING_DATA_SET) String icannReportingDataSet) {
     try {
       return bigQueryConnectionBuilder
           .setExecutorService(MoreExecutors.newDirectExecutorService())
-          .setDatasetId(ICANN_REPORTING_DATA_SET)
+          .setDatasetId(icannReportingDataSet)
           .setOverwrite(true)
           .setPollInterval(Duration.standardSeconds(1))
           .build();
@@ -100,4 +101,10 @@ public final class IcannReportingModule {
       throw new RuntimeException("Could not initialize BigqueryConnection!", e);
     }
   }
+
+  @Provides
+  @Named(ICANN_REPORTING_DATA_SET)
+  static String provideIcannReportingDataSet(TransactionManager tm) {
+    return tm.isOfy() ? "icann_reporting" : "cloud_sql_icann_reporting";
+  }
 }

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

@@ -23,6 +23,7 @@ import com.google.common.collect.ImmutableMap;
 import google.registry.config.RegistryConfig.Config;
 import google.registry.util.SqlTemplate;
 import javax.inject.Inject;
+import javax.inject.Named;
 import org.joda.time.DateTime;
 import org.joda.time.LocalTime;
 import org.joda.time.YearMonth;
@@ -34,9 +35,16 @@ import org.joda.time.format.DateTimeFormatter;
  */
 public final class TransactionsReportingQueryBuilder implements QueryBuilder {
 
-  @Inject @Config("projectId") String projectId;
+  final String projectId;
+  private final String icannReportingDataSet;
 
-  @Inject TransactionsReportingQueryBuilder() {}
+  @Inject
+  TransactionsReportingQueryBuilder(
+      @Config("projectId") String projectId,
+      @Named(ICANN_REPORTING_DATA_SET) String icannReportingDataSet) {
+    this.projectId = projectId;
+    this.icannReportingDataSet = icannReportingDataSet;
+  }
 
   static final String TRANSACTIONS_REPORT_AGGREGATION = "transactions_report_aggregation";
   static final String REGISTRAR_IANA_ID = "registrar_iana_id";
@@ -51,9 +59,7 @@ public final class TransactionsReportingQueryBuilder implements QueryBuilder {
   public String getReportQuery(YearMonth yearMonth) {
     return String.format(
         "#standardSQL\nSELECT * FROM `%s.%s.%s`",
-        projectId,
-        ICANN_REPORTING_DATA_SET,
-        getTableName(TRANSACTIONS_REPORT_AGGREGATION, yearMonth));
+        projectId, icannReportingDataSet, getTableName(TRANSACTIONS_REPORT_AGGREGATION, yearMonth));
   }
 
   /** Sets the month we're doing transactions reporting for, and returns the view query map. */
@@ -151,7 +157,7 @@ public final class TransactionsReportingQueryBuilder implements QueryBuilder {
             .put("PROJECT_ID", projectId)
             .put("DATASTORE_EXPORT_DATA_SET", DATASTORE_EXPORT_DATA_SET)
             .put("REGISTRY_TABLE", "Registry")
-            .put("ICANN_REPORTING_DATA_SET", ICANN_REPORTING_DATA_SET)
+            .put("ICANN_REPORTING_DATA_SET", icannReportingDataSet)
             .put("REGISTRAR_IANA_ID_TABLE", getTableName(REGISTRAR_IANA_ID, yearMonth))
             .put("TOTAL_DOMAINS_TABLE", getTableName(TOTAL_DOMAINS, yearMonth))
             .put("TOTAL_NAMESERVERS_TABLE", getTableName(TOTAL_NAMESERVERS, yearMonth))

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

@@ -36,6 +36,7 @@ import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.secdns.DelegationSignerData;
 import google.registry.model.eppcommon.StatusValue;
 import google.registry.model.host.HostResource;
+import google.registry.tools.soy.DomainRenewSoyInfo;
 import google.registry.tools.soy.UniformRapidSuspensionSoyInfo;
 import java.util.ArrayList;
 import java.util.List;
@@ -43,6 +44,7 @@ import java.util.Optional;
 import java.util.Set;
 import javax.xml.bind.annotation.adapters.HexBinaryAdapter;
 import org.joda.time.DateTime;
+import org.joda.time.format.DateTimeFormat;
 
 /** A command to suspend a domain for the Uniform Rapid Suspension process. */
 @Parameters(separators = " =",
@@ -97,6 +99,13 @@ final class UniformRapidSuspensionCommand extends MutatingEppToolCommand {
       description = "Flag indicating that is is an undo command, which removes locks.")
   private boolean undo;
 
+  @Parameter(
+      names = {"--renew_one_year"},
+      required = true,
+      description = "Flag indicating whether or not the domain will be renewed for a year.",
+      arity = 1)
+  private boolean renewOneYear;
+
   /** Set of existing locks that need to be preserved during undo, sorted for nicer output. */
   ImmutableSortedSet<String> existingLocks;
 
@@ -114,19 +123,20 @@ final class UniformRapidSuspensionCommand extends MutatingEppToolCommand {
     superuser = true;
     DateTime now = DateTime.now(UTC);
     ImmutableSet<String> newHostsSet = ImmutableSet.copyOf(newHosts);
-    Optional<DomainBase> domain = loadByForeignKey(DomainBase.class, domainName, now);
-    checkArgumentPresent(domain, "Domain '%s' does not exist or is deleted", domainName);
+    Optional<DomainBase> domainOpt = loadByForeignKey(DomainBase.class, domainName, now);
+    checkArgumentPresent(domainOpt, "Domain '%s' does not exist or is deleted", domainName);
+    DomainBase domain = domainOpt.get();
     Set<String> missingHosts =
         difference(newHostsSet, checkResourcesExist(HostResource.class, newHosts, now));
     checkArgument(missingHosts.isEmpty(), "Hosts do not exist: %s", missingHosts);
     checkArgument(
         locksToPreserve.isEmpty() || undo,
         "Locks can only be preserved when running with --undo");
-    existingNameservers = getExistingNameservers(domain.get());
-    existingLocks = getExistingLocks(domain.get());
-    existingDsData = getExistingDsData(domain.get());
+    existingNameservers = getExistingNameservers(domain);
+    existingLocks = getExistingLocks(domain);
+    existingDsData = getExistingDsData(domain);
     removeStatuses =
-        (hasClientHold(domain.get()) && !undo)
+        (hasClientHold(domain) && !undo)
             ? ImmutableSet.of(StatusValue.CLIENT_HOLD.getXmlName())
             : ImmutableSet.of();
     ImmutableSet<String> statusesToApply;
@@ -138,6 +148,25 @@ final class UniformRapidSuspensionCommand extends MutatingEppToolCommand {
     } else {
       statusesToApply = URS_LOCKS;
     }
+
+    // trigger renew flow
+    if (renewOneYear) {
+      setSoyTemplate(DomainRenewSoyInfo.getInstance(), DomainRenewSoyInfo.RENEWDOMAIN);
+      addSoyRecord(
+          CLIENT_ID,
+          new SoyMapData(
+              "domainName",
+              domain.getDomainName(),
+              "expirationDate",
+              domain
+                  .getRegistrationExpirationTime()
+                  .toString(DateTimeFormat.forPattern("YYYY-MM-dd")),
+              // period is the number of years to renew the registration for
+              "period",
+              String.valueOf(1)));
+    }
+
+    // trigger update flow
     setSoyTemplate(
         UniformRapidSuspensionSoyInfo.getInstance(),
         UniformRapidSuspensionSoyInfo.UNIFORMRAPIDSUSPENSION);

core/src/main/java/google/registry/tools/javascrap/CreateSyntheticHistoryEntriesAction.java

@@ -21,6 +21,7 @@ import static google.registry.persistence.transaction.TransactionManagerFactory.
 import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
 
 import com.google.appengine.tools.mapreduce.Mapper;
+import com.google.common.base.CaseFormat;
 import com.google.common.collect.ImmutableList;
 import com.googlecode.objectify.Key;
 import google.registry.config.RegistryConfig.Config;
@@ -29,16 +30,12 @@ import google.registry.mapreduce.inputs.EppResourceInputs;
 import google.registry.model.EppResource;
 import google.registry.model.domain.DomainHistory;
 import google.registry.model.reporting.HistoryEntry;
-import google.registry.persistence.transaction.CriteriaQueryBuilder;
 import google.registry.rde.RdeStagingAction;
 import google.registry.request.Action;
 import google.registry.request.Response;
 import google.registry.request.auth.Auth;
 import google.registry.tools.server.GenerateZoneFilesAction;
-import java.util.Optional;
 import javax.inject.Inject;
-import javax.persistence.criteria.CriteriaBuilder;
-import javax.persistence.criteria.CriteriaQuery;
 
 /**
  * A mapreduce that creates synthetic history objects in SQL for all {@link EppResource} objects.
@@ -86,12 +83,12 @@ public class CreateSyntheticHistoryEntriesAction implements Runnable {
   }
 
   /**
-   * The number of shards to run the map-only mapreduce on.
+   * The default number of shards to run the map-only mapreduce on.
    *
-   * <p>This is much lower than the default of 100, or even 10, because we can afford it being slow
-   * and we want to avoid overloading SQL.
+   * <p>This is much lower than the default of 100 because we can afford it being slow and we want
+   * to avoid overloading SQL.
    */
-  private static final int NUM_SHARDS = 3;
+  private static final int NUM_SHARDS = 10;
 
   @Override
   public void run() {
@@ -105,8 +102,9 @@ public class CreateSyntheticHistoryEntriesAction implements Runnable {
         .sendLinkToMapreduceConsole(response);
   }
 
-  // Lifted from HistoryEntryDao
-  private static Optional<? extends HistoryEntry> mostRecentHistoryFromSql(EppResource resource) {
+  // Returns true iff any of the *History objects in SQL contain a representation of this resource
+  // at the point in time that the *History object was created.
+  private static boolean hasHistoryContainingResource(EppResource resource) {
     return jpaTm()
         .transact(
             () -> {
@@ -114,18 +112,24 @@ public class CreateSyntheticHistoryEntriesAction implements Runnable {
               Class<? extends HistoryEntry> historyClass =
                   getHistoryClassFromParent(resource.getClass());
               // The field representing repo ID unfortunately varies by history class
-              String repoIdFieldName = getRepoIdFieldNameFromHistoryClass(historyClass);
-              CriteriaBuilder criteriaBuilder = jpaTm().getEntityManager().getCriteriaBuilder();
-              CriteriaQuery<? extends HistoryEntry> criteriaQuery =
-                  CriteriaQueryBuilder.create(historyClass)
-                      .where(repoIdFieldName, criteriaBuilder::equal, resource.getRepoId())
-                      .orderByDesc("modificationTime")
-                      .build();
-              return jpaTm()
-                  .criteriaQuery(criteriaQuery)
-                  .setMaxResults(1)
-                  .getResultStream()
-                  .findFirst();
+              String repoIdFieldName =
+                  CaseFormat.LOWER_CAMEL.to(
+                      CaseFormat.LOWER_UNDERSCORE,
+                      getRepoIdFieldNameFromHistoryClass(historyClass));
+              // The "history" fields in the *History objects are all prefixed with "history_". If
+              // any of the non-"history_" fields are non-null, that means that that row contains
+              // a representation of that EppResource at that point in time. We use creation_time as
+              // a marker since it's the simplest field and all EppResources will have it.
+              return (boolean)
+                  jpaTm()
+                      .getEntityManager()
+                      .createNativeQuery(
+                          String.format(
+                              "SELECT EXISTS (SELECT 1 FROM \"%s\" WHERE %s = :repoId AND"
+                                  + " creation_time IS NOT NULL)",
+                              historyClass.getSimpleName(), repoIdFieldName))
+                      .setParameter("repoId", resource.getRepoId())
+                      .getSingleResult();
             });
   }
 
@@ -164,10 +168,7 @@ public class CreateSyntheticHistoryEntriesAction implements Runnable {
       EppResource eppResource = auditedOfy().load().key(resourceKey).now();
       // Only save new history entries if the most recent history for this object in SQL does not
       // have the resource at that point in time already
-      Optional<? extends HistoryEntry> maybeHistory = mostRecentHistoryFromSql(eppResource);
-      if (maybeHistory
-          .map(history -> !history.getResourceAtPointInTime().isPresent())
-          .orElse(true)) {
+      if (!hasHistoryContainingResource(eppResource)) {
         ofyTm()
             .transact(
                 () ->

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

@@ -6,9 +6,9 @@
       "name": "registryEnvironment",
       "label": "The Registry environment.",
       "helpText": "The Registry environment, required if environment-specific initialization (such as JPA) is needed on worker VMs.",
-      "is_optional": true,
+      "is_optional": false,
       "regexes": [
-        "^[0-9A-Z_]+$"
+        "^PRODUCTION|SANDBOX|CRASH|QA|ALPHA$"
       ]
     },
     {

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

@@ -2,12 +2,21 @@
   "name": "RDE/BRDA Deposit Generation",
   "description": "An Apache Beam pipeline generates RDE or BRDA deposits and deposits them to GCS with GhostRyde encryption.",
   "parameters": [
+    {
+      "name": "registryEnvironment",
+      "label": "The Registry environment.",
+      "helpText": "The Registry environment, required if environment-specific initialization (such as JPA) is needed on worker VMs.",
+      "is_optional": false,
+      "regexes": [
+        "^PRODUCTION|SANDBOX|CRASH|QA|ALPHA$"
+      ]
+    },
     {
       "name": "pendings",
       "label": "The pendings deposits to generate.",
       "helpText": "A TLD to PendingDeposit map that is serialized and Base64 URL-safe encoded.",
       "regexes": [
-        "A-Za-z0-9\\-_"
+        "[A-Za-z0-9\\-_]+"
       ]
     },
     {
@@ -19,12 +28,12 @@
       ]
     },
     {
-      "name": "gcsBucket",
-      "label": "The GCS bucket that where the resulting files will be  stored.",
+      "name": "rdeStagingBucket",
+      "label": "The GCS bucket that where the resulting files will be stored.",
       "helpText": "Only the bucket name itself, without the leading \"gs://\".",
       "is_optional": false,
       "regexes": [
-        "^[a-zA-Z0-9_\\-]+$"
+        "[a-zA-Z0-9_\\-]+$"
       ]
     },
     {
@@ -32,7 +41,7 @@
       "label": "The PGP public key used to encrypt the RDE/BRDA deposit files.",
       "helpText": "The key is Base64 URL-safe encoded.",
       "regexes": [
-        "A-Za-z0-9\\-_"
+        "[A-Za-z0-9\\-_]+"
       ]
     }
   ]

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

@@ -6,9 +6,9 @@
       "name": "registryEnvironment",
       "label": "The Registry environment.",
       "helpText": "The Registry environment, required if environment-specific initialization (such as JPA) is needed on worker VMs.",
-      "is_optional": true,
+      "is_optional": false,
       "regexes": [
-        "^[0-9A-Z_]+$"
+        "^PRODUCTION|SANDBOX|CRASH|QA|ALPHA$"
       ]
     },
     {

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

@@ -23,7 +23,6 @@ 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;
@@ -50,8 +49,8 @@ import org.junit.jupiter.api.extension.RegisterExtension;
 @DualDatabaseTest
 class WipeOutContactHistoryPiiActionTest {
 
+  private static final int TEST_BATCH_SIZE = 20;
   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")
@@ -115,7 +114,8 @@ class WipeOutContactHistoryPiiActionTest {
   void beforeEach() {
     response = new FakeResponse();
     action =
-        new WipeOutContactHistoryPiiAction(clock, MIN_MONTHS_BEFORE_WIPE_OUT, BATCH_SIZE, response);
+        new WipeOutContactHistoryPiiAction(
+            clock, MIN_MONTHS_BEFORE_WIPE_OUT, TEST_BATCH_SIZE, response);
   }
 
   @TestSqlOnly
@@ -133,10 +133,10 @@ class WipeOutContactHistoryPiiActionTest {
   }
 
   @TestSqlOnly
-  void getAllHistoryEntitiesOlderThan_returnsOnlyPartOfThePersistedEntities() {
+  void getAllHistoryEntitiesOlderThan_returnsOnlyOldEnoughPersistedEntities() {
     ImmutableList<ContactHistory> expectedToBeWipedOut =
         persistLotsOfContactHistoryEntities(
-            40, MIN_MONTHS_BEFORE_WIPE_OUT + 2, 0, defaultContactResource);
+            19, MIN_MONTHS_BEFORE_WIPE_OUT + 2, 0, defaultContactResource);
 
     // persisted entities that should not be part of the actual result
     persistLotsOfContactHistoryEntities(
@@ -145,7 +145,7 @@ class WipeOutContactHistoryPiiActionTest {
     jpaTm()
         .transact(
             () ->
-                Truth8.assertThat(
+                assertThat(
                         action.getNextContactHistoryEntitiesWithPiiBatch(
                             clock.nowUtc().minusMonths(MIN_MONTHS_BEFORE_WIPE_OUT)))
                     .containsExactlyElementsIn(expectedToBeWipedOut));
@@ -175,6 +175,8 @@ class WipeOutContactHistoryPiiActionTest {
         .isEqualTo(0);
 
     assertThat(response.getStatus()).isEqualTo(SC_OK);
+    assertThat(response.getPayload())
+        .isEqualTo("Done. Wiped out PII of 0 ContactHistory entities in total.");
   }
 
   @TestSqlOnly
@@ -197,6 +199,8 @@ class WipeOutContactHistoryPiiActionTest {
     assertAllEntitiesContainPii(DatabaseHelper.loadByEntitiesIfPresent(expectedToBeWipedOut));
 
     action.run();
+    assertThat(response.getPayload())
+        .isEqualTo("Done. Wiped out PII of 20 ContactHistory entities in total.");
 
     // The query should return an empty stream after the wipe out action.
     assertThat(
@@ -216,7 +220,7 @@ class WipeOutContactHistoryPiiActionTest {
   void run_withMultipleBatches_numOfEntitiesAsNonMultipleOfBatchSize_success() {
     int numOfMonthsFromNow = MIN_MONTHS_BEFORE_WIPE_OUT + 2;
     ImmutableList<ContactHistory> expectedToBeWipedOut =
-        persistLotsOfContactHistoryEntities(1234, numOfMonthsFromNow, 0, defaultContactResource);
+        persistLotsOfContactHistoryEntities(56, numOfMonthsFromNow, 0, defaultContactResource);
 
     // The query should return a subset of all persisted data.
     assertThat(
@@ -227,10 +231,12 @@ class WipeOutContactHistoryPiiActionTest {
                             .getNextContactHistoryEntitiesWithPiiBatch(
                                 clock.nowUtc().minusMonths(MIN_MONTHS_BEFORE_WIPE_OUT))
                             .count()))
-        .isEqualTo(BATCH_SIZE);
+        .isEqualTo(TEST_BATCH_SIZE);
 
     assertAllEntitiesContainPii(DatabaseHelper.loadByEntitiesIfPresent(expectedToBeWipedOut));
     action.run();
+    assertThat(response.getPayload())
+        .isEqualTo("Done. Wiped out PII of 56 ContactHistory entities in total.");
 
     // The query should return an empty stream after the wipe out action.
     assertThat(
@@ -250,7 +256,8 @@ class WipeOutContactHistoryPiiActionTest {
   void run_withMultipleBatches_numOfEntitiesAsMultiplesOfBatchSize_success() {
     int numOfMonthsFromNow = MIN_MONTHS_BEFORE_WIPE_OUT + 2;
     ImmutableList<ContactHistory> expectedToBeWipedOut =
-        persistLotsOfContactHistoryEntities(2000, numOfMonthsFromNow, 0, defaultContactResource);
+        persistLotsOfContactHistoryEntities(
+            TEST_BATCH_SIZE * 2, numOfMonthsFromNow, 0, defaultContactResource);
 
     // The query should return a subset of all persisted data.
     assertThat(
@@ -261,10 +268,12 @@ class WipeOutContactHistoryPiiActionTest {
                             .getNextContactHistoryEntitiesWithPiiBatch(
                                 clock.nowUtc().minusMonths(MIN_MONTHS_BEFORE_WIPE_OUT))
                             .count()))
-        .isEqualTo(BATCH_SIZE);
+        .isEqualTo(TEST_BATCH_SIZE);
 
     assertAllEntitiesContainPii(DatabaseHelper.loadByEntitiesIfPresent(expectedToBeWipedOut));
     action.run();
+    assertThat(response.getPayload())
+        .isEqualTo("Done. Wiped out PII of 40 ContactHistory entities in total.");
 
     // The query should return an empty stream after the wipe out action.
     assertThat(

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

@@ -40,13 +40,13 @@ public abstract class BeamActionTestBase {
   protected Dataflow dataflow = mock(Dataflow.class);
   private Projects projects = mock(Projects.class);
   private Locations locations = mock(Locations.class);
-  private FlexTemplates templates = mock(FlexTemplates.class);
+  protected FlexTemplates templates = mock(FlexTemplates.class);
   protected Launch launch = mock(Launch.class);
   private LaunchFlexTemplateResponse launchResponse =
       new LaunchFlexTemplateResponse().setJob(new Job().setId("jobid"));
 
   @BeforeEach
-  void beforeEach() throws Exception {
+  protected void beforeEach() throws Exception {
     when(dataflow.projects()).thenReturn(projects);
     when(projects.locations()).thenReturn(locations);
     when(locations.flexTemplates()).thenReturn(templates);

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

@@ -16,7 +16,7 @@ package google.registry.beam.initsql;
 
 import static com.google.common.base.Preconditions.checkState;
 import static google.registry.model.ofy.ObjectifyService.auditedOfy;
-import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
 
 import com.google.appengine.api.datastore.DatastoreService;
 import com.google.appengine.api.datastore.DatastoreServiceFactory;
@@ -75,7 +75,8 @@ public final class BackupTestStore implements AutoCloseable {
   /** Returns the timestamp of the transaction. */
   long transact(Iterable<Object> deletes, Iterable<Object> newOrUpdated) {
     long timestamp = fakeClock.nowUtc().getMillis();
-    tm().transact(
+    ofyTm()
+        .transact(
             () -> {
               auditedOfy().delete().entities(deletes);
               auditedOfy().save().entities(newOrUpdated);
@@ -91,7 +92,7 @@ public final class BackupTestStore implements AutoCloseable {
   @SafeVarargs
   public final long insertOrUpdate(Object... entities) {
     long timestamp = fakeClock.nowUtc().getMillis();
-    tm().transact(() -> auditedOfy().save().entities(entities).now());
+    ofyTm().transact(() -> auditedOfy().save().entities(entities).now());
     fakeClock.advanceOneMilli();
     return timestamp;
   }
@@ -100,7 +101,7 @@ public final class BackupTestStore implements AutoCloseable {
   @SafeVarargs
   public final long delete(Object... entities) {
     long timestamp = fakeClock.nowUtc().getMillis();
-    tm().transact(() -> auditedOfy().delete().entities(entities).now());
+    ofyTm().transact(() -> auditedOfy().delete().entities(entities).now());
     fakeClock.advanceOneMilli();
     return timestamp;
   }

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

@@ -0,0 +1,64 @@
+// Copyright 2021 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.beam.initsql;
+
+import static com.google.common.truth.Truth.assertThat;
+import static google.registry.beam.initsql.Transforms.repairBadData;
+import static google.registry.model.ofy.ObjectifyService.auditedOfy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
+import static google.registry.testing.DatabaseHelper.createTld;
+import static google.registry.testing.DatabaseHelper.newDomainBase;
+import static google.registry.testing.DatabaseHelper.newHostResource;
+
+import com.google.appengine.api.datastore.Entity;
+import google.registry.model.domain.DomainBase;
+import google.registry.model.host.HostResource;
+import google.registry.testing.AppEngineExtension;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+/** Unit tests for {@link Transforms}. */
+public class TransformsTest {
+
+  @RegisterExtension
+  public final AppEngineExtension appEngine =
+      AppEngineExtension.builder().withDatastoreAndCloudSql().build();
+
+  @BeforeEach
+  void beforeEach() {
+    createTld("tld");
+  }
+
+  @Test
+  void testRepairBadData_canonicalizesDomainName() {
+    DomainBase domain = newDomainBase("foobar.tld");
+    Entity entity = ofyTm().transact(() -> auditedOfy().toEntity(domain));
+    entity.setIndexedProperty("fullyQualifiedDomainName", "FOOBäR.TLD");
+    assertThat(((DomainBase) auditedOfy().toPojo(repairBadData(entity))).getDomainName())
+        .isEqualTo("xn--foobr-jra.tld");
+  }
+
+  @Test
+  void testRepairBadData_canonicalizesHostName() {
+    HostResource host = newHostResource("baz.foobar.tld");
+    Entity entity = ofyTm().transact(() -> auditedOfy().toEntity(host));
+    entity.setIndexedProperty(
+        "fullyQualifiedHostName", "b̴̹͔͓̣̭̫͇͕̻̬̱͇͗͌́̆̋͒a̶̬̖͚̋̈́̽̇͝͠z̵͠.FOOBäR.TLD");
+    assertThat(((HostResource) auditedOfy().toPojo(repairBadData(entity))).getHostName())
+        .isEqualTo(
+            "xn--baz-kdcb2ajgzb4jtg6doej4e6b9am7c7b6c5nd4k7gpa2a9a7dufyewec.xn--foobr-jra.tld");
+  }
+}

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

@@ -17,9 +17,6 @@ package google.registry.beam.invoicing;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.tld.Registry.TldState.GENERAL_AVAILABILITY;
-import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
-import static google.registry.persistence.transaction.TransactionManagerFactory.removeTmOverrideForTest;
-import static google.registry.persistence.transaction.TransactionManagerFactory.setTmForTest;
 import static google.registry.testing.DatabaseHelper.createTld;
 import static google.registry.testing.DatabaseHelper.newRegistry;
 import static google.registry.testing.DatabaseHelper.persistActiveDomain;
@@ -48,6 +45,7 @@ import google.registry.persistence.transaction.JpaTestExtensions.JpaIntegrationT
 import google.registry.testing.DatastoreEntityExtension;
 import google.registry.testing.FakeClock;
 import google.registry.testing.TestDataHelper;
+import google.registry.testing.TmOverrideExtension;
 import google.registry.util.ResourceUtils;
 import java.io.File;
 import java.nio.file.Files;
@@ -77,6 +75,25 @@ import org.junit.jupiter.api.io.TempDir;
 /** Unit tests for {@link InvoicingPipeline}. */
 class InvoicingPipelineTest {
 
+  @RegisterExtension
+  @Order(Order.DEFAULT - 1)
+  final transient DatastoreEntityExtension datastore =
+      new DatastoreEntityExtension().allThreads(true);
+
+  @RegisterExtension
+  final TestPipelineExtension pipeline =
+      TestPipelineExtension.create().enableAbandonedNodeEnforcement(true);
+
+  @RegisterExtension
+  final JpaIntegrationTestExtension database =
+      new JpaTestExtensions.Builder().withClock(new FakeClock()).buildIntegrationTestExtension();
+
+  @RegisterExtension
+  @Order(Order.DEFAULT + 1)
+  TmOverrideExtension tmOverrideExtension = TmOverrideExtension.withJpa();
+
+  @TempDir Path tmpDir;
+
   private static final String BILLING_BUCKET_URL = "billing_bucket";
   private static final String YEAR_MONTH = "2017-10";
   private static final String INVOICE_FILE_PREFIX = "REG-INV";
@@ -225,21 +242,6 @@ class InvoicingPipelineTest {
           "2017-10-01,2018-09-30,456,20.50,USD,10125,1,PURCHASE,bestdomains - test,1,"
               + "RENEW | TLD: test | TERM: 1-year,20.50,USD,116688");
 
-  @RegisterExtension
-  @Order(Order.DEFAULT - 1)
-  final transient DatastoreEntityExtension datastore =
-      new DatastoreEntityExtension().allThreads(true);
-
-  @RegisterExtension
-  final TestPipelineExtension pipeline =
-      TestPipelineExtension.create().enableAbandonedNodeEnforcement(true);
-
-  @RegisterExtension
-  final JpaIntegrationTestExtension database =
-      new JpaTestExtensions.Builder().withClock(new FakeClock()).buildIntegrationTestExtension();
-
-  @TempDir Path tmpDir;
-
   private final InvoicingPipelineOptions options =
       PipelineOptionsFactory.create().as(InvoicingPipelineOptions.class);
 
@@ -261,13 +263,12 @@ class InvoicingPipelineTest {
     String query = InvoicingPipeline.makeQuery("2017-10", "my-project-id");
     assertThat(query)
         .isEqualTo(TestDataHelper.loadFile(this.getClass(), "billing_events_test.sql"));
-    // This is necessary because the TestPipelineExtension verifies that the pipelien is run.
+    // This is necessary because the TestPipelineExtension verifies that the pipeline is run.
     pipeline.run();
   }
 
   @Test
   void testSuccess_fullSqlPipeline() throws Exception {
-    setTmForTest(jpaTm());
     setupCloudSql();
     options.setDatabase("CLOUD_SQL");
     InvoicingPipeline invoicingPipeline = new InvoicingPipeline(options);
@@ -282,18 +283,15 @@ class InvoicingPipelineTest {
                 + "UnitPriceCurrency,PONumber");
     assertThat(overallInvoice.subList(1, overallInvoice.size()))
         .containsExactlyElementsIn(EXPECTED_INVOICE_OUTPUT);
-    removeTmOverrideForTest();
   }
 
   @Test
   void testSuccess_readFromCloudSql() throws Exception {
-    setTmForTest(jpaTm());
     setupCloudSql();
     PCollection<BillingEvent> billingEvents = InvoicingPipeline.readFromCloudSql(options, pipeline);
     billingEvents = billingEvents.apply(new ChangeDomainRepo());
     PAssert.that(billingEvents).containsInAnyOrder(INPUT_EVENTS);
     pipeline.run().waitUntilFinish();
-    removeTmOverrideForTest();
   }
 
   @Test

core/src/test/java/google/registry/beam/rde/RdePipelineTest.java

@@ -22,7 +22,6 @@ import static google.registry.beam.rde.RdePipeline.encodePendings;
 import static google.registry.model.common.Cursor.CursorType.RDE_STAGING;
 import static google.registry.model.rde.RdeMode.FULL;
 import static google.registry.model.rde.RdeMode.THIN;
-import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.rde.RdeResourceType.CONTACT;
 import static google.registry.rde.RdeResourceType.DOMAIN;
@@ -64,14 +63,16 @@ import google.registry.model.tld.Registry;
 import google.registry.persistence.VKey;
 import google.registry.persistence.transaction.JpaTestExtensions;
 import google.registry.persistence.transaction.JpaTestExtensions.JpaIntegrationTestExtension;
-import google.registry.persistence.transaction.TransactionManagerFactory;
 import google.registry.rde.DepositFragment;
 import google.registry.rde.Ghostryde;
 import google.registry.rde.PendingDeposit;
 import google.registry.rde.RdeResourceType;
+import google.registry.testing.CloudTasksHelper;
+import google.registry.testing.CloudTasksHelper.TaskMatcher;
 import google.registry.testing.DatastoreEntityExtension;
 import google.registry.testing.FakeClock;
 import google.registry.testing.FakeKeyringModule;
+import google.registry.testing.TmOverrideExtension;
 import java.io.IOException;
 import java.util.function.Function;
 import java.util.regex.Matcher;
@@ -86,7 +87,6 @@ import org.bouncycastle.openpgp.PGPPrivateKey;
 import org.bouncycastle.openpgp.PGPPublicKey;
 import org.joda.time.DateTime;
 import org.joda.time.Duration;
-import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Order;
 import org.junit.jupiter.api.Test;
@@ -134,6 +134,8 @@ public class RdePipelineTest {
 
   private final GcsUtils gcsUtils = new GcsUtils(LocalStorageHelper.getOptions());
 
+  private final CloudTasksHelper cloudTasksHelper = new CloudTasksHelper();
+
   private final PGPPublicKey encryptionKey =
       new FakeKeyringModule().get().getRdeStagingEncryptionKey();
 
@@ -152,6 +154,10 @@ public class RdePipelineTest {
   final JpaIntegrationTestExtension database =
       new JpaTestExtensions.Builder().withClock(clock).buildIntegrationTestExtension();
 
+  @RegisterExtension
+  @Order(Order.DEFAULT + 1)
+  TmOverrideExtension tmOverrideExtension = TmOverrideExtension.withJpa();
+
   @RegisterExtension
   final TestPipelineExtension pipeline =
       TestPipelineExtension.fromOptions(options).enableAbandonedNodeEnforcement(true);
@@ -160,7 +166,6 @@ public class RdePipelineTest {
 
   @BeforeEach
   void beforeEach() throws Exception {
-    TransactionManagerFactory.setTmForTest(jpaTm());
     loadInitialData();
 
     // Two real registrars have been created by loadInitialData(), named "New Registrar" and "The
@@ -212,14 +217,9 @@ public class RdePipelineTest {
     options.setValidationMode("LENIENT");
     options.setStagingKey(
         BaseEncoding.base64Url().encode(PgpHelper.convertPublicKeyToBytes(encryptionKey)));
-    options.setGcsBucket("gcs-bucket");
+    options.setRdeStagingBucket("gcs-bucket");
     options.setJobName("rde-job");
-    rdePipeline = new RdePipeline(options, gcsUtils);
-  }
-
-  @AfterEach
-  void afterEach() {
-    TransactionManagerFactory.removeTmOverrideForTest();
+    rdePipeline = new RdePipeline(options, gcsUtils, cloudTasksHelper.getTestCloudTasksUtils());
   }
 
   @Test
@@ -314,6 +314,21 @@ public class RdePipelineTest {
     assertThat(loadCursorTime(CursorType.RDE_STAGING))
         .isEquivalentAccordingToCompareTo(now.plus(Duration.standardDays(1)));
     assertThat(loadRevision(now, FULL)).isEqualTo(1);
+    cloudTasksHelper.assertTasksEnqueued(
+        "brda",
+        new TaskMatcher()
+            .url("/_dr/task/brdaCopy")
+            .service("backend")
+            .param("tld", "soy")
+            .param("watermark", now.toString())
+            .param("prefix", "rde-job/"));
+    cloudTasksHelper.assertTasksEnqueued(
+        "rde-upload",
+        new TaskMatcher()
+            .url("/_dr/task/rdeUpload")
+            .service("backend")
+            .param("tld", "soy")
+            .param("prefix", "rde-job/"));
   }
 
   // The GCS folder listing can be a bit flaky, so retry if necessary
@@ -337,6 +352,7 @@ public class RdePipelineTest {
 
     assertThat(loadCursorTime(CursorType.RDE_STAGING)).isEquivalentAccordingToCompareTo(now);
     assertThat(loadRevision(now, FULL)).isEqualTo(0);
+    cloudTasksHelper.assertNoTasksEnqueued("brda", "rde-upload");
   }
 
   private void verifyFiles(

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

@@ -18,8 +18,6 @@ import static com.google.common.collect.ImmutableList.toImmutableList;
 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.persistence.transaction.TransactionManagerFactory.removeTmOverrideForTest;
-import static google.registry.persistence.transaction.TransactionManagerFactory.setTmForTest;
 import static google.registry.testing.AppEngineExtension.makeRegistrar1;
 import static google.registry.testing.DatabaseHelper.createTld;
 import static google.registry.testing.DatabaseHelper.persistActiveContact;
@@ -52,6 +50,7 @@ import google.registry.persistence.transaction.JpaTestExtensions.JpaIntegrationT
 import google.registry.testing.DatastoreEntityExtension;
 import google.registry.testing.FakeClock;
 import google.registry.testing.FakeSleeper;
+import google.registry.testing.TmOverrideExtension;
 import google.registry.util.ResourceUtils;
 import google.registry.util.Retrier;
 import java.io.File;
@@ -129,6 +128,10 @@ class Spec11PipelineTest {
   final JpaIntegrationTestExtension database =
       new JpaTestExtensions.Builder().withClock(new FakeClock()).buildIntegrationTestExtension();
 
+  @RegisterExtension
+  @Order(Order.DEFAULT + 1)
+  TmOverrideExtension tmOverrideExtension = TmOverrideExtension.withJpa();
+
   private final Spec11PipelineOptions options =
       PipelineOptionsFactory.create().as(Spec11PipelineOptions.class);
 
@@ -233,7 +236,6 @@ class Spec11PipelineTest {
   }
 
   private void setupCloudSql() {
-    setTmForTest(jpaTm());
     persistNewRegistrar("TheRegistrar");
     persistNewRegistrar("NewRegistrar");
     Registrar registrar1 =
@@ -273,7 +275,6 @@ class Spec11PipelineTest {
     persistResource(createDomain("no-email.com", "2A4BA9BBC-COM", registrar2, contact2));
     persistResource(
         createDomain("anti-anti-anti-virus.dev", "555666888-DEV", registrar3, contact3));
-    removeTmOverrideForTest();
   }
 
   private void verifySaveToGcs() throws Exception {

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

@@ -696,7 +696,7 @@ public class DomainBaseTest extends EntityTestCase {
             IllegalArgumentException.class, () -> domain.asBuilder().setDomainName("AAA.BBB"));
     assertThat(thrown)
         .hasMessageThat()
-        .contains("Domain name must be in puny-coded, lower-case form");
+        .isEqualTo("Domain name AAA.BBB not in puny-coded, lower-case form");
   }
 
   @Test
@@ -706,7 +706,7 @@ public class DomainBaseTest extends EntityTestCase {
             IllegalArgumentException.class, () -> domain.asBuilder().setDomainName("みんな.みんな"));
     assertThat(thrown)
         .hasMessageThat()
-        .contains("Domain name must be in puny-coded, lower-case form");
+        .isEqualTo("Domain name みんな.みんな not in puny-coded, lower-case form");
   }
 
   @Test

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

@@ -200,7 +200,7 @@ class HostResourceTest extends EntityTestCase {
             IllegalArgumentException.class, () -> host.asBuilder().setHostName("AAA.BBB.CCC"));
     assertThat(thrown)
         .hasMessageThat()
-        .contains("Host name must be in puny-coded, lower-case form");
+        .isEqualTo("Host name AAA.BBB.CCC not in puny-coded, lower-case form");
   }
 
   @TestOfyAndSql
@@ -210,7 +210,7 @@ class HostResourceTest extends EntityTestCase {
             IllegalArgumentException.class, () -> host.asBuilder().setHostName("みんな.みんな.みんな"));
     assertThat(thrown)
         .hasMessageThat()
-        .contains("Host name must be in puny-coded, lower-case form");
+        .isEqualTo("Host name みんな.みんな.みんな not in puny-coded, lower-case form");
   }
 
   @TestOfyAndSql

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

@@ -16,7 +16,6 @@ package google.registry.model.ofy;
 
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.ofy.ObjectifyService.auditedOfy;
-import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatabaseHelper.createTld;
 
 import com.google.appengine.api.datastore.Entity;
@@ -66,7 +65,7 @@ public class CommitLogMutationTest {
     Entity rawEntity = convertToEntityInTxn(someObject);
     // Needs to be in a transaction so that registry-saving-to-entity will work.
     CommitLogMutation mutation =
-        tm().transact(() -> CommitLogMutation.create(manifestKey, someObject));
+        auditedOfy().transact(() -> CommitLogMutation.create(manifestKey, someObject));
     assertThat(Key.create(mutation))
         .isEqualTo(CommitLogMutation.createKey(manifestKey, Key.create(someObject)));
     assertThat(mutation.getEntity()).isEqualTo(rawEntity);
@@ -86,6 +85,6 @@ public class CommitLogMutationTest {
   }
 
   private static Entity convertToEntityInTxn(final ImmutableObject object) {
-    return tm().transact(() -> auditedOfy().save().toEntity(object));
+    return auditedOfy().transact(() -> auditedOfy().save().toEntity(object));
   }
 }

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

@@ -49,8 +49,7 @@ import java.util.logging.Level;
 import java.util.logging.Logger;
 import org.joda.time.DateTime;
 import org.joda.time.Duration;
-import org.junit.jupiter.api.AfterEach;
-import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 import org.junitpioneer.jupiter.RetryingTest;
@@ -74,7 +73,8 @@ public class ReplicateToDatastoreActionTest {
   private ReplicateToDatastoreAction action;
   private FakeResponse response;
 
-  @BeforeEach
+  // TODO(b/197534789): fix these tests and re-add the @BeforeEach
+  // @BeforeEach
   void setUp() {
     resetAction();
     injectExtension.setStaticField(Ofy.class, "clock", fakeClock);
@@ -88,13 +88,15 @@ public class ReplicateToDatastoreActionTest {
     TestObject.beforeDatastoreSaveCallCount = 0;
   }
 
-  @AfterEach
+  // TODO(b/197534789): fix these tests and re-add the @AfterEach
+  // @AfterEach
   void tearDown() {
     DatabaseHelper.removeDatabaseMigrationSchedule();
     fakeClock.disableAutoIncrement();
   }
 
   @RetryingTest(4)
+  @Disabled("b/197534789")
   void testReplication() {
     TestObject foo = TestObject.create("foo");
     TestObject bar = TestObject.create("bar");
@@ -120,6 +122,7 @@ public class ReplicateToDatastoreActionTest {
   }
 
   @RetryingTest(4)
+  @Disabled("b/197534789")
   void testReplayFromLastTxn() {
     TestObject foo = TestObject.create("foo");
     TestObject bar = TestObject.create("bar");
@@ -142,6 +145,7 @@ public class ReplicateToDatastoreActionTest {
   }
 
   @RetryingTest(4)
+  @Disabled("b/197534789")
   void testUnintentionalConcurrency() {
     TestObject foo = TestObject.create("foo");
     TestObject bar = TestObject.create("bar");
@@ -177,6 +181,7 @@ public class ReplicateToDatastoreActionTest {
   }
 
   @RetryingTest(4)
+  @Disabled("b/197534789")
   void testMissingTransactions() {
     // Write a transaction (should have a transaction id of 1).
     TestObject foo = TestObject.create("foo");
@@ -194,6 +199,7 @@ public class ReplicateToDatastoreActionTest {
   }
 
   @Test
+  @Disabled("b/197534789")
   void testMissingTransactions_fullTask() {
     // Write a transaction (should have a transaction id of 1).
     TestObject foo = TestObject.create("foo");
@@ -212,6 +218,7 @@ public class ReplicateToDatastoreActionTest {
   }
 
   @Test
+  @Disabled("b/197534789")
   void testBeforeDatastoreSaveCallback() {
     TestObject testObject = TestObject.create("foo");
     insertInDb(testObject);
@@ -221,6 +228,7 @@ public class ReplicateToDatastoreActionTest {
   }
 
   @Test
+  @Disabled("b/197534789")
   void testNotInMigrationState_doesNothing() {
     // set a schedule that backtracks the current status to DATASTORE_PRIMARY
     DateTime now = fakeClock.nowUtc();
@@ -257,6 +265,7 @@ public class ReplicateToDatastoreActionTest {
   }
 
   @Test
+  @Disabled("b/197534789")
   void testFailure_cannotAcquireLock() {
     RequestStatusChecker requestStatusChecker = mock(RequestStatusChecker.class);
     when(requestStatusChecker.getLogId()).thenReturn("logId");

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

@@ -36,6 +36,7 @@ import google.registry.persistence.VKey;
 import google.registry.persistence.transaction.JpaTestExtensions.JpaUnitTestExtension;
 import google.registry.testing.DatabaseHelper;
 import google.registry.testing.FakeClock;
+import google.registry.testing.TmOverrideExtension;
 import java.io.Serializable;
 import java.math.BigInteger;
 import java.sql.SQLException;
@@ -48,8 +49,7 @@ 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.Order;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
@@ -84,15 +84,9 @@ class JpaTransactionManagerImplTest {
               TestEntity.class, TestCompoundIdEntity.class, TestNamedCompoundIdEntity.class)
           .buildUnitTestExtension();
 
-  @BeforeEach
-  void beforeEach() {
-    TransactionManagerFactory.setTmForTest(jpaTm());
-  }
-
-  @AfterEach
-  void afterEach() {
-    TransactionManagerFactory.removeTmOverrideForTest();
-  }
+  @RegisterExtension
+  @Order(Order.DEFAULT + 1)
+  TmOverrideExtension tmOverrideExtension = TmOverrideExtension.withJpa();
 
   @Test
   void transact_succeeds() {

core/src/test/java/google/registry/rdap/RdapDomainActionTest.java

@@ -45,7 +45,12 @@ import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 
-/** Unit tests for {@link RdapDomainAction}. */
+/**
+ * Unit tests for {@link RdapDomainAction}.
+ *
+ * <p>TODO(b/26872828): The next time we do any work on RDAP, consider adding the APNIC RDAP
+ * conformance checker to the unit test suite.
+ */
 class RdapDomainActionTest extends RdapActionBaseTestCase<RdapDomainAction> {
 
   RdapDomainActionTest() {

core/src/test/java/google/registry/rde/BrdaCopyActionTest.java

@@ -35,26 +35,28 @@ import java.io.File;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.InputStreamReader;
+import java.util.Optional;
 import org.bouncycastle.openpgp.PGPKeyPair;
 import org.bouncycastle.openpgp.PGPPrivateKey;
 import org.bouncycastle.openpgp.PGPPublicKey;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
 
 /** Unit tests for {@link BrdaCopyAction}. */
 public class BrdaCopyActionTest {
 
   private static final ByteSource DEPOSIT_XML = RdeTestData.loadBytes("deposit_full.xml");
+  private static final String STAGE_FILENAME = "lol_2010-10-17_thin_S1_R0";
+  private static final BlobId RYDE_FILE =
+      BlobId.of("tub", String.format("%s.ryde", STAGE_FILENAME));
+  private static final BlobId SIG_FILE = BlobId.of("tub", String.format("%s.sig", STAGE_FILENAME));
 
-  private static final BlobId STAGE_FILE =
-      BlobId.of("keg", "lol_2010-10-17_thin_S1_R0.xml.ghostryde");
-  private static final BlobId STAGE_LENGTH_FILE =
-      BlobId.of("keg", "lol_2010-10-17_thin_S1_R0.xml.length");
-  private static final BlobId RYDE_FILE = BlobId.of("tub", "lol_2010-10-17_thin_S1_R0.ryde");
-  private static final BlobId SIG_FILE = BlobId.of("tub", "lol_2010-10-17_thin_S1_R0.sig");
+  private BlobId stageFile;
+  private BlobId stageLengthFile;
 
   @RegisterExtension
   public final BouncyCastleProviderExtension bouncy = new BouncyCastleProviderExtension();
@@ -87,6 +89,16 @@ public class BrdaCopyActionTest {
   private final GcsUtils gcsUtils = new GcsUtils(LocalStorageHelper.getOptions());
   private final BrdaCopyAction action = new BrdaCopyAction();
 
+  private void runAction(String prefix) throws IOException {
+    stageFile = BlobId.of("keg", String.format("%s%s.xml.ghostryde", prefix, STAGE_FILENAME));
+    stageLengthFile = BlobId.of("keg", String.format("%s%s.xml.length", prefix, STAGE_FILENAME));
+    byte[] xml = DEPOSIT_XML.read();
+    gcsUtils.createFromBytes(stageFile, Ghostryde.encode(xml, encryptKey));
+    gcsUtils.createFromBytes(stageLengthFile, Long.toString(xml.length).getBytes(UTF_8));
+    action.prefix = prefix.isEmpty() ? Optional.empty() : Optional.of(prefix);
+    action.run();
+  }
+
   @BeforeEach
   void beforeEach() throws Exception {
     action.gcsUtils = gcsUtils;
@@ -97,24 +109,22 @@ public class BrdaCopyActionTest {
     action.receiverKey = receiverKey;
     action.signingKey = signingKey;
     action.stagingDecryptionKey = decryptKey;
-
-    byte[] xml = DEPOSIT_XML.read();
-    gcsUtils.createFromBytes(STAGE_FILE, Ghostryde.encode(xml, encryptKey));
-    gcsUtils.createFromBytes(STAGE_LENGTH_FILE, Long.toString(xml.length).getBytes(UTF_8));
   }
 
-  @Test
-  void testRun() {
-    action.run();
-    assertThat(gcsUtils.existsAndNotEmpty(STAGE_FILE)).isTrue();
-    assertThat(gcsUtils.existsAndNotEmpty(RYDE_FILE)).isTrue();
+  @ParameterizedTest
+  @ValueSource(strings = {"", "job-name/"})
+  void testRun(String prefix) throws Exception {
+    runAction(prefix);
+    assertThat(gcsUtils.existsAndNotEmpty(stageFile)).isTrue();
+    assertThat(gcsUtils.existsAndNotEmpty(stageLengthFile)).isTrue();
     assertThat(gcsUtils.existsAndNotEmpty(SIG_FILE)).isTrue();
   }
 
-  @Test
-  void testRun_rydeFormat() throws Exception {
+  @ParameterizedTest
+  @ValueSource(strings = {"", "job-name/"})
+  void testRun_rydeFormat(String prefix) throws Exception {
     assumeTrue(hasCommand("gpg --version"));
-    action.run();
+    runAction(prefix);
 
     File rydeTmp = new File(gpg.getCwd(), "ryde");
     Files.write(gcsUtils.readBytesFrom(RYDE_FILE), rydeTmp);
@@ -158,10 +168,11 @@ public class BrdaCopyActionTest {
         .contains("ID 7F9084EE54E1EB0F");
   }
 
-  @Test
-  void testRun_rydeSignature() throws Exception {
+  @ParameterizedTest
+  @ValueSource(strings = {"", "job-name/"})
+  void testRun_rydeSignature(String prefix) throws Exception {
     assumeTrue(hasCommand("gpg --version"));
-    action.run();
+    runAction(prefix);
 
     File rydeTmp = new File(gpg.getCwd(), "ryde");
     File sigTmp = new File(gpg.getCwd(), "ryde.sig");

core/src/test/java/google/registry/rde/RdeReportActionTest.java

@@ -65,6 +65,7 @@ import google.registry.xml.XmlException;
 import java.io.ByteArrayInputStream;
 import java.net.SocketTimeoutException;
 import java.util.Map;
+import java.util.Optional;
 import org.bouncycastle.openpgp.PGPPublicKey;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
@@ -91,7 +92,8 @@ public class RdeReportActionTest {
   private final URLFetchService urlFetchService = mock(URLFetchService.class);
   private final ArgumentCaptor<HTTPRequest> request = ArgumentCaptor.forClass(HTTPRequest.class);
   private final HTTPResponse httpResponse = mock(HTTPResponse.class);
-
+  private final PGPPublicKey encryptKey =
+      new FakeKeyringModule().get().getRdeStagingEncryptionKey();
   private final GcsUtils gcsUtils = new GcsUtils(LocalStorageHelper.getOptions());
   private final BlobId reportFile =
       BlobId.of("tub", "test_2006-06-06_full_S1_R0-report.xml.ghostryde");
@@ -112,12 +114,12 @@ public class RdeReportActionTest {
     action.timeout = standardSeconds(30);
     action.stagingDecryptionKey = new FakeKeyringModule().get().getRdeStagingDecryptionKey();
     action.runner = runner;
+    action.prefix = Optional.empty();
     return action;
   }
 
   @BeforeEach
   void beforeEach() throws Exception {
-    PGPPublicKey encryptKey = new FakeKeyringModule().get().getRdeStagingEncryptionKey();
     createTld("test");
     persistResource(
         Cursor.create(RDE_REPORT, DateTime.parse("2006-06-06TZ"), Registry.get("test")));
@@ -148,6 +150,37 @@ public class RdeReportActionTest {
     assertThat(response.getContentType()).isEqualTo(PLAIN_TEXT_UTF_8);
     assertThat(response.getPayload()).isEqualTo("OK test 2006-06-06T00:00:00.000Z\n");
 
+    // Verify the HTTP request was correct.
+    assertThat(request.getValue().getMethod()).isSameInstanceAs(PUT);
+    assertThat(request.getValue().getURL().getProtocol()).isEqualTo("https");
+    assertThat(request.getValue().getURL().getPath()).endsWith("/test/20101017001");
+    Map<String, String> headers = mapifyHeaders(request.getValue().getHeaders());
+    assertThat(headers).containsEntry("CONTENT_TYPE", "text/xml");
+    assertThat(headers).containsEntry("AUTHORIZATION", "Basic dGVzdF9yeTpmb28=");
+
+    // Verify the payload XML was the same as what's in testdata/report.xml.
+    XjcRdeReportReport report = parseReport(request.getValue().getPayload());
+    assertThat(report.getId()).isEqualTo("20101017001");
+    assertThat(report.getCrDate()).isEqualTo(DateTime.parse("2010-10-17T00:15:00.0Z"));
+    assertThat(report.getWatermark()).isEqualTo(DateTime.parse("2010-10-17T00:00:00Z"));
+  }
+
+  @TestOfyAndSql
+  void testRunWithLock_withPrefix() throws Exception {
+    when(httpResponse.getResponseCode()).thenReturn(SC_OK);
+    when(httpResponse.getContent()).thenReturn(IIRDEA_GOOD_XML.read());
+    when(urlFetchService.fetch(request.capture())).thenReturn(httpResponse);
+    RdeReportAction action = createAction();
+    action.prefix = Optional.of("job-name/");
+    gcsUtils.delete(reportFile);
+    gcsUtils.createFromBytes(
+        BlobId.of("tub", "job-name/test_2006-06-06_full_S1_R0-report.xml.ghostryde"),
+        Ghostryde.encode(REPORT_XML.read(), encryptKey));
+    action.runWithLock(loadRdeReportCursor());
+    assertThat(response.getStatus()).isEqualTo(200);
+    assertThat(response.getContentType()).isEqualTo(PLAIN_TEXT_UTF_8);
+    assertThat(response.getPayload()).isEqualTo("OK test 2006-06-06T00:00:00.000Z\n");
+
     // Verify the HTTP request was correct.
     assertThat(request.getValue().getMethod()).isSameInstanceAs(PUT);
     assertThat(request.getValue().getURL().getProtocol()).isEqualTo("https");

core/src/test/java/google/registry/rde/RdeStagingActionCloudSqlTest.java

@@ -0,0 +1,263 @@
+// Copyright 2017 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.rde;
+
+import static com.google.common.truth.Truth.assertThat;
+import static google.registry.testing.DatabaseHelper.createTld;
+import static google.registry.testing.DatabaseHelper.persistResource;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+
+import com.google.api.services.dataflow.model.LaunchFlexTemplateRequest;
+import com.google.cloud.storage.contrib.nio.testing.LocalStorageHelper;
+import com.google.common.collect.ImmutableSet;
+import google.registry.beam.BeamActionTestBase;
+import google.registry.gcs.GcsUtils;
+import google.registry.model.tld.Registry;
+import google.registry.request.HttpException.BadRequestException;
+import google.registry.testing.AppEngineExtension;
+import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.FakeClock;
+import google.registry.testing.TestSqlOnly;
+import java.nio.charset.StandardCharsets;
+import java.util.Optional;
+import org.joda.time.DateTime;
+import org.joda.time.DateTimeConstants;
+import org.joda.time.Duration;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+/** Unit tests for {@link RdeStagingAction} in Cloud SQL. */
+@DualDatabaseTest
+public class RdeStagingActionCloudSqlTest extends BeamActionTestBase {
+
+  private final FakeClock clock = new FakeClock();
+  private final GcsUtils gcsUtils = new GcsUtils(LocalStorageHelper.getOptions());
+  private final RdeStagingAction action = new RdeStagingAction();
+
+  @RegisterExtension
+  public final AppEngineExtension extension =
+      AppEngineExtension.builder().withClock(clock).withDatastoreAndCloudSql().build();
+
+  @BeforeEach
+  @Override
+  public void beforeEach() throws Exception {
+    super.beforeEach();
+    action.clock = clock;
+    action.lenient = false;
+    action.projectId = "projectId";
+    action.jobRegion = "jobRegion";
+    action.rdeBucket = "rde-bucket";
+    action.pendingDepositChecker = new PendingDepositChecker();
+    action.pendingDepositChecker.brdaDayOfWeek = DateTimeConstants.TUESDAY;
+    action.pendingDepositChecker.brdaInterval = Duration.standardDays(7);
+    action.pendingDepositChecker.clock = clock;
+    action.pendingDepositChecker.rdeInterval = Duration.standardDays(1);
+    action.gcsUtils = gcsUtils;
+    action.response = response;
+    action.transactionCooldown = Duration.ZERO;
+    action.stagingKeyBytes = "ABCE".getBytes(StandardCharsets.UTF_8);
+    action.directory = Optional.empty();
+    action.modeStrings = ImmutableSet.of();
+    action.tlds = ImmutableSet.of();
+    action.watermarks = ImmutableSet.of();
+    action.revision = Optional.empty();
+    action.dataflow = dataflow;
+  }
+
+  @TestSqlOnly
+  void testRun_modeInNonManualMode_throwsException() {
+    createTldWithEscrowEnabled("lol");
+    clock.setTo(DateTime.parse("2000-01-01TZ"));
+    action.modeStrings = ImmutableSet.of("full");
+    assertThrows(BadRequestException.class, action::run);
+    verifyNoMoreInteractions(dataflow);
+  }
+
+  @TestSqlOnly
+  void testRun_tldInNonManualMode_throwsException() {
+    createTldWithEscrowEnabled("lol");
+    clock.setTo(DateTime.parse("2000-01-01TZ"));
+    action.tlds = ImmutableSet.of("tld");
+    assertThrows(BadRequestException.class, action::run);
+    verifyNoMoreInteractions(dataflow);
+  }
+
+  @TestSqlOnly
+  void testRun_watermarkInNonManualMode_throwsException() {
+    createTldWithEscrowEnabled("lol");
+    clock.setTo(DateTime.parse("2000-01-01TZ"));
+    action.watermarks = ImmutableSet.of(clock.nowUtc());
+    assertThrows(BadRequestException.class, action::run);
+    verifyNoMoreInteractions(dataflow);
+  }
+
+  @TestSqlOnly
+  void testRun_revisionInNonManualMode_throwsException() {
+    createTldWithEscrowEnabled("lol");
+    clock.setTo(DateTime.parse("2000-01-01TZ"));
+    action.revision = Optional.of(42);
+    assertThrows(BadRequestException.class, action::run);
+    verifyNoMoreInteractions(dataflow);
+  }
+
+  @TestSqlOnly
+  void testRun_noTlds_returns204() {
+    action.run();
+    assertThat(response.getStatus()).isEqualTo(204);
+    verifyNoMoreInteractions(dataflow);
+  }
+
+  @TestSqlOnly
+  void testRun_tldWithoutEscrowEnabled_returns204() {
+    createTld("lol");
+    persistResource(Registry.get("lol").asBuilder().setEscrowEnabled(false).build());
+    clock.setTo(DateTime.parse("2000-01-01TZ"));
+    action.run();
+    assertThat(response.getStatus()).isEqualTo(204);
+    verifyNoMoreInteractions(dataflow);
+  }
+
+  @TestSqlOnly
+  void testRun_tldWithEscrowEnabled_launchesPipeline() throws Exception {
+    createTldWithEscrowEnabled("lol");
+    clock.setTo(DateTime.parse("2000-01-01TZ"));
+    action.run();
+    assertThat(response.getStatus()).isEqualTo(200);
+    assertThat(response.getPayload()).contains("Launched RDE pipeline: jobid");
+    verify(templates, times(1))
+        .launch(eq("projectId"), eq("jobRegion"), any(LaunchFlexTemplateRequest.class));
+  }
+
+  @TestSqlOnly
+  void testRun_withinTransactionCooldown_getsExcludedAndReturns204() {
+    createTldWithEscrowEnabled("lol");
+    clock.setTo(DateTime.parse("2000-01-01T00:04:59Z"));
+    action.transactionCooldown = Duration.standardMinutes(5);
+    action.run();
+    assertThat(response.getStatus()).isEqualTo(204);
+    verifyNoMoreInteractions(dataflow);
+  }
+
+  @TestSqlOnly
+  void testRun_afterTransactionCooldown_runsMapReduce() throws Exception {
+    createTldWithEscrowEnabled("lol");
+    clock.setTo(DateTime.parse("2000-01-01T00:05:00Z"));
+    action.transactionCooldown = Duration.standardMinutes(5);
+    action.run();
+    assertThat(response.getStatus()).isEqualTo(200);
+    assertThat(response.getPayload()).contains("Launched RDE pipeline: jobid");
+    verify(templates, times(1))
+        .launch(eq("projectId"), eq("jobRegion"), any(LaunchFlexTemplateRequest.class));
+  }
+
+  @TestSqlOnly
+  void testManualRun_emptyMode_throwsException() {
+    createTldWithEscrowEnabled("lol");
+    clock.setTo(DateTime.parse("2000-01-01TZ"));
+    action.manual = true;
+    action.directory = Optional.of("test/");
+    action.modeStrings = ImmutableSet.of();
+    action.tlds = ImmutableSet.of("lol");
+    action.watermarks = ImmutableSet.of(clock.nowUtc());
+    assertThrows(BadRequestException.class, action::run);
+  }
+
+  @TestSqlOnly
+  void testManualRun_invalidMode_throwsException() {
+    createTldWithEscrowEnabled("lol");
+    clock.setTo(DateTime.parse("2000-01-01TZ"));
+    action.manual = true;
+    action.directory = Optional.of("test/");
+    action.modeStrings = ImmutableSet.of("full", "thing");
+    action.tlds = ImmutableSet.of("lol");
+    action.watermarks = ImmutableSet.of(clock.nowUtc());
+    assertThrows(BadRequestException.class, action::run);
+  }
+
+  @TestSqlOnly
+  void testManualRun_emptyTld_throwsException() {
+    createTldWithEscrowEnabled("lol");
+    clock.setTo(DateTime.parse("2000-01-01TZ"));
+    action.manual = true;
+    action.directory = Optional.of("test/");
+    action.modeStrings = ImmutableSet.of("full");
+    action.tlds = ImmutableSet.of();
+    action.watermarks = ImmutableSet.of(clock.nowUtc());
+    assertThrows(BadRequestException.class, action::run);
+  }
+
+  @TestSqlOnly
+  void testManualRun_emptyWatermark_throwsException() {
+    createTldWithEscrowEnabled("lol");
+    clock.setTo(DateTime.parse("2000-01-01TZ"));
+    action.manual = true;
+    action.directory = Optional.of("test/");
+    action.modeStrings = ImmutableSet.of("full");
+    action.tlds = ImmutableSet.of("lol");
+    action.watermarks = ImmutableSet.of();
+    assertThrows(BadRequestException.class, action::run);
+  }
+
+  @TestSqlOnly
+  void testManualRun_nonDayStartWatermark_throwsException() {
+    createTldWithEscrowEnabled("lol");
+    clock.setTo(DateTime.parse("2000-01-01TZ"));
+    action.manual = true;
+    action.directory = Optional.of("test/");
+    action.modeStrings = ImmutableSet.of("full");
+    action.tlds = ImmutableSet.of("lol");
+    action.watermarks = ImmutableSet.of(DateTime.parse("2001-01-01T01:36:45Z"));
+    assertThrows(BadRequestException.class, action::run);
+  }
+
+  @TestSqlOnly
+  void testManualRun_invalidRevision_throwsException() {
+    createTldWithEscrowEnabled("lol");
+    clock.setTo(DateTime.parse("2000-01-01TZ"));
+    action.manual = true;
+    action.directory = Optional.of("test/");
+    action.modeStrings = ImmutableSet.of("full");
+    action.tlds = ImmutableSet.of("lol");
+    action.watermarks = ImmutableSet.of(DateTime.parse("2001-01-01T00:00:00Z"));
+    action.revision = Optional.of(-1);
+    assertThrows(BadRequestException.class, action::run);
+  }
+
+  @TestSqlOnly
+  void testManualRun_validParameters_runsMapReduce() throws Exception {
+    createTldWithEscrowEnabled("lol");
+    clock.setTo(DateTime.parse("2000-01-01TZ"));
+    action.manual = true;
+    action.directory = Optional.of("test/");
+    action.modeStrings = ImmutableSet.of("full");
+    action.tlds = ImmutableSet.of("lol");
+    action.watermarks = ImmutableSet.of(DateTime.parse("2001-01-01TZ"));
+    action.run();
+    assertThat(response.getStatus()).isEqualTo(200);
+    assertThat(response.getPayload()).contains("Launched RDE pipeline: jobid");
+    verify(templates, times(1))
+        .launch(eq("projectId"), eq("jobRegion"), any(LaunchFlexTemplateRequest.class));
+  }
+
+  private static void createTldWithEscrowEnabled(final String tld) {
+    createTld(tld);
+    persistResource(Registry.get(tld).asBuilder().setEscrowEnabled(true).build());
+  }
+}

core/src/test/java/google/registry/rde/RdeStagingActionDatastoreTest.java

@@ -90,8 +90,8 @@ import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
-/** Unit tests for {@link RdeStagingAction}. */
-public class RdeStagingActionTest extends MapreduceTestCase<RdeStagingAction> {
+/** Unit tests for {@link RdeStagingAction} in Datastore. */
+public class RdeStagingActionDatastoreTest extends MapreduceTestCase<RdeStagingAction> {
 
   private static final BlobId XML_FILE =
       BlobId.of("rde-bucket", "lol_2000-01-01_full_S1_R0.xml.ghostryde");

core/src/test/java/google/registry/rde/RdeTestSuite.java

@@ -27,7 +27,8 @@ import org.junit.runner.RunWith;
   GhostrydeGpgIntegrationTest.class,
   GhostrydeTest.class,
   HostResourceToXjcConverterTest.class,
-  RdeStagingActionTest.class,
+  RdeStagingActionDatastoreTest.class,
+  RdeStagingActionCloudSqlTest.class,
   RdeUploadActionTest.class,
   RdeReportActionTest.class,
   RegistrarToXjcConverterTest.class,

core/src/test/java/google/registry/rde/RdeUploadActionTest.java

@@ -25,8 +25,6 @@ import static google.registry.testing.DatabaseHelper.createTld;
 import static google.registry.testing.DatabaseHelper.persistResource;
 import static google.registry.testing.DatabaseHelper.persistSimpleResource;
 import static google.registry.testing.SystemInfo.hasCommand;
-import static google.registry.testing.TaskQueueHelper.assertNoTasksEnqueued;
-import static google.registry.testing.TaskQueueHelper.assertTasksEnqueued;
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.joda.time.Duration.standardDays;
 import static org.joda.time.Duration.standardHours;
@@ -41,7 +39,6 @@ import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
 import static org.mockito.Mockito.when;
 
-import com.google.appengine.api.taskqueue.QueueFactory;
 import com.google.appengine.api.utils.SystemProperty;
 import com.google.cloud.storage.BlobId;
 import com.google.cloud.storage.contrib.nio.testing.LocalStorageHelper;
@@ -62,6 +59,8 @@ import google.registry.request.HttpException.NoContentException;
 import google.registry.request.RequestParameters;
 import google.registry.testing.AppEngineExtension;
 import google.registry.testing.BouncyCastleProviderExtension;
+import google.registry.testing.CloudTasksHelper;
+import google.registry.testing.CloudTasksHelper.TaskMatcher;
 import google.registry.testing.DualDatabaseTest;
 import google.registry.testing.FakeClock;
 import google.registry.testing.FakeKeyringModule;
@@ -69,17 +68,16 @@ import google.registry.testing.FakeResponse;
 import google.registry.testing.FakeSleeper;
 import google.registry.testing.GpgSystemCommandExtension;
 import google.registry.testing.Lazies;
-import google.registry.testing.TaskQueueHelper.TaskMatcher;
 import google.registry.testing.TestOfyAndSql;
 import google.registry.testing.sftp.SftpServerExtension;
 import google.registry.util.Retrier;
-import google.registry.util.TaskQueueUtils;
 import java.io.File;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.net.Socket;
 import java.net.URI;
+import java.util.Optional;
 import org.bouncycastle.openpgp.PGPPublicKey;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
@@ -100,6 +98,12 @@ public class RdeUploadActionTest {
       BlobId.of("bucket", "tld_2010-10-17_full_S1_R0.xml.length");
   private static final BlobId REPORT_FILE =
       BlobId.of("bucket", "tld_2010-10-17_full_S1_R0-report.xml.ghostryde");
+  private static final BlobId GHOSTRYDE_FILE_WITH_PREFIX =
+      BlobId.of("bucket", "job-name/tld_2010-10-17_full_S1_R0.xml.ghostryde");
+  private static final BlobId LENGTH_FILE_WITH_PREFIX =
+      BlobId.of("bucket", "job-name/tld_2010-10-17_full_S1_R0.xml.length");
+  private static final BlobId REPORT_FILE_WITH_PREFIX =
+      BlobId.of("bucket", "job-name/tld_2010-10-17_full_S1_R0-report.xml.ghostryde");
 
   private static final BlobId GHOSTRYDE_R1_FILE =
       BlobId.of("bucket", "tld_2010-10-17_full_S1_R1.xml.ghostryde");
@@ -109,6 +113,7 @@ public class RdeUploadActionTest {
       BlobId.of("bucket", "tld_2010-10-17_full_S1_R1-report.xml.ghostryde");
 
   private final GcsUtils gcsUtils = new GcsUtils(LocalStorageHelper.getOptions());
+  private final CloudTasksHelper cloudTasksHelper = new CloudTasksHelper();
 
   @RegisterExtension final SftpServerExtension sftpd = new SftpServerExtension();
 
@@ -129,6 +134,8 @@ public class RdeUploadActionTest {
   public final AppEngineExtension appEngine =
       AppEngineExtension.builder().withDatastoreAndCloudSql().withTaskQueue().build();
 
+  private final PGPPublicKey encryptKey =
+      new FakeKeyringModule().get().getRdeStagingEncryptionKey();
   private final FakeResponse response = new FakeResponse();
   private final EscrowTaskRunner runner = mock(EscrowTaskRunner.class);
   private final FakeClock clock = new FakeClock(DateTime.parse("2010-10-17TZ"));
@@ -155,10 +162,10 @@ public class RdeUploadActionTest {
       action.receiverKey = keyring.getRdeReceiverKey();
       action.signingKey = keyring.getRdeSigningKey();
       action.stagingDecryptionKey = keyring.getRdeStagingDecryptionKey();
-      action.reportQueue = QueueFactory.getQueue("rde-report");
       action.runner = runner;
-      action.taskQueueUtils = new TaskQueueUtils(new Retrier(null, 1));
+      action.cloudTasksUtils = cloudTasksHelper.getTestCloudTasksUtils();
       action.retrier = new Retrier(new FakeSleeper(clock), 3);
+      action.prefix = Optional.empty();
       return action;
     }
   }
@@ -181,15 +188,13 @@ public class RdeUploadActionTest {
     SystemProperty.environment.set(SystemProperty.Environment.Value.Development);
 
     createTld("tld");
-    PGPPublicKey encryptKey = new FakeKeyringModule().get().getRdeStagingEncryptionKey();
     gcsUtils.createFromBytes(GHOSTRYDE_FILE, Ghostryde.encode(DEPOSIT_XML.read(), encryptKey));
     gcsUtils.createFromBytes(GHOSTRYDE_R1_FILE, Ghostryde.encode(DEPOSIT_XML.read(), encryptKey));
     gcsUtils.createFromBytes(LENGTH_FILE, Long.toString(DEPOSIT_XML.size()).getBytes(UTF_8));
     gcsUtils.createFromBytes(LENGTH_R1_FILE, Long.toString(DEPOSIT_XML.size()).getBytes(UTF_8));
     gcsUtils.createFromBytes(REPORT_FILE, Ghostryde.encode(REPORT_XML.read(), encryptKey));
     gcsUtils.createFromBytes(REPORT_R1_FILE, Ghostryde.encode(REPORT_XML.read(), encryptKey));
-    tm()
-        .transact(
+    tm().transact(
             () -> {
               RdeRevision.saveRevision("lol", DateTime.parse("2010-10-17TZ"), FULL, 0);
               RdeRevision.saveRevision("tld", DateTime.parse("2010-10-17TZ"), FULL, 0);
@@ -210,11 +215,48 @@ public class RdeUploadActionTest {
     RdeUploadAction action = createAction(null);
     action.tld = "lol";
     action.run();
-    verify(runner).lockRunAndRollForward(
-        action, Registry.get("lol"), standardSeconds(23), CursorType.RDE_UPLOAD, standardDays(1));
-    assertTasksEnqueued("rde-report", new TaskMatcher()
-        .url(RdeReportAction.PATH)
-        .param(RequestParameters.PARAM_TLD, "lol"));
+    verify(runner)
+        .lockRunAndRollForward(
+            action,
+            Registry.get("lol"),
+            standardSeconds(23),
+            CursorType.RDE_UPLOAD,
+            standardDays(1));
+    cloudTasksHelper.assertTasksEnqueued(
+        "rde-report",
+        new TaskMatcher().url(RdeReportAction.PATH).param(RequestParameters.PARAM_TLD, "lol"));
+    verifyNoMoreInteractions(runner);
+  }
+
+  @TestOfyAndSql
+  void testRun_withPrefix() throws Exception {
+    createTld("lol");
+    RdeUploadAction action = createAction(null);
+    action.prefix = Optional.of("job-name/");
+    action.tld = "lol";
+    gcsUtils.delete(GHOSTRYDE_FILE);
+    gcsUtils.createFromBytes(
+        GHOSTRYDE_FILE_WITH_PREFIX, Ghostryde.encode(DEPOSIT_XML.read(), encryptKey));
+    gcsUtils.delete(LENGTH_FILE);
+    gcsUtils.createFromBytes(
+        LENGTH_FILE_WITH_PREFIX, Long.toString(DEPOSIT_XML.size()).getBytes(UTF_8));
+    gcsUtils.delete(REPORT_FILE);
+    gcsUtils.createFromBytes(
+        REPORT_FILE_WITH_PREFIX, Ghostryde.encode(REPORT_XML.read(), encryptKey));
+    action.run();
+    verify(runner)
+        .lockRunAndRollForward(
+            action,
+            Registry.get("lol"),
+            standardSeconds(23),
+            CursorType.RDE_UPLOAD,
+            standardDays(1));
+    cloudTasksHelper.assertTasksEnqueued(
+        "rde-report",
+        new TaskMatcher()
+            .url(RdeReportAction.PATH)
+            .param(RequestParameters.PARAM_TLD, "lol")
+            .param(RdeModule.PARAM_PREFIX, "job-name/"));
     verifyNoMoreInteractions(runner);
   }
 
@@ -231,7 +273,7 @@ public class RdeUploadActionTest {
     assertThat(response.getStatus()).isEqualTo(200);
     assertThat(response.getContentType()).isEqualTo(PLAIN_TEXT_UTF_8);
     assertThat(response.getPayload()).isEqualTo("OK tld 2010-10-17T00:00:00.000Z\n");
-    assertNoTasksEnqueued("rde-upload");
+    cloudTasksHelper.assertNoTasksEnqueued("rde-upload");
     assertThat(folder.list())
         .asList()
         .containsExactly("tld_2010-10-17_full_S1_R0.ryde", "tld_2010-10-17_full_S1_R0.sig");
@@ -262,7 +304,7 @@ public class RdeUploadActionTest {
     assertThat(response.getStatus()).isEqualTo(200);
     assertThat(response.getContentType()).isEqualTo(PLAIN_TEXT_UTF_8);
     assertThat(response.getPayload()).isEqualTo("OK tld 2010-10-17T00:00:00.000Z\n");
-    assertNoTasksEnqueued("rde-upload");
+    cloudTasksHelper.assertNoTasksEnqueued("rde-upload");
     // Assert that both files are written to SFTP and GCS, and that the contents are identical.
     String rydeFilename = "tld_2010-10-17_full_S1_R0.ryde";
     String sigFilename = "tld_2010-10-17_full_S1_R0.sig";
@@ -273,6 +315,41 @@ public class RdeUploadActionTest {
         .isEqualTo(Files.toByteArray(new File(folder, sigFilename)));
   }
 
+  @TestOfyAndSql
+  void testRunWithLock_copiesOnGcs_withPrefix() throws Exception {
+    int port = sftpd.serve("user", "password", folder);
+    URI uploadUrl = URI.create(String.format("sftp://user:password@localhost:%d/", port));
+    DateTime stagingCursor = DateTime.parse("2010-10-18TZ");
+    DateTime uploadCursor = DateTime.parse("2010-10-17TZ");
+    persistResource(Cursor.create(RDE_STAGING, stagingCursor, Registry.get("tld")));
+    RdeUploadAction action = createAction(uploadUrl);
+    action.prefix = Optional.of("job-name/");
+    gcsUtils.delete(GHOSTRYDE_FILE);
+    gcsUtils.createFromBytes(
+        GHOSTRYDE_FILE_WITH_PREFIX, Ghostryde.encode(DEPOSIT_XML.read(), encryptKey));
+    gcsUtils.delete(LENGTH_FILE);
+    gcsUtils.createFromBytes(
+        LENGTH_FILE_WITH_PREFIX, Long.toString(DEPOSIT_XML.size()).getBytes(UTF_8));
+    gcsUtils.delete(REPORT_FILE);
+    gcsUtils.createFromBytes(
+        REPORT_FILE_WITH_PREFIX, Ghostryde.encode(REPORT_XML.read(), encryptKey));
+    action.runWithLock(uploadCursor);
+    assertThat(response.getStatus()).isEqualTo(200);
+    assertThat(response.getContentType()).isEqualTo(PLAIN_TEXT_UTF_8);
+    assertThat(response.getPayload()).isEqualTo("OK tld 2010-10-17T00:00:00.000Z\n");
+    cloudTasksHelper.assertNoTasksEnqueued("rde-upload");
+    // Assert that both files are written to SFTP and GCS, and that the contents are identical.
+    String rydeFilename = "tld_2010-10-17_full_S1_R0.ryde";
+    String rydeGcsFilename = "job-name/tld_2010-10-17_full_S1_R0.ryde";
+    String sigFilename = "tld_2010-10-17_full_S1_R0.sig";
+    String sigGcsFilename = "job-name/tld_2010-10-17_full_S1_R0.sig";
+    assertThat(folder.list()).asList().containsExactly(rydeFilename, sigFilename);
+    assertThat(gcsUtils.readBytesFrom(BlobId.of("bucket", rydeGcsFilename)))
+        .isEqualTo(Files.toByteArray(new File(folder, rydeFilename)));
+    assertThat(gcsUtils.readBytesFrom(BlobId.of("bucket", sigGcsFilename)))
+        .isEqualTo(Files.toByteArray(new File(folder, sigFilename)));
+  }
+
   @TestOfyAndSql
   void testRunWithLock_resend() throws Exception {
     tm().transact(() -> RdeRevision.saveRevision("tld", DateTime.parse("2010-10-17TZ"), FULL, 1));
@@ -285,7 +362,7 @@ public class RdeUploadActionTest {
     assertThat(response.getStatus()).isEqualTo(200);
     assertThat(response.getContentType()).isEqualTo(PLAIN_TEXT_UTF_8);
     assertThat(response.getPayload()).isEqualTo("OK tld 2010-10-17T00:00:00.000Z\n");
-    assertNoTasksEnqueued("rde-upload");
+    cloudTasksHelper.assertNoTasksEnqueued("rde-upload");
     assertThat(folder.list())
         .asList()
         .containsExactly("tld_2010-10-17_full_S1_R1.ryde", "tld_2010-10-17_full_S1_R1.sig");
@@ -327,7 +404,7 @@ public class RdeUploadActionTest {
         .isEqualTo(
             "Waiting on RdeStagingAction for TLD tld to send 2010-10-17T00:00:00.000Z upload; last"
                 + " RDE staging completion was at 1970-01-01T00:00:00.000Z");
-    assertNoTasksEnqueued("rde-upload");
+    cloudTasksHelper.assertNoTasksEnqueued("rde-upload");
     assertThat(folder.list()).isEmpty();
   }
 

core/src/test/java/google/registry/reporting/icann/ActivityReportingQueryBuilderTest.java

@@ -39,18 +39,17 @@ class ActivityReportingQueryBuilderTest {
 
   private final YearMonth yearMonth = new YearMonth(2017, 9);
 
-  private ActivityReportingQueryBuilder getQueryBuilder() {
-    ActivityReportingQueryBuilder queryBuilder = new ActivityReportingQueryBuilder();
-    queryBuilder.projectId = "domain-registry-alpha";
-    queryBuilder.dnsCountQueryCoordinator =
+  private ActivityReportingQueryBuilder createQueryBuilder(String datasetName) {
+    return new ActivityReportingQueryBuilder(
+        "domain-registry-alpha",
+        datasetName,
         new BasicDnsCountQueryCoordinator(
-            new BasicDnsCountQueryCoordinator.Params(null, queryBuilder.projectId));
-    return queryBuilder;
+            new BasicDnsCountQueryCoordinator.Params(null, "domain-registry-alpha")));
   }
 
   @TestOfyOnly
   void testAggregateQueryMatch_datastore() {
-    ActivityReportingQueryBuilder queryBuilder = getQueryBuilder();
+    ActivityReportingQueryBuilder queryBuilder = createQueryBuilder("icann_reporting");
     assertThat(queryBuilder.getReportQuery(yearMonth))
         .isEqualTo(
             "#standardSQL\nSELECT * FROM "
@@ -58,8 +57,8 @@ class ActivityReportingQueryBuilderTest {
   }
 
   @TestSqlOnly
-  void testAggregateQueryMatch_cloud_sql() {
-    ActivityReportingQueryBuilder queryBuilder = getQueryBuilder();
+  void testAggregateQueryMatch_cloudSql() {
+    ActivityReportingQueryBuilder queryBuilder = createQueryBuilder("cloud_sql_icann_reporting");
     assertThat(queryBuilder.getReportQuery(yearMonth))
         .isEqualTo(
             "#standardSQL\n"
@@ -78,7 +77,7 @@ class ActivityReportingQueryBuilderTest {
             ActivityReportingQueryBuilder.WHOIS_COUNTS,
             ActivityReportingQueryBuilder.ACTIVITY_REPORT_AGGREGATION);
 
-    ActivityReportingQueryBuilder queryBuilder = getQueryBuilder();
+    ActivityReportingQueryBuilder queryBuilder = createQueryBuilder("icann_reporting");
     ImmutableMap<String, String> actualQueries = queryBuilder.getViewQueryMap(yearMonth);
     for (String queryName : expectedQueryNames) {
       String actualTableName = String.format("%s_201709", queryName);
@@ -99,7 +98,7 @@ class ActivityReportingQueryBuilderTest {
             ActivityReportingQueryBuilder.WHOIS_COUNTS,
             ActivityReportingQueryBuilder.ACTIVITY_REPORT_AGGREGATION);
 
-    ActivityReportingQueryBuilder queryBuilder = getQueryBuilder();
+    ActivityReportingQueryBuilder queryBuilder = createQueryBuilder("cloud_sql_icann_reporting");
     ImmutableMap<String, String> actualQueries = queryBuilder.getViewQueryMap(yearMonth);
     for (String queryName : expectedQueryNames) {
       String actualTableName = String.format("%s_201709", queryName);

core/src/test/java/google/registry/reporting/icann/IcannReportingModuleTest.java

@@ -18,34 +18,55 @@ import static com.google.common.truth.Truth.assertThat;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+import google.registry.persistence.transaction.JpaTestExtensions;
+import google.registry.persistence.transaction.JpaTestExtensions.JpaUnitTestExtension;
 import javax.servlet.http.HttpServletRequest;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Unit tests for {@link google.registry.reporting.icann.IcannReportingModule}. */
 class IcannReportingModuleTest {
 
-  @Test
-  void testProvideReportTypes() {
-    HttpServletRequest req = mock(HttpServletRequest.class);
+  @RegisterExtension
+  JpaUnitTestExtension jpaExtension = new JpaTestExtensions.Builder().buildUnitTestExtension();
 
+  @Test
+  void testProvideReportTypes_null() {
+    HttpServletRequest req = mock(HttpServletRequest.class);
     when(req.getParameter("reportTypes")).thenReturn(null);
     assertThat(IcannReportingModule.provideReportTypes(req))
         .containsExactly(
             IcannReportingModule.ReportType.ACTIVITY, IcannReportingModule.ReportType.TRANSACTIONS);
+  }
 
+  @Test
+  void testProvideReportTypes_empty() {
+    HttpServletRequest req = mock(HttpServletRequest.class);
     when(req.getParameter("reportTypes")).thenReturn("");
     assertThat(IcannReportingModule.provideReportTypes(req))
         .containsExactly(
             IcannReportingModule.ReportType.ACTIVITY, IcannReportingModule.ReportType.TRANSACTIONS);
+  }
 
+  @Test
+  void testProvideReportTypes_activity() {
+    HttpServletRequest req = mock(HttpServletRequest.class);
     when(req.getParameter("reportTypes")).thenReturn("activity");
     assertThat(IcannReportingModule.provideReportTypes(req))
         .containsExactly(IcannReportingModule.ReportType.ACTIVITY);
+  }
 
+  @Test
+  void testProvideReportTypes_transactions() {
+    HttpServletRequest req = mock(HttpServletRequest.class);
     when(req.getParameter("reportTypes")).thenReturn("transactions");
     assertThat(IcannReportingModule.provideReportTypes(req))
         .containsExactly(IcannReportingModule.ReportType.TRANSACTIONS);
+  }
 
+  @Test
+  void testProvideReportTypes_bothTypes() {
+    HttpServletRequest req = mock(HttpServletRequest.class);
     when(req.getParameter("reportTypes")).thenReturn("activity,transactions");
     assertThat(IcannReportingModule.provideReportTypes(req))
         .containsExactly(

core/src/test/java/google/registry/reporting/icann/IcannReportingStagerTest.java

@@ -54,13 +54,12 @@ class IcannReportingStagerTest {
 
   private IcannReportingStager createStager() {
     IcannReportingStager action = new IcannReportingStager();
-    ActivityReportingQueryBuilder activityBuilder = new ActivityReportingQueryBuilder();
-    activityBuilder.projectId = "test-project";
-    activityBuilder.dnsCountQueryCoordinator = new BasicDnsCountQueryCoordinator(null);
+    ActivityReportingQueryBuilder activityBuilder =
+        new ActivityReportingQueryBuilder(
+            "test-project", "icann_reporting", new BasicDnsCountQueryCoordinator(null));
     action.activityQueryBuilder = activityBuilder;
-    TransactionsReportingQueryBuilder transactionsBuilder = new TransactionsReportingQueryBuilder();
-    transactionsBuilder.projectId = "test-project";
-    action.transactionsQueryBuilder = transactionsBuilder;
+    action.transactionsQueryBuilder =
+        new TransactionsReportingQueryBuilder("test-project", "icann_reporting");
     action.reportingBucket = "test-bucket";
     action.bigquery = bigquery;
     action.gcsUtils = gcsUtils;

core/src/test/java/google/registry/reporting/icann/TransactionsReportingQueryBuilderTest.java

@@ -27,8 +27,8 @@ class TransactionsReportingQueryBuilderTest {
   private final YearMonth yearMonth = new YearMonth(2017, 9);
 
   private TransactionsReportingQueryBuilder getQueryBuilder() {
-    TransactionsReportingQueryBuilder queryBuilder = new TransactionsReportingQueryBuilder();
-    queryBuilder.projectId = "domain-registry-alpha";
+    TransactionsReportingQueryBuilder queryBuilder =
+        new TransactionsReportingQueryBuilder("domain-registry-alpha", "icann_reporting");
     return queryBuilder;
   }
 

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

@@ -16,8 +16,6 @@ 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;
@@ -28,6 +26,7 @@ import google.registry.model.registrar.RegistrarContact;
 import google.registry.persistence.transaction.JpaTestExtensions;
 import google.registry.persistence.transaction.JpaTestExtensions.JpaIntegrationWithCoverageExtension;
 import google.registry.testing.DatastoreEntityExtension;
+import google.registry.testing.TmOverrideExtension;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Order;
 import org.junit.jupiter.api.Test;
@@ -44,13 +43,16 @@ class RegistrarContactTest {
   JpaIntegrationWithCoverageExtension jpa =
       new JpaTestExtensions.Builder().buildIntegrationWithCoverageExtension();
 
+  @RegisterExtension
+  @Order(Order.DEFAULT + 1)
+  TmOverrideExtension tmOverrideExtension = TmOverrideExtension.withJpa();
+
   private Registrar testRegistrar;
 
   private RegistrarContact testRegistrarPoc;
 
   @BeforeEach
   public void beforeEach() {
-    setTmForTest(jpaTm());
     testRegistrar = saveRegistrar("registrarId");
     testRegistrarPoc =
         new RegistrarContact.Builder()

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

@@ -15,7 +15,6 @@
 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;
@@ -29,11 +28,10 @@ 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 google.registry.testing.TmOverrideExtension;
 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,13 +50,16 @@ public class RegistrarDaoTest {
   JpaIntegrationWithCoverageExtension jpa =
       new JpaTestExtensions.Builder().withClock(fakeClock).buildIntegrationWithCoverageExtension();
 
+  @RegisterExtension
+  @Order(Order.DEFAULT + 1)
+  TmOverrideExtension tmOverrideExtension = TmOverrideExtension.withJpa();
+
   private final VKey<Registrar> registrarKey = VKey.createSql(Registrar.class, "registrarId");
 
   private Registrar testRegistrar;
 
   @BeforeEach
   void beforeEach() {
-    TransactionManagerFactory.setTmForTest(jpaTm());
     testRegistrar =
         new Registrar.Builder()
             .setType(Registrar.Type.TEST)
@@ -75,11 +76,6 @@ public class RegistrarDaoTest {
             .build();
   }
 
-  @AfterEach
-  void afterEach() {
-    TransactionManagerFactory.removeTmOverrideForTest();
-  }
-
   @Test
   void saveNew_worksSuccessfully() {
     assertThat(existsInDb(testRegistrar)).isFalse();

core/src/test/java/google/registry/schema/replay/SqlEntityTest.java

@@ -21,10 +21,9 @@ import google.registry.model.registrar.Registrar;
 import google.registry.model.registrar.RegistrarContact;
 import google.registry.model.registrar.RegistrarContact.RegistrarPocId;
 import google.registry.persistence.VKey;
-import google.registry.persistence.transaction.TransactionManagerFactory;
 import google.registry.testing.AppEngineExtension;
 import google.registry.testing.DatastoreEntityExtension;
-import org.junit.jupiter.api.AfterEach;
+import google.registry.testing.TmOverrideExtension;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Order;
 import org.junit.jupiter.api.Test;
@@ -41,17 +40,15 @@ public class SqlEntityTest {
   final AppEngineExtension database =
       new AppEngineExtension.Builder().withCloudSql().withoutCannedData().build();
 
+  @RegisterExtension
+  @Order(Order.DEFAULT + 1)
+  TmOverrideExtension tmOverrideExtension = TmOverrideExtension.withJpa();
+
   @BeforeEach
   void setup() throws Exception {
-    TransactionManagerFactory.setTmForTest(TransactionManagerFactory.jpaTm());
     AppEngineExtension.loadInitialData();
   }
 
-  @AfterEach
-  void teardown() {
-    TransactionManagerFactory.removeTmOverrideForTest();
-  }
-
   @Test
   void getPrimaryKeyString_oneIdColumn() {
     // AppEngineExtension canned data: Registrar1

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

@@ -23,13 +23,8 @@ import static com.google.common.truth.Truth.assertWithMessage;
 import static google.registry.util.DiffUtils.prettyPrintEntityDeepDiff;
 import static java.util.Arrays.asList;
 import static java.util.stream.Collectors.joining;
-import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
 
-import com.google.cloud.tasks.v2.CloudTasksClient;
 import com.google.cloud.tasks.v2.HttpMethod;
-import com.google.cloud.tasks.v2.QueueName;
 import com.google.cloud.tasks.v2.Task;
 import com.google.common.base.Ascii;
 import com.google.common.base.Joiner;
@@ -37,16 +32,20 @@ import com.google.common.collect.ArrayListMultimap;
 import com.google.common.collect.ImmutableMultimap;
 import com.google.common.collect.ImmutableMultiset;
 import com.google.common.collect.LinkedListMultimap;
+import com.google.common.collect.ListMultimap;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Multimap;
+import com.google.common.collect.Multimaps;
 import com.google.common.net.HttpHeaders;
 import com.google.common.net.MediaType;
 import com.google.common.truth.Truth8;
 import google.registry.model.ImmutableObject;
 import google.registry.util.CloudTasksUtils;
 import google.registry.util.Retrier;
+import java.io.Serializable;
 import java.net.URI;
 import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
@@ -54,6 +53,9 @@ import java.util.List;
 import java.util.Map;
 import java.util.NoSuchElementException;
 import java.util.Objects;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Function;
 import java.util.function.Predicate;
 import javax.annotation.Nonnull;
@@ -65,30 +67,37 @@ import javax.annotation.Nonnull;
  * helper methods because we have not yet encountered all the use cases with Cloud Tasks. As more
  * and more Task Queue API usage is migrated to Cloud Tasks we may replicate more methods from the
  * latter.
+ *
+ * <p>Note the use of {@link AtomicInteger} {@code nextInstanceId} here. When a {@link
+ * FakeCloudTasksClient} instance, and by extension the {@link CloudTasksHelper} instance that
+ * contains it is serialized/deserialized, as happens in a Beam pipeline, we to want to push tasks
+ * to the same test task container that the original instance pushes to, so that we can make
+ * assertions on them by accessing the original instance. We cannot make the test task container
+ * itself static because we do not want tasks enqueued in previous tests to interfere with latter
+ * tests, when they run on the same JVM (and therefore share the same static class members). To
+ * solve this we put the test container in a static map whose keys are the instance IDs. An
+ * explicitly created new {@link CloudTasksHelper} (as would be created for a new test method) would
+ * have a new ID allocated to it, and therefore stores its tasks in a distinct container. A
+ * deserialized {@link CloudTasksHelper}, on the other hand, will have the same instance ID and
+ * share the same test class container with its progenitor.
  */
-public class CloudTasksHelper {
+public class CloudTasksHelper implements Serializable {
+
+  private static final long serialVersionUID = -8949359648199614677L;
+  private static final AtomicInteger nextInstanceId = new AtomicInteger(0);
+  protected static ConcurrentMap<Integer, ListMultimap<String, Task>> testTasks =
+      new ConcurrentHashMap<>();
 
   private static final String PROJECT_ID = "test-project";
   private static final String LOCATION_ID = "test-location";
 
   private final Retrier retrier = new Retrier(new FakeSleeper(new FakeClock()), 1);
-  private final LinkedListMultimap<String, Task> testTasks = LinkedListMultimap.create();
-  private final CloudTasksClient mockClient = mock(CloudTasksClient.class);
+  private final int instanceId = nextInstanceId.getAndIncrement();
   private final CloudTasksUtils cloudTasksUtils =
-      new CloudTasksUtils(retrier, PROJECT_ID, LOCATION_ID, () -> mockClient);
+      new CloudTasksUtils(retrier, PROJECT_ID, LOCATION_ID, new FakeCloudTasksClient());
 
   public CloudTasksHelper() {
-    when(mockClient.createTask(any(QueueName.class), any(Task.class)))
-        .thenAnswer(
-            invocation -> {
-              QueueName queue = invocation.getArgument(0);
-              Task task = invocation.getArgument(1);
-              if (task.getName().isEmpty()) {
-                task = task.toBuilder().setName(String.format("test-%d", testTasks.size())).build();
-              }
-              testTasks.put(queue.getQueue(), task);
-              return task;
-            });
+    testTasks.put(instanceId, Multimaps.synchronizedListMultimap(LinkedListMultimap.create()));
   }
 
   public CloudTasksUtils getTestCloudTasksUtils() {
@@ -96,7 +105,7 @@ public class CloudTasksHelper {
   }
 
   public List<Task> getTestTasksFor(String queue) {
-    return testTasks.get(queue);
+    return new ArrayList<>(testTasks.get(instanceId).get(queue));
   }
 
   /**
@@ -160,6 +169,20 @@ public class CloudTasksHelper {
     }
   }
 
+  private class FakeCloudTasksClient extends CloudTasksUtils.SerializableCloudTasksClient {
+
+    private static final long serialVersionUID = 6661964844791720639L;
+
+    @Override
+    public Task enqueue(String projectId, String locationId, String queueName, Task task) {
+      if (task.getName().isEmpty()) {
+        task = task.toBuilder().setName(String.format("test-%d", testTasks.size())).build();
+      }
+      testTasks.get(instanceId).put(queueName, task);
+      return task;
+    }
+  }
+
   /** An adapter to clean up a {@link Task} for ease of matching. */
   private static class MatchableTask extends ImmutableObject {
 

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

@@ -154,7 +154,7 @@ class DualDatabaseTestInvocationContextProvider implements TestTemplateInvocatio
       context.getStore(NAMESPACE).put(ORIGINAL_TM_KEY, tm());
       DatabaseType databaseType =
           (DatabaseType) context.getStore(NAMESPACE).get(INJECTED_TM_SUPPLIER_KEY);
-      TransactionManagerFactory.setTmForTest(databaseType.getTm());
+      TransactionManagerFactory.setTmOverrideForTest(databaseType.getTm());
     }
   }
 

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

@@ -38,7 +38,6 @@ import google.registry.persistence.transaction.Transaction.Delete;
 import google.registry.persistence.transaction.Transaction.Mutation;
 import google.registry.persistence.transaction.Transaction.Update;
 import google.registry.persistence.transaction.TransactionEntity;
-import google.registry.util.RequestStatusChecker;
 import java.io.IOException;
 import java.util.List;
 import java.util.Optional;
@@ -46,7 +45,6 @@ import javax.annotation.Nullable;
 import org.junit.jupiter.api.extension.AfterEachCallback;
 import org.junit.jupiter.api.extension.BeforeEachCallback;
 import org.junit.jupiter.api.extension.ExtensionContext;
-import org.mockito.Mockito;
 
 /**
  * A JUnit extension that replays datastore transactions against postgresql.
@@ -83,11 +81,13 @@ public class ReplayExtension implements BeforeEachCallback, AfterEachCallback {
    * Create a replay extension that replays from SQL to cloud datastore when running in SQL mode.
    */
   public static ReplayExtension createWithDoubleReplay(FakeClock clock) {
-    return new ReplayExtension(
-        clock,
-        true,
-        new ReplicateToDatastoreAction(
-            clock, Mockito.mock(RequestStatusChecker.class), new FakeResponse()));
+    // TODO: use the proper double-replay extension when the tests are not flaky
+    // return new ReplayExtension(
+    //     clock,
+    //     true,
+    //     new ReplicateToDatastoreAction(
+    //         clock, Mockito.mock(RequestStatusChecker.class), new FakeResponse()));
+    return createWithCompare(clock);
   }
 
   @Override

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

@@ -0,0 +1,73 @@
+// 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.testing;
+
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
+
+import google.registry.persistence.transaction.TransactionManager;
+import google.registry.persistence.transaction.TransactionManagerFactory;
+import org.junit.jupiter.api.extension.AfterEachCallback;
+import org.junit.jupiter.api.extension.BeforeEachCallback;
+import org.junit.jupiter.api.extension.ExtensionContext;
+
+/**
+ * JUnit extension for overriding the {@link TransactionManager} in tests.
+ *
+ * <p>You will typically want to run this at <code>@Order(Order.DEFAULT + 1)</code> alongside a
+ * {@link google.registry.persistence.transaction.JpaTransactionManagerExtension} or {@link
+ * DatastoreEntityExtension} with default {@link org.junit.jupiter.api.Order}. The transaction
+ * manager extension needs to run first so that when this override is called it's not trying to use
+ * the default dummy one.
+ *
+ * <p>This extension is incompatible with {@link DualDatabaseTest}. Use either that or this, but not
+ * both.
+ */
+public final class TmOverrideExtension implements BeforeEachCallback, AfterEachCallback {
+
+  private static enum TmOverride {
+    OFY,
+    JPA;
+  }
+
+  private final TmOverride tmOverride;
+
+  private TmOverrideExtension(TmOverride tmOverride) {
+    this.tmOverride = tmOverride;
+  }
+
+  /** Use the {@link google.registry.model.ofy.DatastoreTransactionManager} for all tests. */
+  public static TmOverrideExtension withOfy() {
+    return new TmOverrideExtension(TmOverride.OFY);
+  }
+
+  /**
+   * Use the {@link google.registry.persistence.transaction.JpaTransactionManager} for all tests.
+   */
+  public static TmOverrideExtension withJpa() {
+    return new TmOverrideExtension(TmOverride.JPA);
+  }
+
+  @Override
+  public void beforeEach(ExtensionContext context) {
+    TransactionManagerFactory.setTmOverrideForTest(
+        tmOverride == TmOverride.OFY ? ofyTm() : jpaTm());
+  }
+
+  @Override
+  public void afterEach(ExtensionContext context) {
+    TransactionManagerFactory.removeTmOverrideForTest();
+  }
+}

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

@@ -23,12 +23,15 @@ import static google.registry.testing.DatabaseHelper.persistResource;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
 import com.beust.jcommander.ParameterException;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.secdns.DelegationSignerData;
 import google.registry.model.eppcommon.StatusValue;
 import google.registry.model.host.HostResource;
 import google.registry.persistence.VKey;
 import javax.xml.bind.annotation.adapters.HexBinaryAdapter;
+import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
@@ -40,6 +43,8 @@ class UniformRapidSuspensionCommandTest
   private HostResource ns2;
   private HostResource urs1;
   private HostResource urs2;
+  private DomainBase defaultDomainBase;
+  private ImmutableSet<DelegationSignerData> defaultDsData;
 
   @BeforeEach
   void beforeEach() {
@@ -50,32 +55,36 @@ class UniformRapidSuspensionCommandTest
     ns2 = persistActiveHost("ns2.example.com");
     urs1 = persistActiveHost("urs1.example.com");
     urs2 = persistActiveHost("urs2.example.com");
+    defaultDomainBase = newDomainBase("evil.tld");
+    defaultDsData =
+        ImmutableSet.of(
+            DelegationSignerData.create(1, 2, 3, new HexBinaryAdapter().unmarshal("dead")),
+            DelegationSignerData.create(4, 5, 6, new HexBinaryAdapter().unmarshal("beef")));
   }
 
-  private void persistDomainWithHosts(HostResource... hosts) {
+  private void persistDomainWithHosts(
+      DomainBase domainBase, ImmutableSet<DelegationSignerData> dsData, HostResource... hosts) {
     ImmutableSet.Builder<VKey<HostResource>> hostRefs = new ImmutableSet.Builder<>();
     for (HostResource host : hosts) {
       hostRefs.add(host.createVKey());
     }
-    persistResource(newDomainBase("evil.tld").asBuilder()
-        .setNameservers(hostRefs.build())
-        .setDsData(ImmutableSet.of(
-            DelegationSignerData.create(1, 2, 3, new HexBinaryAdapter().unmarshal("dead")),
-            DelegationSignerData.create(4, 5, 6, new HexBinaryAdapter().unmarshal("beef"))))
-        .build());
+    persistResource(
+        domainBase.asBuilder().setNameservers(hostRefs.build()).setDsData(dsData).build());
   }
 
   @Test
   void testCommand_addsLocksReplacesHostsAndDsDataPrintsUndo() throws Exception {
-    persistDomainWithHosts(ns1, ns2);
+    persistDomainWithHosts(defaultDomainBase, defaultDsData, ns1, ns2);
     runCommandForced(
         "--domain_name=evil.tld",
         "--hosts=urs1.example.com,urs2.example.com",
-        "--dsdata=1 1 1 abcd");
+        "--dsdata=1 1 1 abcd",
+        "--renew_one_year=false");
     eppVerifier
         .expectRegistrarId("CharlestonRoad")
         .expectSuperuser()
-        .verifySent("uniform_rapid_suspension.xml");
+        .verifySent("uniform_rapid_suspension.xml")
+        .verifyNoMoreSent();
     assertInStdout("uniform_rapid_suspension --undo");
     assertInStdout("--domain_name evil.tld");
     assertInStdout("--hosts ns1.example.com,ns2.example.com");
@@ -86,12 +95,16 @@ class UniformRapidSuspensionCommandTest
 
   @Test
   void testCommand_respectsExistingHost() throws Exception {
-    persistDomainWithHosts(urs2, ns1);
-    runCommandForced("--domain_name=evil.tld", "--hosts=urs1.example.com,urs2.example.com");
+    persistDomainWithHosts(defaultDomainBase, defaultDsData, urs2, ns1);
+    runCommandForced(
+        "--domain_name=evil.tld",
+        "--hosts=urs1.example.com,urs2.example.com",
+        "--renew_one_year=false");
     eppVerifier
         .expectRegistrarId("CharlestonRoad")
         .expectSuperuser()
-        .verifySent("uniform_rapid_suspension_existing_host.xml");
+        .verifySent("uniform_rapid_suspension_existing_host.xml")
+        .verifyNoMoreSent();
     assertInStdout("uniform_rapid_suspension --undo ");
     assertInStdout("--domain_name evil.tld");
     assertInStdout("--hosts ns1.example.com,urs2.example.com");
@@ -101,8 +114,11 @@ class UniformRapidSuspensionCommandTest
   @Test
   void testCommand_generatesUndoForUndelegatedDomain() throws Exception {
     persistActiveDomain("evil.tld");
-    runCommandForced("--domain_name=evil.tld", "--hosts=urs1.example.com,urs2.example.com");
-    eppVerifier.verifySentAny();
+    runCommandForced(
+        "--domain_name=evil.tld",
+        "--hosts=urs1.example.com,urs2.example.com",
+        "--renew_one_year=false");
+    eppVerifier.verifySentAny().verifyNoMoreSent();
     assertInStdout("uniform_rapid_suspension --undo");
     assertInStdout("--domain_name evil.tld");
     assertNotInStdout("--locks_to_preserve");
@@ -114,8 +130,8 @@ class UniformRapidSuspensionCommandTest
         newDomainBase("evil.tld").asBuilder()
           .addStatusValue(StatusValue.SERVER_DELETE_PROHIBITED)
           .build());
-    runCommandForced("--domain_name=evil.tld");
-    eppVerifier.verifySentAny();
+    runCommandForced("--domain_name=evil.tld", "--renew_one_year=false");
+    eppVerifier.verifySentAny().verifyNoMoreSent();
     assertInStdout("uniform_rapid_suspension --undo");
     assertInStdout("--domain_name evil.tld");
     assertInStdout("--locks_to_preserve serverDeleteProhibited");
@@ -133,11 +149,13 @@ class UniformRapidSuspensionCommandTest
     runCommandForced(
         "--domain_name=evil.tld",
         "--hosts=urs1.example.com,urs2.example.com",
-        "--dsdata=1 1 1 abcd");
+        "--dsdata=1 1 1 abcd",
+        "--renew_one_year=false");
     eppVerifier
         .expectRegistrarId("CharlestonRoad")
         .expectSuperuser()
-        .verifySent("uniform_rapid_suspension_with_client_hold.xml");
+        .verifySent("uniform_rapid_suspension_with_client_hold.xml")
+        .verifyNoMoreSent();
     assertInStdout("uniform_rapid_suspension --undo");
     assertInStdout("--domain_name evil.tld");
     assertInStdout("--hosts ns1.example.com,ns2.example.com");
@@ -146,54 +164,62 @@ class UniformRapidSuspensionCommandTest
 
   @Test
   void testUndo_removesLocksReplacesHostsAndDsData() throws Exception {
-    persistDomainWithHosts(urs1, urs2);
+    persistDomainWithHosts(defaultDomainBase, defaultDsData, urs1, urs2);
     runCommandForced(
-        "--domain_name=evil.tld", "--undo", "--hosts=ns1.example.com,ns2.example.com");
+        "--domain_name=evil.tld",
+        "--undo",
+        "--hosts=ns1.example.com,ns2.example.com",
+        "--renew_one_year=false");
     eppVerifier
         .expectRegistrarId("CharlestonRoad")
         .expectSuperuser()
-        .verifySent("uniform_rapid_suspension_undo.xml");
+        .verifySent("uniform_rapid_suspension_undo.xml")
+        .verifyNoMoreSent();
     assertNotInStdout("--undo");  // Undo shouldn't print a new undo command.
   }
 
   @Test
   void testUndo_respectsLocksToPreserveFlag() throws Exception {
-    persistDomainWithHosts(urs1, urs2);
+    persistDomainWithHosts(defaultDomainBase, defaultDsData, urs1, urs2);
     runCommandForced(
         "--domain_name=evil.tld",
         "--undo",
         "--locks_to_preserve=serverDeleteProhibited",
-        "--hosts=ns1.example.com,ns2.example.com");
+        "--hosts=ns1.example.com,ns2.example.com",
+        "--renew_one_year=false");
     eppVerifier
         .expectRegistrarId("CharlestonRoad")
         .expectSuperuser()
-        .verifySent("uniform_rapid_suspension_undo_preserve.xml");
+        .verifySent("uniform_rapid_suspension_undo_preserve.xml")
+        .verifyNoMoreSent();
     assertNotInStdout("--undo");  // Undo shouldn't print a new undo command.
   }
 
   @Test
   void testUndo_restoresClientHolds() throws Exception {
-    persistDomainWithHosts(urs1, urs2);
+    persistDomainWithHosts(defaultDomainBase, defaultDsData, urs1, urs2);
     runCommandForced(
         "--domain_name=evil.tld",
         "--undo",
         "--hosts=ns1.example.com,ns2.example.com",
-        "--restore_client_hold");
+        "--restore_client_hold",
+        "--renew_one_year=false");
     eppVerifier
         .expectRegistrarId("CharlestonRoad")
         .expectSuperuser()
-        .verifySent("uniform_rapid_suspension_undo_client_hold.xml");
+        .verifySent("uniform_rapid_suspension_undo_client_hold.xml")
+        .verifyNoMoreSent();
     assertNotInStdout("--undo"); // Undo shouldn't print a new undo command.
   }
 
   @Test
-  void testAutorenews_setToFalsebyDefault() throws Exception {
+  void testAutorenews_setToFalseByDefault() throws Exception {
     persistResource(
         newDomainBase("evil.tld")
             .asBuilder()
             .addStatusValue(StatusValue.SERVER_DELETE_PROHIBITED)
             .build());
-    runCommandForced("--domain_name=evil.tld");
+    runCommandForced("--domain_name=evil.tld", "--renew_one_year=false");
     eppVerifier.verifySentAny();
     assertInStdout("<superuser:autorenews>false</superuser:autorenews>");
   }
@@ -209,11 +235,57 @@ class UniformRapidSuspensionCommandTest
         "--domain_name=evil.tld",
         "--undo",
         "--hosts=ns1.example.com,ns2.example.com",
-        "--restore_client_hold");
+        "--restore_client_hold",
+        "--renew_one_year=false");
     eppVerifier.verifySentAny();
     assertInStdout("<superuser:autorenews>true</superuser:autorenews>");
   }
 
+  @Test
+  void testRenewOneYear_renewFlowIsTriggered() throws Exception {
+    // this test case was written based on an existing test case,
+    // testUndo_removesLocksReplacesHostsAndDsData() but two things were modified to test the
+    // renew workflow, which were:
+    // 1) a different domain that contains creation time and expiration time
+    // 2) renew_one_year is set to true
+
+    persistDomainWithHosts(
+        newDomainBase("evil.tld")
+            .asBuilder()
+            .setCreationTimeForTest(DateTime.parse("2021-10-01T05:01:11Z"))
+            .setRegistrationExpirationTime(DateTime.parse("2022-10-01T05:01:11Z"))
+            .setPersistedCurrentSponsorRegistrarId("CharlestonRoad")
+            .build(),
+        defaultDsData,
+        urs1,
+        urs2);
+
+    runCommandForced(
+        "--domain_name=evil.tld",
+        "--undo",
+        "--hosts=ns1.example.com,ns2.example.com",
+        "--renew_one_year=true");
+
+    // verify if renew flow is triggered
+    eppVerifier
+        .expectRegistrarId("CharlestonRoad")
+        .expectSuperuser()
+        .verifySent(
+            "domain_renew.xml",
+            ImmutableMap.of("DOMAIN", "evil.tld", "EXPDATE", "2022-10-01", "YEARS", "1"));
+
+    // verify if update flow is triggered
+    eppVerifier
+        .expectRegistrarId("CharlestonRoad")
+        .expectSuperuser()
+        .verifySent("uniform_rapid_suspension_undo.xml")
+        .verifyNoMoreSent();
+    assertNotInStdout("--undo"); // Undo shouldn't print a new undo command.
+
+    // verify that no other flows are triggered after the renew and update flows
+    eppVerifier.verifyNoMoreSent();
+  }
+
   @Test
   void testFailure_locksToPreserveWithoutUndo() {
     persistActiveDomain("evil.tld");
@@ -222,7 +294,9 @@ class UniformRapidSuspensionCommandTest
             IllegalArgumentException.class,
             () ->
                 runCommandForced(
-                    "--domain_name=evil.tld", "--locks_to_preserve=serverDeleteProhibited"));
+                    "--domain_name=evil.tld",
+                    "--locks_to_preserve=serverDeleteProhibited",
+                    "--renew_one_year=false"));
     assertThat(thrown).hasMessageThat().contains("--undo");
   }
 
@@ -232,17 +306,29 @@ class UniformRapidSuspensionCommandTest
     ParameterException thrown =
         assertThrows(
             ParameterException.class,
-            () -> runCommandForced("--hosts=urs1.example.com,urs2.example.com"));
+            () ->
+                runCommandForced(
+                    "--hosts=urs1.example.com,urs2.example.com", "--renew_one_year=false"));
     assertThat(thrown).hasMessageThat().contains("--domain_name");
   }
 
+  @Test
+  void testFailure_renewOneYearRequired() {
+    persistActiveDomain("evil.tld");
+    ParameterException thrown =
+        assertThrows(ParameterException.class, () -> runCommandForced("--domain_name=evil.tld"));
+    assertThat(thrown).hasMessageThat().contains("--renew_one_year");
+  }
+
   @Test
   void testFailure_extraFieldInDsData() {
     persistActiveDomain("evil.tld");
     IllegalArgumentException thrown =
         assertThrows(
             IllegalArgumentException.class,
-            () -> runCommandForced("--domain_name=evil.tld", "--dsdata=1 1 1 abc 1"));
+            () ->
+                runCommandForced(
+                    "--domain_name=evil.tld", "--dsdata=1 1 1 abc 1", "--renew_one_year=false"));
     assertThat(thrown)
         .hasMessageThat()
         .contains("dsRecord 1 1 1 abc 1 should have 4 parts, but has 5");
@@ -254,7 +340,9 @@ class UniformRapidSuspensionCommandTest
     IllegalArgumentException thrown =
         assertThrows(
             IllegalArgumentException.class,
-            () -> runCommandForced("--domain_name=evil.tld", "--dsdata=1 1 1"));
+            () ->
+                runCommandForced(
+                    "--domain_name=evil.tld", "--dsdata=1 1 1", "--renew_one_year=false"));
     assertThat(thrown).hasMessageThat().contains("dsRecord 1 1 1 should have 4 parts, but has 3");
   }
 
@@ -264,7 +352,9 @@ class UniformRapidSuspensionCommandTest
     IllegalArgumentException thrown =
         assertThrows(
             IllegalArgumentException.class,
-            () -> runCommandForced("--domain_name=evil.tld", "--dsdata=1,2,3"));
+            () ->
+                runCommandForced(
+                    "--domain_name=evil.tld", "--dsdata=1,2,3", "--renew_one_year=false"));
     assertThat(thrown).hasMessageThat().contains("dsRecord 1 should have 4 parts, but has 1");
   }
 }

release/cloudbuild-nomulus.yaml

@@ -84,6 +84,7 @@ steps:
   - |
     ./release/stage_beam_pipeline.sh \
       beamPipelineCommon \
+      beam_pipeline_common \
       ${TAG_NAME} \
       ${PROJECT_ID} \
       google.registry.beam.initsql.InitSqlPipeline \

release/stage_beam_pipeline.sh

@@ -38,21 +38,22 @@
 
 set -e
 
-if (( "$#" < 5 ||  $(("$#" % 2)) == 0 ));
+if (( "$#" < 6 ||  $(("$#" % 2)) == 1 ));
 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
 

util/src/main/java/google/registry/util/CloudTasksUtils.java

@@ -37,7 +37,7 @@ import com.google.protobuf.ByteString;
 import java.io.Serializable;
 import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
-import javax.inject.Provider;
+import java.util.function.Supplier;
 
 /** Utilities for dealing with Cloud Tasks. */
 public class CloudTasksUtils implements Serializable {
@@ -48,17 +48,14 @@ public class CloudTasksUtils implements Serializable {
   private final Retrier retrier;
   private final String projectId;
   private final String locationId;
-  private final Provider<CloudTasksClient> clientProvider;
+  private final SerializableCloudTasksClient client;
 
   public CloudTasksUtils(
-      Retrier retrier,
-      String projectId,
-      String locationId,
-      Provider<CloudTasksClient> clientProvider) {
+      Retrier retrier, String projectId, String locationId, SerializableCloudTasksClient client) {
     this.retrier = retrier;
     this.projectId = projectId;
     this.locationId = locationId;
-    this.clientProvider = clientProvider;
+    this.client = client;
   }
 
   public Task enqueue(String queue, Task task) {
@@ -69,9 +66,7 @@ public class CloudTasksUtils implements Serializable {
               queue,
               task.getAppEngineHttpRequest().getRelativeUri(),
               task.getAppEngineHttpRequest().getAppEngineRouting().getService());
-          try (CloudTasksClient client = clientProvider.get()) {
-            return client.createTask(QueueName.of(projectId, locationId, queue), task);
-          }
+          return client.enqueue(projectId, locationId, queue, task);
         },
         ApiException.class);
   }
@@ -141,4 +136,28 @@ public class CloudTasksUtils implements Serializable {
   public static Task createGetTask(String path, String service, Multimap<String, String> params) {
     return createTask(path, HttpMethod.GET, service, params);
   }
+
+  public abstract static class SerializableCloudTasksClient implements Serializable {
+    public abstract Task enqueue(String projectId, String locationId, String queueName, Task task);
+  }
+
+  public static class GcpCloudTasksClient extends SerializableCloudTasksClient {
+
+    private static final long serialVersionUID = -5959253033129154037L;
+
+    // Use a supplier so that we can use try-with-resources with the client, which implements
+    // Autocloseable.
+    private final Supplier<CloudTasksClient> clientSupplier;
+
+    public GcpCloudTasksClient(Supplier<CloudTasksClient> clientSupplier) {
+      this.clientSupplier = clientSupplier;
+    }
+
+    @Override
+    public Task enqueue(String projectId, String locationId, String queueName, Task task) {
+      try (CloudTasksClient client = clientSupplier.get()) {
+        return client.createTask(QueueName.of(projectId, locationId, queueName), task);
+      }
+    }
+  }
 }

util/src/main/java/google/registry/util/GoogleCredentialsBundle.java

@@ -22,25 +22,23 @@ import com.google.api.client.http.HttpTransport;
 import com.google.api.client.json.JsonFactory;
 import com.google.auth.http.HttpCredentialsAdapter;
 import com.google.auth.oauth2.GoogleCredentials;
+import java.io.Serializable;
 
 /**
  * Helper class to provide {@link HttpTransport}, {@link JsonFactory} and {@link
  * HttpRequestInitializer} for a given {@link GoogleCredentials}. These classes are normally needed
  * for creating the instance of a GCP client.
  */
-public class GoogleCredentialsBundle {
+public class GoogleCredentialsBundle implements Serializable {
+
+  private static final HttpTransport HTTP_TRANSPORT = Utils.getDefaultTransport();
+  private static final JsonFactory JSON_FACTORY = Utils.getDefaultJsonFactory();
 
   private GoogleCredentials googleCredentials;
-  private HttpTransport httpTransport;
-  private JsonFactory jsonFactory;
-  private HttpRequestInitializer httpRequestInitializer;
 
   private GoogleCredentialsBundle(GoogleCredentials googleCredentials) {
     checkNotNull(googleCredentials);
     this.googleCredentials = googleCredentials;
-    this.httpTransport = Utils.getDefaultTransport();
-    this.jsonFactory = Utils.getDefaultJsonFactory();
-    this.httpRequestInitializer = new HttpCredentialsAdapter(googleCredentials);
   }
 
   /** Creates a {@link GoogleCredentialsBundle} instance from given {@link GoogleCredentials}. */
@@ -55,16 +53,16 @@ public class GoogleCredentialsBundle {
 
   /** Returns the instance of {@link HttpTransport}. */
   public HttpTransport getHttpTransport() {
-    return httpTransport;
+    return HTTP_TRANSPORT;
   }
 
   /** Returns the instance of {@link JsonFactory}. */
   public JsonFactory getJsonFactory() {
-    return jsonFactory;
+    return JSON_FACTORY;
   }
 
   /** Returns the instance of {@link HttpRequestInitializer}. */
   public HttpRequestInitializer getHttpRequestInitializer() {
-    return httpRequestInitializer;
+    return new HttpCredentialsAdapter(googleCredentials);
   }
 }

util/src/main/java/google/registry/util/PlaceholderEnvironment.java

@@ -0,0 +1,82 @@
+// Copyright 2020 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.util;
+
+import com.google.apphosting.api.ApiProxy.Environment;
+import com.google.common.collect.ImmutableMap;
+import java.util.Map;
+
+/** A placeholder GAE environment class that is used when masquerading a thread as a GAE thread. */
+public final class PlaceholderEnvironment implements Environment {
+
+  private static final PlaceholderEnvironment INSTANCE = new PlaceholderEnvironment();
+
+  public static PlaceholderEnvironment get() {
+    return INSTANCE;
+  }
+
+  private PlaceholderEnvironment() {}
+
+  @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();
+  }
+}

util/src/test/java/google/registry/util/CloudTasksUtilsTest.java

@@ -17,18 +17,18 @@ package google.registry.util;
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
-import com.google.cloud.tasks.v2.CloudTasksClient;
 import com.google.cloud.tasks.v2.HttpMethod;
-import com.google.cloud.tasks.v2.QueueName;
 import com.google.cloud.tasks.v2.Task;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.LinkedListMultimap;
 import google.registry.testing.FakeClock;
 import google.registry.testing.FakeSleeper;
+import google.registry.util.CloudTasksUtils.SerializableCloudTasksClient;
 import java.nio.charset.StandardCharsets;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -37,21 +37,18 @@ import org.junit.jupiter.api.Test;
 public class CloudTasksUtilsTest {
   // Use a LinkedListMultimap to preserve order of the inserted entries for assertion.
   private final LinkedListMultimap<String, String> params = LinkedListMultimap.create();
-  private final CloudTasksClient mockClient = mock(CloudTasksClient.class);
+  private final SerializableCloudTasksClient mockClient = mock(SerializableCloudTasksClient.class);
   private final CloudTasksUtils cloudTasksUtils =
       new CloudTasksUtils(
-          new Retrier(new FakeSleeper(new FakeClock()), 1),
-          "project",
-          "location",
-          () -> mockClient);
+          new Retrier(new FakeSleeper(new FakeClock()), 1), "project", "location", mockClient);
 
   @BeforeEach
   void beforeEach() {
     params.put("key1", "val1");
     params.put("key2", "val2");
     params.put("key1", "val3");
-    when(mockClient.createTask(any(QueueName.class), any(Task.class)))
-        .thenAnswer(invocation -> invocation.getArgument(1));
+    when(mockClient.enqueue(anyString(), anyString(), anyString(), any(Task.class)))
+        .thenAnswer(invocation -> invocation.getArgument(3));
   }
 
   @Test
@@ -94,7 +91,7 @@ public class CloudTasksUtilsTest {
   void testSuccess_enqueueTask() {
     Task task = CloudTasksUtils.createGetTask("/the/path", "myservice", params);
     cloudTasksUtils.enqueue("test-queue", task);
-    verify(mockClient).createTask(QueueName.of("project", "location", "test-queue"), task);
+    verify(mockClient).enqueue("project", "location", "test-queue", task);
   }
 
   @Test
@@ -102,8 +99,8 @@ public class CloudTasksUtilsTest {
     Task task1 = CloudTasksUtils.createGetTask("/the/path", "myservice", params);
     Task task2 = CloudTasksUtils.createGetTask("/other/path", "yourservice", params);
     cloudTasksUtils.enqueue("test-queue", task1, task2);
-    verify(mockClient).createTask(QueueName.of("project", "location", "test-queue"), task1);
-    verify(mockClient).createTask(QueueName.of("project", "location", "test-queue"), task2);
+    verify(mockClient).enqueue("project", "location", "test-queue", task1);
+    verify(mockClient).enqueue("project", "location", "test-queue", task2);
   }
 
   @Test
@@ -111,7 +108,7 @@ public class CloudTasksUtilsTest {
     Task task1 = CloudTasksUtils.createGetTask("/the/path", "myservice", params);
     Task task2 = CloudTasksUtils.createGetTask("/other/path", "yourservice", params);
     cloudTasksUtils.enqueue("test-queue", ImmutableList.of(task1, task2));
-    verify(mockClient).createTask(QueueName.of("project", "location", "test-queue"), task1);
-    verify(mockClient).createTask(QueueName.of("project", "location", "test-queue"), task2);
+    verify(mockClient).enqueue("project", "location", "test-queue", task1);
+    verify(mockClient).enqueue("project", "location", "test-queue", task2);
   }
 }