Skip to main content

partiql/sql/
env.rs

1use std::collections::VecDeque;
2
3use indexmap::IndexMap as Map;
4
5use crate::sql::Expr;
6use crate::sql::Selector;
7use crate::sql::SelectorNode;
8use crate::value::PqlValue;
9
10#[derive(Debug, Default, Clone)]
11pub struct Env {
12    data: Map<String, Expr>,
13}
14
15impl Env {
16    pub fn insert(&mut self, alias: &str, expr: &Expr) -> Option<Expr> {
17        self.data.insert(alias.to_string(), expr.to_owned())
18    }
19
20    pub fn insert_from_selector(&mut self, alias: &str, selector: &Selector) -> Option<Expr> {
21        let value = Expr::Selector(selector.to_owned());
22        self.insert(alias, &value)
23    }
24
25    pub fn insert_from_pqlval(&mut self, alias: &str, value: &PqlValue) -> Option<Expr> {
26        let value = Expr::Value(value.to_owned());
27        self.insert(alias, &value)
28    }
29
30    pub fn get(&self, key: &str) -> Option<Expr> {
31        self.data.get(key).map(|e| e.to_owned())
32    }
33
34    pub fn get_mut(&mut self, key: &str) -> Option<&mut Expr> {
35        self.data.get_mut(key)
36    }
37
38    pub fn get_by_selector(&self, selector: &Selector) -> Option<PqlValue> {
39        if let Some((head, tail)) = selector.split_first() {
40            if let Some(expr) = self.get(head.to_string().as_str()) {
41                match expr {
42                    Expr::Value(value) => {
43                        let v = if tail.data.len() > 0 {
44                            value.select_by_selector(&tail)
45                        } else {
46                            Some(value)
47                        };
48                        v
49                    }
50                    _ => todo!(),
51                }
52            } else {
53                todo!()
54            }
55        } else {
56            unreachable!()
57        }
58    }
59
60    pub fn get_as_selector(&self, key: &str) -> Option<Selector> {
61        match self.get(key) {
62            Some(Expr::Selector(selector)) => Some(selector),
63            _ => None,
64        }
65    }
66
67    fn rec_get_full_path(&self, selector: &Selector, trace_path: &mut Selector) {
68        if let Some((first, tail)) = selector.split_first() {
69            if let Some(alias_path) = self.get_as_selector(&first.to_string()) {
70                self.rec_get_full_path(&alias_path, trace_path)
71            } else {
72                (*trace_path)
73                    .data
74                    .push_back(SelectorNode::String(first.to_string()));
75            }
76            if tail.data.len() > 0 {
77                let tail_path = Selector::from(tail);
78                let mut vec_path = tail_path
79                    .to_vec()
80                    .into_iter()
81                    .map(|s| SelectorNode::String(s.to_string()))
82                    .collect::<VecDeque<_>>();
83                (*trace_path).data.append(&mut vec_path);
84            }
85        }
86    }
87
88    pub fn expand_fullpath_as_selector(&self, selector: &Selector) -> Selector {
89        let mut trace_path = Selector::default();
90
91        self.rec_get_full_path(selector, &mut trace_path);
92        trace_path
93    }
94
95    pub fn expand_fullpath(&self, expr: &Expr) -> Expr {
96        expr.expand_fullpath(self)
97    }
98}
99
100#[cfg(test)]
101mod tests {
102    use std::str::FromStr;
103
104    use super::Env;
105    use crate::planner::Drain;
106    use crate::sql::Expr;
107    use crate::sql::Field;
108    use crate::sql::Sql;
109
110    #[test]
111    fn get_full_path() -> anyhow::Result<()> {
112        let sql = Sql::from_str(
113            r#"
114SELECT
115  e.name AS employeeName, p.name AS projectName
116FROM
117  hr.employeesNest AS e, e.projects AS p
118        "#,
119        )?;
120
121        let mut env = Env::default();
122        Drain(sql.from_clause).execute(&mut env);
123
124        assert_eq!(
125            env.expand_fullpath(&Field::from_str("e.name AS employeeName")?.expr)
126                .to_string(),
127            "hr.employeesNest.name",
128        );
129
130        assert_eq!(
131            env.expand_fullpath(&Field::from_str("p.name AS projectName")?.expr)
132                .to_string(),
133            "hr.employeesNest.projects.name",
134        );
135
136        Ok(())
137    }
138
139    #[test]
140    fn test_update_env() -> anyhow::Result<()> {
141        let mut env = Env::default();
142        env.insert("name", &Expr::from("Alice"));
143
144        if let Some(name) = env.get_mut("name") {
145            *name = Expr::from("Bob");
146        }
147
148        assert_eq!(env.get("name"), Some(Expr::from("Bob")));
149
150        Ok(())
151    }
152}