From 4e71421c818439bb326f24dee97f835eb1d1975f Mon Sep 17 00:00:00 2001 From: weiminyu Date: Tue, 8 Jan 2019 14:28:29 -0800 Subject: [PATCH] Support datastore restore in Nomulus tool Two commands are being added: - ImportDatastoreCommand starts an async import operation. User may choose to wait until import completes or quit immediately. - GetOperationStatusCommand checks the status of an operation. It may be used to check the status of an operation started by ImportDatastoreCommand. Both commands communicate with Datastore admin api directly, without going through the Registry server. ------------- Created by MOE: https://github.com/google/moe MOE_MIGRATED_REVID=228400059 --- .../registry/export/CheckBackupAction.java | 3 + .../export/datastore/DatastoreAdmin.java | 56 ++++++- .../export/datastore/EntityFilter.java | 2 - .../registry/export/datastore/Operation.java | 38 ++++- java/google/registry/tools/BUILD | 1 + .../tools/GetOperationStatusCommand.java | 47 ++++++ .../tools/ImportDatastoreCommand.java | 127 +++++++++++++++ java/google/registry/tools/RegistryTool.java | 2 + .../registry/tools/RegistryToolComponent.java | 4 + .../export/datastore/EntityFilterTest.java | 5 - javatests/google/registry/tools/BUILD | 1 + .../tools/GetOperationStatusCommandTest.java | 59 +++++++ .../tools/ImportDatastoreCommandTest.java | 146 ++++++++++++++++++ 13 files changed, 477 insertions(+), 14 deletions(-) create mode 100644 java/google/registry/tools/GetOperationStatusCommand.java create mode 100644 java/google/registry/tools/ImportDatastoreCommand.java create mode 100644 javatests/google/registry/tools/GetOperationStatusCommandTest.java create mode 100644 javatests/google/registry/tools/ImportDatastoreCommandTest.java diff --git a/java/google/registry/export/CheckBackupAction.java b/java/google/registry/export/CheckBackupAction.java index 3ef56da79..1247b1a2b 100644 --- a/java/google/registry/export/CheckBackupAction.java +++ b/java/google/registry/export/CheckBackupAction.java @@ -14,6 +14,7 @@ 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; @@ -131,6 +132,8 @@ public class CheckBackupAction implements Runnable { Set kindsToLoad = ImmutableSet.copyOf(Splitter.on(',').split(kindsToLoadParam)); Operation backup = getExportStatus(); + checkArgument(backup.isExport(), "Expecting an export operation: [%s].", backupName); + if (backup.isProcessing() && backup.getRunningTime(clock).isShorterThan(MAXIMUM_BACKUP_RUNNING_TIME)) { // Backup might still be running, so send a 304 to have the task retry. diff --git a/java/google/registry/export/datastore/DatastoreAdmin.java b/java/google/registry/export/datastore/DatastoreAdmin.java index 4e21495f6..ab1d49f30 100644 --- a/java/google/registry/export/datastore/DatastoreAdmin.java +++ b/java/google/registry/export/datastore/DatastoreAdmin.java @@ -25,6 +25,7 @@ import com.google.api.client.json.GenericJson; import com.google.api.client.json.JsonFactory; import com.google.api.client.util.Key; import com.google.common.base.Strings; +import com.google.common.collect.ImmutableList; import java.util.Collection; import java.util.Optional; @@ -73,12 +74,32 @@ public class DatastoreAdmin extends AbstractGoogleJsonClient { * * * @param outputUrlPrefix the full resource URL of the external storage location - * @param kinds the datastore 'kinds' to be exported + * @param kinds the datastore 'kinds' to be exported. If empty, all kinds will be exported */ public Export export(String outputUrlPrefix, Collection kinds) { return new Export(new ExportRequest(outputUrlPrefix, kinds)); } + /** + * Imports the entire backup specified by {@code backupUrl} back to Cloud Datastore. + * + *

