Merge pull request #281 from Qrlew/support_rereading

Support for parsing CTEs with column aliases

CHANGELOG.md

@@ -6,6 +6,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
+## [0.9.19] - 2024-05-23
+### Added
+- parsing of CTEs with column alias
+
 ## [0.9.18] - 2024-05-23
 ### Added
 - use a column present in the schema as privacy unit weight

Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 authors = ["Nicolas Grislain <[email protected]>"]
 name = "qrlew"
-version = "0.9.18"
+version = "0.9.19"
 edition = "2021"
 description = "Sarus Qrlew Engine"
 documentation = "https://docs.rs/qrlew"

examples/ast.rs

@@ -83,7 +83,7 @@ fn build_ast() -> Result<(), &'static str> {
         },
         query: Box::new(query.clone()),
         from: None,
-        materialized: None
+        materialized: None,
     };
     println!("{}", cte);
     let cte = Cte {
@@ -93,7 +93,7 @@ fn build_ast() -> Result<(), &'static str> {
         },
         query: Box::new(query),
         from: Some(Ident::new("fro")),
-        materialized: None
+        materialized: None,
     };
     println!("{}", cte);
 

src/rewriting/mod.rs

@@ -281,6 +281,17 @@ mod tests {
                 .iter()
                 .map(ToString::to_string)
                 .join("\n");
+
+            // Test re-parsing
+            let re_parsed = parse(dp_query.as_str()).unwrap();
+            let relation = Relation::try_from(re_parsed.with(&relations)).unwrap();
+            let query = ast::Query::from(&relation).to_string();
+            _ = database
+                .query(query.as_str())
+                .unwrap()
+                .iter()
+                .map(ToString::to_string)
+                .join("\n");
         }
     }
 

src/sql/expr.rs

@@ -1109,6 +1109,8 @@ impl<'a> Visitor<'a, Result<Expr>> for TryIntoExprVisitor<'a> {
                 };
                 Expr::unix_timestamp(arg)
             }
+            "greatest" => Expr::greatest(flat_args[0].clone(), flat_args[1].clone()),
+            "least" => Expr::least(flat_args[0].clone(), flat_args[1].clone()),
             // Aggregates
             "min" => Expr::min(flat_args[0].clone()),
             "max" => Expr::max(flat_args[0].clone()),

src/sql/mod.rs

@@ -4,6 +4,7 @@
 //!
 
 pub mod expr;
+pub mod query_aliases;
 pub mod query_names;
 pub mod reader;
 pub mod relation;
@@ -133,6 +134,10 @@ mod tests {
         ] {
             let relation = Relation::try_from(parse(query).unwrap().with(&database.relations())).unwrap();
             relation.display_dot().unwrap();
+            let relation_query: &str = &ast::Query::from(&relation).to_string();
+            println!("QUERY:\n{}", relation_query);
+            let relation = Relation::try_from(parse(relation_query).unwrap().with(&database.relations())).unwrap();
+            relation.display_dot().unwrap();
         }
     }
 
@@ -159,4 +164,137 @@ mod tests {
             assert_eq!(res1, res2);
         }
     }
