sql-parser: select with CTE

Author: Almaz-KG

src/ast/cte.rs

@@ -9,7 +9,7 @@ pub struct WithCteStatement {
 
     pub cte_expressions: Vec<CteExpression>,
 
-    pub statement: Statement,
+    pub statement: Box<Statement>,
 }
 
 /// An AST for a single CTE expression.
@@ -19,7 +19,13 @@ pub struct CteExpression {
 
     pub column_names: Vec<Identifier>,
 
-    pub materialized: Option<bool>,
+    pub materialized: Option<MaterializationType>,
 
-    pub select: Vec<SelectStatement>,
+    pub select: SelectStatement,
+}
+
+#[derive(Debug, PartialEq)]
+pub enum MaterializationType {
+    Materialized,
+    NotMaterialized,
 }

src/ast/mod.rs

@@ -20,6 +20,7 @@ pub use create::{
     CreateIndexStatement, CreateTableStatement, CreateTriggerStatement, CreateViewStatement,
     CreateVirtualTableStatement,
 };
+pub use cte::{CteExpression, MaterializationType, WithCteStatement};
 pub use delete::DeleteStatement;
 pub use drop::{DropIndexStatement, DropTableStatement, DropTriggerStatement, DropViewStatement};
 pub use explain::ExplainStatement;
@@ -94,6 +95,9 @@ pub enum Statement {
     /// SQLite Administrative statements, see [ReindexStatement]
     Reindex(ReindexStatement),
 
+    /// WITH statement, see [WithCteStatement]
+    WithCte(WithCteStatement),
+
     /// EXPLAIN statement (can wrap another statement), see [ExplainStatement]
     Explain(ExplainStatement),
 }

src/parser/cte.rs

@@ -0,0 +1,110 @@
+use crate::{
+    expression::IdentifierParser, CteExpression, Identifier, Keyword, MaterializationType,
+    TokenType, WithCteStatement,
+};
+
+use super::{select::SelectStatementParser, Parser, ParsingError};
+
+pub trait CteStatementParser {
+    fn parse_cte_statement(&mut self) -> Result<WithCteStatement, ParsingError>;
+
+    fn parse_cte_expression(&mut self) -> Result<CteExpression, ParsingError>;
+
+    fn parse_cte_column_names(&mut self) -> Result<Vec<Identifier>, ParsingError>;
+
+    fn parse_cte_materialized(&mut self) -> Result<Option<MaterializationType>, ParsingError>;
+}
+
+impl<'a> CteStatementParser for Parser<'a> {
+    fn parse_cte_statement(&mut self) -> Result<WithCteStatement, ParsingError> {
+        self.consume_as_keyword(Keyword::With)?;
+
+        let recursive = self.consume_as_keyword(Keyword::Recursive).is_ok();
+
+        let mut cte_expressions = Vec::new();
+        loop {
+            let cte_expression = self.parse_cte_expression()?;
+            cte_expressions.push(cte_expression);
+
+            if self.consume_as(TokenType::Comma).is_err() {
+                break;
+            }
+        }
+
+        let statement = self.parse_statement()?;
+
+        Ok(WithCteStatement {
+            recursive,
+            cte_expressions,
+            statement: Box::new(statement),
+        })
+    }
+
+    fn parse_cte_expression(&mut self) -> Result<CteExpression, ParsingError> {
+        let name = self.parse_identifier()?;
+        let column_names = self.parse_cte_column_names()?;
+
+        self.consume_as_keyword(Keyword::As)?;
+
+        let materialized = self.parse_cte_materialized()?;
+
+        self.consume_as(TokenType::LeftParen)?;
+        let select = self.parse_select_statement()?;
+        self.consume_as(TokenType::RightParen)?;
+
+        Ok(CteExpression {
+            name,
+            column_names,
+            materialized,
+            select,
+        })
+    }
+
+    fn parse_cte_column_names(&mut self) -> Result<Vec<Identifier>, ParsingError> {
+        let mut column_names = Vec::new();
+
+        if self.consume_as(TokenType::LeftParen).is_ok() {
+            loop {
+                let column_name = self.parse_identifier()?;
+                column_names.push(column_name);
+
+                if self.consume_as(TokenType::Comma).is_err() {
+                    break;
+                }
+            }
+            self.consume_as(TokenType::RightParen)?;
+        }
+
+        Ok(column_names)
+    }
+
+    fn parse_cte_materialized(&mut self) -> Result<Option<MaterializationType>, ParsingError> {
+        if self.consume_as_keyword(Keyword::Materialized).is_ok() {
+            Ok(Some(MaterializationType::Materialized))
+        } else if self.consume_as_keyword(Keyword::Not).is_ok() {
+            self.consume_as_keyword(Keyword::Materialized)?;
+            Ok(Some(MaterializationType::NotMaterialized))
+        } else {
+            Ok(None)
+        }
+    }
+}
+
+#[cfg(test)]
+pub mod test_utils {
+    use crate::{CteExpression, Identifier, MaterializationType, SelectStatement};
+
+    pub fn cte_expression(
+        name: Identifier,
+        column_names: Vec<Identifier>,
+        materialized: Option<MaterializationType>,
+        select: SelectStatement,
+    ) -> CteExpression {
+        CteExpression {
+            name,
+            column_names,
+            materialized,
+            select,
+        }
+    }
+}

