Support PostgreSQL array subquery constructor (#566)

This commit is contained in:
Alex Qyoun-ae 2022-08-12 22:08:56 +04:00 committed by GitHub
parent 61dc3e9bf9
commit 31ba0012f7
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 101 additions and 2 deletions

View file

@ -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 (")?;

View file

@ -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)?;

View file

@ -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");

View file

@ -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

View file

@ -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'");