DoctrineDataSource.php

<?php declare(strict_types=1);

namespace App\Ublaboo\DataGrid\DataSource;

use Doctrine\ORM\Tools\Pagination\Paginator;
use Nette\Utils\Strings;

/**
 * @method void onDataLoaded(array $result)
 */
class DoctrineDataSource extends \Ublaboo\DataGrid\DataSource\DoctrineDataSource
{
	/**
	 * @param  string  $column
	 * @return string
	 */
	private function checkAliases($column)
	{
		if (Strings::contains($column, '.')) {
			return $column;
		}

		if (!isset($this->root_alias)) {
			$ras = $this->data_source->getRootAliases();
			$this->root_alias = current($ras);
		}

		return $this->root_alias . '.' . $column;
	}

	/**
	 * @return bool
	 */
	private function usePaginator()
	{
		return $this->data_source->getDQLPart('join') || $this->data_source->getDQLPart('groupBy');
	}

	/**
	 * Get count of data
	 * @return int
	 */
	public function getCount()
	{
		if ($this->usePaginator()) {
			return (new Paginator($this->getQuery(), false))->count();
		}
		$data_source = clone $this->data_source;
		$data_source->select(sprintf('COUNT(%s)', $this->checkAliases($this->primary_key)));
		$data_source->resetDQLPart('orderBy');

		return (int) $data_source->getQuery()->getSingleScalarResult();
	}

	/**
	 * @return array
	 */
	public function getData()
	{
		if ($this->usePaginator()) {
			$iterator = (new Paginator($this->getQuery(), false))->getIterator();

			$data = iterator_to_array($iterator);
		} else {
			$data = $this->getQuery()->getResult();
		}

		$this->onDataLoaded($data);

		return $data;
	}
}