Restore ofy keys in DomainTransferData (#838)

* Restore ofy keys in DomainTransferData

Restore composite VKeys correctly in DomainTransferData (they were previously
missing their ofy keys).

* Use "AlsoLoad" to populate history ids

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

@@ -352,9 +352,13 @@ public class DomainContent extends EppResource
         restoreOfyFrom(myKey, autorenewBillingEvent, autorenewBillingEventHistoryId);
     autorenewPollMessage =
         restoreOfyFrom(myKey, autorenewPollMessage, autorenewPollMessageHistoryId);
+
+    if (transferData != null) {
+      transferData.restoreOfyKeys(myKey);
+    }
   }
 
-  private <T> VKey<T> restoreOfyFrom(Key<DomainBase> domainKey, VKey<T> key, Long historyId) {
+  public static <T> VKey<T> restoreOfyFrom(Key<DomainBase> domainKey, VKey<T> key, Long historyId) {
     if (historyId == null) {
       // This is a legacy key (or a null key, in which case this works too)
       return VKey.restoreOfyFrom(key, EntityGroupRoot.class, "per-tld");
@@ -716,7 +720,14 @@ public class DomainContent extends EppResource
             + " use DomainBase instead");
   }
 
-  private static Long getHistoryId(VKey<?> key) {
+  /**
+   * Obtains a history id from the given key.
+   *
+   * <p>The key must be a composite key either of the form domain-key/history-key/long-event-key or
+   * EntityGroupRoot/long-event-key (for legacy keys). In the latter case or for a null key returns
+   * a history id of null.
+   */
+  public static Long getHistoryId(VKey<?> key) {
     if (key == null) {
       return null;
     }

core/src/main/java/google/registry/model/transfer/DomainTransferData.java

@@ -14,12 +14,16 @@
 
 package google.registry.model.transfer;
 
+import com.google.common.annotations.VisibleForTesting;
+import com.googlecode.objectify.Key;
+import com.googlecode.objectify.annotation.AlsoLoad;
 import com.googlecode.objectify.annotation.Embed;
 import com.googlecode.objectify.annotation.Ignore;
 import com.googlecode.objectify.annotation.IgnoreSave;
 import com.googlecode.objectify.annotation.Unindex;
 import com.googlecode.objectify.condition.IfNull;
 import google.registry.model.billing.BillingEvent;
+import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.Period;
 import google.registry.model.domain.Period.Unit;
 import google.registry.model.poll.PollMessage;
@@ -86,6 +90,10 @@ public class DomainTransferData extends TransferData<DomainTransferData.Builder>
   @Column(name = "transfer_billing_event_id")
   VKey<BillingEvent.OneTime> serverApproveBillingEvent;
 
+  @Ignore
+  @Column(name = "transfer_billing_event_history_id")
+  Long serverApproveBillingEventHistoryId;
+
   /**
    * The autorenew billing event that should be associated with this resource after the transfer.
    *
@@ -96,6 +104,10 @@ public class DomainTransferData extends TransferData<DomainTransferData.Builder>
   @Column(name = "transfer_billing_recurrence_id")
   VKey<BillingEvent.Recurring> serverApproveAutorenewEvent;
 
+  @Ignore
+  @Column(name = "transfer_billing_recurrence_history_id")
+  Long serverApproveAutorenewEventHistoryId;
+
   /**
    * The autorenew poll message that should be associated with this resource after the transfer.
    *
@@ -106,11 +118,47 @@ public class DomainTransferData extends TransferData<DomainTransferData.Builder>
   @Column(name = "transfer_autorenew_poll_message_id")
   VKey<PollMessage.Autorenew> serverApproveAutorenewPollMessage;
 
+  @Ignore
+  @Column(name = "transfer_autorenew_poll_message_history_id")
+  Long serverApproveAutorenewPollMessageHistoryId;
+
   @Override
   public Builder copyConstantFieldsToBuilder() {
     return super.copyConstantFieldsToBuilder().setTransferPeriod(this.transferPeriod);
   }
 
+  /**
+   * Restores the set of ofy keys after loading from SQL using the specified {@code rootKey}.
+   *
+   * <p>This is for use by DomainBase/DomainHistory PostLoad methods ONLY.
+   */
+  public void restoreOfyKeys(Key<DomainBase> rootKey) {
+    serverApproveBillingEvent =
+        DomainBase.restoreOfyFrom(
+            rootKey, serverApproveBillingEvent, serverApproveBillingEventHistoryId);
+    serverApproveAutorenewEvent =
+        DomainBase.restoreOfyFrom(
+            rootKey, serverApproveAutorenewEvent, serverApproveAutorenewEventHistoryId);
+    serverApproveAutorenewPollMessage =
+        DomainBase.restoreOfyFrom(
+            rootKey, serverApproveAutorenewPollMessage, serverApproveAutorenewPollMessageHistoryId);
+  }
+
+  private void loadServerApproveBillingEventHistoryId(
+      @AlsoLoad("serverApproveBillingEvent") VKey<BillingEvent.OneTime> val) {
+    serverApproveBillingEventHistoryId = DomainBase.getHistoryId(val);
+  }
+
+  private void loadServerApproveAutorenewEventHistoryId(
+      @AlsoLoad("serverApproveAutorenewEvent") VKey<BillingEvent.Recurring> val) {
+    serverApproveAutorenewEventHistoryId = DomainBase.getHistoryId(val);
+  }
+
+  private void loadServerApproveAutorenewPollMessageHistoryId(
+      @AlsoLoad("serverApproveAutorenewPollMessage") VKey<PollMessage.Autorenew> val) {
+    serverApproveAutorenewPollMessageHistoryId = DomainBase.getHistoryId(val);
+  }
+
   public Period getTransferPeriod() {
     return transferPeriod;
   }
@@ -125,16 +173,34 @@ public class DomainTransferData extends TransferData<DomainTransferData.Builder>
     return serverApproveBillingEvent;
   }
 
+  @VisibleForTesting
+  @Nullable
+  public Long getServerApproveBillingEventHistoryId() {
+    return serverApproveBillingEventHistoryId;
+  }
+
   @Nullable
   public VKey<BillingEvent.Recurring> getServerApproveAutorenewEvent() {
     return serverApproveAutorenewEvent;
   }
 
+  @VisibleForTesting
+  @Nullable
+  public Long getServerApproveAutorenewEventHistoryId() {
+    return serverApproveAutorenewEventHistoryId;
+  }
+
   @Nullable
   public VKey<PollMessage.Autorenew> getServerApproveAutorenewPollMessage() {
     return serverApproveAutorenewPollMessage;
   }
 
+  @VisibleForTesting
+  @Nullable
+  public Long getServerApproveAutorenewPollMessageHistoryId() {
+    return serverApproveAutorenewPollMessageHistoryId;
+  }
+
   @Override
   public boolean isEmpty() {
     return EMPTY.equals(this);
@@ -168,18 +234,24 @@ public class DomainTransferData extends TransferData<DomainTransferData.Builder>
     public Builder setServerApproveBillingEvent(
         VKey<BillingEvent.OneTime> serverApproveBillingEvent) {
       getInstance().serverApproveBillingEvent = serverApproveBillingEvent;
+      getInstance().serverApproveBillingEventHistoryId =
+          DomainBase.getHistoryId(serverApproveBillingEvent);
       return this;
     }
 
     public Builder setServerApproveAutorenewEvent(
         VKey<BillingEvent.Recurring> serverApproveAutorenewEvent) {
       getInstance().serverApproveAutorenewEvent = serverApproveAutorenewEvent;
+      getInstance().serverApproveAutorenewEventHistoryId =
+          DomainBase.getHistoryId(serverApproveAutorenewEvent);
       return this;
     }
 
     public Builder setServerApproveAutorenewPollMessage(
         VKey<PollMessage.Autorenew> serverApproveAutorenewPollMessage) {
       getInstance().serverApproveAutorenewPollMessage = serverApproveAutorenewPollMessage;
+      getInstance().serverApproveAutorenewPollMessageHistoryId =
+          DomainBase.getHistoryId(serverApproveAutorenewPollMessage);
       return this;
     }
   }