Add support for ImportAware in BeanRegistrar

Closes gh-34627

spring-beans/src/main/java/org/springframework/beans/factory/BeanRegistrar.java

@@ -19,18 +19,20 @@ package org.springframework.beans.factory;
 import org.springframework.core.env.Environment;
 
 /**
- * Contract for registering beans programmatically.
+ * Contract for registering beans programmatically, typically imported with an
- *
+ * {@link org.springframework.context.annotation.Import @Import} annotation on
- * <p>Typically imported with an {@link org.springframework.context.annotation.Import @Import}
+ * a {@link org.springframework.context.annotation.Configuration @Configuration}
- * annotation on {@link org.springframework.context.annotation.Configuration @Configuration}
+ * class.
- * classes.
  * <pre class="code">
  * &#064;Configuration
  * &#064;Import(MyBeanRegistrar.class)
  * class MyConfiguration {
  * }</pre>
+ * Can also be applied to an application context via
+ * {@link org.springframework.context.support.GenericApplicationContext#register(BeanRegistrar...)}.
+ *
  *
- * <p>The bean registrar implementation uses {@link BeanRegistry} and {@link Environment}
+ * <p>Bean registrar implementations use {@link BeanRegistry} and {@link Environment}
  * APIs to register beans programmatically in a concise and flexible way.
  * <pre class="code">
  * class MyBeanRegistrar implements BeanRegistrar {
@@ -50,6 +52,10 @@ import org.springframework.core.env.Environment;
  *     }
  * }</pre>
  *
+ * <p>A {@code BeanRegistrar} implementing {@link org.springframework.context.annotation.ImportAware}
+ * can optionally introspect import metadata when used in an import scenario, otherwise the
+ * {@code setImportMetadata} method is simply not being called.
+ *
  * <p>In Kotlin, it is recommended to use {@code BeanRegistrarDsl} instead of
  * implementing {@code BeanRegistrar}.
  *

spring-context/src/main/java/org/springframework/context/annotation/ConfigurationClass.java

@@ -66,7 +66,7 @@ final class ConfigurationClass {
 	private final Map<String, Class<? extends BeanDefinitionReader>> importedResources =
 			new LinkedHashMap<>();
 
-	private final Set<BeanRegistrar> beanRegistrars = new LinkedHashSet<>();
+	private final Map<String, BeanRegistrar> beanRegistrars = new LinkedHashMap<>();
 
 	private final Map<ImportBeanDefinitionRegistrar, AnnotationMetadata> importBeanDefinitionRegistrars =
 			new LinkedHashMap<>();
@@ -222,11 +222,11 @@ final class ConfigurationClass {
 		return this.importedResources;
 	}
 
-	void addBeanRegistrar(BeanRegistrar beanRegistrar) {
+	void addBeanRegistrar(String sourceClassName, BeanRegistrar beanRegistrar) {
-		this.beanRegistrars.add(beanRegistrar);
+		this.beanRegistrars.put(sourceClassName, beanRegistrar);
 	}
 
-	public Set<BeanRegistrar> getBeanRegistrars() {
+	public Map<String, BeanRegistrar> getBeanRegistrars() {
 		return this.beanRegistrars;
 	}
 

spring-context/src/main/java/org/springframework/context/annotation/ConfigurationClassBeanDefinitionReader.java

@@ -404,11 +404,11 @@ class ConfigurationClassBeanDefinitionReader {
 				registrar.registerBeanDefinitions(metadata, this.registry, this.importBeanNameGenerator));
 	}
 
-	private void loadBeanDefinitionsFromBeanRegistrars(Set<BeanRegistrar> registrars) {
+	private void loadBeanDefinitionsFromBeanRegistrars(Map<String, BeanRegistrar> registrars) {
 		Assert.isInstanceOf(ListableBeanFactory.class, this.registry,
 				"Cannot support bean registrars since " + this.registry.getClass().getName() +
 						" does not implement BeanDefinitionRegistry");
-		registrars.forEach(registrar -> registrar.register(new BeanRegistryAdapter(this.registry,
+		registrars.values().forEach(registrar -> registrar.register(new BeanRegistryAdapter(this.registry,
 				(ListableBeanFactory) this.registry, this.environment, registrar.getClass()), this.environment));
 	}
 

spring-context/src/main/java/org/springframework/context/annotation/ConfigurationClassParser.java

@@ -602,7 +602,11 @@ class ConfigurationClassParser {
 					else if (candidate.isAssignable(BeanRegistrar.class)) {
 						Class<?> candidateClass = candidate.loadClass();
 						BeanRegistrar registrar = (BeanRegistrar) BeanUtils.instantiateClass(candidateClass);
-						configClass.addBeanRegistrar(registrar);
+						AnnotationMetadata metadata = currentSourceClass.getMetadata();
+						if (registrar instanceof ImportAware importAware) {
+							importAware.setImportMetadata(metadata);
+						}
+						configClass.addBeanRegistrar(metadata.getClassName(), registrar);
 					}
 					else if (candidate.isAssignable(ImportBeanDefinitionRegistrar.class)) {
 						// Candidate class is an ImportBeanDefinitionRegistrar ->

spring-context/src/main/java/org/springframework/context/annotation/ConfigurationClassPostProcessor.java

@@ -114,6 +114,7 @@ import org.springframework.javapoet.ClassName;
 import org.springframework.javapoet.CodeBlock;
 import org.springframework.javapoet.CodeBlock.Builder;
 import org.springframework.javapoet.MethodSpec;
+import org.springframework.javapoet.NameAllocator;
 import org.springframework.javapoet.ParameterizedTypeName;
 import org.springframework.util.Assert;
 import org.springframework.util.ClassUtils;
@@ -122,6 +123,7 @@ import org.springframework.util.LinkedMultiValueMap;
 import org.springframework.util.MultiValueMap;
 import org.springframework.util.ObjectUtils;
 import org.springframework.util.ReflectionUtils;
+import org.springframework.util.StringUtils;
 
 /**
  * {@link BeanFactoryPostProcessor} used for bootstrapping processing of
@@ -197,7 +199,7 @@ public class ConfigurationClassPostProcessor implements BeanDefinitionRegistryPo
 	@SuppressWarnings("NullAway.Init")
 	private List<PropertySourceDescriptor> propertySourceDescriptors;
 
-	private Set<BeanRegistrar> beanRegistrars = new LinkedHashSet<>();
+	private Map<String, BeanRegistrar> beanRegistrars = new LinkedHashMap<>();
 
 
 	@Override
@@ -443,7 +445,7 @@ public class ConfigurationClassPostProcessor implements BeanDefinitionRegistryPo
 			}
 			this.reader.loadBeanDefinitions(configClasses);
 			for (ConfigurationClass configClass : configClasses) {
-				this.beanRegistrars.addAll(configClass.getBeanRegistrars());
+				this.beanRegistrars.putAll(configClass.getBeanRegistrars());
 			}
 			alreadyParsed.addAll(configClasses);
 			processConfig.tag("classCount", () -> String.valueOf(configClasses.size())).end();
@@ -846,13 +848,13 @@ public class ConfigurationClassPostProcessor implements BeanDefinitionRegistryPo
 
 		private static final String ENVIRONMENT_VARIABLE = "environment";
 
-		private final Set<BeanRegistrar> beanRegistrars;
+		private final Map<String, BeanRegistrar> beanRegistrars;
 
 		private final ConfigurableListableBeanFactory beanFactory;
 
 		private final AotServices<BeanRegistrationAotProcessor> aotProcessors;
 
-		public BeanRegistrarAotContribution(Set<BeanRegistrar> beanRegistrars, ConfigurableListableBeanFactory beanFactory) {
+		public BeanRegistrarAotContribution(Map<String, BeanRegistrar> beanRegistrars, ConfigurableListableBeanFactory beanFactory) {
 			this.beanRegistrars = beanRegistrars;
 			this.beanFactory = beanFactory;
 			this.aotProcessors = AotServices.factoriesAndBeans(this.beanFactory).load(BeanRegistrationAotProcessor.class);
@@ -935,13 +937,32 @@ public class ConfigurationClassPostProcessor implements BeanDefinitionRegistryPo
 
 		private CodeBlock generateRegisterCode() {
 			Builder code = CodeBlock.builder();
-			for (BeanRegistrar beanRegistrar : this.beanRegistrars) {
+			Builder metadataReaderFactoryCode = null;
-				code.addStatement("new $T().register(new $T(($T)$L, $L, $L, $T.class, $L), $L)", beanRegistrar.getClass(),
+			NameAllocator nameAllocator = new NameAllocator();
+			for (Map.Entry<String, BeanRegistrar> beanRegistrarEntry : this.beanRegistrars.entrySet()) {
+				BeanRegistrar beanRegistrar = beanRegistrarEntry.getValue();
+				String beanRegistrarName = nameAllocator.newName(StringUtils.uncapitalize(beanRegistrar.getClass().getSimpleName()));
+				code.addStatement("$T $L = new $T()", beanRegistrar.getClass(), beanRegistrarName, beanRegistrar.getClass());
+				if (beanRegistrar instanceof ImportAware) {
+					if (metadataReaderFactoryCode == null) {
+						metadataReaderFactoryCode = CodeBlock.builder();
+						metadataReaderFactoryCode.addStatement("$T metadataReaderFactory = new $T()",
+								MetadataReaderFactory.class, CachingMetadataReaderFactory.class);
+					}
+					code.beginControlFlow("try")
+							.addStatement("$L.setImportMetadata(metadataReaderFactory.getMetadataReader($S).getAnnotationMetadata())",
+									beanRegistrarName, beanRegistrarEntry.getKey())
+							.nextControlFlow("catch ($T ex)", IOException.class)
+							.addStatement("throw new $T(\"Failed to read metadata for '$L'\", ex)",
+									IllegalStateException.class, beanRegistrarEntry.getKey())
+							.endControlFlow();
+				}
+				code.addStatement("$L.register(new $T(($T)$L, $L, $L, $T.class, $L), $L)", beanRegistrarName,
 						BeanRegistryAdapter.class, BeanDefinitionRegistry.class, BeanFactoryInitializationCode.BEAN_FACTORY_VARIABLE,
 						BeanFactoryInitializationCode.BEAN_FACTORY_VARIABLE, ENVIRONMENT_VARIABLE, beanRegistrar.getClass(),
 						CUSTOMIZER_MAP_VARIABLE, ENVIRONMENT_VARIABLE);
 			}
-			return code.build();
+			return (metadataReaderFactoryCode == null ? code.build() : metadataReaderFactoryCode.add(code.build()).build());
 		}
 
 		private CodeBlock generateInitDestroyMethods(String beanName, AbstractBeanDefinition beanDefinition,

spring-context/src/test/java/org/springframework/context/annotation/ConfigurationClassPostProcessorAotContributionTests.java

@@ -500,6 +500,22 @@ public class ConfigurationClassPostProcessorAotContributionTests {
 			});
 		}
 
+		@Test
+		void applyToWhenIsImportAware() {
+			BeanFactoryInitializationAotContribution contribution = getContribution(CommonAnnotationBeanPostProcessor.class,
+					ImportAwareBeanRegistrarConfiguration.class);
+			assertThat(contribution).isNotNull();
+			contribution.applyTo(generationContext, beanFactoryInitializationCode);
+			compile((initializer, compiled) -> {
+				GenericApplicationContext freshContext = new GenericApplicationContext();
+				initializer.accept(freshContext);
+				freshContext.refresh();
+				assertThat(freshContext.getBean(ClassNameHolder.class).className())
+						.isEqualTo(ImportAwareBeanRegistrarConfiguration.class.getName());
+				freshContext.close();
+			});
+		}
+
 		@SuppressWarnings("unchecked")
 		private void compile(BiConsumer<Consumer<GenericApplicationContext>, Compiled> result) {
 			MethodReference methodReference = beanFactoryInitializationCode.getInitializers().get(0);
@@ -561,6 +577,31 @@ public class ConfigurationClassPostProcessorAotContributionTests {
 			}
 		}
 
+		@Import(ImportAwareBeanRegistrar.class)
+		public static class ImportAwareBeanRegistrarConfiguration {
+		}
+
+		public static class ImportAwareBeanRegistrar implements BeanRegistrar, ImportAware {
+
+			@Nullable
+			private AnnotationMetadata importMetadata;
+
+			@Override
+			public void register(BeanRegistry registry, Environment env) {
+				registry.registerBean(ClassNameHolder.class, spec -> spec.supplier(context ->
+						new ClassNameHolder(this.importMetadata == null ? null : this.importMetadata.getClassName())));
+			}
+
+			@Override
+			public void setImportMetadata(AnnotationMetadata importMetadata) {
+				this.importMetadata = importMetadata;
+			}
+
+			public @Nullable AnnotationMetadata getImportMetadata() {
+				return this.importMetadata;
+			}
+		}
+
 		static class Foo {
 		}
 
@@ -576,6 +617,8 @@ public class ConfigurationClassPostProcessorAotContributionTests {
 
 	}
 
+	public record ClassNameHolder(@Nullable String className) {}
+
 
 	private @Nullable BeanFactoryInitializationAotContribution getContribution(Class<?>... types) {
 		DefaultListableBeanFactory beanFactory = new DefaultListableBeanFactory();

spring-context/src/test/java/org/springframework/context/annotation/beanregistrar/BeanRegistrarConfigurationTests.java

@@ -24,12 +24,14 @@ import org.springframework.beans.factory.config.BeanDefinition;
 import org.springframework.beans.factory.support.RootBeanDefinition;
 import org.springframework.context.annotation.AnnotationConfigApplicationContext;
 import org.springframework.context.testfixture.beans.factory.GenericBeanRegistrar;
+import org.springframework.context.testfixture.beans.factory.ImportAwareBeanRegistrar;
 import org.springframework.context.testfixture.beans.factory.SampleBeanRegistrar.Bar;
 import org.springframework.context.testfixture.beans.factory.SampleBeanRegistrar.Baz;
 import org.springframework.context.testfixture.beans.factory.SampleBeanRegistrar.Foo;
 import org.springframework.context.testfixture.beans.factory.SampleBeanRegistrar.Init;
 import org.springframework.context.testfixture.context.annotation.registrar.BeanRegistrarConfiguration;
 import org.springframework.context.testfixture.context.annotation.registrar.GenericBeanRegistrarConfiguration;
+import org.springframework.context.testfixture.context.annotation.registrar.ImportAwareBeanRegistrarConfiguration;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -82,4 +84,13 @@ public class BeanRegistrarConfigurationTests {
 		assertThat(beanDefinition.getResolvableType().resolveGeneric(0)).isEqualTo(GenericBeanRegistrar.Foo.class);
 	}
 
+	@Test
+	void beanRegistrarWithImportAware() {
+		AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext();
+		context.register(ImportAwareBeanRegistrarConfiguration.class);
+		context.refresh();
+		assertThat(context.getBean(ImportAwareBeanRegistrar.ClassNameHolder.class).className())
+				.isEqualTo(ImportAwareBeanRegistrarConfiguration.class.getName());
+	}
+
 }

spring-context/src/test/java/org/springframework/context/support/GenericApplicationContextTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2024 the original author or authors.
+ * Copyright 2002-2025 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.
@@ -46,6 +46,8 @@ import org.springframework.beans.factory.support.MergedBeanDefinitionPostProcess
 import org.springframework.beans.factory.support.RootBeanDefinition;
 import org.springframework.context.ApplicationContext;
 import org.springframework.context.ApplicationContextAware;
+import org.springframework.context.testfixture.beans.factory.ImportAwareBeanRegistrar;
+import org.springframework.context.testfixture.beans.factory.SampleBeanRegistrar;
 import org.springframework.core.DecoratingProxy;
 import org.springframework.core.env.ConfigurableEnvironment;
 import org.springframework.core.env.Environment;
@@ -627,6 +629,22 @@ class GenericApplicationContextTests {
 		context.close();
 	}
 
+	@Test
+	void beanRegistrar() {
+		GenericApplicationContext context = new GenericApplicationContext();
+		context.register(new SampleBeanRegistrar());
+		context.refresh();
+		assertThat(context.getBean(SampleBeanRegistrar.Bar.class).foo()).isEqualTo(context.getBean(SampleBeanRegistrar.Foo.class));
+	}
+
+	@Test
+	void importAwareBeanRegistrar() {
+		GenericApplicationContext context = new GenericApplicationContext();
+		context.register(new ImportAwareBeanRegistrar());
+		context.refresh();
+		assertThat(context.getBean(ImportAwareBeanRegistrar.ClassNameHolder.class).className()).isNull();
+	}
+
 
 	private MergedBeanDefinitionPostProcessor registerMockMergedBeanDefinitionPostProcessor(GenericApplicationContext context) {
 		MergedBeanDefinitionPostProcessor bpp = mock();

spring-context/src/testFixtures/java/org/springframework/context/testfixture/beans/factory/ImportAwareBeanRegistrar.java

@@ -0,0 +1,48 @@
+/*
+ * Copyright 2002-2025 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.context.testfixture.beans.factory;
+
+import org.jspecify.annotations.Nullable;
+
+import org.springframework.beans.factory.BeanRegistrar;
+import org.springframework.beans.factory.BeanRegistry;
+import org.springframework.context.annotation.ImportAware;
+import org.springframework.core.env.Environment;
+import org.springframework.core.type.AnnotationMetadata;
+
+public class ImportAwareBeanRegistrar implements BeanRegistrar, ImportAware {
+
+	@Nullable
+	private AnnotationMetadata importMetadata;
+
+	@Override
+	public void register(BeanRegistry registry, Environment env) {
+		registry.registerBean(ClassNameHolder.class, spec -> spec.supplier(context ->
+				new ClassNameHolder(this.importMetadata == null ? null : this.importMetadata.getClassName())));
+	}
+
+	@Override
+	public void setImportMetadata(AnnotationMetadata importMetadata) {
+		this.importMetadata = importMetadata;
+	}
+
+	public @Nullable AnnotationMetadata getImportMetadata() {
+		return this.importMetadata;
+	}
+
+	public record ClassNameHolder(@Nullable String className) {}
+}

spring-context/src/testFixtures/java/org/springframework/context/testfixture/context/annotation/registrar/ImportAwareBeanRegistrarConfiguration.java

@@ -0,0 +1,24 @@
+/*
+ * Copyright 2002-2025 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.context.testfixture.context.annotation.registrar;
+
+import org.springframework.context.annotation.Import;
+import org.springframework.context.testfixture.beans.factory.ImportAwareBeanRegistrar;
+
+@Import(ImportAwareBeanRegistrar.class)
+public class ImportAwareBeanRegistrarConfiguration {
+}