Adapt JPA auto-configuration to PersistenceManagedTypes

This commit exposes a PersistenceManagedTypes bean with the entities
to consider in a typical auto-configuration scenario. This allows the
result of the scanning to be optimized AOT, if necessary.

Closes gh-32119

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/orm/jpa/JpaBaseConfiguration.java

@@ -25,11 +25,8 @@ import jakarta.persistence.EntityManagerFactory;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 
-import org.springframework.beans.BeansException;
 import org.springframework.beans.factory.BeanFactory;
-import org.springframework.beans.factory.BeanFactoryAware;
 import org.springframework.beans.factory.ObjectProvider;
-import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
 import org.springframework.boot.autoconfigure.AutoConfigurationPackages;
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
@@ -45,9 +42,12 @@ import org.springframework.boot.orm.jpa.EntityManagerFactoryBuilder;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.context.annotation.Primary;
+import org.springframework.core.io.ResourceLoader;
 import org.springframework.orm.jpa.JpaTransactionManager;
 import org.springframework.orm.jpa.JpaVendorAdapter;
 import org.springframework.orm.jpa.LocalContainerEntityManagerFactoryBean;
+import org.springframework.orm.jpa.persistenceunit.PersistenceManagedTypes;
+import org.springframework.orm.jpa.persistenceunit.PersistenceManagedTypesScanner;
 import org.springframework.orm.jpa.persistenceunit.PersistenceUnitManager;
 import org.springframework.orm.jpa.support.OpenEntityManagerInViewFilter;
 import org.springframework.orm.jpa.support.OpenEntityManagerInViewInterceptor;
@@ -73,7 +73,7 @@ import org.springframework.web.servlet.config.annotation.WebMvcConfigurer;
  */
 @Configuration(proxyBeanMethods = false)
 @EnableConfigurationProperties(JpaProperties.class)
