Adapt to AOT Infrastructure changes in Commons.

mp911de · mp911de · commit a5b5c2e610a3 · 2025-09-24T10:07:42.000+02:00
See spring-projects/spring-data-commons#3267
diff --git a/spring-data-jpa/src/jmh/java/org/springframework/data/jpa/benchmark/AotRepositoryQueryMethodBenchmarks.java b/spring-data-jpa/src/jmh/java/org/springframework/data/jpa/benchmark/AotRepositoryQueryMethodBenchmarks.java
@@ -40,6 +40,7 @@
 import org.springframework.aot.test.generate.TestGenerationContext;
 import org.springframework.beans.factory.support.BeanDefinitionRegistry;
 import org.springframework.beans.factory.support.DefaultBeanNameGenerator;
+import org.springframework.beans.factory.support.DefaultListableBeanFactory;
 import org.springframework.core.env.StandardEnvironment;
 import org.springframework.core.io.DefaultResourceLoader;
 import org.springframework.core.test.tools.TestCompiler;
@@ -56,6 +57,7 @@
 import org.springframework.data.projection.ProjectionFactory;
 import org.springframework.data.projection.SpelAwareProxyProjectionFactory;
 import org.springframework.data.repository.config.AnnotationRepositoryConfigurationSource;
+import org.springframework.data.repository.config.RepositoryConfigurationSource;
 import org.springframework.data.repository.core.RepositoryMetadata;
 import org.springframework.data.repository.core.support.RepositoryComposition;
 import org.springframework.data.repository.core.support.RepositoryFactoryBeanSupport;
