Support named arguments in function invocations (#250)

This commit supports functions with argument names.

the format is :
"Select some_function( a => exp, b => exp2 .. ) FROM table1
OR
"select * from table(function(a => exp)) f;"

see:
https://jakewheat.github.io/sql-overview/sql-2016-foundation-grammar.html#named-argument-assignment-token
or the motivating example from snowflake:
https://docs.snowflake.com/en/sql-reference/functions/flatten.html
This commit is contained in:
eyalleshem 2020-08-02 08:04:55 +03:00 committed by GitHub
parent 580e4b1d64
commit 1cc3bf4099
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 90 additions and 10 deletions

View file

@ -873,12 +873,28 @@ impl fmt::Display for Assignment {
} }
} }
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub enum FunctionArg {
Named { name: Ident, arg: Expr },
Unnamed(Expr),
}
impl fmt::Display for FunctionArg {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self {
FunctionArg::Named { name, arg } => write!(f, "{} => {}", name, arg),
FunctionArg::Unnamed(unnamed_arg) => write!(f, "{}", unnamed_arg),
}
}
}
/// A function call /// A function call
#[derive(Debug, Clone, PartialEq, Eq, Hash)] #[derive(Debug, Clone, PartialEq, Eq, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct Function { pub struct Function {
pub name: ObjectName, pub name: ObjectName,
pub args: Vec<Expr>, pub args: Vec<FunctionArg>,
pub over: Option<WindowSpec>, pub over: Option<WindowSpec>,
// aggregate functions may specify eg `COUNT(DISTINCT x)` // aggregate functions may specify eg `COUNT(DISTINCT x)`
pub distinct: bool, pub distinct: bool,

View file

@ -226,7 +226,7 @@ pub enum TableFactor {
/// Arguments of a table-valued function, as supported by Postgres /// Arguments of a table-valued function, as supported by Postgres
/// and MSSQL. Note that deprecated MSSQL `FROM foo (NOLOCK)` syntax /// and MSSQL. Note that deprecated MSSQL `FROM foo (NOLOCK)` syntax
/// will also be parsed as `args`. /// will also be parsed as `args`.
args: Vec<Expr>, args: Vec<FunctionArg>,
/// MSSQL-specific `WITH (...)` hints such as NOLOCK. /// MSSQL-specific `WITH (...)` hints such as NOLOCK.
with_hints: Vec<Expr>, with_hints: Vec<Expr>,
}, },

View file

