Skip to content

Commit 9d6676c

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 42d8956 commit 9d6676c

File tree

7 files changed

+66
-24
lines changed

7 files changed

+66
-24
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 & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -151,9 +151,11 @@ protected Query doCreateCountQuery(JpaParametersParameterAccessor accessor) {
151151
String queryString = countQuery.get().getQueryString();
152152
EntityManager em = getEntityManager();
153153

154+
String queryStringToUse = potentiallyRewriteQuery(queryString, accessor.getSort(), accessor.getPageable());
155+
154156
Query query = getQueryMethod().isNativeQuery() //
155-
? em.createNativeQuery(queryString) //
156-
: em.createQuery(queryString, Long.class);
157+
? em.createNativeQuery(queryStringToUse) //
158+
: em.createQuery(queryStringToUse, Long.class);
157159

158160
QueryParameterSetter.QueryMetadata metadata = metadataCache.getMetadata(queryString, query);
159161

@@ -184,16 +186,17 @@ protected Query createJpaQuery(String queryString, Sort sort, @Nullable Pageable
184186
ReturnedType returnedType) {
185187

186188
EntityManager em = getEntityManager();
189+
String queryToUse = potentiallyRewriteQuery(queryString, sort, pageable);
187190

188191
if (this.query.hasConstructorExpression() || this.query.isDefaultProjection()) {
189-
return em.createQuery(potentiallyRewriteQuery(queryString, sort, pageable));
192+
return em.createQuery(queryToUse);
190193
}
191194

192195
Class<?> typeToRead = getTypeToRead(returnedType);
193196

194197
return typeToRead == null //
195-
? em.createQuery(potentiallyRewriteQuery(queryString, sort, pageable)) //
196-
: em.createQuery(potentiallyRewriteQuery(queryString, sort, pageable), typeToRead);
198+
? em.createQuery(queryToUse) //
199+
: em.createQuery(queryToUse, typeToRead);
197200
}
198201

199202
/**

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

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -181,11 +181,9 @@ protected RepositoryQuery resolveQuery(JpaQueryMethod method, QueryRewriter quer
181181
getCountQuery(method, namedQueries, em), queryRewriter, valueExpressionDelegate);
182182
}
183183

184-
RepositoryQuery query = NamedQuery.lookupFrom(method, em);
184+
RepositoryQuery query = NamedQuery.lookupFrom(method, em, queryRewriter);
185185

186-
return query != null //
187-
? query //
188-
: NO_QUERY;
186+
return query != null ? query : NO_QUERY;
189187
}
190188

191189
@Nullable

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

Lines changed: 27 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,11 @@
2222

2323
import org.apache.commons.logging.Log;
2424
import org.apache.commons.logging.LogFactory;
25+
26+
import org.springframework.data.domain.Pageable;
27+
import org.springframework.data.domain.Sort;
2528
import org.springframework.data.jpa.provider.QueryExtractor;
29+
import org.springframework.data.jpa.repository.QueryRewriter;
2630
import org.springframework.data.repository.query.Parameters;
2731
import org.springframework.data.repository.query.QueryCreationException;
2832
import org.springframework.data.repository.query.RepositoryQuery;
@@ -53,18 +57,20 @@ final class NamedQuery extends AbstractJpaQuery {
5357
private final boolean namedCountQueryIsPresent;
5458
private final Lazy<DeclaredQuery> declaredQuery;
5559
private final QueryParameterSetter.QueryMetadataCache metadataCache;
60+
private final QueryRewriter queryRewriter;
5661

5762
/**
5863
* Creates a new {@link NamedQuery}.
5964
*/
60-
private NamedQuery(JpaQueryMethod method, EntityManager em) {
65+
private NamedQuery(JpaQueryMethod method, EntityManager em, QueryRewriter queryRewriter) {
6166

6267
super(method, em);
6368

6469
this.queryName = method.getNamedQueryName();
6570
this.countQueryName = method.getNamedCountQueryName();
6671
QueryExtractor extractor = method.getQueryExtractor();
6772
this.countProjection = method.getCountQueryProjection();
73+
this.queryRewriter = queryRewriter;
6874

6975
Parameters<?, ?> parameters = method.getParameters();
7076

@@ -127,9 +133,10 @@ static boolean hasNamedQuery(EntityManager em, String queryName) {
127133
*
128134
* @param method must not be {@literal null}.
129135
* @param em must not be {@literal null}.
136+
* @param queryRewriter must not be {@literal null}.
130137
*/
131138
@Nullable
132-
public static RepositoryQuery lookupFrom(JpaQueryMethod method, EntityManager em) {
139+
public static RepositoryQuery lookupFrom(JpaQueryMethod method, EntityManager em, QueryRewriter queryRewriter) {
133140

134141
String queryName = method.getNamedQueryName();
135142

@@ -147,7 +154,7 @@ public static RepositoryQuery lookupFrom(JpaQueryMethod method, EntityManager em
147154
method.isNativeQuery() ? "NativeQuery" : "Query"));
148155
}
149156

150-
RepositoryQuery query = new NamedQuery(method, em);
157+
RepositoryQuery query = new NamedQuery(method, em, queryRewriter);
151158
if (LOG.isDebugEnabled()) {
152159
LOG.debug(String.format("Found named query '%s'", queryName));
153160
}
@@ -187,6 +194,7 @@ protected TypedQuery<Long> doCreateCountQuery(JpaParametersParameterAccessor acc
187194
} else {
188195

189196
String countQueryString = declaredQuery.get().deriveCountQuery(countProjection).getQueryString();
197+
countQueryString = potentiallyRewriteQuery(countQueryString, accessor.getSort(), accessor.getPageable());
190198
cacheKey = countQueryString;
191199
countQuery = em.createQuery(countQueryString, Long.class);
192200
}
@@ -222,4 +230,20 @@ protected Class<?> getTypeToRead(ReturnedType returnedType) {
222230
? null //
223231
: super.getTypeToRead(returnedType);
224232
}
233+
234+
/**
235+
* Use the {@link QueryRewriter}, potentially rewrite the query, using relevant {@link Sort} and {@link Pageable}
236+
* information.
237+
*
238+
* @param originalQuery
239+
* @param sort
240+
* @param pageable
241+
* @return
242+
*/
243+
private String potentiallyRewriteQuery(String originalQuery, Sort sort, Pageable pageable) {
244+
245+
return pageable.isPaged() //
246+
? queryRewriter.rewrite(originalQuery, pageable) //
247+
: queryRewriter.rewrite(originalQuery, sort);
248+
}
225249
}

‎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
@@ -19,8 +19,10 @@
1919
import static org.assertj.core.api.Assertions.entry;
2020

2121
import java.util.HashMap;
22+
import java.util.LinkedHashSet;
2223
import java.util.List;
2324
import java.util.Map;
25+
import java.util.Set;
2426

2527
import org.junit.jupiter.api.BeforeEach;
2628
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;
@@ -44,7 +47,7 @@
4447
import org.springframework.test.context.junit.jupiter.SpringExtension;
4548

4649
/**
47-
* Unit tests for repository with {@link Query} and {@link QueryRewrite}.
50+
* Unit tests for repository with {@link Query} and {@link QueryRewriter}.
4851
*
4952
* @author Greg Turnquist
5053
* @author Krzysztof Krason
@@ -60,10 +63,12 @@ class JpaQueryRewriteIntegrationTests {
6063
static final String REWRITTEN_QUERY = "rewritten query";
6164
static final String SORT = "sort";
6265
static Map<String, String> results = new HashMap<>();
66+
static Set<String> queries = new LinkedHashSet<>();
6367

6468
@BeforeEach
6569
void setUp() {
6670
results.clear();
71+
repository.deleteAll();
6772
}
6873

6974
@Test
@@ -77,15 +82,15 @@ void nativeQueryShouldHandleRewrites() {
7782
entry(SORT, Sort.unsorted().toString()));
7883
}
7984

80-
@Test
85+
@Test // GH-3801
8186
void nonNativeQueryShouldHandleRewrites() {
8287

83-
repository.findByNonNativeQuery("Matthews");
88+
repository.save(new User("D", "A", "foo@bar"));
8489

85-
assertThat(results).containsExactly( //
86-
entry(ORIGINAL_QUERY, "select original_user_alias from User original_user_alias"), //
87-
entry(REWRITTEN_QUERY, "select rewritten_user_alias from User rewritten_user_alias"), //
88-
entry(SORT, Sort.unsorted().toString()));
90+
repository.findByNonNativeQuery("Matthews", PageRequest.of(0, 1));
91+
92+
assertThat(queries).contains("select original_user_alias from User original_user_alias");
93+
assertThat(queries).contains("select count(original_user_alias) from User original_user_alias");
8994
}
9095

9196
@Test
@@ -169,7 +174,7 @@ public interface UserRepositoryWithRewriter
169174
List<User> findByNativeQuery(String param);
170175

171176
@Query(value = "select original_user_alias from User original_user_alias", queryRewriter = TestQueryRewriter.class)
172-
List<User> findByNonNativeQuery(String param);
177+
Page<User> findByNonNativeQuery(String param, PageRequest pageRequest);
173178

174179
@Query(value = "select original_user_alias from User original_user_alias", queryRewriter = TestQueryRewriter.class)
175180
List<User> findByNonNativeSortedQuery(String param, Sort sort);
@@ -214,6 +219,7 @@ private static String replaceAlias(String query, Sort sort) {
214219
results.put(ORIGINAL_QUERY, query);
215220
results.put(REWRITTEN_QUERY, rewrittenQuery);
216221
results.put(SORT, sort.toString());
222+
queries.add(query);
217223

218224
return rewrittenQuery;
219225
}

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

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
import org.springframework.data.domain.Page;
3737
import org.springframework.data.domain.Pageable;
3838
import org.springframework.data.jpa.provider.QueryExtractor;
39+
import org.springframework.data.jpa.repository.QueryRewriter;
3940
import org.springframework.data.projection.ProjectionFactory;
4041
import org.springframework.data.projection.SpelAwareProxyProjectionFactory;
4142
import org.springframework.data.repository.core.RepositoryMetadata;
@@ -88,7 +89,8 @@ void rejectsPersistenceProviderIfIncapableOfExtractingQueriesAndPagebleBeingUsed
8889
JpaQueryMethod queryMethod = new JpaQueryMethod(method, metadata, projectionFactory, extractor);
8990

9091
when(em.createNamedQuery(queryMethod.getNamedCountQueryName())).thenThrow(new IllegalArgumentException());
91-
assertThatExceptionOfType(QueryCreationException.class).isThrownBy(() -> NamedQuery.lookupFrom(queryMethod, em));
92+
assertThatExceptionOfType(QueryCreationException.class)
93+
.isThrownBy(() -> NamedQuery.lookupFrom(queryMethod, em, QueryRewriter.IdentityQueryRewriter.INSTANCE));
9294
}
9395

9496
@Test // DATAJPA-142
@@ -100,7 +102,8 @@ void doesNotRejectPersistenceProviderIfNamedCountQueryIsAvailable() {
100102

101103
TypedQuery<Long> countQuery = mock(TypedQuery.class);
102104
when(em.createNamedQuery(eq(queryMethod.getNamedCountQueryName()), eq(Long.class))).thenReturn(countQuery);
103-
NamedQuery query = (NamedQuery) NamedQuery.lookupFrom(queryMethod, em);
105+
NamedQuery query = (NamedQuery) NamedQuery.lookupFrom(queryMethod, em,
106+
QueryRewriter.IdentityQueryRewriter.INSTANCE);
104107

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

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

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -176,8 +176,11 @@ public interface UserRepository extends JpaRepository<User, Long> {
176176
Sometimes, no matter how many features you try to apply, it seems impossible to get Spring Data JPA to apply every thing
177177
you'd like to a query before it is sent to the `EntityManager`.
178178

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

182185
.Declare a QueryRewriter using `@Query`
183186
====

0 commit comments

Comments
 (0)