Merge pull request #91 from bburnichon/ticket_66_dbal_memory

Rework DbalReader to reduce memory usage.

src/Ddeboer/DataImport/Reader/DbalReader.php

@@ -2,27 +2,33 @@
 namespace Ddeboer\DataImport\Reader;
 
 use Doctrine\DBAL\Connection;
+use Doctrine\DBAL\Statement;
 
 /**
  * Reads data through the Doctrine DBAL
  */
 class DbalReader implements ReaderInterface
 {
-    /***
-     * @var \Doctrine\DBAL\Connection
-     */
+    /** @var Connection */
     private $connection;
 
-    /**
-     * @var array
-     */
     private $data;
 
-    /**
-     * @var \Doctrine\DBAL\Driver\Statement
-     */
+    /** @var Statement */
     private $stmt;
 
+    /** @var string */
+    private $sql;
+    /** @var array */
+    private $params;
+
+    /** @var integer */
+    private $rowCount;
+    /** @var bool */
+    private $rowCountCalculated = true;
+
+    private $key;
+
     /**
      * Constructor
      *
@@ -33,84 +39,181 @@ class DbalReader implements ReaderInterface
     public function __construct(Connection $connection, $sql, array $params = array())
     {
         $this->connection = $connection;
-        $this->stmt = $this->connection->prepare($sql);
 
-        foreach ($params as $key => $value) {
-            $this->stmt->bindValue($key, $value);
-        }
+        $this->setSql($sql, $params);
+    }
+
+    /**
+     * Do calculate row count?
+     *
+     * @param bool $calculate
+     *
+     * @return $this
+     */
+    public function setRowCountCalculated($calculate = true)
+    {
+        $this->rowCountCalculated = (bool) $calculate;
+
+        return $this;
+    }
+
+    /**
+     * Is row count calculated?
+     *
+     * @return bool
+     */
+    public function isRowCountCalculated()
+    {
+        return $this->rowCountCalculated;
     }
 
     /**
      * {@inheritdoc}
      */
     public function getFields()
     {
-        $this->stmt->execute();
+        if (null === $this->data) {
+            $this->rewind();
+        }
+        if (false === $this->data) {
+            return array();
+        }
+
+        return array_keys((array) $this->data);
+    }
+
+    /**
+     * Set Query string with Parameters
+     *
+     * @param string $sql
+     * @param array  $params
+     *
+     * @return $this
+     */
+    public function setSql($sql, array $params = array())
+    {
+        $this->sql = (string) $sql;
+
+        $this->setSqlParameters($params);
+
+        return $this;
+    }
+
+    /**
+     * Set SQL parameters
+     *
+     * @param array  $params
+     *
+     * @return $this
+     */
+    public function setSqlParameters(array $params)
+    {
+        $this->params = $params;
 
-        return array_keys($this->stmt->fetch(\PDO::FETCH_ASSOC));
+        $this->stmt = null;
+        $this->rowCount = null;
+
+        return $this;
     }
 
     /**
      * {@inheritdoc}
      */
     public function current()
     {
-        return current($this->data);
+        if (null === $this->data) {
+            $this->rewind();
+        }
+
+        return $this->data;
     }
 
     /**
      * {@inheritdoc}
      */
     public function next()
     {
-        next($this->data);
+        $this->key++;
+        $this->data = $this->stmt->fetch(\PDO::FETCH_ASSOC);
     }
 
     /**
      * {@inheritdoc}
      */
     public function key()
     {
-        return key($this->data);
+        return $this->key;
     }
 
     /**
      * {@inheritdoc}
      */
     public function valid()
     {
-        $key = key($this->data);
+        if (null === $this->data) {
+            $this->rewind();
+        }
 
-        return ($key !== null && $key !== false);
+        return (false !== $this->data);
     }
 
     /**
      * {@inheritdoc}
      */
     public function rewind()
     {
-        $this->loadData();
-        reset($this->data);
+        if (null === $this->stmt) {
+            $this->stmt = $this->prepare($this->sql, $this->params);
+        }
+        if (0 !== $this->key) {
+            $this->stmt->execute();
+            $this->data = $this->stmt->fetch(\PDO::FETCH_ASSOC);
+            $this->key = 0;
+        }
     }
 
     /**
      * {@inheritdoc}
      */
     public function count()
     {
-        $this->loadData();
+        if (null === $this->rowCount) {
+            if ($this->rowCountCalculated) {
+                $this->doCalcRowCount();
+            } else {
+                if (null === $this->stmt) {
+                    $this->rewind();
+                }
+                $this->rowCount = $this->stmt->rowCount();
+            }
+        }
+
+        return $this->rowCount;
+    }
+
+    private function doCalcRowCount()
+    {
+        $statement = $this->prepare('SELECT COUNT(*) FROM ('.$this->sql.')', $this->params);
+        $statement->execute();
 
-        return count($this->data);
+        $this->rowCount = (int) $statement->fetchColumn(0);
     }
 
     /**
-     * Load data if it hasn't been loaded yet
+     * Prepare given statement
+     *
+     * @param string $sql
+     * @param array  $params
+     *
+     * @return Statement
      */
-    protected function loadData()
+    private function prepare($sql, array $params)
     {
-        if (null === $this->data) {
-            $this->stmt->execute();
-            $this->data = $this->stmt->fetchAll(\PDO::FETCH_ASSOC);
+        $statement = $this->connection->prepare($sql);
+        foreach ($params as $key => $value) {
+            $statement->bindValue($key, $value);
         }
+
+        return $statement;
     }
 }