From 71f86c9970c1d0c3f2643f270e0dabe05059eb27 Mon Sep 17 00:00:00 2001 From: Michael Muller Date: Wed, 30 Sep 2020 11:15:58 -0400 Subject: [PATCH] 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. --- .../google/registry/persistence/VKey.java | 77 +++++++++++++++++++ .../google/registry/persistence/VKeyTest.java | 57 ++++++++++++++ 2 files changed, 134 insertions(+) diff --git a/core/src/main/java/google/registry/persistence/VKey.java b/core/src/main/java/google/registry/persistence/VKey.java index d12bfdd9d..b8fda844f 100644 --- a/core/src/main/java/google/registry/persistence/VKey.java +++ b/core/src/main/java/google/registry/persistence/VKey.java @@ -113,6 +113,83 @@ public class VKey extends ImmutableObject implements Serializable { return new VKey(kind, Key.create(kind, name), name); } + /** + * Returns a clone with an ofy key restored from {@code ancestors}. + * + *

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. + * + *

For example, to restore the objectify key for + * DomainBase("COM-1234")/HistoryEntry(123)/PollEvent(567), one might use: + * + *

{@code
+   * pollEvent.restoreOfy(DomainBase.class, "COM-1234", HistoryEntry.class, 567)
+   * }
+ * + *

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. + * + *

As a special case, an objectify Key may be used as the first ancestor instead of a Class, + * value pair. + */ + public VKey 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 ofyKey = + sqlKey instanceof Long + ? Key.create(lastKey, getKind(), (Long) sqlKey) + : Key.create(lastKey, getKind(), (String) sqlKey); + + return VKey.create((Class) getKind(), sqlKey, ofyKey); + } + + /** + * Returns a clone of {@code key} with an ofy key restored from {@code ancestors}. + * + *

This is the static form of the method restoreOfy() above. If {@code key} is null, it returns + * null. + */ + public static VKey restoreOfyFrom(@Nullable VKey key, Object... ancestors) { + return key == null ? null : key.restoreOfy(ancestors); + } + /** Returns the type of the entity. */ public Class getKind() { return this.kind; diff --git a/core/src/test/java/google/registry/persistence/VKeyTest.java b/core/src/test/java/google/registry/persistence/VKeyTest.java index ca392cda1..ad8d2d00c 100644 --- a/core/src/test/java/google/registry/persistence/VKeyTest.java +++ b/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 key = VKey.createSql(TestObject.class, "foo"); + VKey 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 {} }