Skip to content

Commit eeca4df

Browse files
mp911demp911de
committed
Detect PersistenceUnitManager and PersistenceManaged types in JpaRepositoryContributor.
See #3875
1 parent 01eb2b8 commit eeca4df

File tree

4 files changed

+220
-71
lines changed

4 files changed

+220
-71
lines changed

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

Lines changed: 78 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -22,74 +22,131 @@
2222
import jakarta.persistence.metamodel.ManagedType;
2323
import jakarta.persistence.metamodel.Metamodel;
2424
import jakarta.persistence.spi.ClassTransformer;
25+
import jakarta.persistence.spi.PersistenceUnitInfo;
2526

27+
import java.net.URL;
28+
import java.util.Collection;
2629
import java.util.List;
2730
import java.util.Map;
2831
import java.util.Set;
32+
import java.util.function.Supplier;
2933

34+
import org.hibernate.cfg.JdbcSettings;
35+
import org.hibernate.dialect.H2Dialect;
36+
import org.hibernate.engine.jdbc.connections.internal.UserSuppliedConnectionProviderImpl;
3037
import org.hibernate.jpa.HibernatePersistenceProvider;
3138
import org.hibernate.jpa.boot.internal.EntityManagerFactoryBuilderImpl;
3239
import org.hibernate.jpa.boot.internal.PersistenceUnitInfoDescriptor;
40+
import org.jspecify.annotations.Nullable;
41+
42+
import org.springframework.data.repository.config.AotRepositoryContext;
3343
import org.springframework.data.util.Lazy;
3444
import org.springframework.instrument.classloading.SimpleThrowawayClassLoader;
3545
import org.springframework.orm.jpa.persistenceunit.MutablePersistenceUnitInfo;
46+
import org.springframework.orm.jpa.persistenceunit.PersistenceManagedTypes;
3647

