diff --git a/crates/pgls_statement_splitter/src/lib.rs b/crates/pgls_statement_splitter/src/lib.rs index 8cde70a3f..c6189986e 100644 --- a/crates/pgls_statement_splitter/src/lib.rs +++ b/crates/pgls_statement_splitter/src/lib.rs @@ -352,6 +352,81 @@ 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 create_as_with_values_cte() { + let create = "CREATE TABLE mappers.alur_zones AS +WITH + zone_ranks (zone_alur, tension_rank) AS ( + VALUES + ('Abis', 1), + ('A', 2), + ('B1', 3), + ('B2', 4), + ('C', 5) + ) +SELECT + code_postal AS zip_code, + commune, + zone_alur +FROM apis.alur_zones;"; + let input = format!( + "{create} +SELECT 2;" + ); + + Tester::from(input.as_str()) + .expect_statements(vec![create, "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 +527,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") diff --git a/crates/pgls_statement_splitter/src/splitter/ddl.rs b/crates/pgls_statement_splitter/src/splitter/ddl.rs index e0d2db64d..7b25f237b 100644 --- a/crates/pgls_statement_splitter/src/splitter/ddl.rs +++ b/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 { diff --git a/crates/pgls_statement_splitter/src/splitter/dml.rs b/crates/pgls_statement_splitter/src/splitter/dml.rs index 22fe23f1c..f6aa5a731 100644 --- a/crates/pgls_statement_splitter/src/splitter/dml.rs +++ b/crates/pgls_statement_splitter/src/splitter/dml.rs @@ -14,6 +14,9 @@ pub(crate) fn cte(p: &mut Splitter) -> SplitterResult { loop { p.expect(SyntaxKind::IDENT)?; + if p.current() == SyntaxKind::L_PAREN { + parenthesis(p)?; + } p.expect(SyntaxKind::AS_KW)?; // Handle optional [NOT] MATERIALIZED hint (PostgreSQL 12+) p.eat(SyntaxKind::NOT_KW)?; @@ -27,17 +30,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 ] ` and