wasi_pg_client/query/
cache.rs

1//! Prepared statement cache with LRU eviction.
2//!
3//! Automatically caches prepared statements to avoid redundant `Parse`
4//! round-trips when the same SQL is executed multiple times.
5
6use std::collections::{HashMap, VecDeque};
7
8use crate::query::prepared::PreparedStatement;
9
10// ---------------------------------------------------------------------------
11// StatementCache
12// ---------------------------------------------------------------------------
13
14/// A least-recently-used (LRU) cache for prepared statements.
15///
16/// The cache is keyed by SQL text. When a statement is looked up that is
17/// not in the cache, it is prepared on the server and stored. If the cache
18/// is at capacity, the least-recently-used statement is evicted (closed on
19/// the server) before the new one is inserted.
20#[derive(Debug, Clone)]
21#[non_exhaustive]
22pub struct StatementCache {
23    cache: HashMap<String, PreparedStatement>,
24    capacity: usize,
25    order: VecDeque<String>,
26}
27
28impl StatementCache {
29    /// Create a new cache with the given capacity.
30    pub fn new(capacity: usize) -> Self {
31        Self {
32            cache: HashMap::with_capacity(capacity),
33            capacity,
34            order: VecDeque::with_capacity(capacity),
35        }
36    }
37
38    /// Returns the number of statements currently cached.
39    pub fn len(&self) -> usize {
40        self.cache.len()
41    }
42
43    /// Returns true if the cache contains no statements.
44    pub fn is_empty(&self) -> bool {
45        self.cache.is_empty()
46    }
47
48    /// Look up a prepared statement by SQL text.
49    ///
50    /// Returns `Some` if the statement is cached. The statement is moved to
51    /// the front of the LRU order.
52    pub fn get(&mut self, sql: &str) -> Option<&PreparedStatement> {
53        if self.cache.contains_key(sql) {
54            // Move to front (most recently used)
55            self.order.retain(|s| s != sql);
56            self.order.push_front(sql.to_string());
57            self.cache.get(sql)
58        } else {
59            None
60        }
61    }
62
63    /// Insert a prepared statement into the cache.
64    ///
65    /// If the cache is at capacity, the LRU entry is removed from the map
66    /// (but **not** closed on the server — the caller must do that if needed).
67    /// Returns the evicted statement, if any.
68    pub fn insert(&mut self, stmt: PreparedStatement) -> Option<PreparedStatement> {
69        let sql = stmt.sql().to_string();
70
71        // Remove old entry if present
72        self.order.retain(|s| s != &sql);
73
74        let evicted = if self.cache.len() >= self.capacity && !self.cache.contains_key(&sql) {
75            self.order
76                .pop_back()
77                .and_then(|old_sql| self.cache.remove(&old_sql))
78        } else {
79            None
80        };
81
82        self.order.push_front(sql.clone());
83        self.cache.insert(sql, stmt);
84        evicted
85    }
86
87    /// Remove a statement from the cache by SQL text.
88    ///
89    /// Returns the removed statement, if any.
90    pub fn remove(&mut self, sql: &str) -> Option<PreparedStatement> {
91        self.order.retain(|s| s != sql);
92        self.cache.remove(sql)
93    }
94
95    /// Clear all entries from the cache.
96    ///
97    /// Returns the evicted statements. The caller is responsible for closing
98    /// them on the server if needed.
99    pub fn clear(&mut self) -> Vec<PreparedStatement> {
100        self.order.clear();
101        self.cache.drain().map(|(_, v)| v).collect()
102    }
103}
104
105// ---------------------------------------------------------------------------
106// Tests
107// ---------------------------------------------------------------------------
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112    use std::sync::Arc;
113
114    fn dummy_stmt(sql: &str, name: &str) -> PreparedStatement {
115        PreparedStatement {
116            name: name.into(),
117            sql: sql.into(),
118            param_types: vec![],
119            columns: Arc::new(vec![]),
120        }
121    }
122
123    #[test]
124    fn test_cache_insert_and_get() {
125        let mut cache = StatementCache::new(2);
126        let stmt = dummy_stmt("SELECT 1", "s1");
127        cache.insert(stmt);
128
129        assert_eq!(cache.len(), 1);
130        assert!(cache.get("SELECT 1").is_some());
131        assert!(cache.get("SELECT 2").is_none());
132    }
133
134    #[test]
135    fn test_cache_lru_eviction() {
136        let mut cache = StatementCache::new(2);
137        cache.insert(dummy_stmt("SELECT 1", "s1"));
138        cache.insert(dummy_stmt("SELECT 2", "s2"));
139
140        // Access SELECT 1 to make it MRU
141        let _ = cache.get("SELECT 1");
142
143        // Insert SELECT 3 — should evict SELECT 2 (LRU)
144        let evicted = cache.insert(dummy_stmt("SELECT 3", "s3"));
145        assert!(evicted.is_some());
146        assert_eq!(evicted.unwrap().sql(), "SELECT 2");
147
148        assert!(cache.get("SELECT 1").is_some());
149        assert!(cache.get("SELECT 2").is_none());
150        assert!(cache.get("SELECT 3").is_some());
151    }
152
153    #[test]
154    fn test_cache_remove() {
155        let mut cache = StatementCache::new(2);
156        cache.insert(dummy_stmt("SELECT 1", "s1"));
157
158        let removed = cache.remove("SELECT 1");
159        assert!(removed.is_some());
160        assert_eq!(removed.unwrap().sql(), "SELECT 1");
161        assert!(cache.is_empty());
162    }
163
164    #[test]
165    fn test_cache_clear() {
166        let mut cache = StatementCache::new(2);
167        cache.insert(dummy_stmt("SELECT 1", "s1"));
168        cache.insert(dummy_stmt("SELECT 2", "s2"));
169
170        let cleared = cache.clear();
171        assert_eq!(cleared.len(), 2);
172        assert!(cache.is_empty());
173    }
174}
wasi_pg_client/query/cache.rs

wasi_pg_client/query/
cache.rs
