Move "WithLongVKey" to BillingEvent subclasses (#821)

When loading the VKeys for the BillingEvents hierarchy, it is necessary to
restore the original concrete class for the type, otherwise we end up with a
different (and incompatible) VKey.

As part of this, convert the cancellation matching billing event to
VKey<Recurring>, which seems like the only thing it actually can be.

core/src/main/java/google/registry/beam/invoicing/sql/billing_events.sql

@@ -55,7 +55,7 @@ FROM (
     FROM
       `%PROJECT_ID%.%DATASTORE_EXPORT_DATA_SET%.%REGISTRY_TABLE%`
     WHERE
-      enableInvoicing IS TRUE) ) AS BillingEvent
+      invoicingEnabled IS TRUE) ) AS BillingEvent
   -- Gather billing ID from registrar table
   -- This is a 'JOIN' as opposed to 'LEFT JOIN' to filter out
   -- non-billable registrars

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

@@ -67,7 +67,6 @@ import org.joda.time.DateTime;
 
 /** A billable event in a domain's lifecycle. */
 @MappedSuperclass
-@WithLongVKey
 public abstract class BillingEvent extends ImmutableObject
     implements Buildable, TransferServerApproveEntity {
 
@@ -292,6 +291,7 @@ public abstract class BillingEvent extends ImmutableObject
         @javax.persistence.Index(columnList = "allocation_token_id")
       })
   @AttributeOverride(name = "id", column = @Column(name = "billing_event_id"))
+  @WithLongVKey
   public static class OneTime extends BillingEvent implements DatastoreAndSqlEntity {
 
     /** The billable value. */
@@ -327,7 +327,7 @@ public abstract class BillingEvent extends ImmutableObject
      * Cancellation}s.
      */
     @Column(name = "cancellation_matching_billing_recurrence_id")
-    VKey<? extends BillingEvent> cancellationMatchingBillingEvent;
+    VKey<Recurring> cancellationMatchingBillingEvent;
 
     /**
      * The {@link AllocationToken} used in the creation of this event, or null if one was not used.
@@ -409,7 +409,7 @@ public abstract class BillingEvent extends ImmutableObject
       }
 
       public Builder setCancellationMatchingBillingEvent(
-          VKey<? extends BillingEvent> cancellationMatchingBillingEvent) {
+          VKey<Recurring> cancellationMatchingBillingEvent) {
         getInstance().cancellationMatchingBillingEvent = cancellationMatchingBillingEvent;
         return this;
       }
@@ -468,6 +468,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
   public static class Recurring extends BillingEvent implements DatastoreAndSqlEntity {
 
     /**
@@ -562,6 +563,7 @@ public abstract class BillingEvent extends ImmutableObject
         @javax.persistence.Index(columnList = "billingTime")
       })
   @AttributeOverride(name = "id", column = @Column(name = "billing_cancellation_id"))
+  @WithLongVKey
   public static class Cancellation extends BillingEvent implements DatastoreAndSqlEntity {
 
     /** The billing time of the charge that is being cancelled. */
@@ -682,6 +684,7 @@ public abstract class BillingEvent extends ImmutableObject
   /** An event representing a modification of an existing one-time billing event. */
   @ReportedOn
   @Entity
