Repository: apache/datafusion-sqlparser-rs (Rust)

Issue:
Support Postgres' `JSON_OBJECT` function `RETURNING` clause
[`json_object()`](https://www.postgresql.org/docs/current/functions-json.html#:~:text=%20json_object%20( ) supports quite broad syntax:

```
json_object ( [ { key_expression { VALUE | ':' } value_expression
[ FORMAT JSON [ ENCODING UTF8 ] ] }[, ...] ] [ { NULL | ABSENT } ON NULL ]
[ { WITH | WITHOUT } UNIQUE [ KEYS ] ]
[ RETURNING data_type [ FORMAT JSON [ ENCODING UTF8 ] ] ])
```

#1545 added support for the `VALUE` keyword between values. I now need support for `RETURNING JSONB`, as [used by Django](https://github.com/django/django/blob/cd0966cd4e37da8e6153cbf57c194dce29caaddc/django/db/models/functions/json.py#L111-L114) on Postgres 16+.



Fix (patch):
diff --git a/src/ast/mod.rs b/src/ast/mod.rs
index 71cb6c5ee..479b24a9c 100644
--- a/src/ast/mod.rs
+++ b/src/ast/mod.rs
@@ -7817,11 +7817,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>
-    /// [`JSON_OBJECT`]: <https://learn.microsoft.com/en-us/sql/t-sql/functions/json-object-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)
+    /// [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 {
@@ -7838,6 +7843,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}")
+            }
         }
     }
 }
@@ -10172,6 +10180,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))]
diff --git a/src/ast/spans.rs b/src/ast/spans.rs
index 39761751f..cbc598e57 100644
--- a/src/ast/spans.rs
+++ b/src/ast/spans.rs
@@ -1791,6 +1791,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(),
         }
     }
 }
diff --git a/src/parser/mod.rs b/src/parser/mod.rs
index 6c559eed4..f793ccea0 100644
--- a/src/parser/mod.rs
+++ b/src/parser/mod.rs
@@ -15573,7 +15573,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:
@@ -15584,11 +15584,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,
@@ -15644,6 +15651,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,
@@ -15652,7 +15665,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)
@@ -15663,6 +15675,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 (