Polishing.

Introduce Environment to AotContext.
Move JSON classes to aot/generate package.

Refine exposed API.

See #3265

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

@@ -30,7 +30,9 @@ import org.springframework.beans.factory.config.BeanDefinition;
 import org.springframework.beans.factory.config.BeanReference;
 import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
 import org.springframework.beans.factory.support.RootBeanDefinition;
-import org.springframework.core.SpringProperties;
+import org.springframework.core.env.Environment;
+import org.springframework.core.env.EnvironmentCapable;
+import org.springframework.core.env.StandardEnvironment;
 import org.springframework.data.util.TypeScanner;
 import org.springframework.util.Assert;
 
@@ -45,17 +47,13 @@ import org.springframework.util.Assert;
  * @author Christoph Strobl
  * @author John Blum
  * @author Mark Paluch
- * @see BeanFactory
  * @since 3.0
+ * @see BeanFactory
  */
-public interface AotContext {
+public interface AotContext extends EnvironmentCapable {
 
 	String GENERATED_REPOSITORIES_ENABLED = "spring.aot.repositories.enabled";
 
-	static boolean aotGeneratedRepositoriesEnabled() {
-		return SpringProperties.getFlag(GENERATED_REPOSITORIES_ENABLED);
-	}
-
 	/**
 	 * Create an {@link AotContext} backed by the given {@link BeanFactory}.
 	 *
@@ -67,7 +65,24 @@ public interface AotContext {
 
 		Assert.notNull(beanFactory, "BeanFactory must not be null");
 
-		return new DefaultAotContext(beanFactory);
+		return new DefaultAotContext(beanFactory, new StandardEnvironment());
+	}
+
+	/**
+	 * Create an {@link AotContext} backed by the given {@link BeanFactory}.
+	 *
+	 * @param beanFactory reference to the {@link BeanFactory}; must not be {@literal null}.
+	 * @return a new instance of {@link AotContext}.
+	 * @param environment reference to the {@link Environment}; must not be {@literal null}.
+	 * @return a new instance of {@link AotContext}.
+	 * @see BeanFactory
+	 */
+	static AotContext from(BeanFactory beanFactory, Environment environment) {
+
+		Assert.notNull(beanFactory, "BeanFactory must not be null");
+		Assert.notNull(environment, "Environment must not be null");
+
+		return new DefaultAotContext(beanFactory, environment);
 	}
 
 	/**

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

@@ -28,6 +28,7 @@ import org.springframework.beans.factory.config.BeanDefinition;
 import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
 import org.springframework.beans.factory.support.DefaultListableBeanFactory;
 import org.springframework.beans.factory.support.RootBeanDefinition;
+import org.springframework.core.env.Environment;
 import org.springframework.util.ClassUtils;
 
 /**
@@ -40,9 +41,12 @@ class DefaultAotContext implements AotContext {
 
 	private final ConfigurableListableBeanFactory factory;
 
-	public DefaultAotContext(BeanFactory beanFactory) {
+	private final Environment environment;
+
+	public DefaultAotContext(BeanFactory beanFactory, Environment environment) {
 		factory = beanFactory instanceof ConfigurableListableBeanFactory cbf ? cbf
 				: new DefaultListableBeanFactory(beanFactory);
+		this.environment = environment;
 	}
 
 	@Override
@@ -50,6 +54,11 @@ class DefaultAotContext implements AotContext {
 		return factory;
 	}
 
+	@Override
+	public Environment getEnvironment() {
+		return environment;
+	}
+
 	@Override
 	public TypeIntrospector introspectType(String typeName) {
 		return new DefaultTypeIntrospector(typeName);
@@ -138,4 +147,5 @@ class DefaultAotContext implements AotContext {
 			return factory.getType(beanName, false);
 		}
 	}
+
 }

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

@@ -31,8 +31,12 @@ import org.springframework.beans.factory.aot.BeanRegistrationAotProcessor;
 import org.springframework.beans.factory.config.ConstructorArgumentValues.ValueHolder;
 import org.springframework.beans.factory.support.RegisteredBean;
 import org.springframework.beans.factory.support.RootBeanDefinition;
+import org.springframework.context.EnvironmentAware;
 import org.springframework.core.ResolvableType;
+import org.springframework.core.env.Environment;
+import org.springframework.core.env.StandardEnvironment;
 import org.springframework.data.domain.ManagedTypes;
+import org.springframework.data.util.Lazy;
 import org.springframework.data.util.QTypeContributor;
 import org.springframework.data.util.TypeContributor;
 import org.springframework.data.util.TypeUtils;
@@ -47,10 +51,11 @@ import org.springframework.util.StringUtils;
  * @author John Blum
  * @since 3.0
  */
-public class ManagedTypesBeanRegistrationAotProcessor implements BeanRegistrationAotProcessor {
+public class ManagedTypesBeanRegistrationAotProcessor implements BeanRegistrationAotProcessor, EnvironmentAware {
 
 	private final Log logger = LogFactory.getLog(getClass());
 	private @Nullable String moduleIdentifier;
+	private Lazy<Environment> environment = Lazy.of(StandardEnvironment::new);
 
 	public void setModuleIdentifier(@Nullable String moduleIdentifier) {
 		this.moduleIdentifier = moduleIdentifier;
@@ -61,6 +66,11 @@ public class ManagedTypesBeanRegistrationAotProcessor implements BeanRegistratio
 		return this.moduleIdentifier;
 	}
 
+	@Override
+	public void setEnvironment(Environment environment) {
+		this.environment = Lazy.of(() -> environment);
+	}
+
 	@Override
 	public @Nullable BeanRegistrationAotContribution processAheadOfTime(RegisteredBean registeredBean) {
 
@@ -69,7 +79,8 @@ public class ManagedTypesBeanRegistrationAotProcessor implements BeanRegistratio
 		}
 
 		BeanFactory beanFactory = registeredBean.getBeanFactory();
-		return contribute(AotContext.from(beanFactory), resolveManagedTypes(registeredBean), registeredBean);
+		return contribute(AotContext.from(beanFactory, this.environment.get()), resolveManagedTypes(registeredBean),
+				registeredBean);
 	}
 
 	private ManagedTypes resolveManagedTypes(RegisteredBean registeredBean) {

src/main/java/org/springframework/data/repository/aot/generate/AotFragmentTarget.java

@@ -17,10 +17,9 @@ package org.springframework.data.repository.aot.generate;
 
 import org.jspecify.annotations.Nullable;
 
-import org.springframework.data.repository.aot.generate.json.JSONException;
-import org.springframework.data.repository.aot.generate.json.JSONObject;
-
 /**
+ * Value object capturing metadata about a fragment target.
+ *
  * @author Mark Paluch
  * @since 4.0
  */
@@ -39,4 +38,5 @@ record AotFragmentTarget(String signature, @Nullable String implementation) {
 
 		return fragment;
 	}
+
 }

src/main/java/org/springframework/data/repository/aot/generate/AotQueryMethodGenerationContext.java

@@ -1,19 +1,3 @@
-/*
- * Copyright 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
- *
- *      http://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.
- */
-
 /*
  * Copyright 2025 the original author or authors.
  *
@@ -21,7 +5,7 @@
  * you may not use this file except in compliance with the License.
  * You may obtain a copy of the License at
  *
- *      http://www.apache.org/licenses/LICENSE-2.0
+ *      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,
@@ -68,7 +52,7 @@ public class AotQueryMethodGenerationContext {
 	private final QueryMethod queryMethod;
 	private final RepositoryInformation repositoryInformation;
 	private final AotRepositoryFragmentMetadata targetTypeMetadata;
-	private final AotRepositoryMethodImplementationMetadata targetMethodMetadata;
+	private final MethodMetadata targetMethodMetadata;
 	private final CodeBlocks codeBlocks;
 
 	AotQueryMethodGenerationContext(RepositoryInformation repositoryInformation, Method method, QueryMethod queryMethod,
@@ -79,7 +63,7 @@ public class AotQueryMethodGenerationContext {
 		this.queryMethod = queryMethod;
 		this.repositoryInformation = repositoryInformation;
 		this.targetTypeMetadata = targetTypeMetadata;
-		this.targetMethodMetadata = new AotRepositoryMethodImplementationMetadata(repositoryInformation, method);
+		this.targetMethodMetadata = new MethodMetadata(repositoryInformation, method);
 		this.codeBlocks = new CodeBlocks(targetTypeMetadata);
 	}
 
@@ -87,7 +71,7 @@ public class AotQueryMethodGenerationContext {
 		return targetTypeMetadata;
 	}
 
-	AotRepositoryMethodImplementationMetadata getTargetMethodMetadata() {
+	MethodMetadata getTargetMethodMetadata() {
 		return targetMethodMetadata;
 	}
 

src/main/java/org/springframework/data/repository/aot/generate/AotRepositoryBuilder.java

@@ -34,8 +34,6 @@ import org.springframework.aot.generate.ClassNameGenerator;
 import org.springframework.aot.generate.Generated;
 import org.springframework.data.projection.ProjectionFactory;
 import org.springframework.data.repository.aot.generate.AotRepositoryFragmentMetadata.ConstructorArgument;
-import org.springframework.data.repository.aot.generate.json.JSONException;
-import org.springframework.data.repository.aot.generate.json.JSONObject;
 import org.springframework.data.repository.core.RepositoryInformation;
 import org.springframework.data.repository.core.support.RepositoryComposition;
 import org.springframework.data.repository.core.support.RepositoryFragment;
@@ -59,7 +57,7 @@ class AotRepositoryBuilder {
 	private final AotRepositoryFragmentMetadata generationMetadata;
 
 	private @Nullable Consumer<AotRepositoryConstructorBuilder> constructorCustomizer;
-	private @Nullable BiFunction<Method, RepositoryInformation, MethodContributor<? extends QueryMethod>> methodContributorFunction;
+	private @Nullable BiFunction<Method, RepositoryInformation, @Nullable MethodContributor<? extends QueryMethod>> methodContributorFunction;
 	private ClassCustomizer customizer;
 
 	private AotRepositoryBuilder(RepositoryInformation repositoryInformation, ProjectionFactory projectionFactory) {
@@ -89,7 +87,7 @@ class AotRepositoryBuilder {
 	}
 
 	public AotRepositoryBuilder withQueryMethodContributor(
-			BiFunction<Method, RepositoryInformation, MethodContributor<? extends QueryMethod>> methodContributorFunction) {
+			BiFunction<Method, RepositoryInformation, @Nullable MethodContributor<? extends QueryMethod>> methodContributorFunction) {
 		this.methodContributorFunction = methodContributorFunction;
 		return this;
 	}
@@ -144,11 +142,7 @@ class AotRepositoryBuilder {
 		AotRepositoryMetadata metadata = new AotRepositoryMetadata(repositoryInformation.getRepositoryInterface().getName(),
 				"", repositoryType, methodMetadata);
 
-		try {
-			return new AotBundle(javaFile, metadata.toJson());
-		} catch (JSONException e) {
-			throw new IllegalStateException(e);
-		}
+		return new AotBundle(javaFile, metadata.toJson());
 	}
 
 	private void contributeMethod(Method method, RepositoryComposition repositoryComposition,
@@ -218,7 +212,7 @@ class AotRepositoryBuilder {
 	public Map<String, TypeName> getAutowireFields() {
 		Map<String, TypeName> autowireFields = new LinkedHashMap<>(generationMetadata.getConstructorArguments().size());
 		for (Map.Entry<String, ConstructorArgument> entry : generationMetadata.getConstructorArguments().entrySet()) {
-			autowireFields.put(entry.getKey(), entry.getValue().getTypeName());
+			autowireFields.put(entry.getKey(), entry.getValue().typeName());
 		}
 		return autowireFields;
 	}

src/main/java/org/springframework/data/repository/aot/generate/AotRepositoryConstructorBuilder.java

@@ -15,7 +15,6 @@
  */
 package org.springframework.data.repository.aot.generate;
 
-import java.util.List;
 import java.util.Map.Entry;
 
 import javax.lang.model.element.Modifier;
@@ -34,6 +33,7 @@ import org.springframework.javapoet.TypeName;
  * @author Mark Paluch
  * @since 4.0
  */
+// TODO: extract constructor contributor in a similar way to MethodContributor.
 public class AotRepositoryConstructorBuilder {
 
 	private final RepositoryInformation repositoryInformation;
@@ -103,7 +103,7 @@ public class AotRepositoryConstructorBuilder {
 		MethodSpec.Builder builder = MethodSpec.constructorBuilder().addModifiers(Modifier.PUBLIC);
 
 		for (Entry<String, ConstructorArgument> parameter : this.metadata.getConstructorArguments().entrySet()) {
-			builder.addParameter(parameter.getValue().getTypeName(), parameter.getKey());
+			builder.addParameter(parameter.getValue().typeName(), parameter.getKey());
 		}
 
 		customizer.customize(repositoryInformation, builder);
@@ -118,19 +118,6 @@ public class AotRepositoryConstructorBuilder {
 		return builder.build();
 	}
 
-	private static TypeName getDefaultStoreRepositoryImplementationType(RepositoryInformation repositoryInformation) {
-
-		ResolvableType resolvableType = ResolvableType.forClass(repositoryInformation.getRepositoryBaseClass());
-		if (resolvableType.hasGenerics()) {
-			List<Class<?>> generics = List.of();
-			if (resolvableType.getGenerics().length == 2) { // TODO: Find some other way to resolve generics
-				generics = List.of(repositoryInformation.getDomainType(), repositoryInformation.getIdType());
-			}
-			return ParameterizedTypeName.get(repositoryInformation.getRepositoryBaseClass(), generics.toArray(Class[]::new));
-		}
-		return TypeName.get(repositoryInformation.getRepositoryBaseClass());
-	}
-
 	/**
 	 * Customizer for the AOT repository constructor.
 	 */

src/main/java/org/springframework/data/repository/aot/generate/AotRepositoryFragmentMetadata.java

@@ -31,6 +31,7 @@ import org.springframework.javapoet.TypeName;
 /**
  * @author Christoph Strobl
  */
+// TODO: Can we make this package-private?
 public class AotRepositoryFragmentMetadata {
 
 	private final ClassName className;
@@ -94,31 +95,11 @@ public class AotRepositoryFragmentMetadata {
 		this.constructorArguments.put(parameterName, new ConstructorArgument(parameterName, type, fieldName));
 	}
 
-	static class ConstructorArgument {
-		String parameterName;
-		@Nullable String fieldName;
-		TypeName typeName;
-
-		public ConstructorArgument(String parameterName,TypeName typeName, String fieldName) {
-			this.parameterName = parameterName;
-			this.fieldName = fieldName;
-			this.typeName = typeName;
-		}
+	public record ConstructorArgument(String parameterName, TypeName typeName, @Nullable String fieldName) {
 
 		boolean isForLocalField() {
 			return fieldName != null;
 		}
 
-		public String getParameterName() {
-			return parameterName;
-		}
-
-		public String getFieldName() {
-			return fieldName;
-		}
-
-		public TypeName getTypeName() {
-			return typeName;
-		}
 	}
 }

src/main/java/org/springframework/data/repository/aot/generate/AotRepositoryMetadata.java

@@ -17,11 +17,9 @@ package org.springframework.data.repository.aot.generate;
 
 import java.util.List;
 
-import org.springframework.data.repository.aot.generate.json.JSONArray;
-import org.springframework.data.repository.aot.generate.json.JSONException;
-import org.springframework.data.repository.aot.generate.json.JSONObject;
-
 /**
+ * Value object capturing metadata about a repository.
+ *
  * @author Mark Paluch
  * @since 4.0
  */
@@ -33,6 +31,12 @@ record AotRepositoryMetadata(String name, String module,
 		IMPERATIVE, REACTIVE
 	}
 
+	/**
+	 * Convert this {@link AotRepositoryMetadata} to a {@link JSONObject}.
+	 *
+	 * @return
+	 * @throws JSONException
+	 */
 	JSONObject toJson() throws JSONException {
 
 		JSONObject metadata = new JSONObject();
@@ -49,6 +53,6 @@ record AotRepositoryMetadata(String name, String module,
 		metadata.put("methods", methods);
 
 		return metadata;
-
 	}
+
 }

src/main/java/org/springframework/data/repository/aot/generate/AotRepositoryMethod.java

@@ -15,18 +15,25 @@
  */
 package org.springframework.data.repository.aot.generate;
 
+import java.util.Map;
+
 import org.jspecify.annotations.Nullable;
 
-import org.springframework.data.repository.aot.generate.json.JSONException;
-import org.springframework.data.repository.aot.generate.json.JSONObject;
-
 /**
+ * Value object capturing metadata about a repository method.
+ *
  * @author Mark Paluch
  * @since 4.0
  */
 record AotRepositoryMethod(String name, String signature, @Nullable QueryMetadata query,
 		@Nullable AotFragmentTarget fragment) {
 
+	/**
+	 * Convert this {@link AotRepositoryMethod} to a {@link JSONObject}.
+	 *
+	 * @return
+	 * @throws JSONException
+	 */
 	public JSONObject toJson() throws JSONException {
 
 		JSONObject method = new JSONObject();
@@ -34,11 +41,23 @@ record AotRepositoryMethod(String name, String signature, @Nullable QueryMetadat
 		method.put("signature", signature());
 
 		if (query() != null) {
-			method.put("query", query().toJson());
+			method.put("query", queryMetadataToJson(query()));
 		} else if (fragment() != null) {
 			method.put("fragment", fragment().toJson());
 		}
 
 		return method;
 	}
+
+	static JSONObject queryMetadataToJson(QueryMetadata queryMetadata) throws JSONException {
+
+		JSONObject query = new JSONObject();
+
+		for (Map.Entry<String, Object> entry : queryMetadata.serialize().entrySet()) {
+			query.put(entry.getKey(), entry.getValue());
+		}
+
+		return query;
+	}
+
 }

src/main/java/org/springframework/data/repository/aot/generate/AotRepositoryMethodBuilder.java

@@ -18,6 +18,8 @@ package org.springframework.data.repository.aot.generate;
 import java.lang.reflect.Method;
 import java.lang.reflect.Parameter;
 import java.lang.reflect.TypeVariable;
+import java.util.function.BiConsumer;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 
 import javax.lang.model.element.Modifier;
@@ -40,12 +42,12 @@ import org.springframework.util.StringUtils;
  * @author Mark Paluch
  * @since 4.0
  */
-public class AotRepositoryMethodBuilder {
+class AotRepositoryMethodBuilder {
 
 	private final AotQueryMethodGenerationContext context;
 
-	private RepositoryMethodContribution contribution = (context) -> CodeBlock.builder().build();
-	private RepositoryMethodCustomizer customizer = (context, body) -> {};
+	private Function<AotQueryMethodGenerationContext, CodeBlock> contribution = (context) -> CodeBlock.builder().build();
+	private BiConsumer<AotQueryMethodGenerationContext, MethodSpec.Builder> customizer = (context, body) -> {};
 
 	AotRepositoryMethodBuilder(AotQueryMethodGenerationContext context) {
 
@@ -70,36 +72,40 @@ public class AotRepositoryMethodBuilder {
 	}
 
 	/**
-	 * Register a {@link RepositoryMethodContribution} for the repository interface that can contribute a query method
-	 * implementation block.
+	 * Register a {@link org.springframework.data.repository.aot.generate.MethodContributor.RepositoryMethodContribution}
+	 * for the repository interface that can contribute a query method implementation block.
 	 *
 	 * @param contribution
 	 * @return
 	 */
-	public AotRepositoryMethodBuilder contribute(RepositoryMethodContribution contribution) {
+	public AotRepositoryMethodBuilder contribute(Function<AotQueryMethodGenerationContext, CodeBlock> contribution) {
 		this.contribution = contribution;
 		return this;
 	}
 
 	/**
-	 * Register a query method customizer that is applied after a successful {@link RepositoryMethodContribution}.
+	 * Register a query method customizer that is applied after a successful
+	 * {@link org.springframework.data.repository.aot.generate.MethodContributor.RepositoryMethodContribution}.
 	 *
 	 * @param customizer
 	 * @return
 	 */
-	public AotRepositoryMethodBuilder customize(RepositoryMethodCustomizer customizer) {
+	public AotRepositoryMethodBuilder customize(
+			BiConsumer<AotQueryMethodGenerationContext, MethodSpec.Builder> customizer) {
 		this.customizer = customizer;
 		return this;
 	}
 
 	/**
-	 * Builds an AOT repository method if {@link RepositoryMethodContribution} can contribute a method.
+	 * Builds an AOT repository method if
+	 * {@link org.springframework.data.repository.aot.generate.MethodContributor.RepositoryMethodContribution} can
+	 * contribute a method.
 	 *
 	 * @return the {@link MethodSpec} or {@literal null}, if the method cannot be contributed.
 	 */
 	public MethodSpec buildMethod() {
 
-		CodeBlock methodBody = contribution.contribute(context);
+		CodeBlock methodBody = contribution.apply(context);
 
 		MethodSpec.Builder builder = MethodSpec.methodBuilder(context.getMethod().getName()).addModifiers(Modifier.PUBLIC);
 		builder.returns(TypeName.get(context.getReturnType().getType()));
@@ -115,25 +121,9 @@ public class AotRepositoryMethodBuilder {
 						.getMethodArguments().values().stream().map(it -> it.type.toString()).collect(Collectors.toList())));
 		context.getTargetMethodMetadata().getMethodArguments().forEach((name, spec) -> builder.addParameter(spec));
 		builder.addCode(methodBody);
-		customizer.customize(context, builder);
+		customizer.accept(context, builder);
 
 		return builder.build();
 	}
 
-	/**
-	 * AOT contribution from a {@link AotRepositoryMethodBuilder} used to contribute a repository query method body.
-	 */
-	public interface RepositoryMethodContribution {
-
-		CodeBlock contribute(AotQueryMethodGenerationContext context);
-	}
-
-	/**
-	 * Customizer for a contributed AOT repository query method.
-	 */
-	public interface RepositoryMethodCustomizer {
-
-		void customize(AotQueryMethodGenerationContext context, MethodSpec.Builder builder);
-
-	}
 }

src/main/java/org/springframework/data/repository/aot/generate/CodeBlocks.java

@@ -25,6 +25,7 @@ import org.springframework.util.StringUtils;
  *
  * @author Christoph Strobl
  */
+@Deprecated(forRemoval = true)
 public class CodeBlocks {
 
 	private final AotRepositoryFragmentMetadata metadata;

src/main/java/org/springframework/data/repository/aot/generate/JSON.java

@@ -14,8 +14,11 @@
  * limitations under the License.
  */
 
-package org.springframework.data.repository.aot.generate.json;
+package org.springframework.data.repository.aot.generate;
 
+import org.jspecify.annotations.NullUnmarked;
+
+@NullUnmarked
 class JSON {
 
 	static double checkDouble(double d) throws JSONException {

src/main/java/org/springframework/data/repository/aot/generate/JSONArray.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package org.springframework.data.repository.aot.generate.json;
+package org.springframework.data.repository.aot.generate;
 
 import java.lang.reflect.Array;
 import java.util.ArrayList;
@@ -22,6 +22,8 @@ import java.util.Collection;
 import java.util.Iterator;
 import java.util.List;
 
+import org.jspecify.annotations.NullUnmarked;
+
 // Note: this class was written without inspecting the non-free org.json source code.
 
 /**
@@ -40,7 +42,8 @@ import java.util.List;
  * should not be subclassed. In particular, self-use by overridable methods is not specified. See <i>Effective Java</i>
  * Item 17, "Design and Document or inheritance or else prohibit it" for further information.
  */
-public class JSONArray {
+@NullUnmarked
+class JSONArray {
 
 	private final List<Object> values;
 

src/main/java/org/springframework/data/repository/aot/generate/JSONException.java

@@ -14,10 +14,12 @@
  * limitations under the License.
  */
 
-package org.springframework.data.repository.aot.generate.json;
+package org.springframework.data.repository.aot.generate;
 
 // Note: this class was written without inspecting the non-free org.json source code.
 
+import org.jspecify.annotations.NullUnmarked;
+
 /**
  * Thrown to indicate a problem with the JSON API. Such problems include:
  * <ul>
@@ -43,7 +45,8 @@ package org.springframework.data.repository.aot.generate.json;
  * }
  * </pre>
  */
-public class JSONException extends Exception {
+@NullUnmarked
+class JSONException extends RuntimeException {
 
 	public JSONException(String s) {
 		super(s);

src/main/java/org/springframework/data/repository/aot/generate/JSONObject.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package org.springframework.data.repository.aot.generate.json;
+package org.springframework.data.repository.aot.generate;
 
 import java.util.ArrayList;
 import java.util.Collection;
@@ -22,6 +22,8 @@ import java.util.Iterator;
 import java.util.LinkedHashMap;
 import java.util.Map;
 
+import org.jspecify.annotations.NullUnmarked;
+
 // Note: this class was written without inspecting the non-free org.json source code.
 
 /**
@@ -63,7 +65,8 @@ import java.util.Map;
  * should not be subclassed. In particular, self-use by overrideable methods is not specified. See <i>Effective Java</i>
  * Item 17, "Design and Document or inheritance or else prohibit it" for further information.
  */
-public class JSONObject {
+@NullUnmarked
+class JSONObject {
 
 	private static final Double NEGATIVE_ZERO = -0d;
 

src/main/java/org/springframework/data/repository/aot/generate/JSONStringer.java

@@ -14,12 +14,14 @@
  * limitations under the License.
  */
 
-package org.springframework.data.repository.aot.generate.json;
+package org.springframework.data.repository.aot.generate;
 
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 
+import org.jspecify.annotations.NullUnmarked;
+
 // Note: this class was written without inspecting the non-free org.json source code.
 
 /**
@@ -53,7 +55,8 @@ import java.util.List;
  * overrideable methods is not specified. See <i>Effective Java</i> Item 17, "Design and Document or inheritance or else
  * prohibit it" for further information.
  */
-public class JSONStringer {
+@NullUnmarked
+class JSONStringer {
 
 	/**
 	 * The output data, containing at most one top-level array or object.

src/main/java/org/springframework/data/repository/aot/generate/JSONTokener.java

@@ -14,10 +14,12 @@
  * limitations under the License.
  */
 
-package org.springframework.data.repository.aot.generate.json;
+package org.springframework.data.repository.aot.generate;
 
 // Note: this class was written without inspecting the non-free org.json source code.
 
+import org.jspecify.annotations.NullUnmarked;
+
 /**
  * Parses a JSON (<a href="https://www.ietf.org/rfc/rfc4627.txt">RFC 4627</a>) encoded string into the corresponding
  * object. Most clients of this class will use only need the {@link #JSONTokener(String) constructor} and
@@ -51,7 +53,8 @@ package org.springframework.data.repository.aot.generate.json;
  * overrideable methods is not specified. See <i>Effective Java</i> Item 17, "Design and Document or inheritance or else
  * prohibit it" for further information.
  */
-public class JSONTokener {
+@NullUnmarked
+class JSONTokener {
 
 	/**
 	 * The input JSON.

src/main/java/org/springframework/data/repository/aot/generate/MethodContributor.java

@@ -21,6 +21,7 @@ import org.jspecify.annotations.NonNull;
 import org.jspecify.annotations.Nullable;
 
 import org.springframework.data.repository.query.QueryMethod;
+import org.springframework.javapoet.CodeBlock;
 import org.springframework.javapoet.MethodSpec;
 
 /**
@@ -49,16 +50,20 @@ public abstract class MethodContributor<M extends QueryMethod> {
 	 */
 	public static <M extends QueryMethod> QueryMethodMetadataContributorBuilder<M> forQueryMethod(M queryMethod) {
 
-		return new QueryMethodMetadataContributorBuilder<M>() {
+		return new QueryMethodMetadataContributorBuilder<>() {
 
 			@Override
 			public MethodContributor<M> metadataOnly(QueryMetadata metadata) {
-				return new MetadataMethodContributor<>(queryMethod, metadata);
+				return new MetadataContributor<>(queryMethod, metadata);
 			}
 
+
 			@Override
 			public QueryMethodContributorBuilder<M> withMetadata(QueryMetadata metadata) {
-				return builderConsumer -> new AotMethodContributor<>(queryMethod, metadata, builderConsumer);
+
+				return contribution -> new GeneratedMethodContributor<>(queryMethod, metadata, b -> {
+					b.contribute(contribution::contribute);
+				});
 			}
 		};
 	}
@@ -86,42 +91,6 @@ public abstract class MethodContributor<M extends QueryMethod> {
 	 */
 	public abstract @Nullable MethodSpec contribute(AotQueryMethodGenerationContext context);
 
-	private static class MetadataMethodContributor<M extends QueryMethod> extends MethodContributor<M> {
-
-		private MetadataMethodContributor(M queryMethod, QueryMetadata metadata) {
-			super(queryMethod, metadata);
-		}
-
-		@Override
-		public @Nullable MethodSpec contribute(AotQueryMethodGenerationContext context) {
-			return null;
-		}
-	}
-
-	private static class AotMethodContributor<M extends QueryMethod> extends MethodContributor<M> {
-
-		private final Consumer<AotRepositoryMethodBuilder> builderConsumer;
-
-		private AotMethodContributor(M queryMethod, QueryMetadata metadata,
-				Consumer<AotRepositoryMethodBuilder> builderConsumer) {
-			super(queryMethod, metadata);
-			this.builderConsumer = builderConsumer;
-		}
-
-		@Override
-		public boolean contributesMethodSpec() {
-			return true;
-		}
-
-		@Override
-		public @NonNull MethodSpec contribute(AotQueryMethodGenerationContext context) {
-
-			AotRepositoryMethodBuilder builder = new AotRepositoryMethodBuilder(context);
-			builderConsumer.accept(builder);
-			return builder.buildMethod();
-		}
-	}
-
 	/**
 	 * Initial builder for a query method contributor. This builder allows returning a {@link MethodContributor} using
 	 * metadata-only (i.e. no code contribution) or a {@link QueryMethodContributorBuilder} accepting code contributions.
@@ -161,18 +130,62 @@ public abstract class MethodContributor<M extends QueryMethod> {
 		 * @param contribution the method contribution that provides the method to be added to the repository.
 		 * @return the method contributor to use.
 		 */
-		default MethodContributor<M> contribute(AotRepositoryMethodBuilder.RepositoryMethodContribution contribution) {
-			return using(builder -> builder.contribute(contribution));
+		MethodContributor<M> contribute(RepositoryMethodContribution contribution);
+
+	}
+
+	/**
+	 * AOT contribution from a {@link AotRepositoryMethodBuilder} used to contribute a repository query method body.
+	 */
+	public interface RepositoryMethodContribution {
+
+		CodeBlock contribute(AotQueryMethodGenerationContext context);
+	}
+
+	/**
+	 * Customizer for a contributed AOT repository query method.
+	 */
+	public interface RepositoryMethodCustomizer {
+
+		void customize(AotQueryMethodGenerationContext context, MethodSpec.Builder builder);
+
+	}
+
+	private static class MetadataContributor<M extends QueryMethod> extends MethodContributor<M> {
+
+		private MetadataContributor(M queryMethod, QueryMetadata metadata) {
+			super(queryMethod, metadata);
 		}
 
-		/**
-		 * Terminal method accepting a consumer that uses {@link AotRepositoryMethodBuilder} to build the method.
-		 *
-		 * @param builderConsumer consumer method being provided with the {@link AotRepositoryMethodBuilder} that provides
-		 *          the method to be added to the repository.
-		 * @return the method contributor to use.
-		 */
-		MethodContributor<M> using(Consumer<AotRepositoryMethodBuilder> builderConsumer);
+		@Override
+		public @Nullable MethodSpec contribute(AotQueryMethodGenerationContext context) {
+			return null;
+		}
+
+	}
+
+	private static class GeneratedMethodContributor<M extends QueryMethod> extends MethodContributor<M> {
+
+		private final Consumer<AotRepositoryMethodBuilder> builderConsumer;
+
+		private GeneratedMethodContributor(M queryMethod, QueryMetadata metadata,
+				Consumer<AotRepositoryMethodBuilder> builderConsumer) {
+			super(queryMethod, metadata);
+			this.builderConsumer = builderConsumer;
+		}
+
+		@Override
+		public boolean contributesMethodSpec() {
+			return true;
+		}
+
+		@Override
+		public @NonNull MethodSpec contribute(AotQueryMethodGenerationContext context) {
+
+			AotRepositoryMethodBuilder builder = new AotRepositoryMethodBuilder(context);
+			builderConsumer.accept(builder);
+			return builder.buildMethod();
+		}
 
 	}
 

src/main/java/org/springframework/data/repository/aot/generate/MethodMetadata.java

@@ -28,15 +28,17 @@ import org.springframework.javapoet.ParameterSpec;
 import org.springframework.javapoet.TypeName;
 
 /**
+ * Metadata about an AOT Repository method.
+ *
  * @author Christoph Strobl
  */
-class AotRepositoryMethodImplementationMetadata {
+class MethodMetadata {
 
 	private final Map<String, ParameterSpec> methodArguments = new LinkedHashMap<>();
 	private final ResolvableType actualReturnType;
 	private final ResolvableType returnType;
 
-	public AotRepositoryMethodImplementationMetadata(RepositoryInformation repositoryInformation, Method method) {
+	public MethodMetadata(RepositoryInformation repositoryInformation, Method method) {
 
 		this.returnType = repositoryInformation.getReturnType(method).toResolvableType();
 		this.actualReturnType = ResolvableType.forType(repositoryInformation.getReturnedDomainClass(method));
@@ -64,7 +66,8 @@ class AotRepositoryMethodImplementationMetadata {
 		this.methodArguments.put(parameterSpec.name, parameterSpec);
 	}
 
-	Map<String, ParameterSpec> getMethodArguments() {
+	public Map<String, ParameterSpec> getMethodArguments() {
 		return methodArguments;
 	}
+
 }

src/main/java/org/springframework/data/repository/aot/generate/QueryMetadata.java

@@ -17,24 +17,17 @@ package org.springframework.data.repository.aot.generate;
 
 import java.util.Map;
 
-import org.springframework.data.repository.aot.generate.json.JSONException;
-import org.springframework.data.repository.aot.generate.json.JSONObject;
-
 /**
+ * Interface providing metadata about a query. Name and multiplicity of keys is subject to a provider's implementation.
+ *
  * @author Mark Paluch
+ * @since 4.0
  */
 public interface QueryMetadata {
 
+	/**
+	 * @return serialize query metadata to a {@link Map} of key/value pairs using simple types (string, numbers).
+	 */
 	Map<String, Object> serialize();
 
-	public default JSONObject toJson() throws JSONException {
-
-		JSONObject query = new JSONObject();
-
-		for (Map.Entry<String, Object> entry : serialize().entrySet()) {
-			query.put(entry.getKey(), entry.getValue());
-		}
-
-		return query;
-	}
 }

src/main/java/org/springframework/data/repository/aot/generate/RepositoryContributor.java

@@ -26,7 +26,6 @@ import org.springframework.aot.hint.MemberCategory;
 import org.springframework.aot.hint.TypeReference;
 import org.springframework.data.projection.ProjectionFactory;
 import org.springframework.data.projection.SpelAwareProxyProjectionFactory;
-import org.springframework.data.repository.aot.generate.json.JSONException;
 import org.springframework.data.repository.config.AotRepositoryContext;
 import org.springframework.data.repository.core.RepositoryInformation;
 import org.springframework.data.repository.query.QueryMethod;

src/main/java/org/springframework/data/repository/aot/generate/json/package-info.java

@@ -1,5 +0,0 @@
-/**
- * Shaded JSON. This source was originally taken from com.vaadin.external.google:android-json which provides a clean
- * room re-implementation of the org.json APIs and does not include the "Do not use for evil" clause.
- */
-package org.springframework.data.repository.aot.generate.json;

src/main/java/org/springframework/data/repository/aot/generate/package-info.java

@@ -1,5 +1,8 @@
 /**
  * Ahead-of-Time (AOT) generation for Spring Data repositories.
+ * <p>
+ * Contains also shaded JSON. This source was originally taken from com.vaadin.external.google:android-json which
+ * provides a clean room re-implementation of the org.json APIs and does not include the "Do not use for evil" clause.
  */
 @org.jspecify.annotations.NullMarked
 package org.springframework.data.repository.aot.generate;

src/main/java/org/springframework/data/repository/config/DefaultAotRepositoryContext.java

@@ -25,6 +25,7 @@ import org.jspecify.annotations.Nullable;
 
 import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
 import org.springframework.core.annotation.MergedAnnotation;
+import org.springframework.core.env.Environment;
 import org.springframework.data.aot.AotContext;
 import org.springframework.data.repository.core.RepositoryInformation;
 import org.springframework.data.util.Lazy;
@@ -64,6 +65,11 @@ class DefaultAotRepositoryContext implements AotRepositoryContext {
 		return getAotContext().getBeanFactory();
 	}
 
+	@Override
+	public Environment getEnvironment() {
+		return getAotContext().getEnvironment();
+	}
+
 	@Override
 	public Set<String> getBasePackages() {
 		return basePackages == null ? Collections.emptySet() : basePackages;

src/main/java/org/springframework/data/repository/config/RepositoryRegistrationAotContribution.java

@@ -21,7 +21,6 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
-import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.function.BiConsumer;
@@ -80,38 +79,38 @@ public class RepositoryRegistrationAotContribution implements BeanRegistrationAo
 
 	private @Nullable BiFunction<AotRepositoryContext, GenerationContext, @Nullable RepositoryContributor> moduleContribution;
 
-	private final RepositoryRegistrationAotProcessor repositoryRegistrationAotProcessor;
+	private final RepositoryRegistrationAotProcessor aotProcessor;
 
 	/**
 	 * Constructs a new instance of the {@link RepositoryRegistrationAotContribution} initialized with the given, required
 	 * {@link RepositoryRegistrationAotProcessor} from which this contribution was created.
 	 *
-	 * @param repositoryRegistrationAotProcessor reference back to the {@link RepositoryRegistrationAotProcessor} from
-	 *          which this contribution was created.
+	 * @param processor reference back to the {@link RepositoryRegistrationAotProcessor} from which this contribution was
+	 *          created.
 	 * @throws IllegalArgumentException if the {@link RepositoryRegistrationAotProcessor} is {@literal null}.
 	 * @see RepositoryRegistrationAotProcessor
 	 */
 	protected RepositoryRegistrationAotContribution(
-			RepositoryRegistrationAotProcessor repositoryRegistrationAotProcessor) {
+			RepositoryRegistrationAotProcessor processor) {
 
-		Assert.notNull(repositoryRegistrationAotProcessor, "RepositoryRegistrationAotProcessor must not be null");
+		Assert.notNull(processor, "RepositoryRegistrationAotProcessor must not be null");
 
-		this.repositoryRegistrationAotProcessor = repositoryRegistrationAotProcessor;
+		this.aotProcessor = processor;
 	}
 
 	/**
 	 * Factory method used to construct a new instance of {@link RepositoryRegistrationAotContribution} initialized with
 	 * the given, required {@link RepositoryRegistrationAotProcessor} from which this contribution was created.
 	 *
-	 * @param repositoryRegistrationAotProcessor reference back to the {@link RepositoryRegistrationAotProcessor} from
-	 *          which this contribution was created.
+	 * @param processor reference back to the {@link RepositoryRegistrationAotProcessor} from which this contribution was
+	 *          created.
 	 * @return a new instance of {@link RepositoryRegistrationAotContribution}.
 	 * @throws IllegalArgumentException if the {@link RepositoryRegistrationAotProcessor} is {@literal null}.
 	 * @see RepositoryRegistrationAotProcessor
 	 */
 	public static RepositoryRegistrationAotContribution fromProcessor(
-			RepositoryRegistrationAotProcessor repositoryRegistrationAotProcessor) {
-		return new RepositoryRegistrationAotContribution(repositoryRegistrationAotProcessor);
+			RepositoryRegistrationAotProcessor processor) {
+		return new RepositoryRegistrationAotContribution(processor);
 	}
 
 	protected ConfigurableListableBeanFactory getBeanFactory() {
@@ -131,7 +130,7 @@ public class RepositoryRegistrationAotContribution implements BeanRegistrationAo
 	}
 
 	protected RepositoryRegistrationAotProcessor getRepositoryRegistrationAotProcessor() {
-		return this.repositoryRegistrationAotProcessor;
+		return this.aotProcessor;
 	}
 
 	public RepositoryInformation getRepositoryInformation() {
@@ -173,7 +172,7 @@ public class RepositoryRegistrationAotContribution implements BeanRegistrationAo
 	 * @return this.
 	 */
 	public RepositoryRegistrationAotContribution withModuleContribution(
-			@Nullable BiFunction<AotRepositoryContext, GenerationContext, RepositoryContributor> moduleContribution) {
+			@Nullable BiFunction<AotRepositoryContext, GenerationContext, @Nullable RepositoryContributor> moduleContribution) {
 		this.moduleContribution = moduleContribution;
 		return this;
 	}
@@ -360,9 +359,6 @@ public class RepositoryRegistrationAotContribution implements BeanRegistrationAo
 				|| ClassUtils.isPrimitiveArray(type); //
 	}
 
-	// TODO What was this meant to be used for? Was this type filter maybe meant to be used in
-	// the TypeContributor.contribute(:Class, :Predicate :GenerationContext) method
-	// used in the contributeType(..) method above?
 	public Predicate<Class<?>> typeFilter() { // like only document ones. // TODO: As in MongoDB?
 		return Predicates.isTrue();
 	}
@@ -372,7 +368,7 @@ public class RepositoryRegistrationAotContribution implements BeanRegistrationAo
 			RepositoryConfiguration<?> repositoryMetadata) {
 
 		DefaultAotRepositoryContext repositoryContext = new DefaultAotRepositoryContext(
-				AotContext.from(this.getBeanFactory()));
+				AotContext.from(getBeanFactory(), getRepositoryRegistrationAotProcessor().getEnvironment()));
 
 		RepositoryFactoryBeanSupport rfbs = bean.getBeanFactory().getBean("&" + bean.getBeanName(),
 				RepositoryFactoryBeanSupport.class);
@@ -385,33 +381,9 @@ public class RepositoryRegistrationAotContribution implements BeanRegistrationAo
 		return repositoryContext;
 	}
 
+	// TODO: Capture Repository Config
 	private Set<Class<? extends Annotation>> resolveIdentifyingAnnotations() {
-
-		Set<Class<? extends Annotation>> identifyingAnnotations = Collections.emptySet();
-
-		try {
-			// TODO: Getting all beans of type RepositoryConfigurationExtensionSupport will have the effect that
-			// if the user is currently operating in multi-store mode, then all identifying annotations from
-			// all stores will be included in the resulting Set.
-			// When using AOT, is multi-store mode allowed? I don't see why not, but does this work correctly
-			// with AOT, ATM?
-			Map<String, RepositoryConfigurationExtensionSupport> repositoryConfigurationExtensionBeans = getBeanFactory()
-					.getBeansOfType(RepositoryConfigurationExtensionSupport.class);
-
-			// repositoryConfigurationExtensionBeans.values().stream()
-			// .map(RepositoryConfigurationExtensionSupport::getIdentifyingAnnotations)
-			// .flatMap(Collection::stream)
-			// .collect(Collectors.toCollection(() -> identifyingAnnotations));
-
-		} catch (Throwable ignore) {
-			// Possible BeansException because no bean exists of type RepositoryConfigurationExtension,
-			// which included non-Singletons and occurred during eager initialization.
-		}
-
-		return identifyingAnnotations;
+		return Collections.emptySet();
 	}
 
-	private RepositoryInformation resolveRepositoryInformation(RepositoryConfiguration<?> repositoryMetadata) {
-		return RepositoryBeanDefinitionReader.readRepositoryInformation(repositoryMetadata, getBeanFactory());
-	}
 }

src/main/java/org/springframework/data/repository/config/RepositoryRegistrationAotProcessor.java

@@ -16,12 +16,9 @@
 package org.springframework.data.repository.config;
 
 import java.lang.annotation.Annotation;
-import java.util.ArrayList;
 import java.util.Collections;
-import java.util.List;
 import java.util.Map;
 import java.util.function.BiConsumer;
-import java.util.function.BiFunction;
 import java.util.function.Predicate;
 import java.util.stream.Stream;
 
@@ -42,14 +39,16 @@ import org.springframework.beans.factory.aot.BeanRegistrationAotProcessor;
 import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
 import org.springframework.beans.factory.support.RegisteredBean;
 import org.springframework.beans.factory.support.RootBeanDefinition;
+import org.springframework.context.EnvironmentAware;
 import org.springframework.core.annotation.MergedAnnotation;
-import org.springframework.data.aot.AotContext;
+import org.springframework.core.env.Environment;
+import org.springframework.core.env.EnvironmentCapable;
+import org.springframework.core.env.StandardEnvironment;
 import org.springframework.data.repository.aot.generate.RepositoryContributor;
 import org.springframework.data.repository.core.RepositoryInformation;
 import org.springframework.data.repository.core.support.RepositoryFactoryBeanSupport;
 import org.springframework.data.util.TypeContributor;
 import org.springframework.util.Assert;
-import org.springframework.util.StringUtils;
 
 /**
  * {@link BeanRegistrationAotProcessor} responsible processing and providing AOT configuration for repositories.
@@ -72,13 +71,16 @@ import org.springframework.util.StringUtils;
  * @author John Blum
  * @since 3.0
  */
-public class RepositoryRegistrationAotProcessor implements BeanRegistrationAotProcessor, BeanFactoryAware {
-
-	private @Nullable ConfigurableListableBeanFactory beanFactory;
+public class RepositoryRegistrationAotProcessor
+		implements BeanRegistrationAotProcessor, BeanFactoryAware, EnvironmentAware, EnvironmentCapable {
 
 	private final Log logger = LogFactory.getLog(getClass());
 
-	private @Nullable Map<String, RepositoryConfiguration<?>> configMap;
+	private @Nullable ConfigurableListableBeanFactory beanFactory;
+
+	private Environment environment = new StandardEnvironment();
+
+	private Map<String, RepositoryConfiguration<?>> configMap = Collections.emptyMap();
 
 	@Override
 	public @Nullable BeanRegistrationAotContribution processAheadOfTime(RegisteredBean bean) {
@@ -113,10 +115,6 @@ public class RepositoryRegistrationAotProcessor implements BeanRegistrationAotPr
 		ReflectiveRuntimeHintsRegistrar registrar = new ReflectiveRuntimeHintsRegistrar();
 		RuntimeHints hints = generationContext.getRuntimeHints();
 
-		List<Class<?>> aggregateRootTypes = new ArrayList<>();
-		aggregateRootTypes.add(information.getDomainType());
-		aggregateRootTypes.addAll(information.getAlternativeDomainTypes());
-
 		Stream.concat(Stream.of(information.getDomainType()), information.getAlternativeDomainTypes().stream())
 				.forEach(it -> registrar.registerRuntimeHints(hints, it));
 	}
@@ -133,12 +131,9 @@ public class RepositoryRegistrationAotProcessor implements BeanRegistrationAotPr
 
 		//TODO: add the hook for customizing bean initialization code here!
 
-		return contribution.withModuleContribution(new BiFunction<AotRepositoryContext, GenerationContext, RepositoryContributor>() {
-			@Override
-			public RepositoryContributor apply(AotRepositoryContext repositoryContext, GenerationContext generationContext) {
-				registerReflectiveForAggregateRoot(repositoryContext, generationContext);
-				return contribute(repositoryContext, generationContext);
-			}
+		return contribution.withModuleContribution((repositoryContext, generationContext) -> {
+			registerReflectiveForAggregateRoot(repositoryContext, generationContext);
+			return contribute(repositoryContext, generationContext);
 		});
 	}
 
@@ -151,40 +146,54 @@ public class RepositoryRegistrationAotProcessor implements BeanRegistrationAotPr
 		this.beanFactory = (ConfigurableListableBeanFactory) beanFactory;
 	}
 
-	protected ConfigurableListableBeanFactory getBeanFactory() {
-
-		Assert.state(this.beanFactory != null, "BeanFactory must not be null");
-
-		return this.beanFactory;
+	@Override
+	public void setEnvironment(Environment environment) {
+		this.environment = environment;
 	}
 
-	public void setConfigMap(@Nullable Map<String, RepositoryConfiguration<?>> configMap) {
+	@Override
+	public Environment getEnvironment() {
+		return this.environment;
+	}
+
+	public void setConfigMap(Map<String, RepositoryConfiguration<?>> configMap) {
 		this.configMap = configMap;
 	}
 
 	public Map<String, RepositoryConfiguration<?>> getConfigMap() {
-		return nullSafeMap(this.configMap);
+		return this.configMap;
 	}
 
-	private <K, V> Map<K, V> nullSafeMap(@Nullable Map<K, V> map) {
-		return map != null ? map : Collections.emptyMap();
+	protected ConfigurableListableBeanFactory getBeanFactory() {
+
+		if (this.beanFactory == null) {
+			throw new IllegalStateException(
+					"No BeanFactory available. Make sure to set the BeanFactory before using this processor.");
+		}
+
+		return this.beanFactory;
 	}
 
 	protected @Nullable RepositoryConfiguration<?> getRepositoryMetadata(RegisteredBean bean) {
-		return getConfigMap().get(nullSafeBeanName(bean));
+		return getConfigMap().get(bean.getBeanName());
 	}
 
-	private String nullSafeBeanName(RegisteredBean bean) {
-
-		String beanName = bean.getBeanName();
-
-		return StringUtils.hasText(beanName) ? beanName : "";
+	protected void contributeType(Class<?> type, GenerationContext generationContext) {
+		TypeContributor.contribute(type, it -> true, generationContext);
 	}
 
 	protected Log getLogger() {
 		return this.logger;
 	}
 
+	protected void logDebug(String message, Object... arguments) {
+		logAt(Log::isDebugEnabled, Log::debug, message, arguments);
+	}
+
+	protected void logTrace(String message, Object... arguments) {
+		logAt(Log::isTraceEnabled, Log::trace, message, arguments);
+	}
+
 	private void logAt(Predicate<Log> logLevelPredicate, BiConsumer<Log, String> logOperation, String message,
 			Object... arguments) {
 
@@ -195,25 +204,14 @@ public class RepositoryRegistrationAotProcessor implements BeanRegistrationAotPr
 		}
 	}
 
-	protected void logDebug(String message, Object... arguments) {
-		logAt(Log::isDebugEnabled, Log::debug, message, arguments);
+	private static boolean isSpringDataManagedAnnotation(MergedAnnotation<?> annotation) {
+
+		return isSpringDataType(annotation.getType())
+				|| annotation.getMetaTypes().stream().anyMatch(RepositoryRegistrationAotProcessor::isSpringDataType);
 	}
 
-	protected void logTrace(String message, Object... arguments) {
-		logAt(Log::isTraceEnabled, Log::trace, message, arguments);
+	private static boolean isSpringDataType(Class<?> type) {
+		return type.getPackageName().startsWith(TypeContributor.DATA_NAMESPACE);
 	}
 
-	private static boolean isSpringDataManagedAnnotation(@Nullable MergedAnnotation<?> annotation) {
-
-		return annotation != null && (isInSpringDataNamespace(annotation.getType())
-				|| annotation.getMetaTypes().stream().anyMatch(RepositoryRegistrationAotProcessor::isInSpringDataNamespace));
-	}
-
-	protected void contributeType(Class<?> type, GenerationContext generationContext) {
-		TypeContributor.contribute(type, it -> true, generationContext);
-	}
-
-	private static boolean isInSpringDataNamespace(Class<?> type) {
-		return type.getPackage().getName().startsWith(TypeContributor.DATA_NAMESPACE);
-	}
 }

src/test/java/org/springframework/data/repository/aot/generate/DummyModuleAotRepositoryContext.java

@@ -22,6 +22,7 @@ import java.util.Set;
 
 import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
 import org.springframework.core.annotation.MergedAnnotation;
+import org.springframework.core.env.Environment;
 import org.springframework.core.io.ClassPathResource;
 import org.springframework.core.test.tools.ClassFile;
 import org.springframework.data.repository.config.AotRepositoryContext;
@@ -47,6 +48,11 @@ class DummyModuleAotRepositoryContext implements AotRepositoryContext {
 		return null;
 	}
 
+	@Override
+	public Environment getEnvironment() {
+		return null;
+	}
+
 	@Override
 	public TypeIntrospector introspectType(String typeName) {
 		return null;