diff options
Diffstat (limited to 'www/wiki/includes/utils/BatchRowIterator.php')
-rw-r--r-- | www/wiki/includes/utils/BatchRowIterator.php | 296 |
1 files changed, 296 insertions, 0 deletions
diff --git a/www/wiki/includes/utils/BatchRowIterator.php b/www/wiki/includes/utils/BatchRowIterator.php new file mode 100644 index 00000000..60720c87 --- /dev/null +++ b/www/wiki/includes/utils/BatchRowIterator.php @@ -0,0 +1,296 @@ +<?php + +use Wikimedia\Rdbms\IDatabase; + +/** + * 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. + * + * 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 + */ +class BatchRowIterator implements RecursiveIterator { + + /** + * @var IDatabase $db The database to read from + */ + protected $db; + + /** + * @var string|array $table The name or names of the table to read from + */ + protected $table; + + /** + * @var array $primaryKey The name of the primary key(s) + */ + protected $primaryKey; + + /** + * @var int $batchSize The number of rows to fetch per iteration + */ + protected $batchSize; + + /** + * @var array $conditions Array of strings containing SQL conditions + * to add to the query + */ + protected $conditions = []; + + /** + * @var array $joinConditions + */ + protected $joinConditions = []; + + /** + * @var array $fetchColumns List of column names to select from the + * table suitable for use with IDatabase::select() + */ + protected $fetchColumns; + + /** + * @var string $orderBy SQL Order by condition generated from $this->primaryKey + */ + protected $orderBy; + + /** + * @var array $current The current iterator value + */ + private $current = []; + + /** + * @var int key 0-indexed number of pages fetched since self::reset() + */ + private $key; + + /** + * @var array Additional query options + */ + protected $options = []; + + /** + * @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 + ) ); + } + } + + /** + * 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() { + $res = $this->db->select( + $this->table, + $this->fetchColumns, + $this->buildConditions(), + __METHOD__, + [ + '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 ); + } +} |