Doctrine_Pager extension that handles multiple queries

Multi.php

<?php
/*
 *  $Id$
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 * This software consists of voluntary contributions made by many individuals
 * and is licensed under the LGPL. For more information, see
 * <http://www.doctrine-project.org>.
 */

/**
 * Doctrine_Pager_Multi
 * This is an extension to Doctrine_Pager which lets you paginate over
 * results from multiple queries. Useful for e.g. searching from different
 * tables.
 *
 * @author      Guilherme Blanco <guilhermeblanco@hotmail.com>
 * @author      Jan Dudek <jd@jandudek.com>
 * @package     Doctrine
 * @subpackage  Pager
 * @license     http://www.opensource.org/licenses/lgpl-license.php LGPL
 */
class Doctrine_Pager_Multi extends Doctrine_Pager
{
    /**
     * @var array $_queries               Array of Doctrine_Query objects related to the pager
     */
    protected $_queries;
    
    /**
     * @var array $_countQueries          Array of Doctrine_Query objects related to the counter of pager
     */
    protected $_countQueries;

    /**
     * @var array $_countQueriesParams    Hold the params to be used by Doctrine_Query counter object of pager
     */
    protected $_countQueriesParams;

    /**
     * @var integer $_counts              Number of results found in each query
     */
    protected $_counts;

    /**
     * __construct
     *
     * @param array $queries   Array of either Doctrine_Query objects or strings 
     * @param int $page        Current page
     * @param int $maxPerPage  Maximum itens per page
     * @return void
     */
    public function __construct($queries, $page, $maxPerPage = 0)
    {
        $this->_setExecuted(false);
        
        foreach ($queries as $i => $query) {
            $this->_setQuery($i, $query);
        }
        $this->_setPage($page);

        $this->setMaxPerPage($maxPerPage);
    }

    /**
     * _initialize
     *
     * Initialize Pager object calculating number of results
     *
     * @param $params  Optional parameters to Doctrine_Query::execute
     * @return void
     */
    protected function _initialize($params = array())
    {
        foreach ($this->_queries as $i => $query) {
            $countQuery = $this->getCountQuery($i);
            $this->_counts[$i] = $countQuery->count($this->getCountQueryParams($i, $params));
        }

        $this->_setNumResults(array_sum($this->_counts));
        $this->_setExecuted(true); // _adjustOffset relies of _executed equals true = getNumResults()

        $this->_adjustOffset();
    }

    /**
     * _adjustOffset
     *
     * Adjusts last page of Doctrine_Pager
     *
     * @return void
     */
    protected function _adjustOffset()
    {
        // Define new total of pages
        $this->_setLastPage(
            max(1, ceil($this->getNumResults() / $this->getMaxPerPage()))
        );
    }

    /**
     * getQuery
     *
     * Returns the Doctrine_Query collector object related to the pager
     *
     * @param int $i             Index in $queries array
     * @return Doctrine_Query    Doctrine_Query object related to the pager
     */
    public function getQuery($i)
    {
        return $this->_queries[$i];
    }

    /**
     * _setQuery
     *
     * Defines the collector query to be used by pager
     *
     * @param int $i             Index in $queries array
     * @param Doctrine_Query     Accepts either a Doctrine_Query object or a string 
     *                           (which does the Doctrine_Query class creation).
     * @return void
     */
    protected function _setQuery($i, $query)
    {
        if (is_string($query)) {
            $query = Doctrine_Query::create()
                ->parseDqlQuery($query);
        }

        $this->_queries[$i] = $query;
    }

    /**
     * getCountQuery
     *
     * Returns the Doctrine_Query object that is used to make the count results to pager
     *
     * @param int $i              Index in $queries array
     * @return Doctrine_Query     Doctrine_Query object related to the pager
     */
    public function getCountQuery($i)
    {
        return (isset($this->_countQueries[$i]) && $this->_countQueries[$i] !== null)
            ? $this->_countQueries[$i] : $this->_queries[$i];
    }

    /**
     * setCountQuery
     *
     * Defines the counter query to be used by pager
     *
     * @param int $i          Index in $queries array
     * @param Doctrine_Query  Accepts either a Doctrine_Query object or a string 
     *                        (which does the Doctrine_Query class creation).
     * @param array           Optional params to be used by counter Doctrine_Query. 
     *                        If not defined, the params passed to execute method will be used.
     * @return void
     */
    public function setCountQuery($i, $query, $params = null)
    {
        if (is_string($query)) {
            $query = Doctrine_Query::create()
                ->parseDqlQuery($query);
        }

        $this->_countQueries[$i] = $query;

        $this->setCountQueryParams($i, $params);

        $this->_setExecuted(false);
    }

    /**
     * getCountQueryParams
     *
     * Returns the params to be used by counter Doctrine_Query
     *
     * @param int $i     Index in $queries array
     * @return array     Doctrine_Query counter params
     */
    public function getCountQueryParams($i, $defaultParams = array())
    {
        return (isset($this->_countQueriesParams[$i]) && $this->_countQueriesParams[$i] !== null)
            ? $this->_countQueriesParams[$i] : $defaultParams;
    }

    /**
     * setCountQueryParams
     *
     * Defines the params to be used by counter Doctrine_Query
     *
     * @param int $i      Index in $queries array
     * @param array       Optional params to be used by counter Doctrine_Query. 
     *                    If not defined, the params passed to execute method will be used.
     * @param boolean     Optional argument that append the query param instead of overriding the existent ones.
     * @return void
     */
    public function setCountQueryParams($i, $params = array(), $append = false)
    {
        if ($append && is_array($this->_countQueriesParams[$i])) {
            $this->_countQueriesParams[$i] = array_merge($this->_countQueriesParams[$i], $params);
        } else {
            if ($params !== null && !is_array($params)) {
                $params = array($params);
            }

            $this->_countQueriesParams[$i] = $params;
        }

        $this->_setExecuted(false);
    }

    /**
     * execute
     *
     * Executes the query, populates the collection and then return it
     *
     * @param $params               Optional parameters to Doctrine_Query::execute
     * @param $hydrationMode        Hydration Mode of Doctrine_Query::execute returned ResultSet.
     * @return array                Results (Doctrine_Records or arrays, depending on $hydrationMode)
     */
    public function execute($params = array(), $hydrationMode = null)
    {
        if ( !$this->getExecuted()) {
            $this->_initialize($params);
        }

        $results = array();
        $offset = $this->getFirstIndice() - 1; // getFirstIndice() result is 1-based
        $limit = $this->getMaxPerPage();
        foreach ($this->_queries as $i => $query) {
            // got enough results, stop querying
            if ($limit <= 0) {
                break;
            }
            // all results from the query where shown on previous pages, skip to next query
            if ($offset > $this->_counts[$i]) {
                $offset -= $this->_counts[$i];
                continue;
            }
            $query->offset($offset);
            $query->limit($limit);
            $limit -= ($this->_counts[$i] - $offset);
            $offset = 0;

            // merge() will convert any Doctrine_Collection to array of Doctrine_Records.
            // This is because Doctrine_Collection cannot contain records from different tables.
            $results = $this->merge($results, $query->execute($params, $hydrationMode));
        }
        
        return $results;
    }
    
    private function merge($first, $second)
    {
        if ($second instanceof Doctrine_Collection) {
            $second = $second->getData();
        }
        return array_merge($first, $second);
    }
}