diff --git a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/JpaCodeBlocks.java b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/JpaCodeBlocks.java index 8dba82710..5dacdd7cb 100644 --- a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/JpaCodeBlocks.java +++ b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/JpaCodeBlocks.java @@ -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, + @Nullable String queryRewriterName, AotQuery query, @Nullable String sqlResultSetMapping, + MergedAnnotation 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); } } diff --git a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/JpaRepositoryContributor.java b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/JpaRepositoryContributor.java index 157e77c5e..53ce0f8cc 100644 --- a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/JpaRepositoryContributor.java +++ b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/JpaRepositoryContributor.java @@ -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()); diff --git a/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/aot/JpaRepositoryContributorIntegrationTests.java b/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/aot/JpaRepositoryContributorIntegrationTests.java index aa2f2f0c5..9cbd86109 100644 --- a/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/aot/JpaRepositoryContributorIntegrationTests.java +++ b/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/aot/JpaRepositoryContributorIntegrationTests.java @@ -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 page = fragment.findAndApplyQueryRewriter(kylo.getEmailAddress(), Pageable.unpaged()); + + assertThat(page).isNotEmpty(); + } + void todo() { // dynamic projections: Not implemented diff --git a/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/aot/UserRepository.java b/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/aot/UserRepository.java index de4ae656d..3d5eeb930 100644 --- a/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/aot/UserRepository.java +++ b/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/aot/UserRepository.java @@ -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 { @Query("select u from User u where u.emailAddress = ?1 AND TYPE(u) = ?2") T findByEmailAddress(String emailAddress, Class 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 findAndApplyQueryRewriter(String emailAddress, Pageable pageable); + interface EmailOnly { String getEmailAddress(); } @@ -236,4 +243,16 @@ interface UserRepository extends CrudRepository { 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"); + } + } } diff --git a/src/main/antora/modules/ROOT/pages/jpa/aot.adoc b/src/main/antora/modules/ROOT/pages/jpa/aot.adoc index 145c19c95..031a75f52 100644 --- a/src/main/antora/modules/ROOT/pages/jpa/aot.adoc +++ b/src/main/antora/modules/ROOT/pages/jpa/aot.adoc @@ -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 `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