Use VKeys instead of Ofy keys in mutating command (#1682)

* Use VKeys instead of Ofy keys in mutating command

* Add createVKey to ImmutableObject

* Use SQL only VKeys

core/src/main/java/google/registry/model/ImmutableObject.java

@@ -27,6 +27,7 @@ import com.google.common.base.Joiner;
 import com.google.common.collect.Maps;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Ignore;
+import google.registry.persistence.VKey;
 import java.lang.annotation.Documented;
 import java.lang.annotation.Retention;
 import java.lang.annotation.Target;
@@ -255,4 +256,8 @@ public abstract class ImmutableObject implements Cloneable {
   public Map<String, Object> toDiffableFieldMap() {
     return (Map<String, Object>) toMapRecursive(this);
   }
+
+  public VKey createVKey() {
+    throw new UnsupportedOperationException("VKey creation is not supported for this entity");
+  }
 }

core/src/main/java/google/registry/model/tld/Registry.java

@@ -293,6 +293,11 @@ public class Registry extends ImmutableObject implements Buildable, UnsafeSerial
     return createVKey(key.getName());
   }
 
+  @Override
+  public VKey<Registry> createVKey() {
+    return VKey.createSql(Registry.class, this.tldStrId);
+  }
+
   /**
    * The name of the pricing engine that this TLD uses.
    *

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

@@ -22,7 +22,6 @@ import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.base.Strings.emptyToNull;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
-import static google.registry.util.DatastoreServiceUtils.getNameOrId;
 import static google.registry.util.DiffUtils.prettyPrintEntityDeepDiff;
 import static java.util.stream.Collectors.joining;
 
@@ -30,7 +29,6 @@ import com.google.common.base.MoreObjects;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
-import com.googlecode.objectify.Key;
 import google.registry.model.ImmutableObject;
 import google.registry.persistence.VKey;
 import java.util.ArrayList;
@@ -77,65 +75,21 @@ public abstract class MutatingCommand extends ConfirmingCommand implements Comma
     final ImmutableObject newEntity;
 
     /** The key that points to the entity being changed. */
-    final VKey<?> key;
+    final VKey<?> vKey;
 
     private EntityChange(ImmutableObject oldEntity, ImmutableObject newEntity) {
       type = ChangeType.get(oldEntity != null, newEntity != null);
       checkArgument(
-          type != ChangeType.UPDATE || Key.create(oldEntity).equals(Key.create(newEntity)),
+          type != ChangeType.UPDATE || oldEntity.createVKey().equals(newEntity.createVKey()),
           "Both entity versions in an update must have the same Key.");
       this.oldEntity = oldEntity;
       this.newEntity = newEntity;
-      ImmutableObject entity = MoreObjects.firstNonNull(oldEntity, newEntity);
+      vKey = MoreObjects.firstNonNull(oldEntity, newEntity).createVKey();
-
-      // This is one of the few cases where it is acceptable to create an asymmetric VKey (using
-      // createOfy()).  We can use this code on datastore-only entities where we can't construct a
-      // SQL key.
-      VKey<?> createdKey;
-      try {
-        createdKey = VKey.from(Key.create(entity));
-      } catch (RuntimeException e) {
-        createdKey = VKey.createOfy(entity.getClass(), Key.create(entity));
-      }
-      key = createdKey;
-    }
-
-    /**
-     * EntityChange constructor that supports Vkey override. A Vkey is a key of an entity. This is a
-     * workaround to handle cases when a SqlEntity instance does not have a primary key before being
-     * persisted.
-     */
-    private EntityChange(
-        @Nullable ImmutableObject oldEntity, @Nullable ImmutableObject newEntity, VKey<?> vkey) {
-      type = ChangeType.get(oldEntity != null, newEntity != null);
-      if (type == ChangeType.UPDATE) {
-        checkArgument(
-            Key.create(oldEntity).equals(Key.create(newEntity)),
-            "Both entity versions in an update must have the same Key.");
-        checkArgument(
-            Key.create(oldEntity).equals(vkey.getOfyKey()),
-            "The Key of the entity must be the same as the OfyKey of the vkey");
-      } else if (type == ChangeType.CREATE) {
-        checkArgument(
-            Key.create(newEntity).equals(vkey.getOfyKey()),
-            "Both entity versions in an update must have the same Key.");
-      } else if (type == ChangeType.DELETE) {
-        checkArgument(
-            Key.create(oldEntity).equals(vkey.getOfyKey()),
-            "The Key of the entity must be the same as the OfyKey of the vkey");
-      }
-      this.oldEntity = oldEntity;
-      this.newEntity = newEntity;
-      key = vkey;
     }
 
     /** Returns a human-readable ID string for the entity being changed. */
     String getEntityId() {
-      return String.format(
+      return String.format("%s@%s", vKey.getKind().getSimpleName(), vKey.getSqlKey().toString());
-          "%s@%s",
-          key.getOfyKey().getKind(),
-          // NB: try name before id, since name defaults to null, whereas id defaults to 0.
-          getNameOrId(key.getOfyKey().getRaw()));
     }
 
     /** Returns a string representation of this entity change. */
