Add support for ImportAware in BeanRegistrar

Closes gh-34627
This commit is contained in:
Sébastien Deleuze 2025-03-21 11:49:15 +01:00
parent 3e788e4ca1
commit 5ce64f47b2
10 changed files with 196 additions and 21 deletions

View File

@ -19,18 +19,20 @@ package org.springframework.beans.factory;
import org.springframework.core.env.Environment;
/**
* Contract for registering beans programmatically.
*
* <p>Typically imported with an {@link org.springframework.context.annotation.Import @Import}
* annotation on {@link org.springframework.context.annotation.Configuration @Configuration}
* classes.
* Contract for registering beans programmatically, typically imported with an
* {@link org.springframework.context.annotation.Import @Import} annotation on
* a {@link org.springframework.context.annotation.Configuration @Configuration}
* class.
* <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}.
*

View File

@ -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) {
this.beanRegistrars.add(beanRegistrar);
void addBeanRegistrar(String sourceClassName, BeanRegistrar beanRegistrar) {
this.beanRegistrars.put(sourceClassName, beanRegistrar);
}
public Set<BeanRegistrar> getBeanRegistrars() {
public Map<String, BeanRegistrar> getBeanRegistrars() {
return this.beanRegistrars;
}

View File

@ -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));
}

View File

@ -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 ->

View File

@ -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) {
code.addStatement("new $T().register(new $T(($T)$L, $L, $L, $T.class, $L), $L)", beanRegistrar.getClass(),
Builder metadataReaderFactoryCode = null;
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,

View File

@ -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();

View File

@ -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());
}
}

View File

@ -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();

View File

@ -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) {}
}

View File

@ -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 {
}