diff --git a/core/src/main/java/google/registry/backup/BackupModule.java b/core/src/main/java/google/registry/backup/BackupModule.java index d599bdcc5..c005a991a 100644 --- a/core/src/main/java/google/registry/backup/BackupModule.java +++ b/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 provideBucketOverride(HttpServletRequest req) { + return extractOptionalParameter(req, BUCKET_OVERRIDE_PARAM); + } + @Provides @Parameter(FROM_TIME_PARAM) static DateTime provideFromTime(HttpServletRequest req) { diff --git a/core/src/main/java/google/registry/backup/GcsDiffFileLister.java b/core/src/main/java/google/registry/backup/GcsDiffFileLister.java index 56e02fc8c..7fd887a8e 100644 --- a/core/src/main/java/google/registry/backup/GcsDiffFileLister.java +++ b/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> 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 listDiffFiles(DateTime fromTime, @Nullable DateTime toTime) { + ImmutableList 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 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) { diff --git a/core/src/main/java/google/registry/backup/ReplayCommitLogsToSqlAction.java b/core/src/main/java/google/registry/backup/ReplayCommitLogsToSqlAction.java index aa184177c..5709ab07e 100644 --- a/core/src/main/java/google/registry/backup/ReplayCommitLogsToSqlAction.java +++ b/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 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; } diff --git a/core/src/main/java/google/registry/backup/RestoreCommitLogsAction.java b/core/src/main/java/google/registry/backup/RestoreCommitLogsAction.java index 2f93667bd..fdf99f69f 100644 --- a/core/src/main/java/google/registry/backup/RestoreCommitLogsAction.java +++ b/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 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 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 diffFiles = diffLister.listDiffFiles(fromTime, toTime); + String gcsBucket = gcsBucketOverride.orElse(defaultGcsBucket); + logger.atInfo().log("Restoring from %s.", gcsBucket); + List diffFiles = diffLister.listDiffFiles(gcsBucket, fromTime, toTime); if (diffFiles.isEmpty()) { logger.atInfo().log("Nothing to restore"); return; diff --git a/core/src/test/java/google/registry/backup/GcsDiffFileListerTest.java b/core/src/test/java/google/registry/backup/GcsDiffFileListerTest.java index bd9ae256d..7d4dd5de0 100644 --- a/core/src/test/java/google/registry/backup/GcsDiffFileListerTest.java +++ b/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 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 { diff --git a/core/src/test/java/google/registry/backup/ReplayCommitLogsToSqlActionTest.java b/core/src/test/java/google/registry/backup/ReplayCommitLogsToSqlActionTest.java index 064f9c4ca..693f015ac 100644 --- a/core/src/test/java/google/registry/backup/ReplayCommitLogsToSqlActionTest.java +++ b/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( diff --git a/core/src/test/java/google/registry/backup/RestoreCommitLogsActionTest.java b/core/src/test/java/google/registry/backup/RestoreCommitLogsActionTest.java index b7001199d..9fd9c3522 100644 --- a/core/src/test/java/google/registry/backup/RestoreCommitLogsActionTest.java +++ b/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(); }