Add support for PostgreSQL '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/dialect/mod.rs

@@ -303,6 +303,11 @@ pub trait Dialect: Debug + Any {
         false
     }
 
+    /// Returns true if the dialect supports `IN` with query bodies
+    fn supports_in_query_body(&self) -> bool {
+        false
+    }
+
     /// Returns true if the dialect supports `BEGIN {DEFERRED | IMMEDIATE | EXCLUSIVE | TRY | CATCH} [TRANSACTION]` statements
     fn supports_start_transaction_modifier(&self) -> bool {
         false

src/dialect/postgresql.rs

@@ -258,4 +258,8 @@ impl Dialect for PostgreSqlDialect {
     fn supports_set_names(&self) -> bool {
         true
     }
+
+    fn supports_in_query_body(&self) -> bool {
+        true
+    }
 }

src/parser/mod.rs

@@ -3744,10 +3744,18 @@ 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 {
-                expr: Box::new(expr),
-                subquery: self.parse_query()?,
-                negated,
+            if self.dialect.supports_in_query_body() {
+                Expr::InSubquery {
+                    expr: Box::new(expr),
+                    subquery: self.parse_query_body(self.dialect.prec_unknown())?,
+                    negated,
+                }
+            } else {
+                Expr::InSubquery {
+                    expr: Box::new(expr),
+                    subquery: Box::new(SetExpr::Query(self.parse_query()?)),
+                    negated,
+                }
             }
         } else {
             Expr::InList {

tests/sqlparser_common.rs

@@ -2224,7 +2224,9 @@ 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: Box::new(SetExpr::Query(Box::new(verified_query(
+                "SELECT segm FROM bar"
+            )))),
             negated: false,
         },
         select.selection.unwrap()

tests/sqlparser_postgres.rs

@@ -4063,6 +4063,13 @@ fn parse_update_in_with_subquery() {
     pg_and_generic().verified_stmt(r#"WITH "result" AS (UPDATE "Hero" SET "name" = 'Captain America', "number_of_movies" = "number_of_movies" + 1 WHERE "secret_identity" = 'Sam Wilson' RETURNING "id", "name", "secret_identity", "number_of_movies") SELECT * FROM "result""#);
 }
 
+#[test]
+fn parse_in_with_set_expression() {
+    pg().verified_stmt(
+        r#"SELECT 1 FROM t WHERE t.id IN ((SELECT id FROM t) UNION (SELECT id FROM t))"#,
+    );
+}
+
 #[test]
 fn parse_create_function() {
     let sql = "CREATE FUNCTION add(INTEGER, INTEGER) RETURNS INTEGER LANGUAGE SQL IMMUTABLE STRICT PARALLEL SAFE AS 'select $1 + $2;'";