Introduce QueryRewriter.

Allow a QueryRewriter to be applied to any query crafted using @query via an additional @QueryRewriter annotation. Also supported directly inside @query.

See #2162.

Author: gregturn

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

@@ -29,7 +29,7 @@
  * @author Oliver Gierke
  * @author Thomas Darimont
  * @author Christoph Strobl
- *
+ * @author Greg Turnquist
  * @see Modifying
  */
 @Retention(RetentionPolicy.RUNTIME)
@@ -45,7 +45,8 @@
 
 	/**
 	 * Defines a special count query that shall be used for pagination queries to lookup the total number of elements for
-	 * a page. If none is configured we will derive the count query from the original query or {@link #countProjection()} query if any.
+	 * a page. If none is configured we will derive the count query from the original query or {@link #countProjection()}
+	 * query if any.
 	 */
 	String countQuery() default "";
 
@@ -70,11 +71,19 @@
 	String name() default "";
 
 	/**
-	 * Returns the name of the {@link jakarta.persistence.NamedQuery} to be used to execute count queries when pagination is
-	 * used. Will default to the named query name configured suffixed by {@code .count}.
+	 * Returns the name of the {@link jakarta.persistence.NamedQuery} to be used to execute count queries when pagination
+	 * is used. Will default to the named query name configured suffixed by {@code .count}.
 	 *
 	 * @see #name()
 	 * @return
 	 */
 	String countName() default "";
+
+	/**
+	 * Define the {@link QueryRewriter} bean that should be applied to this query after the query is full assembled.
+	 *
+	 * @return
+	 * @since 3.0
+	 */
+	Class<? extends QueryRewriter> queryRewriter() default QueryRewriter.NoopQueryRewriter.class;
 }

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

@@ -0,0 +1,64 @@
+/*
+ * Copyright 2008-2022 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.data.jpa.repository;
+
+import org.springframework.data.domain.Pageable;
+import org.springframework.data.domain.Sort;
+
+/**
+ * Callback to rewrite a query right before it's handed to the EntityManager.
+ *
+ * @author Greg Turnquist
+ * @since 3.0
+ */
+@FunctionalInterface
+public interface QueryRewriter {
+
+	/**
+	 * The assembled query and current {@link Sort} settings are offered. This is the query right before it's handed to
+	 * the EntityManager, so everything that Spring Data and tools intends to do has been done. The user is able to make
+	 * any last minute changes.<br/>
+	 * <br/>
+	 * WARNING: No checks are performed before the transformed query is passed to the EntityManager.
+	 * 
+	 * @param query - the assembled generated query, right before it's handed over to the EntityManager.
+	 * @param sort - current {@link Sort} settings provided by the method, or {@link Sort#unsorted()}} if there are none.
+	 * @return alter the query however you like.
+	 */
+	String rewrite(String query, Sort sort);
+
+	/**
+	 * This alternative is used to handle {@link Pageable}-based methods.
+	 * 
+	 * @param query - the assembled generated query, right before it's handed over to the EntityManager.
+	 * @param pageRequest
+	 * @return
+	 */
+	default String rewrite(String query, Pageable pageRequest) {
+		return rewrite(query, pageRequest.getSort());
+	}
+
+	/**
+	 * A {@link QueryRewriter} that doesn't change the query.
+	 */
+	public class NoopQueryRewriter implements QueryRewriter {
+
+		@Override
+		public String rewrite(String query, Sort sort) {
+			return query;
+		}
+	}
+}

spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/cdi/JpaRepositoryBean.java

@@ -15,18 +15,21 @@
  */
 package org.springframework.data.jpa.repository.cdi;
 
-import java.lang.annotation.Annotation;
-import java.util.Optional;
-import java.util.Set;
-
 import jakarta.enterprise.context.spi.CreationalContext;
 import jakarta.enterprise.inject.spi.Bean;
 import jakarta.enterprise.inject.spi.BeanManager;
 import jakarta.persistence.EntityManager;
 
+import java.lang.annotation.Annotation;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.Supplier;
+
+import org.springframework.data.jpa.repository.query.QueryRewriterProvider;
 import org.springframework.data.jpa.repository.support.JpaRepositoryFactory;
 import org.springframework.data.repository.cdi.CdiRepositoryBean;
 import org.springframework.data.repository.config.CustomRepositoryImplementationDetector;
