Skip to content

Commit 3f17014

Browse files
mp911demp911de
committed
Capture @EnableJpaRepositories configuration for AOT processing.
Closes #3838
1 parent 88da07e commit 3f17014

File tree

10 files changed

+191
-47
lines changed

10 files changed

+191
-47
lines changed

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

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525

2626
import org.hibernate.jpa.HibernatePersistenceProvider;
2727
import org.junit.platform.commons.annotation.Testable;
28+
import org.mockito.Mockito;
2829
import org.openjdk.jmh.annotations.Benchmark;
2930
import org.openjdk.jmh.annotations.Fork;
3031
import org.openjdk.jmh.annotations.Level;
@@ -37,16 +38,24 @@
3738
import org.openjdk.jmh.annotations.Warmup;
3839

3940
import org.springframework.aot.test.generate.TestGenerationContext;
41+
import org.springframework.beans.factory.support.BeanDefinitionRegistry;
42+
import org.springframework.beans.factory.support.DefaultBeanNameGenerator;
43+
import org.springframework.core.env.StandardEnvironment;
44+
import org.springframework.core.io.DefaultResourceLoader;
4045
import org.springframework.core.test.tools.TestCompiler;
46+
import org.springframework.core.type.AnnotationMetadata;
4147
import org.springframework.data.domain.Sort;
4248
import org.springframework.data.jpa.benchmark.model.Person;
4349
import org.springframework.data.jpa.benchmark.model.Profile;
4450
import org.springframework.data.jpa.benchmark.repository.PersonRepository;
4551
import org.springframework.data.jpa.repository.aot.JpaRepositoryContributor;
4652
import org.springframework.data.jpa.repository.aot.TestJpaAotRepositoryContext;
53+
import org.springframework.data.jpa.repository.config.EnableJpaRepositories;
54+
import org.springframework.data.jpa.repository.sample.SampleConfig;
4755
import org.springframework.data.jpa.repository.support.JpaRepositoryFactory;
4856
import org.springframework.data.projection.ProjectionFactory;
4957
import org.springframework.data.projection.SpelAwareProxyProjectionFactory;
58+
import org.springframework.data.repository.config.AnnotationRepositoryConfigurationSource;
5059
import org.springframework.data.repository.core.RepositoryMetadata;
5160
import org.springframework.data.repository.core.support.RepositoryComposition;
5261
import org.springframework.data.repository.core.support.RepositoryFactoryBeanSupport;
@@ -74,7 +83,10 @@ public static class BenchmarkParameters {
7483

7584
public static Class<?> aot;
7685
public static TestJpaAotRepositoryContext<PersonRepository> repositoryContext = new TestJpaAotRepositoryContext<>(
77-
PersonRepository.class, null);
86+
PersonRepository.class, null,
87+
new AnnotationRepositoryConfigurationSource(AnnotationMetadata.introspect(SampleConfig.class),
88+
EnableJpaRepositories.class, new DefaultResourceLoader(), new StandardEnvironment(),
89+
Mockito.mock(BeanDefinitionRegistry.class), DefaultBeanNameGenerator.INSTANCE));
7890

