Consider actual query method return type in AOT-generated repository implementations.

We now consider the actual query method return type in AOT generated fragments.
Ensure AOT-generated code uses correct entity class in query methods returning subclass types.

Closes #5178
Original pull request: #5179

Signed-off-by: Tomasz Forys <tomito14@interia.pl>

Author: infoShare

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

@@ -30,6 +30,7 @@
 import org.springframework.data.mongodb.core.query.DiskUse;
 import org.springframework.data.mongodb.repository.Hint;
 import org.springframework.data.mongodb.repository.Meta;
+import org.springframework.data.mongodb.repository.query.MongoEntityMetadata;
 import org.springframework.data.mongodb.repository.query.MongoQueryExecution.PagedExecution;
 import org.springframework.data.mongodb.repository.query.MongoQueryExecution.SlicedExecution;
 import org.springframework.data.mongodb.repository.query.MongoQueryMethod;
@@ -44,6 +45,7 @@
 
 /**
  * @author Christoph Strobl
+ * @author Tomasz Forys
  * @since 5.0
  */
 class QueryBlocks {
@@ -71,27 +73,29 @@ CodeBlock build() {
 
 			MethodReturn methodReturn = context.getMethodReturn();
 			String mongoOpsRef = context.fieldNameOf(MongoOperations.class);
+			MongoEntityMetadata<?> entityMetadata = queryMethod.getEntityInformation();
 
 			Builder builder = CodeBlock.builder();
 
 			boolean isProjecting = context.getReturnedType().isProjecting();
-			Class<?> domainType = context.getRepositoryInformation().getDomainType();
-			Object actualReturnType = queryMethod.getParameters().hasDynamicProjection() || isProjecting
+			boolean hasDynamicProjection = queryMethod.getParameters().hasDynamicProjection();
+			Class<?> queryType = entityMetadata.getCollectionEntity().getType();
+			Class<?> entityType = entityMetadata.getJavaType();
+			Object actualReturnType = hasDynamicProjection || isProjecting
 					? methodReturn.getActualTypeName()
-					: domainType;
+					: entityType;
 
 			builder.add("\n");
 
-			if (queryMethod.getParameters().hasDynamicProjection()) {
+			if (hasDynamicProjection) {
 				builder.addStatement("$T<$T> $L = $L.query($T.class).as($L)", FindWithQuery.class, actualReturnType,
-						context.localVariable("finder"), mongoOpsRef, domainType, context.getDynamicProjectionParameterName());
-			} else if (isProjecting) {
+						context.localVariable("finder"), mongoOpsRef, queryType, context.getDynamicProjectionParameterName());
+			} else if (isProjecting || !queryType.equals(entityType)) {
 				builder.addStatement("$T<$T> $L = $L.query($T.class).as($T.class)", FindWithQuery.class, actualReturnType,
-						context.localVariable("finder"), mongoOpsRef, domainType, actualReturnType);
+						context.localVariable("finder"), mongoOpsRef, queryType, actualReturnType);
 			} else {
-
 				builder.addStatement("$T<$T> $L = $L.query($T.class)", FindWithQuery.class, actualReturnType,
-						context.localVariable("finder"), mongoOpsRef, domainType);
+						context.localVariable("finder"), mongoOpsRef, queryType);
 			}
 
 			String terminatingMethod;

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

@@ -24,6 +24,7 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.List;
+import java.util.Optional;
 import java.util.regex.Pattern;
 
 import org.junit.jupiter.api.Assertions;
@@ -47,6 +48,7 @@
 import org.springframework.data.mongodb.core.geo.GeoJsonPolygon;
 import org.springframework.data.mongodb.core.geo.Sphere;
 import org.springframework.data.mongodb.repository.Hint;
+import org.springframework.data.mongodb.repository.Query;
 import org.springframework.data.mongodb.repository.ReadPreference;
 import org.springframework.data.mongodb.repository.VectorSearch;
 import org.springframework.data.repository.Repository;
@@ -61,6 +63,7 @@
  *
  * @author Christoph Strobl
  * @author Mark Paluch
+ * @author Tomasz Forys
  */
 class QueryMethodContributionUnitTests {
 
@@ -406,6 +409,30 @@ void rendersStreamableReturnTypeForAggregation() throws NoSuchMethodException {
 				.containsSubsequence("return", "Streamable.of(", "getMappedResults())");
 	}
 
+	@Test // GH-5178
+	void rendersAnnotatedQueryUsingSubtypeAsCollectionAndResultType() throws NoSuchMethodException {
+
+		MethodSpec methodSpec = codeOf(UserRepoWithMeta.class, "findSpecialUserByLastname", String.class);
+
+		assertThat(methodSpec.toString()) //
+				.containsSubsequence("FindWithQuery<", ".SpecialUser> finder =")
+				.containsSubsequence(".query(",".SpecialUser.class)")
+				.doesNotContain(".query(example.aot.User.class)")
+				.contains("return finder.matching(filterQuery).all()");
+	}
+
+	@Test // GH-5178
+	void rendersOptionalAnnotatedQueryUsingSubtypeAsCollectionAndResultType() throws NoSuchMethodException {
+
+		MethodSpec methodSpec = codeOf(UserRepoWithMeta.class, "findOptionalSubtypeByUsername", String.class);
+
+		assertThat(methodSpec.toString()) //
+				.containsSubsequence("FindWithQuery<", ".SpecialUser> finder =")
+				.containsSubsequence(".query(",".SpecialUser.class)")
+				.doesNotContain(".query(example.aot.User.class)")
+				.contains("return finder.matching(filterQuery).one()");
+	}
+
 	private static MethodSpec codeOf(Class<?> repository, String methodName, Class<?>... args)
 			throws NoSuchMethodException {
 
@@ -445,6 +472,8 @@ static class TestQueryMethodGenerationContext extends AotQueryMethodGenerationCo
 		}
 	}
 
+	static class SpecialUser extends User {}
+
 	interface UserRepoWithMeta extends Repository<User, String> {
 
 		@Hint("fn-idx")
@@ -460,5 +489,11 @@ interface UserRepoWithMeta extends Repository<User, String> {
 		@VectorSearch(indexName = "embedding.vector_cos", limit = "#{5+5}")
 		SearchResults<User> searchWithLimitAsExpressionByLastnameAndEmbeddingWithinOrderByFirstname(String lastname,
 				Vector vector, Range<Similarity> distance);
+
+		@Query("{ 'lastname' : { '$regex' : '^?0' } }")
+		List<SpecialUser> findSpecialUserByLastname(String lastname);
+
+		@Query("{ 'username' : ?0 }")
+		Optional<SpecialUser> findOptionalSubtypeByUsername(String username);
 	}
 }