Convert ofy -> tm for two more classes (#1049)

* Convert ofy -> tm for two more classes

Convert ofy -> tm for MutatingCommand and DedupeOneTimeBillingEventIdsCommand.

Note that DedupeOneTimeBillingEventIdsCommand will not be needed after
migration, so this conversion is just to remove the ofy uses from the
codebase.  We don't update the test (other than to keep it working) and it
wouldn't currently work in SQL.

* Fixed a test broken by this PR

core/src/main/java/google/registry/persistence/VKey.java

@@ -64,7 +64,7 @@ public class VKey<T> extends ImmutableObject implements Serializable {
   }
 
   /** Creates a {@link VKey} which only contains the ofy primary key. */
-  public static <T> VKey<T> createOfy(Class<T> kind, Key<T> ofyKey) {
+  public static <T> VKey<T> createOfy(Class<? extends T> kind, Key<T> ofyKey) {
     checkArgumentNotNull(kind, "kind must not be null");
     checkArgumentNotNull(ofyKey, "ofyKey must not be null");
     return new VKey<T>(kind, ofyKey, null);

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

@@ -15,13 +15,13 @@
 package google.registry.tools;
 
 import static com.google.common.base.Preconditions.checkState;
-import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.beust.jcommander.Parameters;
-import com.googlecode.objectify.Key;
 import google.registry.model.billing.BillingEvent;
 import google.registry.model.billing.BillingEvent.OneTime;
 import google.registry.model.domain.DomainBase;
+import google.registry.persistence.VKey;
 
 /**
  * Command to dedupe {@link BillingEvent.OneTime} entities having duplicate IDs.
@@ -49,9 +49,9 @@ public class DedupeOneTimeBillingEventIdsCommand extends ReadEntityFromKeyPathCo
 
   @Override
   void process(OneTime entity) {
-    Key<BillingEvent> key = Key.create(entity);
+    VKey<OneTime> key = entity.createVKey();
-    Key<DomainBase> domainKey = getGrandParentAsDomain(key);
+    VKey<DomainBase> domainKey = getGrandParentAsDomain(key);
-    DomainBase domain = ofy().load().key(domainKey).now();
+    DomainBase domain = tm().transact(() -> tm().loadByKey(domainKey));
 
     // 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.
@@ -61,7 +61,7 @@ public class DedupeOneTimeBillingEventIdsCommand extends ReadEntityFromKeyPathCo
         .forEach(
             gracePeriod ->
                 checkState(
-                    !gracePeriod.getOneTimeBillingEvent().getOfyKey().equals(key),
+                    !gracePeriod.getOneTimeBillingEvent().equals(key),
                     "Entity %s is referenced by a grace period in domain %s",
                     key,
                     domainKey));
@@ -71,7 +71,7 @@ public class DedupeOneTimeBillingEventIdsCommand extends ReadEntityFromKeyPathCo
     stageEntityKeyChange(entity, uniqIdBillingEvent);
   }
 
-  private static void assertNotInDomainTransferData(DomainBase domainBase, Key<?> key) {
+  private static void assertNotInDomainTransferData(DomainBase domainBase, VKey<?> key) {
     if (!domainBase.getTransferData().isEmpty()) {
       domainBase
           .getTransferData()
@@ -79,10 +79,10 @@ public class DedupeOneTimeBillingEventIdsCommand extends ReadEntityFromKeyPathCo
           .forEach(
               entityKey ->
                   checkState(
-                      !entityKey.getOfyKey().equals(key),
+                      !entityKey.equals(key),
                       "Entity %s is referenced by the transfer data in domain %s",
                       key,
-                      domainBase.createVKey().getOfyKey()));
+                      domainBase.createVKey()));
     }
   }
 }

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

@@ -21,7 +21,6 @@ import static com.google.common.base.Preconditions.checkNotNull;
 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.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.DatastoreServiceUtils.getNameOrId;
 import static google.registry.util.DiffUtils.prettyPrintEntityDeepDiff;
@@ -33,6 +32,8 @@ 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 google.registry.schema.replay.SqlEntity;
 import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.LinkedHashMap;
@@ -75,7 +76,7 @@ public abstract class MutatingCommand extends ConfirmingCommand implements Comma
     final ImmutableObject newEntity;
 
     /** The key that points to the entity being changed. */
-    final Key<ImmutableObject> key;
+    final VKey<?> key;
 
     public EntityChange(ImmutableObject oldEntity, ImmutableObject newEntity) {
       type = ChangeType.get(oldEntity != null, newEntity != null);
@@ -84,16 +85,24 @@ public abstract class MutatingCommand extends ConfirmingCommand implements Comma
           "Both entity versions in an update must have the same Key.");
       this.oldEntity = oldEntity;
       this.newEntity = newEntity;
-      key = Key.create(MoreObjects.firstNonNull(oldEntity, newEntity));
+      ImmutableObject entity = MoreObjects.firstNonNull(oldEntity, newEntity);
+
+      // This is one of the few cases where it is acceptable to create an asymmetric VKey (using
+      // createOfy()).  We can use this code on DatastoreOnlyEntity's where we can't construct an
+      // SQL key.
+      key =
+          entity instanceof SqlEntity
+              ? VKey.from(Key.create(entity))
+              : VKey.createOfy(entity.getClass(), Key.create(entity));
     }
 
     /** Returns a human-readable ID string for the entity being changed. */
     public String getEntityId() {
       return String.format(
           "%s@%s",
-          key.getKind(),
+          key.getOfyKey().getKind(),
           // NB: try name before id, since name defaults to null, whereas id defaults to 0.
-          getNameOrId(key.getRaw()));
+          getNameOrId(key.getOfyKey().getRaw()));
     }
 
     /** Returns a string representation of this entity change. */
@@ -114,14 +123,12 @@ public abstract class MutatingCommand extends ConfirmingCommand implements Comma
   }
 
   /** Map from entity keys to EntityChange objects representing changes to those entities. */
