Make it possible to create union main part for a CTE

docs/en/reference/query-builder.rst

@@ -369,7 +369,10 @@ or QueryBuilder instances to one of the following methods:
         ->setMaxResults(100);
 
 Common Table Expressions
-~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+`SELECT` main query
+^^^^^^^^^^^^^^^^^^^
 
 To define Common Table Expressions (CTEs) that can be used in select query.
 
@@ -400,6 +403,39 @@ Multiple CTEs can be defined by calling the with method multiple times.
 
 Values of parameters used in a CTE should be defined in the main QueryBuilder.
 
+`UNION` main query part
+^^^^^^^^^^^^^^^^^^^^^^^
+
+To define Common Table Expressions (CTEs) that can be used in union query the union
+api needs to be used instead of using `select()`:
+
+.. code-block:: php
+
+    <?php
+
+    $qb = $connection->createQueryBuilder();
+
+    $baseQueryBuilder = $qb->sub()
+        ->select('id')
+        ->from('table_a');
+
+    $unionPart1 = $qb->sub()
+        ->select('id', $qb->expr()->literal('first') . ' AS value', '1 AS sort')
+        ->from('cte_base')
+        ->where($qb->expr()->eq('id', ':id1'));
+
+    $unionPart2 = $qb->sub()
+        ->select('id', $qb->expr()->literal('second') . ' AS value', '2 AS sort')
+        ->from('cte_base')
+        ->where($qb->expr()->eq('id', ':id2'));
+
+    $qb->with('cte_base', $baseQueryBuilder)
+        ->union($unionPart1)
+        ->addUnion($unionPart2)
+        ->orderBy('sort')
+        ->setParameter('id1', 2)
+        ->setParameter('id2', 1);
+
 Building Expressions
 --------------------
 

src/Query/QueryBuilder.php

@@ -1445,7 +1445,15 @@ private function getSQLForUnion(): string
             );
         }
 
-        return $this->connection->getDatabasePlatform()
+        $databasePlatform = $this->connection->getDatabasePlatform();
+        $unionParts       = [];
+        if (count($this->commonTableExpressions) > 0) {
+            $unionParts[] = $databasePlatform
+                ->createWithSQLBuilder()
+                ->buildSQL(...$this->commonTableExpressions);
+        }
+
+        $unionParts[] = $databasePlatform
             ->createUnionSQLBuilder()
             ->buildSQL(
                 new UnionQuery(
@@ -1454,6 +1462,8 @@ private function getSQLForUnion(): string
                     new Limit($this->maxResults, $this->firstResult),
                 ),
             );
+
+        return implode(' ', $unionParts);
     }
 
     /**

tests/Functional/Query/QueryBuilderTest.php

@@ -528,6 +528,38 @@ public function testSelectWithCTEUnion(): void
         self::assertSame($expectedRows, $qb->executeQuery()->fetchAllAssociative());
     }
 
+    public function testCTEUnionMainQuery(): void
+    {
+        if (! $this->platformSupportsCTEs()) {
+            self::markTestSkipped('The database platform does not support CTE.');
+        }
+
+        $expectedRows = [['id' => 2, 'value' => 'first', 'sort' => 1], ['id' => 1, 'value' => 'second', 'sort' => 2]];
+        $expectedRows = $this->prepareExpectedRows($expectedRows);
+        $qb           = $this->connection->createQueryBuilder();
+
+        $baseQueryBuilder = $qb->sub()
+            ->select('id')
+            ->from('for_update');
+
+        $unionPart1 = $qb->sub()
+            ->select('id', $qb->expr()->literal('first') . ' AS value', '1 AS sort')
+            ->from('cte_base')
+            ->where($qb->expr()->eq('id', '2'));
+
+        $unionPart2 = $qb->sub()
+            ->select('id', $qb->expr()->literal('second') . ' AS value', '2 AS sort')
+            ->from('cte_base')
+            ->where($qb->expr()->eq('id', '1'));
+
+        $qb->with('cte_base', $baseQueryBuilder)
+            ->union($unionPart1)
+            ->addUnion($unionPart2)
+            ->orderBy('sort');
+
+        self::assertSame($expectedRows, $qb->executeQuery()->fetchAllAssociative());
+    }
+
     public function testPlatformDoesNotSupportCTE(): void
     {
         if ($this->platformSupportsCTEs()) {
@@ -549,7 +581,7 @@ public function testPlatformDoesNotSupportCTE(): void
     }
 
     /**
-     * @param array<array<string, int>> $rows
+     * @param array<array<string, int|string>> $rows
      *
      * @return array<array<string, int|string>>
      */