sql_parser: parsing indexed columns in upsert clause

Almaz-KG · Almaz-KG · commit c78ad1039ad8 · 2024-12-19T17:50:45.000+01:00
diff --git a/src/ast/expression.rs b/src/ast/expression.rs
@@ -32,7 +32,7 @@ pub enum Expression {
     Cast(Box<Expression>, DataType),
 
     /// A collate expression (e.g. expression COLLATE collation_name)
-    CollateExpression(Box<Expression>, String),
+    CollateExpression(CollateExpressionStatement),
 
     /// A binary matching expressions (e.g. $expr1 NOT MATCH $expr2)
     BinaryMatchingExpression(Box<Expression>, BinaryMatchingExpression),
@@ -489,6 +489,16 @@ pub struct CaseExpression {
     pub else_expression: Option<Box<Expression>>,
 }
 
+/// A collate expression
+#[derive(Debug, PartialEq, Clone)]
+pub struct CollateExpressionStatement {
+    /// The expression
+    pub expression: Box<Expression>,
+
+    /// The collation name
+    pub collation_name: String,
+}
+
 /// A when expression
 #[derive(Debug, PartialEq, Clone)]
 pub struct WhenExpression {
diff --git a/src/ast/insert.rs b/src/ast/insert.rs
@@ -1,6 +1,6 @@
 use super::{
-    ConflictClause, Expression, Identifier, QualifiedTableName, ReturningClause, SelectStatement,
-    SetClause,
+    ConflictClause, Expression, Identifier, Ordering, QualifiedTableName, ReturningClause,
+    SelectStatement, SetClause,
 };
 
 /// An AST for [INSERT](https://www.sqlite.org/lang_insert.html) SQL statement.
@@ -45,10 +45,18 @@ pub struct UpsertClause {
 /// to update the table with the new values.
 #[derive(Debug, PartialEq)]
 pub struct UpsertConflictTarget {
-    pub columns: Vec<Identifier>,
+    pub columns: Vec<IndexedColumn>,
     pub where_clause: Option<Box<Expression>>,
 }
 
+/// An indexed column
+#[derive(Debug, PartialEq)]
+pub struct IndexedColumn {
+    pub column: Expression,
+
+    pub ordering: Option<Ordering>,
+}
+
 /// The [action](https://www.sqlite.org/lang_upsert.html#upsert_action) to take when a conflict occurs.
 #[derive(Debug, PartialEq)]
 pub enum UpsertAction {
diff --git a/src/parser/expression/collate_expr.rs b/src/parser/expression/collate_expr.rs
@@ -1,4 +1,4 @@
-use crate::{Expression, Keyword, Parser, ParsingError};
+use crate::{CollateExpressionStatement, Expression, Keyword, Parser, ParsingError};
 
 pub trait CollateExpressionParser {
     /// Parse a COLLATE expression, i.e. an expression followed by the COLLATE keyword and a string literal
@@ -18,10 +18,10 @@ impl<'a> CollateExpressionParser for Parser<'a> {
         let name = self.peek_as_string()?;
         self.consume_token()?;
 
-        Ok(Expression::CollateExpression(
-            Box::new(expression),
-            name.to_string(),
-        ))
+        Ok(Expression::CollateExpression(CollateExpressionStatement {
+            expression: Box::new(expression),
+            collation_name: name.to_string(),
+        }))
     }
 }
 
diff --git a/src/parser/expression/mod.rs b/src/parser/expression/mod.rs
@@ -395,8 +395,9 @@ impl<'a> ExpressionParser for Parser<'a> {
 pub(crate) mod test_utils {
     use crate::ast::{Expression, SelectItem};
     use crate::{
-        BinaryOp, DataType, ExistsStatement, Function, FunctionArg, Identifier, LiteralValue,
-        OverClause, Parser, RaiseFunction, SelectStatement, Statement, UnaryOp,
+        BinaryOp, CollateExpressionStatement, DataType, ExistsStatement, Function, FunctionArg,
+        Identifier, LiteralValue, OverClause, Parser, RaiseFunction, SelectStatement, Statement,
+        UnaryOp,
     };
 
     pub fn run_sunny_day_test_with_multiple_expressions(
@@ -516,7 +517,10 @@ pub(crate) mod test_utils {
     }
 
     pub fn collate_expression(expression: Expression, name: String) -> Expression {
-        Expression::CollateExpression(Box::new(expression), name)
+        Expression::CollateExpression(CollateExpressionStatement {
+            expression: Box::new(expression),
+            collation_name: name,
+        })
     }
 
     pub fn expression_list(expressions: Vec<Expression>) -> Expression {
diff --git a/src/parser/insert/mod.rs b/src/parser/insert/mod.rs
@@ -1,7 +1,7 @@
 use crate::{
-    parser::select::SelectStatementParser, ConflictClause, Identifier, InsertStatement,
-    InsertValues, Keyword, QualifiedTableName, TokenType, UpsertAction, UpsertClause,
-    UpsertConflictTarget, UpsertUpdate,
+    parser::select::SelectStatementParser, ConflictClause, Identifier, IndexedColumn,
+    InsertStatement, InsertValues, Keyword, Ordering, QualifiedTableName, TokenType, UpsertAction,
+    UpsertClause, UpsertConflictTarget, UpsertUpdate,
 };
 
 use super::{
@@ -29,6 +29,10 @@ pub trait InsertStatementParser {
     ) -> Result<Option<UpsertConflictTarget>, ParsingError>;
 
     fn parse_upsert_action(&mut self) -> Result<UpsertAction, ParsingError>;
+
+    fn parse_indexed_column(&mut self) -> Result<IndexedColumn, ParsingError>;
+
+    fn parse_ordering(&mut self) -> Result<Option<Ordering>, ParsingError>;
 }
 
 impl<'a> InsertStatementParser for Parser<'a> {
@@ -154,7 +158,7 @@ impl<'a> InsertStatementParser for Parser<'a> {
         if self.consume_as(TokenType::LeftParen).is_ok() {
             let mut columns = vec![];
             loop {
-                columns.push(self.parse_identifier()?);
+                columns.push(self.parse_indexed_column()?);
                 if self.consume_as(TokenType::Comma).is_err() {
                     break;
                 }
@@ -194,6 +198,21 @@ impl<'a> InsertStatementParser for Parser<'a> {
             self.peek_token()?.to_string(),
         ))
     }
+
+    fn parse_indexed_column(&mut self) -> Result<IndexedColumn, ParsingError> {
+        let column = self.parse_expression()?;
+        let ordering = self.parse_ordering()?;
+        Ok(IndexedColumn { column, ordering })
+    }
+
+    fn parse_ordering(&mut self) -> Result<Option<Ordering>, ParsingError> {
+        if self.consume_as_keyword(Keyword::Asc).is_ok() {
+            return Ok(Some(Ordering::Asc));
+        } else if self.consume_as_keyword(Keyword::Desc).is_ok() {
+            return Ok(Some(Ordering::Desc));
+        }
+        Ok(None)
+    }
 }
 
 #[cfg(test)]
@@ -212,29 +231,23 @@ mod test_utils {
             returning_clause: vec![],
         }
     }
-
-    // pub fn insert_statement_with_values(values: InsertValues) -> InsertStatement {
-    //     let mut statement = insert_statement();
-    //     statement.values = values;
-    //     statement
-    // }
 }
 
 #[cfg(test)]
 mod tests_insert_statement {
     use super::test_utils::*;
     use crate::{
         expression::test_utils::{
-            binary_op_expression, expression_list, identifier_expression,
+            binary_op_expression, collate_expression, expression_list, identifier_expression,
             numeric_literal_expression,
         },
         parser::{
             cte::test_utils::cte_expression, select::test_utils::select_from,
             test_utils::run_sunny_day_test,
         },
-        BinaryOp, ConflictClause, CteExpression, FromClause, Identifier, InsertValues,
-        QualifiedTableName, ReturningClause, SetClause, Statement, UpsertAction, UpsertClause,
-        UpsertConflictTarget, UpsertUpdate, WithCteStatement,
+        BinaryOp, ConflictClause, CteExpression, FromClause, Identifier, IndexedColumn,
+        InsertValues, QualifiedTableName, ReturningClause, SetClause, Statement, UpsertAction,
+        UpsertClause, UpsertConflictTarget, UpsertUpdate, WithCteStatement,
     };
 
     #[test]
@@ -299,6 +312,7 @@ mod tests_insert_statement {
             Identifier::Single("col1".to_string()),
             Identifier::Single("col2".to_string()),
         ];
+
         statement.values = InsertValues::Values(vec![
             vec![
                 numeric_literal_expression("1"),
@@ -352,7 +366,7 @@ mod tests_insert_statement {
         let sql = r#"
             INSERT INTO table1 DEFAULT VALUES 
                 ON CONFLICT (col1) DO NOTHING 
-                ON CONFLICT (col2) 
+                ON CONFLICT (col2 COLLATE utf8) 
                     WHERE col2 > 10
                     DO UPDATE SET col2 = col2 + 1
                 ON CONFLICT (col3, col4)
@@ -365,14 +379,23 @@ mod tests_insert_statement {
         expected_statement.upsert_clause = Some(vec![
             UpsertClause {
                 conflict_target: Some(UpsertConflictTarget {
-                    columns: vec![Identifier::Single("col1".to_string())],
+                    columns: vec![IndexedColumn {
+                        column: identifier_expression(&["col1"]),
+                        ordering: None,
+                    }],
                     where_clause: None,
                 }),
                 action: UpsertAction::Nothing,
             },
             UpsertClause {
                 conflict_target: Some(UpsertConflictTarget {
-                    columns: vec![Identifier::Single("col2".to_string())],
+                    columns: vec![IndexedColumn {
+                        column: collate_expression(
+                            identifier_expression(&["col2"]),
+                            "utf8".to_string(),
+                        ),
+                        ordering: None,
+                    }],
                     where_clause: Some(Box::new(binary_op_expression(
                         BinaryOp::GreaterThan,
                         identifier_expression(&["col2"]),
@@ -394,8 +417,14 @@ mod tests_insert_statement {
             UpsertClause {
                 conflict_target: Some(UpsertConflictTarget {
                     columns: vec![
-                        Identifier::Single("col3".to_string()),
-                        Identifier::Single("col4".to_string()),
+                        IndexedColumn {
+                            column: identifier_expression(&["col3"]),
+                            ordering: None,
+                        },
+                        IndexedColumn {
+                            column: identifier_expression(&["col4"]),
+                            ordering: None,
+                        },
                     ],
                     where_clause: Some(Box::new(binary_op_expression(
                         BinaryOp::EqualsEquals,
diff --git a/src/parser/select/mod.rs b/src/parser/select/mod.rs
@@ -1489,12 +1489,14 @@ mod test_select_having_clause {
 #[cfg(test)]
 mod test_select_window_clause {
     use super::test_utils::select_statement_with_window_clause;
-    use crate::expression::test_utils::{identifier_expression, numeric_literal_expression};
+    use crate::expression::test_utils::{
+        collate_expression, identifier_expression, numeric_literal_expression,
+    };
     use crate::parser::test_utils::*;
     use crate::{
-        BetweenFrameSpec, BetweenFrameSpecType, Expression, FrameSpec, FrameSpecExclude,
-        FrameSpecType, FrameType, NamedWindowDefinition, NullsOrdering, Ordering, OrderingTerm,
-        Statement, WindowDefinition,
+        BetweenFrameSpec, BetweenFrameSpecType, FrameSpec, FrameSpecExclude, FrameSpecType,
+        FrameType, NamedWindowDefinition, NullsOrdering, Ordering, OrderingTerm, Statement,
+        WindowDefinition,
     };
 
     #[test]
@@ -1526,8 +1528,8 @@ mod test_select_window_clause {
                         nulls_ordering: None,
                     },
                     OrderingTerm {
-                        expression: Box::new(Expression::CollateExpression(
-                            Box::new(identifier_expression(&["col4"])),
+                        expression: Box::new(collate_expression(
+                            identifier_expression(&["col4"]),
                             "binary".to_string(),
                         )),
                         ordering: Some(Ordering::Asc),
@@ -1623,9 +1625,9 @@ mod test_select_union_clause {
 #[cfg(test)]
 mod test_select_order_by_clause {
     use super::test_utils::select_statement_with_order_by_clause;
-    use crate::expression::test_utils::identifier_expression;
+    use crate::expression::test_utils::{collate_expression, identifier_expression};
     use crate::parser::test_utils::*;
-    use crate::{Expression, NullsOrdering, Ordering, OrderingTerm, Statement};
+    use crate::{NullsOrdering, Ordering, OrderingTerm, Statement};
 
     #[test]
     fn test_select_order_by_clause() {
@@ -1680,16 +1682,16 @@ mod test_select_order_by_clause {
     fn test_select_order_by_clause_with_collation_and_nulls_ordering() {
         let expected_statement = select_statement_with_order_by_clause(vec![
             OrderingTerm {
-                expression: Box::new(Expression::CollateExpression(
-                    Box::new(identifier_expression(&["col1"])),
+                expression: Box::new(collate_expression(
+                    identifier_expression(&["col1"]),
                     "binary".to_string(),
                 )),
                 ordering: Some(Ordering::Asc),
                 nulls_ordering: Some(NullsOrdering::Last),
             },
             OrderingTerm {
-                expression: Box::new(Expression::CollateExpression(
-                    Box::new(identifier_expression(&["col2"])),
+                expression: Box::new(collate_expression(
+                    identifier_expression(&["col2"]),
                     "utf8".to_string(),
                 )),
                 ordering: Some(Ordering::Desc),
