Skip to content

Commit

Permalink
Add support for MSSQL's JSON_ARRAY/JSON_OBJECT expr (apache#1507)
Browse files Browse the repository at this point in the history
Co-authored-by: Ifeanyi Ubah <[email protected]>
  • Loading branch information
gaoqiangz and iffyio authored Nov 18, 2024
1 parent f961efc commit 92be237
Show file tree
Hide file tree
Showing 9 changed files with 617 additions and 44 deletions.
56 changes: 55 additions & 1 deletion src/ast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5449,6 +5449,8 @@ pub enum FunctionArgOperator {
RightArrow,
/// function(arg1 := value1)
Assignment,
/// function(arg1 : value1)
Colon,
}

impl fmt::Display for FunctionArgOperator {
Expand All @@ -5457,6 +5459,7 @@ impl fmt::Display for FunctionArgOperator {
FunctionArgOperator::Equals => f.write_str("="),
FunctionArgOperator::RightArrow => f.write_str("=>"),
FunctionArgOperator::Assignment => f.write_str(":="),
FunctionArgOperator::Colon => f.write_str(":"),
}
}
}
Expand All @@ -5465,11 +5468,22 @@ impl fmt::Display for FunctionArgOperator {
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub enum FunctionArg {
/// `name` is identifier
///
/// Enabled when `Dialect::supports_named_fn_args_with_expr_name` returns 'false'
Named {
name: Ident,
arg: FunctionArgExpr,
operator: FunctionArgOperator,
},
/// `name` is arbitrary expression
///
/// Enabled when `Dialect::supports_named_fn_args_with_expr_name` returns 'true'
ExprNamed {
name: Expr,
arg: FunctionArgExpr,
operator: FunctionArgOperator,
},
Unnamed(FunctionArgExpr),
}

Expand All @@ -5481,6 +5495,11 @@ impl fmt::Display for FunctionArg {
arg,
operator,
} => write!(f, "{name} {operator} {arg}"),
FunctionArg::ExprNamed {
name,
arg,
operator,
} => write!(f, "{name} {operator} {arg}"),
FunctionArg::Unnamed(unnamed_arg) => write!(f, "{unnamed_arg}"),
}
}
Expand Down Expand Up @@ -5619,7 +5638,10 @@ impl fmt::Display for FunctionArgumentList {
}
write!(f, "{}", display_comma_separated(&self.args))?;
if !self.clauses.is_empty() {
write!(f, " {}", display_separated(&self.clauses, " "))?;
if !self.args.is_empty() {
write!(f, " ")?;
}
write!(f, "{}", display_separated(&self.clauses, " "))?;
}
Ok(())
}
Expand Down Expand Up @@ -5661,6 +5683,11 @@ pub enum FunctionArgumentClause {
///
/// [`GROUP_CONCAT`]: https://dev.mysql.com/doc/refman/8.0/en/aggregate-functions.html#function_group-concat
Separator(Value),
/// The json-null-clause to the [`JSON_ARRAY`]/[`JSON_OBJECT`] function in MSSQL.
///
/// [`JSON_ARRAY`]: <https://learn.microsoft.com/en-us/sql/t-sql/functions/json-array-transact-sql?view=sql-server-ver16>
/// [`JSON_OBJECT`]: <https://learn.microsoft.com/en-us/sql/t-sql/functions/json-object-transact-sql?view=sql-server-ver16>
JsonNullClause(JsonNullClause),
}

impl fmt::Display for FunctionArgumentClause {
Expand All @@ -5676,6 +5703,7 @@ impl fmt::Display for FunctionArgumentClause {
FunctionArgumentClause::OnOverflow(on_overflow) => write!(f, "{on_overflow}"),
FunctionArgumentClause::Having(bound) => write!(f, "{bound}"),
FunctionArgumentClause::Separator(sep) => write!(f, "SEPARATOR {sep}"),
FunctionArgumentClause::JsonNullClause(null_clause) => write!(f, "{null_clause}"),
}
}
}
Expand Down Expand Up @@ -7564,6 +7592,32 @@ impl fmt::Display for ShowStatementIn {
}
}

/// MSSQL's json null clause
///
/// ```plaintext
/// <json_null_clause> ::=
/// NULL ON NULL
/// | ABSENT ON NULL
/// ```
///
/// <https://learn.microsoft.com/en-us/sql/t-sql/functions/json-object-transact-sql?view=sql-server-ver16#json_null_clause>
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub enum JsonNullClause {
NullOnNull,
AbsentOnNull,
}

