Spring/spring-data-jpa
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add support for QueryRewriter.

Browse Source 
See #3830
This commit is contained in:
Mark Paluch
2025-04-04 16:06:10 +02:00
parent 9fe7d423ef
commit d3cc7da3d9
5 changed files with 100 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

75
spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/JpaCodeBlocks.java
View File

@@ -35,6 +35,7 @@ import org.springframework.data.domain.Sort;
import org.springframework.data.jpa.repository.Modifying;
import org.springframework.data.jpa.repository.NativeQuery;
import org.springframework.data.jpa.repository.QueryHints;
import org.springframework.data.jpa.repository.QueryRewriter;
import org.springframework.data.jpa.repository.query.DeclaredQuery;
import org.springframework.data.jpa.repository.query.JpaQueryMethod;
import org.springframework.data.jpa.repository.query.ParameterBinding;
@@ -85,6 +86,7 @@ class JpaCodeBlocks {
private @Nullable AotEntityGraph entityGraph;
private @Nullable String sqlResultSetMapping;
private @Nullable Class<?> queryReturnType;
private @Nullable Class<?> queryRewriter = QueryRewriter.IdentityQueryRewriter.class;
private QueryBlockBuilder(AotQueryMethodGenerationContext context, JpaQueryMethod queryMethod) {
this.context = context;
@@ -126,6 +128,11 @@ class JpaCodeBlocks {
return this;
}
public QueryBlockBuilder queryRewriter(@Nullable Class<?> queryRewriter) {
this.queryRewriter = queryRewriter == null ? QueryRewriter.IdentityQueryRewriter.class : queryRewriter;
return this;
}
/**
* Build the query block.
*
@@ -145,12 +152,20 @@ class JpaCodeBlocks {
CodeBlock.Builder builder = CodeBlock.builder();
builder.add("\n");
String queryStringNameVariableName = null;
String queryStringVariableName = null;
String queryRewriterName = null;
if (queries.result() instanceof StringAotQuery && queryRewriter != QueryRewriter.IdentityQueryRewriter.class) {
queryRewriterName = "queryRewriter";
builder.addStatement("$T $L = new $T()", queryRewriter, queryRewriterName, queryRewriter);
}
if (queries != null && queries.result() instanceof StringAotQuery sq) {
queryStringNameVariableName = "%sString".formatted(queryVariableName);
builder.addStatement("$T $L = $S", String.class, queryStringNameVariableName, sq.getQueryString());
queryStringVariableName = "%sString".formatted(queryVariableName);
builder.add(buildQueryString(sq, queryStringVariableName));
}
String countQueryStringNameVariableName = null;
@@ -159,7 +174,7 @@ class JpaCodeBlocks {
if (queryMethod.isPageQuery() && queries.count() instanceof StringAotQuery sq) {
countQueryStringNameVariableName = "count%sString".formatted(StringUtils.capitalize(queryVariableName));
builder.addStatement("$T $L = $S", String.class, countQueryStringNameVariableName, sq.getQueryString());
builder.add(buildQueryString(sq, countQueryStringNameVariableName));
}
String sortParameterName = context.getSortParameterName();
@@ -169,14 +184,14 @@ class JpaCodeBlocks {
if ((StringUtils.hasText(sortParameterName) || StringUtils.hasText(dynamicReturnType))
&& queries.result() instanceof StringAotQuery) {
builder.add(applyRewrite(sortParameterName, dynamicReturnType, queryStringNameVariableName, actualReturnType));
builder.add(applyRewrite(sortParameterName, dynamicReturnType, queryStringVariableName, actualReturnType));
}
if (queries.result().hasExpression() || queries.count().hasExpression()) {
builder.addStatement("class ExpressionMarker{}");
}
builder.add(createQuery(false, queryVariableName, queryStringNameVariableName, queries.result(),
builder.add(createQuery(false, queryVariableName, queryStringVariableName, queryRewriterName, queries.result(),
this.sqlResultSetMapping, this.queryHints, this.entityGraph, this.queryReturnType));
builder.add(applyLimits(queries.result().isExists()));
@@ -187,7 +202,8 @@ class JpaCodeBlocks {
boolean queryHints = this.queryHints.isPresent() && this.queryHints.getBoolean("forCounting");
builder.add(createQuery(true, countQueryVariableName, countQueryStringNameVariableName, queries.count(), null,
builder.add(createQuery(true, countQueryVariableName, countQueryStringNameVariableName, queryRewriterName,
queries.count(), null,
queryHints ? this.queryHints : MergedAnnotation.missing(), null, Long.class));
builder.addStatement("return ($T) $L.getSingleResult()", Long.class, countQueryVariableName);
@@ -199,6 +215,13 @@ class JpaCodeBlocks {
return builder.build();
}
private CodeBlock buildQueryString(StringAotQuery sq, String queryStringVariableName) {
CodeBlock.Builder builder = CodeBlock.builder();
builder.addStatement("$T $L = $S", String.class, queryStringVariableName, sq.getQueryString());
return builder.build();
}
private CodeBlock applyRewrite(@Nullable String sort, @Nullable String dynamicReturnType, String queryString,
Class<?> actualReturnType) {
@@ -268,12 +291,14 @@ class JpaCodeBlocks {
}
private CodeBlock createQuery(boolean count, String queryVariableName, @Nullable String queryStringNameVariableName,
AotQuery query, @Nullable String sqlResultSetMapping, MergedAnnotation<QueryHints> queryHints,
@Nullable String queryRewriterName, AotQuery query, @Nullable String sqlResultSetMapping,
MergedAnnotation<QueryHints> queryHints,
@Nullable AotEntityGraph entityGraph, @Nullable Class<?> queryReturnType) {
Builder builder = CodeBlock.builder();
builder.add(doCreateQuery(count, queryVariableName, queryStringNameVariableName, query, sqlResultSetMapping,
builder.add(doCreateQuery(count, queryVariableName, queryStringNameVariableName, queryRewriterName, query,
sqlResultSetMapping,
queryReturnType));
if (entityGraph != null) {
@@ -306,18 +331,36 @@ class JpaCodeBlocks {
}
private CodeBlock doCreateQuery(boolean count, String queryVariableName,
@Nullable String queryStringNameVariableName, AotQuery query, @Nullable String sqlResultSetMapping,
@Nullable String queryStringName, @Nullable String queryRewriterName, AotQuery query,
@Nullable String sqlResultSetMapping,
@Nullable Class<?> queryReturnType) {
ReturnedType returnedType = context.getReturnedType();
Builder builder = CodeBlock.builder();
String queryStringNameToUse = queryStringName;
if (query instanceof StringAotQuery sq) {
if (StringUtils.hasText(queryRewriterName)) {
queryStringNameToUse = queryStringName + "Rewritten";
if (StringUtils.hasText(context.getPageableParameterName())) {
builder.addStatement("$T $L = $L.rewrite($L, $L)", String.class, queryStringNameToUse, queryRewriterName,
queryStringName, context.getPageableParameterName());
} else if (StringUtils.hasText(context.getSortParameterName())) {
builder.addStatement("$T $L = $L.rewrite($L, $L)", String.class, queryStringNameToUse, queryRewriterName,
queryStringName, context.getSortParameterName());
} else {
builder.addStatement("$T $L = $L.rewrite($L, $T.unsorted())", String.class, queryStringNameToUse,
queryRewriterName, queryStringName, Sort.class);
}
}
if (StringUtils.hasText(sqlResultSetMapping)) {
builder.addStatement("$T $L = this.$L.createNativeQuery($L, $S)", Query.class, queryVariableName,
context.fieldNameOf(EntityManager.class), queryStringNameVariableName, sqlResultSetMapping);
context.fieldNameOf(EntityManager.class), queryStringNameToUse, sqlResultSetMapping);
return builder.build();
}
@@ -327,10 +370,10 @@ class JpaCodeBlocks {
if (queryReturnType != null) {
builder.addStatement("$T $L = this.$L.createNativeQuery($L, $T.class)", Query.class, queryVariableName,
context.fieldNameOf(EntityManager.class), queryStringNameVariableName, queryReturnType);
context.fieldNameOf(EntityManager.class), queryStringNameToUse, queryReturnType);
} else {
builder.addStatement("$T $L = this.$L.createNativeQuery($L)", Query.class, queryVariableName,
context.fieldNameOf(EntityManager.class), queryStringNameVariableName);
context.fieldNameOf(EntityManager.class), queryStringNameToUse);
}
return builder.build();
@@ -339,7 +382,7 @@ class JpaCodeBlocks {
if (sq.hasConstructorExpressionOrDefaultProjection() && !count && returnedType.isProjecting()
&& returnedType.getReturnedType().isInterface()) {
builder.addStatement("$T $L = this.$L.createQuery($L)", Query.class, queryVariableName,
context.fieldNameOf(EntityManager.class), queryStringNameVariableName);
context.fieldNameOf(EntityManager.class), queryStringNameToUse);
} else {
String createQueryMethod = query.isNative() ? "createNativeQuery" : "createQuery";
@@ -347,10 +390,10 @@ class JpaCodeBlocks {
if (!sq.hasConstructorExpressionOrDefaultProjection() && !count && returnedType.isProjecting()
&& returnedType.getReturnedType().isInterface()) {
builder.addStatement("$T $L = this.$L.$L($L, $T.class)", Query.class, queryVariableName,
context.fieldNameOf(EntityManager.class), createQueryMethod, queryStringNameVariableName, Tuple.class);
context.fieldNameOf(EntityManager.class), createQueryMethod, queryStringNameToUse, Tuple.class);
} else {
builder.addStatement("$T $L = this.$L.$L($L)", Query.class, queryVariableName,
context.fieldNameOf(EntityManager.class), createQueryMethod, queryStringNameVariableName);
context.fieldNameOf(EntityManager.class), createQueryMethod, queryStringNameToUse);
}
}

5
spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/JpaRepositoryContributor.java
View File

@@ -81,6 +81,8 @@ public class JpaRepositoryContributor extends RepositoryContributor {
@Override
protected void customizeConstructor(AotRepositoryConstructorBuilder constructorBuilder) {
// TODO: BeanFactoryQueryRewriterProvider if there is a method using QueryRewriters.
constructorBuilder.addParameter("entityManager", EntityManager.class);
constructorBuilder.addParameter("context", RepositoryFactoryBeanSupport.FragmentCreationContext.class);
@@ -149,7 +151,8 @@ public class JpaRepositoryContributor extends RepositoryContributor {
body.add(JpaCodeBlocks.queryBuilder(context, queryMethod).filter(aotQueries)
.queryReturnType(QueriesFactory.getQueryReturnType(aotQueries.result(), returnedType, context))
.nativeQuery(nativeQuery).queryHints(queryHints).entityGraph(aotEntityGraph).build());
.nativeQuery(nativeQuery).queryHints(queryHints).entityGraph(aotEntityGraph)
.queryRewriter(query.isPresent() ? query.getClass("queryRewriter") : null).build());
body.add(
JpaCodeBlocks.executionBuilder(context, queryMethod).modifying(modifying).query(aotQueries.result()).build());

13
spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/aot/JpaRepositoryContributorIntegrationTests.java
View File

@@ -33,6 +33,7 @@ import org.springframework.context.annotation.Configuration;
import org.springframework.data.domain.Limit;
import org.springframework.data.domain.Page;
import org.springframework.data.domain.PageRequest;
import org.springframework.data.domain.Pageable;
import org.springframework.data.domain.Slice;
import org.springframework.data.domain.Sort;
import org.springframework.data.jpa.domain.sample.Role;
@@ -624,6 +625,18 @@ class JpaRepositoryContributorIntegrationTests {
assertThat(result).isInstanceOf(SpecialUser.class);
}
@Test
void shouldApplyQueryRewriter() {
User result = fragment.findAndApplyQueryRewriter(kylo.getEmailAddress());
assertThat(result).isNotNull();
Page<User> page = fragment.findAndApplyQueryRewriter(kylo.getEmailAddress(), Pageable.unpaged());
assertThat(page).isNotEmpty();
}
void todo() {
// dynamic projections: Not implemented

19
spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/aot/UserRepository.java
View File

@@ -32,6 +32,7 @@ import org.springframework.data.jpa.repository.Modifying;
import org.springframework.data.jpa.repository.NativeQuery;
import org.springframework.data.jpa.repository.Query;
import org.springframework.data.jpa.repository.QueryHints;
import org.springframework.data.jpa.repository.QueryRewriter;
import org.springframework.data.repository.CrudRepository;
/**
@@ -229,6 +230,12 @@ interface UserRepository extends CrudRepository<User, Integer> {
@Query("select u from User u where u.emailAddress = ?1 AND TYPE(u) = ?2")
<T extends User> T findByEmailAddress(String emailAddress, Class<T> type);
@Query(value = "select u from PLACEHOLDER u where u.emailAddress = ?1", queryRewriter = MyQueryRewriter.class)
User findAndApplyQueryRewriter(String emailAddress);
@Query(value = "select u from OTHER u where u.emailAddress = ?1", queryRewriter = MyQueryRewriter.class)
Page<User> findAndApplyQueryRewriter(String emailAddress, Pageable pageable);
interface EmailOnly {
String getEmailAddress();
}
@@ -236,4 +243,16 @@ interface UserRepository extends CrudRepository<User, Integer> {
record Names(String firstname, String lastname) {
}
static class MyQueryRewriter implements QueryRewriter {
@Override
public String rewrite(String query, Sort sort) {
return query.replaceAll("PLACEHOLDER", "User");
}
@Override
public String rewrite(String query, Pageable pageRequest) {
return query.replaceAll("OTHER", "User");
}
}
}

7
src/main/antora/modules/ROOT/pages/jpa/aot.adoc
View File

@@ -38,6 +38,9 @@ This optimization moves query method processing from runtime to build-time, whic
The resulting AOT repository fragment follows the naming scheme of `<Repository FQCN>Impl__Aot` and is placed in the same package as the repository interface.
You can find all queries in their String form for generated repository query methods.
NOTE: Consider AOT repository classes an internal optimization.
Do not use them directly in your code as generation and implementation details may change in future releases.
=== Running with AOT Repositories
AOT is a mandatory step to transform a Spring application to a native executable, so it is automatically enabled when running in this mode.
@@ -79,9 +82,9 @@ Mind that using Value Expressions requires expression parsing and contextual inf
* `CrudRepository` and other base interface methods
* Querydsl and Query by Example methods
* Methods whose implementation would be overly complex
** Methods accepting `ScrollPosition (e.g. `Keyset` pagination)
** Methods accepting `ScrollPosition` (e.g. `Keyset` pagination)
** Stored procedure query methods annotated with `@Procedure`
** For now: Dynamic and interface projections
** Dynamic projections
[[aot.repositories.json]]
== Repository Metadata