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

Author: mp911de

spring-data-jpa/src/main/java/org/springframework/data/jpa/provider/PersistenceProvider.java

@@ -30,13 +30,15 @@
 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;
 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 String getCommentHintKey() {
 			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 String getCommentHintKey() {
 		public String getCommentHintValue(String comment) {
 			return "/* " + comment + " */";
 		}
+
 	},
 
 	/**
@@ -197,6 +211,7 @@ public boolean shouldUseAccessorFor(Object entity) {
 		public @Nullable String getCommentHintKey() {
 			return null;
 		}
+
 	};
 
 	private static final @Nullable Class<?> typedParameterValueClass;
@@ -406,6 +421,18 @@ public boolean isPresent() {
 		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 @@ interface Constants {
 
 		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<Object> executeQueryWithResultStream(Query jpaQuery) {
@@ -482,6 +510,7 @@ public void close() {
 				scrollableResults.close();
 			}
 		}
+
 	}
 
 	/**
@@ -531,5 +560,7 @@ public void close() {
 				scrollableCursor.close();
 			}
 		}
+
 	}
+
 }

spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/AbstractJpaQuery.java

@@ -106,7 +106,7 @@ public AbstractJpaQuery(JpaQueryMethod method, EntityManager em) {
 			} 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 JpaQueryMethod getQueryMethod() {
 		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}.
 	 *

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<ParameterBinder> 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 @@ public AbstractStringBasedJpaQuery(JpaQueryMethod method, EntityManager em, Decl
 		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 @@ public AbstractStringBasedJpaQuery(JpaQueryMethod method, EntityManager em, Decl
 				"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

spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/JpaQueryExecution.java

@@ -188,20 +188,47 @@ protected Object doExecute(AbstractJpaQuery query, JpaParametersParameterAccesso
 	 */
 	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) {
 
 			Query query = repositoryQuery.createQuery(accessor);
 
 			return PageableExecutionUtils.getPage(query.getResultList(), accessor.getPageable(),
-					() -> count(repositoryQuery, accessor));
+					() -> count(query, repositoryQuery, accessor));
+		}
+
+		private long count(Query resultQuery, AbstractJpaQuery repositoryQuery, JpaParametersParameterAccessor accessor) {
+
+			if (repositoryQuery.hasDeclaredCountQuery()) {
+				return doCount(repositoryQuery, accessor);
+			}
+
+			return provider.getResultCount(resultQuery, () -> doCount(repositoryQuery, accessor));
 		}
 
-		private long count(AbstractJpaQuery repositoryQuery, JpaParametersParameterAccessor 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();
 		}
 	}
 

spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/NamedQuery.java

@@ -182,6 +182,11 @@ static boolean hasNamedQuery(EntityManager em, String queryName) {
 		return query;
 	}
 
+	@Override
+	public boolean hasDeclaredCountQuery() {
+		return namedCountQueryIsPresent;
+	}
+
 	@Override
 	protected Query doCreateQuery(JpaParametersParameterAccessor accessor) {
 

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);

spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/StoredProcedureJpaQuery.java

@@ -81,6 +81,11 @@ private static boolean useNamedParameters(QueryMethod method) {
 		return false;
 	}
 
+	@Override
+	public boolean hasDeclaredCountQuery() {
+		return false;
+	}
+
 	@Override
 	protected StoredProcedureQuery createQuery(JpaParametersParameterAccessor accessor) {
 		return applyHints(doCreateQuery(accessor), getQueryMethod());

spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/AbstractJpaQueryTests.java

@@ -230,6 +230,11 @@ protected Query doCreateQuery(JpaParametersParameterAccessor accessor) {
 			return query;
 		}
 
+		@Override
+		public boolean hasDeclaredCountQuery() {
+			return true;
+		}
+
 		@Override
 		protected TypedQuery<Long> doCreateCountQuery(JpaParametersParameterAccessor accessor) {
 			return countQuery;

spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/JpaQueryExecutionUnitTests.java

@@ -40,6 +40,7 @@
 
 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 @@ void pagedExecutionRetrievesObjectsForPageableOutOfRange() throws Exception {
 		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 @@ void pagedExecutionShouldNotGenerateCountQueryIfQueryReportedNoResults() throws
 		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 @@ void pagedExecutionShouldUseCountFromResultIfOffsetIsZeroAndResultsWithinPageSiz
 		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 @@ void pagedExecutionShouldUseCountFromResultWithOffsetAndResultsWithinPageSize()
 		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 @@ void pagedExecutionShouldUseRequestCountFromResultWithOffsetAndResultsHitLowerPa
 		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 @@ void pagedExecutionShouldUseRequestCountFromResultWithOffsetAndResultsHitUpperPa
 		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) }));
 

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`
 ====