feat: add remaining dml statements

Author: psteinroe

crates/pg_statement_splitter/src/lib.rs

@@ -1,6 +1,6 @@
-///! Postgres Statement Splitter
-///!
-///! This crate provides a function to split a SQL source string into individual statements.
+//! Postgres Statement Splitter
+//!
+//! This crate provides a function to split a SQL source string into individual statements.
 mod parser;
 mod syntax_error;
 
@@ -36,7 +36,13 @@ mod tests {
 
     impl Tester {
         fn expect_statements(&self, expected: Vec<&str>) {
-            assert_eq!(self.parse.ranges.len(), expected.len());
+            assert_eq!(
+                self.parse.ranges.len(),
+                expected.len(),
+                "Expected {} statements, got {}",
+                expected.len(),
+                self.parse.ranges.len()
+            );
 
             for (range, expected) in self.parse.ranges.iter().zip(expected.iter()) {
                 assert_eq!(*expected, self.input[*range].to_string());
@@ -65,4 +71,46 @@ mod tests {
             "select 3",
         ]);
     }
+
+    #[test]
+    fn insert_into() {
+        Tester::from("randomness\ninsert into tbl (id) values (1)\nselect 3").expect_statements(
+            vec!["randomness", "insert into tbl (id) values (1)", "select 3"],
+        );
+    }
+
+    #[test]
+    fn update() {
+        Tester::from("more randomness\nupdate tbl set col = '1'\n\nselect 3").expect_statements(
+            vec!["more randomness", "update tbl set col = '1'", "select 3"],
+        );
+    }
+
+    #[test]
+    fn delete_from() {
+        Tester::from("more randomness\ndelete from test where id = 1\n\nselect 3")
+            .expect_statements(vec![
+                "more randomness",
+                "delete from test where id = 1",
+                "select 3",
+            ]);
+    }
+
+    #[test]
+    fn unknown() {
+        Tester::from("random stuff\n\nmore randomness\n\nselect 3").expect_statements(vec![
+            "random stuff",
+            "more randomness",
+            "select 3",
+        ]);
+    }
+
+    #[test]
+    fn unknown_2() {
+        Tester::from("random stuff\nselect 1\n\nselect 3").expect_statements(vec![
+            "random stuff",
+            "select 1",
+            "select 3",
+        ]);
+    }
 }

crates/pg_statement_splitter/src/parser.rs

@@ -40,14 +40,7 @@ impl Parser {
         // we dont care about whitespace tokens, except for double newlines
         // to make everything simpler, we just filter them out
         // the token holds the text range, so we dont need to worry about that
-        let tokens = lex(sql)
-            .iter()
-            .filter(|t| {
-                return !WHITESPACE_TOKENS.contains(&t.kind)
-                    || (t.kind == SyntaxKind::Newline && t.text.chars().count() > 1);
-            })
-            .cloned()
-            .collect::<Vec<_>>();
+        let tokens = lex(sql);
 
         let eof_token = Token::eof(usize::from(
             tokens
@@ -104,9 +97,25 @@ impl Parser {
     pub fn close_stmt(&mut self) {
         assert!(self.next_pos > 0);
 
+        // go back the positions until we find the first relevant token
+        let mut end_token_pos = self.next_pos - 1;
+        loop {
+            let token = self.tokens.get(end_token_pos);
+
+            if end_token_pos == 0 || token.is_none() {
+                break;
+            }
+
+            if !is_irrelevant_token(token.unwrap()) {
+                break;
+            }
+
+            end_token_pos -= 1;
+        }
+
         self.ranges.push((
             self.current_stmt_start.expect("Expected active statement"),
-            self.next_pos - 1,
+            end_token_pos,
         ));
 
         self.current_stmt_start = None;

crates/pg_statement_splitter/src/parser/common.rs

@@ -2,7 +2,7 @@ use pg_lexer::{SyntaxKind, Token, TokenType};
 
 use super::{
     data::at_statement_start,
-    dml::{cte, select},
+    dml::{cte, delete, insert, select, update},
     Parser,
 };
 
@@ -16,7 +16,9 @@ pub fn source(p: &mut Parser) {
                 break;
             }
             Token {
-                token_type: TokenType::Whitespace | TokenType::NoKeyword,
+                // we might want to ignore TokenType::NoKeyword here too
+                // but this will lead to invalid statements to not being picked up
+                token_type: TokenType::Whitespace,
                 ..
             } => {
                 p.advance();
@@ -38,20 +40,16 @@ pub(crate) fn statement(p: &mut Parser) {
             select(p);
         }
         SyntaxKind::Insert => {
-            todo!();
-            // insert(p);
+            insert(p);
         }
         SyntaxKind::Update => {
-            todo!();
-            // update(p);
+            update(p);
         }
         SyntaxKind::DeleteP => {
-            todo!();
-            // delete(p);
+            delete(p);
         }
-        t => {
-            panic!("stmt: Unknown start token {:?}", t);
-            // unknown(p);
+        _ => {
+            unknown(p);
         }
     }
     p.close_stmt();

crates/pg_statement_splitter/src/parser/dml.rs

@@ -26,3 +26,23 @@ pub(crate) fn select(p: &mut Parser) {
 
     unknown(p);
 }
+
+pub(crate) fn insert(p: &mut Parser) {
+    p.expect(SyntaxKind::Insert);
+    p.expect(SyntaxKind::Into);
+
+    unknown(p);
+}
+
+pub(crate) fn update(p: &mut Parser) {
+    p.expect(SyntaxKind::Update);
+
+    unknown(p);
+}
+
+pub(crate) fn delete(p: &mut Parser) {
+    p.expect(SyntaxKind::DeleteP);
+    p.expect(SyntaxKind::From);
+
+    unknown(p);
+}