+
+    #[test]
+    fn test_parsing_complex_cte_query_with_column_aliases() {
+        let database = postgresql::test_database();
+        let query: &str = r#"
+        WITH
+        "map_3kfp" ("a", "d") AS (
+            SELECT
+            "a" AS "a",
+            "d" AS "d"
+            FROM
+            "table_1"
+        ),
+        "map_l_fx" ("x", "y", "z") AS (
+            SELECT
+            "x" AS "x",
+            "y" AS "y",
+            "z" AS "z"
+            FROM
+            "table_2"
+        ),
+        "join_pgtf" (
+            "field_yc23",
+            "field_yep7",
+            "field_dzgn",
+            "field_pwch",
+            "field_w_o_"
+        ) AS (
+            SELECT
+            *
+            FROM
+            "map_3kfp" AS "_LEFT_"
+            JOIN "map_l_fx" AS "_RIGHT_" ON ("_LEFT_"."d") = ("_RIGHT_"."x")
+        ),
+        "join_uvlf" (
+            "field_l5nj",
+            "field_bn3m",
+            "field_m0fl",
+            "field_v3e7",
+            "field_z8oi",
+            "field_dzgn",
+            "field_pwch",
+            "field_w_o_"
+        ) AS (
+            SELECT
+            *
+            FROM
+            "join_pgtf" AS "_LEFT_"
+            JOIN "table_2" AS "_RIGHT_" ON ("_LEFT_"."field_yep7") = ("_RIGHT_"."x")
+        ),
+        "map_hfqo" (
+            "a",
+            "d",
+            "field_m0fl",
+            "field_v3e7",
+            "field_z8oi",
+            "field_dzgn",
+            "field_pwch",
+            "field_w_o_"
+        ) AS (
+            SELECT
+            "field_l5nj" AS "a",
+            "field_bn3m" AS "d",
+            "field_m0fl" AS "field_m0fl",
+            "field_v3e7" AS "field_v3e7",
+            "field_z8oi" AS "field_z8oi",
+            "field_dzgn" AS "field_dzgn",
+            "field_pwch" AS "field_pwch",
+            "field_w_o_" AS "field_w_o_"
+            FROM
+            "join_uvlf"
+        ),
+        "map_m5ph" (
+            "a",
+            "d",
+            "field_m0fl",
+            "field_v3e7",
+            "field_z8oi",
+            "field_dzgn",
+            "field_pwch",
+            "field_w_o_"
+        ) AS (
+            SELECT
+            "a" AS "a",
+            "d" AS "d",
+            "field_m0fl" AS "field_m0fl",
+            "field_v3e7" AS "field_v3e7",
+            "field_z8oi" AS "field_z8oi",
+            "field_dzgn" AS "field_dzgn",
+            "field_pwch" AS "field_pwch",
+            "field_w_o_" AS "field_w_o_"
+            FROM
+            "map_hfqo"
+            ORDER BY
+            "a" ASC
+            LIMIT
+            10
+        )
+        SELECT
+        *
+        FROM
+        "map_m5ph"
+        LIMIT
+        10
+        "#;
+        let query = parse(query).unwrap();
+        println!("QUERY: {}", query);
+        let binding = database.relations();
+        let qwr = query.with(&binding);
+        let relation = Relation::try_from(qwr).unwrap();
+        relation.display_dot().unwrap();
+    }
+
+    #[test]
+    fn test_parsing_simple_cte_query_with_column_aliases() {
+        let database = postgresql::test_database();
+        let query: &str = r#"
+        WITH
+        "tab1" ("my_new_name_a", "my_new_name_b") AS (
+            SELECT * FROM "table_1"
+        ),
+        "tab2" ("my_new_new_name_a") AS (
+            SELECT "my_new_name_a" FROM "tab1"
+        )
+        SELECT "my_new_new_name_a" FROM "tab2"
+        "#;
+        let query = parse(query).unwrap();
+        println!("QUERY: {}", query);
+        let binding = database.relations();
+        let qwr = query.with(&binding);
+        let relation = Relation::try_from(qwr).unwrap();
+        relation.display_dot().unwrap();
+    }
 }

src/sql/query_aliases.rs

