includes/utils/BatchRowIterator.php

   1 <?php
   2
   3 use Wikimedia\Rdbms\IDatabase;
   4
   5 /**
   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.
  10  *
  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.
  15  *
  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.
  20  *
  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
  25  *
  26  * @file
  27  * @ingroup Maintenance
  28  */
  29 class BatchRowIterator implements RecursiveIterator {
  30
  31         /**
  32          * @var IDatabase $db The database to read from
  33          */
  34         protected $db;
  35
  36         /**
  37          * @var string|array $table The name or names of the table to read from
  38          */
  39         protected $table;
  40
  41         /**
  42          * @var array $primaryKey The name of the primary key(s)
  43          */
  44         protected $primaryKey;
  45
  46         /**
  47          * @var integer $batchSize The number of rows to fetch per iteration
  48          */
  49         protected $batchSize;
  50
  51         /**
  52          * @var array $conditions Array of strings containing SQL conditions
  53          *  to add to the query
  54          */
  55         protected $conditions = [];
  56
  57         /**
  58          * @var array $joinConditions
  59          */
  60         protected $joinConditions = [];
  61
  62         /**
  63          * @var array $fetchColumns List of column names to select from the
  64          *  table suitable for use with IDatabase::select()
  65          */
  66         protected $fetchColumns;
  67
  68         /**
  69          * @var string $orderBy SQL Order by condition generated from $this->primaryKey
  70          */
  71         protected $orderBy;
  72
  73         /**
  74          * @var array $current The current iterator value
  75          */
  76         private $current = [];
  77
  78         /**
  79          * @var integer key 0-indexed number of pages fetched since self::reset()
  80          */
  81         private $key;
  82
  83         /**
  84          * @var array Additional query options
  85          */
  86         protected $options = [];
  87
  88         /**
  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 integer $batchSize The number of rows to fetch per iteration
  93          * @throws InvalidArgumentException
  94          */
  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.' );
  98                 }
  99                 $this->db = $db;
 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;
 105         }
 106
 107         /**
 108          * @param array $conditions Query conditions suitable for use with
 109          *  IDatabase::select
 110          */
 111         public function addConditions( array $conditions ) {
 112                 $this->conditions = array_merge( $this->conditions, $conditions );
 113         }
 114
 115         /**
 116          * @param array $options Query options suitable for use with
 117          *  IDatabase::select
 118          */
 119         public function addOptions( array $options ) {
 120                 $this->options = array_merge( $this->options, $options );
 121         }
 122
 123         /**
 124          * @param array $conditions Query join conditions suitable for use
 125          *  with IDatabase::select
 126          */
 127         public function addJoinConditions( array $conditions ) {
 128                 $this->joinConditions = array_merge( $this->joinConditions, $conditions );
 129         }
 130
 131         /**
 132          * @param array $columns List of column names to select from the
 133          *  table suitable for use with IDatabase::select()
 134          */
 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;
 139                 } else {
 140                         $this->fetchColumns = array_unique( array_merge(
 141                                 $this->primaryKey,
 142                                 $columns
 143                         ) );
 144                 }
 145         }
 146
 147         /**
 148          * Extracts the primary key(s) from a database row.
 149          *
 150          * @param stdClass $row An individual database row from this iterator
 151          * @return array Map of primary key column to value within the row
 152          */
 153         public function extractPrimaryKeys( $row ) {
 154                 $pk = [];
 155                 foreach ( $this->primaryKey as $alias => $column ) {
 156                         $name = is_numeric( $alias ) ? $column : $alias;
 157                         $pk[$name] = $row->{$name};
 158                 }
 159                 return $pk;
 160         }
 161
 162         /**
 163          * @return array The most recently fetched set of rows from the database
 164          */
 165         public function current() {
 166                 return $this->current;
 167         }
 168
 169         /**
 170          * @return integer 0-indexed count of the page number fetched
 171          */
 172         public function key() {
 173                 return $this->key;
 174         }
 175
 176         /**
 177          * Reset the iterator to the begining of the table.
 178          */
 179         public function rewind() {
 180                 $this->key = -1; // self::next() will turn this into 0
 181                 $this->current = [];
 182                 $this->next();
 183         }
 184
 185         /**
 186          * @return bool True when the iterator is in a valid state
 187          */
 188         public function valid() {
 189                 return (bool)$this->current;
 190         }
 191
 192         /**
 193          * @return bool True when this result set has rows
 194          */
 195         public function hasChildren() {
 196                 return $this->current && count( $this->current );
 197         }
 198
 199         /**
 200          * @return RecursiveIterator
 201          */
 202         public function getChildren() {
 203                 return new NotRecursiveIterator( new ArrayIterator( $this->current ) );
 204         }
 205
 206         /**
 207          * Fetch the next set of rows from the database.
 208          */
 209         public function next() {
 210                 $res = $this->db->select(
 211                         $this->table,
 212                         $this->fetchColumns,
 213                         $this->buildConditions(),
 214                         __METHOD__,
 215                         [
 216                                 'LIMIT' => $this->batchSize,
 217                                 'ORDER BY' => $this->orderBy,
 218                         ] + $this->options,
 219                         $this->joinConditions
 220                 );
 221
 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 );
 226                 $this->key++;
 227         }
 228
 229         /**
 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
 234          *
 235          * Example output:
 236          *     [ '( foo = 42 AND bar > 7 ) OR ( foo > 42 )' ]
 237          *
 238          * @return array The SQL conditions necessary to select the next set
 239          *  of rows in the batched query
 240          */
 241         protected function buildConditions() {
 242                 if ( !$this->current ) {
 243                         return $this->conditions;
 244                 }
 245
 246                 $maxRow = end( $this->current );
 247                 $maximumValues = [];
 248                 foreach ( $this->primaryKey as $alias => $column ) {
 249                         $name = is_numeric( $alias ) ? $column : $alias;
 250                         $maximumValues[$column] = $this->db->addQuotes( $maxRow->{$name} );
 251                 }
 252
 253                 $pkConditions = [];
 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
 260                 //   col1 > 4
 261                 while ( $maximumValues ) {
 262                         $pkConditions[] = $this->buildGreaterThanCondition( $maximumValues );
 263                         array_pop( $maximumValues );
 264                 }
 265
 266                 $conditions = $this->conditions;
 267                 $conditions[] = sprintf( '( %s )', implode( ' ) OR ( ', $pkConditions ) );
 268
 269                 return $conditions;
 270         }
 271
 272         /**
 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
 277          *
 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
 283          */
 284         protected function buildGreaterThanCondition( array $quotedMaximumValues ) {
 285                 $keys = array_keys( $quotedMaximumValues );
 286                 $lastColumn = end( $keys );
 287                 $lastValue = array_pop( $quotedMaximumValues );
 288                 $conditions = [];
 289                 foreach ( $quotedMaximumValues as $column => $value ) {
 290                         $conditions[] = "$column = $value";
 291                 }
 292                 $conditions[] = "$lastColumn > $lastValue";
 293
 294                 return implode( ' AND ', $conditions );
 295         }
 296 }