Remove static methods in back up actions (#1481)

* Remove static methods in back up actions

* Remove BigqueryPollJob helper class

* Add schedule time in task comparison

* Change payload type from byte[] to ByteString

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

@@ -14,18 +14,22 @@
 
 package google.registry.export;
 
-import static google.registry.export.CheckBackupAction.enqueuePollTask;
 import static google.registry.request.Action.Method.POST;
 
+import com.google.common.base.Joiner;
+import com.google.common.collect.ImmutableMultimap;
 import com.google.common.flogger.FluentLogger;
 import google.registry.config.RegistryConfig;
 import google.registry.export.datastore.DatastoreAdmin;
 import google.registry.export.datastore.Operation;
 import google.registry.model.annotations.DeleteAfterMigration;
 import google.registry.request.Action;
+import google.registry.request.Action.Service;
 import google.registry.request.HttpException.InternalServerErrorException;
 import google.registry.request.Response;
 import google.registry.request.auth.Auth;
+import google.registry.util.Clock;
+import google.registry.util.CloudTasksUtils;
 import javax.inject.Inject;
 
 /**
@@ -59,6 +63,8 @@ public class BackupDatastoreAction implements Runnable {
 
   @Inject DatastoreAdmin datastoreAdmin;
   @Inject Response response;
+  @Inject Clock clock;
+  @Inject CloudTasksUtils cloudTasksUtils;
 
   @Inject
   BackupDatastoreAction() {}
@@ -73,8 +79,19 @@ public class BackupDatastoreAction implements Runnable {
               .execute();
 
       String backupName = backup.getName();
-      // Enqueue a poll task to monitor the backup and load REPORTING-related kinds into bigquery.
-      enqueuePollTask(backupName, AnnotatedEntities.getReportingKinds());
+      // Enqueue a poll task to monitor the backup for completion and load reporting-related kinds
+      // into bigquery.
+      cloudTasksUtils.enqueue(
+          CheckBackupAction.QUEUE,
+          cloudTasksUtils.createPostTaskWithDelay(
+              CheckBackupAction.PATH,
+              Service.BACKEND.toString(),
+              ImmutableMultimap.of(
+                  CheckBackupAction.CHECK_BACKUP_NAME_PARAM,
+                  backupName,
+                  CheckBackupAction.CHECK_BACKUP_KINDS_TO_LOAD_PARAM,
+                  Joiner.on(',').join(AnnotatedEntities.getReportingKinds())),
+              CheckBackupAction.POLL_COUNTDOWN));
       String message =
           String.format(
               "Datastore backup started with name: %s\nSaving to %s",

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

@@ -14,18 +14,14 @@
 
 package google.registry.export;
 
-import static com.google.appengine.api.taskqueue.QueueFactory.getQueue;
-import static com.google.appengine.api.taskqueue.TaskOptions.Builder.withUrl;
 import static google.registry.bigquery.BigqueryUtils.toJobReferenceString;
 
 import com.google.api.services.bigquery.Bigquery;
 import com.google.api.services.bigquery.model.Job;
 import com.google.api.services.bigquery.model.JobReference;
-import com.google.appengine.api.taskqueue.Queue;
-import com.google.appengine.api.taskqueue.TaskHandle;
-import com.google.appengine.api.taskqueue.TaskOptions;
-import com.google.appengine.api.taskqueue.TaskOptions.Method;
+import com.google.cloud.tasks.v2.Task;
 import com.google.common.flogger.FluentLogger;
+import com.google.protobuf.ByteString;
 import dagger.Lazy;
 import google.registry.request.Action;
 import google.registry.request.Header;
@@ -33,12 +29,10 @@ import google.registry.request.HttpException.BadRequestException;
 import google.registry.request.HttpException.NotModifiedException;
 import google.registry.request.Payload;
 import google.registry.request.auth.Auth;
-import google.registry.util.TaskQueueUtils;
+import google.registry.util.CloudTasksUtils;
 import java.io.ByteArrayInputStream;
-import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.ObjectInputStream;
-import java.io.ObjectOutputStream;
 import javax.inject.Inject;
 import org.joda.time.Duration;
 
@@ -67,11 +61,14 @@ public class BigqueryPollJobAction implements Runnable {
   static final Duration POLL_COUNTDOWN = Duration.standardSeconds(20);
 
   @Inject Bigquery bigquery;
-  @Inject TaskQueueUtils taskQueueUtils;
+  @Inject CloudTasksUtils cloudTasksUtils;
+
   @Inject @Header(CHAINED_TASK_QUEUE_HEADER) Lazy<String> chainedQueueName;
   @Inject @Header(PROJECT_ID_HEADER) String projectId;
   @Inject @Header(JOB_ID_HEADER) String jobId;
-  @Inject @Payload byte[] payload;
+
+  @Inject @Payload ByteString payload;
+
   @Inject BigqueryPollJobAction() {}
 
   @Override
@@ -79,20 +76,25 @@ public class BigqueryPollJobAction implements Runnable {
     boolean jobOutcome =
         checkJobOutcome(); // Throws a NotModifiedException if the job hasn't completed.
     // If the job failed, do not enqueue the next step.
-    if (!jobOutcome || payload == null || payload.length == 0) {
+    if (!jobOutcome || payload == null || payload.size() == 0) {
       return;
     }
     // If there is a payload, it's a chained task, so enqueue it.
-    TaskOptions task;
+    Task task;
     try {
-      task = (TaskOptions) new ObjectInputStream(new ByteArrayInputStream(payload)).readObject();
+      task =
+          (Task)
+              new ObjectInputStream(new ByteArrayInputStream(payload.toByteArray())).readObject();
     } catch (ClassNotFoundException | IOException e) {
       throw new BadRequestException("Cannot deserialize task from payload", e);
     }
-    String taskName = taskQueueUtils.enqueue(getQueue(chainedQueueName.get()), task).getName();
+    Task enqueuedTask = cloudTasksUtils.enqueue(chainedQueueName.get(), task);
     logger.atInfo().log(
         "Added chained task %s for %s to queue %s: %s",
-        taskName, task.getUrl(), chainedQueueName.get(), task);
+        enqueuedTask.getName(),
+        enqueuedTask.getAppEngineHttpRequest().getRelativeUri(),
+        chainedQueueName.get(),
+        enqueuedTask);
   }
 
   /**
@@ -124,51 +126,4 @@ public class BigqueryPollJobAction implements Runnable {
     logger.atInfo().log("Bigquery job succeeded - %s.", jobRefString);
     return true;
   }
-
-
-  /** Helper class to enqueue a bigquery poll job. */
-  public static class BigqueryPollJobEnqueuer {
-
-    private final TaskQueueUtils taskQueueUtils;
-
-    @Inject
-    BigqueryPollJobEnqueuer(TaskQueueUtils taskQueueUtils) {
-      this.taskQueueUtils = taskQueueUtils;
-    }
-
-    /** Enqueue a task to poll for the success or failure of the referenced BigQuery job. */
-    public TaskHandle enqueuePollTask(JobReference jobRef) {
-      return taskQueueUtils.enqueue(
-          getQueue(QUEUE), createCommonPollTask(jobRef).method(Method.GET));
-    }
-
-    /**
-     * Enqueue a task to poll for the success or failure of the referenced BigQuery job and to
-     * launch the provided task in the specified queue if the job succeeds.
-     */
-    public TaskHandle enqueuePollTask(
-        JobReference jobRef, TaskOptions chainedTask, Queue chainedTaskQueue) throws IOException {
-      // Serialize the chainedTask into a byte array to put in the task payload.
-      ByteArrayOutputStream taskBytes = new ByteArrayOutputStream();
-      new ObjectOutputStream(taskBytes).writeObject(chainedTask);
-      return taskQueueUtils.enqueue(
-          getQueue(QUEUE),
-          createCommonPollTask(jobRef)
-              .method(Method.POST)
-              .header(CHAINED_TASK_QUEUE_HEADER, chainedTaskQueue.getQueueName())
-              .payload(taskBytes.toByteArray()));
-    }
-
-    /**
-     * Enqueue a task to poll for the success or failure of the referenced BigQuery job and to
-     * launch the provided task in the specified queue if the job succeeds.
-     */
-    private static TaskOptions createCommonPollTask(JobReference jobRef) {
-      // Omit host header so that task will be run on the current backend/module.
-      return withUrl(PATH)
-          .countdownMillis(POLL_COUNTDOWN.getMillis())
-          .header(PROJECT_ID_HEADER, jobRef.getProjectId())
-          .header(JOB_ID_HEADER, jobRef.getJobId());
-    }
-  }
 }

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

