Mysql: Add support for := operator (apache#1779)

Author: barsela1

src/ast/operator.rs

@@ -321,6 +321,9 @@ pub enum BinaryOperator {
     /// `~=` Same as? (PostgreSQL/Redshift geometric operator)
     /// See <https://www.postgresql.org/docs/9.5/functions-geometry.html>
     TildeEq,
+    /// ':=' Assignment Operator
+    /// See <https://dev.mysql.com/doc/refman/8.4/en/assignment-operators.html#operator_assign-value>
+    Assignment,
 }
 
 impl fmt::Display for BinaryOperator {
@@ -394,6 +397,7 @@ impl fmt::Display for BinaryOperator {
             BinaryOperator::QuestionDoublePipe => f.write_str("?||"),
             BinaryOperator::At => f.write_str("@"),
             BinaryOperator::TildeEq => f.write_str("~="),
+            BinaryOperator::Assignment => f.write_str(":="),
         }
     }
 }

src/dialect/mod.rs

@@ -620,7 +620,8 @@ pub trait Dialect: Debug + Any {
             Token::Word(w) if w.keyword == Keyword::OPERATOR => Ok(p!(Between)),
             Token::Word(w) if w.keyword == Keyword::DIV => Ok(p!(MulDivModOp)),
             Token::Period => Ok(p!(Period)),
-            Token::Eq
+            Token::Assignment
+            | Token::Eq
             | Token::Lt
             | Token::LtEq
             | Token::Neq

src/parser/mod.rs

@@ -3233,6 +3233,7 @@ impl<'a> Parser<'a> {
         let regular_binary_operator = match &tok.token {
             Token::Spaceship => Some(BinaryOperator::Spaceship),
             Token::DoubleEq => Some(BinaryOperator::Eq),
+            Token::Assignment => Some(BinaryOperator::Assignment),
             Token::Eq => Some(BinaryOperator::Eq),
             Token::Neq => Some(BinaryOperator::NotEq),
             Token::Gt => Some(BinaryOperator::Gt),

tests/sqlparser_mysql.rs

@@ -3483,3 +3483,107 @@ fn parse_match_against_with_alias() {
         _ => unreachable!(),
     }
 }
+
+#[test]
+fn test_variable_assignment_using_colon_equal() {
+    let sql_select = "SELECT @price := price, @tax := price * 0.1 FROM products WHERE id = 1";
+    let stmt = mysql().verified_stmt(sql_select);
+    match stmt {
+        Statement::Query(query) => {
+            let select = query.body.as_select().unwrap();
+
+            assert_eq!(
+                select.projection,
+                vec![
+                    SelectItem::UnnamedExpr(Expr::BinaryOp {
+                        left: Box::new(Expr::Identifier(Ident {
+                            value: "@price".to_string(),
+                            quote_style: None,
+                            span: Span::empty(),
+                        })),
+                        op: BinaryOperator::Assignment,
+                        right: Box::new(Expr::Identifier(Ident {
+                            value: "price".to_string(),
+                            quote_style: None,
+                            span: Span::empty(),
+                        })),
+                    }),
+                    SelectItem::UnnamedExpr(Expr::BinaryOp {
+                        left: Box::new(Expr::Identifier(Ident {
+                            value: "@tax".to_string(),
+                            quote_style: None,
+                            span: Span::empty(),
+                        })),
+                        op: BinaryOperator::Assignment,
+                        right: Box::new(Expr::BinaryOp {
+                            left: Box::new(Expr::Identifier(Ident {
+                                value: "price".to_string(),
+                                quote_style: None,
+                                span: Span::empty(),
+                            })),
+                            op: BinaryOperator::Multiply,
+                            right: Box::new(Expr::Value(
+                                (test_utils::number("0.1")).with_empty_span()
+                            )),
+                        }),
+                    }),
+                ]
+            );
+
+            assert_eq!(
+                select.selection,
+                Some(Expr::BinaryOp {
+                    left: Box::new(Expr::Identifier(Ident {
+                        value: "id".to_string(),
+                        quote_style: None,
+                        span: Span::empty(),
+                    })),
+                    op: BinaryOperator::Eq,
+                    right: Box::new(Expr::Value((test_utils::number("1")).with_empty_span())),
+                })
+            );
+        }
+        _ => panic!("Unexpected statement {stmt}"),
+    }
+
+    let sql_update =
+        "UPDATE products SET price = @new_price := price * 1.1 WHERE category = 'Books'";
+    let stmt = mysql().verified_stmt(sql_update);
+
+    match stmt {
+        Statement::Update { assignments, .. } => {
+            assert_eq!(
+                assignments,
+                vec![Assignment {
+                    target: AssignmentTarget::ColumnName(ObjectName(vec![
+                        ObjectNamePart::Identifier(Ident {
+                            value: "price".to_string(),
+                            quote_style: None,
+                            span: Span::empty(),
+                        })
+                    ])),
+                    value: Expr::BinaryOp {
+                        left: Box::new(Expr::Identifier(Ident {
+                            value: "@new_price".to_string(),
+                            quote_style: None,
+                            span: Span::empty(),
+                        })),
+                        op: BinaryOperator::Assignment,
+                        right: Box::new(Expr::BinaryOp {
+                            left: Box::new(Expr::Identifier(Ident {
+                                value: "price".to_string(),
+                                quote_style: None,
+                                span: Span::empty(),
+                            })),
+                            op: BinaryOperator::Multiply,
+                            right: Box::new(Expr::Value(
+                                (test_utils::number("1.1")).with_empty_span()
+                            )),
+                        }),
+                    },
+                }]
+            )
+        }
+        _ => panic!("Unexpected statement {stmt}"),
+    }
+}