+  @WithLongVKey
   public static class Modification extends BillingEvent implements DatastoreAndSqlEntity {
 
     /** The change in cost that should be applied to the original billing event. */

core/src/main/java/google/registry/model/contact/ContactHistory.java

@@ -17,8 +17,11 @@ package google.registry.model.contact;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.EntitySubclass;
 import google.registry.model.EppResource;
+import google.registry.model.ImmutableObject;
+import google.registry.model.contact.ContactHistory.ContactHistoryId;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.persistence.VKey;
+import java.io.Serializable;
 import java.util.Optional;
 import javax.annotation.Nullable;
 import javax.persistence.Access;
@@ -28,6 +31,7 @@ import javax.persistence.Entity;
 import javax.persistence.GeneratedValue;
 import javax.persistence.GenerationType;
 import javax.persistence.Id;
+import javax.persistence.IdClass;
 import javax.persistence.PostLoad;
 
 /**
@@ -47,13 +51,13 @@ import javax.persistence.PostLoad;
     })
 @EntitySubclass
 @Access(AccessType.FIELD)
+@IdClass(ContactHistoryId.class)
 public class ContactHistory extends HistoryEntry {
 
   // Store ContactBase instead of ContactResource so we don't pick up its @Id
   @Nullable ContactBase contactBase;
 
-  @Column(nullable = false)
-  VKey<ContactResource> contactRepoId;
+  @Id String contactRepoId;
 
   @Id
   @GeneratedValue(strategy = GenerationType.SEQUENCE, generator = "TempHistorySequenceGenerator")
@@ -76,7 +80,14 @@ public class ContactHistory extends HistoryEntry {
 
   /** The key to the {@link ContactResource} this is based off of. */
   public VKey<ContactResource> getContactRepoId() {
-    return contactRepoId;
+    return VKey.create(
+        ContactResource.class, contactRepoId, Key.create(ContactResource.class, contactRepoId));
+  }
+
+  /** Creates a {@link VKey} instance for this entity. */
+  public VKey<ContactHistory> createVKey() {
+    return VKey.create(
+        ContactHistory.class, new ContactHistoryId(contactRepoId, getId()), Key.create(this));
   }
 
   @PostLoad
@@ -87,10 +98,65 @@ public class ContactHistory extends HistoryEntry {
       contactBase = null;
     }
     // Fill in the full, symmetric, parent repo ID key
-    Key<ContactResource> parentKey =
-        Key.create(ContactResource.class, (String) contactRepoId.getSqlKey());
-    parent = parentKey;
-    contactRepoId = VKey.create(ContactResource.class, contactRepoId.getSqlKey(), parentKey);
+    parent = Key.create(ContactResource.class, contactRepoId);
+  }
+
+  /** Class to represent the composite primary key of {@link ContactHistory} entity. */
+  static class ContactHistoryId extends ImmutableObject implements Serializable {
+
+    private String contactRepoId;
+
+    private Long id;
+
+    /** Hibernate requires this default constructor. */
+    private ContactHistoryId() {}
+
+    ContactHistoryId(String contactRepoId, long id) {
+      this.contactRepoId = contactRepoId;
+      this.id = id;
+    }
+
+    /**
+     * Returns the contact repository id.
+     *
+     * <p>This method is private because it is only used by Hibernate.
+     */
+    @SuppressWarnings("unused")
+    private String getContactRepoId() {
+      return contactRepoId;
+    }
+
+    /**
+     * Returns the history revision id.
+     *
+     * <p>This method is private because it is only used by Hibernate.
+     */
+    @SuppressWarnings("unused")
+    private long getId() {
+      return id;
+    }
+
+    /**
+     * Sets the contact repository id.
+     *
+     * <p>This method is private because it is only used by Hibernate and should not be used
+     * externally to keep immutability.
+     */
+    @SuppressWarnings("unused")
+    private void setContactRepoId(String contactRepoId) {
+      this.contactRepoId = contactRepoId;
+    }
+
+    /**
+     * Sets the history revision id.
+     *
+     * <p>This method is private because it is only used by Hibernate and should not be used
+     * externally to keep immutability.
+     */
+    @SuppressWarnings("unused")
+    private void setId(long id) {
+      this.id = id;
+    }
   }
 
   @Override
@@ -111,9 +177,9 @@ public class ContactHistory extends HistoryEntry {
       return this;
     }
 