@@ -82,11 +84,16 @@ public class AotRepositoryQueryMethodBenchmarks {
 	public static class BenchmarkParameters {
 
 		public static Class<?> aot;
-		public static TestJpaAotRepositoryContext<PersonRepository> repositoryContext = new TestJpaAotRepositoryContext<>(
-				PersonRepository.class, null,
-				new AnnotationRepositoryConfigurationSource(AnnotationMetadata.introspect(SampleConfig.class),
-						EnableJpaRepositories.class, new DefaultResourceLoader(), new StandardEnvironment(),
-						Mockito.mock(BeanDefinitionRegistry.class), DefaultBeanNameGenerator.INSTANCE));
+		public static TestJpaAotRepositoryContext<PersonRepository> repositoryContext;
+
+		static {
+			RepositoryConfigurationSource configurationSource = new AnnotationRepositoryConfigurationSource(
+					AnnotationMetadata.introspect(SampleConfig.class), EnableJpaRepositories.class, new DefaultResourceLoader(),
+					new StandardEnvironment(), Mockito.mock(BeanDefinitionRegistry.class), DefaultBeanNameGenerator.INSTANCE);
+
+			repositoryContext = new TestJpaAotRepositoryContext<>(new DefaultListableBeanFactory(), PersonRepository.class,
+					null, configurationSource);
+		}
 
 		EntityManager entityManager;
 		RepositoryComposition.RepositoryFragments fragments;
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
@@ -58,6 +58,7 @@
 import org.springframework.core.io.ResourceLoader;
 import org.springframework.dao.DataAccessException;
 import org.springframework.dao.annotation.PersistenceExceptionTranslationPostProcessor;
+import org.springframework.data.aot.AotContext;
 import org.springframework.data.jpa.repository.JpaRepository;
 import org.springframework.data.jpa.repository.aot.JpaRepositoryContributor;
 import org.springframework.data.jpa.repository.support.DefaultJpaContext;
@@ -374,10 +375,22 @@ static boolean isActive(@Nullable ClassLoader classLoader) {
 	public static class JpaRepositoryRegistrationAotProcessor extends RepositoryRegistrationAotProcessor {
 
 		public static final String USE_ENTITY_MANAGER = "spring.aot.jpa.repositories.use-entitymanager";
+
 		private static final String MODULE_NAME = "jpa";
 
-		protected @Nullable JpaRepositoryContributor contribute(AotRepositoryContext repositoryContext,
+		@Override
+		protected void configureTypeContributions(AotRepositoryContext repositoryContext,
 				GenerationContext generationContext) {
+			super.configureTypeContributions(repositoryContext, generationContext);
+		}
+
+		@Override
+		protected void configureTypeContribution(Class<?> type, AotContext aotContext) {
+			aotContext.typeConfiguration(type, config -> config.contributeAccessors().forQuerydsl());
+		}
+
+		@Override
+		protected @Nullable JpaRepositoryContributor contributeAotRepository(AotRepositoryContext repositoryContext) {
 
 			if (!repositoryContext.isGeneratedRepositoriesEnabled(MODULE_NAME)) {
 				return null;
diff --git a/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/aot/AotContributionIntegrationTests.java b/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/aot/AotContributionIntegrationTests.java
@@ -30,10 +30,8 @@
 import org.springframework.context.aot.ApplicationContextAotGenerator;
 import org.springframework.core.io.InputStreamResource;
 import org.springframework.core.io.InputStreamSource;
-import org.springframework.data.aot.AotContext;
 import org.springframework.data.jpa.repository.config.EnableJpaRepositories;
 import org.springframework.data.jpa.repository.config.InfrastructureConfig;
-import org.springframework.mock.env.MockPropertySource;
 
 /**
  * Integration tests for AOT processing.
@@ -70,8 +68,6 @@ void shouldGenerateMetadataForBaseRepositoryAndQuerydslFragment() throws IOExcep
 	private static TestGenerationContext generate(Class<?>... configurationClasses) {
 
 		AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext();
-		context.getEnvironment().getPropertySources()
-				.addFirst(new MockPropertySource().withProperty(AotContext.GENERATED_REPOSITORIES_ENABLED, "true"));
 		context.register(configurationClasses);
 
 		ApplicationContextAotGenerator generator = new ApplicationContextAotGenerator();
diff --git a/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/aot/AotFragmentTestConfigurationSupport.java b/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/aot/AotFragmentTestConfigurationSupport.java
@@ -47,6 +47,7 @@
 import org.springframework.data.projection.ProjectionFactory;
 import org.springframework.data.projection.SpelAwareProxyProjectionFactory;
 import org.springframework.data.repository.config.AnnotationRepositoryConfigurationSource;
+import org.springframework.data.repository.config.RepositoryConfigurationSource;
 import org.springframework.data.repository.core.RepositoryMetadata;
 import org.springframework.data.repository.core.support.RepositoryComposition;
 import org.springframework.data.repository.core.support.RepositoryFactoryBeanSupport;
@@ -69,7 +70,8 @@ public class AotFragmentTestConfigurationSupport implements BeanFactoryPostProce
 
 	private final Class<?> repositoryInterface;
 	private final boolean registerFragmentFacade;
-	private final TestJpaAotRepositoryContext<?> repositoryContext;
+	private final Class<?>[] additionalFragments;
+	private final RepositoryConfigurationSource configSource;
 
 	public AotFragmentTestConfigurationSupport(Class<?> repositoryInterface) {
 		this(repositoryInterface, SampleConfig.class, true);
@@ -82,22 +84,22 @@ public AotFragmentTestConfigurationSupport(Class<?> repositoryInterface, Class<?
 	public AotFragmentTestConfigurationSupport(Class<?> repositoryInterface, Class<?> configClass,
 			boolean registerFragmentFacade, Class<?>... additionalFragments) {
 		this.repositoryInterface = repositoryInterface;
-
-		RepositoryComposition composition = RepositoryComposition
-				.of((List) Arrays.stream(additionalFragments).map(RepositoryFragment::structural).toList());
-		this.repositoryContext = new TestJpaAotRepositoryContext<>(repositoryInterface, composition,
-				new AnnotationRepositoryConfigurationSource(AnnotationMetadata.introspect(configClass),
-						EnableJpaRepositories.class, new DefaultResourceLoader(), new StandardEnvironment(),
-						Mockito.mock(BeanDefinitionRegistry.class), DefaultBeanNameGenerator.INSTANCE));
 		this.registerFragmentFacade = registerFragmentFacade;
+		this.additionalFragments = additionalFragments;
+		this.configSource = new AnnotationRepositoryConfigurationSource(AnnotationMetadata.introspect(configClass),
+				EnableJpaRepositories.class, new DefaultResourceLoader(), new StandardEnvironment(),
+				Mockito.mock(BeanDefinitionRegistry.class), DefaultBeanNameGenerator.INSTANCE);
 	}
 
 	@Override
 	public void postProcessBeanFactory(ConfigurableListableBeanFactory beanFactory) throws BeansException {
 
 		TestGenerationContext generationContext = new TestGenerationContext(repositoryInterface);
+		RepositoryComposition composition = RepositoryComposition
+				.of((List) Arrays.stream(additionalFragments).map(RepositoryFragment::structural).toList());
 
-		repositoryContext.setBeanFactory(beanFactory);
+		TestJpaAotRepositoryContext<?> repositoryContext = new TestJpaAotRepositoryContext<>(beanFactory,
+				repositoryInterface, composition, configSource);
 
 		new JpaRepositoryContributor(repositoryContext).contribute(generationContext);
 
diff --git a/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/aot/TestJpaAotRepositoryContext.java b/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/aot/TestJpaAotRepositoryContext.java
@@ -19,24 +19,20 @@
 import jakarta.persistence.MappedSuperclass;
 
 import java.lang.annotation.Annotation;
-import java.util.Collection;
-import java.util.List;
 import java.util.Set;
-import java.util.function.Consumer;
 
 import org.jspecify.annotations.Nullable;
 
-import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
+import org.springframework.beans.factory.BeanFactory;
 import org.springframework.core.annotation.MergedAnnotation;
-import org.springframework.core.env.Environment;
-import org.springframework.core.env.StandardEnvironment;
-import org.springframework.data.aot.AotTypeConfiguration;
+import org.springframework.data.aot.AotContext;
 import org.springframework.data.jpa.domain.sample.Role;
 import org.springframework.data.jpa.domain.sample.SpecialUser;
 import org.springframework.data.jpa.domain.sample.User;
 import org.springframework.data.jpa.repository.support.JpaRepositoryFragmentsContributor;
 import org.springframework.data.jpa.repository.support.SimpleJpaRepository;
 import org.springframework.data.repository.config.AotRepositoryContext;
+import org.springframework.data.repository.config.AotRepositoryContextSupport;
 import org.springframework.data.repository.config.AotRepositoryInformation;
 import org.springframework.data.repository.config.RepositoryConfigurationSource;
 import org.springframework.data.repository.core.RepositoryInformation;
@@ -49,15 +45,16 @@
  *
  * @author Christoph Strobl
  */
-public class TestJpaAotRepositoryContext<T> implements AotRepositoryContext {
+public class TestJpaAotRepositoryContext<T> extends AotRepositoryContextSupport {
 
 	private final AotRepositoryInformation repositoryInformation;
 	private final Class<T> repositoryInterface;
 	private final RepositoryConfigurationSource configurationSource;
-	private @Nullable ConfigurableListableBeanFactory beanFactory;
 
-	public TestJpaAotRepositoryContext(Class<T> repositoryInterface, @Nullable RepositoryComposition composition,
+	public TestJpaAotRepositoryContext(BeanFactory beanFactory, Class<T> repositoryInterface,
+			@Nullable RepositoryComposition composition,
 			RepositoryConfigurationSource configurationSource) {
+		super(AotContext.from(beanFactory));
 		this.repositoryInterface = repositoryInterface;
 		this.configurationSource = configurationSource;
 
@@ -69,45 +66,6 @@ public TestJpaAotRepositoryContext(Class<T> repositoryInterface, @Nullable Repos
 				composition.append(fragments).getFragments().stream().toList());
 	}
 
-	public Class<T> getRepositoryInterface() {
-		return repositoryInterface;
-	}
-
-	@Override
-	public ConfigurableListableBeanFactory getBeanFactory() {
-		return beanFactory;
-	}
-
-	@Override
-	public Environment getEnvironment() {
-		return new StandardEnvironment();
-	}
-
-	@Override
-	public TypeIntrospector introspectType(String typeName) {
-		return null;
-	}
-
-	@Override
-	public IntrospectedBeanDefinition introspectBeanDefinition(String beanName) {
-		return null;
-	}
-
-	@Override
-	public void typeConfiguration(Class<?> type, Consumer<AotTypeConfiguration> configurationConsumer) {
-
-	}
-
-	@Override
-	public Collection<AotTypeConfiguration> typeConfigurations() {
-		return List.of();
-	}
-
-	@Override
-	public String getBeanName() {
-		return "dummyRepository";
-	}
-
 	@Override
 	public String getModuleName() {
 		return "JPA";
@@ -118,11 +76,6 @@ public RepositoryConfigurationSource getConfigurationSource() {
 		return configurationSource;
 	}
 
-	@Override
-	public Set<String> getBasePackages() {
-		return Set.of("org.springframework.data.dummy.repository.aot");
-	}
-
 	@Override
 	public Set<Class<? extends Annotation>> getIdentifyingAnnotations() {
 		return Set.of(Entity.class, MappedSuperclass.class);
@@ -143,12 +96,4 @@ public Set<Class<?>> getResolvedTypes() {
 		return Set.of(User.class, SpecialUser.class, Role.class);
 	}
 
-	@Override
-	public Set<Class<?>> getUserDomainTypes() {
-		return Set.of();
-	}
-
-	public void setBeanFactory(ConfigurableListableBeanFactory beanFactory) {
-		this.beanFactory = beanFactory;
-	}
 }
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
