fix: with ordinality

Author: psteinroe

crates/pglt_statement_splitter/src/lib.rs

@@ -233,6 +233,11 @@ values ('insert', new.id, now());",
             ]);
     }
 
+    #[test]
+    fn with_ordinality() {
+        Tester::from("insert into table (col) select 1 from other t cross join lateral jsonb_array_elements(t.buttons) with ordinality as a(b, nr) where t.buttons is not null;").expect_statements(vec!["insert into table (col) select 1 from other t cross join lateral jsonb_array_elements(t.buttons) with ordinality as a(b, nr) where t.buttons is not null;"]);
+    }
+
     #[test]
     fn unknown() {
         Tester::from("random stuff\n\nmore randomness\n\nselect 3").expect_statements(vec![

crates/pglt_statement_splitter/src/parser.rs

@@ -85,7 +85,12 @@ impl Parser {
 
     /// Start statement
     pub fn start_stmt(&mut self) {
-        assert!(self.current_stmt_start.is_none());
+        assert!(
+            self.current_stmt_start.is_none(),
+            "cannot start statement without statement {:?} at {:?}",
+            self.current_stmt_start,
+            self.tokens.get(self.current_stmt_start.unwrap()).unwrap()
+        );
         self.current_stmt_start = Some(self.next_pos);
     }
 
@@ -142,6 +147,21 @@ impl Parser {
         }
     }
 
+    /// Look ahead to the next relevant token *after the next token*
+    fn look_ahead(&self) -> Option<&Token> {
+        // we need to look ahead to the next relevant token
+        let mut look_ahead_pos = self.next_pos + 1;
+        loop {
+            let token = self.tokens.get(look_ahead_pos)?;
+
+            if !is_irrelevant_token(token) {
+                return Some(token);
+            }
+
+            look_ahead_pos += 1;
+        }
+    }
+
     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;

crates/pglt_statement_splitter/src/parser/common.rs

@@ -169,6 +169,19 @@ pub(crate) fn unknown(p: &mut Parser, exclude: &[SyntaxKind]) {
                     }
                     p.advance();
                 }
+                Some(SyntaxKind::With) => {
+                    let next = p.look_ahead().map(|t| t.kind);
+                    if [
+                        // WITH ORDINALITY should not start a new statement
+                        SyntaxKind::Ordinality,
+                    ]
+                    .iter()
+                    .all(|x| Some(x) != next.as_ref())
+                    {
+                        break;
+                    }
+                    p.advance();
+                }
                 Some(_) => {
                     break;
                 }