Skip to content

Commit b509828

Browse files
yeomin4242yeomin4242
committed
Limit AOT query rewriter support to declared beans.
Resolve declared query rewriters through BeanFactory during generated fragment construction and store them in fields. Remove the changes to AotRepositoryFragmentSupport and the QueryRewriterProvider SPI. Preserve direct no-args construction where supported and add regression coverage for bean-backed rewriters with constructors that declare checked exceptions. Closes spring-projects#3984 Signed-off-by: YeongJae Min <whereismysejong@naver.com>
1 parent 332d8c3 commit b509828

14 files changed

Lines changed: 155 additions & 203 deletions

‎spring-data-jpa/src/jmh/java/org/springframework/data/jpa/benchmark/AotRepositoryQueryMethodBenchmarks.java‎

Lines changed: 9 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,6 @@
3838
import org.openjdk.jmh.annotations.Warmup;
3939

4040
import org.springframework.aot.test.generate.TestGenerationContext;
41-
import org.springframework.beans.factory.BeanFactory;
4241
import org.springframework.beans.factory.support.BeanDefinitionRegistry;
4342
import org.springframework.beans.factory.support.DefaultBeanNameGenerator;
4443
import org.springframework.beans.factory.support.DefaultListableBeanFactory;
@@ -50,7 +49,6 @@
5049
import org.springframework.data.jpa.benchmark.model.Person;
5150
import org.springframework.data.jpa.benchmark.model.Profile;
5251
import org.springframework.data.jpa.benchmark.repository.PersonRepository;
53-
import org.springframework.data.jpa.repository.aot.AotRepositoryFragmentSupport;
5452
import org.springframework.data.jpa.repository.aot.JpaRepositoryContributor;
5553
import org.springframework.data.jpa.repository.aot.TestJpaAotRepositoryContext;
5654
import org.springframework.data.jpa.repository.config.EnableJpaRepositories;
@@ -71,7 +69,6 @@
7169
/**
7270
* @author Christoph Strobl
7371
* @author Mark Paluch
74-
* @author YeongJae Min
7572
*/
7673
@Testable
7774
@Fork(1)
@@ -95,7 +92,7 @@ public static class BenchmarkParameters {
9592
new StandardEnvironment(), Mockito.mock(BeanDefinitionRegistry.class), DefaultBeanNameGenerator.INSTANCE);
9693

9794
repositoryContext = new TestJpaAotRepositoryContext<>(new DefaultListableBeanFactory(), PersonRepository.class,
98-
RepositoryComposition.empty(), configurationSource);
95+
null, configurationSource);
9996
}
10097

10198
EntityManager entityManager;
@@ -135,20 +132,21 @@ public void doSetup() {
135132
new JpaRepositoryContributor(repositoryContext, entityManager.getEntityManagerFactory())
136133
.contribute(generationContext);
137134

138-
generationContext.writeGeneratedContent();
139-
140135
TestCompiler.forSystem().withCompilerOptions("-parameters").with(generationContext).compile(compiled -> {
141-
this.aot = compiled.getAllCompiledClasses().stream()
142-
.filter(AotRepositoryFragmentSupport.class::isAssignableFrom).findFirst()
143-
.orElseThrow(() -> new IllegalStateException("No generated AOT repository fragment found"));
136+
137+
try {
138+
this.aot = compiled.getClassLoader().loadClass(PersonRepository.class.getName() + "Impl__Aot");
139+
} catch (Exception e) {
140+
throw new RuntimeException(e);
141+
}
144142
});
145143
}
146144

