<commit_before><?php
namespace Psa\ElasticIndex\Shell;

use Cake\Console\Shell;
use Cake\Core\Configure;
use Cake\Datasource\ConnectionManager;
use Cake\ElasticSearch\TypeRegistry;
use Cake\ORM\TableRegistry;
use Exception;

class ElasticIndexShell extends Shell {

    /**
     * Indexable tables
     *
     * @var array
     */
    protected $_indexableTables = [];

    /**
     * Counter
     *
     * @var int
     */
    protected $_counter = 0;

    /**
     * {@inheritdoc}
     */
    public function initialize()
    {
        parent::initialize();
        $this->_indexableTables = (array)Configure::read('ElasticIndex.indexableTables');
    }

    /**
     * Main entry point
     *
     * @return void
     */
    public function main()
    {
        $this->chooseModel();
    }

    /**
     * Let the user choose a model to index.
     *
     * @return void
     */
    public function chooseModel()
    {
        if (empty($this->_indexableTables)) {
            $this->out(__d('elastic_index', 'No tables that could be indexed were found.'));
            return;
        }

        $this->out(__d('elastic_index', 'Choose a model to index:'));
        $this->nl(2);

        foreach ($this->_indexableTables as $key => $table) {
            $this->out('[' . $key . ']' . ' ' . $table);
        }

        $selection = $this->in(__d('elastic_index', 'Choose the table you want to index:'));
        $this->params['table'] = $this->_indexableTables[$selection];

        $this->build();
    }

    /**
     * Gets a table object of a table to index.
     *
     * If the table doesn't have the behavior loaded it will be loaded.
     *
     * @param string $tableName
     * @return \Cake\ORM\Table
     */
    protected function _getTable($tableName = null)
    {
        if ($tableName === null) {
            $tableName = $this->param('table');
        }

        if (empty($tableName)) {
            $this->abort('No --table option provided!');
        }

        try {
            $table = TableRegistry::get($tableName);
            if (!in_array('ElasticIndex', $table->behaviors()->loaded())) {
                $table->addBehavior('Psa/ElasticIndex.ElasticIndex');
            }

            return $table;
        } catch (\Exception $e) {
            $this->printException($e);
        }
    }

    /**
     * Gets table name(s) from the input
     *
     * @return array
     */
    protected function _getTablesFromInput()
    {
        $table = $this->param('table');
        if (empty($table)) {
            return $this->_indexableTables;
        }

        return explode(',', $table);
    }

    /**
     * (Re-)Builds the whole index for a given table.
     *
     * @return void
     */
    public function build()
    {
        $tables = $this->_getTablesFromInput();
        foreach ($tables as $table) {
            $this->out(__d('elastic_index', 'Start building index for table "{0}".', [$table]));
            $this->_buildIndex($table);
            $this->out(__d('elastic_index', 'Finished building index for table "{0}".', [$table]));
        }

        $this->out(__d('elastic_index', 'Done indexing.'));
    }

    /**
     * The actual index building method that iterates over the table data.
     *
     * @param string $tableName
     * @return void
     */
    protected function _buildIndex($tableName)
    {
        $table = $this->_getTable($tableName);

        $query = $table->find();
        if ($table->hasFinder('indexDataCount')) {
            $query->find('indexDataCount');
        }

        $total = $query
            ->all()
            ->count();

        $offset = $this->param('offset');
        $limit = $this->param('limit');

        if ($offset > $total) {
            $this->abort(sprintf('Offset (%d) is bigger than the total number (%d) of records', $offset, $total));
        }

        if ($offset > 0) {
            $this->out(sprintf('Going to process %d records of %d.', $total - $offset, $total));
        } else {
            $this->out(sprintf('Going to process %d records.', $total));
        }

        $this->helper('progress')->output([
            'total' => $total,
            'callback' => function ($progress) use ($total, $table, $offset, $limit) {
                while ($offset <= $total) {
                    $this->_process($table, $offset, $limit);
                    $offset = $offset + $limit;
                    $progress->increment($limit);
                    $progress->draw();
                }

                return;
            }
        ]);
    }

