Detect PersistenceUnitManager and PersistenceManaged types in JpaRepositoryContributor.

See #3875
This commit is contained in:
Mark Paluch
2025-05-13 15:59:41 +02:00
parent 01eb2b833d
commit eeca4df36f
4 changed files with 237 additions and 88 deletions

View File

@@ -22,85 +22,51 @@ import jakarta.persistence.metamodel.EntityType;
import jakarta.persistence.metamodel.ManagedType;
import jakarta.persistence.metamodel.Metamodel;
import jakarta.persistence.spi.ClassTransformer;
import jakarta.persistence.spi.PersistenceUnitInfo;
import java.net.URL;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.function.Supplier;
import org.hibernate.cfg.JdbcSettings;
import org.hibernate.dialect.H2Dialect;
import org.hibernate.engine.jdbc.connections.internal.UserSuppliedConnectionProviderImpl;
import org.hibernate.jpa.HibernatePersistenceProvider;
import org.hibernate.jpa.boot.internal.EntityManagerFactoryBuilderImpl;
import org.hibernate.jpa.boot.internal.PersistenceUnitInfoDescriptor;
import org.jspecify.annotations.Nullable;
import org.springframework.data.repository.config.AotRepositoryContext;
import org.springframework.data.util.Lazy;
import org.springframework.instrument.classloading.SimpleThrowawayClassLoader;
import org.springframework.orm.jpa.persistenceunit.MutablePersistenceUnitInfo;
import org.springframework.orm.jpa.persistenceunit.PersistenceManagedTypes;
/**
* AOT metamodel implementation that uses Hibernate to build the metamodel.
*
* @author Christoph Strobl
* @author Mark Paluch
* @since 4.0
*/
class AotMetamodel implements Metamodel {
private final String persistenceUnit;
private final Set<Class<?>> managedTypes;
private final Lazy<EntityManagerFactory> entityManagerFactory = Lazy.of(this::init);
private final Lazy<Metamodel> metamodel = Lazy.of(() -> entityManagerFactory.get().getMetamodel());
private final Lazy<EntityManager> entityManager = Lazy.of(() -> entityManagerFactory.get().createEntityManager());
private final Lazy<EntityManagerFactory> entityManagerFactory;
private final Lazy<EntityManager> entityManager = Lazy.of(() -> getEntityManagerFactory().createEntityManager());
public AotMetamodel(Set<Class<?>> managedTypes) {
this("AotMetamodel", managedTypes);
public AotMetamodel(AotRepositoryContext repositoryContext) {
this(repositoryContext.getResolvedTypes().stream().map(Class::getName)
.filter(name -> !name.startsWith("jakarta.persistence")).toList(), null);
}
private AotMetamodel(String persistenceUnit, Set<Class<?>> managedTypes) {
this.persistenceUnit = persistenceUnit;
this.managedTypes = managedTypes;
public AotMetamodel(PersistenceManagedTypes managedTypes) {
this(managedTypes.getManagedClassNames(), managedTypes.getPersistenceUnitRootUrl());
}
public static AotMetamodel hibernateModel(Class<?>... types) {
return new AotMetamodel(Set.of(types));
}
public static AotMetamodel hibernateModel(String persistenceUnit, Class<?>... types) {
return new AotMetamodel(persistenceUnit, Set.of(types));
}
public <X> EntityType<X> entity(Class<X> cls) {
return metamodel.get().entity(cls);
}
@Override
public EntityType<?> entity(String s) {
return metamodel.get().entity(s);
}
public <X> ManagedType<X> managedType(Class<X> cls) {
return metamodel.get().managedType(cls);
}
public <X> EmbeddableType<X> embeddable(Class<X> cls) {
return metamodel.get().embeddable(cls);
}
public Set<ManagedType<?>> getManagedTypes() {
return metamodel.get().getManagedTypes();
}
public Set<EntityType<?>> getEntities() {
return metamodel.get().getEntities();
}
public Set<EmbeddableType<?>> getEmbeddables() {
return metamodel.get().getEmbeddables();
}
public EntityManager entityManager() {
return entityManager.get();
}
public EntityManagerFactory getEntityManagerFactory() {
return entityManagerFactory.get();
}
EntityManagerFactory init() {
public AotMetamodel(Collection<String> managedTypes, @Nullable URL persistenceUnitRootUrl) {
MutablePersistenceUnitInfo persistenceUnitInfo = new MutablePersistenceUnitInfo() {
@Override
@@ -113,19 +79,82 @@ class AotMetamodel implements Metamodel {
// just ignore it
}
};
persistenceUnitInfo.setPersistenceUnitName("AotMetaModel");
persistenceUnitInfo.setPersistenceUnitName(persistenceUnit);
this.managedTypes.stream().map(Class::getName).forEach(persistenceUnitInfo::addManagedClassName);
this.entityManagerFactory = init(() -> {
persistenceUnitInfo.setPersistenceProviderClassName(HibernatePersistenceProvider.class.getName());
managedTypes.stream().forEach(persistenceUnitInfo::addManagedClassName);
return new EntityManagerFactoryBuilderImpl(new PersistenceUnitInfoDescriptor(persistenceUnitInfo) {
@Override
public List<String> getManagedClassNames() {
return persistenceUnitInfo.getManagedClassNames();
}
}, Map.of("hibernate.dialect", "org.hibernate.dialect.H2Dialect", "hibernate.boot.allow_jdbc_metadata_access",
"false")).build();
persistenceUnitInfo.setPersistenceProviderClassName(HibernatePersistenceProvider.class.getName());
return new PersistenceUnitInfoDescriptor(persistenceUnitInfo) {
@Override
public List<String> getManagedClassNames() {
return persistenceUnitInfo.getManagedClassNames();
}
@Override
public URL getPersistenceUnitRootUrl() {
return persistenceUnitRootUrl;
}
};
});
}
public AotMetamodel(PersistenceUnitInfo unitInfo) {
this.entityManagerFactory = init(() -> new PersistenceUnitInfoDescriptor(unitInfo));
}
static Lazy<EntityManagerFactory> init(Supplier<PersistenceUnitInfoDescriptor> unitInfo) {
return Lazy.of(() -> new EntityManagerFactoryBuilderImpl(unitInfo.get(),
Map.of(JdbcSettings.DIALECT, H2Dialect.class.getName(), //
JdbcSettings.ALLOW_METADATA_ON_BOOT, "false", //
JdbcSettings.CONNECTION_PROVIDER, new UserSuppliedConnectionProviderImpl()))
.build());
}
private Metamodel getMetamodel() {
return getEntityManagerFactory().getMetamodel();
}
public <X> EntityType<X> entity(Class<X> cls) {
return getMetamodel().entity(cls);
}
@Override
public EntityType<?> entity(String s) {
return getMetamodel().entity(s);
}
public <X> ManagedType<X> managedType(Class<X> cls) {
return getMetamodel().managedType(cls);
}
public <X> EmbeddableType<X> embeddable(Class<X> cls) {
return getMetamodel().embeddable(cls);
}
public Set<ManagedType<?>> getManagedTypes() {
return getMetamodel().getManagedTypes();
}
public Set<EntityType<?>> getEntities() {
return getMetamodel().getEntities();
}
public Set<EmbeddableType<?>> getEmbeddables() {
return getMetamodel().getEmbeddables();
}
public EntityManager entityManager() {
return entityManager.get();
}
public EntityManagerFactory getEntityManagerFactory() {
return entityManagerFactory.get();
}
}

View File

@@ -17,10 +17,11 @@ package org.springframework.data.jpa.repository.aot;
import jakarta.persistence.EntityManager;
import jakarta.persistence.EntityManagerFactory;
import jakarta.persistence.metamodel.Metamodel;
import jakarta.persistence.spi.PersistenceUnitInfo;
import java.lang.reflect.Method;
import java.util.Map;
import java.util.stream.Collectors;
import org.jspecify.annotations.Nullable;
@@ -49,6 +50,7 @@ import org.springframework.data.repository.query.ReturnedType;
import org.springframework.data.util.TypeInformation;
import org.springframework.javapoet.CodeBlock;
import org.springframework.javapoet.TypeName;
import org.springframework.orm.jpa.persistenceunit.PersistenceManagedTypes;
import org.springframework.util.ClassUtils;
import org.springframework.util.StringUtils;
@@ -64,31 +66,43 @@ import org.springframework.util.StringUtils;
*/
public class JpaRepositoryContributor extends RepositoryContributor {
private final Metamodel metamodel;
private final PersistenceProvider persistenceProvider;
private final QueriesFactory queriesFactory;
private final EntityGraphLookup entityGraphLookup;
public JpaRepositoryContributor(AotRepositoryContext repositoryContext) {
this(repositoryContext, new AotMetamodel(repositoryContext));
}
super(repositoryContext);
public JpaRepositoryContributor(AotRepositoryContext repositoryContext, PersistenceUnitInfo unitInfo) {
this(repositoryContext, new AotMetamodel(unitInfo));
}
AotMetamodel amm = new AotMetamodel(repositoryContext.getResolvedTypes().stream()
.filter(it -> !it.getName().startsWith("jakarta.persistence")).collect(Collectors.toSet()));
this.persistenceProvider = PersistenceProvider.fromEntityManagerFactory(amm.getEntityManagerFactory());
this.queriesFactory = new QueriesFactory(amm.getEntityManagerFactory(), amm);
this.entityGraphLookup = new EntityGraphLookup(amm.getEntityManagerFactory());
public JpaRepositoryContributor(AotRepositoryContext repositoryContext, PersistenceManagedTypes managedTypes) {
this(repositoryContext, new AotMetamodel(managedTypes));
}
public JpaRepositoryContributor(AotRepositoryContext repositoryContext, EntityManagerFactory entityManagerFactory) {
super(repositoryContext);
this.metamodel = entityManagerFactory.getMetamodel();
this.persistenceProvider = PersistenceProvider.fromEntityManagerFactory(entityManagerFactory);
this.queriesFactory = new QueriesFactory(entityManagerFactory);
this.entityGraphLookup = new EntityGraphLookup(entityManagerFactory);
}
private JpaRepositoryContributor(AotRepositoryContext repositoryContext, AotMetamodel metamodel) {
super(repositoryContext);
this.metamodel = metamodel;
this.persistenceProvider = PersistenceProvider.fromEntityManagerFactory(metamodel.getEntityManagerFactory());
this.queriesFactory = new QueriesFactory(metamodel.getEntityManagerFactory(), metamodel);
this.entityGraphLookup = new EntityGraphLookup(metamodel.getEntityManagerFactory());
}
@Override
protected void customizeClass(AotRepositoryClassBuilder classBuilder) {
classBuilder.customize(builder -> builder.superclass(TypeName.get(AotRepositoryFragmentSupport.class)));
@@ -203,6 +217,10 @@ public class JpaRepositoryContributor extends RepositoryContributor {
});
}
public Metamodel getMetamodel() {
return metamodel;
}
record StoredProcedureMetadata(String procedure) implements QueryMetadata {
@Override

View File

@@ -22,6 +22,7 @@ import jakarta.persistence.EntityManagerFactory;
import jakarta.persistence.MappedSuperclass;
import jakarta.persistence.PersistenceContext;
import jakarta.persistence.PersistenceUnit;
import jakarta.persistence.spi.PersistenceUnitInfo;
import java.lang.annotation.Annotation;
import java.util.Arrays;
@@ -35,8 +36,11 @@ import java.util.Optional;
import java.util.Set;
import org.jspecify.annotations.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.aot.generate.GenerationContext;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.beans.factory.aot.BeanRegistrationAotProcessor;
import org.springframework.beans.factory.config.BeanDefinition;
import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
@@ -46,6 +50,7 @@ import org.springframework.beans.factory.support.BeanDefinitionRegistry;
import org.springframework.beans.factory.support.RootBeanDefinition;
import org.springframework.context.annotation.AnnotationConfigUtils;
import org.springframework.core.annotation.AnnotationAttributes;
import org.springframework.core.env.Environment;
import org.springframework.core.io.ResourceLoader;
import org.springframework.dao.DataAccessException;
import org.springframework.dao.annotation.PersistenceExceptionTranslationPostProcessor;
@@ -57,13 +62,14 @@ import org.springframework.data.jpa.repository.support.EntityManagerBeanDefiniti
import org.springframework.data.jpa.repository.support.JpaEvaluationContextExtension;
import org.springframework.data.jpa.repository.support.JpaRepositoryFactoryBean;
import org.springframework.data.jpa.repository.support.SimpleJpaRepository;
import org.springframework.data.repository.aot.generate.RepositoryContributor;
import org.springframework.data.repository.config.AnnotationRepositoryConfigurationSource;
import org.springframework.data.repository.config.AotRepositoryContext;
import org.springframework.data.repository.config.RepositoryConfigurationExtensionSupport;
import org.springframework.data.repository.config.RepositoryConfigurationSource;
import org.springframework.data.repository.config.RepositoryRegistrationAotProcessor;
import org.springframework.data.repository.config.XmlRepositoryConfigurationSource;
import org.springframework.orm.jpa.persistenceunit.PersistenceManagedTypes;
import org.springframework.orm.jpa.persistenceunit.PersistenceUnitManager;
import org.springframework.orm.jpa.support.PersistenceAnnotationBeanPostProcessor;
import org.springframework.util.ClassUtils;
import org.springframework.util.StringUtils;
@@ -90,6 +96,7 @@ public class JpaRepositoryConfigExtension extends RepositoryConfigurationExtensi
private static final String ENABLE_DEFAULT_TRANSACTIONS_ATTRIBUTE = "enableDefaultTransactions";
private static final String JPA_METAMODEL_CACHE_CLEANUP_CLASSNAME = "org.springframework.data.jpa.util.JpaMetamodelCacheCleanup";
private static final String ESCAPE_CHARACTER_PROPERTY = "escapeCharacter";
private static final Logger log = LoggerFactory.getLogger(JpaRepositoryConfigExtension.class);
private final Map<Object, String> entityManagerRefs = new LinkedHashMap<>();
@@ -327,28 +334,60 @@ public class JpaRepositoryConfigExtension extends RepositoryConfigurationExtensi
String GENERATED_REPOSITORIES_JPA_USE_ENTITY_MANAGER = "spring.aot.jpa.repositories.use-entitymanager";
protected @Nullable RepositoryContributor contribute(AotRepositoryContext repositoryContext,
protected @Nullable JpaRepositoryContributor contribute(AotRepositoryContext repositoryContext,
GenerationContext generationContext) {
Environment environment = repositoryContext.getEnvironment();
boolean enabled = Boolean.parseBoolean(
repositoryContext.getEnvironment().getProperty(AotContext.GENERATED_REPOSITORIES_ENABLED, "false"));
environment.getProperty(AotContext.GENERATED_REPOSITORIES_ENABLED, "false"));
if (!enabled) {
return null;
}
ConfigurableListableBeanFactory beanFactory = repositoryContext.getBeanFactory();
boolean useEntityManager = Boolean.parseBoolean(
repositoryContext.getEnvironment().getProperty(GENERATED_REPOSITORIES_JPA_USE_ENTITY_MANAGER, "false"));
environment.getProperty(GENERATED_REPOSITORIES_JPA_USE_ENTITY_MANAGER, "false"));
if (useEntityManager) {
ConfigurableListableBeanFactory beanFactory = repositoryContext.getBeanFactory();
EntityManagerFactory emf = beanFactory.getBeanProvider(EntityManagerFactory.class).getIfAvailable();
ObjectProvider<PersistenceUnitManager> unitManagerProvider = beanFactory
.getBeanProvider(PersistenceUnitManager.class);
PersistenceUnitManager unitManager = unitManagerProvider.getIfAvailable();
if (emf != null) {
return new JpaRepositoryContributor(repositoryContext, emf);
if (unitManager != null) {
log.debug("Using PersistenceUnitManager for AOT repository generation");
return new JpaRepositoryContributor(repositoryContext, unitManager.obtainDefaultPersistenceUnitInfo());
}
log.debug("Using EntityManager for AOT repository generation");
EntityManagerFactory emf = beanFactory.getBean(EntityManagerFactory.class);
return new JpaRepositoryContributor(repositoryContext, emf);
}
ObjectProvider<PersistenceManagedTypes> managedTypesProvider = beanFactory
.getBeanProvider(PersistenceManagedTypes.class);
PersistenceManagedTypes managedTypes = managedTypesProvider.getIfAvailable();
if (managedTypes != null) {
log.debug("Using PersistenceManagedTypes for AOT repository generation");
return new JpaRepositoryContributor(repositoryContext, managedTypes);
}
ObjectProvider<PersistenceUnitInfo> infoProvider = beanFactory.getBeanProvider(PersistenceUnitInfo.class);
PersistenceUnitInfo unitInfo = infoProvider.getIfAvailable();
if (unitInfo != null) {
log.debug("Using PersistenceUnitInfo for AOT repository generation");
return new JpaRepositoryContributor(repositoryContext, unitInfo);
}
log.debug("Using scanned types for AOT repository generation");
return new JpaRepositoryContributor(repositoryContext);
}
}

View File

@@ -18,24 +18,39 @@ package org.springframework.data.jpa.repository.config;
import static org.assertj.core.api.Assertions.*;
import jakarta.persistence.Entity;
import jakarta.persistence.Id;
import java.lang.annotation.Annotation;
import java.net.URL;
import java.util.Collections;
import java.util.List;
import java.util.Set;
import org.jspecify.annotations.Nullable;
import org.junit.jupiter.api.Test;
import org.springframework.aot.generate.ClassNameGenerator;
import org.springframework.aot.generate.DefaultGenerationContext;
import org.springframework.aot.generate.GenerationContext;
import org.springframework.aot.generate.InMemoryGeneratedFiles;
import org.springframework.aot.hint.predicate.RuntimeHintsPredicates;
import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
import org.springframework.context.support.AbstractApplicationContext;
import org.springframework.context.support.GenericApplicationContext;
import org.springframework.core.annotation.MergedAnnotation;
import org.springframework.core.env.Environment;
import org.springframework.core.env.StandardEnvironment;
import org.springframework.data.aot.AotContext;
import org.springframework.data.jpa.repository.aot.JpaRepositoryContributor;
import org.springframework.data.jpa.repository.support.SimpleJpaRepository;
import org.springframework.data.repository.Repository;
import org.springframework.data.repository.config.AotRepositoryContext;
import org.springframework.data.repository.config.AotRepositoryInformation;
import org.springframework.data.repository.core.RepositoryInformation;
import org.springframework.data.repository.core.support.AbstractRepositoryMetadata;
import org.springframework.javapoet.ClassName;
import org.springframework.mock.env.MockPropertySource;
import org.springframework.orm.jpa.persistenceunit.PersistenceManagedTypes;
/**
* @author Christoph Strobl
@@ -49,7 +64,7 @@ class JpaRepositoryRegistrationAotProcessorUnitTests {
new InMemoryGeneratedFiles());
new JpaRepositoryConfigExtension.JpaRepositoryRegistrationAotProcessor()
.contribute(new DummyAotRepositoryContext() {
.contribute(new DummyAotRepositoryContext(null) {
@Override
public Set<Class<?>> getResolvedTypes() {
return Collections.singleton(Person.class);
@@ -66,7 +81,7 @@ class JpaRepositoryRegistrationAotProcessorUnitTests {
new InMemoryGeneratedFiles());
new JpaRepositoryConfigExtension.JpaRepositoryRegistrationAotProcessor()
.contribute(new DummyAotRepositoryContext() {
.contribute(new DummyAotRepositoryContext(null) {
@Override
public Set<MergedAnnotation<Annotation>> getResolvedAnnotations() {
@@ -79,10 +94,57 @@ class JpaRepositoryRegistrationAotProcessorUnitTests {
assertThat(RuntimeHintsPredicates.reflection().onType(Entity.class)).rejects(ctx.getRuntimeHints());
}
static class Person {}
@Test // GH-3838
void repositoryProcessorShouldConsiderPersistenceManagedTypes() {
GenerationContext ctx = new DefaultGenerationContext(new ClassNameGenerator(ClassName.OBJECT),
new InMemoryGeneratedFiles());
GenericApplicationContext context = new GenericApplicationContext();
context.registerBean(PersistenceManagedTypes.class, () -> {
return new PersistenceManagedTypes() {
@Override
public List<String> getManagedClassNames() {
return List.of(Person.class.getName());
}
@Override
public List<String> getManagedPackages() {
return List.of();
}
@Override
public @Nullable URL getPersistenceUnitRootUrl() {
return null;
}
};
});
context.getEnvironment().getPropertySources()
.addFirst(new MockPropertySource().withProperty(AotContext.GENERATED_REPOSITORIES_ENABLED, "true"));
JpaRepositoryContributor contributor = new JpaRepositoryConfigExtension.JpaRepositoryRegistrationAotProcessor()
.contribute(new DummyAotRepositoryContext(context), ctx);
assertThat(contributor.getMetamodel().managedType(Person.class)).isNotNull();
}
@Entity
static class Person {
@Id Long id;
}
interface PersonRepository extends Repository<Person, Long> {}
static class DummyAotRepositoryContext implements AotRepositoryContext {
private final @Nullable AbstractApplicationContext applicationContext;
DummyAotRepositoryContext(@Nullable AbstractApplicationContext applicationContext) {
this.applicationContext = applicationContext;
}
@Override
public String getBeanName() {
return "jpaRepository";
@@ -105,7 +167,8 @@ class JpaRepositoryRegistrationAotProcessorUnitTests {
@Override
public RepositoryInformation getRepositoryInformation() {
return null;
return new AotRepositoryInformation(AbstractRepositoryMetadata.getMetadata(PersonRepository.class),
SimpleJpaRepository.class, List.of());
}
@Override
@@ -120,12 +183,12 @@ class JpaRepositoryRegistrationAotProcessorUnitTests {
@Override
public ConfigurableListableBeanFactory getBeanFactory() {
return null;
return applicationContext != null ? applicationContext.getBeanFactory() : null;
}
@Override
public Environment getEnvironment() {
return new StandardEnvironment();
return applicationContext == null ? new StandardEnvironment() : applicationContext.getEnvironment();
}
@Override