diff --git a/crates/pgls_statement_splitter/src/lib.rs b/crates/pgls_statement_splitter/src/lib.rs index 8cde70a3f..c51ccbcea 100644 --- a/crates/pgls_statement_splitter/src/lib.rs +++ b/crates/pgls_statement_splitter/src/lib.rs @@ -352,6 +352,23 @@ END;", .expect_statements(vec!["insert into tbl (id) select 1", "select 3"]); } + #[test] + fn insert_with_cte() { + let insert = "INSERT INTO target (id) +WITH source AS ( + SELECT 1 AS id +) +SELECT id FROM source;"; + let input = format!( + "{insert} +SELECT 2;" + ); + + Tester::from(input.as_str()) + .expect_statements(vec![insert, "SELECT 2;"]) + .assert_no_errors(); + } + #[test] fn c_style_comments() { Tester::from("/* this is a test */\nselect 1").expect_statements(vec!["select 1"]); diff --git a/crates/pgls_statement_splitter/src/splitter/dml.rs b/crates/pgls_statement_splitter/src/splitter/dml.rs index 22fe23f1c..bd70f3245 100644 --- a/crates/pgls_statement_splitter/src/splitter/dml.rs +++ b/crates/pgls_statement_splitter/src/splitter/dml.rs @@ -82,7 +82,7 @@ pub(crate) fn insert(p: &mut Splitter) -> SplitterResult { p.expect(SyntaxKind::INSERT_KW)?; p.expect(SyntaxKind::INTO_KW)?; - unknown(p, &[SyntaxKind::SELECT_KW]) + unknown(p, &[SyntaxKind::WITH_KW, SyntaxKind::SELECT_KW]) } pub(crate) fn update(p: &mut Splitter) -> SplitterResult {