diff --git a/core/lib/Drupal/Core/Database/Query/Select.php b/core/lib/Drupal/Core/Database/Query/Select.php index 2d3159c36e7..22aca8f3347 100644 --- a/core/lib/Drupal/Core/Database/Query/Select.php +++ b/core/lib/Drupal/Core/Database/Query/Select.php @@ -582,7 +582,25 @@ class Select extends Query implements SelectInterface { return $this; } + /** + * Implements SelectInterface::countQuery(). + */ public function countQuery() { + $count = $this->prepareCountQuery(); + + $query = $this->connection->select($count); + $query->addExpression('COUNT(*)'); + + return $query; + } + + /** + * Prepares a count query from the current query object. + * + * @return Drupal\Core\Database\Query\Select + * A new query object ready to have COUNT(*) performed on it. + */ + protected function prepareCountQuery() { // Create our new query object that we will mutate into a count query. $count = clone($this); @@ -630,10 +648,13 @@ class Select extends Query implements SelectInterface { $count->distinct = FALSE; } - $query = $this->connection->select($count); - $query->addExpression('COUNT(*)'); + // If there are any dependent queries to UNION, prepare each of those for + // the count query also. + foreach ($count->union as &$union) { + $union['query'] = $union['query']->prepareCountQuery(); + } - return $query; + return $count; } public function __toString() { diff --git a/core/modules/system/lib/Drupal/system/Tests/Database/SelectTest.php b/core/modules/system/lib/Drupal/system/Tests/Database/SelectTest.php index d487ba27f26..af4485de60e 100644 --- a/core/modules/system/lib/Drupal/system/Tests/Database/SelectTest.php +++ b/core/modules/system/lib/Drupal/system/Tests/Database/SelectTest.php @@ -281,6 +281,28 @@ class SelectTest extends DatabaseTestBase { $this->assertEqual($names[2], 'Ringo', 'Third query returned correct name.'); } + /** + * Tests that we can get a count query for a UNION Select query. + */ + function testUnionCount() { + $query_1 = db_select('test', 't') + ->fields('t', array('name', 'age')) + ->condition('age', array(27, 28), 'IN'); + + $query_2 = db_select('test', 't') + ->fields('t', array('name', 'age')) + ->condition('age', 28); + + $query_1->union($query_2, 'ALL'); + $names = $query_1->execute()->fetchCol(); + + $query_3 = $query_1->countQuery(); + $count = $query_3->execute()->fetchField(); + + // Ensure the counts match. + $this->assertEqual(count($names), $count, "The count query's result matched the number of rows in the UNION query."); + } + /** * Test that random ordering of queries works. *