fix to_tree for queries with ROLLUP

Author: StpMax

mindsdb_sql_parser/ast/select/identifier.py

@@ -95,7 +95,8 @@ def parts_to_str(self):
 
     def to_tree(self, *args, level=0, **kwargs):
         alias_str = f', alias={self.alias.to_tree()}' if self.alias else ''
-        return indent(level) + f'Identifier(parts={[str(i) for i in self.parts]}{alias_str})'
+        with_rollup_str = ', with_rollup=True' if self.with_rollup else ''
+        return indent(level) + f'Identifier(parts={[str(i) for i in self.parts]}{alias_str}{with_rollup_str})'
 
     def get_string(self, *args, **kwargs):
         return self.parts_to_str()

mindsdb_sql_parser/ast/show.py

@@ -31,12 +31,12 @@ def to_tree(self, *args, level=0, **kwargs):
         ind = indent(level)
         ind1 = indent(level+1)
         category_str = f'{ind1}category={repr(self.category)},'
-        from_str = f'\n{ind1}from={self.from_table.to_string()},' if self.from_table else ''
-        in_str = f'\n{ind1}in={self.in_table.to_tree(level=level + 2)},' if self.in_table else ''
+        from_str = f'\n{ind1}from_table={self.from_table.to_tree(level=level + 2)},' if self.from_table else ''
+        in_str = f'\n{ind1}in_table={self.in_table.to_tree(level=level + 2)},' if self.in_table else ''
         where_str = f'\n{ind1}where=\n{self.where.to_tree(level=level+2)},' if self.where else ''
-        name_str = f'\n{ind1}name={self.name},' if self.name else ''
-        like_str = f'\n{ind1}like={self.like},' if self.like else ''
-        modes_str = f'\n{ind1}modes=[{",".join(self.modes)}],' if self.modes else ''
+        name_str = f'\n{ind1}name={repr(self.name)},' if self.name else ''
+        like_str = f'\n{ind1}like={repr(self.like)},' if self.like else ''
+        modes_str = f'\n{ind1}modes=[{",".join([repr(m) for m in self.modes])}],' if self.modes else ''
         out_str = f'{ind}Show(' \
                   f'{category_str}' \
                   f'{name_str}' \

tests/test_mysql/test_mysql_parser.py

@@ -1,27 +1,29 @@
 from mindsdb_sql_parser import parse_sql
-from mindsdb_sql_parser.ast import Select, Identifier, BinaryOperation, Star
-from mindsdb_sql_parser.ast import Variable, Function
 from mindsdb_sql_parser.parser import Show
+from mindsdb_sql_parser.ast import Select, Identifier, BinaryOperation, Star, Variable, Function, ASTNode
+
+
+def compare_ast(parsed: ASTNode, expected: ASTNode, sql: str) -> None:
+    assert parsed.to_tree() == expected.to_tree()
+    assert str(parsed).lower() == sql.lower()
+    assert str(parsed) == str(expected)
+    assert str(eval(parsed.to_tree())) == str(parsed)
 
 
 class TestMySQLParser:
     def test_select_variable(self):
         sql = 'SELECT @version'
         ast = parse_sql(sql)
         expected_ast = Select(targets=[Variable('version')])
-        assert ast.to_tree() == expected_ast.to_tree()
-        assert str(ast).lower() == sql.lower()
-        assert str(ast) == str(expected_ast)
+        compare_ast(ast, expected_ast, sql)
 
         sql = 'SELECT @@version'
         ast = parse_sql(sql)
         expected_ast = Select(targets=[Variable('version', is_system_var=True)])
-        assert ast.to_tree() == expected_ast.to_tree()
-        assert str(ast).lower() == sql.lower()
-        assert str(ast) == str(expected_ast)
+        compare_ast(ast, expected_ast, sql)
 
     def test_select_varialbe_complex(self):
-        sql = f"""SELECT * FROM tab1 WHERE column1 in (SELECT column2 + @variable FROM t2)"""
+        sql = """SELECT * FROM tab1 WHERE column1 in (SELECT column2 + @variable FROM t2)"""
         ast = parse_sql(sql)
         expected_ast = Select(targets=[Star()],
                               from_table=Identifier('tab1'),
@@ -36,10 +38,7 @@ def test_select_varialbe_complex(self):
                                                                parentheses=True)
                                                     )
                                                     ))
-
-        assert ast.to_tree() == expected_ast.to_tree()
-        assert str(ast).lower() == sql.lower()
-        assert str(ast) == str(expected_ast)
+        compare_ast(ast, expected_ast, sql)
 
     def test_show_index(self):
         sql = "SHOW INDEX FROM `predictors`"
@@ -48,10 +47,7 @@ def test_show_index(self):
             category='INDEX',
             from_table=Identifier('`predictors`')
         )
-
-        assert str(ast).lower() == sql.lower()
-        assert str(ast) == str(expected_ast)
-        assert ast.to_tree() == expected_ast.to_tree()
+        compare_ast(ast, expected_ast, sql)
 
     def test_show_index_from_db(self):
         sql = "SHOW INDEX FROM `predictors` FROM db"
@@ -60,10 +56,7 @@ def test_show_index_from_db(self):
             category='INDEX',
             from_table=Identifier('db.`predictors`'),
         )
-
-        # assert str(ast).lower() == sql.lower()
-        assert str(ast) == str(expected_ast)
-        assert ast.to_tree() == expected_ast.to_tree()
+        compare_ast(ast, expected_ast, sql)
 
     def test_with_rollup(self):
         sql = "SELECT country, SUM(sales) FROM booksales GROUP BY country WITH ROLLUP"
@@ -77,7 +70,23 @@ def test_with_rollup(self):
             from_table=Identifier('booksales'),
             group_by=[Identifier('country', with_rollup=True)]
         )
+        compare_ast(ast, expected_ast, sql)
+
+    def test_with_rollup_multiple_columns(self):
+        """Test WITH ROLLUP with multiple GROUP BY columns"""
+        sql = "SELECT year, country, SUM(sales) FROM booksales GROUP BY year, country WITH ROLLUP"
 
-        assert str(ast).lower() == sql.lower()
-        assert str(ast) == str(expected_ast)
-        assert ast.to_tree() == expected_ast.to_tree()
+        ast = parse_sql(sql)
+        expected_ast = Select(
+            targets=[
+                Identifier('year'),
+                Identifier('country'),
+                Function(op='SUM', args=[Identifier('sales')])
+            ],
+            from_table=Identifier('booksales'),
+            group_by=[
+                Identifier('year'),
+                Identifier('country', with_rollup=True)
+            ]
+        )
+        compare_ast(ast, expected_ast, sql)