Retrieve Environment from RegisteredBean.

Author: christophstrobl

src/main/java/org/springframework/data/aot/ManagedTypesBeanRegistrationAotProcessor.java

@@ -33,13 +33,15 @@
 import org.springframework.context.EnvironmentAware;
 import org.springframework.core.ResolvableType;
 import org.springframework.core.env.Environment;
+import org.springframework.core.env.EnvironmentCapable;
 import org.springframework.core.env.StandardEnvironment;
 import org.springframework.data.domain.ManagedTypes;
 import org.springframework.data.util.Lazy;
 import org.springframework.data.util.TypeCollector;
 import org.springframework.data.util.TypeContributor;
 import org.springframework.data.util.TypeUtils;
 import org.springframework.util.ClassUtils;
+import org.springframework.util.ObjectUtils;
 import org.springframework.util.StringUtils;
 
 /**
@@ -54,7 +56,8 @@ public class ManagedTypesBeanRegistrationAotProcessor implements BeanRegistratio
 
 	private final Log logger = LogFactory.getLog(getClass());
 	private @Nullable String moduleIdentifier;
-	private Lazy<Environment> environment = Lazy.of(StandardEnvironment::new);
+	private static final Lazy<Environment> DEFAULT_ENVIRONMENT = Lazy.of(StandardEnvironment::new);
+	private @Nullable Environment environment = null;
 
 	public void setModuleIdentifier(@Nullable String moduleIdentifier) {
 		this.moduleIdentifier = moduleIdentifier;
@@ -67,7 +70,7 @@ public String getModuleIdentifier() {
 
 	@Override
 	public void setEnvironment(Environment environment) {
-		this.environment = Lazy.of(environment);
+		this.environment = environment;
 	}
 
 	@Override
@@ -77,7 +80,7 @@ public void setEnvironment(Environment environment) {
 			return null;
 		}
 
-		DefaultAotContext aotContext = new DefaultAotContext(registeredBean.getBeanFactory(), environment.get());
+		DefaultAotContext aotContext = new DefaultAotContext(registeredBean.getBeanFactory(), getConfiguredEnvironmentOrTryToResolveOne(registeredBean));
 		return contribute(aotContext, resolveManagedTypes(registeredBean), registeredBean);
 	}
 
@@ -183,4 +186,21 @@ protected boolean matchesByType(@Nullable Class<?> beanType) {
 	protected boolean matchesPrefix(@Nullable String beanName) {
 		return StringUtils.startsWithIgnoreCase(beanName, getModuleIdentifier());
 	}
+
+	protected Environment getConfiguredEnvironmentOrTryToResolveOne(RegisteredBean registeredBean) {
+
+		if (this.environment != null) {
+			return this.environment;
+		}
+
+		if (registeredBean.getBeanFactory() instanceof EnvironmentCapable ec) {
+			return ec.getEnvironment();
+		} else {
+			String[] beanNamesForType = registeredBean.getBeanFactory().getBeanNamesForType(Environment.class);
+			if (!ObjectUtils.isEmpty(beanNamesForType)) {
+				return registeredBean.getBeanFactory().getBean(beanNamesForType[0], Environment.class);
+			}
+		}
+		return DEFAULT_ENVIRONMENT.get();
+	}
 }

src/test/java/org/springframework/data/aot/ManagedTypesBeanRegistrationAotProcessorUnitTests.java

@@ -41,10 +41,12 @@
 import org.springframework.beans.factory.support.RootBeanDefinition;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.core.env.Environment;
 import org.springframework.data.aot.ManagedTypesRegistrationAotContribution.ManagedTypesInstanceCodeFragment;
 import org.springframework.data.domain.ManagedTypes;
 import org.springframework.javapoet.MethodSpec;
 import org.springframework.javapoet.MethodSpec.Builder;
+import org.springframework.mock.env.MockEnvironment;
 import org.springframework.test.util.ReflectionTestUtils;
 
 /**
@@ -198,15 +200,15 @@ void generatesInstanceSupplierCodeFragmentToAvoidDuplicateInvocations() {
 	@Test // GH-2680
 	void generatesInstanceSupplierCodeFragmentToAvoidDuplicateInvocationsForEmptyManagedTypes() {
 
-		beanFactory.registerBeanDefinition("commons.managed-types", BeanDefinitionBuilder.rootBeanDefinition(EmptyManagedTypes.class).getBeanDefinition());
+		beanFactory.registerBeanDefinition("commons.managed-types",
+				BeanDefinitionBuilder.rootBeanDefinition(EmptyManagedTypes.class).getBeanDefinition());
 		RegisteredBean registeredBean = RegisteredBean.of(beanFactory, "commons.managed-types");
 
 		BeanRegistrationAotContribution contribution = createPostProcessor("commons")
 				.processAheadOfTime(RegisteredBean.of(beanFactory, "commons.managed-types"));
 
 		AotTestCodeContributionBuilder.withContextFor(this.getClass()).writeContentFor(contribution).compile(it -> {
 
-
 			InstanceSupplier<ManagedTypes> types = ReflectionTestUtils
 					.invokeMethod(it.getAllCompiledClasses().iterator().next(), "instance");
 			try {
@@ -266,6 +268,40 @@ void canGenerateCodeReturnsFalseIfNoFactoryMethodPresent() {
 		assertThat(fragment.canGenerateCode()).isFalse();
 	}
 
+	@Test // GH-3414
+	void usesConfiguredEnvironment() {
+
+		MockEnvironment env = spy(new MockEnvironment());
+		ManagedTypesBeanRegistrationAotProcessor processor = createPostProcessor("commons");
+		processor.setEnvironment(env);
+
+		beanFactory.registerBeanDefinition("commons.managed-types", managedTypesDefinition);
+
+		BeanRegistrationAotContribution contribution = processor
+				.processAheadOfTime(RegisteredBean.of(beanFactory, "commons.managed-types"));
+
+		contribution.applyTo(new TestGenerationContext(Object.class), null);
+
+		verify(env).getProperty(eq("spring.aot.data.accessors.enabled"), eq(Boolean.class), eq(true));
+	}
+
+	@Test // GH-3414
+	void usesUsesEnvironmentFromBeanIfNotSet() {
+
+		MockEnvironment env = spy(new MockEnvironment());
+		ManagedTypesBeanRegistrationAotProcessor processor = createPostProcessor("commons");
+
+		beanFactory.registerBeanDefinition("commons.managed-types", managedTypesDefinition);
+		beanFactory.registerBeanDefinition("environment", new RootBeanDefinition(Environment.class, () -> env));
+
+		BeanRegistrationAotContribution contribution = processor
+				.processAheadOfTime(RegisteredBean.of(beanFactory, "commons.managed-types"));
+
+		contribution.applyTo(new TestGenerationContext(Object.class), null);
+
+		verify(env).getProperty(eq("spring.aot.data.accessors.enabled"), eq(Boolean.class), eq(true));
+	}
+
 	private ManagedTypesBeanRegistrationAotProcessor createPostProcessor(String moduleIdentifier) {
 		ManagedTypesBeanRegistrationAotProcessor postProcessor = new ManagedTypesBeanRegistrationAotProcessor();
 		postProcessor.setModuleIdentifier(moduleIdentifier);