Skip to content

Commit 6165783

Browse files
adamchainzadamchainz
committed
Add support for PostgreSQL JSON function 'RETURNING' clauses
1 parent 60a5c8d commit 6165783

File tree

4 files changed

+142
-7
lines changed

4 files changed

+142
-7
lines changed

‎src/ast/mod.rs

Lines changed: 30 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7802,11 +7802,16 @@ pub enum FunctionArgumentClause {
78027802
///
78037803
/// [`GROUP_CONCAT`]: https://dev.mysql.com/doc/refman/8.0/en/aggregate-functions.html#function_group-concat
78047804
Separator(Value),
7805-
/// The json-null-clause to the [`JSON_ARRAY`]/[`JSON_OBJECT`] function in MSSQL.
7805+
/// The `ON NULL` clause for some JSON functions in MSSQL and PostgreSQL
78067806
///
7807-
/// [`JSON_ARRAY`]: <https://learn.microsoft.com/en-us/sql/t-sql/functions/json-array-transact-sql?view=sql-server-ver16>
7808-
/// [`JSON_OBJECT`]: <https://learn.microsoft.com/en-us/sql/t-sql/functions/json-object-transact-sql?view=sql-server-ver16>
7807+
/// [MSSQL `JSON_ARRAY`](https://learn.microsoft.com/en-us/sql/t-sql/functions/json-array-transact-sql?view=sql-server-ver16)
7808+
/// [MSSQL `JSON_OBJECT`](https://learn.microsoft.com/en-us/sql/t-sql/functions/json-object-transact-sql?view=sql-server-ver16>)
7809+
/// [PostgreSQL JSON functions](https://www.postgresql.org/docs/current/functions-json.html#FUNCTIONS-JSON-PROCESSING)
78097810
JsonNullClause(JsonNullClause),
7811+
/// The `RETURNING` clause for some JSON functions in PostgreSQL
7812+
///
7813+
/// [`JSON_OBJECT`](https://www.postgresql.org/docs/current/functions-json.html#:~:text=json_object)
7814+
JsonReturningClause(JsonReturningClause),
78107815
}
78117816

78127817
impl fmt::Display for FunctionArgumentClause {
@@ -7823,6 +7828,9 @@ impl fmt::Display for FunctionArgumentClause {
78237828
FunctionArgumentClause::Having(bound) => write!(f, "{bound}"),
78247829
FunctionArgumentClause::Separator(sep) => write!(f, "SEPARATOR {sep}"),
78257830
FunctionArgumentClause::JsonNullClause(null_clause) => write!(f, "{null_clause}"),
7831+
FunctionArgumentClause::JsonReturningClause(returning_clause) => {
7832+
write!(f, "{returning_clause}")
7833+
}
78267834
}
78277835
}
78287836
}
@@ -10071,6 +10079,25 @@ impl Display for JsonNullClause {
1007110079
}
1007210080
}
1007310081

10082+
/// PostgreSQL JSON function RETURNING clause
10083+
///
10084+
/// Example:
10085+
/// ```sql
10086+
/// JSON_OBJECT('a': 1 RETURNING jsonb)
10087+
/// ```
10088+
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
10089+
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
10090+
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
10091+
pub struct JsonReturningClause {
10092+
pub data_type: DataType,
10093+
}
10094+
10095+
impl Display for JsonReturningClause {
10096+
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
10097+
write!(f, "RETURNING {}", self.data_type)
10098+
}
10099+
}
10100+
1007410101
/// rename object definition
1007510102
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
1007610103
#[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
@@ -1779,6 +1779,7 @@ impl Spanned for FunctionArgumentClause {
17791779
FunctionArgumentClause::Having(HavingBound(_kind, expr)) => expr.span(),
17801780
FunctionArgumentClause::Separator(value) => value.span(),
17811781
FunctionArgumentClause::JsonNullClause(_) => Span::empty(),
1782+
FunctionArgumentClause::JsonReturningClause(_) => Span::empty(),
17821783
}
17831784
}
17841785
}

‎src/parser/mod.rs

Lines changed: 24 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15307,7 +15307,7 @@ impl<'a> Parser<'a> {
1530715307
Ok(TableFunctionArgs { args, settings })
1530815308
}
1530915309

