Apply count distinct to the criteria query when specified in the pagi…

…ng count query (#2766)
micronaut-projects · Jan 30, 2024 · 47fc3cd · 47fc3cd
1 parent a7d77c0
commit 47fc3cd
Show file tree

Hide file tree

Showing 7 changed files with 46 additions and 8 deletions.
diff --git a/...jpa/src/test/groovy/io/micronaut/data/hibernate/JpaSpecificationCrudRepositorySpec.groovy b/...jpa/src/test/groovy/io/micronaut/data/hibernate/JpaSpecificationCrudRepositorySpec.groovy
@@ -102,6 +102,10 @@ class JpaSpecificationCrudRepositorySpec extends Specification {
         results.size() == 4
         results.every({ it instanceof Person})
 
+        def pageReq = Pageable.from(0, 2, Sort.of(Sort.Order.asc("age")))
+        def page = crudRepository.findAll(JpaSpecificationCrudRepository.Specifications.ageGreaterThanThirty(true), pageReq)
+        page.totalSize <= 4
+
         def sorted = crudRepository.findAll(JpaSpecificationCrudRepository.Specifications.ageGreaterThanThirty(), Sort.of(Sort.Order.asc("age")))
 
         sorted.first().name == "James"

diff --git a/...bernate-jpa/src/test/java/io/micronaut/data/hibernate/JpaSpecificationCrudRepository.java b/...bernate-jpa/src/test/java/io/micronaut/data/hibernate/JpaSpecificationCrudRepository.java
@@ -72,9 +72,18 @@ public interface JpaSpecificationCrudRepository extends CrudRepository<Person, L
 
     class Specifications {
         public static Specification<Person> ageGreaterThanThirty() {
-            return (root, query, criteriaBuilder) -> criteriaBuilder.greaterThan(
+            return ageGreaterThanThirty(false);
+        }
+
+        public static Specification<Person> ageGreaterThanThirty(boolean countDistinct) {
+            return (root, query, criteriaBuilder) -> {
+                if (countDistinct && query.getResultType() == Long.class) {
+                    query.distinct(true);
+                }
+                return criteriaBuilder.greaterThan(
                     root.get("age"), 30
-            );
+                );
+            };
         }
 
         public static Specification<Person> nameEquals(String name) {

diff --git a/...hibernate/reactive/repository/jpa/intercept/ReactiveFindPageSpecificationInterceptor.java b/...hibernate/reactive/repository/jpa/intercept/ReactiveFindPageSpecificationInterceptor.java
@@ -97,7 +97,11 @@ protected Publisher<?> interceptPublisher(RepositoryMethodKey methodKey, MethodI
                 if (countPredicate != null) {
                     countQuery.where(countPredicate);
                 }
-                countQuery.select(criteriaBuilder.count(countRoot));
+                if (countQuery.isDistinct()) {
+                    countQuery.select(criteriaBuilder.countDistinct(countRoot));
+                } else {
+                    countQuery.select(criteriaBuilder.count(countRoot));
+                }
                 return Mono.fromCompletionStage(() -> session.createQuery(countQuery).getSingleResult())
                     .map(total -> Page.of(results, pageable, total));
             });

diff --git a/...ain/java/io/micronaut/data/jpa/repository/intercept/FindPageSpecificationInterceptor.java b/...ain/java/io/micronaut/data/jpa/repository/intercept/FindPageSpecificationInterceptor.java
@@ -107,7 +107,11 @@ public Page intercept(RepositoryMethodKey methodKey, MethodInvocationContext<Obj
             if (countPredicate != null) {
                 countQuery.where(countPredicate);
             }
-            countQuery.select(criteriaBuilder.count(countRoot));
+            if (countQuery.isDistinct()) {
+                countQuery.select(criteriaBuilder.countDistinct(countRoot));
+            } else {
+                countQuery.select(criteriaBuilder.count(countRoot));
+            }
             Long singleResult = entityManager.createQuery(countQuery).getSingleResult();
 
             return Page.of(

diff --git a/...rc/main/java/io/micronaut/data/spring/jpa/intercept/FindPageSpecificationInterceptor.java b/...rc/main/java/io/micronaut/data/spring/jpa/intercept/FindPageSpecificationInterceptor.java
@@ -101,9 +101,13 @@ public Object intercept(RepositoryMethodKey methodKey, MethodInvocationContext<O
                     final List<Object> results = typedQuery.getResultList();
                     final CriteriaQuery<Long> countQuery = criteriaBuilder.createQuery(Long.class);
                     final Root<?> countRoot = countQuery.from(getRequiredRootEntity(context));
-                    final Predicate countPredicate = specification.toPredicate(countRoot, query, criteriaBuilder);
+                    final Predicate countPredicate = specification.toPredicate(countRoot, countQuery, criteriaBuilder);
                     countQuery.where(countPredicate);
-                    countQuery.select(criteriaBuilder.count(countRoot));
+                    if (countQuery.isDistinct()) {
+                        countQuery.select(criteriaBuilder.countDistinct(countRoot));
+                    } else {
+                        countQuery.select(criteriaBuilder.count(countRoot));
+                    }
 
                     return new PageImpl<>(
                             results,

diff --git a/...jpa/src/test/groovy/io/micronaut/data/spring/hibernate/SpringCrudRepositoryJpaSpec.groovy b/...jpa/src/test/groovy/io/micronaut/data/spring/hibernate/SpringCrudRepositoryJpaSpec.groovy
@@ -81,6 +81,10 @@ class SpringCrudRepositoryJpaSpec extends Specification implements H2Properties
         results.size() == 4
         results.every({ it instanceof Person})
 
+        def pageReq = PageRequest.of(0, 2, Sort.by("age"))
+        def page = crudRepository.findAll(SpringCrudRepository.Specifications.ageGreaterThanThirty(true), pageReq)
+        page.totalElements <= 4
+
         def sorted = crudRepository.findAll(SpringCrudRepository.Specifications.ageGreaterThanThirty(), Sort.by("age"))
 
         sorted.first().name == "James"

diff --git a/...ing-jpa/src/test/java/io/micronaut/data/spring/hibernate/spring/SpringCrudRepository.java b/...ing-jpa/src/test/java/io/micronaut/data/spring/hibernate/spring/SpringCrudRepository.java
@@ -73,9 +73,18 @@ public interface SpringCrudRepository extends CrudRepository<Person, Long>, JpaS
 
     class Specifications {
         public static Specification<Person> ageGreaterThanThirty() {
-            return (root, query, criteriaBuilder) -> criteriaBuilder.greaterThan(
+            return ageGreaterThanThirty(false);
+        }
+
+        public static Specification<Person> ageGreaterThanThirty(boolean countDistinct) {
+            return (root, query, criteriaBuilder) -> {
+                if (countDistinct && query.getResultType() == Long.class) {
+                    query.distinct(true);
+                }
+                return criteriaBuilder.greaterThan(
                     root.get("age"), 30
-            );
+                );
+            };
         }
 
         public static Specification<Person> nameEquals(String name) {