From 651e6c59f45aaa7d6e96afeff4fd811f1fa527d7 Mon Sep 17 00:00:00 2001
From: psteinroe <philipp@steinroetter.com>
Date: Mon, 14 Apr 2025 20:34:14 +0200
Subject: [PATCH 1/2] fix: plpglsql cmds

---
 crates/pgt_statement_splitter/src/lib.rs      | 18 ++++++
 crates/pgt_statement_splitter/src/parser.rs   |  3 +
 .../src/parser/common.rs                      | 57 ++++++++++++++++++-
 3 files changed, 77 insertions(+), 1 deletion(-)

diff --git a/crates/pgt_statement_splitter/src/lib.rs b/crates/pgt_statement_splitter/src/lib.rs
index 06440da1..3fa67213 100644
--- a/crates/pgt_statement_splitter/src/lib.rs
+++ b/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")
diff --git a/crates/pgt_statement_splitter/src/parser.rs b/crates/pgt_statement_splitter/src/parser.rs
index c94fe245..183abd97 100644
--- a/crates/pgt_statement_splitter/src/parser.rs
+++ b/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
diff --git a/crates/pgt_statement_splitter/src/parser/common.rs b/crates/pgt_statement_splitter/src/parser/common.rs
index 1a355f08..09071fd7 100644
--- a/crates/pgt_statement_splitter/src/parser/common.rs
+++ b/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,13 @@ pub fn source(p: &mut Parser) {
             } => {
                 p.advance();
             }
+            Token {
+                kind: SyntaxKind::Ascii92,
+                ..
+            } => {
+                println!("current {:?}", p.current());
+                plpgsql_command(p);
+            }
             _ => {
                 statement(p);
             }
@@ -87,6 +94,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 +150,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,
                 ..

From 1289763a4e486262d15b17f7a9ce2a9ebf66f870 Mon Sep 17 00:00:00 2001
From: psteinroe <philipp@steinroetter.com>
Date: Mon, 14 Apr 2025 20:36:00 +0200
Subject: [PATCH 2/2] cleanup

---
 crates/pgt_statement_splitter/src/parser/common.rs | 1 -
 1 file changed, 1 deletion(-)

diff --git a/crates/pgt_statement_splitter/src/parser/common.rs b/crates/pgt_statement_splitter/src/parser/common.rs
index 09071fd7..a353791b 100644
--- a/crates/pgt_statement_splitter/src/parser/common.rs
+++ b/crates/pgt_statement_splitter/src/parser/common.rs
@@ -28,7 +28,6 @@ pub fn source(p: &mut Parser) {
                 kind: SyntaxKind::Ascii92,
                 ..
             } => {
-                println!("current {:?}", p.current());
                 plpgsql_command(p);
             }
             _ => {