Add a nomulus command to resave entity with unique id (#656)

* Add a nomulus command to resave entity with unique id

* Remove PollMessage

* Remove logic for PollMessage and resolve comments

* Resolve comments

core/src/main/java/google/registry/tools/MutatingCommand.java

@@ -143,10 +143,14 @@ public abstract class MutatingCommand extends ConfirmingCommand implements Comma
   protected String execute() throws Exception {
     for (final List<EntityChange> batch : getCollatedEntityChangeBatches()) {
       tm().transact(() -> batch.forEach(this::executeChange));
+      postBatchExecute();
     }
     return String.format("Updated %d entities.\n", changedEntitiesMap.size());
   }
 
+  /** Performs any execution step after each batch. */
+  protected void postBatchExecute() {}
+
   private void executeChange(EntityChange change) {
     // Load the key of the entity to mutate and double-check that it hasn't been
     // modified from the version that existed when the change was prepared.

core/src/main/java/google/registry/tools/RegistryTool.java

@@ -99,6 +99,7 @@ public final class RegistryTool {
           .put("remove_ip_address", RemoveIpAddressCommand.class)
           .put("renew_domain", RenewDomainCommand.class)
           .put("resave_entities", ResaveEntitiesCommand.class)
+          .put("resave_entities_with_unique_id", ResaveEntitiesWithUniqueIdCommand.class)
           .put("resave_environment_entities", ResaveEnvironmentEntitiesCommand.class)
           .put("resave_epp_resource", ResaveEppResourceCommand.class)
           .put("send_escrow_report_to_icann", SendEscrowReportToIcannCommand.class)

core/src/main/java/google/registry/tools/ResaveEntitiesWithUniqueIdCommand.java

@@ -0,0 +1,196 @@
+// Copyright 2020 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.checkState;
+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.api.datastore.KeyFactory;
+import com.google.common.base.Splitter;
+import com.google.common.io.CharStreams;
+import com.google.common.io.Files;
+import com.googlecode.objectify.Key;
+import google.registry.model.ImmutableObject;
+import google.registry.model.billing.BillingEvent;
+import google.registry.model.domain.DomainBase;
+import google.registry.util.NonFinalForTesting;
+import java.io.File;
+import java.io.InputStream;
+import java.io.InputStreamReader;
+import java.util.List;
+
+/**
+ * Command to resave entities with a unique id.
+ *
+ * <p>This command is used to address the duplicate id issue we found for certain {@link
+ * BillingEvent.OneTime} entities. The command reassigns an application wide unique id to the
+ * problematic entity and resaves it, it also resaves the entity having reference to the problematic
+ * entity with the updated id.
+ *
+ * <p>To use this command, you will need to provide the path to a file containing a list of strings
+ * representing the literal of Objectify key for the problematic entities. An example key literal
+ * is:
+ *
+ * <pre>
+ * "DomainBase", "111111-TEST", "HistoryEntry", 2222222, "OneTime", 3333333
+ * </pre>
+ *
+ * <p>Note that the double quotes are part of the key literal. The key literal can be retrieved from
+ * the column <code>__key__.path</code> in BigQuery.
+ */
+@Parameters(separators = " =", commandDescription = "Resave entities with a unique id.")
+public class ResaveEntitiesWithUniqueIdCommand extends MutatingCommand {
+
+  @Parameter(
+      names = "--key_paths_file",
+      description =
+          "Key paths file name, each line in the file should be a key literal. An example key"
+              + " literal is: \"DomainBase\", \"111111-TEST\", \"HistoryEntry\", 2222222,"
+              + " \"OneTime\", 3333333")
+  File keyPathsFile;
+
+  @NonFinalForTesting private static InputStream stdin = System.in;
+
+  private String keyChangeMessage;
+
+  @Override
+  protected void init() throws Exception {
+    List<String> keyPaths =
+        keyPathsFile == null
+            ? CharStreams.readLines(new InputStreamReader(stdin, UTF_8))
+            : Files.readLines(keyPathsFile, UTF_8);
+    for (String keyPath : keyPaths) {
+      Key<?> untypedKey = parseKeyPath(keyPath);
+      Object entity = ofy().load().key(untypedKey).now();
+      if (entity == null) {
+        System.err.println(
+            String.format(
+                "Entity %s read from %s doesn't exist in Datastore! Skipping.",
+                untypedKey,
+                keyPathsFile == null ? "STDIN" : "File " + keyPathsFile.getAbsolutePath()));
+        continue;
+      }
+      if (entity instanceof BillingEvent.OneTime) {
+        resaveBillingEvent((BillingEvent.OneTime) entity);
+      } else {
+        throw new IllegalArgumentException("Unsupported entity key: " + untypedKey);
+      }
+      flushTransaction();
+    }
+  }
+
+  @Override
+  protected void postBatchExecute() {
+    System.out.println(keyChangeMessage);
+  }
+
+  private void deleteOldAndSaveNewEntity(ImmutableObject oldEntity, ImmutableObject newEntity) {
+    stageEntityChange(oldEntity, null);
+    stageEntityChange(null, newEntity);
+  }
+
+  private void resaveBillingEvent(BillingEvent.OneTime billingEvent) {
+    Key<BillingEvent> key = Key.create(billingEvent);
+    Key<DomainBase> domainKey = getGrandParentAsDomain(key);
+    DomainBase domain = ofy().load().key(domainKey).now();
+
+    // The BillingEvent.OneTime entity to be resaved should be the billing event created a few
+    // years ago, so they should not be referenced from TransferData and GracePeriod in the domain.
+    assertNotInDomainTransferData(domain, key);
+    domain
+        .getGracePeriods()
+        .forEach(
+            gracePeriod ->
+                checkState(
+                    !gracePeriod.getOneTimeBillingEvent().getOfyKey().equals(key),
+                    "Entity %s is referenced by a grace period in domain %s",
+                    key,
+                    domainKey));
+
+    // By setting id to 0L, Buildable.build() will assign an application wide unique id to it.
+    BillingEvent.OneTime uniqIdBillingEvent = billingEvent.asBuilder().setId(0L).build();
+    deleteOldAndSaveNewEntity(billingEvent, uniqIdBillingEvent);
+    keyChangeMessage =
+        String.format("Old Entity Key: %s New Entity Key: %s", key, Key.create(uniqIdBillingEvent));
+  }
+
+  private static boolean isKind(Key<?> key, Class<?> clazz) {
+    return key.getKind().equals(Key.getKind(clazz));
+  }
+
+  static Key<?> parseKeyPath(String keyPath) {
+    List<String> keyComponents = Splitter.on(',').splitToList(keyPath);
+    checkState(
+        keyComponents.size() > 0 && keyComponents.size() % 2 == 0,
+        "Invalid number of key components");
+    com.google.appengine.api.datastore.Key rawKey = null;
+    for (int i = 0, j = 1; j < keyComponents.size(); i += 2, j += 2) {
+      String kindLiteral = keyComponents.get(i).trim();
+      String idOrNameLiteral = keyComponents.get(j).trim();
+      rawKey = createDatastoreKey(rawKey, kindLiteral, idOrNameLiteral);
+    }
+    return Key.create(rawKey);
+  }
+
+  private static com.google.appengine.api.datastore.Key createDatastoreKey(
+      com.google.appengine.api.datastore.Key parent, String kindLiteral, String idOrNameLiteral) {
+    if (isLiteralString(idOrNameLiteral)) {
+      return KeyFactory.createKey(parent, removeQuotes(kindLiteral), removeQuotes(idOrNameLiteral));
+    } else {
+      return KeyFactory.createKey(
+          parent, removeQuotes(kindLiteral), Long.parseLong(idOrNameLiteral));
+    }
+  }
+
+  private static boolean isLiteralString(String raw) {
+    return raw.charAt(0) == '"' && raw.charAt(raw.length() - 1) == '"';
+  }
+
+  private static String removeQuotes(String literal) {
+    return literal.substring(1, literal.length() - 1);
+  }
+
+  private static Key<DomainBase> getGrandParentAsDomain(Key<?> key) {
+    Key<?> grandParent;
+    try {
+      grandParent = key.getParent().getParent();
+    } catch (Throwable e) {
+      throw new IllegalArgumentException("Error retrieving grand parent key", e);
+    }
+    if (!isKind(grandParent, DomainBase.class)) {
+      throw new IllegalArgumentException(
+          String.format("Expected a Key<DomainBase> but got %s", grandParent));
+    }
+    return (Key<DomainBase>) grandParent;
+  }
+
+  private static void assertNotInDomainTransferData(DomainBase domainBase, Key<?> key) {
+    if (!domainBase.getTransferData().isEmpty()) {
+      domainBase
+          .getTransferData()
+          .getServerApproveEntities()
+          .forEach(
+              entityKey ->
+                  checkState(
+                      !entityKey.getOfyKey().equals(key),
+                      "Entity %s is referenced by the transfer data in domain %s",
+                      key,
+                      domainBase.createVKey().getOfyKey()));
+    }
+  }
+}

core/src/test/java/google/registry/tools/ResaveEntitiesWithUniqueIdCommandTest.java

@@ -0,0 +1,146 @@
+// Copyright 2020 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.createTld;
+import static google.registry.testing.DatastoreHelper.persistActiveDomain;
+import static google.registry.testing.DatastoreHelper.persistResource;
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.joda.money.CurrencyUnit.USD;
+import static org.junit.Assert.assertThrows;
+
+import com.google.common.collect.ImmutableSet;
+import com.googlecode.objectify.Key;
+import google.registry.model.EppResource;
+import google.registry.model.billing.BillingEvent;
+import google.registry.model.billing.BillingEvent.Reason;
+import google.registry.model.domain.DomainBase;
+import google.registry.model.domain.Period;
+import google.registry.model.eppcommon.Trid;
+import google.registry.model.poll.PollMessage;
+import google.registry.model.reporting.HistoryEntry;
+import google.registry.model.transfer.DomainTransferData;
+import org.joda.money.Money;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+/** Unit tests for {@link ResaveEntitiesWithUniqueIdCommand}. */
+class ResaveEntitiesWithUniqueIdCommandTest
+    extends CommandTestCase<ResaveEntitiesWithUniqueIdCommand> {
+
+  DomainBase domain;
+  HistoryEntry historyEntry;
+  PollMessage.Autorenew autorenewToResave;
+  BillingEvent.OneTime billingEventToResave;
+
+  @BeforeEach
+  void setUp() {
+    createTld("foobar");
+    domain = persistActiveDomain("foo.foobar");
+    historyEntry = persistHistoryEntry(domain);
+    autorenewToResave = persistAutorenewPollMessage(historyEntry);
+    billingEventToResave = persistBillingEvent(historyEntry);
+  }
+
+  @Test
+  void resaveBillingEvent_succeeds() throws Exception {
+    runCommand(
+        "--force",
+        "--key_paths_file",
+        writeToNamedTmpFile("keypath.txt", getKeyPathLiteral(billingEventToResave)));
+
+    int count = 0;
+    for (BillingEvent.OneTime billingEvent :
+        ofy().load().type(BillingEvent.OneTime.class).ancestor(historyEntry)) {
+      count++;
+      assertThat(billingEvent.getId()).isNotEqualTo(billingEventToResave.getId());
+      assertThat(billingEvent.asBuilder().setId(billingEventToResave.getId()).build())
+          .isEqualTo(billingEventToResave);
+    }
+    assertThat(count).isEqualTo(1);
+  }
+
+  @Test
+  void resaveBillingEvent_failsWhenReferredByDomain() throws Exception {
+    persistResource(
+        domain
+            .asBuilder()
+            .setTransferData(
+                new DomainTransferData.Builder()
+                    .setServerApproveEntities(ImmutableSet.of(billingEventToResave.createVKey()))
+                    .build())
+            .build());
+
+    assertThrows(
+        IllegalStateException.class,
+        () ->
+            runCommand(
+                "--force",
+                "--key_paths_file",
+                writeToNamedTmpFile("keypath.txt", getKeyPathLiteral(billingEventToResave))));
+  }
+
+  private PollMessage.Autorenew persistAutorenewPollMessage(HistoryEntry historyEntry) {
+    return persistResource(
+        new PollMessage.Autorenew.Builder()
+            .setClientId("TheRegistrar")
+            .setEventTime(fakeClock.nowUtc())
+            .setMsg("Test poll message")
+            .setParent(historyEntry)
+            .setAutorenewEndTime(fakeClock.nowUtc().plusDays(365))
+            .setTargetId("foobar.foo")
+            .build());
+  }
+
+  private BillingEvent.OneTime persistBillingEvent(HistoryEntry historyEntry) {
+    return persistResource(
+        new BillingEvent.OneTime.Builder()
+            .setClientId("a registrar")
+            .setTargetId("foo.tld")
+            .setParent(historyEntry)
+            .setReason(Reason.CREATE)
+            .setFlags(ImmutableSet.of(BillingEvent.Flag.ANCHOR_TENANT))
+            .setPeriodYears(2)
+            .setCost(Money.of(USD, 1))
+            .setEventTime(fakeClock.nowUtc())
+            .setBillingTime(fakeClock.nowUtc().plusDays(5))
+            .build());
+  }
+
+  private HistoryEntry persistHistoryEntry(EppResource parent) {
+    return persistResource(
+        new HistoryEntry.Builder()
+            .setParent(parent)
+            .setType(HistoryEntry.Type.DOMAIN_CREATE)
+            .setPeriod(Period.create(1, Period.Unit.YEARS))
+            .setXmlBytes("<xml></xml>".getBytes(UTF_8))
+            .setModificationTime(fakeClock.nowUtc())
+            .setClientId("foo")
+            .setTrid(Trid.create("ABC-123", "server-trid"))
+            .setBySuperuser(false)
+            .setReason("reason")
+            .setRequestedByRegistrar(false)
+            .build());
+  }
+
+  private static String getKeyPathLiteral(Object entity) {
+    Key<?> key = Key.create(entity);
+    return String.format(
+        "\"DomainBase\", \"%s\", \"HistoryEntry\", %s, \"%s\", %s",
+        key.getParent().getParent().getName(), key.getParent().getId(), key.getKind(), key.getId());
+  }
+}