sql-parser: parse with CTEs

Author: Almaz-KG

src/ast/cte.rs

@@ -1,3 +1,5 @@
+use std::fmt::Display;
+
 use super::{Identifier, SelectStatement, Statement};
 
 /// An AST for [WITH](https://www.sqlite.org/lang_with.html) SQL statement.
@@ -24,8 +26,17 @@ pub struct CteExpression {
     pub select: SelectStatement,
 }
 
-#[derive(Debug, PartialEq)]
+#[derive(Debug, PartialEq, Clone)]
 pub enum MaterializationType {
     Materialized,
     NotMaterialized,
 }
+
+impl Display for MaterializationType {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            MaterializationType::Materialized => write!(f, "MATERIALIZED"),
+            MaterializationType::NotMaterialized => write!(f, "NOT MATERIALIZED"),
+        }
+    }
+}

src/ast/expression.rs

@@ -100,6 +100,15 @@ pub enum Identifier {
     NameWithWildcard(String),
 }
 
+impl From<&str> for Identifier {
+    fn from(s: &str) -> Self {
+        if s == "*" {
+            return Identifier::Wildcard;
+        }
+        Identifier::Single(s.to_string())
+    }
+}
+
 /// A binary operation
 #[derive(Debug, PartialEq, Clone)]
 pub enum BinaryOp {

src/parser/cte.rs

@@ -108,3 +108,115 @@ pub mod test_utils {
         }
     }
 }
+
+#[cfg(test)]
+pub mod test_cte_statement_parser {
+    use super::super::select::test_utils::select_star_from;
+    use super::test_utils::cte_expression;
+    use crate::{
+        parser::test_utils::run_sunny_day_test, Identifier, MaterializationType, Statement,
+        WithCteStatement,
+    };
+
+    #[test]
+    fn test_cte_statement_with_recursive() {
+        let expected_statement = Statement::WithCte(WithCteStatement {
+            recursive: true,
+            cte_expressions: vec![cte_expression(
+                Identifier::from("cte_1"),
+                vec![],
+                None,
+                select_star_from(Identifier::from("cte_table")),
+            )],
+            statement: Box::new(Statement::Select(select_star_from(Identifier::from(
+                "table_1",
+            )))),
+        });
+
+        run_sunny_day_test(
+            "WITH RECURSIVE cte_1 AS (SELECT * FROM cte_table) SELECT * FROM table_1",
+            expected_statement,
+        );
+    }
+
+    #[test]
+    fn test_cte_statement_with_multiple_ctes() {
+        let expected_statement = Statement::WithCte(WithCteStatement {
+            recursive: false,
+            cte_expressions: vec![
+                cte_expression(
+                    Identifier::from("cte_1"),
+                    vec![],
+                    None,
+                    select_star_from(Identifier::from("cte_table1")),
+                ),
+                cte_expression(
+                    Identifier::from("cte_2"),
+                    vec![],
+                    None,
+                    select_star_from(Identifier::from("cte_1")),
+                ),
+            ],
+            statement: Box::new(Statement::Select(select_star_from(Identifier::from(
+                "cte_2",
+            )))),
+        });
+
+        run_sunny_day_test(
+            "WITH cte_1 AS (SELECT * FROM cte_table1), cte_2 AS (SELECT * FROM cte_1) SELECT * FROM cte_2",
+            expected_statement,
+        );
+    }
+
+    #[test]
+    fn test_cte_statement_with_materialized() {
+        let materialization_types = vec![
+            MaterializationType::Materialized,
+            MaterializationType::NotMaterialized,
+        ];
+
+        for materialization_type in materialization_types {
+            let expected_statement = Statement::WithCte(WithCteStatement {
+                recursive: false,
+                cte_expressions: vec![cte_expression(
+                    Identifier::from("cte_1"),
+                    vec![],
+                    Some(materialization_type.clone()),
+                    select_star_from(Identifier::from("cte_table")),
+                )],
+                statement: Box::new(Statement::Select(select_star_from(Identifier::from(
+                    "cte_1",
+                )))),
+            });
+
+            run_sunny_day_test(
+                &format!(
+                    "WITH cte_1 AS {} (SELECT * FROM cte_table) SELECT * FROM cte_1",
+                    materialization_type
+                ),
+                expected_statement,
+            );
+        }
+    }
+
+    #[test]
+    fn test_cte_statement_with_column_names() {
+        let expected_statement = Statement::WithCte(WithCteStatement {
+            recursive: false,
+            cte_expressions: vec![cte_expression(
+                Identifier::from("cte_1"),
+                vec![Identifier::from("col1"), Identifier::from("col2")],
+                None,
+                select_star_from(Identifier::from("cte_table")),
+            )],
+            statement: Box::new(Statement::Select(select_star_from(Identifier::from(
+                "cte_1",
+            )))),
+        });
+
+        run_sunny_day_test(
+            "WITH cte_1 (col1, col2) AS (SELECT * FROM cte_table) SELECT * FROM cte_1",
+            expected_statement,
+        );
+    }
+}

src/parser/select/mod.rs

@@ -474,7 +474,7 @@ impl<'a> SelectStatementParser for Parser<'a> {
 }
 
 #[cfg(test)]
-mod test_utils {
+pub mod test_utils {
     use crate::{
         CteExpression, DistinctType, Expression, Identifier, LimitClause, NamedWindowDefinition,
         OrderingTerm, QualifiedTableName, Select, SelectFrom, SelectItem, SelectStatement,
@@ -503,6 +503,10 @@ mod test_utils {
         })
     }
 
+    pub fn select_star_from(table_name: Identifier) -> SelectStatement {
+        select_statement_with_from(SelectFrom::Table(QualifiedTableName::from(table_name)))
+    }
+
     pub fn select_statement_with_where_clause(where_clause: Expression) -> SelectStatement {
         SelectStatement::Select(Select {
             distinct_type: DistinctType::None,