Fix CTE query expressions (#761)

* Fix cte * Fix tests * Add `ExpressionInterface` as alias type and update tests * Update psalm * Update tests * Add line to CHANGELOG.md
yiisoft · Oct 14, 2023 · 4f1dbb9 · 4f1dbb9
1 parent 071f46a
commit 4f1dbb9
Show file tree

Hide file tree

Showing 7 changed files with 130 additions and 15 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -8,6 +8,7 @@
 - Bug #756: Fix `Quoter::quoteSql()` for SQL containing table with prefix (@Tigrov)
 - Bug #756: Fix `Quoter::getTableNameParts()` for cases when different quotes for tables and columns (@Tigrov)
 - Bug #756: Fix `Quoter::quoteTableName()` for sub-query with alias (@Tigrov)
+- Bug #761: Quote aliases of CTE in `WITH` queries (@Tigrov)
 - Chg #765: Deprecate `SchemaInterface::TYPE_JSONB` (@Tigrov)
 
 ## 1.1.1 August 16, 2023

diff --git a/src/Query/Query.php b/src/Query/Query.php
@@ -662,7 +662,7 @@ public function where(array|string|ExpressionInterface|null $condition, array $p
         return $this;
     }
 
-    public function withQuery(QueryInterface|string $query, string $alias, bool $recursive = false): static
+    public function withQuery(QueryInterface|string $query, ExpressionInterface|string $alias, bool $recursive = false): static
     {
         $this->withQueries[] = ['query' => $query, 'alias' => $alias, 'recursive' => $recursive];
         return $this;

diff --git a/src/Query/QueryPartsInterface.php b/src/Query/QueryPartsInterface.php
@@ -645,10 +645,11 @@ public function where(array|string|ExpressionInterface|null $condition, array $p
      * Prepends an SQL statement using `WITH` syntax.
      *
      * @param QueryInterface|string $query The SQL statement to append using `UNION`.
-     * @param string $alias The query alias in `WITH` construction.
+     * @param ExpressionInterface|string $alias The query alias in `WITH` construction.
+     * To specify the alias in plain SQL, you may pass an instance of {@see ExpressionInterface}.
      * @param bool $recursive Its `true` if using `WITH RECURSIVE` and `false` if using `WITH`.
      */
-    public function withQuery(QueryInterface|string $query, string $alias, bool $recursive = false): static;
+    public function withQuery(QueryInterface|string $query, ExpressionInterface|string $alias, bool $recursive = false): static;
 
     /**
      * Specifies the `WITH` query clause for the query.

diff --git a/src/QueryBuilder/AbstractDQLQueryBuilder.php b/src/QueryBuilder/AbstractDQLQueryBuilder.php
@@ -411,7 +411,7 @@ public function buildWithQueries(array $withs, array &$params): string
         $recursive = false;
         $result = [];
 
-        /** @psalm-var array<array-key, array{query:string|Query, alias:string, recursive:bool}> $withs */
+        /** @psalm-var array{query:string|Query, alias:ExpressionInterface|string, recursive:bool}[] $withs */
         foreach ($withs as $with) {
             if ($with['recursive']) {
                 $recursive = true;
@@ -421,7 +421,9 @@ public function buildWithQueries(array $withs, array &$params): string
                 [$with['query'], $params] = $this->build($with['query'], $params);
             }
 
-            $result[] = $with['alias'] . ' AS (' . $with['query'] . ')';
+            $quotedAlias = $this->quoteCteAlias($with['alias']);
+
+            $result[] = $quotedAlias . ' AS (' . $with['query'] . ')';
         }
 
         return 'WITH ' . ($recursive ? 'RECURSIVE ' : '') . implode(', ', $result);
@@ -610,4 +612,32 @@ private function quoteTableNames(array $tables, array &$params): array
 
         return $tables;
     }
+
+    /**
+     * Quotes an alias of Common Table Expressions (CTE)
+     *
+     * @param ExpressionInterface|string $name The alias name with or without column names to quote.
+     *
+     * @return string The quoted alias.
+     */
+    private function quoteCteAlias(ExpressionInterface|string $name): string
+    {
+        if ($name instanceof ExpressionInterface) {
+            return $this->buildExpression($name);
+        }
+
+        if (!str_contains($name, '(')) {
+            return $this->quoter->quoteTableName($name);
+        }
+
+        if (!str_ends_with($name, ')')) {
+            return $name;
+        }
+
+        /** @psalm-suppress PossiblyUndefinedArrayOffset */
+        [$name, $columns] = explode('(', substr($name, 0, -1), 2);
+        $name = trim($name);
+
+        return $this->quoter->quoteTableName($name) . '(' . $this->buildColumns($columns) . ')';
+    }
 }
diff --git a/tests/AbstractQueryBuilderTest.php b/tests/AbstractQueryBuilderTest.php
@@ -668,7 +668,7 @@ public function testBuildWithQueries(): void
         $this->assertSame(
             DbHelper::replaceQuotes(
                 <<<SQL
-                WITH cte AS (SELECT * FROM [[admin_profile]])
+                WITH [[cte]] AS (SELECT * FROM [[admin_profile]])
                 SQL,
                 $db->getDriverName(),
             ),
@@ -1131,7 +1131,7 @@ public function testBuildWithQuery(): void
         $this->assertSame(
             DbHelper::replaceQuotes(
                 <<<SQL
-                WITH a1 AS (SELECT [[id]] FROM [[t1]] WHERE expr = 1), a2 AS ((SELECT [[id]] FROM [[t2]] INNER JOIN [[a1]] ON t2.id = a1.id WHERE expr = 2) UNION ( SELECT [[id]] FROM [[t3]] WHERE expr = 3 )) SELECT * FROM [[a2]]
+                WITH [[a1]] AS (SELECT [[id]] FROM [[t1]] WHERE expr = 1), [[a2]] AS ((SELECT [[id]] FROM [[t2]] INNER JOIN [[a1]] ON t2.id = a1.id WHERE expr = 2) UNION ( SELECT [[id]] FROM [[t3]] WHERE expr = 3 )) SELECT * FROM [[a2]]
                 SQL,
                 $db->getDriverName(),
             ),
@@ -1157,15 +1157,40 @@ public function testBuildWithQueryRecursive(): void
 
         [$sql, $params] = $qb->build($query);
 
-        $this->assertSame(
-            DbHelper::replaceQuotes(
-                <<<SQL
-                WITH RECURSIVE a1 AS (SELECT [[id]] FROM [[t1]] WHERE expr = 1) SELECT * FROM [[a1]]
-                SQL,
-                $db->getDriverName(),
-            ),
-            $sql,
+        $expected = DbHelper::replaceQuotes(
+            <<<SQL
+            WITH RECURSIVE [[a1]] AS (SELECT [[id]] FROM [[t1]] WHERE expr = 1) SELECT * FROM [[a1]]
+            SQL,
+            $db->getDriverName(),
         );
+
+        if (in_array($db->getDriverName(), ['oci', 'sqlsrv'], true)) {
+            $expected = str_replace('WITH RECURSIVE ', 'WITH ', $expected);
+        }
+
+        $this->assertSame($expected, $sql);
+        $this->assertSame([], $params);
+    }
+
+    /** @dataProvider \Yiisoft\Db\Tests\Provider\QueryBuilderProvider::cteAliases */
+    public function testBuildWithQueryAlias($alias, $expected)
+    {
+        $db = $this->getConnection();
+        $qb = $db->getQueryBuilder();
+
+        $withQuery = (new Query($db))->from('t');
+        $query = (new Query($db))->withQuery($withQuery, $alias)->from('t');
+
+        [$sql, $params] = $qb->build($query);
+
+        $expectedSql = DbHelper::replaceQuotes(
+            <<<SQL
+            WITH $expected AS (SELECT * FROM [[t]]) SELECT * FROM [[t]]
+            SQL,
+            $db->getDriverName(),
+        );
+
+        $this->assertSame($expectedSql, $sql);
         $this->assertSame([], $params);
     }
 

diff --git a/tests/Common/CommonQueryTest.php b/tests/Common/CommonQueryTest.php
@@ -4,6 +4,7 @@
 
 namespace Yiisoft\Db\Tests\Common;
 
+use Yiisoft\Db\Expression\Expression;
 use Yiisoft\Db\Query\Query;
 use Yiisoft\Db\Tests\AbstractQueryTest;
 
@@ -34,4 +35,50 @@ public function testColumnIndexByWithClosure()
 
         $db->close();
     }
+
+    public function testWithQuery()
+    {
+        $db = $this->getConnection(true);
+
+        $with = (new Query($db))
+            ->distinct()
+            ->select(['status'])
+            ->from('customer');
+
+        $query = (new Query($db))
+            ->withQuery($with, 'statuses')
+            ->from('statuses');
+
+        $this->assertEquals(2, $query->count());
+
+        $db->close();
+    }
+
+    public function testWithQueryRecursive()
+    {
+        $db = $this->getConnection();
+        $quoter = $db->getQuoter();
+        $isOracle = $db->getDriverName() === 'oci';
+
+        /** Sum 1 to 10 equals 55 */
+        $quotedName = $quoter->quoteColumnName('n');
+        $union = (new Query($db))
+            ->select(new Expression($quotedName . ' + 1'))
+            ->from('t')
+            ->where(['<', 'n', 10]);
+
+        $with = (new Query($db))
+            ->select(new Expression('1'))
+            ->from($isOracle ? new Expression('DUAL') : [])
+            ->union($union, true);
+
+        $sum = (new Query($db))
+            ->withQuery($with, 't(n)', true)
+            ->from('t')
+            ->sum($quotedName);
+
+        $this->assertEquals(55, $sum);
+
+        $db->close();
+    }
 }
diff --git a/tests/Provider/QueryBuilderProvider.php b/tests/Provider/QueryBuilderProvider.php
@@ -1231,4 +1231,15 @@ public static function upsert(): array
             ],
         ];
     }
+
+    public static function cteAliases(): array
+    {
+        return [
+            'simple' => ['a', '[[a]]'],
+            'with one column' => ['a(b)', '[[a]]([[b]])'],
+            'with columns' => ['a(b,c,d)', '[[a]]([[b]], [[c]], [[d]])'],
+            'with extra space' => ['a(b,c,d) ', 'a(b,c,d) '],
+            'expression' => [new Expression('a(b,c,d)'), 'a(b,c,d)'],
+        ];
+    }
 }