Code Coverage
 
Classes and Traits
Functions and Methods
Lines
Total
0.00% covered (danger)
0.00%
0 / 1
50.00% covered (danger)
50.00%
8 / 16
CRAP
85.14% covered (warning)
85.14%
63 / 74
BatchRowIterator
0.00% covered (danger)
0.00%
0 / 1
50.00% covered (danger)
50.00%
8 / 16
30.58
85.14% covered (warning)
85.14%
63 / 74
 __construct
0.00% covered (danger)
0.00%
0 / 1
2.01
88.89% covered (warning)
88.89%
8 / 9
 addConditions
100.00% covered (success)
100.00%
1 / 1
1
100.00% covered (success)
100.00%
2 / 2
 addOptions
0.00% covered (danger)
0.00%
0 / 1
2
0.00% covered (danger)
0.00%
0 / 2
 addJoinConditions
0.00% covered (danger)
0.00%
0 / 1
2
0.00% covered (danger)
0.00%
0 / 2
 setFetchColumns
100.00% covered (success)
100.00%
1 / 1
3
100.00% covered (success)
100.00%
6 / 6
 setCaller
0.00% covered (danger)
0.00%
0 / 1
2
0.00% covered (danger)
0.00%
0 / 2
 extractPrimaryKeys
100.00% covered (success)
100.00%
1 / 1
3
100.00% covered (success)
100.00%
5 / 5
 current
100.00% covered (success)
100.00%
1 / 1
1
100.00% covered (success)
100.00%
1 / 1
 key
0.00% covered (danger)
0.00%
0 / 1
2
0.00% covered (danger)
0.00%
0 / 1
 rewind
100.00% covered (success)
100.00%
1 / 1
1
100.00% covered (success)
100.00%
4 / 4
 valid
100.00% covered (success)
100.00%
1 / 1
1
100.00% covered (success)
100.00%
1 / 1
 hasChildren
0.00% covered (danger)
0.00%
0 / 1
6
0.00% covered (danger)
0.00%
0 / 1
 getChildren
0.00% covered (danger)
0.00%
0 / 1
2
0.00% covered (danger)
0.00%
0 / 1
 next
0.00% covered (danger)
0.00%
0 / 1
2.00
93.33% covered (success)
93.33%
14 / 15
 buildConditions
100.00% covered (success)
100.00%
1 / 1
5
100.00% covered (success)
100.00%
14 / 14
 buildGreaterThanCondition
100.00% covered (success)
100.00%
1 / 1
2
100.00% covered (success)
100.00%
8 / 8
<?php
use Wikimedia\Rdbms\IDatabase;
/**
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License along
 * with this program; if not, write to the Free Software Foundation, Inc.,
 * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 * http://www.gnu.org/copyleft/gpl.html
 *
 * @file
 * @ingroup Maintenance
 */
/**
 * Allows iterating a large number of rows in batches transparently.
 * By default when iterated over returns the full query result as an
 * array of rows.  Can be wrapped in RecursiveIteratorIterator to
 * collapse those arrays into a single stream of rows queried in batches.
 *
 * @newable
 */
