Skip to content

Commit 926b03a

Browse files
alex-dukhnoalex-dukhno
authored
Add parsing for PostgreSQL math operators (apache#267)
1 parent 2f71324 commit 926b03a

6 files changed

Lines changed: 171 additions & 6 deletions

File tree

‎src/ast/mod.rs‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -282,7 +282,13 @@ impl fmt::Display for Expr {
282282
high
283283
),
284284
Expr::BinaryOp { left, op, right } => write!(f, "{} {} {}", left, op, right),
285-
Expr::UnaryOp { op, expr } => write!(f, "{} {}", op, expr),
285+
Expr::UnaryOp { op, expr } => {
286+
if op == &UnaryOperator::PGPostfixFactorial {
287+
write!(f, "{}{}", expr, op)
288+
} else {
289+
write!(f, "{} {}", op, expr)
290+
}
291+
}
286292
Expr::Cast { expr, data_type } => write!(f, "CAST({} AS {})", expr, data_type),
287293
Expr::Extract { field, expr } => write!(f, "EXTRACT({} FROM {})", field, expr),
288294
Expr::Collate { expr, collation } => write!(f, "{} COLLATE {}", expr, collation),

‎src/ast/operator.rs‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,18 @@ pub enum UnaryOperator {
2121
Plus,
2222
Minus,
2323
Not,
24+
/// Bitwise Not, e.g. `~9` (PostgreSQL-specific)
25+
PGBitwiseNot,
26+
/// Square root, e.g. `|/9` (PostgreSQL-specific)
27+
PGSquareRoot,
28+
/// Cube root, e.g. `||/27` (PostgreSQL-specific)
29+
PGCubeRoot,
30+
/// Factorial, e.g. `9!` (PostgreSQL-specific)
31+
PGPostfixFactorial,
32+
/// Factorial, e.g. `!!9` (PostgreSQL-specific)
33+
PGPrefixFactorial,
34+
/// Absolute value, e.g. `@ -9` (PostgreSQL-specific)
35+
PGAbs,
2436
}
2537

2638
impl fmt::Display for UnaryOperator {
@@ -29,6 +41,12 @@ impl fmt::Display for UnaryOperator {
2941
UnaryOperator::Plus => "+",
3042
UnaryOperator::Minus => "-",
3143
UnaryOperator::Not => "NOT",
44+
UnaryOperator::PGBitwiseNot => "~",
45+
UnaryOperator::PGSquareRoot => "|/",
46+
UnaryOperator::PGCubeRoot => "||/",
47+
UnaryOperator::PGPostfixFactorial => "!",
48+
UnaryOperator::PGPrefixFactorial => "!!",
49+
UnaryOperator::PGAbs => "@",
3250
})
3351
}
3452
}
@@ -56,6 +74,9 @@ pub enum BinaryOperator {
5674
BitwiseOr,
5775
BitwiseAnd,
5876
BitwiseXor,
77+
PGBitwiseXor,
78+
PGBitwiseShiftLeft,
79+
PGBitwiseShiftRight,
5980
}
6081

6182
impl fmt::Display for BinaryOperator {
@@ -80,6 +101,9 @@ impl fmt::Display for BinaryOperator {
80101
BinaryOperator::BitwiseOr => "|",
81102
BinaryOperator::BitwiseAnd => "&",
82103
BinaryOperator::BitwiseXor => "^",
104+
BinaryOperator::PGBitwiseXor => "#",
105+
BinaryOperator::PGBitwiseShiftLeft => "<<",
106+
BinaryOperator::PGBitwiseShiftRight => ">>",
83107
})
84108
}
85109
}

