added parsing for PostgreSQL operations

Author: alex-dukhno

src/ast/mod.rs

@@ -191,7 +191,11 @@ pub enum Expr {
         right: Box<Expr>,
     },
     /// Unary operation e.g. `NOT foo`
-    UnaryOp { op: UnaryOperator, expr: Box<Expr> },
+    UnaryOp {
+        op: UnaryOperator,
+        expr: Box<Expr>,
+        infix: bool,
+    },
     /// CAST an expression to a different data type e.g. `CAST(foo AS VARCHAR(123))`
     Cast {
         expr: Box<Expr>,
@@ -282,7 +286,13 @@ impl fmt::Display for Expr {
                 high
             ),
             Expr::BinaryOp { left, op, right } => write!(f, "{} {} {}", left, op, right),
-            Expr::UnaryOp { op, expr } => write!(f, "{} {}", op, expr),
+            Expr::UnaryOp { op, expr, infix } => {
+                if *infix {
+                    write!(f, "{}{}", expr, op)
+                } else {
+                    write!(f, "{} {}", op, expr)
+                }
+            }
             Expr::Cast { expr, data_type } => write!(f, "CAST({} AS {})", expr, data_type),
             Expr::Extract { field, expr } => write!(f, "EXTRACT({} FROM {})", field, expr),
             Expr::Collate { expr, collation } => write!(f, "{} COLLATE {}", expr, collation),

src/ast/operator.rs

@@ -21,6 +21,12 @@ pub enum UnaryOperator {
     Plus,
     Minus,
     Not,
+    PGBitwiseNot,
+    PGSqrt,
+    PGCbrt,
+    PGFactorial,
+    PGInfixFactorial,
+    PGAbs,
 }
 
 impl fmt::Display for UnaryOperator {
@@ -29,6 +35,12 @@ impl fmt::Display for UnaryOperator {
             UnaryOperator::Plus => "+",
             UnaryOperator::Minus => "-",
             UnaryOperator::Not => "NOT",
+            UnaryOperator::PGBitwiseNot => "~",
+            UnaryOperator::PGSqrt => "|/",
+            UnaryOperator::PGCbrt => "||/",
+            UnaryOperator::PGFactorial => "!",
+            UnaryOperator::PGInfixFactorial => "!!",
+            UnaryOperator::PGAbs => "@",
         })
     }
 }
@@ -56,6 +68,9 @@ pub enum BinaryOperator {
     BitwiseOr,
     BitwiseAnd,
     BitwiseXor,
+    PGBitwiseXor,
+    PGBitwiseShiftLeft,
+    PGBitwiseShiftRight,
 }
 
 impl fmt::Display for BinaryOperator {
@@ -80,6 +95,9 @@ impl fmt::Display for BinaryOperator {
             BinaryOperator::BitwiseOr => "|",
             BinaryOperator::BitwiseAnd => "&",
             BinaryOperator::BitwiseXor => "^",
+            BinaryOperator::PGBitwiseXor => "#",
+            BinaryOperator::PGBitwiseShiftLeft => "<<",
+            BinaryOperator::PGBitwiseShiftRight => ">>",
         })
     }
 }

src/parser.rs

@@ -250,6 +250,7 @@ impl<'a> Parser<'a> {
                 Keyword::NOT => Ok(Expr::UnaryOp {
                     op: UnaryOperator::Not,
                     expr: Box::new(self.parse_subexpr(Self::UNARY_NOT_PREC)?),
+                    infix: false,
                 }),
                 // Here `w` is a word, check if it's a part of a multi-part
                 // identifier, a function call, or a simple identifier:
@@ -283,6 +284,31 @@ impl<'a> Parser<'a> {
                 },
             }, // End of Token::Word
             Token::Mult => Ok(Expr::Wildcard),
