Apply QueryRewriter to count queries as well.

We now use QueryRewriter to post-process count queries as well. Previously, only the actual result query has been processed.

Closes #3801

Author: mp911de

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

@@ -112,7 +112,6 @@ public AbstractStringBasedJpaQuery(JpaQueryMethod method, EntityManager em, Decl
 		});
 
 		this.countParameterBinder = Lazy.of(() -> this.createBinder(this.countQuery.get()));
-
 		this.queryRewriter = queryConfiguration.getQueryRewriter(method);
 
 		JpaParameters parameters = method.getParameters();

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

@@ -169,7 +169,7 @@ protected RepositoryQuery resolveQuery(JpaQueryMethod method, JpaQueryConfigurat
 						configuration);
 			}
 
-			RepositoryQuery query = NamedQuery.lookupFrom(method, em, configuration.getSelector());
+			RepositoryQuery query = NamedQuery.lookupFrom(method, em, configuration);
 
 			return query != null ? query : NO_QUERY;
 		}

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

@@ -56,19 +56,20 @@ final class NamedQuery extends AbstractJpaQuery {
 	private final @Nullable String countProjection;
 	private final boolean namedCountQueryIsPresent;
 	private final Lazy<EntityQuery> entityQuery;
+	private final QueryRewriter queryRewriter;
 
 	/**
 	 * Creates a new {@link NamedQuery}.
 	 */
-	private NamedQuery(JpaQueryMethod method, EntityManager em, QueryEnhancerSelector selector, QueryRewriter queryRewriter) {
+	private NamedQuery(JpaQueryMethod method, EntityManager em, JpaQueryConfiguration queryConfiguration) {
 
 		super(method, em);
 
 		this.queryName = method.getNamedQueryName();
 		this.countQueryName = method.getNamedCountQueryName();
 		QueryExtractor extractor = method.getQueryExtractor();
 		this.countProjection = method.getCountQueryProjection();
-		this.queryRewriter = queryRewriter;
+		this.queryRewriter = queryConfiguration.getQueryRewriter(method);
 
 		Parameters<?, ?> parameters = method.getParameters();
 
@@ -104,7 +105,7 @@ private NamedQuery(JpaQueryMethod method, EntityManager em, QueryEnhancerSelecto
 			declaredQuery = DeclaredQuery.jpqlQuery(queryString);
 		}
 
-		this.entityQuery = Lazy.of(() -> EntityQuery.create(declaredQuery, selector));
+		this.entityQuery = Lazy.of(() -> EntityQuery.create(declaredQuery, queryConfiguration.getSelector()));
 	}
 
 	/**
@@ -138,9 +139,10 @@ static boolean hasNamedQuery(EntityManager em, String queryName) {
 	 * @param method must not be {@literal null}.
 	 * @param em must not be {@literal null}.
 	 * @param selector must not be {@literal null}.
+	 * @param queryConfiguration must not be {@literal null}.
 	 */
 	public static @Nullable RepositoryQuery lookupFrom(JpaQueryMethod method, EntityManager em,
-			QueryEnhancerSelector selector) {
+			JpaQueryConfiguration queryConfiguration) {
 
 		String queryName = method.getNamedQueryName();
 
@@ -158,7 +160,7 @@ static boolean hasNamedQuery(EntityManager em, String queryName) {
 					method.isNativeQuery() ? "NativeQuery" : "Query"));
 		}
 
-		RepositoryQuery query = new NamedQuery(method, em, selector);
+		RepositoryQuery query = new NamedQuery(method, em, queryConfiguration);
 		if (LOG.isDebugEnabled()) {
 			LOG.debug(String.format("Found named query '%s'", queryName));
 		}
@@ -193,6 +195,7 @@ protected TypedQuery<Long> doCreateCountQuery(JpaParametersParameterAccessor acc
 		} else {
 
 			String countQueryString = entityQuery.get().deriveCountQuery(countProjection).getQueryString();
+			countQueryString = potentiallyRewriteQuery(countQueryString, accessor.getSort(), accessor.getPageable());
 			countQuery = em.createQuery(countQueryString, Long.class);
 		}
 
@@ -235,9 +238,9 @@ protected TypedQuery<Long> doCreateCountQuery(JpaParametersParameterAccessor acc
 	 * @param pageable
 	 * @return
 	 */
-	private String potentiallyRewriteQuery(String originalQuery, Sort sort, Pageable pageable) {
+	private String potentiallyRewriteQuery(String originalQuery, Sort sort, @Nullable Pageable pageable) {
 
-		return pageable.isPaged() //
+		return pageable != null && pageable.isPaged() //
 				? queryRewriter.rewrite(originalQuery, pageable) //
 				: queryRewriter.rewrite(originalQuery, sort);
 	}

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

@@ -41,6 +41,7 @@
 import org.springframework.data.projection.SpelAwareProxyProjectionFactory;
 import org.springframework.data.repository.core.RepositoryMetadata;
 import org.springframework.data.repository.query.QueryCreationException;
+import org.springframework.data.repository.query.ValueExpressionDelegate;
 import org.springframework.data.util.TypeInformation;
 
 /**
@@ -55,6 +56,9 @@
 @MockitoSettings(strictness = Strictness.LENIENT)
 class NamedQueryUnitTests {
 
+	private static final JpaQueryConfiguration CONFIG = new JpaQueryConfiguration(QueryRewriterProvider.simple(),
+			QueryEnhancerSelector.DEFAULT_SELECTOR, ValueExpressionDelegate.create(), EscapeCharacter.DEFAULT);
+
 	@Mock RepositoryMetadata metadata;
 	@Mock QueryExtractor extractor;
 	@Mock EntityManager em;
@@ -89,7 +93,8 @@ void rejectsPersistenceProviderIfIncapableOfExtractingQueriesAndPagebleBeingUsed
 		JpaQueryMethod queryMethod = new JpaQueryMethod(method, metadata, projectionFactory, extractor);
 
 		when(em.createNamedQuery(queryMethod.getNamedCountQueryName())).thenThrow(new IllegalArgumentException());
-		assertThatExceptionOfType(QueryCreationException.class).isThrownBy(() -> NamedQuery.lookupFrom(queryMethod, em, QueryEnhancerSelector.DEFAULT_SELECTOR, QueryRewriter.IdentityQueryRewriter.INSTANCE));
+		assertThatExceptionOfType(QueryCreationException.class)
+				.isThrownBy(() -> NamedQuery.lookupFrom(queryMethod, em, CONFIG));
 	}
 
 	@Test // DATAJPA-142
@@ -101,7 +106,7 @@ void doesNotRejectPersistenceProviderIfNamedCountQueryIsAvailable() {
 
 		TypedQuery<Long> countQuery = mock(TypedQuery.class);
 		when(em.createNamedQuery(eq(queryMethod.getNamedCountQueryName()), eq(Long.class))).thenReturn(countQuery);
-		NamedQuery query = (NamedQuery) NamedQuery.lookupFrom(queryMethod, em, QueryEnhancerSelector.DEFAULT_SELECTOR, QueryRewriter.IdentityQueryRewriter.INSTANCE);
+		NamedQuery query = (NamedQuery) NamedQuery.lookupFrom(queryMethod, em, CONFIG);
 
 		query.doCreateCountQuery(new JpaParametersParameterAccessor(queryMethod.getParameters(), new Object[1]));
 		verify(em, times(1)).createNamedQuery(queryMethod.getNamedCountQueryName(), Long.class);

src/main/antora/modules/ROOT/pages/jpa/query-methods.adoc

@@ -376,6 +376,8 @@ 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`.
 
 .Declare a QueryRewriter using `@Query`
 ====