Fix vkey reconstruction for PollMessage (#823)

* Fix vkey reconstruction for PollMessage

* Add foreign key

* Rebase on HEAD

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

@@ -38,8 +38,6 @@ import javax.persistence.CascadeType;
 import javax.persistence.Column;
 import javax.persistence.ElementCollection;
 import javax.persistence.Entity;
-import javax.persistence.GeneratedValue;
-import javax.persistence.GenerationType;
 import javax.persistence.Id;
 import javax.persistence.IdClass;
 import javax.persistence.Index;
@@ -128,7 +126,6 @@ public class DomainHistory extends HistoryEntry {
   }
 
   @Id
-  @GeneratedValue(strategy = GenerationType.SEQUENCE, generator = "TempHistorySequenceGenerator")
   @Column(name = "historyRevisionId")
   @Access(AccessType.PROPERTY)
   @Override

core/src/main/java/google/registry/model/poll/PollMessage.java

@@ -27,14 +27,18 @@ import com.googlecode.objectify.annotation.EntitySubclass;
 import com.googlecode.objectify.annotation.Id;
 import com.googlecode.objectify.annotation.Ignore;
 import com.googlecode.objectify.annotation.Index;
+import com.googlecode.objectify.annotation.OnLoad;
 import com.googlecode.objectify.annotation.Parent;
 import google.registry.model.Buildable;
 import google.registry.model.EppResource;
 import google.registry.model.ImmutableObject;
 import google.registry.model.annotations.ExternalMessagingName;
 import google.registry.model.annotations.ReportedOn;
+import google.registry.model.contact.ContactResource;
+import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.DomainRenewData;
 import google.registry.model.eppoutput.EppResponse.ResponseData;
+import google.registry.model.host.HostResource;
 import google.registry.model.poll.PendingActionNotificationResponse.ContactPendingActionNotificationResponse;
 import google.registry.model.poll.PendingActionNotificationResponse.DomainPendingActionNotificationResponse;
 import google.registry.model.poll.PendingActionNotificationResponse.HostPendingActionNotificationResponse;
@@ -54,10 +58,9 @@ import javax.persistence.Column;
 import javax.persistence.DiscriminatorColumn;
 import javax.persistence.DiscriminatorValue;
 import javax.persistence.Embedded;
-import javax.persistence.GeneratedValue;
-import javax.persistence.GenerationType;
 import javax.persistence.Inheritance;
 import javax.persistence.InheritanceType;
+import javax.persistence.PostLoad;
 import javax.persistence.Transient;
 import org.joda.time.DateTime;
 
@@ -98,7 +101,6 @@ public abstract class PollMessage extends ImmutableObject
   /** Entity id. */
   @Id
   @javax.persistence.Id
-  @GeneratedValue(strategy = GenerationType.IDENTITY)
   @Column(name = "poll_message_id")
   Long id;
 
@@ -124,11 +126,11 @@ public abstract class PollMessage extends ImmutableObject
 
   @Ignore String hostRepoId;
 
-  @Ignore Long domainRevisionId;
+  @Ignore Long domainHistoryRevisionId;
 
-  @Ignore Long contactRevisionId;
+  @Ignore Long contactHistoryRevisionId;
 
-  @Ignore Long hostRevisionId;
+  @Ignore Long hostHistoryRevisionId;
 
   public Key<HistoryEntry> getParentKey() {
     return parent;
@@ -152,6 +154,34 @@ public abstract class PollMessage extends ImmutableObject
 
   public abstract ImmutableList<ResponseData> getResponseData();
 
+  @PostLoad
+  void postLoad() {
+    if (domainRepoId != null) {
+      parent =
+          Key.create(
+              Key.create(DomainBase.class, domainRepoId),
+              HistoryEntry.class,
+              domainHistoryRevisionId);
+    } else if (contactRepoId != null) {
+      parent =
+          Key.create(
+              Key.create(ContactResource.class, contactRepoId),
+              HistoryEntry.class,
+              contactHistoryRevisionId);
+    } else if (hostHistoryRevisionId != null) {
+      parent =
+          Key.create(
+              Key.create(HostResource.class, hostRepoId),
+              HistoryEntry.class,
+              hostHistoryRevisionId);
+    }
+  }
+
+  @OnLoad
+  void onLoad() {
+    setSqlForeignKeys(this);
+  }
+
   @Override
   public abstract VKey<? extends PollMessage> createVKey();
 
@@ -217,10 +247,30 @@ public abstract class PollMessage extends ImmutableObject
       checkArgumentNotNull(instance.clientId, "clientId must be specified");
       checkArgumentNotNull(instance.eventTime, "eventTime must be specified");
       checkArgumentNotNull(instance.parent, "parent must be specified");
+      checkArgumentNotNull(instance.parent.getParent(), "parent.getParent() must be specified");
+      setSqlForeignKeys(instance);
       return super.build();
     }
   }
 
