Add parsing for PostgreSQL math operators (#267)

Author: alex-dukhno

src/ast/mod.rs

@@ -282,7 +282,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 } => {
+                if op == &UnaryOperator::PGPostfixFactorial {
+                    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,18 @@ pub enum UnaryOperator {
     Plus,
     Minus,
     Not,
+    /// Bitwise Not, e.g. `~9` (PostgreSQL-specific)
+    PGBitwiseNot,
+    /// Square root, e.g. `|/9` (PostgreSQL-specific)
+    PGSquareRoot,
+    /// Cube root, e.g. `||/27` (PostgreSQL-specific)
+    PGCubeRoot,
+    /// Factorial, e.g. `9!` (PostgreSQL-specific)
+    PGPostfixFactorial,
+    /// Factorial, e.g. `!!9` (PostgreSQL-specific)
+    PGPrefixFactorial,
+    /// Absolute value, e.g. `@ -9` (PostgreSQL-specific)
+    PGAbs,
 }
 
 impl fmt::Display for UnaryOperator {
@@ -29,6 +41,12 @@ impl fmt::Display for UnaryOperator {
             UnaryOperator::Plus => "+",
             UnaryOperator::Minus => "-",
             UnaryOperator::Not => "NOT",
+            UnaryOperator::PGBitwiseNot => "~",
+            UnaryOperator::PGSquareRoot => "|/",
+            UnaryOperator::PGCubeRoot => "||/",
+            UnaryOperator::PGPostfixFactorial => "!",
+            UnaryOperator::PGPrefixFactorial => "!!",
+            UnaryOperator::PGAbs => "@",
         })
     }
 }
@@ -56,6 +74,9 @@ pub enum BinaryOperator {
     BitwiseOr,
     BitwiseAnd,
     BitwiseXor,
+    PGBitwiseXor,
+    PGBitwiseShiftLeft,
+    PGBitwiseShiftRight,
 }
 
 impl fmt::Display for BinaryOperator {
@@ -80,6 +101,9 @@ impl fmt::Display for BinaryOperator {
             BinaryOperator::BitwiseOr => "|",
             BinaryOperator::BitwiseAnd => "&",
             BinaryOperator::BitwiseXor => "^",
+            BinaryOperator::PGBitwiseXor => "#",
+            BinaryOperator::PGBitwiseShiftLeft => "<<",
+            BinaryOperator::PGBitwiseShiftRight => ">>",
         })
     }
 }

src/parser.rs

@@ -294,6 +294,26 @@ impl<'a> Parser<'a> {
                     expr: Box::new(self.parse_subexpr(Self::PLUS_MINUS_PREC)?),
                 })
             }
+            tok @ Token::DoubleExclamationMark
+            | tok @ Token::PGSquareRoot
+            | tok @ Token::PGCubeRoot
+            | tok @ Token::AtSign
+            | tok @ Token::Tilde
+                if dialect_of!(self is PostgreSqlDialect) =>
+            {
+                let op = match tok {
+                    Token::DoubleExclamationMark => UnaryOperator::PGPrefixFactorial,
+                    Token::PGSquareRoot => UnaryOperator::PGSquareRoot,
+                    Token::PGCubeRoot => UnaryOperator::PGCubeRoot,
+                    Token::AtSign => UnaryOperator::PGAbs,
+                    Token::Tilde => UnaryOperator::PGBitwiseNot,
+                    _ => unreachable!(),
+                };
+                Ok(Expr::UnaryOp {
+                    op,
+                    expr: Box::new(self.parse_subexpr(Self::PLUS_MINUS_PREC)?),
+                })
+            }
             Token::Number(_)
             | Token::SingleQuotedString(_)
             | Token::NationalStringLiteral(_)