    protected function _getRecords($table, $offset, $limit) {
        $query = $table->find();
        if ($table->hasFinder('indexData')) {
            $query->find('indexData');
        }

        return $query
            ->offset($offset)<commit_msg>Working on an alternative bulk getting method<commit_after><?php
namespace Psa\ElasticIndex\Shell;

use Cake\Chronos\Chronos;
use Cake\Collection\Collection;
use Cake\Console\Shell;
use Cake\Core\Configure;
use Cake\Datasource\ConnectionManager;
use Cake\ElasticSearch\TypeRegistry;
use Cake\ORM\TableRegistry;
use Cake\Utility\Hash;
use Exception;

class ElasticIndexShell extends Shell {

    /**
     * Indexable tables
     *
     * @var array
     */
    protected $_indexableTables = [];

    /**
     * Counter
     *
     * @var int
     */
    protected $_counter = 0;

    /**
     * Start time of a process
     *
     * @var \DateTime|null
     */
    protected $_startTime = null;

    /**
     * {@inheritdoc}
     */
    public function initialize()
    {
        parent::initialize();
        $this->_indexableTables = (array)Configure::read('ElasticIndex.indexableTables');
    }

    /**
     * Main entry point
     *
     * @return void
     */
    public function main()
    {
        $this->chooseModel();
    }

    /**
     * Let the user choose a model to index.
     *
     * @return void
     */
    public function chooseModel()
    {
        if (empty($this->_indexableTables)) {
            $this->out(__d('elastic_index', 'No tables that could be indexed were found.'));
            return;
        }

        $this->out(__d('elastic_index', 'Choose a model to index:'));
        $this->nl(2);

        foreach ($this->_indexableTables as $key => $table) {
            $this->out('[' . $key . ']' . ' ' . $table);
        }

        $selection = $this->in(__d('elastic_index', 'Choose the table you want to index:'));
        $this->params['table'] = $this->_indexableTables[$selection];

        $this->build();
    }

    /**
     * Gets a table object of a table to index.
     *
     * If the table doesn't have the behavior loaded it will be loaded.
     *
     * @param string $tableName
     * @return \Cake\ORM\Table
     */
    protected function _getTable($tableName = null)
    {
        if ($tableName === null) {
            $tableName = $this->param('table');
        }

        if (empty($tableName)) {
            $this->abort('No --table option provided!');
        }

        try {
            $table = TableRegistry::get($tableName);
            if (!in_array('ElasticIndex', $table->behaviors()->loaded())) {
                $table->addBehavior('Psa/ElasticIndex.ElasticIndex');
            }

            return $table;
        } catch (\Exception $e) {
            $this->printException($e);
        }
    }

    /**
     * Gets table name(s) from the input
     *
     * @return array
     */
    protected function _getTablesFromInput()
    {
        $table = $this->param('table');
        if (empty($table)) {
            return $this->_indexableTables;
        }

        return explode(',', $table);
    }

    /**
     * (Re-)Builds the whole index for a given table.
     *
     * @return void
     */
    public function build()
    {
        $tables = $this->_getTablesFromInput();
        foreach ($tables as $table) {
            $this->out(__d('elastic_index', 'Start building index for table "{0}".', [$table]));
            $this->_buildIndex($table);
            $this->out(__d('elastic_index', 'Finished building index for table "{0}".', [$table]));
        }

        $this->out(__d('elastic_index', 'Done indexing.'));
    }

    /**
     * The actual index building method that iterates over the table data.
     *
     * @param string $tableName
     * @return void
     */
    protected function _buildIndex($tableName)
    {
        $table = $this->_getTable($tableName);

        $query = $table->find();
        if ($table->hasFinder('indexDataCount')) {
            $query->find('indexDataCount');
        }

        $total = $query
            ->all()
            ->count();

        $offset = $this->param('offset');
        $limit = $this->param('limit');

        if ($offset > $total) {
            $this->abort(sprintf('Offset (%d) is bigger than the total number (%d) of records', $offset, $total));
        }

        if ($offset > 0) {
            $this->out(sprintf('Going to process %d records of %d.', $total - $offset, $total));
        } else {
            $this->out(sprintf('Going to process %d records.', $total));
        }

        $this->_startTime = Chronos::now();

        //$this->_getRecordIds($table, $offset, $limit);

        $this->helper('progress')->output([
            'total' => $total,
            'callback' => function ($progress) use ($total, $table, $offset, $limit) {
                while ($offset <= $total) {
                    $this->_process($table, $offset, $limit);
                    $offset = $offset + $limit;
                    $progress->increment($limit);
                    $progress->draw();
                }

                return;
            }
        ]);

        $this->_showEndTime();
    }

    protected function _showEndTime() {
        $seconds = Chronos::now()->diffInSeconds($this->_startTime);
        $minutes = Chronos::now()->diffInMinutes($this->_startTime);
        if ($seconds > 60) {
            $seconds = $seconds - ($minutes * 60);
            $this->out($minutes . ' minutes and ' . $seconds . ' seconds');
        } else {
            $this->out($seconds . ' seconds');
        }
    }

    protected function _getRecordIds($table) {
        $query = $table->find();
//        if ($table->hasFinder('indexIds')) {
//            $query->find('indexIds');
//        }

        $field = (string)$table->getPrimaryKey();

        $stmt = ConnectionManager::get('default2')
            ->newQuery()
            ->select([
                'id'
            ])
            ->orderDesc('id')
            ->from('wa_profiles')
            ->bufferResults(false)
            ->execute();

        $f = function () use ($stmt) {
            while ($row = $stmt->fetch('assoc')) {
                yield $row;
            }
        };

        $limit = $this->param('limit');
        $collection = (new Collection($f()))
            ->chunk($limit)
            ->each(function($chunck) use ($table) {

                $query = $table->find();
                if ($table->hasFinder('indexData')) {
                    $query->find('indexData');
                }

                $ids = Hash::combine($chunck, '{n}.id', '{n}.id');

                $results = $query
                    ->where([
                        'Profiles.id IN' => $ids
                    ])
                    ->all()
                    ->toList();

                $table->saveIndexDocuments($results);
            });
    }

    protected function _getRecords($table, $offset, $limit) {
        $query = $table->find();
        if ($table->hasFinder('indexData')) {
            $query->find('indexData');
        }

        return $query
            ->offset($offset)
{'fraction_non_alphanumeric': 0.1592829705505762, 'fraction_numerical': 0.0012804097311139564, 'mean_word_length': 1.7515265382808831, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 135, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4099619', 'n_tokens_mistral': 3394, 'n_tokens_neox': 3251, 'n_words': 740}