-  private final LinkedHashMap<Key<ImmutableObject>, EntityChange> changedEntitiesMap =
+  private final LinkedHashMap<VKey<?>, EntityChange> changedEntitiesMap = new LinkedHashMap<>();
-      new LinkedHashMap<>();
 
   /** A set of resource keys for which new transactions should be created after. */
-  private final Set<Key<ImmutableObject>> transactionBoundaries = new HashSet<>();
+  private final Set<VKey<?>> transactionBoundaries = new HashSet<>();
 
-  @Nullable
+  @Nullable private VKey<?> lastAddedKey;
-  private Key<ImmutableObject> lastAddedKey;
 
   /**
    * Initializes the command.
@@ -154,20 +161,24 @@ 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.
-    ImmutableObject existingEntity = ofy().load().key(change.key).now();
+    Optional<?> existingEntity = tm().loadByKeyIfPresent(change.key);
     checkState(
-        Objects.equals(change.oldEntity, existingEntity),
+        Objects.equals(change.oldEntity, existingEntity.orElse(null)),
         "Entity changed since init() was called.\n%s",
         prettyPrintEntityDeepDiff(
             (change.oldEntity == null) ? ImmutableMap.of() : change.oldEntity.toDiffableFieldMap(),
-            (existingEntity == null) ? ImmutableMap.of() : existingEntity.toDiffableFieldMap()));
+            existingEntity.isPresent()
+                ? ((ImmutableObject) existingEntity.get()).toDiffableFieldMap()
+                : ImmutableMap.of()));
     switch (change.type) {
-      case CREATE: // Fall through.
+      case CREATE:
+        tm().insert(change.newEntity);
+        return;
       case UPDATE:
-        ofy().save().entity(change.newEntity).now();
+        tm().update(change.newEntity);
         return;
       case DELETE:
-        ofy().delete().key(change.key).now();
+        tm().delete(change.key);
         return;
     }
     throw new UnsupportedOperationException("Unknown entity change type: " + change.type);

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

@@ -26,6 +26,7 @@ import com.google.common.io.Files;
 import com.googlecode.objectify.Key;
 import google.registry.model.ImmutableObject;
 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;
@@ -148,4 +149,18 @@ abstract class ReadEntityFromKeyPathCommand<T> extends MutatingCommand {
     }
     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/test/java/google/registry/tools/DedupeOneTimeBillingEventIdsCommandTest.java

@@ -15,8 +15,8 @@
 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.DatabaseHelper.createTld;
+import static google.registry.testing.DatabaseHelper.loadAllOf;
 import static google.registry.testing.DatabaseHelper.persistActiveDomain;
 import static google.registry.testing.DatabaseHelper.persistResource;
 import static java.nio.charset.StandardCharsets.UTF_8;
@@ -38,7 +38,13 @@ import org.joda.money.Money;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
-/** Unit tests for {@link DedupeOneTimeBillingEventIdsCommand}. */
+/**
+ * Unit tests for {@link DedupeOneTimeBillingEventIdsCommand}.
+ *
+ * <p>Note that these are _not_ dual database tests even though the action has been converted. The
+ * dedupe was strictly a one-time event that needed to be done prior to moving to SQL. It should no
+ * longer be necessary and we may want to simply remove the command.
+ */
 class DedupeOneTimeBillingEventIdsCommandTest
     extends CommandTestCase<DedupeOneTimeBillingEventIdsCommand> {
 
@@ -64,8 +70,7 @@ class DedupeOneTimeBillingEventIdsCommandTest
         writeToNamedTmpFile("keypath.txt", getKeyPathLiteral(billingEventToResave)));
 
     int count = 0;
