Add support for 'IN <SetExpression>'

Author: adamchainz

src/ast/mod.rs

@@ -724,7 +724,7 @@ pub enum Expr {
     /// `[ NOT ] IN (SELECT ...)`
     InSubquery {
         expr: Box<Expr>,
-        subquery: Box<Query>,
+        subquery: Box<SetExpr>,
         negated: bool,
     },
     /// `[ NOT ] IN UNNEST(array_expression)`

src/parser/mod.rs

@@ -3742,23 +3742,21 @@ impl<'a> Parser<'a> {
             });
         }
         self.expect_token(&Token::LParen)?;
-        let in_op = if self.parse_keyword(Keyword::SELECT) || self.parse_keyword(Keyword::WITH) {
-            self.prev_token();
-            Expr::InSubquery {
+        let in_op = match self.maybe_parse(|p| p.parse_query_body(p.dialect.prec_unknown()))? {
+            Some(subquery) => Expr::InSubquery {
                 expr: Box::new(expr),
-                subquery: self.parse_query()?,
+                subquery,
                 negated,
-            }
-        } else {
-            Expr::InList {
+            },
+            None => Expr::InList {
                 expr: Box::new(expr),
                 list: if self.dialect.supports_in_empty_list() {
                     self.parse_comma_separated0(Parser::parse_expr, Token::RParen)?
                 } else {
                     self.parse_comma_separated(Parser::parse_expr)?
                 },
                 negated,
-            }
+            },
         };
         self.expect_token(&Token::RParen)?;
         Ok(in_op)

tests/sqlparser_common.rs

@@ -2224,7 +2224,21 @@ fn parse_in_subquery() {
     assert_eq!(
         Expr::InSubquery {
             expr: Box::new(Expr::Identifier(Ident::new("segment"))),
-            subquery: Box::new(verified_query("SELECT segm FROM bar")),
+            subquery: verified_query("SELECT segm FROM bar").body,
+            negated: false,
+        },
+        select.selection.unwrap()
+    );
+}
+
+#[test]
+fn parse_in_union() {
+    let sql = "SELECT * FROM customers WHERE segment IN ((SELECT segm FROM bar) UNION (SELECT segm FROM bar2))";
+    let select = verified_only_select(sql);
+    assert_eq!(
+        Expr::InSubquery {
+            expr: Box::new(Expr::Identifier(Ident::new("segment"))),
+            subquery: verified_query("(SELECT segm FROM bar) UNION (SELECT segm FROM bar2)").body,
             negated: false,
         },
         select.selection.unwrap()