Use raw EntityManager to load during beforeSqlSave (#1253)

If we use the transaction manager methods, JpaTransactionManagerImpl
will attempt to detach the EppResource in question that we're loading --
this fails because that entity has been saved in the same transaction
already. We don't need detaching during these methods (it's just for
resource population) so we can use the raw loads to get around it.

core/src/main/java/google/registry/backup/ReplayCommitLogsToSqlAction.java

@@ -239,7 +239,7 @@ public class ReplayCommitLogsToSqlAction implements Runnable {
           .toSqlEntity()
           .ifPresent(
               sqlEntity -> {
-                ReplaySpecializer.beforeSqlSave(sqlEntity);
+                sqlEntity.beforeSqlSaveOnReplay();
                 jpaTm().put(sqlEntity);
               });
     } else {

core/src/main/java/google/registry/model/contact/ContactHistory.java

@@ -136,9 +136,11 @@ public class ContactHistory extends HistoryEntry implements SqlEntity {
   }
 
   // Used to fill out the contactBase field during asynchronous replay
-  public static void beforeSqlSave(ContactHistory contactHistory) {
+  @Override
-    contactHistory.contactBase =
+  public void beforeSqlSaveOnReplay() {
-        jpaTm().loadByKey(VKey.createSql(ContactResource.class, contactHistory.getContactRepoId()));
+    if (contactBase == null) {
+      contactBase = jpaTm().getEntityManager().find(ContactResource.class, getContactRepoId());
+    }
   }
 
   /** Class to represent the composite primary key of {@link ContactHistory} entity. */

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

@@ -294,9 +294,26 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
   }
 
   // Used to fill out the domainContent field during asynchronous replay
-  public static void beforeSqlSave(DomainHistory domainHistory) {
+  @Override
-    domainHistory.domainContent =
+  public void beforeSqlSaveOnReplay() {
-        jpaTm().loadByKey(VKey.createSql(DomainBase.class, domainHistory.getDomainRepoId()));
+    if (domainContent == null) {
+      domainContent = jpaTm().getEntityManager().find(DomainBase.class, getDomainRepoId());
+      fillAuxiliaryFieldsFromDomain(this);
+    }
+  }
+
+  private static void fillAuxiliaryFieldsFromDomain(DomainHistory domainHistory) {
+    if (domainHistory.domainContent != null) {
+      domainHistory.nsHosts = nullToEmptyImmutableCopy(domainHistory.domainContent.nsHosts);
+      domainHistory.dsDataHistories =
+          nullToEmptyImmutableCopy(domainHistory.domainContent.getDsData()).stream()
+              .map(dsData -> DomainDsDataHistory.createFrom(domainHistory.id, dsData))
+              .collect(toImmutableSet());
+      domainHistory.gracePeriodHistories =
+          nullToEmptyImmutableCopy(domainHistory.domainContent.getGracePeriods()).stream()
+              .map(gracePeriod -> GracePeriodHistory.createFrom(domainHistory.id, gracePeriod))
+              .collect(toImmutableSet());
+    }
   }
 
   /** Class to represent the composite primary key of {@link DomainHistory} entity. */
@@ -391,17 +408,7 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
       // Note that we cannot assert that instance.domainContent is not null here because this
       // builder is also used to convert legacy HistoryEntry objects to DomainHistory, when
       // domainContent is not available.
-      if (instance.domainContent != null) {
+      fillAuxiliaryFieldsFromDomain(instance);
-        instance.nsHosts = nullToEmptyImmutableCopy(instance.domainContent.nsHosts);
-        instance.dsDataHistories =
-            nullToEmptyImmutableCopy(instance.domainContent.getDsData()).stream()
-                .map(dsData -> DomainDsDataHistory.createFrom(instance.id, dsData))
-                .collect(toImmutableSet());
-        instance.gracePeriodHistories =
-            nullToEmptyImmutableCopy(instance.domainContent.getGracePeriods()).stream()
-                .map(gracePeriod -> GracePeriodHistory.createFrom(instance.id, gracePeriod))
-                .collect(toImmutableSet());
-      }
       return instance;
     }
   }

core/src/main/java/google/registry/model/host/HostHistory.java

