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/QueryRewriter.java

@@ -26,6 +26,10 @@
  * and tools intends to do has been done. You can customize the query to apply final changes. Rewriting can only make
  * use of already existing contextual data. That is, adding or replacing query text or reuse of bound parameters. Query
  * rewriting must not add additional bindable parameters as these cannot be materialized.
+ * <p>
+ * 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
+ * {@code SelectionQuery}.
  *
  * @author Greg Turnquist
  * @author Mark Paluch
@@ -71,4 +75,5 @@ public String rewrite(String query, Sort sort) {
 			return query;
 		}
 	}
+
 }

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

@@ -151,9 +151,11 @@ protected Query doCreateCountQuery(JpaParametersParameterAccessor accessor) {
 		String queryString = countQuery.get().getQueryString();
 		EntityManager em = getEntityManager();
 
+		String queryStringToUse = potentiallyRewriteQuery(queryString, accessor.getSort(), accessor.getPageable());
+
 		Query query = getQueryMethod().isNativeQuery() //
-				? em.createNativeQuery(queryString) //
-				: em.createQuery(queryString, Long.class);
+				? em.createNativeQuery(queryStringToUse) //
+				: em.createQuery(queryStringToUse, Long.class);
 
 		QueryParameterSetter.QueryMetadata metadata = metadataCache.getMetadata(queryString, query);
 
@@ -184,16 +186,17 @@ protected Query createJpaQuery(String queryString, Sort sort, @Nullable Pageable
 			ReturnedType returnedType) {
 
 		EntityManager em = getEntityManager();
+		String queryToUse = potentiallyRewriteQuery(queryString, sort, pageable);
 
 		if (this.query.hasConstructorExpression() || this.query.isDefaultProjection()) {
-			return em.createQuery(potentiallyRewriteQuery(queryString, sort, pageable));
+			return em.createQuery(queryToUse);
 		}
 
 		Class<?> typeToRead = getTypeToRead(returnedType);
 
 		return typeToRead == null //
-				? em.createQuery(potentiallyRewriteQuery(queryString, sort, pageable)) //
-				: em.createQuery(potentiallyRewriteQuery(queryString, sort, pageable), typeToRead);
+				? em.createQuery(queryToUse) //
+				: em.createQuery(queryToUse, typeToRead);
 	}
 
 	/**

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

@@ -181,11 +181,9 @@ protected RepositoryQuery resolveQuery(JpaQueryMethod method, QueryRewriter quer
 						getCountQuery(method, namedQueries, em), queryRewriter, valueExpressionDelegate);
 			}
 
-			RepositoryQuery query = NamedQuery.lookupFrom(method, em);
+			RepositoryQuery query = NamedQuery.lookupFrom(method, em, queryRewriter);
 
-			return query != null //
-					? query //
-					: NO_QUERY;
+			return query != null ? query : NO_QUERY;
 		}
 
 		@Nullable

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

@@ -22,7 +22,11 @@
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
+
+import org.springframework.data.domain.Pageable;
+import org.springframework.data.domain.Sort;
 import org.springframework.data.jpa.provider.QueryExtractor;
+import org.springframework.data.jpa.repository.QueryRewriter;
 import org.springframework.data.repository.query.Parameters;
 import org.springframework.data.repository.query.QueryCreationException;
 import org.springframework.data.repository.query.RepositoryQuery;
@@ -53,18 +57,20 @@ final class NamedQuery extends AbstractJpaQuery {
 	private final boolean namedCountQueryIsPresent;
 	private final Lazy<DeclaredQuery> declaredQuery;
 	private final QueryParameterSetter.QueryMetadataCache metadataCache;
+	private final QueryRewriter queryRewriter;
 
 	/**
 	 * Creates a new {@link NamedQuery}.
 	 */
-	private NamedQuery(JpaQueryMethod method, EntityManager em) {
+	private NamedQuery(JpaQueryMethod method, EntityManager em, QueryRewriter queryRewriter) {
 
 		super(method, em);
 
 		this.queryName = method.getNamedQueryName();
 		this.countQueryName = method.getNamedCountQueryName();
 		QueryExtractor extractor = method.getQueryExtractor();
 		this.countProjection = method.getCountQueryProjection();
+		this.queryRewriter = queryRewriter;
 
 		Parameters<?, ?> parameters = method.getParameters();
 
@@ -127,9 +133,10 @@ static boolean hasNamedQuery(EntityManager em, String queryName) {
 	 *
 	 * @param method must not be {@literal null}.
 	 * @param em must not be {@literal null}.
+	 * @param queryRewriter must not be {@literal null}.
 	 */
 	@Nullable
-	public static RepositoryQuery lookupFrom(JpaQueryMethod method, EntityManager em) {
+	public static RepositoryQuery lookupFrom(JpaQueryMethod method, EntityManager em, QueryRewriter queryRewriter) {
 
 		String queryName = method.getNamedQueryName();
 
@@ -147,7 +154,7 @@ public static RepositoryQuery lookupFrom(JpaQueryMethod method, EntityManager em
 					method.isNativeQuery() ? "NativeQuery" : "Query"));
 		}
 
-		RepositoryQuery query = new NamedQuery(method, em);
+		RepositoryQuery query = new NamedQuery(method, em, queryRewriter);
 		if (LOG.isDebugEnabled()) {
 			LOG.debug(String.format("Found named query '%s'", queryName));
 		}
@@ -187,6 +194,7 @@ protected TypedQuery<Long> doCreateCountQuery(JpaParametersParameterAccessor acc
 		} else {
 
 			String countQueryString = declaredQuery.get().deriveCountQuery(countProjection).getQueryString();
+			countQueryString = potentiallyRewriteQuery(countQueryString, accessor.getSort(), accessor.getPageable());
 			cacheKey = countQueryString;
 			countQuery = em.createQuery(countQueryString, Long.class);
 		}
@@ -222,4 +230,20 @@ protected Class<?> getTypeToRead(ReturnedType returnedType) {
 				? null //
 				: super.getTypeToRead(returnedType);
 	}
+
+	/**
+	 * Use the {@link QueryRewriter}, potentially rewrite the query, using relevant {@link Sort} and {@link Pageable}
+	 * information.
+	 *
+	 * @param originalQuery
+	 * @param sort
+	 * @param pageable
+	 * @return
+	 */
+	private String potentiallyRewriteQuery(String originalQuery, Sort sort, Pageable pageable) {
+
+		return pageable.isPaged() //
+				? queryRewriter.rewrite(originalQuery, pageable) //
+				: queryRewriter.rewrite(originalQuery, sort);
+	}
 }

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

@@ -19,8 +19,10 @@
 import static org.assertj.core.api.Assertions.entry;
 
 import java.util.HashMap;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -31,6 +33,7 @@
 import org.springframework.context.annotation.Configuration;
 import org.springframework.context.annotation.FilterType;
 import org.springframework.context.annotation.ImportResource;
+import org.springframework.data.domain.Page;
 import org.springframework.data.domain.PageRequest;
 import org.springframework.data.domain.Pageable;
 import org.springframework.data.domain.Sort;
@@ -44,7 +47,7 @@
 import org.springframework.test.context.junit.jupiter.SpringExtension;
 
 /**
- * Unit tests for repository with {@link Query} and {@link QueryRewrite}.
+ * Unit tests for repository with {@link Query} and {@link QueryRewriter}.
  *
  * @author Greg Turnquist
  * @author Krzysztof Krason
@@ -60,10 +63,12 @@ class JpaQueryRewriteIntegrationTests {
 	static final String REWRITTEN_QUERY = "rewritten query";
 	static final String SORT = "sort";
 	static Map<String, String> results = new HashMap<>();
+	static Set<String> queries = new LinkedHashSet<>();
 
 	@BeforeEach
 	void setUp() {
 		results.clear();
+		repository.deleteAll();
 	}
 
 	@Test
@@ -77,15 +82,15 @@ void nativeQueryShouldHandleRewrites() {
 				entry(SORT, Sort.unsorted().toString()));
 	}
 
-	@Test
+	@Test // GH-3801
 	void nonNativeQueryShouldHandleRewrites() {
 
-		repository.findByNonNativeQuery("Matthews");
+		repository.save(new User("D", "A", "foo@bar"));
 
-		assertThat(results).containsExactly( //
-				entry(ORIGINAL_QUERY, "select original_user_alias from User original_user_alias"), //
-				entry(REWRITTEN_QUERY, "select rewritten_user_alias from User rewritten_user_alias"), //
-				entry(SORT, Sort.unsorted().toString()));
+		repository.findByNonNativeQuery("Matthews", PageRequest.of(0, 1));
+
+		assertThat(queries).contains("select original_user_alias from User original_user_alias");
+		assertThat(queries).contains("select count(original_user_alias) from User original_user_alias");
 	}
 
 	@Test
@@ -169,7 +174,7 @@ public interface UserRepositoryWithRewriter
 		List<User> findByNativeQuery(String param);
 
 		@Query(value = "select original_user_alias from User original_user_alias", queryRewriter = TestQueryRewriter.class)
-		List<User> findByNonNativeQuery(String param);
+		Page<User> findByNonNativeQuery(String param, PageRequest pageRequest);
 
 		@Query(value = "select original_user_alias from User original_user_alias", queryRewriter = TestQueryRewriter.class)
 		List<User> findByNonNativeSortedQuery(String param, Sort sort);
@@ -214,6 +219,7 @@ private static String replaceAlias(String query, Sort sort) {
 		results.put(ORIGINAL_QUERY, query);
 		results.put(REWRITTEN_QUERY, rewrittenQuery);
 		results.put(SORT, sort.toString());
+		queries.add(query);
 
 		return rewrittenQuery;
 	}

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

@@ -36,6 +36,7 @@
 import org.springframework.data.domain.Page;
 import org.springframework.data.domain.Pageable;
 import org.springframework.data.jpa.provider.QueryExtractor;
+import org.springframework.data.jpa.repository.QueryRewriter;
 import org.springframework.data.projection.ProjectionFactory;
 import org.springframework.data.projection.SpelAwareProxyProjectionFactory;
 import org.springframework.data.repository.core.RepositoryMetadata;
@@ -88,7 +89,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));
+		assertThatExceptionOfType(QueryCreationException.class)
+				.isThrownBy(() -> NamedQuery.lookupFrom(queryMethod, em, QueryRewriter.IdentityQueryRewriter.INSTANCE));
 	}
 
 	@Test // DATAJPA-142
@@ -100,7 +102,8 @@ 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);
+		NamedQuery query = (NamedQuery) NamedQuery.lookupFrom(queryMethod, em,
+				QueryRewriter.IdentityQueryRewriter.INSTANCE);
 
 		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

@@ -176,8 +176,11 @@ public interface UserRepository extends JpaRepository<User, Long> {
 Sometimes, no matter how many features you try to apply, it seems impossible to get Spring Data JPA to apply every thing
 you'd like to a query before it is sent to the `EntityManager`.
 
-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.
+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`
 ====