Merge pull request #62 from mindsdb/fix-CONN-1340

Fix CTE with UNION

Author: StpMax

mindsdb_sql_parser/parser.py

@@ -1106,14 +1106,15 @@ def select(self, p):
         select.cte = p.ctes
         return select
 
-    @_('ctes COMMA identifier cte_columns_or_nothing AS LPAREN select RPAREN')
+    @_('ctes COMMA identifier cte_columns_or_nothing AS LPAREN select RPAREN',
+       'ctes COMMA identifier cte_columns_or_nothing AS LPAREN union RPAREN')
     def ctes(self, p):
         ctes = p.ctes
         ctes = ctes + [
             CommonTableExpression(
                 name=p.identifier,
                 columns=p.cte_columns_or_nothing,
-                query=p.select)
+                query=p[6])
         ]
         return ctes
 

tests/test_base_sql/test_select_common_table_expression.py

@@ -87,3 +87,48 @@ def test_cte_nested(self):
         assert str(ast).lower() == sql.lower()
         assert str(ast) == str(expected_ast)
         assert ast.to_tree() == expected_ast.to_tree()
+
+    def test_cte_union(self):
+        sql = """
+            WITH ta AS (
+                SELECT 'a' AS a
+                UNION
+                SELECT 'b' AS a
+            ), tb AS (
+                SELECT 'c' AS a
+                UNION
+                SELECT 'd' AS a
+            )
+            SELECT a FROM ta
+            UNION
+            SELECT a FROM tb
+        """
+        ast = parse_sql(sql)
+
+        expected_ast = Union(
+            left=Select(
+                cte=[
+                    CommonTableExpression(
+                        name=Identifier('ta'),
+                        query=Union(
+                            left=Select(targets=[Constant('a', alias=Identifier('a'))]),
+                            right=Select(targets=[Constant('b', alias=Identifier('a'))])
+                        )
+                    ),
+                    CommonTableExpression(
+                        name=Identifier('tb'),
+                        query=Union(
+                            left=Select(targets=[Constant('c', alias=Identifier('a'))]),
+                            right=Select(targets=[Constant('d', alias=Identifier('a'))])
+                        )
+                    ),
+                ],
+                targets=[Identifier('a')],
+                from_table=Identifier('ta')
+            ),
+            right=Select(targets=[Identifier('a')], from_table=Identifier('tb'))
+        )
+
+        assert (' '.join(str(ast).split())).lower() == (' '.join(sql.split())).lower()
+        assert str(ast) == str(expected_ast)
+        assert ast.to_tree() == expected_ast.to_tree()