From d73a86de7be55589ee3c064b37e781353e0bba28 Mon Sep 17 00:00:00 2001 From: Yanming Zhou Date: Mon, 29 Apr 2024 15:30:52 +0800 Subject: [PATCH] Use provider built-in result count to reuse query if possible Hibernate introduce `SelectionQuery::getResultCount` since 6.5.0, It may be JPA Query method in the future. --- .../repository/query/AbstractJpaQuery.java | 3 ++- .../repository/query/JpaQueryExecution.java | 23 +++++++++++++++++-- .../query/JpaQueryExecutionUnitTests.java | 13 ++++++----- 3 files changed, 30 insertions(+), 9 deletions(-) diff --git a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/AbstractJpaQuery.java b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/AbstractJpaQuery.java index 205c585025..236ac9f59a 100644 --- a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/AbstractJpaQuery.java +++ b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/AbstractJpaQuery.java @@ -62,6 +62,7 @@ * @author Сергей Цыпанов * @author Wonchul Heo * @author Julia Lee + * @author Yanming Zhou */ public abstract class AbstractJpaQuery implements RepositoryQuery { @@ -99,7 +100,7 @@ public AbstractJpaQuery(JpaQueryMethod method, EntityManager em) { } else if (method.isSliceQuery()) { return new SlicedExecution(); } else if (method.isPageQuery()) { - return new PagedExecution(); + return new PagedExecution(em); } else if (method.isModifyingQuery()) { return null; } else { diff --git a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/JpaQueryExecution.java b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/JpaQueryExecution.java index 59f161293d..9c7f751342 100644 --- a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/JpaQueryExecution.java +++ b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/JpaQueryExecution.java @@ -24,7 +24,9 @@ import java.util.Collection; import java.util.List; import java.util.Optional; +import java.util.function.LongSupplier; +import org.hibernate.query.SelectionQuery; import org.springframework.core.convert.ConversionService; import org.springframework.core.convert.support.ConfigurableConversionService; import org.springframework.core.convert.support.DefaultConversionService; @@ -57,6 +59,7 @@ * @author Jens Schauder * @author Gabriel Basilio * @author Greg Turnquist + * @author Yanming Zhou */ public abstract class JpaQueryExecution { @@ -195,14 +198,30 @@ protected Object doExecute(AbstractJpaQuery query, JpaParametersParameterAccesso */ static class PagedExecution extends JpaQueryExecution { + private final EntityManager em; + + public PagedExecution(EntityManager em) { + + Assert.notNull(em, "The EntityManager must not be null"); + this.em = em; + } + @Override @SuppressWarnings("unchecked") protected Object doExecute(AbstractJpaQuery repositoryQuery, JpaParametersParameterAccessor accessor) { Query query = repositoryQuery.createQuery(accessor); - return PageableExecutionUtils.getPage(query.getResultList(), accessor.getPageable(), - () -> count(repositoryQuery, accessor)); + LongSupplier count = () -> { + boolean userDefinedCountQuery = (repositoryQuery instanceof AbstractStringBasedJpaQuery); + if (!userDefinedCountQuery && PersistenceProvider.fromEntityManager(em) == PersistenceProvider.HIBERNATE) { + // Hibernate introduce SelectionQuery::getResultCount since 6.5.0 + return ((SelectionQuery) query).getResultCount(); + } + return count(repositoryQuery, accessor); + }; + + return PageableExecutionUtils.getPage(query.getResultList(), accessor.getPageable(), count); } private long count(AbstractJpaQuery repositoryQuery, JpaParametersParameterAccessor accessor) { diff --git a/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/JpaQueryExecutionUnitTests.java b/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/JpaQueryExecutionUnitTests.java index 078cf1249c..65f40cf17e 100644 --- a/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/JpaQueryExecutionUnitTests.java +++ b/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/JpaQueryExecutionUnitTests.java @@ -76,6 +76,7 @@ public static void sampleMethod(Pageable pageable) {} @BeforeEach void setUp() { + when(em.getDelegate()).thenReturn(em); // for PersistenceProvider.fromEntityManager(em) when(query.executeUpdate()).thenReturn(0); when(jpaQuery.createQuery(Mockito.any(JpaParametersParameterAccessor.class))).thenReturn(query); when(jpaQuery.getQueryMethod()).thenReturn(method); @@ -183,7 +184,7 @@ void pagedExecutionRetrievesObjectsForPageableOutOfRange() throws Exception { when(jpaQuery.createQuery(Mockito.any())).thenReturn(query); when(countQuery.getResultList()).thenReturn(Arrays.asList(20L)); - PagedExecution execution = new PagedExecution(); + PagedExecution execution = new PagedExecution(em); execution.doExecute(jpaQuery, new JpaParametersParameterAccessor(parameters, new Object[] { PageRequest.of(2, 10) })); @@ -199,7 +200,7 @@ void pagedExecutionShouldNotGenerateCountQueryIfQueryReportedNoResults() throws when(jpaQuery.createQuery(Mockito.any())).thenReturn(query); when(query.getResultList()).thenReturn(Arrays.asList(0L)); - PagedExecution execution = new PagedExecution(); + PagedExecution execution = new PagedExecution(em); execution.doExecute(jpaQuery, new JpaParametersParameterAccessor(parameters, new Object[] { PageRequest.of(0, 10) })); @@ -215,7 +216,7 @@ void pagedExecutionShouldUseCountFromResultIfOffsetIsZeroAndResultsWithinPageSiz when(jpaQuery.createQuery(Mockito.any())).thenReturn(query); when(query.getResultList()).thenReturn(Arrays.asList(new Object(), new Object(), new Object(), new Object())); - PagedExecution execution = new PagedExecution(); + PagedExecution execution = new PagedExecution(em); execution.doExecute(jpaQuery, new JpaParametersParameterAccessor(parameters, new Object[] { PageRequest.of(0, 10) })); @@ -230,7 +231,7 @@ void pagedExecutionShouldUseCountFromResultWithOffsetAndResultsWithinPageSize() when(jpaQuery.createQuery(Mockito.any())).thenReturn(query); when(query.getResultList()).thenReturn(Arrays.asList(new Object(), new Object(), new Object(), new Object())); - PagedExecution execution = new PagedExecution(); + PagedExecution execution = new PagedExecution(em); execution.doExecute(jpaQuery, new JpaParametersParameterAccessor(parameters, new Object[] { PageRequest.of(5, 10) })); @@ -247,7 +248,7 @@ void pagedExecutionShouldUseRequestCountFromResultWithOffsetAndResultsHitLowerPa when(jpaQuery.createCountQuery(Mockito.any())).thenReturn(query); when(countQuery.getResultList()).thenReturn(Arrays.asList(20L)); - PagedExecution execution = new PagedExecution(); + PagedExecution execution = new PagedExecution(em); execution.doExecute(jpaQuery, new JpaParametersParameterAccessor(parameters, new Object[] { PageRequest.of(4, 4) })); @@ -264,7 +265,7 @@ void pagedExecutionShouldUseRequestCountFromResultWithOffsetAndResultsHitUpperPa when(jpaQuery.createCountQuery(Mockito.any())).thenReturn(query); when(countQuery.getResultList()).thenReturn(Arrays.asList(20L)); - PagedExecution execution = new PagedExecution(); + PagedExecution execution = new PagedExecution(em); execution.doExecute(jpaQuery, new JpaParametersParameterAccessor(parameters, new Object[] { PageRequest.of(4, 4) }));