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<?> value();
+}

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

@@ -0,0 +1,35 @@
+/*
+ * 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 {
+
+	String rewrite(String query, Sort sort);
+
+	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,11 @@
 import jakarta.persistence.EntityManager;
 import jakarta.persistence.Query;
 
+import java.lang.reflect.Constructor;
+import java.lang.reflect.InvocationTargetException;
+
+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 +40,7 @@
  * @author David Madden
  * @author Mark Paluch
  * @author Diego Krupitza
+ * @author Greg Turnquist
  */
 abstract class AbstractStringBasedJpaQuery extends AbstractJpaQuery {
 
@@ -43,6 +49,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 +81,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.");
@@ -148,7 +156,30 @@ 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)) //
+				: em.createQuery(potentiallyRewriteQuery(queryString), typeToRead);
+	}
+
+	@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) {
+		return this.queryRewriter == null //
+				? originalQuery //
+				: this.queryRewriter.rewrite(originalQuery, Sort.unsorted());
 	}
 }

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;
@@ -430,4 +432,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

@@ -35,6 +35,7 @@
  * @author Oliver Gierke
  * @author Jens Schauder
  * @author Mark Paluch
+ * @author Greg Turnquist
  */
 final class NativeJpaQuery extends AbstractStringBasedJpaQuery {
 
@@ -65,7 +66,8 @@ protected Query createJpaQuery(String queryString, 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))
+				: em.createNativeQuery(potentiallyRewriteQuery(queryString), type);
 	}
 
 	@Nullable

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;