@@ -195,7 +149,7 @@ public abstract class MutatingCommand extends ConfirmingCommand implements Comma
   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.
-    Optional<?> existingEntity = tm().loadByKeyIfPresent(change.key);
+    Optional<?> existingEntity = tm().loadByKeyIfPresent(change.vKey);
     checkState(
         Objects.equals(change.oldEntity, existingEntity.orElse(null)),
         "Entity changed since init() was called.\n%s",
@@ -212,7 +166,7 @@ public abstract class MutatingCommand extends ConfirmingCommand implements Comma
         tm().update(change.newEntity);
         return;
       case DELETE:
-        tm().delete(change.key);
+        tm().delete(change.vKey);
         return;
     }
     throw new UnsupportedOperationException("Unknown entity change type: " + change.type);
@@ -227,7 +181,7 @@ public abstract class MutatingCommand extends ConfirmingCommand implements Comma
     ArrayList<EntityChange> nextBatch = new ArrayList<>();
     for (EntityChange change : changedEntitiesMap.values()) {
       nextBatch.add(change);
-      if (transactionBoundaries.contains(change.key)) {
+      if (transactionBoundaries.contains(change.vKey)) {
         batches.add(ImmutableList.copyOf(nextBatch));
         nextBatch.clear();
       }
@@ -249,30 +203,11 @@ public abstract class MutatingCommand extends ConfirmingCommand implements Comma
       @Nullable ImmutableObject oldEntity, @Nullable ImmutableObject newEntity) {
     EntityChange change = new EntityChange(oldEntity, newEntity);
     checkArgument(
-        !changedEntitiesMap.containsKey(change.key),
+        !changedEntitiesMap.containsKey(change.vKey),
         "Cannot apply multiple changes for the same entity: %s",
         change.getEntityId());
-    changedEntitiesMap.put(change.key, change);
+    changedEntitiesMap.put(change.vKey, change);
-    lastAddedKey = change.key;
+    lastAddedKey = change.vKey;
-  }
-
-  /**
-   * Stages an entity change which will be applied by execute(), with the support of Vkey override.
-   * It supports cases of SqlEntity instances that do not have primary keys before being persisted.
-   *
-   * @param oldEntity the existing version of the entity, or null to create a new entity
-   * @param newEntity the new version of the entity to save, or null to delete the entity
-   * @param vkey the key of the entity
-   */
-  protected void stageEntityChange(
-      @Nullable ImmutableObject oldEntity, @Nullable ImmutableObject newEntity, VKey vkey) {
-    EntityChange change = new EntityChange(oldEntity, newEntity, vkey);
-    checkArgument(
-        !changedEntitiesMap.containsKey(change.key),
-        "Cannot apply multiple changes for the same entity: %s",
-        change.getEntityId());
-    changedEntitiesMap.put(change.key, change);
-    lastAddedKey = change.key;
   }
 
   /**

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

@@ -1,166 +0,0 @@
-// 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.auditedOfy;
-import static java.nio.charset.StandardCharsets.UTF_8;
-
-import com.beust.jcommander.Parameter;
-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.annotations.DeleteAfterMigration;
-import google.registry.model.domain.DomainBase;
-import google.registry.persistence.VKey;
-import google.registry.util.NonFinalForTesting;
-import google.registry.util.TypeUtils.TypeInstantiator;
-import java.io.File;
-import java.io.InputStream;
-import java.io.InputStreamReader;
-import java.util.List;
-
-/**
- * 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.
- */
-@DeleteAfterMigration
-abstract class ReadEntityFromKeyPathCommand<T> 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 StringBuilder changeMessage = new StringBuilder();
-
-  abstract void process(T entity);
-
-  @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 = auditedOfy().load().key(untypedKey).now();
-      if (entity == null) {
-        System.err.printf(
-            "Entity %s read from %s doesn't exist in Datastore! Skipping.%n",
-            untypedKey, keyPathsFile == null ? "STDIN" : "File " + keyPathsFile.getAbsolutePath());
-        continue;
-      }
-      Class<T> clazz = new TypeInstantiator<T>(getClass()) {}.getExactType();
-      if (clazz.isInstance(entity)) {
-        process((T) entity);
-      } else {
-        throw new IllegalArgumentException("Unsupported entity key: " + untypedKey);
-      }
-      flushTransaction();
-    }
-  }
-
-  @Override
-  protected void postBatchExecute() {
-    System.out.println(changeMessage);
-  }
-
-  void stageEntityKeyChange(ImmutableObject oldEntity, ImmutableObject newEntity) {
-    stageEntityChange(oldEntity, null);
-    stageEntityChange(null, newEntity);
-    appendChangeMessage(
-        String.format(
-            "Changed entity key from: %s to: %s", Key.create(oldEntity), Key.create(newEntity)));
-  }
-
-  void appendChangeMessage(String message) {
-    changeMessage.append(message);
-  }
-
-  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);
-  }
-
-  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;
-  }
-
-  static VKey<DomainBase> getGrandParentAsDomain(VKey<?> key) {
-    Key<DomainBase> grandParent;
-    try {
-      grandParent = key.getOfyKey().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 VKey.create(DomainBase.class, grandParent.getName(), grandParent);
-  }
-}

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

