Repository: apache/datafusion-sqlparser-rs (rust)

Issue:
PostgreSQL dialect fails to parse `IN` with union
Minimal example:

```sql
SELECT 1 FROM t WHERE t.id IN ((SELECT id FROM t) UNION (SELECT id FROM t))
```

Currently this fails with:

```rs
ParserError("Expected: ), found: UNION at Line: 1, Column: 51")
```

Pull request description:
Add support for parenthesized subquery as `IN` predicate
Fixes #1792.

This is my attempt to hack in support with my minimal Rust skills. I stopped when I realized that my change of `InSubquery` to use `SetExpr` for its subquery broke many tests.

Patch:
diff --git a/src/ast/mod.rs b/src/ast/mod.rs
index 4031936..72fb999 100644
--- a/src/ast/mod.rs
+++ b/src/ast/mod.rs
@@ -728,7 +728,7 @@ pub enum Expr {
     /// `[ NOT ] IN (SELECT ...)`
     InSubquery {
         expr: Box<Expr>,
-        subquery: Box<Query>,
+        subquery: Box<SetExpr>,
         negated: bool,
     },
     /// `[ NOT ] IN UNNEST(array_expression)`
diff --git a/src/parser/mod.rs b/src/parser/mod.rs
index bb989fe..d43e5d0 100644
--- a/src/parser/mod.rs
+++ b/src/parser/mod.rs
@@ -3756,15 +3756,13 @@ 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)?
@@ -3772,7 +3770,7 @@ impl<'a> Parser<'a> {
                     self.parse_comma_separated(Parser::parse_expr)?
                 },
                 negated,
-            }
+            },
         };
         self.expect_token(&Token::RParen)?;
         Ok(in_op)


Test patch:
diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs
index 14716dd..be848a6 100644
--- a/tests/sqlparser_common.rs
+++ b/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()