+import org.springframework.data.repository.core.support.RepositoryFactorySupport;
 import org.springframework.util.Assert;
 
 /**
@@ -41,6 +44,7 @@
 class JpaRepositoryBean<T> extends CdiRepositoryBean<T> {
 
 	private final Bean<EntityManager> entityManagerBean;
+	private final QueryRewriterProvider queryRewriterProvider;
 
 	/**
 	 * Constructs a {@link JpaRepositoryBean}.
@@ -58,13 +62,20 @@ class JpaRepositoryBean<T> extends CdiRepositoryBean<T> {
 
 		Assert.notNull(entityManagerBean, "EntityManager bean must not be null!");
 		this.entityManagerBean = entityManagerBean;
+		this.queryRewriterProvider = new QueryRewriterBeanManagerProvider(beanManager);
 	}
 
 	@Override
 	protected T create(CreationalContext<T> creationalContext, Class<T> repositoryType) {
 
 		EntityManager entityManager = getDependencyInstance(entityManagerBean, EntityManager.class);
 
-		return create(() -> new JpaRepositoryFactory(entityManager), repositoryType);
+		Supplier<RepositoryFactorySupport> repositoryFactorySupportSupplier = () -> {
+			JpaRepositoryFactory jpaRepositoryFactory = new JpaRepositoryFactory(entityManager);
+			jpaRepositoryFactory.setQueryRewriterProvider(queryRewriterProvider);
+			return jpaRepositoryFactory;
+		};
+
+		return create(repositoryFactorySupportSupplier, repositoryType);
 	}
 }

spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/cdi/QueryRewriterBeanManagerProvider.java

@@ -0,0 +1,55 @@
+/*
+ * Copyright 2008-2022 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.data.jpa.repository.cdi;
+
+import jakarta.enterprise.context.spi.CreationalContext;
+import jakarta.enterprise.inject.spi.Bean;
+import jakarta.enterprise.inject.spi.BeanManager;
+
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.springframework.data.jpa.repository.QueryRewriter;
+import org.springframework.data.jpa.repository.query.QueryRewriterProvider;
+
+/**
+ * A {@link BeanManager}-based {@link QueryRewriterProvider}.
+ *
+ * @author Greg Turnquist
+ * @since 3.0
+ */
+public class QueryRewriterBeanManagerProvider extends QueryRewriterProvider {
+
+	private static final Log LOGGER = LogFactory.getLog(QueryRewriterBeanManagerProvider.class);
+
+	private final BeanManager beanManager;
+
+	public QueryRewriterBeanManagerProvider(BeanManager beanManager) {
+		this.beanManager = beanManager;
+	}
+
+	@Override
+	protected QueryRewriter extractQueryRewriterBean(Class<? extends QueryRewriter> queryRewriter) {
+
+		try {
+			Bean<QueryRewriter> bean = (Bean<QueryRewriter>) beanManager.getBeans(queryRewriter).iterator().next();
+			CreationalContext<QueryRewriter> context = beanManager.createCreationalContext(bean);
+			return (QueryRewriter) beanManager.getReference(bean, queryRewriter, context);
+		} catch (Exception e) {
+			LOGGER.error(e.toString());
+			return null;
+		}
+	}
+}

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

@@ -18,6 +18,13 @@
 import jakarta.persistence.EntityManager;
 import jakarta.persistence.Query;
 
+import java.util.function.Supplier;
+
+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.repository.QueryRewriter;
 import org.springframework.data.repository.query.QueryMethodEvaluationContextProvider;
 import org.springframework.data.repository.query.ResultProcessor;
 import org.springframework.data.repository.query.ReturnedType;
@@ -35,14 +42,18 @@
  * @author David Madden
  * @author Mark Paluch
  * @author Diego Krupitza
