Support PostgreSQL array subquery constructor (apache#566)

Author: MazterQyou

src/ast/mod.rs

@@ -376,6 +376,8 @@ pub enum Expr {
     /// A parenthesized subquery `(SELECT ...)`, used in expression like
     /// `SELECT (subquery) AS x` or `WHERE (subquery) = x`
     Subquery(Box<Query>),
+    /// An array subquery constructor, e.g. `SELECT ARRAY(SELECT 1 UNION SELECT 2)`
+    ArraySubquery(Box<Query>),
     /// The `LISTAGG` function `SELECT LISTAGG(...) WITHIN GROUP (ORDER BY ...)`
     ListAgg(ListAgg),
     /// The `GROUPING SETS` expr.
@@ -573,6 +575,7 @@ impl fmt::Display for Expr {
                 subquery
             ),
             Expr::Subquery(s) => write!(f, "({})", s),
+            Expr::ArraySubquery(s) => write!(f, "ARRAY({})", s),
             Expr::ListAgg(listagg) => write!(f, "{}", listagg),
             Expr::GroupingSets(sets) => {
                 write!(f, "GROUPING SETS (")?;

src/parser.rs

@@ -449,11 +449,18 @@ impl<'a> Parser<'a> {
                 Keyword::TRIM => self.parse_trim_expr(),
                 Keyword::INTERVAL => self.parse_literal_interval(),
                 Keyword::LISTAGG => self.parse_listagg_expr(),
-                // Treat ARRAY[1,2,3] as an array [1,2,3], otherwise try as function call
+                // Treat ARRAY[1,2,3] as an array [1,2,3], otherwise try as subquery or a function call
                 Keyword::ARRAY if self.peek_token() == Token::LBracket => {
                     self.expect_token(&Token::LBracket)?;
                     self.parse_array_expr(true)
                 }
+                Keyword::ARRAY
+                    if self.peek_token() == Token::LParen
+                        && !dialect_of!(self is ClickHouseDialect) =>
+                {
+                    self.expect_token(&Token::LParen)?;
+                    self.parse_array_subquery()
+                }
                 Keyword::NOT => self.parse_not(),
                 // Here `w` is a word, check if it's a part of a multi-part
                 // identifier, a function call, or a simple identifier:
@@ -927,6 +934,13 @@ impl<'a> Parser<'a> {
         }
     }
 
+    // Parses an array constructed from a subquery
+    pub fn parse_array_subquery(&mut self) -> Result<Expr, ParserError> {
+        let query = self.parse_query()?;
+        self.expect_token(&Token::RParen)?;
+        Ok(Expr::ArraySubquery(Box::new(query)))
+    }
+
     /// Parse a SQL LISTAGG expression, e.g. `LISTAGG(...) WITHIN GROUP (ORDER BY ...)`.
     pub fn parse_listagg_expr(&mut self) -> Result<Expr, ParserError> {
         self.expect_token(&Token::LParen)?;

tests/sqpparser_clickhouse.rs renamed to tests/sqlparser_clickhouse.rs

@@ -121,6 +121,25 @@ fn parse_array_expr() {
     )
 }
 
+#[test]
+fn parse_array_fn() {
+    let sql = "SELECT array(x1, x2) FROM foo";
+    let select = clickhouse().verified_only_select(sql);
+    assert_eq!(
+        &Expr::Function(Function {
+            name: ObjectName(vec![Ident::new("array")]),
+            args: vec![
+                FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Identifier(Ident::new("x1")))),
+                FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Identifier(Ident::new("x2")))),
+            ],
+            over: None,
+            distinct: false,
+            special: false,
+        }),
+        expr_from_projection(only(&select.projection))
+    );
+}
+
 #[test]
 fn parse_kill() {
     let stmt = clickhouse().verified_stmt("KILL MUTATION 5");

tests/sqlparser_common.rs

@@ -2597,7 +2597,7 @@ fn parse_bad_constraint() {
 
 #[test]
 fn parse_scalar_function_in_projection() {
-    let names = vec!["sqrt", "array", "foo"];
+    let names = vec!["sqrt", "foo"];
 
     for function_name in names {
         // like SELECT sqrt(id) FROM foo

tests/sqlparser_postgres.rs

@@ -1241,6 +1241,69 @@ fn parse_array_index_expr() {
     );
 }
 
+#[test]
+fn parse_array_subquery_expr() {
+    let sql = "SELECT ARRAY(SELECT 1 UNION SELECT 2)";
+    let select = pg().verified_only_select(sql);
+    assert_eq!(
+        &Expr::ArraySubquery(Box::new(Query {
+            with: None,
+            body: Box::new(SetExpr::SetOperation {
+                op: SetOperator::Union,
+                all: false,
+                left: Box::new(SetExpr::Select(Box::new(Select {
+                    distinct: false,
+                    top: None,
+                    projection: vec![SelectItem::UnnamedExpr(Expr::Value(Value::Number(
+                        #[cfg(not(feature = "bigdecimal"))]
+                        "1".to_string(),
+                        #[cfg(feature = "bigdecimal")]
+                        bigdecimal::BigDecimal::from(1),
+                        false,
+                    )))],
+                    into: None,
+                    from: vec![],
+                    lateral_views: vec![],
+                    selection: None,
+                    group_by: vec![],
+                    cluster_by: vec![],
+                    distribute_by: vec![],
+                    sort_by: vec![],
+                    having: None,
+                    qualify: None,
+                }))),
+                right: Box::new(SetExpr::Select(Box::new(Select {
+                    distinct: false,
+                    top: None,
+                    projection: vec![SelectItem::UnnamedExpr(Expr::Value(Value::Number(
+                        #[cfg(not(feature = "bigdecimal"))]
+                        "2".to_string(),
+                        #[cfg(feature = "bigdecimal")]
+                        bigdecimal::BigDecimal::from(2),
+                        false,
+                    )))],
+                    into: None,
+                    from: vec![],
+                    lateral_views: vec![],
+                    selection: None,
+                    group_by: vec![],
+                    cluster_by: vec![],
+                    distribute_by: vec![],
+                    sort_by: vec![],
+                    having: None,
+                    qualify: None,
+                }))),
+            }),
+            order_by: vec![],
+            limit: None,
+            offset: None,
+            fetch: None,
+            lock: None,
+        })),
+        expr_from_projection(only(&select.projection)),
+    );
+}
+
 #[test]
 fn test_transaction_statement() {
     let statement = pg().verified_stmt("SET TRANSACTION SNAPSHOT '000003A1-1'");