From 3e1f190b45ee9c848dc65be5d5ef9ecadf5cffba Mon Sep 17 00:00:00 2001
From: Massimiliano Arione <garakkio@gmail.com>
Date: Fri, 24 Jul 2015 16:59:40 +0200
Subject: [PATCH] add optional QueryBuilder

---
 src/DoctrineReader.php | 32 ++++++++++++++++++++++++--------
 1 file changed, 24 insertions(+), 8 deletions(-)

diff --git a/src/DoctrineReader.php b/src/DoctrineReader.php
index 943b80e..48d419c 100644
--- a/src/DoctrineReader.php
+++ b/src/DoctrineReader.php
@@ -5,6 +5,8 @@
 use Doctrine\Common\Persistence\ObjectManager;
 use Doctrine\ORM\Internal\Hydration\IterableResult;
 use Doctrine\ORM\Query;
+use Doctrine\ORM\QueryBuilder;
+use Doctrine\ORM\Tools\Pagination\Paginator;
 use Port\Reader\CountableReader;
 
 /**
@@ -29,14 +31,32 @@ class DoctrineReader implements CountableReader
      */
     protected $iterableResult;
 
+    /**
+     * @var QueryBuilder
+     */
+    protected $queryBuilder;
+
     /**
      * @param ObjectManager $objectManager
      * @param string        $objectName    e.g. YourBundle:YourEntity
+     * @param QueryBuilder  $queryBuilder
      */
-    public function __construct(ObjectManager $objectManager, $objectName)
+    public function __construct(ObjectManager $objectManager, $objectName, QueryBuilder $queryBuilder = null)
     {
         $this->objectManager = $objectManager;
         $this->objectName = $objectName;
+        if (is_null($queryBuilder)) {
+            $queryBuilder = $objectManager->getRepository($objectName)->createQueryBuilder(substr($objectName, 0, 1));
+        }
+        $this->queryBuilder = $queryBuilder;
+    }
+
+    /**
+     * @param QueryBuilder $queryBuilder
+     */
+    public function setQueryBuilder(QueryBuilder $queryBuilder)
+    {
+        $this->queryBuilder = $queryBuilder;
     }
 
     /**
@@ -86,9 +106,7 @@ public function valid()
     public function rewind()
     {
         if (!$this->iterableResult) {
-            $query = $this->objectManager->createQuery(
-                sprintf('SELECT o FROM %s o', $this->objectName)
-            );
+            $query = $this->queryBuilder->getQuery();
             $this->iterableResult = $query->iterate([], Query::HYDRATE_ARRAY);
         }
 
@@ -100,10 +118,8 @@ public function rewind()
      */
     public function count()
     {
-        $query = $this->objectManager->createQuery(
-            sprintf('SELECT COUNT(o) FROM %s o', $this->objectName)
-        );
+        $paginator = new Paginator($this->queryBuilder->getQuery());
 
-        return $query->getSingleScalarResult();
+        return count($paginator);
     }
 }