diff --git a/java/google/registry/tools/RegistryTool.java b/java/google/registry/tools/RegistryTool.java index 5beb36bd3..f3f24211b 100644 --- a/java/google/registry/tools/RegistryTool.java +++ b/java/google/registry/tools/RegistryTool.java @@ -17,6 +17,7 @@ package google.registry.tools; import com.google.common.collect.ImmutableMap; import google.registry.tools.javascrap.PopulateNullRegistrarFieldsCommand; import google.registry.tools.javascrap.RemoveIpAddressCommand; +import google.registry.tools.javascrap.ResaveAllocationTokensCommand; /** Container class to create and run remote commands against a Datastore instance. */ public final class RegistryTool { @@ -96,6 +97,7 @@ public final class RegistryTool { .put("registrar_contact", RegistrarContactCommand.class) .put("remove_ip_address", RemoveIpAddressCommand.class) .put("renew_domain", RenewDomainCommand.class) + .put("resave_allocation_tokens", ResaveAllocationTokensCommand.class) .put("resave_entities", ResaveEntitiesCommand.class) .put("resave_environment_entities", ResaveEnvironmentEntitiesCommand.class) .put("resave_epp_resource", ResaveEppResourceCommand.class) diff --git a/java/google/registry/tools/javascrap/ResaveAllocationTokensCommand.java b/java/google/registry/tools/javascrap/ResaveAllocationTokensCommand.java new file mode 100644 index 000000000..b9c89c9de --- /dev/null +++ b/java/google/registry/tools/javascrap/ResaveAllocationTokensCommand.java @@ -0,0 +1,43 @@ +// 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.javascrap; + +import static com.google.common.collect.Lists.partition; +import static google.registry.model.ofy.ObjectifyService.ofy; + +import com.beust.jcommander.Parameters; +import com.googlecode.objectify.Key; +import google.registry.model.domain.token.AllocationToken; +import google.registry.tools.MutatingCommand; +import java.util.List; + +/** Scrap tool to load and save all {@link AllocationToken}s to populate on-load attributes. */ +@Parameters(commandDescription = "Load and resave all allocation tokens") +public final class ResaveAllocationTokensCommand extends MutatingCommand { + + private static final int BATCH_SIZE = 20; + + @Override + protected void init() { + List> tokenKeys = ofy().load().type(AllocationToken.class).keys().list(); + for (List> partitionedKeys : partition(tokenKeys, BATCH_SIZE)) { + for (Key key : partitionedKeys) { + AllocationToken token = ofy().load().key(key).now(); + stageEntityChange(token, token); + } + flushTransaction(); + } + } +}