fix(splitter): keep create-as CTEs attached

crates/pgls_statement_splitter/src/lib.rs

@@ -352,6 +352,54 @@ END;",
             .expect_statements(vec!["insert into tbl (id) select 1", "select 3"]);
     }
 
+    #[test]
+    fn create_as_with_cte() {
+        let create_statements = [
+            "CREATE TABLE target AS",
+            "CREATE MATERIALIZED VIEW target AS",
+            "CREATE VIEW target AS",
+        ];
+
+        for create_as in create_statements {
+            let create = format!(
+                "{create_as}
+WITH cte AS (
+    SELECT 1 AS id
+)
+SELECT id FROM cte;"
+            );
+            let input = format!(
+                "{create}
+SELECT 2;"
+            );
+
+            Tester::from(input.as_str())
+                .expect_statements(vec![create.as_str(), "SELECT 2;"])
+                .assert_no_errors();
+        }
+    }
+
+    #[test]
+    fn create_view_with_options_as_cte() {
+        Tester::from(
+            "CREATE VIEW target WITH (security_invoker) AS
+WITH cte AS (
+    SELECT 1 AS id
+)
+SELECT id FROM cte;
+SELECT 2;",
+        )
+        .expect_statements(vec![
+            "CREATE VIEW target WITH (security_invoker) AS
+WITH cte AS (
+    SELECT 1 AS id
+)
+SELECT id FROM cte;",
+            "SELECT 2;",
+        ])
+        .assert_no_errors();
+    }
+
     #[test]
     fn c_style_comments() {
         Tester::from("/* this is a test */\nselect 1").expect_statements(vec!["select 1"]);
@@ -452,6 +500,25 @@ LIMIT
             .expect_statements(vec!["with test as (select 1 as id) select * from test;"]);
     }
 
+    #[test]
+    fn with_cte_followed_by_statement() {
+        Tester::from(
+            "WITH cte AS (
+    SELECT 1 AS id
+)
+SELECT id FROM cte;
+SELECT 2;",
+        )
+        .expect_statements(vec![
+            "WITH cte AS (
+    SELECT 1 AS id
+)
+SELECT id FROM cte;",
+            "SELECT 2;",
+        ])
+        .assert_no_errors();
+    }
+
     #[test]
     fn case() {
         Tester::from("select case when select 2 then 1 else 0 end")

crates/pgls_statement_splitter/src/splitter/ddl.rs

@@ -2,12 +2,31 @@ use pgls_lexer::SyntaxKind;
 
 use crate::splitter::common::SplitterResult;
 
-use super::{Splitter, common::unknown};
+use super::{
+    Splitter,
+    common::{parenthesis, unknown},
+    dml::cte,
+};
 
 pub(crate) fn create(p: &mut Splitter) -> SplitterResult {
     p.expect(SyntaxKind::CREATE_KW)?;
 
-    unknown(p, &[SyntaxKind::WITH_KW])
+    loop {
+        unknown(p, &[])?;
+
+        if p.current() != SyntaxKind::WITH_KW {
+            return Ok(());
+        }
+
+        if p.look_back(true) == Some(SyntaxKind::AS_KW) {
+            return cte(p);
+        }
+
+        p.expect(SyntaxKind::WITH_KW)?;
+        if p.current() == SyntaxKind::L_PAREN {
+            parenthesis(p)?;
+        }
+    }
 }
 
 pub(crate) fn alter(p: &mut Splitter) -> SplitterResult {

crates/pgls_statement_splitter/src/splitter/dml.rs

@@ -27,17 +27,13 @@ pub(crate) fn cte(p: &mut Splitter) -> SplitterResult {
         }
     }
 
-    unknown(
-        p,
-        &[
-            SyntaxKind::SELECT_KW,
-            SyntaxKind::INSERT_KW,
-            SyntaxKind::UPDATE_KW,
-            SyntaxKind::DELETE_KW,
-            SyntaxKind::MERGE_KW,
-        ],
-    )?;
-    Ok(())
+    match p.current() {
+        SyntaxKind::SELECT_KW => select(p),
+        SyntaxKind::INSERT_KW => insert(p),
+        SyntaxKind::UPDATE_KW => update(p),
+        SyntaxKind::DELETE_KW => delete(p),
+        _ => unknown(p, &[]),
+    }
 }
 
 /// `EXPLAIN [ ANALYZE ] [ VERBOSE ] <statement>` and