Add support for UNION ALL and UNION DISTINCT (#89)

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.

Author: Zheness

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');
     }
 
     /**

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);

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)));