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

pom.xml

@@ -27,7 +27,7 @@
 		<jsqlparser>4.3</jsqlparser>
 		<mysql-connector-java>8.0.23</mysql-connector-java>
 		<postgresql>42.2.19</postgresql>
-		<springdata.commons>3.0.0-SNAPSHOT</springdata.commons>
+		<springdata.commons>3.0.0-gh-2162-SNAPSHOT</springdata.commons>
 		<vavr>0.10.3</vavr>
 
 		<hibernate.groupId>org.hibernate</hibernate.groupId>

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

@@ -29,7 +29,6 @@
  * @author Oliver Gierke
  * @author Thomas Darimont
  * @author Christoph Strobl
- *
  * @see Modifying
  */
 @Retention(RetentionPolicy.RUNTIME)
@@ -45,7 +44,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 +70,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 "";
+
+	/**
+	 * Returns the class of 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.NoopRewriter.class;
 }

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

@@ -0,0 +1,39 @@
+/*
+ * 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 java.lang.annotation.Documented;
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * Annotation to indicate source of {@link QueryRewriter}.
+ *
+ * @author Greg Turnquist
+ * @since 3.0
+ */
+@Retention(RetentionPolicy.RUNTIME)
+@Target({ ElementType.METHOD, ElementType.ANNOTATION_TYPE })
+@Documented
+public @interface QueryRewrite {
+
+	/**
+	 * Define the {@link QueryRewriter} to callback.
+	 */
+	Class<? extends QueryRewriter> value();
+}

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 via {@link QueryRewrite} on a method with {@link Query}.
+ *
+ * @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 this 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()}} is 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 NoopRewriter implements QueryRewriter {
+
+		@Override
+		public String rewrite(String query, Sort sort) {
+			return query;
+		}
+	}
+}

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

@@ -18,6 +18,15 @@
 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.beans.BeansException;
+import org.springframework.beans.factory.BeanFactory;
+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 +44,19 @@
  * @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;
+	private final BeanFactory beanFactory;
 
 	/**
 	 * Creates a new {@link AbstractStringBasedJpaQuery} from the given {@link JpaQueryMethod}, {@link EntityManager} and
@@ -57,7 +71,7 @@ abstract class AbstractStringBasedJpaQuery extends AbstractJpaQuery {
 	 */
 	public AbstractStringBasedJpaQuery(JpaQueryMethod method, EntityManager em, String queryString,
 			@Nullable String countQueryString, QueryMethodEvaluationContextProvider evaluationContextProvider,
-			SpelExpressionParser parser) {
+			SpelExpressionParser parser, BeanFactory beanFactory) {
 
 		super(method, em);
 
@@ -74,6 +88,8 @@ public AbstractStringBasedJpaQuery(JpaQueryMethod method, EntityManager em, Stri
 				method.isNativeQuery());
 
 		this.parser = parser;
+		this.beanFactory = beanFactory;
+		this.queryRewriterSupplier = () -> findQueryRewriter(method);
 
 		Assert.isTrue(method.isNativeQuery() || !query.usesJdbcStyleParameters(),
 				"JDBC style parameters (?) are not supported for JPA queries.");
@@ -86,7 +102,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 +154,65 @@ 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);
+	}
+
+	/**
+	 * Using the {@link org.springframework.data.jpa.repository.QueryRewrite} annotation, look for a {@link QueryRewriter}
+	 * and instantiate one.
+	 * 
+	 * @param method - {@link JpaQueryMethod} that has the annotation details
+	 * @return a {@link QueryRewriter for the method or {@code null}
+	 */
+	@Nullable
+	private QueryRewriter findQueryRewriter(JpaQueryMethod method) {
+
+		Class<? extends QueryRewriter> queryRewriter = method.getQueryRewriter();
+
+		if (queryRewriter == null) {
+			return null;
+		}
+
+		try {
+			return beanFactory.getBean(queryRewriter);
+		} catch (BeansException e) {
+			LOGGER.error(e.toString());
+			return null;
+		}
+	}
+
+	/**
+	 * 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

@@ -17,6 +17,7 @@
 
 import jakarta.persistence.EntityManager;
 
+import org.springframework.beans.factory.BeanFactory;
 import org.springframework.data.repository.query.QueryMethodEvaluationContextProvider;
 import org.springframework.data.repository.query.RepositoryQuery;
 import org.springframework.expression.spel.standard.SpelExpressionParser;
@@ -45,12 +46,12 @@ enum JpaQueryFactory {
 	 * @return
 	 */
 	AbstractJpaQuery fromMethodWithQueryString(JpaQueryMethod method, EntityManager em, String queryString,
-			@Nullable String countQueryString,
-			QueryMethodEvaluationContextProvider evaluationContextProvider) {
+			@Nullable String countQueryString, QueryMethodEvaluationContextProvider evaluationContextProvider,
+			BeanFactory beanFactory) {
 
 		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, beanFactory)
+				: new SimpleJpaQuery(method, em, queryString, countQueryString, evaluationContextProvider, PARSER, beanFactory);
 	}
 
 	/**