fix: handle insert update delete and select within unknown

Author: psteinroe

crates/pg_statement_splitter/src/lib.rs

@@ -72,6 +72,34 @@ mod tests {
         ]);
     }
 
+    #[test]
+    fn insert_with_select() {
+        Tester::from("\ninsert into tbl (id) select 1\n\nselect 3")
+            .expect_statements(vec!["insert into tbl (id) select 1", "select 3"]);
+    }
+
+    #[test]
+    fn case() {
+        Tester::from("select case when select 2 then 1 else 0 end")
+            .expect_statements(vec!["select case when select 2 then 1 else 0 end"]);
+    }
+
+    #[test]
+    fn create_rule() {
+        Tester::from(
+            "create rule log_employee_insert as
+on insert to employees
+do also insert into employee_log (action, employee_id, log_time)
+values ('insert', new.id, now());",
+        )
+        .expect_statements(vec![
+            "create rule log_employee_insert as
+on insert to employees
+do also insert into employee_log (action, employee_id, log_time)
+values ('insert', new.id, now());",
+        ]);
+    }
+
     #[test]
     fn insert_into() {
         Tester::from("randomness\ninsert into tbl (id) values (1)\nselect 3").expect_statements(

crates/pg_statement_splitter/src/parser.rs

@@ -1,5 +1,6 @@
 mod common;
 mod data;
+mod ddl;
 mod dml;
 
 pub use common::source;
@@ -143,6 +144,24 @@ impl Parser {
         }
     }
 
+    fn look_back(&self) -> Option<&Token> {
+        // we need to look back to the last relevant token
+        let mut look_back_pos = self.next_pos - 1;
+        loop {
+            let token = self.tokens.get(look_back_pos);
+
+            if look_back_pos == 0 || token.is_none() {
+                return None;
+            }
+
+            if !is_irrelevant_token(token.unwrap()) {
+                return token;
+            }
+
+            look_back_pos -= 1;
+        }
+    }
+
     /// checks if the current token is of `kind` and advances if true
     /// returns true if the current token is of `kind`
     pub fn eat(&mut self, kind: SyntaxKind) -> bool {
@@ -164,7 +183,7 @@ impl Parser {
 
     /// collects an SyntaxError with an `error` message at the current position
     fn error_at(&mut self, error: String) {
-        todo!();
+        todo!("{error}");
     }
 }
 

crates/pg_statement_splitter/src/parser/common.rs

@@ -2,6 +2,7 @@ use pg_lexer::{SyntaxKind, Token, TokenType};
 
 use super::{
     data::at_statement_start,
+    ddl::create,
     dml::{cte, delete, insert, select, update},
     Parser,
 };
@@ -48,6 +49,9 @@ pub(crate) fn statement(p: &mut Parser) {
         SyntaxKind::DeleteP => {
             delete(p);
         }
+        SyntaxKind::Create => {
+            create(p);
+        }
         _ => {
             unknown(p);
         }
@@ -71,6 +75,22 @@ pub(crate) fn parenthesis(p: &mut Parser) {
     }
 }
 
+pub(crate) fn case(p: &mut Parser) {
+    p.expect(SyntaxKind::Case);
+
+    loop {
+        match p.peek().kind {
+            SyntaxKind::EndP => {
+                p.advance();
+                break;
+            }
+            _ => {
+                p.advance();
+            }
+        }
+    }
+}
+
 pub(crate) fn unknown(p: &mut Parser) {
     loop {
         match p.peek() {
@@ -87,19 +107,52 @@ pub(crate) fn unknown(p: &mut Parser) {
             } => {
                 break;
             }
+            Token {
+                kind: SyntaxKind::Case,
+                ..
+            } => {
+                case(p);
+            }
             Token {
                 kind: SyntaxKind::Ascii40,
                 ..
             } => {
                 parenthesis(p);
             }
-            t => {
-                if at_statement_start(t.kind) {
+            t => match at_statement_start(t.kind) {
+                Some(SyntaxKind::Select) => {
+                    // we need to check for `as` here to not break on `select as`
+                    if p.look_back().map(|t| t.kind) != Some(SyntaxKind::As) {
+                        break;
+                    }
+                    p.advance();
+                }
+                Some(SyntaxKind::Insert) | Some(SyntaxKind::Update) | Some(SyntaxKind::DeleteP) => {
+                    let prev = p.look_back().map(|t| t.kind);
+                    if [
+                        // for create trigger
+                        SyntaxKind::After,
+                        // for create rule
+                        SyntaxKind::On,
+                        // for create rule
+                        SyntaxKind::Also,
+                        // for create rule
+                        SyntaxKind::Instead,
+                    ]
+                    .iter()
+                    .all(|x| Some(x) != prev.as_ref())
+                    {
+                        break;
+                    }
+                    p.advance();
+                }
+                Some(_) => {
                     break;
                 }
-
-                p.advance();
-            }
+                None => {
+                    p.advance();
+                }
+            },
         }
     }
 }

crates/pg_statement_splitter/src/parser/data.rs

@@ -1,5 +1,7 @@
 use pg_lexer::SyntaxKind;
 
+// All tokens listed here must be explicitly handled in the `unknown` function to ensure that we do
+// not break in the middle of another statement that contains a statement start token.
 static STATEMENT_START_TOKENS: &[SyntaxKind] = &[
     SyntaxKind::With,
     SyntaxKind::Select,
@@ -9,6 +11,6 @@ static STATEMENT_START_TOKENS: &[SyntaxKind] = &[
     SyntaxKind::Create,
 ];
 
-pub(crate) fn at_statement_start(kind: SyntaxKind) -> bool {
-    STATEMENT_START_TOKENS.contains(&kind)
+pub(crate) fn at_statement_start(kind: SyntaxKind) -> Option<SyntaxKind> {
+    STATEMENT_START_TOKENS.iter().find(|&x| x == &kind).cloned()
 }

crates/pg_statement_splitter/src/parser/ddl.rs

@@ -0,0 +1,9 @@
+use pg_lexer::SyntaxKind;
+
+use super::{common::unknown, Parser};
+
+pub(crate) fn create(p: &mut Parser) {
+    p.expect(SyntaxKind::Create);
+
+    unknown(p);
+}