From 4476510acca0e6a04e2935ef31882b6ee9147185 Mon Sep 17 00:00:00 2001 From: Gus Brodman Date: Tue, 3 Nov 2026 14:47:06 -0500 Subject: [PATCH] Use getResultList() in DetachingTypedQuery When possible, we want to delegate to TypedQuery.getResultList instead of TypedQuery.getResultStream. This is because the stream hydrates and `@PostLoad`s entities (including eager-loading collections) one at a time, which means we can't batch. Retrieving the list hydrates all the results into the persistence context before initializing child collections, allowing us to batch the child collection loads across different loaded entities. --- .../JpaTransactionManagerImpl.java | 22 +++--- .../JpaTransactionManagerImplTest.java | 72 ++++++++++++++++++- 2 files changed, 85 insertions(+), 9 deletions(-) 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; + } + } }