-public abstract class JpaBaseConfiguration implements BeanFactoryAware {
+public abstract class JpaBaseConfiguration {
 
 	private final DataSource dataSource;
 
@@ -81,8 +81,6 @@ public abstract class JpaBaseConfiguration implements BeanFactoryAware {
 
 	private final JtaTransactionManager jtaTransactionManager;
 
-	private ConfigurableListableBeanFactory beanFactory;
-
 	protected JpaBaseConfiguration(DataSource dataSource, JpaProperties properties,
 			ObjectProvider<JtaTransactionManager> jtaTransactionManager) {
 		this.dataSource = dataSource;
@@ -128,11 +126,12 @@ public abstract class JpaBaseConfiguration implements BeanFactoryAware {
 	@Bean
 	@Primary
 	@ConditionalOnMissingBean({ LocalContainerEntityManagerFactoryBean.class, EntityManagerFactory.class })
-	public LocalContainerEntityManagerFactoryBean entityManagerFactory(EntityManagerFactoryBuilder factoryBuilder) {
+	public LocalContainerEntityManagerFactoryBean entityManagerFactory(EntityManagerFactoryBuilder factoryBuilder,
+			PersistenceManagedTypes persistenceManagedTypes) {
 		Map<String, Object> vendorProperties = getVendorProperties();
 		customizeVendorProperties(vendorProperties);
-		return factoryBuilder.dataSource(this.dataSource).packages(getPackagesToScan()).properties(vendorProperties)
-				.mappingResources(getMappingResources()).jta(isJta()).build();
+		return factoryBuilder.dataSource(this.dataSource).managedTypes(persistenceManagedTypes)
+				.properties(vendorProperties).mappingResources(getMappingResources()).jta(isJta()).build();
 	}
 
 	protected abstract AbstractJpaVendorAdapter createJpaVendorAdapter();
@@ -147,14 +146,6 @@ public abstract class JpaBaseConfiguration implements BeanFactoryAware {
 	protected void customizeVendorProperties(Map<String, Object> vendorProperties) {
 	}
 
-	protected String[] getPackagesToScan() {
-		List<String> packages = EntityScanPackages.get(this.beanFactory).getPackageNames();
-		if (packages.isEmpty() && AutoConfigurationPackages.has(this.beanFactory)) {
-			packages = AutoConfigurationPackages.get(this.beanFactory);
-		}
-		return StringUtils.toStringArray(packages);
-	}
-
 	private String[] getMappingResources() {
 		List<String> mappingResources = this.properties.getMappingResources();
 		return (!ObjectUtils.isEmpty(mappingResources) ? StringUtils.toStringArray(mappingResources) : null);
@@ -192,9 +183,26 @@ public abstract class JpaBaseConfiguration implements BeanFactoryAware {
 		return this.dataSource;
 	}
 
-	@Override
-	public void setBeanFactory(BeanFactory beanFactory) throws BeansException {
-		this.beanFactory = (ConfigurableListableBeanFactory) beanFactory;
+	@Configuration(proxyBeanMethods = false)
+	@ConditionalOnMissingBean({ LocalContainerEntityManagerFactoryBean.class, EntityManagerFactory.class })
+	static class PersistenceManagedTypesConfiguration {
+
+		@Bean
+		@Primary
+		@ConditionalOnMissingBean
+		PersistenceManagedTypes persistenceManagedTypes(BeanFactory beanFactory, ResourceLoader resourceLoader) {
+			String[] packagesToScan = getPackagesToScan(beanFactory);
+			return new PersistenceManagedTypesScanner(resourceLoader).scan(packagesToScan);
+		}
+
+		private static String[] getPackagesToScan(BeanFactory beanFactory) {
+			List<String> packages = EntityScanPackages.get(beanFactory).getPackageNames();
+			if (packages.isEmpty() && AutoConfigurationPackages.has(beanFactory)) {
+				packages = AutoConfigurationPackages.get(beanFactory);
+			}
+			return StringUtils.toStringArray(packages);
+		}
+
 	}
 
 	@Configuration(proxyBeanMethods = false)

spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/orm/jpa/AbstractJpaAutoConfigurationTests.java

@@ -23,13 +23,16 @@ import java.util.UUID;
 
 import javax.sql.DataSource;
 
+import jakarta.persistence.EntityManager;
 import jakarta.persistence.EntityManagerFactory;
+import jakarta.persistence.metamodel.ManagedType;
 import jakarta.persistence.spi.PersistenceUnitInfo;
 import org.hibernate.engine.transaction.jta.platform.internal.NoJtaPlatform;
 import org.junit.jupiter.api.Test;
 
 import org.springframework.boot.autoconfigure.AutoConfigurations;
 import org.springframework.boot.autoconfigure.TestAutoConfigurationPackage;
+import org.springframework.boot.autoconfigure.data.jpa.country.Country;
 import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
 import org.springframework.boot.autoconfigure.jdbc.DataSourceTransactionManagerAutoConfiguration;
 import org.springframework.boot.autoconfigure.orm.jpa.test.City;
@@ -49,6 +52,7 @@ import org.springframework.orm.jpa.JpaTransactionManager;
 import org.springframework.orm.jpa.JpaVendorAdapter;
 import org.springframework.orm.jpa.LocalContainerEntityManagerFactoryBean;
 import org.springframework.orm.jpa.persistenceunit.DefaultPersistenceUnitManager;
+import org.springframework.orm.jpa.persistenceunit.PersistenceManagedTypes;
 import org.springframework.orm.jpa.persistenceunit.PersistenceUnitManager;
 import org.springframework.orm.jpa.support.OpenEntityManagerInViewFilter;
 import org.springframework.orm.jpa.support.OpenEntityManagerInViewInterceptor;
@@ -112,6 +116,7 @@ abstract class AbstractJpaAutoConfigurationTests {
 			assertThat(context).hasSingleBean(DataSource.class);
 			assertThat(context).hasSingleBean(JpaTransactionManager.class);
 			assertThat(context).hasSingleBean(EntityManagerFactory.class);
+			assertThat(context).hasSingleBean(PersistenceManagedTypes.class);
 		});
 	}
 
@@ -121,6 +126,7 @@ abstract class AbstractJpaAutoConfigurationTests {
 			assertThat(context).getBeans(DataSource.class).hasSize(2);
 			assertThat(context).hasSingleBean(JpaTransactionManager.class);
 			assertThat(context).hasSingleBean(EntityManagerFactory.class);
+			assertThat(context).hasSingleBean(PersistenceManagedTypes.class);
 		});
 	}
 
@@ -225,6 +231,28 @@ abstract class AbstractJpaAutoConfigurationTests {
 		});
 	}
 
+	@Test
+	void defaultPersistenceManagedTypes() {
+		this.contextRunner.run((context) -> {
+			assertThat(context).hasSingleBean(PersistenceManagedTypes.class);
+			EntityManager entityManager = context.getBean(EntityManagerFactory.class).createEntityManager();
+			assertThat(entityManager.getMetamodel().getManagedTypes().stream().map(ManagedType::getJavaType)
+					.toArray(Class<?>[]::new)).contains(City.class).doesNotContain(Country.class);
+		});
+	}
+
+	@Test
+	void customPersistenceManagedTypes() {
+		this.contextRunner
+				.withBean(PersistenceManagedTypes.class, () -> PersistenceManagedTypes.of(Country.class.getName()))
+				.run((context) -> {
+					assertThat(context).hasSingleBean(PersistenceManagedTypes.class);
+					EntityManager entityManager = context.getBean(EntityManagerFactory.class).createEntityManager();
+					assertThat(entityManager.getMetamodel().getManagedTypes().stream().map(ManagedType::getJavaType)
+							.toArray(Class<?>[]::new)).contains(Country.class).doesNotContain(City.class);
+				});
+	}
+
 	@Test
 	void customPersistenceUnitManager() {
 		this.contextRunner.withUserConfiguration(TestConfigurationWithCustomPersistenceUnitManager.class)

spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/orm/jpa/domain/country/Country.java

@@ -0,0 +1,56 @@
+/*
+ * Copyright 2012-2022 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.boot.autoconfigure.orm.jpa.domain.country;
+
+import java.io.Serializable;
+
+import jakarta.persistence.Column;
+import jakarta.persistence.Entity;
+import jakarta.persistence.GeneratedValue;
+import jakarta.persistence.Id;
+import org.hibernate.envers.Audited;
+
+@Entity
+public class Country implements Serializable {
+
+	private static final long serialVersionUID = 1L;
+
+	@Id
+	@GeneratedValue
+	private Long id;
+
+	@Audited
+	@Column
+	private String name;
+
+	public Long getId() {
+		return this.id;
+	}
+
+	public void setId(Long id) {
+		this.id = id;
+	}
+
+	public String getName() {
+		return this.name;
+	}
+
+	public void setName(String name) {
+		this.name = name;
+	}
+
+}

spring-boot-project/spring-boot/src/main/java/org/springframework/boot/orm/jpa/EntityManagerFactoryBuilder.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2021 the original author or authors.
+ * Copyright 2012-2022 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -28,6 +28,7 @@ import javax.sql.DataSource;
 import org.springframework.core.task.AsyncTaskExecutor;
 import org.springframework.orm.jpa.JpaVendorAdapter;
 import org.springframework.orm.jpa.LocalContainerEntityManagerFactoryBean;
+import org.springframework.orm.jpa.persistenceunit.PersistenceManagedTypes;
 import org.springframework.orm.jpa.persistenceunit.PersistenceUnitManager;
 import org.springframework.orm.jpa.persistenceunit.PersistenceUnitPostProcessor;
 import org.springframework.util.ClassUtils;
@@ -125,6 +126,8 @@ public class EntityManagerFactoryBuilder {
 
 		private DataSource dataSource;
 
+		private PersistenceManagedTypes managedTypes;
+
 		private String[] packagesToScan;
 
 		private String persistenceUnit;
@@ -139,10 +142,22 @@ public class EntityManagerFactoryBuilder {
 			this.dataSource = dataSource;
 		}
 
+		/**
+		 * The persistence managed types, providing both the managed entities and packages
+		 * the entity manager should consider.
+		 * @param managedTypes managed types.
+		 * @return the builder for fluent usage
+		 */
+		public Builder managedTypes(PersistenceManagedTypes managedTypes) {
+			this.managedTypes = managedTypes;
+			return this;
+		}
+
 		/**
 		 * The names of packages to scan for {@code @Entity} annotations.
 		 * @param packagesToScan packages to scan
 		 * @return the builder for fluent usage
+		 * @see #managedTypes(PersistenceManagedTypes)
 		 */
 		public Builder packages(String... packagesToScan) {
 			this.packagesToScan = packagesToScan;
@@ -153,6 +168,7 @@ public class EntityManagerFactoryBuilder {
 		 * The classes whose packages should be scanned for {@code @Entity} annotations.
 		 * @param basePackageClasses the classes to use
 		 * @return the builder for fluent usage
+		 * @see #managedTypes(PersistenceManagedTypes)
 		 */
 		public Builder packages(Class<?>... basePackageClasses) {
 			Set<String> packages = new HashSet<>();
@@ -233,7 +249,12 @@ public class EntityManagerFactoryBuilder {
 			else {
 				entityManagerFactoryBean.setDataSource(this.dataSource);
 			}
-			entityManagerFactoryBean.setPackagesToScan(this.packagesToScan);
+			if (this.managedTypes != null) {
+				entityManagerFactoryBean.setManagedTypes(this.managedTypes);
+			}
+			else {
+				entityManagerFactoryBean.setPackagesToScan(this.packagesToScan);
+			}
 			entityManagerFactoryBean.getJpaPropertyMap().putAll(EntityManagerFactoryBuilder.this.jpaProperties);
 			entityManagerFactoryBean.getJpaPropertyMap().putAll(this.properties);
 			if (!ObjectUtils.isEmpty(this.mappingResources)) {