@@ -0,0 +1,148 @@
+//! Data structures and visitor to collect aliases associated to CTEs query.
+
+use crate::{ast, sql::visitor::Visitor, visitor::Visited};
+use colored::Colorize;
+use itertools::Itertools;
+use std::{
+    collections::BTreeMap,
+    fmt,
+    iter::Iterator,
+    ops::{Deref, DerefMut},
+};
+
+/// A mapping between a Query the column aliases associated to it.
+/// It is used to correctly associate aliases to queries in CTEs with column aliases (e.g. WITH tab (a, b, c) AS (SELECT...)).
+#[derive(Clone, Debug, Hash, PartialEq, Eq)]
+pub struct QueryAliases<'a>(BTreeMap<&'a ast::Query, Option<&'a Vec<ast::Ident>>>);
+
+impl<'a> QueryAliases<'a> {
+    /// Build a new QueryAliases object
+    pub fn new() -> Self {
+        QueryAliases(BTreeMap::new())
+    }
+}
+
+impl<'a> Deref for QueryAliases<'a> {
+    type Target = BTreeMap<&'a ast::Query, Option<&'a Vec<ast::Ident>>>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl<'a> DerefMut for QueryAliases<'a> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
+impl<'a> IntoIterator for QueryAliases<'a> {
+    type Item = <BTreeMap<&'a ast::Query, Option<&'a Vec<ast::Ident>>> as IntoIterator>::Item;
+    type IntoIter =
+        <BTreeMap<&'a ast::Query, Option<&'a Vec<ast::Ident>>> as IntoIterator>::IntoIter;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.0.into_iter()
+    }
+}
+
+impl<'a> fmt::Display for QueryAliases<'a> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "Query Aliases\n{}",
+            self.0
+                .iter()
+                .map(|(q, r)| match r {
+                    Some(r) => format!(
+                        "{} -> {}",
+                        format!("{q}").blue(),
+                        format!("({})", r.iter().join(", ")).green()
+                    ),
+                    None => format!(
+                        "{} -> {}",
+                        format!("{q}").blue(),
+                        format!("?").bold().green()
+                    ),
+                })
+                .join(",\n")
+        )
+    }
+}
+
+/// A visitor to build a reference free query tree
+#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
+pub struct IntoQueryAliasesVisitor;
+
+impl<'a> Visitor<'a, QueryAliases<'a>> for IntoQueryAliasesVisitor {
+    fn query(
+        &self,
+        query: &'a ast::Query,
+        _dependencies: Visited<'a, ast::Query, QueryAliases<'a>>,
+    ) -> QueryAliases<'a> {
+        let mut query_aliases = QueryAliases::new();
+        if let Some(with) = &query.with {
+            for cte in &with.cte_tables {
+                if !cte.alias.columns.is_empty() {
+                    query_aliases.insert(cte.query.as_ref(), Some(cte.alias.columns.as_ref()));
+                }
+            }
+        }
+        query_aliases
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use colored::Colorize;
+
+    use super::*;
+    use crate::{sql::relation, visitor::Acceptor as _};
+
+    #[test]
+    fn test_query_aliases() {
+        let query_1 = relation::parse("select * from table_1").unwrap();
+        let query_2 = relation::parse("select * from table_2").unwrap();
+        let aliases_query_2: Vec<ast::Ident> = vec!["aa".into(), "bb".into(), "cc".into()];
+        let query_3 = relation::parse("select * from table_3").unwrap();
+        let aliases_query_3: Vec<ast::Ident> = vec!["aaa".into(), "bbb".into(), "ccc".into()];
+
+        println!("query_1 = {}", query_1.to_string().blue());
+        let mut query_aliases_1 = QueryAliases::new();
+        let mut query_aliases_2 = QueryAliases::new();
+        let mut query_aliases_3 = QueryAliases::new();
+
+        query_aliases_2.insert(&query_1, None);
+        query_aliases_2.insert(&query_2, Some(&aliases_query_2));
+        println!("query_aliases_2: {query_aliases_2}");
+
+        query_aliases_3.insert(&query_3, Some(&aliases_query_3));
+        query_aliases_3.insert(&query_2, None);
+        query_aliases_3.insert(&query_1, None);
+        println!("query_aliases_3: {query_aliases_3}");
+
+        query_aliases_1.extend(query_aliases_2);
+        query_aliases_1.extend(query_aliases_3);
+        println!("query_aliases_1: {query_aliases_1}");
+    }
+
+    const COMPLEX_QUERY: &str = "
+        WITH
+        view_1 (a, b, c) as (select * from schema.table_1),
+        view_2 (aa, bb, cc) as (select * from view_1)
+        SELECT 2*new_a, b+1
+        FROM (
+            SELECT view_2.cc as new_a, right.d as b FROM (SELECT * FROM view_2) LEFT OUTER JOIN schema.table_2 as right ON view_1.a = table_2.id
+            );
+        select * from table_1;
+    ";
+
+    #[test]
+    fn test_query_aliases_visitor() {
+        let query = relation::parse(COMPLEX_QUERY).unwrap();
+        println!("Query = {}", query.to_string().blue());
+        let visitor = IntoQueryAliasesVisitor;
+        let query_aliases = query.accept(visitor);
+        println!("{}", query_aliases);
+    }
+}