Spring/spring-data-jpa
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Apply QueryRewriter to count queries as well.

Browse Source 
We now use QueryRewriter to post-process count queries as well. Previously, only the actual result query has been processed.

Closes #3801
This commit is contained in:
Mark Paluch
2025-03-19 10:06:27 +01:00
parent 42d8956ca5
commit 9d6676c893
7 changed files with 66 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/QueryRewriter.java
View File

@@ -26,6 +26,10 @@ import org.springframework.data.domain.Sort;
* and tools intends to do has been done. You can customize the query to apply final changes. Rewriting can only make
* use of already existing contextual data. That is, adding or replacing query text or reuse of bound parameters. Query
* rewriting must not add additional bindable parameters as these cannot be materialized.
* <p>
* Query rewriting applies to the actual query and, when applicable, to count queries. Count queries are optimized and
* therefore, either not necessary or a count is obtained through other means, such as derived from a Hibernate
* {@code SelectionQuery}.
*
* @author Greg Turnquist
* @author Mark Paluch
@@ -71,4 +75,5 @@ public interface QueryRewriter {
return query;
}
}
}

13
spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/AbstractStringBasedJpaQuery.java
View File

@@ -151,9 +151,11 @@ abstract class AbstractStringBasedJpaQuery extends AbstractJpaQuery {
String queryString = countQuery.get().getQueryString();
EntityManager em = getEntityManager();
String queryStringToUse = potentiallyRewriteQuery(queryString, accessor.getSort(), accessor.getPageable());
Query query = getQueryMethod().isNativeQuery() //
? em.createNativeQuery(queryString) //
: em.createQuery(queryString, Long.class);
? em.createNativeQuery(queryStringToUse) //
: em.createQuery(queryStringToUse, Long.class);
QueryParameterSetter.QueryMetadata metadata = metadataCache.getMetadata(queryString, query);
@@ -184,16 +186,17 @@ abstract class AbstractStringBasedJpaQuery extends AbstractJpaQuery {
ReturnedType returnedType) {
EntityManager em = getEntityManager();
String queryToUse = potentiallyRewriteQuery(queryString, sort, pageable);
if (this.query.hasConstructorExpression() || this.query.isDefaultProjection()) {
return em.createQuery(potentiallyRewriteQuery(queryString, sort, pageable));
return em.createQuery(queryToUse);
}
Class<?> typeToRead = getTypeToRead(returnedType);
return typeToRead == null //
? em.createQuery(potentiallyRewriteQuery(queryString, sort, pageable)) //
: em.createQuery(potentiallyRewriteQuery(queryString, sort, pageable), typeToRead);
? em.createQuery(queryToUse) //
: em.createQuery(queryToUse, typeToRead);
}
/**

6
spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/JpaQueryLookupStrategy.java
View File

@@ -181,11 +181,9 @@ public final class JpaQueryLookupStrategy {
getCountQuery(method, namedQueries, em), queryRewriter, valueExpressionDelegate);
}
RepositoryQuery query = NamedQuery.lookupFrom(method, em);
RepositoryQuery query = NamedQuery.lookupFrom(method, em, queryRewriter);
return query != null //
? query //
: NO_QUERY;
return query != null ? query : NO_QUERY;
}
@Nullable

30
spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/NamedQuery.java
View File

@@ -22,7 +22,11 @@ import jakarta.persistence.TypedQuery;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.data.domain.Pageable;
import org.springframework.data.domain.Sort;
import org.springframework.data.jpa.provider.QueryExtractor;
import org.springframework.data.jpa.repository.QueryRewriter;
import org.springframework.data.repository.query.Parameters;
import org.springframework.data.repository.query.QueryCreationException;
import org.springframework.data.repository.query.RepositoryQuery;
@@ -53,11 +57,12 @@ final class NamedQuery extends AbstractJpaQuery {
private final boolean namedCountQueryIsPresent;
private final Lazy<DeclaredQuery> declaredQuery;
private final QueryParameterSetter.QueryMetadataCache metadataCache;
private final QueryRewriter queryRewriter;
/**
* Creates a new {@link NamedQuery}.
*/
private NamedQuery(JpaQueryMethod method, EntityManager em) {
private NamedQuery(JpaQueryMethod method, EntityManager em, QueryRewriter queryRewriter) {
super(method, em);
@@ -65,6 +70,7 @@ final class NamedQuery extends AbstractJpaQuery {
this.countQueryName = method.getNamedCountQueryName();
QueryExtractor extractor = method.getQueryExtractor();
this.countProjection = method.getCountQueryProjection();
this.queryRewriter = queryRewriter;
Parameters<?, ?> parameters = method.getParameters();
@@ -127,9 +133,10 @@ final class NamedQuery extends AbstractJpaQuery {
*
* @param method must not be {@literal null}.
* @param em must not be {@literal null}.
* @param queryRewriter must not be {@literal null}.
*/
@Nullable
public static RepositoryQuery lookupFrom(JpaQueryMethod method, EntityManager em) {
public static RepositoryQuery lookupFrom(JpaQueryMethod method, EntityManager em, QueryRewriter queryRewriter) {
String queryName = method.getNamedQueryName();
@@ -147,7 +154,7 @@ final class NamedQuery extends AbstractJpaQuery {
method.isNativeQuery() ? "NativeQuery" : "Query"));
}
RepositoryQuery query = new NamedQuery(method, em);
RepositoryQuery query = new NamedQuery(method, em, queryRewriter);
if (LOG.isDebugEnabled()) {
LOG.debug(String.format("Found named query '%s'", queryName));
}
@@ -187,6 +194,7 @@ final class NamedQuery extends AbstractJpaQuery {
} else {
String countQueryString = declaredQuery.get().deriveCountQuery(countProjection).getQueryString();
countQueryString = potentiallyRewriteQuery(countQueryString, accessor.getSort(), accessor.getPageable());
cacheKey = countQueryString;
countQuery = em.createQuery(countQueryString, Long.class);
}
@@ -222,4 +230,20 @@ final class NamedQuery extends AbstractJpaQuery {
? null //
: super.getTypeToRead(returnedType);
}
/**
* Use the {@link QueryRewriter}, potentially rewrite the query, using relevant {@link Sort} and {@link Pageable}
* information.
*
* @param originalQuery
* @param sort
* @param pageable
* @return
*/
private String potentiallyRewriteQuery(String originalQuery, Sort sort, Pageable pageable) {
return pageable.isPaged() //
? queryRewriter.rewrite(originalQuery, pageable) //
: queryRewriter.rewrite(originalQuery, sort);
}
}

22
spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/JpaQueryRewriteIntegrationTests.java
View File

@@ -19,8 +19,10 @@ import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.entry;
import java.util.HashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
@@ -31,6 +33,7 @@ import org.springframework.context.annotation.ComponentScan;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.FilterType;
import org.springframework.context.annotation.ImportResource;
import org.springframework.data.domain.Page;
import org.springframework.data.domain.PageRequest;
import org.springframework.data.domain.Pageable;
import org.springframework.data.domain.Sort;
@@ -44,7 +47,7 @@ import org.springframework.test.context.ContextConfiguration;
import org.springframework.test.context.junit.jupiter.SpringExtension;
/**
* Unit tests for repository with {@link Query} and {@link QueryRewrite}.
* Unit tests for repository with {@link Query} and {@link QueryRewriter}.
*
* @author Greg Turnquist
* @author Krzysztof Krason
@@ -60,10 +63,12 @@ class JpaQueryRewriteIntegrationTests {
static final String REWRITTEN_QUERY = "rewritten query";
static final String SORT = "sort";
static Map<String, String> results = new HashMap<>();
static Set<String> queries = new LinkedHashSet<>();
@BeforeEach
void setUp() {
results.clear();
repository.deleteAll();
}
@Test
@@ -77,15 +82,15 @@ class JpaQueryRewriteIntegrationTests {
entry(SORT, Sort.unsorted().toString()));
}
@Test
@Test // GH-3801
void nonNativeQueryShouldHandleRewrites() {
repository.findByNonNativeQuery("Matthews");
repository.save(new User("D", "A", "foo@bar"));
assertThat(results).containsExactly( //
entry(ORIGINAL_QUERY, "select original_user_alias from User original_user_alias"), //
entry(REWRITTEN_QUERY, "select rewritten_user_alias from User rewritten_user_alias"), //
entry(SORT, Sort.unsorted().toString()));
repository.findByNonNativeQuery("Matthews", PageRequest.of(0, 1));
assertThat(queries).contains("select original_user_alias from User original_user_alias");
assertThat(queries).contains("select count(original_user_alias) from User original_user_alias");
}
@Test
@@ -169,7 +174,7 @@ class JpaQueryRewriteIntegrationTests {
List<User> findByNativeQuery(String param);
@Query(value = "select original_user_alias from User original_user_alias", queryRewriter = TestQueryRewriter.class)
List<User> findByNonNativeQuery(String param);
Page<User> findByNonNativeQuery(String param, PageRequest pageRequest);
@Query(value = "select original_user_alias from User original_user_alias", queryRewriter = TestQueryRewriter.class)
List<User> findByNonNativeSortedQuery(String param, Sort sort);
@@ -214,6 +219,7 @@ class JpaQueryRewriteIntegrationTests {
results.put(ORIGINAL_QUERY, query);
results.put(REWRITTEN_QUERY, rewrittenQuery);
results.put(SORT, sort.toString());
queries.add(query);
return rewrittenQuery;
}

7
spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/NamedQueryUnitTests.java
View File

@@ -36,6 +36,7 @@ import org.mockito.quality.Strictness;
import org.springframework.data.domain.Page;
import org.springframework.data.domain.Pageable;
import org.springframework.data.jpa.provider.QueryExtractor;
import org.springframework.data.jpa.repository.QueryRewriter;
import org.springframework.data.projection.ProjectionFactory;
import org.springframework.data.projection.SpelAwareProxyProjectionFactory;
import org.springframework.data.repository.core.RepositoryMetadata;
@@ -88,7 +89,8 @@ class NamedQueryUnitTests {
JpaQueryMethod queryMethod = new JpaQueryMethod(method, metadata, projectionFactory, extractor);
when(em.createNamedQuery(queryMethod.getNamedCountQueryName())).thenThrow(new IllegalArgumentException());
assertThatExceptionOfType(QueryCreationException.class).isThrownBy(() -> NamedQuery.lookupFrom(queryMethod, em));
assertThatExceptionOfType(QueryCreationException.class)
.isThrownBy(() -> NamedQuery.lookupFrom(queryMethod, em, QueryRewriter.IdentityQueryRewriter.INSTANCE));
}
@Test // DATAJPA-142
@@ -100,7 +102,8 @@ class NamedQueryUnitTests {
TypedQuery<Long> countQuery = mock(TypedQuery.class);
when(em.createNamedQuery(eq(queryMethod.getNamedCountQueryName()), eq(Long.class))).thenReturn(countQuery);
NamedQuery query = (NamedQuery) NamedQuery.lookupFrom(queryMethod, em);
NamedQuery query = (NamedQuery) NamedQuery.lookupFrom(queryMethod, em,
QueryRewriter.IdentityQueryRewriter.INSTANCE);
query.doCreateCountQuery(new JpaParametersParameterAccessor(queryMethod.getParameters(), new Object[1]));
verify(em, times(1)).createNamedQuery(queryMethod.getNamedCountQueryName(), Long.class);

7
src/main/antora/modules/ROOT/pages/jpa/query-methods.adoc
View File

@@ -176,8 +176,11 @@ public interface UserRepository extends JpaRepository<User, Long> {
Sometimes, no matter how many features you try to apply, it seems impossible to get Spring Data JPA to apply every thing
you'd like to a query before it is sent to the `EntityManager`.
You have the ability to get your hands on the query, right before it's sent to the `EntityManager` and "rewrite" it. That is,
you can make any alterations at the last moment.
You have the ability to get your hands on the query, right before it's sent to the `EntityManager` and "rewrite" it.
That is, you can make any alterations at the last moment.
Query rewriting applies to the actual query and, when applicable, to count queries.
Count queries are optimized and therefore, either not necessary or a count is obtained through other means, such as derived from a Hibernate `SelectionQuery`.
.Declare a QueryRewriter using `@Query`
====