Add domain name support to AllocationToken entities

The design doc is at []
The next step will be to tie this into the domain create flow, and if the domain
name is on a reserved list, allow it to be created if the token is specified that
has the given domain name on it.

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

java/google/registry/model/domain/token/AllocationToken.java

@@ -21,12 +21,14 @@ import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Entity;
 import com.googlecode.objectify.annotation.Id;
+import com.googlecode.objectify.annotation.Index;
 import google.registry.model.BackupGroupRoot;
 import google.registry.model.Buildable;
 import google.registry.model.CreateAutoTimestamp;
 import google.registry.model.annotations.ReportedOn;
 import google.registry.model.reporting.HistoryEntry;
 import java.util.Optional;
+import javax.annotation.Nullable;
 import org.joda.time.DateTime;
 
 /** An entity representing an allocation token. */
@@ -40,6 +42,9 @@ public class AllocationToken extends BackupGroupRoot implements Buildable {
   /** The key of the history entry for which the token was used. Null if not yet used. */
   Key<HistoryEntry> redemptionHistoryEntry;
 
+  /** The fully-qualified domain name that this token is limited to, if any. */
+  @Nullable @Index String domainName;
+
   /** When this token was created. */
   CreateAutoTimestamp creationTime = CreateAutoTimestamp.create(null);
 
@@ -55,6 +60,10 @@ public class AllocationToken extends BackupGroupRoot implements Buildable {
     return redemptionHistoryEntry != null;
   }
 
+  public Optional<String> getDomainName() {
+    return Optional.ofNullable(domainName);
+  }
+
   public Optional<DateTime> getCreationTime() {
     return Optional.ofNullable(creationTime.getTimestamp());
   }
@@ -86,6 +95,11 @@ public class AllocationToken extends BackupGroupRoot implements Buildable {
       return this;
     }
 
+    public Builder setDomainName(@Nullable String domainName) {
+      getInstance().domainName = domainName;
+      return this;
+    }
+
     public Builder setCreationTime(DateTime creationTime) {
       checkState(
           getInstance().creationTime.getTimestamp() == null, "creationTime can only be set once");

java/google/registry/tools/GenerateAllocationTokensCommand.java

@@ -14,22 +14,31 @@
 
 package google.registry.tools;
 
+import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
+import static com.google.common.collect.Queues.newArrayDeque;
 import static com.google.common.collect.Sets.difference;
 import static google.registry.model.ofy.ObjectifyService.ofy;
+import static java.nio.charset.StandardCharsets.UTF_8;
 
 import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
 import com.google.appengine.tools.remoteapi.RemoteApiException;
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Joiner;
+import com.google.common.base.Splitter;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.io.Files;
 import com.googlecode.objectify.Key;
 import google.registry.model.domain.token.AllocationToken;
 import google.registry.tools.Command.RemoteApiCommand;
 import google.registry.util.NonFinalForTesting;
 import google.registry.util.Retrier;
 import google.registry.util.StringGenerator;
+import java.io.File;
+import java.io.IOException;
 import java.util.Collection;
+import java.util.Deque;
 import javax.inject.Inject;
 import javax.inject.Named;
 
@@ -50,11 +59,16 @@ public class GenerateAllocationTokensCommand implements RemoteApiCommand {
 
   @Parameter(
     names = {"-n", "--number"},
-    description = "The number of tokens to generate",
+    description = "The number of tokens to generate"
-    required = true
   )
   private long numTokens;
 
+  @Parameter(
+      names = {"-d", "--domain_names_file"},
+      description = "A file with a list of newline-delimited domain names to create tokens for"
+  )
+  private String domainNamesFile;
+
   @Parameter(
     names = {"-l", "--length"},
     description = "The length of each token, exclusive of the prefix (if specified); defaults to 12"
@@ -62,7 +76,7 @@ public class GenerateAllocationTokensCommand implements RemoteApiCommand {
   private int tokenLength = 12;
 
   @Parameter(
-      names = {"-d", "--dry_run"},
+      names = {"--dry_run"},
       description = "Do not actually persist the tokens; defaults to false")
   boolean dryRun;
 
@@ -70,16 +84,41 @@ public class GenerateAllocationTokensCommand implements RemoteApiCommand {
   @Inject Retrier retrier;
 
   private static final int BATCH_SIZE = 20;
+  private static final Joiner SKIP_NULLS = Joiner.on(", ").skipNulls();
 
   @Override
-  public void run() {
+  public void run() throws IOException {
+    checkArgument(
+        (numTokens > 0) ^ (domainNamesFile != null),
+        "Must specify either --number or --domain_names_file, but not both");
+
+    Deque<String> domainNames;
+    if (domainNamesFile == null) {
+      domainNames = null;
+    } else {
+      domainNames =
+          newArrayDeque(
+              Splitter.on('\n')
+                  .omitEmptyStrings()
+                  .trimResults()
+                  .split(Files.asCharSource(new File(domainNamesFile), UTF_8).read()));
+      numTokens = domainNames.size();
+    }
+
     int tokensSaved = 0;
     do {
       ImmutableSet<AllocationToken> tokens =
           generateTokens(BATCH_SIZE)
               .stream()
               .limit(numTokens - tokensSaved)
-              .map(t -> new AllocationToken.Builder().setToken(t).build())
+              .map(
+                  t -> {
+                    AllocationToken.Builder token = new AllocationToken.Builder().setToken(t);
+                    if (domainNames != null) {
+                      token.setDomainName(domainNames.removeFirst());
+                    }
+                    return token.build();
+                  })
               .collect(toImmutableSet());
       // Wrap in a retrier to deal with transient 404 errors (thrown as RemoteApiExceptions).
       tokensSaved += retrier.callWithRetry(() -> saveTokens(tokens), RemoteApiException.class);
@@ -90,7 +129,8 @@ public class GenerateAllocationTokensCommand implements RemoteApiCommand {
   int saveTokens(final ImmutableSet<AllocationToken> tokens) {
     Collection<AllocationToken> savedTokens =
         dryRun ? tokens : ofy().transact(() -> ofy().save().entities(tokens).now().values());
-    savedTokens.stream().map(AllocationToken::getToken).forEach(System.out::println);
+    savedTokens.forEach(
+        t -> System.out.println(SKIP_NULLS.join(t.getDomainName().orElse(null), t.getToken())));
     return savedTokens.size();
   }
 

javatests/google/registry/model/domain/token/AllocationTokenTest.java

@@ -36,6 +36,7 @@ public class AllocationTokenTest extends EntityTestCase {
             new AllocationToken.Builder()
                 .setToken("abc123")
                 .setRedemptionHistoryEntry(Key.create(HistoryEntry.class, 1L))
+                .setDomainName("foo.example")
                 .setCreationTime(DateTime.parse("2010-11-12T05:00:00Z"))
                 .build());
     assertThat(ofy().load().entity(token).now()).isEqualTo(token);
@@ -44,11 +45,14 @@ public class AllocationTokenTest extends EntityTestCase {
   @Test
   public void testIndexing() throws Exception {
     verifyIndexing(
-        new AllocationToken.Builder()
+        persistResource(
-            .setToken("abc123")
+            new AllocationToken.Builder()
-            .setCreationTime(DateTime.parse("2010-11-12T05:00:00Z"))
+                .setToken("abc123")
-            .build(),
+                .setDomainName("blahdomain.fake")
-        "token");
+                .setCreationTime(DateTime.parse("2010-11-12T05:00:00Z"))
+                .build()),
+        "token",
+        "domainName");
   }
 
   @Test

javatests/google/registry/model/testdata/schema.txt

@@ -307,6 +307,7 @@ class google.registry.model.domain.token.AllocationToken {
   com.googlecode.objectify.Key<google.registry.model.reporting.HistoryEntry> redemptionHistoryEntry;
   google.registry.model.CreateAutoTimestamp creationTime;
   google.registry.model.UpdateAutoTimestamp updateTimestamp;
+  java.lang.String domainName;
 }
 class google.registry.model.eppcommon.AuthInfo$PasswordAuth {
   java.lang.String repoId;

javatests/google/registry/tools/GenerateAllocationTokensCommandTest.java

@@ -18,12 +18,13 @@ 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 static java.nio.charset.StandardCharsets.UTF_8;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.spy;
 
-import com.beust.jcommander.ParameterException;
 import com.google.appengine.tools.remoteapi.RemoteApiException;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.io.Files;
 import com.googlecode.objectify.Key;
 import google.registry.model.domain.token.AllocationToken;
 import google.registry.model.reporting.HistoryEntry;
@@ -33,6 +34,7 @@ import google.registry.testing.FakeClock;
 import google.registry.testing.FakeSleeper;
 import google.registry.util.Retrier;
 import google.registry.util.StringGenerator.Alphabets;
+import java.io.File;
 import java.util.function.Function;
 import javax.annotation.Nullable;
 import org.joda.time.DateTime;
@@ -54,7 +56,7 @@ public class GenerateAllocationTokensCommandTest
   @Test
   public void testSuccess_oneToken() throws Exception {
     runCommand("--prefix", "blah", "--number", "1", "--length", "9");
-    assertAllocationTokens(createToken("blah123456789", null));
+    assertAllocationTokens(createToken("blah123456789", null, null));
     assertInStdout("blah123456789");
   }
 
@@ -62,16 +64,16 @@ public class GenerateAllocationTokensCommandTest
   public void testSuccess_threeTokens() throws Exception {
     runCommand("--prefix", "foo", "--number", "3", "--length", "10");
     assertAllocationTokens(
-        createToken("foo123456789A", null),
+        createToken("foo123456789A", null, null),
-        createToken("fooBCDEFGHJKL", null),
+        createToken("fooBCDEFGHJKL", null, null),
-        createToken("fooMNPQRSTUVW", null));
+        createToken("fooMNPQRSTUVW", null, null));
     assertInStdout("foo123456789A\nfooBCDEFGHJKL\nfooMNPQRSTUVW");
   }
 
   @Test
   public void testSuccess_defaults() throws Exception {
     runCommand("--number", "1");
-    assertAllocationTokens(createToken("123456789ABC", null));
+    assertAllocationTokens(createToken("123456789ABC", null, null));
     assertInStdout("123456789ABC");
   }
 
@@ -85,7 +87,7 @@ public class GenerateAllocationTokensCommandTest
         .when(spyCommand)
         .saveTokens(Mockito.any());
     runCommand("--number", "1");
-    assertAllocationTokens(createToken("123456789ABC", null));
+    assertAllocationTokens(createToken("123456789ABC", null, null));
     assertInStdout("123456789ABC");
   }
 
@@ -94,7 +96,7 @@ public class GenerateAllocationTokensCommandTest
     AllocationToken existingToken =
         persistResource(new AllocationToken.Builder().setToken("DEADBEEF123456789ABC").build());
     runCommand("--number", "1", "--prefix", "DEADBEEF");
-    assertAllocationTokens(existingToken, createToken("DEADBEEFDEFGHJKLMNPQ", null));
+    assertAllocationTokens(existingToken, createToken("DEADBEEFDEFGHJKLMNPQ", null, null));
     assertInStdout("DEADBEEFDEFGHJKLMNPQ");
   }
 
@@ -116,10 +118,39 @@ public class GenerateAllocationTokensCommandTest
   }
 
   @Test
-  public void testFailure_mustSpecifyNumberOfTokens() {
+  public void testSuccess_domainNames() throws Exception {
-    ParameterException thrown =
+    File domainNamesFile = tmpDir.newFile("domain_names.txt");
-        assertThrows(ParameterException.class, () -> runCommand("--prefix", "FEET"));
+    Files.asCharSink(domainNamesFile, UTF_8).write("foo1.tld\nboo2.tld\nbaz9.tld\n");
-    assertThat(thrown).hasMessageThat().contains("The following option is required: -n, --number");
+    runCommand("--domain_names_file", domainNamesFile.getPath());
+    assertAllocationTokens(
+        createToken("123456789ABC", null, "foo1.tld"),
+        createToken("DEFGHJKLMNPQ", null, "boo2.tld"),
+        createToken("RSTUVWXYZabc", null, "baz9.tld"));
+    assertInStdout("foo1.tld, 123456789ABC\nboo2.tld, DEFGHJKLMNPQ\nbaz9.tld, RSTUVWXYZabc");
+  }
+
+  @Test
+  public void testFailure_mustSpecifyNumberOfTokensOrDomainsFile() {
+    IllegalArgumentException thrown =
+        assertThrows(IllegalArgumentException.class, () -> runCommand("--prefix", "FEET"));
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo("Must specify either --number or --domain_names_file, but not both");
+  }
+
+  @Test
+  public void testFailure_mustNotSpecifyBothNumberOfTokensAndDomainsFile() {
+    IllegalArgumentException thrown =
+        assertThrows(
+            IllegalArgumentException.class,
+            () ->
+                runCommand(
+                    "--prefix", "FEET",
+                    "--number", "999",
+                    "--domain_names_file", "/path/to/blaaaaah"));
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo("Must specify either --number or --domain_names_file, but not both");
   }
 
   private void assertAllocationTokens(AllocationToken... expectedTokens) {
@@ -142,11 +173,14 @@ public class GenerateAllocationTokensCommandTest
   }
 
   private AllocationToken createToken(
-      String token, @Nullable Key<HistoryEntry> redemptionHistoryEntry) {
+      String token,
+      @Nullable Key<HistoryEntry> redemptionHistoryEntry,
+      @Nullable String domainName) {
     AllocationToken.Builder builder = new AllocationToken.Builder().setToken(token);
     if (redemptionHistoryEntry != null) {
       builder.setRedemptionHistoryEntry(redemptionHistoryEntry);
     }
+    builder.setDomainName(domainName);
     return builder.build();
   }
 }