From af18ae4f76f13f8e4e4cf7eeca3a576d5523a262 Mon Sep 17 00:00:00 2001 From: psteinroe Date: Wed, 3 Sep 2025 22:50:29 +0200 Subject: [PATCH 1/3] fix: statement splitter --- crates/pgt_statement_splitter/src/lib.rs | 91 ++++++++++++++++--- crates/pgt_statement_splitter/src/splitter.rs | 9 ++ .../src/splitter/common.rs | 2 + .../src/splitter/dml.rs | 1 + .../src/workspace/server.tests.rs | 40 ++++++++ 5 files changed, 132 insertions(+), 11 deletions(-) diff --git a/crates/pgt_statement_splitter/src/lib.rs b/crates/pgt_statement_splitter/src/lib.rs index 02ca1b304..ad3537c5a 100644 --- a/crates/pgt_statement_splitter/src/lib.rs +++ b/crates/pgt_statement_splitter/src/lib.rs @@ -62,6 +62,32 @@ mod tests { } impl Tester { + fn assert_single_statement(&self) -> &Self { + assert_eq!( + self.result.ranges.len(), + 1, + "Expected a single statement for input {}, got {}: {:?}", + self.input, + self.result.ranges.len(), + self.result + .ranges + .iter() + .map(|r| &self.input[*r]) + .collect::>() + ); + self + } + + fn assert_no_errors(&self) -> &Self { + assert!( + self.result.errors.is_empty(), + "Expected no errors, got {}: {:?}", + self.result.errors.len(), + self.result.errors + ); + self + } + fn expect_statements(&self, expected: Vec<&str>) -> &Self { assert_eq!( self.result.ranges.len(), @@ -114,6 +140,16 @@ mod tests { ); } + #[test] + fn test_for_no_key_update() { + Tester::from( + "SELECT 1 FROM assessments AS a WHERE a.id = $assessment_id FOR NO KEY UPDATE;", + ) + .expect_statements(vec![ + "SELECT 1 FROM assessments AS a WHERE a.id = $assessment_id FOR NO KEY UPDATE;", + ]); + } + #[test] fn test_crash_eof() { Tester::from("CREATE INDEX \"idx_analytics_read_ratio\" ON \"public\".\"message\" USING \"btree\" (\"inbox_id\", \"timestamp\") INCLUDE (\"status\") WHERE (\"is_inbound\" = false and channel_type not in ('postal'', 'sms'));") @@ -241,19 +277,52 @@ mod tests { } #[test] - fn trigger_instead_of() { + fn with_recursive() { Tester::from( - "CREATE OR REPLACE TRIGGER my_trigger - INSTEAD OF INSERT ON my_table - FOR EACH ROW - EXECUTE FUNCTION my_table_trigger_fn();", + " +WITH RECURSIVE + template_questions AS ( + -- non-recursive term that finds the ID of the template question (if any) for question_id + SELECT + tq.id, + tq.qid, + tq.course_id, + tq.template_directory + FROM + questions AS q + JOIN questions AS tq ON ( + tq.qid = q.template_directory + AND tq.course_id = q.course_id + ) + WHERE + q.id = $question_id + AND tq.deleted_at IS NULL + -- required UNION for a recursive WITH statement + UNION + -- recursive term that references template_questions again + SELECT + tq.id, + tq.qid, + tq.course_id, + tq.template_directory + FROM + template_questions AS q + JOIN questions AS tq ON ( + tq.qid = q.template_directory + AND tq.course_id = q.course_id + ) + WHERE + tq.deleted_at IS NULL + ) +SELECT + id +FROM + template_questions +LIMIT + 100;", ) - .expect_statements(vec![ - "CREATE OR REPLACE TRIGGER my_trigger - INSTEAD OF INSERT ON my_table - FOR EACH ROW - EXECUTE FUNCTION my_table_trigger_fn();", - ]); + .assert_single_statement() + .assert_no_errors(); } #[test] diff --git a/crates/pgt_statement_splitter/src/splitter.rs b/crates/pgt_statement_splitter/src/splitter.rs index cfb4716d5..9061999ea 100644 --- a/crates/pgt_statement_splitter/src/splitter.rs +++ b/crates/pgt_statement_splitter/src/splitter.rs @@ -102,6 +102,15 @@ impl<'a> Splitter<'a> { self.lexed.kind(self.current_pos) } + fn eat(&mut self, kind: SyntaxKind) -> bool { + if self.current() == kind { + self.advance(); + true + } else { + false + } + } + fn kind(&self, idx: usize) -> SyntaxKind { self.lexed.kind(idx) } diff --git a/crates/pgt_statement_splitter/src/splitter/common.rs b/crates/pgt_statement_splitter/src/splitter/common.rs index 786c24788..ff7cd372b 100644 --- a/crates/pgt_statement_splitter/src/splitter/common.rs +++ b/crates/pgt_statement_splitter/src/splitter/common.rs @@ -224,6 +224,8 @@ pub(crate) fn unknown(p: &mut Splitter, exclude: &[SyntaxKind]) { SyntaxKind::COMMA, // Do update in INSERT stmt SyntaxKind::DO_KW, + // FOR NO KEY UPDATE + SyntaxKind::KEY_KW, ] .iter() .all(|x| Some(x) != prev.as_ref()) diff --git a/crates/pgt_statement_splitter/src/splitter/dml.rs b/crates/pgt_statement_splitter/src/splitter/dml.rs index 9c8333016..acfbebfc9 100644 --- a/crates/pgt_statement_splitter/src/splitter/dml.rs +++ b/crates/pgt_statement_splitter/src/splitter/dml.rs @@ -7,6 +7,7 @@ use super::{ pub(crate) fn cte(p: &mut Splitter) { p.expect(SyntaxKind::WITH_KW); + p.eat(SyntaxKind::RECURSIVE_KW); loop { p.expect(SyntaxKind::IDENT); diff --git a/crates/pgt_workspace/src/workspace/server.tests.rs b/crates/pgt_workspace/src/workspace/server.tests.rs index 33520cbf6..a41977877 100644 --- a/crates/pgt_workspace/src/workspace/server.tests.rs +++ b/crates/pgt_workspace/src/workspace/server.tests.rs @@ -386,6 +386,46 @@ async fn test_positional_params(test_db: PgPool) { assert_eq!(diagnostics.len(), 0, "Expected no diagnostic"); } +#[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] +async fn test_named_params(_test_db: PgPool) { + let conf = PartialConfiguration::init(); + + let workspace = get_test_workspace(Some(conf)).expect("Unable to create test workspace"); + + let path = PgTPath::new("test.sql"); + + let content = r#" +SELECT + 1 +FROM + assessments AS a +WHERE + a.id = $assessment_id +FOR NO KEY UPDATE; + "#; + + workspace + .open_file(OpenFileParams { + path: path.clone(), + content: content.into(), + version: 1, + }) + .expect("Unable to open test file"); + + let diagnostics = workspace + .pull_diagnostics(crate::workspace::PullDiagnosticsParams { + path: path.clone(), + categories: RuleCategories::all(), + max_diagnostics: 100, + only: vec![], + skip: vec![], + }) + .expect("Unable to pull diagnostics") + .diagnostics; + + assert_eq!(diagnostics.len(), 0, "Expected no diagnostic"); +} + #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")] async fn test_cstyle_comments(test_db: PgPool) { let mut conf = PartialConfiguration::init(); From d2d69d6effdd8747a2961a4ba4a2f746ef6dff16 Mon Sep 17 00:00:00 2001 From: psteinroe Date: Wed, 3 Sep 2025 22:54:32 +0200 Subject: [PATCH 2/3] progress --- crates/pgt_statement_splitter/src/lib.rs | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/crates/pgt_statement_splitter/src/lib.rs b/crates/pgt_statement_splitter/src/lib.rs index ad3537c5a..b8d9a451b 100644 --- a/crates/pgt_statement_splitter/src/lib.rs +++ b/crates/pgt_statement_splitter/src/lib.rs @@ -276,6 +276,22 @@ mod tests { Tester::from("/* this is a test */\nselect 1").expect_statements(vec!["select 1"]); } + #[test] + fn trigger_instead_of() { + Tester::from( + "CREATE OR REPLACE TRIGGER my_trigger + INSTEAD OF INSERT ON my_table + FOR EACH ROW + EXECUTE FUNCTION my_table_trigger_fn();", + ) + .expect_statements(vec![ + "CREATE OR REPLACE TRIGGER my_trigger + INSTEAD OF INSERT ON my_table + FOR EACH ROW + EXECUTE FUNCTION my_table_trigger_fn();", + ]); + } + #[test] fn with_recursive() { Tester::from( From 0ba78f3d9d2821516414f8c12d21ccdae875288b Mon Sep 17 00:00:00 2001 From: psteinroe Date: Wed, 3 Sep 2025 23:25:26 +0200 Subject: [PATCH 3/3] progress --- crates/pgt_workspace/src/workspace/server.tests.rs | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/crates/pgt_workspace/src/workspace/server.tests.rs b/crates/pgt_workspace/src/workspace/server.tests.rs index a41977877..ec47d21ba 100644 --- a/crates/pgt_workspace/src/workspace/server.tests.rs +++ b/crates/pgt_workspace/src/workspace/server.tests.rs @@ -423,7 +423,14 @@ FOR NO KEY UPDATE; .expect("Unable to pull diagnostics") .diagnostics; - assert_eq!(diagnostics.len(), 0, "Expected no diagnostic"); + assert_eq!( + diagnostics + .iter() + .filter(|d| d.category().is_some_and(|c| c.name() == "syntax")) + .count(), + 0, + "Expected no syntax diagnostic" + ); } #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")]