Skip to content

Commit 57cab9a

Browse files
mp911demp911de
committed
Apply QueryRewriter to count queries as well.
We now use QueryRewriter to post-process count queries as well. Previously, only the actual result query has been processed. Closes #3801
1 parent a24597a commit 57cab9a

File tree

7 files changed

+64
-21
lines changed

7 files changed

+64
-21
lines changed

‎spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/QueryRewriter.java

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,10 @@
2626
* and tools intends to do has been done. You can customize the query to apply final changes. Rewriting can only make
2727
* use of already existing contextual data. That is, adding or replacing query text or reuse of bound parameters. Query
2828
* rewriting must not add additional bindable parameters as these cannot be materialized.
29+
* <p>
30+
* Query rewriting applies to the actual query and, when applicable, to count queries. Count queries are optimized and
31+
* therefore, either not necessary or a count is obtained through other means, such as derived from a Hibernate
32+
* {@code SelectionQuery}.
2933
*
3034
* @author Greg Turnquist
3135
* @author Mark Paluch
@@ -71,4 +75,5 @@ public String rewrite(String query, Sort sort) {
7175
return query;
7276
}
7377
}
78+
7479
}

‎spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/AbstractStringBasedJpaQuery.java

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -105,7 +105,6 @@ public AbstractStringBasedJpaQuery(JpaQueryMethod method, EntityManager em, Decl
105105
});
106106

107107
this.countParameterBinder = Lazy.of(() -> this.createBinder(this.countQuery.get()));
108-
109108
this.queryRewriter = queryConfiguration.getQueryRewriter(method);
110109