7991
EntityManager entityManager;
8092
RepositoryComposition.RepositoryFragments fragments;

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

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -58,8 +58,16 @@ class AotMetamodel implements Metamodel {
5858
private final Lazy<EntityManager> entityManager = Lazy.of(() -> getEntityManagerFactory().createEntityManager());
5959

6060
public AotMetamodel(AotRepositoryContext repositoryContext) {
61-
this(repositoryContext.getResolvedTypes().stream().map(Class::getName)
62-
.filter(name -> !name.startsWith("jakarta.persistence")).toList(), null);
61+
this(repositoryContext.getResolvedTypes().stream().filter(AotMetamodel::isJakartaAnnotated).map(Class::getName)
62+
.toList(), null);
63+
}
64+
65+
private static boolean isJakartaAnnotated(Class<?> cls) {
66+
67+
return cls.isAnnotationPresent(jakarta.persistence.Entity.class)
68+
|| cls.isAnnotationPresent(jakarta.persistence.Embeddable.class)
69+
|| cls.isAnnotationPresent(jakarta.persistence.MappedSuperclass.class)
70+
|| cls.isAnnotationPresent(jakarta.persistence.Converter.class);
6371
}
6472

6573
public AotMetamodel(PersistenceManagedTypes managedTypes) {

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

Lines changed: 25 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -22,9 +22,11 @@
2222

2323
import java.lang.reflect.Method;
2424
import java.util.Map;
25+
import java.util.Optional;
2526

2627
import org.jspecify.annotations.Nullable;
2728

29+
import org.springframework.beans.BeanUtils;
2830
import org.springframework.core.annotation.AnnotatedElementUtils;
2931
import org.springframework.core.annotation.MergedAnnotation;
3032
import org.springframework.core.annotation.MergedAnnotations;
@@ -70,6 +72,7 @@ public class JpaRepositoryContributor extends RepositoryContributor {
7072
private final PersistenceProvider persistenceProvider;
7173
private final QueriesFactory queriesFactory;
7274
private final EntityGraphLookup entityGraphLookup;
75+
private final AotRepositoryContext context;
7376

7477
public JpaRepositoryContributor(AotRepositoryContext repositoryContext) {
7578
this(repositoryContext, new AotMetamodel(repositoryContext));
@@ -87,19 +90,22 @@ public JpaRepositoryContributor(AotRepositoryContext repositoryContext, EntityMa
8790

8891
super(repositoryContext);
8992

93+
this.context = repositoryContext;
9094
this.metamodel = entityManagerFactory.getMetamodel();
9195
this.persistenceProvider = PersistenceProvider.fromEntityManagerFactory(entityManagerFactory);
92-
this.queriesFactory = new QueriesFactory(entityManagerFactory);
96+
this.queriesFactory = new QueriesFactory(repositoryContext.getConfigurationSource(), entityManagerFactory);
9397
this.entityGraphLookup = new EntityGraphLookup(entityManagerFactory);
9498
}
9599

96100
private JpaRepositoryContributor(AotRepositoryContext repositoryContext, AotMetamodel metamodel) {
97101

98102
super(repositoryContext);
99103

104+
this.context = repositoryContext;
100105
this.metamodel = metamodel;
101106
this.persistenceProvider = PersistenceProvider.fromEntityManagerFactory(metamodel.getEntityManagerFactory());
102-
this.queriesFactory = new QueriesFactory(metamodel.getEntityManagerFactory(), metamodel);
107+
this.queriesFactory = new QueriesFactory(repositoryContext.getConfigurationSource(),
108+
metamodel.getEntityManagerFactory(), metamodel);
103109
this.entityGraphLookup = new EntityGraphLookup(metamodel.getEntityManagerFactory());
104110
}
105111

@@ -111,25 +117,36 @@ protected void customizeClass(AotRepositoryClassBuilder classBuilder) {
111117
@Override
112118
protected void customizeConstructor(AotRepositoryConstructorBuilder constructorBuilder) {
113119

114-
// TODO: BeanFactoryQueryRewriterProvider if there is a method using QueryRewriters.
115-
116120
constructorBuilder.addParameter("entityManager", EntityManager.class);
117121
constructorBuilder.addParameter("context", RepositoryFactoryBeanSupport.FragmentCreationContext.class);
118122

119-
// TODO: Pick up the configured QueryEnhancerSelector
123+
Optional<Class<QueryEnhancerSelector>> queryEnhancerSelector = getQueryEnhancerSelectorClass();
124+
120125
constructorBuilder.customize(builder -> {
121-
builder.addStatement("super($T.DEFAULT_SELECTOR, context)", QueryEnhancerSelector.class);
126+
127+
if (queryEnhancerSelector.isPresent()) {
128+
builder.addStatement("super(new T$(), context)", queryEnhancerSelector.get());
129+
} else {
130+
builder.addStatement("super($T.DEFAULT_SELECTOR, context)", QueryEnhancerSelector.class);
131+
}
122132
});
123133
}
124134

135+
@SuppressWarnings({ "rawtypes", "unchecked" })
136+
private Optional<Class<QueryEnhancerSelector>> getQueryEnhancerSelectorClass() {
137+
return (Optional) context.getConfigurationSource().getAttribute("queryEnhancerSelector", Class.class)
138+
.filter(it -> !it.equals(QueryEnhancerSelector.DefaultQueryEnhancerSelector.class));
139+
}
140+
125141
@Override
126142
protected @Nullable MethodContributor<? extends QueryMethod> contributeQueryMethod(Method method) {
127143

128144
JpaQueryMethod queryMethod = new JpaQueryMethod(method, getRepositoryInformation(), getProjectionFactory(),
129145
persistenceProvider);
130146

131-
// meh!
132-
QueryEnhancerSelector selector = QueryEnhancerSelector.DEFAULT_SELECTOR;
147+
Optional<Class<QueryEnhancerSelector>> queryEnhancerSelectorClass = getQueryEnhancerSelectorClass();
148+
QueryEnhancerSelector selector = queryEnhancerSelectorClass.map(BeanUtils::instantiateClass)
149+
.orElse(QueryEnhancerSelector.DEFAULT_SELECTOR);
133150

134151
// no stored procedures for now.
135152
if (queryMethod.isProcedureQuery()) {
@@ -183,7 +200,6 @@ protected void customizeConstructor(AotRepositoryConstructorBuilder constructorB
183200
TypeInformation<?> returnType = getRepositoryInformation().getReturnType(method);
184201

185202
boolean returnsCount = JpaCodeBlocks.QueryExecutionBlockBuilder.returnsModifying(returnType.getType());
186-
187203
boolean isVoid = ClassUtils.isVoidType(returnType.getType());
188204

189205
if (!returnsCount && !isVoid) {

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

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import java.util.Arrays;
2525
import java.util.List;
2626
import java.util.Map;
27+
import java.util.Optional;
2728
import java.util.function.Function;
2829
import java.util.function.UnaryOperator;
2930

@@ -37,6 +38,7 @@
3738
import org.springframework.data.jpa.repository.query.*;
3839
import org.springframework.data.jpa.repository.support.JpqlQueryTemplates;
3940
import org.springframework.data.repository.aot.generate.AotQueryMethodGenerationContext;
41+
import org.springframework.data.repository.config.RepositoryConfigurationSource;
4042
import org.springframework.data.repository.core.RepositoryInformation;
4143
import org.springframework.data.repository.query.ReturnedType;
4244
import org.springframework.data.repository.query.parser.PartTree;
@@ -53,14 +55,21 @@ class QueriesFactory {
5355

5456
private final EntityManagerFactory entityManagerFactory;
5557
private final Metamodel metamodel;
58+
private final EscapeCharacter escapeCharacter;
59+
private final JpqlQueryTemplates templates = JpqlQueryTemplates.UPPER;
5660

57-
public QueriesFactory(EntityManagerFactory entityManagerFactory) {
58-
this(entityManagerFactory, entityManagerFactory.getMetamodel());
61+
public QueriesFactory(RepositoryConfigurationSource configurationSource, EntityManagerFactory entityManagerFactory) {
62+
this(configurationSource, entityManagerFactory, entityManagerFactory.getMetamodel());
5963
}
6064

61-
public QueriesFactory(EntityManagerFactory entityManagerFactory, Metamodel metamodel) {
65+
public QueriesFactory(RepositoryConfigurationSource configurationSource, EntityManagerFactory entityManagerFactory,
66+
Metamodel metamodel) {
67+
6268
this.metamodel = metamodel;
6369
this.entityManagerFactory = entityManagerFactory;
70+
71+
Optional<Character> escapeCharacter = configurationSource.getAttribute("escapeCharacter", Character.class);
72+
this.escapeCharacter = escapeCharacter.map(EscapeCharacter::of).orElse(EscapeCharacter.DEFAULT);
6473
}
6574

6675
/**
@@ -77,8 +86,7 @@ public AotQueries createQueries(RepositoryInformation repositoryInformation, Mer
7786
QueryEnhancerSelector selector, JpaQueryMethod queryMethod, ReturnedType returnedType) {
7887

7988
if (query.isPresent() && StringUtils.hasText(query.getString("value"))) {
80-
return buildStringQuery(repositoryInformation.getDomainType(), returnedType, selector, query,
81-
queryMethod);
89+
return buildStringQuery(repositoryInformation.getDomainType(), returnedType, selector, query, queryMethod);
8290
}
8391

8492
TypedQueryReference<?> namedQuery = getNamedQuery(returnedType, queryMethod.getNamedQueryName());
@@ -201,9 +209,6 @@ private AotQueries buildPartTreeQuery(ReturnedType returnedType, RepositoryInfor
201209
MergedAnnotation<Query> query, JpaQueryMethod queryMethod) {
202210

203211
PartTree partTree = new PartTree(queryMethod.getName(), repositoryInformation.getDomainType());
204-
// TODO make configurable
205-
JpqlQueryTemplates templates = JpqlQueryTemplates.UPPER;
206-
207212
AotQuery aotQuery = createQuery(partTree, returnedType, queryMethod.getParameters(), templates);
208213

209214
if (query.isPresent() && StringUtils.hasText(query.getString("countQuery"))) {
@@ -222,8 +227,7 @@ private AotQueries buildPartTreeQuery(ReturnedType returnedType, RepositoryInfor
222227
private AotQuery createQuery(PartTree partTree, ReturnedType returnedType, JpaParameters parameters,
223228
JpqlQueryTemplates templates) {
224229

225-
ParameterMetadataProvider metadataProvider = new ParameterMetadataProvider(parameters, EscapeCharacter.DEFAULT,
226-
templates);
230+
ParameterMetadataProvider metadataProvider = new ParameterMetadataProvider(parameters, escapeCharacter, templates);
227231
JpaQueryCreator queryCreator = new JpaQueryCreator(partTree, false, returnedType, metadataProvider, templates,
228232
metamodel);
229233

@@ -234,8 +238,7 @@ private AotQuery createQuery(PartTree partTree, ReturnedType returnedType, JpaPa
234238
private AotQuery createCountQuery(PartTree partTree, ReturnedType returnedType, JpaParameters parameters,
235239
JpqlQueryTemplates templates) {
236240

237-
ParameterMetadataProvider metadataProvider = new ParameterMetadataProvider(parameters, EscapeCharacter.DEFAULT,
238-
templates);
241+
ParameterMetadataProvider metadataProvider = new ParameterMetadataProvider(parameters, escapeCharacter, templates);
239242
JpaQueryCreator queryCreator = new JpaCountQueryCreator(partTree, returnedType, metadataProvider, templates,
240243
metamodel);
241244

‎spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/config/JpaRepositoryConfigExtension.java‎

Lines changed: 13 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,6 @@
6969
import org.springframework.data.repository.config.RepositoryRegistrationAotProcessor;
7070
import org.springframework.data.repository.config.XmlRepositoryConfigurationSource;
7171
import org.springframework.orm.jpa.persistenceunit.PersistenceManagedTypes;
72-
import org.springframework.orm.jpa.persistenceunit.PersistenceUnitManager;
7372
import org.springframework.orm.jpa.support.PersistenceAnnotationBeanPostProcessor;
7473
import org.springframework.util.ClassUtils;
7574
import org.springframework.util.StringUtils;
@@ -339,38 +338,34 @@ public static class JpaRepositoryRegistrationAotProcessor extends RepositoryRegi
339338

340339
Environment environment = repositoryContext.getEnvironment();
341340

342-
boolean enabled = Boolean.parseBoolean(
343-
environment.getProperty(AotContext.GENERATED_REPOSITORIES_ENABLED, "false"));
341+
boolean enabled = Boolean
342+
.parseBoolean(environment.getProperty(AotContext.GENERATED_REPOSITORIES_ENABLED, "false"));
344343
if (!enabled) {
345344
return null;
346345
}
347346

348347
ConfigurableListableBeanFactory beanFactory = repositoryContext.getBeanFactory();
349348

350-
boolean useEntityManager = Boolean.parseBoolean(
351-
environment.getProperty(GENERATED_REPOSITORIES_JPA_USE_ENTITY_MANAGER, "false"));
349+
boolean useEntityManager = Boolean
350+
.parseBoolean(environment.getProperty(GENERATED_REPOSITORIES_JPA_USE_ENTITY_MANAGER, "false"));
352351

353352
if (useEntityManager) {
354353

355-
ObjectProvider<PersistenceUnitManager> unitManagerProvider = beanFactory
356-
.getBeanProvider(PersistenceUnitManager.class);
357-
PersistenceUnitManager unitManager = unitManagerProvider.getIfAvailable();
354+
Optional<String> entityManagerFactoryRef = repositoryContext.getConfigurationSource()
355+
.getAttribute("entityManagerFactoryRef");
358356

359-
if (unitManager != null) {
357+
log.debug(
358+
"Using EntityManager '%s' for AOT repository generation".formatted(entityManagerFactoryRef.orElse("")));
360359

361-
log.debug("Using PersistenceUnitManager for AOT repository generation");
362-
return new JpaRepositoryContributor(repositoryContext, unitManager.obtainDefaultPersistenceUnitInfo());
363-
}
364-
365-
log.debug("Using EntityManager for AOT repository generation");
366-
367-
EntityManagerFactory emf = beanFactory.getBean(EntityManagerFactory.class);
360+
EntityManagerFactory emf = entityManagerFactoryRef
361+
.map(it -> beanFactory.getBean(it, EntityManagerFactory.class))
362+
.orElseGet(() -> beanFactory.getBean(EntityManagerFactory.class));
368363
return new JpaRepositoryContributor(repositoryContext, emf);
369364
}
370365

371366
ObjectProvider<PersistenceManagedTypes> managedTypesProvider = beanFactory
372367
.getBeanProvider(PersistenceManagedTypes.class);
373-
PersistenceManagedTypes managedTypes = managedTypesProvider.getIfAvailable();
368+
PersistenceManagedTypes managedTypes = managedTypesProvider.getIfUnique();
374369

375370
if (managedTypes != null) {
376371

@@ -379,7 +374,7 @@ public static class JpaRepositoryRegistrationAotProcessor extends RepositoryRegi
379374
}
380375

381376
ObjectProvider<PersistenceUnitInfo> infoProvider = beanFactory.getBeanProvider(PersistenceUnitInfo.class);
382-
PersistenceUnitInfo unitInfo = infoProvider.getIfAvailable();
377+
PersistenceUnitInfo unitInfo = infoProvider.getIfUnique();
383378

384379
if (unitInfo != null) {
385380

‎spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/aot/AotFragmentTestConfigurationSupport.java‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@
2121
import java.lang.reflect.Method;
2222
import java.lang.reflect.Proxy;
2323

24+
import org.mockito.Mockito;
25+
2426
import org.springframework.aot.test.generate.TestGenerationContext;
2527
import org.springframework.beans.BeansException;
2628
import org.springframework.beans.factory.config.BeanDefinition;
@@ -29,11 +31,18 @@
2931
import org.springframework.beans.factory.support.AbstractBeanDefinition;
3032
import org.springframework.beans.factory.support.BeanDefinitionBuilder;
3133
import org.springframework.beans.factory.support.BeanDefinitionRegistry;
34+
import org.springframework.beans.factory.support.DefaultBeanNameGenerator;
3235
import org.springframework.context.annotation.Bean;
3336
import org.springframework.context.annotation.ImportResource;
37+
import org.springframework.core.env.StandardEnvironment;
38+
import org.springframework.core.io.DefaultResourceLoader;
3439
import org.springframework.core.test.tools.TestCompiler;
40+
import org.springframework.core.type.AnnotationMetadata;
41+
import org.springframework.data.jpa.repository.config.EnableJpaRepositories;
42+
import org.springframework.data.jpa.repository.sample.SampleConfig;
3543
import org.springframework.data.projection.ProjectionFactory;
3644
import org.springframework.data.projection.SpelAwareProxyProjectionFactory;
45+
import org.springframework.data.repository.config.AnnotationRepositoryConfigurationSource;
3746
import org.springframework.data.repository.core.RepositoryMetadata;
3847
import org.springframework.data.repository.core.support.RepositoryFactoryBeanSupport;
3948
import org.springframework.data.repository.query.ValueExpressionDelegate;
@@ -56,8 +65,15 @@ class AotFragmentTestConfigurationSupport implements BeanFactoryPostProcessor {
5665
private final TestJpaAotRepositoryContext<?> repositoryContext;
5766

5867
public AotFragmentTestConfigurationSupport(Class<?> repositoryInterface) {
68+
this(repositoryInterface, SampleConfig.class);
69+
}
70+
71+
public AotFragmentTestConfigurationSupport(Class<?> repositoryInterface, Class<?> configClass) {
5972
this.repositoryInterface = repositoryInterface;
60-
this.repositoryContext = new TestJpaAotRepositoryContext<>(repositoryInterface, null);
73+
this.repositoryContext = new TestJpaAotRepositoryContext<>(repositoryInterface, null,
74+
new AnnotationRepositoryConfigurationSource(AnnotationMetadata.introspect(configClass),
75+
EnableJpaRepositories.class, new DefaultResourceLoader(), new StandardEnvironment(),
76+
Mockito.mock(BeanDefinitionRegistry.class), DefaultBeanNameGenerator.INSTANCE));
6177
}
6278

6379
@Override

0 commit comments

Comments
 (0)