BatchEntityIterator.php

<?php

namespace Sidus\BatchIteratorBundle\Command;

use Doctrine\Common\Persistence\ObjectManager;
use Doctrine\ORM\Query;
use Doctrine\ORM\QueryBuilder;
use Symfony\Bundle\FrameworkBundle\Command\ContainerAwareCommand;
use Symfony\Component\Console\Helper\ProgressBar;
use Symfony\Component\Console\Output\OutputInterface;

abstract class EntityIteratorCommand extends ContainerAwareCommand
{
    /**
     * @var OutputInterface
     */
    protected $output;

    /**
     * @var ProgressBar
     */
    protected $progress;

    protected $processedEntities = [];

    /**
     *
     * @param ObjectManager $em
     * @param QueryBuilder $qb
     * @param callable $callback
     * @param int $batchSize
     * @param bool $withProgress
     * @param bool $detach
     */
    protected function batchIterate(ObjectManager $em, QueryBuilder $qb, callable $callback, $batchSize = 100, $withProgress = true, $detach = true)
    {
        $q = $qb->getQuery();
        $alias = $qb->getRootAliases()[0];
        $total = $qb->select("count({$alias}.id)")->getQuery()->getSingleScalarResult();
        $iterableResult = $q->iterate();
        if ($this->output && $withProgress) {
            $this->progress = new ProgressBar($this->output, $total);
            $this->progress->setBarWidth(100);
            $this->progress->start();
        }
        $saved = 0;
        $i = 0;

        foreach ($iterableResult AS $entity) {
            if (is_array($entity)) {
                $entity = $entity[0];
            }
            if ($callback($entity)) {
                $saved++;
            }
            $this->processedEntities[] = $entity;
            $i++;

            // If batchsize reached but no entity was saved, let's clear the em anyway
            if (0 === $saved && $i > $batchSize) {
                if ($detach) {
                    unset($entity);
                    $this->clearEm($em);
                }
                $i = 0;
            }
            if ($saved > $batchSize) {
                unset($entity);
                $this->flushEm($em, $detach);
                $saved = 0;
                $i = 0;
            }
            $this->progress->advance();
        }

        $this->flushEm($em, $detach);
        $this->progress->finish();
        $this->progress = null;
    }

    protected function flushEm(ObjectManager $em, $detach = true)
    {
        $this->writeln("Flushing entity manager", OutputInterface::VERBOSITY_VERBOSE);
        $em->flush(); // Executes all updates.
        if ($detach) {
            $this->clearEm($em);
        }
    }

    /**
     * Detaches all processed objects from Doctrine, keeps the rest
     * @param ObjectManager $em
     */
    protected function clearEm(ObjectManager $em)
    {
        foreach ($this->processedEntities as $entity) {
            $em->detach($entity);
        }
        $this->processedEntities = [];
        gc_collect_cycles();
    }

    protected function writeln($message, $verbosity = OutputInterface::VERBOSITY_NORMAL)
    {
        if ($this->output && $verbosity <= $this->output->getVerbosity()) {
            if ($this->progress) {
                $this->output->writeln("\n");
            }
            $this->output->writeln($message);
        }
    }
}