Skip to content

Commit 5ce64f4

Browse files
sdeleuzesdeleuze
committed
Add support for ImportAware in BeanRegistrar
Closes gh-34627
1 parent 3e788e4 commit 5ce64f4

File tree

10 files changed

+196
-21
lines changed

10 files changed

+196
-21
lines changed

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

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -19,18 +19,20 @@
1919
import org.springframework.core.env.Environment;
2020

2121
/**
22-
* Contract for registering beans programmatically.
23-
*
24-
* <p>Typically imported with an {@link org.springframework.context.annotation.Import @Import}
25-
* annotation on {@link org.springframework.context.annotation.Configuration @Configuration}
26-
* classes.
22+
* Contract for registering beans programmatically, typically imported with an
23+
* {@link org.springframework.context.annotation.Import @Import} annotation on
24+
* a {@link org.springframework.context.annotation.Configuration @Configuration}
25+
* class.
2726
* <pre class="code">
2827
* &#064;Configuration
2928
* &#064;Import(MyBeanRegistrar.class)
3029
* class MyConfiguration {
3130
* }</pre>
31+
* Can also be applied to an application context via
32+
* {@link org.springframework.context.support.GenericApplicationContext#register(BeanRegistrar...)}.
33+
*
3234
*
33-
* <p>The bean registrar implementation uses {@link BeanRegistry} and {@link Environment}
35+
* <p>Bean registrar implementations use {@link BeanRegistry} and {@link Environment}
3436
* APIs to register beans programmatically in a concise and flexible way.
3537
* <pre class="code">
3638
* class MyBeanRegistrar implements BeanRegistrar {
@@ -50,6 +52,10 @@
5052
* }
5153
* }</pre>
5254
*
55+
* <p>A {@code BeanRegistrar} implementing {@link org.springframework.context.annotation.ImportAware}
56+
* can optionally introspect import metadata when used in an import scenario, otherwise the
57+
* {@code setImportMetadata} method is simply not being called.
58+
*
5359
* <p>In Kotlin, it is recommended to use {@code BeanRegistrarDsl} instead of
5460
* implementing {@code BeanRegistrar}.
5561
*

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ final class ConfigurationClass {
6666
private final Map<String, Class<? extends BeanDefinitionReader>> importedResources =
6767
new LinkedHashMap<>();
6868

69-
private final Set<BeanRegistrar> beanRegistrars = new LinkedHashSet<>();
69+
private final Map<String, BeanRegistrar> beanRegistrars = new LinkedHashMap<>();
7070

7171
private final Map<ImportBeanDefinitionRegistrar, AnnotationMetadata> importBeanDefinitionRegistrars =
7272
new LinkedHashMap<>();
@@ -222,11 +222,11 @@ Map<String, Class<? extends BeanDefinitionReader>> getImportedResources() {
222222
return this.importedResources;
223223
}
224224

225-
void addBeanRegistrar(BeanRegistrar beanRegistrar) {
226-
this.beanRegistrars.add(beanRegistrar);
225+
void addBeanRegistrar(String sourceClassName, BeanRegistrar beanRegistrar) {
226+
this.beanRegistrars.put(sourceClassName, beanRegistrar);
227227
}
228228

229-
public Set<BeanRegistrar> getBeanRegistrars() {
229+
public Map<String, BeanRegistrar> getBeanRegistrars() {
230230
return this.beanRegistrars;
231231
}
232232

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -404,11 +404,11 @@ private void loadBeanDefinitionsFromImportBeanDefinitionRegistrars(Map<ImportBea
404404
registrar.registerBeanDefinitions(metadata, this.registry, this.importBeanNameGenerator));
405405
}
406406

407-
private void loadBeanDefinitionsFromBeanRegistrars(Set<BeanRegistrar> registrars) {
407+
private void loadBeanDefinitionsFromBeanRegistrars(Map<String, BeanRegistrar> registrars) {
408408
Assert.isInstanceOf(ListableBeanFactory.class, this.registry,
409409
"Cannot support bean registrars since " + this.registry.getClass().getName() +
410410
" does not implement BeanDefinitionRegistry");
411-
registrars.forEach(registrar -> registrar.register(new BeanRegistryAdapter(this.registry,
411+
registrars.values().forEach(registrar -> registrar.register(new BeanRegistryAdapter(this.registry,
412412
(ListableBeanFactory) this.registry, this.environment, registrar.getClass()), this.environment));
413413
}
414414

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -602,7 +602,11 @@ private void processImports(ConfigurationClass configClass, SourceClass currentS
602602
else if (candidate.isAssignable(BeanRegistrar.class)) {
603603
Class<?> candidateClass = candidate.loadClass();
604604
BeanRegistrar registrar = (BeanRegistrar) BeanUtils.instantiateClass(candidateClass);
605-
configClass.addBeanRegistrar(registrar);
605+
AnnotationMetadata metadata = currentSourceClass.getMetadata();
606+
if (registrar instanceof ImportAware importAware) {
607+
importAware.setImportMetadata(metadata);
608+
}
609+
configClass.addBeanRegistrar(metadata.getClassName(), registrar);
606610
}
607611
else if (candidate.isAssignable(ImportBeanDefinitionRegistrar.class)) {
608612
// Candidate class is an ImportBeanDefinitionRegistrar ->

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

Lines changed: 28 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -114,6 +114,7 @@
114114
import org.springframework.javapoet.CodeBlock;
115115
import org.springframework.javapoet.CodeBlock.Builder;
116116
import org.springframework.javapoet.MethodSpec;
117+
import org.springframework.javapoet.NameAllocator;
117118
import org.springframework.javapoet.ParameterizedTypeName;
118119
import org.springframework.util.Assert;
119120
import org.springframework.util.ClassUtils;
@@ -122,6 +123,7 @@
122123
import org.springframework.util.MultiValueMap;
123124
import org.springframework.util.ObjectUtils;
124125
import org.springframework.util.ReflectionUtils;
126+
import org.springframework.util.StringUtils;
125127

126128
/**
127129
* {@link BeanFactoryPostProcessor} used for bootstrapping processing of
@@ -197,7 +199,7 @@ public class ConfigurationClassPostProcessor implements BeanDefinitionRegistryPo
197199
@SuppressWarnings("NullAway.Init")
198200
private List<PropertySourceDescriptor> propertySourceDescriptors;
199201

200-
private Set<BeanRegistrar> beanRegistrars = new LinkedHashSet<>();
202+
private Map<String, BeanRegistrar> beanRegistrars = new LinkedHashMap<>();
201203

202204

203205
@Override
@@ -443,7 +445,7 @@ else if (ConfigurationClassUtils.checkConfigurationClassCandidate(beanDef, this.
443445
}
444446
this.reader.loadBeanDefinitions(configClasses);
445447
for (ConfigurationClass configClass : configClasses) {
446-
this.beanRegistrars.addAll(configClass.getBeanRegistrars());
448+
this.beanRegistrars.putAll(configClass.getBeanRegistrars());
447449
}
448450
alreadyParsed.addAll(configClasses);
449451
processConfig.tag("classCount", () -> String.valueOf(configClasses.size())).end();
@@ -846,13 +848,13 @@ private static class BeanRegistrarAotContribution implements BeanFactoryInitiali
846848

847849
private static final String ENVIRONMENT_VARIABLE = "environment";
848850

849-
private final Set<BeanRegistrar> beanRegistrars;
851+
private final Map<String, BeanRegistrar> beanRegistrars;
850852

851853
private final ConfigurableListableBeanFactory beanFactory;
852854

853855
private final AotServices<BeanRegistrationAotProcessor> aotProcessors;
854856

855-
public BeanRegistrarAotContribution(Set<BeanRegistrar> beanRegistrars, ConfigurableListableBeanFactory beanFactory) {
857+
public BeanRegistrarAotContribution(Map<String, BeanRegistrar> beanRegistrars, ConfigurableListableBeanFactory beanFactory) {
856858
this.beanRegistrars = beanRegistrars;
857859
this.beanFactory = beanFactory;
858860
this.aotProcessors = AotServices.factoriesAndBeans(this.beanFactory).load(BeanRegistrationAotProcessor.class);
@@ -935,13 +937,32 @@ private CodeBlock generateCustomizerMap() {
935937

936938
private CodeBlock generateRegisterCode() {
937939
Builder code = CodeBlock.builder();
938-
for (BeanRegistrar beanRegistrar : this.beanRegistrars) {
939-
code.addStatement("new $T().register(new $T(($T)$L, $L, $L, $T.class, $L), $L)", beanRegistrar.getClass(),
940+
Builder metadataReaderFactoryCode = null;
941+
NameAllocator nameAllocator = new NameAllocator();
942+
for (Map.Entry<String, BeanRegistrar> beanRegistrarEntry : this.beanRegistrars.entrySet()) {
943+
BeanRegistrar beanRegistrar = beanRegistrarEntry.getValue();
944+
String beanRegistrarName = nameAllocator.newName(StringUtils.uncapitalize(beanRegistrar.getClass().getSimpleName()));
945+
code.addStatement("$T $L = new $T()", beanRegistrar.getClass(), beanRegistrarName, beanRegistrar.getClass());
946+
if (beanRegistrar instanceof ImportAware) {
947+
if (metadataReaderFactoryCode == null) {
948+
metadataReaderFactoryCode = CodeBlock.builder();
949+
metadataReaderFactoryCode.addStatement("$T metadataReaderFactory = new $T()",
950+
MetadataReaderFactory.class, CachingMetadataReaderFactory.class);
951+
}
952+
code.beginControlFlow("try")
953+
.addStatement("$L.setImportMetadata(metadataReaderFactory.getMetadataReader($S).getAnnotationMetadata())",
954+
beanRegistrarName, beanRegistrarEntry.getKey())
955+
.nextControlFlow("catch ($T ex)", IOException.class)
956+
.addStatement("throw new $T(\"Failed to read metadata for '$L'\", ex)",
957+
IllegalStateException.class, beanRegistrarEntry.getKey())
958+
.endControlFlow();
959+
}
960+
code.addStatement("$L.register(new $T(($T)$L, $L, $L, $T.class, $L), $L)", beanRegistrarName,
940961
BeanRegistryAdapter.class, BeanDefinitionRegistry.class, BeanFactoryInitializationCode.BEAN_FACTORY_VARIABLE,
941962
BeanFactoryInitializationCode.BEAN_FACTORY_VARIABLE, ENVIRONMENT_VARIABLE, beanRegistrar.getClass(),
942963
CUSTOMIZER_MAP_VARIABLE, ENVIRONMENT_VARIABLE);
943964
}
944-
return code.build();
965+
return (metadataReaderFactoryCode == null ? code.build() : metadataReaderFactoryCode.add(code.build()).build());
945966
}
946967

947968
private CodeBlock generateInitDestroyMethods(String beanName, AbstractBeanDefinition beanDefinition,

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

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -500,6 +500,22 @@ void applyToWhenHasPostConstructAnnotationPostProcessed() {
500500
});
501501
}
502502

503+
@Test
504+
void applyToWhenIsImportAware() {
505+
BeanFactoryInitializationAotContribution contribution = getContribution(CommonAnnotationBeanPostProcessor.class,
506+
ImportAwareBeanRegistrarConfiguration.class);
507+
assertThat(contribution).isNotNull();
508+
contribution.applyTo(generationContext, beanFactoryInitializationCode);
509+
compile((initializer, compiled) -> {
510+
GenericApplicationContext freshContext = new GenericApplicationContext();
511+
initializer.accept(freshContext);
512+
freshContext.refresh();
513+
assertThat(freshContext.getBean(ClassNameHolder.class).className())
514+
.isEqualTo(ImportAwareBeanRegistrarConfiguration.class.getName());
515+
freshContext.close();
516+
});
517+
}
518+
503519
@SuppressWarnings("unchecked")
504520
private void compile(BiConsumer<Consumer<GenericApplicationContext>, Compiled> result) {
505521
MethodReference methodReference = beanFactoryInitializationCode.getInitializers().get(0);
@@ -561,6 +577,31 @@ public void register(BeanRegistry registry, Environment env) {
561577
}
562578
}
563579

580+
@Import(ImportAwareBeanRegistrar.class)
581+
public static class ImportAwareBeanRegistrarConfiguration {
582+
}
583+
584+
public static class ImportAwareBeanRegistrar implements BeanRegistrar, ImportAware {
585+
586+
@Nullable
587+
private AnnotationMetadata importMetadata;
588+
589+
@Override
590+
public void register(BeanRegistry registry, Environment env) {
591+
registry.registerBean(ClassNameHolder.class, spec -> spec.supplier(context ->
592+
new ClassNameHolder(this.importMetadata == null ? null : this.importMetadata.getClassName())));
593+
}
594+
595+
@Override
596+
public void setImportMetadata(AnnotationMetadata importMetadata) {
597+
this.importMetadata = importMetadata;
598+
}
599+
600+
public @Nullable AnnotationMetadata getImportMetadata() {
601+
return this.importMetadata;
602+
}
603+
}
604+
564605
static class Foo {
565606
}
566607

@@ -576,6 +617,8 @@ void postConstruct() {
576617

577618
}
578619

620+
public record ClassNameHolder(@Nullable String className) {}
621+
579622

580623
private @Nullable BeanFactoryInitializationAotContribution getContribution(Class<?>... types) {
581624
DefaultListableBeanFactory beanFactory = new DefaultListableBeanFactory();

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,12 +24,14 @@
2424
import org.springframework.beans.factory.support.RootBeanDefinition;
2525
import org.springframework.context.annotation.AnnotationConfigApplicationContext;
2626
import org.springframework.context.testfixture.beans.factory.GenericBeanRegistrar;
27+
import org.springframework.context.testfixture.beans.factory.ImportAwareBeanRegistrar;
2728
import org.springframework.context.testfixture.beans.factory.SampleBeanRegistrar.Bar;
2829
import org.springframework.context.testfixture.beans.factory.SampleBeanRegistrar.Baz;
2930
import org.springframework.context.testfixture.beans.factory.SampleBeanRegistrar.Foo;
3031
import org.springframework.context.testfixture.beans.factory.SampleBeanRegistrar.Init;
3132
import org.springframework.context.testfixture.context.annotation.registrar.BeanRegistrarConfiguration;
3233
import org.springframework.context.testfixture.context.annotation.registrar.GenericBeanRegistrarConfiguration;
34+
import org.springframework.context.testfixture.context.annotation.registrar.ImportAwareBeanRegistrarConfiguration;
3335

3436
import static org.assertj.core.api.Assertions.assertThat;
3537
import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -82,4 +84,13 @@ void beanRegistrarWithTargetType() {
8284
assertThat(beanDefinition.getResolvableType().resolveGeneric(0)).isEqualTo(GenericBeanRegistrar.Foo.class);
8385
}
8486

87+
@Test
88+
void beanRegistrarWithImportAware() {
89+
AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext();
90+
context.register(ImportAwareBeanRegistrarConfiguration.class);
91+
context.refresh();
92+
assertThat(context.getBean(ImportAwareBeanRegistrar.ClassNameHolder.class).className())
93+
.isEqualTo(ImportAwareBeanRegistrarConfiguration.class.getName());
94+
}
95+
8596
}

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

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2002-2024 the original author or authors.
2+
* Copyright 2002-2025 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -46,6 +46,8 @@
4646
import org.springframework.beans.factory.support.RootBeanDefinition;
4747
import org.springframework.context.ApplicationContext;
4848
import org.springframework.context.ApplicationContextAware;
49+
import org.springframework.context.testfixture.beans.factory.ImportAwareBeanRegistrar;
50+
import org.springframework.context.testfixture.beans.factory.SampleBeanRegistrar;
4951
import org.springframework.core.DecoratingProxy;
5052
import org.springframework.core.env.ConfigurableEnvironment;
5153
import org.springframework.core.env.Environment;
@@ -627,6 +629,22 @@ public Class<?> determineBeanType(Class<?> beanClass, String beanName) throws Be
627629
context.close();
628630
}
629631

632+
@Test
633+
void beanRegistrar() {
634+
GenericApplicationContext context = new GenericApplicationContext();
635+
context.register(new SampleBeanRegistrar());
636+
context.refresh();
637+
assertThat(context.getBean(SampleBeanRegistrar.Bar.class).foo()).isEqualTo(context.getBean(SampleBeanRegistrar.Foo.class));
638+
}
639+
640+
@Test
641+
void importAwareBeanRegistrar() {
642+
GenericApplicationContext context = new GenericApplicationContext();
643+
context.register(new ImportAwareBeanRegistrar());
644+
context.refresh();
645+
assertThat(context.getBean(ImportAwareBeanRegistrar.ClassNameHolder.class).className()).isNull();
646+
}
647+
630648

631649
private MergedBeanDefinitionPostProcessor registerMockMergedBeanDefinitionPostProcessor(GenericApplicationContext context) {
632650
MergedBeanDefinitionPostProcessor bpp = mock();

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

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
/*
2+
* Copyright 2002-2025 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.context.testfixture.beans.factory;
18+
19+
import org.jspecify.annotations.Nullable;
20+
21+
import org.springframework.beans.factory.BeanRegistrar;
22+
import org.springframework.beans.factory.BeanRegistry;
23+
import org.springframework.context.annotation.ImportAware;
24+
import org.springframework.core.env.Environment;
25+
import org.springframework.core.type.AnnotationMetadata;
26+
27+
public class ImportAwareBeanRegistrar implements BeanRegistrar, ImportAware {
28+
29+
@Nullable
30+
private AnnotationMetadata importMetadata;
31+
32+
@Override
33+
public void register(BeanRegistry registry, Environment env) {
34+
registry.registerBean(ClassNameHolder.class, spec -> spec.supplier(context ->
35+
new ClassNameHolder(this.importMetadata == null ? null : this.importMetadata.getClassName())));
36+
}
37+
38+
@Override
39+
public void setImportMetadata(AnnotationMetadata importMetadata) {
40+
this.importMetadata = importMetadata;
41+
}
42+
43+
public @Nullable AnnotationMetadata getImportMetadata() {
44+
return this.importMetadata;
45+
}
46+
47+
public record ClassNameHolder(@Nullable String className) {}
48+
}

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

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
/*
2+
* Copyright 2002-2025 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.context.testfixture.context.annotation.registrar;
18+
19+
import org.springframework.context.annotation.Import;
20+
import org.springframework.context.testfixture.beans.factory.ImportAwareBeanRegistrar;
21+
22+
@Import(ImportAwareBeanRegistrar.class)
23+
public class ImportAwareBeanRegistrarConfiguration {
24+
}

0 commit comments

Comments
 (0)