Skip to content

Commit 7cdb62e

Browse files
adamchainzadamchainz
committed
Add support for PostgreSQL JSON function 'RETURNING' clauses
1 parent 779dcf9 commit 7cdb62e

File tree

4 files changed

+139
-7
lines changed

4 files changed

+139
-7
lines changed

‎src/ast/mod.rs‎

Lines changed: 30 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7817,11 +7817,16 @@ pub enum FunctionArgumentClause {
78177817
///
78187818
/// [`GROUP_CONCAT`]: https://dev.mysql.com/doc/refman/8.0/en/aggregate-functions.html#function_group-concat
78197819
Separator(Value),
7820-
/// The json-null-clause to the [`JSON_ARRAY`]/[`JSON_OBJECT`] function in MSSQL.
7820+
/// The `ON NULL` clause for some JSON functions.
78217821
///
7822-
/// [`JSON_ARRAY`]: <https://learn.microsoft.com/en-us/sql/t-sql/functions/json-array-transact-sql?view=sql-server-ver16>
7823-
/// [`JSON_OBJECT`]: <https://learn.microsoft.com/en-us/sql/t-sql/functions/json-object-transact-sql?view=sql-server-ver16>
7822+
/// [MSSQL `JSON_ARRAY`](https://learn.microsoft.com/en-us/sql/t-sql/functions/json-array-transact-sql?view=sql-server-ver16)
7823+
/// [MSSQL `JSON_OBJECT`](https://learn.microsoft.com/en-us/sql/t-sql/functions/json-object-transact-sql?view=sql-server-ver16>)
7824+
/// [PostgreSQL JSON functions](https://www.postgresql.org/docs/current/functions-json.html#FUNCTIONS-JSON-PROCESSING)
78247825
JsonNullClause(JsonNullClause),
7826+
/// The `RETURNING` clause for some JSON functions in PostgreSQL
7827+
///
7828+
/// [`JSON_OBJECT`](https://www.postgresql.org/docs/current/functions-json.html#:~:text=json_object)
7829+
JsonReturningClause(JsonReturningClause),
78257830
}
78267831

78277832
impl fmt::Display for FunctionArgumentClause {
@@ -7838,6 +7843,9 @@ impl fmt::Display for FunctionArgumentClause {
78387843
FunctionArgumentClause::Having(bound) => write!(f, "{bound}"),
78397844
FunctionArgumentClause::Separator(sep) => write!(f, "SEPARATOR {sep}"),
78407845
FunctionArgumentClause::JsonNullClause(null_clause) => write!(f, "{null_clause}"),
7846+
FunctionArgumentClause::JsonReturningClause(returning_clause) => {
7847+
write!(f, "{returning_clause}")
7848+
}
78417849
}
78427850
}
78437851
}
@@ -10172,6 +10180,25 @@ impl Display for JsonNullClause {
1017210180
}
1017310181
}
1017410182

10183+
/// PostgreSQL JSON function RETURNING clause
10184+
///
10185+
/// Example:
10186+
/// ```sql
10187+
/// JSON_OBJECT('a': 1 RETURNING jsonb)
10188+
/// ```
10189+
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
10190+
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
10191+
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
10192+
pub struct JsonReturningClause {
10193+
pub data_type: DataType,
10194+
}
10195+
10196+
impl Display for JsonReturningClause {
10197+
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
10198+
write!(f, "RETURNING {}", self.data_type)
10199+
}
10200+
}
10201+
1017510202
/// rename object definition
1017610203
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
1017710204
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]

‎src/ast/spans.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1791,6 +1791,7 @@ impl Spanned for FunctionArgumentClause {
17911791
FunctionArgumentClause::Having(HavingBound(_kind, expr)) => expr.span(),
17921792
FunctionArgumentClause::Separator(value) => value.span(),
17931793
FunctionArgumentClause::JsonNullClause(_) => Span::empty(),
1794+
FunctionArgumentClause::JsonReturningClause(_) => Span::empty(),
17941795
}
17951796
}
17961797
}