+ * @author Greg Turnquist
  */
 abstract class AbstractStringBasedJpaQuery extends AbstractJpaQuery {
 
+	private static final Log LOGGER = LogFactory.getLog(AbstractStringBasedJpaQuery.class);
+
 	private final DeclaredQuery query;
 	private final DeclaredQuery countQuery;
 	private final QueryMethodEvaluationContextProvider evaluationContextProvider;
 	private final SpelExpressionParser parser;
 	private final QueryParameterSetter.QueryMetadataCache metadataCache = new QueryParameterSetter.QueryMetadataCache();
+	private final Supplier<QueryRewriter> queryRewriterSupplier;
 
 	/**
 	 * Creates a new {@link AbstractStringBasedJpaQuery} from the given {@link JpaQueryMethod}, {@link EntityManager} and
@@ -57,7 +68,7 @@ abstract class AbstractStringBasedJpaQuery extends AbstractJpaQuery {
 	 */
 	public AbstractStringBasedJpaQuery(JpaQueryMethod method, EntityManager em, String queryString,
 			@Nullable String countQueryString, QueryMethodEvaluationContextProvider evaluationContextProvider,
-			SpelExpressionParser parser) {
+			SpelExpressionParser parser, QueryRewriterProvider queryRewriterProvider) {
 
 		super(method, em);
 
@@ -74,6 +85,7 @@ public AbstractStringBasedJpaQuery(JpaQueryMethod method, EntityManager em, Stri
 				method.isNativeQuery());
 
 		this.parser = parser;
+		this.queryRewriterSupplier = queryRewriterProvider.of(method);
 
 		Assert.isTrue(method.isNativeQuery() || !query.usesJdbcStyleParameters(),
 				"JDBC style parameters (?) are not supported for JPA queries.");
@@ -86,7 +98,8 @@ public Query doCreateQuery(JpaParametersParameterAccessor accessor) {
 				.applySorting(accessor.getSort(), query.getAlias());
 		ResultProcessor processor = getQueryMethod().getResultProcessor().withDynamicProjection(accessor);
 
-		Query query = createJpaQuery(sortedQueryString, processor.getReturnedType());
+		Query query = createJpaQuery(sortedQueryString, accessor.getSort(), accessor.getPageable(),
+				processor.getReturnedType());
 
 		QueryParameterSetter.QueryMetadata metadata = metadataCache.getMetadata(sortedQueryString, query);
 
@@ -137,18 +150,41 @@ public DeclaredQuery getCountQuery() {
 	 * Creates an appropriate JPA query from an {@link EntityManager} according to the current {@link AbstractJpaQuery}
 	 * type.
 	 */
-	protected Query createJpaQuery(String queryString, ReturnedType returnedType) {
+	protected Query createJpaQuery(String queryString, Sort sort, @Nullable Pageable pageable,
+			ReturnedType returnedType) {
 
 		EntityManager em = getEntityManager();
 
 		if (this.query.hasConstructorExpression() || this.query.isDefaultProjection()) {
-			return em.createQuery(queryString);
+			return em.createQuery(potentiallyRewriteQuery(queryString, sort, pageable));
 		}
 
 		Class<?> typeToRead = getTypeToRead(returnedType);
 
 		return typeToRead == null //
-				? em.createQuery(queryString) //
-				: em.createQuery(queryString, typeToRead);
+				? em.createQuery(potentiallyRewriteQuery(queryString, sort, pageable)) //
+				: em.createQuery(potentiallyRewriteQuery(queryString, sort, pageable), typeToRead);
+	}
+
+	/**
+	 * Use the {@link QueryRewriter}, potentially rewrite the query, using relevant {@link Sort} and {@link Pageable}
+	 * information.
+	 * 
+	 * @param originalQuery
+	 * @param sort
+	 * @param pageable
+	 * @return
+	 */
+	protected String potentiallyRewriteQuery(String originalQuery, Sort sort, @Nullable Pageable pageable) {
+
+		QueryRewriter queryRewriter = this.queryRewriterSupplier.get();
+
+		if (queryRewriter == null) {
+			return originalQuery;
+		}
+
+		return pageable != null && pageable.isPaged() //
+				? queryRewriter.rewrite(originalQuery, pageable) //
+				: queryRewriter.rewrite(originalQuery, sort);
 	}
 }

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

@@ -45,12 +45,14 @@ enum JpaQueryFactory {
 	 * @return
 	 */
 	AbstractJpaQuery fromMethodWithQueryString(JpaQueryMethod method, EntityManager em, String queryString,
-			@Nullable String countQueryString,
-			QueryMethodEvaluationContextProvider evaluationContextProvider) {
+			@Nullable String countQueryString, QueryMethodEvaluationContextProvider evaluationContextProvider,
+			QueryRewriterProvider queryRewriterProvider) {
 
 		return method.isNativeQuery()
-				? new NativeJpaQuery(method, em, queryString, countQueryString, evaluationContextProvider, PARSER)
-				: new SimpleJpaQuery(method, em, queryString, countQueryString, evaluationContextProvider, PARSER);
+				? new NativeJpaQuery(method, em, queryString, countQueryString, evaluationContextProvider, PARSER,
+				queryRewriterProvider)
+				: new SimpleJpaQuery(method, em, queryString, countQueryString, evaluationContextProvider, PARSER,
+				queryRewriterProvider);
 	}
 
 	/**