Restore commit logs from other project (#1236)

* Restore commit logs from other project

Allow non-production projects to restore commit logs from another
project. This feature can be used to duplicate a realistic testing
environment.

An optional parameter is added that can override the default commit log
location.

Tested successfully in QA.

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

@@ -18,8 +18,10 @@ import static com.google.appengine.api.ThreadManager.currentRequestThreadFactory
 import static com.google.common.util.concurrent.MoreExecutors.listeningDecorator;
 import static google.registry.backup.ExportCommitLogDiffAction.LOWER_CHECKPOINT_TIME_PARAM;
 import static google.registry.backup.ExportCommitLogDiffAction.UPPER_CHECKPOINT_TIME_PARAM;
+import static google.registry.backup.RestoreCommitLogsAction.BUCKET_OVERRIDE_PARAM;
 import static google.registry.backup.RestoreCommitLogsAction.FROM_TIME_PARAM;
 import static google.registry.backup.RestoreCommitLogsAction.TO_TIME_PARAM;
+import static google.registry.request.RequestParameters.extractOptionalParameter;
 import static google.registry.request.RequestParameters.extractRequiredDatetimeParameter;
 import static google.registry.request.RequestParameters.extractRequiredParameter;
 import static java.util.concurrent.Executors.newFixedThreadPool;
@@ -32,6 +34,7 @@ import google.registry.cron.CommitLogFanoutAction;
 import google.registry.request.HttpException.BadRequestException;
 import google.registry.request.Parameter;
 import java.lang.annotation.Documented;
+import java.util.Optional;
 import javax.inject.Qualifier;
 import javax.servlet.http.HttpServletRequest;
 import org.joda.time.DateTime;
@@ -75,6 +78,12 @@ public final class BackupModule {
     return extractRequiredDatetimeParameter(req, UPPER_CHECKPOINT_TIME_PARAM);
   }
 
+  @Provides
+  @Parameter(BUCKET_OVERRIDE_PARAM)
+  static Optional<String> provideBucketOverride(HttpServletRequest req) {
+    return extractOptionalParameter(req, BUCKET_OVERRIDE_PARAM);
+  }
+
   @Provides
   @Parameter(FROM_TIME_PARAM)
   static DateTime provideFromTime(HttpServletRequest req) {

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

@@ -32,7 +32,6 @@ import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.ListeningExecutorService;
 import google.registry.backup.BackupModule.Backups;
-import google.registry.config.RegistryConfig.Config;
 import java.io.IOException;
 import java.util.Iterator;
 import java.util.Map;
@@ -47,16 +46,16 @@ class GcsDiffFileLister {
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
 
   @Inject GcsService gcsService;
-  @Inject @Config("commitLogGcsBucket") String gcsBucket;
   @Inject @Backups ListeningExecutorService executor;
   @Inject GcsDiffFileLister() {}
 
   /**
    * Traverses the sequence of diff files backwards from checkpointTime and inserts the file
-   * metadata into "sequence".  Returns true if a complete sequence was discovered, false if one or
+   * metadata into "sequence". Returns true if a complete sequence was discovered, false if one or
    * more files are missing.
    */
   private boolean constructDiffSequence(
+      String gcsBucket,
       Map<DateTime, ListenableFuture<GcsFileMetadata>> upperBoundTimesToMetadata,
       DateTime fromTime,
       DateTime lastTime,
@@ -69,7 +68,7 @@ class GcsDiffFileLister {
       } else {
         String filename = DIFF_FILE_PREFIX + checkpointTime;
         logger.atInfo().log("Patching GCS list; discovered file: %s", filename);
-        metadata = getMetadata(filename);
+        metadata = getMetadata(gcsBucket, filename);
 
         // If we hit a gap, quit.
         if (metadata == null) {
@@ -87,7 +86,8 @@ class GcsDiffFileLister {
     return true;
   }
 
-  ImmutableList<GcsFileMetadata> listDiffFiles(DateTime fromTime, @Nullable DateTime toTime) {
+  ImmutableList<GcsFileMetadata> listDiffFiles(
+      String gcsBucket, DateTime fromTime, @Nullable DateTime toTime) {
     logger.atInfo().log("Requested restore from time: %s", fromTime);
     if (toTime != null) {
       logger.atInfo().log("  Until time: %s", toTime);
@@ -111,7 +111,8 @@ class GcsDiffFileLister {
       final String filename = listItems.next().getName();
       DateTime upperBoundTime = DateTime.parse(filename.substring(DIFF_FILE_PREFIX.length()));
       if (isInRange(upperBoundTime, fromTime, toTime)) {
-        upperBoundTimesToMetadata.put(upperBoundTime, executor.submit(() -> getMetadata(filename)));
+        upperBoundTimesToMetadata.put(
+            upperBoundTime, executor.submit(() -> getMetadata(gcsBucket, filename)));
         lastUpperBoundTime = latestOf(upperBoundTime, lastUpperBoundTime);
       }
     }
@@ -130,8 +131,9 @@ class GcsDiffFileLister {
     // may be missing files at the end).
     TreeMap<DateTime, GcsFileMetadata> sequence = new TreeMap<>();
     logger.atInfo().log("Restoring until: %s", lastUpperBoundTime);
-    boolean inconsistentFileSet = !constructDiffSequence(
-        upperBoundTimesToMetadata, fromTime, lastUpperBoundTime, sequence);
+    boolean inconsistentFileSet =
+        !constructDiffSequence(
+            gcsBucket, upperBoundTimesToMetadata, fromTime, lastUpperBoundTime, sequence);
 
     // Verify that all of the elements in the original set are represented in the sequence.  If we
     // find anything that's not represented, construct a sequence for it.
@@ -143,7 +145,7 @@ class GcsDiffFileLister {
           break;
         }
         if (!sequence.containsKey(key)) {
-          constructDiffSequence(upperBoundTimesToMetadata, fromTime, key, sequence);
+          constructDiffSequence(gcsBucket, upperBoundTimesToMetadata, fromTime, key, sequence);
           checkForMoreExtraDiffs = true;
           inconsistentFileSet = true;
           break;
@@ -175,7 +177,7 @@ class GcsDiffFileLister {
     return DateTime.parse(metadata.getOptions().getUserMetadata().get(LOWER_BOUND_CHECKPOINT));
   }
 
-  private GcsFileMetadata getMetadata(String filename) {
+  private GcsFileMetadata getMetadata(String gcsBucket, String filename) {
     try {
       return gcsService.getMetadata(new GcsFilename(gcsBucket, filename));
     } catch (IOException e) {

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

@@ -30,6 +30,7 @@ import com.google.appengine.tools.cloudstorage.GcsService;
 import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableList;
 import com.google.common.flogger.FluentLogger;
+import google.registry.config.RegistryConfig.Config;
 import google.registry.model.common.DatabaseMigrationStateSchedule;
 import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState;
 import google.registry.model.common.DatabaseMigrationStateSchedule.ReplayDirection;
@@ -82,6 +83,10 @@ public class ReplayCommitLogsToSqlAction implements Runnable {
   @Inject GcsDiffFileLister diffLister;
   @Inject Clock clock;
 
+  @Inject
+  @Config("commitLogGcsBucket")
+  String gcsBucket;
+
   /** If true, will exit after logging the commit log files that would otherwise be replayed. */
   @Inject
   @Parameter(DRY_RUN_PARAM)
@@ -154,7 +159,7 @@ public class ReplayCommitLogsToSqlAction implements Runnable {
     // If there's an inconsistent file set, this will throw IllegalStateException and the job
     // will try later -- this is likely because an export hasn't finished yet.
     ImmutableList<GcsFileMetadata> commitLogFiles =
-        diffLister.listDiffFiles(fromTime, /* current time */ null);
+        diffLister.listDiffFiles(gcsBucket, fromTime, /* current time */ null);
     logger.atInfo().log("Found %d new commit log files to process.", commitLogFiles.size());
     return commitLogFiles;
   }

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

@@ -26,6 +26,7 @@ import com.google.appengine.api.datastore.EntityTranslator;
 import com.google.appengine.tools.cloudstorage.GcsFileMetadata;
 import com.google.appengine.tools.cloudstorage.GcsService;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Lists;
 import com.google.common.collect.PeekingIterator;
 import com.google.common.collect.Streams;
@@ -33,6 +34,7 @@ import com.google.common.flogger.FluentLogger;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.Result;
 import com.googlecode.objectify.util.ResultNow;
+import google.registry.config.RegistryConfig.Config;
 import google.registry.config.RegistryEnvironment;
 import google.registry.model.ImmutableObject;
 import google.registry.model.ofy.CommitLogBucket;
@@ -50,6 +52,7 @@ import java.nio.channels.Channels;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Stream;
 import javax.inject.Inject;
@@ -72,27 +75,41 @@ public class RestoreCommitLogsAction implements Runnable {
   static final String DRY_RUN_PARAM = "dryRun";
   static final String FROM_TIME_PARAM = "fromTime";
   static final String TO_TIME_PARAM = "toTime";
+  static final String BUCKET_OVERRIDE_PARAM = "gcsBucket";
+
+  private static final ImmutableSet<RegistryEnvironment> FORBIDDEN_ENVIRONMENTS =
+      ImmutableSet.of(RegistryEnvironment.PRODUCTION, RegistryEnvironment.SANDBOX);
 
   @Inject GcsService gcsService;
   @Inject @Parameter(DRY_RUN_PARAM) boolean dryRun;
   @Inject @Parameter(FROM_TIME_PARAM) DateTime fromTime;
   @Inject @Parameter(TO_TIME_PARAM) DateTime toTime;
+
+  @Inject
+  @Parameter(BUCKET_OVERRIDE_PARAM)
+  Optional<String> gcsBucketOverride;
+
   @Inject DatastoreService datastoreService;
   @Inject GcsDiffFileLister diffLister;
+
+  @Inject
+  @Config("commitLogGcsBucket")
+  String defaultGcsBucket;
+
   @Inject Retrier retrier;
   @Inject RestoreCommitLogsAction() {}
 
   @Override
   public void run() {
     checkArgument(
-        RegistryEnvironment.get() == RegistryEnvironment.ALPHA
-            || RegistryEnvironment.get() == RegistryEnvironment.CRASH
-            || RegistryEnvironment.get() == RegistryEnvironment.UNITTEST,
-        "DO NOT RUN ANYWHERE ELSE EXCEPT ALPHA, CRASH OR TESTS.");
+        !FORBIDDEN_ENVIRONMENTS.contains(RegistryEnvironment.get()),
+        "DO NOT RUN IN PRODUCTION OR SANDBOX.");
     if (dryRun) {
       logger.atInfo().log("Running in dryRun mode");
     }
-    List<GcsFileMetadata> diffFiles = diffLister.listDiffFiles(fromTime, toTime);
+    String gcsBucket = gcsBucketOverride.orElse(defaultGcsBucket);
+    logger.atInfo().log("Restoring from %s.", gcsBucket);
+    List<GcsFileMetadata> diffFiles = diffLister.listDiffFiles(gcsBucket, fromTime, toTime);
     if (diffFiles.isEmpty()) {
       logger.atInfo().log("Nothing to restore");
       return;

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

@@ -65,7 +65,6 @@ public class GcsDiffFileListerTest {
   @BeforeEach
   void beforeEach() throws Exception {
     diffLister.gcsService = gcsService;
-    diffLister.gcsBucket = GCS_BUCKET;
     diffLister.executor = newDirectExecutorService();
     for (int i = 0; i < 5; i++) {
       gcsService.createOrReplace(
@@ -87,7 +86,7 @@ public class GcsDiffFileListerTest {
   }
 
   private Iterable<DateTime> listDiffFiles(DateTime fromTime, DateTime toTime) {
-    return extractTimesFromDiffFiles(diffLister.listDiffFiles(fromTime, toTime));
+    return extractTimesFromDiffFiles(diffLister.listDiffFiles(GCS_BUCKET, fromTime, toTime));
   }
 
   private void addGcsFile(int fileAge, int prevAge) throws IOException {

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

@@ -17,7 +17,6 @@ package google.registry.backup;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.util.concurrent.MoreExecutors.newDirectExecutorService;
-import static google.registry.backup.RestoreCommitLogsActionTest.GCS_BUCKET;
 import static google.registry.backup.RestoreCommitLogsActionTest.createCheckpoint;
 import static google.registry.backup.RestoreCommitLogsActionTest.saveDiffFile;
 import static google.registry.backup.RestoreCommitLogsActionTest.saveDiffFileNotToRestore;
@@ -128,9 +127,9 @@ public class ReplayCommitLogsToSqlActionTest {
     action.response = response;
     action.requestStatusChecker = requestStatusChecker;
     action.clock = fakeClock;
+    action.gcsBucket = "gcs bucket";
     action.diffLister = new GcsDiffFileLister();
     action.diffLister.gcsService = gcsService;
-    action.diffLister.gcsBucket = GCS_BUCKET;
     action.diffLister.executor = newDirectExecutorService();
     ofyTm()
         .transact(

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

@@ -55,6 +55,7 @@ import java.nio.ByteBuffer;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
+import java.util.Optional;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -83,9 +84,10 @@ public class RestoreCommitLogsActionTest {
     action.datastoreService = DatastoreServiceFactory.getDatastoreService();
     action.fromTime = now.minusMillis(1);
     action.retrier = new Retrier(new FakeSleeper(new FakeClock()), 1);
+    action.defaultGcsBucket = GCS_BUCKET;
+    action.gcsBucketOverride = Optional.empty();
     action.diffLister = new GcsDiffFileLister();
     action.diffLister.gcsService = gcsService;
-    action.diffLister.gcsBucket = GCS_BUCKET;
     action.diffLister.executor = newDirectExecutorService();
   }