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))
-            .setAutorenewPollMessage(Key.create(autorenewPollMessage))
+            .setAutorenewBillingEvent(autorenewBillingEvent.createVKey())
+            .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 =
-              ofy().load().key(gracePeriod.getOneTimeBillingEvent()).now();
+          BillingEvent.OneTime oneTime = tm().load(gracePeriod.getOneTimeBillingEvent());
           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);
+              .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))
-            .setAutorenewPollMessage(Key.create(newAutorenewPollMessage))
+            .setAutorenewBillingEvent(newAutorenewEvent.createVKey())
+            .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());
-    ofy().delete().key(existingDomain.getDeletePollMessage());
+    tm().saveNewOrUpdateAll(entitiesToSave.build());
+    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))
-        .setAutorenewPollMessage(Key.create(autorenewPollMessage))
+        .setAutorenewBillingEvent(autorenewEvent.createVKey())
+        .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))
-            .setAutorenewPollMessage(Key.create(gainingClientAutorenewPollMessage))
+            .setAutorenewBillingEvent(autorenewEvent.createVKey())
+            .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(
-                  transferData.getServerApproveAutorenewEvent() == null
-                      ? null
-                      : transferData.getServerApproveAutorenewEvent().getOfyKey())
-              .setAutorenewPollMessage(
-                  transferData.getServerApproveAutorenewPollMessage() == null
-                      ? null
-                      : transferData.getServerApproveAutorenewPollMessage().getOfyKey());
+              .setAutorenewBillingEvent(transferData.getServerApproveAutorenewEvent())
+              .setAutorenewPollMessage(transferData.getServerApproveAutorenewPollMessage());
       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
-                        ? null
-                        : transferData.getServerApproveBillingEvent().getOfyKey())));
+                    transferData.getServerApproveBillingEvent())));
       } 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(
-        Key<BillingEvent.Recurring> autorenewBillingEvent) {
+    public Builder setAutorenewBillingEvent(VKey<BillingEvent.Recurring> autorenewBillingEvent) {
       getInstance().autorenewBillingEvent = autorenewBillingEvent;
       return this;
     }
 
-    public Builder setAutorenewPollMessage(
-        Key<PollMessage.Autorenew> autorenewPollMessage) {
+    public Builder setAutorenewPollMessage(VKey<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
-   * {@code billingEventRecurring}) or for redemption grace periods (since deletes have no cost).
+   * null if not applicable. Not set for autorenew grace periods (which instead use the field {@code
+   * 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 Key<BillingEvent.OneTime> getOneTimeBillingEvent() {
+  public VKey<BillingEvent.OneTime> getOneTimeBillingEvent() {
     return billingEventOneTime;
   }
 
@@ -100,16 +100,16 @@ 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;
   }
 
   private static GracePeriod createInternal(
-       GracePeriodStatus type,
-       DateTime expirationTime,
-       String clientId,
-       @Nullable Key<BillingEvent.OneTime> billingEventOneTime,
-       @Nullable Key<BillingEvent.Recurring> billingEventRecurring) {
+      GracePeriodStatus type,
+      DateTime expirationTime,
+      String clientId,
+      @Nullable VKey<BillingEvent.OneTime> billingEventOneTime,
+      @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
-   * need to avoid loading the BillingEvent from Datastore.  This method should typically be
-   * called only from test code to explicitly construct GracePeriods.
+   * <p>Normal callers should always use {@link #forBillingEvent} instead, assuming they do not need
+   * to avoid loading the BillingEvent from Datastore. This method should typically be called only
+   * 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))
-            .setAutorenewPollMessage(Key.create(newAutorenewPollMessage))
+            .setAutorenewBillingEvent(newAutorenewEvent.createVKey())
+            .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