src/parser/mod.rs

@@ -2,6 +2,7 @@ use std::iter::Peekable;
 
 mod alter;
 mod column_definition;
+mod cte;
 mod drop;
 mod errors;
 pub(crate) mod expression;
@@ -15,12 +16,12 @@ mod test_utils;
 
 use crate::{Keyword, SelectStatement, Statement, Token, TokenType, Tokenizer};
 use alter::AlterTableStatementParser;
+use cte::CteStatementParser;
 use drop::DropStatementParser;
 pub use errors::*;
 use select::{SelectStatementParser, ValuesStatementParser};
 use sqlite::SQLite3StatementParser;
 use trx::TransactionStatementParser;
-
 /// A parser for SQLite SQL statements
 pub struct Parser<'a> {
     tokenizer: Peekable<Tokenizer<'a>>,
@@ -240,7 +241,7 @@ impl<'a> Parser<'a> {
             }
             Keyword::Values => ValuesStatementParser::parse_values_statement(self)
                 .map(|stmt| Statement::Select(SelectStatement::Values(stmt))),
-
+            Keyword::With => CteStatementParser::parse_cte_statement(self).map(Statement::WithCte),
             Keyword::Alter => AlterTableStatementParser::parse_alter_table_statement(self)
                 .map(Statement::AlterTable),
             keyword => Err(ParsingError::UnexpectedKeyword(keyword)),

src/parser/select/mod.rs

@@ -476,9 +476,9 @@ impl<'a> SelectStatementParser for Parser<'a> {
 #[cfg(test)]
 mod test_utils {
     use crate::{
-        DistinctType, Expression, Identifier, LimitClause, NamedWindowDefinition, OrderingTerm,
-        QualifiedTableName, Select, SelectFrom, SelectItem, SelectStatement, UnionStatement,
-        UnionStatementType,
+        CteExpression, DistinctType, Expression, Identifier, LimitClause, NamedWindowDefinition,
+        OrderingTerm, QualifiedTableName, Select, SelectFrom, SelectItem, SelectStatement,
+        Statement, UnionStatement, UnionStatementType, WithCteStatement,
     };
 
     pub fn select_statement_with_columns(
@@ -613,6 +613,28 @@ mod test_utils {
             ..Default::default()
         })
     }
+
+    pub fn select_statement_with_cte_clause(
+        recursive: bool,
+        expressions: Vec<CteExpression>,
+    ) -> Statement {
+        Statement::WithCte(WithCteStatement {
+            recursive,
+            cte_expressions: expressions,
+            statement: Box::new(Statement::Select(SelectStatement::Select(Select {
+                distinct_type: DistinctType::None,
+                columns: vec![SelectItem::Expression(Expression::Identifier(
+                    Identifier::Wildcard,
+                ))],
+                from: Some(SelectFrom::Table(QualifiedTableName {
+                    table_id: Identifier::Single("table_1".to_string()),
+                    alias: None,
+                    indexed_type: None,
+                })),
+                ..Default::default()
+            }))),
+        })
+    }
 }
 
 #[cfg(test)]
@@ -1826,3 +1848,61 @@ mod test_select_limit_clause {
         );
     }
 }
+
+#[cfg(test)]
+mod test_select_cte_clause {
+    use super::super::cte::test_utils::cte_expression;
+    use super::test_utils::{select_statement_with_cte_clause, select_statement_with_from};
+    use crate::parser::test_utils::*;
+    use crate::{Identifier, QualifiedTableName, SelectFrom};
+
+    #[test]
+    fn test_select_cte_clause() {
+        let expected_statement = select_statement_with_cte_clause(
+            true,
+            vec![cte_expression(
+                Identifier::Single("cte_1".to_string()),
+                vec![],
+                None,
+                select_statement_with_from(SelectFrom::Table(QualifiedTableName::from(
+                    Identifier::Single("cte_table".to_string()),
+                ))),
+            )],
+        );
+
+        run_sunny_day_test(
+            "WITH RECURSIVE cte_1 AS (SELECT * FROM cte_table) SELECT * FROM table_1",
+            expected_statement,
+        );
+    }
+
+    #[test]
+    fn test_select_with_multiple_ctes() {
+        let expected_statement = select_statement_with_cte_clause(
+            true,
+            vec![
+                cte_expression(
+                    Identifier::Single("cte_1".to_string()),
+                    vec![],
+                    None,
+                    select_statement_with_from(SelectFrom::Table(QualifiedTableName::from(
+                        Identifier::Single("cte_table1".to_string()),
+                    ))),
+                ),
+                cte_expression(
+                    Identifier::Single("cte_2".to_string()),
+                    vec![],
+                    None,
+                    select_statement_with_from(SelectFrom::Table(QualifiedTableName::from(
+                        Identifier::Single("cte_table2".to_string()),
+                    ))),
+                ),
+            ],
+        );
+
+        run_sunny_day_test(
+            "WITH RECURSIVE cte_1 AS (SELECT * FROM cte_table1), cte_2 AS (SELECT * FROM cte_table2) SELECT * FROM table_1",
+            expected_statement,
+        );
+    }
+}