Add a command to update allocation tokens

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

java/google/registry/tools/RegistryTool.java

@@ -107,6 +107,7 @@ public final class RegistryTool {
           .put("uniform_rapid_suspension", UniformRapidSuspensionCommand.class)
           .put("unlock_domain", UnlockDomainCommand.class)
           .put("unrenew_domain", UnrenewDomainCommand.class)
+          .put("update_allocation_tokens", UpdateAllocationTokensCommand.class)
           .put("update_cursors", UpdateCursorsCommand.class)
           .put("update_domain", UpdateDomainCommand.class)
           .put("update_kms_keyring", UpdateKmsKeyringCommand.class)

java/google/registry/tools/UpdateAllocationTokensCommand.java

@@ -0,0 +1,174 @@
+// 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 static com.google.common.collect.ImmutableList.toImmutableList;
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
+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.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableSortedMap;
+import com.googlecode.objectify.Key;
+import google.registry.model.domain.token.AllocationToken;
+import google.registry.model.domain.token.AllocationToken.TokenStatus;
+import google.registry.tools.params.TransitionListParameter.TokenStatusTransitions;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.function.Function;
+import org.joda.time.DateTime;
+
+/** Command to update existing {@link AllocationToken}s. */
+@Parameters(
+    separators = " =",
+    commandDescription =
+        "Updates AllocationTokens with the given prefix, modifying some subset of their allowed"
+            + " client IDs, allowed TLDs, discount fraction, or status transitions")
+public final class UpdateAllocationTokensCommand extends ConfirmingCommand
+    implements CommandWithRemoteApi {
+
+  @Parameter(
+      names = {"-p", "--prefix"},
+      description =
+          "Update all allocation tokens with this prefix; otherwise use '--tokens' to specify"
+              + " exact tokens(s) to update")
+  private String prefix;
+
+  @Parameter(
+      names = {"--tokens"},
+      description =
+          "Comma-separated list of tokens to update; otherwise use '--prefix' to update all tokens"
+              + " with a given prefix")
+  private List<String> tokens;
+
+  @Parameter(
+      names = {"--dry_run"},
+      description = "Do not actually update the tokens; defaults to false")
+  private boolean dryRun;
+
+  @Parameter(
+      names = {"--allowed_client_ids"},
+      description =
+          "Comma-separated list of allowed client IDs. Use the empty string to clear the"
+              + " existing list.")
+  private List<String> allowedClientIds;
+
+  @Parameter(
+      names = {"--allowed_tlds"},
+      description =
+          "Comma-separated list of allowed TLDs. Use the empty string to clear the"
+              + " existing list.")
+  private List<String> allowedTlds;
+
+  @Parameter(
+      names = {"--discount_fraction"},
+      description =
+          "A discount off the base price for the first year between 0.0 and 1.0. Default is 0.0,"
+              + " i.e. no discount.")
+  private Double discountFraction;
+
+  @Parameter(
+      names = "--token_status_transitions",
+      converter = TokenStatusTransitions.class,
+      validateWith = TokenStatusTransitions.class,
+      description =
+          "Comma-delimited list of token status transitions effective on specific dates, of the"
+              + " form <time>=<status>[,<time>=<status>]* where each status represents the status.")
+  private ImmutableSortedMap<DateTime, TokenStatus> tokenStatusTransitions;
+
+  private static final int BATCH_SIZE = 20;
+  private static final Joiner JOINER = Joiner.on(", ");
+
+  private ImmutableSet<AllocationToken> tokensToSave;
+
+  @Override
+  public void init() {
+    // A single entry with the empty string means that the user passed an empty argument to the
+    // lists, so we should clear them
+    if (ImmutableList.of("").equals(allowedClientIds)) {
+      allowedClientIds = ImmutableList.of();
+    }
+    if (ImmutableList.of("").equals(allowedTlds)) {
+      allowedTlds = ImmutableList.of();
+    }
+
+    ImmutableList<Key<AllocationToken>> keysToUpdate = getKeysToUpdate();
+    tokensToSave =
+        ofy().load().keys(keysToUpdate).values().stream()
+            .collect(toImmutableMap(Function.identity(), this::updateToken))
+            .entrySet()
+            .stream()
+            .filter(entry -> !entry.getKey().equals(entry.getValue())) // only update changed tokens
+            .map(Map.Entry::getValue)
+            .collect(toImmutableSet());
+  }
+
+  @Override
+  protected String execute() {
+    long numUpdated =
+        stream(partition(tokensToSave, BATCH_SIZE))
+            .mapToLong(batch -> ofy().transact(() -> saveBatch(batch)))
+            .sum();
+    return String.format("Updated %d tokens in total.", numUpdated);
+  }
+
+  private ImmutableList<Key<AllocationToken>> getKeysToUpdate() {
+    checkArgument(
+        tokens == null ^ prefix == null, "Must provide one of --tokens or --prefix, not both");
+    if (tokens != null) {
+      return tokens.stream()
+          .map(token -> Key.create(AllocationToken.class, token))
+          .collect(toImmutableList());
+    } else {
+      return ofy().load().type(AllocationToken.class).keys().list().stream()
+          .filter(key -> key.getName().startsWith(prefix))
+          .collect(toImmutableList());
+    }
+  }
+
+  private AllocationToken updateToken(AllocationToken original) {
+    AllocationToken.Builder builder = original.asBuilder();
+    Optional.ofNullable(allowedClientIds)
+        .ifPresent(clientIds -> builder.setAllowedClientIds(ImmutableSet.copyOf(clientIds)));
+    Optional.ofNullable(allowedTlds)
+        .ifPresent(tlds -> builder.setAllowedTlds(ImmutableSet.copyOf(tlds)));
+    Optional.ofNullable(discountFraction).ifPresent(builder::setDiscountFraction);
+    Optional.ofNullable(tokenStatusTransitions).ifPresent(builder::setTokenStatusTransitions);
+    return builder.build();
+  }
+
+  private long saveBatch(List<AllocationToken> batch) {
+    if (!dryRun) {
+      ofy().save().entities(batch);
+    }
+    System.out.printf(
+        "%s tokens: %s\n",
+        dryRun ? "Would update" : "Updated",
+        JOINER.join(
+            tokensToSave.stream()
+                .map(AllocationToken::getToken)
+                .sorted()
+                .collect(toImmutableSet())));
+    return batch.size();
+  }
+}

javatests/google/registry/tools/UpdateAllocationTokensCommandTest.java

@@ -0,0 +1,203 @@
+// 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.model.domain.token.AllocationToken.TokenStatus.CANCELLED;
+import static google.registry.model.domain.token.AllocationToken.TokenStatus.ENDED;
+import static google.registry.model.domain.token.AllocationToken.TokenStatus.NOT_STARTED;
+import static google.registry.model.domain.token.AllocationToken.TokenStatus.VALID;
+import static google.registry.model.domain.token.AllocationToken.TokenType.SINGLE_USE;
+import static google.registry.model.domain.token.AllocationToken.TokenType.UNLIMITED_USE;
+import static google.registry.testing.DatastoreHelper.persistResource;
+import static google.registry.testing.JUnitBackports.assertThrows;
+import static google.registry.util.DateTimeUtils.START_OF_TIME;
+import static org.joda.time.DateTimeZone.UTC;
+
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableSortedMap;
+import google.registry.model.domain.token.AllocationToken;
+import google.registry.model.domain.token.AllocationToken.TokenStatus;
+import org.joda.time.DateTime;
+import org.junit.Test;
+
+public class UpdateAllocationTokensCommandTest
+    extends CommandTestCase<UpdateAllocationTokensCommand> {
+
+  @Test
+  public void testUpdateTlds_setTlds() throws Exception {
+    AllocationToken token =
+        persistResource(builderWithPromo().setAllowedTlds(ImmutableSet.of("toRemove")).build());
+    runCommandForced("--prefix", "token", "--allowed_tlds", "tld,example");
+    assertThat(reloadResource(token).getAllowedTlds()).containsExactly("tld", "example");
+  }
+
+  @Test
+  public void testUpdateTlds_clearTlds() throws Exception {
+    AllocationToken token =
+        persistResource(builderWithPromo().setAllowedTlds(ImmutableSet.of("toRemove")).build());
+    runCommandForced("--prefix", "token", "--allowed_tlds", "");
+    assertThat(reloadResource(token).getAllowedTlds()).isEmpty();
+  }
+
+  @Test
+  public void testUpdateClientIds_setClientIds() throws Exception {
+    AllocationToken token =
+        persistResource(
+            builderWithPromo().setAllowedClientIds(ImmutableSet.of("toRemove")).build());
+    runCommandForced("--prefix", "token", "--allowed_client_ids", "clientone,clienttwo");
+    assertThat(reloadResource(token).getAllowedClientIds())
+        .containsExactly("clientone", "clienttwo");
+  }
+
+  @Test
+  public void testUpdateClientIds_clearClientIds() throws Exception {
+    AllocationToken token =
+        persistResource(
+            builderWithPromo().setAllowedClientIds(ImmutableSet.of("toRemove")).build());
+    runCommandForced("--prefix", "token", "--allowed_client_ids", "");
+    assertThat(reloadResource(token).getAllowedClientIds()).isEmpty();
+  }
+
+  @Test
+  public void testUpdateDiscountFraction() throws Exception {
+    AllocationToken token = persistResource(builderWithPromo().setDiscountFraction(0.5).build());
+    runCommandForced("--prefix", "token", "--discount_fraction", "0.15");
+    assertThat(reloadResource(token).getDiscountFraction()).isEqualTo(0.15);
+  }
+
+  @Test
+  public void testUpdateStatusTransitions() throws Exception {
+    DateTime now = DateTime.now(UTC);
+    AllocationToken token = persistResource(builderWithPromo().build());
+    runCommandForced(
+        "--prefix",
+        "token",
+        "--token_status_transitions",
+        String.format(
+            "\"%s=NOT_STARTED,%s=VALID,%s=CANCELLED\"", START_OF_TIME, now.minusDays(1), now));
+    token = reloadResource(token);
+    assertThat(token.getTokenStatusTransitions().toValueMap())
+        .containsExactly(START_OF_TIME, NOT_STARTED, now.minusDays(1), VALID, now, CANCELLED);
+  }
+
+  @Test
+  public void testUpdateStatusTransitions_badTransitions() {
+    DateTime now = DateTime.now(UTC);
+    persistResource(builderWithPromo().build());
+    IllegalArgumentException thrown =
+        assertThrows(
+            IllegalArgumentException.class,
+            () ->
+                runCommandForced(
+                    "--prefix",
+                    "token",
+                    "--token_status_transitions",
+                    String.format(
+                        "\"%s=NOT_STARTED,%s=ENDED,%s=VALID\"",
+                        START_OF_TIME, now.minusDays(1), now)));
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo("tokenStatusTransitions map cannot transition from NOT_STARTED to ENDED.");
+  }
+
+  @Test
+  public void testUpdate_onlyWithPrefix() throws Exception {
+    AllocationToken token =
+        persistResource(builderWithPromo().setAllowedTlds(ImmutableSet.of("tld")).build());
+    AllocationToken otherToken =
+        persistResource(
+            new AllocationToken.Builder()
+                .setToken("otherToken")
+                .setTokenType(SINGLE_USE)
+                .setAllowedTlds(ImmutableSet.of("tld"))
+                .build());
+    runCommandForced("--prefix", "other", "--allowed_tlds", "");
+    assertThat(reloadResource(token).getAllowedTlds()).containsExactly("tld");
+    assertThat(reloadResource(otherToken).getAllowedTlds()).isEmpty();
+  }
+
+  @Test
+  public void testUpdate_onlyTokensProvided() throws Exception {
+    AllocationToken firstToken =
+        persistResource(builderWithPromo().setAllowedTlds(ImmutableSet.of("tld")).build());
+    AllocationToken secondToken =
+        persistResource(
+            new AllocationToken.Builder()
+                .setToken("secondToken")
+                .setTokenType(SINGLE_USE)
+                .setAllowedTlds(ImmutableSet.of("tld"))
+                .build());
+    AllocationToken thirdToken =
+        persistResource(
+            new AllocationToken.Builder()
+                .setToken("thirdToken")
+                .setTokenType(SINGLE_USE)
+                .setAllowedTlds(ImmutableSet.of("tld"))
+                .build());
+    runCommandForced("--tokens", "secondToken,thirdToken", "--allowed_tlds", "");
+    assertThat(reloadResource(firstToken).getAllowedTlds()).containsExactly("tld");
+    assertThat(reloadResource(secondToken).getAllowedTlds()).isEmpty();
+    assertThat(reloadResource(thirdToken).getAllowedTlds()).isEmpty();
+  }
+
+  @Test
+  public void testDoNothing() throws Exception {
+    AllocationToken token =
+        persistResource(
+            builderWithPromo()
+                .setAllowedClientIds(ImmutableSet.of("clientid"))
+                .setAllowedTlds(ImmutableSet.of("tld"))
+                .setDiscountFraction(0.15)
+                .build());
+    runCommandForced("--prefix", "token");
+    AllocationToken reloaded = reloadResource(token);
+    assertThat(reloaded.getAllowedTlds()).isEqualTo(token.getAllowedTlds());
+    assertThat(reloaded.getAllowedClientIds()).isEqualTo(token.getAllowedClientIds());
+    assertThat(reloaded.getDiscountFraction()).isEqualTo(token.getDiscountFraction());
+  }
+
+  @Test
+  public void testFailure_bothTokensAndPrefix() {
+    assertThat(
+            assertThrows(
+                IllegalArgumentException.class,
+                () -> runCommandForced("--prefix", "token", "--tokens", "token")))
+        .hasMessageThat()
+        .isEqualTo("Must provide one of --tokens or --prefix, not both");
+  }
+
+  @Test
+  public void testFailure_neitherTokensNorPrefix() {
+    assertThat(
+            assertThrows(
+                IllegalArgumentException.class, () -> runCommandForced("--allowed_tlds", "tld")))
+        .hasMessageThat()
+        .isEqualTo("Must provide one of --tokens or --prefix, not both");
+  }
+
+  private static AllocationToken.Builder builderWithPromo() {
+    DateTime now = DateTime.now(UTC);
+    return new AllocationToken.Builder()
+        .setToken("token")
+        .setTokenType(UNLIMITED_USE)
+        .setTokenStatusTransitions(
+            ImmutableSortedMap.<DateTime, TokenStatus>naturalOrder()
+                .put(START_OF_TIME, NOT_STARTED)
+                .put(now.minusDays(1), VALID)
+                .put(now.plusDays(1), ENDED)
+                .build());
+  }
+}