Skip to content

Commit 332d8c3

Browse files
yeomin4242yeomin4242
committed
Resolve QueryRewriter beans in AOT repositories
Signed-off-by: YeongJae Min <whereismysejong@naver.com>
1 parent e087d54 commit 332d8c3

12 files changed

Lines changed: 236 additions & 26 deletions

File tree

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

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

4040
import org.springframework.aot.test.generate.TestGenerationContext;
41+
import org.springframework.beans.factory.BeanFactory;
4142
import org.springframework.beans.factory.support.BeanDefinitionRegistry;
4243
import org.springframework.beans.factory.support.DefaultBeanNameGenerator;
4344
import org.springframework.beans.factory.support.DefaultListableBeanFactory;
@@ -49,6 +50,7 @@
4950
import org.springframework.data.jpa.benchmark.model.Person;
5051
import org.springframework.data.jpa.benchmark.model.Profile;
5152
import org.springframework.data.jpa.benchmark.repository.PersonRepository;
53+
import org.springframework.data.jpa.repository.aot.AotRepositoryFragmentSupport;
5254
import org.springframework.data.jpa.repository.aot.JpaRepositoryContributor;
5355
import org.springframework.data.jpa.repository.aot.TestJpaAotRepositoryContext;
5456
import org.springframework.data.jpa.repository.config.EnableJpaRepositories;
@@ -69,6 +71,7 @@
6971
/**
7072
* @author Christoph Strobl
7173
* @author Mark Paluch
74+
* @author YeongJae Min
7275
*/
7376
@Testable
7477
@Fork(1)
@@ -92,7 +95,7 @@ public static class BenchmarkParameters {
9295
new StandardEnvironment(), Mockito.mock(BeanDefinitionRegistry.class), DefaultBeanNameGenerator.INSTANCE);
9396

9497
repositoryContext = new TestJpaAotRepositoryContext<>(new DefaultListableBeanFactory(), PersonRepository.class,
95-
null, configurationSource);
98+
RepositoryComposition.empty(), configurationSource);
9699
}
97100

98101
EntityManager entityManager;
@@ -132,21 +135,20 @@ public void doSetup() {
132135
new JpaRepositoryContributor(repositoryContext, entityManager.getEntityManagerFactory())
133136
.contribute(generationContext);
134137

135-
TestCompiler.forSystem().withCompilerOptions("-parameters").with(generationContext).compile(compiled -> {
138+
generationContext.writeGeneratedContent();
136139

137-
try {
138-
this.aot = compiled.getClassLoader().loadClass(PersonRepository.class.getName() + "Impl__Aot");
139-
} catch (Exception e) {
140-
throw new RuntimeException(e);
141-
}
140+
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"));
142144
});
143145
}
144146

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