-    for (BillingEvent.OneTime billingEvent :
+    for (BillingEvent.OneTime billingEvent : loadAllOf(BillingEvent.OneTime.class)) {
-        ofy().load().type(BillingEvent.OneTime.class).ancestor(historyEntry)) {
       count++;
       assertThat(billingEvent.getId()).isNotEqualTo(billingEventToResave.getId());
       assertThat(billingEvent.asBuilder().setId(billingEventToResave.getId()).build())

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

@@ -14,18 +14,23 @@
 
 package google.registry.tools;
 
+import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.truth.Truth.assertThat;
-import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatabaseHelper.createTld;
 import static google.registry.testing.DatabaseHelper.deleteResource;
+import static google.registry.testing.DatabaseHelper.loadByEntity;
 import static google.registry.testing.DatabaseHelper.persistActiveHost;
 import static google.registry.testing.DatabaseHelper.persistNewRegistrar;
 import static google.registry.testing.DatabaseHelper.persistResource;
 import static org.joda.time.DateTimeZone.UTC;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
+import com.google.common.collect.ImmutableList;
+import com.googlecode.objectify.Key;
 import google.registry.model.host.HostResource;
 import google.registry.model.registrar.Registrar;
+import google.registry.persistence.VKey;
 import google.registry.testing.AppEngineExtension;
 import java.util.Arrays;
 import org.joda.time.DateTime;
@@ -103,15 +108,19 @@ public class MutatingCommandTest {
             + "blockPremiumNames: false -> true\n");
     String results = command.execute();
     assertThat(results).isEqualTo("Updated 4 entities.\n");
-    assertThat(ofy().load().entity(host1).now()).isEqualTo(newHost1);
+    assertThat(loadByEntity(host1)).isEqualTo(newHost1);
-    assertThat(ofy().load().entity(host2).now()).isEqualTo(newHost2);
+    assertThat(loadByEntity(host2)).isEqualTo(newHost2);
-    assertThat(ofy().load().entity(registrar1).now()).isEqualTo(newRegistrar1);
+    assertThat(loadByEntity(registrar1)).isEqualTo(newRegistrar1);
-    assertThat(ofy().load().entity(registrar2).now()).isEqualTo(newRegistrar2);
+    assertThat(loadByEntity(registrar2)).isEqualTo(newRegistrar2);
   }
 
   @Test
   void testSuccess_create() throws Exception {
-    ofy().deleteWithoutBackup().entities(Arrays.asList(host1, host2, registrar1, registrar2)).now();
+    ImmutableList<VKey<?>> keys =
+        Arrays.asList(host1, host2, registrar1, registrar2).stream()
+            .map(entity -> VKey.from(Key.create(entity)))
+            .collect(toImmutableList());
+    tm().transact(() -> tm().deleteWithoutBackup(keys));
     MutatingCommand command = new MutatingCommand() {
       @Override
       protected void init() {
@@ -137,10 +146,10 @@ public class MutatingCommandTest {
             + newRegistrar2 + "\n");
     String results = command.execute();
     assertThat(results).isEqualTo("Updated 4 entities.\n");
-    assertThat(ofy().load().entity(newHost1).now()).isEqualTo(newHost1);
+    assertThat(loadByEntity(newHost1)).isEqualTo(newHost1);
-    assertThat(ofy().load().entity(newHost2).now()).isEqualTo(newHost2);
+    assertThat(loadByEntity(newHost2)).isEqualTo(newHost2);
-    assertThat(ofy().load().entity(newRegistrar1).now()).isEqualTo(newRegistrar1);
+    assertThat(loadByEntity(newRegistrar1)).isEqualTo(newRegistrar1);
-    assertThat(ofy().load().entity(newRegistrar2).now()).isEqualTo(newRegistrar2);
+    assertThat(loadByEntity(newRegistrar2)).isEqualTo(newRegistrar2);
   }
 
   @Test
@@ -170,10 +179,10 @@ public class MutatingCommandTest {
             + registrar2 + "\n");
     String results = command.execute();
     assertThat(results).isEqualTo("Updated 4 entities.\n");
-    assertThat(ofy().load().entity(host1).now()).isNull();
+    assertThat(loadByEntity(host1)).isNull();
-    assertThat(ofy().load().entity(host2).now()).isNull();
+    assertThat(loadByEntity(host2)).isNull();
-    assertThat(ofy().load().entity(registrar1).now()).isNull();
+    assertThat(loadByEntity(registrar1)).isNull();
-    assertThat(ofy().load().entity(registrar2).now()).isNull();
+    assertThat(loadByEntity(registrar2)).isNull();
   }
 
   @Test
@@ -196,8 +205,8 @@ public class MutatingCommandTest {
             + "[no changes]\n");
     String results = command.execute();
     assertThat(results).isEqualTo("Updated 2 entities.\n");
-    assertThat(ofy().load().entity(host1).now()).isEqualTo(host1);
+    assertThat(loadByEntity(host1)).isEqualTo(host1);
-    assertThat(ofy().load().entity(registrar1).now()).isEqualTo(registrar1);
+    assertThat(loadByEntity(registrar1)).isEqualTo(registrar1);
   }
 
   @Test
@@ -230,10 +239,10 @@ public class MutatingCommandTest {
             + "blockPremiumNames: false -> true\n");
     String results = command.execute();
     assertThat(results).isEqualTo("Updated 4 entities.\n");
-    assertThat(ofy().load().entity(host1).now()).isNull();
+    assertThat(loadByEntity(host1)).isNull();
-    assertThat(ofy().load().entity(host2).now()).isEqualTo(newHost2);
+    assertThat(loadByEntity(host2)).isEqualTo(newHost2);
-    assertThat(ofy().load().entity(registrar1).now()).isNull();
+    assertThat(loadByEntity(registrar1)).isNull();
-    assertThat(ofy().load().entity(registrar2).now()).isEqualTo(newRegistrar2);
+    assertThat(loadByEntity(registrar2)).isEqualTo(newRegistrar2);
   }
 
   @Test
@@ -271,11 +280,11 @@ public class MutatingCommandTest {
 
     IllegalStateException thrown = assertThrows(IllegalStateException.class, command::execute);
     assertThat(thrown).hasMessageThat().contains("Entity changed since init() was called.");
-    assertThat(ofy().load().entity(host1).now()).isNull();
+    assertThat(loadByEntity(host1)).isNull();
-    assertThat(ofy().load().entity(host2).now()).isEqualTo(newHost2);
+    assertThat(loadByEntity(host2)).isEqualTo(newHost2);
     // These two shouldn't've changed.
-    assertThat(ofy().load().entity(registrar1).now()).isEqualTo(registrar1);
+    assertThat(loadByEntity(registrar1)).isEqualTo(registrar1);
-    assertThat(ofy().load().entity(registrar2).now()).isEqualTo(registrar2);
+    assertThat(loadByEntity(registrar2)).isEqualTo(registrar2);
   }
 
   @Test