3 use Wikimedia\Rdbms\IDatabase;
6 * Allows iterating a large number of rows in batches transparently.
7 * By default when iterated over returns the full query result as an
8 * array of rows. Can be wrapped in RecursiveIteratorIterator to
9 * collapse those arrays into a single stream of rows queried in batches.
11 * This program is free software; you can redistribute it and/or modify
12 * it under the terms of the GNU General Public License as published by
13 * the Free Software Foundation; either version 2 of the License, or
14 * (at your option) any later version.
16 * This program is distributed in the hope that it will be useful,
17 * but WITHOUT ANY WARRANTY; without even the implied warranty of
18 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
19 * GNU General Public License for more details.
21 * You should have received a copy of the GNU General Public License along
22 * with this program; if not, write to the Free Software Foundation, Inc.,
23 * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
24 * http://www.gnu.org/copyleft/gpl.html
27 * @ingroup Maintenance
29 class BatchRowIterator implements RecursiveIterator {
32 * @var IDatabase $db The database to read from
37 * @var string|array $table The name or names of the table to read from
42 * @var array $primaryKey The name of the primary key(s)
44 protected $primaryKey;
47 * @var int $batchSize The number of rows to fetch per iteration
52 * @var array $conditions Array of strings containing SQL conditions
55 protected $conditions = [];
58 * @var array $joinConditions
60 protected $joinConditions = [];
63 * @var array $fetchColumns List of column names to select from the
64 * table suitable for use with IDatabase::select()
66 protected $fetchColumns;
69 * @var string $orderBy SQL Order by condition generated from $this->primaryKey
74 * @var array $current The current iterator value
76 private $current = [];
79 * @var int key 0-indexed number of pages fetched since self::reset()
84 * @var array Additional query options
86 protected $options = [];
89 * @param IDatabase $db The database to read from
90 * @param string|array $table The name or names of the table to read from
91 * @param string|array $primaryKey The name or names of the primary key columns
92 * @param int $batchSize The number of rows to fetch per iteration
93 * @throws InvalidArgumentException
95 public function __construct( IDatabase $db, $table, $primaryKey, $batchSize ) {
96 if ( $batchSize < 1 ) {
97 throw new InvalidArgumentException( 'Batch size must be at least 1 row.' );
100 $this->table = $table;
101 $this->primaryKey = (array)$primaryKey;
102 $this->fetchColumns = $this->primaryKey;
103 $this->orderBy = implode( ' ASC,', $this->primaryKey ) . ' ASC';
104 $this->batchSize = $batchSize;
108 * @param array $conditions Query conditions suitable for use with
111 public function addConditions( array $conditions ) {
112 $this->conditions = array_merge( $this->conditions, $conditions );
116 * @param array $options Query options suitable for use with
119 public function addOptions( array $options ) {
120 $this->options = array_merge( $this->options, $options );
124 * @param array $conditions Query join conditions suitable for use
125 * with IDatabase::select
127 public function addJoinConditions( array $conditions ) {
128 $this->joinConditions = array_merge( $this->joinConditions, $conditions );
132 * @param array $columns List of column names to select from the
133 * table suitable for use with IDatabase::select()
135 public function setFetchColumns( array $columns ) {
136 // If it's not the all column selector merge in the primary keys we need
137 if ( count( $columns ) === 1 && reset( $columns ) === '*' ) {
138 $this->fetchColumns = $columns;
140 $this->fetchColumns = array_unique( array_merge(
148 * Extracts the primary key(s) from a database row.
150 * @param stdClass $row An individual database row from this iterator
151 * @return array Map of primary key column to value within the row
153 public function extractPrimaryKeys( $row ) {
155 foreach ( $this->primaryKey as $alias => $column ) {
156 $name = is_numeric( $alias ) ? $column : $alias;
157 $pk[$name] = $row->{$name};
163 * @return array The most recently fetched set of rows from the database
165 public function current() {
166 return $this->current;
170 * @return int 0-indexed count of the page number fetched
172 public function key() {
177 * Reset the iterator to the begining of the table.
179 public function rewind() {
180 $this->key = -1; // self::next() will turn this into 0
186 * @return bool True when the iterator is in a valid state
188 public function valid() {
189 return (bool)$this->current;
193 * @return bool True when this result set has rows
195 public function hasChildren() {
196 return $this->current && count( $this->current );
200 * @return RecursiveIterator
202 public function getChildren() {
203 return new NotRecursiveIterator( new ArrayIterator( $this->current ) );
207 * Fetch the next set of rows from the database.
209 public function next() {
210 $res = $this->db->select(
213 $this->buildConditions(),
216 'LIMIT' => $this->batchSize,
217 'ORDER BY' => $this->orderBy,
219 $this->joinConditions
222 // The iterator is converted to an array because in addition to
223 // returning it in self::current() we need to use the end value
224 // in self::buildConditions()
225 $this->current = iterator_to_array( $res );
230 * Uses the primary key list and the maximal result row from the
231 * previous iteration to build an SQL condition sufficient for
232 * selecting the next page of results. All except the final key use
233 * `=` conditions while the final key uses a `>` condition
236 * [ '( foo = 42 AND bar > 7 ) OR ( foo > 42 )' ]
238 * @return array The SQL conditions necessary to select the next set
239 * of rows in the batched query
241 protected function buildConditions() {
242 if ( !$this->current ) {
243 return $this->conditions;
246 $maxRow = end( $this->current );
248 foreach ( $this->primaryKey as $alias => $column ) {
249 $name = is_numeric( $alias ) ? $column : $alias;
250 $maximumValues[$column] = $this->db->addQuotes( $maxRow->{$name} );
254 // For example: If we have 3 primary keys
255 // first run through will generate
256 // col1 = 4 AND col2 = 7 AND col3 > 1
257 // second run through will generate
258 // col1 = 4 AND col2 > 7
259 // and the final run through will generate
261 while ( $maximumValues ) {
262 $pkConditions[] = $this->buildGreaterThanCondition( $maximumValues );
263 array_pop( $maximumValues );
266 $conditions = $this->conditions;
267 $conditions[] = sprintf( '( %s )', implode( ' ) OR ( ', $pkConditions ) );
273 * Given an array of column names and their maximum value generate
274 * an SQL condition where all keys except the last match $quotedMaximumValues
275 * exactly and the last column is greater than the matching value in
276 * $quotedMaximumValues
278 * @param array $quotedMaximumValues The maximum values quoted with
279 * $this->db->addQuotes()
280 * @return string An SQL condition that will select rows where all
281 * columns match the maximum value exactly except the last column
282 * which must be greater than the provided maximum value
284 protected function buildGreaterThanCondition( array $quotedMaximumValues ) {
285 $keys = array_keys( $quotedMaximumValues );
286 $lastColumn = end( $keys );
287 $lastValue = array_pop( $quotedMaximumValues );
289 foreach ( $quotedMaximumValues as $column => $value ) {
290 $conditions[] = "$column = $value";
292 $conditions[] = "$lastColumn > $lastValue";
294 return implode( ' AND ', $conditions );