mirror of
https://github.com/google/nomulus.git
synced 2025-05-14 16:37:13 +02:00
Add a command to update allocation tokens
------------- Created by MOE: https://github.com/google/moe MOE_MIGRATED_REVID=246824258
This commit is contained in:
gbrodman
jianglai
parent
df7e9a1225
commit
3483f348dc
3 changed files with 378 additions and 0 deletions
|
|
@ -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)
|
||||
|
|
|
|||
174
java/google/registry/tools/UpdateAllocationTokensCommand.java
Normal file
174
java/google/registry/tools/UpdateAllocationTokensCommand.java
Normal file
|
|
@ -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();
|
||||
}
|
||||
}
|
||||
|
|
@ -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());
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue