Merge pull request #28 from appsignal/sanitize-then-else

Add THEN and ELSE as logical operators
appsignal · Jul 27, 2023 · 9573053 · 9573053
2 parents 3ea4f71 + b9d9424
commit 9573053
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 3 deletions.
diff --git a/src/lexer.rs b/src/lexer.rs
@@ -315,6 +315,12 @@ impl SqlLexer {
                         s if s.eq_ignore_ascii_case("not") => {
                             Token::Operator(Operator::Logical(LogicalOperator::Not))
                         }
+                        s if s.eq_ignore_ascii_case("then") => {
+                            Token::Operator(Operator::Logical(LogicalOperator::Then))
+                        }
+                        s if s.eq_ignore_ascii_case("else") => {
+                            Token::Operator(Operator::Logical(LogicalOperator::Else))
+                        }
                         s if s.eq_ignore_ascii_case("like") => {
                             Token::Operator(Operator::Logical(LogicalOperator::Like))
                         }
@@ -626,7 +632,7 @@ mod tests {
 
     #[test]
     fn test_logical_operators_uppercase() {
-        let sql = "IN NOT LIKE ILIKE RLIKE GLOB MATCH REGEXP".to_string();
+        let sql = "IN NOT LIKE ILIKE RLIKE GLOB MATCH REGEXP THEN ELSE".to_string();
         let lexer = SqlLexer::new(sql);
 
         let expected = vec![
@@ -645,14 +651,18 @@ mod tests {
             Token::Operator(Operator::Logical(LogicalOperator::Match)),
             Token::Space,
             Token::Operator(Operator::Logical(LogicalOperator::Regexp)),
+            Token::Space,
+            Token::Operator(Operator::Logical(LogicalOperator::Then)),
+            Token::Space,
+            Token::Operator(Operator::Logical(LogicalOperator::Else)),
         ];
 
         assert_eq!(lexer.lex().tokens, expected);
     }
 
     #[test]
     fn test_logical_operators_lowercase() {
-        let sql = "in not like rlike glob match regexp".to_string();
+        let sql = "in not like rlike glob match regexp then else".to_string();
         let lexer = SqlLexer::new(sql);
 
         let expected = vec![
@@ -669,14 +679,18 @@ mod tests {
             Token::Operator(Operator::Logical(LogicalOperator::Match)),
             Token::Space,
             Token::Operator(Operator::Logical(LogicalOperator::Regexp)),
+            Token::Space,
+            Token::Operator(Operator::Logical(LogicalOperator::Then)),
+            Token::Space,
+            Token::Operator(Operator::Logical(LogicalOperator::Else)),
         ];
 
         assert_eq!(lexer.lex().tokens, expected);
     }
 
     #[test]
     fn test_logical_operators_mixedcase() {
-        let sql = "In Not Like Rlike Glob Match Regexp".to_string();
+        let sql = "In Not Like Rlike Glob Match Regexp tHen ElsE".to_string();
         let lexer = SqlLexer::new(sql);
 
         let expected = vec![
@@ -693,6 +707,10 @@ mod tests {
             Token::Operator(Operator::Logical(LogicalOperator::Match)),
             Token::Space,
             Token::Operator(Operator::Logical(LogicalOperator::Regexp)),
+            Token::Space,
+            Token::Operator(Operator::Logical(LogicalOperator::Then)),
+            Token::Space,
+            Token::Operator(Operator::Logical(LogicalOperator::Else)),
         ];
 
         assert_eq!(lexer.lex().tokens, expected);

diff --git a/src/lib.rs b/src/lib.rs
@@ -57,6 +57,8 @@ pub enum LogicalOperator {
     Glob,   // GLOB
     Match,  // MATCH
     Regexp, // REGEXP
+    Then,   // THEN
+    Else,   // ELSE
 }
 
 #[derive(Debug, PartialEq)]

diff --git a/src/sanitizer.rs b/src/sanitizer.rs
@@ -381,6 +381,17 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_case_then_else_subquery() {
+        assert_eq!(
+            sanitize_string(
+                "CASE WHEN NOT EXISTS (SELECT * FROM `table` WHERE `id` = 1) THEN 1 ELSE '0' END;"
+                    .to_string()
+            ),
+            "CASE WHEN NOT EXISTS (SELECT * FROM `table` WHERE `id` = ?) THEN ? ELSE ? END;"
+        );
+    }
+
     #[test]
     fn test_select_array() {
         assert_eq!(

diff --git a/src/writer.rs b/src/writer.rs
@@ -33,6 +33,8 @@ impl SqlWriter {
                 Token::Operator(Operator::Logical(LogicalOperator::Rlike)) => out.push_str("RLIKE"),
                 Token::Operator(Operator::Logical(LogicalOperator::Glob)) => out.push_str("GLOB"),
                 Token::Operator(Operator::Logical(LogicalOperator::Match)) => out.push_str("MATCH"),
+                Token::Operator(Operator::Logical(LogicalOperator::Then)) => out.push_str("THEN"),
+                Token::Operator(Operator::Logical(LogicalOperator::Else)) => out.push_str("ELSE"),
                 Token::Operator(Operator::Logical(LogicalOperator::Regexp)) => {
                     out.push_str("REGEXP")
                 }