Refactoring.

Introduce AotRepositoryFragmentSupport, adopt to FragmentCreationContext.
Reduce visibility. Refactor CodeBlocks builder. Simplify query rewriting and use base class methods.
Use typed verifier through a JDK proxy to avoid reflective frontend.
Revise testing to a plain old Spring test but testing the AOT fragment through its interface by forwarding reflective calls to the AOT fragment.
Refactor AotQuery into AotQueries to support a wider range of possible queries.

See #3830

Author: mp911de

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

@@ -37,7 +37,7 @@
 /**
  * @author Christoph Strobl
  */
-public class AotMetaModel implements Metamodel {
+class AotMetaModel implements Metamodel {
 
 	private final String persistenceUnit;
 	private final Set<Class<?>> managedTypes;
@@ -105,7 +105,7 @@ public ClassLoader getNewTempClassLoader() {
 
 			@Override
 			public void addTransformer(ClassTransformer classTransformer) {
-				// just ingnore it
+				// just ignore it
 			}
 		};
 

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

@@ -0,0 +1,58 @@
+/*
+ * 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
+ *
+ *      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.jpa.repository.aot.generated;
+
+import jakarta.validation.constraints.Null;
+
+import org.springframework.data.jpa.repository.query.DeclaredQuery;
+import org.springframework.data.jpa.repository.query.QueryEnhancer;
+import org.springframework.data.jpa.repository.query.QueryEnhancerSelector;
+import org.springframework.util.StringUtils;
+
+/**
+ * Value object capturing queries used for repository query methods.
+ *
+ * @author Mark Paluch
+ * @since 4.0
+ */
+record AotQueries(AotQuery result, AotQuery count) {
+
+	/**
+	 * Derive a count query from the given query.
+	 */
+	public static AotQueries from(StringAotQuery query, @Null String countProjection, QueryEnhancerSelector selector) {
+
+		QueryEnhancer queryEnhancer = selector.select(query.getQuery()).create(query.getQuery());
+
+		String derivedCountQuery = queryEnhancer
+				.createCountQueryFor(StringUtils.hasText(countProjection) ? countProjection : null);
+
+		DeclaredQuery countQuery = query.getQuery().rewrite(derivedCountQuery);
+		return new AotQueries(query, StringAotQuery.of(countQuery));
+	}
+
+	/**
+	 * Create new {@code AotQueries} for the given queries.
+	 */
+	public static AotQueries from(AotQuery result, AotQuery count) {
+		return new AotQueries(result, count);
+	}
+
+	public boolean isNative() {
+		return result().isNative();
+	}
+
+}

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

@@ -0,0 +1,61 @@
+/*
+ * 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
+ *
+ *      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.jpa.repository.aot.generated;
+
+import java.util.List;
+
+import org.springframework.data.domain.Limit;
+import org.springframework.data.jpa.repository.query.ParameterBinding;
+
+/**
+ * AOT query value object along with its parameter bindings.
+ *
+ * @author Christoph Strobl
+ * @author Mark Paluch
+ * @since 4.0
+ */
+abstract class AotQuery {
+
+	private final List<ParameterBinding> parameterBindings;
+
+	AotQuery(List<ParameterBinding> parameterBindings) {
+		this.parameterBindings = parameterBindings;
+	}
+
+	/**
+	 * @return whether the query is a {@link jakarta.persistence.EntityManager#createNativeQuery native} one.
+	 */
+	public abstract boolean isNative();
+
+	public List<ParameterBinding> getParameterBindings() {
+		return parameterBindings;
+	}
+
+	/**
+	 * @return the preliminary query limit.
+	 */
+	public Limit getLimit() {
+		return Limit.unlimited();
+	}
+
+	/**
+	 * @return whether the query is limited (e.g. {@code findTop10By}).
+	 */
+	public boolean isLimited() {
+		return getLimit().isLimited();
+	}
+
+}

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

@@ -17,46 +17,18 @@
 
 import jakarta.persistence.metamodel.Metamodel;
 
-import org.springframework.data.jpa.repository.Query;
-import org.springframework.data.jpa.repository.query.EscapeCharacter;
-import org.springframework.data.jpa.repository.query.JpaParameters;
-import org.springframework.data.jpa.repository.query.JpaQueryCreator;
-import org.springframework.data.jpa.repository.query.ParameterMetadataProvider;
-import org.springframework.data.jpa.repository.support.JpqlQueryTemplates;
-import org.springframework.data.repository.aot.generate.AotRepositoryMethodGenerationContext;
-import org.springframework.data.repository.query.ParametersSource;
-import org.springframework.data.repository.query.ReturnedType;
-import org.springframework.data.repository.query.parser.PartTree;
-
 /**
  * @author Christoph Strobl
  * @since 2025/01
  */
-public class AotQueryCreator {
+class AotQueryCreator {
 
 	Metamodel metamodel;
 
 	public AotQueryCreator(Metamodel metamodel) {
 		this.metamodel = metamodel;
 	}
 
-	AotStringQuery createQuery(PartTree partTree, ReturnedType returnedType,
-			AotRepositoryMethodGenerationContext context) {
-
-		ParametersSource parametersSource = ParametersSource.of(context.getRepositoryInformation(), context.getMethod());
-		JpaParameters parameters = new JpaParameters(parametersSource);
-		ParameterMetadataProvider metadataProvider = new ParameterMetadataProvider(parameters, EscapeCharacter.DEFAULT,
-				JpqlQueryTemplates.UPPER);
 
-		JpaQueryCreator queryCreator = new JpaQueryCreator(partTree, returnedType, metadataProvider,
-				JpqlQueryTemplates.UPPER, metamodel);
-		AotStringQuery query = AotStringQuery.bindable(queryCreator.createQuery(), metadataProvider.getBindings());
-
-		if (partTree.isLimiting()) {
-			query.setLimit(partTree.getResultLimit());
-		}
-		query.setCountQuery(context.annotationValue(Query.class, "countQuery"));
-		return query;
-	}
 
 }

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

@@ -0,0 +1,118 @@
+/*
+ * 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
+ *
+ *      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.jpa.repository.aot.generated;
+
+import java.lang.reflect.Method;
+
+import org.jspecify.annotations.Nullable;
+
+import org.springframework.data.domain.Sort;
+import org.springframework.data.expression.ValueEvaluationContextProvider;
+import org.springframework.data.expression.ValueExpression;
+import org.springframework.data.jpa.repository.query.DeclaredQuery;
+import org.springframework.data.jpa.repository.query.JpaParameters;
+import org.springframework.data.jpa.repository.query.QueryEnhancer;
+import org.springframework.data.jpa.repository.query.QueryEnhancerSelector;
+import org.springframework.data.projection.ProjectionFactory;
+import org.springframework.data.repository.core.RepositoryMetadata;
+import org.springframework.data.repository.core.support.RepositoryFactoryBeanSupport;
+import org.springframework.data.repository.query.ParametersSource;
+import org.springframework.data.repository.query.ReturnedType;
+import org.springframework.data.repository.query.ValueExpressionDelegate;
+import org.springframework.util.ConcurrentLruCache;
+
+/**
+ * @author Mark Paluch
+ */
+public class AotRepositoryFragmentSupport {
+
+	private final RepositoryMetadata repositoryMetadata;
+
+	private final ValueExpressionDelegate valueExpressions;
+
+	private final ProjectionFactory projectionFactory;
+
+	private final ConcurrentLruCache<DeclaredQuery, QueryEnhancer> enhancers;
+
+	private final ConcurrentLruCache<String, ValueExpression> expressions;
+
+	private final ConcurrentLruCache<Method, ValueEvaluationContextProvider> contextProviders;
+
+	protected AotRepositoryFragmentSupport(QueryEnhancerSelector selector,
+			RepositoryFactoryBeanSupport.FragmentCreationContext context) {
+		this(selector, context.getRepositoryMetadata(), context.getValueExpressionDelegate(),
+				context.getProjectionFactory());
+	}
+
+	protected AotRepositoryFragmentSupport(QueryEnhancerSelector selector, RepositoryMetadata repositoryMetadata,
+			ValueExpressionDelegate valueExpressions, ProjectionFactory projectionFactory) {
+
+		this.repositoryMetadata = repositoryMetadata;
+		this.valueExpressions = valueExpressions;
+		this.projectionFactory = projectionFactory;
+		this.enhancers = new ConcurrentLruCache<>(32, query -> selector.select(query).create(query));
+		this.expressions = new ConcurrentLruCache<>(32, valueExpressions::parse);
+		this.contextProviders = new ConcurrentLruCache<>(32, it -> valueExpressions
+				.createValueContextProvider(new JpaParameters(ParametersSource.of(repositoryMetadata, it))));
+	}
+
+	/**
+	 * Rewrite a {@link DeclaredQuery} to apply {@link Sort} and {@link Class} projection.
+	 *
+	 * @param query
+	 * @param sort
+	 * @param returnedType
+	 * @return
+	 */
+	protected String rewriteQuery(DeclaredQuery query, Sort sort, Class<?> returnedType) {
+
+		QueryEnhancer queryStringEnhancer = this.enhancers.get(query);
+		return queryStringEnhancer.rewrite(new DefaultQueryRewriteInformation(sort,
+				ReturnedType.of(returnedType, repositoryMetadata.getDomainType(), projectionFactory)));
+	}
+
+	/**
+	 * Evaluate a Value Expression.
+	 *
+	 * @param method
+	 * @param expressionString
+	 * @param args
+	 * @return
+	 */
+	protected @Nullable Object evaluateExpression(Method method, String expressionString, Object... args) {
+
+		ValueExpression expression = this.expressions.get(expressionString);
+		ValueEvaluationContextProvider contextProvider = this.contextProviders.get(method);
+
+		return expression.evaluate(contextProvider.getEvaluationContext(args, expression.getExpressionDependencies()));
+	}
+
+	private record DefaultQueryRewriteInformation(Sort sort,
+			ReturnedType returnedType) implements QueryEnhancer.QueryRewriteInformation {
+
+		@Override
+		public Sort getSort() {
+			return sort();
+		}
+
+		@Override
+		public ReturnedType getReturnedType() {
+			return returnedType();
+		}
+
+	}
+
+}