diff --git a/java/google/registry/tools/DeleteAllocationTokensCommand.java b/java/google/registry/tools/DeleteAllocationTokensCommand.java new file mode 100644 index 000000000..aaf02f7f1 --- /dev/null +++ b/java/google/registry/tools/DeleteAllocationTokensCommand.java @@ -0,0 +1,107 @@ +// Copyright 2018 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.collect.ImmutableSet.toImmutableSet; +import static com.google.common.collect.Iterables.partition; +import static com.google.common.collect.Streams.stream; +import static google.registry.model.ofy.ObjectifyService.ofy; + +import com.beust.jcommander.Parameter; +import com.beust.jcommander.Parameters; +import com.google.common.base.Joiner; +import com.google.common.collect.ImmutableSet; +import com.googlecode.objectify.Key; +import com.googlecode.objectify.cmd.Query; +import google.registry.model.domain.token.AllocationToken; +import java.util.List; + +/** + * Command to delete unused {@link AllocationToken}s. + * + *

Allocation tokens that have been redeemed cannot be deleted. To delete a single allocation + * token, specify the entire token as the prefix. + */ +@Parameters( + separators = " =", + commandDescription = "Deletes the unused AllocationTokens with a given prefix.") +final class DeleteAllocationTokensCommand extends ConfirmingCommand + implements CommandWithRemoteApi { + + @Parameter( + names = {"-p", "--prefix"}, + description = "Allocation token prefix; if blank, deletes all unused tokens", + required = true) + private String prefix; + + @Parameter( + names = {"--with_domains"}, + description = "Allow deletion of allocation tokens with specified domains; defaults to false") + boolean withDomains; + + @Parameter( + names = {"--dry_run"}, + description = "Do not actually delete the tokens; defaults to false") + boolean dryRun; + + private static final int BATCH_SIZE = 20; + private static final Joiner JOINER = Joiner.on(", "); + + private ImmutableSet> tokensToDelete; + + @Override + public void init() { + Query query = + ofy().load().type(AllocationToken.class).filter("redemptionHistoryEntry", null); + tokensToDelete = + query.keys().list().stream() + .filter(key -> key.getName().startsWith(prefix)) + .collect(toImmutableSet()); + } + + @Override + protected String prompt() { + return String.format( + "Found %d unused tokens starting with '%s' to delete.", tokensToDelete.size(), prefix); + } + + @Override + protected String execute() { + long numDeleted = + stream(partition(tokensToDelete, BATCH_SIZE)) + .mapToLong(batch -> ofy().transact(() -> deleteBatch(batch))) + .sum(); + return String.format("Deleted %d tokens in total.", numDeleted); + } + + /** Deletes a (filtered) batch of AllocationTokens and returns how many were deleted. */ + private long deleteBatch(List> batch) { + // Load the tokens in the same transaction as they are deleted to verify they weren't redeemed + // since the query ran. This also filters out per-domain tokens if they're not to be deleted. + ImmutableSet tokensToDelete = + ofy().load().keys(batch).values().stream() + .filter(t -> withDomains || !t.getDomainName().isPresent()) + .filter(t -> !t.isRedeemed()) + .collect(toImmutableSet()); + if (!dryRun) { + ofy().delete().entities(tokensToDelete); + } + System.out.printf( + "%s tokens: %s\n", + dryRun ? "Would delete" : "Deleted", + JOINER.join(batch.stream().map(Key::getName).collect(toImmutableSet()))); + return tokensToDelete.size(); + } +} diff --git a/java/google/registry/tools/GenerateAllocationTokensCommand.java b/java/google/registry/tools/GenerateAllocationTokensCommand.java index 2daf39be6..f0cc11927 100644 --- a/java/google/registry/tools/GenerateAllocationTokensCommand.java +++ b/java/google/registry/tools/GenerateAllocationTokensCommand.java @@ -43,13 +43,13 @@ import javax.inject.Inject; import javax.inject.Named; /** Command to generate and persist {@link AllocationToken}s. */ -@NonFinalForTesting @Parameters( separators = " =", commandDescription = "Generates and persists the given number of AllocationTokens, printing each token to stdout." ) -public class GenerateAllocationTokensCommand implements CommandWithRemoteApi { +@NonFinalForTesting +class GenerateAllocationTokensCommand implements CommandWithRemoteApi { @Parameter( names = {"-p", "--prefix"}, diff --git a/java/google/registry/tools/RegistryTool.java b/java/google/registry/tools/RegistryTool.java index a16b59433..ac47090d9 100644 --- a/java/google/registry/tools/RegistryTool.java +++ b/java/google/registry/tools/RegistryTool.java @@ -48,6 +48,7 @@ public final class RegistryTool { .put("create_sandbox_tld", CreateSandboxTldCommand.class) .put("create_tld", CreateTldCommand.class) .put("curl", CurlCommand.class) + .put("delete_allocation_tokens", DeleteAllocationTokensCommand.class) .put("delete_domain", DeleteDomainCommand.class) .put("delete_host", DeleteHostCommand.class) .put("delete_premium_list", DeletePremiumListCommand.class) diff --git a/javatests/google/registry/tools/DeleteAllocationTokensCommandTest.java b/javatests/google/registry/tools/DeleteAllocationTokensCommandTest.java new file mode 100644 index 000000000..0ce485f23 --- /dev/null +++ b/javatests/google/registry/tools/DeleteAllocationTokensCommandTest.java @@ -0,0 +1,140 @@ +// Copyright 2018 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.model.ofy.ObjectifyService.ofy; +import static google.registry.testing.DatastoreHelper.persistResource; +import static google.registry.testing.JUnitBackports.assertThrows; + +import com.beust.jcommander.ParameterException; +import com.googlecode.objectify.Key; +import google.registry.model.domain.token.AllocationToken; +import google.registry.model.reporting.HistoryEntry; +import java.util.Collection; +import javax.annotation.Nullable; +import org.junit.Before; +import org.junit.Test; + +/** Unit tests for {@link DeleteAllocationTokensCommand}. */ +public class DeleteAllocationTokensCommandTest + extends CommandTestCase { + + private AllocationToken preRed1; + private AllocationToken preRed2; + private AllocationToken preNot1; + private AllocationToken preNot2; + private AllocationToken othrRed; + private AllocationToken othrNot; + + @Before + public void init() { + preRed1 = persistToken("prefix12345AA", null, true); + preRed2 = persistToken("prefixgh8907a", null, true); + preNot1 = persistToken("prefix2978204", null, false); + preNot2 = persistToken("prefix8ZZZhs8", null, false); + othrRed = persistToken("h97987sasdfhh", null, true); + othrNot = persistToken("asdgfho7HASDS", null, false); + } + + @Test + public void test_deleteAllUnredeemedTokens_whenEmptyPrefixSpecified() throws Exception { + runCommandForced("--prefix", ""); + assertThat(reloadTokens(preNot1, preNot2, othrNot)).isEmpty(); + assertThat(reloadTokens(preRed1, preRed2, othrRed)).containsExactly(preRed1, preRed2, othrRed); + } + + @Test + public void test_deleteOnlyUnredeemedTokensWithPrefix() throws Exception { + runCommandForced("--prefix", "prefix"); + assertThat(reloadTokens(preNot1, preNot2)).isEmpty(); + assertThat(reloadTokens(preRed1, preRed2, othrRed, othrNot)) + .containsExactly(preRed1, preRed2, othrRed, othrNot); + } + + @Test + public void test_deleteSingleAllocationToken() throws Exception { + runCommandForced("--prefix", "asdgfho7HASDS"); + assertThat(reloadTokens(othrNot)).isEmpty(); + assertThat(reloadTokens(preRed1, preRed2, preNot1, preNot2, othrRed)) + .containsExactly(preRed1, preRed2, preNot1, preNot2, othrRed); + } + + @Test + public void test_deleteTokensWithNonExistentPrefix_doesNothing() throws Exception { + runCommandForced("--prefix", "nonexistent"); + assertThat(reloadTokens(preRed1, preRed2, preNot1, preNot2, othrRed, othrNot)) + .containsExactly(preRed1, preRed2, preNot1, preNot2, othrRed, othrNot); + } + + @Test + public void test_dryRun_deletesNothing() throws Exception { + runCommandForced("--prefix", "", "--dry_run"); + assertThat(reloadTokens(preRed1, preRed2, preNot1, preNot2, othrRed, othrNot)) + .containsExactly(preRed1, preRed2, preNot1, preNot2, othrRed, othrNot); + } + + @Test + public void test_defaultOptions_doesntDeletePerDomainTokens() throws Exception { + AllocationToken preDom1 = persistToken("prefixasdfg897as", "foo.bar", false); + AllocationToken preDom2 = persistToken("prefix98HAZXadbn", "foo.bar", true); + runCommandForced("--prefix", "prefix"); + assertThat(reloadTokens(preNot1, preNot2)).isEmpty(); + assertThat(reloadTokens(preRed1, preRed2, preDom1, preDom2, othrRed, othrNot)) + .containsExactly(preRed1, preRed2, preDom1, preDom2, othrRed, othrNot); + } + + @Test + public void test_withDomains_doesDeletePerDomainTokens() throws Exception { + AllocationToken preDom1 = persistToken("prefixasdfg897as", "foo.bar", false); + AllocationToken preDom2 = persistToken("prefix98HAZXadbn", "foo.bar", true); + runCommandForced("--prefix", "prefix", "--with_domains"); + assertThat(reloadTokens(preNot1, preNot2, preDom1)).isEmpty(); + assertThat(reloadTokens(preRed1, preRed2, preDom2, othrRed, othrNot)) + .containsExactly(preRed1, preRed2, preDom2, othrRed, othrNot); + } + + @Test + public void test_batching() throws Exception { + for (int i = 0; i < 50; i++) { + persistToken(String.format("batch%2d", i), null, i % 2 == 0); + } + assertThat(ofy().load().type(AllocationToken.class).count()).isEqualTo(56); + runCommandForced("--prefix", "batch"); + assertThat(ofy().load().type(AllocationToken.class).count()).isEqualTo(56 - 25); + } + + @Test + public void test_prefixIsRequired() { + ParameterException thrown = assertThrows(ParameterException.class, () -> runCommandForced()); + assertThat(thrown) + .hasMessageThat() + .isEqualTo("The following option is required: -p, --prefix "); + } + + private static AllocationToken persistToken( + String token, @Nullable String domainName, boolean redeemed) { + AllocationToken.Builder builder = + new AllocationToken.Builder().setToken(token).setDomainName(domainName); + if (redeemed) { + builder.setRedemptionHistoryEntry(Key.create(HistoryEntry.class, 1051L)); + } + return persistResource(builder.build()); + } + + private static Collection reloadTokens(AllocationToken ... tokens) { + return ofy().load().entities(tokens).values(); + } +}