Skip to content

Commit

Permalink
Support PostgreSQL array subquery constructor (#566)
Browse files Browse the repository at this point in the history
  • Loading branch information
MazterQyou committed Aug 12, 2022
1 parent 61dc3e9 commit 31ba001
Show file tree
Hide file tree
Showing 5 changed files with 101 additions and 2 deletions.
3 changes: 3 additions & 0 deletions src/ast/mod.rs
Expand Up @@ -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.
Expand Down Expand Up @@ -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 (")?;
Expand Down
16 changes: 15 additions & 1 deletion src/parser.rs
Expand Up @@ -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:
Expand Down Expand Up @@ -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)?;
Expand Down
19 changes: 19 additions & 0 deletions tests/sqpparser_clickhouse.rs → tests/sqlparser_clickhouse.rs
Expand Up @@ -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");
Expand Down
2 changes: 1 addition & 1 deletion tests/sqlparser_common.rs
Expand Up @@ -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
Expand Down
63 changes: 63 additions & 0 deletions tests/sqlparser_postgres.rs
Expand Up @@ -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'");
Expand Down

0 comments on commit 31ba001

Please sign in to comment.