From 837efabc7b180dedaca26e47359aa75859acc5a6 Mon Sep 17 00:00:00 2001
From: psteinroe <philipp@steinroetter.com>
Date: Sat, 29 Mar 2025 13:42:36 +0100
Subject: [PATCH] fix: cte

---
 crates/pgt_statement_splitter/src/lib.rs        |  6 ++++++
 crates/pgt_statement_splitter/src/parser/dml.rs | 13 +++++++++++--
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/crates/pgt_statement_splitter/src/lib.rs b/crates/pgt_statement_splitter/src/lib.rs
index 5418139a..ee182087 100644
--- a/crates/pgt_statement_splitter/src/lib.rs
+++ b/crates/pgt_statement_splitter/src/lib.rs
@@ -142,6 +142,12 @@ mod tests {
             .expect_statements(vec!["insert into tbl (id) select 1", "select 3"]);
     }
 
+    #[test]
+    fn with_cte() {
+        Tester::from("with test as (select 1 as id) select * from test;")
+            .expect_statements(vec!["with test as (select 1 as id) select * from test;"]);
+    }
+
     #[test]
     fn case() {
         Tester::from("select case when select 2 then 1 else 0 end")
diff --git a/crates/pgt_statement_splitter/src/parser/dml.rs b/crates/pgt_statement_splitter/src/parser/dml.rs
index e758f3d7..a45f6c40 100644
--- a/crates/pgt_statement_splitter/src/parser/dml.rs
+++ b/crates/pgt_statement_splitter/src/parser/dml.rs
@@ -2,7 +2,7 @@ use pgt_lexer::SyntaxKind;
 
 use super::{
     Parser,
-    common::{parenthesis, statement, unknown},
+    common::{parenthesis, unknown},
 };
 
 pub(crate) fn cte(p: &mut Parser) {
@@ -18,7 +18,16 @@ pub(crate) fn cte(p: &mut Parser) {
         }
     }
 
-    statement(p);
+    unknown(
+        p,
+        &[
+            SyntaxKind::Select,
+            SyntaxKind::Insert,
+            SyntaxKind::Update,
+            SyntaxKind::DeleteP,
+            SyntaxKind::Merge,
+        ],
+    );
 }
 
 pub(crate) fn select(p: &mut Parser) {