Add CTE support to select in QueryBuilder

Author: nio-dtp

src/Query/QueryBuilder.php

@@ -26,6 +26,7 @@
 use function is_object;
 use function key;
 use function method_exists;
+use function sprintf;
 use function strtoupper;
 use function substr;
 use function ucfirst;
@@ -82,6 +83,7 @@ class QueryBuilder
         'orderBy'    => [],
         'values'     => [],
         'for_update' => null,
+        'with_cte'   => [],
     ];
 
     /**
@@ -425,7 +427,7 @@ public function getSQL()
                 break;
 
             case self::SELECT:
-                $sql = $this->getSQLForSelect();
+                $sql = $this->getCommonTableExpressions() . $this->getSQLForSelect();
                 break;
         }
 
@@ -1491,6 +1493,29 @@ public function resetOrderBy(): self
         return $this;
     }
 
+    /**
+     * Adds a CTE to a select query.
+     *
+     * @param string $name       The name of the CTE.
+     * @param self $queryBuilder The query builder to use within the CTE.
+     *
+     * @return $this This QueryBuilder instance.
+     */
+    public function with(string $name, self $queryBuilder): self
+    {
+        if ($queryBuilder->hasCommonTableExpressions()) {
+            throw new QueryException('CTEs cannot be nested.');
+        }
+        $this->sqlParts['with_cte'][$name] = $queryBuilder;
+
+        return $this;
+    }
+
+    private function hasCommonTableExpressions(): bool
+    {
+        return 0 < count($this->sqlParts['with_cte']);
+    }
+
     /** @throws Exception */
     private function getSQLForSelect(): string
     {
@@ -1511,6 +1536,21 @@ private function getSQLForSelect(): string
             );
     }
 
+    private function getCommonTableExpressions(): string
+    {
+        if (!$this->hasCommonTableExpressions()) {
+            return '';
+        }
+        /** @var array<string,self> $cteParts */
+        $cteParts = $this->sqlParts['with_cte'];
+        $cte = [];
+        foreach ($cteParts as $name => $queryBuilder) {
+            $cte[] = sprintf('%s AS (%s)', $name, $queryBuilder->getSQL());
+        }
+
+        return 'WITH ' . implode(', ', $cte) . ' ';
+    }
+
     /**
      * @return string[]
      *

tests/Query/QueryBuilderTest.php

@@ -1030,6 +1030,46 @@ public function testSelectAllWithoutTableAlias(): void
         self::assertEquals('SELECT * FROM users', (string) $qb);
     }
 
+    public function testSelectWithCTE(): void
+    {
+        $qbWith = new QueryBuilder($this->conn);
+        $qbWith->select('id', 'name')
+            ->from('users')
+            ->where('name LIKE :name');
+
+        $qb = new QueryBuilder($this->conn);
+        $qb->with('filtered_users', $qbWith)
+            ->select('*')
+            ->from('filtered_users', 'fu');
+
+        self::assertEquals(
+            'WITH filtered_users AS (SELECT id, name FROM users WHERE name LIKE :name)'
+            . ' SELECT * FROM filtered_users fu',
+            $qb->getSQL(),
+        );
+    }
+
+    public function testSelectNestedCTE(): void
+    {
+        $qbNested = new QueryBuilder($this->conn);
+        $qbNested->select('*')
+            ->from('country');
+
+        $qbWith = new QueryBuilder($this->conn);
+        $qbWith->select('id', 'name')
+            ->with('countries', $qbNested)
+            ->from('users')
+            ->where('name LIKE :name');
+
+        self::expectException(QueryException::class);
+        self::expectExceptionMessage('CTEs cannot be nested.');
+
+        $qb = new QueryBuilder($this->conn);
+        $qb->with('filtered_users', $qbWith)
+            ->select('*')
+            ->from('filtered_users', 'fu');
+    }
+
     public function testGetParameterType(): void
     {
         $qb = new QueryBuilder($this->conn);