147145
try {
148146
RepositoryFactoryBeanSupport.FragmentCreationContext creationContext = getCreationContext(repositoryContext);
149147
fragments = RepositoryComposition.RepositoryFragments
150-
.just(aot.getConstructor(EntityManager.class, RepositoryFactoryBeanSupport.FragmentCreationContext.class,
151-
BeanFactory.class).newInstance(entityManager, creationContext, new DefaultListableBeanFactory()));
148+
.just(aot.getConstructor(EntityManager.class, RepositoryFactoryBeanSupport.FragmentCreationContext.class)
149+
.newInstance(entityManager, creationContext));
152150

153151
this.repositoryProxy = createRepository(fragments);
154152
} catch (Exception e) {

‎spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/AotRepositoryFragmentSupport.java‎

Lines changed: 0 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -24,13 +24,11 @@
2424
import java.util.Collection;
2525
import java.util.List;
2626
import java.util.Optional;
27-
import java.util.function.Supplier;
2827
import java.util.function.UnaryOperator;
2928
import java.util.stream.Stream;
3029

3130
import org.jspecify.annotations.Nullable;
3231

33-
import org.springframework.beans.factory.BeanFactory;
3432
import org.springframework.core.CollectionFactory;
3533
import org.springframework.core.convert.ConversionService;
3634
import org.springframework.core.convert.support.ConfigurableConversionService;
@@ -39,14 +37,11 @@
3937
import org.springframework.data.domain.Sort;
4038
import org.springframework.data.expression.ValueEvaluationContextProvider;
4139
import org.springframework.data.expression.ValueExpression;
42-
import org.springframework.data.jpa.repository.QueryRewriter;
43-
import org.springframework.data.jpa.repository.query.BeanFactoryQueryRewriterProvider;
4440
import org.springframework.data.jpa.repository.query.DeclaredQuery;
4541
import org.springframework.data.jpa.repository.query.JpaParameters;
4642
import org.springframework.data.jpa.repository.query.JpaResultConverters;
4743
import org.springframework.data.jpa.repository.query.QueryEnhancer;
4844
import org.springframework.data.jpa.repository.query.QueryEnhancerSelector;
49-
import org.springframework.data.jpa.repository.query.QueryRewriterProvider;
5045
import org.springframework.data.jpa.util.TupleBackedMap;
5146
import org.springframework.data.projection.ProjectionFactory;
5247
import org.springframework.data.repository.core.RepositoryMetadata;
@@ -61,7 +56,6 @@
6156
* Support class for JPA AOT repository fragments.
6257
*
6358
* @author Mark Paluch
64-
* @author YeongJae Min
6559
* @since 4.0
6660
*/
6761
public class AotRepositoryFragmentSupport {
@@ -91,33 +85,18 @@ public class AotRepositoryFragmentSupport {
9185

9286
private final Lazy<ConcurrentLruCache<Method, ValueEvaluationContextProvider>> contextProviders;
9387

94-
private final QueryRewriterProvider queryRewriterProvider;
95-
9688
protected AotRepositoryFragmentSupport(QueryEnhancerSelector selector,
9789
RepositoryFactoryBeanSupport.FragmentCreationContext context) {
9890
this(selector, context.getRepositoryMetadata(), context.getValueExpressionDelegate(),
9991
context.getProjectionFactory());
10092
}
10193

102-
protected AotRepositoryFragmentSupport(QueryEnhancerSelector selector,
103-
RepositoryFactoryBeanSupport.FragmentCreationContext context, BeanFactory beanFactory) {
104-
this(selector, context.getRepositoryMetadata(), context.getValueExpressionDelegate(), context.getProjectionFactory(),
105-
new BeanFactoryQueryRewriterProvider(beanFactory));
106-
}
107-
10894
protected AotRepositoryFragmentSupport(QueryEnhancerSelector selector, RepositoryMetadata repositoryMetadata,
10995
ValueExpressionDelegate valueExpressions, ProjectionFactory projectionFactory) {
110-
this(selector, repositoryMetadata, valueExpressions, projectionFactory, QueryRewriterProvider.simple());
111-
}
112-
113-
private AotRepositoryFragmentSupport(QueryEnhancerSelector selector, RepositoryMetadata repositoryMetadata,
114-
ValueExpressionDelegate valueExpressions, ProjectionFactory projectionFactory,
115-
QueryRewriterProvider queryRewriterProvider) {
11696

11797
this.repositoryMetadata = repositoryMetadata;
11898
this.valueExpressions = valueExpressions;
11999
this.projectionFactory = projectionFactory;
120-
this.queryRewriterProvider = queryRewriterProvider;
121100
this.enhancers = Lazy.of(() -> new ConcurrentLruCache<>(32, query -> selector.select(query).create(query)));
122101
this.expressions = Lazy.of(() -> new ConcurrentLruCache<>(32, valueExpressions::parse));
123102
this.contextProviders = Lazy.of(() -> new ConcurrentLruCache<>(32, it -> valueExpressions
@@ -139,15 +118,6 @@ protected String rewriteQuery(DeclaredQuery query, Sort sort, Class<?> returnedT
139118
ReturnedType.of(returnedType, repositoryMetadata.getDomainType(), projectionFactory)));
140119
}
141120

142-
protected QueryRewriter getQueryRewriter(Class<? extends QueryRewriter> queryRewriter) {
143-
return this.queryRewriterProvider.getQueryRewriter(queryRewriter);
144-
}
145-
146-
protected QueryRewriter getQueryRewriter(Class<? extends QueryRewriter> queryRewriter,
147-
Supplier<QueryRewriter> fallback) {
148-
return this.queryRewriterProvider.getQueryRewriter(queryRewriter, fallback);
149-
}
150-
151121
/**
152122
* Evaluate a Value Expression.
153123
*

‎spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/JpaCodeBlocks.java‎

Lines changed: 4 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -21,9 +21,6 @@
2121
import jakarta.persistence.QueryHint;
2222
import jakarta.persistence.Tuple;
2323

24-
import java.lang.reflect.Constructor;
25-
import java.lang.reflect.Modifier;
26-
import java.util.Arrays;
2724
import java.util.Collection;
2825
import java.util.List;
2926
import java.util.function.LongSupplier;
@@ -45,7 +42,6 @@
4542
import org.springframework.data.jpa.repository.Modifying;
4643
import org.springframework.data.jpa.repository.NativeQuery;
4744
import org.springframework.data.jpa.repository.QueryHints;
48-
import org.springframework.data.jpa.repository.QueryRewriter;
4945
import org.springframework.data.jpa.repository.query.DeclaredQuery;
5046
import org.springframework.data.jpa.repository.query.JpaQueryMethod;
5147
import org.springframework.data.jpa.repository.query.ParameterBinding;
@@ -105,7 +101,7 @@ static class QueryBlockBuilder {
105101
private @Nullable AotEntityGraph entityGraph;
106102
private @Nullable String sqlResultSetMapping;
107103
private @Nullable Class<?> queryReturnType;
108-
private @Nullable Class<?> queryRewriter = QueryRewriter.IdentityQueryRewriter.class;
104+
private @Nullable String queryRewriterField;
109105

110106
private QueryBlockBuilder(AotQueryMethodGenerationContext context, JpaQueryMethod queryMethod) {
111107

@@ -156,8 +152,8 @@ public QueryBlockBuilder queryReturnType(@Nullable Class<?> queryReturnType) {
156152
return this;
157153
}
158154

159-
public QueryBlockBuilder queryRewriter(@Nullable Class<?> queryRewriter) {
160-
this.queryRewriter = queryRewriter == null ? QueryRewriter.IdentityQueryRewriter.class : queryRewriter;
155+
public QueryBlockBuilder queryRewriterField(@Nullable String queryRewriterField) {
156+
this.queryRewriterField = queryRewriterField;
161157
return this;
162158
}
163159

@@ -181,19 +177,7 @@ public CodeBlock build() {
181177
CodeBlock.Builder builder = CodeBlock.builder();
182178

183179
String queryStringVariableName = null;
184-
String queryRewriterName = null;
185-
186-
if (queries.result() instanceof StringAotQuery && queryRewriter != QueryRewriter.IdentityQueryRewriter.class) {
187-
188-
queryRewriterName = context.localVariable("queryRewriter");
189-
if (hasAccessibleNoArgsConstructor(queryRewriter)) {
190-
builder.addStatement("$T $L = getQueryRewriter($T.class, () -> new $T())", QueryRewriter.class,
191-
queryRewriterName, queryRewriter, queryRewriter);
192-
} else {
193-
builder.addStatement("$T $L = getQueryRewriter($T.class)", QueryRewriter.class, queryRewriterName,
194-
queryRewriter);
195-
}
196-
}
180+
String queryRewriterName = queryRewriterField == null ? null : "this." + queryRewriterField;
197181

198182
if (queries.result() instanceof StringAotQuery sq) {
199183

@@ -261,35 +245,6 @@ private CodeBlock buildQueryString(StringAotQuery sq, String queryStringVariable
261245
return builder.build();
262246
}
263247

264-
private boolean hasAccessibleNoArgsConstructor(Class<?> type) {
265-
266-
try {
267-
Constructor<?> constructor = type.getDeclaredConstructor();
268-
boolean samePackage = type.getPackageName()
269-
.equals(context.getRepositoryInformation().getRepositoryInterface().getPackageName());
270-
271-
return isAccessible(type, samePackage) && isAccessible(constructor, samePackage);
272-
} catch (NoSuchMethodException e) {
273-
return false;
274-
}
275-
}
276-
277-
private static boolean isAccessible(Class<?> type, boolean samePackage) {
278-
279-
if (Modifier.isPrivate(type.getModifiers()) || (!Modifier.isPublic(type.getModifiers()) && !samePackage)) {
280-
return false;
281-
}
282-
283-
Class<?> enclosingClass = type.getEnclosingClass();
284-
return enclosingClass == null || isAccessible(enclosingClass, samePackage);
285-
}
286-
287-
private static boolean isAccessible(Constructor<?> constructor, boolean samePackage) {
288-
289-
return Modifier.isPublic(constructor.getModifiers())
290-
|| (!Modifier.isPrivate(constructor.getModifiers()) && samePackage);
291-
}
292-
293248
private CodeBlock applyRewrite(@Nullable String sort, @Nullable String dynamicReturnType, boolean isProjecting,
294249
String queryString) {
295250

‎spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/JpaRepositoryContributor.java‎

Lines changed: 67 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,11 @@
2121
import jakarta.persistence.PersistenceUnitUtil;
2222
import jakarta.persistence.metamodel.Metamodel;
2323

24+
import java.lang.reflect.Constructor;
2425
import java.lang.reflect.Method;
26+
import java.lang.reflect.Modifier;
27+
import java.util.Arrays;
28+
import java.util.LinkedHashMap;
2529
import java.util.Map;
2630
import java.util.Optional;
2731
import java.util.function.Function;
@@ -42,6 +46,7 @@
4246
import org.springframework.data.jpa.repository.NativeQuery;
4347
import org.springframework.data.jpa.repository.Query;
4448
import org.springframework.data.jpa.repository.QueryHints;
49+
import org.springframework.data.jpa.repository.QueryRewriter;
4550
import org.springframework.data.jpa.repository.query.JpaEntityMetadata;
4651
import org.springframework.data.jpa.repository.query.JpaParameters;
4752
import org.springframework.data.jpa.repository.query.JpaQueryMethod;
@@ -95,6 +100,7 @@ public class JpaRepositoryContributor extends RepositoryContributor {
95100
private final PersistenceProvider persistenceProvider;
96101
private final QueriesFactory queriesFactory;
97102
private final EntityGraphLookup entityGraphLookup;
103+
private final Map<Class<? extends QueryRewriter>, String> queryRewriterFields = new LinkedHashMap<>();
98104

99105
public JpaRepositoryContributor(AotRepositoryContext repositoryContext) {
100106
this(repositoryContext, AotEntityManagerFactoryCreator.from(repositoryContext).getEntityManagerFactory());
@@ -122,6 +128,8 @@ public EntityManagerFactory getEntityManagerFactory() {
122128
protected void customizeClass(AotRepositoryClassBuilder classBuilder) {
123129
classBuilder.customize(builder -> {
124130
builder.superclass(TypeName.get(AotRepositoryFragmentSupport.class));
131+
queryRewriterFields.values().forEach(name -> builder.addField(QueryRewriter.class, name,
132+
javax.lang.model.element.Modifier.PRIVATE, javax.lang.model.element.Modifier.FINAL));
125133

126134
// suppress compiler warnings for the many @Deprecated(since = "4.0", forRemoval = true) instances
127135
if (JpaDetector.isJpa4Present()) {
@@ -145,21 +153,66 @@ protected void customizeConstructor(AotRepositoryConstructorBuilder constructorB
145153
});
146154

147155
constructorBuilder.addParameter("context", RepositoryFactoryBeanSupport.FragmentCreationContext.class);
148-
constructorBuilder.addParameter("beanFactory", BeanFactory.class,
149-
customizer -> customizer.origin(context -> context.getBeanFactory()));
156+
157+
if (getRepositoryInformation().getQueryMethods().stream().anyMatch(this::hasQueryRewriter)) {
158+
constructorBuilder.addParameter("beanFactory", BeanFactory.class,
159+
customizer -> customizer.origin(context -> context.getBeanFactory()));
160+
}
150161

151162
Optional<Class<QueryEnhancerSelector>> queryEnhancerSelector = getQueryEnhancerSelectorClass();
152163

153164
constructorBuilder.customize(builder -> {
154165

155166
if (queryEnhancerSelector.isPresent()) {
156-
builder.addStatement("super(new $T(), context, beanFactory)", queryEnhancerSelector.get());
167+
builder.addStatement("super(new $T(), context)", queryEnhancerSelector.get());
157168
} else {
158-
builder.addStatement("super($T.DEFAULT_SELECTOR, context, beanFactory)", QueryEnhancerSelector.class);
169+
builder.addStatement("super($T.DEFAULT_SELECTOR, context)", QueryEnhancerSelector.class);
159170
}
171+
172+
queryRewriterFields.forEach((type, field) -> {
173+
if (hasAccessibleNoArgsConstructor(type)) {
174+
builder.addStatement("this.$L = beanFactory.getBeanProvider($T.class).getIfAvailable(() -> new $T())",
175+
field, type, type);
176+
} else {
177+
builder.addStatement("this.$L = beanFactory.getBean($T.class)", field, type);
178+
}
179+
});
160180
});
161181
}
162182

183+
private boolean hasQueryRewriter(Method method) {
184+
185+
MergedAnnotation<Query> query = MergedAnnotations.from(method).get(Query.class);
186+
return query.isPresent() && query.getClass("queryRewriter") != QueryRewriter.IdentityQueryRewriter.class;
187+
}
188+
189+
private boolean hasAccessibleNoArgsConstructor(Class<?> type) {
190+
191+
try {
192+
Constructor<?> constructor = type.getDeclaredConstructor();
193+
boolean samePackage = type.getPackageName()
194+
.equals(getRepositoryInformation().getRepositoryInterface().getPackageName());
195+
196+
return !Modifier.isAbstract(type.getModifiers()) && isAccessible(type, samePackage)
197+
&& (Modifier.isPublic(constructor.getModifiers())
198+
|| (!Modifier.isPrivate(constructor.getModifiers()) && samePackage))
199+
&& Arrays.stream(constructor.getExceptionTypes())
200+
.allMatch(it -> RuntimeException.class.isAssignableFrom(it) || Error.class.isAssignableFrom(it));
201+
} catch (NoSuchMethodException e) {
202+
return false;
203+
}
204+
}
205+
206+
private static boolean isAccessible(Class<?> type, boolean samePackage) {
207+
208+
if (Modifier.isPrivate(type.getModifiers()) || (!Modifier.isPublic(type.getModifiers()) && !samePackage)) {
209+
return false;
210+
}
211+
212+
Class<?> enclosingClass = type.getEnclosingClass();
213+
return enclosingClass == null || isAccessible(enclosingClass, samePackage);
214+
}
215+
163216
private @Nullable String getEntityManagerFactoryRef() {
164217
return context.getConfigurationSource().getAttribute("entityManagerFactoryRef")
165218
.filter(it -> !"entityManagerFactory".equals(it)).orElse(null);
@@ -243,6 +296,15 @@ private Optional<Class<QueryEnhancerSelector>> getQueryEnhancerSelectorClass() {
243296
}
244297
}
245298

299+
String queryRewriterField = null;
300+
if (aotQueries.result() instanceof StringAotQuery && hasQueryRewriter(method)) {
301+
Class<? extends QueryRewriter> type = query.getClass("queryRewriter").asSubclass(QueryRewriter.class);
302+
queryRewriterField = queryRewriterFields.computeIfAbsent(type,
303+
it -> "queryRewriter" + queryRewriterFields.size());
304+
}
305+
306+
String rewriterField = queryRewriterField;
307+
246308
return MethodContributor.forQueryMethod(queryMethod).withMetadata(aotQueries.toMetadata(queryMethod.isPageQuery()))
247309
.contribute(context -> {
248310

@@ -261,7 +323,7 @@ private Optional<Class<QueryEnhancerSelector>> getQueryEnhancerSelectorClass() {
261323
.queryReturnType(QueriesFactory.getQueryReturnType(aotQueries.result(), returnedType, context))
262324
.nativeQuery(nativeQuery).queryHints(queryHints).entityGraph(aotEntityGraph)
263325
.lockMode(lock.isPresent() ? lock.getEnum("value", LockModeType.class) : null)
264-
.queryRewriter(query.isPresent() ? query.getClass("queryRewriter") : null).build());
326+
.queryRewriterField(rewriterField).build());
265327

266328
body.add(JpaCodeBlocks.executionBuilder(context, queryMethod).modifying(modifying).query(aotQueries.result())
267329
.build());

0 commit comments

Comments
 (0)