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}