diff --git a/crates/pglt_statement_splitter/src/lib.rs b/crates/pglt_statement_splitter/src/lib.rs index 7d74db037..1e0bb6202 100644 --- a/crates/pglt_statement_splitter/src/lib.rs +++ b/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![ diff --git a/crates/pglt_statement_splitter/src/parser.rs b/crates/pglt_statement_splitter/src/parser.rs index f89f28dda..c3351b911 100644 --- a/crates/pglt_statement_splitter/src/parser.rs +++ b/crates/pglt_statement_splitter/src/parser.rs @@ -85,7 +85,11 @@ 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 within statement at {:?}", + self.tokens.get(self.current_stmt_start.unwrap()) + ); self.current_stmt_start = Some(self.next_pos); } @@ -142,6 +146,21 @@ impl Parser { } } + /// Look ahead to the next relevant 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; diff --git a/crates/pglt_statement_splitter/src/parser/common.rs b/crates/pglt_statement_splitter/src/parser/common.rs index f156be0c5..debf8e4c8 100644 --- a/crates/pglt_statement_splitter/src/parser/common.rs +++ b/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; }