111110
JpaParameters parameters = method.getParameters();
@@ -162,9 +161,11 @@ protected Query doCreateCountQuery(JpaParametersParameterAccessor accessor) {
162161
String queryString = countQuery.get().getQueryString();
163162
EntityManager em = getEntityManager();
164163

164+
String queryStringToUse = potentiallyRewriteQuery(queryString, accessor.getSort(), accessor.getPageable());
165+
165166
Query query = getQueryMethod().isNativeQuery() //
166-
? em.createNativeQuery(queryString) //
167-
: em.createQuery(queryString, Long.class);
167+
? em.createNativeQuery(queryStringToUse) //
168+
: em.createQuery(queryStringToUse, Long.class);
168169

169170
countParameterBinder.get().bind(new QueryParameterSetter.BindableQuery(query), accessor,
170171
QueryParameterSetter.ErrorHandling.LENIENT);
@@ -194,16 +195,17 @@ protected Query createJpaQuery(QueryProvider query, Sort sort, @Nullable Pageabl
194195
ReturnedType returnedType) {
195196

196197
EntityManager em = getEntityManager();
198+
String queryToUse = potentiallyRewriteQuery(query.getQueryString(), sort, pageable);
197199

198200
if (this.query.hasConstructorExpression() || this.query.isDefaultProjection()) {
199-
return em.createQuery(potentiallyRewriteQuery(query.getQueryString(), sort, pageable));
201+
return em.createQuery(queryToUse);
200202
}
201203

202204
Class<?> typeToRead = getTypeToRead(returnedType);
203205

204206
return typeToRead == null //
205-
? em.createQuery(potentiallyRewriteQuery(query.getQueryString(), sort, pageable)) //
206-
: em.createQuery(potentiallyRewriteQuery(query.getQueryString(), sort, pageable), typeToRead);
207+
? em.createQuery(queryToUse) //
208+
: em.createQuery(queryToUse, typeToRead);
207209
}
208210

209211
/**

‎spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/JpaQueryLookupStrategy.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -169,7 +169,7 @@ protected RepositoryQuery resolveQuery(JpaQueryMethod method, JpaQueryConfigurat
169169
configuration);
170170
}
171171

172-
RepositoryQuery query = NamedQuery.lookupFrom(method, em, configuration.getSelector());
172+
RepositoryQuery query = NamedQuery.lookupFrom(method, em, configuration);
173173

174174
return query != null ? query : NO_QUERY;
175175
}

‎spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/NamedQuery.java

Lines changed: 27 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,10 @@
2424
import org.apache.commons.logging.LogFactory;
2525
import org.jspecify.annotations.Nullable;
2626

27+
import org.springframework.data.domain.Pageable;
28+
import org.springframework.data.domain.Sort;
2729
import org.springframework.data.jpa.provider.QueryExtractor;
30+
import org.springframework.data.jpa.repository.QueryRewriter;
2831
import org.springframework.data.repository.query.Parameters;
2932
import org.springframework.data.repository.query.QueryCreationException;
3033
import org.springframework.data.repository.query.RepositoryQuery;
@@ -53,18 +56,20 @@ final class NamedQuery extends AbstractJpaQuery {
5356
private final @Nullable String countProjection;
5457
private final boolean namedCountQueryIsPresent;
5558
private final Lazy<EntityQuery> entityQuery;
59+
private final QueryRewriter queryRewriter;
5660

5761
/**
5862
* Creates a new {@link NamedQuery}.
5963
*/
60-
private NamedQuery(JpaQueryMethod method, EntityManager em, QueryEnhancerSelector selector) {
64+
private NamedQuery(JpaQueryMethod method, EntityManager em, JpaQueryConfiguration queryConfiguration) {
6165

6266
super(method, em);
6367

6468
this.queryName = method.getNamedQueryName();
6569
this.countQueryName = method.getNamedCountQueryName();
6670
QueryExtractor extractor = method.getQueryExtractor();
6771
this.countProjection = method.getCountQueryProjection();
72+
this.queryRewriter = queryConfiguration.getQueryRewriter(method);
6873

6974
Parameters<?, ?> parameters = method.getParameters();
7075

@@ -100,7 +105,7 @@ private NamedQuery(JpaQueryMethod method, EntityManager em, QueryEnhancerSelecto
100105
declaredQuery = DeclaredQuery.jpqlQuery(queryString);
101106
}
102107

103-
this.entityQuery = Lazy.of(() -> EntityQuery.create(declaredQuery, selector));
108+
this.entityQuery = Lazy.of(() -> EntityQuery.create(declaredQuery, queryConfiguration.getSelector()));
104109
}
105110

106111
/**
@@ -134,9 +139,10 @@ static boolean hasNamedQuery(EntityManager em, String queryName) {
134139
* @param method must not be {@literal null}.
135140
* @param em must not be {@literal null}.
136141
* @param selector must not be {@literal null}.
142+
* @param queryConfiguration must not be {@literal null}.
137143
*/
138144
public static @Nullable RepositoryQuery lookupFrom(JpaQueryMethod method, EntityManager em,
139-
QueryEnhancerSelector selector) {
145+
JpaQueryConfiguration queryConfiguration) {
140146

141147
String queryName = method.getNamedQueryName();
142148

@@ -154,7 +160,7 @@ static boolean hasNamedQuery(EntityManager em, String queryName) {
154160
method.isNativeQuery() ? "NativeQuery" : "Query"));
155161
}
156162

157-
RepositoryQuery query = new NamedQuery(method, em, selector);
163+
RepositoryQuery query = new NamedQuery(method, em, queryConfiguration);
158164
if (LOG.isDebugEnabled()) {
159165
LOG.debug(String.format("Found named query '%s'", queryName));
160166
}
@@ -189,6 +195,7 @@ protected TypedQuery<Long> doCreateCountQuery(JpaParametersParameterAccessor acc
189195
} else {
190196

191197
String countQueryString = entityQuery.get().deriveCountQuery(countProjection).getQueryString();
198+
countQueryString = potentiallyRewriteQuery(countQueryString, accessor.getSort(), accessor.getPageable());
192199
countQuery = em.createQuery(countQueryString, Long.class);
193200
}
194201

@@ -221,4 +228,20 @@ protected TypedQuery<Long> doCreateCountQuery(JpaParametersParameterAccessor acc
221228
? null //
222229
: super.getTypeToRead(returnedType);
223230
}
231+
232+
/**
233+
* Use the {@link QueryRewriter}, potentially rewrite the query, using relevant {@link Sort} and {@link Pageable}
234+
* information.
235+
*
236+
* @param originalQuery
237+
* @param sort
238+
* @param pageable
239+
* @return
240+
*/
241+
protected String potentiallyRewriteQuery(String originalQuery, Sort sort, @Nullable Pageable pageable) {
242+
243+
return pageable != null && pageable.isPaged() //
244+
? queryRewriter.rewrite(originalQuery, pageable) //
245+
: queryRewriter.rewrite(originalQuery, sort);
246+
}
224247
}