-    public Builder setContactRepoId(VKey<ContactResource> contactRepoId) {
+    public Builder setContactRepoId(String contactRepoId) {
       getInstance().contactRepoId = contactRepoId;
-      contactRepoId.maybeGetOfyKey().ifPresent(parent -> getInstance().parent = parent);
+      getInstance().parent = Key.create(ContactResource.class, contactRepoId);
       return this;
     }
 
@@ -121,8 +187,7 @@ public class ContactHistory extends HistoryEntry {
     @Override
     public Builder setParent(Key<? extends EppResource> parent) {
       super.setParent(parent);
-      getInstance().contactRepoId =
-          VKey.create(ContactResource.class, parent.getName(), (Key<ContactResource>) parent);
+      getInstance().contactRepoId = parent.getName();
       return this;
     }
   }

core/src/main/java/google/registry/model/contact/ContactResource.java

@@ -51,7 +51,6 @@ public class ContactResource extends ContactBase
 
   @Override
   public VKey<ContactResource> createVKey() {
-    // TODO(mmuller): create symmetric keys if we can ever reload both sides.
     return VKey.create(ContactResource.class, getRepoId(), Key.create(this));
   }
 

core/src/main/java/google/registry/model/domain/DomainHistory.java

@@ -156,8 +156,10 @@ public class DomainHistory extends HistoryEntry {
     return VKey.create(DomainBase.class, domainRepoId, Key.create(DomainBase.class, domainRepoId));
   }
 
+  /** Creates a {@link VKey} instance for this entity. */
   public VKey<DomainHistory> createVKey() {
-    return VKey.createSql(DomainHistory.class, new DomainHistoryId(domainRepoId, getId()));
+    return VKey.create(
+        DomainHistory.class, new DomainHistoryId(domainRepoId, getId()), Key.create(this));
   }
 
   @PostLoad
@@ -188,19 +190,45 @@ public class DomainHistory extends HistoryEntry {
       this.id = id;
     }
 
-    String getDomainRepoId() {
+    /**
+     * Returns the domain repository id.
+     *
+     * <p>This method is private because it is only used by Hibernate.
+     */
+    @SuppressWarnings("unused")
+    private String getDomainRepoId() {
       return domainRepoId;
     }
 
-    void setDomainRepoId(String domainRepoId) {
-      this.domainRepoId = domainRepoId;
-    }
-
-    long getId() {
+    /**
+     * Returns the history revision id.
+     *
+     * <p>This method is private because it is only used by Hibernate.
+     */
+    @SuppressWarnings("unused")
+    private long getId() {
       return id;
     }
 
-    void setId(long id) {
+    /**
+     * Sets the domain repository id.
+     *
+     * <p>This method is private because it is only used by Hibernate and should not be used
+     * externally to keep immutability.
+     */
+    @SuppressWarnings("unused")
+    private void setDomainRepoId(String domainRepoId) {
+      this.domainRepoId = domainRepoId;
+    }
+
+    /**
+     * Sets the history revision id.
+     *
+     * <p>This method is private because it is only used by Hibernate and should not be used
+     * externally to keep immutability.
+     */
+    @SuppressWarnings("unused")
+    private void setId(long id) {
       this.id = id;
     }
   }

core/src/main/java/google/registry/model/host/HostHistory.java

@@ -14,11 +14,15 @@
 
 package google.registry.model.host;
 
+
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.EntitySubclass;
 import google.registry.model.EppResource;
+import google.registry.model.ImmutableObject;
+import google.registry.model.host.HostHistory.HostHistoryId;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.persistence.VKey;
+import java.io.Serializable;
 import java.util.Optional;
 import javax.annotation.Nullable;
 import javax.persistence.Access;
@@ -28,6 +32,7 @@ import javax.persistence.Entity;
 import javax.persistence.GeneratedValue;
 import javax.persistence.GenerationType;
 import javax.persistence.Id;
+import javax.persistence.IdClass;
 import javax.persistence.PostLoad;
 
 /**
@@ -48,13 +53,13 @@ import javax.persistence.PostLoad;
     })
 @EntitySubclass
 @Access(AccessType.FIELD)
+@IdClass(HostHistoryId.class)
 public class HostHistory extends HistoryEntry {
 
   // Store HostBase instead of HostResource so we don't pick up its @Id
   @Nullable HostBase hostBase;
 
-  @Column(nullable = false)
-  VKey<HostResource> hostRepoId;
+  @Id String hostRepoId;
 
   @Id
   @GeneratedValue(strategy = GenerationType.SEQUENCE, generator = "TempHistorySequenceGenerator")
@@ -77,7 +82,12 @@ public class HostHistory extends HistoryEntry {
 
   /** The key to the {@link google.registry.model.host.HostResource} this is based off of. */
   public VKey<HostResource> getHostRepoId() {
-    return hostRepoId;
+    return VKey.create(HostResource.class, hostRepoId, Key.create(HostResource.class, hostRepoId));
+  }
+
+  /** Creates a {@link VKey} instance for this entity. */
+  public VKey<HostHistory> createVKey() {
+    return VKey.create(HostHistory.class, new HostHistoryId(hostRepoId, getId()), Key.create(this));
   }
 
   @PostLoad
@@ -88,9 +98,65 @@ public class HostHistory extends HistoryEntry {
       hostBase = null;
     }
     // Fill in the full, symmetric, parent repo ID key
-    Key<HostResource> parentKey = Key.create(HostResource.class, (String) hostRepoId.getSqlKey());
-    parent = parentKey;
-    hostRepoId = VKey.create(HostResource.class, hostRepoId.getSqlKey(), parentKey);
+    parent = Key.create(HostResource.class, hostRepoId);
+  }
+
+  /** Class to represent the composite primary key of {@link HostHistory} entity. */
+  static class HostHistoryId extends ImmutableObject implements Serializable {
+
+    private String hostRepoId;
+
+    private Long id;
+
+    /** Hibernate requires this default constructor. */
+    private HostHistoryId() {}
+
+    HostHistoryId(String hostRepoId, long id) {
+      this.hostRepoId = hostRepoId;
+      this.id = id;
+    }
+
+    /**
+     * Returns the host repository id.
+     *
+     * <p>This method is private because it is only used by Hibernate.
+     */
+    @SuppressWarnings("unused")
+    private String getHostRepoId() {
+      return hostRepoId;
+    }
+
+    /**
+     * Returns the history revision id.
+     *
+     * <p>This method is private because it is only used by Hibernate.
+     */
+    @SuppressWarnings("unused")
+    private long getId() {
+      return id;
+    }
+
+    /**
+     * Sets the host repository id.
+     *
+     * <p>This method is private because it is only used by Hibernate and should not be used
+     * externally to keep immutability.
+     */
+    @SuppressWarnings("unused")
+    private void setHostRepoId(String hostRepoId) {
+      this.hostRepoId = hostRepoId;
+    }
+
+    /**
+     * Sets the history revision id.
+     *
+     * <p>This method is private because it is only used by Hibernate and should not be used
+     * externally to keep immutability.
+     */
+    @SuppressWarnings("unused")
+    private void setId(long id) {
+      this.id = id;
+    }
   }
 
   @Override
@@ -111,9 +177,9 @@ public class HostHistory extends HistoryEntry {
       return this;
     }
 
-    public Builder setHostRepoId(VKey<HostResource> hostRepoId) {
+    public Builder setHostRepoId(String hostRepoId) {
       getInstance().hostRepoId = hostRepoId;
-      hostRepoId.maybeGetOfyKey().ifPresent(parent -> getInstance().parent = parent);
+      getInstance().parent = Key.create(HostResource.class, hostRepoId);
       return this;
     }
 
@@ -121,8 +187,7 @@ public class HostHistory extends HistoryEntry {
     @Override
     public Builder setParent(Key<? extends EppResource> parent) {
       super.setParent(parent);
-      getInstance().hostRepoId =
-          VKey.create(HostResource.class, parent.getName(), (Key<HostResource>) parent);
+      getInstance().hostRepoId = parent.getName();
       return this;
     }
   }

core/src/main/java/google/registry/model/reporting/HistoryEntry.java

@@ -308,19 +308,10 @@ public class HistoryEntry extends ImmutableObject implements Buildable, Datastor
           new DomainHistory.Builder().copyFrom(this).setDomainRepoId(parent.getName()).build();
     } else if (parentKind.equals(getKind(HostResource.class))) {
       resultEntity =
-          new HostHistory.Builder()
-              .copyFrom(this)
-              .setHostRepoId(
-                  VKey.create(HostResource.class, parent.getName(), (Key<HostResource>) parent))
-              .build();
+          new HostHistory.Builder().copyFrom(this).setHostRepoId(parent.getName()).build();
     } else if (parentKind.equals(getKind(ContactResource.class))) {
       resultEntity =
-          new ContactHistory.Builder()
-              .copyFrom(this)
-              .setContactRepoId(
-                  VKey.create(
-                      ContactResource.class, parent.getName(), (Key<ContactResource>) parent))
-              .build();
+          new ContactHistory.Builder().copyFrom(this).setContactRepoId(parent.getName()).build();
     } else {
       throw new IllegalStateException(
           String.format("Unknown kind of HistoryEntry parent %s", parentKind));

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/main/resources/META-INF/persistence.xml

@@ -94,7 +94,10 @@
     <class>google.registry.persistence.converter.ZonedDateTimeConverter</class>
 
     <!-- Generated converters for VKey -->
-    <class>google.registry.model.billing.VKeyConverter_BillingEvent</class>
+    <class>google.registry.model.billing.VKeyConverter_Cancellation</class>
+    <class>google.registry.model.billing.VKeyConverter_Modification</class>
+    <class>google.registry.model.billing.VKeyConverter_OneTime</class>
+    <class>google.registry.model.billing.VKeyConverter_Recurring</class>
     <class>google.registry.model.contact.VKeyConverter_ContactResource</class>
     <class>google.registry.model.domain.VKeyConverter_DomainBase</class>
     <class>google.registry.model.domain.token.VKeyConverter_AllocationToken</class>

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

@@ -47,14 +47,13 @@ public class ContactHistoryTest extends EntityTestCase {
     jpaTm().transact(() -> jpaTm().insert(contact));
     VKey<ContactResource> contactVKey = contact.createVKey();
     ContactResource contactFromDb = jpaTm().transact(() -> jpaTm().load(contactVKey));
-    ContactHistory contactHistory = createContactHistory(contactFromDb, contactVKey);
+    ContactHistory contactHistory = createContactHistory(contactFromDb, contact.getRepoId());
     contactHistory.id = null;
     jpaTm().transact(() -> jpaTm().insert(contactHistory));
     jpaTm()
         .transact(
             () -> {
-              ContactHistory fromDatabase =
-                  jpaTm().load(VKey.createSql(ContactHistory.class, contactHistory.getId()));
+              ContactHistory fromDatabase = jpaTm().load(contactHistory.createVKey());
               assertContactHistoriesEqual(fromDatabase, contactHistory);
               assertThat(fromDatabase.getContactRepoId().getSqlKey())
                   .isEqualTo(contactHistory.getContactRepoId().getSqlKey());
@@ -70,15 +69,17 @@ public class ContactHistoryTest extends EntityTestCase {
     VKey<ContactResource> contactVKey = contact.createVKey();
     ContactResource contactFromDb = jpaTm().transact(() -> jpaTm().load(contactVKey));
     ContactHistory contactHistory =
-        createContactHistory(contactFromDb, contactVKey).asBuilder().setContactBase(null).build();
+        createContactHistory(contactFromDb, contact.getRepoId())
+            .asBuilder()
+            .setContactBase(null)
+            .build();
     contactHistory.id = null;
     jpaTm().transact(() -> jpaTm().insert(contactHistory));
 
     jpaTm()
         .transact(
             () -> {
-              ContactHistory fromDatabase =
-                  jpaTm().load(VKey.createSql(ContactHistory.class, contactHistory.getId()));
+              ContactHistory fromDatabase = jpaTm().load(contactHistory.createVKey());
               assertContactHistoriesEqual(fromDatabase, contactHistory);
               assertThat(fromDatabase.getContactRepoId().getSqlKey())
                   .isEqualTo(contactHistory.getContactRepoId().getSqlKey());
@@ -94,14 +95,13 @@ public class ContactHistoryTest extends EntityTestCase {
     VKey<ContactResource> contactVKey = contact.createVKey();
     ContactResource contactFromDb = tm().transact(() -> tm().load(contactVKey));
     fakeClock.advanceOneMilli();
-    ContactHistory contactHistory = createContactHistory(contactFromDb, contactVKey);
+    ContactHistory contactHistory = createContactHistory(contactFromDb, contact.getRepoId());
     tm().transact(() -> tm().insert(contactHistory));
 
     // retrieving a HistoryEntry or a ContactHistory with the same key should return the same object
     // note: due to the @EntitySubclass annotation. all Keys for ContactHistory objects will have
     // type HistoryEntry
-    VKey<ContactHistory> contactHistoryVKey =
-        VKey.createOfy(ContactHistory.class, Key.create(contactHistory));
+    VKey<ContactHistory> contactHistoryVKey = contactHistory.createVKey();
     VKey<HistoryEntry> historyEntryVKey =
         VKey.createOfy(HistoryEntry.class, Key.create(contactHistory.asHistoryEntry()));
     ContactHistory hostHistoryFromDb = tm().transact(() -> tm().load(contactHistoryVKey));
@@ -110,8 +110,7 @@ public class ContactHistoryTest extends EntityTestCase {
     assertThat(hostHistoryFromDb).isEqualTo(historyEntryFromDb);
   }
 
-  private ContactHistory createContactHistory(
-      ContactBase contact, VKey<ContactResource> contactVKey) {
+  private ContactHistory createContactHistory(ContactBase contact, String contactRepoId) {
     return new ContactHistory.Builder()
         .setType(HistoryEntry.Type.HOST_CREATE)
         .setXmlBytes("<xml></xml>".getBytes(UTF_8))
@@ -122,14 +121,14 @@ public class ContactHistoryTest extends EntityTestCase {
         .setReason("reason")
         .setRequestedByRegistrar(true)
         .setContactBase(contact)
-        .setContactRepoId(contactVKey)
+        .setContactRepoId(contactRepoId)
         .build();
   }
 
   static void assertContactHistoriesEqual(ContactHistory one, ContactHistory two) {
     assertAboutImmutableObjects()
         .that(one)
-        .isEqualExceptFields(two, "contactBase", "contactRepoId", "parent");
+        .isEqualExceptFields(two, "contactBase", "contactRepoId");
     assertAboutImmutableObjects()
         .that(one.getContactBase().orElse(null))
         .isEqualExceptFields(two.getContactBase().orElse(null), "repoId");

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

@@ -121,8 +121,7 @@ public class DomainHistoryTest extends EntityTestCase {
     // retrieving a HistoryEntry or a DomainHistory with the same key should return the same object
     // note: due to the @EntitySubclass annotation. all Keys for DomainHistory objects will have
     // type HistoryEntry
-    VKey<DomainHistory> domainHistoryVKey =
-        VKey.createOfy(DomainHistory.class, Key.create(domainHistory));
+    VKey<DomainHistory> domainHistoryVKey = domainHistory.createVKey();
     VKey<HistoryEntry> historyEntryVKey =
         VKey.createOfy(HistoryEntry.class, Key.create(domainHistory.asHistoryEntry()));
     DomainHistory domainHistoryFromDb = tm().transact(() -> tm().load(domainHistoryVKey));
@@ -155,7 +154,7 @@ public class DomainHistoryTest extends EntityTestCase {
     assertAboutImmutableObjects()
         .that(one)
         .isEqualExceptFields(
-            two, "domainContent", "domainRepoId", "parent", "nsHosts", "domainTransactionRecords");
+            two, "domainContent", "domainRepoId", "nsHosts", "domainTransactionRecords");
     assertThat(one.getDomainContent().map(DomainContent::getDomainName))
         .isEqualTo(two.getDomainContent().map(DomainContent::getDomainName));
     // NB: the record's ID gets reset by Hibernate, causing the hash code to differ so we have to

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

@@ -48,14 +48,13 @@ public class HostHistoryTest extends EntityTestCase {
     VKey<HostResource> hostVKey =
         VKey.create(HostResource.class, "host1", Key.create(HostResource.class, "host1"));
     HostResource hostFromDb = jpaTm().transact(() -> jpaTm().load(hostVKey));
-    HostHistory hostHistory = createHostHistory(hostFromDb, hostVKey);
+    HostHistory hostHistory = createHostHistory(hostFromDb, host.getRepoId());
     hostHistory.id = null;
     jpaTm().transact(() -> jpaTm().insert(hostHistory));
     jpaTm()
         .transact(
             () -> {
-              HostHistory fromDatabase =
-                  jpaTm().load(VKey.createSql(HostHistory.class, hostHistory.getId()));
+              HostHistory fromDatabase = jpaTm().load(hostHistory.createVKey());
               assertHostHistoriesEqual(fromDatabase, hostHistory);
               assertThat(fromDatabase.getHostRepoId().getSqlKey())
                   .isEqualTo(hostHistory.getHostRepoId().getSqlKey());
@@ -68,20 +67,16 @@ public class HostHistoryTest extends EntityTestCase {
     HostResource host = newHostResourceWithRoid("ns1.example.com", "host1");
     jpaTm().transact(() -> jpaTm().insert(host));
 
-    VKey<HostResource> hostVKey =
-        VKey.create(HostResource.class, "host1", Key.create(HostResource.class, "host1"));
-    HostResource hostFromDb = jpaTm().transact(() -> jpaTm().load(hostVKey));
-
+    HostResource hostFromDb = jpaTm().transact(() -> jpaTm().load(host.createVKey()));
     HostHistory hostHistory =
-        createHostHistory(hostFromDb, hostVKey).asBuilder().setHostBase(null).build();
+        createHostHistory(hostFromDb, host.getRepoId()).asBuilder().setHostBase(null).build();
     hostHistory.id = null;
     jpaTm().transact(() -> jpaTm().insert(hostHistory));
 
     jpaTm()
         .transact(
             () -> {
-              HostHistory fromDatabase =
-                  jpaTm().load(VKey.createSql(HostHistory.class, hostHistory.getId()));
+              HostHistory fromDatabase = jpaTm().load(hostHistory.createVKey());
               assertHostHistoriesEqual(fromDatabase, hostHistory);
               assertThat(fromDatabase.getHostRepoId().getSqlKey())
                   .isEqualTo(hostHistory.getHostRepoId().getSqlKey());
@@ -97,14 +92,14 @@ public class HostHistoryTest extends EntityTestCase {
     VKey<HostResource> hostVKey =
         VKey.create(HostResource.class, "host1", Key.create(HostResource.class, "host1"));
     HostResource hostFromDb = tm().transact(() -> tm().load(hostVKey));
-    HostHistory hostHistory = createHostHistory(hostFromDb, hostVKey);
+    HostHistory hostHistory = createHostHistory(hostFromDb, host.getRepoId());
     fakeClock.advanceOneMilli();
     tm().transact(() -> tm().insert(hostHistory));
 
     // retrieving a HistoryEntry or a HostHistory with the same key should return the same object
     // note: due to the @EntitySubclass annotation. all Keys for HostHistory objects will have
     // type HistoryEntry
-    VKey<HostHistory> hostHistoryVKey = VKey.createOfy(HostHistory.class, Key.create(hostHistory));
+    VKey<HostHistory> hostHistoryVKey = hostHistory.createVKey();
     VKey<HistoryEntry> historyEntryVKey =
         VKey.createOfy(HistoryEntry.class, Key.create(hostHistory.asHistoryEntry()));
     HostHistory hostHistoryFromDb = tm().transact(() -> tm().load(hostHistoryVKey));
@@ -120,7 +115,7 @@ public class HostHistoryTest extends EntityTestCase {
         .isEqualExceptFields(two.getHostBase().orElse(null), "repoId");
   }
 
-  private HostHistory createHostHistory(HostBase hostBase, VKey<HostResource> hostVKey) {
+  private HostHistory createHostHistory(HostBase hostBase, String hostRepoId) {
     return new HostHistory.Builder()
         .setType(HistoryEntry.Type.HOST_CREATE)
         .setXmlBytes("<xml></xml>".getBytes(UTF_8))
@@ -131,7 +126,7 @@ public class HostHistoryTest extends EntityTestCase {
         .setReason("reason")
         .setRequestedByRegistrar(true)
         .setHostBase(hostBase)
-        .setHostRepoId(hostVKey)
+        .setHostRepoId(hostRepoId)
         .build();
   }
 }

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

@@ -88,11 +88,7 @@ public class LegacyHistoryObjectTest extends EntityTestCase {
               jpaTm().insert(legacyContactHistory);
             });
     ContactHistory legacyHistoryFromSql =
-        jpaTm()
-            .transact(
-                () ->
-                    jpaTm()
-                        .load(VKey.createSql(ContactHistory.class, legacyContactHistory.getId())));
+        jpaTm().transact(() -> jpaTm().load(legacyContactHistory.createVKey()));
     assertAboutImmutableObjects()
         .that(legacyContactHistory)
         .isEqualExceptFields(legacyHistoryFromSql);
@@ -172,9 +168,7 @@ public class LegacyHistoryObjectTest extends EntityTestCase {
               jpaTm().insert(legacyHostHistory);
             });
     HostHistory legacyHistoryFromSql =
-        jpaTm()
-            .transact(
-                () -> jpaTm().load(VKey.createSql(HostHistory.class, legacyHostHistory.getId())));
+        jpaTm().transact(() -> jpaTm().load(legacyHostHistory.createVKey()));
     assertAboutImmutableObjects().that(legacyHostHistory).isEqualExceptFields(legacyHistoryFromSql);
     // can't compare hostRepoId directly since it doesn't save the ofy key in SQL
     assertThat(legacyHostHistory.getHostRepoId().getSqlKey())

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 {}
 }

core/src/test/resources/google/registry/beam/invoicing/billing_events_test.sql

@@ -55,7 +55,7 @@ FROM (
     FROM
       `my-project-id.latest_datastore_export.Registry`
     WHERE
-      enableInvoicing IS TRUE) ) AS BillingEvent
+      invoicingEnabled IS TRUE) ) AS BillingEvent
   -- Gather billing ID from registrar table
   -- This is a 'JOIN' as opposed to 'LEFT JOIN' to filter out
   -- non-billable registrars

core/src/test/resources/google/registry/model/schema.txt

@@ -41,7 +41,7 @@ class google.registry.model.billing.BillingEvent$OneTime {
   @Id java.lang.Long id;
   @Parent com.googlecode.objectify.Key<google.registry.model.reporting.HistoryEntry> parent;
   google.registry.model.billing.BillingEvent$Reason reason;
-  google.registry.persistence.VKey<? extends google.registry.model.billing.BillingEvent> cancellationMatchingBillingEvent;
+  google.registry.persistence.VKey<google.registry.model.billing.BillingEvent$Recurring> cancellationMatchingBillingEvent;
   google.registry.persistence.VKey<google.registry.model.domain.token.AllocationToken> allocationToken;
   java.lang.Integer periodYears;
   java.lang.String clientId;
@@ -131,9 +131,9 @@ class google.registry.model.contact.ContactHistory {
   google.registry.model.domain.Period period;
   google.registry.model.eppcommon.Trid trid;
   google.registry.model.reporting.HistoryEntry$Type type;
-  google.registry.persistence.VKey<google.registry.model.contact.ContactResource> contactRepoId;
   java.lang.Boolean requestedByRegistrar;
   java.lang.String clientId;
+  java.lang.String contactRepoId;
   java.lang.String otherClientId;
   java.lang.String reason;
   java.util.Set<google.registry.model.reporting.DomainTransactionRecord> domainTransactionRecords;
@@ -401,9 +401,9 @@ class google.registry.model.host.HostHistory {
   google.registry.model.eppcommon.Trid trid;
   google.registry.model.host.HostBase hostBase;
   google.registry.model.reporting.HistoryEntry$Type type;
-  google.registry.persistence.VKey<google.registry.model.host.HostResource> hostRepoId;
   java.lang.Boolean requestedByRegistrar;
   java.lang.String clientId;
+  java.lang.String hostRepoId;
   java.lang.String otherClientId;
   java.lang.String reason;
   java.util.Set<google.registry.model.reporting.DomainTransactionRecord> domainTransactionRecords;

db/src/main/resources/sql/flyway.txt

@@ -56,3 +56,4 @@ V55__domain_history_fields.sql
 V56__rename_host_table.sql
 V57__history_null_content.sql
 V58__drop_default_value_and_sequences_for_billing_event.sql
+V59__use_composite_primary_key_for_contact_and_host_history_table.sql

db/src/main/resources/sql/flyway/V59__use_composite_primary_key_for_contact_and_host_history_table.sql

@@ -0,0 +1,23 @@
+-- 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.
+
+alter table "ContactHistory" drop constraint "ContactHistory_pkey";
+
+alter table "ContactHistory"
+    add constraint "ContactHistory_pkey" primary key (contact_repo_id, history_revision_id);
+
+alter table "HostHistory" drop constraint "HostHistory_pkey";
+
+alter table "HostHistory"
+    add constraint "HostHistory_pkey" primary key (host_repo_id, history_revision_id);

db/src/main/resources/sql/schema/db-schema.sql.generated

@@ -151,7 +151,8 @@ create sequence temp_history_id_sequence start 1 increment 50;
     );
 
     create table "ContactHistory" (
-       history_revision_id int8 not null,
+       contact_repo_id text not null,
+        history_revision_id int8 not null,
         history_by_superuser boolean not null,
         history_registrar_id text,
         history_modification_time timestamptz not null,
@@ -215,8 +216,7 @@ create sequence temp_history_id_sequence start 1 increment 50;
         last_epp_update_time timestamptz,
         statuses text[],
         update_timestamp timestamptz,
-        contact_repo_id text not null,
-        primary key (history_revision_id)
+        primary key (contact_repo_id, history_revision_id)
     );
 
     create table "Cursor" (
@@ -400,7 +400,8 @@ create sequence temp_history_id_sequence start 1 increment 50;
     );
 
     create table "HostHistory" (
-       history_revision_id int8 not null,
+       host_repo_id text not null,
+        history_revision_id int8 not null,
         history_by_superuser boolean not null,
         history_registrar_id text,
         history_modification_time timestamptz not null,
@@ -423,8 +424,7 @@ create sequence temp_history_id_sequence start 1 increment 50;
         last_epp_update_time timestamptz,
         statuses text[],
         update_timestamp timestamptz,
-        host_repo_id text not null,
-        primary key (history_revision_id)
+        primary key (host_repo_id, history_revision_id)
     );
 
     create table "Lock" (

db/src/main/resources/sql/schema/nomulus.golden.sql

@@ -1084,7 +1084,7 @@ ALTER TABLE ONLY public."ClaimsList"
 --
 
 ALTER TABLE ONLY public."ContactHistory"
-    ADD CONSTRAINT "ContactHistory_pkey" PRIMARY KEY (history_revision_id);
+    ADD CONSTRAINT "ContactHistory_pkey" PRIMARY KEY (contact_repo_id, history_revision_id);
 
 
 --
@@ -1140,7 +1140,7 @@ ALTER TABLE ONLY public."GracePeriod"
 --
 
 ALTER TABLE ONLY public."HostHistory"
-    ADD CONSTRAINT "HostHistory_pkey" PRIMARY KEY (history_revision_id);
+    ADD CONSTRAINT "HostHistory_pkey" PRIMARY KEY (host_repo_id, history_revision_id);
 
 
 --