Merge pull request #81 from mindsdb/union-in-in

Fix insert from select with using

Author: ea-rus

mindsdb_sql_parser/__about__.py

@@ -1,6 +1,6 @@
 __title__ = 'mindsdb_sql_parser'
 __package_name__ = 'mindsdb_sql_parser'
-__version__ = '0.13.5'
+__version__ = '0.13.6'
 __description__ = "Mindsdb SQL parser"
 __email__ = "jorge@mindsdb.com"
 __author__ = 'MindsDB Inc'

mindsdb_sql_parser/parser.py

@@ -210,11 +210,11 @@ def create_agent(self, p):
             params=params,
             if_not_exists=p.if_not_exists_or_empty
         )
-    
+
     @_('DROP AGENT if_exists_or_empty identifier')
     def drop_agent(self, p):
         return DropAgent(name=p.identifier, if_exists=p.if_exists_or_empty)
-    
+
     @_('UPDATE AGENT identifier SET kw_parameter_list')
     @_('ALTER AGENT identifier USING kw_parameter_list')
     def update_agent(self, p):
@@ -667,11 +667,13 @@ def update(self, p):
                       from_select=p.select)
 
     # INSERT
-    @_('INSERT INTO identifier LPAREN column_list RPAREN union',
+    @_('INSERT INTO identifier LPAREN column_list RPAREN select',
+       'INSERT INTO identifier LPAREN column_list RPAREN union',
+       'INSERT INTO identifier select',
        'INSERT INTO identifier union')
     def insert(self, p):
         columns = getattr(p, 'column_list', None)
-        query = p.union
+        query = p.select if hasattr(p, 'select') else p.union
         return Insert(table=p.identifier, columns=columns, from_select=query)
 
     @_('INSERT INTO identifier LPAREN column_list RPAREN VALUES expr_list_set',
@@ -1084,38 +1086,49 @@ def database_engine(self, p):
         return {'identifier':p.identifier, 'engine':engine, 'if_not_exists':p.if_not_exists_or_empty}
 
     # Combining
-    @_('union UNION select',
+    @_('select UNION select',
+       'union UNION select',
+       'select UNION ALL select',
        'union UNION ALL select',
+       'select UNION DISTINCT select',
        'union UNION DISTINCT select')
     def union(self, p):
         unique = not hasattr(p, 'ALL')
         distinct_key = hasattr(p, 'DISTINCT')
         return Union(left=p[0], right=p[-1], unique=unique, distinct_key=distinct_key)
 
-    @_('union INTERSECT select',
+    @_('select INTERSECT select',
+       'union INTERSECT select',
+       'select INTERSECT ALL select',
        'union INTERSECT ALL select',
+       'select INTERSECT DISTINCT select',
        'union INTERSECT DISTINCT select')
     def union(self, p):
         unique = not hasattr(p, 'ALL')
         distinct_key = hasattr(p, 'DISTINCT')
         return Intersect(left=p[0], right=p[-1], unique=unique, distinct_key=distinct_key)
 
-    @_('union EXCEPT select',
+    @_('select EXCEPT select',
+       'union EXCEPT select',
+       'select EXCEPT ALL select',
        'union EXCEPT ALL select',
+       'select EXCEPT DISTINCT select',
        'union EXCEPT DISTINCT select')
     def union(self, p):
         unique = not hasattr(p, 'ALL')
         distinct_key = hasattr(p, 'DISTINCT')
         return Except(left=p[0], right=p[-1], unique=unique, distinct_key=distinct_key)
 
     # tableau
-    @_('LPAREN union RPAREN')
-    def union(self, p):
+    @_('LPAREN select RPAREN')
+    def select(self, p):
         return p[1]
 
-    @_('select')
+    @_('LPAREN union RPAREN')
     def union(self, p):
-        return p[0]
+        node = p[1]
+        node.parentheses = True
+        return node
 
     # WITH
     @_('ctes select')
@@ -1124,7 +1137,8 @@ def select(self, p):
         select.cte = p.ctes
         return select
 
-    @_('ctes COMMA identifier cte_columns_or_nothing AS LPAREN union 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 + [
@@ -1135,7 +1149,8 @@ def ctes(self, p):
         ]
         return ctes
 
-    @_('WITH identifier cte_columns_or_nothing AS LPAREN union RPAREN')
+    @_('WITH identifier cte_columns_or_nothing AS LPAREN select RPAREN',
+       'WITH identifier cte_columns_or_nothing AS LPAREN union RPAREN')
     def ctes(self, p):
         return [
             CommonTableExpression(
@@ -1508,11 +1523,12 @@ def window(self, p):
 
     # OPERATIONS
 
-    @_('LPAREN union RPAREN')
+    @_('LPAREN select RPAREN',
+       'LPAREN union RPAREN')
     def expr(self, p):
-        union = p.union
-        union.parentheses = True
-        return union
+        node = p[1]
+        node.parentheses = True
+        return node
 
     @_('LPAREN expr RPAREN')
     def expr(self, p):

tests/test_base_sql/test_union.py

@@ -80,3 +80,33 @@ def test_union_alias(self):
         assert ast.to_tree() == expected_ast.to_tree()
         assert str(ast) == str(expected_ast)
 
+    def test_union_alias_parens(self):
+        sql = f'''
+        SELECT * FROM (
+           (
+            SELECT col1 FROM tab1
+            UNION
+            SELECT col1 FROM tab2
+           )
+           UNION
+            SELECT col1 FROM tab3
+        ) AS tt
+        '''
+        ast = parse_sql(sql)
+
+        expected_ast = Select(
+            targets=[Star()],
+            from_table=Union(
+                left=Union(
+                    left=Select(targets=[Identifier('col1')], from_table=Identifier('tab1')),
+                    right=Select(targets=[Identifier('col1')], from_table=Identifier('tab2')),
+                    parentheses=True
+                ),
+                right=Select(targets=[Identifier('col1')], from_table=Identifier('tab3')),
+                parentheses=True,
+                alias=Identifier('tt')
+            )
+        )
+
+        assert str(ast) == str(expected_ast)
+        assert ast.to_tree() == expected_ast.to_tree()

tests/test_mindsdb/test_selects.py

@@ -71,7 +71,7 @@ def test_select_using(self):
            using p1=1, p2='2', p3=column
         """
         ast = parse_sql(sql)
-        expected_ast = Select(
+        expected_select = Select(
             targets=[Identifier('status')],
             from_table=Identifier('tbl1'),
             group_by=[Constant(1)],
@@ -82,6 +82,23 @@ def test_select_using(self):
             }
         )
 
+        assert str(ast) == str(expected_select)
+        assert ast.to_tree() == expected_select.to_tree()
+
+        sql = """
+           insert into int2.table2
+            (
+             SELECT status FROM tbl1
+             group by 1
+            )
+           using p1=1, p2='2', p3=column
+        """
+        ast = parse_sql(sql)
+        expected_ast = Insert(
+            table=Identifier('int2.table2'),
+            from_select=expected_select
+        )
+
         assert str(ast) == str(expected_ast)
         assert ast.to_tree() == expected_ast.to_tree()