qrlew/sql/
query_names.rs

1//! Data structures and visitor to collect object names and the queries they refer to.
2
3use crate::{
4    ast,
5    expr::Identifier,
6    hierarchy::Path,
7    sql::visitor::{TableWithJoins, Visitor},
8    visitor::Visited,
9};
10use colored::Colorize;
11use itertools::Itertools;
12use std::{
13    collections::BTreeMap,
14    fmt,
15    iter::Iterator,
16    ops::{Deref, DerefMut},
17};
18
19/// A mapping between an ObjectName in a Query and the Query referred by the Name (when available)
20#[derive(Clone, Debug, Hash, PartialEq, Eq)]
21pub struct QueryNames<'a>(BTreeMap<(&'a ast::Query, ast::ObjectName), Option<&'a ast::Query>>);
22
23impl<'a> QueryNames<'a> {
24    /// Build a new QueryNames object
25    pub fn new() -> Self {
26        QueryNames(BTreeMap::new())
27    }
28
29    /// Set all unresolved names
30    pub fn set(&mut self, name: ast::ObjectName, referred: &'a ast::Query) -> &mut Self {
31        for ((_, n), r) in self.iter_mut() {
32            if (*n == name) && r.is_none() {
33                *r = Some(referred);
34            }
35        }
36        self
37    }
38
39    /// Return all the names and referred queries in a query
40    pub fn name_referred(
41        &self,
42        query: &'a ast::Query,
43    ) -> impl Iterator<Item = (&ast::ObjectName, &'a ast::Query)> {
44        self.iter()
45            .filter_map(move |((q, n), r)| if *q == query { Some((n, (*r)?)) } else { None })
46    }
47}
48
49impl<'a> Deref for QueryNames<'a> {
50    type Target = BTreeMap<(&'a ast::Query, ast::ObjectName), Option<&'a ast::Query>>;
51
52    fn deref(&self) -> &Self::Target {
53        &self.0
54    }
55}
56
57impl<'a> DerefMut for QueryNames<'a> {
58    fn deref_mut(&mut self) -> &mut Self::Target {
59        &mut self.0
60    }
61}
62
63impl<'a> IntoIterator for QueryNames<'a> {
64    type Item =
65        <BTreeMap<(&'a ast::Query, ast::ObjectName), Option<&'a ast::Query>> as IntoIterator>::Item;
66    type IntoIter = <BTreeMap<(&'a ast::Query, ast::ObjectName), Option<&'a ast::Query>> as IntoIterator>::IntoIter;
67
68    fn into_iter(self) -> Self::IntoIter {
69        self.0.into_iter()
70    }
71}
72
73impl<'a> fmt::Display for QueryNames<'a> {
74    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
75        write!(
76            f,
77            "Query Names\n{}",
78            self.0
79                .iter()
80                .map(|((q, n), r)| match r {
81                    Some(r) => format!(
82                        "{} | {} -> {}",
83                        format!("{q}").blue(),
84                        format!("{n}").red(),
85                        format!("{r}").green()
86                    ),
87                    None => format!(
88                        "{} | {} -> {}",
89                        format!("{q}").blue(),
90                        format!("{n}").red(),
91                        format!("?").bold().green()
92                    ),
93                })
94                .join(",\n")
95        )
96    }
97}
98
99/// A visitor to build a reference free query tree
100#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
101pub struct IntoQueryNamesVisitor;
102
103fn names_from_set_expr<'a>(set_expr: &'a ast::SetExpr) -> Vec<&'a ast::ObjectName> {
104    match set_expr {
105        ast::SetExpr::Select(select) => select
106            .from
107            .iter()
108            .flat_map(|table_with_joins| TableWithJoins::new(table_with_joins).names())
109            .collect(),
110        ast::SetExpr::SetOperation { left, right, .. } => names_from_set_expr(left.as_ref())
111            .into_iter()
112            .chain(names_from_set_expr(right.as_ref()))
113            .collect(),
114        _ => todo!(),
115    }
116}
117
118impl<'a> Visitor<'a, QueryNames<'a>> for IntoQueryNamesVisitor {
119    fn query(
120        &self,
121        query: &'a ast::Query,
122        dependencies: Visited<'a, ast::Query, QueryNames<'a>>,
123    ) -> QueryNames<'a> {
124        let mut query_names = QueryNames::new();
125        // Add all elemdnts in dependencies
126        for (_, dependency) in dependencies {
127            query_names.extend(dependency);
128        }
129        // Add reference elements from current query
130        for name in names_from_set_expr(query.body.as_ref()) {
131            query_names.insert((query, name.clone()), None);
132        }
133        // Set names
134        if let Some(with) = &query.with {
135            for cte in &with.cte_tables {
136                query_names.set(
137                    ast::ObjectName(vec![cte.alias.name.clone()]),
138                    cte.query.as_ref(),
139                );
140            }
141        }
142        query_names
143    }
144}
145
146/// Implement conversion from ObjectName to Path
147impl Path for ast::ObjectName {
148    fn path(self) -> Vec<String> {
149        self.0.path()
150    }
151}
152
153/// Implement conversion to Identifier
154impl From<&ast::ObjectName> for Identifier {
155    fn from(value: &ast::ObjectName) -> Self {
156        value.0.iter().map(|i| i.value.clone()).collect()
157    }
158}
159
160#[cfg(test)]
161mod tests {
162    use colored::Colorize;
163
164    use super::*;
165    use crate::{sql::relation, visitor::Acceptor as _};
166
167    #[test]
168    fn test_query_names() {
169        let query_1 = relation::parse("select * from table_1").unwrap();
170        let query_2 = relation::parse("select * from table_2").unwrap();
171        let query_3 = relation::parse("select * from table_3").unwrap();
172        let query_4 = relation::parse("select * from table_4").unwrap();
173        let name_1 = ast::ObjectName(vec!["name_1".into()]);
174        let name_2 = ast::ObjectName(vec!["name_2".into()]);
175        let name_3 = ast::ObjectName(vec!["name_3".into()]);
176        println!("query_1 = {}", query_1.to_string().blue());
177        let mut query_names_1 = QueryNames::new();
178        let mut query_names_2 = QueryNames::new();
179        let mut query_names_3 = QueryNames::new();
180        query_names_2.insert((&query_1, name_1), None);
181        query_names_2.insert((&query_1, name_2), Some(&query_2));
182        query_names_3.insert((&query_1, name_3.clone()), Some(&query_3));
183        query_names_3.insert((&query_2, name_3.clone()), None);
184        query_names_3.insert((&query_3, name_3.clone()), None);
185        query_names_1.extend(query_names_2);
186        query_names_1.extend(query_names_3);
187        println!("BEFORE: {query_names_1}");
188        query_names_1.set(name_3.clone(), &query_4);
189        println!("AFTER: {query_names_1}");
190    }
191
192    const COMPLEX_QUERY: &str = "
193        with view_1 as (select * from schema.table_1),
194        view_2 as (select * from view_1)
195        select 2*a, b+1 from (select view_2.c as a, right.d as b from (select * from view_2) LEFT OUTER JOIN schema.table_2 as right ON view_1.id = table_2.id);
196        select * from table_1;
197    ";
198
199    #[test]
200    fn test_query_names_visitor() {
201        let query = relation::parse(COMPLEX_QUERY).unwrap();
202        println!("Query = {}", query.to_string().blue());
203        let visitor = IntoQueryNamesVisitor;
204        let query_names = query.accept(visitor);
205        println!("{}", query_names);
206    }
207}
qrlew/sql/query_names.rs

qrlew/sql/
query_names.rs