@@ -658,6 +678,15 @@ impl<'a> Parser<'a> {
             Token::Caret => Some(BinaryOperator::BitwiseXor),
             Token::Ampersand => Some(BinaryOperator::BitwiseAnd),
             Token::Div => Some(BinaryOperator::Divide),
+            Token::ShiftLeft if dialect_of!(self is PostgreSqlDialect) => {
+                Some(BinaryOperator::PGBitwiseShiftLeft)
+            }
+            Token::ShiftRight if dialect_of!(self is PostgreSqlDialect) => {
+                Some(BinaryOperator::PGBitwiseShiftRight)
+            }
+            Token::Sharp if dialect_of!(self is PostgreSqlDialect) => {
+                Some(BinaryOperator::PGBitwiseXor)
+            }
             Token::Word(w) => match w.keyword {
                 Keyword::AND => Some(BinaryOperator::And),
                 Keyword::OR => Some(BinaryOperator::Or),
@@ -707,6 +736,12 @@ 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::PGPostfixFactorial,
+                expr: Box::new(expr),
+            })
         } else {
             // Can only happen if `get_next_precedence` got out of sync with this function
             panic!("No infix parser for token {:?}", tok)
@@ -785,11 +820,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 `#` used for PostgreSQL Bitwise XOR operator
+    Sharp,
+    /// Tilde `~` used for PostgreSQL Bitwise NOT operator
+    Tilde,
+    /// `<<`, a bitwise shift left operator in PostgreSQL
+    ShiftLeft,
+    /// `>>`, a bitwise shift right operator in PostgreSQL
+    ShiftRight,
+    /// Exclamation Mark `!` used for PostgreSQL factorial operator
+    ExclamationMark,
+    /// Double Exclamation Mark `!!` used for PostgreSQL prefix factorial operator
+    DoubleExclamationMark,
+    /// AtSign `@` used for PostgreSQL abs operator
+    AtSign,
+    /// `|/`, a square root math operator in PostgreSQL
+    PGSquareRoot,
+    /// `||/` , a cube root math operator in PostgreSQL
+    PGCubeRoot,
 }
 
 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::AtSign => f.write_str("@"),
+            Token::ShiftLeft => f.write_str("<<"),
+            Token::ShiftRight => f.write_str(">>"),
+            Token::PGSquareRoot => f.write_str("|/"),
+            Token::PGCubeRoot => 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::PGSquareRoot),
+                        Some('|') => {
+                            chars.next(); // consume the second '|'
+                            match chars.peek() {
+                                Some('/') => self.consume_and_return(chars, Token::PGCubeRoot),
+                                _ => 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::AtSign),
                 other => self.consume_and_return(chars, Token::Char(other)),
             },
             None => Ok(None),

tests/sqlparser_common.rs

@@ -343,7 +343,7 @@ 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"))),
             })],
             over: None,
             distinct: true,

tests/sqlparser_postgres.rs

@@ -551,6 +551,65 @@ fn parse_prepare() {
     );
 }
 
+#[test]
+fn parse_pg_bitwise_binary_ops() {
+    let bitwise_ops = &[
+        ("#", BinaryOperator::PGBitwiseXor),
+        (">>", BinaryOperator::PGBitwiseShiftRight),
+        ("<<", BinaryOperator::PGBitwiseShiftLeft),
+    ];
+
+    for (str_op, op) in bitwise_ops {
+        let select = pg().verified_only_select(&format!("SELECT a {} b", &str_op));
+        assert_eq!(
+            SelectItem::UnnamedExpr(Expr::BinaryOp {
+                left: Box::new(Expr::Identifier(Ident::new("a"))),
+                op: op.clone(),
+                right: Box::new(Expr::Identifier(Ident::new("b"))),
+            }),
+            select.projection[0]
+        );
+    }
+}
+
+#[test]
+fn parse_pg_unary_ops() {
+    let pg_unary_ops = &[
+        ("~", UnaryOperator::PGBitwiseNot),
+        ("|/", UnaryOperator::PGSquareRoot),
+        ("||/", UnaryOperator::PGCubeRoot),
+        ("!!", UnaryOperator::PGPrefixFactorial),
+        ("@", UnaryOperator::PGAbs),
+    ];
+
+    for (str_op, op) in pg_unary_ops {
+        let select = pg().verified_only_select(&format!("SELECT {} a", &str_op));
+        assert_eq!(
+            SelectItem::UnnamedExpr(Expr::UnaryOp {
+                op: op.clone(),
+                expr: Box::new(Expr::Identifier(Ident::new("a"))),
+            }),
+            select.projection[0]
+        );
+    }
+}
+
+#[test]
+fn parse_pg_postfix_factorial() {
+    let postfix_factorial = &[("!", UnaryOperator::PGPostfixFactorial)];
+
+    for (str_op, op) in postfix_factorial {
+        let select = pg().verified_only_select(&format!("SELECT a{}", &str_op));
+        assert_eq!(
+            SelectItem::UnnamedExpr(Expr::UnaryOp {
+                op: op.clone(),
+                expr: Box::new(Expr::Identifier(Ident::new("a"))),
+            }),
+            select.projection[0]
+        );
+    }
+}
+
 fn pg() -> TestedDialects {
     TestedDialects {
         dialects: vec![Box::new(PostgreSqlDialect {})],