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

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

@@ -21,6 +21,7 @@ import static google.registry.testing.SqlHelper.assertThrowForeignKeyViolation;
 import static google.registry.testing.SqlHelper.saveRegistrar;
 import static google.registry.util.DateTimeUtils.END_OF_TIME;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
+import static org.joda.money.CurrencyUnit.USD;
 import static org.joda.time.DateTimeZone.UTC;
 import static org.junit.jupiter.api.Assertions.fail;
 
@@ -42,12 +43,14 @@ import google.registry.model.host.HostResource;
 import google.registry.model.poll.PollMessage;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.transfer.ContactTransferData;
+import google.registry.model.transfer.DomainTransferData;
 import google.registry.persistence.VKey;
 import google.registry.persistence.transaction.JpaTestRules;
 import google.registry.persistence.transaction.JpaTestRules.JpaIntegrationWithCoverageExtension;
 import google.registry.testing.DatastoreEntityExtension;
 import google.registry.testing.FakeClock;
 import java.util.Arrays;
+import org.joda.money.Money;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Order;
@@ -445,6 +448,7 @@ public class DomainBaseSqlTest {
             () -> {
               historyEntry =
                   new DomainHistory.Builder()
+                      .setId(100L)
                       .setType(HistoryEntry.Type.DOMAIN_CREATE)
                       .setPeriod(Period.create(1, Period.Unit.YEARS))
                       .setModificationTime(DateTime.now(UTC))
@@ -458,6 +462,7 @@ public class DomainBaseSqlTest {
                       .build();
               BillingEvent.Recurring billEvent =
                   new BillingEvent.Recurring.Builder()
+                      .setId(200L)
                       .setReason(Reason.RENEW)
                       .setFlags(ImmutableSet.of(Flag.AUTO_RENEW))
                       .setTargetId("example.com")
@@ -470,16 +475,37 @@ public class DomainBaseSqlTest {
                       .build();
               PollMessage.Autorenew autorenewPollMessage =
                   new PollMessage.Autorenew.Builder()
+                      .setId(300L)
                       .setClientId("registrar1")
                       .setEventTime(DateTime.now(UTC).plusYears(1))
                       .setParent(historyEntry)
                       .build();
               PollMessage.OneTime deletePollMessage =
                   new PollMessage.OneTime.Builder()
+                      .setId(400L)
                       .setClientId("registrar1")
                       .setEventTime(DateTime.now(UTC).plusYears(1))
                       .setParent(historyEntry)
                       .build();
+              BillingEvent.OneTime oneTimeBillingEvent =
+                  new BillingEvent.OneTime.Builder()
+                      .setId(500L)
+                      // Use SERVER_STATUS so we don't have to add a period.
+                      .setReason(Reason.SERVER_STATUS)
+                      .setTargetId("example.com")
+                      .setClientId("registrar1")
+                      .setDomainRepoId("4-COM")
+                      .setBillingTime(DateTime.now(UTC))
+                      .setCost(Money.of(USD, 100))
+                      .setEventTime(DateTime.now(UTC).plusYears(1))
+                      .setParent(historyEntry)
+                      .build();
+              DomainTransferData transferData =
+                  new DomainTransferData.Builder()
+                      .setServerApproveBillingEvent(oneTimeBillingEvent.createVKey())
+                      .setServerApproveAutorenewEvent(billEvent.createVKey())
+                      .setServerApproveAutorenewPollMessage(autorenewPollMessage.createVKey())
+                      .build();
 
               jpaTm().insert(contact);
               jpaTm().insert(contact2);
@@ -490,12 +516,14 @@ public class DomainBaseSqlTest {
                       .setAutorenewBillingEvent(billEvent.createVKey())
                       .setAutorenewPollMessage(autorenewPollMessage.createVKey())
                       .setDeletePollMessage(deletePollMessage.createVKey())
+                      .setTransferData(transferData)
                       .build();
               historyEntry = historyEntry.asBuilder().setDomainContent(domain).build();
               jpaTm().insert(historyEntry);
               jpaTm().insert(autorenewPollMessage);
               jpaTm().insert(billEvent);
               jpaTm().insert(deletePollMessage);
+              jpaTm().insert(oneTimeBillingEvent);
               jpaTm().insert(domain);
             });
 
@@ -516,6 +544,15 @@ public class DomainBaseSqlTest {
         .isEqualTo(domain.getAutorenewBillingEvent());
     assertThat(persistedHistoryEntry.getDomainContent().get().getDeletePollMessage())
         .isEqualTo(domain.getDeletePollMessage());
+    DomainTransferData persistedTransferData =
+        persistedHistoryEntry.getDomainContent().get().getTransferData();
+    DomainTransferData originalTransferData = domain.getTransferData();
+    assertThat(persistedTransferData.getServerApproveBillingEvent())
+        .isEqualTo(originalTransferData.getServerApproveBillingEvent());
+    assertThat(persistedTransferData.getServerApproveAutorenewEvent())
+        .isEqualTo(originalTransferData.getServerApproveAutorenewEvent());
+    assertThat(persistedTransferData.getServerApproveAutorenewPollMessage())
+        .isEqualTo(originalTransferData.getServerApproveAutorenewPollMessage());
   }
 
   @Test
@@ -525,6 +562,7 @@ public class DomainBaseSqlTest {
             () -> {
               historyEntry =
                   new DomainHistory.Builder()
+                      .setId(100L)
                       .setType(HistoryEntry.Type.DOMAIN_CREATE)
                       .setPeriod(Period.create(1, Period.Unit.YEARS))
                       .setModificationTime(DateTime.now(UTC))
@@ -538,6 +576,7 @@ public class DomainBaseSqlTest {
                       .build();
               BillingEvent.Recurring billEvent =
                   new BillingEvent.Recurring.Builder()
+                      .setId(200L)
                       .setReason(Reason.RENEW)
                       .setFlags(ImmutableSet.of(Flag.AUTO_RENEW))
                       .setTargetId("example.com")
@@ -550,16 +589,41 @@ public class DomainBaseSqlTest {
                       .build();
               PollMessage.Autorenew autorenewPollMessage =
                   new PollMessage.Autorenew.Builder()
+                      .setId(300L)
                       .setClientId("registrar1")
                       .setEventTime(DateTime.now(UTC).plusYears(1))
                       .setParent(historyEntry)
                       .build();
               PollMessage.OneTime deletePollMessage =
                   new PollMessage.OneTime.Builder()
+                      .setId(400L)
                       .setClientId("registrar1")
                       .setEventTime(DateTime.now(UTC).plusYears(1))
                       .setParent(historyEntry)
                       .build();
+              BillingEvent.OneTime oneTimeBillingEvent =
+                  new BillingEvent.OneTime.Builder()
+                      .setId(500L)
+                      // Use SERVER_STATUS so we don't have to add a period.
+                      .setReason(Reason.SERVER_STATUS)
+                      .setTargetId("example.com")
+                      .setClientId("registrar1")
+                      .setDomainRepoId("4-COM")
+                      .setBillingTime(DateTime.now(UTC))
+                      .setCost(Money.of(USD, 100))
+                      .setEventTime(DateTime.now(UTC).plusYears(1))
+                      .setParent(historyEntry)
+                      .build();
+              DomainTransferData transferData =
+                  new DomainTransferData.Builder()
+                      .setServerApproveBillingEvent(
+                          createLegacyVKey(BillingEvent.OneTime.class, oneTimeBillingEvent.getId()))
+                      .setServerApproveAutorenewEvent(
+                          createLegacyVKey(BillingEvent.Recurring.class, billEvent.getId()))
+                      .setServerApproveAutorenewPollMessage(
+                          createLegacyVKey(
+                              PollMessage.Autorenew.class, autorenewPollMessage.getId()))
+                      .build();
 
               jpaTm().insert(contact);
               jpaTm().insert(contact2);
@@ -570,15 +634,18 @@ public class DomainBaseSqlTest {
                       .setAutorenewBillingEvent(
                           createLegacyVKey(BillingEvent.Recurring.class, billEvent.getId()))
                       .setAutorenewPollMessage(
-                          createLegacyVKey(PollMessage.Autorenew.class, deletePollMessage.getId()))
+                          createLegacyVKey(
+                              PollMessage.Autorenew.class, autorenewPollMessage.getId()))
                       .setDeletePollMessage(
-                          createLegacyVKey(PollMessage.OneTime.class, autorenewPollMessage.getId()))
+                          createLegacyVKey(PollMessage.OneTime.class, deletePollMessage.getId()))
+                      .setTransferData(transferData)
                       .build();
               historyEntry = historyEntry.asBuilder().setDomainContent(domain).build();
               jpaTm().insert(historyEntry);
               jpaTm().insert(autorenewPollMessage);
               jpaTm().insert(billEvent);
               jpaTm().insert(deletePollMessage);
+              jpaTm().insert(oneTimeBillingEvent);
               jpaTm().insert(domain);
             });
 
@@ -599,6 +666,15 @@ public class DomainBaseSqlTest {
         .isEqualTo(domain.getAutorenewBillingEvent());
     assertThat(persistedHistoryEntry.getDomainContent().get().getDeletePollMessage())
         .isEqualTo(domain.getDeletePollMessage());
+    DomainTransferData persistedTransferData =
+        persistedHistoryEntry.getDomainContent().get().getTransferData();
+    DomainTransferData originalTransferData = domain.getTransferData();
+    assertThat(persistedTransferData.getServerApproveBillingEvent())
+        .isEqualTo(originalTransferData.getServerApproveBillingEvent());
+    assertThat(persistedTransferData.getServerApproveAutorenewEvent())
+        .isEqualTo(originalTransferData.getServerApproveAutorenewEvent());
+    assertThat(persistedTransferData.getServerApproveAutorenewPollMessage())
+        .isEqualTo(originalTransferData.getServerApproveAutorenewPollMessage());
   }
 
   private <T> VKey<T> createLegacyVKey(Class<T> clazz, long id) {

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

@@ -67,6 +67,7 @@ public class DomainBaseTest extends EntityTestCase {
   private DomainBase domain;
   private VKey<BillingEvent.OneTime> oneTimeBillKey;
   private VKey<BillingEvent.Recurring> recurringBillKey;
+  private Key<HistoryEntry> historyEntryKey;
 
   @BeforeEach
   void setUp() {
@@ -94,7 +95,7 @@ public class DomainBaseTest extends EntityTestCase {
                     .setRepoId("3-COM")
                     .build())
             .createVKey();
-    Key<HistoryEntry> historyEntryKey =
+    historyEntryKey =
         Key.create(
             persistResource(new HistoryEntry.Builder().setParent(domainKey.getOfyKey()).build()));
     oneTimeBillKey = VKey.from(Key.create(historyEntryKey, BillingEvent.OneTime.class, 1));
@@ -164,10 +165,26 @@ public class DomainBaseTest extends EntityTestCase {
 
   @Test
   void testPersistence() {
+    // Note that this only verifies that the value stored under the foreign key is the same as that
+    // stored under the primary key ("domain" is the domain loaded from the datastore, not the
+    // original domain object).
     assertThat(loadByForeignKey(DomainBase.class, domain.getForeignKey(), fakeClock.nowUtc()))
         .hasValue(domain);
   }
 
+  @Test
+  void testVKeyRestoration() {
+    assertThat(domain.deletePollMessageHistoryId).isEqualTo(historyEntryKey.getId());
+    assertThat(domain.autorenewBillingEventHistoryId).isEqualTo(historyEntryKey.getId());
+    assertThat(domain.autorenewPollMessageHistoryId).isEqualTo(historyEntryKey.getId());
+    assertThat(domain.getTransferData().getServerApproveBillingEventHistoryId())
+        .isEqualTo(historyEntryKey.getId());
+    assertThat(domain.getTransferData().getServerApproveAutorenewEventHistoryId())
+        .isEqualTo(historyEntryKey.getId());
+    assertThat(domain.getTransferData().getServerApproveAutorenewPollMessageHistoryId())
+        .isEqualTo(historyEntryKey.getId());
+  }
+
   @Test
   void testIndexing() throws Exception {
     verifyIndexing(

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

@@ -61,3 +61,4 @@ V60__remove_pollmessage_sequence.sql
 V61__domain_hist_columns.sql
 V62__disable_key_auto_generation_for_history_tables.sql
 V63__add_schema_for_ds_data.sql
+V64__transfer_history_columns.sql

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

@@ -0,0 +1,21 @@
+-- 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 transfer_billing_recurrence_history_id int8;
+ALTER TABLE "Domain" ADD COLUMN transfer_autorenew_poll_message_history_id int8;
+ALTER TABLE "Domain" ADD COLUMN transfer_billing_event_history_id int8;
+
+ALTER TABLE "DomainHistory" ADD COLUMN transfer_billing_recurrence_history_id int8;
+ALTER TABLE "DomainHistory" ADD COLUMN transfer_autorenew_poll_message_history_id int8;
+ALTER TABLE "DomainHistory" ADD COLUMN transfer_billing_event_history_id int8;

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

@@ -271,8 +271,11 @@
         tld text,
         transfer_billing_cancellation_id int8,
         transfer_billing_recurrence_id int8,
+        transfer_billing_recurrence_history_id int8,
         transfer_autorenew_poll_message_id int8,
+        transfer_autorenew_poll_message_history_id int8,
         transfer_billing_event_id int8,
+        transfer_billing_event_history_id int8,
         transfer_renew_period_unit text,
         transfer_renew_period_value int4,
         transfer_registration_expiration_time timestamptz,
@@ -326,8 +329,11 @@
         tld text,
         transfer_billing_cancellation_id int8,
         transfer_billing_recurrence_id int8,
+        transfer_billing_recurrence_history_id int8,
         transfer_autorenew_poll_message_id int8,
+        transfer_autorenew_poll_message_history_id int8,
         transfer_billing_event_id int8,
+        transfer_billing_event_history_id int8,
         transfer_renew_period_unit text,
         transfer_renew_period_value int4,
         transfer_registration_expiration_time timestamptz,

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

@@ -369,7 +369,10 @@ CREATE TABLE public."Domain" (
     autorenew_end_time timestamp with time zone,
     billing_recurrence_history_id bigint,
     autorenew_poll_message_history_id bigint,
-    deletion_poll_message_history_id bigint
+    deletion_poll_message_history_id bigint,
+    transfer_billing_recurrence_history_id bigint,
+    transfer_autorenew_poll_message_history_id bigint,
+    transfer_billing_event_history_id bigint
 );
 
 
@@ -439,7 +442,10 @@ CREATE TABLE public."DomainHistory" (
     history_period_value integer,
     billing_recurrence_history_id bigint,
     autorenew_poll_message_history_id bigint,
-    deletion_poll_message_history_id bigint
+    deletion_poll_message_history_id bigint,
+    transfer_billing_recurrence_history_id bigint,
+    transfer_autorenew_poll_message_history_id bigint,
+    transfer_billing_event_history_id bigint
 );