Always detach entities during load (#1116)

* Always detach entities during load

The mutations on non-transient fields that we do in some of the PostLoad
methods have been causing the objects to be marked as "dirty", and hibernate
has been quietly persisting them during transaction commit.

By detaching the entities on load, we avoid any possibility of this, which
works in our case because we treat all of our model objects as immutable
during normal use.

There is another mixed blessing to this: lazy loading won't work on these
objects once they are detached from a session, meaning that all fields must be
lazy loaded up front.  This is unfortunate in that we don't always need those
lazy-loaded fields and there is a performance cost to loading them, but it is
also useful in that objects will now be complete when used outseide of the
transaction that loaded them (prior to this, an attempt to access a
lazy-loaded field after its transaction closed would have caused an error at
runtime).

* Changes requested in review

* A few improvements to test logic

* Deal with premature detachment of mutated objects

* Add unit tests, use a more specific exception

* Changes for review

- Deal with DomainDeleteFlow, which appears to be the only case in the
  codebase where we're doing a load-after-save.
- Display the object that is being loaded after save in the exception message.
- Add a TODO for figuring out why Eager loads aren't working as expected.

* Move the recurring billing event into a parameter

* Changes for review and rebase error fix

* Remove initialization of list entries

Remove initialization of list entries that we want to be lazy loaded (premium,
reserved, and claims lists).

* Post-rebase cleanups

core/src/main/java/google/registry/flows/domain/DomainDeleteFlow.java

@@ -251,7 +251,9 @@ public final class DomainDeleteFlow implements TransactionalFlow {
         buildDomainHistory(newDomain, registry, now, durationUntilDelete, inAddGracePeriod);
     updateForeignKeyIndexDeletionTime(newDomain);
     handlePendingTransferOnDelete(existingDomain, newDomain, now, domainHistory);
-    // Close the autorenew billing event and poll message. This may delete the poll message.
+    // Close the autorenew billing event and poll message. This may delete the poll message.  Store
+    // the updated recurring billing event, we'll need it later and can't reload it.
+    BillingEvent.Recurring recurringBillingEvent =
         updateAutorenewRecurrenceEndTime(existingDomain, now);
     // If there's a pending transfer, the gaining client's autorenew billing
     // event and poll message will already have been deleted in
@@ -275,7 +277,8 @@ public final class DomainDeleteFlow implements TransactionalFlow {
                     newDomain.getDeletionTime().isAfter(now)
                         ? SUCCESS_WITH_ACTION_PENDING
                         : SUCCESS)
-                .setResponseExtensions(getResponseExtensions(existingDomain, now))
+                .setResponseExtensions(
+                    getResponseExtensions(recurringBillingEvent, existingDomain, now))
                 .build());
     persistEntityChanges(entityChanges);
     return responseBuilder
@@ -377,7 +380,7 @@ public final class DomainDeleteFlow implements TransactionalFlow {
 
   @Nullable
   private ImmutableList<FeeTransformResponseExtension> getResponseExtensions(
-      DomainBase existingDomain, DateTime now) {
+      BillingEvent.Recurring recurringBillingEvent, DomainBase existingDomain, DateTime now) {
     FeeTransformResponseExtension.Builder feeResponseBuilder = getDeleteResponseBuilder();
     if (feeResponseBuilder == null) {
       return ImmutableList.of();
@@ -385,7 +388,7 @@ public final class DomainDeleteFlow implements TransactionalFlow {
     ImmutableList.Builder<Credit> creditsBuilder = new ImmutableList.Builder<>();
     for (GracePeriod gracePeriod : existingDomain.getGracePeriods()) {
       if (gracePeriod.hasBillingEvent()) {
-        Money cost = getGracePeriodCost(gracePeriod, now);
+        Money cost = getGracePeriodCost(recurringBillingEvent, gracePeriod, now);
         creditsBuilder.add(Credit.create(
             cost.negated().getAmount(), FeeType.CREDIT, gracePeriod.getType().getXmlName()));
         feeResponseBuilder.setCurrency(checkNotNull(cost.getCurrencyUnit()));
@@ -398,12 +401,12 @@ public final class DomainDeleteFlow implements TransactionalFlow {
     return ImmutableList.of(feeResponseBuilder.setCredits(credits).build());
   }
 
-  private Money getGracePeriodCost(GracePeriod gracePeriod, DateTime now) {
+  private Money getGracePeriodCost(
+      BillingEvent.Recurring recurringBillingEvent, GracePeriod gracePeriod, DateTime now) {
     if (gracePeriod.getType() == GracePeriodStatus.AUTO_RENEW) {
+      // If we updated the autorenew billing event, reuse it.
       DateTime autoRenewTime =
-          tm().loadByKey(checkNotNull(gracePeriod.getRecurringBillingEvent()))
+          recurringBillingEvent.getRecurrenceTimeOfYear().getLastInstanceBeforeOrAt(now);
-              .getRecurrenceTimeOfYear()
-              .getLastInstanceBeforeOrAt(now);
       return getDomainRenewCost(targetId, autoRenewTime, 1);
     }
     return tm().loadByKey(checkNotNull(gracePeriod.getOneTimeBillingEvent())).getCost();

core/src/main/java/google/registry/flows/domain/DomainFlowUtils.java

@@ -517,8 +517,10 @@ public class DomainFlowUtils {
    * <p>This may end up deleting the poll message (if closing the message interval) or recreating it
    * (if opening the message interval). This may cause an autorenew billing event to have an end
    * time earlier than its event time (i.e. if it's being ended before it was ever triggered).
+   *
+   * <p>Returns the new autorenew recurring billing event.
    */
-  public static void updateAutorenewRecurrenceEndTime(DomainBase domain, DateTime newEndTime) {
+  public static Recurring updateAutorenewRecurrenceEndTime(DomainBase domain, DateTime newEndTime) {
     Optional<PollMessage.Autorenew> autorenewPollMessage =
         tm().loadByKeyIfPresent(domain.getAutorenewPollMessage());
 
@@ -545,8 +547,13 @@ public class DomainFlowUtils {
       tm().put(updatedAutorenewPollMessage);
     }
 
-    Recurring recurring = tm().loadByKey(domain.getAutorenewBillingEvent());
+    Recurring recurring =
-    tm().put(recurring.asBuilder().setRecurrenceEndTime(newEndTime).build());
+        tm().loadByKey(domain.getAutorenewBillingEvent())
+            .asBuilder()
+            .setRecurrenceEndTime(newEndTime)
+            .build();
+    tm().put(recurring);
+    return recurring;
   }
 
   /**

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

@@ -36,7 +36,9 @@ import javax.persistence.Index;
 import javax.persistence.JoinColumn;
 import javax.persistence.JoinTable;
 import javax.persistence.OneToMany;
+import javax.persistence.PostLoad;
 import javax.persistence.Table;
+import org.hibernate.Hibernate;
 import org.joda.time.DateTime;
 
 /**
@@ -78,6 +80,8 @@ public class DomainBase extends DomainContent
     return super.getRepoId();
   }
 
+  // It seems like this should be FetchType.EAGER, but for some reason when we do that we get a lazy
+  // load error during the load of a domain.
   @ElementCollection
   @JoinTable(
       name = "DomainHost",
@@ -139,6 +143,16 @@ public class DomainBase extends DomainContent
     return dsData;
   }
 
+  /** Post-load method to eager load the collections. */
+  @PostLoad
+  @Override
+  protected void postLoad() {
+    super.postLoad();
+    // TODO(b/188044616): Determine why Eager loading doesn't work here.
+    Hibernate.initialize(dsData);
+    Hibernate.initialize(gracePeriods);
+  }
+
   @Override
   public VKey<DomainBase> createVKey() {
     return VKey.create(DomainBase.class, getRepoId(), Key.create(this));

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

@@ -338,8 +338,7 @@ public class DomainContent extends EppResource
   }
 
   @PostLoad
-  @SuppressWarnings("UnusedMethod")
+  protected void postLoad() {
-  private void postLoad() {
     // Reconstitute the contact list.
     ImmutableSet.Builder<DesignatedContact> contactsBuilder = new ImmutableSet.Builder<>();
 

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

@@ -54,6 +54,7 @@ import javax.persistence.JoinTable;
 import javax.persistence.OneToMany;
 import javax.persistence.PostLoad;
 import javax.persistence.Table;
+import org.hibernate.Hibernate;
 
 /**
  * A persisted history entry representing an EPP modification to a domain.
@@ -261,6 +262,12 @@ public class DomainHistory extends HistoryEntry implements SqlEntity {
         }
       }
     }
+
+    // TODO(b/188044616): Determine why Eager loading doesn't work here.
+    Hibernate.initialize(domainTransactionRecords);
+    Hibernate.initialize(nsHosts);
+    Hibernate.initialize(dsDataHistories);
+    Hibernate.initialize(gracePeriodHistories);
   }
 
   // In Datastore, save as a HistoryEntry object regardless of this object's type

core/src/main/java/google/registry/model/ofy/DatastoreTransactionManager.java

@@ -297,8 +297,9 @@ public class DatastoreTransactionManager implements TransactionManager {
   }
 
   @Override
-  public void delete(Object entity) {
+  public <T> T delete(T entity) {
     syncIfTransactionless(getOfy().delete().entity(toDatastoreEntity(entity)));
+    return entity;
   }
 
   @Override

core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java

@@ -44,13 +44,17 @@ import google.registry.util.Clock;
 import google.registry.util.Retrier;
 import google.registry.util.SystemSleeper;
 import java.lang.reflect.Field;
+import java.util.Collections;
+import java.util.IdentityHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.NoSuchElementException;
 import java.util.Optional;
+import java.util.Set;
 import java.util.function.Supplier;
 import java.util.stream.Stream;
 import java.util.stream.StreamSupport;
+import javax.annotation.Nullable;
 import javax.persistence.EntityManager;
 import javax.persistence.EntityManagerFactory;
 import javax.persistence.EntityTransaction;
@@ -269,8 +273,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     assertInTransaction();
     // Necessary due to the changes in HistoryEntry representation during the migration to SQL
     Object toPersist = toSqlEntity(entity);
-    getEntityManager().persist(toPersist);
+    transactionInfo.get().insertObject(toPersist);
-    transactionInfo.get().addUpdate(toPersist);
   }
 
   @Override
@@ -299,8 +302,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     assertInTransaction();
     // Necessary due to the changes in HistoryEntry representation during the migration to SQL
     Object toPersist = toSqlEntity(entity);
-    getEntityManager().merge(toPersist);
+    transactionInfo.get().updateObject(toPersist);
-    transactionInfo.get().addUpdate(toPersist);
   }
 
   @Override
@@ -339,8 +341,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     checkArgument(exists(entity), "Given entity does not exist");
     // Necessary due to the changes in HistoryEntry representation during the migration to SQL
     Object toPersist = toSqlEntity(entity);
-    getEntityManager().merge(toPersist);
+    transactionInfo.get().updateObject(toPersist);
-    transactionInfo.get().addUpdate(toPersist);
   }
 
   @Override
@@ -397,7 +398,8 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   public <T> Optional<T> loadByKeyIfPresent(VKey<T> key) {
     checkArgumentNotNull(key, "key must be specified");
     assertInTransaction();
-    return Optional.ofNullable(getEntityManager().find(key.getKind(), key.getSqlKey()));
+    return Optional.ofNullable(getEntityManager().find(key.getKind(), key.getSqlKey()))
+        .map(this::detach);
   }
 
   @Override
@@ -411,7 +413,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
         .map(
             key ->
                 new SimpleEntry<VKey<? extends T>, T>(
-                    key, getEntityManager().find(key.getKind(), key.getSqlKey())))
+                    key, detach(getEntityManager().find(key.getKind(), key.getSqlKey()))))
         .filter(entry -> entry.getValue() != null)
         .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue));
   }
@@ -433,7 +435,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     if (result == null) {
       throw new NoSuchElementException(key.toString());
     }
-    return result;
+    return detach(result);
   }
 
   @Override
@@ -472,12 +474,11 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   public <T> ImmutableList<T> loadAllOf(Class<T> clazz) {
     checkArgumentNotNull(clazz, "clazz must be specified");
     assertInTransaction();
-    return ImmutableList.copyOf(
+    return getEntityManager()
-        getEntityManager()
+        .createQuery(String.format("FROM %s", getEntityType(clazz).getName()), clazz)
-            .createQuery(
+        .getResultStream()
-                String.format("SELECT entity FROM %s entity", getEntityType(clazz).getName()),
+        .map(this::detach)
-                clazz)
+        .collect(toImmutableList());
-            .getResultList());
   }
 
   @Override
@@ -524,18 +525,21 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   }
 
   @Override
-  public void delete(Object entity) {
+  public <T> T delete(T entity) {
     checkArgumentNotNull(entity, "entity must be specified");
     if (isEntityOfIgnoredClass(entity)) {
-      return;
+      return entity;
     }
     assertInTransaction();
     entity = toSqlEntity(entity);
-    Object managedEntity = entity;
+    T managedEntity = entity;
     if (!getEntityManager().contains(entity)) {
+      // We don't add the entity to "objectsToSave": once deleted, the object should never be
+      // returned as a result of the query or lookup.
       managedEntity = getEntityManager().merge(entity);
     }
     getEntityManager().remove(managedEntity);
+    return managedEntity;
   }
 
   @Override
@@ -555,7 +559,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
 
   @Override
   public <T> QueryComposer<T> createQueryComposer(Class<T> entity) {
-    return new JpaQueryComposerImpl<T>(entity, getEntityManager());
+    return new JpaQueryComposerImpl<T>(entity);
   }
 
   @Override
@@ -658,6 +662,21 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     }
   }
 
+  /** Detach the entity, suitable for use in Optional.map(). */
+  @Nullable
+  private <T> T detach(@Nullable T entity) {
+    if (entity != null) {
+
+      // If the entity was previously persisted or merged, we have to throw an exception.
+      if (transactionInfo.get().willSave(entity)) {
+        throw new IllegalStateException("Inserted/updated object reloaded: " + entity);
+      }
+
+      getEntityManager().detach(entity);
+    }
+    return entity;
+  }
+
   private static class TransactionInfo {
     EntityManager entityManager;
     boolean inTransaction = false;
@@ -666,6 +685,12 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     // Serializable representation of the transaction to be persisted in the Transaction table.
     Transaction.Builder contentsBuilder;
 
+    // The set of entity objects that have been either persisted (via insert()) or merged (via
+    // put()/update()).  If the entity manager returns these as a result of a find() or query
+    // operation, we can not detach them -- detaching removes them from the transaction and causes
+    // them to not be saved to the database -- so we throw an exception instead.
+    Set<Object> objectsToSave = Collections.newSetFromMap(new IdentityHashMap<Object, Boolean>());
+
     /** Start a new transaction. */
     private void start(Clock clock) {
       checkArgumentNotNull(clock);
@@ -680,6 +705,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
       inTransaction = false;
       transactionTime = null;
       contentsBuilder = null;
+      objectsToSave = Collections.newSetFromMap(new IdentityHashMap<Object, Boolean>());
       if (entityManager != null) {
         // Close this EntityManager just let the connection pool be able to reuse it, it doesn't
         // close the underlying database connection.
@@ -708,25 +734,40 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
         }
       }
     }
+
+    /** Does the full "update" on an object including all internal housekeeping. */
+    private void updateObject(Object object) {
+      Object merged = entityManager.merge(object);
+      objectsToSave.add(merged);
+      addUpdate(object);
     }
 
-  private static class JpaQueryComposerImpl<T> extends QueryComposer<T> {
+    /** Does the full "insert" on a new object including all internal housekeeping. */
+    private void insertObject(Object object) {
+      entityManager.persist(object);
+      objectsToSave.add(object);
+      addUpdate(object);
+    }
+
+    /** Returns true if the object has been persisted/merged and will be saved on commit. */
+    private boolean willSave(Object object) {
+      return objectsToSave.contains(object);
+    }
+  }
+
+  private class JpaQueryComposerImpl<T> extends QueryComposer<T> {
 
     private static final int DEFAULT_FETCH_SIZE = 1000;
 
-    EntityManager em;
-
     private int fetchSize = DEFAULT_FETCH_SIZE;
 
-    private boolean autoDetachOnLoad = true;
+    JpaQueryComposerImpl(Class<T> entityClass) {
-
-    JpaQueryComposerImpl(Class<T> entityClass, EntityManager em) {
       super(entityClass);
-      this.em = em;
     }
 
     private TypedQuery<T> buildQuery() {
-      CriteriaQueryBuilder<T> queryBuilder = CriteriaQueryBuilder.create(em, entityClass);
+      CriteriaQueryBuilder<T> queryBuilder =
+          CriteriaQueryBuilder.create(getEntityManager(), entityClass);
       return addCriteria(queryBuilder);
     }
 
@@ -739,13 +780,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
         queryBuilder.orderByAsc(orderBy);
       }
 
-      return em.createQuery(queryBuilder.build());
+      return getEntityManager().createQuery(queryBuilder.build());
-    }
-
-    @Override
-    public QueryComposer<T> withAutoDetachOnLoad(boolean autoDetachOnLoad) {
-      this.autoDetachOnLoad = autoDetachOnLoad;
-      return this;
     }
 
     @Override
@@ -758,12 +793,12 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     @Override
     public Optional<T> first() {
       List<T> results = buildQuery().setMaxResults(1).getResultList();
-      return results.size() > 0 ? Optional.of(maybeDetachEntity(results.get(0))) : Optional.empty();
+      return results.size() > 0 ? Optional.of(detach(results.get(0))) : Optional.empty();
     }
 
     @Override
     public T getSingleResult() {
-      return maybeDetachEntity(buildQuery().getSingleResult());
+      return detach(buildQuery().getSingleResult());
     }
 
     @Override
@@ -777,27 +812,21 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
       } else {
         logger.atWarning().log("Query implemention does not support result streaming.");
       }
-      return query.getResultStream().map(this::maybeDetachEntity);
+      return query.getResultStream().map(JpaTransactionManagerImpl.this::detach);
     }
 
     @Override
     public long count() {
-      CriteriaQueryBuilder<Long> queryBuilder = CriteriaQueryBuilder.createCount(em, entityClass);
+      CriteriaQueryBuilder<Long> queryBuilder =
+          CriteriaQueryBuilder.createCount(getEntityManager(), entityClass);
       return addCriteria(queryBuilder).getSingleResult();
     }
 
     @Override
     public ImmutableList<T> list() {
       return buildQuery().getResultList().stream()
-          .map(this::maybeDetachEntity)
+          .map(JpaTransactionManagerImpl.this::detach)
           .collect(ImmutableList.toImmutableList());
     }
-
-    private T maybeDetachEntity(T entity) {
-      if (autoDetachOnLoad) {
-        em.detach(entity);
-      }
-      return entity;
-    }
   }
 }

core/src/main/java/google/registry/persistence/transaction/TransactionManager.java

@@ -262,8 +262,14 @@ public interface TransactionManager {
   /** Deletes the set of entities by their key id. */
   void delete(Iterable<? extends VKey<?>> keys);
 
-  /** Deletes the given entity from the database. */
+  /**
-  void delete(Object entity);
+   * Deletes the given entity from the database.
+   *
+   * <p>This returns the deleted entity, which may not necessarily be the same as the original
+   * entity passed in, as it may be a) converted to a different type of object more appropriate to
+   * the database type or b) merged with an object managed by the database entity manager.
+   */
+  <T> T delete(T entity);
 
   /**
    * Deletes the entity by its id without writing commit logs if the underlying database is

core/src/test/java/google/registry/persistence/EntityCallbacksListenerTest.java

@@ -58,7 +58,7 @@ class EntityCallbacksListenerTest {
             .transact(
                 () -> {
                   TestEntity merged = jpaTm().getEntityManager().merge(testUpdate);
-                  merged.foo++;
+                  merged.nonTransientField++;
                   jpaTm().getEntityManager().flush();
                   return merged;
                 });
@@ -74,8 +74,7 @@ class EntityCallbacksListenerTest {
             .transact(
                 () -> {
                   TestEntity removed = jpaTm().loadByKey(VKey.createSql(TestEntity.class, "id"));
-                  jpaTm().delete(removed);
+                  return jpaTm().delete(removed);
-                  return removed;
                 });
     checkAll(testRemove, 0, 0, 1, 1);
   }
@@ -99,6 +98,24 @@ class EntityCallbacksListenerTest {
     assertThat(hasMethodAnnotatedWithEmbedded(ViolationEntity.class)).isTrue();
   }
 
+  @Test
+  void verifyCallbacksNotCalledOnCommit() {
+    TestEntity testEntity = new TestEntity();
+    jpaTm().transact(() -> jpaTm().insert(testEntity));
+
+    TestEntity testLoad =
+        jpaTm().transact(() -> jpaTm().loadByKey(VKey.createSql(TestEntity.class, "id")));
+    assertThat(testLoad.entityPreUpdate).isEqualTo(0);
+
+    testLoad = jpaTm().transact(() -> jpaTm().loadByKey(VKey.createSql(TestEntity.class, "id")));
+
+    // Verify that post-load happened but pre-update didn't.
+    assertThat(testLoad.entityPostLoad).isEqualTo(1);
+    assertThat(testLoad.entityPreUpdate).isEqualTo(0);
+    // since we didn't save the non-transient field, should only be 1
+    assertThat(testLoad.nonTransientField).isEqualTo(1);
+  }
+
   private static boolean hasMethodAnnotatedWithEmbedded(Class<?> entityType) {
     boolean result = false;
     Class<?> parentType = entityType.getSuperclass();
@@ -152,15 +169,22 @@ class EntityCallbacksListenerTest {
   @Entity(name = "TestEntity")
   private static class TestEntity extends ParentEntity {
     @Id String name = "id";
-    int foo = 0;
+    int nonTransientField = 0;
 
     @Transient int entityPostLoad = 0;
+    @Transient int entityPreUpdate = 0;
 
     @Embedded EntityEmbedded entityEmbedded = new EntityEmbedded();
 
     @PostLoad
     void entityPostLoad() {
       entityPostLoad++;
+      nonTransientField++;
+    }
+
+    @PreUpdate
+    void entityPreUpdate() {
+      entityPreUpdate++;
     }
   }
 

core/src/test/java/google/registry/persistence/transaction/JpaTransactionManagerImplTest.java

@@ -14,8 +14,10 @@
 
 package google.registry.persistence.transaction;
 
+import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.testing.DatabaseHelper.assertDetached;
 import static google.registry.testing.TestDataHelper.fileClassPath;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
@@ -25,9 +27,11 @@ import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import google.registry.model.ImmutableObject;
 import google.registry.persistence.VKey;
 import google.registry.persistence.transaction.JpaTestRules.JpaUnitTestExtension;
+import google.registry.testing.DatabaseHelper;
 import google.registry.testing.FakeClock;
 import java.io.Serializable;
 import java.math.BigInteger;
@@ -384,7 +388,7 @@ class JpaTransactionManagerImplTest {
   void load_succeeds() {
     assertThat(jpaTm().transact(() -> jpaTm().exists(theEntity))).isFalse();
     jpaTm().transact(() -> jpaTm().insert(theEntity));
-    TestEntity persisted = jpaTm().transact(() -> jpaTm().loadByKey(theEntityKey));
+    TestEntity persisted = jpaTm().transact(() -> assertDetached(jpaTm().loadByKey(theEntityKey)));
     assertThat(persisted.name).isEqualTo("theEntity");
     assertThat(persisted.data).isEqualTo("foo");
   }
@@ -397,11 +401,20 @@ class JpaTransactionManagerImplTest {
         () -> jpaTm().transact(() -> jpaTm().loadByKey(theEntityKey)));
   }
 
+  @Test
+  void loadByEntity_succeeds() {
+    jpaTm().transact(() -> jpaTm().insert(theEntity));
+    TestEntity persisted = jpaTm().transact(() -> assertDetached(jpaTm().loadByEntity(theEntity)));
+    assertThat(persisted.name).isEqualTo("theEntity");
+    assertThat(persisted.data).isEqualTo("foo");
+  }
+
   @Test
   void maybeLoad_succeeds() {
     assertThat(jpaTm().transact(() -> jpaTm().exists(theEntity))).isFalse();
     jpaTm().transact(() -> jpaTm().insert(theEntity));
-    TestEntity persisted = jpaTm().transact(() -> jpaTm().loadByKeyIfPresent(theEntityKey).get());
+    TestEntity persisted =
+        jpaTm().transact(() -> assertDetached(jpaTm().loadByKeyIfPresent(theEntityKey).get()));
     assertThat(persisted.name).isEqualTo("theEntity");
     assertThat(persisted.data).isEqualTo("foo");
   }
@@ -417,17 +430,81 @@ class JpaTransactionManagerImplTest {
   void loadCompoundIdEntity_succeeds() {
     assertThat(jpaTm().transact(() -> jpaTm().exists(compoundIdEntity))).isFalse();
     jpaTm().transact(() -> jpaTm().insert(compoundIdEntity));
-    TestCompoundIdEntity persisted = jpaTm().transact(() -> jpaTm().loadByKey(compoundIdEntityKey));
+    TestCompoundIdEntity persisted =
+        jpaTm().transact(() -> assertDetached(jpaTm().loadByKey(compoundIdEntityKey)));
     assertThat(persisted.name).isEqualTo("compoundIdEntity");
     assertThat(persisted.age).isEqualTo(10);
     assertThat(persisted.data).isEqualTo("foo");
   }
 
+  @Test
+  void loadByKeysIfPresent() {
+    jpaTm().transact(() -> jpaTm().insert(theEntity));
+    jpaTm()
+        .transact(
+            () -> {
+              ImmutableMap<VKey<? extends TestEntity>, TestEntity> results =
+                  jpaTm()
+                      .loadByKeysIfPresent(
+                          ImmutableList.of(
+                              theEntityKey, VKey.createSql(TestEntity.class, "does-not-exist")));
+
+              assertThat(results).containsExactly(theEntityKey, theEntity);
+              assertDetached(results.get(theEntityKey));
+            });
+  }
+
+  @Test
+  void loadByKeys_succeeds() {
+    jpaTm().transact(() -> jpaTm().insert(theEntity));
+    jpaTm()
+        .transact(
+            () -> {
+              ImmutableMap<VKey<? extends TestEntity>, TestEntity> results =
+                  jpaTm().loadByKeysIfPresent(ImmutableList.of(theEntityKey));
+              assertThat(results).containsExactly(theEntityKey, theEntity);
+              assertDetached(results.get(theEntityKey));
+            });
+  }
+
+  @Test
+  void loadByEntitiesIfPresent_succeeds() {
+    jpaTm().transact(() -> jpaTm().insert(theEntity));
+    jpaTm()
+        .transact(
+            () -> {
+              ImmutableList<TestEntity> results =
+                  jpaTm()
+                      .loadByEntitiesIfPresent(
+                          ImmutableList.of(theEntity, new TestEntity("does-not-exist", "bar")));
+              assertThat(results).containsExactly(theEntity);
+              assertDetached(results.get(0));
+            });
+  }
+
+  @Test
+  void loadByEntities_succeeds() {
+    jpaTm().transact(() -> jpaTm().insert(theEntity));
+    jpaTm()
+        .transact(
+            () -> {
+              ImmutableList<TestEntity> results =
+                  jpaTm().loadByEntities(ImmutableList.of(theEntity));
+              assertThat(results).containsExactly(theEntity);
+              assertDetached(results.get(0));
+            });
+  }
+
   @Test
   void loadAll_succeeds() {
     jpaTm().transact(() -> jpaTm().insertAll(moreEntities));
     ImmutableList<TestEntity> persisted =
-        jpaTm().transact(() -> jpaTm().loadAllOf(TestEntity.class));
+        jpaTm()
+            .transact(
+                () ->
+                    jpaTm().loadAllOf(TestEntity.class).stream()
+                        .map(DatabaseHelper::assertDetached)
+                        .collect(toImmutableList()));
     assertThat(persisted).containsExactlyElementsIn(moreEntities);
   }
 
@@ -462,6 +539,55 @@ class JpaTransactionManagerImplTest {
         () -> jpaTm().transact(() -> jpaTm().assertDelete(theEntityKey)));
   }
 
+  @Test
+  void loadAfterInsert_fails() {
+    assertThat(
+            assertThrows(
+                IllegalStateException.class,
+                () ->
+                    jpaTm()
+                        .transact(
+                            () -> {
+                              jpaTm().insert(theEntity);
+                              jpaTm().loadByKey(theEntityKey);
+                            })))
+        .hasMessageThat()
+        .contains("Inserted/updated object reloaded: ");
+  }
+
+  @Test
+  void loadAfterUpdate_fails() {
+    jpaTm().transact(() -> jpaTm().insert(theEntity));
+    assertThat(
+            assertThrows(
+                IllegalStateException.class,
+                () ->
+                    jpaTm()
+                        .transact(
+                            () -> {
+                              jpaTm().update(theEntity);
+                              jpaTm().loadByKey(theEntityKey);
+                            })))
+        .hasMessageThat()
+        .contains("Inserted/updated object reloaded: ");
+  }
+
+  @Test
+  void loadAfterPut_fails() {
+    assertThat(
+            assertThrows(
+                IllegalStateException.class,
+                () ->
+                    jpaTm()
+                        .transact(
+                            () -> {
+                              jpaTm().put(theEntity);
+                              jpaTm().loadByKey(theEntityKey);
+                            })))
+        .hasMessageThat()
+        .contains("Inserted/updated object reloaded: ");
+  }
+
   private void insertPerson(int age) {
     jpaTm()
         .getEntityManager()

core/src/test/java/google/registry/persistence/transaction/QueryComposerTest.java

@@ -26,6 +26,7 @@ import com.googlecode.objectify.annotation.Id;
 import com.googlecode.objectify.annotation.Index;
 import google.registry.model.ImmutableObject;
 import google.registry.testing.AppEngineExtension;
+import google.registry.testing.DatabaseHelper;
 import google.registry.testing.DualDatabaseTest;
 import google.registry.testing.FakeClock;
 import google.registry.testing.TestOfyAndSql;
@@ -70,20 +71,13 @@ public class QueryComposerTest {
 
   @TestOfyAndSql
   public void testFirstQueries() {
-    assertThat(
-            transactIfJpaTm(
-                () ->
-                    tm().createQueryComposer(TestEntity.class)
-                        .where("name", Comparator.EQ, "bravo")
-                        .first()
-                        .get()))
-        .isEqualTo(bravo);
     assertThat(
             transactIfJpaTm(
                 () ->
                     tm().createQueryComposer(TestEntity.class)
                         .where("name", Comparator.GT, "bravo")
                         .first()
+                        .map(DatabaseHelper::assertDetached)
                         .get()))
         .isEqualTo(charlie);
     assertThat(
@@ -92,6 +86,7 @@ public class QueryComposerTest {
                     tm().createQueryComposer(TestEntity.class)
                         .where("name", Comparator.GTE, "charlie")
                         .first()
+                        .map(DatabaseHelper::assertDetached)
                         .get()))
         .isEqualTo(charlie);
     assertThat(
@@ -100,6 +95,7 @@ public class QueryComposerTest {
                     tm().createQueryComposer(TestEntity.class)
                         .where("name", Comparator.LT, "bravo")
                         .first()
+                        .map(DatabaseHelper::assertDetached)
                         .get()))
         .isEqualTo(alpha);
     assertThat(
@@ -108,6 +104,7 @@ public class QueryComposerTest {
                     tm().createQueryComposer(TestEntity.class)
                         .where("name", Comparator.LTE, "alpha")
                         .first()
+                        .map(DatabaseHelper::assertDetached)
                         .get()))
         .isEqualTo(alpha);
   }
@@ -128,9 +125,10 @@ public class QueryComposerTest {
     assertThat(
             transactIfJpaTm(
                 () ->
+                    DatabaseHelper.assertDetached(
                         tm().createQueryComposer(TestEntity.class)
                             .where("name", Comparator.EQ, "alpha")
-                        .getSingleResult()))
+                            .getSingleResult())))
         .isEqualTo(alpha);
   }
 
@@ -176,6 +174,7 @@ public class QueryComposerTest {
                         .createQueryComposer(TestEntity.class)
                         .where("name", Comparator.GT, "alpha")
                         .stream()
+                        .map(DatabaseHelper::assertDetached)
                         .collect(toImmutableList())))
         .containsExactly(bravo, charlie);
     assertThat(
@@ -185,6 +184,7 @@ public class QueryComposerTest {
                         .createQueryComposer(TestEntity.class)
                         .where("name", Comparator.GTE, "bravo")
                         .stream()
+                        .map(DatabaseHelper::assertDetached)
                         .collect(toImmutableList())))
         .containsExactly(bravo, charlie);
     assertThat(
@@ -194,6 +194,7 @@ public class QueryComposerTest {
                         .createQueryComposer(TestEntity.class)
                         .where("name", Comparator.LT, "charlie")
                         .stream()
+                        .map(DatabaseHelper::assertDetached)
                         .collect(toImmutableList())))
         .containsExactly(alpha, bravo);
     assertThat(
@@ -203,6 +204,7 @@ public class QueryComposerTest {
                         .createQueryComposer(TestEntity.class)
                         .where("name", Comparator.LTE, "bravo")
                         .stream()
+                        .map(DatabaseHelper::assertDetached)
                         .collect(toImmutableList())))
         .containsExactly(alpha, bravo);
   }
@@ -226,6 +228,7 @@ public class QueryComposerTest {
                     tm().createQueryComposer(TestEntity.class)
                         .where("val", Comparator.EQ, 2)
                         .first()
+                        .map(DatabaseHelper::assertDetached)
                         .get()))
         .isEqualTo(bravo);
   }
@@ -240,6 +243,7 @@ public class QueryComposerTest {
                         .where("val", Comparator.GT, 1)
                         .orderBy("val")
                         .stream()
+                        .map(DatabaseHelper::assertDetached)
                         .collect(toImmutableList())))
         .containsExactly(bravo, alpha);
   }

core/src/test/java/google/registry/persistence/transaction/TransactionManagerTest.java

@@ -395,6 +395,17 @@ public class TransactionManagerTest {
         .isEqualTo("Expected at most one entity of type TestEntity, found at least two");
   }
 
+  @TestOfyAndSql
+  void mutatedObjectNotPersisted() {
+    tm().transact(() -> tm().insert(theEntity));
+    tm().transact(
+            () -> {
+              TestEntity e = tm().loadByKey(theEntity.key());
+              e.data = "some other data!";
+            });
+    assertThat(tm().transact(() -> tm().loadByKey(theEntity.key())).data).isEqualTo("foo");
+  }
+
   private static void assertEntityExists(TestEntity entity) {
     assertThat(tm().transact(() -> tm().exists(entity))).isTrue();
   }

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

@@ -110,7 +110,6 @@ import google.registry.model.transfer.TransferStatus;
 import google.registry.persistence.VKey;
 import google.registry.schema.tld.PremiumListDao;
 import google.registry.tmch.LordnTaskUtils;
-import java.lang.reflect.Field;
 import java.util.List;
 import java.util.Map;
 import java.util.NoSuchElementException;
@@ -119,7 +118,6 @@ import java.util.Set;
 import java.util.function.Consumer;
 import java.util.function.Function;
 import javax.annotation.Nullable;
-import org.hibernate.Hibernate;
 import org.joda.money.CurrencyUnit;
 import org.joda.money.Money;
 import org.joda.time.DateTime;
@@ -1065,9 +1063,6 @@ public class DatabaseHelper {
     // Force the session to be cleared so that when we read it back, we read from Datastore
     // and not from the transaction's session cache.
     tm().clearSessionCache();
-    for (R resource : resources) {
-      transactIfJpaTm(() -> tm().loadByEntity(resource));
-    }
   }
 
   /**
@@ -1284,9 +1279,6 @@ public class DatabaseHelper {
                 ImmutableList.Builder<Object> result = new ImmutableList.Builder<>();
                 for (Class<?> entityClass : entityClasses) {
                   for (Object object : jpaTm().loadAllOf(entityClass)) {
-                    // Initialize and detach the objects so they can be used for comparison later
-                    initializeHibernateObject(object);
-                    jpaTm().getEntityManager().detach(object);
                     result.add(object);
                   }
                 }
@@ -1295,20 +1287,6 @@ public class DatabaseHelper {
     }
   }
 
-  /**
-   * Initializes all fields in a Hibernate proxy object so it can be used outside of a transaction.
-   */
-  private static void initializeHibernateObject(Object object) {
-    for (Field field : object.getClass().getDeclaredFields()) {
-      field.setAccessible(true);
-      try {
-        Hibernate.initialize(field.get(object));
-      } catch (IllegalAccessException e) {
-        // shouldn't happen since we set the field to be accessible
-      }
-    }
-  }
-
   /**
    * Loads (i.e. reloads) the specified entity from the DB.
    *
@@ -1373,5 +1351,17 @@ public class DatabaseHelper {
     return transactIfJpaTm(() -> tm().loadByKeysIfPresent(keys));
   }
 
+  /**
+   * In JPA mode, assert that the given entity is detached from the current entity manager.
+   *
+   * <p>Returns the original entity object.
+   */
+  public static <T> T assertDetached(T entity) {
+    if (!tm().isOfy()) {
+      assertThat(jpaTm().getEntityManager().contains(entity)).isFalse();
+    }
+    return entity;
+  }
+
   private DatabaseHelper() {}
 }