@@ -99,9 +99,7 @@ public final class RegistryTool {
           .put("pending_escrow", PendingEscrowCommand.class)
           .put("registrar_contact", RegistrarContactCommand.class)
           .put("renew_domain", RenewDomainCommand.class)
-          .put("resave_entities", ResaveEntitiesCommand.class)
           .put("resave_environment_entities", ResaveEnvironmentEntitiesCommand.class)
-          .put("resave_epp_resource", ResaveEppResourceCommand.class)
           .put("save_sql_credential", SaveSqlCredentialCommand.class)
           .put("send_escrow_report_to_icann", SendEscrowReportToIcannCommand.class)
           .put("set_num_instances", SetNumInstancesCommand.class)

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

@@ -1,55 +0,0 @@
-// Copyright 2017 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.collect.Lists.partition;
-import static google.registry.model.ofy.ObjectifyService.auditedOfy;
-
-import com.beust.jcommander.Parameter;
-import com.beust.jcommander.Parameters;
-import google.registry.model.ImmutableObject;
-import google.registry.persistence.VKey;
-import java.util.List;
-
-/**
- * A command to load and resave an entity by websafe key.
- *
- * <p>This triggers @OnSave changes. If the entity was directly edited in the Datastore viewer, this
- * can be used to make sure that the commit logs reflect the new state.
- */
-@Parameters(
-    separators = " =",
-    commandDescription = "Load and resave entities by websafe key")
-public final class ResaveEntitiesCommand extends MutatingCommand {
-
-  /** The number of resaves to do in a single transaction. */
-  private static final int BATCH_SIZE = 10;
-
-  // TODO(b/207376744): figure out if there's a guide that shows how a websafe key should look like
-  @Parameter(description = "Websafe keys", required = true)
-  List<String> mainParameters;
-
-  @Override
-  protected void init() {
-    for (List<String> batch : partition(mainParameters, BATCH_SIZE)) {
-      for (String websafeKey : batch) {
-        ImmutableObject entity =
-            (ImmutableObject) auditedOfy().load().key(VKey.create(websafeKey).getOfyKey()).now();
-        stageEntityChange(entity, entity);
-      }
-      flushTransaction();
-    }
-  }
-}

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

@@ -1,62 +0,0 @@
-// Copyright 2017 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 google.registry.persistence.transaction.TransactionManagerFactory.tm;
-import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
-import static org.joda.time.DateTimeZone.UTC;
-
-import com.beust.jcommander.Parameter;
-import com.beust.jcommander.Parameters;
-import google.registry.model.EppResource;
-import google.registry.persistence.VKey;
-import google.registry.tools.CommandUtilities.ResourceType;
-import org.joda.time.DateTime;
-
-/**
- * A command to load and resave an {@link EppResource} by foreign key.
- *
- * <p>This triggers @OnSave changes. If the entity was directly edited in the Datastore viewer, this
- * can be used to make sure that the commit logs reflect the new state.
- */
-@Parameters(
-    separators = " =",
-    commandDescription = "Load and resave EPP resources by foreign key")
-public final class ResaveEppResourceCommand extends MutatingCommand {
-
-  @Parameter(
-    names = "--type",
-    description = "Resource type.")
-  protected ResourceType type;
-
-  @Parameter(
-    names = "--id",
-    description = "Foreign key of the resource.")
-  protected String uniqueId;
-
-  @Override
-  protected void init() {
-    VKey<? extends EppResource> resourceKey =
-        checkArgumentNotNull(
-            type.getKey(uniqueId, DateTime.now(UTC)),
-            "Could not find active resource of type %s: %s",
-            type,
-            uniqueId);
-    // Load the resource directly to bypass running cloneProjectedAtTime() automatically, which can
-    // cause stageEntityChange() to fail due to implicit projection changes.
-    EppResource resource = tm().loadByKey(resourceKey);
-    stageEntityChange(resource, resource);
-  }
-}

core/src/test/java/google/registry/persistence/BillingVKeyTest.java

@@ -106,9 +106,8 @@ class BillingVKeyTest {
       return billingRecurrenceVKey.createVKey();
     }
 
-    VKey<BillingVKeyTestEntity> createVKey() {
+    public VKey<BillingVKeyTestEntity> createVKey() {
-      return VKey.create(
+      return VKey.createSql(BillingVKeyTestEntity.class, id);
-          BillingVKeyTestEntity.class, id, Key.create(parent, BillingVKeyTestEntity.class, id));
     }
   }
 }

core/src/test/java/google/registry/persistence/DomainHistoryVKeyTest.java

@@ -96,8 +96,8 @@ class DomainHistoryVKeyTest {
       this.domainHistoryVKey = domainHistoryVKey;
     }
 
-    VKey<TestEntity> createVKey() {
+    public VKey<TestEntity> createVKey() {
-      return VKey.create(TestEntity.class, id, Key.create(parent, TestEntity.class, id));
+      return VKey.createSql(TestEntity.class, id);
     }
   }
 }