Add support for PostgreSQL JSON function 'RETURNING' clauses (#2001)

src/ast/mod.rs

@@ -7822,11 +7822,16 @@ 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.
+    /// The `ON NULL` clause for some JSON functions.
     ///
-    /// [`JSON_ARRAY`]: <https://learn.microsoft.com/en-us/sql/t-sql/functions/json-array-transact-sql?view=sql-server-ver16>
+    /// [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>
+    /// [MSSQL `JSON_OBJECT`](https://learn.microsoft.com/en-us/sql/t-sql/functions/json-object-transact-sql?view=sql-server-ver16>)
+    /// [PostgreSQL JSON functions](https://www.postgresql.org/docs/current/functions-json.html#FUNCTIONS-JSON-PROCESSING)
     JsonNullClause(JsonNullClause),
+    /// The `RETURNING` clause for some JSON functions in PostgreSQL
+    ///
+    /// [`JSON_OBJECT`](https://www.postgresql.org/docs/current/functions-json.html#:~:text=json_object)
+    JsonReturningClause(JsonReturningClause),
 }
 
 impl fmt::Display for FunctionArgumentClause {
@@ -7843,6 +7848,9 @@ impl fmt::Display for FunctionArgumentClause {
             FunctionArgumentClause::Having(bound) => write!(f, "{bound}"),
             FunctionArgumentClause::Separator(sep) => write!(f, "SEPARATOR {sep}"),
             FunctionArgumentClause::JsonNullClause(null_clause) => write!(f, "{null_clause}"),
+            FunctionArgumentClause::JsonReturningClause(returning_clause) => {
+                write!(f, "{returning_clause}")
+            }
         }
     }
 }
@@ -10177,6 +10185,25 @@ impl Display for JsonNullClause {
     }
 }
 
+/// PostgreSQL JSON function RETURNING clause
+///
+/// Example:
+/// ```sql
+/// JSON_OBJECT('a': 1 RETURNING jsonb)
+/// ```
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct JsonReturningClause {
+    pub data_type: DataType,
+}
+
+impl Display for JsonReturningClause {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "RETURNING {}", self.data_type)
+    }
+}
+
 /// rename object definition
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]

src/ast/spans.rs

@@ -1792,6 +1792,7 @@ impl Spanned for FunctionArgumentClause {
             FunctionArgumentClause::Having(HavingBound(_kind, expr)) => expr.span(),
             FunctionArgumentClause::Separator(value) => value.span(),
             FunctionArgumentClause::JsonNullClause(_) => Span::empty(),
+            FunctionArgumentClause::JsonReturningClause(_) => Span::empty(),
         }
     }
 }

src/parser/mod.rs

@@ -15578,7 +15578,7 @@ impl<'a> Parser<'a> {
         Ok(TableFunctionArgs { args, settings })
     }
 
-    /// Parses a potentially empty list of arguments to a window function
+    /// Parses a potentially empty list of arguments to a function
     /// (including the closing parenthesis).
     ///
     /// Examples:
