Refine Repository Composition retrieval during AOT processing

pom.xml

@@ -5,7 +5,7 @@
 
 	<groupId>org.springframework.data</groupId>
 	<artifactId>spring-data-commons</artifactId>
-	<version>4.0.0-SNAPSHOT</version>
+	<version>4.0.x-GH-3279-SNAPSHOT</version>
 
 	<name>Spring Data Core</name>
 	<description>Core Spring concepts underpinning every Spring Data module.</description>

src/main/java/org/springframework/data/mapping/Parameter.java

@@ -118,6 +118,7 @@ public boolean hasName() {
 	 * @since 3.5
 	 * @see org.springframework.core.ParameterNameDiscoverer
 	 */
+	@SuppressWarnings("NullAway")
 	public String getRequiredName() {
 
 		if (!hasName()) {

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

@@ -22,14 +22,13 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.function.BiFunction;
-import java.util.function.Consumer;
 
 import javax.lang.model.element.Modifier;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.jspecify.annotations.Nullable;
+
 import org.springframework.aot.generate.ClassNameGenerator;
 import org.springframework.aot.generate.Generated;
 import org.springframework.data.projection.ProjectionFactory;
@@ -41,6 +40,7 @@
 import org.springframework.javapoet.ClassName;
 import org.springframework.javapoet.FieldSpec;
 import org.springframework.javapoet.JavaFile;
+import org.springframework.javapoet.MethodSpec;
 import org.springframework.javapoet.TypeName;
 import org.springframework.javapoet.TypeSpec;
 
@@ -53,16 +53,19 @@
 class AotRepositoryBuilder {
 
 	private final RepositoryInformation repositoryInformation;
+	private final String moduleName;
 	private final ProjectionFactory projectionFactory;
 	private final AotRepositoryFragmentMetadata generationMetadata;
 
-	private @Nullable Consumer<AotRepositoryConstructorBuilder> constructorCustomizer;
-	private @Nullable BiFunction<Method, RepositoryInformation, @Nullable MethodContributor<? extends QueryMethod>> methodContributorFunction;
+	private @Nullable ConstructorCustomizer constructorCustomizer;
+	private @Nullable MethodContributorFactory methodContributorFactory;
 	private ClassCustomizer customizer;
 
-	private AotRepositoryBuilder(RepositoryInformation repositoryInformation, ProjectionFactory projectionFactory) {
+	private AotRepositoryBuilder(RepositoryInformation repositoryInformation, String moduleName,
+			ProjectionFactory projectionFactory) {
 
 		this.repositoryInformation = repositoryInformation;
+		this.moduleName = moduleName;
 		this.projectionFactory = projectionFactory;
 
 		this.generationMetadata = new AotRepositoryFragmentMetadata(className());
@@ -74,54 +77,69 @@ private AotRepositoryBuilder(RepositoryInformation repositoryInformation, Projec
 		this.customizer = (info, metadata, builder) -> {};
 	}
 
-	public static <M extends QueryMethod> AotRepositoryBuilder forRepository(RepositoryInformation repositoryInformation,
+	/**
+	 * Create a new {@code AotRepositoryBuilder} for the given {@link RepositoryInformation}.
+	 *
+	 * @param information must not be {@literal null}.
+	 * @param moduleName must not be {@literal null}.
+	 * @param projectionFactory must not be {@literal null}.
+	 * @return
+	 */
+	public static AotRepositoryBuilder forRepository(RepositoryInformation information, String moduleName,
 			ProjectionFactory projectionFactory) {
-		return new AotRepositoryBuilder(repositoryInformation, projectionFactory);
+		return new AotRepositoryBuilder(information, moduleName, projectionFactory);
 	}
 
-	public AotRepositoryBuilder withConstructorCustomizer(
-			Consumer<AotRepositoryConstructorBuilder> constructorCustomizer) {
+	/**
+	 * Configure a {@link ClassCustomizer} customizer.
+	 *
+	 * @param classCustomizer must not be {@literal null}.
+	 * @return {@code this}.
+	 */
+	public AotRepositoryBuilder withClassCustomizer(ClassCustomizer classCustomizer) {
 
-		this.constructorCustomizer = constructorCustomizer;
+		this.customizer = classCustomizer;
 		return this;
 	}
 
-	public AotRepositoryBuilder withQueryMethodContributor(
-			BiFunction<Method, RepositoryInformation, @Nullable MethodContributor<? extends QueryMethod>> methodContributorFunction) {
-		this.methodContributorFunction = methodContributorFunction;
+	/**
+	 * Configure a {@link AotRepositoryConstructorBuilder} customizer.
+	 *
+	 * @param constructorCustomizer must not be {@literal null}.
+	 * @return {@code this}.
+	 */
+	public AotRepositoryBuilder withConstructorCustomizer(ConstructorCustomizer constructorCustomizer) {
+
+		this.constructorCustomizer = constructorCustomizer;
 		return this;
 	}
 
-	public AotRepositoryBuilder withClassCustomizer(ClassCustomizer classCustomizer) {
+	/**
+	 * Configure a {@link MethodContributor} factory.
+	 *
+	 * @param methodContributorFactory must not be {@literal null}.
+	 * @return {@code this}.
+	 */
+	public AotRepositoryBuilder withQueryMethodContributor(MethodContributorFactory methodContributorFactory) {
 
-		this.customizer = classCustomizer;
+		this.methodContributorFactory = methodContributorFactory;
 		return this;
 	}
 
 	public AotBundle build() {
 
+		List<AotRepositoryMethod> methodMetadata = new ArrayList<>();
+		RepositoryComposition repositoryComposition = repositoryInformation.getRepositoryComposition();
+
 		// start creating the type
 		TypeSpec.Builder builder = TypeSpec.classBuilder(this.generationMetadata.getTargetTypeName()) //
 				.addModifiers(Modifier.PUBLIC) //
 				.addAnnotation(Generated.class) //
-				.addJavadoc("AOT generated repository implementation for {@link $T}.\n",
+				.addJavadoc("AOT generated $L repository implementation for {@link $T}.\n", moduleName,
 						repositoryInformation.getRepositoryInterface());
 
 		// create the constructor
-		AotRepositoryConstructorBuilder constructorBuilder = new AotRepositoryConstructorBuilder(repositoryInformation,
-				generationMetadata);
-		if (constructorCustomizer != null) {
-			constructorCustomizer.accept(constructorBuilder);
-		}
-
-		builder.addMethod(constructorBuilder.buildConstructor());
-
-		List<AotRepositoryMethod> methodMetadata = new ArrayList<>();
-		AotRepositoryMetadata.RepositoryType repositoryType = repositoryInformation.isReactiveRepository()
-				? AotRepositoryMetadata.RepositoryType.REACTIVE
-				: AotRepositoryMetadata.RepositoryType.IMPERATIVE;
-
-		RepositoryComposition repositoryComposition = repositoryInformation.getRepositoryComposition();
+		builder.addMethod(buildConstructor());
 
 		Arrays.stream(repositoryInformation.getRepositoryInterface().getMethods())
 				.sorted(Comparator.<Method, String> comparing(it -> {
@@ -136,35 +154,58 @@ public AotBundle build() {
 
 		// finally customize the file itself
 		this.customizer.customize(repositoryInformation, generationMetadata, builder);
+
 		JavaFile javaFile = JavaFile.builder(packageName(), builder.build()).build();
+		AotRepositoryMetadata metadata = getAotRepositoryMetadata(methodMetadata);
+
+		return new AotBundle(javaFile, metadata);
+	}
 
-		// TODO: module identifier
-		AotRepositoryMetadata metadata = new AotRepositoryMetadata(repositoryInformation.getRepositoryInterface().getName(),
-				"", repositoryType, methodMetadata);
+	private MethodSpec buildConstructor() {
 
-		return new AotBundle(javaFile, metadata.toJson());
+		AotRepositoryConstructorBuilder constructorBuilder = new AotRepositoryConstructorBuilder(repositoryInformation,
+				generationMetadata);
+
+		if (constructorCustomizer != null) {
+			constructorCustomizer.customize(constructorBuilder);
+		}
+
+		return constructorBuilder.buildConstructor();
+	}
+
+	private AotRepositoryMetadata getAotRepositoryMetadata(List<AotRepositoryMethod> methodMetadata) {
+
+		AotRepositoryMetadata.RepositoryType repositoryType = repositoryInformation.isReactiveRepository()
+				? AotRepositoryMetadata.RepositoryType.REACTIVE
+				: AotRepositoryMetadata.RepositoryType.IMPERATIVE;
+
+		String jsonModuleName = moduleName != null ? moduleName.replaceAll("Reactive", "").trim() : null;
+
+		return new AotRepositoryMetadata(repositoryInformation.getRepositoryInterface().getName(), jsonModuleName,
+				repositoryType, methodMetadata);
 	}
 
 	private void contributeMethod(Method method, RepositoryComposition repositoryComposition,
 			List<AotRepositoryMethod> methodMetadata, TypeSpec.Builder builder) {
 
-		if (repositoryInformation.isCustomMethod(method) || repositoryInformation.isBaseClassMethod(method)) {
+		if (repositoryInformation.isCustomMethod(method)
+				|| (repositoryInformation.isBaseClassMethod(method) && !repositoryInformation.isQueryMethod(method))) {
 
 			RepositoryFragment<?> fragment = repositoryComposition.findFragment(method);
 
 			if (fragment != null) {
 				methodMetadata.add(getFragmentMetadata(method, fragment));
+				return;
 			}
-			return;
 		}
 
 		if (method.isBridge() || method.isDefault() || java.lang.reflect.Modifier.isStatic(method.getModifiers())) {
 			return;
 		}
 
-		if (repositoryInformation.isQueryMethod(method) && methodContributorFunction != null) {
+		if (repositoryInformation.isQueryMethod(method) && methodContributorFactory != null) {
 
-			MethodContributor<? extends QueryMethod> contributor = methodContributorFunction.apply(method,
+			MethodContributor<? extends QueryMethod> contributor = methodContributorFactory.create(method,
 					repositoryInformation);
 
 			if (contributor != null) {
@@ -186,8 +227,7 @@ private void contributeMethod(Method method, RepositoryComposition repositoryCom
 	private AotRepositoryMethod getFragmentMetadata(Method method, RepositoryFragment<?> fragment) {
 
 		String signature = fragment.getSignatureContributor().getName();
-		String implementation = fragment.getImplementation().map(it -> it.getClass().getName()).orElse(null);
-
+		String implementation = fragment.getImplementationClass().map(Class::getName).orElse(null);
 		AotFragmentTarget fragmentTarget = new AotFragmentTarget(signature, implementation);
 
 		return new AotRepositoryMethod(method.getName(), method.toGenericString(), null, fragmentTarget);
@@ -231,17 +271,51 @@ public ProjectionFactory getProjectionFactory() {
 	public interface ClassCustomizer {
 
 		/**
-		 * Apply customization ot the AOT repository fragment class after it has been defined..
+		 * Apply customization ot the AOT repository fragment class after it has been defined.
 		 *
-		 * @param information
-		 * @param metadata
-		 * @param builder
+		 * @param information repository information.
+		 * @param metadata metadata of the AOT repository fragment.
+		 * @param builder the actual builder.
 		 */
 		void customize(RepositoryInformation information, AotRepositoryFragmentMetadata metadata, TypeSpec.Builder builder);
 
 	}
 
-	record AotBundle(JavaFile javaFile, JSONObject metadata) {
+	/**
+	 * Customizer interface to customize the AOT repository fragment constructor through
+	 * {@link AotRepositoryConstructorBuilder}.
+	 */
+	public interface ConstructorCustomizer {
+
+		/**
+		 * Apply customization ot the AOT repository fragment constructor.
+		 *
+		 * @param constructorBuilder the builder to be customized.
+		 */
+		void customize(AotRepositoryConstructorBuilder constructorBuilder);
+
+	}
+
+	/**
+	 * Factory interface to conditionally create {@link MethodContributor} instances. An implementation may decide whether
+	 * to return a {@link MethodContributor} or {@literal null}, if no method (code or metadata) should be contributed.
+	 */
+	public interface MethodContributorFactory {
+
+		/**
+		 * Apply customization ot the AOT repository fragment constructor.
+		 *
+		 * @param method the method to be contributed.
+		 * @param information repository information.
+		 * @return the {@link MethodContributor} to be used. Can be {@literal null} if the method and method metadata should
+		 *         not be contributed.
+		 */
+		@Nullable
+		MethodContributor<? extends QueryMethod> create(Method method, RepositoryInformation information);
+
+	}
+
+	record AotBundle(JavaFile javaFile, AotRepositoryMetadata metadata) {
 	}
 
 }

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

@@ -36,7 +36,7 @@ public abstract class MethodContributor<M extends QueryMethod> {
 	private final M queryMethod;
 	private final QueryMetadata metadata;
 
-	private MethodContributor(M queryMethod, QueryMetadata metadata) {
+	MethodContributor(M queryMethod, QueryMetadata metadata) {
 		this.queryMethod = queryMethod;
 		this.metadata = metadata;
 	}

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

@@ -32,33 +32,48 @@
 import org.springframework.javapoet.JavaFile;
 import org.springframework.javapoet.TypeName;
 import org.springframework.javapoet.TypeSpec;
-import org.springframework.util.StringUtils;
 
 /**
  * Contributor for AOT repository fragments.
  *
  * @author Christoph Strobl
  * @author Mark Paluch
+ * @since 4.0
  */
 public class RepositoryContributor {
 
 	private static final Log logger = LogFactory.getLog(RepositoryContributor.class);
 
 	private final AotRepositoryBuilder builder;
 
+	/**
+	 * Create a new {@code RepositoryContributor} for the given {@link AotRepositoryContext}.
+	 *
+	 * @param repositoryContext
+	 */
 	public RepositoryContributor(AotRepositoryContext repositoryContext) {
 		this.builder = AotRepositoryBuilder.forRepository(repositoryContext.getRepositoryInformation(),
-				createProjectionFactory());
+				repositoryContext.getModuleName(), createProjectionFactory());
 	}
 
+	/**
+	 * @return a new {@link ProjectionFactory} to be used with the AOT repository builder. The actual instance should be
+	 *         accessed through {@link #getProjectionFactory()}.
+	 */
 	protected ProjectionFactory createProjectionFactory() {
 		return new SpelAwareProxyProjectionFactory();
 	}
 
+	/**
+	 * @return the used {@link ProjectionFactory}.
+	 */
 	protected ProjectionFactory getProjectionFactory() {
 		return builder.getProjectionFactory();
 	}
 
+	/**
+	 * @return the used {@link RepositoryInformation}.
+	 */
 	protected RepositoryInformation getRepositoryInformation() {
 		return builder.getRepositoryInformation();
 	}
@@ -73,13 +88,10 @@ public java.util.Map<String, TypeName> requiredArgs() {
 
 	public void contribute(GenerationContext generationContext) {
 
-		// TODO: do we need - generationContext.withName("spring-data");
-
-		builder.withClassCustomizer(this::customizeClass);
-		builder.withConstructorCustomizer(this::customizeConstructor);
-		builder.withQueryMethodContributor(this::contributeQueryMethod);
-
-		AotRepositoryBuilder.AotBundle aotBundle = builder.build();
+		AotRepositoryBuilder.AotBundle aotBundle = builder.withClassCustomizer(this::customizeClass) //
+				.withConstructorCustomizer(this::customizeConstructor) //
+				.withQueryMethodContributor(this::contributeQueryMethod) //
+				.build();
 
 		Class<?> repositoryInterface = getRepositoryInformation().getRepositoryInterface();
 		String repositoryJsonFileName = getRepositoryJsonFileName(repositoryInterface);
@@ -89,7 +101,7 @@ public void contribute(GenerationContext generationContext) {
 		String repositoryJson;
 
 		try {
-			repositoryJson = aotBundle.metadata().toString(2);
+			repositoryJson = aotBundle.metadata().toJson().toString(2);
 		} catch (JSONException e) {
 			throw new RuntimeException(e);
 		}
@@ -118,11 +130,7 @@ public void contribute(GenerationContext generationContext) {
 	}
 
 	private static String getRepositoryJsonFileName(Class<?> repositoryInterface) {
-
-		String repositoryJsonName = repositoryInterface.getSimpleName() + ".json";
-		String repositoryJsonPath = repositoryInterface.getPackageName().replace('.', '/');
-
-		return StringUtils.hasText(repositoryJsonPath) ? repositoryJsonPath + "/" + repositoryJsonName : repositoryJsonName;
+		return repositoryInterface.getName().replace('.', '/') + ".json";
 	}
 
 	/**