‎src/parser/mod.rs‎

Lines changed: 26 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15573,7 +15573,7 @@ impl<'a> Parser<'a> {
1557315573
Ok(TableFunctionArgs { args, settings })
1557415574
}
1557515575

15576-
/// Parses a potentially empty list of arguments to a window function
15576+
/// Parses a potentially empty list of arguments to a function
1557715577
/// (including the closing parenthesis).
1557815578
///
1557915579
/// Examples:
@@ -15584,11 +15584,18 @@ impl<'a> Parser<'a> {
1558415584
fn parse_function_argument_list(&mut self) -> Result<FunctionArgumentList, ParserError> {
1558515585
let mut clauses = vec![];
1558615586

15587-
// For MSSQL empty argument list with json-null-clause case, e.g. `JSON_ARRAY(NULL ON NULL)`
15587+
// Handle clauses that may exist with an empty argument list
15588+
1558815589
if let Some(null_clause) = self.parse_json_null_clause() {
1558915590
clauses.push(FunctionArgumentClause::JsonNullClause(null_clause));
1559015591
}
1559115592

15593+
if let Some(json_returning_clause) = self.maybe_parse_json_returning_clause()? {
15594+
clauses.push(FunctionArgumentClause::JsonReturningClause(
15595+
json_returning_clause,
15596+
));
15597+
}
15598+
1559215599
if self.consume_token(&Token::RParen) {
1559315600
return Ok(FunctionArgumentList {
1559415601
duplicate_treatment: None,
@@ -15644,6 +15651,12 @@ impl<'a> Parser<'a> {
1564415651
clauses.push(FunctionArgumentClause::JsonNullClause(null_clause));
1564515652
}
1564615653

15654+
if let Some(json_returning_clause) = self.maybe_parse_json_returning_clause()? {
15655+
clauses.push(FunctionArgumentClause::JsonReturningClause(
15656+
json_returning_clause,
15657+
));
15658+
}
15659+
1564715660
self.expect_token(&Token::RParen)?;
1564815661
Ok(FunctionArgumentList {
1564915662
duplicate_treatment,
@@ -15652,7 +15665,6 @@ impl<'a> Parser<'a> {
1565215665
})
1565315666
}
1565415667

15655-
/// Parses MSSQL's json-null-clause
1565615668
fn parse_json_null_clause(&mut self) -> Option<JsonNullClause> {
1565715669
if self.parse_keywords(&[Keyword::ABSENT, Keyword::ON, Keyword::NULL]) {
1565815670
Some(JsonNullClause::AbsentOnNull)
@@ -15663,6 +15675,17 @@ impl<'a> Parser<'a> {
1566315675
}
1566415676
}
1566515677

15678+
fn maybe_parse_json_returning_clause(
15679+
&mut self,
15680+
) -> Result<Option<JsonReturningClause>, ParserError> {
15681+
if self.parse_keyword(Keyword::RETURNING) {
15682+
let data_type = self.parse_data_type()?;
15683+
Ok(Some(JsonReturningClause { data_type }))
15684+
} else {
15685+
Ok(None)
15686+
}
15687+
}
15688+
1566615689
fn parse_duplicate_treatment(&mut self) -> Result<Option<DuplicateTreatment>, ParserError> {
1566715690
let loc = self.peek_token().span.start;
1566815691
match (

‎tests/sqlparser_postgres.rs‎

Lines changed: 82 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3351,7 +3351,31 @@ fn test_json() {
33513351
}
33523352

33533353
#[test]
3354-
fn test_fn_arg_with_value_operator() {
3354+
fn json_object_colon_syntax() {
3355+
match pg().verified_expr("JSON_OBJECT('name' : 'value')") {
3356+
Expr::Function(Function {
3357+
args: FunctionArguments::List(FunctionArgumentList { args, .. }),
3358+
..
3359+
}) => {
3360+
assert!(
3361+
matches!(
3362+
&args[..],
3363+
&[FunctionArg::ExprNamed {
3364+
operator: FunctionArgOperator::Colon,
3365+
..
3366+
}]
3367+
),
3368+
"Invalid function argument: {args:?}"
3369+
);
3370+
}
3371+
other => panic!(
3372+
"Expected: JSON_OBJECT('name' : 'value') to be parsed as a function, but got {other:?}"
3373+
),
3374+
}
3375+
}
3376+
3377+
#[test]
3378+
fn json_object_value_syntax() {
33553379
match pg().verified_expr("JSON_OBJECT('name' VALUE 'value')") {
33563380
Expr::Function(Function { args: FunctionArguments::List(FunctionArgumentList { args, .. }), .. }) => {
33573381
assert!(matches!(
@@ -3363,6 +3387,63 @@ fn test_fn_arg_with_value_operator() {
33633387
}
33643388
}
33653389

3390+
#[test]
3391+
fn parse_json_object() {
3392+
let sql = "JSON_OBJECT('name' VALUE 'value' NULL ON NULL)";
3393+
let expr = pg().verified_expr(sql);
3394+
assert!(
3395+
matches!(
3396+
expr.clone(),
3397+
Expr::Function(Function {
3398+
name: ObjectName(parts),
3399+
args: FunctionArguments::List(FunctionArgumentList { args, clauses, .. }),
3400+
..
3401+
}) if parts == vec![ObjectNamePart::Identifier(Ident::new("JSON_OBJECT"))]
3402+
&& matches!(
3403+
&args[..],
3404+
&[FunctionArg::ExprNamed { operator: FunctionArgOperator::Value, .. }]
3405+
)
3406+
&& clauses == vec![FunctionArgumentClause::JsonNullClause(JsonNullClause::NullOnNull)]
3407+
),
3408+
"Failed to parse JSON_OBJECT with expected structure, got: {expr:?}"
3409+
);
3410+
3411+
let sql = "JSON_OBJECT('name' VALUE 'value' RETURNING JSONB)";
3412+
let expr = pg().verified_expr(sql);
3413+
assert!(
3414+
matches!(
3415+
expr.clone(),
3416+
Expr::Function(Function {
3417+
name: ObjectName(parts),
3418+
args: FunctionArguments::List(FunctionArgumentList { args, clauses, .. }),
3419+
..
3420+
}) if parts == vec![ObjectNamePart::Identifier(Ident::new("JSON_OBJECT"))]
3421+
&& matches!(
3422+
&args[..],
3423+
&[FunctionArg::ExprNamed { operator: FunctionArgOperator::Value, .. }]
3424+
)
3425+
&& clauses == vec![FunctionArgumentClause::JsonReturningClause(JsonReturningClause { data_type: DataType::JSONB })]
3426+
),
3427+
"Failed to parse JSON_OBJECT with expected structure, got: {expr:?}"
3428+
);
3429+
3430+
let sql = "JSON_OBJECT(RETURNING JSONB)";
3431+
let expr = pg().verified_expr(sql);
3432+
assert!(
3433+
matches!(
3434+
expr.clone(),
3435+
Expr::Function(Function {
3436+
name: ObjectName(parts),
3437+
args: FunctionArguments::List(FunctionArgumentList { args, clauses, .. }),
3438+
..
3439+
}) if parts == vec![ObjectNamePart::Identifier(Ident::new("JSON_OBJECT"))]
3440+
&& args.is_empty()
3441+
&& clauses == vec![FunctionArgumentClause::JsonReturningClause(JsonReturningClause { data_type: DataType::JSONB })]
3442+
),
3443+
"Failed to parse JSON_OBJECT with expected structure, got: {expr:?}"
3444+
);
3445+
}
3446+
33663447
#[test]
33673448
fn parse_json_table_is_not_reserved() {
33683449
// JSON_TABLE is not a reserved keyword in PostgreSQL, even though it is in SQL:2023

0 commit comments

Comments
 (0)