+            Token::Tilde => Ok(Expr::UnaryOp {
+                op: UnaryOperator::PGBitwiseNot,
+                expr: Box::new(self.parse_subexpr(0)?),
+                infix: false,
+            }),
+            Token::DoubleExclamationMark => Ok(Expr::UnaryOp {
+                op: UnaryOperator::PGInfixFactorial,
+                expr: Box::new(self.parse_subexpr(0)?),
+                infix: false,
+            }),
+            Token::SquareRoot => Ok(Expr::UnaryOp {
+                op: UnaryOperator::PGSqrt,
+                expr: Box::new(self.parse_subexpr(0)?),
+                infix: false,
+            }),
+            Token::CubeRoot => Ok(Expr::UnaryOp {
+                op: UnaryOperator::PGCbrt,
+                expr: Box::new(self.parse_subexpr(0)?),
+                infix: false,
+            }),
+            Token::Ampersat => Ok(Expr::UnaryOp {
+                op: UnaryOperator::PGAbs,
+                expr: Box::new(self.parse_subexpr(0)?),
+                infix: false,
+            }),
             tok @ Token::Minus | tok @ Token::Plus => {
                 let op = if tok == Token::Plus {
                     UnaryOperator::Plus
@@ -292,6 +318,7 @@ impl<'a> Parser<'a> {
                 Ok(Expr::UnaryOp {
                     op,
                     expr: Box::new(self.parse_subexpr(Self::PLUS_MINUS_PREC)?),
+                    infix: false,
                 })
             }
             Token::Number(_)
@@ -658,6 +685,9 @@ impl<'a> Parser<'a> {
             Token::Caret => Some(BinaryOperator::BitwiseXor),
             Token::Ampersand => Some(BinaryOperator::BitwiseAnd),
             Token::Div => Some(BinaryOperator::Divide),
+            Token::ShiftLeft => Some(BinaryOperator::PGBitwiseShiftLeft),
+            Token::ShiftRight => Some(BinaryOperator::PGBitwiseShiftRight),
+            Token::Sharp => Some(BinaryOperator::PGBitwiseXor),
             Token::Word(w) => match w.keyword {
                 Keyword::AND => Some(BinaryOperator::And),
                 Keyword::OR => Some(BinaryOperator::Or),
@@ -707,6 +737,13 @@ impl<'a> Parser<'a> {
             }
         } else if Token::DoubleColon == tok {
             self.parse_pg_cast(expr)
+        } else if Token::ExclamationMark == tok {
+            // PostgreSQL factorial operation
+            Ok(Expr::UnaryOp {
+                op: UnaryOperator::PGFactorial,
+                expr: Box::new(expr),
+                infix: true,
+            })
         } else {
             // Can only happen if `get_next_precedence` got out of sync with this function
             panic!("No infix parser for token {:?}", tok)
@@ -785,11 +822,12 @@ impl<'a> Parser<'a> {
             Token::Word(w) if w.keyword == Keyword::LIKE => Ok(Self::BETWEEN_PREC),
             Token::Eq | Token::Lt | Token::LtEq | Token::Neq | Token::Gt | Token::GtEq => Ok(20),
             Token::Pipe => Ok(21),
-            Token::Caret => Ok(22),
+            Token::Caret | Token::Sharp | Token::ShiftRight | Token::ShiftLeft => Ok(22),
             Token::Ampersand => Ok(23),
             Token::Plus | Token::Minus => Ok(Self::PLUS_MINUS_PREC),
             Token::Mult | Token::Div | Token::Mod | Token::StringConcat => Ok(40),
             Token::DoubleColon => Ok(50),
+            Token::ExclamationMark => Ok(50),
             _ => Ok(0),
         }
     }

src/tokenizer.rs

@@ -54,7 +54,7 @@ pub enum Token {
     Neq,
     /// Less Than operator `<`
     Lt,
-    /// Greater han operator `>`
+    /// Greater Than operator `>`
     Gt,
     /// Less Than Or Equals operator `<=`
     LtEq,
@@ -102,6 +102,24 @@ pub enum Token {
     RBrace,
     /// Right Arrow `=>`
     RArrow,
+    /// Sharp `#` use for PostgreSQL Bitwise XOR operator
+    Sharp,
+    /// Tilde `~` use for PostgreSQL Bitwise NOT operator
+    Tilde,
+    /// Bitwise left operator `<<` use for PostgreSQL
+    ShiftLeft,
+    /// Bitwise right operator `>>` use for PostgreSQL
+    ShiftRight,
+    /// Exclamation Mark `!` use for PostgreSQL factorial operator
+    ExclamationMark,
+    /// Exclamation Mark `!!` use for PostgreSQL prefix factorial operator
+    DoubleExclamationMark,
+    /// Ampersat `@` use for PostgreSQL abs operator
+    Ampersat,
+    /// PostgreSQL square root math operator
+    SquareRoot,
+    /// PostgreSQL cube root math operator
+    CubeRoot,
 }
 
 impl fmt::Display for Token {
@@ -143,6 +161,15 @@ impl fmt::Display for Token {
             Token::LBrace => f.write_str("{"),
             Token::RBrace => f.write_str("}"),
             Token::RArrow => f.write_str("=>"),
+            Token::Sharp => f.write_str("#"),
+            Token::ExclamationMark => f.write_str("!"),
+            Token::DoubleExclamationMark => f.write_str("!!"),
+            Token::Tilde => f.write_str("~"),
+            Token::Ampersat => f.write_str("@"),
+            Token::ShiftLeft => f.write_str("<<"),
+            Token::ShiftRight => f.write_str(">>"),
+            Token::SquareRoot => f.write_str("|/"),
+            Token::CubeRoot => f.write_str("||/"),
         }
     }
 }
@@ -406,7 +433,14 @@ impl<'a> Tokenizer<'a> {
                 '|' => {
                     chars.next(); // consume the '|'
                     match chars.peek() {
-                        Some('|') => self.consume_and_return(chars, Token::StringConcat),
+                        Some('/') => self.consume_and_return(chars, Token::SquareRoot),
+                        Some('|') => {
+                            chars.next(); // consume the second '|'
+                            match chars.peek() {
+                                Some('/') => self.consume_and_return(chars, Token::CubeRoot),
+                                _ => Ok(Some(Token::StringConcat)),
+                            }
+                        }
                         // Bitshift '|' operator
                         _ => Ok(Some(Token::Pipe)),
                     }
@@ -423,21 +457,24 @@ impl<'a> Tokenizer<'a> {
                     chars.next(); // consume
                     match chars.peek() {
                         Some('=') => self.consume_and_return(chars, Token::Neq),
-                        _ => self.tokenizer_error("Expected to see '=' after '!' character"),
+                        Some('!') => self.consume_and_return(chars, Token::DoubleExclamationMark),
+                        _ => Ok(Some(Token::ExclamationMark)),
                     }
                 }
                 '<' => {
                     chars.next(); // consume
                     match chars.peek() {
                         Some('=') => self.consume_and_return(chars, Token::LtEq),
                         Some('>') => self.consume_and_return(chars, Token::Neq),
+                        Some('<') => self.consume_and_return(chars, Token::ShiftLeft),
                         _ => Ok(Some(Token::Lt)),
                     }
                 }
                 '>' => {
                     chars.next(); // consume
                     match chars.peek() {
                         Some('=') => self.consume_and_return(chars, Token::GtEq),
+                        Some('>') => self.consume_and_return(chars, Token::ShiftRight),
                         _ => Ok(Some(Token::Gt)),
                     }
                 }
@@ -464,6 +501,9 @@ impl<'a> Tokenizer<'a> {
                         comment,
                     })))
                 }
+                '~' => self.consume_and_return(chars, Token::Tilde),
+                '#' => self.consume_and_return(chars, Token::Sharp),
+                '@' => self.consume_and_return(chars, Token::Ampersat),
                 other => self.consume_and_return(chars, Token::Char(other)),
             },
             None => Ok(None),