@@ -16,18 +16,14 @@ package google.registry.export;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.Sets.intersection;
-import static google.registry.export.UploadDatastoreBackupAction.enqueueUploadBackupTask;
 import static google.registry.request.Action.Method.GET;
 import static google.registry.request.Action.Method.POST;
 import static javax.servlet.http.HttpServletResponse.SC_NOT_FOUND;
 
 import com.google.api.client.googleapis.json.GoogleJsonResponseException;
-import com.google.appengine.api.taskqueue.QueueFactory;
-import com.google.appengine.api.taskqueue.TaskHandle;
-import com.google.appengine.api.taskqueue.TaskOptions;
-import com.google.appengine.api.taskqueue.TaskOptions.Method;
 import com.google.common.base.Joiner;
 import com.google.common.base.Splitter;
+import com.google.common.collect.ImmutableMultimap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Sets;
 import com.google.common.flogger.FluentLogger;
@@ -35,6 +31,7 @@ import google.registry.export.datastore.DatastoreAdmin;
 import google.registry.export.datastore.Operation;
 import google.registry.model.annotations.DeleteAfterMigration;
 import google.registry.request.Action;
+import google.registry.request.Action.Service;
 import google.registry.request.HttpException;
 import google.registry.request.HttpException.BadRequestException;
 import google.registry.request.HttpException.InternalServerErrorException;
