Skip to content

Commit fa05077

Browse files
adamchainzayman-sigma
adamchainz
authored and
ayman-sigma
committed
Add support for PostgreSQL JSON function 'RETURNING' clauses (apache#2001)
1 parent cfd8bf1 commit fa05077

4 files changed

Lines changed: 139 additions & 7 deletions

File tree

‎src/ast/mod.rs‎

Lines changed: 30 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7851,11 +7851,16 @@ pub enum FunctionArgumentClause {
78517851
///
78527852
/// [`GROUP_CONCAT`]: https://dev.mysql.com/doc/refman/8.0/en/aggregate-functions.html#function_group-concat
78537853
Separator(Value),
7854-
/// The json-null-clause to the [`JSON_ARRAY`]/[`JSON_OBJECT`] function in MSSQL.
7854+
/// The `ON NULL` clause for some JSON functions.
78557855
///
7856-
/// [`JSON_ARRAY`]: <https://learn.microsoft.com/en-us/sql/t-sql/functions/json-array-transact-sql?view=sql-server-ver16>
7857-
/// [`JSON_OBJECT`]: <https://learn.microsoft.com/en-us/sql/t-sql/functions/json-object-transact-sql?view=sql-server-ver16>
7856+
/// [MSSQL `JSON_ARRAY`](https://learn.microsoft.com/en-us/sql/t-sql/functions/json-array-transact-sql?view=sql-server-ver16)
7857+
/// [MSSQL `JSON_OBJECT`](https://learn.microsoft.com/en-us/sql/t-sql/functions/json-object-transact-sql?view=sql-server-ver16>)
7858+
/// [PostgreSQL JSON functions](https://www.postgresql.org/docs/current/functions-json.html#FUNCTIONS-JSON-PROCESSING)
78587859
JsonNullClause(JsonNullClause),
7860+
/// The `RETURNING` clause for some JSON functions in PostgreSQL
7861+
///
7862+
/// [`JSON_OBJECT`](https://www.postgresql.org/docs/current/functions-json.html#:~:text=json_object)
7863+
JsonReturningClause(JsonReturningClause),
78597864
}
78607865

78617866
impl fmt::Display for FunctionArgumentClause {
@@ -7872,6 +7877,9 @@ impl fmt::Display for FunctionArgumentClause {
78727877
FunctionArgumentClause::Having(bound) => write!(f, "{bound}"),
78737878
FunctionArgumentClause::Separator(sep) => write!(f, "SEPARATOR {sep}"),
78747879
FunctionArgumentClause::JsonNullClause(null_clause) => write!(f, "{null_clause}"),
7880+
FunctionArgumentClause::JsonReturningClause(returning_clause) => {
7881+
write!(f, "{returning_clause}")
7882+
}
78757883
}
78767884
}
78777885
}
@@ -10206,6 +10214,25 @@ impl Display for JsonNullClause {
1020610214
}
1020710215
}
1020810216

10217+
/// PostgreSQL JSON function RETURNING clause
10218+
///
10219+
/// Example:
10220+
/// ```sql
10221+
/// JSON_OBJECT('a': 1 RETURNING jsonb)
10222+
/// ```
10223+
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
10224+
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
10225+
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
10226+
pub struct JsonReturningClause {
10227+
pub data_type: DataType,
10228+
}
10229+
10230+
impl Display for JsonReturningClause {
10231+
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
10232+
write!(f, "RETURNING {}", self.data_type)
10233+
}
10234+
}
10235+
1020910236
/// rename object definition
1021010237
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
1021110238
#[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
@@ -1797,6 +1797,7 @@ impl Spanned for FunctionArgumentClause {
17971797
FunctionArgumentClause::Having(HavingBound(_kind, expr)) => expr.span(),
17981798
FunctionArgumentClause::Separator(value) => value.span(),
17991799
FunctionArgumentClause::JsonNullClause(_) => Span::empty(),
1800+
FunctionArgumentClause::JsonReturningClause(_) => Span::empty(),
18001801
}
18011802
}
18021803
}

‎src/parser/mod.rs‎

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

15613-
/// Parses a potentially empty list of arguments to a window function
15613+
/// Parses a potentially empty list of arguments to a function
1561415614
/// (including the closing parenthesis).
1561515615
///
1561615616
/// Examples:
@@ -15621,11 +15621,18 @@ impl<'a> Parser<'a> {
1562115621
fn parse_function_argument_list(&mut self) -> Result<FunctionArgumentList, ParserError> {
1562215622
let mut clauses = vec![];
1562315623

15624-
// For MSSQL empty argument list with json-null-clause case, e.g. `JSON_ARRAY(NULL ON NULL)`
15624+
// Handle clauses that may exist with an empty argument list
15625+
1562515626
if let Some(null_clause) = self.parse_json_null_clause() {
1562615627
clauses.push(FunctionArgumentClause::JsonNullClause(null_clause));
1562715628
}
1562815629

15630+
if let Some(json_returning_clause) = self.maybe_parse_json_returning_clause()? {
15631+
clauses.push(FunctionArgumentClause::JsonReturningClause(
15632+
json_returning_clause,
15633+
));
15634+
}
15635+
1562915636
if self.consume_token(&Token::RParen) {
1563015637
return Ok(FunctionArgumentList {
1563115638
duplicate_treatment: None,
@@ -15681,6 +15688,12 @@ impl<'a> Parser<'a> {
1568115688
clauses.push(FunctionArgumentClause::JsonNullClause(null_clause));
1568215689
}
1568315690

15691+
if let Some(json_returning_clause) = self.maybe_parse_json_returning_clause()? {
15692+
clauses.push(FunctionArgumentClause::JsonReturningClause(
15693+
json_returning_clause,
15694+
));
15695+
}
15696+
1568415697
self.expect_token(&Token::RParen)?;
1568515698
Ok(FunctionArgumentList {
1568615699
duplicate_treatment,
@@ -15689,7 +15702,6 @@ impl<'a> Parser<'a> {
1568915702
})
1569015703
}
1569115704

15692-
/// Parses MSSQL's json-null-clause
1569315705
fn parse_json_null_clause(&mut self) -> Option<JsonNullClause> {
1569415706
if self.parse_keywords(&[Keyword::ABSENT, Keyword::ON, Keyword::NULL]) {
1569515707
Some(JsonNullClause::AbsentOnNull)
@@ -15700,6 +15712,17 @@ impl<'a> Parser<'a> {
1570015712
}
1570115713
}
1570215714

15715+
fn maybe_parse_json_returning_clause(
15716+
&mut self,
15717+
) -> Result<Option<JsonReturningClause>, ParserError> {
15718+
if self.parse_keyword(Keyword::RETURNING) {
15719+
let data_type = self.parse_data_type()?;
15720+
Ok(Some(JsonReturningClause { data_type }))
15721+
} else {
15722+
Ok(None)
15723+
}
15724+
}
15725+
1570315726
fn parse_duplicate_treatment(&mut self) -> Result<Option<DuplicateTreatment>, ParserError> {
1570415727
let loc = self.peek_token().span.start;
1570515728
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)