Merge pull request #14 from mindsdb/support-with-rollup

Support WITH ROLLUP

Author: ea-rus

mindsdb_sql_parser/ast/select/identifier.py

@@ -39,7 +39,7 @@ def get_reserved_words():
 
 
 class Identifier(ASTNode):
-    def __init__(self, path_str=None, parts=None, is_outer=False, *args, **kwargs):
+    def __init__(self, path_str=None, parts=None, is_outer=False, with_rollup=False, *args, **kwargs):
         super().__init__(*args, **kwargs)
         assert path_str or parts, "Either path_str or parts must be provided for an Identifier"
         assert not (path_str and parts), "Provide either path_str or parts, but not both"
@@ -56,6 +56,7 @@ def __init__(self, path_str=None, parts=None, is_outer=False, *args, **kwargs):
         self.is_quoted: List[bool] = is_quoted
         # used to define type of implicit join in oracle
         self.is_outer: bool = is_outer
+        self.with_rollup: bool = with_rollup
 
     @classmethod
     def from_path_str(self, value, *args, **kwargs):

mindsdb_sql_parser/ast/select/select.py

@@ -100,7 +100,7 @@ def to_tree(self, *args, level=0, **kwargs):
         return out_str
 
     def get_string(self, *args, **kwargs):
-
+        from mindsdb_sql_parser.ast.select.identifier import Identifier
         out_str = ''
         if self.cte is not None:
             cte_str = ', '.join([out.to_string() for out in self.cte])
@@ -126,7 +126,13 @@ def get_string(self, *args, **kwargs):
             out_str += f' WHERE {self.where.to_string()}'
 
         if self.group_by is not None:
-            group_by_str = ', '.join([out.to_string() for out in self.group_by])
+            parts = []
+            for item in self.group_by:
+                part = item.to_string()
+                if isinstance(item, Identifier) and item.with_rollup:
+                    part += ' WITH ROLLUP'
+                parts.append(part)
+            group_by_str = ', '.join(parts)
             out_str += f' GROUP BY {group_by_str}'
 
         if self.having is not None:
@@ -147,7 +153,7 @@ def get_string(self, *args, **kwargs):
             out_str += f' {self.mode}'
 
         if self.using is not None:
-            from mindsdb_sql_parser.ast.select.identifier import Identifier
+
 
             using_ar = []
             for key, value in self.using.items():

mindsdb_sql_parser/lexer.py

@@ -51,7 +51,7 @@ class MindsDBLexer(Lexer):
         WITH, SELECT, DISTINCT, FROM, WHERE, AS,
         LIMIT, OFFSET, ASC, DESC, NULLS_FIRST, NULLS_LAST,
         GROUP_BY, HAVING, ORDER_BY,
-        STAR, FOR, UPDATE,
+        STAR, FOR, UPDATE, ROLLUP,
 
         JOIN, INNER, OUTER, CROSS, LEFT, RIGHT, ON, ASOF, LATERAL, IS_OUTER,
 
@@ -233,6 +233,7 @@ class MindsDBLexer(Lexer):
     # FETCH FIRST and FETCH NEXT are the same
     FETCH_FIRST = r'\bFETCH[\s](FIRST|NEXT)\b'
     ROWS_ONLY = r'\bROWS[\s]ONLY\b'
+    ROLLUP = r'\bROLLUP\b'
 
     JOIN = r'\bJOIN\b'
     INNER = r'\bINNER\b'

mindsdb_sql_parser/parser.py

@@ -1775,6 +1775,12 @@ def identifier(self, p):
     def string(self, p):
         return p[0]
 
+    @_('identifier WITH ROLLUP')
+    def identifier(self, p):
+        value = p[0]
+        value.with_rollup = True
+        return value
+
     @_('identifier IS_OUTER')
     def identifier(self, p):
         value = p[0]
@@ -1899,6 +1905,9 @@ def function_name(self, p):
        'EVALUATE',
        'DETECTION',
        'JOB',
+       'ASOF',
+       'LATERAL',
+       'ROLLUP',
        # fixme
        # 'EXCEPT',
        # 'INTERSECT',

tests/test_mysql/test_mysql_parser.py

@@ -1,6 +1,6 @@
 from mindsdb_sql_parser import parse_sql
 from mindsdb_sql_parser.ast import Select, Identifier, BinaryOperation, Star
-from mindsdb_sql_parser.ast import Variable
+from mindsdb_sql_parser.ast import Variable, Function
 from mindsdb_sql_parser.parser import Show
 
 
@@ -64,3 +64,20 @@ def test_show_index_from_db(self):
         # assert str(ast).lower() == sql.lower()
         assert str(ast) == str(expected_ast)
         assert ast.to_tree() == expected_ast.to_tree()
+
+    def test_with_rollup(self):
+        sql = "SELECT country, SUM(sales) FROM booksales GROUP BY country WITH ROLLUP"
+
+        ast = parse_sql(sql)
+        expected_ast = Select(
+            targets=[
+                Identifier('country'),
+                Function(op='SUM', args=[Identifier('sales')])
+            ],
+            from_table=Identifier('booksales'),
+            group_by=[Identifier('country', with_rollup=True)]
+        )
+
+        assert str(ast).lower() == sql.lower()
+        assert str(ast) == str(expected_ast)
+        assert ast.to_tree() == expected_ast.to_tree()