Fix AOT vs reflective behavior failures in AotUserRepositoryTests.

Closes #3951

Author: mp911de

spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/AotQueries.java

@@ -39,15 +39,7 @@ record AotQueries(AotQuery result, AotQuery count) {
 	/**
 	 * Derive a count query from the given query.
 	 */
-	public static AotQueries from(StringAotQuery query, @Nullable String countProjection,
-			QueryEnhancerSelector selector) {
-		return from(query, StringAotQuery::getQuery, countProjection, selector);
-	}
-
-	/**
-	 * Derive a count query from the given query.
-	 */
-	public static <T extends AotQuery> AotQueries from(T query, Function<T, DeclaredQuery> queryMapper,
+	public static <T extends AotQuery> AotQueries withDerivedCountQuery(T query, Function<T, DeclaredQuery> queryMapper,
 			@Nullable String countProjection, QueryEnhancerSelector selector) {
 
 		DeclaredQuery underlyingQuery = queryMapper.apply(query);
@@ -56,8 +48,7 @@ public static <T extends AotQuery> AotQueries from(T query, Function<T, Declared
 		String derivedCountQuery = queryEnhancer
 				.createCountQueryFor(StringUtils.hasText(countProjection) ? countProjection : null);
 
-		DeclaredQuery countQuery = underlyingQuery.rewrite(derivedCountQuery);
-		return new AotQueries(query, StringAotQuery.of(countQuery));
+		return new AotQueries(query, StringAotQuery.of(underlyingQuery.rewrite(derivedCountQuery)));
 	}
 
 	/**

spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/AotQuery.java

@@ -18,6 +18,7 @@
 import java.util.List;
 
 import org.springframework.data.domain.Limit;
+import org.springframework.data.jpa.repository.query.EntityQuery;
 import org.springframework.data.jpa.repository.query.ParameterBinding;
 
 /**
@@ -35,6 +36,10 @@ abstract class AotQuery {
 		this.parameterBindings = parameterBindings;
 	}
 
+	static boolean hasConstructorExpressionOrDefaultProjection(EntityQuery query) {
+		return query.hasConstructorExpression() || query.isDefaultProjection();
+	}
+
 	/**
 	 * @return whether the query is a {@link jakarta.persistence.EntityManager#createNativeQuery native} one.
 	 */
@@ -89,4 +94,10 @@ public boolean hasExpression() {
 		return false;
 	}
 
+	/**
+	 * @return {@literal true} if query is expected to return the declared method type directly; {@literal false} if the
+	 *         result requires projection post-processing. See also {@code NativeJpaQuery#getTypeToQueryFor}.
+	 */
+	public abstract boolean hasConstructorExpressionOrDefaultProjection();
+
 }

spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/JpaCodeBlocks.java

@@ -96,6 +96,7 @@ static class QueryBlockBuilder {
 		private @Nullable Class<?> queryRewriter = QueryRewriter.IdentityQueryRewriter.class;
 
 		private QueryBlockBuilder(AotQueryMethodGenerationContext context, JpaQueryMethod queryMethod) {
+
 			this.context = context;
 			this.queryMethod = queryMethod;
 			this.queryVariableName = context.localVariable("query");
@@ -294,14 +295,16 @@ private CodeBlock applyLimits(boolean exists, String pageable) {
 				return builder.build();
 			}
 
+			if (queries != null && queries.result() instanceof StringAotQuery sq && sq.hasPagingExpression()) {
+				return builder.build();
+			}
+
 			String limit = context.getLimitParameterName();
 
 			if (StringUtils.hasText(limit)) {
 				builder.beginControlFlow("if ($L.isLimited())", limit);
 				builder.addStatement("$L.setMaxResults($L.max())", queryVariableName, limit);
 				builder.endControlFlow();
-			} else if (queries != null && queries.result().isLimited()) {
-				builder.addStatement("$L.setMaxResults($L)", queryVariableName, queries.result().getLimit().max());
 			}
 
 			if (StringUtils.hasText(pageable)) {
@@ -316,6 +319,20 @@ private CodeBlock applyLimits(boolean exists, String pageable) {
 				builder.endControlFlow();
 			}
 
+			if (queries.result().isLimited()) {
+
+				int max = queries.result().getLimit().max();
+
+				builder.beginControlFlow("if ($L.getMaxResults() != $T.MAX_VALUE)", queryVariableName, Integer.class);
+				builder.beginControlFlow("if ($1L.getMaxResults() > $2L && $1L.getFirstResult() > 0)", queryVariableName, max);
+				builder.addStatement("$1L.setFirstResult($1L.getFirstResult() - ($1L.getMaxResults() - $2L))",
+						queryVariableName, max);
+				builder.endControlFlow();
+				builder.endControlFlow();
+
+				builder.addStatement("$1L.setMaxResults($2L)", queryVariableName, max);
+			}
+
 			return builder.build();
 		}
 
@@ -484,11 +501,12 @@ private CodeBlock doCreateQuery(boolean count, String queryVariableName,
 
 			if (query instanceof NamedAotQuery nq) {
 
-				if (!count && returnedType.isProjecting() && returnedType.getReturnedType().isInterface()) {
-					builder.addStatement("$T $L = this.$L.createNamedQuery($S)", Query.class, queryVariableName,
-							context.fieldNameOf(EntityManager.class), nq.getName());
-					return builder.build();
-				} else if (queryReturnType != null) {
+				if (!count && !nq.hasConstructorExpressionOrDefaultProjection() && returnedType.isProjecting()
+						&& returnedType.getReturnedType().isInterface()) {
+					queryReturnType = Tuple.class;
+				}
+
+				if (queryReturnType != null) {
 
 					builder.addStatement("$T $L = this.$L.createNamedQuery($S, $T.class)", Query.class, queryVariableName,
 							context.fieldNameOf(EntityManager.class), nq.getName(), queryReturnType);

spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/JpaRepositoryContributor.java

@@ -183,8 +183,8 @@ private Optional<Class<QueryEnhancerSelector>> getQueryEnhancerSelectorClass() {
 
 		MergedAnnotation<Query> query = MergedAnnotations.from(method).get(Query.class);
 
-		AotQueries aotQueries = queriesFactory.createQueries(getRepositoryInformation(), query, selector, queryMethod,
-				returnedType);
+		AotQueries aotQueries = queriesFactory.createQueries(getRepositoryInformation(), returnedType, selector, query,
+				queryMethod);
 
 		// no KeysetScrolling for now.
 		if (parameters.hasScrollPositionParameter()) {

spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/NamedAotQuery.java

@@ -15,11 +15,8 @@
  */
 package org.springframework.data.jpa.repository.aot;
 
-import java.util.List;
-
 import org.springframework.data.jpa.repository.query.DeclaredQuery;
-import org.springframework.data.jpa.repository.query.ParameterBinding;
-import org.springframework.data.jpa.repository.query.PreprocessedQuery;
+import org.springframework.data.jpa.repository.query.EntityQuery;
 
 /**
  * Value object to describe a named AOT query.
@@ -31,20 +28,20 @@ class NamedAotQuery extends AotQuery {
 
 	private final String name;
 	private final DeclaredQuery query;
+	private final boolean constructorExpressionOrDefaultProjection;
 
-	private NamedAotQuery(String name, DeclaredQuery queryString, List<ParameterBinding> parameterBindings) {
-		super(parameterBindings);
+	public NamedAotQuery(String name, EntityQuery entityQuery) {
+		super(entityQuery.getParameterBindings());
 		this.name = name;
-		this.query = queryString;
+		this.query = entityQuery.getQuery();
+		this.constructorExpressionOrDefaultProjection = AotQuery.hasConstructorExpressionOrDefaultProjection(entityQuery);
 	}
 
 	/**
 	 * Creates a new {@code NamedAotQuery}.
 	 */
-	public static NamedAotQuery named(String namedQuery, DeclaredQuery queryString) {
-
-		PreprocessedQuery parsed = PreprocessedQuery.parse(queryString);
-		return new NamedAotQuery(namedQuery, queryString, parsed.getBindings());
+	public static NamedAotQuery named(String namedQuery, EntityQuery query) {
+		return new NamedAotQuery(namedQuery, query);
 	}
 
 	public String getName() {
@@ -64,4 +61,9 @@ public boolean isNative() {
 		return query.isNative();
 	}
 
+	@Override
+	public boolean hasConstructorExpressionOrDefaultProjection() {
+		return constructorExpressionOrDefaultProjection;
+	}
+
 }

spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/QueriesFactory.java

@@ -112,28 +112,28 @@ private NamedQueries getNamedQueries(@Nullable RepositoryConfigurationSource con
 	 * Creates the {@link AotQueries} used within a specific {@link JpaQueryMethod}.
 	 *
 	 * @param repositoryInformation
-	 * @param query
+	 * @param returnedType
 	 * @param selector
+	 * @param query
 	 * @param queryMethod
-	 * @param returnedType
 	 * @return
 	 */
-	public AotQueries createQueries(RepositoryInformation repositoryInformation, MergedAnnotation<Query> query,
-			QueryEnhancerSelector selector, JpaQueryMethod queryMethod, ReturnedType returnedType) {
+	public AotQueries createQueries(RepositoryInformation repositoryInformation, ReturnedType returnedType,
+			QueryEnhancerSelector selector, MergedAnnotation<Query> query, JpaQueryMethod queryMethod) {
 
 		if (query.isPresent() && StringUtils.hasText(query.getString("value"))) {
 			return buildStringQuery(repositoryInformation.getDomainType(), returnedType, selector, query, queryMethod);
 		}
 
 		String queryName = queryMethod.getNamedQueryName();
-		if (hasNamedQuery(queryName, returnedType)) {
+		if (hasNamedQuery(returnedType, queryName)) {
 			return buildNamedQuery(returnedType, selector, queryName, query, queryMethod);
 		}
 
-		return buildPartTreeQuery(returnedType, repositoryInformation, query, queryMethod);
+		return buildPartTreeQuery(repositoryInformation, returnedType, selector, query, queryMethod);
 	}
 
-	private boolean hasNamedQuery(String queryName, ReturnedType returnedType) {
+	private boolean hasNamedQuery(ReturnedType returnedType, String queryName) {
 		return namedQueries.hasQuery(queryName) || getNamedQuery(returnedType, queryName) != null;
 	}
 
@@ -142,19 +142,15 @@ private AotQueries buildStringQuery(Class<?> domainType, ReturnedType returnedTy
 
 		UnaryOperator<String> operator = s -> s.replaceAll("#\\{#entityName}", domainType.getName());
 		boolean isNative = query.getBoolean("nativeQuery");
-		Function<String, StringAotQuery> queryFunction = isNative ? StringAotQuery::nativeQuery : StringAotQuery::jpqlQuery;
+		Function<String, DeclaredQuery> queryFunction = isNative ? DeclaredQuery::nativeQuery : DeclaredQuery::jpqlQuery;
 		queryFunction = operator.andThen(queryFunction);
 
 		String queryString = query.getString("value");
 
-		StringAotQuery aotStringQuery = queryFunction.apply(queryString);
+		EntityQuery entityQuery = EntityQuery.create(queryFunction.apply(queryString), selector);
+		StringAotQuery aotStringQuery = StringAotQuery.of(entityQuery);
 		String countQuery = query.getString("countQuery");
 
-		EntityQuery entityQuery = EntityQuery.create(aotStringQuery.getQuery(), selector);
-		if (entityQuery.hasConstructorExpression() || entityQuery.isDefaultProjection()) {
-			aotStringQuery = aotStringQuery.withConstructorExpressionOrDefaultProjection();
-		}
-
 		if (returnedType.isProjecting() && returnedType.hasInputProperties()
 				&& !returnedType.getReturnedType().isInterface()) {
 
@@ -174,38 +170,38 @@ public ReturnedType getReturnedType() {
 		}
 
 		if (StringUtils.hasText(countQuery)) {
-			return AotQueries.from(aotStringQuery, queryFunction.apply(countQuery));
+			return AotQueries.from(aotStringQuery, StringAotQuery.of(queryFunction.apply(countQuery)));
 		}
 
-		if (hasNamedQuery(queryMethod.getNamedCountQueryName(), returnedType)) {
+		if (hasNamedQuery(returnedType, queryMethod.getNamedCountQueryName())) {
 			return AotQueries.from(aotStringQuery,
-					createNamedAotQuery(returnedType, queryMethod.getNamedCountQueryName(), queryMethod, isNative));
+					createNamedAotQuery(returnedType, selector, queryMethod.getNamedCountQueryName(), queryMethod, isNative));
 		}
 
 		String countProjection = query.getString("countProjection");
-		return AotQueries.from(aotStringQuery, countProjection, selector);
+		return AotQueries.withDerivedCountQuery(aotStringQuery, StringAotQuery::getQuery, countProjection, selector);
 	}
 
-	private AotQueries buildNamedQuery(ReturnedType returnedType, QueryEnhancerSelector selector,
-			String queryName, MergedAnnotation<Query> query, JpaQueryMethod queryMethod) {
+	private AotQueries buildNamedQuery(ReturnedType returnedType, QueryEnhancerSelector selector, String queryName,
+			MergedAnnotation<Query> query, JpaQueryMethod queryMethod) {
 
 		boolean nativeQuery = query.isPresent() && query.getBoolean("nativeQuery");
-		AotQuery aotQuery = createNamedAotQuery(returnedType, queryName, queryMethod, nativeQuery);
-
+		AotQuery aotQuery = createNamedAotQuery(returnedType, selector, queryName, queryMethod, nativeQuery);
 		String countQuery = query.isPresent() ? query.getString("countQuery") : null;
 
 		if (StringUtils.hasText(countQuery)) {
 			return AotQueries.from(aotQuery,
-					aotQuery.isNative() ? StringAotQuery.nativeQuery(countQuery) : StringAotQuery.jpqlQuery(countQuery));
+					StringAotQuery
+							.of(aotQuery.isNative() ? DeclaredQuery.nativeQuery(countQuery) : DeclaredQuery.jpqlQuery(countQuery)));
 		}
 
-		if (hasNamedQuery(queryMethod.getNamedCountQueryName(), returnedType)) {
+		if (hasNamedQuery(returnedType, queryMethod.getNamedCountQueryName())) {
 			return AotQueries.from(aotQuery,
-					createNamedAotQuery(returnedType, queryMethod.getNamedCountQueryName(), queryMethod, nativeQuery));
+					createNamedAotQuery(returnedType, selector, queryMethod.getNamedCountQueryName(), queryMethod, nativeQuery));
 		}
 
 		String countProjection = query.isPresent() ? query.getString("countProjection") : null;
-		return AotQueries.from(aotQuery, it -> {
+		return AotQueries.withDerivedCountQuery(aotQuery, it -> {
 
 			if (it instanceof StringAotQuery sq) {
 				return sq.getQuery();
@@ -215,25 +211,26 @@ private AotQueries buildNamedQuery(ReturnedType returnedType, QueryEnhancerSelec
 		}, countProjection, selector);
 	}
 
-	private AotQuery createNamedAotQuery(ReturnedType returnedType, String queryName, JpaQueryMethod queryMethod,
-			boolean isNative) {
+	private AotQuery createNamedAotQuery(ReturnedType returnedType, QueryEnhancerSelector selector, String queryName,
+			JpaQueryMethod queryMethod, boolean isNative) {
 
 		if (namedQueries.hasQuery(queryName)) {
 
 			String queryString = namedQueries.getQuery(queryName);
-			return StringAotQuery.named(queryName,
-					isNative ? DeclaredQuery.nativeQuery(queryString) : DeclaredQuery.jpqlQuery(queryString));
+
+			DeclaredQuery query = isNative ? DeclaredQuery.nativeQuery(queryString) : DeclaredQuery.jpqlQuery(queryString);
+			return StringAotQuery.named(queryName, EntityQuery.create(query, selector));
 		}
 
 		TypedQueryReference<?> namedQuery = getNamedQuery(returnedType, queryName);
 
 		Assert.state(namedQuery != null, "Native named query must not be null");
 
-		return createNamedAotQuery(namedQuery, queryMethod, isNative);
+		return createNamedAotQuery(namedQuery, selector, isNative, queryMethod);
 	}
 
-	private AotQuery createNamedAotQuery(TypedQueryReference<?> namedQuery, JpaQueryMethod queryMethod,
-			boolean isNative) {
+	private AotQuery createNamedAotQuery(TypedQueryReference<?> namedQuery, QueryEnhancerSelector selector,
+			boolean isNative, JpaQueryMethod queryMethod) {
 
 		QueryExtractor queryExtractor = queryMethod.getQueryExtractor();
 		String queryString = queryExtractor.extractQueryString(namedQuery);
@@ -244,8 +241,9 @@ private AotQuery createNamedAotQuery(TypedQueryReference<?> namedQuery, JpaQuery
 
 		Assert.hasText(queryString, () -> "Cannot extract Query from named query [%s]".formatted(namedQuery.getName()));
 
-		return NamedAotQuery.named(namedQuery.getName(),
-				isNative ? DeclaredQuery.nativeQuery(queryString) : DeclaredQuery.jpqlQuery(queryString));
+		DeclaredQuery query = isNative ? DeclaredQuery.nativeQuery(queryString) : DeclaredQuery.jpqlQuery(queryString);
+
+		return NamedAotQuery.named(namedQuery.getName(), EntityQuery.create(query, selector));
 	}
 
 	private @Nullable TypedQueryReference<?> getNamedQuery(ReturnedType returnedType, String queryName) {
@@ -266,19 +264,20 @@ private AotQuery createNamedAotQuery(TypedQueryReference<?> namedQuery, JpaQuery
 		return null;
 	}
 
-	private AotQueries buildPartTreeQuery(ReturnedType returnedType, RepositoryInformation repositoryInformation,
+	private AotQueries buildPartTreeQuery(RepositoryInformation repositoryInformation, ReturnedType returnedType,
+			QueryEnhancerSelector selector,
 			MergedAnnotation<Query> query, JpaQueryMethod queryMethod) {
 
 		PartTree partTree = new PartTree(queryMethod.getName(), repositoryInformation.getDomainType());
 		AotQuery aotQuery = createQuery(partTree, returnedType, queryMethod.getParameters(), templates);
 
 		if (query.isPresent() && StringUtils.hasText(query.getString("countQuery"))) {
-			return AotQueries.from(aotQuery, StringAotQuery.jpqlQuery(query.getString("countQuery")));
+			return AotQueries.from(aotQuery, StringAotQuery.of(DeclaredQuery.jpqlQuery(query.getString("countQuery"))));
 		}
 
-		if (hasNamedQuery(queryMethod.getNamedCountQueryName(), returnedType)) {
+		if (hasNamedQuery(returnedType, queryMethod.getNamedCountQueryName())) {
 			return AotQueries.from(aotQuery,
-					createNamedAotQuery(returnedType, queryMethod.getNamedCountQueryName(), queryMethod, false));
+					createNamedAotQuery(returnedType, selector, queryMethod.getNamedCountQueryName(), queryMethod, false));
 		}
 
 		AotQuery partTreeCountQuery = createCountQuery(partTree, returnedType, queryMethod.getParameters(), templates);
@@ -318,19 +317,21 @@ private AotQuery createCountQuery(PartTree partTree, ReturnedType returnedType,
 
 		Class<?> result = queryForEntity ? returnedType.getDomainType() : null;
 
-		if (query instanceof StringAotQuery sq && sq.hasConstructorExpressionOrDefaultProjection()) {
-			return result;
-		}
-
 		if (returnedType.isProjecting()) {
 
 			if (returnedType.getReturnedType().isInterface()) {
+
+				if (query.hasConstructorExpressionOrDefaultProjection()) {
+					return result;
+				}
+
 				return Tuple.class;
 			}
 
 			return returnedType.getReturnedType();
 		}
 
+
 		return result;
 	}