‎src/parser.rs‎

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -294,6 +294,26 @@ impl<'a> Parser<'a> {
294294
expr: Box::new(self.parse_subexpr(Self::PLUS_MINUS_PREC)?),
295295
})
296296
}
297+
tok @ Token::DoubleExclamationMark
298+
| tok @ Token::PGSquareRoot
299+
| tok @ Token::PGCubeRoot
300+
| tok @ Token::AtSign
301+
| tok @ Token::Tilde
302+
if dialect_of!(self is PostgreSqlDialect) =>
303+
{
304+
let op = match tok {
305+
Token::DoubleExclamationMark => UnaryOperator::PGPrefixFactorial,
306+
Token::PGSquareRoot => UnaryOperator::PGSquareRoot,
307+
Token::PGCubeRoot => UnaryOperator::PGCubeRoot,
308+
Token::AtSign => UnaryOperator::PGAbs,
309+
Token::Tilde => UnaryOperator::PGBitwiseNot,
310+
_ => unreachable!(),
311+
};
312+
Ok(Expr::UnaryOp {
313+
op,
314+
expr: Box::new(self.parse_subexpr(Self::PLUS_MINUS_PREC)?),
315+
})
316+
}
297317
Token::Number(_)
298318
| Token::SingleQuotedString(_)
299319
| Token::NationalStringLiteral(_)
@@ -658,6 +678,15 @@ impl<'a> Parser<'a> {
658678
Token::Caret => Some(BinaryOperator::BitwiseXor),
659679
Token::Ampersand => Some(BinaryOperator::BitwiseAnd),
660680
Token::Div => Some(BinaryOperator::Divide),
681+
Token::ShiftLeft if dialect_of!(self is PostgreSqlDialect) => {
682+
Some(BinaryOperator::PGBitwiseShiftLeft)
683+
}
684+
Token::ShiftRight if dialect_of!(self is PostgreSqlDialect) => {
685+
Some(BinaryOperator::PGBitwiseShiftRight)
686+
}
687+
Token::Sharp if dialect_of!(self is PostgreSqlDialect) => {
688+
Some(BinaryOperator::PGBitwiseXor)
689+
}
661690
Token::Word(w) => match w.keyword {
662691
Keyword::AND => Some(BinaryOperator::And),
663692
Keyword::OR => Some(BinaryOperator::Or),
@@ -707,6 +736,12 @@ impl<'a> Parser<'a> {
707736
}
708737
} else if Token::DoubleColon == tok {
709738
self.parse_pg_cast(expr)
739+
} else if Token::ExclamationMark == tok {
740+
// PostgreSQL factorial operation
741+
Ok(Expr::UnaryOp {
742+
op: UnaryOperator::PGPostfixFactorial,
743+
expr: Box::new(expr),
744+
})
710745
} else {
711746
// Can only happen if `get_next_precedence` got out of sync with this function
712747
panic!("No infix parser for token {:?}", tok)
@@ -785,11 +820,12 @@ impl<'a> Parser<'a> {
785820
Token::Word(w) if w.keyword == Keyword::LIKE => Ok(Self::BETWEEN_PREC),
786821
Token::Eq | Token::Lt | Token::LtEq | Token::Neq | Token::Gt | Token::GtEq => Ok(20),
787822
Token::Pipe => Ok(21),
788-
Token::Caret => Ok(22),
823+
Token::Caret | Token::Sharp | Token::ShiftRight | Token::ShiftLeft => Ok(22),
789824
Token::Ampersand => Ok(23),
790825
Token::Plus | Token::Minus => Ok(Self::PLUS_MINUS_PREC),
791826
Token::Mult | Token::Div | Token::Mod | Token::StringConcat => Ok(40),
792827
Token::DoubleColon => Ok(50),
828+
Token::ExclamationMark => Ok(50),
793829
_ => Ok(0),
794830
}
795831
}

‎src/tokenizer.rs‎

Lines changed: 43 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ pub enum Token {
5454
Neq,
5555
/// Less Than operator `<`
5656
Lt,
57-
/// Greater han operator `>`
57+
/// Greater Than operator `>`
5858
Gt,
5959
/// Less Than Or Equals operator `<=`
6060
LtEq,
@@ -102,6 +102,24 @@ pub enum Token {
102102
RBrace,
103103
/// Right Arrow `=>`
104104
RArrow,
105+
/// Sharp `#` used for PostgreSQL Bitwise XOR operator
106+
Sharp,
107+
/// Tilde `~` used for PostgreSQL Bitwise NOT operator
108+
Tilde,
109+
/// `<<`, a bitwise shift left operator in PostgreSQL
110+
ShiftLeft,
111+
/// `>>`, a bitwise shift right operator in PostgreSQL
112+
ShiftRight,
113+
/// Exclamation Mark `!` used for PostgreSQL factorial operator
114+
ExclamationMark,
115+
/// Double Exclamation Mark `!!` used for PostgreSQL prefix factorial operator
116+
DoubleExclamationMark,
117+
/// AtSign `@` used for PostgreSQL abs operator
118+
AtSign,
119+
/// `|/`, a square root math operator in PostgreSQL
120+
PGSquareRoot,
121+
/// `||/` , a cube root math operator in PostgreSQL
122+
PGCubeRoot,
105123
}
106124

107125
impl fmt::Display for Token {
@@ -143,6 +161,15 @@ impl fmt::Display for Token {
143161
Token::LBrace => f.write_str("{"),
144162
Token::RBrace => f.write_str("}"),
145163
Token::RArrow => f.write_str("=>"),
164+
Token::Sharp => f.write_str("#"),
165+
Token::ExclamationMark => f.write_str("!"),
166+
Token::DoubleExclamationMark => f.write_str("!!"),
167+
Token::Tilde => f.write_str("~"),
168+
Token::AtSign => f.write_str("@"),
169+
Token::ShiftLeft => f.write_str("<<"),
170+
Token::ShiftRight => f.write_str(">>"),
171+
Token::PGSquareRoot => f.write_str("|/"),
172+
Token::PGCubeRoot => f.write_str("||/"),
146173
}
147174
}
148175
}
@@ -406,7 +433,14 @@ impl<'a> Tokenizer<'a> {
406433
'|' => {
407434
chars.next(); // consume the '|'
408435
match chars.peek() {
409-
Some('|') => self.consume_and_return(chars, Token::StringConcat),
436+
Some('/') => self.consume_and_return(chars, Token::PGSquareRoot),
437+
Some('|') => {
438+
chars.next(); // consume the second '|'
439+
match chars.peek() {
440+
Some('/') => self.consume_and_return(chars, Token::PGCubeRoot),
441+
_ => Ok(Some(Token::StringConcat)),
442+
}
443+
}
410444
// Bitshift '|' operator
411445
_ => Ok(Some(Token::Pipe)),
412446
}
@@ -423,21 +457,24 @@ impl<'a> Tokenizer<'a> {
423457
chars.next(); // consume
424458
match chars.peek() {
425459
Some('=') => self.consume_and_return(chars, Token::Neq),
426-
_ => self.tokenizer_error("Expected to see '=' after '!' character"),
460+
Some('!') => self.consume_and_return(chars, Token::DoubleExclamationMark),
461+
_ => Ok(Some(Token::ExclamationMark)),
427462
}
428463
}
429464
'<' => {
430465
chars.next(); // consume
431466
match chars.peek() {
432467
Some('=') => self.consume_and_return(chars, Token::LtEq),
433468
Some('>') => self.consume_and_return(chars, Token::Neq),
469+
Some('<') => self.consume_and_return(chars, Token::ShiftLeft),
434470
_ => Ok(Some(Token::Lt)),
435471
}
436472
}
437473
'>' => {
438474
chars.next(); // consume
439475
match chars.peek() {
440476
Some('=') => self.consume_and_return(chars, Token::GtEq),
477+
Some('>') => self.consume_and_return(chars, Token::ShiftRight),
441478
_ => Ok(Some(Token::Gt)),
442479
}
443480
}
@@ -464,6 +501,9 @@ impl<'a> Tokenizer<'a> {
464501
comment,
465502
})))
466503
}
504+
'~' => self.consume_and_return(chars, Token::Tilde),
505+
'#' => self.consume_and_return(chars, Token::Sharp),
506+
'@' => self.consume_and_return(chars, Token::AtSign),
467507
other => self.consume_and_return(chars, Token::Char(other)),
468508
},
469509
None => Ok(None),