@@ -45,6 +42,7 @@ import google.registry.request.RequestMethod;
 import google.registry.request.Response;
 import google.registry.request.auth.Auth;
 import google.registry.util.Clock;
+import google.registry.util.CloudTasksUtils;
 import java.io.IOException;
 import java.util.Set;
 import javax.inject.Inject;
@@ -83,6 +81,7 @@ public class CheckBackupAction implements Runnable {
   @Inject DatastoreAdmin datastoreAdmin;
   @Inject Clock clock;
   @Inject Response response;
+  @Inject CloudTasksUtils cloudTasksUtils;
   @Inject @RequestMethod Action.Method requestMethod;
 
   @Inject
@@ -175,21 +174,22 @@ public class CheckBackupAction implements Runnable {
     if (exportedKindsToLoad.isEmpty()) {
       message += "no kinds to load into BigQuery.";
     } else {
-      enqueueUploadBackupTask(backupId, backup.getExportFolderUrl(), exportedKindsToLoad);
+      /** Enqueue a task for starting a backup load. */
+      cloudTasksUtils.enqueue(
+          UploadDatastoreBackupAction.QUEUE,
+          cloudTasksUtils.createPostTask(
+              UploadDatastoreBackupAction.PATH,
+              Service.BACKEND.toString(),
+              ImmutableMultimap.of(
+                  UploadDatastoreBackupAction.UPLOAD_BACKUP_ID_PARAM,
+                  backupId,
+                  UploadDatastoreBackupAction.UPLOAD_BACKUP_FOLDER_PARAM,
+                  backup.getExportFolderUrl(),
+                  UploadDatastoreBackupAction.UPLOAD_BACKUP_KINDS_PARAM,
+                  Joiner.on(',').join(exportedKindsToLoad))));
       message += "BigQuery load task enqueued.";
     }
     logger.atInfo().log(message);
     response.setPayload(message);
   }
-
-  /** Enqueue a poll task to monitor the named backup for completion. */
-  static TaskHandle enqueuePollTask(String backupId, ImmutableSet<String> kindsToLoad) {
-    return QueueFactory.getQueue(QUEUE)
-        .add(
-            TaskOptions.Builder.withUrl(PATH)
-                .method(Method.POST)
-                .countdownMillis(POLL_COUNTDOWN.getMillis())
-                .param(CHECK_BACKUP_NAME_PARAM, backupId)
-                .param(CHECK_BACKUP_KINDS_TO_LOAD_PARAM, Joiner.on(',').join(kindsToLoad)));
-  }
 }

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

@@ -21,8 +21,6 @@ import com.google.api.services.bigquery.Bigquery;
 import com.google.api.services.bigquery.model.Table;
 import com.google.api.services.bigquery.model.TableReference;
 import com.google.api.services.bigquery.model.ViewDefinition;
