fix: plpglsql cmds (#336)

Author: psteinroe

crates/pgt_statement_splitter/src/lib.rs

@@ -136,6 +136,24 @@ mod tests {
             )]);
     }
 
+    #[test]
+    fn command_between_not_starting() {
+        Tester::from("select 1\n      \\com test\nselect 2")
+            .expect_statements(vec!["select 1", "select 2"]);
+    }
+
+    #[test]
+    fn command_between() {
+        Tester::from("select 1\n\\com test\nselect 2")
+            .expect_statements(vec!["select 1", "select 2"]);
+    }
+
+    #[test]
+    fn command_standalone() {
+        Tester::from("select 1\n\n\\com test\n\nselect 2")
+            .expect_statements(vec!["select 1", "select 2"]);
+    }
+
     #[test]
     fn insert_with_select() {
         Tester::from("\ninsert into tbl (id) select 1\n\nselect 3")

crates/pgt_statement_splitter/src/parser.rs

@@ -113,6 +113,9 @@ impl Parser {
         }
     }
 
+    /// Advances the parser to the next relevant token and returns it.
+    ///
+    /// NOTE: This will skip irrelevant tokens.
     fn advance(&mut self) -> &Token {
         // can't reuse any `find_next_relevant` logic because of Mr. Borrow Checker
         let (pos, token) = self

crates/pgt_statement_splitter/src/parser/common.rs

@@ -1,4 +1,4 @@
-use pgt_lexer::{SyntaxKind, Token, TokenType};
+use pgt_lexer::{SyntaxKind, Token, TokenType, WHITESPACE_TOKENS};
 
 use super::{
     Parser,
@@ -24,6 +24,12 @@ pub fn source(p: &mut Parser) {
             } => {
                 p.advance();
             }
+            Token {
+                kind: SyntaxKind::Ascii92,
+                ..
+            } => {
+                plpgsql_command(p);
+            }
             _ => {
                 statement(p);
             }
@@ -87,6 +93,24 @@ pub(crate) fn parenthesis(p: &mut Parser) {
     }
 }
 
+pub(crate) fn plpgsql_command(p: &mut Parser) {
+    p.expect(SyntaxKind::Ascii92);
+
+    loop {
+        match p.current().kind {
+            SyntaxKind::Newline => {
+                p.advance();
+                break;
+            }
+            _ => {
+                // advance the parser to the next token without ignoring irrelevant tokens
+                // we would skip a newline with `advance()`
+                p.current_pos += 1;
+            }
+        }
+    }
+}
+
 pub(crate) fn case(p: &mut Parser) {
     p.expect(SyntaxKind::Case);
 
@@ -125,6 +149,36 @@ pub(crate) fn unknown(p: &mut Parser, exclude: &[SyntaxKind]) {
             } => {
                 case(p);
             }
+            Token {
+                kind: SyntaxKind::Ascii92,
+                ..
+            } => {
+                // pgsql commands e.g.
+                //
+                // ```
+                // \if test
+                // ```
+                //
+                // we wait for "\" and check if the previous token is a newline
+
+                // newline is a whitespace, but we do not want to ignore it here
+                let irrelevant = WHITESPACE_TOKENS
+                    .iter()
+                    .filter(|t| **t != SyntaxKind::Newline)
+                    .collect::<Vec<_>>();
+
+                // go back from the current position without ignoring irrelevant tokens
+                if p.tokens
+                    .iter()
+                    .take(p.current_pos)
+                    .rev()
+                    .find(|t| !irrelevant.contains(&&t.kind))
+                    .is_some_and(|t| t.kind == SyntaxKind::Newline)
+                {
+                    break;
+                }
+                p.advance();
+            }
             Token {
                 kind: SyntaxKind::Ascii40,
                 ..