qrlew/sql/
visitor.rs

1//! An Acceptor and Visitor implementation for ast::Query
2
3use crate::{
4    ast,
5    visitor::{self, Acceptor, Dependencies, Visited},
6};
7
8use std::iter::Iterator;
9
10/// A type to hold queries and relations with their aliases
11/// A Table with its dependencies (a table can be a simple name refenrencing another Table)
12#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
13enum TableWithAlias<'a> {
14    // Subqueries in the FROM or queries in CTEs
15    Query(&'a ast::Query, Option<&'a ast::TableAlias>),
16    // References to named objects, possibly aliased
17    Name(&'a ast::ObjectName, Option<&'a ast::TableAlias>),
18}
19
20/// A Wrapper, newtype pattern, to add conversion methods to an existing object
21pub struct TableWithJoins<'a>(&'a ast::TableWithJoins);
22
23impl<'a> TableWithJoins<'a> {
24    pub fn new(table_with_joins: &'a ast::TableWithJoins) -> Self {
25        TableWithJoins(table_with_joins)
26    }
27
28    fn tables_with_aliases(self) -> impl Iterator<Item = TableWithAlias<'a>> {
29        match &self.0.relation {
30            ast::TableFactor::Derived {
31                subquery, alias, ..
32            } => Some(TableWithAlias::Query(subquery.as_ref(), alias.as_ref())),
33            ast::TableFactor::Table { name, alias, .. } => {
34                Some(TableWithAlias::Name(name, alias.as_ref()))
35            }
36            _ => None,
37        }
38        .into_iter()
39        .chain(
40            // Then get the joins
41            self.0.joins.iter().filter_map(|join| match &join.relation {
42                ast::TableFactor::Derived {
43                    subquery, alias, ..
44                } => Some(TableWithAlias::Query(subquery.as_ref(), alias.as_ref())),
45                ast::TableFactor::Table { name, alias, .. } => {
46                    Some(TableWithAlias::Name(name, alias.as_ref()))
47                }
48                _ => None,
49            }),
50        )
51    }
52
53    /// Iterate over queries
54    pub fn queries(self) -> impl Iterator<Item = &'a ast::Query> {
55        self.tables_with_aliases().filter_map(|t| match t {
56            TableWithAlias::Query(q, _) => Some(q),
57            _ => None,
58        })
59    }
60
61    /// Iterate over names
62    pub fn names(self) -> impl Iterator<Item = &'a ast::ObjectName> {
63        self.tables_with_aliases().filter_map(|t| match t {
64            TableWithAlias::Name(n, _) => Some(n),
65            _ => None,
66        })
67    }
68}
69
70fn queries_from_set_expr<'a>(set_expr: &'a ast::SetExpr) -> Vec<&'a ast::Query> {
71    match set_expr {
72        ast::SetExpr::Select(select) => select
73            .from
74            .iter()
75            .flat_map(|table_with_joins| TableWithJoins(table_with_joins).queries())
76            .collect(),
77        ast::SetExpr::SetOperation { .. } => vec![],
78        ast::SetExpr::Values(_values) => todo!(),
79        _ => todo!(), // Not implemented
80    }
81}
82
83/// Implement the Acceptor trait
84impl<'a> Acceptor<'a> for ast::Query {
85    fn dependencies(&'a self) -> Dependencies<'a, Self> {
86        let mut dependencies = Dependencies::empty();
87        // Add CTEs subqueries
88        dependencies.extend(
89            self.with
90                .iter()
91                .flat_map(|with| with.cte_tables.iter().map(|cte| cte.query.as_ref())),
92        );
93        // Add subqueries from the body
94        dependencies.extend(queries_from_set_expr(self.body.as_ref()));
95        dependencies
96    }
97}
98
99/// A Visitor for Queries
100/// Dependencies are overrdden because references have to be unfold after a first pass to build the links (see IntoQueryNamesVisitor)
101pub trait Visitor<'a, T: Clone> {
102    /// The query visitor can affect the set of sub-queries it will visit by overriding this method
103    fn dependencies(&self, acceptor: &'a ast::Query) -> Dependencies<'a, ast::Query> {
104        acceptor.dependencies()
105    }
106    fn query(&self, query: &'a ast::Query, dependencies: Visited<'a, ast::Query, T>) -> T;
107}
108
109/// Unpack the visited queries of visitor::Visitor to ease the writing of Visitor
110impl<'a, T: Clone, V: Visitor<'a, T>> visitor::Visitor<'a, ast::Query, T> for V {
111    fn visit(&self, acceptor: &'a ast::Query, dependencies: Visited<'a, ast::Query, T>) -> T {
112        self.query(acceptor, dependencies)
113    }
114    /// We override the dependencies to allow visitor control on the acceptors' dependencies
115    fn dependencies(&self, acceptor: &'a ast::Query) -> Dependencies<'a, ast::Query> {
116        self.dependencies(acceptor)
117    }
118}