From b66c96d363f0804c8142abada4d55e4bed3742c6 Mon Sep 17 00:00:00 2001 From: Mark Paluch Date: Fri, 31 Jan 2025 09:58:20 +0100 Subject: [PATCH] Use `SelectionQuery.getResultCount()` for count queries if possible. We now use Hibernate's built-in mechanism to obtain the result count if there is an enclosing transaction. Without the transaction, the session is being closed and we cannot run the query. Closes #3456 --- .../jpa/provider/PersistenceProvider.java | 31 +++++++++++++++++ .../repository/query/AbstractJpaQuery.java | 11 ++++++- .../query/AbstractStringBasedJpaQuery.java | 7 ++-- .../repository/query/JpaQueryExecution.java | 33 +++++++++++++++++-- .../data/jpa/repository/query/NamedQuery.java | 5 +++ .../repository/query/PartTreeJpaQuery.java | 5 +++ .../query/StoredProcedureJpaQuery.java | 5 +++ .../query/AbstractJpaQueryTests.java | 5 +++ .../query/JpaQueryExecutionUnitTests.java | 13 ++++---- .../modules/ROOT/pages/jpa/query-methods.adoc | 2 +- 10 files changed, 104 insertions(+), 13 deletions(-) diff --git a/spring-data-jpa/src/main/java/org/springframework/data/jpa/provider/PersistenceProvider.java b/spring-data-jpa/src/main/java/org/springframework/data/jpa/provider/PersistenceProvider.java index 4d604b452..9755f19f0 100644 --- a/spring-data-jpa/src/main/java/org/springframework/data/jpa/provider/PersistenceProvider.java +++ b/spring-data-jpa/src/main/java/org/springframework/data/jpa/provider/PersistenceProvider.java @@ -30,6 +30,7 @@ import java.util.Collections; import java.util.List; import java.util.NoSuchElementException; import java.util.Set; +import java.util.function.LongSupplier; import org.eclipse.persistence.config.QueryHints; import org.eclipse.persistence.jpa.JpaQuery; @@ -37,6 +38,7 @@ import org.eclipse.persistence.queries.ScrollableCursor; import org.hibernate.ScrollMode; import org.hibernate.ScrollableResults; import org.hibernate.proxy.HibernateProxy; +import org.hibernate.query.SelectionQuery; import org.jspecify.annotations.Nullable; import org.springframework.data.util.CloseableIterator; @@ -117,6 +119,17 @@ public enum PersistenceProvider implements QueryExtractor, ProxyIdAccessor, Quer return "org.hibernate.comment"; } + @Override + public long getResultCount(Query resultQuery, LongSupplier countSupplier) { + + if (TransactionSynchronizationManager.isActualTransactionActive() + && resultQuery instanceof SelectionQuery sq) { + return sq.getResultCount(); + } + + return super.getResultCount(resultQuery, countSupplier); + } + }, /** @@ -160,6 +173,7 @@ public enum PersistenceProvider implements QueryExtractor, ProxyIdAccessor, Quer public String getCommentHintValue(String comment) { return "/* " + comment + " */"; } + }, /** @@ -197,6 +211,7 @@ public enum PersistenceProvider implements QueryExtractor, ProxyIdAccessor, Quer public @Nullable String getCommentHintKey() { return null; } + }; private static final @Nullable Class typedParameterValueClass; @@ -406,6 +421,18 @@ public enum PersistenceProvider implements QueryExtractor, ProxyIdAccessor, Quer return this.present; } + /** + * Obtain the result count from a {@link Query} returning the result or fall back to {@code countSupplier} if the + * query does not provide the result count. + * + * @param resultQuery the query that has returned {@link Query#getResultList()} + * @param countSupplier fallback supplier to provide the count if the query does not provide it. + * @return the result count. + */ + public long getResultCount(Query resultQuery, LongSupplier countSupplier) { + return countSupplier.getAsLong(); + } + /** * Holds the PersistenceProvider specific interface names. * @@ -427,6 +454,7 @@ public enum PersistenceProvider implements QueryExtractor, ProxyIdAccessor, Quer String HIBERNATE_JPA_METAMODEL_TYPE = "org.hibernate.metamodel.model.domain.JpaMetamodel"; String ECLIPSELINK_JPA_METAMODEL_TYPE = "org.eclipse.persistence.internal.jpa.metamodel.MetamodelImpl"; + } public CloseableIterator executeQueryWithResultStream(Query jpaQuery) { @@ -482,6 +510,7 @@ public enum PersistenceProvider implements QueryExtractor, ProxyIdAccessor, Quer scrollableResults.close(); } } + } /** @@ -531,5 +560,7 @@ public enum PersistenceProvider implements QueryExtractor, ProxyIdAccessor, Quer scrollableCursor.close(); } } + } + } diff --git a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/AbstractJpaQuery.java b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/AbstractJpaQuery.java index 4e672ccc8..ef604e1f5 100644 --- a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/AbstractJpaQuery.java +++ b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/AbstractJpaQuery.java @@ -106,7 +106,7 @@ public abstract class AbstractJpaQuery implements RepositoryQuery { } else if (method.isSliceQuery()) { return new SlicedExecution(); } else if (method.isPageQuery()) { - return new PagedExecution(); + return new PagedExecution(this.provider); } else if (method.isModifyingQuery()) { return null; } else { @@ -120,6 +120,15 @@ public abstract class AbstractJpaQuery implements RepositoryQuery { return method; } + /** + * Returns {@literal true} if the query has a dedicated count query associated with it or {@literal false} if the + * count query shall be derived. + * + * @return {@literal true} if the query has a dedicated count query {@literal false} if the * count query is derived. + * @since 3.5 + */ + public abstract boolean hasDeclaredCountQuery(); + /** * Returns the {@link EntityManager}. * diff --git a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/AbstractStringBasedJpaQuery.java b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/AbstractStringBasedJpaQuery.java index 64b21e8cb..c288d4a35 100644 --- a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/AbstractStringBasedJpaQuery.java +++ b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/AbstractStringBasedJpaQuery.java @@ -62,6 +62,7 @@ abstract class AbstractStringBasedJpaQuery extends AbstractJpaQuery { private final QuerySortRewriter querySortRewriter; private final Lazy countParameterBinder; private final ValueEvaluationContextProvider valueExpressionContextProvider; + private final boolean hasDeclaredCountQuery; /** * Creates a new {@link AbstractStringBasedJpaQuery} from the given {@link JpaQueryMethod}, {@link EntityManager} and @@ -101,6 +102,7 @@ abstract class AbstractStringBasedJpaQuery extends AbstractJpaQuery { this.valueExpressionContextProvider = valueExpressionDelegate.createValueContextProvider(method.getParameters()); this.query = TemplatedQuery.create(query, method.getEntityInformation(), queryConfiguration); + this.hasDeclaredCountQuery = countQuery != null; this.countQuery = Lazy.of(() -> { @@ -130,8 +132,9 @@ abstract class AbstractStringBasedJpaQuery extends AbstractJpaQuery { "JDBC style parameters (?) are not supported for JPA queries"); } - private DeclaredQuery createQuery(String queryString, boolean nativeQuery) { - return nativeQuery ? DeclaredQuery.nativeQuery(queryString) : DeclaredQuery.jpqlQuery(queryString); + @Override + public boolean hasDeclaredCountQuery() { + return hasDeclaredCountQuery; } @Override diff --git a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/JpaQueryExecution.java b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/JpaQueryExecution.java index 338a2204e..be0a09bc4 100644 --- a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/JpaQueryExecution.java +++ b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/JpaQueryExecution.java @@ -188,6 +188,12 @@ public abstract class JpaQueryExecution { */ static class PagedExecution extends JpaQueryExecution { + private final PersistenceProvider provider; + + PagedExecution(PersistenceProvider provider) { + this.provider = provider; + } + @Override @SuppressWarnings("unchecked") protected Object doExecute(AbstractJpaQuery repositoryQuery, JpaParametersParameterAccessor accessor) { @@ -195,13 +201,34 @@ public abstract class JpaQueryExecution { Query query = repositoryQuery.createQuery(accessor); return PageableExecutionUtils.getPage(query.getResultList(), accessor.getPageable(), - () -> count(repositoryQuery, accessor)); + () -> count(query, repositoryQuery, accessor)); } - private long count(AbstractJpaQuery repositoryQuery, JpaParametersParameterAccessor accessor) { + private long count(Query resultQuery, AbstractJpaQuery repositoryQuery, JpaParametersParameterAccessor accessor) { + + if (repositoryQuery.hasDeclaredCountQuery()) { + return doCount(repositoryQuery, accessor); + } + + return provider.getResultCount(resultQuery, () -> doCount(repositoryQuery, accessor)); + } + + long doCount(AbstractJpaQuery repositoryQuery, JpaParametersParameterAccessor accessor) { List totals = repositoryQuery.createCountQuery(accessor).getResultList(); - return (totals.size() == 1 ? CONVERSION_SERVICE.convert(totals.get(0), Long.class) : totals.size()); + + if (totals.size() == 1) { + Object result = totals.get(0); + + if (result instanceof Number n) { + return n.longValue(); + } + + return CONVERSION_SERVICE.convert(result, Long.class); + } + + // group by count + return totals.size(); } } diff --git a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/NamedQuery.java b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/NamedQuery.java index a38bf9eaa..125ec40c6 100644 --- a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/NamedQuery.java +++ b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/NamedQuery.java @@ -182,6 +182,11 @@ final class NamedQuery extends AbstractJpaQuery { return query; } + @Override + public boolean hasDeclaredCountQuery() { + return namedCountQueryIsPresent; + } + @Override protected Query doCreateQuery(JpaParametersParameterAccessor accessor) { diff --git a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/PartTreeJpaQuery.java b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/PartTreeJpaQuery.java index 66dac4792..bf254c46b 100644 --- a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/PartTreeJpaQuery.java +++ b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/PartTreeJpaQuery.java @@ -112,6 +112,11 @@ public class PartTreeJpaQuery extends AbstractJpaQuery { } } + @Override + public boolean hasDeclaredCountQuery() { + return false; + } + @Override public Query doCreateQuery(JpaParametersParameterAccessor accessor) { return queryPreparer.createQuery(accessor); diff --git a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/StoredProcedureJpaQuery.java b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/StoredProcedureJpaQuery.java index 3423c71e4..caece33d0 100644 --- a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/StoredProcedureJpaQuery.java +++ b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/StoredProcedureJpaQuery.java @@ -81,6 +81,11 @@ class StoredProcedureJpaQuery extends AbstractJpaQuery { return false; } + @Override + public boolean hasDeclaredCountQuery() { + return false; + } + @Override protected StoredProcedureQuery createQuery(JpaParametersParameterAccessor accessor) { return applyHints(doCreateQuery(accessor), getQueryMethod()); diff --git a/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/AbstractJpaQueryTests.java b/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/AbstractJpaQueryTests.java index 8728e0322..fdcbabf84 100644 --- a/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/AbstractJpaQueryTests.java +++ b/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/AbstractJpaQueryTests.java @@ -230,6 +230,11 @@ class AbstractJpaQueryTests { return query; } + @Override + public boolean hasDeclaredCountQuery() { + return true; + } + @Override protected TypedQuery doCreateCountQuery(JpaParametersParameterAccessor accessor) { return countQuery; diff --git a/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/JpaQueryExecutionUnitTests.java b/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/JpaQueryExecutionUnitTests.java index e8907f16f..c7d160d6d 100644 --- a/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/JpaQueryExecutionUnitTests.java +++ b/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/JpaQueryExecutionUnitTests.java @@ -40,6 +40,7 @@ import org.mockito.quality.Strictness; import org.springframework.data.domain.PageRequest; import org.springframework.data.domain.Pageable; +import org.springframework.data.jpa.provider.PersistenceProvider; import org.springframework.data.jpa.provider.QueryExtractor; import org.springframework.data.jpa.repository.Modifying; import org.springframework.data.jpa.repository.query.JpaQueryExecution.ModifyingExecution; @@ -183,7 +184,7 @@ class JpaQueryExecutionUnitTests { when(jpaQuery.createQuery(Mockito.any())).thenReturn(query); when(countQuery.getResultList()).thenReturn(Arrays.asList(20L)); - PagedExecution execution = new PagedExecution(); + PagedExecution execution = new PagedExecution(PersistenceProvider.GENERIC_JPA); execution.doExecute(jpaQuery, new JpaParametersParameterAccessor(parameters, new Object[] { PageRequest.of(2, 10) })); @@ -199,7 +200,7 @@ class JpaQueryExecutionUnitTests { when(jpaQuery.createQuery(Mockito.any())).thenReturn(query); when(query.getResultList()).thenReturn(Arrays.asList(0L)); - PagedExecution execution = new PagedExecution(); + PagedExecution execution = new PagedExecution(PersistenceProvider.GENERIC_JPA); execution.doExecute(jpaQuery, new JpaParametersParameterAccessor(parameters, new Object[] { PageRequest.of(0, 10) })); @@ -215,7 +216,7 @@ class JpaQueryExecutionUnitTests { when(jpaQuery.createQuery(Mockito.any())).thenReturn(query); when(query.getResultList()).thenReturn(Arrays.asList(new Object(), new Object(), new Object(), new Object())); - PagedExecution execution = new PagedExecution(); + PagedExecution execution = new PagedExecution(PersistenceProvider.GENERIC_JPA); execution.doExecute(jpaQuery, new JpaParametersParameterAccessor(parameters, new Object[] { PageRequest.of(0, 10) })); @@ -230,7 +231,7 @@ class JpaQueryExecutionUnitTests { when(jpaQuery.createQuery(Mockito.any())).thenReturn(query); when(query.getResultList()).thenReturn(Arrays.asList(new Object(), new Object(), new Object(), new Object())); - PagedExecution execution = new PagedExecution(); + PagedExecution execution = new PagedExecution(PersistenceProvider.GENERIC_JPA); execution.doExecute(jpaQuery, new JpaParametersParameterAccessor(parameters, new Object[] { PageRequest.of(5, 10) })); @@ -247,7 +248,7 @@ class JpaQueryExecutionUnitTests { when(jpaQuery.createCountQuery(Mockito.any())).thenReturn(query); when(countQuery.getResultList()).thenReturn(Arrays.asList(20L)); - PagedExecution execution = new PagedExecution(); + PagedExecution execution = new PagedExecution(PersistenceProvider.GENERIC_JPA); execution.doExecute(jpaQuery, new JpaParametersParameterAccessor(parameters, new Object[] { PageRequest.of(4, 4) })); @@ -264,7 +265,7 @@ class JpaQueryExecutionUnitTests { when(jpaQuery.createCountQuery(Mockito.any())).thenReturn(query); when(countQuery.getResultList()).thenReturn(Arrays.asList(20L)); - PagedExecution execution = new PagedExecution(); + PagedExecution execution = new PagedExecution(PersistenceProvider.GENERIC_JPA); execution.doExecute(jpaQuery, new JpaParametersParameterAccessor(parameters, new Object[] { PageRequest.of(4, 4) })); diff --git a/src/main/antora/modules/ROOT/pages/jpa/query-methods.adoc b/src/main/antora/modules/ROOT/pages/jpa/query-methods.adoc index 5513309f4..b947fca73 100644 --- a/src/main/antora/modules/ROOT/pages/jpa/query-methods.adoc +++ b/src/main/antora/modules/ROOT/pages/jpa/query-methods.adoc @@ -294,7 +294,7 @@ Sometimes, no matter how many features you try to apply, it seems impossible to You have the ability to get your hands on the query, right before it's sent to the `EntityManager` and "rewrite" it. That is, you can make any alterations at the last moment. Query rewriting applies to the actual query and, when applicable, to count queries. -Count queries are optimized and therefore, either not necessary or a count is obtained through other means, such as derived from a Hibernate `SelectionQuery`. +Count queries are optimized and therefore, either not necessary or a count is obtained through other means, such as derived from a Hibernate `SelectionQuery` if there is an enclosing transaction. .Declare a QueryRewriter using `@Query` ====