impl Display for JsonNullClause {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
JsonNullClause::NullOnNull => write!(f, "NULL ON NULL"),
JsonNullClause::AbsentOnNull => write!(f, "ABSENT ON NULL"),
}
}
}

#[cfg(test)]
mod tests {
use super::*;
Expand Down
4 changes: 4 additions & 0 deletions src/dialect/duckdb.rs
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,10 @@ impl Dialect for DuckDbDialect {
true
}

fn supports_named_fn_args_with_assignment_operator(&self) -> bool {
true
}

// DuckDB uses this syntax for `STRUCT`s.
//
// https://duckdb.org/docs/sql/data_types/struct.html#creating-structs
Expand Down
4 changes: 4 additions & 0 deletions src/dialect/generic.rs
Original file line number Diff line number Diff line change
Expand Up @@ -119,4 +119,8 @@ impl Dialect for GenericDialect {
fn supports_load_extension(&self) -> bool {
true
}

fn supports_named_fn_args_with_assignment_operator(&self) -> bool {
true
}
}
25 changes: 24 additions & 1 deletion src/dialect/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -231,11 +231,34 @@ pub trait Dialect: Debug + Any {
false
}

/// Returns true if the dialect supports named arguments of the form FUN(a = '1', b = '2').
/// Returns true if the dialect supports named arguments of the form `FUN(a = '1', b = '2')`.
fn supports_named_fn_args_with_eq_operator(&self) -> bool {
false
}

/// Returns true if the dialect supports named arguments of the form `FUN(a : '1', b : '2')`.
fn supports_named_fn_args_with_colon_operator(&self) -> bool {
false
}

/// Returns true if the dialect supports named arguments of the form `FUN(a := '1', b := '2')`.
fn supports_named_fn_args_with_assignment_operator(&self) -> bool {
false
}

/// Returns true if the dialect supports named arguments of the form `FUN(a => '1', b => '2')`.
fn supports_named_fn_args_with_rarrow_operator(&self) -> bool {
true
}

/// Returns true if dialect supports argument name as arbitrary expression.
/// e.g. `FUN(LOWER('a'):'1', b:'2')`
/// Such function arguments are represented in the AST by the `FunctionArg::ExprNamed` variant,
/// otherwise use the `FunctionArg::Named` variant (compatible reason).
fn supports_named_fn_args_with_expr_name(&self) -> bool {
false
}

/// Returns true if the dialect supports identifiers starting with a numeric
/// prefix such as tables named `59901_user_login`
fn supports_numeric_prefix(&self) -> bool {
Expand Down
12 changes: 12 additions & 0 deletions src/dialect/mssql.rs
Original file line number Diff line number Diff line change
Expand Up @@ -66,4 +66,16 @@ impl Dialect for MsSqlDialect {
fn supports_methods(&self) -> bool {
true
}

fn supports_named_fn_args_with_colon_operator(&self) -> bool {
true
}

fn supports_named_fn_args_with_expr_name(&self) -> bool {
true
}

fn supports_named_fn_args_with_rarrow_operator(&self) -> bool {
false
}
}
1 change: 1 addition & 0 deletions src/keywords.rs
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,7 @@ macro_rules! define_keywords {
define_keywords!(
ABORT,
ABS,
ABSENT,
ABSOLUTE,
ACCESS,
ACCOUNT,
Expand Down
115 changes: 74 additions & 41 deletions src/parser/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11321,45 +11321,58 @@ impl<'a> Parser<'a> {
}

pub fn parse_function_args(&mut self) -> Result<FunctionArg, ParserError> {
if self.peek_nth_token(1) == Token::RArrow {
let name = self.parse_identifier(false)?;

self.expect_token(&Token::RArrow)?;
let arg = self.parse_wildcard_expr()?.into();

Ok(FunctionArg::Named {
name,
arg,
operator: FunctionArgOperator::RightArrow,
})
} else if self.dialect.supports_named_fn_args_with_eq_operator()
&& self.peek_nth_token(1) == Token::Eq
{
let name = self.parse_identifier(false)?;

self.expect_token(&Token::Eq)?;
let arg = self.parse_wildcard_expr()?.into();

Ok(FunctionArg::Named {
name,
arg,
operator: FunctionArgOperator::Equals,
})
} else if dialect_of!(self is DuckDbDialect | GenericDialect)
&& self.peek_nth_token(1) == Token::Assignment
{
let name = self.parse_identifier(false)?;

self.expect_token(&Token::Assignment)?;
let arg = self.parse_expr()?.into();

Ok(FunctionArg::Named {
name,
arg,
operator: FunctionArgOperator::Assignment,
})
let arg = if self.dialect.supports_named_fn_args_with_expr_name() {
self.maybe_parse(|p| {
let name = p.parse_expr()?;
let operator = p.parse_function_named_arg_operator()?;
let arg = p.parse_wildcard_expr()?.into();
Ok(FunctionArg::ExprNamed {
name,
arg,
operator,
})
})?
} else {
Ok(FunctionArg::Unnamed(self.parse_wildcard_expr()?.into()))
self.maybe_parse(|p| {
let name = p.parse_identifier(false)?;
let operator = p.parse_function_named_arg_operator()?;
let arg = p.parse_wildcard_expr()?.into();
Ok(FunctionArg::Named {
name,
arg,
operator,
})
})?
};
if let Some(arg) = arg {
return Ok(arg);
}
Ok(FunctionArg::Unnamed(self.parse_wildcard_expr()?.into()))
}

fn parse_function_named_arg_operator(&mut self) -> Result<FunctionArgOperator, ParserError> {
let tok = self.next_token();
match tok.token {
Token::RArrow if self.dialect.supports_named_fn_args_with_rarrow_operator() => {
Ok(FunctionArgOperator::RightArrow)
}
Token::Eq if self.dialect.supports_named_fn_args_with_eq_operator() => {
Ok(FunctionArgOperator::Equals)
}
Token::Assignment
if self
.dialect
.supports_named_fn_args_with_assignment_operator() =>
{
Ok(FunctionArgOperator::Assignment)
}
Token::Colon if self.dialect.supports_named_fn_args_with_colon_operator() => {
Ok(FunctionArgOperator::Colon)
}
_ => {
self.prev_token();
self.expected("argument operator", tok)
}
}
}

Expand Down Expand Up @@ -11403,19 +11416,24 @@ impl<'a> Parser<'a> {
/// FIRST_VALUE(x IGNORE NULL);
/// ```
fn parse_function_argument_list(&mut self) -> Result<FunctionArgumentList, ParserError> {
let mut clauses = vec![];

// For MSSQL empty argument list with json-null-clause case, e.g. `JSON_ARRAY(NULL ON NULL)`
if let Some(null_clause) = self.parse_json_null_clause() {
clauses.push(FunctionArgumentClause::JsonNullClause(null_clause));
}

if self.consume_token(&Token::RParen) {
return Ok(FunctionArgumentList {
duplicate_treatment: None,
args: vec![],
clauses: vec![],
clauses,
});
}

let duplicate_treatment = self.parse_duplicate_treatment()?;
let args = self.parse_comma_separated(Parser::parse_function_args)?;

let mut clauses = vec![];

if self.dialect.supports_window_function_null_treatment_arg() {
if let Some(null_treatment) = self.parse_null_treatment()? {
clauses.push(FunctionArgumentClause::IgnoreOrRespectNulls(null_treatment));
Expand Down Expand Up @@ -11456,6 +11474,10 @@ impl<'a> Parser<'a> {
clauses.push(FunctionArgumentClause::OnOverflow(on_overflow));
}

if let Some(null_clause) = self.parse_json_null_clause() {
clauses.push(FunctionArgumentClause::JsonNullClause(null_clause));
}

self.expect_token(&Token::RParen)?;
Ok(FunctionArgumentList {
duplicate_treatment,
Expand All @@ -11464,6 +11486,17 @@ impl<'a> Parser<'a> {
})
}

/// Parses MSSQL's json-null-clause
fn parse_json_null_clause(&mut self) -> Option<JsonNullClause> {
if self.parse_keywords(&[Keyword::ABSENT, Keyword::ON, Keyword::NULL]) {
Some(JsonNullClause::AbsentOnNull)
} else if self.parse_keywords(&[Keyword::NULL, Keyword::ON, Keyword::NULL]) {
Some(JsonNullClause::NullOnNull)
} else {
None
}
}

fn parse_duplicate_treatment(&mut self) -> Result<Option<DuplicateTreatment>, ParserError> {
let loc = self.peek_token().location;
match (
Expand Down
3 changes: 2 additions & 1 deletion tests/sqlparser_common.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4402,8 +4402,9 @@ fn parse_explain_query_plan() {

#[test]
fn parse_named_argument_function() {
let dialects = all_dialects_where(|d| d.supports_named_fn_args_with_rarrow_operator());
let sql = "SELECT FUN(a => '1', b => '2') FROM foo";
let select = verified_only_select(sql);
let select = dialects.verified_only_select(sql);

assert_eq!(
&Expr::Function(Function {
Expand Down
Loading

0 comments on commit 92be237

Please sign in to comment.