diff --git a/core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java b/core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java index d2c0704c14f..2acce3363d3 100644 --- a/core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java +++ b/core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java @@ -622,11 +622,19 @@ public Stream loadAllOfStream(Class clazz) { @Override public ImmutableList loadAllOfSorted(Class clazz, String... sortFields) { - return loadAllOfSortedStream(clazz, sortFields).collect(toImmutableList()); + // Use getResultList() instead of getResultStream() so Hibernate hydrates all entities into + // the PersistenceContext and BatchFetchQueue before initializing collections and detaching. + return createLoadAllQuery(clazz, sortFields).getResultList().stream() + .map(this::detach) + .collect(toImmutableList()); } @Override public Stream loadAllOfSortedStream(Class clazz, String... sortFields) { + return createLoadAllQuery(clazz, sortFields).getResultStream().map(this::detach); + } + + private TypedQuery createLoadAllQuery(Class clazz, String... sortFields) { checkArgumentNotNull(clazz, "clazz must be specified"); checkArgumentNotNull(sortFields, "sortFields must not be null"); assertInTransaction(); @@ -642,10 +650,7 @@ public Stream loadAllOfSortedStream(Class clazz, String... sortFields) queryString.append(" ORDER BY "); queryString.append(String.join(", ", sortFields)); } - return getEntityManager() - .createQuery(queryString.toString(), clazz) - .getResultStream() - .map(this::detach); + return getEntityManager().createQuery(queryString.toString(), clazz); } @Override @@ -1034,13 +1039,14 @@ public TypedQuery setCacheRetrieveMode(CacheRetrieveMode mode) { @Override public T getSingleResultOrNull() { - return delegate.getSingleResultOrNull(); + return detachIfEntity(delegate.getSingleResultOrNull()); } @Override public List getResultList() { - return delegate - .getResultStream() + // Delegate to getResultList() rather than getResultStream() so Hibernate can batch-fetch + // child collections across all loaded entities before each entity is detached. + return delegate.getResultList().stream() .map(JpaTransactionManagerImpl.this::detachIfEntity) .collect(toImmutableList()); } diff --git a/core/src/test/java/google/registry/persistence/transaction/JpaTransactionManagerImplTest.java b/core/src/test/java/google/registry/persistence/transaction/JpaTransactionManagerImplTest.java index 18f5277fbef..a6c628842b5 100644 --- a/core/src/test/java/google/registry/persistence/transaction/JpaTransactionManagerImplTest.java +++ b/core/src/test/java/google/registry/persistence/transaction/JpaTransactionManagerImplTest.java @@ -38,6 +38,7 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; import google.registry.config.RegistryConfig; import google.registry.model.ImmutableObject; import google.registry.persistence.PersistenceModule.TransactionIsolationLevel; @@ -45,15 +46,20 @@ import google.registry.persistence.transaction.JpaTestExtensions.JpaUnitTestExtension; import google.registry.testing.DatabaseHelper; import google.registry.testing.FakeClock; +import jakarta.persistence.ElementCollection; import jakarta.persistence.Entity; import jakarta.persistence.EntityManager; +import jakarta.persistence.FetchType; import jakarta.persistence.Id; import jakarta.persistence.IdClass; import jakarta.persistence.OptimisticLockException; import jakarta.persistence.PersistenceException; import jakarta.persistence.RollbackException; import java.io.Serializable; +import java.util.List; import java.util.NoSuchElementException; +import java.util.Set; +import java.util.function.Supplier; import org.hibernate.SessionFactory; import org.hibernate.stat.Statistics; import org.junit.jupiter.api.Test; @@ -89,7 +95,10 @@ class JpaTransactionManagerImplTest { .withInitScript(fileClassPath(getClass(), "test_schema.sql")) .withClock(fakeClock) .withEntityClass( - TestEntity.class, TestCompoundIdEntity.class, TestNamedCompoundIdEntity.class) + TestEntity.class, + TestCompoundIdEntity.class, + TestNamedCompoundIdEntity.class, + TestCollectionEntity.class) .buildUnitTestExtension(); @Test @@ -960,6 +969,52 @@ void query_detachesResults() { "FROM TestEntity WHERE name = 'entity1'", TestEntity.class) .getSingleResult())) .isFalse()); + tm().transact( + () -> { + TestEntity result = + tm().query("FROM TestEntity WHERE name = 'entity1'", TestEntity.class) + .getSingleResultOrNull(); + assertThat(result).isEqualTo(moreEntities.get(0)); + assertThat(tm().getEntityManager().contains(result)).isFalse(); + assertThat( + tm().query("FROM TestEntity WHERE name = 'nonexistent'", TestEntity.class) + .getSingleResultOrNull()) + .isNull(); + }); + } + + @Test + void queryAndLoadAllOf_batchFetchEagerCollections() { + ImmutableList entities = + ImmutableList.of( + new TestCollectionEntity("c1", ImmutableSet.of("a", "b")), + new TestCollectionEntity("c2", ImmutableSet.of("c")), + new TestCollectionEntity("c3", ImmutableSet.of("d", "e"))); + persistResources(entities); + + for (Supplier> querySupplier : + ImmutableList.>>of( + () -> + tm().query("FROM TestCollectionEntity ORDER BY name", TestCollectionEntity.class) + .getResultList(), + () -> tm().loadAllOfSorted(TestCollectionEntity.class, "name"))) { + tm().transact( + () -> { + Statistics stats = + tm().getEntityManager() + .getEntityManagerFactory() + .unwrap(SessionFactory.class) + .getStatistics(); + stats.setStatisticsEnabled(true); + stats.clear(); + + List loaded = querySupplier.get(); + // 1 query for TestCollectionEntity + 1 batched query for all 3 eager collections. + assertThat(stats.getPrepareStatementCount()).isEqualTo(2); + loaded.forEach(DatabaseHelper::assertDetachedFromEntityManager); + assertThat(loaded).containsExactlyElementsIn(entities).inOrder(); + }); + } } @Test @@ -1168,4 +1223,19 @@ private void setAgeField(int ageField) { this.ageField = ageField; } } + + @Entity(name = "TestCollectionEntity") + private static class TestCollectionEntity extends ImmutableObject { + @Id private String name; + + @ElementCollection(fetch = FetchType.EAGER) + private Set items; + + private TestCollectionEntity() {} + + private TestCollectionEntity(String name, Set items) { + this.name = name; + this.items = items; + } + } }