fix(splitter): do not split EXPLAIN from the explained statement (#759)

Author: amogiska

crates/pgls_statement_splitter/src/lib.rs

@@ -133,6 +133,39 @@ mod tests {
         }
     }
 
+    #[test]
+    fn explain_statements() {
+        let single_statements = vec![
+            "EXPLAIN SELECT 1;",
+            "explain select 1;",
+            "EXPLAIN ANALYZE VERBOSE SELECT * FROM contact;",
+            "EXPLAIN ANALYSE SELECT * FROM contact;",
+            "EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) UPDATE contact SET name = 'x';",
+            "EXPLAIN INSERT INTO contact (id) VALUES (1);",
+            "EXPLAIN DELETE FROM contact WHERE id = 1;",
+            "EXPLAIN WITH x AS (SELECT 1) SELECT * FROM x;",
+            "EXPLAIN CREATE TABLE contact_copy AS SELECT * FROM contact;",
+            "EXPLAIN VALUES (1);",
+        ];
+
+        for stmt in single_statements {
+            Tester::from(stmt).expect_statements(vec![stmt]);
+        }
+
+        Tester::from("EXPLAIN SELECT 1; SELECT 2;")
+            .expect_statements(vec!["EXPLAIN SELECT 1;", "SELECT 2;"]);
+
+        Tester::from("explain select 1\nselect 2")
+            .expect_statements(vec!["explain select 1", "select 2"]);
+
+        Tester::from("explain select 1\n\nselect 2")
+            .expect_statements(vec!["explain select 1", "select 2"]);
+
+        // explain is an unreserved keyword and remains usable as an identifier
+        Tester::from("SELECT explain FROM contact;")
+            .expect_statements(vec!["SELECT explain FROM contact;"]);
+    }
+
     #[test]
     fn begin_commit() {
         Tester::from(

crates/pgls_statement_splitter/src/splitter/common.rs

@@ -7,7 +7,7 @@ use super::{
     Splitter,
     data::at_statement_start,
     ddl::{alter, create},
-    dml::{cte, delete, insert, select, update},
+    dml::{cte, delete, explain, insert, select, update},
 };
 
 #[derive(Debug)]
@@ -58,6 +58,7 @@ pub(crate) fn statement(p: &mut Splitter) -> SplitterResult {
         SyntaxKind::DELETE_KW => delete(p),
         SyntaxKind::CREATE_KW => create(p),
         SyntaxKind::ALTER_KW => alter(p),
+        SyntaxKind::EXPLAIN_KW => explain(p),
         _ => unknown(p, &[]),
     };
 

crates/pgls_statement_splitter/src/splitter/dml.rs

@@ -5,6 +5,7 @@ use crate::splitter::common::SplitterResult;
 use super::{
     Splitter,
     common::{parenthesis, unknown},
+    ddl::create,
 };
 
 pub(crate) fn cte(p: &mut Splitter) -> SplitterResult {
@@ -39,6 +40,38 @@ pub(crate) fn cte(p: &mut Splitter) -> SplitterResult {
     Ok(())
 }
 
+/// `EXPLAIN [ ANALYZE ] [ VERBOSE ] <statement>` and
+/// `EXPLAIN ( option [, ...] ) <statement>`
+///
+/// EXPLAIN is a prefix to the statement it explains, so after consuming the
+/// prefix we delegate to the splitter function of the inner statement instead
+/// of treating its leading keyword as a new statement start.
+pub(crate) fn explain(p: &mut Splitter) -> SplitterResult {
+    p.expect(SyntaxKind::EXPLAIN_KW)?;
+
+    // EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON, ...) <statement>
+    if p.current() == SyntaxKind::L_PAREN {
+        parenthesis(p)?;
+    } else {
+        // legacy EXPLAIN [ANALYZE | ANALYSE] [VERBOSE] <statement>
+        p.eat(SyntaxKind::ANALYZE_KW)?;
+        p.eat(SyntaxKind::ANALYSE_KW)?;
+        p.eat(SyntaxKind::VERBOSE_KW)?;
+    }
+
+    match p.current() {
+        SyntaxKind::WITH_KW => cte(p),
+        SyntaxKind::SELECT_KW => select(p),
+        SyntaxKind::INSERT_KW => insert(p),
+        SyntaxKind::UPDATE_KW => update(p),
+        SyntaxKind::DELETE_KW => delete(p),
+        // EXPLAIN CREATE TABLE AS / CREATE MATERIALIZED VIEW AS
+        SyntaxKind::CREATE_KW => create(p),
+        // MERGE, VALUES, EXECUTE, DECLARE
+        _ => unknown(p, &[]),
+    }
+}
+
 pub(crate) fn select(p: &mut Splitter) -> SplitterResult {
     p.expect(SyntaxKind::SELECT_KW)?;
 

crates/pgls_statement_splitter/tests/data/explain_statements__5.sql

@@ -0,0 +1,9 @@
+explain select 1 from contact;
+
+explain analyze select * from contact where id = 1;
+
+explain (analyze, buffers, format json) update contact set name = 'x' where id = 1;
+
+explain with x as (select 1) select * from x;
+
+select 1;