1use alloc::collections::{BTreeMap, VecDeque};
20use alloc::string::String;
21use alloc::vec::Vec;
22
23use spg_sql::ast::{Expr, FromClause, FromJoin, SelectItem, SelectStatement, Statement, TableRef};
24use spg_storage::ColumnSchema;
25
26pub(crate) const PLAN_CACHE_MAX_ENTRIES: usize = 256;
31
32#[derive(Debug, Clone)]
36pub struct PreparedPlan {
37 pub stmt: Statement,
38 pub statistics_version: u64,
42 pub source_tables: Vec<String>,
45 pub describe_columns: Vec<ColumnSchema>,
48}
49
50#[derive(Debug, Clone)]
51pub struct PlanCache {
52 entries: BTreeMap<String, PreparedPlan>,
56 lru: VecDeque<String>,
60 max_entries: usize,
65}
66
67impl Default for PlanCache {
68 fn default() -> Self {
69 Self {
70 entries: BTreeMap::new(),
71 lru: VecDeque::new(),
72 max_entries: PLAN_CACHE_MAX_ENTRIES,
73 }
74 }
75}
76
77impl PlanCache {
78 pub fn new() -> Self {
79 Self::default()
80 }
81
82 pub fn set_max_entries(&mut self, n: usize) {
88 self.max_entries = n.max(1).min(PLAN_CACHE_MAX_ENTRIES);
89 }
90
91 pub fn max_entries(&self) -> usize {
92 self.max_entries
93 }
94
95 pub fn len(&self) -> usize {
96 self.entries.len()
97 }
98
99 pub fn is_empty(&self) -> bool {
100 self.entries.is_empty()
101 }
102
103 pub fn get_snapshot(&self, sql: &str) -> Option<&PreparedPlan> {
107 self.entries.get(sql)
108 }
109
110 pub fn get(&mut self, sql: &str) -> Option<&PreparedPlan> {
113 if !self.entries.contains_key(sql) {
114 return None;
115 }
116 if let Some(idx) = self.lru.iter().position(|k| k == sql) {
117 let key = self.lru.remove(idx).expect("idx came from position()");
118 self.lru.push_back(key);
119 }
120 self.entries.get(sql)
121 }
122
123 pub fn insert(&mut self, sql: String, plan: PreparedPlan) {
126 if self.entries.contains_key(&sql) {
127 if let Some(idx) = self.lru.iter().position(|k| k == &sql) {
128 let key = self.lru.remove(idx).expect("idx came from position()");
129 self.lru.push_back(key);
130 }
131 self.entries.insert(sql, plan);
132 return;
133 }
134 if self.entries.len() >= self.max_entries {
135 if let Some(oldest) = self.lru.pop_front() {
136 self.entries.remove(&oldest);
137 }
138 }
139 self.lru.push_back(sql.clone());
140 self.entries.insert(sql, plan);
141 }
142
143 pub fn clear(&mut self) {
144 self.entries.clear();
145 self.lru.clear();
146 }
147
148 pub fn evict(&mut self, sql: &str) -> Option<PreparedPlan> {
151 let plan = self.entries.remove(sql)?;
152 if let Some(idx) = self.lru.iter().position(|k| k == sql) {
153 self.lru.remove(idx);
154 }
155 Some(plan)
156 }
157
158 pub fn evict_referencing(&mut self, table: &str) -> usize {
161 let to_evict: Vec<String> = self
162 .entries
163 .iter()
164 .filter_map(|(k, p)| {
165 if p.source_tables.iter().any(|t| t == table) {
166 Some(k.clone())
167 } else {
168 None
169 }
170 })
171 .collect();
172 let n = to_evict.len();
173 for k in to_evict {
174 self.entries.remove(&k);
175 if let Some(idx) = self.lru.iter().position(|x| x == &k) {
176 self.lru.remove(idx);
177 }
178 }
179 n
180 }
181}
182
183pub fn collect_source_tables(stmt: &Statement) -> Vec<String> {
188 let mut out: Vec<String> = Vec::new();
189 match stmt {
190 Statement::Select(s) => collect_from_select(s, &mut out),
191 Statement::Insert(s) => push_unique(&mut out, &s.table),
192 Statement::Update(s) => {
193 push_unique(&mut out, &s.table);
194 if let Some(w) = &s.where_ {
195 collect_expr(w, &mut out);
196 }
197 }
198 Statement::Delete(s) => {
199 push_unique(&mut out, &s.table);
200 if let Some(w) = &s.where_ {
201 collect_expr(w, &mut out);
202 }
203 }
204 Statement::Explain(inner) => {
205 collect_from_select(&inner.inner, &mut out);
206 }
207 _ => {}
208 }
209 out.sort();
210 out.dedup();
211 out
212}
213
214fn collect_from_select(s: &SelectStatement, out: &mut Vec<String>) {
215 if let Some(from) = &s.from {
216 collect_from_clause(from, out);
217 }
218 if let Some(w) = &s.where_ {
219 collect_expr(w, out);
220 }
221 if let Some(h) = &s.having {
222 collect_expr(h, out);
223 }
224 for item in &s.items {
225 if let SelectItem::Expr { expr, .. } = item {
226 collect_expr(expr, out);
227 }
228 }
229 for (_, peer) in &s.unions {
230 collect_from_select(peer, out);
231 }
232}
233
234fn collect_from_clause(from: &FromClause, out: &mut Vec<String>) {
235 collect_table_ref(&from.primary, out);
236 for j in &from.joins {
237 collect_from_join(j, out);
238 }
239}
240
241fn collect_from_join(j: &FromJoin, out: &mut Vec<String>) {
242 collect_table_ref(&j.table, out);
243 if let Some(on) = &j.on {
244 collect_expr(on, out);
245 }
246}
247
248fn collect_table_ref(t: &TableRef, out: &mut Vec<String>) {
249 push_unique(out, &t.name);
250}
251
252fn collect_expr(e: &Expr, out: &mut Vec<String>) {
253 match e {
254 Expr::ScalarSubquery(inner) => collect_from_select(inner, out),
255 Expr::Exists { subquery, .. } => collect_from_select(subquery, out),
256 Expr::InSubquery { expr, subquery, .. } => {
257 collect_expr(expr, out);
258 collect_from_select(subquery, out);
259 }
260 Expr::Binary { lhs, rhs, .. } => {
261 collect_expr(lhs, out);
262 collect_expr(rhs, out);
263 }
264 Expr::Unary { expr, .. } => collect_expr(expr, out),
265 Expr::Cast { expr, .. } => collect_expr(expr, out),
266 Expr::IsNull { expr, .. } => collect_expr(expr, out),
267 Expr::Like { expr, pattern, .. } => {
268 collect_expr(expr, out);
269 collect_expr(pattern, out);
270 }
271 Expr::FunctionCall { args, .. } => {
272 for a in args {
273 collect_expr(a, out);
274 }
275 }
276 Expr::WindowFunction {
277 args,
278 partition_by,
279 order_by,
280 ..
281 } => {
282 for a in args {
283 collect_expr(a, out);
284 }
285 for p in partition_by {
286 collect_expr(p, out);
287 }
288 for (o, _) in order_by {
289 collect_expr(o, out);
290 }
291 }
292 Expr::Extract { source, .. } => collect_expr(source, out),
293 Expr::Array(items) => {
294 for elem in items {
295 collect_expr(elem, out);
296 }
297 }
298 Expr::ArraySubscript { target, index } => {
299 collect_expr(target, out);
300 collect_expr(index, out);
301 }
302 Expr::AnyAll { expr, array, .. } => {
303 collect_expr(expr, out);
304 collect_expr(array, out);
305 }
306 Expr::Literal(_) | Expr::Column(_) | Expr::Placeholder(_) => {}
307 }
308}
309
310fn push_unique(out: &mut Vec<String>, s: &str) {
311 if !out.iter().any(|x| x == s) {
312 out.push(String::from(s));
313 }
314}
315
316#[cfg(test)]
319mod tests {
320 use super::*;
321 use alloc::string::ToString;
322 use spg_sql::parser::parse_statement;
323
324 fn dummy_plan(version: u64, tables: &[&str]) -> PreparedPlan {
325 let stmt = parse_statement("SELECT 1").expect("trivial SELECT parses");
326 PreparedPlan {
327 stmt,
328 statistics_version: version,
329 source_tables: tables.iter().map(|s| s.to_string()).collect(),
330 describe_columns: Vec::new(),
331 }
332 }
333
334 #[test]
335 fn new_cache_is_empty() {
336 let cache = PlanCache::new();
337 assert!(cache.is_empty());
338 assert_eq!(cache.len(), 0);
339 }
340
341 #[test]
342 fn insert_then_get_returns_the_plan() {
343 let mut cache = PlanCache::new();
344 cache.insert("SELECT 1".into(), dummy_plan(0, &["t"]));
345 assert_eq!(cache.len(), 1);
346 let plan = cache.get("SELECT 1").expect("hit");
347 assert_eq!(plan.source_tables, alloc::vec!["t".to_string()]);
348 }
349
350 #[test]
351 fn miss_returns_none() {
352 let mut cache = PlanCache::new();
353 cache.insert("SELECT 1".into(), dummy_plan(0, &[]));
354 assert!(cache.get("SELECT 2").is_none());
355 }
356
357 #[test]
358 fn replace_overwrites_existing_entry() {
359 let mut cache = PlanCache::new();
360 cache.insert("SELECT 1".into(), dummy_plan(1, &["a"]));
361 cache.insert("SELECT 1".into(), dummy_plan(2, &["b"]));
362 assert_eq!(cache.len(), 1);
363 let plan = cache.get("SELECT 1").expect("hit");
364 assert_eq!(plan.statistics_version, 2);
365 }
366
367 #[test]
368 fn lru_evicts_oldest_at_cap() {
369 let mut cache = PlanCache::new();
370 for i in 0..PLAN_CACHE_MAX_ENTRIES {
371 cache.insert(alloc::format!("SELECT {i}"), dummy_plan(i as u64, &[]));
372 }
373 assert_eq!(cache.len(), PLAN_CACHE_MAX_ENTRIES);
374 cache.insert("SELECT new".into(), dummy_plan(999, &[]));
375 assert_eq!(cache.len(), PLAN_CACHE_MAX_ENTRIES);
376 assert!(cache.get("SELECT 0").is_none());
377 assert!(cache.get("SELECT new").is_some());
378 }
379
380 #[test]
381 fn get_promotes_lru_position() {
382 let mut cache = PlanCache::new();
383 cache.insert("a".into(), dummy_plan(0, &[]));
384 cache.insert("b".into(), dummy_plan(0, &[]));
385 cache.insert("c".into(), dummy_plan(0, &[]));
386 let _ = cache.get("a");
388 for i in 0..(PLAN_CACHE_MAX_ENTRIES - 3) {
391 cache.insert(alloc::format!("filler{i}"), dummy_plan(0, &[]));
392 }
393 cache.insert("trigger".into(), dummy_plan(0, &[]));
394 assert!(
395 cache.get("a").is_some(),
396 "a was MRU after get(); should survive"
397 );
398 assert!(cache.get("b").is_none(), "b should be evicted");
399 }
400
401 #[test]
402 fn clear_drops_everything() {
403 let mut cache = PlanCache::new();
404 cache.insert("a".into(), dummy_plan(0, &[]));
405 cache.insert("b".into(), dummy_plan(0, &[]));
406 cache.clear();
407 assert!(cache.is_empty());
408 assert!(cache.get("a").is_none());
409 }
410
411 #[test]
412 fn evict_referencing_drops_only_matching_plans() {
413 let mut cache = PlanCache::new();
414 cache.insert("a".into(), dummy_plan(0, &["users"]));
415 cache.insert("b".into(), dummy_plan(0, &["orders"]));
416 cache.insert("c".into(), dummy_plan(0, &["users", "orders"]));
417 let n = cache.evict_referencing("users");
418 assert_eq!(n, 2);
419 assert!(cache.get("a").is_none());
420 assert!(cache.get("b").is_some());
421 assert!(cache.get("c").is_none());
422 }
423
424 #[test]
425 fn collect_source_tables_from_simple_select() {
426 let stmt = parse_statement("SELECT a, b FROM t1 WHERE x = 1").expect("parses");
427 let tables = collect_source_tables(&stmt);
428 assert_eq!(tables, alloc::vec!["t1".to_string()]);
429 }
430
431 #[test]
432 fn collect_source_tables_from_join() {
433 let stmt =
434 parse_statement("SELECT * FROM t1 JOIN t2 ON t1.a = t2.b JOIN t3 ON t2.c = t3.d")
435 .expect("parses");
436 let tables = collect_source_tables(&stmt);
437 assert_eq!(
438 tables,
439 alloc::vec!["t1".to_string(), "t2".to_string(), "t3".to_string()]
440 );
441 }
442
443 #[test]
444 fn collect_source_tables_dedupes_self_join() {
445 let stmt = parse_statement("SELECT * FROM t1 a JOIN t1 b ON a.x = b.y").expect("parses");
446 let tables = collect_source_tables(&stmt);
447 assert_eq!(tables, alloc::vec!["t1".to_string()]);
448 }
449}