use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::collections::{HashMap, VecDeque};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tracing::{debug, info};
#[derive(Debug, Clone)]
pub struct StatementCacheConfig {
pub max_size: usize,
pub ttl_secs: u64,
pub enable_fingerprinting: bool,
}
impl Default for StatementCacheConfig {
fn default() -> Self {
Self {
max_size: 1000,
ttl_secs: 3600, enable_fingerprinting: true,
}
}
}
#[derive(Debug, Clone)]
struct CachedStatement {
sql: String,
#[allow(dead_code)]
fingerprint: String,
cached_at: Instant,
reuse_count: u64,
last_accessed: Instant,
}
impl CachedStatement {
fn new(sql: String, fingerprint: String) -> Self {
let now = Instant::now();
Self {
sql,
fingerprint,
cached_at: now,
reuse_count: 0,
last_accessed: now,
}
}
fn is_expired(&self, ttl: Duration) -> bool {
self.cached_at.elapsed() > ttl
}
fn touch(&mut self) {
self.reuse_count += 1;
self.last_accessed = Instant::now();
}
}
#[derive(Debug, Clone)]
pub struct StatementCache {
config: StatementCacheConfig,
cache: Arc<RwLock<HashMap<String, CachedStatement>>>,
lru_queue: Arc<RwLock<VecDeque<String>>>,
stats: Arc<RwLock<CacheStats>>,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct CacheStats {
pub hits: u64,
pub misses: u64,
pub evictions: u64,
pub expirations: u64,
pub cached_statements: usize,
}
impl CacheStats {
pub fn hit_rate(&self) -> f64 {
let total = self.hits + self.misses;
if total == 0 {
0.0
} else {
self.hits as f64 / total as f64
}
}
pub fn reset(&mut self) {
self.hits = 0;
self.misses = 0;
self.evictions = 0;
self.expirations = 0;
}
}
impl Default for StatementCache {
fn default() -> Self {
Self::new(StatementCacheConfig::default())
}
}
impl StatementCache {
pub fn new(config: StatementCacheConfig) -> Self {
Self {
config,
cache: Arc::new(RwLock::new(HashMap::new())),
lru_queue: Arc::new(RwLock::new(VecDeque::new())),
stats: Arc::new(RwLock::new(CacheStats::default())),
}
}
pub fn fingerprint(&self, sql: &str) -> String {
if !self.config.enable_fingerprinting {
return sql.to_string();
}
let normalized = normalize_sql(sql);
let mut hasher = Sha256::new();
hasher.update(normalized.as_bytes());
let result = hasher.finalize();
hex::encode(result.as_slice())
}
pub fn get(&self, sql: &str) -> Option<String> {
let fingerprint = self.fingerprint(sql);
let mut cache = self.cache.write();
let mut stats = self.stats.write();
if let Some(stmt) = cache.get_mut(&fingerprint) {
if stmt.is_expired(Duration::from_secs(self.config.ttl_secs)) {
cache.remove(&fingerprint);
self.remove_from_lru(&fingerprint);
stats.expirations += 1;
stats.misses += 1;
debug!(fingerprint = %fingerprint, "Statement expired");
return None;
}
stmt.touch();
stats.hits += 1;
self.update_lru(&fingerprint);
debug!(
fingerprint = %fingerprint,
reuse_count = stmt.reuse_count,
"Statement cache hit"
);
Some(stmt.sql.clone())
} else {
stats.misses += 1;
debug!(fingerprint = %fingerprint, "Statement cache miss");
None
}
}
pub fn put(&self, sql: String) {
let fingerprint = self.fingerprint(&sql);
let mut cache = self.cache.write();
let mut stats = self.stats.write();
while cache.len() >= self.config.max_size {
if let Some(oldest) = self.lru_queue.write().pop_back() {
cache.remove(&oldest);
stats.evictions += 1;
debug!(fingerprint = %oldest, "Statement evicted from cache");
} else {
break;
}
}
let stmt = CachedStatement::new(sql, fingerprint.clone());
cache.insert(fingerprint.clone(), stmt);
self.lru_queue.write().push_front(fingerprint.clone());
stats.cached_statements = cache.len();
debug!(
fingerprint = %fingerprint,
cache_size = cache.len(),
"Statement cached"
);
}
pub fn clear(&self) {
self.cache.write().clear();
self.lru_queue.write().clear();
self.stats.write().cached_statements = 0;
info!("Statement cache cleared");
}
pub fn stats(&self) -> CacheStats {
self.stats.read().clone()
}
pub fn remove_expired(&self) {
let ttl = Duration::from_secs(self.config.ttl_secs);
let mut cache = self.cache.write();
let mut stats = self.stats.write();
let expired: Vec<String> = cache
.iter()
.filter(|(_, stmt)| stmt.is_expired(ttl))
.map(|(fp, _)| fp.clone())
.collect();
for fp in &expired {
cache.remove(fp);
self.remove_from_lru(fp);
stats.expirations += 1;
}
stats.cached_statements = cache.len();
if !expired.is_empty() {
info!(count = expired.len(), "Removed expired statements");
}
}
fn update_lru(&self, fingerprint: &str) {
let mut queue = self.lru_queue.write();
if let Some(pos) = queue.iter().position(|fp| fp == fingerprint) {
queue.remove(pos);
}
queue.push_front(fingerprint.to_string());
}
fn remove_from_lru(&self, fingerprint: &str) {
let mut queue = self.lru_queue.write();
if let Some(pos) = queue.iter().position(|fp| fp == fingerprint) {
queue.remove(pos);
}
}
}
fn normalize_sql(sql: &str) -> String {
let normalized = sql
.split_whitespace()
.collect::<Vec<_>>()
.join(" ")
.to_lowercase();
replace_literals(&normalized)
}
static RE_LITERAL_STRING: std::sync::LazyLock<regex::Regex> = std::sync::LazyLock::new(|| {
regex::Regex::new(r"'[^']*'")
.expect("static regex pattern for SQL string literals is always valid")
});
static RE_LITERAL_NUMBER: std::sync::LazyLock<regex::Regex> = std::sync::LazyLock::new(|| {
regex::Regex::new(r"\b\d+\b")
.expect("static regex pattern for SQL numeric literals is always valid")
});
fn replace_literals(sql: &str) -> String {
let sql = RE_LITERAL_STRING.replace_all(sql, "?");
let sql = RE_LITERAL_NUMBER.replace_all(&sql, "?");
sql.to_string()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_statement_cache_config_default() {
let config = StatementCacheConfig::default();
assert_eq!(config.max_size, 1000);
assert_eq!(config.ttl_secs, 3600);
assert!(config.enable_fingerprinting);
}
#[test]
fn test_fingerprint_generation() {
let cache = StatementCache::default();
let sql1 = "SELECT * FROM users WHERE id = 1";
let sql2 = "SELECT * FROM users WHERE id = 2";
let fp1 = cache.fingerprint(sql1);
let fp2 = cache.fingerprint(sql2);
assert_eq!(fp1, fp2);
}
#[test]
fn test_cache_put_and_get() {
let cache = StatementCache::default();
let sql = "SELECT * FROM users WHERE id = 1".to_string();
cache.put(sql.clone());
let result = cache.get(&sql);
assert!(result.is_some());
assert_eq!(result.unwrap(), sql);
}
#[test]
fn test_cache_miss() {
let cache = StatementCache::default();
let sql = "SELECT * FROM users WHERE id = 1";
let result = cache.get(sql);
assert!(result.is_none());
}
#[test]
fn test_cache_stats() {
let cache = StatementCache::default();
let sql = "SELECT * FROM users WHERE id = 1".to_string();
cache.get(&sql);
cache.put(sql.clone());
cache.get(&sql);
cache.get(&sql);
let stats = cache.stats();
assert_eq!(stats.hits, 2);
assert_eq!(stats.misses, 1);
assert_eq!(stats.cached_statements, 1);
}
#[test]
fn test_cache_hit_rate() {
let stats = CacheStats {
hits: 80,
misses: 20,
evictions: 0,
expirations: 0,
cached_statements: 100,
};
assert!((stats.hit_rate() - 0.8).abs() < 0.001);
}
#[test]
fn test_cache_eviction() {
let config = StatementCacheConfig {
max_size: 3,
enable_fingerprinting: false, ..Default::default()
};
let cache = StatementCache::new(config);
cache.put("SELECT 1".to_string());
cache.put("SELECT 2".to_string());
cache.put("SELECT 3".to_string());
cache.put("SELECT 4".to_string());
let stats = cache.stats();
assert_eq!(stats.cached_statements, 3);
assert_eq!(stats.evictions, 1);
}
#[test]
fn test_cache_clear() {
let cache = StatementCache::default();
cache.put("SELECT 1".to_string());
cache.put("SELECT 2".to_string());
cache.clear();
let stats = cache.stats();
assert_eq!(stats.cached_statements, 0);
}
#[test]
fn test_normalize_sql() {
let sql = "SELECT * FROM users WHERE id = 1";
let normalized = normalize_sql(sql);
assert_eq!(normalized, "select * from users where id = ?");
}
#[test]
fn test_replace_literals() {
let sql = "SELECT * FROM users WHERE id = 123 AND name = 'john'";
let replaced = replace_literals(sql);
assert!(replaced.contains("?"));
assert!(!replaced.contains("123"));
assert!(!replaced.contains("john"));
}
#[test]
fn test_statement_reuse_count() {
let cache = StatementCache::default();
let sql = "SELECT * FROM users WHERE id = 1".to_string();
cache.put(sql.clone());
for _ in 0..5 {
cache.get(&sql);
}
let fingerprint = cache.fingerprint(&sql);
let cached_cache = cache.cache.read();
let stmt = cached_cache.get(&fingerprint).unwrap();
assert_eq!(stmt.reuse_count, 5);
}
#[test]
fn test_stats_reset() {
let mut stats = CacheStats {
hits: 100,
misses: 50,
evictions: 10,
expirations: 5,
cached_statements: 200,
};
stats.reset();
assert_eq!(stats.hits, 0);
assert_eq!(stats.misses, 0);
assert_eq!(stats.evictions, 0);
assert_eq!(stats.expirations, 0);
}
}