1use {
2 super::context::Context,
3 crate::{
4 ast::{ColumnDef, ColumnUniqueOption, Expr, Function, Query, TableAlias, TableFactor},
5 data::Schema,
6 },
7 std::sync::Arc,
8};
9
10pub trait Planner<'a> {
11 fn get_schema(&self, name: &str) -> Option<&'a Schema>;
12
13 fn query(&self, outer_context: Option<Arc<Context<'a>>>, query: Query) -> Query;
14
15 fn subquery_expr(&self, outer_context: Option<Arc<Context<'a>>>, expr: Expr) -> Expr {
16 match expr {
17 Expr::Identifier(_)
18 | Expr::CompoundIdentifier { .. }
19 | Expr::Literal(_)
20 | Expr::TypedString { .. } => expr,
21 Expr::IsNull(expr) => Expr::IsNull(Box::new(self.subquery_expr(outer_context, *expr))),
22 Expr::IsNotNull(expr) => {
23 Expr::IsNotNull(Box::new(self.subquery_expr(outer_context, *expr)))
24 }
25 Expr::InList {
26 expr,
27 list,
28 negated,
29 } => {
30 let list = list
31 .into_iter()
32 .map(|expr| self.subquery_expr(outer_context.as_ref().map(Arc::clone), expr))
33 .collect();
34 let expr = Box::new(self.subquery_expr(outer_context, *expr));
35
36 Expr::InList {
37 expr,
38 list,
39 negated,
40 }
41 }
42 Expr::Subquery(query) => Expr::Subquery(Box::new(self.query(outer_context, *query))),
43 Expr::Exists { subquery, negated } => Expr::Exists {
44 subquery: Box::new(self.query(outer_context, *subquery)),
45 negated,
46 },
47 Expr::InSubquery {
48 expr,
49 subquery,
50 negated,
51 } => {
52 let expr =
53 Box::new(self.subquery_expr(outer_context.as_ref().map(Arc::clone), *expr));
54 let subquery = Box::new(self.query(outer_context, *subquery));
55
56 Expr::InSubquery {
57 expr,
58 subquery,
59 negated,
60 }
61 }
62 Expr::Between {
63 expr,
64 negated,
65 low,
66 high,
67 } => {
68 let expr =
69 Box::new(self.subquery_expr(outer_context.as_ref().map(Arc::clone), *expr));
70 let low =
71 Box::new(self.subquery_expr(outer_context.as_ref().map(Arc::clone), *low));
72 let high = Box::new(self.subquery_expr(outer_context, *high));
73
74 Expr::Between {
75 expr,
76 negated,
77 low,
78 high,
79 }
80 }
81 Expr::Like {
82 expr,
83 negated,
84 pattern,
85 } => {
86 let expr =
87 Box::new(self.subquery_expr(outer_context.as_ref().map(Arc::clone), *expr));
88 let pattern =
89 Box::new(self.subquery_expr(outer_context.as_ref().map(Arc::clone), *pattern));
90
91 Expr::Like {
92 expr,
93 negated,
94 pattern,
95 }
96 }
97 Expr::ILike {
98 expr,
99 negated,
100 pattern,
101 } => {
102 let expr =
103 Box::new(self.subquery_expr(outer_context.as_ref().map(Arc::clone), *expr));
104 let pattern =
105 Box::new(self.subquery_expr(outer_context.as_ref().map(Arc::clone), *pattern));
106
107 Expr::ILike {
108 expr,
109 negated,
110 pattern,
111 }
112 }
113 Expr::BinaryOp { left, op, right } => Expr::BinaryOp {
114 left: Box::new(self.subquery_expr(outer_context.as_ref().map(Arc::clone), *left)),
115 op,
116 right: Box::new(self.subquery_expr(outer_context, *right)),
117 },
118 Expr::UnaryOp { op, expr } => Expr::UnaryOp {
119 op,
120 expr: Box::new(self.subquery_expr(outer_context, *expr)),
121 },
122 Expr::Nested(expr) => Expr::Nested(Box::new(self.subquery_expr(outer_context, *expr))),
123 Expr::Case {
124 operand,
125 when_then,
126 else_result,
127 } => {
128 let operand = operand.map(|expr| {
129 Box::new(self.subquery_expr(outer_context.as_ref().map(Arc::clone), *expr))
130 });
131 let when_then = when_then
132 .into_iter()
133 .map(|(when, then)| {
134 let when = self.subquery_expr(outer_context.as_ref().map(Arc::clone), when);
135 let then = self.subquery_expr(outer_context.as_ref().map(Arc::clone), then);
136
137 (when, then)
138 })
139 .collect();
140 let else_result =
141 else_result.map(|expr| Box::new(self.subquery_expr(outer_context, *expr)));
142
143 Expr::Case {
144 operand,
145 when_then,
146 else_result,
147 }
148 }
149 Expr::ArrayIndex { obj, indexes } => {
150 let indexes = indexes
151 .into_iter()
152 .map(|expr| self.subquery_expr(outer_context.as_ref().map(Arc::clone), expr))
153 .collect();
154 let obj = Box::new(self.subquery_expr(outer_context, *obj));
155 Expr::ArrayIndex { obj, indexes }
156 }
157 Expr::Array { elem } => {
158 let elem = elem
159 .into_iter()
160 .map(|expr| self.subquery_expr(outer_context.as_ref().map(Arc::clone), expr))
161 .collect();
162 Expr::Array { elem }
163 }
164 Expr::Interval {
165 expr,
166 leading_field,
167 last_field,
168 } => Expr::Interval {
169 expr: Box::new(self.subquery_expr(outer_context, *expr)),
170 leading_field,
171 last_field,
172 },
173 Expr::Function(func) => match *func {
174 Function::Cast { expr, data_type } => Expr::Function(Box::new(Function::Cast {
175 expr: self.subquery_expr(outer_context, expr),
176 data_type,
177 })),
178 Function::Extract { field, expr } => Expr::Function(Box::new(Function::Extract {
179 field,
180 expr: self.subquery_expr(outer_context, expr),
181 })),
182 _ => Expr::Function(func),
183 },
184 Expr::Aggregate(_) => expr,
185 }
186 }
187
188 fn update_context(
189 &self,
190 next: Option<Arc<Context<'a>>>,
191 table_factor: &TableFactor,
192 ) -> Option<Arc<Context<'a>>> {
193 let (name, alias) = match table_factor {
194 TableFactor::Table { name, alias, .. } => {
195 let alias = alias.as_ref().map(|TableAlias { name, .. }| name.clone());
196
197 (name, alias)
198 }
199 TableFactor::Derived { .. }
200 | TableFactor::Series { .. }
201 | TableFactor::Dictionary { .. } => return next,
202 };
203
204 let column_defs = match self.get_schema(name) {
205 Some(Schema { column_defs, .. }) => column_defs,
206 None => return next,
207 };
208
209 let column_defs = match column_defs {
210 Some(column_defs) => column_defs,
211 None => return next,
212 };
213
214 let columns = column_defs
215 .iter()
216 .map(|ColumnDef { name, .. }| name.as_str())
217 .collect::<Vec<_>>();
218
219 let primary_key = column_defs
220 .iter()
221 .find_map(|ColumnDef { name, unique, .. }| {
222 (unique == &Some(ColumnUniqueOption { is_primary: true })).then_some(name.as_str())
223 });
224
225 let context = Context::new(
226 alias.unwrap_or_else(|| name.to_owned()),
227 columns,
228 primary_key,
229 next,
230 );
231 Some(Arc::new(context))
232 }
233}