A successful backup restores deleted entities and reverts updates to existing entities since + * the backup time. However, it does not affect newly added entities. + */ + public Import importBackup(String backupUrl) { + return new Import(new ImportRequest(backupUrl, ImmutableList.of())); + } + + /** + * Imports the backup specified by {@code backupUrl} back to Cloud Datastore. Only entities whose + * types are included in {@code kinds} are imported. + * + * @see #importBackup(String) + */ + public Import importBackup(String backupUrl, Collection kinds) { + return new Import(new ImportRequest(backupUrl, kinds)); + } + /** * Returns a {@link Get} request that retrieves the details of an export or import {@link * Operation}. @@ -158,6 +179,20 @@ public class DatastoreAdmin extends AbstractGoogleJsonClient { } } + /** A request to restore an backup to a Cloud Datastore database. */ + public class Import extends DatastoreAdminRequest { + + Import(ImportRequest importRequest) { + super( + DatastoreAdmin.this, + "POST", + "projects/{projectId}:import", + importRequest, + Operation.class); + set("projectId", projectId); + } + } + /** A request to retrieve details of an export or import operation. */ public class Get extends DatastoreAdminRequest { @@ -216,8 +251,27 @@ public class DatastoreAdmin extends AbstractGoogleJsonClient { ExportRequest(String outputUrlPrefix, Collection kinds) { checkNotNull(outputUrlPrefix, "outputUrlPrefix"); + checkArgument(!kinds.isEmpty(), "kinds must not be empty"); this.outputUrlPrefix = outputUrlPrefix; this.entityFilter = new EntityFilter(kinds); } } + + /** + * Model object that describes the JSON content in an export request. + * + *

Please note that some properties defined in the API are excluded, e.g., {@code databaseId} + * (not supported by Cloud Datastore) and labels (not used by Domain Registry). + */ + @SuppressWarnings("unused") + static class ImportRequest extends GenericJson { + @Key private final String inputUrl; + @Key private final EntityFilter entityFilter; + + ImportRequest(String inputUrl, Collection kinds) { + checkNotNull(inputUrl, "outputUrlPrefix"); + this.inputUrl = inputUrl; + this.entityFilter = new EntityFilter(kinds); + } + } } diff --git a/java/google/registry/export/datastore/EntityFilter.java b/java/google/registry/export/datastore/EntityFilter.java index ce261b49d..d4d530236 100644 --- a/java/google/registry/export/datastore/EntityFilter.java +++ b/java/google/registry/export/datastore/EntityFilter.java @@ -14,7 +14,6 @@ package google.registry.export.datastore; -import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Preconditions.checkNotNull; import com.google.api.client.json.GenericJson; @@ -39,7 +38,6 @@ public class EntityFilter extends GenericJson { EntityFilter(Collection kinds) { checkNotNull(kinds, "kinds"); - checkArgument(!kinds.isEmpty(), "kinds must not be empty"); this.kinds = ImmutableList.copyOf(kinds); } diff --git a/java/google/registry/export/datastore/Operation.java b/java/google/registry/export/datastore/Operation.java index 871cee365..0c27d3a88 100644 --- a/java/google/registry/export/datastore/Operation.java +++ b/java/google/registry/export/datastore/Operation.java @@ -15,10 +15,10 @@ package google.registry.export.datastore; import static com.google.common.base.Preconditions.checkState; +import static com.google.common.base.Strings.isNullOrEmpty; import com.google.api.client.json.GenericJson; import com.google.api.client.util.Key; -import com.google.common.base.Strings; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; import google.registry.export.datastore.DatastoreAdmin.Get; @@ -52,6 +52,14 @@ public class Operation extends GenericJson { return name; } + public boolean isExport() { + return !isNullOrEmpty(getExportFolderUrl()); + } + + public boolean isImport() { + return !isNullOrEmpty(getMetadata().getInputUrl()); + } + public boolean isDone() { return done; } @@ -88,7 +96,9 @@ public class Operation extends GenericJson { /** * Returns the URL to the GCS folder that holds the exported data. This folder is created by * Datastore and is under the {@code outputUrlPrefix} set to {@linkplain - * DatastoreAdmin#export(String, List) the export request}. + * DatastoreAdmin#export(String, java.util.Collection) the export request}. + * + * @throws IllegalStateException if this is not an export operation */ public String getExportFolderUrl() { return getMetadata().getOutputUrlPrefix(); @@ -98,6 +108,8 @@ public class Operation extends GenericJson { * Returns the last segment of the {@linkplain #getExportFolderUrl() export folder URL} which can * be used as unique identifier of this export operation. This is a better ID than the {@linkplain * #getName() operation name}, which is opaque. + * + * @throws IllegalStateException if this is not an export operation */ public String getExportId() { String exportFolderUrl = getExportFolderUrl(); @@ -138,12 +150,12 @@ public class Operation extends GenericJson { public CommonMetadata() {} String getOperationType() { - checkState(!Strings.isNullOrEmpty(operationType), "operationType may not be null or empty"); + checkState(!isNullOrEmpty(operationType), "operationType may not be null or empty"); return operationType; } String getState() { - checkState(!Strings.isNullOrEmpty(state), "state may not be null or empty"); + checkState(!isNullOrEmpty(state), "state may not be null or empty"); return state; } @@ -165,6 +177,7 @@ public class Operation extends GenericJson { @Key private Progress progressEntities; @Key private Progress progressBytes; @Key private EntityFilter entityFilter; + @Key private String inputUrl; @Key private String outputUrlPrefix; public Metadata() {} @@ -186,9 +199,22 @@ public class Operation extends GenericJson { return entityFilter; } + public String getInputUrl() { + return checkUrls().inputUrl; + } + public String getOutputUrlPrefix() { - checkState(!Strings.isNullOrEmpty(outputUrlPrefix), "outputUrlPrefix"); - return outputUrlPrefix; + return checkUrls().outputUrlPrefix; + } + + Metadata checkUrls() { + checkState( + isNullOrEmpty(inputUrl) || isNullOrEmpty(outputUrlPrefix), + "inputUrl and outputUrlPrefix must not be both present"); + checkState( + !isNullOrEmpty(inputUrl) || !isNullOrEmpty(outputUrlPrefix), + "inputUrl and outputUrlPrefix must not be both missing"); + return this; } } diff --git a/java/google/registry/tools/BUILD b/java/google/registry/tools/BUILD index 2535fd9b3..4f91842ba 100644 --- a/java/google/registry/tools/BUILD +++ b/java/google/registry/tools/BUILD @@ -45,6 +45,7 @@ java_library( "//java/google/registry/dns/writer/clouddns", "//java/google/registry/dns/writer/dnsupdate", "//java/google/registry/export", + "//java/google/registry/export/datastore", "//java/google/registry/flows", "//java/google/registry/gcs", "//java/google/registry/keyring", diff --git a/java/google/registry/tools/GetOperationStatusCommand.java b/java/google/registry/tools/GetOperationStatusCommand.java new file mode 100644 index 000000000..9a7c9ee91 --- /dev/null +++ b/java/google/registry/tools/GetOperationStatusCommand.java @@ -0,0 +1,47 @@ +// Copyright 2019 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.tools; + +import static com.google.common.base.Preconditions.checkArgument; + +import com.beust.jcommander.Parameter; +import com.beust.jcommander.Parameters; +import com.google.api.client.json.JsonFactory; +import com.google.api.client.json.jackson2.JacksonFactory; +import com.google.common.base.Strings; +import google.registry.export.datastore.DatastoreAdmin; +import java.util.List; +import javax.inject.Inject; + +/** Command to get the status of a Datastore operation, e.g., an import or export. */ +@Parameters(separators = " =", commandDescription = "Get status of a Datastore operation.") +public class GetOperationStatusCommand implements Command { + + private static final JsonFactory JSON_FACTORY = JacksonFactory.getDefaultInstance(); + + @Parameter(description = "Name of the Datastore import or export operation.") + private List mainParameters; + + @Inject DatastoreAdmin datastoreAdmin; + + @Override + public void run() throws Exception { + checkArgument( + mainParameters.size() == 1, "Requires exactly one argument: the name of the operation."); + String operationName = mainParameters.get(0); + checkArgument(!Strings.isNullOrEmpty(operationName), "Missing operation name."); + System.out.println(JSON_FACTORY.toPrettyString(datastoreAdmin.get(operationName).execute())); + } +} diff --git a/java/google/registry/tools/ImportDatastoreCommand.java b/java/google/registry/tools/ImportDatastoreCommand.java new file mode 100644 index 000000000..ca4c537d7 --- /dev/null +++ b/java/google/registry/tools/ImportDatastoreCommand.java @@ -0,0 +1,127 @@ +// Copyright 2019 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.tools; + +import static com.google.common.base.Preconditions.checkArgument; + +import com.beust.jcommander.Parameter; +import com.beust.jcommander.Parameters; +import com.google.common.base.Ascii; +import com.google.common.collect.ImmutableList; +import google.registry.export.datastore.DatastoreAdmin; +import google.registry.export.datastore.Operation; +import java.util.List; +import java.util.concurrent.TimeUnit; +import javax.annotation.Nullable; +import javax.inject.Inject; +import org.joda.time.Duration; + +/** + * Command that imports an earlier backup into Datastore. + * + *

This command is part of the Datastore restore process. Please refer to the playbook for + * the entire process. + */ +@Parameters(separators = " =", commandDescription = "Imports a backup of the Datastore.") +public class ImportDatastoreCommand extends ConfirmingCommand { + + @Parameter(names = "--backup_url", description = "URL to the backup on GCS to be imported.") + private String backupUrl; + + @Nullable + @Parameter( + names = "--kinds", + description = "List of entity kinds to be imported. Default is to import all.") + private List kinds = ImmutableList.of(); + + @Parameter( + names = "--async", + description = "If true, command will launch import operation and quit.") + private boolean async; + + @Parameter( + names = "--poll_interval", + description = + "Polling interval while waiting for completion synchronously. " + + "Value is in ISO-8601 format, e.g., PT10S for 10 seconds.") + private Duration pollingInterval = Duration.standardSeconds(30); + + @Parameter( + names = "--confirm_production_import", + description = "Set this option to 'PRODUCTION' to confirm import in production environment.") + private String confirmProductionImport = ""; + + @Inject DatastoreAdmin datastoreAdmin; + + @Override + protected String execute() throws Exception { + RegistryToolEnvironment currentEnvironment = RegistryToolEnvironment.get(); + + // Extra confirmation for running in production + checkArgument( + !currentEnvironment.equals(RegistryToolEnvironment.PRODUCTION) + || confirmProductionImport.equals("PRODUCTION"), + "The confirm_production_import option must be set when restoring production environment."); + + Operation importOperation = datastoreAdmin.importBackup(backupUrl, kinds).execute(); + + String statusCommand = + String.format( + "nomulus -e %s get_operation_status %s", + Ascii.toLowerCase(currentEnvironment.name()), importOperation.getName()); + + if (async) { + return String.format( + "Datastore import started. Run this command to check its progress:\n%s", + statusCommand); + } + + System.out.println( + "Waiting for import to complete.\n" + + "You may press Ctrl-C at any time, and use this command to check progress:\n" + + statusCommand); + while (importOperation.isProcessing()) { + waitInteractively(pollingInterval); + + importOperation = datastoreAdmin.get(importOperation.getName()).execute(); + + System.out.printf("\n%s\n", importOperation.getProgress()); + } + return String.format( + "\nDatastore import %s %s.", + importOperation.getName(), importOperation.isSuccessful() ? "succeeded" : "failed"); + } + + @Override + protected String prompt() { + return "\nThis command is an intermediate step in the Datastore restore process.\n\n" + + "Please read and understand the playbook entry at\n" + + " http://playbooks/domain_registry/procedures/backup-restore-testing.md\n" + + "before proceeding.\n"; + } + + /** Prints dots to console at regular interval while waiting. */ + private static void waitInteractively(Duration pollingInterval) throws InterruptedException { + int sleepSeconds = 2; + long iterations = (pollingInterval.getStandardSeconds() + sleepSeconds - 1) / sleepSeconds; + + for (int i = 0; i < iterations; i++) { + TimeUnit.SECONDS.sleep(sleepSeconds); + System.out.print('.'); + System.out.flush(); + } + } +} diff --git a/java/google/registry/tools/RegistryTool.java b/java/google/registry/tools/RegistryTool.java index 28acc36c9..f2e42c8d0 100644 --- a/java/google/registry/tools/RegistryTool.java +++ b/java/google/registry/tools/RegistryTool.java @@ -75,6 +75,7 @@ public final class RegistryTool { .put("get_history_entries", GetHistoryEntriesCommand.class) .put("get_host", GetHostCommand.class) .put("get_keyring_secret", GetKeyringSecretCommand.class) + .put("get_operation_status", GetOperationStatusCommand.class) .put("get_registrar", GetRegistrarCommand.class) .put("get_resource_by_key", GetResourceByKeyCommand.class) .put("get_routing_map", GetRoutingMapCommand.class) @@ -83,6 +84,7 @@ public final class RegistryTool { .put("get_tld", GetTldCommand.class) .put("ghostryde", GhostrydeCommand.class) .put("hash_certificate", HashCertificateCommand.class) + .put("import_datastore", ImportDatastoreCommand.class) .put("list_cursors", ListCursorsCommand.class) .put("list_domains", ListDomainsCommand.class) .put("list_hosts", ListHostsCommand.class) diff --git a/java/google/registry/tools/RegistryToolComponent.java b/java/google/registry/tools/RegistryToolComponent.java index 606cd36b9..a5a1e24f0 100644 --- a/java/google/registry/tools/RegistryToolComponent.java +++ b/java/google/registry/tools/RegistryToolComponent.java @@ -22,6 +22,7 @@ import google.registry.config.RegistryConfig.ConfigModule; import google.registry.dns.writer.VoidDnsWriterModule; import google.registry.dns.writer.clouddns.CloudDnsWriterModule; import google.registry.dns.writer.dnsupdate.DnsUpdateWriterModule; +import google.registry.export.datastore.DatastoreAdminModule; import google.registry.keyring.KeyringModule; import google.registry.keyring.api.DummyKeyringModule; import google.registry.keyring.api.KeyModule; @@ -56,6 +57,7 @@ import javax.inject.Singleton; BigqueryModule.class, ConfigModule.class, CloudDnsWriterModule.class, + DatastoreAdminModule.class, DatastoreServiceModule.class, DummyKeyringModule.class, DnsUpdateWriterModule.class, @@ -92,7 +94,9 @@ interface RegistryToolComponent { void inject(GenerateDnsReportCommand command); void inject(GenerateEscrowDepositCommand command); void inject(GetKeyringSecretCommand command); + void inject(GetOperationStatusCommand command); void inject(GhostrydeCommand command); + void inject(ImportDatastoreCommand command); void inject(ListCursorsCommand command); void inject(LoadSnapshotCommand command); void inject(LockDomainCommand command); diff --git a/javatests/google/registry/export/datastore/EntityFilterTest.java b/javatests/google/registry/export/datastore/EntityFilterTest.java index ba56ea9c9..33dc74f68 100644 --- a/javatests/google/registry/export/datastore/EntityFilterTest.java +++ b/javatests/google/registry/export/datastore/EntityFilterTest.java @@ -37,11 +37,6 @@ public class EntityFilterTest { assertThrows(NullPointerException.class, () -> new EntityFilter(null)); } - @Test - public void testEntityFilter_create_emptyKinds() { - assertThrows(IllegalArgumentException.class, () -> new EntityFilter(ImmutableList.of())); - } - @Test public void testEntityFilter_marshall() throws IOException { EntityFilter entityFilter = diff --git a/javatests/google/registry/tools/BUILD b/javatests/google/registry/tools/BUILD index c80ce057d..90bb08b59 100644 --- a/javatests/google/registry/tools/BUILD +++ b/javatests/google/registry/tools/BUILD @@ -19,6 +19,7 @@ java_library( deps = [ "//java/google/registry/backup", "//java/google/registry/config", + "//java/google/registry/export/datastore", "//java/google/registry/flows", "//java/google/registry/keyring/api", "//java/google/registry/model", diff --git a/javatests/google/registry/tools/GetOperationStatusCommandTest.java b/javatests/google/registry/tools/GetOperationStatusCommandTest.java new file mode 100644 index 000000000..4fc08aa78 --- /dev/null +++ b/javatests/google/registry/tools/GetOperationStatusCommandTest.java @@ -0,0 +1,59 @@ +// Copyright 2019 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.tools; + +import static com.google.common.truth.Truth.assertThat; +import static google.registry.testing.JUnitBackports.assertThrows; +import static org.mockito.Mockito.when; + +import google.registry.export.datastore.DatastoreAdmin; +import google.registry.export.datastore.DatastoreAdmin.Get; +import google.registry.export.datastore.Operation; +import java.io.IOException; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.mockito.ArgumentCaptor; +import org.mockito.Captor; +import org.mockito.Mock; + +/** Unit tests for {@link GetOperationStatusCommand}. */ +@RunWith(JUnit4.class) +public class GetOperationStatusCommandTest extends CommandTestCase { + + @Mock private DatastoreAdmin datastoreAdmin; + @Mock private Get getRequest; + @Captor ArgumentCaptor operationName; + + @Before + public void setup() throws IOException { + command.datastoreAdmin = datastoreAdmin; + + when(datastoreAdmin.get(operationName.capture())).thenReturn(getRequest); + when(getRequest.execute()).thenReturn(new Operation()); + } + + @Test + public void test_success() throws Exception { + runCommand("projects/project-id/operations/HASH"); + assertThat(operationName.getValue()).isEqualTo("projects/project-id/operations/HASH"); + } + + @Test + public void test_failure_tooManyNames() { + assertThrows(IllegalArgumentException.class, () -> runCommand("a", "b")); + } +} diff --git a/javatests/google/registry/tools/ImportDatastoreCommandTest.java b/javatests/google/registry/tools/ImportDatastoreCommandTest.java new file mode 100644 index 000000000..5f6d1bc7b --- /dev/null +++ b/javatests/google/registry/tools/ImportDatastoreCommandTest.java @@ -0,0 +1,146 @@ +// Copyright 2019 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.tools; + +import static com.google.common.truth.Truth.assertThat; +import static google.registry.testing.JUnitBackports.assertThrows; +import static org.mockito.Matchers.anyString; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import google.registry.export.datastore.DatastoreAdmin; +import google.registry.export.datastore.DatastoreAdmin.Get; +import google.registry.export.datastore.DatastoreAdmin.Import; +import google.registry.export.datastore.Operation; +import java.util.Collection; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.mockito.ArgumentCaptor; +import org.mockito.Captor; +import org.mockito.Mock; + +/** Unit tests for {@link ImportDatastoreCommand}. */ +@RunWith(JUnit4.class) +public class ImportDatastoreCommandTest extends CommandTestCase { + + @Captor ArgumentCaptor backupUrl; + @Captor ArgumentCaptor> kinds; + @Captor ArgumentCaptor operationName; + + @Mock private DatastoreAdmin datastoreAdmin; + @Mock private Import importRequest; + @Mock private Get getRequest; + @Mock private Operation importOperation; + @Mock private Operation getOperation; + + @Before + public void setup() throws Exception { + command.datastoreAdmin = datastoreAdmin; + + when(datastoreAdmin.importBackup(backupUrl.capture(), kinds.capture())) + .thenReturn(importRequest); + when(importRequest.execute()).thenReturn(importOperation); + when(importOperation.getName()).thenReturn("opName"); + + when(datastoreAdmin.get(operationName.capture())).thenReturn(getRequest); + when(getRequest.execute()).thenReturn(getOperation); + } + + @Test + public void test_importAllKinds_immediateSuccess() throws Exception { + runCommandForced( + "--poll_interval", "PT0.001S", + "--backup_url", "gs://bucket/export-id/export-id.overall_export_metadata"); + assertThat(backupUrl.getValue()) + .isEqualTo("gs://bucket/export-id/export-id.overall_export_metadata"); + assertThat(kinds.getValue()).isEmpty(); + verify(datastoreAdmin, never()).get(anyString()); + } + + @Test + public void test_importSomeKinds_immediateSuccess() throws Exception { + runCommandForced( + "--poll_interval", + "PT0.001S", + "--backup_url", + "gs://bucket/export-id/export-id.overall_export_metadata", + "--kinds", + "Registrar", + "--kinds", + "Registry"); + assertThat(backupUrl.getValue()) + .isEqualTo("gs://bucket/export-id/export-id.overall_export_metadata"); + assertThat(kinds.getValue()).containsExactly("Registrar", "Registry"); + verify(datastoreAdmin, never()).get(anyString()); + } + + @Test + public void test_delayedSuccess_sync() throws Exception { + when(importOperation.isProcessing()).thenReturn(true); + when(getOperation.isProcessing()).thenReturn(true).thenReturn(false); + when(getOperation.isSuccessful()).thenReturn(true); + runCommandForced( + "--poll_interval", "PT0.001S", + "--backup_url", "gs://bucket/export-id/export-id.overall_export_metadata"); + + verify(datastoreAdmin, times(1)).get("opName"); + } + + @Test + public void test_delayedSuccess_async() throws Exception { + when(importOperation.isProcessing()).thenReturn(false); + when(getOperation.isProcessing()).thenReturn(true).thenReturn(false); + when(getOperation.isSuccessful()).thenReturn(true); + runCommandForced( + "--poll_interval", + "PT0.001S", + "--backup_url", + "gs://bucket/export-id/export-id.overall_export_metadata", + "--async"); + + verify(datastoreAdmin, never()).get("opName"); + } + + @Test + public void test_failure_notAllowedInProduction() { + assertThrows( + IllegalArgumentException.class, + () -> + runCommandInEnvironment( + RegistryToolEnvironment.PRODUCTION, + "--force", + "--poll_interval", + "PT0.001S", + "--backup_url", + "gs://bucket/export-id/export-id.overall_export_metadata")); + } + + @Test + public void test_success_runInProduction() throws Exception { + runCommandInEnvironment( + RegistryToolEnvironment.PRODUCTION, + "--force", + "--confirm_production_import", + "PRODUCTION", + "--poll_interval", + "PT0.001S", + "--backup_url", + "gs://bucket/export-id/export-id.overall_export_metadata"); + } +}