Fix quadratic INSERT INTO ... VALUES sanitisation performance

src/bin/bench.rs

@@ -0,0 +1,22 @@
+fn long_insert_into_values_query(rows: usize) -> String {
+    let mut query = r#"INSERT INTO "table_name" ("one","two","three") VALUES "#.to_owned();
+    for i in 0..rows {
+        query.push_str(&format!("({}, {}, {}),", i, i + 1, i + 2));
+    }
+    query.push_str("(0, 0, 0);");
+    query
+}
+
+extern crate sql_lexer;
+
+fn main() {
+    for rows in [200, 2000, 20000] {
+        println!("Sanitising long_insert_into_values_query with {rows} rows");
+        let query = long_insert_into_values_query(rows);
+        let start = std::time::Instant::now();
+        let output = sql_lexer::sanitize_string(query);
+        let elapsed = start.elapsed();
+        println!("Output: {}", output);
+        println!("Elapsed: {:?}", elapsed);
+    }
+}

src/lib.rs

@@ -137,8 +137,12 @@ pub enum Token {
     SquareBracketClose,
     Colon,
     Semicolon,
+    // Used by the sanitizer to replace values
     Placeholder,
+    // Used by the sanitizer to replace repeated value lists
     Ellipsis,
+    // Used by the sanitizer to remove tokens
+    None,
     Null,
     True,
     False,

src/sanitizer.rs

@@ -97,14 +97,17 @@ impl SqlSanitizer {
                         if keep {
                             kept += 1;
                         } else {
+                            let mut removed = 1;
                             self.remove(pos + kept);
                             while kept > 0 {
                                 kept -= 1;
+                                removed += 1;
                                 self.remove(pos + kept);
                             }
+                            pos += removed;
                         }
                     }
-                    self.sql.tokens.insert(start_pos, Token::Ellipsis);
+                    self.ellipsis(start_pos);
                 }
                 Token::ParentheseClose | Token::SquareBracketClose => state = State::Default,
                 Token::Dot if state == State::JoinOn => (),
@@ -146,23 +149,28 @@ impl SqlSanitizer {
                                 match self.sql.tokens[pos] {
                                     Token::ParentheseClose => break,
                                     Token::SquareBracketClose => break,
-                                    _ => self.remove(pos),
+                                    _ => {
+                                        self.remove(pos);
+                                        pos += 1;
+                                    }
                                 }
                             }
-                            self.sql.tokens.insert(start_pos, Token::Placeholder);
+                            self.placeholder(start_pos);
                         }
                         _ => (),
                     }
                 }
                 // Remove comments
                 Token::Comment(_) => {
-                    self.sql.tokens.remove(pos);
+                    self.remove(pos);
                     if self.sql.tokens.get(pos - 1) == Some(&Token::Space) {
-                        self.sql.tokens.remove(pos - 1);
+                        self.remove(pos - 1);
                     }
                 }
                 // Spaces don't influence the state by default
                 Token::Space => (),
+                // Non-tokens don't influence the state
+                Token::None => (),
                 // Keep state the same if we're in a insert values or keyword scope state
                 _ if state == State::InsertValues || state == State::KeywordScopeStarted => (),
                 // Reset state to default if there were no matches
@@ -176,12 +184,17 @@ impl SqlSanitizer {
     }
 
     fn remove(&mut self, position: usize) {
-        self.sql.tokens.remove(position);
+        self.sql.tokens[position] = Token::None;
     }
 
+    // Replaces the token at position `position` with a placeholder.
     fn placeholder(&mut self, position: usize) {
-        self.sql.tokens.remove(position);
-        self.sql.tokens.insert(position, Token::Placeholder);
+        self.sql.tokens[position] = Token::Placeholder;
+    }
+
+    // Replaces the token at position `position` with an ellipsis.
+    fn ellipsis(&mut self, position: usize) {
+        self.sql.tokens[position] = Token::Ellipsis;
     }
 }
 

src/writer.rs

@@ -171,6 +171,7 @@ impl SqlWriter {
                 Token::Semicolon => out.push(';'),
                 Token::Placeholder => out.push('?'),
                 Token::Ellipsis => out.push_str("..."),
+                Token::None => {}
                 Token::Null => out.push_str("NULL"),
                 Token::True => out.push_str("TRUE"),
                 Token::False => out.push_str("FALSE"),