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

UNION can be an expression

Author: ea-rus

mindsdb_sql_parser/__about__.py

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

mindsdb_sql_parser/parser.py

@@ -667,13 +667,11 @@ def update(self, p):
                       from_select=p.select)
 
     # INSERT
-    @_('INSERT INTO identifier LPAREN column_list RPAREN select',
-       'INSERT INTO identifier LPAREN column_list RPAREN union',
-       'INSERT INTO identifier select',
+    @_('INSERT INTO identifier LPAREN column_list RPAREN union',
        'INSERT INTO identifier union')
     def insert(self, p):
         columns = getattr(p, 'column_list', None)
-        query = p.select if hasattr(p, 'select') else p.union
+        query = p.union
         return Insert(table=p.identifier, columns=columns, from_select=query)
 
     @_('INSERT INTO identifier LPAREN column_list RPAREN VALUES expr_list_set',
@@ -1086,33 +1084,24 @@ def database_engine(self, p):
         return {'identifier':p.identifier, 'engine':engine, 'if_not_exists':p.if_not_exists_or_empty}
 
     # Combining
-    @_('select UNION select',
-       'union UNION select',
-       'select UNION ALL select',
+    @_('union UNION 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)
 
-    @_('select INTERSECT select',
-       'union INTERSECT select',
-       'select INTERSECT ALL select',
+    @_('union INTERSECT 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)
 
-    @_('select EXCEPT select',
-       'union EXCEPT select',
-       'select EXCEPT ALL select',
+    @_('union EXCEPT select',
        'union EXCEPT ALL select',
-       'select EXCEPT DISTINCT select',
        'union EXCEPT DISTINCT select')
     def union(self, p):
         unique = not hasattr(p, 'ALL')
@@ -1121,19 +1110,21 @@ def union(self, p):
 
     # tableau
     @_('LPAREN select RPAREN')
-    @_('LPAREN union RPAREN')
     def select(self, p):
         return p[1]
 
+    @_('select')
+    def union(self, p):
+        return p[0]
+
     # WITH
     @_('ctes select')
     def select(self, p):
         select = p.select
         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 union RPAREN')
+    @_('ctes COMMA identifier cte_columns_or_nothing AS LPAREN union RPAREN')
     def ctes(self, p):
         ctes = p.ctes
         ctes = ctes + [
@@ -1144,8 +1135,7 @@ def ctes(self, p):
         ]
         return ctes
 
-    @_('WITH identifier cte_columns_or_nothing AS LPAREN select RPAREN',
-       'WITH identifier cte_columns_or_nothing AS LPAREN union RPAREN')
+    @_('WITH identifier cte_columns_or_nothing AS LPAREN union RPAREN')
     def ctes(self, p):
         return [
             CommonTableExpression(
@@ -1518,11 +1508,11 @@ def window(self, p):
 
     # OPERATIONS
 
-    @_('LPAREN select RPAREN')
+    @_('LPAREN union RPAREN')
     def expr(self, p):
-        select = p.select
-        select.parentheses = True
-        return select
+        union = p.union
+        union.parentheses = True
+        return union
 
     @_('LPAREN expr RPAREN')
     def expr(self, p):

tests/test_base_sql/test_select_structure.py

@@ -877,6 +877,33 @@ def test_is_not_precedence(self):
         assert str(ast) == str(expected_ast)
         assert ast.to_tree() == expected_ast.to_tree()
 
+    def test_in_union(self):
+        sql = f'''
+          SELECT * FROM tbl1 
+          where col in (
+            select id from tbl2
+            union 
+            select id from tbl3
+          )
+        '''
+        ast = parse_sql(sql)
+
+        expected_ast = Select(
+            targets=[Star()],
+            from_table=Identifier('tbl1'),
+            where=BinaryOperation(op='in', args=(
+                Identifier('col'),
+                Union(
+                    left=Select(targets=[Identifier('id')], from_table=Identifier('tbl2')),
+                    right=Select(targets=[Identifier('id')], from_table=Identifier('tbl3')),
+                    parentheses=True
+                )
+            ))
+        )
+
+        assert str(ast) == str(expected_ast)
+        assert ast.to_tree() == expected_ast.to_tree()
+
 
 class TestSelectStructureNoSqlite:
     def test_select_from_plugins(self):