Add a command to remove Registry 1.0 key in DomainBase (#900)

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

@@ -45,10 +45,10 @@ import google.registry.model.domain.DomainBase;
 @Parameters(
     separators = " =",
     commandDescription = "Dedupe BillingEvent.OneTime entities with duplicate IDs.")
-public class DedupeOneTimeBillingEventIdsCommand extends DedupeEntityIdsCommand<OneTime> {
+public class DedupeOneTimeBillingEventIdsCommand extends ReadEntityFromKeyPathCommand<OneTime> {
 
   @Override
-  void dedupe(OneTime entity) {
+  void process(OneTime entity) {
     Key<BillingEvent> key = Key.create(entity);
     Key<DomainBase> domainKey = getGrandParentAsDomain(key);
     DomainBase domain = ofy().load().key(domainKey).now();

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

@@ -54,11 +54,10 @@ import java.util.Set;
 @Parameters(
     separators = " =",
     commandDescription = "Dedupe BillingEvent.Recurring entities with duplicate IDs.")
-public class DedupeRecurringBillingEventIdsCommand
-    extends DedupeEntityIdsCommand<BillingEvent.Recurring> {
+public class DedupeRecurringBillingEventIdsCommand extends ReadEntityFromKeyPathCommand<Recurring> {
 
   @Override
-  void dedupe(Recurring recurring) {
+  void process(Recurring recurring) {
     // Loads the associated DomainBase and BillingEvent.OneTime entities that
     // may have reference to this BillingEvent.Recurring entity.
     Key<DomainBase> domainKey = getGrandParentAsDomain(Key.create(recurring));

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

@@ -33,8 +33,13 @@ import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.util.List;
 
-/** Base Command to dedupe entities with duplicate IDs. */
-abstract class DedupeEntityIdsCommand<T> extends MutatingCommand {
+/**
+ * Base Command to read entities from Datastore by their key paths retrieved from BigQuery.
+ *
+ * <p>The key path is the value of column __key__.path of the entity's BigQuery table. Its value is
+ * converted from the entity's key.
+ */
+abstract class ReadEntityFromKeyPathCommand<T> extends MutatingCommand {
 
   @Parameter(
       names = "--key_paths_file",
@@ -48,7 +53,7 @@ abstract class DedupeEntityIdsCommand<T> extends MutatingCommand {
 
   private StringBuilder changeMessage = new StringBuilder();
 
-  abstract void dedupe(T entity);
+  abstract void process(T entity);
 
   @Override
   protected void init() throws Exception {
@@ -69,7 +74,7 @@ abstract class DedupeEntityIdsCommand<T> extends MutatingCommand {
       }
       Class<T> clazz = new TypeInstantiator<T>(getClass()) {}.getExactType();
       if (clazz.isInstance(entity)) {
-        dedupe((T) entity);
+        process((T) entity);
       } else {
         throw new IllegalArgumentException("Unsupported entity key: " + untypedKey);
       }

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

@@ -99,6 +99,7 @@ public final class RegistryTool {
           .put("populate_null_registrar_fields", PopulateNullRegistrarFieldsCommand.class)
           .put("registrar_contact", RegistrarContactCommand.class)
           .put("remove_ip_address", RemoveIpAddressCommand.class)
+          .put("remove_registry_one_key", RemoveRegistryOneKeyCommand.class)
           .put("renew_domain", RenewDomainCommand.class)
           .put("resave_entities", ResaveEntitiesCommand.class)
           .put("resave_environment_entities", ResaveEnvironmentEntitiesCommand.class)

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

@@ -0,0 +1,69 @@
+// 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 com.beust.jcommander.Parameters;
+import com.googlecode.objectify.Key;
+import google.registry.model.domain.DomainBase;
+import google.registry.persistence.VKey;
+import javax.annotation.Nullable;
+import org.joda.time.DateTime;
+
+/** Command to remove the Registry 1.0 key in {@link DomainBase} entity. */
+@Parameters(separators = " =", commandDescription = "Remove .")
+public class RemoveRegistryOneKeyCommand extends ReadEntityFromKeyPathCommand<DomainBase> {
+
+  @Override
+  void process(DomainBase entity) {
+    // Assert that the DomainBase entity must be deleted before 2017-08-01(most of the problematic
+    // entities were deleted before 2017, though there are still a few entities deleted in 2017-07).
+    // This is because we finished the Registry 2.0 migration in 2017 and should not generate any
+    // Registry 1.0 key after it.
+    if (!entity.getDeletionTime().isBefore(DateTime.parse("2017-08-01T00:00:00Z"))) {
+      throw new IllegalStateException(
+          String.format(
+              "Entity's deletion time %s is not before 2017-08-01T00:00:00Z",
+              entity.getDeletionTime()));
+    }
+    boolean hasChange = false;
+    DomainBase.Builder domainBuilder = entity.asBuilder();
+    // We only found the registry 1.0 key existed in fields autorenewBillingEvent,
+    // autorenewPollMessage and deletePollMessage so we just need to check these fields for each
+    // entity.
+    if (isRegistryOneKey(entity.getAutorenewBillingEvent())) {
+      domainBuilder.setAutorenewBillingEvent(null);
+      hasChange = true;
+    }
+    if (isRegistryOneKey(entity.getAutorenewPollMessage())) {
+      domainBuilder.setAutorenewPollMessage(null);
+      hasChange = true;
+    }
+    if (isRegistryOneKey(entity.getDeletePollMessage())) {
+      domainBuilder.setDeletePollMessage(null);
+      hasChange = true;
+    }
+    if (hasChange) {
+      stageEntityChange(entity, domainBuilder.build());
+    }
+  }
+
+  private static boolean isRegistryOneKey(@Nullable VKey<?> vKey) {
+    if (vKey == null || vKey.getOfyKey() == null || vKey.getOfyKey().getParent() == null) {
+      return false;
+    }
+    Key<?> parentKey = vKey.getOfyKey().getParent();
+    return parentKey.getKind().equals("EntityGroupRoot") && parentKey.getName().equals("per-tld");
+  }
+}

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

@@ -0,0 +1,122 @@
+// 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.ImmutableObjectSubject.immutableObjectCorrespondence;
+import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.testing.DatabaseHelper.createTld;
+import static google.registry.testing.DatabaseHelper.newDomainBase;
+import static google.registry.testing.DatabaseHelper.persistResource;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.truth.Correspondence;
+import com.googlecode.objectify.Key;
+import google.registry.model.ImmutableObject;
+import google.registry.model.billing.BillingEvent;
+import google.registry.model.common.EntityGroupRoot;
+import google.registry.model.domain.DomainBase;
+import google.registry.model.poll.PollMessage;
+import google.registry.model.reporting.HistoryEntry;
+import google.registry.persistence.VKey;
+import org.joda.time.DateTime;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+/** Unit test for {@link RemoveRegistryOneKeyCommand}. */
+public class RemoveRegistryOneKeyCommandTest extends CommandTestCase<RemoveRegistryOneKeyCommand> {
+  DomainBase domain;
+  HistoryEntry historyEntry;
+
+  @BeforeEach
+  void beforeEach() {
+    createTld("foobar");
+    domain =
+        newDomainBase("foo.foobar")
+            .asBuilder()
+            .setDeletionTime(DateTime.parse("2016-01-01T00:00:00Z"))
+            .setAutorenewBillingEvent(createRegistryOneVKey(BillingEvent.Recurring.class, 100L))
+            .setAutorenewPollMessage(createRegistryOneVKey(PollMessage.Autorenew.class, 200L))
+            .setDeletePollMessage(createRegistryOneVKey(PollMessage.OneTime.class, 300L))
+            .build();
+  }
+
+  @Test
+  void removeRegistryOneKeyInDomainBase_succeeds() throws Exception {
+    DomainBase origin = persistResource(domain);
+
+    runCommand(
+        "--force",
+        "--key_paths_file",
+        writeToNamedTmpFile("keypath.txt", getKeyPathLiteral(domain)));
+
+    DomainBase persisted = ofy().load().key(domain.createVKey().getOfyKey()).now();
+    assertThat(ImmutableList.of(persisted))
+        .comparingElementsUsing(getDomainBaseCorrespondence())
+        .containsExactly(origin);
+    assertThat(persisted.getAutorenewBillingEvent()).isNull();
+    assertThat(persisted.getAutorenewPollMessage()).isNull();
+    assertThat(persisted.getDeletePollMessage()).isNull();
+  }
+
+  @Test
+  void removeRegistryOneKeyInDomainBase_notModifyRegistryTwoKey() throws Exception {
+    DomainBase origin =
+        persistResource(
+            domain
+                .asBuilder()
+                .setAutorenewBillingEvent(
+                    createRegistryTwoVKey(BillingEvent.Recurring.class, domain, 300L))
+                .build());
+
+    runCommand(
+        "--force",
+        "--key_paths_file",
+        writeToNamedTmpFile("keypath.txt", getKeyPathLiteral(domain)));
+
+    DomainBase persisted = ofy().load().key(domain.createVKey().getOfyKey()).now();
+    assertThat(ImmutableList.of(persisted))
+        .comparingElementsUsing(getDomainBaseCorrespondence())
+        .containsExactly(origin);
+    assertThat(persisted.getAutorenewBillingEvent())
+        .isEqualTo(createRegistryTwoVKey(BillingEvent.Recurring.class, domain, 300L));
+    assertThat(persisted.getAutorenewPollMessage()).isNull();
+    assertThat(persisted.getDeletePollMessage()).isNull();
+  }
+
+  private static String getKeyPathLiteral(Object entity) {
+    Key<?> key = Key.create(entity);
+    return String.format("\"DomainBase\", \"%s\"", key.getName());
+  }
+
+  private static <T> VKey<T> createRegistryOneVKey(Class<T> clazz, long id) {
+    Key<?> parent = Key.create(EntityGroupRoot.class, "per-tld");
+    return VKey.create(clazz, id, Key.create(parent, clazz, id));
+  }
+
+  private static <T> VKey<T> createRegistryTwoVKey(Class<T> clazz, DomainBase domain, long id) {
+    Key<?> parent = Key.create(domain.createVKey().getOfyKey(), HistoryEntry.class, 1000L);
+    return VKey.create(clazz, id, Key.create(parent, clazz, id));
+  }
+
+  private static Correspondence<ImmutableObject, ImmutableObject> getDomainBaseCorrespondence() {
+    return immutableObjectCorrespondence(
+        "revisions",
+        "updateTimestamp",
+        "autorenewBillingEvent",
+        "autorenewPollMessage",
+        "deletePollMessage");
+  }
+}