@@ -15589,11 +15589,18 @@ impl<'a> Parser<'a> {
     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)`
+        // Handle clauses that may exist with an empty argument list
+
         if let Some(null_clause) = self.parse_json_null_clause() {
             clauses.push(FunctionArgumentClause::JsonNullClause(null_clause));
         }
 
+        if let Some(json_returning_clause) = self.maybe_parse_json_returning_clause()? {
+            clauses.push(FunctionArgumentClause::JsonReturningClause(
+                json_returning_clause,
+            ));
+        }
+
         if self.consume_token(&Token::RParen) {
             return Ok(FunctionArgumentList {
                 duplicate_treatment: None,
@@ -15649,6 +15656,12 @@ impl<'a> Parser<'a> {
             clauses.push(FunctionArgumentClause::JsonNullClause(null_clause));
         }
 
+        if let Some(json_returning_clause) = self.maybe_parse_json_returning_clause()? {
+            clauses.push(FunctionArgumentClause::JsonReturningClause(
+                json_returning_clause,
+            ));
+        }
+
         self.expect_token(&Token::RParen)?;
         Ok(FunctionArgumentList {
             duplicate_treatment,
@@ -15657,7 +15670,6 @@ 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)
@@ -15668,6 +15680,17 @@ impl<'a> Parser<'a> {
         }
     }
 
+    fn maybe_parse_json_returning_clause(
+        &mut self,
+    ) -> Result<Option<JsonReturningClause>, ParserError> {
+        if self.parse_keyword(Keyword::RETURNING) {
+            let data_type = self.parse_data_type()?;
+            Ok(Some(JsonReturningClause { data_type }))
+        } else {
+            Ok(None)
+        }
+    }
+
     fn parse_duplicate_treatment(&mut self) -> Result<Option<DuplicateTreatment>, ParserError> {
         let loc = self.peek_token().span.start;
         match (

tests/sqlparser_postgres.rs

@@ -3351,7 +3351,31 @@ fn test_json() {
 }
 
 #[test]
-fn test_fn_arg_with_value_operator() {
+fn json_object_colon_syntax() {
+    match pg().verified_expr("JSON_OBJECT('name' : 'value')") {
+        Expr::Function(Function {
+            args: FunctionArguments::List(FunctionArgumentList { args, .. }),
+            ..
+        }) => {
+            assert!(
+                matches!(
+                    &args[..],
+                    &[FunctionArg::ExprNamed {
+                        operator: FunctionArgOperator::Colon,
+                        ..
+                    }]
+                ),
+                "Invalid function argument: {args:?}"
+            );
+        }
+        other => panic!(
+            "Expected: JSON_OBJECT('name' : 'value') to be parsed as a function, but got {other:?}"
+        ),
+    }
+}
+
+#[test]
+fn json_object_value_syntax() {
     match pg().verified_expr("JSON_OBJECT('name' VALUE 'value')") {
         Expr::Function(Function { args: FunctionArguments::List(FunctionArgumentList { args, .. }), .. }) => {
             assert!(matches!(
@@ -3363,6 +3387,63 @@ fn test_fn_arg_with_value_operator() {
     }
 }
 
+#[test]
+fn parse_json_object() {
+    let sql = "JSON_OBJECT('name' VALUE 'value' NULL ON NULL)";
+    let expr = pg().verified_expr(sql);
+    assert!(
+        matches!(
+            expr.clone(),
+            Expr::Function(Function {
+                name: ObjectName(parts),
+                args: FunctionArguments::List(FunctionArgumentList { args, clauses, .. }),
+                ..
+            }) if parts == vec![ObjectNamePart::Identifier(Ident::new("JSON_OBJECT"))]
+                && matches!(
+                    &args[..],
+                    &[FunctionArg::ExprNamed { operator: FunctionArgOperator::Value, .. }]
+                )
+                && clauses == vec![FunctionArgumentClause::JsonNullClause(JsonNullClause::NullOnNull)]
+        ),
+        "Failed to parse JSON_OBJECT with expected structure, got: {expr:?}"
+    );
+
+    let sql = "JSON_OBJECT('name' VALUE 'value' RETURNING JSONB)";
+    let expr = pg().verified_expr(sql);
+    assert!(
+        matches!(
+            expr.clone(),
+            Expr::Function(Function {
+                name: ObjectName(parts),
+                args: FunctionArguments::List(FunctionArgumentList { args, clauses, .. }),
+                ..
+            }) if parts == vec![ObjectNamePart::Identifier(Ident::new("JSON_OBJECT"))]
+                && matches!(
+                    &args[..],
+                    &[FunctionArg::ExprNamed { operator: FunctionArgOperator::Value, .. }]
+                )
+                && clauses == vec![FunctionArgumentClause::JsonReturningClause(JsonReturningClause { data_type: DataType::JSONB })]
+        ),
+        "Failed to parse JSON_OBJECT with expected structure, got: {expr:?}"
+    );
+
+    let sql = "JSON_OBJECT(RETURNING JSONB)";
+    let expr = pg().verified_expr(sql);
+    assert!(
+        matches!(
+            expr.clone(),
+            Expr::Function(Function {
+                name: ObjectName(parts),
+                args: FunctionArguments::List(FunctionArgumentList { args, clauses, .. }),
+                ..
+            }) if parts == vec![ObjectNamePart::Identifier(Ident::new("JSON_OBJECT"))]
+                && args.is_empty()
+                && clauses == vec![FunctionArgumentClause::JsonReturningClause(JsonReturningClause { data_type: DataType::JSONB })]
+        ),
+        "Failed to parse JSON_OBJECT with expected structure, got: {expr:?}"
+    );
+}
+
 #[test]
 fn parse_json_table_is_not_reserved() {
     // JSON_TABLE is not a reserved keyword in PostgreSQL, even though it is in SQL:2023