Finish RDE pipeline implementation in SQL mode (#1330)

This PR adds the final step in RDE pipeline (enqueueing the next action
  to Cloud Tasks) and makes some necessary changes, namely by making all
  CloudTasksUtils related classes serializable, so that they can be used
  on Beam.

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

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/rde/RdePipelineTest.java

@@ -69,6 +69,8 @@ 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;
@@ -134,6 +136,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();
 
@@ -212,9 +216,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);
+    rdePipeline = new RdePipeline(options, gcsUtils, cloudTasksHelper.getTestCloudTasksUtils());
   }
 
   @AfterEach
@@ -314,6 +318,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 +356,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/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/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 {