@@ -136,9 +136,11 @@ public class HostHistory extends HistoryEntry implements SqlEntity {
   }
 
   // Used to fill out the hostBase field during asynchronous replay
-  public static void beforeSqlSave(HostHistory hostHistory) {
+  @Override
-    hostHistory.hostBase =
+  public void beforeSqlSaveOnReplay() {
-        jpaTm().loadByKey(VKey.createSql(HostResource.class, hostHistory.getHostRepoId()));
+    if (hostBase == null) {
+      hostBase = jpaTm().getEntityManager().find(HostResource.class, getHostRepoId());
+    }
   }
 
   /** Class to represent the composite primary key of {@link HostHistory} entity. */

core/src/main/java/google/registry/schema/replay/ReplaySpecializer.java

@@ -28,17 +28,11 @@ import java.lang.reflect.Method;
 public class ReplaySpecializer {
 
   public static void beforeSqlDelete(VKey<?> key) {
-    invokeMethod(key.getKind(), "beforeSqlDelete", key);
+    String methodName = "beforeSqlDelete";
-  }
+    Class<?> clazz = key.getKind();
-
-  public static void beforeSqlSave(SqlEntity sqlEntity) {
-    invokeMethod(sqlEntity.getClass(), "beforeSqlSave", sqlEntity);
-  }
-
-  private static <T> void invokeMethod(Class<T> clazz, String methodName, Object argument) {
     try {
-      Method method = clazz.getMethod(methodName, argument.getClass());
+      Method method = clazz.getMethod(methodName, VKey.class);
-      method.invoke(null, argument);
+      method.invoke(null, key);
     } catch (NoSuchMethodException e) {
       // Ignore, this just means that the class doesn't need this hook.
     } catch (IllegalAccessException e) {

core/src/main/java/google/registry/schema/replay/SqlEntity.java

@@ -26,4 +26,7 @@ import java.util.Optional;
 public interface SqlEntity {
 
   Optional<DatastoreEntity> toDatastoreEntity();
+
+  /** A method that will ber called before the object is saved to SQL in asynchronous replay. */
+  default void beforeSqlSaveOnReplay() {}
 }

core/src/main/java/google/registry/tools/javascrap/CreateSyntheticHistoryEntriesAction.java

@@ -53,7 +53,7 @@ import javax.inject.Inject;
  * However, since this is meant to be run during the Datastore-primary, SQL-secondary stage of the
  * migration, we want to make sure that we are using the most up-to-date version of the data. The
  * resource field of the history objects will be populated during asynchronous migration, e.g. in
- * {@link DomainHistory#beforeSqlSave(DomainHistory)}.
+ * {@link DomainHistory#beforeSqlSaveOnReplay}.
  */
 @Action(
     service = Action.Service.BACKEND,

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

@@ -18,6 +18,7 @@ import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.ImmutableObjectSubject.assertAboutImmutableObjects;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.testing.DatabaseHelper.newContactResource;
 import static google.registry.testing.DatabaseHelper.newContactResourceWithRoid;
 import static google.registry.testing.SqlHelper.saveRegistrar;
 import static java.nio.charset.StandardCharsets.UTF_8;
@@ -47,7 +48,7 @@ public class ContactHistoryTest extends EntityTestCase {
     jpaTm().transact(() -> jpaTm().insert(contact));
     VKey<ContactResource> contactVKey = contact.createVKey();
     ContactResource contactFromDb = jpaTm().transact(() -> jpaTm().loadByKey(contactVKey));
-    ContactHistory contactHistory = createContactHistory(contactFromDb, contact.getRepoId());
+    ContactHistory contactHistory = createContactHistory(contactFromDb);
     jpaTm().transact(() -> jpaTm().insert(contactHistory));
     jpaTm()
         .transact(
@@ -67,10 +68,7 @@ public class ContactHistoryTest extends EntityTestCase {
     VKey<ContactResource> contactVKey = contact.createVKey();
     ContactResource contactFromDb = jpaTm().transact(() -> jpaTm().loadByKey(contactVKey));
     ContactHistory contactHistory =
-        createContactHistory(contactFromDb, contact.getRepoId())
+        createContactHistory(contactFromDb).asBuilder().setContact(null).build();
-            .asBuilder()
-            .setContact(null)
-            .build();
     jpaTm().transact(() -> jpaTm().insert(contactHistory));
 
     jpaTm()
@@ -91,7 +89,7 @@ public class ContactHistoryTest extends EntityTestCase {
     VKey<ContactResource> contactVKey = contact.createVKey();
     ContactResource contactFromDb = tm().transact(() -> tm().loadByKey(contactVKey));
     fakeClock.advanceOneMilli();
-    ContactHistory contactHistory = createContactHistory(contactFromDb, contact.getRepoId());
+    ContactHistory contactHistory = createContactHistory(contactFromDb);
     tm().transact(() -> tm().insert(contactHistory));
 
     // retrieving a HistoryEntry or a ContactHistory with the same key should return the same object
@@ -106,7 +104,38 @@ public class ContactHistoryTest extends EntityTestCase {
     assertThat(hostHistoryFromDb).isEqualTo(historyEntryFromDb);
   }
 
-  private ContactHistory createContactHistory(ContactBase contact, String contactRepoId) {
+  @Test
+  void testBeforeSqlSave_afterContactPersisted() {
+    saveRegistrar("TheRegistrar");
+    ContactResource contactResource = newContactResource("contactId");
+    ContactHistory contactHistory =
+        new ContactHistory.Builder()
+            .setType(HistoryEntry.Type.HOST_CREATE)
+            .setXmlBytes("<xml></xml>".getBytes(UTF_8))
+            .setModificationTime(fakeClock.nowUtc())
+            .setClientId("TheRegistrar")
+            .setTrid(Trid.create("ABC-123", "server-trid"))
+            .setBySuperuser(false)
+            .setReason("reason")
+            .setRequestedByRegistrar(true)
+            .setContactRepoId(contactResource.getRepoId())
+            .build();
+    jpaTm()
+        .transact(
+            () -> {
+              jpaTm().put(contactResource);
+              contactHistory.beforeSqlSaveOnReplay();
+              jpaTm().put(contactHistory);
+            });
+    jpaTm()
+        .transact(
+            () ->
+                assertAboutImmutableObjects()
+                    .that(jpaTm().loadByEntity(contactResource))
+                    .hasFieldsEqualTo(jpaTm().loadByEntity(contactHistory).getContactBase().get()));
+  }
+
+  private ContactHistory createContactHistory(ContactBase contact) {
     return new ContactHistory.Builder()
         .setType(HistoryEntry.Type.HOST_CREATE)
         .setXmlBytes("<xml></xml>".getBytes(UTF_8))
@@ -117,7 +146,7 @@ public class ContactHistoryTest extends EntityTestCase {
         .setReason("reason")
         .setRequestedByRegistrar(true)
         .setContact(contact)
-        .setContactRepoId(contactRepoId)
+        .setContactRepoId(contact.getRepoId())
         .build();
   }
 

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

@@ -190,6 +190,44 @@ public class DomainHistoryTest extends EntityTestCase {
     jpaTm().transact(() -> jpaTm().put(domainHistoryFromDb2));
   }
 
+  @TestSqlOnly
+  void testBeforeSqlSave_afterDomainPersisted() {
+    DomainBase domain = createDomainWithContactsAndHosts();
+    DomainHistory historyWithoutResource =
+        new DomainHistory.Builder()
+            .setType(HistoryEntry.Type.DOMAIN_CREATE)
+            .setXmlBytes("<xml></xml>".getBytes(UTF_8))
+            .setModificationTime(fakeClock.nowUtc())
+            .setClientId("TheRegistrar")
+            .setTrid(Trid.create("ABC-123", "server-trid"))
+            .setBySuperuser(false)
+            .setReason("reason")
+            .setRequestedByRegistrar(true)
+            .setDomainRepoId(domain.getRepoId())
+            .setOtherClientId("otherClient")
+            .setPeriod(Period.create(1, Period.Unit.YEARS))
+            .build();
+    jpaTm()
+        .transact(
+            () -> {
+              jpaTm()
+                  .put(
+                      domain
+                          .asBuilder()
+                          .setPersistedCurrentSponsorClientId("NewRegistrar")
+                          .build());
+              historyWithoutResource.beforeSqlSaveOnReplay();
+              jpaTm().put(historyWithoutResource);
+            });
+    jpaTm()
+        .transact(
+            () ->
+                assertAboutImmutableObjects()
+                    .that(jpaTm().loadByEntity(domain))
+                    .hasFieldsEqualTo(
+                        jpaTm().loadByEntity(historyWithoutResource).getDomainContent().get()));
+  }
+
   static DomainBase createDomainWithContactsAndHosts() {
     createTld("tld");
     HostResource host = newHostResourceWithRoid("ns1.example.com", "host1");

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

@@ -18,6 +18,7 @@ import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.ImmutableObjectSubject.assertAboutImmutableObjects;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.testing.DatabaseHelper.newHostResource;
 import static google.registry.testing.DatabaseHelper.newHostResourceWithRoid;
 import static google.registry.testing.SqlHelper.saveRegistrar;
 import static java.nio.charset.StandardCharsets.UTF_8;
@@ -48,7 +49,7 @@ public class HostHistoryTest extends EntityTestCase {
     VKey<HostResource> hostVKey =
         VKey.create(HostResource.class, "host1", Key.create(HostResource.class, "host1"));
     HostResource hostFromDb = jpaTm().transact(() -> jpaTm().loadByKey(hostVKey));
-    HostHistory hostHistory = createHostHistory(hostFromDb, host.getRepoId());
+    HostHistory hostHistory = createHostHistory(hostFromDb);
     jpaTm().transact(() -> jpaTm().insert(hostHistory));
     jpaTm()
         .transact(
@@ -66,8 +67,7 @@ public class HostHistoryTest extends EntityTestCase {
     jpaTm().transact(() -> jpaTm().insert(host));
 
     HostResource hostFromDb = jpaTm().transact(() -> jpaTm().loadByKey(host.createVKey()));
-    HostHistory hostHistory =
+    HostHistory hostHistory = createHostHistory(hostFromDb).asBuilder().setHost(null).build();
-        createHostHistory(hostFromDb, host.getRepoId()).asBuilder().setHost(null).build();
     jpaTm().transact(() -> jpaTm().insert(hostHistory));
 
     jpaTm()
@@ -88,7 +88,7 @@ public class HostHistoryTest extends EntityTestCase {
     VKey<HostResource> hostVKey =
         VKey.create(HostResource.class, "host1", Key.create(HostResource.class, "host1"));
     HostResource hostFromDb = tm().transact(() -> tm().loadByKey(hostVKey));
-    HostHistory hostHistory = createHostHistory(hostFromDb, host.getRepoId());
+    HostHistory hostHistory = createHostHistory(hostFromDb);
     fakeClock.advanceOneMilli();
     tm().transact(() -> tm().insert(hostHistory));
 
@@ -104,6 +104,37 @@ public class HostHistoryTest extends EntityTestCase {
     assertThat(hostHistoryFromDb).isEqualTo(historyEntryFromDb);
   }
 
+  @Test
+  void testBeforeSqlSave_afterHostPersisted() {
+    saveRegistrar("TheRegistrar");
+    HostResource hostResource = newHostResource("ns1.example.tld");
+    HostHistory hostHistory =
+        new HostHistory.Builder()
+            .setType(HistoryEntry.Type.HOST_CREATE)
+            .setXmlBytes("<xml></xml>".getBytes(UTF_8))
+            .setModificationTime(fakeClock.nowUtc())
+            .setClientId("TheRegistrar")
+            .setTrid(Trid.create("ABC-123", "server-trid"))
+            .setBySuperuser(false)
+            .setReason("reason")
+            .setRequestedByRegistrar(true)
+            .setHostRepoId(hostResource.getRepoId())
+            .build();
+    jpaTm()
+        .transact(
+            () -> {
+              jpaTm().put(hostResource);
+              hostHistory.beforeSqlSaveOnReplay();
+              jpaTm().put(hostHistory);
+            });
+    jpaTm()
+        .transact(
+            () ->
+                assertAboutImmutableObjects()
+                    .that(jpaTm().loadByEntity(hostResource))
+                    .hasFieldsEqualTo(jpaTm().loadByEntity(hostHistory).getHostBase().get()));
+  }
+
   private void assertHostHistoriesEqual(HostHistory one, HostHistory two) {
     assertAboutImmutableObjects().that(one).isEqualExceptFields(two, "hostBase");
     assertAboutImmutableObjects()
@@ -111,7 +142,7 @@ public class HostHistoryTest extends EntityTestCase {
         .isEqualExceptFields(two.getHostBase().orElse(null), "repoId");
   }
 
-  private HostHistory createHostHistory(HostBase hostBase, String hostRepoId) {
+  private HostHistory createHostHistory(HostBase hostBase) {
     return new HostHistory.Builder()
         .setType(HistoryEntry.Type.HOST_CREATE)
         .setXmlBytes("<xml></xml>".getBytes(UTF_8))
@@ -122,7 +153,7 @@ public class HostHistoryTest extends EntityTestCase {
         .setReason("reason")
         .setRequestedByRegistrar(true)
         .setHost(hostBase)
-        .setHostRepoId(hostRepoId)
+        .setHostRepoId(hostBase.getRepoId())
         .build();
   }
 }

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

@@ -74,7 +74,7 @@ public class LegacyHistoryObjectTest extends EntityTestCase {
     assertAboutImmutableObjects()
         .that(legacyHistoryEntry)
         .isEqualExceptFields(fromObjectify, "contactBase", "contactRepoId");
-    assertThat(fromObjectify instanceof ContactHistory).isTrue();
+    assertThat(fromObjectify).isInstanceOf(ContactHistory.class);
     ContactHistory legacyContactHistory = (ContactHistory) fromObjectify;
 
     // Next, save that from-Datastore object in SQL and verify we can load it back in
@@ -124,7 +124,7 @@ public class LegacyHistoryObjectTest extends EntityTestCase {
             "nsHosts",
             "dsDataHistories",
             "gracePeriodHistories");
-    assertThat(fromObjectify instanceof DomainHistory).isTrue();
+    assertThat(fromObjectify).isInstanceOf(DomainHistory.class);
     DomainHistory legacyDomainHistory = (DomainHistory) fromObjectify;
 
     // Next, save that from-Datastore object in SQL and verify we can load it back in
@@ -137,15 +137,7 @@ public class LegacyHistoryObjectTest extends EntityTestCase {
               // Don't compare nsHosts directly because one is null and the other is empty
               assertAboutImmutableObjects()
                   .that(legacyDomainHistory)
-                  .isEqualExceptFields(
+                  .isEqualExceptFields(legacyHistoryFromSql, "nsHosts");
-                      // NB: period, transaction records, and other client ID are added in #794
-                      legacyHistoryFromSql,
-                      "period",
-                      "domainTransactionRecords",
-                      "otherClientId",
-                      "nsHosts",
-                      "dsDataHistories",
-                      "gracePeriodHistories");
               assertThat(nullToEmpty(legacyDomainHistory.getNsHosts()))
                   .isEqualTo(nullToEmpty(legacyHistoryFromSql.getNsHosts()));
             });
@@ -174,7 +166,7 @@ public class LegacyHistoryObjectTest extends EntityTestCase {
     assertAboutImmutableObjects()
         .that(legacyHistoryEntry)
         .isEqualExceptFields(fromObjectify, "hostBase", "hostRepoId");
-    assertThat(fromObjectify instanceof HostHistory).isTrue();
+    assertThat(fromObjectify).isInstanceOf(HostHistory.class);
     HostHistory legacyHostHistory = (HostHistory) fromObjectify;
 
     // Next, save that from-Datastore object in SQL and verify we can load it back in

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

@@ -75,7 +75,8 @@ public class TestObject extends ImmutableObject implements DatastoreAndSqlEntity
     beforeSqlDeleteCallCount++;
   }
 
-  public static void beforeSqlSave(TestObject testObject) {
+  @Override
+  public void beforeSqlSaveOnReplay() {
     beforeSqlSaveCallCount++;
   }