‎spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/JpaQueryRewriteIntegrationTests.java

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,10 @@
1818
import static org.assertj.core.api.Assertions.*;
1919

2020
import java.util.HashMap;
21+
import java.util.LinkedHashSet;
2122
import java.util.List;
2223
import java.util.Map;
24+
import java.util.Set;
2325

2426
import org.junit.jupiter.api.BeforeEach;
2527
import org.junit.jupiter.api.Test;
@@ -31,6 +33,7 @@
3133
import org.springframework.context.annotation.Configuration;
3234
import org.springframework.context.annotation.FilterType;
3335
import org.springframework.context.annotation.ImportResource;
36+
import org.springframework.data.domain.Page;
3437
import org.springframework.data.domain.PageRequest;
3538
import org.springframework.data.domain.Pageable;
3639
import org.springframework.data.domain.Sort;
@@ -47,7 +50,7 @@
4750
import org.springframework.test.util.ReflectionTestUtils;
4851

4952
/**
50-
* Unit tests for repository with {@link Query} and {@link QueryRewrite}.
53+
* Unit tests for repository with {@link Query} and {@link QueryRewriter}.
5154
*
5255
* @author Greg Turnquist
5356
* @author Krzysztof Krason
@@ -64,10 +67,12 @@ class JpaQueryRewriteIntegrationTests {
6467
static final String REWRITTEN_QUERY = "rewritten query";
6568
static final String SORT = "sort";
6669
static Map<String, String> results = new HashMap<>();
70+
static Set<String> queries = new LinkedHashSet<>();
6771

6872
@BeforeEach
6973
void setUp() {
7074
results.clear();
75+
repository.deleteAll();
7176
}
7277

7378
@Test
@@ -89,15 +94,15 @@ void nativeQueryShouldHandleRewrites() {
8994
entry(SORT, Sort.unsorted().toString()));
9095
}
9196

92-
@Test
97+
@Test // GH-3801
9398
void nonNativeQueryShouldHandleRewrites() {
9499

95-
repository.findByNonNativeQuery("Matthews");
100+
repository.save(new User("D", "A", "foo@bar"));
96101

97-
assertThat(results).containsExactly( //
98-
entry(ORIGINAL_QUERY, "select original_user_alias from User original_user_alias"), //
99-
entry(REWRITTEN_QUERY, "select rewritten_user_alias from User rewritten_user_alias"), //
100-
entry(SORT, Sort.unsorted().toString()));
102+
repository.findByNonNativeQuery("Matthews", PageRequest.of(0, 1));
103+
104+
assertThat(queries).contains("select original_user_alias from User original_user_alias");
105+
assertThat(queries).contains("select count(original_user_alias) from User original_user_alias");
101106
}
102107

103108
@Test
@@ -181,7 +186,7 @@ public interface UserRepositoryWithRewriter
181186
List<User> findByNativeQuery(String param);
182187

183188
@Query(value = "select original_user_alias from User original_user_alias", queryRewriter = TestQueryRewriter.class)
184-
List<User> findByNonNativeQuery(String param);
189+
Page<User> findByNonNativeQuery(String param, PageRequest pageRequest);
185190

186191
@Query(value = "select original_user_alias from User original_user_alias", queryRewriter = TestQueryRewriter.class)
187192
List<User> findByNonNativeSortedQuery(String param, Sort sort);
@@ -226,6 +231,7 @@ private static String replaceAlias(String query, Sort sort) {
226231
results.put(ORIGINAL_QUERY, query);
227232
results.put(REWRITTEN_QUERY, rewrittenQuery);
228233
results.put(SORT, sort.toString());
234+
queries.add(query);
229235

230236
return rewrittenQuery;
231237
}

‎spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/NamedQueryUnitTests.java

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@
4040
import org.springframework.data.projection.SpelAwareProxyProjectionFactory;
4141
import org.springframework.data.repository.core.RepositoryMetadata;
4242
import org.springframework.data.repository.query.QueryCreationException;
43+
import org.springframework.data.repository.query.ValueExpressionDelegate;
4344
import org.springframework.data.util.TypeInformation;
4445

4546
/**
@@ -54,6 +55,9 @@
5455
@MockitoSettings(strictness = Strictness.LENIENT)
5556
class NamedQueryUnitTests {
5657

58+
private static final JpaQueryConfiguration CONFIG = new JpaQueryConfiguration(QueryRewriterProvider.simple(),
59+
QueryEnhancerSelector.DEFAULT_SELECTOR, ValueExpressionDelegate.create(), EscapeCharacter.DEFAULT);
60+
5761
@Mock RepositoryMetadata metadata;
5862
@Mock QueryExtractor extractor;
5963
@Mock EntityManager em;
@@ -88,7 +92,8 @@ void rejectsPersistenceProviderIfIncapableOfExtractingQueriesAndPagebleBeingUsed
8892
JpaQueryMethod queryMethod = new JpaQueryMethod(method, metadata, projectionFactory, extractor);
8993

9094
when(em.createNamedQuery(queryMethod.getNamedCountQueryName())).thenThrow(new IllegalArgumentException());
91-
assertThatExceptionOfType(QueryCreationException.class).isThrownBy(() -> NamedQuery.lookupFrom(queryMethod, em, QueryEnhancerSelector.DEFAULT_SELECTOR));
95+
assertThatExceptionOfType(QueryCreationException.class)
96+
.isThrownBy(() -> NamedQuery.lookupFrom(queryMethod, em, CONFIG));
9297
}
9398

9499
@Test // DATAJPA-142
@@ -100,7 +105,7 @@ void doesNotRejectPersistenceProviderIfNamedCountQueryIsAvailable() {
100105

101106
TypedQuery<Long> countQuery = mock(TypedQuery.class);
102107
when(em.createNamedQuery(eq(queryMethod.getNamedCountQueryName()), eq(Long.class))).thenReturn(countQuery);
103-
NamedQuery query = (NamedQuery) NamedQuery.lookupFrom(queryMethod, em, QueryEnhancerSelector.DEFAULT_SELECTOR);
108+
NamedQuery query = (NamedQuery) NamedQuery.lookupFrom(queryMethod, em, CONFIG);
104109

105110
query.doCreateCountQuery(new JpaParametersParameterAccessor(queryMethod.getParameters(), new Object[1]));
106111
verify(em, times(1)).createNamedQuery(queryMethod.getNamedCountQueryName(), Long.class);

‎src/main/antora/modules/ROOT/pages/jpa/query-methods.adoc

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -293,6 +293,8 @@ Sometimes, no matter how many features you try to apply, it seems impossible to
293293

294294
You have the ability to get your hands on the query, right before it's sent to the `EntityManager` and "rewrite" it.
295295
That is, you can make any alterations at the last moment.
296+
Query rewriting applies to the actual query and, when applicable, to count queries.
297+
Count queries are optimized and therefore, either not necessary or a count is obtained through other means, such as derived from a Hibernate `SelectionQuery`.
296298

297299
.Declare a QueryRewriter using `@Query`
298300
====

0 commit comments

Comments
 (0)