Code Coverage
 
Lines
Functions and Methods
Classes and Traits
Total
84.21% covered (warning)
84.21%
48 / 57
46.67% covered (danger)
46.67%
7 / 15
CRAP
0.00% covered (danger)
0.00%
0 / 1
BatchRowIterator
84.21% covered (warning)
84.21%
48 / 57
46.67% covered (danger)
46.67%
7 / 15
28.66
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.33% covered (success)
93.33%
14 / 15
0.00% covered (danger)
0.00%
0 / 1
3.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     */
104    public function __construct( IReadableDatabase $db, $table, $primaryKey, $batchSize ) {
105        if ( $batchSize < 1 ) {
106            throw new InvalidArgumentException( 'Batch size must be at least 1 row.' );
107        }
108        $this->db = $db;
109        $this->table = $table;
110        $this->primaryKey = (array)$primaryKey;
111        $this->fetchColumns = $this->primaryKey;
112        $this->orderBy = implode( ' ASC,', $this->primaryKey ) . ' ASC';
113        $this->batchSize = $batchSize;
114    }
115
116    /**
117     * @param array $conditions Query conditions suitable for use with
118     *  IDatabase::select
119     */
120    public function addConditions( array $conditions ) {
121        $this->conditions = array_merge( $this->conditions, $conditions );
122    }
123
124    /**
125     * @param array $options Query options suitable for use with
126     *  IDatabase::select
127     */
128    public function addOptions( array $options ) {
129        $this->options = array_merge( $this->options, $options );
130    }
131
132    /**
133     * @param array $conditions Query join conditions suitable for use
134     *  with IDatabase::select
135     */
136    public function addJoinConditions( array $conditions ) {
137        $this->joinConditions = array_merge( $this->joinConditions, $conditions );
138    }
139
140    /**
141     * @param array $columns List of column names to select from the
142     *  table suitable for use with IDatabase::select()
143     */
144    public function setFetchColumns( array $columns ) {
145        // If it's not the all column selector merge in the primary keys we need
146        if ( count( $columns ) === 1 && reset( $columns ) === '*' ) {
147            $this->fetchColumns = $columns;
148        } else {
149            $this->fetchColumns = array_unique( array_merge(
150                $this->primaryKey,
151                $columns
152            ) );
153        }
154    }
155
156    /**
157     * Use ->setCaller( __METHOD__ ) to indicate which code is using this
158     * class. Only used in debugging output.
159     * @since 1.36
160     *
161     * @param string $caller
162     * @return self
163     */
164    public function setCaller( $caller ) {
165        $this->caller = $caller;
166
167        return $this;
168    }
169
170    /**
171     * Extracts the primary key(s) from a database row.
172     *
173     * @param stdClass $row An individual database row from this iterator
174     * @return array Map of primary key column to value within the row
175     */
176    public function extractPrimaryKeys( $row ) {
177        $pk = [];
178        foreach ( $this->primaryKey as $alias => $column ) {
179            $name = is_numeric( $alias ) ? $column : $alias;
180            $pk[$name] = $row->{$name};
181        }
182        return $pk;
183    }
184
185    /**
186     * @return array The most recently fetched set of rows from the database
187     */
188    public function current(): array {
189        return $this->current;
190    }
191
192    /**
193     * @return int 0-indexed count of the page number fetched
194     */
195    public function key(): int {
196        return $this->key;
197    }
198
199    /**
200     * Reset the iterator to the beginning of the table.
201     */
202    public function rewind(): void {
203        $this->key = -1; // self::next() will turn this into 0
204        $this->current = [];
205        $this->next();
206    }
207
208    /**
209     * @return bool True when the iterator is in a valid state
210     */
211    public function valid(): bool {
212        return (bool)$this->current;
213    }
214
215    /**
216     * @return bool True when this result set has rows
217     */
218    public function hasChildren(): bool {
219        return $this->current && count( $this->current );
220    }
221
222    /**
223     * @return null|RecursiveIterator
224     */
225    public function getChildren(): ?RecursiveIterator {
226        return new NotRecursiveIterator( new ArrayIterator( $this->current ) );
227    }
228
229    /**
230     * Fetch the next set of rows from the database.
231     */
232    public function next(): void {
233        $caller = __METHOD__;
234        if ( (string)$this->caller !== '' ) {
235            $caller .= " (for {$this->caller})";
236        }
237
238        $res = $this->db->newSelectQueryBuilder()
239            ->tables( is_array( $this->table ) ? $this->table : [ $this->table ] )
240            ->fields( $this->fetchColumns )
241            ->where( $this->buildConditions() )
242            ->caller( $caller )
243            ->limit( $this->batchSize )
244            ->orderBy( $this->orderBy )
245            ->options( $this->options )
246            ->joinConds( $this->joinConditions )
247            ->fetchResultSet();
248
249        // The iterator is converted to an array because in addition to
250        // returning it in self::current() we need to use the end value
251        // in self::buildConditions()
252        $this->current = iterator_to_array( $res );
253        $this->key++;
254    }
255
256    /**
257     * Uses the primary key list and the maximal result row from the
258     * previous iteration to build an SQL condition sufficient for
259     * selecting the next page of results.
260     *
261     * @return array The SQL conditions necessary to select the next set
262     *  of rows in the batched query
263     */
264    protected function buildConditions() {
265        if ( !$this->current ) {
266            return $this->conditions;
267        }
268
269        $maxRow = end( $this->current );
270        $maximumValues = [];
271        foreach ( $this->primaryKey as $alias => $column ) {
272            $name = is_numeric( $alias ) ? $column : $alias;
273            $maximumValues[$column] = $maxRow->$name;
274        }
275
276        $conditions = $this->conditions;
277        $conditions[] = $this->db->buildComparison( '>', $maximumValues );
278
279        return $conditions;
280    }
281}