15310-
/// Parses a potentially empty list of arguments to a window function
15310+
/// Parses a potentially empty list of arguments to a function
1531115311
/// (including the closing parenthesis).
1531215312
///
1531315313
/// Examples:
@@ -15318,11 +15318,18 @@ impl<'a> Parser<'a> {
1531815318
fn parse_function_argument_list(&mut self) -> Result<FunctionArgumentList, ParserError> {
1531915319
let mut clauses = vec![];
1532015320

15321-
// For MSSQL empty argument list with json-null-clause case, e.g. `JSON_ARRAY(NULL ON NULL)`
15321+
// Handle clauses that may exist with an empty argument list
15322+
1532215323
if let Some(null_clause) = self.parse_json_null_clause() {
1532315324
clauses.push(FunctionArgumentClause::JsonNullClause(null_clause));
1532415325
}
1532515326

15327+
if let Some(json_returning_clause) = self.parse_json_returning_clause()? {
15328+
clauses.push(FunctionArgumentClause::JsonReturningClause(
15329+
json_returning_clause,
15330+
));
15331+
}
15332+
1532615333
if self.consume_token(&Token::RParen) {
1532715334
return Ok(FunctionArgumentList {
1532815335
duplicate_treatment: None,
@@ -15378,6 +15385,12 @@ impl<'a> Parser<'a> {
1537815385
clauses.push(FunctionArgumentClause::JsonNullClause(null_clause));
1537915386
}
1538015387

15388+
if let Some(json_returning_clause) = self.parse_json_returning_clause()? {
15389+
clauses.push(FunctionArgumentClause::JsonReturningClause(
15390+
json_returning_clause,
15391+
));
15392+
}
15393+
1538115394
self.expect_token(&Token::RParen)?;
1538215395
Ok(FunctionArgumentList {
1538315396
duplicate_treatment,
@@ -15386,7 +15399,6 @@ impl<'a> Parser<'a> {
1538615399
})
1538715400
}
1538815401

15389-
/// Parses MSSQL's json-null-clause
1539015402
fn parse_json_null_clause(&mut self) -> Option<JsonNullClause> {
1539115403
if self.parse_keywords(&[Keyword::ABSENT, Keyword::ON, Keyword::NULL]) {
1539215404
Some(JsonNullClause::AbsentOnNull)
@@ -15397,6 +15409,15 @@ impl<'a> Parser<'a> {
1539715409
}
1539815410
}
1539915411

15412+
fn parse_json_returning_clause(&mut self) -> Result<Option<JsonReturningClause>, ParserError> {
15413+
if self.parse_keyword(Keyword::RETURNING) {
15414+
let data_type = self.parse_data_type()?;
15415+
Ok(Some(JsonReturningClause { data_type }))
15416+
} else {
15417+
Ok(None)
15418+
}
15419+
}
15420+
1540015421
fn parse_duplicate_treatment(&mut self) -> Result<Option<DuplicateTreatment>, ParserError> {
1540115422
let loc = self.peek_token().span.start;
1540215423
match (

‎tests/sqlparser_postgres.rs

Lines changed: 87 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,68 @@ fn test_fn_arg_with_value_operator() {
33633387
}
33643388
}
33653389

3390+
#[test]
3391+
fn json_object_with_null_clause() {
3392+
match pg().verified_expr("JSON_OBJECT('name' VALUE 'value' NULL ON NULL)") {
3393+
Expr::Function(Function { args: FunctionArguments::List(FunctionArgumentList { args, clauses, .. }), .. }) => {
3394+
assert!(matches!(
3395+
&args[..],
3396+
&[FunctionArg::ExprNamed { operator: FunctionArgOperator::Value, .. }]
3397+
), "Invalid function argument: {args:?}");
3398+
assert_eq!(
3399+
clauses,
3400+
vec![FunctionArgumentClause::JsonNullClause(JsonNullClause::NullOnNull)],
3401+
"Expected: NULL ON NULL to be parsed as a null treatment clause, but got {clauses:?}"
3402+
);
3403+
}
3404+
other => panic!("Expected: JSON_OBJECT('name' VALUE 'value' NULL ON NULL) to be parsed as a function, but got {other:?}"),
3405+
}
3406+
}
3407+
3408+
#[test]
3409+
fn json_object_with_returning_clause() {
3410+
match pg().verified_expr("JSON_OBJECT('name' VALUE 'value' RETURNING JSONB)") {
3411+
Expr::Function(Function { args: FunctionArguments::List(FunctionArgumentList { args, clauses, .. }), .. }) => {
3412+
assert!(matches!(
3413+
&args[..],
3414+
&[FunctionArg::ExprNamed { operator: FunctionArgOperator::Value, .. }]
3415+
), "Invalid function argument: {args:?}");
3416+
assert_eq!(
3417+
clauses,
3418+
vec![FunctionArgumentClause::JsonReturningClause(JsonReturningClause {
3419+
data_type: DataType::JSONB
3420+
})],
3421+
"Expected: RETURNING JSONB to be parsed as a returning clause, but got {clauses:?}"
3422+
);
3423+
}
3424+
other => panic!("Expected: JSON_OBJECT('name' VALUE 'value' RETURNING jsonb) to be parsed as a function, but got {other:?}"),
3425+
}
3426+
}
3427+
3428+
#[test]
3429+
fn json_object_only_returning_clause() {
3430+
match pg().verified_expr("JSON_OBJECT(RETURNING JSONB)") {
3431+
Expr::Function(Function {
3432+
args: FunctionArguments::List(FunctionArgumentList { args, clauses, .. }),
3433+
..
3434+
}) => {
3435+
assert!(args.is_empty(), "Expected no arguments, got: {args:?}");
3436+
assert_eq!(
3437+
clauses,
3438+
vec![FunctionArgumentClause::JsonReturningClause(
3439+
JsonReturningClause {
3440+
data_type: DataType::JSONB
3441+
}
3442+
)],
3443+
"Expected: RETURNING JSONB to be parsed as a returning clause, but got {clauses:?}"
3444+
);
3445+
}
3446+
other => panic!(
3447+
"Expected: JSON_OBJECT(RETURNING jsonb) to be parsed as a function, but got {other:?}"
3448+
),
3449+
}
3450+
}
3451+
33663452
#[test]
33673453
fn parse_json_table_is_not_reserved() {
33683454
// JSON_TABLE is not a reserved keyword in PostgreSQL, even though it is in SQL:2023

0 commit comments

Comments
 (0)