+  private static void setSqlForeignKeys(PollMessage pollMessage) {
+    String grandparentKind = pollMessage.parent.getParent().getKind();
+    String repoId = pollMessage.parent.getParent().getName();
+    long historyRevisionId = pollMessage.parent.getId();
+    if (Key.getKind(DomainBase.class).equals(grandparentKind)) {
+      pollMessage.domainRepoId = repoId;
+      pollMessage.domainHistoryRevisionId = historyRevisionId;
+    } else if (Key.getKind(ContactResource.class).equals(grandparentKind)) {
+      pollMessage.contactRepoId = repoId;
+      pollMessage.contactHistoryRevisionId = historyRevisionId;
+    } else if (Key.getKind(HostResource.class).equals(grandparentKind)) {
+      pollMessage.hostRepoId = repoId;
+      pollMessage.hostHistoryRevisionId = historyRevisionId;
+    } else {
+      throw new IllegalArgumentException("Unknown grandparent kind: " + grandparentKind);
+    }
+  }
+
   /**
    * A one-time poll message.
    *

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

@@ -59,7 +59,6 @@ public class DomainHistoryTest extends EntityTestCase {
   void testPersistence() {
     DomainBase domain = createDomainWithContactsAndHosts();
     DomainHistory domainHistory = createDomainHistory(domain);
-    domainHistory.id = null;
     jpaTm().transact(() -> jpaTm().insert(domainHistory));
 
     jpaTm()
@@ -77,7 +76,6 @@ public class DomainHistoryTest extends EntityTestCase {
     DomainBase domain = createDomainWithContactsAndHosts();
     DomainHistory domainHistory =
         createDomainHistory(domain).asBuilder().setDomainContent(null).build();
-    domainHistory.id = null;
     jpaTm().transact(() -> jpaTm().insert(domainHistory));
 
     jpaTm()

core/src/test/java/google/registry/model/poll/PollMessageTest.java

@@ -18,12 +18,15 @@ import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.testing.DatastoreHelper.createTld;
-import static google.registry.testing.DatastoreHelper.persistActiveDomain;
+import static google.registry.testing.DatastoreHelper.newDomainBase;
+import static google.registry.testing.DatastoreHelper.persistActiveContact;
 import static google.registry.testing.DatastoreHelper.persistResource;
 import static google.registry.testing.SqlHelper.saveRegistrar;
 import static java.nio.charset.StandardCharsets.UTF_8;
 
 import google.registry.model.EntityTestCase;
+import google.registry.model.contact.ContactResource;
+import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.Period;
 import google.registry.model.eppcommon.Trid;
 import google.registry.model.reporting.HistoryEntry;
@@ -34,6 +37,7 @@ import org.junit.jupiter.api.Test;
 /** Unit tests for {@link PollMessage}. */
 public class PollMessageTest extends EntityTestCase {
 
+  private DomainBase domain;
   private HistoryEntry historyEntry;
   private PollMessage.OneTime oneTime;
   private PollMessage.Autorenew autoRenew;
@@ -45,15 +49,17 @@ public class PollMessageTest extends EntityTestCase {
   @BeforeEach
   void setUp() {
     createTld("foobar");
+    ContactResource contact = persistActiveContact("contact1234");
+    domain = persistResource(newDomainBase("foo.foobar", contact));
     historyEntry =
         persistResource(
             new HistoryEntry.Builder()
-                .setParent(persistActiveDomain("foo.foobar"))
+                .setParent(domain)
                 .setType(HistoryEntry.Type.DOMAIN_CREATE)
                 .setPeriod(Period.create(1, Period.Unit.YEARS))
                 .setXmlBytes("<xml></xml>".getBytes(UTF_8))
                 .setModificationTime(fakeClock.nowUtc())
-                .setClientId("foo")
+                .setClientId("TheRegistrar")
                 .setTrid(Trid.create("ABC-123", "server-trid"))
                 .setBySuperuser(false)
                 .setReason("reason")
@@ -75,50 +81,46 @@ public class PollMessageTest extends EntityTestCase {
             .setAutorenewEndTime(fakeClock.nowUtc().plusDays(365))
             .setTargetId("foobar.foo")
             .build();
-    // TODO(shicong): Remove these two lines after we use symmetric vkey and change the cloud sql
+    jpaTm()
-    //  schema
+        .transact(
-    oneTime.id = null;
+            () -> {
-    autoRenew.id = null;
+              saveRegistrar("TheRegistrar");
+              jpaTm().insert(contact);
+              jpaTm().insert(domain);
+              jpaTm().insert(historyEntry.toChildHistoryEntity());
+            });
   }
 
   @Test
   void testCloudSqlPersistenceOneTime() {
-    saveRegistrar("TheRegistrar");
     jpaTm().transact(() -> jpaTm().insert(oneTime));
     PollMessage.OneTime persisted =
         jpaTm().transact(() -> jpaTm().load(VKey.createSql(PollMessage.OneTime.class, oneTime.id)));
-    persisted.parent = oneTime.parent;
     assertThat(persisted).isEqualTo(oneTime);
   }
 
   @Test
   void testCloudSqlPersistenceAutorenew() {
-    saveRegistrar("TheRegistrar");
     jpaTm().transact(() -> jpaTm().insert(autoRenew));
     PollMessage.Autorenew persisted =
         jpaTm()
             .transact(
                 () -> jpaTm().load(VKey.createSql(PollMessage.Autorenew.class, autoRenew.id)));
-    persisted.parent = autoRenew.parent;
     assertThat(persisted).isEqualTo(autoRenew);
   }
 
   @Test
   void testCloudSqlSupportForPolymorphicVKey() {
-    saveRegistrar("TheRegistrar");
-
     jpaTm().transact(() -> jpaTm().insert(oneTime));
     PollMessage persistedOneTime =
         jpaTm().transact(() -> jpaTm().load(VKey.createSql(PollMessage.class, oneTime.getId())));
     assertThat(persistedOneTime).isInstanceOf(PollMessage.OneTime.class);
-    persistedOneTime.parent = oneTime.parent;
     assertThat(persistedOneTime).isEqualTo(oneTime);
 
     jpaTm().transact(() -> jpaTm().insert(autoRenew));
     PollMessage persistedAutoRenew =
         jpaTm().transact(() -> jpaTm().load(VKey.createSql(PollMessage.class, autoRenew.getId())));
     assertThat(persistedAutoRenew).isInstanceOf(PollMessage.Autorenew.class);
-    persistedAutoRenew.parent = oneTime.parent;
     assertThat(persistedAutoRenew).isEqualTo(autoRenew);
   }
 

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

@@ -57,3 +57,4 @@ V56__rename_host_table.sql
 V57__history_null_content.sql
 V58__drop_default_value_and_sequences_for_billing_event.sql
 V59__use_composite_primary_key_for_contact_and_host_history_table.sql
+V60__remove_pollmessage_sequence.sql

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

@@ -0,0 +1,36 @@
+-- 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 "PollMessage" alter column poll_message_id drop default;
+
+drop sequence "PollMessage_poll_message_id_seq";
+
+alter table "PollMessage" rename column "domain_revision_id" to "domain_history_revision_id";
+alter table "PollMessage" rename column "contact_revision_id" to "contact_history_revision_id";
+alter table "PollMessage" rename column "host_revision_id" to "host_history_revision_id";
+
+alter table if exists "PollMessage"
+   add constraint fk_poll_message_domain_history
+   foreign key (domain_repo_id, domain_history_revision_id)
+   references "DomainHistory";
+
+alter table if exists "PollMessage"
+   add constraint fk_poll_message_contact_history
+   foreign key (contact_repo_id, contact_history_revision_id)
+   references "ContactHistory";
+
+alter table if exists "PollMessage"
+   add constraint fk_poll_message_host_history
+   foreign key (host_repo_id, host_history_revision_id)
+   references "HostHistory";

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

@@ -438,15 +438,15 @@ create sequence temp_history_id_sequence start 1 increment 50;
 
     create table "PollMessage" (
        type text not null,
-        poll_message_id  bigserial not null,
+        poll_message_id int8 not null,
         registrar_id text not null,
+        contact_history_revision_id int8,
         contact_repo_id text,
-        contact_revision_id int8,
+        domain_history_revision_id int8,
         domain_repo_id text,
-        domain_revision_id int8,
         event_time timestamptz not null,
+        host_history_revision_id int8,
         host_repo_id text,
-        host_revision_id int8,
         message text,
         transfer_response_contact_id text,
         transfer_response_domain_expiration_time timestamptz,

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

@@ -589,12 +589,12 @@ CREATE TABLE public."PollMessage" (
     poll_message_id bigint NOT NULL,
     registrar_id text NOT NULL,
     contact_repo_id text,
-    contact_revision_id bigint,
+    contact_history_revision_id bigint,
     domain_repo_id text,
-    domain_revision_id bigint,
+    domain_history_revision_id bigint,
     event_time timestamp with time zone NOT NULL,
     host_repo_id text,
-    host_revision_id bigint,
+    host_history_revision_id bigint,
     message text,
     transfer_response_contact_id text,
     transfer_response_domain_expiration_time timestamp with time zone,
@@ -614,25 +614,6 @@ CREATE TABLE public."PollMessage" (
 );
 
 
---
--- Name: PollMessage_poll_message_id_seq; Type: SEQUENCE; Schema: public; Owner: -
---
-
-CREATE SEQUENCE public."PollMessage_poll_message_id_seq"
-    START WITH 1
-    INCREMENT BY 1
-    NO MINVALUE
-    NO MAXVALUE
-    CACHE 1;
-
-
---
--- Name: PollMessage_poll_message_id_seq; Type: SEQUENCE OWNED BY; Schema: public; Owner: -
---
-
-ALTER SEQUENCE public."PollMessage_poll_message_id_seq" OWNED BY public."PollMessage".poll_message_id;
-
-
 --
 -- Name: PremiumEntry; Type: TABLE; Schema: public; Owner: -
 --
@@ -989,13 +970,6 @@ ALTER TABLE ONLY public."DomainTransactionRecord" ALTER COLUMN id SET DEFAULT ne
 ALTER TABLE ONLY public."GracePeriod" ALTER COLUMN id SET DEFAULT nextval('public."GracePeriod_id_seq"'::regclass);
 
 
---
--- Name: PollMessage poll_message_id; Type: DEFAULT; Schema: public; Owner: -
---
-
-ALTER TABLE ONLY public."PollMessage" ALTER COLUMN poll_message_id SET DEFAULT nextval('public."PollMessage_poll_message_id_seq"'::regclass);
-
-
 --
 -- Name: PremiumList revision_id; Type: DEFAULT; Schema: public; Owner: -
 --
@@ -1901,6 +1875,14 @@ ALTER TABLE ONLY public."HostHistory"
     ADD CONSTRAINT fk_hosthistory_host FOREIGN KEY (host_repo_id) REFERENCES public."Host"(repo_id);
 
 
+--
+-- Name: PollMessage fk_poll_message_contact_history; Type: FK CONSTRAINT; Schema: public; Owner: -
+--
+
+ALTER TABLE ONLY public."PollMessage"
+    ADD CONSTRAINT fk_poll_message_contact_history FOREIGN KEY (contact_repo_id, contact_history_revision_id) REFERENCES public."ContactHistory"(contact_repo_id, history_revision_id);
+
+
 --
 -- Name: PollMessage fk_poll_message_contact_repo_id; Type: FK CONSTRAINT; Schema: public; Owner: -
 --
@@ -1909,6 +1891,14 @@ ALTER TABLE ONLY public."PollMessage"
     ADD CONSTRAINT fk_poll_message_contact_repo_id FOREIGN KEY (contact_repo_id) REFERENCES public."Contact"(repo_id);
 
 
+--
+-- Name: PollMessage fk_poll_message_domain_history; Type: FK CONSTRAINT; Schema: public; Owner: -
+--
+
+ALTER TABLE ONLY public."PollMessage"
+    ADD CONSTRAINT fk_poll_message_domain_history FOREIGN KEY (domain_repo_id, domain_history_revision_id) REFERENCES public."DomainHistory"(domain_repo_id, history_revision_id);
+
+
 --
 -- Name: PollMessage fk_poll_message_domain_repo_id; Type: FK CONSTRAINT; Schema: public; Owner: -
 --
@@ -1917,6 +1907,14 @@ ALTER TABLE ONLY public."PollMessage"
     ADD CONSTRAINT fk_poll_message_domain_repo_id FOREIGN KEY (domain_repo_id) REFERENCES public."Domain"(repo_id);
 
 
+--
+-- Name: PollMessage fk_poll_message_host_history; Type: FK CONSTRAINT; Schema: public; Owner: -
+--
+
+ALTER TABLE ONLY public."PollMessage"
+    ADD CONSTRAINT fk_poll_message_host_history FOREIGN KEY (host_repo_id, host_history_revision_id) REFERENCES public."HostHistory"(host_repo_id, history_revision_id);
+
+
 --
 -- Name: PollMessage fk_poll_message_host_repo_id; Type: FK CONSTRAINT; Schema: public; Owner: -
 --