Code Coverage
 
Lines
Functions and Methods
Classes and Traits
Total
84.48% covered (warning)
84.48%
49 / 58
46.67% covered (danger)
46.67%
7 / 15
CRAP
0.00% covered (danger)
0.00%
0 / 1
BatchRowIterator
84.48% covered (warning)
84.48%
49 / 58
46.67% covered (danger)
46.67%
7 / 15
27.34
0.00% covered (danger)
0.00%
0 / 1
 __construct
87.50% covered (warning)
87.50%
7 / 8
0.00% covered (danger)
0.00%
0 / 1
2.01
 addConditions
100.00% covered (success)
100.00%
1 / 1
100.00% covered (success)
100.00%
1 / 1
1
 addOptions
0.00% covered (danger)
0.00%
0 / 1
0.00% covered (danger)
0.00%
0 / 1
2
 addJoinConditions
0.00% covered (danger)
0.00%
0 / 1
0.00% covered (danger)
0.00%
0 / 1
2
 setFetchColumns
100.00% covered (success)
100.00%
6 / 6
100.00% covered (success)
100.00%
1 / 1
3
 setCaller
0.00% covered (danger)
0.00%
0 / 2
0.00% covered (danger)
0.00%
0 / 1
2
 extractPrimaryKeys
100.00% covered (success)
100.00%
5 / 5
100.00% covered (success)
100.00%
1 / 1
3
 current
100.00% covered (success)
100.00%
1 / 1
100.00% covered (success)
100.00%
1 / 1
1
 key
0.00% covered (danger)
0.00%
0 / 1
0.00% covered (danger)
0.00%
0 / 1
2
 rewind
100.00% covered (success)
100.00%
3 / 3
100.00% covered (success)
100.00%
1 / 1
1
 valid
100.00% covered (success)
100.00%
1 / 1
100.00% covered (success)
100.00%
1 / 1
1
 hasChildren
0.00% covered (danger)
0.00%
0 / 1
0.00% covered (danger)
0.00%
0 / 1
6
 getChildren
0.00% covered (danger)
0.00%
0 / 1
0.00% covered (danger)
0.00%
0 / 1
2
 next
93.75% covered (success)
93.75%
15 / 16
0.00% covered (danger)
0.00%
0 / 1
2.00
 buildConditions
