1use std::{
8 collections::{HashMap, HashSet},
9 sync::{
10 atomic::{AtomicUsize, Ordering},
11 RwLock,
12 },
13};
14
15use vibesql_storage::Row;
16
17use super::{CacheStats, QuerySignature};
18use crate::schema::CombinedSchema;
19
20#[derive(Clone)]
22struct CachedResult {
23 rows: Vec<Row>,
25 schema: CombinedSchema,
27 tables: HashSet<String>,
29}
30
31pub struct QueryResultCache {
36 cache: RwLock<HashMap<QuerySignature, CachedResult>>,
37 max_size: usize,
38 hits: AtomicUsize,
39 misses: AtomicUsize,
40 evictions: AtomicUsize,
41}
42
43impl QueryResultCache {
44 pub fn new(max_size: usize) -> Self {
46 Self {
47 cache: RwLock::new(HashMap::new()),
48 max_size,
49 hits: AtomicUsize::new(0),
50 misses: AtomicUsize::new(0),
51 evictions: AtomicUsize::new(0),
52 }
53 }
54
55 pub fn get(&self, signature: &QuerySignature) -> Option<(Vec<Row>, CombinedSchema)> {
57 let cache = self.cache.read().unwrap();
58 if let Some(entry) = cache.get(signature) {
59 self.hits.fetch_add(1, Ordering::Relaxed);
60 Some((entry.rows.clone(), entry.schema.clone()))
61 } else {
62 self.misses.fetch_add(1, Ordering::Relaxed);
63 None
64 }
65 }
66
67 pub fn insert(
69 &self,
70 signature: QuerySignature,
71 rows: Vec<Row>,
72 schema: CombinedSchema,
73 tables: HashSet<String>,
74 ) {
75 let entry = CachedResult { rows, schema, tables };
76 let mut cache = self.cache.write().unwrap();
77
78 if cache.len() >= self.max_size {
80 if let Some(key) = cache.keys().next().cloned() {
81 cache.remove(&key);
82 self.evictions.fetch_add(1, Ordering::Relaxed);
83 }
84 }
85
86 cache.insert(signature, entry);
87 }
88
89 pub fn contains(&self, signature: &QuerySignature) -> bool {
91 self.cache.read().unwrap().contains_key(signature)
92 }
93
94 pub fn clear(&self) {
96 self.cache.write().unwrap().clear();
97 }
98
99 pub fn invalidate_table(&self, table: &str) {
103 let mut cache = self.cache.write().unwrap();
104 cache.retain(|_, entry| !entry.tables.iter().any(|t| t.eq_ignore_ascii_case(table)));
105 }
106
107 pub fn invalidate_all(&self) {
109 self.clear();
110 }
111
112 pub fn stats(&self) -> CacheStats {
114 let cache = self.cache.read().unwrap();
115 let hits = self.hits.load(Ordering::Relaxed);
116 let misses = self.misses.load(Ordering::Relaxed);
117 let total = hits + misses;
118 let hit_rate = if total > 0 { hits as f64 / total as f64 } else { 0.0 };
119
120 CacheStats {
121 hits,
122 misses,
123 evictions: self.evictions.load(Ordering::Relaxed),
124 size: cache.len(),
125 hit_rate,
126 }
127 }
128
129 pub fn max_size(&self) -> usize {
131 self.max_size
132 }
133}
134
135#[cfg(test)]
136mod tests {
137 use vibesql_types::SqlValue;
138
139 use super::*;
140
141 fn make_test_row(values: Vec<SqlValue>) -> Row {
142 Row::new(values)
143 }
144
145 fn make_test_schema() -> CombinedSchema {
146 use vibesql_catalog::{ColumnSchema, TableSchema};
147 use vibesql_types::DataType;
148
149 let columns = vec![
150 ColumnSchema {
151 name: "id".to_string(),
152 data_type: DataType::Integer,
153 nullable: false,
154 default_value: None,
155 generated_expr: None,
156 is_exact_integer_type: false,
157 collation: None,
158 },
159 ColumnSchema {
160 name: "name".to_string(),
161 data_type: DataType::Varchar { max_length: Some(255) },
162 nullable: true,
163 default_value: None,
164 generated_expr: None,
165 is_exact_integer_type: false,
166 collation: None,
167 },
168 ];
169
170 let schema = TableSchema::new("users".to_string(), columns);
171 CombinedSchema::from_table("users".to_string(), schema)
172 }
173
174 #[test]
175 fn test_cache_hit() {
176 let cache = QueryResultCache::new(10);
177 let sig = QuerySignature::from_sql("SELECT * FROM users");
178 let rows = vec![
179 make_test_row(vec![
180 SqlValue::Integer(1),
181 SqlValue::Varchar(arcstr::ArcStr::from("Alice")),
182 ]),
183 make_test_row(vec![
184 SqlValue::Integer(2),
185 SqlValue::Varchar(arcstr::ArcStr::from("Bob")),
186 ]),
187 ];
188 let schema = make_test_schema();
189 let mut tables = HashSet::new();
190 tables.insert("users".to_string());
191
192 cache.insert(sig.clone(), rows.clone(), schema.clone(), tables);
193 let result = cache.get(&sig);
194
195 assert!(result.is_some());
196 let (cached_rows, _cached_schema) = result.unwrap();
197 assert_eq!(cached_rows.len(), 2);
198 assert_eq!(cached_rows[0].values.len(), 2);
199 }
200
201 #[test]
202 fn test_cache_miss() {
203 let cache = QueryResultCache::new(10);
204 let sig = QuerySignature::from_sql("SELECT * FROM users");
205
206 let result = cache.get(&sig);
207 assert!(result.is_none());
208 }
209
210 #[test]
211 fn test_lru_eviction() {
212 let cache = QueryResultCache::new(2);
213 let schema = make_test_schema();
214 let rows = vec![make_test_row(vec![SqlValue::Integer(1)])];
215 let tables = HashSet::new();
216
217 let sig1 = QuerySignature::from_sql("SELECT * FROM users");
218 let sig2 = QuerySignature::from_sql("SELECT * FROM orders");
219 let sig3 = QuerySignature::from_sql("SELECT * FROM products");
220
221 cache.insert(sig1, rows.clone(), schema.clone(), tables.clone());
222 cache.insert(sig2, rows.clone(), schema.clone(), tables.clone());
223 assert_eq!(cache.stats().size, 2);
224
225 cache.insert(sig3, rows, schema, tables);
226 assert_eq!(cache.stats().size, 2);
227 assert_eq!(cache.stats().evictions, 1);
228 }
229
230 #[test]
231 fn test_cache_clear() {
232 let cache = QueryResultCache::new(10);
233 let sig = QuerySignature::from_sql("SELECT * FROM users");
234 let rows = vec![make_test_row(vec![SqlValue::Integer(1)])];
235 let schema = make_test_schema();
236 let tables = HashSet::new();
237
238 cache.insert(sig.clone(), rows, schema, tables);
239 assert!(cache.contains(&sig));
240
241 cache.clear();
242 assert!(!cache.contains(&sig));
243 }
244
245 #[test]
246 fn test_table_invalidation() {
247 let cache = QueryResultCache::new(10);
248 let sig = QuerySignature::from_sql("SELECT * FROM users WHERE id = 1");
249 let rows = vec![make_test_row(vec![SqlValue::Integer(1)])];
250 let schema = make_test_schema();
251 let mut tables = HashSet::new();
252 tables.insert("users".to_string());
253
254 cache.insert(sig.clone(), rows, schema, tables);
255 assert!(cache.contains(&sig));
256
257 cache.invalidate_table("users");
258 assert!(!cache.contains(&sig));
259 }
260
261 #[test]
262 fn test_table_invalidation_case_insensitive() {
263 let cache = QueryResultCache::new(10);
264 let sig = QuerySignature::from_sql("SELECT * FROM users");
265 let rows = vec![make_test_row(vec![SqlValue::Integer(1)])];
266 let schema = make_test_schema();
267 let mut tables = HashSet::new();
268 tables.insert("users".to_string());
269
270 cache.insert(sig.clone(), rows, schema, tables);
271 assert!(cache.contains(&sig));
272
273 cache.invalidate_table("USERS");
275 assert!(!cache.contains(&sig));
276 }
277
278 #[test]
279 fn test_cache_stats() {
280 let cache = QueryResultCache::new(10);
281 let sig = QuerySignature::from_sql("SELECT * FROM users");
282 let rows = vec![make_test_row(vec![SqlValue::Integer(1)])];
283 let schema = make_test_schema();
284 let tables = HashSet::new();
285
286 cache.insert(sig.clone(), rows, schema, tables);
287
288 cache.get(&sig);
290 cache.get(&sig);
291
292 let other_sig = QuerySignature::from_sql("SELECT * FROM orders");
294 cache.get(&other_sig);
295
296 let stats = cache.stats();
297 assert_eq!(stats.hits, 2);
298 assert_eq!(stats.misses, 1);
299 assert!((stats.hit_rate - 2.0 / 3.0).abs() < 0.01);
300 }
301}