Add nomulus command for deleting AllocationTokens

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=214459480

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.
+ *
+ * <p>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<Key<AllocationToken>> tokensToDelete;
+
+  @Override
+  public void init() {
+    Query<AllocationToken> 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<Key<AllocationToken>> 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<AllocationToken> 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();
+  }
+}

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"},

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)

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<DeleteAllocationTokensCommand> {
+
+  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<AllocationToken> reloadTokens(AllocationToken ... tokens) {
+    return ofy().load().entities(tokens).values();
+  }
+}