100.00% covered (success)
100.00%
10 / 10
100.00% covered (success)
100.00%
1 / 1
4
1<?php
2
3use Wikimedia\Rdbms\IReadableDatabase;
4
5/**
6 * This program is free software; you can redistribute it and/or modify
7 * it under the terms of the GNU General Public License as published by
8 * the Free Software Foundation; either version 2 of the License, or
9 * (at your option) any later version.
10 *
11 * This program is distributed in the hope that it will be useful,
12 * but WITHOUT ANY WARRANTY; without even the implied warranty of
13 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14 * GNU General Public License for more details.
15 *
16 * You should have received a copy of the GNU General Public License along
17 * with this program; if not, write to the Free Software Foundation, Inc.,
18 * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
19 * http://www.gnu.org/copyleft/gpl.html
20 *
21 * @file
22 * @ingroup Maintenance
23 */
24
25/**
26 * Allows iterating a large number of rows in batches transparently.
27 * By default when iterated over returns the full query result as an
28 * array of rows.  Can be wrapped in RecursiveIteratorIterator to
29 * collapse those arrays into a single stream of rows queried in batches.
30 *
31 * @newable
32 */
33class BatchRowIterator implements RecursiveIterator {
34
35    /**
36     * @var IReadableDatabase
37     */
38    protected $db;
39
40    /**
41     * @var string|array The name or names of the table to read from
42     */
43    protected $table;
44
45    /**
46     * @var array The name of the primary key(s)
47     */
48    protected $primaryKey;
49
50    /**
51     * @var int The number of rows to fetch per iteration
52     */
53    protected $batchSize;
54
55    /**
56     * @var array Array of strings containing SQL conditions to add to the query
57     */
58    protected $conditions = [];
59
60    /**
61     * @var array
62     */
63    protected $joinConditions = [];
64
65    /**
66     * @var array List of column names to select from the table suitable for use
67     *  with IDatabase::select()
68     */
69    protected $fetchColumns;
70
71    /**
72     * @var string SQL Order by condition generated from $this->primaryKey
73     */
74    protected $orderBy;
75
76    /**
77     * @var array The current iterator value
78     */
79    private $current = [];
80
81    /**
82     * @var int 0-indexed number of pages fetched since self::reset()
83     */
84    private $key = -1;
85
86    /**
87     * @var array Additional query options
88     */
89    protected $options = [];
90
91    /**
92     * @var string|null For debugging which method is using this class.
93     */
94    protected $caller;
95
96    /**
97     * @stable to call
98     *
99     * @param IReadableDatabase $db
100     * @param string|array $table The name or names of the table to read from
101     * @param string|array $primaryKey The name or names of the primary key columns
102     * @param int $batchSize The number of rows to fetch per iteration
103     * @throws InvalidArgumentException
104     */
105    public function __construct( IReadableDatabase $db, $table, $primaryKey, $batchSize ) {
106        if ( $batchSize < 1 ) {
107            throw new InvalidArgumentException( 'Batch size must be at least 1 row.' );
108        }
109        $this->db = $db;
110        $this->table = $table;
111        $this->primaryKey = (array)$primaryKey;
112        $this->fetchColumns = $this->primaryKey;
113        $this->orderBy = implode( ' ASC,', $this->primaryKey ) . ' ASC';
114        $this->batchSize = $batchSize;
115    }
116
117    /**
118     * @param array $conditions Query conditions suitable for use with
119     *  IDatabase::select
120     */
121    public function addConditions( array $conditions ) {
122        $this->conditions = array_merge( $this->conditions, $conditions );
123    }
124
125    /**
126     * @param array $options Query options suitable for use with
127     *  IDatabase::select
128     */
129    public function addOptions( array $options ) {
130        $this->options = array_merge( $this->options, $options );
131    }
132
133    /**
134     * @param array $conditions Query join conditions suitable for use
135     *  with IDatabase::select
136     */
137    public function addJoinConditions( array $conditions ) {
138        $this->joinConditions = array_merge( $this->joinConditions, $conditions );
139    }
140
141    /**
142     * @param array $columns List of column names to select from the
143     *  table suitable for use with IDatabase::select()
144     */
145    public function setFetchColumns( array $columns ) {
146        // If it's not the all column selector merge in the primary keys we need
147        if ( count( $columns ) === 1 && reset( $columns ) === '*' ) {
148            $this->fetchColumns = $columns;
149        } else {
150            $this->fetchColumns = array_unique( array_merge(
151                $this->primaryKey,
152                $columns
153            ) );
154        }
155    }
156
157    /**
158     * Use ->setCaller( __METHOD__ ) to indicate which code is using this
159     * class. Only used in debugging output.
160     * @since 1.36
161     *
162     * @param string $caller
163     * @return self
164     */
165    public function setCaller( $caller ) {
166        $this->caller = $caller;
167
168        return $this;
169    }
170
171    /**
172     * Extracts the primary key(s) from a database row.
173     *
174     * @param stdClass $row An individual database row from this iterator
175     * @return array Map of primary key column to value within the row
176     */
177    public function extractPrimaryKeys( $row ) {
178        $pk = [];
179        foreach ( $this->primaryKey as $alias => $column ) {
180            $name = is_numeric( $alias ) ? $column : $alias;
181            $pk[$name] = $row->{$name};
182        }
183        return $pk;
184    }
185
186    /**
187     * @return array The most recently fetched set of rows from the database
188     */
189    public function current(): array {
190        return $this->current;
191    }
192
193    /**
194     * @return int 0-indexed count of the page number fetched
195     */
196    public function key(): int {
197        return $this->key;
198    }
199
200    /**
201     * Reset the iterator to the beginning of the table.
202     */
203    public function rewind(): void {
204        $this->key = -1; // self::next() will turn this into 0
205        $this->current = [];
206        $this->next();
207    }
208
209    /**
210     * @return bool True when the iterator is in a valid state
211     */
212    public function valid(): bool {
213        return (bool)$this->current;
214    }
215
216    /**
217     * @return bool True when this result set has rows
218     */
219    public function hasChildren(): bool {
220        return $this->current && count( $this->current );
221    }
222
223    /**
224     * @return null|RecursiveIterator
225     */
226    public function getChildren(): ?RecursiveIterator {
227        return new NotRecursiveIterator( new ArrayIterator( $this->current ) );
228    }
229
230    /**
231     * Fetch the next set of rows from the database.
232     */
233    public function next(): void {
234        $caller = __METHOD__;
235        if ( (string)$this->caller !== '' ) {
236            $caller .= " (for {$this->caller})";
237        }
238
239        $res = $this->db->select(
240            $this->table,
241            $this->fetchColumns,
242            $this->buildConditions(),
243            $caller,
244            [
245                'LIMIT' => $this->batchSize,
246                'ORDER BY' => $this->orderBy,
247            ] + $this->options,
248            $this->joinConditions
249        );
250
251        // The iterator is converted to an array because in addition to
252        // returning it in self::current() we need to use the end value
253        // in self::buildConditions()
254        $this->current = iterator_to_array( $res );
255        $this->key++;
256    }
257
258    /**
259     * Uses the primary key list and the maximal result row from the
260     * previous iteration to build an SQL condition sufficient for
261     * selecting the next page of results.
262     *
263     * @return array The SQL conditions necessary to select the next set
264     *  of rows in the batched query
265     */
266    protected function buildConditions() {
267        if ( !$this->current ) {
268            return $this->conditions;
269        }
270
271        $maxRow = end( $this->current );
272        $maximumValues = [];
273        foreach ( $this->primaryKey as $alias => $column ) {
274            $name = is_numeric( $alias ) ? $column : $alias;
275            $maximumValues[$column] = $maxRow->$name;
276        }
277
278        $conditions = $this->conditions;
279        $conditions[] = $this->db->buildComparison( '>', $maximumValues );
280
281        return $conditions;
282    }
283}