Resurrect symmetric vkeys when possible. (#899)

* Resurrect symmetric vkeys when possible.

AbstractVKeyConverter had never been updated to generate symmetric VKeys for
simple (i.e. non-composite) VKey types.  Update it to default to generating a
VKey with a simple Ofy key.

With this change, composite VKeys must specify the "compositeKey = true" field
in the With*VKey annotations.  Doing so creates an asymmetric (SQL only) VKey
that can triggers higher-level logic to generate a corrected VKey containing
the composite Key.

core/src/main/java/google/registry/model/billing/BillingEvent.java

@@ -290,7 +290,7 @@ public abstract class BillingEvent extends ImmutableObject
         @javax.persistence.Index(columnList = "allocationToken")
       })
   @AttributeOverride(name = "id", column = @Column(name = "billing_event_id"))
-  @WithLongVKey
+  @WithLongVKey(compositeKey = true)
   public static class OneTime extends BillingEvent implements DatastoreAndSqlEntity {
 
     /** The billable value. */
@@ -464,7 +464,7 @@ public abstract class BillingEvent extends ImmutableObject
         @javax.persistence.Index(columnList = "recurrence_time_of_year")
       })
   @AttributeOverride(name = "id", column = @Column(name = "billing_recurrence_id"))
-  @WithLongVKey
+  @WithLongVKey(compositeKey = true)
   public static class Recurring extends BillingEvent implements DatastoreAndSqlEntity {
 
     /**
@@ -559,7 +559,7 @@ public abstract class BillingEvent extends ImmutableObject
         @javax.persistence.Index(columnList = "billingTime")
       })
   @AttributeOverride(name = "id", column = @Column(name = "billing_cancellation_id"))
-  @WithLongVKey
+  @WithLongVKey(compositeKey = true)
   public static class Cancellation extends BillingEvent implements DatastoreAndSqlEntity {
 
     /** The billing time of the charge that is being cancelled. */
@@ -680,7 +680,7 @@ public abstract class BillingEvent extends ImmutableObject
   /** An event representing a modification of an existing one-time billing event. */
   @ReportedOn
   @Entity
-  @WithLongVKey
+  @WithLongVKey(compositeKey = true)
   public static class Modification extends BillingEvent implements DatastoreOnlyEntity {
 
     /** The change in cost that should be applied to the original billing event. */

core/src/main/java/google/registry/model/poll/PollMessage.java

@@ -279,7 +279,7 @@ public abstract class PollMessage extends ImmutableObject
   @EntitySubclass(index = false)
   @javax.persistence.Entity
   @DiscriminatorValue("ONE_TIME")
-  @WithLongVKey
+  @WithLongVKey(compositeKey = true)
   public static class OneTime extends PollMessage {
 
     // Response data. Objectify cannot persist a base class type, so we must have a separate field
@@ -432,7 +432,7 @@ public abstract class PollMessage extends ImmutableObject
   @EntitySubclass(index = false)
   @javax.persistence.Entity
   @DiscriminatorValue("AUTORENEW")
-  @WithLongVKey
+  @WithLongVKey(compositeKey = true)
   public static class Autorenew extends PollMessage {
 
     /** The target id of the autorenew event. */

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

@@ -31,4 +31,12 @@ public @interface WithLongVKey {
    * class name will be "VKeyConverter_" concatenated with the suffix.
    */
   String classNameSuffix() default "";
+
+  /**
+   * Set to true if this is a composite vkey.
+   *
+   * <p>For composite VKeys, we don't attempt to define an objectify key when loading from SQL: the
+   * enclosing class has to take care of that.
+   */
+  boolean compositeKey() default false;
 }

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

@@ -31,4 +31,12 @@ public @interface WithStringVKey {
    * class name will be "VKeyConverter_" concatenated with the suffix.
    */
   String classNameSuffix() default "";
+
+  /**
+   * Set to true if this is a composite vkey.
+   *
+   * <p>For composite VKeys, we don't attempt to define an objectify key when loading from SQL: the
+   * enclosing class has to take care of that.
+   */
+  boolean compositeKey() default false;
 }

core/src/main/java/google/registry/persistence/converter/VKeyConverter.java

@@ -14,6 +14,7 @@
 
 package google.registry.persistence.converter;
 
+import com.googlecode.objectify.Key;
 import google.registry.persistence.VKey;
 import javax.annotation.Nullable;
 import javax.persistence.AttributeConverter;
@@ -29,7 +30,28 @@ public abstract class VKeyConverter<T, C> implements AttributeConverter<VKey<? e
   @Override
   @Nullable
   public VKey<? extends T> convertToEntityAttribute(@Nullable C dbData) {
-    return dbData == null ? null : VKey.createSql(getAttributeClass(), dbData);
+    if (dbData == null) {
+      return null;
+    }
+    Class<? extends T> clazz = getAttributeClass();
+    Key ofyKey = null;
+    if (!hasCompositeOfyKey()) {
+      // If this isn't a composite key, we can create the Ofy key from the SQL key.
+      ofyKey =
+          dbData instanceof String
+              ? Key.create(clazz, (String) dbData)
+              : Key.create(clazz, (Long) dbData);
+      return VKey.create(clazz, dbData, ofyKey);
+    } else {
+      // We don't know how to create the Ofy key and probably don't have everything necessary to do
+      // it anyway, so just create an asymmetric key - the containing object will have to convert it
+      // into a symmetric key.
+      return VKey.createSql(clazz, dbData);
+    }
+  }
+
+  protected boolean hasCompositeOfyKey() {
+    return false;
   }
 
   /** Returns the class of the attribute. */

core/src/test/java/google/registry/model/domain/DomainBaseSqlTest.java

@@ -84,10 +84,10 @@ public class DomainBaseSqlTest {
     saveRegistrar("registrar1");
     saveRegistrar("registrar2");
     saveRegistrar("registrar3");
-    contactKey = VKey.createSql(ContactResource.class, "contact_id1");
+    contactKey = createKey(ContactResource.class, "contact_id1");
-    contact2Key = VKey.createSql(ContactResource.class, "contact_id2");
+    contact2Key = createKey(ContactResource.class, "contact_id2");
 
-    host1VKey = VKey.createSql(HostResource.class, "host1");
+    host1VKey = createKey(HostResource.class, "host1");
 
     domain =
         new DomainBase.Builder()
@@ -696,6 +696,10 @@ public class DomainBaseSqlTest {
     assertThat(domain.getGracePeriods()).isEqualTo(gracePeriods);
   }
 
+  private <T> VKey<T> createKey(Class<T> clazz, String name) {
+    return VKey.create(clazz, name, Key.create(clazz, name));
+  }
+
   private <T> VKey<T> createLegacyVKey(Class<T> clazz, long id) {
     return VKey.create(
         clazz, id, Key.create(Key.create(EntityGroupRoot.class, "per-tld"), clazz, id));

core/src/test/java/google/registry/model/history/DomainHistoryTest.java

@@ -14,7 +14,6 @@
 
 package google.registry.model.history;
 
-import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.ImmutableObjectSubject.assertAboutImmutableObjects;
 import static google.registry.model.ImmutableObjectSubject.immutableObjectCorrespondence;
@@ -92,10 +91,7 @@ public class DomainHistoryTest extends EntityTestCase {
               assertDomainHistoriesEqual(fromDatabase, domainHistory);
               assertThat(fromDatabase.getParentVKey()).isEqualTo(domainHistory.getParentVKey());
               assertThat(fromDatabase.getNsHosts())
-                  .containsExactlyElementsIn(
+                  .containsExactlyElementsIn(domainHistory.getNsHosts());
-                      domainHistory.getNsHosts().stream()
-                          .map(key -> VKey.createSql(HostResource.class, key.getSqlKey()))
-                          .collect(toImmutableSet()));
             });
   }
 

core/src/test/java/google/registry/persistence/converter/LongVKeyConverterTest.java

@@ -19,8 +19,7 @@ import static google.registry.persistence.transaction.TransactionManagerFactory.
 
 import google.registry.persistence.VKey;
 import google.registry.persistence.WithLongVKey;
-import google.registry.persistence.transaction.JpaTestRules;
+import google.registry.testing.AppEngineExtension;
-import google.registry.persistence.transaction.JpaTestRules.JpaUnitTestExtension;
 import javax.persistence.Entity;
 import javax.persistence.Id;
 import org.junit.jupiter.api.Test;
@@ -30,33 +29,56 @@ import org.junit.jupiter.api.extension.RegisterExtension;
 public class LongVKeyConverterTest {
 
   @RegisterExtension
-  public final JpaUnitTestExtension jpaExtension =
+  public final AppEngineExtension appEngineExtension =
-      new JpaTestRules.Builder()
+      new AppEngineExtension.Builder()
-          .withEntityClass(TestEntity.class, VKeyConverter_LongType.class)
+          .withDatastoreAndCloudSql()
-          .buildUnitTestRule();
+          .withoutCannedData()
+          .withJpaUnitTestEntities(
+              TestLongEntity.class,
+              VKeyConverter_LongType.class,
+              VKeyConverter_CompositeLongType.class)
+          .withOfyTestEntities(TestLongEntity.class, CompositeKeyTestLongEntity.class)
+          .build();
 
   @Test
   void testRoundTrip() {
-    TestEntity original = new TestEntity(VKey.createSql(TestEntity.class, 10L));
+    TestLongEntity original =
+        new TestLongEntity(
+            VKey.createSql(TestLongEntity.class, 10L),
+            VKey.createSql(CompositeKeyTestLongEntity.class, 20L));
     jpaTm().transact(() -> jpaTm().getEntityManager().persist(original));
 
-    TestEntity retrieved =
+    TestLongEntity retrieved =
-        jpaTm().transact(() -> jpaTm().getEntityManager().find(TestEntity.class, "id"));
+        jpaTm().transact(() -> jpaTm().getEntityManager().find(TestLongEntity.class, "id"));
     assertThat(retrieved.number.getSqlKey()).isEqualTo(10L);
+    assertThat(retrieved.number.getOfyKey().getId()).isEqualTo(10L);
+
+    assertThat(retrieved.composite.getSqlKey()).isEqualTo(20L);
+    assertThat(retrieved.composite.maybeGetOfyKey().isPresent()).isFalse();
   }
 
-  @Entity(name = "TestEntity")
+  @Entity(name = "TestLongEntity")
+  @com.googlecode.objectify.annotation.Entity
   @WithLongVKey(classNameSuffix = "LongType")
-  static class TestEntity {
+  static class TestLongEntity {
-    @Id String id = "id";
+    @com.googlecode.objectify.annotation.Id @Id String id = "id";
 
-    VKey<TestEntity> number;
+    VKey<TestLongEntity> number;
+    VKey<CompositeKeyTestLongEntity> composite;
 
-    TestEntity(VKey<TestEntity> number) {
+    TestLongEntity(VKey<TestLongEntity> number, VKey<CompositeKeyTestLongEntity> composite) {
       this.number = number;
+      this.composite = composite;
     }
 
     /** Default constructor, needed for hibernate. */
-    public TestEntity() {}
+    public TestLongEntity() {}
+  }
+
+  @Entity(name = "CompositeKeyTestLongEntity")
+  @com.googlecode.objectify.annotation.Entity
+  @WithLongVKey(classNameSuffix = "CompositeLongType", compositeKey = true)
+  static class CompositeKeyTestLongEntity {
+    @com.googlecode.objectify.annotation.Id @Id String id = "id";
   }
 }

core/src/test/java/google/registry/persistence/converter/StringVKeyConverterTest.java

@@ -19,8 +19,7 @@ import static google.registry.persistence.transaction.TransactionManagerFactory.
 
 import google.registry.persistence.VKey;
 import google.registry.persistence.WithStringVKey;
-import google.registry.persistence.transaction.JpaTestRules;
+import google.registry.testing.AppEngineExtension;
-import google.registry.persistence.transaction.JpaTestRules.JpaUnitTestExtension;
 import javax.persistence.Entity;
 import javax.persistence.Id;
 import org.junit.jupiter.api.Test;
@@ -30,36 +29,61 @@ import org.junit.jupiter.api.extension.RegisterExtension;
 public class StringVKeyConverterTest {
 
   @RegisterExtension
-  public final JpaUnitTestExtension jpaExtension =
+  public final AppEngineExtension appEngineExtension =
-      new JpaTestRules.Builder()
+      new AppEngineExtension.Builder()
-          .withEntityClass(TestEntity.class, VKeyConverter_StringType.class)
+          .withDatastoreAndCloudSql()
-          .buildUnitTestRule();
+          .withoutCannedData()
+          .withJpaUnitTestEntities(
+              TestStringEntity.class,
+              VKeyConverter_StringType.class,
+              VKeyConverter_CompositeStringType.class)
+          .withOfyTestEntities(TestStringEntity.class, CompositeKeyTestStringEntity.class)
+          .build();
 
   @Test
   void testRoundTrip() {
-    TestEntity original =
+    TestStringEntity original =
-        new TestEntity("TheRealSpartacus", VKey.createSql(TestEntity.class, "ImSpartacus!"));
+        new TestStringEntity(
+            "TheRealSpartacus",
+            VKey.createSql(TestStringEntity.class, "ImSpartacus!"),
+            VKey.createSql(CompositeKeyTestStringEntity.class, "NoImSpartacus!"));
     jpaTm().transact(() -> jpaTm().getEntityManager().persist(original));
 
-    TestEntity retrieved =
+    TestStringEntity retrieved =
         jpaTm()
-            .transact(() -> jpaTm().getEntityManager().find(TestEntity.class, "TheRealSpartacus"));
+            .transact(
+                () -> jpaTm().getEntityManager().find(TestStringEntity.class, "TheRealSpartacus"));
     assertThat(retrieved.other.getSqlKey()).isEqualTo("ImSpartacus!");
+    assertThat(retrieved.other.getOfyKey().getName()).isEqualTo("ImSpartacus!");
+
+    assertThat(retrieved.composite.getSqlKey()).isEqualTo("NoImSpartacus!");
+    assertThat(retrieved.composite.maybeGetOfyKey().isPresent()).isFalse();
   }
 
-  @Entity(name = "TestEntity")
+  @Entity(name = "TestStringEntity")
+  @com.googlecode.objectify.annotation.Entity
   @WithStringVKey(classNameSuffix = "StringType")
-  static class TestEntity {
+  static class TestStringEntity {
-    @Id String id;
+    @com.googlecode.objectify.annotation.Id @Id String id;
 
-    VKey<TestEntity> other;
+    VKey<TestStringEntity> other;
+    VKey<CompositeKeyTestStringEntity> composite;
 
-    TestEntity(String id, VKey<TestEntity> other) {
+    TestStringEntity(
+        String id, VKey<TestStringEntity> other, VKey<CompositeKeyTestStringEntity> composite) {
       this.id = id;
       this.other = other;
+      this.composite = composite;
     }
 
     /** Default constructor, needed for hibernate. */
-    public TestEntity() {}
+    public TestStringEntity() {}
+  }
+
+  @Entity(name = "CompositeKeyTestStringEntity")
+  @com.googlecode.objectify.annotation.Entity
+  @WithStringVKey(classNameSuffix = "CompositeStringType", compositeKey = true)
+  static class CompositeKeyTestStringEntity {
+    @com.googlecode.objectify.annotation.Id @Id String id = "id";
   }
 }

processor/build.gradle

@@ -18,6 +18,12 @@ plugins {
 
 dependencies {
     def deps = rootProject.dependencyMap
+
+    // Custom-built objectify jar at commit ecd5165, included in Nomulus
+    // release.
+    compile files(
+        "${rootDir}/third_party/objectify/v4_1/objectify-4.1.3.jar")
+
     compile deps['com.google.code.findbugs:jsr305']
     compile deps['com.google.guava:guava']
     compile deps['com.squareup:javapoet']

processor/src/main/java/google/registry/processors/AbstractVKeyProcessor.java

@@ -27,11 +27,14 @@ import com.squareup.javapoet.TypeSpec;
 import java.io.IOException;
 import java.io.UncheckedIOException;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import javax.annotation.processing.AbstractProcessor;
 import javax.annotation.processing.RoundEnvironment;
 import javax.lang.model.element.AnnotationMirror;
+import javax.lang.model.element.AnnotationValue;
 import javax.lang.model.element.Element;
+import javax.lang.model.element.ExecutableElement;
 import javax.lang.model.element.Modifier;
 import javax.lang.model.element.TypeElement;
 import javax.lang.model.type.DeclaredType;
@@ -47,9 +50,12 @@ import javax.persistence.Converter;
 public abstract class AbstractVKeyProcessor extends AbstractProcessor {
 
   private static final String CONVERTER_CLASS_NAME_TEMP = "VKeyConverter_%s";
-  // The method with same name should be defined in StringVKey and LongVKey
+  // The method with same name should be defined in WithStringVKey and WithLongVKey
   private static final String CLASS_NAME_SUFFIX_KEY = "classNameSuffix";
 
+  // Method in WithStringVKey and WithLongVKey to indicate that this is a composite key.
+  private static final String COMPOSITE_KEY_KEY = "compositeKey";
+
   abstract Class<?> getSqlColumnType();
 
   abstract String getAnnotationSimpleName();
@@ -76,18 +82,18 @@ public abstract class AbstractVKeyProcessor extends AbstractProcessor {
                         actualAnnotation.size() == 1,
                         String.format(
                             "type can have only 1 %s annotation", getAnnotationSimpleName()));
-                    String converterClassNameSuffix =
+                    String converterClassNameSuffix = "";
-                        actualAnnotation.get(0).getElementValues().entrySet().stream()
+                    boolean hasCompositeOfyKey = false;
-                            .filter(
+                    for (Map.Entry<? extends ExecutableElement, ? extends AnnotationValue> entry :
-                                entry ->
+                        actualAnnotation.get(0).getElementValues().entrySet()) {
-                                    entry
+                      String keyName = entry.getKey().getSimpleName().toString();
-                                        .getKey()
+                      Object value = entry.getValue().getValue();
-                                        .getSimpleName()
+                      if (keyName.equals(CLASS_NAME_SUFFIX_KEY)) {
-                                        .toString()
+                        converterClassNameSuffix = ((String) value).trim();
-                                        .equals(CLASS_NAME_SUFFIX_KEY))
+                      } else if (keyName.equals(COMPOSITE_KEY_KEY)) {
-                            .map(entry -> ((String) entry.getValue().getValue()).trim())
+                        hasCompositeOfyKey = (Boolean) value;
-                            .findFirst()
+                      }
-                            .orElse("");
+                    }
                     if (converterClassNameSuffix.isEmpty()) {
                       converterClassNameSuffix =
                           getTypeUtils().asElement(entityType).getSimpleName().toString();
@@ -97,7 +103,8 @@ public abstract class AbstractVKeyProcessor extends AbstractProcessor {
                       createJavaFile(
                               getPackageName(annotatedTypeElement),
                               String.format(CONVERTER_CLASS_NAME_TEMP, converterClassNameSuffix),
-                              entityType)
+                              entityType,
+                              hasCompositeOfyKey)
                           .writeTo(processingEnv.getFiler());
                     } catch (IOException e) {
                       throw new UncheckedIOException(e);
@@ -108,7 +115,10 @@ public abstract class AbstractVKeyProcessor extends AbstractProcessor {
   }
 
   private JavaFile createJavaFile(
-      String packageName, String converterClassName, TypeMirror entityTypeMirror) {
+      String packageName,
+      String converterClassName,
+      TypeMirror entityTypeMirror,
+      boolean hasCompositeOfyKey) {
     TypeName entityType = ClassName.get(entityTypeMirror);
 
     ParameterizedTypeName attributeConverter =
@@ -127,7 +137,7 @@ public abstract class AbstractVKeyProcessor extends AbstractProcessor {
             .addStatement("return $T.class", entityType)
             .build();
 
-    TypeSpec vKeyConverter =
+    TypeSpec.Builder classBuilder =
         TypeSpec.classBuilder(converterClassName)
             .addAnnotation(
                 AnnotationSpec.builder(ClassName.get(Converter.class))
@@ -135,9 +145,24 @@ public abstract class AbstractVKeyProcessor extends AbstractProcessor {
                     .build())
             .addModifiers(Modifier.FINAL)
             .superclass(attributeConverter)
-            .addMethod(getAttributeClass)
+            .addMethod(getAttributeClass);
+
+    // If this is a converter for a composite vkey type, generate an override for the default
+    // {@link google.registry.persistence.VKeyConverter.hasCompositeOfyKey()} method, which returns
+    // false.
+    if (hasCompositeOfyKey) {
+      MethodSpec hasCompositeOfyKeyMethod =
+          MethodSpec.methodBuilder("hasCompositeOfyKey")
+              .addAnnotation(Override.class)
+              .addModifiers(Modifier.PROTECTED)
+              .returns(boolean.class)
+              .addStatement("return true", entityType)
               .build();
 
+      classBuilder.addMethod(hasCompositeOfyKeyMethod);
+    }
+
+    TypeSpec vKeyConverter = classBuilder.build();
     return JavaFile.builder(packageName, vKeyConverter).build();
   }