Use explicitly configured MongoOperations for AOT fragment bootstrap.

This commit makes sure to use a configured MongoOperations reference to bootstrap AOT generated repository implementations. Previously mongoTemplateRef set via EnableMongoRepositories had not been taken into account.

Author: christophstrobl

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/repository/aot/MongoRepositoryContributor.java

@@ -25,7 +25,7 @@
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.jspecify.annotations.Nullable;
-
+import org.springframework.beans.factory.config.RuntimeBeanReference;
 import org.springframework.core.annotation.AnnotatedElementUtils;
 import org.springframework.core.io.support.PathMatchingResourcePatternResolver;
 import org.springframework.data.mapping.model.SimpleTypeHolder;
@@ -72,6 +72,7 @@ public class MongoRepositoryContributor extends RepositoryContributor {
 	private final SimpleTypeHolder simpleTypeHolder;
 	private final MongoMappingContext mappingContext;
 	private final NamedQueries namedQueries;
+	private final @Nullable String mongoOperationsRef;
 
 	public MongoRepositoryContributor(AotRepositoryContext repositoryContext) {
 
@@ -81,7 +82,9 @@ public MongoRepositoryContributor(AotRepositoryContext repositoryContext) {
 		if (classLoader == null) {
 			classLoader = getClass().getClassLoader();
 		}
-		namedQueries = getNamedQueries(repositoryContext.getConfigurationSource(), classLoader);
+
+		this.namedQueries = getNamedQueries(repositoryContext.getConfigurationSource(), classLoader);
+		this.mongoOperationsRef = getMongoTemplateRef(repositoryContext.getConfigurationSource());
 
 		// avoid Java Time (JSR-310) Type introspection
 		MongoCustomConversions mongoCustomConversions = MongoCustomConversions
@@ -97,6 +100,14 @@ public MongoRepositoryContributor(AotRepositoryContext repositoryContext) {
 		this.queryCreator = new AotQueryCreator(this.mappingContext);
 	}
 
+	private @Nullable String getMongoTemplateRef(@Nullable RepositoryConfigurationSource configSource) {
+		if (configSource == null) {
+			return null;
+		}
+
+		return configSource.getAttribute("mongoTemplateRef").filter(it -> !"mongoTemplate".equals(it)).orElse(null);
+	}
+
 	@SuppressWarnings("NullAway")
 	private NamedQueries getNamedQueries(@Nullable RepositoryConfigurationSource configSource, ClassLoader classLoader) {
 
@@ -132,9 +143,15 @@ protected void customizeClass(AotRepositoryClassBuilder classBuilder) {
 	@Override
 	protected void customizeConstructor(AotRepositoryConstructorBuilder constructorBuilder) {
 
-		constructorBuilder.addParameter("operations", MongoOperations.class);
-		constructorBuilder.addParameter("context", RepositoryFactoryBeanSupport.FragmentCreationContext.class,
-				false);
+		constructorBuilder.addParameter("operations", MongoOperations.class, customizer -> {
+
+			customizer.bindToField()
+					.origin(StringUtils.hasText(this.mongoOperationsRef)
+							? new RuntimeBeanReference(this.mongoOperationsRef, MongoOperations.class)
+							: new RuntimeBeanReference(MongoOperations.class));
+		});
+
+		constructorBuilder.addParameter("context", RepositoryFactoryBeanSupport.FragmentCreationContext.class, false);
 
 		constructorBuilder.customize((builder) -> {
 			builder.addStatement("super(operations, context)");

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/repository/aot/MongoRepositoryContributorConfigurationTests.java

@@ -0,0 +1,88 @@
+/*
+ * Copyright 2025-present 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
+ *
+ *      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,
+ * 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.
+ */
+package org.springframework.data.mongodb.repository.aot;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+
+import org.junit.jupiter.api.Test;
+import org.springframework.aot.generate.GeneratedFiles.Kind;
+import org.springframework.aot.test.generate.TestGenerationContext;
+import org.springframework.context.annotation.AnnotationConfigApplicationContext;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.context.aot.ApplicationContextAotGenerator;
+import org.springframework.core.io.InputStreamResource;
+import org.springframework.core.io.InputStreamSource;
+import org.springframework.data.mongodb.repository.config.EnableMongoRepositories;
+
+/**
+ * @author Christoph Strobl
+ */
+class MongoRepositoryContributorConfigurationTests {
+
+	@Configuration
+	@EnableMongoRepositories(basePackages = "example.aot")
+	static class ConfigurationWithoutAnythingSpecial {
+
+	}
+
+	@Configuration
+	@EnableMongoRepositories(mongoTemplateRef = "template-2", basePackages = "example.aot")
+	static class ConfigurationWithTemplateRef {
+
+	}
+
+	@Test // GH-5107
+	void usesPrimaryMongoOperationsBeanReferenceByDefault() throws IOException {
+
+		TestGenerationContext testContext = generate(ConfigurationWithoutAnythingSpecial.class);
+		InputStreamSource file = testContext.getGeneratedFiles().getGeneratedFile(Kind.SOURCE,
+				"example/aot/UserRepository__BeanDefinitions.java");
+
+		InputStreamResource isr = new InputStreamResource(file);
+		String sourceCode = isr.getContentAsString(StandardCharsets.UTF_8);
+
+		assertThat(sourceCode).contains("operations = beanFactory.getBean(MongoOperations.class)");
+	}
+
+	@Test // GH-5107
+	void shouldConsiderMongoTemplateReferenceIfPresent() throws IOException {
+
+		TestGenerationContext testContext = generate(ConfigurationWithTemplateRef.class);
+		InputStreamSource file = testContext.getGeneratedFiles().getGeneratedFile(Kind.SOURCE,
+				"example/aot/UserRepository__BeanDefinitions.java");
+
+		InputStreamResource isr = new InputStreamResource(file);
+		String sourceCode = isr.getContentAsString(StandardCharsets.UTF_8);
+
+		assertThat(sourceCode).contains("operations = beanFactory.getBean(\"template-2\", MongoOperations.class)");
+	}
+
+	private static TestGenerationContext generate(Class<?>... configurationClasses) {
+
+		AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext();
+		context.register(configurationClasses);
+
+		ApplicationContextAotGenerator generator = new ApplicationContextAotGenerator();
+
+		TestGenerationContext generationContext = new TestGenerationContext();
+		generator.processAheadOfTime(context, generationContext);
+		generationContext.writeGeneratedContent();
+		return generationContext;
+	}
+}