Introduce QueryRewriter.

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

See #2162.

Author: gregturn

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,53 @@
+/*
+ * 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/>
+	 * <p/>
+	 * 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());
+	}
+}

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

@@ -18,6 +18,12 @@
 import jakarta.persistence.EntityManager;
 import jakarta.persistence.Query;
 
+import java.lang.reflect.Constructor;
+import java.lang.reflect.InvocationTargetException;
+
+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,6 +41,7 @@
  * @author David Madden
  * @author Mark Paluch
  * @author Diego Krupitza
+ * @author Greg Turnquist
  */
 abstract class AbstractStringBasedJpaQuery extends AbstractJpaQuery {
 
@@ -43,6 +50,7 @@ abstract class AbstractStringBasedJpaQuery extends AbstractJpaQuery {
 	private final QueryMethodEvaluationContextProvider evaluationContextProvider;
 	private final SpelExpressionParser parser;
 	private final QueryParameterSetter.QueryMetadataCache metadataCache = new QueryParameterSetter.QueryMetadataCache();
+	private final QueryRewriter queryRewriter;
 
 	/**
 	 * Creates a new {@link AbstractStringBasedJpaQuery} from the given {@link JpaQueryMethod}, {@link EntityManager} and
@@ -74,6 +82,7 @@ public AbstractStringBasedJpaQuery(JpaQueryMethod method, EntityManager em, Stri
 				method.isNativeQuery());
 
 		this.parser = parser;
+		this.queryRewriter = findQueryRewriter(method);
 
 		Assert.isTrue(method.isNativeQuery() || !query.usesJdbcStyleParameters(),
 				"JDBC style parameters (?) are not supported for JPA queries.");
@@ -86,7 +95,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,7 +147,8 @@ 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();
 
@@ -148,7 +159,42 @@ protected Query createJpaQuery(String queryString, ReturnedType returnedType) {
 		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);
+	}
+
+	/**
+	 * Useing 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
+	protected QueryRewriter findQueryRewriter(JpaQueryMethod method) {
+
+		Class<? extends QueryRewriter> queryRewriter = method.getQueryRewriter();
+
+		if (queryRewriter == null) {
+			return null;
+		}
+
+		try {
+			return (QueryRewriter) ((Constructor<?>) queryRewriter.getDeclaredConstructor()).newInstance();
+		} catch (InstantiationException | IllegalAccessException | InvocationTargetException | NoSuchMethodException e) {
+			System.out.println(e);
+			return null;
+		}
+	}
+
+	protected String potentiallyRewriteQuery(String originalQuery, Sort sort, @Nullable Pageable pageable) {
+
+		if (this.queryRewriter == null) {
+			return originalQuery;
+		}
+
+		return pageable != null && pageable.isPaged() //
+				? this.queryRewriter.rewrite(originalQuery, pageable) //
+				: this.queryRewriter.rewrite(originalQuery, sort);
 	}
 }

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

@@ -15,6 +15,9 @@
  */
 package org.springframework.data.jpa.repository.query;
 
+import jakarta.persistence.LockModeType;
+import jakarta.persistence.QueryHint;
+
 import java.lang.annotation.Annotation;
 import java.lang.reflect.Method;
 import java.util.Arrays;
@@ -24,9 +27,6 @@
 import java.util.Optional;
 import java.util.Set;
 
-import jakarta.persistence.LockModeType;
-import jakarta.persistence.QueryHint;
-
 import org.springframework.core.annotation.AnnotatedElementUtils;
 import org.springframework.core.annotation.AnnotationUtils;
 import org.springframework.data.jpa.provider.QueryExtractor;
@@ -35,6 +35,8 @@
 import org.springframework.data.jpa.repository.Modifying;
 import org.springframework.data.jpa.repository.Query;
 import org.springframework.data.jpa.repository.QueryHints;
+import org.springframework.data.jpa.repository.QueryRewrite;
+import org.springframework.data.jpa.repository.QueryRewriter;
 import org.springframework.data.projection.ProjectionFactory;
 import org.springframework.data.repository.core.RepositoryMetadata;
 import org.springframework.data.repository.query.Parameter;
@@ -57,6 +59,7 @@
  * @author Mark Paluch
  * @author Сергей Цыпанов
  * @author Réda Housni Alaoui
+ * @author Greg Turnquist
  */
 public class JpaQueryMethod extends QueryMethod {
 
@@ -430,4 +433,14 @@ StoredProcedureAttributes getProcedureAttributes() {
 		return storedProcedureAttributes;
 	}
 
+	/**
+	 * Returns the {@link QueryRewriter} (if there is one). NOTE: This is only used for native (@Query) methods.
+	 *
+	 * @return
+	 * @since 3.0
+	 */
+	@Nullable
+	Class<? extends QueryRewriter> getQueryRewriter() {
+		return getMergedOrDefaultAnnotationValue("value", QueryRewrite.class, Class.class);
+	}
 }

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

@@ -19,6 +19,8 @@
 import jakarta.persistence.Query;
 import jakarta.persistence.Tuple;
 
+import org.springframework.data.domain.Pageable;
+import org.springframework.data.domain.Sort;
 import org.springframework.data.repository.query.Parameters;
 import org.springframework.data.repository.query.QueryMethodEvaluationContextProvider;
 import org.springframework.data.repository.query.RepositoryQuery;
@@ -35,6 +37,7 @@
  * @author Oliver Gierke
  * @author Jens Schauder
  * @author Mark Paluch
+ * @author Greg Turnquist
  */
 final class NativeJpaQuery extends AbstractStringBasedJpaQuery {
 
@@ -60,12 +63,13 @@ public NativeJpaQuery(JpaQueryMethod method, EntityManager em, String queryStrin
 	}
 
 	@Override
-	protected Query createJpaQuery(String queryString, ReturnedType returnedType) {
+	protected Query createJpaQuery(String queryString, Sort sort, Pageable pageable, ReturnedType returnedType) {
 
 		EntityManager em = getEntityManager();
 		Class<?> type = getTypeToQueryFor(returnedType);
 
-		return type == null ? em.createNativeQuery(queryString) : em.createNativeQuery(queryString, type);
+		return type == null ? em.createNativeQuery(potentiallyRewriteQuery(queryString, sort, pageable))
+				: em.createNativeQuery(potentiallyRewriteQuery(queryString, sort, pageable), type);
 	}
 
 	@Nullable

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

@@ -18,16 +18,16 @@
 import static org.mockito.ArgumentMatchers.*;
 import static org.mockito.Mockito.*;
 
-import java.lang.reflect.Method;
-import java.util.Set;
-
 import jakarta.persistence.EntityManager;
 import jakarta.persistence.PersistenceContext;
 import jakarta.persistence.Tuple;
 
+import java.lang.reflect.Method;
+import java.util.Set;
+
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
-
+import org.springframework.data.domain.Sort;
 import org.springframework.data.jpa.domain.sample.Role;
 import org.springframework.data.jpa.domain.sample.User;
 import org.springframework.data.jpa.provider.PersistenceProvider;
@@ -60,10 +60,11 @@ void createsNormalQueryForJpaManagedReturnTypes() throws Exception {
 		when(mock.getMetamodel()).thenReturn(em.getMetamodel());
 
 		JpaQueryMethod method = getMethod("findRolesByEmailAddress", String.class);
-		AbstractStringBasedJpaQuery jpaQuery = new SimpleJpaQuery(method, mock,
-				null, QueryMethodEvaluationContextProvider.DEFAULT, new SpelExpressionParser());
+		AbstractStringBasedJpaQuery jpaQuery = new SimpleJpaQuery(method, mock, null,
+				QueryMethodEvaluationContextProvider.DEFAULT, new SpelExpressionParser());
 
-		jpaQuery.createJpaQuery(method.getAnnotatedQuery(), method.getResultProcessor().getReturnedType());
+		jpaQuery.createJpaQuery(method.getAnnotatedQuery(), Sort.unsorted(), null,
+				method.getResultProcessor().getReturnedType());
 
 		verify(mock, times(1)).createQuery(anyString());
 		verify(mock, times(0)).createQuery(anyString(), eq(Tuple.class));

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

@@ -19,21 +19,20 @@
 import static org.mockito.ArgumentMatchers.*;
 import static org.mockito.Mockito.*;
 
-import java.lang.reflect.Method;
-import java.util.List;
-
 import jakarta.persistence.EntityManager;
 import jakarta.persistence.EntityManagerFactory;
 import jakarta.persistence.metamodel.Metamodel;
 
+import java.lang.reflect.Method;
+import java.util.List;
+
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.Mock;
 import org.mockito.junit.jupiter.MockitoExtension;
 import org.mockito.junit.jupiter.MockitoSettings;
 import org.mockito.quality.Strictness;
-
 import org.springframework.data.domain.Page;
 import org.springframework.data.domain.Pageable;
 import org.springframework.data.domain.Sort;