3748
/**
49+
* AOT metamodel implementation that uses Hibernate to build the metamodel.
50+
*
3851
* @author Christoph Strobl
52+
* @author Mark Paluch
3953
* @since 4.0
4054
*/
4155
class AotMetamodel implements Metamodel {
4256

43-
private final String persistenceUnit;
44-
private final Set<Class<?>> managedTypes;
45-
private final Lazy<EntityManagerFactory> entityManagerFactory = Lazy.of(this::init);
46-
private final Lazy<Metamodel> metamodel = Lazy.of(() -> entityManagerFactory.get().getMetamodel());
47-
private final Lazy<EntityManager> entityManager = Lazy.of(() -> entityManagerFactory.get().createEntityManager());
57+
private final Lazy<EntityManagerFactory> entityManagerFactory;
58+
private final Lazy<EntityManager> entityManager = Lazy.of(() -> getEntityManagerFactory().createEntityManager());
4859

49-
public AotMetamodel(Set<Class<?>> managedTypes) {
50-
this("AotMetamodel", managedTypes);
60+
public AotMetamodel(AotRepositoryContext repositoryContext) {
61+
this(repositoryContext.getResolvedTypes().stream().map(Class::getName)
62+
.filter(name -> !name.startsWith("jakarta.persistence")).toList(), null);
5163
}
5264

53-
private AotMetamodel(String persistenceUnit, Set<Class<?>> managedTypes) {
54-
this.persistenceUnit = persistenceUnit;
55-
this.managedTypes = managedTypes;
65+
public AotMetamodel(PersistenceManagedTypes managedTypes) {
66+
this(managedTypes.getManagedClassNames(), managedTypes.getPersistenceUnitRootUrl());
5667
}
5768

58-
public static AotMetamodel hibernateModel(Class<?>... types) {
59-
return new AotMetamodel(Set.of(types));
69+
public AotMetamodel(Collection<String> managedTypes, @Nullable URL persistenceUnitRootUrl) {
70+
71+
MutablePersistenceUnitInfo persistenceUnitInfo = new MutablePersistenceUnitInfo() {
72+
@Override
73+
public ClassLoader getNewTempClassLoader() {
74+
return new SimpleThrowawayClassLoader(this.getClass().getClassLoader());
75+
}
76+
77+
@Override
78+
public void addTransformer(ClassTransformer classTransformer) {
79+
// just ignore it
80+
}
81+
};
82+
persistenceUnitInfo.setPersistenceUnitName("AotMetaModel");
83+
84+
this.entityManagerFactory = init(() -> {
85+
86+
managedTypes.stream().forEach(persistenceUnitInfo::addManagedClassName);
87+
88+
persistenceUnitInfo.setPersistenceProviderClassName(HibernatePersistenceProvider.class.getName());
89+
90+
return new PersistenceUnitInfoDescriptor(persistenceUnitInfo) {
91+
92+
@Override
93+
public List<String> getManagedClassNames() {
94+
return persistenceUnitInfo.getManagedClassNames();
95+
}
96+
97+
@Override
98+
public URL getPersistenceUnitRootUrl() {
99+
return persistenceUnitRootUrl;
100+
}
101+
102+
};
103+
});
60104
}
61105

62-
public static AotMetamodel hibernateModel(String persistenceUnit, Class<?>... types) {
63-
return new AotMetamodel(persistenceUnit, Set.of(types));
106+
public AotMetamodel(PersistenceUnitInfo unitInfo) {
107+
this.entityManagerFactory = init(() -> new PersistenceUnitInfoDescriptor(unitInfo));
108+
}
109+
110+
static Lazy<EntityManagerFactory> init(Supplier<PersistenceUnitInfoDescriptor> unitInfo) {
111+
112+
return Lazy.of(() -> new EntityManagerFactoryBuilderImpl(unitInfo.get(),
113+
Map.of(JdbcSettings.DIALECT, H2Dialect.class.getName(), //
114+
JdbcSettings.ALLOW_METADATA_ON_BOOT, "false", //
115+
JdbcSettings.CONNECTION_PROVIDER, new UserSuppliedConnectionProviderImpl()))
116+
.build());
117+
}
118+
119+
private Metamodel getMetamodel() {
120+
return getEntityManagerFactory().getMetamodel();
64121
}
65122

66123
public <X> EntityType<X> entity(Class<X> cls) {
67-
return metamodel.get().entity(cls);
124+
return getMetamodel().entity(cls);
68125
}
69126

70127
@Override
71128
public EntityType<?> entity(String s) {
72-
return metamodel.get().entity(s);
129+
return getMetamodel().entity(s);
73130
}
74131

75132
public <X> ManagedType<X> managedType(Class<X> cls) {
76-
return metamodel.get().managedType(cls);
133+
return getMetamodel().managedType(cls);
77134
}
78135

79136
public <X> EmbeddableType<X> embeddable(Class<X> cls) {
80-
return metamodel.get().embeddable(cls);
137+
return getMetamodel().embeddable(cls);
81138
}
82139

83140
public Set<ManagedType<?>> getManagedTypes() {
84-
return metamodel.get().getManagedTypes();
141+
return getMetamodel().getManagedTypes();
85142
}
86143

87144
public Set<EntityType<?>> getEntities() {
88-
return metamodel.get().getEntities();
145+
return getMetamodel().getEntities();
89146
}
90147

91148
public Set<EmbeddableType<?>> getEmbeddables() {
92-
return metamodel.get().getEmbeddables();
149+
return getMetamodel().getEmbeddables();
93150
}
94151

95152
public EntityManager entityManager() {
@@ -100,32 +157,4 @@ public EntityManagerFactory getEntityManagerFactory() {
100157
return entityManagerFactory.get();
101158
}
102159

103-
EntityManagerFactory init() {
104-
105-
MutablePersistenceUnitInfo persistenceUnitInfo = new MutablePersistenceUnitInfo() {
106-
@Override
107-
public ClassLoader getNewTempClassLoader() {
108-
return new SimpleThrowawayClassLoader(this.getClass().getClassLoader());
109-
}
110-
111-
@Override
112-
public void addTransformer(ClassTransformer classTransformer) {
113-
// just ignore it
114-
}
115-
};
116-
117-
persistenceUnitInfo.setPersistenceUnitName(persistenceUnit);
118-
this.managedTypes.stream().map(Class::getName).forEach(persistenceUnitInfo::addManagedClassName);
119-
120-
persistenceUnitInfo.setPersistenceProviderClassName(HibernatePersistenceProvider.class.getName());
121-
122-
return new EntityManagerFactoryBuilderImpl(new PersistenceUnitInfoDescriptor(persistenceUnitInfo) {
123-
@Override
124-
public List<String> getManagedClassNames() {
125-
return persistenceUnitInfo.getManagedClassNames();
126-
}
127-
}, Map.of("hibernate.dialect", "org.hibernate.dialect.H2Dialect", "hibernate.boot.allow_jdbc_metadata_access",
128-
"false")).build();
129-
}
130-
131160
}

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

Lines changed: 26 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -17,10 +17,11 @@
1717

1818
import jakarta.persistence.EntityManager;
1919
import jakarta.persistence.EntityManagerFactory;
20+
import jakarta.persistence.metamodel.Metamodel;
21+
import jakarta.persistence.spi.PersistenceUnitInfo;
2022

2123
import java.lang.reflect.Method;
2224
import java.util.Map;
23-
import java.util.stream.Collectors;
2425

2526
import org.jspecify.annotations.Nullable;
2627

@@ -49,6 +50,7 @@
4950
import org.springframework.data.util.TypeInformation;
5051
import org.springframework.javapoet.CodeBlock;
5152
import org.springframework.javapoet.TypeName;
53+
import org.springframework.orm.jpa.persistenceunit.PersistenceManagedTypes;
5254
import org.springframework.util.ClassUtils;
5355
import org.springframework.util.StringUtils;
5456

@@ -64,31 +66,43 @@
6466
*/
6567
public class JpaRepositoryContributor extends RepositoryContributor {
6668

69+
private final Metamodel metamodel;
6770
private final PersistenceProvider persistenceProvider;
6871
private final QueriesFactory queriesFactory;
6972
private final EntityGraphLookup entityGraphLookup;
7073

7174
public JpaRepositoryContributor(AotRepositoryContext repositoryContext) {
75+
this(repositoryContext, new AotMetamodel(repositoryContext));
76+
}
7277

73-
super(repositoryContext);
74-
75-
AotMetamodel amm = new AotMetamodel(repositoryContext.getResolvedTypes().stream()
76-
.filter(it -> !it.getName().startsWith("jakarta.persistence")).collect(Collectors.toSet()));
78+
public JpaRepositoryContributor(AotRepositoryContext repositoryContext, PersistenceUnitInfo unitInfo) {
79+
this(repositoryContext, new AotMetamodel(unitInfo));
80+
}
7781

78-
this.persistenceProvider = PersistenceProvider.fromEntityManagerFactory(amm.getEntityManagerFactory());
79-
this.queriesFactory = new QueriesFactory(amm.getEntityManagerFactory(), amm);
80-
this.entityGraphLookup = new EntityGraphLookup(amm.getEntityManagerFactory());
82+
public JpaRepositoryContributor(AotRepositoryContext repositoryContext, PersistenceManagedTypes managedTypes) {
83+
this(repositoryContext, new AotMetamodel(managedTypes));
8184
}
8285

8386
public JpaRepositoryContributor(AotRepositoryContext repositoryContext, EntityManagerFactory entityManagerFactory) {
8487

8588
super(repositoryContext);
8689

90+
this.metamodel = entityManagerFactory.getMetamodel();
8791
this.persistenceProvider = PersistenceProvider.fromEntityManagerFactory(entityManagerFactory);
8892
this.queriesFactory = new QueriesFactory(entityManagerFactory);
8993
this.entityGraphLookup = new EntityGraphLookup(entityManagerFactory);
9094
}
9195

96+
private JpaRepositoryContributor(AotRepositoryContext repositoryContext, AotMetamodel metamodel) {
97+
98+
super(repositoryContext);
99+
100+
this.metamodel = metamodel;
101+
this.persistenceProvider = PersistenceProvider.fromEntityManagerFactory(metamodel.getEntityManagerFactory());
102+
this.queriesFactory = new QueriesFactory(metamodel.getEntityManagerFactory(), metamodel);
103+
this.entityGraphLookup = new EntityGraphLookup(metamodel.getEntityManagerFactory());
104+
}
105+
92106
@Override
93107
protected void customizeClass(AotRepositoryClassBuilder classBuilder) {
94108
classBuilder.customize(builder -> builder.superclass(TypeName.get(AotRepositoryFragmentSupport.class)));
@@ -203,6 +217,10 @@ protected void customizeConstructor(AotRepositoryConstructorBuilder constructorB
203217
});
204218
}
205219

220+
public Metamodel getMetamodel() {
221+
return metamodel;
222+
}
223+
206224
record StoredProcedureMetadata(String procedure) implements QueryMetadata {
207225

208226
@Override

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

Lines changed: 47 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import jakarta.persistence.MappedSuperclass;
2323
import jakarta.persistence.PersistenceContext;
2424
import jakarta.persistence.PersistenceUnit;
25+
import jakarta.persistence.spi.PersistenceUnitInfo;
2526

2627
import java.lang.annotation.Annotation;
2728
import java.util.Arrays;
@@ -35,8 +36,11 @@
3536
import java.util.Set;
3637

3738
import org.jspecify.annotations.Nullable;
39+
import org.slf4j.Logger;
40+
import org.slf4j.LoggerFactory;
3841

3942
import org.springframework.aot.generate.GenerationContext;
43+
import org.springframework.beans.factory.ObjectProvider;
4044
import org.springframework.beans.factory.aot.BeanRegistrationAotProcessor;
4145
import org.springframework.beans.factory.config.BeanDefinition;
4246
import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
@@ -46,6 +50,7 @@
4650
import org.springframework.beans.factory.support.RootBeanDefinition;
4751
import org.springframework.context.annotation.AnnotationConfigUtils;
4852
import org.springframework.core.annotation.AnnotationAttributes;
53+
import org.springframework.core.env.Environment;
4954
import org.springframework.core.io.ResourceLoader;
5055
import org.springframework.dao.DataAccessException;
5156
import org.springframework.dao.annotation.PersistenceExceptionTranslationPostProcessor;
@@ -57,13 +62,14 @@
5762
import org.springframework.data.jpa.repository.support.JpaEvaluationContextExtension;
5863
import org.springframework.data.jpa.repository.support.JpaRepositoryFactoryBean;
5964
import org.springframework.data.jpa.repository.support.SimpleJpaRepository;
60-
import org.springframework.data.repository.aot.generate.RepositoryContributor;
6165
import org.springframework.data.repository.config.AnnotationRepositoryConfigurationSource;
6266
import org.springframework.data.repository.config.AotRepositoryContext;
6367
import org.springframework.data.repository.config.RepositoryConfigurationExtensionSupport;
6468
import org.springframework.data.repository.config.RepositoryConfigurationSource;
6569
import org.springframework.data.repository.config.RepositoryRegistrationAotProcessor;
6670
import org.springframework.data.repository.config.XmlRepositoryConfigurationSource;
71+
import org.springframework.orm.jpa.persistenceunit.PersistenceManagedTypes;
72+
import org.springframework.orm.jpa.persistenceunit.PersistenceUnitManager;
6773
import org.springframework.orm.jpa.support.PersistenceAnnotationBeanPostProcessor;
6874
import org.springframework.util.ClassUtils;
6975
import org.springframework.util.StringUtils;
@@ -90,6 +96,7 @@ public class JpaRepositoryConfigExtension extends RepositoryConfigurationExtensi
9096
private static final String ENABLE_DEFAULT_TRANSACTIONS_ATTRIBUTE = "enableDefaultTransactions";
9197
private static final String JPA_METAMODEL_CACHE_CLEANUP_CLASSNAME = "org.springframework.data.jpa.util.JpaMetamodelCacheCleanup";
9298
private static final String ESCAPE_CHARACTER_PROPERTY = "escapeCharacter";
99+
private static final Logger log = LoggerFactory.getLogger(JpaRepositoryConfigExtension.class);
93100

94101
private final Map<Object, String> entityManagerRefs = new LinkedHashMap<>();
95102

@@ -327,28 +334,60 @@ public static class JpaRepositoryRegistrationAotProcessor extends RepositoryRegi
327334

328335
String GENERATED_REPOSITORIES_JPA_USE_ENTITY_MANAGER = "spring.aot.jpa.repositories.use-entitymanager";
329336

330-
protected @Nullable RepositoryContributor contribute(AotRepositoryContext repositoryContext,
337+
protected @Nullable JpaRepositoryContributor contribute(AotRepositoryContext repositoryContext,
331338
GenerationContext generationContext) {
332339

340+
Environment environment = repositoryContext.getEnvironment();
341+
333342
boolean enabled = Boolean.parseBoolean(
334-
repositoryContext.getEnvironment().getProperty(AotContext.GENERATED_REPOSITORIES_ENABLED, "false"));
343+
environment.getProperty(AotContext.GENERATED_REPOSITORIES_ENABLED, "false"));
335344
if (!enabled) {
336345
return null;
337346
}
338347

348+
ConfigurableListableBeanFactory beanFactory = repositoryContext.getBeanFactory();
349+
339350
boolean useEntityManager = Boolean.parseBoolean(
340-
repositoryContext.getEnvironment().getProperty(GENERATED_REPOSITORIES_JPA_USE_ENTITY_MANAGER, "false"));
351+
environment.getProperty(GENERATED_REPOSITORIES_JPA_USE_ENTITY_MANAGER, "false"));
341352

342353
if (useEntityManager) {
343354

344-
ConfigurableListableBeanFactory beanFactory = repositoryContext.getBeanFactory();
345-
EntityManagerFactory emf = beanFactory.getBeanProvider(EntityManagerFactory.class).getIfAvailable();
355+
ObjectProvider<PersistenceUnitManager> unitManagerProvider = beanFactory
356+
.getBeanProvider(PersistenceUnitManager.class);
357+
PersistenceUnitManager unitManager = unitManagerProvider.getIfAvailable();
346358

347-
if (emf != null) {
348-
return new JpaRepositoryContributor(repositoryContext, emf);
359+
if (unitManager != null) {
360+
361+
log.debug("Using PersistenceUnitManager for AOT repository generation");
362+
return new JpaRepositoryContributor(repositoryContext, unitManager.obtainDefaultPersistenceUnitInfo());
349363
}
364+
365+
log.debug("Using EntityManager for AOT repository generation");
366+
367+
EntityManagerFactory emf = beanFactory.getBean(EntityManagerFactory.class);
368+
return new JpaRepositoryContributor(repositoryContext, emf);
369+
}
370+
371+
ObjectProvider<PersistenceManagedTypes> managedTypesProvider = beanFactory
372+
.getBeanProvider(PersistenceManagedTypes.class);
373+
PersistenceManagedTypes managedTypes = managedTypesProvider.getIfAvailable();
374+
375+
if (managedTypes != null) {
376+
377+
log.debug("Using PersistenceManagedTypes for AOT repository generation");
378+
return new JpaRepositoryContributor(repositoryContext, managedTypes);
379+
}
380+
381+
ObjectProvider<PersistenceUnitInfo> infoProvider = beanFactory.getBeanProvider(PersistenceUnitInfo.class);
382+
PersistenceUnitInfo unitInfo = infoProvider.getIfAvailable();
383+
384+
if (unitInfo != null) {
385+
386+
log.debug("Using PersistenceUnitInfo for AOT repository generation");
387+
return new JpaRepositoryContributor(repositoryContext, unitInfo);
350388
}
351389

390+
log.debug("Using scanned types for AOT repository generation");
352391
return new JpaRepositoryContributor(repositoryContext);
353392
}
354393
}

0 commit comments

Comments
 (0)