BatchRowIterator.php

Go to the documentation of this file.

<?php

use Wikimedia\Rdbms\IDatabase;

class BatchRowIterator implements RecursiveIterator {

    protected $db;

    protected $table;

    protected $primaryKey;

    protected $batchSize;

    protected $conditions = [];

    protected $joinConditions = [];

    protected $fetchColumns;

    protected $orderBy;

    private $current = [];

    private $key;

    protected $options = [];

    public function __construct( IDatabase $db, $table, $primaryKey, $batchSize ) {
        if ( $batchSize < 1 ) {
            throw new InvalidArgumentException( 'Batch size must be at least 1 row.' );
        }
        $this->db = $db;
        $this->table = $table;
        $this->primaryKey = (array)$primaryKey;
        $this->fetchColumns = $this->primaryKey;
        $this->orderBy = implode( ' ASC,', $this->primaryKey ) . ' ASC';
        $this->batchSize = $batchSize;
    }

    public function addConditions( array $conditions ) {
        $this->conditions = array_merge( $this->conditions, $conditions );
    }

    public function addOptions( array $options ) {
        $this->options = array_merge( $this->options, $options );
    }

    public function addJoinConditions( array $conditions ) {
        $this->joinConditions = array_merge( $this->joinConditions, $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 ) === '*' ) {
            $this->fetchColumns = $columns;
        } else {
            $this->fetchColumns = array_unique( array_merge(
                $this->primaryKey,
                $columns
            ) );
        }
    }

    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() {
        return $this->current;
    }

    public function key() {
        return $this->key;
    }

    public function rewind() {
        $this->key = -1; // self::next() will turn this into 0
        $this->current = [];
        $this->next();
    }

    public function valid() {
        return (bool)$this->current;
    }

    public function hasChildren() {
        return $this->current && count( $this->current );
    }

    public function getChildren() {
        return new NotRecursiveIterator( new ArrayIterator( $this->current ) );
    }

    public function next() {
        $res = $this->db->select(
            $this->table,
            $this->fetchColumns,
            $this->buildConditions(),
            __METHOD__,
            [
                'LIMIT' => $this->batchSize,
                'ORDER BY' => $this->orderBy,
            ] + $this->options,
            $this->joinConditions
        );

        // 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 $this->conditions;
        }

        $maxRow = end( $this->current );
        $maximumValues = [];
        foreach ( $this->primaryKey as $alias => $column ) {
            $name = is_numeric( $alias ) ? $column : $alias;
            $maximumValues[$column] = $this->db->addQuotes( $maxRow->{$name} );
        }

        $pkConditions = [];
        // For example: If we have 3 primary keys
        // first run through will generate
        //   col1 = 4 AND col2 = 7 AND col3 > 1
        // second run through will generate
        //   col1 = 4 AND col2 > 7
        // and the final run through will generate
        //   col1 > 4
        while ( $maximumValues ) {
            $pkConditions[] = $this->buildGreaterThanCondition( $maximumValues );
            array_pop( $maximumValues );
        }

        $conditions = $this->conditions;
        $conditions[] = sprintf( '( %s )', implode( ' ) OR ( ', $pkConditions ) );

        return $conditions;
    }

    protected function buildGreaterThanCondition( array $quotedMaximumValues ) {
        $keys = array_keys( $quotedMaximumValues );
        $lastColumn = end( $keys );
        $lastValue = array_pop( $quotedMaximumValues );
        $conditions = [];
        foreach ( $quotedMaximumValues as $column => $value ) {
            $conditions[] = "$column = $value";
        }
        $conditions[] = "$lastColumn > $lastValue";

        return implode( ' AND ', $conditions );
    }
}