Skip to main content

prax_query/
cache.rs

1//! Query caching and prepared statement management.
2//!
3//! This module provides utilities for caching SQL queries and managing
4//! prepared statements to improve performance.
5//!
6//! # Query Cache
7//!
8//! The `QueryCache` stores recently executed queries by their hash,
9//! allowing fast lookup of previously built SQL strings.
10//!
11//! ```rust
12//! use prax_query::cache::QueryCache;
13//!
14//! let cache = QueryCache::new(1000);
15//!
16//! // Cache a query
17//! cache.insert("users_by_id", "SELECT * FROM users WHERE id = $1");
18//!
19//! // Retrieve later
20//! if let Some(sql) = cache.get("users_by_id") {
21//!     println!("Cached SQL: {}", sql);
22//! }
23//! ```
24
25use std::borrow::Cow;
26use std::collections::HashMap;
27use std::hash::{Hash, Hasher};
28use std::sync::atomic::{AtomicU64, Ordering};
29use std::sync::{Arc, RwLock};
30use tracing::debug;
31
32/// A thread-safe cache for SQL queries.
33///
34/// Uses a simple LRU-like eviction strategy when the cache is full.
35#[derive(Debug)]
36pub struct QueryCache {
37    /// Maximum number of entries in the cache.
38    max_size: usize,
39    /// The cached queries.
40    cache: RwLock<HashMap<QueryKey, CachedQuery>>,
41    /// Statistics about cache usage.
42    ///
43    /// Atomic counters so the `get()` hot path can record hits/misses
44    /// without taking the stats write lock while still holding the
45    /// cache read lock (the previous nested-lock pattern was a
46    /// contention hotspot under concurrent reads).
47    stats: AtomicCacheStats,
48}
49
50/// Atomic backing for [`CacheStats`].
51///
52/// `record_*` methods use `Relaxed` because the counters are
53/// strictly monotonic and don't synchronize with anything else; the
54/// only consumer is `QueryCache::stats()` which snapshots them into a
55/// regular [`CacheStats`] value.
56#[derive(Debug, Default)]
57struct AtomicCacheStats {
58    hits: AtomicU64,
59    misses: AtomicU64,
60    evictions: AtomicU64,
61    insertions: AtomicU64,
62}
63
64impl AtomicCacheStats {
65    #[inline]
66    fn record_hit(&self) {
67        self.hits.fetch_add(1, Ordering::Relaxed);
68    }
69
70    #[inline]
71    fn record_miss(&self) {
72        self.misses.fetch_add(1, Ordering::Relaxed);
73    }
74
75    #[inline]
76    fn record_eviction(&self) {
77        self.evictions.fetch_add(1, Ordering::Relaxed);
78    }
79
80    #[inline]
81    fn record_insertion(&self) {
82        self.insertions.fetch_add(1, Ordering::Relaxed);
83    }
84
85    fn snapshot(&self) -> CacheStats {
86        CacheStats {
87            hits: self.hits.load(Ordering::Relaxed),
88            misses: self.misses.load(Ordering::Relaxed),
89            evictions: self.evictions.load(Ordering::Relaxed),
90            insertions: self.insertions.load(Ordering::Relaxed),
91        }
92    }
93
94    fn reset(&self) {
95        self.hits.store(0, Ordering::Relaxed);
96        self.misses.store(0, Ordering::Relaxed);
97        self.evictions.store(0, Ordering::Relaxed);
98        self.insertions.store(0, Ordering::Relaxed);
99    }
100}
101
102/// A key for looking up cached queries.
103#[derive(Debug, Clone, PartialEq, Eq, Hash)]
104pub struct QueryKey {
105    /// The unique identifier for this query type.
106    key: Cow<'static, str>,
107}
108
109impl QueryKey {
110    /// Create a new query key from a static string.
111    #[inline]
112    pub const fn new(key: &'static str) -> Self {
113        Self {
114            key: Cow::Borrowed(key),
115        }
116    }
117
118    /// Create a new query key from an owned string.
119    #[inline]
120    pub fn owned(key: String) -> Self {
121        Self {
122            key: Cow::Owned(key),
123        }
124    }
125}
126
127impl From<&'static str> for QueryKey {
128    fn from(s: &'static str) -> Self {
129        Self::new(s)
130    }
131}
132
133impl From<String> for QueryKey {
134    fn from(s: String) -> Self {
135        Self::owned(s)
136    }
137}
138
139/// A cached SQL query.
140#[derive(Debug)]
141pub struct CachedQuery {
142    /// The SQL string.
143    pub sql: String,
144    /// The number of parameters expected.
145    pub param_count: usize,
146    /// Number of times this query has been accessed.
147    ///
148    /// Atomic so hits can be recorded under the cache's read lock
149    /// (interior mutability, same pattern as `SqlTemplate::touch`).
150    access_count: AtomicU64,
151}
152
153impl Clone for CachedQuery {
154    fn clone(&self) -> Self {
155        Self {
156            sql: self.sql.clone(),
157            param_count: self.param_count,
158            access_count: AtomicU64::new(self.access_count.load(Ordering::Relaxed)),
159        }
160    }
161}
162
163impl CachedQuery {
164    /// Create a new cached query.
165    pub fn new(sql: impl Into<String>, param_count: usize) -> Self {
166        Self {
167            sql: sql.into(),
168            param_count,
169            access_count: AtomicU64::new(0),
170        }
171    }
172
173    /// Record a cache hit for LRU eviction.
174    #[inline]
175    fn record_access(&self) {
176        self.access_count.fetch_add(1, Ordering::Relaxed);
177    }
178
179    /// Get the SQL string.
180    #[inline]
181    pub fn sql(&self) -> &str {
182        &self.sql
183    }
184
185    /// Get the expected parameter count.
186    #[inline]
187    pub fn param_count(&self) -> usize {
188        self.param_count
189    }
190}
191
192/// Statistics about cache usage.
193#[derive(Debug, Default, Clone)]
194pub struct CacheStats {
195    /// Number of cache hits.
196    pub hits: u64,
197    /// Number of cache misses.
198    pub misses: u64,
199    /// Number of evictions.
200    pub evictions: u64,
201    /// Number of insertions.
202    pub insertions: u64,
203}
204
205impl CacheStats {
206    /// Calculate the hit rate.
207    #[inline]
208    pub fn hit_rate(&self) -> f64 {
209        let total = self.hits + self.misses;
210        if total == 0 {
211            0.0
212        } else {
213            self.hits as f64 / total as f64
214        }
215    }
216}
217
218impl QueryCache {
219    /// Create a new query cache with the given maximum size.
220    pub fn new(max_size: usize) -> Self {
221        tracing::info!(max_size, "QueryCache initialized");
222        Self {
223            max_size,
224            cache: RwLock::new(HashMap::with_capacity(max_size)),
225            stats: AtomicCacheStats::default(),
226        }
227    }
228
229    /// Insert a query into the cache.
230    pub fn insert(&self, key: impl Into<QueryKey>, sql: impl Into<String>) {
231        let key = key.into();
232        let sql = sql.into();
233        let param_count = count_placeholders(&sql);
234        debug!(key = ?key.key, sql_len = sql.len(), param_count, "QueryCache::insert()");
235
236        let mut cache = self.cache.write().unwrap();
237
238        // Evict if full
239        if cache.len() >= self.max_size && !cache.contains_key(&key) {
240            self.evict_lru(&mut cache);
241            self.stats.record_eviction();
242            debug!("QueryCache evicted entry");
243        }
244
245        cache.insert(key, CachedQuery::new(sql, param_count));
246        self.stats.record_insertion();
247    }
248
249    /// Insert a query with known parameter count.
250    pub fn insert_with_params(
251        &self,
252        key: impl Into<QueryKey>,
253        sql: impl Into<String>,
254        param_count: usize,
255    ) {
256        let key = key.into();
257        let sql = sql.into();
258
259        let mut cache = self.cache.write().unwrap();
260
261        // Evict if full
262        if cache.len() >= self.max_size && !cache.contains_key(&key) {
263            self.evict_lru(&mut cache);
264            self.stats.record_eviction();
265        }
266
267        cache.insert(key, CachedQuery::new(sql, param_count));
268        self.stats.record_insertion();
269    }
270
271    /// Get a query from the cache.
272    pub fn get(&self, key: impl Into<QueryKey>) -> Option<String> {
273        let key = key.into();
274
275        let cache = self.cache.read().unwrap();
276        if let Some(entry) = cache.get(&key) {
277            // Atomic counters — no write-lock conflict with the read
278            // lock we still hold on `cache`.
279            entry.record_access();
280            self.stats.record_hit();
281            debug!(key = ?key.key, "QueryCache hit");
282            return Some(entry.sql.clone());
283        }
284        drop(cache);
285
286        self.stats.record_miss();
287        debug!(key = ?key.key, "QueryCache miss");
288        None
289    }
290
291    /// Get a cached query entry (includes metadata).
292    pub fn get_entry(&self, key: impl Into<QueryKey>) -> Option<CachedQuery> {
293        let key = key.into();
294
295        let cache = self.cache.read().unwrap();
296        if let Some(entry) = cache.get(&key) {
297            entry.record_access();
298            self.stats.record_hit();
299            return Some(entry.clone());
300        }
301        drop(cache);
302
303        self.stats.record_miss();
304        None
305    }
306
307    /// Get or compute a query.
308    ///
309    /// If the query is cached, returns the cached version.
310    /// Otherwise, computes it using the provided function and caches it.
311    pub fn get_or_insert<F>(&self, key: impl Into<QueryKey>, f: F) -> String
312    where
313        F: FnOnce() -> String,
314    {
315        let key = key.into();
316
317        // Try to get from cache
318        if let Some(sql) = self.get(key.clone()) {
319            return sql;
320        }
321
322        // Compute and insert
323        let sql = f();
324        self.insert(key, sql.clone());
325        sql
326    }
327
328    /// Check if a key exists in the cache.
329    pub fn contains(&self, key: impl Into<QueryKey>) -> bool {
330        let key = key.into();
331        let cache = self.cache.read().unwrap();
332        cache.contains_key(&key)
333    }
334
335    /// Remove a query from the cache.
336    pub fn remove(&self, key: impl Into<QueryKey>) -> Option<String> {
337        let key = key.into();
338        let mut cache = self.cache.write().unwrap();
339        cache.remove(&key).map(|e| e.sql)
340    }
341
342    /// Clear the entire cache.
343    pub fn clear(&self) {
344        let mut cache = self.cache.write().unwrap();
345        cache.clear();
346    }
347
348    /// Get the current number of cached queries.
349    pub fn len(&self) -> usize {
350        let cache = self.cache.read().unwrap();
351        cache.len()
352    }
353
354    /// Check if the cache is empty.
355    pub fn is_empty(&self) -> bool {
356        self.len() == 0
357    }
358
359    /// Get the maximum cache size.
360    pub fn max_size(&self) -> usize {
361        self.max_size
362    }
363
364    /// Get cache statistics.
365    pub fn stats(&self) -> CacheStats {
366        self.stats.snapshot()
367    }
368
369    /// Reset cache statistics.
370    pub fn reset_stats(&self) {
371        self.stats.reset();
372    }
373
374    /// Evict the least recently used entries.
375    fn evict_lru(&self, cache: &mut HashMap<QueryKey, CachedQuery>) {
376        // Evict entries with the lowest access count: 25% of the cache,
377        // but always at least one so a full cache never grows past
378        // max_size (len < 4 would otherwise evict nothing).
379        if cache.is_empty() {
380            return;
381        }
382        let to_evict = (cache.len() / 4).max(1);
383
384        // Snapshot (access_count, key) pairs — only the keys that are
385        // actually evicted get cloned. A partial select partitions the
386        // `to_evict` lowest-count entries up front (O(n) average
387        // instead of an O(n log n) full sort).
388        let mut entries: Vec<_> = cache
389            .iter()
390            .map(|(k, v)| (v.access_count.load(Ordering::Relaxed), k))
391            .collect();
392        entries.select_nth_unstable_by(to_evict - 1, |a, b| a.0.cmp(&b.0));
393
394        let evict_keys: Vec<QueryKey> = entries[..to_evict]
395            .iter()
396            .map(|(_, k)| (*k).clone())
397            .collect();
398        for key in evict_keys {
399            cache.remove(&key);
400        }
401    }
402}
403
404impl Default for QueryCache {
405    fn default() -> Self {
406        Self::new(1000)
407    }
408}
409
410/// Count the number of parameter placeholders in a SQL string.
411fn count_placeholders(sql: &str) -> usize {
412    let mut count = 0;
413    let mut chars = sql.chars().peekable();
414
415    while let Some(c) = chars.next() {
416        if c == '$' {
417            // PostgreSQL-style: $1, $2, etc.
418            let mut num = String::new();
419            while let Some(&d) = chars.peek() {
420                if d.is_ascii_digit() {
421                    num.push(d);
422                    chars.next();
423                } else {
424                    break;
425                }
426            }
427            if !num.is_empty()
428                && let Ok(n) = num.parse::<usize>()
429            {
430                count = count.max(n);
431            }
432        } else if c == '?' {
433            // MySQL/SQLite-style
434            count += 1;
435        }
436    }
437
438    count
439}
440
441/// A query hash for fast lookup.
442#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
443pub struct QueryHash(u64);
444
445impl QueryHash {
446    /// Compute a hash for the given SQL query.
447    pub fn new(sql: &str) -> Self {
448        let mut hasher = std::collections::hash_map::DefaultHasher::new();
449        sql.hash(&mut hasher);
450        Self(hasher.finish())
451    }
452
453    /// Get the raw hash value.
454    #[inline]
455    pub fn value(&self) -> u64 {
456        self.0
457    }
458}
459
460/// Common query patterns for caching.
461pub mod patterns {
462    use super::QueryKey;
463
464    /// Query key for SELECT by ID.
465    #[inline]
466    pub fn select_by_id(table: &str) -> QueryKey {
467        QueryKey::owned(format!("select_by_id:{}", table))
468    }
469
470    /// Query key for SELECT all.
471    #[inline]
472    pub fn select_all(table: &str) -> QueryKey {
473        QueryKey::owned(format!("select_all:{}", table))
474    }
475
476    /// Query key for INSERT.
477    #[inline]
478    pub fn insert(table: &str, columns: usize) -> QueryKey {
479        QueryKey::owned(format!("insert:{}:{}", table, columns))
480    }
481
482    /// Query key for UPDATE by ID.
483    #[inline]
484    pub fn update_by_id(table: &str, columns: usize) -> QueryKey {
485        QueryKey::owned(format!("update_by_id:{}:{}", table, columns))
486    }
487
488    /// Query key for DELETE by ID.
489    #[inline]
490    pub fn delete_by_id(table: &str) -> QueryKey {
491        QueryKey::owned(format!("delete_by_id:{}", table))
492    }
493
494    /// Query key for COUNT.
495    #[inline]
496    pub fn count(table: &str) -> QueryKey {
497        QueryKey::owned(format!("count:{}", table))
498    }
499
500    /// Query key for COUNT with filter.
501    #[inline]
502    pub fn count_filtered(table: &str, filter_hash: u64) -> QueryKey {
503        QueryKey::owned(format!("count:{}:{}", table, filter_hash))
504    }
505}
506
507// =============================================================================
508// High-Performance SQL Template Cache
509// =============================================================================
510
511/// A high-performance SQL template cache optimized for repeated queries.
512///
513/// Unlike `QueryCache` which stores full SQL strings, `SqlTemplateCache` stores
514/// template structures with pre-computed placeholder positions for very fast
515/// instantiation.
516///
517/// # Performance
518///
519/// - Cache lookup: O(1) hash lookup, ~5-10ns
520/// - Template instantiation: O(n) where n is parameter count
521/// - Thread-safe with minimal contention (parking_lot RwLock)
522///
523/// # Examples
524///
525/// ```rust
526/// use prax_query::cache::SqlTemplateCache;
527///
528/// let cache = SqlTemplateCache::new(1000);
529///
530/// // Register a template
531/// let template = cache.register("users_by_id", "SELECT * FROM users WHERE id = $1");
532///
533/// // Instant retrieval (~5ns)
534/// let sql = cache.get("users_by_id");
535/// ```
536#[derive(Debug)]
537pub struct SqlTemplateCache {
538    /// Maximum number of templates.
539    max_size: usize,
540    /// Cached templates (using Arc for cheap cloning).
541    templates: parking_lot::RwLock<HashMap<u64, Arc<SqlTemplate>>>,
542    /// String key to hash lookup.
543    key_index: parking_lot::RwLock<HashMap<Cow<'static, str>, u64>>,
544    /// Statistics.
545    stats: parking_lot::RwLock<CacheStats>,
546}
547
548/// A pre-parsed SQL template for fast instantiation.
549#[derive(Debug)]
550pub struct SqlTemplate {
551    /// The complete SQL string (for direct use).
552    pub sql: Arc<str>,
553    /// Pre-computed hash for fast lookup.
554    pub hash: u64,
555    /// Number of parameters.
556    pub param_count: usize,
557    /// Access timestamp for LRU.
558    last_access: std::sync::atomic::AtomicU64,
559}
560
561impl Clone for SqlTemplate {
562    fn clone(&self) -> Self {
563        use std::sync::atomic::Ordering;
564        Self {
565            sql: Arc::clone(&self.sql),
566            hash: self.hash,
567            param_count: self.param_count,
568            last_access: std::sync::atomic::AtomicU64::new(
569                self.last_access.load(Ordering::Relaxed),
570            ),
571        }
572    }
573}
574
575impl SqlTemplate {
576    /// Create a new SQL template.
577    pub fn new(sql: impl AsRef<str>) -> Self {
578        let sql_str = sql.as_ref();
579        let param_count = count_placeholders(sql_str);
580        let hash = {
581            let mut hasher = std::collections::hash_map::DefaultHasher::new();
582            sql_str.hash(&mut hasher);
583            hasher.finish()
584        };
585
586        Self {
587            sql: Arc::from(sql_str),
588            hash,
589            param_count,
590            last_access: std::sync::atomic::AtomicU64::new(0),
591        }
592    }
593
594    /// Get the SQL string as a reference.
595    #[inline(always)]
596    pub fn sql(&self) -> &str {
597        &self.sql
598    }
599
600    /// Get the SQL string as an Arc (zero-copy clone).
601    #[inline(always)]
602    pub fn sql_arc(&self) -> Arc<str> {
603        Arc::clone(&self.sql)
604    }
605
606    /// Touch the template to update LRU access time.
607    #[inline]
608    fn touch(&self) {
609        use std::sync::atomic::Ordering;
610        use std::time::{SystemTime, UNIX_EPOCH};
611        let now = SystemTime::now()
612            .duration_since(UNIX_EPOCH)
613            .map(|d| d.as_secs())
614            .unwrap_or(0);
615        self.last_access.store(now, Ordering::Relaxed);
616    }
617}
618
619impl SqlTemplateCache {
620    /// Create a new template cache with the given maximum size.
621    pub fn new(max_size: usize) -> Self {
622        tracing::info!(max_size, "SqlTemplateCache initialized");
623        Self {
624            max_size,
625            templates: parking_lot::RwLock::new(HashMap::with_capacity(max_size)),
626            key_index: parking_lot::RwLock::new(HashMap::with_capacity(max_size)),
627            stats: parking_lot::RwLock::new(CacheStats::default()),
628        }
629    }
630
631    /// Register a SQL template with a string key.
632    ///
633    /// Returns the template for immediate use.
634    #[inline]
635    pub fn register(
636        &self,
637        key: impl Into<Cow<'static, str>>,
638        sql: impl AsRef<str>,
639    ) -> Arc<SqlTemplate> {
640        let key = key.into();
641        let template = Arc::new(SqlTemplate::new(sql));
642        let hash = template.hash;
643
644        let mut templates = self.templates.write();
645        let mut key_index = self.key_index.write();
646        let mut stats = self.stats.write();
647
648        // Evict if full
649        if templates.len() >= self.max_size {
650            self.evict_lru_internal(&mut templates, &mut key_index);
651            stats.evictions += 1;
652        }
653
654        key_index.insert(key, hash);
655        templates.insert(hash, Arc::clone(&template));
656        stats.insertions += 1;
657
658        debug!(hash, "SqlTemplateCache::register()");
659        template
660    }
661
662    /// Register a template by hash (for pre-computed hashes).
663    #[inline]
664    pub fn register_by_hash(&self, hash: u64, sql: impl AsRef<str>) -> Arc<SqlTemplate> {
665        let template = Arc::new(SqlTemplate::new(sql));
666
667        let mut templates = self.templates.write();
668        let mut stats = self.stats.write();
669
670        if templates.len() >= self.max_size {
671            let mut key_index = self.key_index.write();
672            self.evict_lru_internal(&mut templates, &mut key_index);
673            stats.evictions += 1;
674        }
675
676        templates.insert(hash, Arc::clone(&template));
677        stats.insertions += 1;
678
679        template
680    }
681
682    /// Get a template by string key (returns Arc for zero-copy).
683    ///
684    /// # Performance
685    ///
686    /// This is the fastest way to get cached SQL:
687    /// - Hash lookup: ~5ns
688    /// - Returns Arc<SqlTemplate> (no allocation)
689    #[inline]
690    pub fn get(&self, key: &str) -> Option<Arc<SqlTemplate>> {
691        let hash = {
692            let key_index = self.key_index.read();
693            match key_index.get(key) {
694                Some(&h) => h,
695                None => {
696                    drop(key_index); // Release read lock before write
697                    let mut stats = self.stats.write();
698                    stats.misses += 1;
699                    return None;
700                }
701            }
702        };
703
704        let templates = self.templates.read();
705        if let Some(template) = templates.get(&hash) {
706            template.touch();
707            let mut stats = self.stats.write();
708            stats.hits += 1;
709            return Some(Arc::clone(template));
710        }
711
712        let mut stats = self.stats.write();
713        stats.misses += 1;
714        None
715    }
716
717    /// Get a template by pre-computed hash (fastest path).
718    ///
719    /// # Performance
720    ///
721    /// ~3-5ns for cache hit with pre-computed hash.
722    #[inline(always)]
723    pub fn get_by_hash(&self, hash: u64) -> Option<Arc<SqlTemplate>> {
724        let templates = self.templates.read();
725        if let Some(template) = templates.get(&hash) {
726            template.touch();
727            // Skip stats update for maximum performance
728            return Some(Arc::clone(template));
729        }
730        None
731    }
732
733    /// Get the SQL string directly (convenience method).
734    #[inline]
735    pub fn get_sql(&self, key: &str) -> Option<Arc<str>> {
736        self.get(key).map(|t| t.sql_arc())
737    }
738
739    /// Get or compute a template.
740    #[inline]
741    pub fn get_or_register<F>(&self, key: impl Into<Cow<'static, str>>, f: F) -> Arc<SqlTemplate>
742    where
743        F: FnOnce() -> String,
744    {
745        let key = key.into();
746
747        // Fast path: check if exists
748        if let Some(template) = self.get(&key) {
749            return template;
750        }
751
752        // Slow path: compute and register
753        let sql = f();
754        self.register(key, sql)
755    }
756
757    /// Check if a key exists.
758    #[inline]
759    pub fn contains(&self, key: &str) -> bool {
760        let key_index = self.key_index.read();
761        key_index.contains_key(key)
762    }
763
764    /// Get cache statistics.
765    pub fn stats(&self) -> CacheStats {
766        self.stats.read().clone()
767    }
768
769    /// Get the number of cached templates.
770    pub fn len(&self) -> usize {
771        self.templates.read().len()
772    }
773
774    /// Check if the cache is empty.
775    pub fn is_empty(&self) -> bool {
776        self.len() == 0
777    }
778
779    /// Clear the cache.
780    pub fn clear(&self) {
781        self.templates.write().clear();
782        self.key_index.write().clear();
783    }
784
785    /// Evict least recently used templates (internal, assumes locks held).
786    fn evict_lru_internal(
787        &self,
788        templates: &mut HashMap<u64, Arc<SqlTemplate>>,
789        key_index: &mut HashMap<Cow<'static, str>, u64>,
790    ) {
791        use std::sync::atomic::Ordering;
792
793        let to_evict = templates.len() / 4;
794        if to_evict == 0 {
795            return;
796        }
797
798        // Find templates with oldest access times (partial select —
799        // only the `to_evict` oldest need to be up front).
800        let mut entries: Vec<_> = templates
801            .iter()
802            .map(|(&hash, t)| (hash, t.last_access.load(Ordering::Relaxed)))
803            .collect();
804        entries.select_nth_unstable_by(to_evict - 1, |a, b| a.1.cmp(&b.1));
805
806        // Collect the set of hashes to evict, remove their templates, then
807        // prune key_index in a single pass (avoids O(evicted * key_index.len())).
808        let evicted: std::collections::HashSet<u64> = entries
809            .into_iter()
810            .take(to_evict)
811            .map(|(hash, _)| {
812                templates.remove(&hash);
813                hash
814            })
815            .collect();
816        key_index.retain(|_, h| !evicted.contains(h));
817    }
818}
819
820impl Default for SqlTemplateCache {
821    fn default() -> Self {
822        Self::new(1000)
823    }
824}
825
826// =============================================================================
827// Global Template Cache (for zero-overhead repeated queries)
828// =============================================================================
829
830/// Global SQL template cache for maximum performance.
831///
832/// Use this for queries that are repeated many times with only parameter changes.
833/// The global cache avoids the overhead of passing cache references around.
834///
835/// # Examples
836///
837/// ```rust
838/// use prax_query::cache::{global_template_cache, register_global_template};
839///
840/// // Pre-register common queries at startup
841/// register_global_template("users_by_id", "SELECT * FROM users WHERE id = $1");
842///
843/// // Later, get the cached SQL (~5ns)
844/// if let Some(template) = global_template_cache().get("users_by_id") {
845///     println!("SQL: {}", template.sql());
846/// }
847/// ```
848static GLOBAL_TEMPLATE_CACHE: std::sync::OnceLock<SqlTemplateCache> = std::sync::OnceLock::new();
849
850/// Get the global SQL template cache.
851#[inline(always)]
852pub fn global_template_cache() -> &'static SqlTemplateCache {
853    GLOBAL_TEMPLATE_CACHE.get_or_init(|| SqlTemplateCache::new(10000))
854}
855
856/// Register a template in the global cache.
857#[inline]
858pub fn register_global_template(
859    key: impl Into<Cow<'static, str>>,
860    sql: impl AsRef<str>,
861) -> Arc<SqlTemplate> {
862    global_template_cache().register(key, sql)
863}
864
865/// Get a template from the global cache.
866#[inline(always)]
867pub fn get_global_template(key: &str) -> Option<Arc<SqlTemplate>> {
868    global_template_cache().get(key)
869}
870
871/// Pre-compute a query hash for repeated lookups.
872///
873/// Use this when you have a query key that will be used many times.
874/// Computing the hash once and using `get_by_hash` is faster than
875/// string key lookups.
876#[inline]
877pub fn precompute_query_hash(key: &str) -> u64 {
878    let mut hasher = std::collections::hash_map::DefaultHasher::new();
879    key.hash(&mut hasher);
880    hasher.finish()
881}
882
883#[cfg(test)]
884mod tests {
885    use super::*;
886
887    #[test]
888    fn test_query_cache_basic() {
889        let cache = QueryCache::new(10);
890
891        cache.insert("users_by_id", "SELECT * FROM users WHERE id = $1");
892        assert!(cache.contains("users_by_id"));
893
894        let sql = cache.get("users_by_id");
895        assert_eq!(sql, Some("SELECT * FROM users WHERE id = $1".to_string()));
896    }
897
898    #[test]
899    fn test_query_cache_get_or_insert() {
900        let cache = QueryCache::new(10);
901
902        let sql1 = cache.get_or_insert("test", || "SELECT 1".to_string());
903        assert_eq!(sql1, "SELECT 1");
904
905        let sql2 = cache.get_or_insert("test", || "SELECT 2".to_string());
906        assert_eq!(sql2, "SELECT 1"); // Should return cached value
907    }
908
909    #[test]
910    fn test_query_cache_stats() {
911        let cache = QueryCache::new(10);
912
913        cache.insert("test", "SELECT 1");
914        cache.get("test"); // Hit
915        cache.get("test"); // Hit
916        cache.get("missing"); // Miss
917
918        let stats = cache.stats();
919        assert_eq!(stats.hits, 2);
920        assert_eq!(stats.misses, 1);
921        assert_eq!(stats.insertions, 1);
922    }
923
924    #[test]
925    fn test_query_cache_evicts_least_accessed() {
926        let cache = QueryCache::new(3);
927
928        cache.insert("a", "SELECT 'a'");
929        cache.insert("b", "SELECT 'b'");
930        cache.insert("c", "SELECT 'c'");
931
932        // Access "a" repeatedly so it has the highest access count.
933        for _ in 0..5 {
934            cache.get("a");
935        }
936
937        // Cache is full; this insert must evict exactly one of the
938        // never-accessed entries ("b" or "c") and retain "a".
939        cache.insert("d", "SELECT 'd'");
940
941        assert!(cache.len() <= cache.max_size());
942        assert!(cache.contains("a"));
943        assert!(cache.contains("d"));
944        let retained = cache.contains("b") as usize + cache.contains("c") as usize;
945        assert_eq!(retained, 1, "exactly one of b/c should be evicted");
946
947        // Repeated inserts at capacity never grow beyond max_size.
948        for i in 0..8 {
949            cache.insert(format!("extra_{i}"), "SELECT 1");
950            assert!(cache.len() <= cache.max_size());
951        }
952    }
953
954    #[test]
955    fn test_count_placeholders_postgres() {
956        assert_eq!(count_placeholders("SELECT * FROM users WHERE id = $1"), 1);
957        assert_eq!(
958            count_placeholders("SELECT * FROM users WHERE id = $1 AND name = $2"),
959            2
960        );
961        assert_eq!(count_placeholders("SELECT * FROM users WHERE id = $10"), 10);
962    }
963
964    #[test]
965    fn test_count_placeholders_mysql() {
966        assert_eq!(count_placeholders("SELECT * FROM users WHERE id = ?"), 1);
967        assert_eq!(
968            count_placeholders("SELECT * FROM users WHERE id = ? AND name = ?"),
969            2
970        );
971    }
972
973    #[test]
974    fn test_query_hash() {
975        let hash1 = QueryHash::new("SELECT * FROM users");
976        let hash2 = QueryHash::new("SELECT * FROM users");
977        let hash3 = QueryHash::new("SELECT * FROM posts");
978
979        assert_eq!(hash1, hash2);
980        assert_ne!(hash1, hash3);
981    }
982
983    #[test]
984    fn test_patterns() {
985        let key = patterns::select_by_id("users");
986        assert!(key.key.starts_with("select_by_id:"));
987    }
988
989    // =========================================================================
990    // SqlTemplateCache Tests
991    // =========================================================================
992
993    #[test]
994    fn test_sql_template_cache_basic() {
995        let cache = SqlTemplateCache::new(100);
996
997        let template = cache.register("users_by_id", "SELECT * FROM users WHERE id = $1");
998        assert_eq!(template.sql(), "SELECT * FROM users WHERE id = $1");
999        assert_eq!(template.param_count, 1);
1000    }
1001
1002    #[test]
1003    fn test_sql_template_cache_get() {
1004        let cache = SqlTemplateCache::new(100);
1005
1006        cache.register("test_query", "SELECT * FROM test WHERE x = $1");
1007
1008        let result = cache.get("test_query");
1009        assert!(result.is_some());
1010        assert_eq!(result.unwrap().sql(), "SELECT * FROM test WHERE x = $1");
1011
1012        let missing = cache.get("nonexistent");
1013        assert!(missing.is_none());
1014    }
1015
1016    #[test]
1017    fn test_sql_template_cache_get_by_hash() {
1018        let cache = SqlTemplateCache::new(100);
1019
1020        let template = cache.register("fast_query", "SELECT 1");
1021        let hash = template.hash;
1022
1023        // Get by hash should be very fast
1024        let result = cache.get_by_hash(hash);
1025        assert!(result.is_some());
1026        assert_eq!(result.unwrap().sql(), "SELECT 1");
1027    }
1028
1029    #[test]
1030    fn test_sql_template_cache_get_or_register() {
1031        let cache = SqlTemplateCache::new(100);
1032
1033        let t1 = cache.get_or_register("computed", || "SELECT * FROM computed".to_string());
1034        assert_eq!(t1.sql(), "SELECT * FROM computed");
1035
1036        // Second call should return cached version
1037        let t2 = cache.get_or_register("computed", || panic!("Should not be called"));
1038        assert_eq!(t2.sql(), "SELECT * FROM computed");
1039        assert_eq!(t1.hash, t2.hash);
1040    }
1041
1042    #[test]
1043    fn test_sql_template_cache_stats() {
1044        let cache = SqlTemplateCache::new(100);
1045
1046        cache.register("q1", "SELECT 1");
1047        cache.get("q1"); // Hit
1048        cache.get("q1"); // Hit
1049        cache.get("missing"); // Miss
1050
1051        let stats = cache.stats();
1052        assert_eq!(stats.hits, 2);
1053        assert_eq!(stats.misses, 1);
1054        assert_eq!(stats.insertions, 1);
1055    }
1056
1057    #[test]
1058    fn test_global_template_cache() {
1059        // Register in global cache
1060        let template = register_global_template("global_test", "SELECT * FROM global");
1061        assert_eq!(template.sql(), "SELECT * FROM global");
1062
1063        // Retrieve from global cache
1064        let result = get_global_template("global_test");
1065        assert!(result.is_some());
1066        assert_eq!(result.unwrap().sql(), "SELECT * FROM global");
1067    }
1068
1069    #[test]
1070    fn test_precompute_query_hash() {
1071        let hash1 = precompute_query_hash("test_key");
1072        let hash2 = precompute_query_hash("test_key");
1073        let hash3 = precompute_query_hash("other_key");
1074
1075        assert_eq!(hash1, hash2);
1076        assert_ne!(hash1, hash3);
1077    }
1078
1079    #[test]
1080    fn test_execution_plan_cache() {
1081        let cache = ExecutionPlanCache::new(100);
1082
1083        // Register a plan
1084        let plan = cache.register(
1085            "users_by_email",
1086            "SELECT * FROM users WHERE email = $1",
1087            PlanHint::IndexScan("users_email_idx".into()),
1088        );
1089        assert_eq!(plan.sql.as_ref(), "SELECT * FROM users WHERE email = $1");
1090
1091        // Get cached plan
1092        let result = cache.get("users_by_email");
1093        assert!(result.is_some());
1094        assert!(matches!(result.unwrap().hint, PlanHint::IndexScan(_)));
1095    }
1096}
1097
1098// ============================================================================
1099// Execution Plan Caching
1100// ============================================================================
1101
1102/// Hints for query execution optimization.
1103///
1104/// These hints can be used by database engines to optimize query execution.
1105/// Different databases support different hints - the engine implementation
1106/// decides how to apply them.
1107#[derive(Debug, Clone, Default)]
1108pub enum PlanHint {
1109    /// No specific hint.
1110    #[default]
1111    None,
1112    /// Force use of a specific index.
1113    IndexScan(String),
1114    /// Force a sequential scan (for analytics queries).
1115    SeqScan,
1116    /// Enable parallel execution.
1117    Parallel(u32),
1118    /// Cache this query's execution plan.
1119    CachePlan,
1120    /// Set a timeout for this query.
1121    Timeout(std::time::Duration),
1122    /// Custom database-specific hint.
1123    Custom(String),
1124}
1125
1126/// A cached execution plan with optimization hints.
1127#[derive(Debug)]
1128pub struct ExecutionPlan {
1129    /// The SQL query.
1130    pub sql: Arc<str>,
1131    /// Pre-computed hash for fast lookup.
1132    pub hash: u64,
1133    /// Execution hint.
1134    pub hint: PlanHint,
1135    /// Estimated cost (if available from EXPLAIN).
1136    pub estimated_cost: Option<f64>,
1137    /// Number of times this plan has been used.
1138    use_count: std::sync::atomic::AtomicU64,
1139    /// Average execution time in microseconds.
1140    avg_execution_us: std::sync::atomic::AtomicU64,
1141}
1142
1143/// Compute a hash for a string.
1144fn compute_hash(s: &str) -> u64 {
1145    let mut hasher = std::collections::hash_map::DefaultHasher::new();
1146    s.hash(&mut hasher);
1147    hasher.finish()
1148}
1149
1150impl ExecutionPlan {
1151    /// Create a new execution plan.
1152    pub fn new(sql: impl AsRef<str>, hint: PlanHint) -> Self {
1153        let sql_str = sql.as_ref();
1154        Self {
1155            sql: Arc::from(sql_str),
1156            hash: compute_hash(sql_str),
1157            hint,
1158            estimated_cost: None,
1159            use_count: std::sync::atomic::AtomicU64::new(0),
1160            avg_execution_us: std::sync::atomic::AtomicU64::new(0),
1161        }
1162    }
1163
1164    /// Create with estimated cost.
1165    pub fn with_cost(sql: impl AsRef<str>, hint: PlanHint, cost: f64) -> Self {
1166        let sql_str = sql.as_ref();
1167        Self {
1168            sql: Arc::from(sql_str),
1169            hash: compute_hash(sql_str),
1170            hint,
1171            estimated_cost: Some(cost),
1172            use_count: std::sync::atomic::AtomicU64::new(0),
1173            avg_execution_us: std::sync::atomic::AtomicU64::new(0),
1174        }
1175    }
1176
1177    /// Record an execution with timing.
1178    pub fn record_execution(&self, duration_us: u64) {
1179        let old_count = self
1180            .use_count
1181            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1182        let old_avg = self
1183            .avg_execution_us
1184            .load(std::sync::atomic::Ordering::Relaxed);
1185
1186        // Update running average
1187        let new_avg = if old_count == 0 {
1188            duration_us
1189        } else {
1190            // Weighted average: (old_avg * old_count + new_value) / (old_count + 1)
1191            (old_avg * old_count + duration_us) / (old_count + 1)
1192        };
1193
1194        self.avg_execution_us
1195            .store(new_avg, std::sync::atomic::Ordering::Relaxed);
1196    }
1197
1198    /// Get the use count.
1199    pub fn use_count(&self) -> u64 {
1200        self.use_count.load(std::sync::atomic::Ordering::Relaxed)
1201    }
1202
1203    /// Get the average execution time in microseconds.
1204    pub fn avg_execution_us(&self) -> u64 {
1205        self.avg_execution_us
1206            .load(std::sync::atomic::Ordering::Relaxed)
1207    }
1208}
1209
1210/// Cache for query execution plans.
1211///
1212/// This cache stores not just SQL strings but also execution hints and
1213/// performance metrics for each query, enabling adaptive optimization.
1214///
1215/// # Example
1216///
1217/// ```rust
1218/// use prax_query::cache::{ExecutionPlanCache, PlanHint};
1219///
1220/// let cache = ExecutionPlanCache::new(1000);
1221///
1222/// // Register a plan with an index hint
1223/// let plan = cache.register(
1224///     "find_user_by_email",
1225///     "SELECT * FROM users WHERE email = $1",
1226///     PlanHint::IndexScan("idx_users_email".into()),
1227/// );
1228///
1229/// // Get the plan later
1230/// if let Some(plan) = cache.get("find_user_by_email") {
1231///     println!("Using plan with hint: {:?}", plan.hint);
1232/// }
1233/// ```
1234#[derive(Debug)]
1235pub struct ExecutionPlanCache {
1236    /// Maximum number of plans to cache.
1237    max_size: usize,
1238    /// Cached plans.
1239    plans: parking_lot::RwLock<HashMap<u64, Arc<ExecutionPlan>>>,
1240    /// Key to hash lookup.
1241    key_index: parking_lot::RwLock<HashMap<Cow<'static, str>, u64>>,
1242}
1243
1244impl ExecutionPlanCache {
1245    /// Create a new execution plan cache.
1246    pub fn new(max_size: usize) -> Self {
1247        Self {
1248            max_size,
1249            plans: parking_lot::RwLock::new(HashMap::with_capacity(max_size / 2)),
1250            key_index: parking_lot::RwLock::new(HashMap::with_capacity(max_size / 2)),
1251        }
1252    }
1253
1254    /// Register a new execution plan.
1255    pub fn register(
1256        &self,
1257        key: impl Into<Cow<'static, str>>,
1258        sql: impl AsRef<str>,
1259        hint: PlanHint,
1260    ) -> Arc<ExecutionPlan> {
1261        let key = key.into();
1262        let plan = Arc::new(ExecutionPlan::new(sql, hint));
1263        let hash = plan.hash;
1264
1265        let mut plans = self.plans.write();
1266        let mut key_index = self.key_index.write();
1267
1268        // Evict if at capacity
1269        if plans.len() >= self.max_size && !plans.contains_key(&hash) {
1270            // Simple eviction: remove least used
1271            if let Some((&evict_hash, _)) = plans.iter().min_by_key(|(_, p)| p.use_count()) {
1272                plans.remove(&evict_hash);
1273                key_index.retain(|_, &mut v| v != evict_hash);
1274            }
1275        }
1276
1277        plans.insert(hash, Arc::clone(&plan));
1278        key_index.insert(key, hash);
1279
1280        plan
1281    }
1282
1283    /// Register a plan with estimated cost.
1284    pub fn register_with_cost(
1285        &self,
1286        key: impl Into<Cow<'static, str>>,
1287        sql: impl AsRef<str>,
1288        hint: PlanHint,
1289        cost: f64,
1290    ) -> Arc<ExecutionPlan> {
1291        let key = key.into();
1292        let plan = Arc::new(ExecutionPlan::with_cost(sql, hint, cost));
1293        let hash = plan.hash;
1294
1295        let mut plans = self.plans.write();
1296        let mut key_index = self.key_index.write();
1297
1298        if plans.len() >= self.max_size
1299            && !plans.contains_key(&hash)
1300            && let Some((&evict_hash, _)) = plans.iter().min_by_key(|(_, p)| p.use_count())
1301        {
1302            plans.remove(&evict_hash);
1303            key_index.retain(|_, &mut v| v != evict_hash);
1304        }
1305
1306        plans.insert(hash, Arc::clone(&plan));
1307        key_index.insert(key, hash);
1308
1309        plan
1310    }
1311
1312    /// Get a cached execution plan.
1313    pub fn get(&self, key: &str) -> Option<Arc<ExecutionPlan>> {
1314        let hash = {
1315            let key_index = self.key_index.read();
1316            *key_index.get(key)?
1317        };
1318
1319        self.plans.read().get(&hash).cloned()
1320    }
1321
1322    /// Get a plan by its hash.
1323    pub fn get_by_hash(&self, hash: u64) -> Option<Arc<ExecutionPlan>> {
1324        self.plans.read().get(&hash).cloned()
1325    }
1326
1327    /// Get or create a plan.
1328    pub fn get_or_register<F>(
1329        &self,
1330        key: impl Into<Cow<'static, str>>,
1331        sql_fn: F,
1332        hint: PlanHint,
1333    ) -> Arc<ExecutionPlan>
1334    where
1335        F: FnOnce() -> String,
1336    {
1337        let key = key.into();
1338
1339        // Fast path: check if exists
1340        if let Some(plan) = self.get(key.as_ref()) {
1341            return plan;
1342        }
1343
1344        // Slow path: create and register
1345        self.register(key, sql_fn(), hint)
1346    }
1347
1348    /// Record execution timing for a plan.
1349    pub fn record_execution(&self, key: &str, duration_us: u64) {
1350        if let Some(plan) = self.get(key) {
1351            plan.record_execution(duration_us);
1352        }
1353    }
1354
1355    /// Get plans sorted by average execution time (slowest first).
1356    pub fn slowest_queries(&self, limit: usize) -> Vec<Arc<ExecutionPlan>> {
1357        let plans = self.plans.read();
1358        let mut sorted: Vec<_> = plans.values().cloned().collect();
1359        sorted.sort_by_key(|a| std::cmp::Reverse(a.avg_execution_us()));
1360        sorted.truncate(limit);
1361        sorted
1362    }
1363
1364    /// Get plans sorted by use count (most used first).
1365    pub fn most_used(&self, limit: usize) -> Vec<Arc<ExecutionPlan>> {
1366        let plans = self.plans.read();
1367        let mut sorted: Vec<_> = plans.values().cloned().collect();
1368        sorted.sort_by_key(|a| std::cmp::Reverse(a.use_count()));
1369        sorted.truncate(limit);
1370        sorted
1371    }
1372
1373    /// Clear all cached plans.
1374    pub fn clear(&self) {
1375        self.plans.write().clear();
1376        self.key_index.write().clear();
1377    }
1378
1379    /// Get the number of cached plans.
1380    pub fn len(&self) -> usize {
1381        self.plans.read().len()
1382    }
1383
1384    /// Check if the cache is empty.
1385    pub fn is_empty(&self) -> bool {
1386        self.plans.read().is_empty()
1387    }
1388}