@@ -586,6 +626,7 @@ mod tests {
     use super::super::dialect::GenericDialect;
     use super::super::dialect::MsSqlDialect;
     use super::*;
+    use crate::dialect::PostgreSqlDialect;
 
     #[test]
     fn tokenize_select_1() {
@@ -958,6 +999,65 @@ mod tests {
         compare(expected, tokens);
     }
 
+    #[test]
+    fn tokenize_postgresql_bitwise_operations() {
+        let sql = String::from("SELECT ~one << two # three >> four");
+        let dialect = PostgreSqlDialect {};
+        let mut tokenizer = Tokenizer::new(&dialect, &sql);
+        let tokens = tokenizer.tokenize().unwrap();
+
+        let expected = vec![
+            Token::make_keyword("SELECT"),
+            Token::Whitespace(Whitespace::Space),
+            Token::Tilde,
+            Token::make_word("one", None),
+            Token::Whitespace(Whitespace::Space),
+            Token::ShiftLeft,
+            Token::Whitespace(Whitespace::Space),
+            Token::make_word("two", None),
+            Token::Whitespace(Whitespace::Space),
+            Token::Sharp,
+            Token::Whitespace(Whitespace::Space),
+            Token::make_word("three", None),
+            Token::Whitespace(Whitespace::Space),
+            Token::ShiftRight,
+            Token::Whitespace(Whitespace::Space),
+            Token::make_word("four", None),
+        ];
+
+        compare(expected, tokens);
+    }
+
+    #[test]
+    fn tokenize_postgresql_math_operations() {
+        let sql = String::from("SELECT !!5 5! @-6 |/4 ||/8");
+        let dialect = PostgreSqlDialect {};
+        let mut tokenizer = Tokenizer::new(&dialect, &sql);
+        let tokens = tokenizer.tokenize().unwrap();
+
+        let expected = vec![
+            Token::make_keyword("SELECT"),
+            Token::Whitespace(Whitespace::Space),
+            Token::DoubleExclamationMark,
+            Token::Number("5".to_string()),
+            Token::Whitespace(Whitespace::Space),
+            Token::Number("5".to_string()),
+            Token::ExclamationMark,
+            Token::Whitespace(Whitespace::Space),
+            Token::Ampersat,
+            Token::Minus,
+            Token::Number("6".to_string()),
+            Token::Whitespace(Whitespace::Space),
+            Token::SquareRoot,
+            Token::Number("4".to_string()),
+            Token::Whitespace(Whitespace::Space),
+            Token::CubeRoot,
+            Token::Number("8".to_string()),
+        ];
+
+        compare(expected, tokens);
+    }
+
     fn compare(expected: Vec<Token>, actual: Vec<Token>) {
         //println!("------------------------------");
         //println!("tokens   = {:?}", actual);

tests/sqlparser_common.rs

@@ -343,7 +343,8 @@ fn parse_select_count_distinct() {
             name: ObjectName(vec![Ident::new("COUNT")]),
             args: vec![FunctionArg::Unnamed(Expr::UnaryOp {
                 op: UnaryOperator::Plus,
-                expr: Box::new(Expr::Identifier(Ident::new("x")))
+                expr: Box::new(Expr::Identifier(Ident::new("x"))),
+                infix: false,
             })],
             over: None,
             distinct: true,
@@ -506,11 +507,13 @@ fn parse_unary_math() {
             left: Box::new(UnaryOp {
                 op: UnaryOperator::Minus,
                 expr: Box::new(Identifier(Ident::new("a"))),
+                infix: false,
             }),
             op: BinaryOperator::Plus,
             right: Box::new(UnaryOp {
                 op: UnaryOperator::Minus,
                 expr: Box::new(Identifier(Ident::new("b"))),
+                infix: false,
             }),
         },
         verified_expr(sql)
@@ -565,6 +568,7 @@ fn parse_not_precedence() {
                 high: Box::new(Expr::Value(number("2"))),
                 negated: true,
             }),
+            infix: false,
         },
     );
 
@@ -579,6 +583,7 @@ fn parse_not_precedence() {
                 op: BinaryOperator::NotLike,
                 right: Box::new(Expr::Value(Value::SingleQuotedString("b".into()))),
             }),
+            infix: false,
         },
     );
 
@@ -593,6 +598,7 @@ fn parse_not_precedence() {
                 list: vec![Expr::Value(Value::SingleQuotedString("a".into()))],
                 negated: true,
             }),
+            infix: false,
         },
     );
 }
@@ -2604,6 +2610,7 @@ fn parse_exists_subquery() {
         Expr::UnaryOp {
             op: UnaryOperator::Not,
             expr: Box::new(Expr::Exists(Box::new(expected_inner))),
+            infix: false,
         },
         select.selection.unwrap(),
     );