‎tests/sqlparser_common.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -343,7 +343,7 @@ fn parse_select_count_distinct() {
343343
name: ObjectName(vec![Ident::new("COUNT")]),
344344
args: vec![FunctionArg::Unnamed(Expr::UnaryOp {
345345
op: UnaryOperator::Plus,
346-
expr: Box::new(Expr::Identifier(Ident::new("x")))
346+
expr: Box::new(Expr::Identifier(Ident::new("x"))),
347347
})],
348348
over: None,
349349
distinct: true,

‎tests/sqlparser_postgres.rs‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -551,6 +551,65 @@ fn parse_prepare() {
551551
);
552552
}
553553

554+
#[test]
555+
fn parse_pg_bitwise_binary_ops() {
556+
let bitwise_ops = &[
557+
("#", BinaryOperator::PGBitwiseXor),
558+
(">>", BinaryOperator::PGBitwiseShiftRight),
559+
("<<", BinaryOperator::PGBitwiseShiftLeft),
560+
];
561+
562+
for (str_op, op) in bitwise_ops {
563+
let select = pg().verified_only_select(&format!("SELECT a {} b", &str_op));
564+
assert_eq!(
565+
SelectItem::UnnamedExpr(Expr::BinaryOp {
566+
left: Box::new(Expr::Identifier(Ident::new("a"))),
567+
op: op.clone(),
568+
right: Box::new(Expr::Identifier(Ident::new("b"))),
569+
}),
570+
select.projection[0]
571+
);
572+
}
573+
}
574+
575+
#[test]
576+
fn parse_pg_unary_ops() {
577+
let pg_unary_ops = &[
578+
("~", UnaryOperator::PGBitwiseNot),
579+
("|/", UnaryOperator::PGSquareRoot),
580+
("||/", UnaryOperator::PGCubeRoot),
581+
("!!", UnaryOperator::PGPrefixFactorial),
582+
("@", UnaryOperator::PGAbs),
583+
];
584+
585+
for (str_op, op) in pg_unary_ops {
586+
let select = pg().verified_only_select(&format!("SELECT {} a", &str_op));
587+
assert_eq!(
588+
SelectItem::UnnamedExpr(Expr::UnaryOp {
589+
op: op.clone(),
590+
expr: Box::new(Expr::Identifier(Ident::new("a"))),
591+
}),
592+
select.projection[0]
593+
);
594+
}
595+
}
596+
597+
#[test]
598+
fn parse_pg_postfix_factorial() {
599+
let postfix_factorial = &[("!", UnaryOperator::PGPostfixFactorial)];
600+
601+
for (str_op, op) in postfix_factorial {
602+
let select = pg().verified_only_select(&format!("SELECT a{}", &str_op));
603+
assert_eq!(
604+
SelectItem::UnnamedExpr(Expr::UnaryOp {
605+
op: op.clone(),
606+
expr: Box::new(Expr::Identifier(Ident::new("a"))),
607+
}),
608+
select.projection[0]
609+
);
610+
}
611+
}
612+
554613
fn pg() -> TestedDialects {
555614
TestedDialects {
556615
dialects: vec![Box::new(PostgreSqlDialect {})],

0 commit comments

Comments
 (0)