-
Notifications
You must be signed in to change notification settings - Fork 2.5k
/
CountOutputWalker.php
142 lines (123 loc) · 5.09 KB
/
CountOutputWalker.php
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
<?php
declare(strict_types=1);
namespace Doctrine\ORM\Tools\Pagination;
use Doctrine\DBAL\Platforms\AbstractPlatform;
use Doctrine\ORM\Mapping\AssociationMetadata;
use Doctrine\ORM\Mapping\FieldMetadata;
use Doctrine\ORM\Query;
use Doctrine\ORM\Query\AST\SelectStatement;
use Doctrine\ORM\Query\ParserResult;
use Doctrine\ORM\Query\ResultSetMapping;
use Doctrine\ORM\Query\SqlWalker;
use RuntimeException;
use function array_diff;
use function array_keys;
use function count;
use function implode;
use function reset;
use function sprintf;
/**
* Wraps the query in order to accurately count the root objects.
*
* Given a DQL like `SELECT u FROM User u` it will generate an SQL query like:
* SELECT COUNT(*) (SELECT DISTINCT <id> FROM (<original SQL>))
*
* Works with composite keys but cannot deal with queries that have multiple
* root entities (e.g. `SELECT f, b from Foo, Bar`)
*/
class CountOutputWalker extends SqlWalker
{
/** @var AbstractPlatform */
private $platform;
/** @var ResultSetMapping */
private $rsm;
/** @var mixed[][] */
private $queryComponents;
/**
* Stores various parameters that are otherwise unavailable
* because Doctrine\ORM\Query\SqlWalker keeps everything private without
* accessors.
*
* @param Query $query
* @param ParserResult $parserResult
* @param mixed[][] $queryComponents
*/
public function __construct($query, $parserResult, array $queryComponents)
{
$this->platform = $query->getEntityManager()->getConnection()->getDatabasePlatform();
$this->rsm = $parserResult->getResultSetMapping();
$this->queryComponents = $queryComponents;
parent::__construct($query, $parserResult, $queryComponents);
}
/**
* Walks down a SelectStatement AST node, wrapping it in a COUNT (SELECT DISTINCT).
*
* Note that the ORDER BY clause is not removed. Many SQL implementations (e.g. MySQL)
* are able to cache subqueries. By keeping the ORDER BY clause intact, the limitSubQuery
* that will most likely be executed next can be read from the native SQL cache.
*
* @return string
*
* @throws RuntimeException
*/
public function walkSelectStatement(SelectStatement $AST)
{
if ($this->platform->getName() === 'mssql') {
$AST->orderByClause = null;
}
$sql = parent::walkSelectStatement($AST);
if ($AST->groupByClause) {
return sprintf(
'SELECT %s AS dctrn_count FROM (%s) dctrn_table',
$this->platform->getCountExpression('*'),
$sql
);
}
// Find out the SQL alias of the identifier column of the root entity
// It may be possible to make this work with multiple root entities but that
// would probably require issuing multiple queries or doing a UNION SELECT
// so for now, It's not supported.
// Get the root entity and alias from the AST fromClause
$from = $AST->fromClause->identificationVariableDeclarations;
if (count($from) > 1) {
throw new RuntimeException('Cannot count query which selects two FROM components, cannot make distinction');
}
$fromRoot = reset($from);
$rootAlias = $fromRoot->rangeVariableDeclaration->aliasIdentificationVariable;
$rootClass = $this->queryComponents[$rootAlias]['metadata'];
$rootIdentifier = $rootClass->identifier;
// For every identifier, find out the SQL alias by combing through the ResultSetMapping
$sqlIdentifier = [];
foreach ($rootIdentifier as $identifier) {
$property = $rootClass->getProperty($identifier);
if ($property instanceof FieldMetadata) {
foreach (array_keys($this->rsm->fieldMappings, $identifier, true) as $alias) {
if ($this->rsm->columnOwnerMap[$alias] === $rootAlias) {
$sqlIdentifier[$identifier] = $alias;
}
}
} elseif ($property instanceof AssociationMetadata) {
$joinColumns = $property->getJoinColumns();
$joinColumn = reset($joinColumns);
foreach (array_keys($this->rsm->metaMappings, $joinColumn->getColumnName(), true) as $alias) {
if ($this->rsm->columnOwnerMap[$alias] === $rootAlias) {
$sqlIdentifier[$identifier] = $alias;
}
}
}
}
if (count($rootIdentifier) !== count($sqlIdentifier)) {
throw new RuntimeException(sprintf(
'Not all identifier properties can be found in the ResultSetMapping: %s',
implode(', ', array_diff($rootIdentifier, array_keys($sqlIdentifier)))
));
}
// Build the counter query
return sprintf(
'SELECT %s AS dctrn_count FROM (SELECT DISTINCT %s FROM (%s) dctrn_result) dctrn_table',
$this->platform->getCountExpression('*'),
implode(', ', $sqlIdentifier),
$sql
);
}
}