From cd4013fe1d845cab1ddd4859e8413b09ad589700 Mon Sep 17 00:00:00 2001
From: Oliver Gierke <info@olivergierke.de>
Date: Mon, 14 Feb 2011 21:44:03 +0100
Subject: [PATCH] Fixed bug in total page number calculation due to additional
 root adding.

Using the initially created root to create the count projection now. Added more tests to verify correct number of elements and pages.
---
 .../jpa/repository/query/JpaCountQueryCreator.java | 14 ++++++--------
 .../data/jpa/repository/query/JpaQueryCreator.java |  4 ++--
 .../jpa/repository/query/PartTreeJpaQuery.java     |  4 ++--
 .../jpa/repository/UserRepositoryFinderTests.java  |  4 +++-
 4 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/src/main/java/org/springframework/data/jpa/repository/query/JpaCountQueryCreator.java b/src/main/java/org/springframework/data/jpa/repository/query/JpaCountQueryCreator.java
index eae9b6415..7728e84e6 100644
--- a/src/main/java/org/springframework/data/jpa/repository/query/JpaCountQueryCreator.java
+++ b/src/main/java/org/springframework/data/jpa/repository/query/JpaCountQueryCreator.java
@@ -19,6 +19,7 @@ import javax.persistence.EntityManager;
 import javax.persistence.criteria.CriteriaBuilder;
 import javax.persistence.criteria.CriteriaQuery;
 import javax.persistence.criteria.Predicate;
+import javax.persistence.criteria.Root;
 
 import org.springframework.data.domain.Sort;
 import org.springframework.data.repository.query.SimpleParameterAccessor;
@@ -32,9 +33,6 @@ import org.springframework.data.repository.query.parser.PartTree;
  */
 public class JpaCountQueryCreator extends JpaQueryCreator {
 
-    private final Class<?> domainClass;
-
-
     /**
      * Creates a new {@link JpaCountQueryCreator}.
      * 
@@ -48,7 +46,6 @@ public class JpaCountQueryCreator extends JpaQueryCreator {
             EntityManager em) {
 
         super(tree, parameters, domainClass, em);
-        this.domainClass = domainClass;
     }
 
 
@@ -56,16 +53,17 @@ public class JpaCountQueryCreator extends JpaQueryCreator {
      * (non-Javadoc)
      * 
      * @see
-     * org.springframework.data.jpa.repository.query.JpaQueryCreator#finalize
+     * org.springframework.data.jpa.repository.query.JpaQueryCreator#complete
      * (javax.persistence.criteria.Predicate,
      * org.springframework.data.domain.Sort,
      * javax.persistence.criteria.CriteriaQuery,
-     * javax.persistence.criteria.CriteriaBuilder)
+     * javax.persistence.criteria.CriteriaBuilder,
+     * javax.persistence.criteria.Root)
      */
     @Override
     protected CriteriaQuery<Object> complete(Predicate predicate, Sort sort,
-            CriteriaQuery<Object> query, CriteriaBuilder builder) {
+            CriteriaQuery<Object> query, CriteriaBuilder builder, Root<?> root) {
 
-        return query.select(builder.count(query.from(domainClass)));
+        return query.select(builder.count(root));
     }
 }
diff --git a/src/main/java/org/springframework/data/jpa/repository/query/JpaQueryCreator.java b/src/main/java/org/springframework/data/jpa/repository/query/JpaQueryCreator.java
index 591c1b7c9..75fefff2b 100644
--- a/src/main/java/org/springframework/data/jpa/repository/query/JpaQueryCreator.java
+++ b/src/main/java/org/springframework/data/jpa/repository/query/JpaQueryCreator.java
@@ -125,7 +125,7 @@ public class JpaQueryCreator extends
     protected final CriteriaQuery<Object> complete(Predicate predicate,
             Sort sort) {
 
-        return complete(predicate, sort, query, builder);
+        return complete(predicate, sort, query, builder, root);
     }
 
 
@@ -140,7 +140,7 @@ public class JpaQueryCreator extends
      * @return
      */
     protected CriteriaQuery<Object> complete(Predicate predicate, Sort sort,
-            CriteriaQuery<Object> query, CriteriaBuilder builder) {
+            CriteriaQuery<Object> query, CriteriaBuilder builder, Root<?> root) {
 
         return this.query.select(root).where(predicate)
                 .orderBy(QueryUtils.toOrders(sort, root, builder));
diff --git a/src/main/java/org/springframework/data/jpa/repository/query/PartTreeJpaQuery.java b/src/main/java/org/springframework/data/jpa/repository/query/PartTreeJpaQuery.java
index 99947a945..ee0a61c85 100644
--- a/src/main/java/org/springframework/data/jpa/repository/query/PartTreeJpaQuery.java
+++ b/src/main/java/org/springframework/data/jpa/repository/query/PartTreeJpaQuery.java
@@ -89,11 +89,11 @@ public class PartTreeJpaQuery extends AbstractJpaQuery {
      */
     public Query createCountQuery(Object[] parameters) {
 
-        CriteriaQuery<Object> createQuery =
+        CriteriaQuery<Object> query =
                 new JpaCountQueryCreator(tree, new SimpleParameterAccessor(
                         getParameters(), parameters), domainClass,
                         getEntityManager()).createQuery();
-        return getEntityManager().createQuery(createQuery);
+        return getEntityManager().createQuery(query);
     }
 
 
diff --git a/src/test/java/org/springframework/data/jpa/repository/UserRepositoryFinderTests.java b/src/test/java/org/springframework/data/jpa/repository/UserRepositoryFinderTests.java
index 57d3749cd..93d3bdcce 100644
--- a/src/test/java/org/springframework/data/jpa/repository/UserRepositoryFinderTests.java
+++ b/src/test/java/org/springframework/data/jpa/repository/UserRepositoryFinderTests.java
@@ -119,7 +119,9 @@ public class UserRepositoryFinderTests {
 
         Page<User> page =
                 userRepository.findByFirstname(new PageRequest(0, 1), "foobar");
-        assertEquals(1, page.getNumberOfElements());
+        assertThat(page.getNumberOfElements(), is(1));
+        assertThat(page.getTotalElements(), is(2L));
+        assertThat(page.getTotalPages(), is(2));
     }