add support for FOR ORDINALITY and NESTED in JSON_TABLE (apache#1493)

lovasoa · web-flow · commit fc0e13b80ea7 · 2024-11-06T16:04:13.000-05:00
diff --git a/src/ast/mod.rs b/src/ast/mod.rs
@@ -54,13 +54,14 @@ pub use self::query::{
     ExceptSelectItem, ExcludeSelectItem, ExprWithAlias, Fetch, ForClause, ForJson, ForXml,
     FormatClause, GroupByExpr, GroupByWithModifier, IdentWithAlias, IlikeSelectItem, Interpolate,
     InterpolateExpr, Join, JoinConstraint, JoinOperator, JsonTableColumn,
-    JsonTableColumnErrorHandling, LateralView, LockClause, LockType, MatchRecognizePattern,
-    MatchRecognizeSymbol, Measure, NamedWindowDefinition, NamedWindowExpr, NonBlock, Offset,
-    OffsetRows, OrderBy, OrderByExpr, PivotValueSource, ProjectionSelect, Query, RenameSelectItem,
-    RepetitionQuantifier, ReplaceSelectElement, ReplaceSelectItem, RowsPerMatch, Select,
-    SelectInto, SelectItem, SetExpr, SetOperator, SetQuantifier, Setting, SymbolDefinition, Table,
-    TableAlias, TableFactor, TableFunctionArgs, TableVersion, TableWithJoins, Top, TopQuantity,
-    ValueTableMode, Values, WildcardAdditionalOptions, With, WithFill,
+    JsonTableColumnErrorHandling, JsonTableNamedColumn, JsonTableNestedColumn, LateralView,
+    LockClause, LockType, MatchRecognizePattern, MatchRecognizeSymbol, Measure,
+    NamedWindowDefinition, NamedWindowExpr, NonBlock, Offset, OffsetRows, OrderBy, OrderByExpr,
+    PivotValueSource, ProjectionSelect, Query, RenameSelectItem, RepetitionQuantifier,
+    ReplaceSelectElement, ReplaceSelectItem, RowsPerMatch, Select, SelectInto, SelectItem, SetExpr,
+    SetOperator, SetQuantifier, Setting, SymbolDefinition, Table, TableAlias, TableFactor,
+    TableFunctionArgs, TableVersion, TableWithJoins, Top, TopQuantity, ValueTableMode, Values,
+    WildcardAdditionalOptions, With, WithFill,
 };
 
 pub use self::trigger::{
diff --git a/src/ast/query.rs b/src/ast/query.rs
@@ -2286,19 +2286,84 @@ impl fmt::Display for ForJson {
 }
 
 /// A single column definition in MySQL's `JSON_TABLE` table valued function.
+///
+/// See
+/// - [MySQL's JSON_TABLE documentation](https://dev.mysql.com/doc/refman/8.0/en/json-table-functions.html#function_json-table)
+/// - [Oracle's JSON_TABLE documentation](https://docs.oracle.com/en/database/oracle/oracle-database/21/sqlrf/JSON_TABLE.html)
+/// - [MariaDB's JSON_TABLE documentation](https://mariadb.com/kb/en/json_table/)
+///
 /// ```sql
 /// SELECT *
 /// FROM JSON_TABLE(
 ///     '["a", "b"]',
 ///     '$[*]' COLUMNS (
-///         value VARCHAR(20) PATH '$'
+///         name FOR ORDINALITY,
+///         value VARCHAR(20) PATH '$',
+///         NESTED PATH '$[*]' COLUMNS (
+///             value VARCHAR(20) PATH '$'
+///         )
 ///     )
 /// ) AS jt;
 /// ```
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
-pub struct JsonTableColumn {
+pub enum JsonTableColumn {
+    /// A named column with a JSON path
+    Named(JsonTableNamedColumn),
+    /// The FOR ORDINALITY column, which is a special column that returns the index of the current row in a JSON array.
+    ForOrdinality(Ident),
+    /// A set of nested columns, which extracts data from a nested JSON array.
+    Nested(JsonTableNestedColumn),
+}
+
+impl fmt::Display for JsonTableColumn {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            JsonTableColumn::Named(json_table_named_column) => {
+                write!(f, "{json_table_named_column}")
+            }
+            JsonTableColumn::ForOrdinality(ident) => write!(f, "{} FOR ORDINALITY", ident),
+            JsonTableColumn::Nested(json_table_nested_column) => {
+                write!(f, "{json_table_nested_column}")
+            }
+        }
+    }
+}
+
+/// A nested column in a JSON_TABLE column list
+///
+/// See <https://mariadb.com/kb/en/json_table/#nested-paths>
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub struct JsonTableNestedColumn {
+    pub path: Value,
+    pub columns: Vec<JsonTableColumn>,
+}
+
+impl fmt::Display for JsonTableNestedColumn {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "NESTED PATH {} COLUMNS ({})",
+            self.path,
+            display_comma_separated(&self.columns)
+        )
+    }
+}
+
+/// A single column definition in MySQL's `JSON_TABLE` table valued function.
+///
+/// See <https://mariadb.com/kb/en/json_table/#path-columns>
+///
+/// ```sql
+///         value VARCHAR(20) PATH '$'
+/// ```
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub struct JsonTableNamedColumn {
     /// The name of the column to be extracted.
     pub name: Ident,
     /// The type of the column to be extracted.
@@ -2313,7 +2378,7 @@ pub struct JsonTableColumn {
     pub on_error: Option<JsonTableColumnErrorHandling>,
 }
 
-impl fmt::Display for JsonTableColumn {
+impl fmt::Display for JsonTableNamedColumn {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         write!(
             f,
diff --git a/src/parser/mod.rs b/src/parser/mod.rs
@@ -10466,7 +10466,23 @@ impl<'a> Parser<'a> {
     /// Parses MySQL's JSON_TABLE column definition.
     /// For example: `id INT EXISTS PATH '$' DEFAULT '0' ON EMPTY ERROR ON ERROR`
     pub fn parse_json_table_column_def(&mut self) -> Result<JsonTableColumn, ParserError> {
+        if self.parse_keyword(Keyword::NESTED) {
+            let _has_path_keyword = self.parse_keyword(Keyword::PATH);
+            let path = self.parse_value()?;
+            self.expect_keyword(Keyword::COLUMNS)?;
+            let columns = self.parse_parenthesized(|p| {
+                p.parse_comma_separated(Self::parse_json_table_column_def)
+            })?;
+            return Ok(JsonTableColumn::Nested(JsonTableNestedColumn {
+                path,
+                columns,
+            }));
+        }
         let name = self.parse_identifier(false)?;
+        if self.parse_keyword(Keyword::FOR) {
+            self.expect_keyword(Keyword::ORDINALITY)?;
+            return Ok(JsonTableColumn::ForOrdinality(name));
+        }
         let r#type = self.parse_data_type()?;
         let exists = self.parse_keyword(Keyword::EXISTS);
         self.expect_keyword(Keyword::PATH)?;
@@ -10481,14 +10497,14 @@ impl<'a> Parser<'a> {
                 on_error = Some(error_handling);
             }
         }
-        Ok(JsonTableColumn {
+        Ok(JsonTableColumn::Named(JsonTableNamedColumn {
             name,
             r#type,
             path,
             exists,
             on_empty,
             on_error,
-        })
+        }))
     }
 
     fn parse_json_table_column_error_handling(
diff --git a/tests/sqlparser_mysql.rs b/tests/sqlparser_mysql.rs
@@ -2773,6 +2773,12 @@ fn parse_json_table() {
         r#"SELECT * FROM JSON_TABLE('[1,2]', '$[*]' COLUMNS(x INT PATH '$' ERROR ON EMPTY)) AS t"#,
     );
     mysql().verified_only_select(r#"SELECT * FROM JSON_TABLE('[1,2]', '$[*]' COLUMNS(x INT PATH '$' ERROR ON EMPTY DEFAULT '0' ON ERROR)) AS t"#);
+    mysql().verified_only_select(
+        r#"SELECT jt.* FROM JSON_TABLE('["Alice", "Bob", "Charlie"]', '$[*]' COLUMNS(row_num FOR ORDINALITY, name VARCHAR(50) PATH '$')) AS jt"#,
+    );
+    mysql().verified_only_select(
+        r#"SELECT * FROM JSON_TABLE('[ {"a": 1, "b": [11,111]}, {"a": 2, "b": [22,222]}, {"a":3}]', '$[*]' COLUMNS(a INT PATH '$.a', NESTED PATH '$.b[*]' COLUMNS (b INT PATH '$'))) AS jt"#,
+    );
     assert_eq!(
         mysql()
             .verified_only_select(
@@ -2784,14 +2790,14 @@ fn parse_json_table() {
             json_expr: Expr::Value(Value::SingleQuotedString("[1,2]".to_string())),
             json_path: Value::SingleQuotedString("$[*]".to_string()),
             columns: vec![
-                JsonTableColumn {
+                JsonTableColumn::Named(JsonTableNamedColumn {
                     name: Ident::new("x"),
                     r#type: DataType::Int(None),
                     path: Value::SingleQuotedString("$".to_string()),
                     exists: false,
                     on_empty: Some(JsonTableColumnErrorHandling::Default(Value::SingleQuotedString("0".to_string()))),
                     on_error: Some(JsonTableColumnErrorHandling::Null),
-                },
+                }),
             ],
             alias: Some(TableAlias {
                 name: Ident::new("t"),