class BatchRowIterator implements RecursiveIterator {
    /**
     * @var IDatabase The database to read from
     */
    protected $db;
    /**
     * @var string|array The name or names of the table to read from
     */
    protected $table;
    /**
     * @var array The name of the primary key(s)
     */
    protected $primaryKey;
    /**
     * @var int The number of rows to fetch per iteration
     */
    protected $batchSize;
    /**
     * @var array Array of strings containing SQL conditions to add to the query
     */
    protected $conditions = [];
    /**
     * @var array
     */
    protected $joinConditions = [];
    /**
     * @var array List of column names to select from the table suitable for use
     *  with IDatabase::select()
     */
    protected $fetchColumns;
    /**
     * @var string SQL Order by condition generated from $this->primaryKey
     */
    protected $orderBy;
    /**
     * @var array The current iterator value
     */
    private $current = [];
    /**
     * @var int 0-indexed number of pages fetched since self::reset()
     */
    private $key;
    /**
     * @var array Additional query options
     */
    protected $options = [];
    /**
     * @var string|null For debugging which method is using this class.
     */
    protected $caller;
    /**
     * @stable to call
     *
     * @param IDatabase $db The database to read from
     * @param string|array $table The name or names of the table to read from
     * @param string|array $primaryKey The name or names of the primary key columns
     * @param int $batchSize The number of rows to fetch per iteration
     * @throws InvalidArgumentException
     */
    public function __construct( IDatabase $db, $table, $primaryKey, $batchSize ) {
        if ( $batchSize < 1 ) {
            throw new InvalidArgumentException( 'Batch size must be at least 1 row.' );
        }
        $this->db = $db;
        $this->table = $table;
        $this->primaryKey = (array)$primaryKey;
        $this->fetchColumns = $this->primaryKey;
        $this->orderBy = implode( ' ASC,', $this->primaryKey ) . ' ASC';
        $this->batchSize = $batchSize;
    }
    /**
     * @param array $conditions Query conditions suitable for use with
     *  IDatabase::select
     */
    public function addConditions( array $conditions ) {
        $this->conditions = array_merge( $this->conditions, $conditions );
    }
    /**
     * @param array $options Query options suitable for use with
     *  IDatabase::select
     */
    public function addOptions( array $options ) {
        $this->options = array_merge( $this->options, $options );
    }
    /**
     * @param array $conditions Query join conditions suitable for use
     *  with IDatabase::select
     */
    public function addJoinConditions( array $conditions ) {
        $this->joinConditions = array_merge( $this->joinConditions, $conditions );
    }
    /**
     * @param array $columns List of column names to select from the
     *  table suitable for use with IDatabase::select()
     */
    public function setFetchColumns( array $columns ) {
        // If it's not the all column selector merge in the primary keys we need
        if ( count( $columns ) === 1 && reset( $columns ) === '*' ) {
            $this->fetchColumns = $columns;
        } else {
            $this->fetchColumns = array_unique( array_merge(
                $this->primaryKey,
                $columns
            ) );
        }
    }
    /**
     * Use ->setCaller( __METHOD__ ) to indicate which code is using this
     * class. Only used in debugging output.
     * @since 1.36
     *
     * @param string $caller
     * @return self
     */
    public function setCaller( $caller ) {
        $this->caller = $caller;
        return $this;
    }
    /**
     * Extracts the primary key(s) from a database row.
     *
     * @param stdClass $row An individual database row from this iterator
     * @return array Map of primary key column to value within the row
     */
    public function extractPrimaryKeys( $row ) {
        $pk = [];
        foreach ( $this->primaryKey as $alias => $column ) {
            $name = is_numeric( $alias ) ? $column : $alias;
            $pk[$name] = $row->{$name};
        }
        return $pk;
    }
    /**
     * @return array The most recently fetched set of rows from the database
     */
    public function current() {
        return $this->current;
    }
    /**
     * @return int 0-indexed count of the page number fetched
     */
    public function key() {
        return $this->key;
    }
    /**
     * Reset the iterator to the begining of the table.
     */
    public function rewind() {
        $this->key = -1; // self::next() will turn this into 0
        $this->current = [];
        $this->next();
    }
    /**
     * @return bool True when the iterator is in a valid state
     */
    public function valid() {
        return (bool)$this->current;
    }
    /**
     * @return bool True when this result set has rows
     */
    public function hasChildren() {
        return $this->current && count( $this->current );
    }
    /**
     * @return RecursiveIterator
     */
    public function getChildren() {
        return new NotRecursiveIterator( new ArrayIterator( $this->current ) );
    }
    /**
     * Fetch the next set of rows from the database.
     */
    public function next() {
        $caller = __METHOD__;
        if ( (string)$this->caller !== '' ) {
            $caller .= " (for {$this->caller})";
        }
        $res = $this->db->select(
            $this->table,
            $this->fetchColumns,
            $this->buildConditions(),
            $caller,
            [
                'LIMIT' => $this->batchSize,
                'ORDER BY' => $this->orderBy,
            ] + $this->options,
            $this->joinConditions
        );
        // The iterator is converted to an array because in addition to
        // returning it in self::current() we need to use the end value
        // in self::buildConditions()
        $this->current = iterator_to_array( $res );
        $this->key++;
    }
    /**
     * Uses the primary key list and the maximal result row from the
     * previous iteration to build an SQL condition sufficient for
     * selecting the next page of results.  All except the final key use
     * `=` conditions while the final key uses a `>` condition
     *
     * Example output:
     *     [ '( foo = 42 AND bar > 7 ) OR ( foo > 42 )' ]
     *
     * @return array The SQL conditions necessary to select the next set
     *  of rows in the batched query
     */
    protected function buildConditions() {
        if ( !$this->current ) {
            return $this->conditions;
        }
        $maxRow = end( $this->current );
        $maximumValues = [];
        foreach ( $this->primaryKey as $alias => $column ) {
            $name = is_numeric( $alias ) ? $column : $alias;
            $maximumValues[$column] = $this->db->addQuotes( $maxRow->{$name} );
        }
        $pkConditions = [];
        // For example: If we have 3 primary keys
        // first run through will generate
        //   col1 = 4 AND col2 = 7 AND col3 > 1
        // second run through will generate
        //   col1 = 4 AND col2 > 7
        // and the final run through will generate
        //   col1 > 4
        while ( $maximumValues ) {
            $pkConditions[] = $this->buildGreaterThanCondition( $maximumValues );
            array_pop( $maximumValues );
        }
        $conditions = $this->conditions;
        $conditions[] = sprintf( '( %s )', implode( ' ) OR ( ', $pkConditions ) );
        return $conditions;
    }
    /**
     * Given an array of column names and their maximum value  generate
     * an SQL condition where all keys except the last match $quotedMaximumValues
     * exactly and the last column is greater than the matching value in
     * $quotedMaximumValues
     *
     * @param array $quotedMaximumValues The maximum values quoted with
     *  $this->db->addQuotes()
     * @return string An SQL condition that will select rows where all
     *  columns match the maximum value exactly except the last column
     *  which must be greater than the provided maximum value
     */
    protected function buildGreaterThanCondition( array $quotedMaximumValues ) {
        $keys = array_keys( $quotedMaximumValues );
        $lastColumn = end( $keys );
        $lastValue = array_pop( $quotedMaximumValues );
        $conditions = [];
        foreach ( $quotedMaximumValues as $column => $value ) {
            $conditions[] = "$column = $value";
        }
        $conditions[] = "$lastColumn > $lastValue";
        return implode( ' AND ', $conditions );
    }
}