From 68ad7c3ba4479f83608ac298afb18eda39abf678 Mon Sep 17 00:00:00 2001
From: Zheness <maxprudhomme@gmail.com>
Date: Fri, 1 Mar 2024 13:16:22 +0100
Subject: [PATCH] Add support for UNION ALL and UNION DISTINCT

The code now supports UNION ALL and UNION DISTINCT statements in addition to the standard UNION. The MagicQueryTest.php file has been updated with added test cases to reflect these changes. Union.php and StatementFactory.php were modified to handle the additional logic required for the new union types.
---
 src/SQLParser/Query/StatementFactory.php | 15 ++++++++++-----
 src/SQLParser/Query/Union.php            | 12 ++++++++++--
 tests/Mouf/Database/MagicQueryTest.php   |  6 ++++++
 3 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/src/SQLParser/Query/StatementFactory.php b/src/SQLParser/Query/StatementFactory.php
index ff50022..9447f17 100644
--- a/src/SQLParser/Query/StatementFactory.php
+++ b/src/SQLParser/Query/StatementFactory.php
@@ -92,11 +92,16 @@ public static function toObject(array $desc)
             }
 
             return $select;
-        } elseif (isset($desc['UNION'])) {
+        }
+        // UNION and UNION DISTINCT have similar behavior
+        if (isset($desc['UNION']) || isset($desc['UNION ALL']) || isset($desc['UNION DISTINCT'])) {
+            $isUnionAll = isset($desc['UNION ALL']);
+            $unionStatement =  $desc['UNION'] ?? ($desc['UNION ALL'] ?? $desc['UNION DISTINCT']);
+
             /** @var Select[] $selects */
-            $selects = array_map([self::class, 'toObject'], $desc['UNION']);
+            $selects = array_map([self::class, 'toObject'], $unionStatement);
 
-            $union = new Union($selects);
+            $union = new Union($selects, $isUnionAll);
 
             if (isset($desc['0']) && isset($desc['0']['ORDER'])) {
                 $order = NodeFactory::mapArrayToNodeObjectList($desc['0']['ORDER']);
@@ -105,9 +110,9 @@ public static function toObject(array $desc)
             }
 
             return $union;
-        } else {
-            throw new \BadMethodCallException('Unknown query');
         }
+
+        throw new \BadMethodCallException('Unknown query');
     }
 
     /**
diff --git a/src/SQLParser/Query/Union.php b/src/SQLParser/Query/Union.php
index f92cf25..5478c5a 100644
--- a/src/SQLParser/Query/Union.php
+++ b/src/SQLParser/Query/Union.php
@@ -25,13 +25,19 @@ class Union implements StatementInterface, NodeInterface
      */
     private $selects;
 
+    /**
+     * @var bool
+     */
+    private $isUnionAll;
+
     /**
      * Union constructor.
      * @param Select[] $selects
      */
-    public function __construct(array $selects)
+    public function __construct(array $selects, bool $isUnionAll)
     {
         $this->selects = $selects;
+        $this->isUnionAll = $isUnionAll;
     }
 
     /** @var NodeInterface[]|NodeInterface */
@@ -105,7 +111,9 @@ public function toSql(array $parameters, AbstractPlatform $platform, int $indent
             return $select->toSql($parameters, $platform, $indent, $conditionsMode, $extrapolateParameters);
         }, $this->selects);
 
-        $sql = '(' . implode(') UNION (', $selectsSql) . ')';
+        $unionStatement = $this->isUnionAll ? 'UNION ALL' : 'UNION';
+
+        $sql = '(' . implode(') ' . $unionStatement . ' (', $selectsSql) . ')';
 
         if (!empty($this->order)) {
             $order = NodeFactory::toSql($this->order, $platform, $parameters, ',', false, $indent + 2, $conditionsMode, $extrapolateParameters);
diff --git a/tests/Mouf/Database/MagicQueryTest.php b/tests/Mouf/Database/MagicQueryTest.php
index 7aef086..64b1967 100644
--- a/tests/Mouf/Database/MagicQueryTest.php
+++ b/tests/Mouf/Database/MagicQueryTest.php
@@ -188,6 +188,12 @@ public function testStandardSelect()
         $sql = 'SELECT a FROM users UNION SELECT a FROM users';
         $this->assertEquals('(SELECT a FROM users) UNION (SELECT a FROM users)', self::simplifySql($magicQuery->build($sql)));
 
+        $sql = 'SELECT a FROM users UNION ALL SELECT a FROM users';
+        $this->assertEquals('(SELECT a FROM users) UNION ALL (SELECT a FROM users)', self::simplifySql($magicQuery->build($sql)));
+
+        $sql = 'SELECT a FROM users UNION DISTINCT SELECT a FROM users';
+        $this->assertEquals('(SELECT a FROM users) UNION (SELECT a FROM users)', self::simplifySql($magicQuery->build($sql)));
+
         $sql = 'SELECT a FROM users u, users u2';
         $this->assertEquals('SELECT a FROM users u CROSS JOIN users u2', self::simplifySql($magicQuery->build($sql)));