@ -2203,11 +2203,24 @@ impl Parser {
Ok(Assignment { id, value }) Ok(Assignment { id, value })
} }
pub fn parse_optional_args(&mut self) -> Result<Vec<Expr>, ParserError> { fn parse_function_args(&mut self) -> Result<FunctionArg, ParserError> {
if self.peek_nth_token(1) == Token::RArrow {
let name = self.parse_identifier()?;
self.expect_token(&Token::RArrow)?;
let arg = self.parse_expr()?;
Ok(FunctionArg::Named { name, arg })
} else {
Ok(FunctionArg::Unnamed(self.parse_expr()?))
}
}
pub fn parse_optional_args(&mut self) -> Result<Vec<FunctionArg>, ParserError> {
if self.consume_token(&Token::RParen) { if self.consume_token(&Token::RParen) {
Ok(vec![]) Ok(vec![])
} else { } else {
let args = self.parse_comma_separated(Parser::parse_expr)?; let args = self.parse_comma_separated(Parser::parse_function_args)?;
self.expect_token(&Token::RParen)?; self.expect_token(&Token::RParen)?;
Ok(args) Ok(args)
} }

View file

@ -99,6 +99,8 @@ pub enum Token {
LBrace, LBrace,
/// Right brace `}` /// Right brace `}`
RBrace, RBrace,
/// Right Arrow `=>`
RArrow,
} }
impl fmt::Display for Token { impl fmt::Display for Token {
@ -139,6 +141,7 @@ impl fmt::Display for Token {
Token::Pipe => f.write_str("|"), Token::Pipe => f.write_str("|"),
Token::LBrace => f.write_str("{"), Token::LBrace => f.write_str("{"),
Token::RBrace => f.write_str("}"), Token::RBrace => f.write_str("}"),
Token::RArrow => f.write_str("=>"),
} }
} }
} }
@ -400,7 +403,13 @@ impl<'a> Tokenizer<'a> {
_ => Ok(Some(Token::Pipe)), _ => Ok(Some(Token::Pipe)),
} }
} }
'=' => self.consume_and_return(chars, Token::Eq), '=' => {
chars.next(); // consume
match chars.peek() {
Some('>') => self.consume_and_return(chars, Token::RArrow),
_ => Ok(Some(Token::Eq)),
}
}
'.' => self.consume_and_return(chars, Token::Period), '.' => self.consume_and_return(chars, Token::Period),
'!' => { '!' => {
chars.next(); // consume chars.next(); // consume
@ -766,6 +775,23 @@ mod tests {
compare(expected, tokens); compare(expected, tokens);
} }
#[test]
fn tokenize_right_arrow() {
let sql = String::from("FUNCTION(key=>value)");
let dialect = GenericDialect {};
let mut tokenizer = Tokenizer::new(&dialect, &sql);
let tokens = tokenizer.tokenize().unwrap();
let expected = vec![
Token::make_word("FUNCTION", None),
Token::LParen,
Token::make_word("key", None),
Token::RArrow,
Token::make_word("value", None),
Token::RParen,
];
compare(expected, tokens);
}
#[test] #[test]
fn tokenize_is_null() { fn tokenize_is_null() {
let sql = String::from("a IS NULL"); let sql = String::from("a IS NULL");

View file

@ -325,7 +325,7 @@ fn parse_select_count_wildcard() {
assert_eq!( assert_eq!(
&Expr::Function(Function { &Expr::Function(Function {
name: ObjectName(vec![Ident::new("COUNT")]), name: ObjectName(vec![Ident::new("COUNT")]),
args: vec![Expr::Wildcard], args: vec![FunctionArg::Unnamed(Expr::Wildcard)],
over: None, over: None,
distinct: false, distinct: false,
}), }),
@ -340,10 +340,10 @@ fn parse_select_count_distinct() {
assert_eq!( assert_eq!(
&Expr::Function(Function { &Expr::Function(Function {
name: ObjectName(vec![Ident::new("COUNT")]), name: ObjectName(vec![Ident::new("COUNT")]),
args: vec![Expr::UnaryOp { args: vec![FunctionArg::Unnamed(Expr::UnaryOp {
op: UnaryOperator::Plus, op: UnaryOperator::Plus,
expr: Box::new(Expr::Identifier(Ident::new("x"))) expr: Box::new(Expr::Identifier(Ident::new("x")))
}], })],
over: None, over: None,
distinct: true, distinct: true,
}), }),
@ -883,7 +883,7 @@ fn parse_select_having() {
Some(Expr::BinaryOp { Some(Expr::BinaryOp {
left: Box::new(Expr::Function(Function { left: Box::new(Expr::Function(Function {
name: ObjectName(vec![Ident::new("COUNT")]), name: ObjectName(vec![Ident::new("COUNT")]),
args: vec![Expr::Wildcard], args: vec![FunctionArg::Unnamed(Expr::Wildcard)],
over: None, over: None,
distinct: false distinct: false
})), })),
@ -1589,7 +1589,32 @@ fn parse_scalar_function_in_projection() {
assert_eq!( assert_eq!(
&Expr::Function(Function { &Expr::Function(Function {
name: ObjectName(vec![Ident::new("sqrt")]), name: ObjectName(vec![Ident::new("sqrt")]),
args: vec![Expr::Identifier(Ident::new("id"))], args: vec![FunctionArg::Unnamed(Expr::Identifier(Ident::new("id")))],
over: None,
distinct: false,
}),
expr_from_projection(only(&select.projection))
);
}
#[test]
fn parse_named_argument_function() {
let sql = "SELECT FUN(a => '1', b => '2') FROM foo";
let select = verified_only_select(sql);
assert_eq!(
&Expr::Function(Function {
name: ObjectName(vec![Ident::new("FUN")]),
args: vec![
FunctionArg::Named {
name: Ident::new("a"),
arg: Expr::Value(Value::SingleQuotedString("1".to_owned()))
},
FunctionArg::Named {
name: Ident::new("b"),
arg: Expr::Value(Value::SingleQuotedString("2".to_owned()))
},
],
over: None, over: None,
distinct: false, distinct: false,
}), }),