151153
this.repositoryProxy = createRepository(fragments);
152154
} catch (Exception e) {

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

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

3031
import org.jspecify.annotations.Nullable;
3132

33+
import org.springframework.beans.factory.BeanFactory;
3234
import org.springframework.core.CollectionFactory;
3335
import org.springframework.core.convert.ConversionService;
3436
import org.springframework.core.convert.support.ConfigurableConversionService;
@@ -37,11 +39,14 @@
3739
import org.springframework.data.domain.Sort;
3840
import org.springframework.data.expression.ValueEvaluationContextProvider;
3941
import org.springframework.data.expression.ValueExpression;
42+
import org.springframework.data.jpa.repository.QueryRewriter;
43+
import org.springframework.data.jpa.repository.query.BeanFactoryQueryRewriterProvider;
4044
import org.springframework.data.jpa.repository.query.DeclaredQuery;
4145
import org.springframework.data.jpa.repository.query.JpaParameters;
4246
import org.springframework.data.jpa.repository.query.JpaResultConverters;
4347
import org.springframework.data.jpa.repository.query.QueryEnhancer;
4448
import org.springframework.data.jpa.repository.query.QueryEnhancerSelector;
49+
import org.springframework.data.jpa.repository.query.QueryRewriterProvider;
4550
import org.springframework.data.jpa.util.TupleBackedMap;
4651
import org.springframework.data.projection.ProjectionFactory;
4752
import org.springframework.data.repository.core.RepositoryMetadata;
@@ -56,6 +61,7 @@
5661
* Support class for JPA AOT repository fragments.
5762
*
5863
* @author Mark Paluch
64+
* @author YeongJae Min
5965
* @since 4.0
6066
*/
6167
public class AotRepositoryFragmentSupport {
@@ -85,18 +91,33 @@ public class AotRepositoryFragmentSupport {
8591

8692
private final Lazy<ConcurrentLruCache<Method, ValueEvaluationContextProvider>> contextProviders;
8793

94+
private final QueryRewriterProvider queryRewriterProvider;
95+
8896
protected AotRepositoryFragmentSupport(QueryEnhancerSelector selector,
8997
RepositoryFactoryBeanSupport.FragmentCreationContext context) {
9098
this(selector, context.getRepositoryMetadata(), context.getValueExpressionDelegate(),
9199
context.getProjectionFactory());
92100
}
93101

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+
94108
protected AotRepositoryFragmentSupport(QueryEnhancerSelector selector, RepositoryMetadata repositoryMetadata,
95109
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) {
96116

97117
this.repositoryMetadata = repositoryMetadata;
98118
this.valueExpressions = valueExpressions;
99119
this.projectionFactory = projectionFactory;
120+
this.queryRewriterProvider = queryRewriterProvider;
100121
this.enhancers = Lazy.of(() -> new ConcurrentLruCache<>(32, query -> selector.select(query).create(query)));
101122
this.expressions = Lazy.of(() -> new ConcurrentLruCache<>(32, valueExpressions::parse));
102123
this.contextProviders = Lazy.of(() -> new ConcurrentLruCache<>(32, it -> valueExpressions
@@ -118,6 +139,15 @@ protected String rewriteQuery(DeclaredQuery query, Sort sort, Class<?> returnedT
118139
ReturnedType.of(returnedType, repositoryMetadata.getDomainType(), projectionFactory)));
119140
}
120141

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+
121151
/**
122152
* Evaluate a Value Expression.
123153
*

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

Lines changed: 40 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,9 @@
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;
2427
import java.util.Collection;
2528
import java.util.List;
2629
import java.util.function.LongSupplier;
@@ -67,6 +70,7 @@
6770
* @author Christoph Strobl
6871
* @author Mark Paluch
6972
* @author JungKyun Ryu
73+
* @author YeongJae Min
7074
* @since 4.0
7175
*/
7276
class JpaCodeBlocks {
@@ -182,7 +186,13 @@ public CodeBlock build() {
182186
if (queries.result() instanceof StringAotQuery && queryRewriter != QueryRewriter.IdentityQueryRewriter.class) {
183187

184188
queryRewriterName = context.localVariable("queryRewriter");
185-
builder.addStatement("$T $L = new $T()", queryRewriter, queryRewriterName, 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+
}
186196
}
187197

188198
if (queries.result() instanceof StringAotQuery sq) {
@@ -251,6 +261,35 @@ private CodeBlock buildQueryString(StringAotQuery sq, String queryStringVariable
251261
return builder.build();
252262
}
253263

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+
254293
private CodeBlock applyRewrite(@Nullable String sort, @Nullable String dynamicReturnType, boolean isProjecting,
255294
String queryString) {
256295

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

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828

2929
import org.jspecify.annotations.Nullable;
3030
import org.springframework.beans.BeanUtils;
31+
import org.springframework.beans.factory.BeanFactory;
3132
import org.springframework.beans.factory.config.RuntimeBeanReference;
3233
import org.springframework.core.annotation.AnnotatedElementUtils;
3334
import org.springframework.core.annotation.MergedAnnotation;
@@ -75,6 +76,7 @@
7576
* @author Christoph Strobl
7677
* @author Mark Paluch
7778
* @author JungKyun Ryu
79+
* @author YeongJae Min
7880
* @since 4.0
7981
*/
8082
public class JpaRepositoryContributor extends RepositoryContributor {
@@ -143,15 +145,17 @@ protected void customizeConstructor(AotRepositoryConstructorBuilder constructorB
143145
});
144146

145147
constructorBuilder.addParameter("context", RepositoryFactoryBeanSupport.FragmentCreationContext.class);
148+
constructorBuilder.addParameter("beanFactory", BeanFactory.class,
149+
customizer -> customizer.origin(context -> context.getBeanFactory()));
146150

147151
Optional<Class<QueryEnhancerSelector>> queryEnhancerSelector = getQueryEnhancerSelectorClass();
148152

149153
constructorBuilder.customize(builder -> {
150154

151155
if (queryEnhancerSelector.isPresent()) {
152-
builder.addStatement("super(new $T(), context)", queryEnhancerSelector.get());
156+
builder.addStatement("super(new $T(), context, beanFactory)", queryEnhancerSelector.get());
153157
} else {
154-
builder.addStatement("super($T.DEFAULT_SELECTOR, context)", QueryEnhancerSelector.class);
158+
builder.addStatement("super($T.DEFAULT_SELECTOR, context, beanFactory)", QueryEnhancerSelector.class);
155159
}
156160
});
157161
}

‎spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/cdi/BeanManagerQueryRewriterProvider.java‎

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import jakarta.enterprise.inject.spi.BeanManager;
2121

2222
import java.util.Iterator;
23+
import java.util.function.Supplier;
2324

2425
import org.springframework.beans.BeanUtils;
2526
import org.springframework.data.jpa.repository.QueryRewriter;
@@ -33,6 +34,7 @@
3334
*
3435
* @author Greg Turnquist
3536
* @author Mark Paluch
37+
* @author YeongJae Min
3638
* @since 3.0
3739
*/
3840
public class BeanManagerQueryRewriterProvider implements QueryRewriterProvider {
@@ -44,10 +46,20 @@ public BeanManagerQueryRewriterProvider(BeanManager beanManager) {
4446
}
4547

4648
@Override
47-
@SuppressWarnings("unchecked")
4849
public QueryRewriter getQueryRewriter(JpaQueryMethod method) {
50+
return getQueryRewriter(method.getQueryRewriter());
51+
}
52+
53+
@Override
54+
public QueryRewriter getQueryRewriter(Class<? extends QueryRewriter> queryRewriter) {
55+
return getQueryRewriter(queryRewriter, () -> BeanUtils.instantiateClass(queryRewriter));
56+
}
57+
58+
@Override
59+
@SuppressWarnings("unchecked")
60+
public QueryRewriter getQueryRewriter(Class<? extends QueryRewriter> queryRewriter,
61+
Supplier<QueryRewriter> fallback) {
4962

50-
Class<? extends QueryRewriter> queryRewriter = method.getQueryRewriter();
5163
if (queryRewriter == QueryRewriter.IdentityQueryRewriter.class) {
5264
return QueryRewriter.IdentityQueryRewriter.INSTANCE;
5365
}
@@ -64,7 +76,7 @@ public QueryRewriter getQueryRewriter(JpaQueryMethod method) {
6476
return new DelegatingQueryRewriter(rewriter);
6577
}
6678

67-
return BeanUtils.instantiateClass(queryRewriter);
79+
return fallback.get();
6880
}
6981

7082
}

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

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@
1515
*/
1616
package org.springframework.data.jpa.repository.query;
1717

18+
import java.util.function.Supplier;
19+
1820
import org.springframework.beans.BeanUtils;
1921
import org.springframework.beans.factory.BeanFactory;
2022
import org.springframework.data.jpa.repository.QueryRewriter;
@@ -25,6 +27,7 @@
2527
*
2628
* @author Greg Turnquist
2729
* @author Mark Paluch
30+
* @author YeongJae Min
2831
* @since 3.0
2932
*/
3033
public class BeanFactoryQueryRewriterProvider implements QueryRewriterProvider {
@@ -36,16 +39,26 @@ public BeanFactoryQueryRewriterProvider(BeanFactory beanFactory) {
3639
}
3740

3841
@Override
39-
@SuppressWarnings("unchecked")
4042
public QueryRewriter getQueryRewriter(JpaQueryMethod method) {
43+
return getQueryRewriter(method.getQueryRewriter());
44+
}
45+
46+
@Override
47+
public QueryRewriter getQueryRewriter(Class<? extends QueryRewriter> queryRewriter) {
48+
return getQueryRewriter(queryRewriter, () -> BeanUtils.instantiateClass(queryRewriter));
49+
}
50+
51+
@Override
52+
@SuppressWarnings("unchecked")
53+
public QueryRewriter getQueryRewriter(Class<? extends QueryRewriter> queryRewriter,
54+
Supplier<QueryRewriter> fallback) {
4155

42-
Class<? extends QueryRewriter> queryRewriter = method.getQueryRewriter();
4356
if (queryRewriter == QueryRewriter.IdentityQueryRewriter.class) {
4457
return QueryRewriter.IdentityQueryRewriter.INSTANCE;
4558
}
4659

4760
Lazy<QueryRewriter> rewriter = Lazy.of(() -> beanFactory.getBeanProvider((Class<QueryRewriter>) queryRewriter)
48-
.getIfAvailable(() -> BeanUtils.instantiateClass(queryRewriter)));
61+
.getIfAvailable(fallback));
4962

5063
return new DelegatingQueryRewriter(rewriter);
5164
}

0 commit comments

Comments
 (0)