--- /dev/null
+<?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 );
+ }
+}
scripts.mit.edu / autoinstallsdev / mediawiki.git / blobdiff