use std::{
collections::{HashMap, HashSet},
sync::{
atomic::{AtomicUsize, Ordering},
RwLock,
},
};
use super::QuerySignature;
#[derive(Clone, Debug)]
pub struct CacheStats {
pub hits: usize,
pub misses: usize,
pub evictions: usize,
pub size: usize,
pub hit_rate: f64,
}
#[derive(Clone)]
struct CachedEntry {
sql: String,
tables: HashSet<String>,
}
pub struct QueryPlanCache {
cache: RwLock<HashMap<QuerySignature, CachedEntry>>,
max_size: usize,
hits: AtomicUsize,
misses: AtomicUsize,
evictions: AtomicUsize,
}
impl QueryPlanCache {
pub fn new(max_size: usize) -> Self {
Self {
cache: RwLock::new(HashMap::new()),
max_size,
hits: AtomicUsize::new(0),
misses: AtomicUsize::new(0),
evictions: AtomicUsize::new(0),
}
}
pub fn get(&self, signature: &QuerySignature) -> Option<String> {
let cache = self.cache.read().unwrap();
if let Some(entry) = cache.get(signature) {
self.hits.fetch_add(1, Ordering::Relaxed);
Some(entry.sql.clone())
} else {
self.misses.fetch_add(1, Ordering::Relaxed);
None
}
}
pub fn insert(&self, signature: QuerySignature, sql: String) {
self.insert_with_tables(signature, sql, HashSet::new());
}
pub fn insert_with_tables(
&self,
signature: QuerySignature,
sql: String,
tables: HashSet<String>,
) {
let entry = CachedEntry { sql, tables };
let mut cache = self.cache.write().unwrap();
if cache.len() >= self.max_size {
if let Some(key) = cache.keys().next().cloned() {
cache.remove(&key);
self.evictions.fetch_add(1, Ordering::Relaxed);
}
}
cache.insert(signature, entry);
}
pub fn contains(&self, signature: &QuerySignature) -> bool {
self.cache.read().unwrap().contains_key(signature)
}
pub fn clear(&self) {
self.cache.write().unwrap().clear();
}
pub fn invalidate_table(&self, table: &str) {
let mut cache = self.cache.write().unwrap();
cache.retain(|_, entry| !entry.tables.iter().any(|t| t.eq_ignore_ascii_case(table)));
}
pub fn stats(&self) -> CacheStats {
let cache = self.cache.read().unwrap();
let hits = self.hits.load(Ordering::Relaxed);
let misses = self.misses.load(Ordering::Relaxed);
let total = hits + misses;
let hit_rate = if total > 0 { hits as f64 / total as f64 } else { 0.0 };
CacheStats {
hits,
misses,
evictions: self.evictions.load(Ordering::Relaxed),
size: cache.len(),
hit_rate,
}
}
pub fn max_size(&self) -> usize {
self.max_size
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_cache_hit() {
let cache = QueryPlanCache::new(10);
let sig = QuerySignature::from_sql("SELECT * FROM users");
let sql = "select * from users".to_string();
cache.insert(sig.clone(), sql.clone());
let result = cache.get(&sig);
assert!(result.is_some());
assert_eq!(result.unwrap(), sql);
}
#[test]
fn test_cache_miss() {
let cache = QueryPlanCache::new(10);
let sig = QuerySignature::from_sql("SELECT * FROM users");
let result = cache.get(&sig);
assert!(result.is_none());
}
#[test]
fn test_lru_eviction() {
let cache = QueryPlanCache::new(2);
let sig1 = QuerySignature::from_sql("SELECT * FROM users");
let sig2 = QuerySignature::from_sql("SELECT * FROM orders");
let sig3 = QuerySignature::from_sql("SELECT * FROM products");
cache.insert(sig1, "select * from users".to_string());
cache.insert(sig2, "select * from orders".to_string());
assert_eq!(cache.stats().size, 2);
cache.insert(sig3, "select * from products".to_string());
assert_eq!(cache.stats().size, 2);
assert_eq!(cache.stats().evictions, 1);
}
#[test]
fn test_cache_clear() {
let cache = QueryPlanCache::new(10);
let sig = QuerySignature::from_sql("SELECT * FROM users");
cache.insert(sig.clone(), "select * from users".to_string());
assert!(cache.contains(&sig));
cache.clear();
assert!(!cache.contains(&sig));
}
#[test]
fn test_table_invalidation() {
let cache = QueryPlanCache::new(10);
let sig = QuerySignature::from_sql("SELECT * FROM users WHERE id = 1");
let mut tables = std::collections::HashSet::new();
tables.insert("users".to_string());
cache.insert_with_tables(
sig.clone(),
"select * from users where id = 1".to_string(),
tables,
);
assert!(cache.contains(&sig));
cache.invalidate_table("users");
assert!(!cache.contains(&sig));
}
#[test]
fn test_cache_stats() {
let cache = QueryPlanCache::new(10);
let sig = QuerySignature::from_sql("SELECT * FROM users");
cache.insert(sig.clone(), "select * from users".to_string());
cache.get(&sig);
cache.get(&sig);
let other_sig = QuerySignature::from_sql("SELECT * FROM orders");
cache.get(&other_sig);
let stats = cache.stats();
assert_eq!(stats.hits, 2);
assert_eq!(stats.misses, 1);
assert!((stats.hit_rate - 2.0 / 3.0).abs() < 0.01);
}
}