php Yii2批量插入助手。支持“忽略”和“重复密钥更新”策略

Posted 2021-05-04

<?php

namespace console\components\Helpers\Traits;

use yii\base\InvalidConfigException;
use yii\db\Connection;
use yii\db\QueryBuilder;

trait BatchInserterTrait
{
    public static $BI_STRATEGY_INSERT_IGNORE           = 1;
    public static $BI_STRATEGY_INSERT_DUPLICATE_UPDATE = 2;

    private $_bi_buffer = [];
    private $_bi_table;
    private $_bi_columns;
    private $_bi_batch_size;
    private $_bi_strategy;
    private $_bi_db;
    private $_bi_on_duplicate_expression;

    /**
     * @param $config
     *
     * @throws InvalidConfigException
     */
    protected function configureBatchInserter($config)
    {
        $this->_bi_strategy = $config['strategy'] ?? self::$BI_STRATEGY_INSERT_IGNORE;
        $this->_bi_batch_size = $config['batch_size'] ?? 1000;
        $this->_bi_on_duplicate_expression = $config['on_duplicate'] ?? '';

        $this->_bi_table = $config['table'] ?? null;
        $this->_bi_columns = $config['columns'] ?? null;
        $this->_bi_db = $config['db'] ?? null;

        if (!$this->_bi_table)
        {
            throw new \yii\base\InvalidConfigException('No table specified');
        }

        if (!$this->_bi_columns)
        {
            throw new \yii\base\InvalidConfigException('No columns specified');
        }

        if (!$this->_bi_db)
        {
            throw new \yii\base\InvalidConfigException('No db specified');
        }

        if (
            $this->_bi_strategy === self::$BI_STRATEGY_INSERT_DUPLICATE_UPDATE
            && empty($this->_bi_on_duplicate_expression)
        )
        {
            throw new InvalidConfigException('Provide \'on_duplicate\' expression for DUPLICATE KEY UPDATE strategy');
        }
    }

    protected function storeDataForBatchInsert($data)
    {
        $this->_bi_buffer[] = $data;
    }

    protected function ensureBatchInsert($force = false)
    {
        $res = false;

        if (count($this->_bi_buffer) >= $this->_bi_batch_size || $force)
        {
            if ($this->_bi_strategy === self::$BI_STRATEGY_INSERT_DUPLICATE_UPDATE)
            {
                $res = $this->batchInsertDuplicateUpdate(
                    $this->_bi_db,
                    $this->_bi_table,
                    $this->_bi_columns,
                    $this->_bi_buffer
                );
            }
            else
            {
                $res = $this->batchInsertIgnore(
                    $this->_bi_db,
                    $this->_bi_table,
                    $this->_bi_columns,
                    $this->_bi_buffer
                );
            }

            $this->_bi_buffer = [];
        }

        return $res;
    }

    private function batchInsertDuplicateUpdate(Connection $db, string $table, array $columns, array $rows)
    {
        if (empty($rows))
        {
            return 0;
        }

        $sql = (new QueryBuilder($db))
            ->batchInsert($table, $columns, $rows);

        $sql .= ' ON DUPLICATE KEY UPDATE ' . $this->_bi_on_duplicate_expression;

        return $db->createCommand($sql)->execute();
    }

    private function batchInsertIgnore(Connection $db, string $table, array $columns, array $rows)
    {
        if (empty($rows))
        {
            return 0;
        }

        $sql = (new QueryBuilder($db))
            ->batchInsert($table, $columns, $rows);
        $sql = 'INSERT IGNORE' . mb_substr( $sql, strlen( 'INSERT' ) );

        return $db->createCommand($sql)->execute();
    }
}

<?php
class Controller
{
    use TableIteratorTrait;
    use BatchInserterTrait;

    public function actionTest()
    {
        /** @var Connection $db */
        $db = \Yii::$app->db;

        $iterator = $this->tableIterator([
            'db' => $db,
            'table' => 'table_name_to_read',
            'key' => 'id',
        ]);

        $this->configureBatchInserter([
            'db' => $db,
            'table' => 'table_name_to_insert',
            // 'strategy' => BatchInserterTrait::$BI_STRATEGY_INSERT_IGNORE,
            'strategy' => BatchInserterTrait::$BI_STRATEGY_INSERT_DUPLICATE_UPDATE,
            'on_duplicate' => 'some_field = 123', // required if strategy is duplicate key update
            'batch_size' => 1000,
            'columns' => ['id', 'some_field'],
        ]);

        foreach ($iterator as $key => $row)
        {
            // store values for columns (order does matter)
            $this->storeDataForBatchInsert([$row['song_id'], RestrictedSong::REASON_NO_HITS]);
            
            // insert batch_size amount of rows in single query
            $res = $this->ensureBatchInsert();

        }

        // force insert last rows
        $res = $this->ensureBatchInsert(true);
    }