Add VKey.restoreOfy() method for fixing ofy keys (#820)

Add a restoreOfy() instance method and a restoreOfyFrom() static method to
assist in restoring the objectify key for classes that have composite keys
that do not restore automatically.

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

@@ -113,6 +113,83 @@ public class VKey<T> extends ImmutableObject implements Serializable {
     return new VKey<T>(kind, Key.create(kind, name), name);
   }
 
+  /**
+   * Returns a clone with an ofy key restored from {@code ancestors}.
+   *
+   * <p>The arguments should generally consist of pairs of Class and value, where the Class is the
+   * kind of the ancestor key and the value is either a String or a Long.
+   *
+   * <p>For example, to restore the objectify key for
+   * DomainBase("COM-1234")/HistoryEntry(123)/PollEvent(567), one might use:
+   *
+   * <pre>{@code
+   * pollEvent.restoreOfy(DomainBase.class, "COM-1234", HistoryEntry.class, 567)
+   * }</pre>
+   *
+   * <p>The final key id or name is obtained from the SQL key. It is assumed that this value must be
+   * either a long integer or a {@code String} and that this proper identifier for the objectify
+   * key.
+   *
+   * <p>As a special case, an objectify Key may be used as the first ancestor instead of a Class,
+   * value pair.
+   */
+  public VKey<T> restoreOfy(Object... ancestors) {
+    Class lastClass = null;
+    Key<?> lastKey = null;
+    for (Object ancestor : ancestors) {
+      if (ancestor instanceof Class) {
+        if (lastClass != null) {
+          throw new IllegalArgumentException(ancestor + " used as a key value.");
+        }
+        lastClass = (Class) ancestor;
+        continue;
+      } else if (ancestor instanceof Key) {
+        if (lastKey != null) {
+          throw new IllegalArgumentException(
+              "Objectify keys may only be used for the first argument");
+        }
+        lastKey = (Key) ancestor;
+        continue;
+      }
+
+      // The argument should be a value.
+      if (lastClass == null) {
+        throw new IllegalArgumentException("Argument " + ancestor + " should be a class.");
+      }
+      if (ancestor instanceof Long) {
+        lastKey = Key.create(lastKey, lastClass, (Long) ancestor);
+      } else if (ancestor instanceof String) {
+        lastKey = Key.create(lastKey, lastClass, (String) ancestor);
+      } else {
+        throw new IllegalArgumentException("Key value " + ancestor + " must be a string or long.");
+      }
+      lastClass = null;
+    }
+
+    // Make sure we didn't end up with a dangling class with no value.
+    if (lastClass != null) {
+      throw new IllegalArgumentException("Missing value for last key of type " + lastClass);
+    }
+
+    Object sqlKey = getSqlKey();
+    Key<T> ofyKey =
+        sqlKey instanceof Long
+            ? Key.create(lastKey, getKind(), (Long) sqlKey)
+            : Key.create(lastKey, getKind(), (String) sqlKey);
+
+    return VKey.create((Class<T>) getKind(), sqlKey, ofyKey);
+  }
+
+  /**
+   * Returns a clone of {@code key} with an ofy key restored from {@code ancestors}.
+   *
+   * <p>This is the static form of the method restoreOfy() above. If {@code key} is null, it returns
+   * null.
+   */
+  public static <T> VKey<T> restoreOfyFrom(@Nullable VKey<T> key, Object... ancestors) {
+    return key == null ? null : key.restoreOfy(ancestors);
+  }
+
   /** Returns the type of the entity. */
   public Class<? extends T> getKind() {
     return this.kind;

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

@@ -18,6 +18,7 @@ import static com.google.common.truth.Truth8.assertThat;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
 import com.googlecode.objectify.Key;
+import com.googlecode.objectify.annotation.Entity;
 import google.registry.model.billing.BillingEvent.OneTime;
 import google.registry.model.registrar.RegistrarContact;
 import google.registry.testing.AppEngineExtension;
@@ -59,4 +60,60 @@ class VKeyTest {
             () -> VKey.create(RegistrarContact.class, "fake@example.com"));
     assertThat(thrown).hasMessageThat().contains("BackupGroupRoot");
   }
+
+  @Test
+  void testRestoreOfy() {
+    assertThat(VKey.restoreOfyFrom(null, TestObject.class, 100)).isNull();
+
+    VKey<TestObject> key = VKey.createSql(TestObject.class, "foo");
+    VKey<TestObject> restored = key.restoreOfy(TestObject.class, "bar");
+    assertThat(restored.getOfyKey())
+        .isEqualTo(Key.create(Key.create(TestObject.class, "bar"), TestObject.class, "foo"));
+    assertThat(restored.getSqlKey()).isEqualTo("foo");
+
+    assertThat(VKey.restoreOfyFrom(key).getOfyKey()).isEqualTo(Key.create(TestObject.class, "foo"));
+
+    restored = key.restoreOfy(OtherObject.class, "baz", TestObject.class, "bar");
+    assertThat(restored.getOfyKey())
+        .isEqualTo(
+            Key.create(
+                Key.create(Key.create(OtherObject.class, "baz"), TestObject.class, "bar"),
+                TestObject.class,
+                "foo"));
+
+    // Verify that we can use a key as the first argument.
+    restored = key.restoreOfy(Key.create(TestObject.class, "bar"));
+    assertThat(restored.getOfyKey())
+        .isEqualTo(Key.create(Key.create(TestObject.class, "bar"), TestObject.class, "foo"));
+
+    // Verify that we get an exception when a key is not the first argument.
+    IllegalArgumentException thrown =
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> key.restoreOfy(TestObject.class, "foo", Key.create(TestObject.class, "bar")));
+    assertThat(thrown)
+        .hasMessageThat()
+        .contains("Objectify keys may only be used for the first argument");
+
+    // Verify other exception cases.
+    thrown =
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> key.restoreOfy(TestObject.class, TestObject.class));
+    assertThat(thrown)
+        .hasMessageThat()
+        .contains("class google.registry.testing.TestObject used as a key value.");
+
+    thrown =
+        assertThrows(IllegalArgumentException.class, () -> key.restoreOfy(TestObject.class, 1.5));
+    assertThat(thrown).hasMessageThat().contains("Key value 1.5 must be a string or long.");
+
+    thrown = assertThrows(IllegalArgumentException.class, () -> key.restoreOfy(TestObject.class));
+    assertThat(thrown)
+        .hasMessageThat()
+        .contains("Missing value for last key of type class google.registry.testing.TestObject");
+  }
+
+  @Entity
+  static class OtherObject {}
 }