From eeca4df36f98f94eeaa27aa12efd5ab5dfda2997 Mon Sep 17 00:00:00 2001 From: Mark Paluch Date: Tue, 13 May 2025 15:59:41 +0200 Subject: [PATCH] Detect PersistenceUnitManager and PersistenceManaged types in JpaRepositoryContributor. See #3875 --- .../data/jpa/repository/aot/AotMetamodel.java | 161 +++++++++++------- .../aot/JpaRepositoryContributor.java | 34 +++- .../config/JpaRepositoryConfigExtension.java | 55 +++++- ...toryRegistrationAotProcessorUnitTests.java | 75 +++++++- 4 files changed, 237 insertions(+), 88 deletions(-) diff --git a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/AotMetamodel.java b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/AotMetamodel.java index 2b3f49bb2..ee8ecd7f5 100644 --- a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/AotMetamodel.java +++ b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/AotMetamodel.java @@ -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> managedTypes; - private final Lazy entityManagerFactory = Lazy.of(this::init); - private final Lazy metamodel = Lazy.of(() -> entityManagerFactory.get().getMetamodel()); - private final Lazy entityManager = Lazy.of(() -> entityManagerFactory.get().createEntityManager()); + private final Lazy entityManagerFactory; + private final Lazy entityManager = Lazy.of(() -> getEntityManagerFactory().createEntityManager()); - public AotMetamodel(Set> 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> 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 EntityType entity(Class cls) { - return metamodel.get().entity(cls); - } - - @Override - public EntityType entity(String s) { - return metamodel.get().entity(s); - } - - public ManagedType managedType(Class cls) { - return metamodel.get().managedType(cls); - } - - public EmbeddableType embeddable(Class cls) { - return metamodel.get().embeddable(cls); - } - - public Set> getManagedTypes() { - return metamodel.get().getManagedTypes(); - } - - public Set> getEntities() { - return metamodel.get().getEntities(); - } - - public Set> getEmbeddables() { - return metamodel.get().getEmbeddables(); - } - - public EntityManager entityManager() { - return entityManager.get(); - } - - public EntityManagerFactory getEntityManagerFactory() { - return entityManagerFactory.get(); - } - - EntityManagerFactory init() { + public AotMetamodel(Collection 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 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 getManagedClassNames() { + return persistenceUnitInfo.getManagedClassNames(); + } + + @Override + public URL getPersistenceUnitRootUrl() { + return persistenceUnitRootUrl; + } + + }; + }); + } + + public AotMetamodel(PersistenceUnitInfo unitInfo) { + this.entityManagerFactory = init(() -> new PersistenceUnitInfoDescriptor(unitInfo)); + } + + static Lazy init(Supplier 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 EntityType entity(Class cls) { + return getMetamodel().entity(cls); + } + + @Override + public EntityType entity(String s) { + return getMetamodel().entity(s); + } + + public ManagedType managedType(Class cls) { + return getMetamodel().managedType(cls); + } + + public EmbeddableType embeddable(Class cls) { + return getMetamodel().embeddable(cls); + } + + public Set> getManagedTypes() { + return getMetamodel().getManagedTypes(); + } + + public Set> getEntities() { + return getMetamodel().getEntities(); + } + + public Set> getEmbeddables() { + return getMetamodel().getEmbeddables(); + } + + public EntityManager entityManager() { + return entityManager.get(); + } + + public EntityManagerFactory getEntityManagerFactory() { + return entityManagerFactory.get(); } } diff --git a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/JpaRepositoryContributor.java b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/JpaRepositoryContributor.java index 1dcb10809..41bf6ff6e 100644 --- a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/JpaRepositoryContributor.java +++ b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/JpaRepositoryContributor.java @@ -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 diff --git a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/config/JpaRepositoryConfigExtension.java b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/config/JpaRepositoryConfigExtension.java index ce3218593..360b449be 100644 --- a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/config/JpaRepositoryConfigExtension.java +++ b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/config/JpaRepositoryConfigExtension.java @@ -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 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 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 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 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); } } diff --git a/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/config/JpaRepositoryRegistrationAotProcessorUnitTests.java b/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/config/JpaRepositoryRegistrationAotProcessorUnitTests.java index 44c260dcb..d75531c50 100644 --- a/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/config/JpaRepositoryRegistrationAotProcessorUnitTests.java +++ b/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/config/JpaRepositoryRegistrationAotProcessorUnitTests.java @@ -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> 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> 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 getManagedClassNames() { + return List.of(Person.class.getName()); + } + + @Override + public List 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 {} 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