Skip to content

Commit 9f515bf

Browse files
adamchainzadamchainz
authored
Add support for PostgreSQL JSON function 'RETURNING' clauses (#2001)
1 parent cffff30 commit 9f515bf

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
@@ -7822,11 +7822,16 @@ pub enum FunctionArgumentClause {
78227822
///
78237823
/// [`GROUP_CONCAT`]: https://dev.mysql.com/doc/refman/8.0/en/aggregate-functions.html#function_group-concat
78247824
Separator(Value),
7825-
/// The json-null-clause to the [`JSON_ARRAY`]/[`JSON_OBJECT`] function in MSSQL.
7825+
/// The `ON NULL` clause for some JSON functions.
78267826
///
7827-
/// [`JSON_ARRAY`]: <https://learn.microsoft.com/en-us/sql/t-sql/functions/json-array-transact-sql?view=sql-server-ver16>
7828-
/// [`JSON_OBJECT`]: <https://learn.microsoft.com/en-us/sql/t-sql/functions/json-object-transact-sql?view=sql-server-ver16>
7827+
/// [MSSQL `JSON_ARRAY`](https://learn.microsoft.com/en-us/sql/t-sql/functions/json-array-transact-sql?view=sql-server-ver16)
7828+
/// [MSSQL `JSON_OBJECT`](https://learn.microsoft.com/en-us/sql/t-sql/functions/json-object-transact-sql?view=sql-server-ver16>)
7829+
/// [PostgreSQL JSON functions](https://www.postgresql.org/docs/current/functions-json.html#FUNCTIONS-JSON-PROCESSING)
78297830
JsonNullClause(JsonNullClause),
7831+
/// The `RETURNING` clause for some JSON functions in PostgreSQL
7832+
///
7833+
/// [`JSON_OBJECT`](https://www.postgresql.org/docs/current/functions-json.html#:~:text=json_object)
7834+
JsonReturningClause(JsonReturningClause),
78307835
}
78317836

78327837
impl fmt::Display for FunctionArgumentClause {
@@ -7843,6 +7848,9 @@ impl fmt::Display for FunctionArgumentClause {
78437848
FunctionArgumentClause::Having(bound) => write!(f, "{bound}"),
78447849
FunctionArgumentClause::Separator(sep) => write!(f, "SEPARATOR {sep}"),
78457850
FunctionArgumentClause::JsonNullClause(null_clause) => write!(f, "{null_clause}"),
7851+
FunctionArgumentClause::JsonReturningClause(returning_clause) => {
7852+
write!(f, "{returning_clause}")
7853+
}
78467854
}
78477855
}
78487856
}
@@ -10177,6 +10185,25 @@ impl Display for JsonNullClause {
1017710185
}
1017810186
}
1017910187

10188+
/// PostgreSQL JSON function RETURNING clause
10189+
///
10190+
/// Example:
10191+
/// ```sql
10192+
/// JSON_OBJECT('a': 1 RETURNING jsonb)
10193+
/// ```
10194+
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
10195+
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
10196+
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
10197+
pub struct JsonReturningClause {
10198+
pub data_type: DataType,
10199+
}
10200+
10201+
impl Display for JsonReturningClause {
10202+
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
10203+
write!(f, "RETURNING {}", self.data_type)
10204+
}
10205+
}
10206+
1018010207
/// rename object definition
1018110208
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
1018210209
#[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
@@ -1792,6 +1792,7 @@ impl Spanned for FunctionArgumentClause {
17921792
FunctionArgumentClause::Having(HavingBound(_kind, expr)) => expr.span(),
17931793
FunctionArgumentClause::Separator(value) => value.span(),
17941794
FunctionArgumentClause::JsonNullClause(_) => Span::empty(),
1795+
FunctionArgumentClause::JsonReturningClause(_) => Span::empty(),
17951796
}
17961797
}
17971798
}

‎src/parser/mod.rs‎

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

15581-
/// Parses a potentially empty list of arguments to a window function
15581+
/// Parses a potentially empty list of arguments to a function
1558215582
/// (including the closing parenthesis).
1558315583
///
1558415584
/// Examples:
@@ -15589,11 +15589,18 @@ impl<'a> Parser<'a> {
1558915589
fn parse_function_argument_list(&mut self) -> Result<FunctionArgumentList, ParserError> {
1559015590
let mut clauses = vec![];
1559115591

15592-
// For MSSQL empty argument list with json-null-clause case, e.g. `JSON_ARRAY(NULL ON NULL)`
15592+
// Handle clauses that may exist with an empty argument list
15593+
1559315594
if let Some(null_clause) = self.parse_json_null_clause() {
1559415595
clauses.push(FunctionArgumentClause::JsonNullClause(null_clause));
1559515596
}
1559615597

15598+
if let Some(json_returning_clause) = self.maybe_parse_json_returning_clause()? {
15599+
clauses.push(FunctionArgumentClause::JsonReturningClause(
15600+
json_returning_clause,
15601+
));
15602+
}
15603+
1559715604
if self.consume_token(&Token::RParen) {
1559815605
return Ok(FunctionArgumentList {
1559915606
duplicate_treatment: None,
@@ -15649,6 +15656,12 @@ impl<'a> Parser<'a> {
1564915656
clauses.push(FunctionArgumentClause::JsonNullClause(null_clause));
1565015657
}
1565115658

15659+
if let Some(json_returning_clause) = self.maybe_parse_json_returning_clause()? {
15660+
clauses.push(FunctionArgumentClause::JsonReturningClause(
15661+
json_returning_clause,
15662+
));
15663+
}
15664+
1565215665
self.expect_token(&Token::RParen)?;
1565315666
Ok(FunctionArgumentList {
1565415667
duplicate_treatment,
@@ -15657,7 +15670,6 @@ impl<'a> Parser<'a> {
1565715670
})
1565815671
}
1565915672

15660-
/// Parses MSSQL's json-null-clause
1566115673
fn parse_json_null_clause(&mut self) -> Option<JsonNullClause> {
1566215674
if self.parse_keywords(&[Keyword::ABSENT, Keyword::ON, Keyword::NULL]) {
1566315675
Some(JsonNullClause::AbsentOnNull)
@@ -15668,6 +15680,17 @@ impl<'a> Parser<'a> {
1566815680
}
1566915681
}
1567015682

15683+
fn maybe_parse_json_returning_clause(
15684+
&mut self,
15685+
) -> Result<Option<JsonReturningClause>, ParserError> {
15686+
if self.parse_keyword(Keyword::RETURNING) {
15687+
let data_type = self.parse_data_type()?;
15688+
Ok(Some(JsonReturningClause { data_type }))
15689+
} else {
15690+
Ok(None)
15691+
}
15692+
}
15693+
1567115694
fn parse_duplicate_treatment(&mut self) -> Result<Option<DuplicateTreatment>, ParserError> {
1567215695
let loc = self.peek_token().span.start;
1567315696
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)