Polishing.

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

Refine exposed API.

See #3265

Author: mp911de

Diff for: src/main/java/org/springframework/data/aot/AotContext.java

@@ -30,7 +30,9 @@
 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,29 +47,42 @@
  * @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}.
+	 *
+	 * @param beanFactory reference to the {@link BeanFactory}; must not be {@literal null}.
+	 * @return a new instance of {@link AotContext}.
+	 * @see BeanFactory
+	 */
+	static AotContext from(BeanFactory beanFactory) {
+
+		Assert.notNull(beanFactory, "BeanFactory must not be null");
+
+		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) {
+	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);
+		return new DefaultAotContext(beanFactory, environment);
 	}
 
 	/**

Diff for: src/main/java/org/springframework/data/aot/DefaultAotContext.java

@@ -28,6 +28,7 @@
 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,16 +41,24 @@ 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
 	public ConfigurableListableBeanFactory getBeanFactory() {
 		return factory;
 	}
 
+	@Override
+	public Environment getEnvironment() {
+		return environment;
+	}
+
 	@Override
 	public TypeIntrospector introspectType(String typeName) {
 		return new DefaultTypeIntrospector(typeName);
@@ -138,4 +147,5 @@ public RootBeanDefinition getRootBeanDefinition() throws NoSuchBeanDefinitionExc
 			return factory.getType(beanName, false);
 		}
 	}
+
 }

Diff for: src/main/java/org/springframework/data/aot/ManagedTypesBeanRegistrationAotProcessor.java

@@ -31,8 +31,12 @@
 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 @@
  * @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 String getModuleIdentifier() {
 		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 String getModuleIdentifier() {
 		}
 
 		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) {

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

@@ -17,10 +17,9 @@
 
 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 @@ public JSONObject toJson() throws JSONException {
 
 		return fragment;
 	}
+
 }

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

@@ -1,27 +1,11 @@
-/*
- * 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.
  *
  * 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
+ *      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,15 +63,15 @@ 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);
 	}
 
 	AotRepositoryFragmentMetadata getTargetTypeMetadata() {
 		return targetTypeMetadata;
 	}
 
-	AotRepositoryMethodImplementationMetadata getTargetMethodMetadata() {
+	MethodMetadata getTargetMethodMetadata() {
 		return targetMethodMetadata;
 	}
 

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

@@ -34,8 +34,6 @@
 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 @@ public AotRepositoryBuilder withConstructorCustomizer(
 	}
 
 	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 @@ public AotBundle build() {
 		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 @@ private String typeName() {
 	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;
 	}

Diff for: 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 @@
  * @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 @@ MethodSpec buildConstructor() {
 		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 @@ MethodSpec buildConstructor() {
 		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.
 	 */

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

@@ -31,6 +31,7 @@
 /**
  * @author Christoph Strobl
  */
+// TODO: Can we make this package-private?
 public class AotRepositoryFragmentMetadata {
 
 	private final ClassName className;
@@ -94,31 +95,11 @@ public void addConstructorArgument(String parameterName, TypeName type, @Nullabl
 		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;
-		}
 	}
 }

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

@@ -17,11 +17,9 @@
 
 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 @@ enum RepositoryType {
 		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 @@ JSONObject toJson() throws JSONException {
 		metadata.put("methods", methods);
 
 		return metadata;
-
 	}
+
 }