1use 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#[derive(Debug)]
36pub struct QueryCache {
37 max_size: usize,
39 cache: RwLock<HashMap<QueryKey, CachedQuery>>,
41 stats: AtomicCacheStats,
48}
49
50#[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#[derive(Debug, Clone, PartialEq, Eq, Hash)]
104pub struct QueryKey {
105 key: Cow<'static, str>,
107}
108
109impl QueryKey {
110 #[inline]
112 pub const fn new(key: &'static str) -> Self {
113 Self {
114 key: Cow::Borrowed(key),
115 }
116 }
117
118 #[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#[derive(Debug)]
141pub struct CachedQuery {
142 pub sql: String,
144 pub param_count: usize,
146 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 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 #[inline]
175 fn record_access(&self) {
176 self.access_count.fetch_add(1, Ordering::Relaxed);
177 }
178
179 #[inline]
181 pub fn sql(&self) -> &str {
182 &self.sql
183 }
184
185 #[inline]
187 pub fn param_count(&self) -> usize {
188 self.param_count
189 }
190}
191
192#[derive(Debug, Default, Clone)]
194pub struct CacheStats {
195 pub hits: u64,
197 pub misses: u64,
199 pub evictions: u64,
201 pub insertions: u64,
203}
204
205impl CacheStats {
206 #[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 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 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 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 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 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 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 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 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 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 if let Some(sql) = self.get(key.clone()) {
319 return sql;
320 }
321
322 let sql = f();
324 self.insert(key, sql.clone());
325 sql
326 }
327
328 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 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 pub fn clear(&self) {
344 let mut cache = self.cache.write().unwrap();
345 cache.clear();
346 }
347
348 pub fn len(&self) -> usize {
350 let cache = self.cache.read().unwrap();
351 cache.len()
352 }
353
354 pub fn is_empty(&self) -> bool {
356 self.len() == 0
357 }
358
359 pub fn max_size(&self) -> usize {
361 self.max_size
362 }
363
364 pub fn stats(&self) -> CacheStats {
366 self.stats.snapshot()
367 }
368
369 pub fn reset_stats(&self) {
371 self.stats.reset();
372 }
373
374 fn evict_lru(&self, cache: &mut HashMap<QueryKey, CachedQuery>) {
376 if cache.is_empty() {
380 return;
381 }
382 let to_evict = (cache.len() / 4).max(1);
383
384 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
410fn 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 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 count += 1;
435 }
436 }
437
438 count
439}
440
441#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
443pub struct QueryHash(u64);
444
445impl QueryHash {
446 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 #[inline]
455 pub fn value(&self) -> u64 {
456 self.0
457 }
458}
459
460pub mod patterns {
462 use super::QueryKey;
463
464 #[inline]
466 pub fn select_by_id(table: &str) -> QueryKey {
467 QueryKey::owned(format!("select_by_id:{}", table))
468 }
469
470 #[inline]
472 pub fn select_all(table: &str) -> QueryKey {
473 QueryKey::owned(format!("select_all:{}", table))
474 }
475
476 #[inline]
478 pub fn insert(table: &str, columns: usize) -> QueryKey {
479 QueryKey::owned(format!("insert:{}:{}", table, columns))
480 }
481
482 #[inline]
484 pub fn update_by_id(table: &str, columns: usize) -> QueryKey {
485 QueryKey::owned(format!("update_by_id:{}:{}", table, columns))
486 }
487
488 #[inline]
490 pub fn delete_by_id(table: &str) -> QueryKey {
491 QueryKey::owned(format!("delete_by_id:{}", table))
492 }
493
494 #[inline]
496 pub fn count(table: &str) -> QueryKey {
497 QueryKey::owned(format!("count:{}", table))
498 }
499
500 #[inline]
502 pub fn count_filtered(table: &str, filter_hash: u64) -> QueryKey {
503 QueryKey::owned(format!("count:{}:{}", table, filter_hash))
504 }
505}
506
507#[derive(Debug)]
537pub struct SqlTemplateCache {
538 max_size: usize,
540 templates: parking_lot::RwLock<HashMap<u64, Arc<SqlTemplate>>>,
542 key_index: parking_lot::RwLock<HashMap<Cow<'static, str>, u64>>,
544 stats: parking_lot::RwLock<CacheStats>,
546}
547
548#[derive(Debug)]
550pub struct SqlTemplate {
551 pub sql: Arc<str>,
553 pub hash: u64,
555 pub param_count: usize,
557 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 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 #[inline(always)]
596 pub fn sql(&self) -> &str {
597 &self.sql
598 }
599
600 #[inline(always)]
602 pub fn sql_arc(&self) -> Arc<str> {
603 Arc::clone(&self.sql)
604 }
605
606 #[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 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 #[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 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 #[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 #[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); 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 #[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 return Some(Arc::clone(template));
729 }
730 None
731 }
732
733 #[inline]
735 pub fn get_sql(&self, key: &str) -> Option<Arc<str>> {
736 self.get(key).map(|t| t.sql_arc())
737 }
738
739 #[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 if let Some(template) = self.get(&key) {
749 return template;
750 }
751
752 let sql = f();
754 self.register(key, sql)
755 }
756
757 #[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 pub fn stats(&self) -> CacheStats {
766 self.stats.read().clone()
767 }
768
769 pub fn len(&self) -> usize {
771 self.templates.read().len()
772 }
773
774 pub fn is_empty(&self) -> bool {
776 self.len() == 0
777 }
778
779 pub fn clear(&self) {
781 self.templates.write().clear();
782 self.key_index.write().clear();
783 }
784
785 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 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 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
826static GLOBAL_TEMPLATE_CACHE: std::sync::OnceLock<SqlTemplateCache> = std::sync::OnceLock::new();
849
850#[inline(always)]
852pub fn global_template_cache() -> &'static SqlTemplateCache {
853 GLOBAL_TEMPLATE_CACHE.get_or_init(|| SqlTemplateCache::new(10000))
854}
855
856#[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#[inline(always)]
867pub fn get_global_template(key: &str) -> Option<Arc<SqlTemplate>> {
868 global_template_cache().get(key)
869}
870
871#[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"); }
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"); cache.get("test"); cache.get("missing"); 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 for _ in 0..5 {
934 cache.get("a");
935 }
936
937 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 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 #[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 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 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"); cache.get("q1"); cache.get("missing"); 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 let template = register_global_template("global_test", "SELECT * FROM global");
1061 assert_eq!(template.sql(), "SELECT * FROM global");
1062
1063 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 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 let result = cache.get("users_by_email");
1093 assert!(result.is_some());
1094 assert!(matches!(result.unwrap().hint, PlanHint::IndexScan(_)));
1095 }
1096}
1097
1098#[derive(Debug, Clone, Default)]
1108pub enum PlanHint {
1109 #[default]
1111 None,
1112 IndexScan(String),
1114 SeqScan,
1116 Parallel(u32),
1118 CachePlan,
1120 Timeout(std::time::Duration),
1122 Custom(String),
1124}
1125
1126#[derive(Debug)]
1128pub struct ExecutionPlan {
1129 pub sql: Arc<str>,
1131 pub hash: u64,
1133 pub hint: PlanHint,
1135 pub estimated_cost: Option<f64>,
1137 use_count: std::sync::atomic::AtomicU64,
1139 avg_execution_us: std::sync::atomic::AtomicU64,
1141}
1142
1143fn 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 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 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 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 let new_avg = if old_count == 0 {
1188 duration_us
1189 } else {
1190 (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 pub fn use_count(&self) -> u64 {
1200 self.use_count.load(std::sync::atomic::Ordering::Relaxed)
1201 }
1202
1203 pub fn avg_execution_us(&self) -> u64 {
1205 self.avg_execution_us
1206 .load(std::sync::atomic::Ordering::Relaxed)
1207 }
1208}
1209
1210#[derive(Debug)]
1235pub struct ExecutionPlanCache {
1236 max_size: usize,
1238 plans: parking_lot::RwLock<HashMap<u64, Arc<ExecutionPlan>>>,
1240 key_index: parking_lot::RwLock<HashMap<Cow<'static, str>, u64>>,
1242}
1243
1244impl ExecutionPlanCache {
1245 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 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 if plans.len() >= self.max_size && !plans.contains_key(&hash) {
1270 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 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 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 pub fn get_by_hash(&self, hash: u64) -> Option<Arc<ExecutionPlan>> {
1324 self.plans.read().get(&hash).cloned()
1325 }
1326
1327 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 if let Some(plan) = self.get(key.as_ref()) {
1341 return plan;
1342 }
1343
1344 self.register(key, sql_fn(), hint)
1346 }
1347
1348 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 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 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 pub fn clear(&self) {
1375 self.plans.write().clear();
1376 self.key_index.write().clear();
1377 }
1378
1379 pub fn len(&self) -> usize {
1381 self.plans.read().len()
1382 }
1383
1384 pub fn is_empty(&self) -> bool {
1386 self.plans.read().is_empty()
1387 }
1388}