Skip to content

Commit 52e453c

Browse files
mp911demp911de
committed
Add query hint support.
See #3830
1 parent 5561fd1 commit 52e453c

File tree

4 files changed

+159
-95
lines changed

4 files changed

+159
-95
lines changed

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

Lines changed: 143 additions & 92 deletions
Original file line numberDiff line numberDiff line change
@@ -17,13 +17,16 @@
1717

1818
import jakarta.persistence.EntityManager;
1919
import jakarta.persistence.Query;
20+
import jakarta.persistence.QueryHint;
2021

2122
import java.util.List;
2223
import java.util.Optional;
2324
import java.util.function.LongSupplier;
2425
import java.util.regex.Pattern;
2526

27+
import org.springframework.core.annotation.MergedAnnotation;
2628
import org.springframework.data.domain.SliceImpl;
29+
import org.springframework.data.jpa.repository.QueryHints;
2730
import org.springframework.data.jpa.repository.query.DeclaredQuery;
2831
import org.springframework.data.jpa.repository.query.ParameterBinding;
2932
import org.springframework.data.repository.aot.generate.AotRepositoryMethodGenerationContext;
@@ -37,93 +40,21 @@
3740

3841
/**
3942
* @author Christoph Strobl
40-
* @since 2025/01
43+
* @author Mark Paluch
44+
* @since 4.0
4145
*/
42-
public class JpaCodeBlocks {
46+
class JpaCodeBlocks {
4347

4448
private static final Pattern PARAMETER_BINDING_PATTERN = Pattern.compile("\\?(\\d+)");
4549

46-
static QueryBlockBuilder queryBlockBuilder(AotRepositoryMethodGenerationContext context) {
50+
public static QueryBlockBuilder queryBuilder(AotRepositoryMethodGenerationContext context) {
4751
return new QueryBlockBuilder(context);
4852
}
4953

50-
static QueryExecutionBlockBuilder queryExecutionBlockBuilder(AotRepositoryMethodGenerationContext context) {
54+
static QueryExecutionBlockBuilder executionBuilder(AotRepositoryMethodGenerationContext context) {
5155
return new QueryExecutionBlockBuilder(context);
5256
}
5357

54-
static class QueryExecutionBlockBuilder {
55-
56-
AotRepositoryMethodGenerationContext context;
57-
private String queryVariableName = "query";
58-
59-
public QueryExecutionBlockBuilder(AotRepositoryMethodGenerationContext context) {
60-
this.context = context;
61-
}
62-
63-
QueryExecutionBlockBuilder referencing(String queryVariableName) {
64-
65-
this.queryVariableName = queryVariableName;
66-
return this;
67-
}
68-
69-
CodeBlock build() {
70-
71-
Builder builder = CodeBlock.builder();
72-
73-
boolean isProjecting = context.getActualReturnType() != null
74-
&& !ObjectUtils.nullSafeEquals(TypeName.get(context.getRepositoryInformation().getDomainType()),
75-
context.getActualReturnType());
76-
Object actualReturnType = isProjecting ? context.getActualReturnType()
77-
: context.getRepositoryInformation().getDomainType();
78-
79-
builder.add("\n");
80-
81-
if (context.isDeleteMethod()) {
82-
83-
builder.addStatement("$T<$T> resultList = $L.getResultList()", List.class, actualReturnType, queryVariableName);
84-
builder.addStatement("resultList.forEach($L::remove)", context.fieldNameOf(EntityManager.class));
85-
if (context.returnsSingleValue()) {
86-
if (ClassUtils.isAssignable(Number.class, context.getMethod().getReturnType())) {
87-
builder.addStatement("return $T.valueOf(resultList.size())", context.getMethod().getReturnType());
88-
} else {
89-
builder.addStatement("return resultList.isEmpty() ? null : resultList.iterator().next()");
90-
}
91-
} else {
92-
builder.addStatement("return resultList");
93-
}
94-
} else if (context.isExistsMethod()) {
95-
builder.addStatement("return !$L.getResultList().isEmpty()", queryVariableName);
96-
} else {
97-
98-
if (context.returnsSingleValue()) {
99-
if (context.returnsOptionalValue()) {
100-
builder.addStatement("return $T.ofNullable(($T) $L.getSingleResultOrNull())", Optional.class,
101-
actualReturnType, queryVariableName);
102-
} else {
103-
builder.addStatement("return ($T) $L.getSingleResultOrNull()", context.getReturnType(), queryVariableName);
104-
}
105-
} else if (context.returnsPage()) {
106-
builder.addStatement("return $T.getPage(($T<$T>) $L.getResultList(), $L, countAll)",
107-
PageableExecutionUtils.class, List.class, actualReturnType, queryVariableName,
108-
context.getPageableParameterName());
109-
} else if (context.returnsSlice()) {
110-
builder.addStatement("$T<$T> resultList = $L.getResultList()", List.class, actualReturnType,
111-
queryVariableName);
112-
builder.addStatement("boolean hasNext = $L.isPaged() && resultList.size() > $L.getPageSize()",
113-
context.getPageableParameterName(), context.getPageableParameterName());
114-
builder.addStatement(
115-
"return new $T<>(hasNext ? resultList.subList(0, $L.getPageSize()) : resultList, $L, hasNext)",
116-
SliceImpl.class, context.getPageableParameterName(), context.getPageableParameterName());
117-
} else {
118-
builder.addStatement("return ($T) query.getResultList()", context.getReturnType());
119-
}
120-
}
121-
122-
return builder.build();
123-
124-
}
125-
}
126-
12758
/**
12859
* Builder for the actual query code block.
12960
*/
@@ -132,23 +63,35 @@ static class QueryBlockBuilder {
13263
private final AotRepositoryMethodGenerationContext context;
13364
private String queryVariableName = "query";
13465
private AotQueries queries;
66+
private MergedAnnotation<QueryHints> queryHints = MergedAnnotation.missing();
13567

136-
public QueryBlockBuilder(AotRepositoryMethodGenerationContext context) {
68+
private QueryBlockBuilder(AotRepositoryMethodGenerationContext context) {
13769
this.context = context;
13870
}
13971

140-
QueryBlockBuilder usingQueryVariableName(String queryVariableName) {
72+
public QueryBlockBuilder usingQueryVariableName(String queryVariableName) {
14173

14274
this.queryVariableName = queryVariableName;
14375
return this;
14476
}
14577

146-
QueryBlockBuilder filter(AotQueries query) {
78+
public QueryBlockBuilder filter(AotQueries query) {
14779
this.queries = query;
14880
return this;
14981
}
15082

151-
CodeBlock build() {
83+
public QueryBlockBuilder queryHints(MergedAnnotation<QueryHints> queryHints) {
84+
85+
this.queryHints = queryHints;
86+
return this;
87+
}
88+
89+
/**
90+
* Build the query block.
91+
*
92+
* @return
93+
*/
94+
public CodeBlock build() {
15295

15396
boolean isProjecting = context.getActualReturnType() != null
15497
&& !ObjectUtils.nullSafeEquals(TypeName.get(context.getRepositoryInformation().getDomainType()),
@@ -172,8 +115,7 @@ CodeBlock build() {
172115
countQuyerVariableName = "count%s".formatted(StringUtils.capitalize(queryVariableName));
173116

174117
StringAotQuery countQuery = (StringAotQuery) queries.count();
175-
builder.addStatement("$T $L = $S", String.class, countQueryStringNameVariableName,
176-
countQuery.getQueryString());
118+
builder.addStatement("$T $L = $S", String.class, countQueryStringNameVariableName, countQuery.getQueryString());
177119
}
178120

179121
// sorting
@@ -185,17 +127,21 @@ CodeBlock build() {
185127
}
186128

187129
if (StringUtils.hasText(sortParameterName)) {
188-
applySorting(builder, sortParameterName, queryStringNameVariableName, actualReturnType);
130+
builder.add(applySorting(sortParameterName, queryStringNameVariableName, actualReturnType));
189131
}
190132

191-
addQueryBlock(builder, queryVariableName, queryStringNameVariableName, queries.result());
133+
builder.add(createQuery(queryVariableName, queryStringNameVariableName, queries.result(), queryHints));
192134

193-
applyLimits(builder);
135+
builder.add(applyLimits());
194136

195137
if (StringUtils.hasText(countQueryStringNameVariableName)) {
196138

197139
builder.beginControlFlow("$T $L = () ->", LongSupplier.class, "countAll");
198-
addQueryBlock(builder, countQuyerVariableName, countQueryStringNameVariableName, queries.count());
140+
141+
boolean queryHints = this.queryHints.isPresent() && this.queryHints.getBoolean("forCounting");
142+
143+
builder.add(createQuery(countQuyerVariableName, countQueryStringNameVariableName, queries.count(),
144+
queryHints ? this.queryHints : MergedAnnotation.missing()));
199145
builder.addStatement("return ($T) $L.getSingleResult()", Long.class, countQuyerVariableName);
200146

201147
// end control flow does not work well with lambdas
@@ -206,8 +152,9 @@ CodeBlock build() {
206152
return builder.build();
207153
}
208154

209-
private void applySorting(Builder builder, String sort, String queryString, Object actualReturnType) {
155+
private CodeBlock applySorting(String sort, String queryString, Object actualReturnType) {
210156

157+
Builder builder = CodeBlock.builder();
211158
builder.beginControlFlow("if ($L.isSorted())", sort);
212159

213160
if (queries.isNative()) {
@@ -221,14 +168,18 @@ private void applySorting(Builder builder, String sort, String queryString, Obje
221168
builder.addStatement("$L = rewriteQuery(declaredQuery, $L, $T.class)", queryString, sort, actualReturnType);
222169

223170
builder.endControlFlow();
171+
172+
return builder.build();
224173
}
225174

226-
private void applyLimits(Builder builder) {
175+
private CodeBlock applyLimits() {
176+
177+
Builder builder = CodeBlock.builder();
227178

228179
if (context.isExistsMethod()) {
229180
builder.addStatement("$L.setMaxResults(1)", queryVariableName);
230181

231-
return;
182+
return builder.build();
232183
}
233184

234185
String limit = context.getLimitParameterName();
@@ -254,15 +205,24 @@ private void applyLimits(Builder builder) {
254205
}
255206
builder.endControlFlow();
256207
}
208+
209+
return builder.build();
257210
}
258211

259-
private void addQueryBlock(Builder builder, String queryVariableName, String queryStringNameVariableName,
260-
AotQuery query) {
212+
private CodeBlock createQuery(String queryVariableName, String queryStringNameVariableName, AotQuery query,
213+
MergedAnnotation<QueryHints> queryHints) {
214+
215+
Builder builder = CodeBlock.builder();
261216

262217
builder.addStatement("$T $L = this.$L.$L($L)", Query.class, queryVariableName,
263218
context.fieldNameOf(EntityManager.class), query.isNative() ? "createNativeQuery" : "createQuery",
264219
queryStringNameVariableName);
265220

221+
if (queryHints.isPresent()) {
222+
builder.add(applyHints(queryVariableName, queryHints));
223+
builder.add("\n");
224+
}
225+
266226
for (ParameterBinding binding : query.getParameterBindings()) {
267227

268228
Object prepare = binding.prepare("s");
@@ -287,6 +247,97 @@ private void addQueryBlock(Builder builder, String queryVariableName, String que
287247
}
288248
}
289249
}
250+
251+
return builder.build();
290252
}
253+
254+
private CodeBlock applyHints(String queryVariableName, MergedAnnotation<QueryHints> queryHints) {
255+
256+
Builder hintsBuilder = CodeBlock.builder();
257+
MergedAnnotation<QueryHint>[] values = queryHints.getAnnotationArray("value", QueryHint.class);
258+
259+
for (MergedAnnotation<QueryHint> hint : values) {
260+
hintsBuilder.addStatement("$L.setHint($S, $S)", queryVariableName, hint.getString("name"),
261+
hint.getString("value"));
262+
}
263+
264+
return hintsBuilder.build();
265+
}
266+
291267
}
268+
269+
static class QueryExecutionBlockBuilder {
270+
271+
private final AotRepositoryMethodGenerationContext context;
272+
private String queryVariableName = "query";
273+
274+
private QueryExecutionBlockBuilder(AotRepositoryMethodGenerationContext context) {
275+
this.context = context;
276+
}
277+
278+
public QueryExecutionBlockBuilder referencing(String queryVariableName) {
279+
280+
this.queryVariableName = queryVariableName;
281+
return this;
282+
}
283+
284+
public CodeBlock build() {
285+
286+
Builder builder = CodeBlock.builder();
287+
288+
boolean isProjecting = context.getActualReturnType() != null
289+
&& !ObjectUtils.nullSafeEquals(TypeName.get(context.getRepositoryInformation().getDomainType()),
290+
context.getActualReturnType());
291+
Object actualReturnType = isProjecting ? context.getActualReturnType()
292+
: context.getRepositoryInformation().getDomainType();
293+
builder.add("\n");
294+
295+
if (context.isDeleteMethod()) {
296+
297+
builder.addStatement("$T<$T> resultList = $L.getResultList()", List.class, actualReturnType, queryVariableName);
298+
builder.addStatement("resultList.forEach($L::remove)", context.fieldNameOf(EntityManager.class));
299+
if (context.returnsSingleValue()) {
300+
if (ClassUtils.isAssignable(Number.class, context.getMethod().getReturnType())) {
301+
builder.addStatement("return $T.valueOf(resultList.size())", context.getMethod().getReturnType());
302+
} else {
303+
builder.addStatement("return resultList.isEmpty() ? null : resultList.iterator().next()");
304+
}
305+
} else {
306+
builder.addStatement("return resultList");
307+
}
308+
} else if (context.isExistsMethod()) {
309+
builder.addStatement("return !$L.getResultList().isEmpty()", queryVariableName);
310+
} else {
311+
312+
if (context.returnsSingleValue()) {
313+
if (context.returnsOptionalValue()) {
314+
builder.addStatement("return $T.ofNullable(($T) $L.getSingleResultOrNull())", Optional.class,
315+
actualReturnType, queryVariableName);
316+
} else {
317+
builder.addStatement("return ($T) $L.getSingleResultOrNull()", context.getReturnType(), queryVariableName);
318+
}
319+
} else if (context.returnsPage()) {
320+
builder.addStatement("return $T.getPage(($T<$T>) $L.getResultList(), $L, countAll)",
321+
PageableExecutionUtils.class, List.class, actualReturnType, queryVariableName,
322+
context.getPageableParameterName());
323+
} else if (context.returnsSlice()) {
324+
builder.addStatement("$T<$T> resultList = $L.getResultList()", List.class, actualReturnType,
325+
queryVariableName);
326+
builder.addStatement("boolean hasNext = $L.isPaged() && resultList.size() > $L.getPageSize()",
327+
context.getPageableParameterName(), context.getPageableParameterName());
328+
builder.addStatement(
329+
"return new $T<>(hasNext ? resultList.subList(0, $L.getPageSize()) : resultList, $L, hasNext)",
330+
SliceImpl.class, context.getPageableParameterName(), context.getPageableParameterName());
331+
} else {
332+
builder.addStatement("return ($T) query.getResultList()", context.getReturnType());
333+
}
334+
}
335+
336+
return builder.build();
337+
338+
}
339+
340+
}
341+
342+
292343
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -125,8 +125,8 @@ protected AotRepositoryMethodBuilder contributeRepositoryMethod(
125125
aotQueries = buildPartTreeQuery(context, query);
126126
}
127127

128-
body.addCode(JpaCodeBlocks.queryBlockBuilder(context).filter(aotQueries).build());
129-
body.addCode(JpaCodeBlocks.queryExecutionBlockBuilder(context).build());
128+
body.addCode(JpaCodeBlocks.queryBuilder(context).filter(aotQueries).queryHints(queryHints).build());
129+
body.addCode(JpaCodeBlocks.executionBuilder(context).build());
130130
});
131131
}
132132

‎spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/aot/generated/JpaRepositoryContributorIntegrationTests.java‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -297,6 +297,12 @@ void testDerivedFinderReturningListOfProjections() {
297297
"kylo@new-empire.com", "luke@jedi.org", "vader@empire.com");
298298
}
299299

300+
@Test
301+
void shouldApplyQueryHints() {
302+
assertThatIllegalArgumentException().isThrownBy(() -> fragment.findHintedByLastname("Skywalker"))
303+
.withMessageContaining("No enum constant jakarta.persistence.CacheStoreMode.foo");
304+
}
305+
300306
@Test
301307
void testDerivedFinderReturningPageOfProjections() {
302308

@@ -344,8 +350,9 @@ void todo() {
344350

345351
// interface projections
346352
// named queries
353+
// dynamic projections
354+
// class type parameter
347355

348-
// query hints
349356
// entity graphs
350357
// native queries
351358
// delete

0 commit comments

Comments
 (0)