Code Coverage
 
Lines
Functions and Methods
Classes and Traits
Total
80.77% covered (warning)
80.77%
42 / 52
53.33% covered (warning)
53.33%
8 / 15
 CRAP
0.00% covered (danger)
0.00%
0 / 1
BatchRowIterator
82.35% covered (warning)
82.35%
42 / 51
53.33% covered (warning)
53.33%
8 / 15
 29.72
0.00% covered (danger)
0.00%
0 / 1
 __construct
81.82% covered (warning)
81.82%
9 / 11
0.00% covered (danger)
0.00%
0 / 1
 4.10
 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%
7 / 7
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
100.00% covered (success)
100.00%
7 / 7
100.00% covered (success)
100.00%
1 / 1
 1
 buildConditions
100.00% covered (success)
100.00%
8 / 8
100.00% covered (success)
100.00%
1 / 1
 4
1<?php
2
3namespace MediaWiki\Utils;
4
5use ArrayIterator;
6use InvalidArgumentException;
7use RecursiveIterator;
8use Wikimedia\Iterators\NotRecursiveIterator;
9use Wikimedia\Rdbms\IReadableDatabase;
10use Wikimedia\Rdbms\SelectQueryBuilder;
11
12/**
13 * @license GPL-2.0-or-later
14 * @file
15 * @ingroup Maintenance
16 */
17
18/**
19 * Allows iterating a large number of rows in batches transparently.
20 * By default when iterated over returns the full query result as an
21 * array of rows.  Can be wrapped in RecursiveIteratorIterator to
22 * collapse those arrays into a single stream of rows queried in batches.
23 *
24 * @newable
25 */
26class BatchRowIterator implements RecursiveIterator {
27
28    protected IReadableDatabase $db;
29
30    /**
31     * @var string[] The name of the primary key(s)
32     */
33    protected $primaryKey;
34
35    /**
36     * @var int The number of rows to fetch per iteration
37     */
38    protected $batchSize;
39
40    /**
41     * @var \stdClass[] The current iterator value
42     */
43    private $current = [];
44
45    /**
46     * @var int 0-indexed number of pages fetched since self::reset()
47     */
48    private $key = -1;
49
50    /**
51     * Underlying database query builder, may be mutated before iteration begins.
52     */
53    public SelectQueryBuilder $sqb;
54
55    /**
56     * @stable to call
57     *
58     * @param IReadableDatabase $db
59     * @param SelectQueryBuilder|string|array $sqb The query to split into batches (or table name/names)
60     * @param string|string[] $primaryKey The name or names of the primary key columns
61     * @param int $batchSize The number of rows to fetch per iteration
62     */
63    public function __construct( IReadableDatabase $db, $sqb, $primaryKey, $batchSize ) {
64        if ( $batchSize < 1 ) {
65            throw new InvalidArgumentException( 'Batch size must be at least 1 row.' );
66        }
67        $this->db = $db;
68        $this->primaryKey = (array)$primaryKey;
69        $this->batchSize = $batchSize;
70
71        if ( $sqb instanceof SelectQueryBuilder ) {
72            $this->sqb = $sqb;
73        } else {
74            $this->sqb = $db->newSelectQueryBuilder()
75                ->tables( is_array( $sqb ) ? $sqb : [ $sqb ] )
76                ->caller( __METHOD__ );
77        }
78        $this->sqb->fields( $this->primaryKey );
79    }
80
81    /**
82     * @deprecated since 1.44 Use the SelectQueryBuilder object directly
83     * @param array $conditions Query conditions suitable for use with
84     *  IDatabase::select
85     */
86    public function addConditions( array $conditions ) {
87        $this->sqb->conds( $conditions );
88    }
89
90    /**
91     * @deprecated since 1.44 Use the SelectQueryBuilder object directly
92     * @param array $options Query options suitable for use with
93     *  IDatabase::select
94     */
95    public function addOptions( array $options ) {
96        $this->sqb->options( $options );
97    }
98
99    /**
100     * @deprecated since 1.44 Use the SelectQueryBuilder object directly
101     * @param array $conditions Query join conditions suitable for use
102     *  with IDatabase::select
103     */
104    public function addJoinConditions( array $conditions ) {
105        $this->sqb->joinConds( $conditions );
106    }
107
108    /**
109     * @deprecated since 1.44 Use the SelectQueryBuilder object directly
110     * @param array $columns List of column names to select from the
111     *  table suitable for use with IDatabase::select()
112     */
113    public function setFetchColumns( array $columns ) {
114        // If it's not the all column selector merge in the primary keys we need
115        if ( count( $columns ) === 1 && reset( $columns ) === '*' ) {
116            $fetchColumns = $columns;
117        } else {
118            $fetchColumns = array_unique( array_merge(
119                $this->primaryKey,
120                $columns
121            ) );
122        }
123        $this->sqb->clearFields()->fields( $fetchColumns );
124    }
125
126    /**
127     * Use ->setCaller( __METHOD__ ) to indicate which code is using this
128     * class. Only used in debugging output.
129     * @since 1.36
130     * @deprecated since 1.44 Use the SelectQueryBuilder object directly
131     *
132     * @param string $caller
133     * @return self
134     */
135    public function setCaller( $caller ) {
136        $this->sqb->caller( $caller );
137
138        return $this;
139    }
140
141    /**
142     * Extracts the primary key(s) from a database row.
143     *
144     * @param \stdClass $row An individual database row from this iterator
145     * @return array Map of primary key column to value within the row
146     */
147    public function extractPrimaryKeys( $row ) {
148        $pk = [];
149        foreach ( $this->primaryKey as $alias => $column ) {
150            $name = is_numeric( $alias ) ? $column : $alias;
151            $pk[$name] = $row->{$name};
152        }
153        return $pk;
154    }
155
156    /**
157     * @return \stdClass[] The most recently fetched set of rows from the database
158     */
159    public function current(): array {
160        return $this->current;
161    }
162
163    /**
164     * @return int 0-indexed count of the page number fetched
165     */
166    public function key(): int {
167        return $this->key;
168    }
169
170    /**
171     * Reset the iterator to the beginning of the table.
172     */
173    public function rewind(): void {
174        $this->key = -1; // self::next() will turn this into 0
175        $this->current = [];
176        $this->next();
177    }
178
179    /**
180     * @return bool True when the iterator is in a valid state
181     */
182    public function valid(): bool {
183        return (bool)$this->current;
184    }
185
186    /**
187     * @return bool True when this result set has rows
188     */
189    public function hasChildren(): bool {
190        return $this->current && count( $this->current );
191    }
192
193    public function getChildren(): ?RecursiveIterator {
194        return new NotRecursiveIterator( new ArrayIterator( $this->current ) );
195    }
196
197    /**
198     * Fetch the next set of rows from the database.
199     */
200    public function next(): void {
201        $res = ( clone $this->sqb )
202            ->andWhere( $this->buildConditions() )
203            ->limit( $this->batchSize )
204            ->orderBy( $this->primaryKey, SelectQueryBuilder::SORT_ASC )
205            ->fetchResultSet();
206
207        // The iterator is converted to an array because in addition to
208        // returning it in self::current() we need to use the end value
209        // in self::buildConditions()
210        $this->current = iterator_to_array( $res );
211        $this->key++;
212    }
213
214    /**
215     * Uses the primary key list and the maximal result row from the
216     * previous iteration to build an SQL condition sufficient for
217     * selecting the next page of results.
218     *
219     * @return array The SQL conditions necessary to select the next set
220     *  of rows in the batched query
221     */
222    protected function buildConditions() {
223        if ( !$this->current ) {
224            return [];
225        }
226
227        $maxRow = end( $this->current );
228        $maximumValues = [];
229        foreach ( $this->primaryKey as $alias => $column ) {
230            $name = is_numeric( $alias ) ? $column : $alias;
231            $maximumValues[$column] = $maxRow->$name;
232        }
233
234        return [ $this->db->buildComparison( '>', $maximumValues ) ];
235    }
236}
237
238/** @deprecated class alias since 1.46 */
239class_alias( BatchRowIterator::class, 'BatchRowIterator' );