BatchRowIterator.php

Go to the documentation of this file.

<?php
 
use Wikimedia\Rdbms\IReadableDatabase;
use Wikimedia\Rdbms\SelectQueryBuilder;
 
class BatchRowIterator implements RecursiveIterator {
 
    protected IReadableDatabase $db;
 
    protected $primaryKey;
 
    protected $batchSize;
 
    private $current = [];
 
    private $key = -1;
 
    public SelectQueryBuilder $sqb;
 
    public function __construct( IReadableDatabase $db, $sqb, $primaryKey, $batchSize ) {
        if ( $batchSize < 1 ) {
            throw new InvalidArgumentException( 'Batch size must be at least 1 row.' );
        }
        $this->db = $db;
        $this->primaryKey = (array)$primaryKey;
        $this->batchSize = $batchSize;
 
        if ( $sqb instanceof SelectQueryBuilder ) {
            $this->sqb = $sqb;
        } else {
            $this->sqb = $db->newSelectQueryBuilder()
                ->tables( is_array( $sqb ) ? $sqb : [ $sqb ] )
                ->caller( __CLASS__ );
        }
        $this->sqb->fields( $this->primaryKey );
    }
 
    public function addConditions( array $conditions ) {
        $this->sqb->conds( $conditions );
    }
 
    public function addOptions( array $options ) {
        $this->sqb->options( $options );
    }
 
    public function addJoinConditions( array $conditions ) {
        $this->sqb->joinConds( $conditions );
    }
 
    public function setFetchColumns( array $columns ) {
        // If it's not the all column selector merge in the primary keys we need
        if ( count( $columns ) === 1 && reset( $columns ) === '*' ) {
            $fetchColumns = $columns;
        } else {
            $fetchColumns = array_unique( array_merge(
                $this->primaryKey,
                $columns
            ) );
        }
        $this->sqb->clearFields()->fields( $fetchColumns );
    }
 
    public function setCaller( $caller ) {
        $this->sqb->caller( $caller );
 
        return $this;
    }
 
    public function extractPrimaryKeys( $row ) {
        $pk = [];
        foreach ( $this->primaryKey as $alias => $column ) {
            $name = is_numeric( $alias ) ? $column : $alias;
            $pk[$name] = $row->{$name};
        }
        return $pk;
    }
 
    public function current(): array {
        return $this->current;
    }
 
    public function key(): int {
        return $this->key;
    }
 
    public function rewind(): void {
        $this->key = -1; // self::next() will turn this into 0
        $this->current = [];
        $this->next();
    }
 
    public function valid(): bool {
        return (bool)$this->current;
    }
 
    public function hasChildren(): bool {
        return $this->current && count( $this->current );
    }
 
    public function getChildren(): ?RecursiveIterator {
        return new NotRecursiveIterator( new ArrayIterator( $this->current ) );
    }
 
    public function next(): void {
        $res = ( clone $this->sqb )
            ->andWhere( $this->buildConditions() )
            ->limit( $this->batchSize )
            ->orderBy( $this->primaryKey, SelectQueryBuilder::SORT_ASC )
            ->fetchResultSet();
 
        // The iterator is converted to an array because in addition to
        // returning it in self::current() we need to use the end value
        // in self::buildConditions()
        $this->current = iterator_to_array( $res );
        $this->key++;
    }
 
    protected function buildConditions() {
        if ( !$this->current ) {
            return [];
        }
 
        $maxRow = end( $this->current );
        $maximumValues = [];
        foreach ( $this->primaryKey as $alias => $column ) {
            $name = is_numeric( $alias ) ? $column : $alias;
            $maximumValues[$column] = $maxRow->$name;
        }
 
        return [ $this->db->buildComparison( '>', $maximumValues ) ];
    }
}