Use IndexColumn in all index definitions (#1900)

src/ast/ddl.rs

@@ -32,9 +32,9 @@ use crate::ast::value::escape_single_quote_string;
 use crate::ast::{
     display_comma_separated, display_separated, CommentDef, CreateFunctionBody,
     CreateFunctionUsing, DataType, Expr, FunctionBehavior, FunctionCalledOnNull,
-    FunctionDeterminismSpecifier, FunctionParallel, Ident, MySQLColumnPosition, ObjectName,
+    FunctionDeterminismSpecifier, FunctionParallel, Ident, IndexColumn, MySQLColumnPosition,
-    OperateFunctionArg, OrderByExpr, ProjectionSelect, SequenceOptions, SqlOption, Tag, Value,
+    ObjectName, OperateFunctionArg, OrderByExpr, ProjectionSelect, SequenceOptions, SqlOption, Tag,
-    ValueWithSpan,
+    Value, ValueWithSpan,
 };
 use crate::keywords::Keyword;
 use crate::tokenizer::Token;
@@ -979,7 +979,7 @@ pub enum TableConstraint {
         /// [1]: IndexType
         index_type: Option<IndexType>,
         /// Identifiers of the columns that are unique.
-        columns: Vec<Ident>,
+        columns: Vec<IndexColumn>,
         index_options: Vec<IndexOption>,
         characteristics: Option<ConstraintCharacteristics>,
         /// Optional Postgres nulls handling: `[ NULLS [ NOT ] DISTINCT ]`
@@ -1015,7 +1015,7 @@ pub enum TableConstraint {
         /// [1]: IndexType
         index_type: Option<IndexType>,
         /// Identifiers of the columns that form the primary key.
-        columns: Vec<Ident>,
+        columns: Vec<IndexColumn>,
         index_options: Vec<IndexOption>,
         characteristics: Option<ConstraintCharacteristics>,
     },
@@ -1060,7 +1060,7 @@ pub enum TableConstraint {
         /// [1]: IndexType
         index_type: Option<IndexType>,
         /// Referred column identifier list.
-        columns: Vec<Ident>,
+        columns: Vec<IndexColumn>,
     },
     /// MySQLs [fulltext][1] definition. Since the [`SPATIAL`][2] definition is exactly the same,
     /// and MySQL displays both the same way, it is part of this definition as well.
@@ -1083,7 +1083,7 @@ pub enum TableConstraint {
         /// Optional index name.
         opt_index_name: Option<Ident>,
         /// Referred column identifier list.
-        columns: Vec<Ident>,
+        columns: Vec<IndexColumn>,
     },
 }
 

src/ast/spans.rs

@@ -28,16 +28,17 @@ use super::{
     ConstraintCharacteristics, CopySource, CreateIndex, CreateTable, CreateTableOptions, Cte,
     Delete, DoUpdate, ExceptSelectItem, ExcludeSelectItem, Expr, ExprWithAlias, Fetch, FromTable,
     Function, FunctionArg, FunctionArgExpr, FunctionArgumentClause, FunctionArgumentList,
-    FunctionArguments, GroupByExpr, HavingBound, IfStatement, IlikeSelectItem, Insert, Interpolate,
+    FunctionArguments, GroupByExpr, HavingBound, IfStatement, IlikeSelectItem, IndexColumn, Insert,
-    InterpolateExpr, Join, JoinConstraint, JoinOperator, JsonPath, JsonPathElem, LateralView,
+    Interpolate, InterpolateExpr, Join, JoinConstraint, JoinOperator, JsonPath, JsonPathElem,
-    LimitClause, MatchRecognizePattern, Measure, NamedParenthesizedList, NamedWindowDefinition,
+    LateralView, LimitClause, MatchRecognizePattern, Measure, NamedParenthesizedList,
-    ObjectName, ObjectNamePart, Offset, OnConflict, OnConflictAction, OnInsert, OpenStatement,
+    NamedWindowDefinition, ObjectName, ObjectNamePart, Offset, OnConflict, OnConflictAction,
-    OrderBy, OrderByExpr, OrderByKind, Partition, PivotValueSource, ProjectionSelect, Query,
+    OnInsert, OpenStatement, OrderBy, OrderByExpr, OrderByKind, Partition, PivotValueSource,
-    RaiseStatement, RaiseStatementValue, ReferentialAction, RenameSelectItem, ReplaceSelectElement,
+    ProjectionSelect, Query, RaiseStatement, RaiseStatementValue, ReferentialAction,
-    ReplaceSelectItem, Select, SelectInto, SelectItem, SetExpr, SqlOption, Statement, Subscript,
+    RenameSelectItem, ReplaceSelectElement, ReplaceSelectItem, Select, SelectInto, SelectItem,
-    SymbolDefinition, TableAlias, TableAliasColumnDef, TableConstraint, TableFactor, TableObject,
+    SetExpr, SqlOption, Statement, Subscript, SymbolDefinition, TableAlias, TableAliasColumnDef,
-    TableOptionsClustered, TableWithJoins, UpdateTableFromKind, Use, Value, Values, ViewColumnDef,
+    TableConstraint, TableFactor, TableObject, TableOptionsClustered, TableWithJoins,
-    WhileStatement, WildcardAdditionalOptions, With, WithFill,
+    UpdateTableFromKind, Use, Value, Values, ViewColumnDef, WhileStatement,
+    WildcardAdditionalOptions, With, WithFill,
 };
 
 /// Given an iterator of spans, return the [Span::union] of all spans.
@@ -650,7 +651,7 @@ impl Spanned for TableConstraint {
                 name.iter()
                     .map(|i| i.span)
                     .chain(index_name.iter().map(|i| i.span))
-                    .chain(columns.iter().map(|i| i.span))
+                    .chain(columns.iter().map(|i| i.span()))
                     .chain(characteristics.iter().map(|i| i.span())),
             ),
             TableConstraint::PrimaryKey {
@@ -664,7 +665,7 @@ impl Spanned for TableConstraint {
                 name.iter()
                     .map(|i| i.span)
                     .chain(index_name.iter().map(|i| i.span))
-                    .chain(columns.iter().map(|i| i.span))
+                    .chain(columns.iter().map(|i| i.span()))
                     .chain(characteristics.iter().map(|i| i.span())),
             ),
             TableConstraint::ForeignKey {
@@ -700,7 +701,7 @@ impl Spanned for TableConstraint {
             } => union_spans(
                 name.iter()
                     .map(|i| i.span)
-                    .chain(columns.iter().map(|i| i.span)),
+                    .chain(columns.iter().map(|i| i.span())),
             ),
             TableConstraint::FulltextOrSpatial {
                 fulltext: _,
@@ -711,7 +712,7 @@ impl Spanned for TableConstraint {
                 opt_index_name
                     .iter()
                     .map(|i| i.span)
-                    .chain(columns.iter().map(|i| i.span)),
+                    .chain(columns.iter().map(|i| i.span())),
             ),
         }
     }
@@ -745,6 +746,12 @@ impl Spanned for CreateIndex {
     }
 }
 
+impl Spanned for IndexColumn {
+    fn span(&self) -> Span {
+        self.column.span()
+    }
+}
+
 impl Spanned for CaseStatement {
     fn span(&self) -> Span {
         let CaseStatement {

src/parser/mod.rs

@@ -6868,9 +6868,7 @@ impl<'a> Parser<'a> {
             None
         };
 
-        self.expect_token(&Token::LParen)?;
+        let columns = self.parse_parenthesized_index_column_list()?;
-        let columns = self.parse_comma_separated(Parser::parse_create_index_expr)?;
-        self.expect_token(&Token::RParen)?;
 
         let include = if self.parse_keyword(Keyword::INCLUDE) {
             self.expect_token(&Token::LParen)?;
@@ -8070,7 +8068,7 @@ impl<'a> Parser<'a> {
                 let index_name = self.parse_optional_ident()?;
                 let index_type = self.parse_optional_using_then_index_type()?;
 
-                let columns = self.parse_parenthesized_column_list(Mandatory, false)?;
+                let columns = self.parse_parenthesized_index_column_list()?;
                 let index_options = self.parse_index_options()?;
                 let characteristics = self.parse_constraint_characteristics()?;
                 Ok(Some(TableConstraint::Unique {
@@ -8092,7 +8090,7 @@ impl<'a> Parser<'a> {
                 let index_name = self.parse_optional_ident()?;
                 let index_type = self.parse_optional_using_then_index_type()?;
 
-                let columns = self.parse_parenthesized_column_list(Mandatory, false)?;
+                let columns = self.parse_parenthesized_index_column_list()?;
                 let index_options = self.parse_index_options()?;
                 let characteristics = self.parse_constraint_characteristics()?;
                 Ok(Some(TableConstraint::PrimaryKey {
@@ -8170,7 +8168,7 @@ impl<'a> Parser<'a> {
                 };
 
                 let index_type = self.parse_optional_using_then_index_type()?;
-                let columns = self.parse_parenthesized_column_list(Mandatory, false)?;
+                let columns = self.parse_parenthesized_index_column_list()?;
 
                 Ok(Some(TableConstraint::Index {
                     display_as_key,
@@ -8199,7 +8197,7 @@ impl<'a> Parser<'a> {
 
                 let opt_index_name = self.parse_optional_ident()?;
 
-                let columns = self.parse_parenthesized_column_list(Mandatory, false)?;
+                let columns = self.parse_parenthesized_index_column_list()?;
 
                 Ok(Some(TableConstraint::FulltextOrSpatial {
                     fulltext,
@@ -10601,6 +10599,14 @@ impl<'a> Parser<'a> {
         self.parse_parenthesized_column_list_inner(optional, allow_empty, |p| p.parse_identifier())
     }
 
+    /// Parses a parenthesized comma-separated list of index columns, which can be arbitrary
+    /// expressions with ordering information (and an opclass in some dialects).
+    fn parse_parenthesized_index_column_list(&mut self) -> Result<Vec<IndexColumn>, ParserError> {
+        self.parse_parenthesized_column_list_inner(Mandatory, false, |p| {
+            p.parse_create_index_expr()
+        })
+    }
+
     /// Parses a parenthesized comma-separated list of qualified, possibly quoted identifiers.
     /// For example: `(db1.sc1.tbl1.col1, db1.sc1.tbl1."col 2", ...)`
     pub fn parse_parenthesized_qualified_column_list(
@@ -16527,6 +16533,20 @@ mod tests {
             }};
         }
 
+        fn mk_expected_col(name: &str) -> IndexColumn {
+            IndexColumn {
+                column: OrderByExpr {
+                    expr: Expr::Identifier(name.into()),
+                    options: OrderByOptions {
+                        asc: None,
+                        nulls_first: None,
+                    },
+                    with_fill: None,
+                },
+                operator_class: None,
+            }
+        }
+
         let dialect =
             TestedDialects::new(vec![Box::new(GenericDialect {}), Box::new(MySqlDialect {})]);
 
@@ -16537,7 +16557,7 @@ mod tests {
                 display_as_key: false,
                 name: None,
                 index_type: None,
-                columns: vec![Ident::new("c1")],
+                columns: vec![mk_expected_col("c1")],
             }
         );
 
@@ -16548,7 +16568,7 @@ mod tests {
                 display_as_key: true,
                 name: None,
                 index_type: None,
-                columns: vec![Ident::new("c1")],
+                columns: vec![mk_expected_col("c1")],
             }
         );
 
@@ -16559,7 +16579,7 @@ mod tests {
                 display_as_key: false,
                 name: Some(Ident::with_quote('\'', "index")),
                 index_type: None,
-                columns: vec![Ident::new("c1"), Ident::new("c2")],
+                columns: vec![mk_expected_col("c1"), mk_expected_col("c2")],
             }
         );
 
@@ -16570,7 +16590,7 @@ mod tests {
                 display_as_key: false,
                 name: None,
                 index_type: Some(IndexType::BTree),
-                columns: vec![Ident::new("c1")],
+                columns: vec![mk_expected_col("c1")],
             }
         );
 
@@ -16581,7 +16601,7 @@ mod tests {
                 display_as_key: false,
                 name: None,
                 index_type: Some(IndexType::Hash),
-                columns: vec![Ident::new("c1")],
+                columns: vec![mk_expected_col("c1")],
             }
         );
 
@@ -16592,7 +16612,7 @@ mod tests {
                 display_as_key: false,
                 name: Some(Ident::new("idx_name")),
                 index_type: Some(IndexType::BTree),
-                columns: vec![Ident::new("c1")],
+                columns: vec![mk_expected_col("c1")],
             }
         );
 
@@ -16603,7 +16623,7 @@ mod tests {
                 display_as_key: false,
                 name: Some(Ident::new("idx_name")),
                 index_type: Some(IndexType::Hash),
-                columns: vec![Ident::new("c1")],
+                columns: vec![mk_expected_col("c1")],
             }
         );
     }

src/test_utils.rs

@@ -448,3 +448,47 @@ pub fn call(function: &str, args: impl IntoIterator<Item = Expr>) -> Expr {
         within_group: vec![],
     })
 }
+
+/// Gets the first index column (mysql calls it a key part) of the first index found in a
+/// [`Statement::CreateIndex`], [`Statement::CreateTable`], or [`Statement::AlterTable`].
+pub fn index_column(stmt: Statement) -> Expr {
+    match stmt {
+        Statement::CreateIndex(CreateIndex { columns, .. }) => {
+            columns.first().unwrap().column.expr.clone()
+        }
+        Statement::CreateTable(CreateTable { constraints, .. }) => {
+            match constraints.first().unwrap() {
+                TableConstraint::Index { columns, .. } => {
+                    columns.first().unwrap().column.expr.clone()
+                }
+                TableConstraint::Unique { columns, .. } => {
+                    columns.first().unwrap().column.expr.clone()
+                }
+                TableConstraint::PrimaryKey { columns, .. } => {
+                    columns.first().unwrap().column.expr.clone()
+                }
+                TableConstraint::FulltextOrSpatial { columns, .. } => {
+                    columns.first().unwrap().column.expr.clone()
+                }
+                _ => panic!("Expected an index, unique, primary, full text, or spatial constraint (foreign key does not support general key part expressions)"),
+            }
+        }
+        Statement::AlterTable { operations, .. } => match operations.first().unwrap() {
+            AlterTableOperation::AddConstraint(TableConstraint::Index { columns, .. }) => {
+                columns.first().unwrap().column.expr.clone()
+            }
+            AlterTableOperation::AddConstraint(TableConstraint::Unique { columns, .. }) => {
+                columns.first().unwrap().column.expr.clone()
+            }
+            AlterTableOperation::AddConstraint(TableConstraint::PrimaryKey { columns, .. }) => {
+                columns.first().unwrap().column.expr.clone()
+            }
+            AlterTableOperation::AddConstraint(TableConstraint::FulltextOrSpatial {
+                columns,
+                ..
+            }) => columns.first().unwrap().column.expr.clone(),
+            _ => panic!("Expected an index, unique, primary, full text, or spatial constraint (foreign key does not support general key part expressions)"),
+        },
+        _ => panic!("Expected CREATE INDEX, ALTER TABLE, or CREATE TABLE, got: {stmt:?}"),
+    }
+}

tests/sqlparser_mysql.rs

@@ -670,6 +670,20 @@ fn table_constraint_unique_primary_ctor(
     characteristics: Option<ConstraintCharacteristics>,
     unique_index_type_display: Option<KeyOrIndexDisplay>,
 ) -> TableConstraint {
+    let columns = columns
+        .into_iter()
+        .map(|ident| IndexColumn {
+            column: OrderByExpr {
+                expr: Expr::Identifier(ident),
+                options: OrderByOptions {
+                    asc: None,
+                    nulls_first: None,
+                },
+                with_fill: None,
+            },
+            operator_class: None,
+        })
+        .collect();
     match unique_index_type_display {
         Some(index_type_display) => TableConstraint::Unique {
             name,
@@ -795,6 +809,67 @@ fn parse_create_table_primary_and_unique_key_with_index_options() {
     }
 }
 
+#[test]
+fn parse_prefix_key_part() {
+    let expected = vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::value(
+        number("10"),
+    )))];
+    for sql in [
+        "CREATE INDEX idx_index ON t(textcol(10))",
+        "ALTER TABLE tab ADD INDEX idx_index (textcol(10))",
+        "ALTER TABLE tab ADD PRIMARY KEY (textcol(10))",
+        "ALTER TABLE tab ADD UNIQUE KEY (textcol(10))",
+        "ALTER TABLE tab ADD UNIQUE KEY (textcol(10))",
+        "ALTER TABLE tab ADD FULLTEXT INDEX (textcol(10))",
+        "CREATE TABLE t (textcol TEXT, INDEX idx_index (textcol(10)))",
+    ] {
+        match index_column(mysql_and_generic().verified_stmt(sql)) {
+            Expr::Function(Function {
+                name,
+                args: FunctionArguments::List(FunctionArgumentList { args, .. }),
+                ..
+            }) => {
+                assert_eq!(name.to_string(), "textcol");
+                assert_eq!(args, expected);
+            }
+            expr => panic!("unexpected expression {expr} for {sql}"),
+        }
+    }
+}
+
+#[test]
+fn test_functional_key_part() {
+    assert_eq!(
+        index_column(
+            mysql_and_generic()
+                .verified_stmt("CREATE INDEX idx_index ON t((col COLLATE utf8mb4_bin) DESC)")
+        ),
+        Expr::Nested(Box::new(Expr::Collate {
+            expr: Box::new(Expr::Identifier("col".into())),
+            collation: ObjectName(vec![sqlparser::ast::ObjectNamePart::Identifier(
+                Ident::new("utf8mb4_bin")
+            )]),
+        }))
+    );
+    assert_eq!(
+        index_column(mysql_and_generic().verified_stmt(
+            r#"CREATE TABLE t (jsoncol JSON, PRIMARY KEY ((CAST(col ->> '$.id' AS UNSIGNED)) ASC))"#
+        )),
+        Expr::Nested(Box::new(Expr::Cast {
+            kind: CastKind::Cast,
+            expr: Box::new(Expr::BinaryOp {
+                left: Box::new(Expr::Identifier(Ident::new("col"))),
+                op: BinaryOperator::LongArrow,
+                right: Box::new(Expr::Value(
+                    Value::SingleQuotedString("$.id".to_string()).with_empty_span()
+                )),
+            }),
+            data_type: DataType::Unsigned,
+            format: None,
+        })),
+    );
+}
+
 #[test]
 fn parse_create_table_primary_and_unique_key_with_index_type() {
     let sqls = ["UNIQUE", "PRIMARY KEY"].map(|key_ty| {