-import com.google.appengine.api.taskqueue.TaskOptions;
-import com.google.appengine.api.taskqueue.TaskOptions.Method;
 import com.google.common.flogger.FluentLogger;
 import google.registry.bigquery.CheckedBigquery;
 import google.registry.config.RegistryConfig.Config;
@@ -85,17 +83,6 @@ public class UpdateSnapshotViewAction implements Runnable {
   @Inject
   UpdateSnapshotViewAction() {}
 
-  /** Create a task for updating a snapshot view. */
-  static TaskOptions createViewUpdateTask(
-      String datasetId, String tableId, String kindName, String viewName) {
-    return TaskOptions.Builder.withUrl(PATH)
-        .method(Method.POST)
-        .param(UPDATE_SNAPSHOT_DATASET_ID_PARAM, datasetId)
-        .param(UPDATE_SNAPSHOT_TABLE_ID_PARAM, tableId)
-        .param(UPDATE_SNAPSHOT_KIND_PARAM, kindName)
-        .param(UPDATE_SNAPSHOT_VIEWNAME_PARAM, viewName);
-  }
-
   @Override
   public void run() {
     try {

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

@@ -14,9 +14,7 @@
 
 package google.registry.export;
 
-import static com.google.appengine.api.taskqueue.QueueFactory.getQueue;
 import static com.google.common.base.MoreObjects.firstNonNull;
-import static google.registry.export.UpdateSnapshotViewAction.createViewUpdateTask;
 import static google.registry.request.Action.Method.POST;
 
 import com.google.api.services.bigquery.Bigquery;
@@ -25,27 +23,33 @@ import com.google.api.services.bigquery.model.JobConfiguration;
 import com.google.api.services.bigquery.model.JobConfigurationLoad;
 import com.google.api.services.bigquery.model.JobReference;
 import com.google.api.services.bigquery.model.TableReference;
-import com.google.appengine.api.taskqueue.TaskHandle;
-import com.google.appengine.api.taskqueue.TaskOptions;
-import com.google.appengine.api.taskqueue.TaskOptions.Method;
+import com.google.cloud.tasks.v2.Task;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Joiner;
 import com.google.common.base.Splitter;
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableMultimap;
 import com.google.common.flogger.FluentLogger;
+import com.google.common.net.HttpHeaders;
+import com.google.common.net.MediaType;
+import com.google.protobuf.ByteString;
+import com.google.protobuf.util.Timestamps;
 import google.registry.bigquery.BigqueryUtils.SourceFormat;
 import google.registry.bigquery.BigqueryUtils.WriteDisposition;
 import google.registry.bigquery.CheckedBigquery;
 import google.registry.config.RegistryConfig.Config;
-import google.registry.export.BigqueryPollJobAction.BigqueryPollJobEnqueuer;
 import google.registry.model.annotations.DeleteAfterMigration;
 import google.registry.request.Action;
+import google.registry.request.Action.Service;
 import google.registry.request.HttpException.BadRequestException;
 import google.registry.request.HttpException.InternalServerErrorException;
 import google.registry.request.Parameter;
 import google.registry.request.auth.Auth;
+import google.registry.util.Clock;
+import google.registry.util.CloudTasksUtils;
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
+import java.io.ObjectOutputStream;
 import javax.inject.Inject;
 
 /** Action to load a Datastore backup from Google Cloud Storage into BigQuery. */
@@ -75,7 +79,9 @@ public class UploadDatastoreBackupAction implements Runnable {
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
 
   @Inject CheckedBigquery checkedBigquery;
-  @Inject BigqueryPollJobEnqueuer bigqueryPollEnqueuer;
+  @Inject CloudTasksUtils cloudTasksUtils;
+  @Inject Clock clock;
+
   @Inject @Config("projectId") String projectId;
 
   @Inject
@@ -93,18 +99,6 @@ public class UploadDatastoreBackupAction implements Runnable {
   @Inject
   UploadDatastoreBackupAction() {}
 
-  /** Enqueue a task for starting a backup load. */
-  public static TaskHandle enqueueUploadBackupTask(
-      String backupId, String gcsFile, ImmutableSet<String> kinds) {
-    return getQueue(QUEUE)
-        .add(
-            TaskOptions.Builder.withUrl(PATH)
-                .method(Method.POST)
-                .param(UPLOAD_BACKUP_ID_PARAM, backupId)
-                .param(UPLOAD_BACKUP_FOLDER_PARAM, gcsFile)
-                .param(UPLOAD_BACKUP_KINDS_PARAM, Joiner.on(',').join(kinds)));
-  }
-
   @Override
   public void run() {
     try {
@@ -142,12 +136,46 @@ public class UploadDatastoreBackupAction implements Runnable {
       Job job = makeLoadJob(jobRef, sourceUri, tableId);
       bigquery.jobs().insert(projectId, job).execute();
 
-      // Enqueue a task to check on the load job's completion, and if it succeeds, to update a
-      // well-known view in BigQuery to point at the newly loaded backup table for this kind.
-      bigqueryPollEnqueuer.enqueuePollTask(
-          jobRef,
-          createViewUpdateTask(BACKUP_DATASET, tableId, kindName, LATEST_BACKUP_VIEW_NAME),
-          getQueue(UpdateSnapshotViewAction.QUEUE));
+      // Serialize the chainedTask into a byte array to put in the task payload.
+      ByteArrayOutputStream taskBytes = new ByteArrayOutputStream();
+      new ObjectOutputStream(taskBytes)
+          .writeObject(
+              cloudTasksUtils.createPostTask(
+                  UpdateSnapshotViewAction.PATH,
+                  Service.BACKEND.toString(),
+                  ImmutableMultimap.of(
+                      UpdateSnapshotViewAction.UPDATE_SNAPSHOT_DATASET_ID_PARAM,
+                      BACKUP_DATASET,
+                      UpdateSnapshotViewAction.UPDATE_SNAPSHOT_TABLE_ID_PARAM,
+                      tableId,
+                      UpdateSnapshotViewAction.UPDATE_SNAPSHOT_KIND_PARAM,
+                      kindName,
+                      UpdateSnapshotViewAction.UPDATE_SNAPSHOT_VIEWNAME_PARAM,
+                      LATEST_BACKUP_VIEW_NAME)));
+
+      // Enqueues a task to poll for the success or failure of the referenced BigQuery job and to
+      // launch the provided task in the specified queue if the job succeeds.
+      cloudTasksUtils.enqueue(
+          BigqueryPollJobAction.QUEUE,
+          Task.newBuilder()
+              .setAppEngineHttpRequest(
+                  cloudTasksUtils
+                      .createPostTask(BigqueryPollJobAction.PATH, Service.BACKEND.toString(), null)
+                      .getAppEngineHttpRequest()
+                      .toBuilder()
+                      .putHeaders(BigqueryPollJobAction.PROJECT_ID_HEADER, jobRef.getProjectId())
+                      .putHeaders(BigqueryPollJobAction.JOB_ID_HEADER, jobRef.getJobId())
+                      .putHeaders(
+                          BigqueryPollJobAction.CHAINED_TASK_QUEUE_HEADER,
+                          UpdateSnapshotViewAction.QUEUE)
+                      // need to include CONTENT_TYPE in header when body is not empty
+                      .putHeaders(HttpHeaders.CONTENT_TYPE, MediaType.FORM_DATA.toString())
+                      .setBody(ByteString.copyFrom(taskBytes.toByteArray()))
+                      .build())
+              .setScheduleTime(
+                  Timestamps.fromMillis(
+                      clock.nowUtc().plus(BigqueryPollJobAction.POLL_COUNTDOWN).getMillis()))
+              .build());
 
       builder.append(String.format(" - %s:%s\n", projectId, jobId));
       logger.atInfo().log("Submitted load job %s:%s.", projectId, jobId);

core/src/main/java/google/registry/request/RequestModule.java

@@ -30,6 +30,7 @@ import com.google.common.collect.ImmutableSet;
 import com.google.common.io.ByteStreams;
 import com.google.common.io.CharStreams;
 import com.google.common.net.MediaType;
+import com.google.protobuf.ByteString;
 import dagger.Module;
 import dagger.Provides;
 import google.registry.model.common.DatabaseMigrationStateSchedule.PrimaryDatabase;
@@ -184,6 +185,16 @@ public final class RequestModule {
     }
   }
 
+  @Provides
+  @Payload
+  static ByteString providePayloadAsByteString(HttpServletRequest req) {
+    try {
+      return ByteString.copyFrom(ByteStreams.toByteArray(req.getInputStream()));
+    } catch (IOException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
   @Provides
   static LockHandler provideLockHandler(LockHandlerImpl lockHandler) {
     return lockHandler;