fix: statement splitter (#496)

fixes #494

Author: psteinroe

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::<Vec<_>>()
+            );
+            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'));")
@@ -256,6 +292,55 @@ mod tests {
         ]);
     }
 
+    #[test]
+    fn with_recursive() {
+        Tester::from(
+            "
+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;",
+        )
+        .assert_single_statement()
+        .assert_no_errors();
+    }
+
     #[test]
     fn with_check() {
         Tester::from("create policy employee_insert on journey_execution for insert to authenticated with check ((select private.organisation_id()) = organisation_id);")

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)
     }

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())

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);

crates/pgt_workspace/src/workspace/server.tests.rs

@@ -579,6 +579,53 @@ async fn test_disable_typecheck(test_db: PgPool) {
     );
 }
 
+#[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
+            .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")]
 async fn test_cstyle_comments(test_db: PgPool) {
     let mut conf = PartialConfiguration::init();