mirrors/datafusion-sqlparse
1
0
Fork 0
mirror of https://github.com/apache/datafusion-sqlparser-rs.git synced 2025-08-23 23:44:07 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

Add parsing for PostgreSQL math operators (#267)

Browse source
This commit is contained in:
Alex Dukhno 2020-09-30 05:29:31 +03:00 • committed by GitHub
parent 2f71324c33
commit 926b03a31d
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 171 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

8
src/ast/mod.rs
View file

@ -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),

24
src/ast/operator.rs
View file

@ -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 => ">>",
})
}
}

38
src/parser.rs
View file

@ -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),
}
}

46
src/tokenizer.rs
View file

@ -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,7 +457,8 @@ 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)),
}
}
'<' => {
@ -431,6 +466,7 @@ impl<'a> Tokenizer<'a> {
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)),
}
}
@ -438,6 +474,7 @@ impl<'a> Tokenizer<'a> {
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),

2
tests/sqlparser_common.rs
View file

@ -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,

59
tests/sqlparser_postgres.rs
View file

@ -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 {})],