From 7bb3a4ad24b28bdfcec6408f16565220a06becbc Mon Sep 17 00:00:00 2001 From: Max Stepanov Date: Tue, 5 Aug 2025 16:07:15 +0300 Subject: [PATCH 1/2] fix cte with union --- mindsdb_sql_parser/parser.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/mindsdb_sql_parser/parser.py b/mindsdb_sql_parser/parser.py index 6403265..043b839 100644 --- a/mindsdb_sql_parser/parser.py +++ b/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 From 2e38e1a5bae1b29c003aed8b1688c3ae0e277ea9 Mon Sep 17 00:00:00 2001 From: Max Stepanov Date: Tue, 5 Aug 2025 16:27:30 +0300 Subject: [PATCH 2/2] test --- .../test_select_common_table_expression.py | 45 +++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/tests/test_base_sql/test_select_common_table_expression.py b/tests/test_base_sql/test_select_common_table_expression.py index 01c1e16..a92848c 100644 --- a/tests/test_base_sql/test_select_common_table_expression.py +++ b/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() \ No newline at end of file