Skip to content

Commit

Permalink
Support PostgreSQL array subquery constructor
Browse files Browse the repository at this point in the history
  • Loading branch information
MazterQyou committed Aug 10, 2022
1 parent 231370a commit c6edf5f
Show file tree
Hide file tree
Showing 4 changed files with 82 additions and 2 deletions.
3 changes: 3 additions & 0 deletions src/ast/mod.rs
Expand Up @@ -355,6 +355,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 @@ -486,6 +488,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 @@ -434,11 +434,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 dialect_of!(self is PostgreSqlDialect | GenericDialect)
&& self.peek_token() == Token::LParen =>
{
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 @@ -910,6 +917,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
2 changes: 1 addition & 1 deletion tests/sqlparser_common.rs
Expand Up @@ -2518,7 +2518,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 c6edf5f

Please sign in to comment.