Add remaining columns to Domain's SQL schema (#702)

core/src/main/java/google/registry/flows/domain/DomainCreateFlow.java

@@ -347,8 +347,8 @@ public class DomainCreateFlow implements TransactionalFlow {
             .setRepoId(repoId)
             .setIdnTableName(validateDomainNameWithIdnTables(domainName))
             .setRegistrationExpirationTime(registrationExpirationTime)
-            .setAutorenewBillingEvent(Key.create(autorenewBillingEvent))
+            .setAutorenewBillingEvent(autorenewBillingEvent.createVKey())
-            .setAutorenewPollMessage(Key.create(autorenewPollMessage))
+            .setAutorenewPollMessage(autorenewPollMessage.createVKey())
             .setLaunchNotice(hasClaimsNotice ? launchCreate.get().getNotice() : null)
             .setSmdId(signedMarkId)
             .setDsData(secDnsCreate.isPresent() ? secDnsCreate.get().getDsData() : null)

core/src/main/java/google/registry/flows/domain/DomainDeleteFlow.java

@@ -31,7 +31,6 @@ import static google.registry.model.ResourceTransferUtils.handlePendingTransferO
 import static google.registry.model.ResourceTransferUtils.updateForeignKeyIndexDeletionTime;
 import static google.registry.model.eppoutput.Result.Code.SUCCESS;
 import static google.registry.model.eppoutput.Result.Code.SUCCESS_WITH_ACTION_PENDING;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.reporting.DomainTransactionRecord.TransactionReportField.ADD_FIELDS;
 import static google.registry.model.reporting.DomainTransactionRecord.TransactionReportField.RENEW_FIELDS;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
@@ -209,7 +208,7 @@ public final class DomainDeleteFlow implements TransactionalFlow {
       PollMessage.OneTime deletePollMessage =
           createDeletePollMessage(existingDomain, historyEntry, deletionTime);
       entitiesToSave.add(deletePollMessage);
-      builder.setDeletePollMessage(Key.create(deletePollMessage));
+      builder.setDeletePollMessage(deletePollMessage.createVKey());
     }
 
     // Cancel any grace periods that were still active, and set the expiration time accordingly.
@@ -222,8 +221,7 @@ public final class DomainDeleteFlow implements TransactionalFlow {
         if (gracePeriod.getOneTimeBillingEvent() != null) {
           // Take the amount of amount of registration time being refunded off the expiration time.
           // This can be either add grace periods or renew grace periods.
-          BillingEvent.OneTime oneTime =
+          BillingEvent.OneTime oneTime = tm().load(gracePeriod.getOneTimeBillingEvent());
-              ofy().load().key(gracePeriod.getOneTimeBillingEvent()).now();
           newExpirationTime = newExpirationTime.minusYears(oneTime.getPeriodYears());
         } else if (gracePeriod.getRecurringBillingEvent() != null) {
           // Take 1 year off the registration if in the autorenew grace period (no need to load the
@@ -370,12 +368,12 @@ public final class DomainDeleteFlow implements TransactionalFlow {
   private Money getGracePeriodCost(GracePeriod gracePeriod, DateTime now) {
     if (gracePeriod.getType() == GracePeriodStatus.AUTO_RENEW) {
       DateTime autoRenewTime =
-          ofy().load().key(checkNotNull(gracePeriod.getRecurringBillingEvent())).now()
+          tm().load(checkNotNull(gracePeriod.getRecurringBillingEvent()))
               .getRecurrenceTimeOfYear()
               .getLastInstanceBeforeOrAt(now);
       return getDomainRenewCost(targetId, autoRenewTime, 1);
     }
-    return ofy().load().key(checkNotNull(gracePeriod.getOneTimeBillingEvent())).now().getCost();
+    return tm().load(checkNotNull(gracePeriod.getOneTimeBillingEvent())).getCost();
   }
 
   @Nullable

core/src/main/java/google/registry/flows/domain/DomainFlowUtils.java

@@ -517,14 +517,14 @@ public class DomainFlowUtils {
    */
   public static void updateAutorenewRecurrenceEndTime(DomainBase domain, DateTime newEndTime) {
     Optional<PollMessage.Autorenew> autorenewPollMessage =
-        Optional.ofNullable(ofy().load().key(domain.getAutorenewPollMessage()).now());
+        tm().maybeLoad(domain.getAutorenewPollMessage());
 
     // Construct an updated autorenew poll message. If the autorenew poll message no longer exists,
     // create a new one at the same id. This can happen if a transfer was requested on a domain
     // where all autorenew poll messages had already been delivered (this would cause the poll
     // message to be deleted), and then subsequently the transfer was canceled, rejected, or deleted
     // (which would cause the poll message to be recreated here).
-    Key<PollMessage.Autorenew> existingAutorenewKey = domain.getAutorenewPollMessage();
+    Key<PollMessage.Autorenew> existingAutorenewKey = domain.getAutorenewPollMessage().getOfyKey();
     PollMessage.Autorenew updatedAutorenewPollMessage =
         autorenewPollMessage.isPresent()
             ? autorenewPollMessage.get().asBuilder().setAutorenewEndTime(newEndTime).build()
@@ -542,7 +542,7 @@ public class DomainFlowUtils {
       ofy().save().entity(updatedAutorenewPollMessage);
     }
 
-    Recurring recurring = ofy().load().key(domain.getAutorenewBillingEvent()).now();
+    Recurring recurring = tm().load(domain.getAutorenewBillingEvent());
     ofy().save().entity(recurring.asBuilder().setRecurrenceEndTime(newEndTime).build());
   }
 

core/src/main/java/google/registry/flows/domain/DomainRenewFlow.java

@@ -33,7 +33,6 @@ import static google.registry.util.DateTimeUtils.leapSafeAddYears;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
-import com.googlecode.objectify.Key;
 import google.registry.flows.EppException;
 import google.registry.flows.EppException.ParameterValueRangeErrorException;
 import google.registry.flows.ExtensionManager;
@@ -181,8 +180,8 @@ public final class DomainRenewFlow implements TransactionalFlow {
             .setLastEppUpdateTime(now)
             .setLastEppUpdateClientId(clientId)
             .setRegistrationExpirationTime(newExpirationTime)
-            .setAutorenewBillingEvent(Key.create(newAutorenewEvent))
+            .setAutorenewBillingEvent(newAutorenewEvent.createVKey())
-            .setAutorenewPollMessage(Key.create(newAutorenewPollMessage))
+            .setAutorenewPollMessage(newAutorenewPollMessage.createVKey())
             .addGracePeriod(
                 GracePeriod.forBillingEvent(GracePeriodStatus.RENEW, explicitRenewEvent))
             .build();

core/src/main/java/google/registry/flows/domain/DomainRestoreRequestFlow.java

@@ -26,7 +26,6 @@ import static google.registry.flows.domain.DomainFlowUtils.verifyNotReserved;
 import static google.registry.flows.domain.DomainFlowUtils.verifyPremiumNameIsNotBlocked;
 import static google.registry.flows.domain.DomainFlowUtils.verifyRegistrarIsActive;
 import static google.registry.model.ResourceTransferUtils.updateForeignKeyIndexDeletionTime;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.DateTimeUtils.END_OF_TIME;
 
@@ -174,8 +173,8 @@ public final class DomainRestoreRequestFlow implements TransactionalFlow  {
             existingDomain, newExpirationTime, autorenewEvent, autorenewPollMessage, now, clientId);
     updateForeignKeyIndexDeletionTime(newDomain);
     entitiesToSave.add(newDomain, historyEntry, autorenewEvent, autorenewPollMessage);
-    ofy().save().entities(entitiesToSave.build());
+    tm().saveNewOrUpdateAll(entitiesToSave.build());
-    ofy().delete().key(existingDomain.getDeletePollMessage());
+    tm().delete(existingDomain.getDeletePollMessage());
     dnsQueue.addDomainRefreshTask(existingDomain.getDomainName());
     return responseBuilder
         .setExtensions(createResponseExtensions(feesAndCredits, feeUpdate, isExpired))
@@ -232,8 +231,8 @@ public final class DomainRestoreRequestFlow implements TransactionalFlow  {
         .setStatusValues(null)
         .setGracePeriods(null)
         .setDeletePollMessage(null)
-        .setAutorenewBillingEvent(Key.create(autorenewEvent))
+        .setAutorenewBillingEvent(autorenewEvent.createVKey())
-        .setAutorenewPollMessage(Key.create(autorenewPollMessage))
+        .setAutorenewPollMessage(autorenewPollMessage.createVKey())
         .setLastEppUpdateTime(now)
         .setLastEppUpdateClientId(clientId)
         .build();

core/src/main/java/google/registry/flows/domain/DomainTransferApproveFlow.java

@@ -186,8 +186,8 @@ public final class DomainTransferApproveFlow implements TransactionalFlow {
                     .setTransferredRegistrationExpirationTime(newExpirationTime)
                     .build())
             .setRegistrationExpirationTime(newExpirationTime)
-            .setAutorenewBillingEvent(Key.create(autorenewEvent))
+            .setAutorenewBillingEvent(autorenewEvent.createVKey())
-            .setAutorenewPollMessage(Key.create(gainingClientAutorenewPollMessage))
+            .setAutorenewPollMessage(gainingClientAutorenewPollMessage.createVKey())
             // Remove all the old grace periods and add a new one for the transfer.
             .setGracePeriods(
                 billingEvent.isPresent()

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

@@ -630,9 +630,9 @@ public abstract class BillingEvent extends ImmutableObject
           .setParent(historyEntry);
       // Set the grace period's billing event using the appropriate Cancellation builder method.
       if (gracePeriod.getOneTimeBillingEvent() != null) {
-        builder.setOneTimeEventKey(VKey.from(gracePeriod.getOneTimeBillingEvent()));
+        builder.setOneTimeEventKey(gracePeriod.getOneTimeBillingEvent());
       } else if (gracePeriod.getRecurringBillingEvent() != null) {
-        builder.setRecurringEventKey(VKey.from(gracePeriod.getRecurringBillingEvent()));
+        builder.setRecurringEventKey(gracePeriod.getRecurringBillingEvent());
       }
       return builder.build();
     }

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

@@ -227,7 +227,8 @@ public class DomainBase extends EppResource
    * refer to a {@link PollMessage} timed to when the domain is fully deleted. If the domain is
    * restored, the message should be deleted.
    */
-  @Transient Key<PollMessage.OneTime> deletePollMessage;
+  @Column(name = "deletion_poll_message_id")
+  VKey<PollMessage.OneTime> deletePollMessage;
 
   /**
    * The recurring billing event associated with this domain's autorenewals.
@@ -237,7 +238,8 @@ public class DomainBase extends EppResource
    * {@link #registrationExpirationTime} is changed the recurrence should be closed, a new one
    * should be created, and this field should be updated to point to the new one.
    */
-  @Transient Key<BillingEvent.Recurring> autorenewBillingEvent;
+  @Column(name = "billing_recurrence_id")
+  VKey<BillingEvent.Recurring> autorenewBillingEvent;
 
   /**
    * The recurring poll message associated with this domain's autorenewals.
@@ -247,7 +249,8 @@ public class DomainBase extends EppResource
    * {@link #registrationExpirationTime} is changed the recurrence should be closed, a new one
    * should be created, and this field should be updated to point to the new one.
    */
-  @Transient Key<PollMessage.Autorenew> autorenewPollMessage;
+  @Column(name = "autorenew_poll_message_id")
+  VKey<PollMessage.Autorenew> autorenewPollMessage;
 
   /** The unexpired grace periods for this domain (some of which may not be active yet). */
   @Transient @ElementCollection Set<GracePeriod> gracePeriods;
@@ -316,15 +319,15 @@ public class DomainBase extends EppResource
     return registrationExpirationTime;
   }
 
-  public Key<PollMessage.OneTime> getDeletePollMessage() {
+  public VKey<PollMessage.OneTime> getDeletePollMessage() {
     return deletePollMessage;
   }
 
-  public Key<BillingEvent.Recurring> getAutorenewBillingEvent() {
+  public VKey<BillingEvent.Recurring> getAutorenewBillingEvent() {
     return autorenewBillingEvent;
   }
 
-  public Key<PollMessage.Autorenew> getAutorenewPollMessage() {
+  public VKey<PollMessage.Autorenew> getAutorenewPollMessage() {
     return autorenewPollMessage;
   }
 
@@ -453,14 +456,8 @@ public class DomainBase extends EppResource
               .setRegistrationExpirationTime(expirationDate)
               // Set the speculatively-written new autorenew events as the domain's autorenew
               // events.
-              .setAutorenewBillingEvent(
+              .setAutorenewBillingEvent(transferData.getServerApproveAutorenewEvent())
-                  transferData.getServerApproveAutorenewEvent() == null
+              .setAutorenewPollMessage(transferData.getServerApproveAutorenewPollMessage());
-                      ? null
-                      : transferData.getServerApproveAutorenewEvent().getOfyKey())
-              .setAutorenewPollMessage(
-                  transferData.getServerApproveAutorenewPollMessage() == null
-                      ? null
-                      : transferData.getServerApproveAutorenewPollMessage().getOfyKey());
       if (transferData.getTransferPeriod().getValue() == 1) {
         // Set the grace period using a key to the prescheduled transfer billing event.  Not using
         // GracePeriod.forBillingEvent() here in order to avoid the actual Datastore fetch.
@@ -471,9 +468,7 @@ public class DomainBase extends EppResource
                     transferExpirationTime.plus(
                         Registry.get(getTld()).getTransferGracePeriodLength()),
                     transferData.getGainingClientId(),
-                    transferData.getServerApproveBillingEvent() == null
+                    transferData.getServerApproveBillingEvent())));
-                        ? null
-                        : transferData.getServerApproveBillingEvent().getOfyKey())));
       } else {
         // There won't be a billing event, so we don't need a grace period
         builder.setGracePeriods(ImmutableSet.of());
@@ -801,19 +796,17 @@ public class DomainBase extends EppResource
       return this;
     }
 
-    public Builder setDeletePollMessage(Key<PollMessage.OneTime> deletePollMessage) {
+    public Builder setDeletePollMessage(VKey<PollMessage.OneTime> deletePollMessage) {
       getInstance().deletePollMessage = deletePollMessage;
       return this;
     }
 
-    public Builder setAutorenewBillingEvent(
+    public Builder setAutorenewBillingEvent(VKey<BillingEvent.Recurring> autorenewBillingEvent) {
-        Key<BillingEvent.Recurring> autorenewBillingEvent) {
       getInstance().autorenewBillingEvent = autorenewBillingEvent;
       return this;
     }
 
-    public Builder setAutorenewPollMessage(
+    public Builder setAutorenewPollMessage(VKey<PollMessage.Autorenew> autorenewPollMessage) {
-        Key<PollMessage.Autorenew> autorenewPollMessage) {
       getInstance().autorenewPollMessage = autorenewPollMessage;
       return this;
     }

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

@@ -17,12 +17,13 @@ package google.registry.model.domain;
 import static com.google.common.base.Preconditions.checkArgument;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 
-import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Embed;
 import com.googlecode.objectify.annotation.Ignore;
 import google.registry.model.ImmutableObject;
 import google.registry.model.billing.BillingEvent;
+import google.registry.model.billing.BillingEvent.Recurring;
 import google.registry.model.domain.rgp.GracePeriodStatus;
+import google.registry.persistence.VKey;
 import javax.annotation.Nullable;
 import javax.persistence.Column;
 import javax.persistence.GeneratedValue;
@@ -57,18 +58,18 @@ public class GracePeriod extends ImmutableObject {
 
   /**
    * The one-time billing event corresponding to the action that triggered this grace period, or
-   * null if not applicable.  Not set for autorenew grace periods (which instead use the field
+   * null if not applicable. Not set for autorenew grace periods (which instead use the field {@code
-   * {@code billingEventRecurring}) or for redemption grace periods (since deletes have no cost).
+   * billingEventRecurring}) or for redemption grace periods (since deletes have no cost).
    */
   // NB: Would @IgnoreSave(IfNull.class), but not allowed for @Embed collections.
-  Key<BillingEvent.OneTime> billingEventOneTime = null;
+  VKey<BillingEvent.OneTime> billingEventOneTime = null;
 
   /**
    * The recurring billing event corresponding to the action that triggered this grace period, if
    * applicable - i.e. if the action was an autorenew - or null in all other cases.
    */
   // NB: Would @IgnoreSave(IfNull.class), but not allowed for @Embed collections.
-  Key<BillingEvent.Recurring> billingEventRecurring = null;
+  VKey<BillingEvent.Recurring> billingEventRecurring = null;
 
   public GracePeriodStatus getType() {
     return type;
@@ -91,8 +92,7 @@ public class GracePeriod extends ImmutableObject {
    * Returns the one time billing event. The value will only be non-null if the type of this grace
    * period is not AUTO_RENEW.
    */
-
+  public VKey<BillingEvent.OneTime> getOneTimeBillingEvent() {
-  public Key<BillingEvent.OneTime> getOneTimeBillingEvent() {
     return billingEventOneTime;
   }
 
@@ -100,7 +100,7 @@ public class GracePeriod extends ImmutableObject {
    * Returns the recurring billing event. The value will only be non-null if the type of this grace
    * period is AUTO_RENEW.
    */
-  public Key<BillingEvent.Recurring> getRecurringBillingEvent() {
+  public VKey<BillingEvent.Recurring> getRecurringBillingEvent() {
     return billingEventRecurring;
   }
 
@@ -108,8 +108,8 @@ public class GracePeriod extends ImmutableObject {
       GracePeriodStatus type,
       DateTime expirationTime,
       String clientId,
-       @Nullable Key<BillingEvent.OneTime> billingEventOneTime,
+      @Nullable VKey<BillingEvent.OneTime> billingEventOneTime,
-       @Nullable Key<BillingEvent.Recurring> billingEventRecurring) {
+      @Nullable VKey<BillingEvent.Recurring> billingEventRecurring) {
     checkArgument((billingEventOneTime == null) || (billingEventRecurring == null),
         "A grace period can have at most one billing event");
     checkArgument(
@@ -127,15 +127,15 @@ public class GracePeriod extends ImmutableObject {
   /**
    * Creates a GracePeriod for an (optional) OneTime billing event.
    *
-   * <p>Normal callers should always use {@link #forBillingEvent} instead, assuming they do not
+   * <p>Normal callers should always use {@link #forBillingEvent} instead, assuming they do not need
-   * need to avoid loading the BillingEvent from Datastore.  This method should typically be
+   * to avoid loading the BillingEvent from Datastore. This method should typically be called only
-   * called only from test code to explicitly construct GracePeriods.
+   * from test code to explicitly construct GracePeriods.
    */
   public static GracePeriod create(
       GracePeriodStatus type,
       DateTime expirationTime,
       String clientId,
-      @Nullable Key<BillingEvent.OneTime> billingEventOneTime) {
+      @Nullable VKey<BillingEvent.OneTime> billingEventOneTime) {
     return createInternal(type, expirationTime, clientId, billingEventOneTime, null);
   }
 
@@ -144,7 +144,7 @@ public class GracePeriod extends ImmutableObject {
       GracePeriodStatus type,
       DateTime expirationTime,
       String clientId,
-      Key<BillingEvent.Recurring> billingEventRecurring) {
+      VKey<Recurring> billingEventRecurring) {
     checkArgumentNotNull(billingEventRecurring, "billingEventRecurring cannot be null");
     return createInternal(type, expirationTime, clientId, null, billingEventRecurring);
   }
@@ -159,6 +159,6 @@ public class GracePeriod extends ImmutableObject {
   public static GracePeriod forBillingEvent(
       GracePeriodStatus type, BillingEvent.OneTime billingEvent) {
     return create(
-        type, billingEvent.getBillingTime(), billingEvent.getClientId(), Key.create(billingEvent));
+        type, billingEvent.getBillingTime(), billingEvent.getClientId(), billingEvent.createVKey());
   }
 }

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

@@ -31,7 +31,6 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableMultimap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Sets;
-import com.googlecode.objectify.Key;
 import google.registry.model.billing.BillingEvent;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.Period;
@@ -224,8 +223,8 @@ class UnrenewDomainCommand extends ConfirmingCommand implements CommandWithRemot
             .setRegistrationExpirationTime(newExpirationTime)
             .setLastEppUpdateTime(now)
             .setLastEppUpdateClientId(domain.getCurrentSponsorClientId())
-            .setAutorenewBillingEvent(Key.create(newAutorenewEvent))
+            .setAutorenewBillingEvent(newAutorenewEvent.createVKey())
-            .setAutorenewPollMessage(Key.create(newAutorenewPollMessage))
+            .setAutorenewPollMessage(newAutorenewPollMessage.createVKey())
             .build();
     // In order to do it'll need to write out a new HistoryEntry (likely of type SYNTHETIC), a new
     // autorenew billing event and poll message, and a new one time poll message at the present time

core/src/test/java/google/registry/flows/FlowTestCase.java

@@ -22,6 +22,7 @@ import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
 import static google.registry.model.eppcommon.EppXmlTransformer.marshal;
 import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatastoreHelper.POLL_MESSAGE_ID_STRIPPER;
 import static google.registry.testing.DatastoreHelper.getPollMessages;
 import static google.registry.testing.DatastoreHelper.stripBillingEventId;
@@ -196,12 +197,9 @@ public abstract class FlowTestCase<F extends Flow> {
     assertWithMessage("Billing event is present for grace period: " + gracePeriod)
         .that(gracePeriod.hasBillingEvent())
         .isTrue();
-    return ofy()
+    return tm().load(
-        .load()
-        .key(
             firstNonNull(
-                gracePeriod.getOneTimeBillingEvent(), gracePeriod.getRecurringBillingEvent()))
+                gracePeriod.getOneTimeBillingEvent(), gracePeriod.getRecurringBillingEvent()));
-        .now();
   }
 
   /**

core/src/test/java/google/registry/flows/domain/DomainCreateFlowTest.java

@@ -31,6 +31,7 @@ import static google.registry.model.registry.Registry.TldState.GENERAL_AVAILABIL
 import static google.registry.model.registry.Registry.TldState.PREDELEGATION;
 import static google.registry.model.registry.Registry.TldState.QUIET_PERIOD;
 import static google.registry.model.registry.Registry.TldState.START_DATE_SUNRISE;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.pricing.PricingEngineProxy.isDomainPremium;
 import static google.registry.testing.DatastoreHelper.assertBillingEvents;
 import static google.registry.testing.DatastoreHelper.assertPollMessagesForResource;
@@ -257,8 +258,7 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
     HistoryEntry historyEntry = getHistoryEntries(domain).get(0);
     assertAboutDomains()
         .that(domain)
-        .hasRegistrationExpirationTime(
+        .hasRegistrationExpirationTime(tm().load(domain.getAutorenewBillingEvent()).getEventTime())
-            ofy().load().key(domain.getAutorenewBillingEvent()).now().getEventTime())
         .and()
         .hasOnlyOneHistoryEntryWhich()
         .hasType(HistoryEntry.Type.DOMAIN_CREATE)

core/src/test/java/google/registry/flows/domain/DomainDeleteFlowTest.java

@@ -34,6 +34,7 @@ import static google.registry.model.reporting.DomainTransactionRecord.Transactio
 import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_CREATE;
 import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_DELETE;
 import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_TRANSFER_REQUEST;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatastoreHelper.assertBillingEvents;
 import static google.registry.testing.DatastoreHelper.createTld;
 import static google.registry.testing.DatastoreHelper.getOnlyHistoryEntryOfType;
@@ -139,8 +140,8 @@ class DomainDeleteFlowTest extends ResourceFlowTestCase<DomainDeleteFlow, Domain
         persistResource(
             domain
                 .asBuilder()
-                .setAutorenewBillingEvent(Key.create(autorenewBillingEvent))
+                .setAutorenewBillingEvent(autorenewBillingEvent.createVKey())
-                .setAutorenewPollMessage(Key.create(autorenewPollMessage))
+                .setAutorenewPollMessage(autorenewPollMessage.createVKey())
                 .build());
     assertTransactionalFlow(true);
   }
@@ -196,9 +197,9 @@ class DomainDeleteFlowTest extends ResourceFlowTestCase<DomainDeleteFlow, Domain
                             GracePeriodStatus.AUTO_RENEW,
                             A_MONTH_AGO.plusDays(45),
                             "TheRegistrar",
-                            Key.create(autorenewBillingEvent))))
+                            autorenewBillingEvent.createVKey())))
-                .setAutorenewBillingEvent(Key.create(autorenewBillingEvent))
+                .setAutorenewBillingEvent(autorenewBillingEvent.createVKey())
-                .setAutorenewPollMessage(Key.create(autorenewPollMessage))
+                .setAutorenewPollMessage(autorenewPollMessage.createVKey())
                 .build());
     assertTransactionalFlow(true);
   }
@@ -436,9 +437,9 @@ class DomainDeleteFlowTest extends ResourceFlowTestCase<DomainDeleteFlow, Domain
     DateTime deletionTime = domain.getDeletionTime();
     assertThat(getPollMessages("TheRegistrar", deletionTime.minusMinutes(1))).isEmpty();
     assertThat(getPollMessages("TheRegistrar", deletionTime)).hasSize(1);
-    assertThat(domain.getDeletePollMessage())
+    assertThat(domain.getDeletePollMessage().getOfyKey())
-        .isEqualTo(Key.create(getOnlyPollMessage("TheRegistrar")));
+        .isEqualTo(getOnlyPollMessage("TheRegistrar").createVKey().getOfyKey());
-    PollMessage.OneTime deletePollMessage = ofy().load().key(domain.getDeletePollMessage()).now();
+    PollMessage.OneTime deletePollMessage = tm().load(domain.getDeletePollMessage());
     assertThat(deletePollMessage.getMsg()).isEqualTo(expectedMessage);
   }
 
@@ -472,10 +473,7 @@ class DomainDeleteFlowTest extends ResourceFlowTestCase<DomainDeleteFlow, Domain
     // Modify the autorenew poll message so that it has unacked messages in the past. This should
     // prevent it from being deleted when the domain is deleted.
     persistResource(
-        ofy()
+        tm().load(reloadResourceByForeignKey().getAutorenewPollMessage())
-            .load()
-            .key(reloadResourceByForeignKey().getAutorenewPollMessage())
-            .now()
             .asBuilder()
             .setEventTime(A_MONTH_FROM_NOW.minusYears(3))
             .build());

core/src/test/java/google/registry/flows/domain/DomainInfoFlowTest.java

@@ -44,6 +44,7 @@ import google.registry.flows.domain.DomainFlowUtils.CurrencyUnitMismatchExceptio
 import google.registry.flows.domain.DomainFlowUtils.FeeChecksDontSupportPhasesException;
 import google.registry.flows.domain.DomainFlowUtils.RestoresAreAlwaysForOneYearException;
 import google.registry.flows.domain.DomainFlowUtils.TransfersAreAlwaysForOneYearException;
+import google.registry.model.billing.BillingEvent;
 import google.registry.model.billing.BillingEvent.Recurring;
 import google.registry.model.contact.ContactAuthInfo;
 import google.registry.model.contact.ContactResource;
@@ -59,6 +60,8 @@ import google.registry.model.eppcommon.StatusValue;
 import google.registry.model.host.HostResource;
 import google.registry.model.ofy.RequestCapturingAsyncDatastoreService;
 import google.registry.model.registry.Registry;
+import google.registry.model.reporting.HistoryEntry;
+import google.registry.persistence.VKey;
 import google.registry.testing.AppEngineRule;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
@@ -324,16 +327,17 @@ class DomainInfoFlowTest extends ResourceFlowTestCase<DomainInfoFlow, DomainBase
   @Test
   void testSuccess_autoRenewGracePeriod() throws Exception {
     persistTestEntities(false);
+    Key<HistoryEntry> historyEntry =
+        Key.create(domain.createVKey().getOfyKey(), HistoryEntry.class, 67890);
+    VKey<BillingEvent.Recurring> recurringVKey =
+        VKey.from(Key.create(historyEntry, Recurring.class, 12345));
     // Add an AUTO_RENEW grace period to the saved resource.
     persistResource(
         domain
             .asBuilder()
             .addGracePeriod(
                 GracePeriod.createForRecurring(
-                    GracePeriodStatus.AUTO_RENEW,
+                    GracePeriodStatus.AUTO_RENEW, clock.nowUtc().plusDays(1), "foo", recurringVKey))
-                    clock.nowUtc().plusDays(1),
-                    "foo",
-                    Key.create(Recurring.class, 12345)))
             .build());
     doSuccessfulTest("domain_info_response_autorenewperiod.xml", false);
   }

core/src/test/java/google/registry/flows/domain/DomainRenewFlowTest.java

@@ -16,7 +16,7 @@ package google.registry.flows.domain;
 
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.flows.domain.DomainTransferFlowTestCase.persistWithPendingTransfer;
-import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatastoreHelper.assertBillingEvents;
 import static google.registry.testing.DatastoreHelper.createTld;
 import static google.registry.testing.DatastoreHelper.getOnlyHistoryEntryOfType;
@@ -38,7 +38,6 @@ import static org.junit.Assert.assertThrows;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
-import com.googlecode.objectify.Key;
 import google.registry.flows.EppException;
 import google.registry.flows.FlowUtils.UnknownCurrencyEppException;
 import google.registry.flows.ResourceFlowTestCase;
@@ -137,8 +136,8 @@ class DomainRenewFlowTest extends ResourceFlowTestCase<DomainRenewFlow, DomainBa
             .asBuilder()
             .setRegistrationExpirationTime(expirationTime)
             .setStatusValues(ImmutableSet.copyOf(statusValues))
-            .setAutorenewBillingEvent(Key.create(autorenewEvent))
+            .setAutorenewBillingEvent(autorenewEvent.createVKey())
-            .setAutorenewPollMessage(Key.create(autorenewPollMessage))
+            .setAutorenewPollMessage(autorenewPollMessage.createVKey())
             .build());
     clock.advanceOneMilli();
   }
@@ -175,7 +174,7 @@ class DomainRenewFlowTest extends ResourceFlowTestCase<DomainRenewFlow, DomainBa
     DomainBase domain = reloadResourceByForeignKey();
     HistoryEntry historyEntryDomainRenew =
         getOnlyHistoryEntryOfType(domain, HistoryEntry.Type.DOMAIN_RENEW);
-    assertThat(ofy().load().key(domain.getAutorenewBillingEvent()).now().getEventTime())
+    assertThat(tm().load(domain.getAutorenewBillingEvent()).getEventTime())
         .isEqualTo(newExpiration);
     assertAboutDomains()
         .that(domain)
@@ -470,10 +469,7 @@ class DomainRenewFlowTest extends ResourceFlowTestCase<DomainRenewFlow, DomainBa
     persistDomain();
     // Modify the autorenew poll message so that it has an undelivered message in the past.
     persistResource(
-        ofy()
+        tm().load(reloadResourceByForeignKey().getAutorenewPollMessage())
-            .load()
-            .key(reloadResourceByForeignKey().getAutorenewPollMessage())
-            .now()
             .asBuilder()
             .setEventTime(expirationTime.minusYears(1))
             .build());

core/src/test/java/google/registry/flows/domain/DomainRestoreRequestFlowTest.java

@@ -15,7 +15,7 @@
 package google.registry.flows.domain;
 
 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.DatastoreHelper.assertBillingEvents;
 import static google.registry.testing.DatastoreHelper.createTld;
 import static google.registry.testing.DatastoreHelper.getOnlyHistoryEntryOfType;
@@ -38,7 +38,6 @@ import static org.junit.Assert.assertThrows;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
-import com.googlecode.objectify.Key;
 import google.registry.flows.EppException;
 import google.registry.flows.EppException.UnimplementedExtensionException;
 import google.registry.flows.FlowUtils.UnknownCurrencyEppException;
@@ -116,13 +115,13 @@ class DomainRestoreRequestFlowTest
                     GracePeriodStatus.REDEMPTION, clock.nowUtc().plusDays(1), "foo", null))
             .setStatusValues(ImmutableSet.of(StatusValue.PENDING_DELETE))
             .setDeletePollMessage(
-                Key.create(
                 persistResource(
                         new PollMessage.OneTime.Builder()
                             .setClientId("TheRegistrar")
                             .setEventTime(clock.nowUtc().plusDays(5))
                             .setParent(historyEntry)
-                            .build())))
+                            .build())
+                    .createVKey())
             .build());
     clock.advanceOneMilli();
   }
@@ -146,7 +145,7 @@ class DomainRestoreRequestFlowTest
     DomainBase domain = reloadResourceByForeignKey();
     HistoryEntry historyEntryDomainRestore =
         getOnlyHistoryEntryOfType(domain, HistoryEntry.Type.DOMAIN_RESTORE);
-    assertThat(ofy().load().key(domain.getAutorenewBillingEvent()).now().getEventTime())
+    assertThat(tm().load(domain.getAutorenewBillingEvent()).getEventTime())
         .isEqualTo(expirationTime);
     assertAboutDomains()
         .that(domain)
@@ -214,7 +213,7 @@ class DomainRestoreRequestFlowTest
     DomainBase domain = reloadResourceByForeignKey();
     HistoryEntry historyEntryDomainRestore =
         getOnlyHistoryEntryOfType(domain, HistoryEntry.Type.DOMAIN_RESTORE);
-    assertThat(ofy().load().key(domain.getAutorenewBillingEvent()).now().getEventTime())
+    assertThat(tm().load(domain.getAutorenewBillingEvent()).getEventTime())
         .isEqualTo(newExpirationTime);
     assertAboutDomains()
         .that(domain)

core/src/test/java/google/registry/flows/domain/DomainTransferApproveFlowTest.java

@@ -16,12 +16,12 @@ package google.registry.flows.domain;
 
 import static com.google.common.collect.MoreCollectors.onlyElement;
 import static com.google.common.truth.Truth.assertThat;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.reporting.DomainTransactionRecord.TransactionReportField.NET_ADDS_4_YR;
 import static google.registry.model.reporting.DomainTransactionRecord.TransactionReportField.TRANSFER_SUCCESSFUL;
 import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_CREATE;
 import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_TRANSFER_APPROVE;
 import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_TRANSFER_REQUEST;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatastoreHelper.assertBillingEventsForResource;
 import static google.registry.testing.DatastoreHelper.createTld;
 import static google.registry.testing.DatastoreHelper.deleteResource;
@@ -42,7 +42,6 @@ import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.collect.Ordering;
 import com.google.common.collect.Streams;
-import com.googlecode.objectify.Key;
 import google.registry.flows.EppException;
 import google.registry.flows.ResourceFlowUtils.BadAuthInfoForResourceException;
 import google.registry.flows.ResourceFlowUtils.ResourceDoesNotExistException;
@@ -192,7 +191,7 @@ class DomainTransferApproveFlowTest
     assertAboutHistoryEntries().that(historyEntryTransferApproved).hasOtherClientId("NewRegistrar");
     assertTransferApproved(domain, originalTransferData);
     assertAboutDomains().that(domain).hasRegistrationExpirationTime(expectedExpirationTime);
-    assertThat(ofy().load().key(domain.getAutorenewBillingEvent()).now().getEventTime())
+    assertThat(tm().load(domain.getAutorenewBillingEvent()).getEventTime())
         .isEqualTo(expectedExpirationTime);
     // The poll message (in the future) to the losing registrar for implicit ack should be gone.
     assertThat(getPollMessages(domain, "TheRegistrar", clock.nowUtc().plusMonths(1))).isEmpty();
@@ -403,7 +402,7 @@ class DomainTransferApproveFlowTest
             .setEventTime(clock.nowUtc()) // The cancellation happens at the moment of transfer.
             .setBillingTime(
                 oldExpirationTime.plus(Registry.get("tld").getAutoRenewGracePeriodLength()))
-            .setRecurringEventKey(VKey.from(domain.getAutorenewBillingEvent())));
+            .setRecurringEventKey(domain.getAutorenewBillingEvent()));
   }
 
   @Test
@@ -634,7 +633,7 @@ class DomainTransferApproveFlowTest
   @Test
   void testSuccess_superuserExtension_transferPeriodZero_autorenewGraceActive() throws Exception {
     DomainBase domain = reloadResourceByForeignKey();
-    Key<Recurring> existingAutorenewEvent = domain.getAutorenewBillingEvent();
+    VKey<Recurring> existingAutorenewEvent = domain.getAutorenewBillingEvent();
     // Set domain to have auto-renewed just before the transfer request, so that it will have an
     // active autorenew grace period spanning the entire transfer window.
     DateTime autorenewTime = clock.nowUtc().minusDays(1);

core/src/test/java/google/registry/flows/domain/DomainTransferRequestFlowTest.java

@@ -318,8 +318,7 @@ class DomainTransferRequestFlowTest
             BillingEvent.class),
         Sets.union(expectedServeApproveBillingEvents, extraBillingEvents));
     // The domain's autorenew billing event should still point to the losing client's event.
-    BillingEvent.Recurring domainAutorenewEvent =
+    BillingEvent.Recurring domainAutorenewEvent = tm().load(domain.getAutorenewBillingEvent());
-        ofy().load().key(domain.getAutorenewBillingEvent()).now();
     assertThat(domainAutorenewEvent.getClientId()).isEqualTo("TheRegistrar");
     assertThat(domainAutorenewEvent.getRecurrenceEndTime()).isEqualTo(implicitTransferTime);
     // The original grace periods should remain untouched.
@@ -446,12 +445,7 @@ class DomainTransferRequestFlowTest
         .hasLastEppUpdateTime(implicitTransferTime)
         .and()
         .hasLastEppUpdateClientId("NewRegistrar");
-    assertThat(
+    assertThat(tm().load(domainAfterAutomaticTransfer.getAutorenewBillingEvent()).getEventTime())
-            ofy()
-                .load()
-                .key(domainAfterAutomaticTransfer.getAutorenewBillingEvent())
-                .now()
-                .getEventTime())
         .isEqualTo(expectedExpirationTime);
     // And after the expected grace time, the grace period should be gone.
     DomainBase afterGracePeriod =
@@ -955,7 +949,7 @@ class DomainTransferRequestFlowTest
   void testSuccess_superuserExtension_zeroPeriod_autorenewGraceActive() throws Exception {
     eppRequestSource = EppRequestSource.TOOL;
     setupDomain("example", "tld");
-    Key<BillingEvent.Recurring> existingAutorenewEvent = domain.getAutorenewBillingEvent();
+    VKey<BillingEvent.Recurring> existingAutorenewEvent = domain.getAutorenewBillingEvent();
     // Set domain to have auto-renewed just before the transfer request, so that it will have an
     // active autorenew grace period spanning the entire transfer window.
     DateTime autorenewTime = clock.nowUtc().minusDays(1);
@@ -1110,7 +1104,7 @@ class DomainTransferRequestFlowTest
   @Test
   void testSuccess_autorenewGraceActive_throughoutTransferWindow() throws Exception {
     setupDomain("example", "tld");
-    Key<BillingEvent.Recurring> existingAutorenewEvent = domain.getAutorenewBillingEvent();
+    VKey<BillingEvent.Recurring> existingAutorenewEvent = domain.getAutorenewBillingEvent();
     // Set domain to have auto-renewed just before the transfer request, so that it will have an
     // active autorenew grace period spanning the entire transfer window.
     DateTime autorenewTime = clock.nowUtc().minusDays(1);
@@ -1142,13 +1136,13 @@ class DomainTransferRequestFlowTest
             .setEventTime(clock.nowUtc().plus(Registry.get("tld").getAutomaticTransferLength()))
             .setBillingTime(autorenewTime.plus(Registry.get("tld").getAutoRenewGracePeriodLength()))
             // The cancellation should refer to the old autorenew billing event.
-            .setRecurringEventKey(VKey.from(existingAutorenewEvent)));
+            .setRecurringEventKey(existingAutorenewEvent));
   }
 
   @Test
   void testSuccess_autorenewGraceActive_onlyAtAutomaticTransferTime() throws Exception {
     setupDomain("example", "tld");
-    Key<BillingEvent.Recurring> existingAutorenewEvent = domain.getAutorenewBillingEvent();
+    VKey<BillingEvent.Recurring> existingAutorenewEvent = domain.getAutorenewBillingEvent();
     // Set domain to expire in 1 day, so that it will be in the autorenew grace period by the
     // automatic transfer time, even though it isn't yet.
     DateTime expirationTime = clock.nowUtc().plusDays(1);
@@ -1170,7 +1164,7 @@ class DomainTransferRequestFlowTest
             .setBillingTime(
                 expirationTime.plus(Registry.get("tld").getAutoRenewGracePeriodLength()))
             // The cancellation should refer to the old autorenew billing event.
-            .setRecurringEventKey(VKey.from(existingAutorenewEvent)));
+            .setRecurringEventKey(existingAutorenewEvent));
   }
 
   @Test

core/src/test/java/google/registry/model/billing/BillingEventTest.java

@@ -396,12 +396,13 @@ public class BillingEventTest extends EntityTestCase {
 
   @Test
   void testSuccess_cancellation_forGracePeriod_withRecurring() {
-    BillingEvent.Cancellation newCancellation = BillingEvent.Cancellation.forGracePeriod(
+    BillingEvent.Cancellation newCancellation =
+        BillingEvent.Cancellation.forGracePeriod(
             GracePeriod.createForRecurring(
                 GracePeriodStatus.AUTO_RENEW,
                 now.plusYears(1).plusDays(45),
                 "a registrar",
-            Key.create(recurring)),
+                recurring.createVKey()),
             historyEntry2,
             "foo.tld");
     // Set ID to be the same to ignore for the purposes of comparison.

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

@@ -64,19 +64,19 @@ import org.junit.jupiter.api.Test;
 public class DomainBaseTest extends EntityTestCase {
 
   private DomainBase domain;
-  private Key<BillingEvent.OneTime> oneTimeBillKey;
+  private VKey<BillingEvent.OneTime> oneTimeBillKey;
-  private Key<BillingEvent.Recurring> recurringBillKey;
+  private VKey<BillingEvent.Recurring> recurringBillKey;
-  private Key<DomainBase> domainKey;
+  private VKey<DomainBase> domainKey;
 
   @BeforeEach
   void setUp() {
     createTld("com");
-    domainKey = Key.create(null, DomainBase.class, "4-COM");
+    domainKey = VKey.from(Key.create(null, DomainBase.class, "4-COM"));
     VKey<HostResource> hostKey =
         persistResource(
                 new HostResource.Builder()
                     .setHostName("ns1.example.com")
-                    .setSuperordinateDomain(VKey.from(domainKey))
+                    .setSuperordinateDomain(domainKey)
                     .setRepoId("1-COM")
                     .build())
             .createVKey();
@@ -95,13 +95,14 @@ public class DomainBaseTest extends EntityTestCase {
                     .build())
             .createVKey();
     Key<HistoryEntry> historyEntryKey =
-        Key.create(persistResource(new HistoryEntry.Builder().setParent(domainKey).build()));
+        Key.create(
-    oneTimeBillKey = Key.create(historyEntryKey, BillingEvent.OneTime.class, 1);
+            persistResource(new HistoryEntry.Builder().setParent(domainKey.getOfyKey()).build()));
-    recurringBillKey = Key.create(historyEntryKey, BillingEvent.Recurring.class, 2);
+    oneTimeBillKey = VKey.from(Key.create(historyEntryKey, BillingEvent.OneTime.class, 1));
-    Key<PollMessage.Autorenew> autorenewPollKey =
+    recurringBillKey = VKey.from(Key.create(historyEntryKey, BillingEvent.Recurring.class, 2));
-        Key.create(historyEntryKey, PollMessage.Autorenew.class, 3);
+    VKey<PollMessage.Autorenew> autorenewPollKey =
-    Key<PollMessage.OneTime> onetimePollKey =
+        VKey.from(Key.create(historyEntryKey, PollMessage.Autorenew.class, 3));
-        Key.create(historyEntryKey, PollMessage.OneTime.class, 1);
+    VKey<PollMessage.OneTime> onetimePollKey =
+        VKey.from(Key.create(historyEntryKey, PollMessage.OneTime.class, 1));
     // Set up a new persisted domain entity.
     domain =
         persistResource(
@@ -138,13 +139,10 @@ public class DomainBaseTest extends EntityTestCase {
                             .setLosingClientId("losing")
                             .setPendingTransferExpirationTime(fakeClock.nowUtc())
                             .setServerApproveEntities(
-                                ImmutableSet.of(
+                                ImmutableSet.of(oneTimeBillKey, recurringBillKey, autorenewPollKey))
-                                    VKey.from(oneTimeBillKey),
+                            .setServerApproveBillingEvent(oneTimeBillKey)
-                                    VKey.from(recurringBillKey),
+                            .setServerApproveAutorenewEvent(recurringBillKey)
-                                    VKey.from(autorenewPollKey)))
+                            .setServerApproveAutorenewPollMessage(autorenewPollKey)
-                            .setServerApproveBillingEvent(VKey.from(oneTimeBillKey))
-                            .setServerApproveAutorenewEvent(VKey.from(recurringBillKey))
-                            .setServerApproveAutorenewPollMessage(VKey.from(autorenewPollKey))
                             .setTransferRequestTime(fakeClock.nowUtc().plusDays(1))
                             .setTransferStatus(TransferStatus.SERVER_APPROVED)
                             .setTransferRequestTrid(Trid.create("client-trid", "server-trid"))
@@ -327,7 +325,7 @@ public class DomainBaseTest extends EntityTestCase {
   private void assertTransferred(
       DomainBase domain,
       DateTime newExpirationTime,
-      Key<BillingEvent.Recurring> newAutorenewEvent) {
+      VKey<BillingEvent.Recurring> newAutorenewEvent) {
     assertThat(domain.getTransferData().getTransferStatus())
         .isEqualTo(TransferStatus.SERVER_APPROVED);
     assertThat(domain.getCurrentSponsorClientId()).isEqualTo("winner");
@@ -377,8 +375,8 @@ public class DomainBaseTest extends EntityTestCase {
             .build();
     DomainBase afterTransfer = domain.cloneProjectedAtTime(fakeClock.nowUtc().plusDays(1));
     DateTime newExpirationTime = oldExpirationTime.plusYears(1);
-    Key<BillingEvent.Recurring> serverApproveAutorenewEvent =
+    VKey<BillingEvent.Recurring> serverApproveAutorenewEvent =
-        domain.getTransferData().getServerApproveAutorenewEvent().getOfyKey();
+        domain.getTransferData().getServerApproveAutorenewEvent();
     assertTransferred(afterTransfer, newExpirationTime, serverApproveAutorenewEvent);
     assertThat(afterTransfer.getGracePeriods())
         .containsExactly(
@@ -389,7 +387,7 @@ public class DomainBaseTest extends EntityTestCase {
                     .plusDays(1)
                     .plus(Registry.get("com").getTransferGracePeriodLength()),
                 "winner",
-                Key.create(transferBillingEvent)));
+                transferBillingEvent.createVKey()));
     // If we project after the grace period expires all should be the same except the grace period.
     DomainBase afterGracePeriod =
         domain.cloneProjectedAtTime(
@@ -747,8 +745,8 @@ public class DomainBaseTest extends EntityTestCase {
             .setPendingTransferExpirationTime(transferExpirationTime)
             .setTransferStatus(TransferStatus.PENDING)
             .setGainingClientId("TheRegistrar")
-            .setServerApproveAutorenewEvent(VKey.from(recurringBillKey))
+            .setServerApproveAutorenewEvent(recurringBillKey)
-            .setServerApproveBillingEvent(VKey.from(oneTimeBillKey))
+            .setServerApproveBillingEvent(oneTimeBillKey)
             .build();
     domain =
         persistResource(

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

@@ -24,6 +24,7 @@ import google.registry.model.billing.BillingEvent.Reason;
 import google.registry.model.billing.BillingEvent.Recurring;
 import google.registry.model.domain.rgp.GracePeriodStatus;
 import google.registry.model.reporting.HistoryEntry;
+import google.registry.persistence.VKey;
 import google.registry.testing.AppEngineRule;
 import org.joda.money.CurrencyUnit;
 import org.joda.money.Money;
@@ -46,12 +47,14 @@ public class GracePeriodTest {
 
   @BeforeEach
   void before() {
-    onetime = new BillingEvent.OneTime.Builder()
+    onetime =
+        new BillingEvent.OneTime.Builder()
             .setEventTime(now)
             .setBillingTime(now.plusDays(1))
             .setClientId("TheRegistrar")
             .setCost(Money.of(CurrencyUnit.USD, 42))
-      .setParent(Key.create(HistoryEntry.class, 12345))
+            .setParent(
+                Key.create(Key.create(DomainBase.class, "domain"), HistoryEntry.class, 12345))
             .setReason(Reason.CREATE)
             .setPeriodYears(1)
             .setTargetId("foo.google")
@@ -62,7 +65,7 @@ public class GracePeriodTest {
   void testSuccess_forBillingEvent() {
     GracePeriod gracePeriod = GracePeriod.forBillingEvent(GracePeriodStatus.ADD, onetime);
     assertThat(gracePeriod.getType()).isEqualTo(GracePeriodStatus.ADD);
-    assertThat(gracePeriod.getOneTimeBillingEvent()).isEqualTo(Key.create(onetime));
+    assertThat(gracePeriod.getOneTimeBillingEvent()).isEqualTo(onetime.createVKey());
     assertThat(gracePeriod.getRecurringBillingEvent()).isNull();
     assertThat(gracePeriod.getClientId()).isEqualTo("TheRegistrar");
     assertThat(gracePeriod.getExpirationTime()).isEqualTo(now.plusDays(1));
@@ -100,7 +103,7 @@ public class GracePeriodTest {
                     GracePeriodStatus.RENEW,
                     now.plusDays(1),
                     "TheRegistrar",
-                    Key.create(Recurring.class, 12345)));
+                    VKey.create(Recurring.class, 12345)));
     assertThat(thrown).hasMessageThat().contains("autorenew");
   }
 }

core/src/test/java/google/registry/rde/DomainBaseToXjcConverterTest.java

@@ -31,7 +31,6 @@ import static org.joda.money.CurrencyUnit.USD;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.net.InetAddresses;
-import com.googlecode.objectify.Key;
 import google.registry.model.billing.BillingEvent;
 import google.registry.model.billing.BillingEvent.Flag;
 import google.registry.model.billing.BillingEvent.Reason;
@@ -302,7 +301,6 @@ public class DomainBaseToXjcConverterTest {
                     StatusValue.CLIENT_TRANSFER_PROHIBITED,
                     StatusValue.SERVER_UPDATE_PROHIBITED))
             .setAutorenewBillingEvent(
-                Key.create(
                 persistResource(
                         new BillingEvent.Recurring.Builder()
                             .setReason(Reason.RENEW)
@@ -312,9 +310,9 @@ public class DomainBaseToXjcConverterTest {
                             .setEventTime(END_OF_TIME)
                             .setRecurrenceEndTime(END_OF_TIME)
                             .setParent(historyEntry)
-                            .build())))
+                            .build())
+                    .createVKey())
             .setAutorenewPollMessage(
-                Key.create(
                 persistResource(
                         new PollMessage.Autorenew.Builder()
                             .setTargetId("lol")
@@ -323,7 +321,8 @@ public class DomainBaseToXjcConverterTest {
                             .setAutorenewEndTime(END_OF_TIME)
                             .setMsg("Domain was auto-renewed.")
                             .setParent(historyEntry)
-                            .build())))
+                            .build())
+                    .createVKey())
             .setTransferData(
                 new DomainTransferData.Builder()
                     .setGainingClientId("gaining")

core/src/test/java/google/registry/rde/RdeFixtures.java

@@ -26,7 +26,6 @@ import static org.joda.money.CurrencyUnit.USD;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.net.InetAddresses;
-import com.googlecode.objectify.Key;
 import google.registry.model.billing.BillingEvent;
 import google.registry.model.billing.BillingEvent.Flag;
 import google.registry.model.billing.BillingEvent.Reason;
@@ -147,7 +146,6 @@ final class RdeFixtures {
                     StatusValue.CLIENT_TRANSFER_PROHIBITED,
                     StatusValue.SERVER_UPDATE_PROHIBITED))
             .setAutorenewBillingEvent(
-                Key.create(
                 persistResource(
                         new BillingEvent.Recurring.Builder()
                             .setReason(Reason.RENEW)
@@ -157,9 +155,9 @@ final class RdeFixtures {
                             .setEventTime(END_OF_TIME)
                             .setRecurrenceEndTime(END_OF_TIME)
                             .setParent(historyEntry)
-                            .build())))
+                            .build())
+                    .createVKey())
             .setAutorenewPollMessage(
-                Key.create(
                 persistSimpleResource(
                         new PollMessage.Autorenew.Builder()
                             .setTargetId(tld)
@@ -168,7 +166,8 @@ final class RdeFixtures {
                             .setAutorenewEndTime(END_OF_TIME)
                             .setMsg("Domain was auto-renewed.")
                             .setParent(historyEntry)
-                            .build())))
+                            .build())
+                    .createVKey())
             .setTransferData(
                 new DomainTransferData.Builder()
                     .setGainingClientId("gaining")

core/src/test/java/google/registry/testing/DatastoreHelper.java

@@ -545,8 +545,8 @@ public class DatastoreHelper {
     return persistResource(
         domain
             .asBuilder()
-            .setAutorenewBillingEvent(Key.create(autorenewEvent))
+            .setAutorenewBillingEvent(autorenewEvent.createVKey())
-            .setAutorenewPollMessage(Key.create(autorenewPollMessage))
+            .setAutorenewPollMessage(autorenewPollMessage.createVKey())
             .build());
   }
 
@@ -588,13 +588,13 @@ public class DatastoreHelper {
                 .build());
     // Modify the existing autorenew event to reflect the pending transfer.
     persistResource(
-        ofy().load().key(domain.getAutorenewBillingEvent()).now().asBuilder()
+        tm().load(domain.getAutorenewBillingEvent())
+            .asBuilder()
             .setRecurrenceEndTime(expirationTime)
             .build());
     // Update the end time of the existing autorenew poll message. We must delete it if it has no
     // events left in it.
-    PollMessage.Autorenew autorenewPollMessage =
+    PollMessage.Autorenew autorenewPollMessage = tm().load(domain.getAutorenewPollMessage());
-        ofy().load().key(domain.getAutorenewPollMessage()).now();
     if (autorenewPollMessage.getEventTime().isBefore(expirationTime)) {
       persistResource(
           autorenewPollMessage.asBuilder()

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

@@ -20,6 +20,7 @@ import static google.registry.model.eppcommon.StatusValue.PENDING_DELETE;
 import static google.registry.model.eppcommon.StatusValue.PENDING_TRANSFER;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.reporting.HistoryEntry.Type.SYNTHETIC;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatastoreHelper.assertBillingEventsEqual;
 import static google.registry.testing.DatastoreHelper.assertPollMessagesEqual;
 import static google.registry.testing.DatastoreHelper.createTld;
@@ -119,7 +120,7 @@ public class UnrenewDomainCommandTest extends CommandTestCase<UnrenewDomainComma
     HistoryEntry synthetic = getOnlyHistoryEntryOfType(domain, SYNTHETIC);
 
     assertBillingEventsEqual(
-        ofy().load().key(domain.getAutorenewBillingEvent()).now(),
+        tm().load(domain.getAutorenewBillingEvent()),
         new BillingEvent.Recurring.Builder()
             .setParent(synthetic)
             .setReason(Reason.RENEW)
@@ -148,7 +149,8 @@ public class UnrenewDomainCommandTest extends CommandTestCase<UnrenewDomainComma
                 .build()));
 
     // Check that fields on domain were updated correctly.
-    assertThat(domain.getAutorenewPollMessage().getParent()).isEqualTo(Key.create(synthetic));
+    assertThat(domain.getAutorenewPollMessage().getOfyKey().getParent())
+        .isEqualTo(Key.create(synthetic));
     assertThat(domain.getRegistrationExpirationTime()).isEqualTo(newExpirationTime);
     assertThat(domain.getLastEppUpdateTime()).isEqualTo(unrenewTime);
     assertThat(domain.getLastEppUpdateClientId()).isEqualTo("TheRegistrar");

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

@@ -164,14 +164,14 @@ class google.registry.model.domain.DomainAuthInfo {
 class google.registry.model.domain.DomainBase {
   @Id java.lang.String repoId;
   com.google.common.collect.ImmutableSortedMap<org.joda.time.DateTime, com.googlecode.objectify.Key<google.registry.model.ofy.CommitLogManifest>> revisions;
-  com.googlecode.objectify.Key<google.registry.model.billing.BillingEvent$Recurring> autorenewBillingEvent;
-  com.googlecode.objectify.Key<google.registry.model.poll.PollMessage$Autorenew> autorenewPollMessage;
-  com.googlecode.objectify.Key<google.registry.model.poll.PollMessage$OneTime> deletePollMessage;
   google.registry.model.CreateAutoTimestamp creationTime;
   google.registry.model.UpdateAutoTimestamp updateTimestamp;
   google.registry.model.domain.DomainAuthInfo authInfo;
   google.registry.model.domain.launch.LaunchNotice launchNotice;
   google.registry.model.transfer.DomainTransferData transferData;
+  google.registry.persistence.VKey<google.registry.model.billing.BillingEvent$Recurring> autorenewBillingEvent;
+  google.registry.persistence.VKey<google.registry.model.poll.PollMessage$Autorenew> autorenewPollMessage;
+  google.registry.persistence.VKey<google.registry.model.poll.PollMessage$OneTime> deletePollMessage;
   java.lang.String creationClientId;
   java.lang.String currentSponsorClientId;
   java.lang.String fullyQualifiedDomainName;
@@ -191,9 +191,9 @@ class google.registry.model.domain.DomainBase {
   org.joda.time.DateTime registrationExpirationTime;
 }
 class google.registry.model.domain.GracePeriod {
-  com.googlecode.objectify.Key<google.registry.model.billing.BillingEvent$OneTime> billingEventOneTime;
-  com.googlecode.objectify.Key<google.registry.model.billing.BillingEvent$Recurring> billingEventRecurring;
   google.registry.model.domain.rgp.GracePeriodStatus type;
+  google.registry.persistence.VKey<google.registry.model.billing.BillingEvent$OneTime> billingEventOneTime;
+  google.registry.persistence.VKey<google.registry.model.billing.BillingEvent$Recurring> billingEventRecurring;
   java.lang.String clientId;
   org.joda.time.DateTime expirationTime;
 }

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

@@ -0,0 +1,32 @@
+-- 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 "Domain" ADD COLUMN billing_recurrence_id int8;
+ALTER TABLE "Domain" ADD COLUMN autorenew_poll_message_id int8;
+ALTER TABLE "Domain" ADD COLUMN deletion_poll_message_id int8;
+
+ALTER TABLE IF EXISTS "Domain"
+   ADD CONSTRAINT fk_domain_billing_recurrence_id
+   FOREIGN KEY (billing_recurrence_id)
+   REFERENCES "BillingEvent";
+
+ALTER TABLE IF EXISTS "Domain"
+   ADD CONSTRAINT fk_domain_autorenew_poll_message_id
+   FOREIGN KEY (autorenew_poll_message_id)
+   REFERENCES "PollMessage";
+
+ALTER TABLE IF EXISTS "Domain"
+   ADD CONSTRAINT fk_domain_deletion_poll_message_id
+   FOREIGN KEY (deletion_poll_message_id)
+   REFERENCES "PollMessage";

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

@@ -232,7 +232,10 @@ create sequence history_id_sequence start 1 increment 1;
         admin_contact text,
         auth_info_repo_id text,
         auth_info_value text,
+        billing_recurrence_id int8,
+        autorenew_poll_message_id int8,
         billing_contact text,
+        deletion_poll_message_id int8,
         domain_name text,
         idn_table_name text,
         last_transfer_time timestamptz,
@@ -272,8 +275,8 @@ create sequence history_id_sequence start 1 increment 1;
 
     create table "GracePeriod" (
        id  bigserial not null,
-        billing_event_one_time bytea,
+        billing_event_one_time int8,
-        billing_event_recurring bytea,
+        billing_event_recurring int8,
         registrar_id text,
         expiration_time timestamptz,
         type int4,

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

@@ -398,7 +398,10 @@ CREATE TABLE public."Domain" (
     transfer_pending_expiration_time timestamp with time zone,
     transfer_request_time timestamp with time zone,
     transfer_status text,
-    update_timestamp timestamp with time zone
+    update_timestamp timestamp with time zone,
+    billing_recurrence_id bigint,
+    autorenew_poll_message_id bigint,
+    deletion_poll_message_id bigint
 );
 
 
@@ -1452,6 +1455,14 @@ ALTER TABLE ONLY public."Domain"
     ADD CONSTRAINT fk_domain_admin_contact FOREIGN KEY (admin_contact) REFERENCES public."Contact"(repo_id);
 
 
+--
+-- Name: Domain fk_domain_autorenew_poll_message_id; Type: FK CONSTRAINT; Schema: public; Owner: -
+--
+
+ALTER TABLE ONLY public."Domain"
+    ADD CONSTRAINT fk_domain_autorenew_poll_message_id FOREIGN KEY (autorenew_poll_message_id) REFERENCES public."PollMessage"(poll_message_id);
+
+
 --
 -- Name: Domain fk_domain_billing_contact; Type: FK CONSTRAINT; Schema: public; Owner: -
 --
@@ -1460,6 +1471,22 @@ ALTER TABLE ONLY public."Domain"
     ADD CONSTRAINT fk_domain_billing_contact FOREIGN KEY (billing_contact) REFERENCES public."Contact"(repo_id);
 
 
+--
+-- Name: Domain fk_domain_billing_recurrence_id; Type: FK CONSTRAINT; Schema: public; Owner: -
+--
+
+ALTER TABLE ONLY public."Domain"
+    ADD CONSTRAINT fk_domain_billing_recurrence_id FOREIGN KEY (billing_recurrence_id) REFERENCES public."BillingEvent"(billing_event_id);
+
+
+--
+-- Name: Domain fk_domain_deletion_poll_message_id; Type: FK CONSTRAINT; Schema: public; Owner: -
+--
+
+ALTER TABLE ONLY public."Domain"
+    ADD CONSTRAINT fk_domain_deletion_poll_message_id FOREIGN KEY (deletion_poll_message_id) REFERENCES public."PollMessage"(poll_message_id);
+
+
 --
 -- Name: Domain fk_domain_registrant_contact; Type: FK CONSTRAINT; Schema: public; Owner: -
 --