1use std::collections::HashMap;
46use std::hash::Hash;
47use std::sync::RwLock;
48
49#[derive(Debug, Clone)]
55pub struct GraphEdge {
56 pub parent_field: String,
58 pub relation: String,
60 pub sub_graph: Option<Box<EntityGraph>>,
62}
63
64#[derive(Debug, Clone, Default)]
78pub struct EntityGraph {
79 edges: Vec<GraphEdge>,
81}
82
83impl EntityGraph {
84 pub fn new() -> Self {
86 Self { edges: Vec::new() }
87 }
88
89 pub fn add_edge(
91 &mut self,
92 parent_field: impl Into<String>,
93 relation: impl Into<String>,
94 ) -> &mut Self {
95 self.edges.push(GraphEdge {
96 parent_field: parent_field.into(),
97 relation: relation.into(),
98 sub_graph: None,
99 });
100 self
101 }
102
103 pub fn add_edge_with_graph(
105 &mut self,
106 parent_field: impl Into<String>,
107 relation: impl Into<String>,
108 sub_graph: EntityGraph,
109 ) -> &mut Self {
110 self.edges.push(GraphEdge {
111 parent_field: parent_field.into(),
112 relation: relation.into(),
113 sub_graph: Some(Box::new(sub_graph)),
114 });
115 self
116 }
117
118 pub fn edges(&self) -> &[GraphEdge] {
120 &self.edges
121 }
122
123 pub fn edge_count(&self) -> usize {
125 self.edges.len()
126 }
127
128 pub fn relations_of(&self, parent_field: &str) -> Vec<&GraphEdge> {
130 self.edges
131 .iter()
132 .filter(|e| e.parent_field == parent_field)
133 .collect()
134 }
135
136 pub fn all_relations(&self) -> Vec<String> {
138 let mut rels: Vec<String> = self.edges.iter().map(|e| e.relation.clone()).collect();
139 rels.sort();
140 rels.dedup();
141 rels
142 }
143
144 pub fn all_parent_fields(&self) -> Vec<String> {
146 let mut fields: Vec<String> = self.edges.iter().map(|e| e.parent_field.clone()).collect();
147 fields.sort();
148 fields.dedup();
149 fields
150 }
151
152 pub fn is_empty(&self) -> bool {
154 self.edges.is_empty()
155 }
156
157 pub fn all_relations_recursive(&self) -> Vec<String> {
159 let mut result = Vec::new();
160 for edge in &self.edges {
161 result.push(edge.relation.clone());
162 if let Some(sub) = &edge.sub_graph {
163 result.extend(sub.all_relations_recursive());
164 }
165 }
166 result.sort();
167 result.dedup();
168 result
169 }
170
171 pub fn detect_cycles(&self) -> Result<(), Vec<String>> {
211 let mut adj: std::collections::HashMap<String, Vec<String>> =
213 std::collections::HashMap::new();
214 self.collect_edges_recursive(&mut adj);
215
216 for neighbors in adj.values_mut() {
218 neighbors.sort();
219 }
220
221 let mut visited = std::collections::HashSet::new();
223 let mut visiting = std::collections::HashSet::new();
224 let mut path = Vec::new();
225
226 let mut sorted_nodes: Vec<String> = adj.keys().cloned().collect();
228 sorted_nodes.sort();
229 for node in &sorted_nodes {
230 if !visited.contains(node) {
231 dfs_cycle_detect(node, &adj, &mut visited, &mut visiting, &mut path)?;
232 }
233 }
234 Ok(())
235 }
236
237 fn collect_edges_recursive(&self, adj: &mut std::collections::HashMap<String, Vec<String>>) {
242 for edge in &self.edges {
243 adj.entry(edge.parent_field.clone())
244 .or_default()
245 .push(edge.relation.clone());
246 if let Some(sub) = &edge.sub_graph {
247 sub.collect_edges_recursive(adj);
248 }
249 }
250 }
251
252 pub fn detect_duplicate_edges(&self) -> Result<(), Vec<(String, String)>> {
258 let mut seen = std::collections::HashSet::new();
259 let mut duplicates = Vec::new();
260 for edge in &self.edges {
261 let key = (edge.parent_field.clone(), edge.relation.clone());
262 if !seen.insert(key.clone()) {
263 duplicates.push((edge.parent_field.clone(), edge.relation.clone()));
264 }
265 }
266 if duplicates.is_empty() {
267 Ok(())
268 } else {
269 Err(duplicates)
270 }
271 }
272
273 pub fn validate(&self) -> Result<(), String> {
277 if let Err(cycle) = self.detect_cycles() {
279 return Err(format!(
280 "EntityGraph 循环引用检测失败:{}",
281 cycle.join(" → ")
282 ));
283 }
284 if let Err(duplicates) = self.detect_duplicate_edges() {
286 let dup_str: Vec<String> = duplicates
287 .iter()
288 .map(|(p, r)| format!("({}->{})", p, r))
289 .collect();
290 return Err(format!(
291 "EntityGraph 重复边检测失败:{}",
292 dup_str.join(", ")
293 ));
294 }
295 Ok(())
296 }
297}
298
299fn dfs_cycle_detect(
304 node: &str,
305 adj: &std::collections::HashMap<String, Vec<String>>,
306 visited: &mut std::collections::HashSet<String>,
307 visiting: &mut std::collections::HashSet<String>,
308 path: &mut Vec<String>,
309) -> Result<(), Vec<String>> {
310 if visiting.contains(node) {
312 let cycle_start = path.iter().position(|n| n == node).unwrap_or(0);
313 let mut cycle = path[cycle_start..].to_vec();
314 cycle.push(node.to_string());
315 return Err(cycle);
316 }
317 if visited.contains(node) {
319 return Ok(());
320 }
321
322 visiting.insert(node.to_string());
324 path.push(node.to_string());
325
326 if let Some(neighbors) = adj.get(node) {
328 for neighbor in neighbors {
329 dfs_cycle_detect(neighbor, adj, visited, visiting, path)?;
330 }
331 }
332
333 visiting.remove(node);
335 visited.insert(node.to_string());
336 path.pop();
337 Ok(())
338}
339
340#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
346pub enum BatchStrategy {
347 #[default]
351 In,
352 Join,
356 Subquery,
360}
361
362impl BatchStrategy {
363 pub fn name(&self) -> &'static str {
365 match self {
366 BatchStrategy::In => "in",
367 BatchStrategy::Join => "join",
368 BatchStrategy::Subquery => "subquery",
369 }
370 }
371
372 pub fn render_in_clause(column: &str, placeholders: usize) -> String {
376 if placeholders == 0 {
377 return format!("{} IN ()", column);
378 }
379 let marks: Vec<&str> = vec!["?"; placeholders];
380 format!("{} IN ({})", column, marks.join(", "))
381 }
382}
383
384#[derive(Debug, Clone, Copy)]
392pub struct BatchSizeConfig {
393 pub size: usize,
395 pub strategy: BatchStrategy,
397}
398
399impl Default for BatchSizeConfig {
400 fn default() -> Self {
401 Self {
402 size: 100,
403 strategy: BatchStrategy::In,
404 }
405 }
406}
407
408impl BatchSizeConfig {
409 pub fn new(size: usize, strategy: BatchStrategy) -> Self {
411 Self { size, strategy }
412 }
413
414 pub fn with_size(size: usize) -> Self {
416 Self {
417 size,
418 strategy: BatchStrategy::In,
419 }
420 }
421
422 pub fn batch_count(&self, total: usize) -> usize {
437 if total == 0 {
438 0
439 } else {
440 total.div_ceil(self.size)
441 }
442 }
443
444 pub fn batch_range(&self, batch_index: usize, total: usize) -> std::ops::Range<usize> {
457 let start = batch_index * self.size;
458 let end = (start + self.size).min(total);
459 start..end
460 }
461}
462
463pub type BatchLoaderFn<K, V> = Box<dyn Fn(&[K]) -> HashMap<K, V> + Send + Sync>;
469
470pub struct BatchLoader<K, V>
494where
495 K: Hash + Eq + Clone + Send + Sync,
496 V: Clone + Send + Sync,
497{
498 batch_size: usize,
500 loader: BatchLoaderFn<K, V>,
502 cache: RwLock<HashMap<K, V>>,
504}
505
506impl<K, V> BatchLoader<K, V>
507where
508 K: Hash + Eq + Clone + Send + Sync,
509 V: Clone + Send + Sync,
510{
511 pub fn new(batch_size: usize, loader: BatchLoaderFn<K, V>) -> Self {
517 Self {
518 batch_size,
519 loader,
520 cache: RwLock::new(HashMap::new()),
521 }
522 }
523
524 pub fn load_many(&self, keys: &[K]) -> HashMap<K, V> {
530 let mut result: HashMap<K, V> = HashMap::new();
531
532 let mut to_load: Vec<K> = Vec::new();
534 if let Ok(cached) = self.cache.read() {
535 for k in keys {
536 if let Some(v) = cached.get(k) {
537 result.insert(k.clone(), v.clone());
538 } else {
539 to_load.push(k.clone());
540 }
541 }
542 } else {
543 to_load.extend(keys.iter().cloned());
544 }
545
546 if to_load.is_empty() {
547 return result;
548 }
549
550 let batch_size = self.batch_size.max(1);
552 let mut all_loaded: HashMap<K, V> = HashMap::new();
553 for chunk in to_load.chunks(batch_size) {
554 let loaded = (self.loader)(chunk);
555 all_loaded.extend(loaded);
556 }
557
558 if let Ok(mut cache) = self.cache.write() {
560 for (k, v) in &all_loaded {
561 cache.insert(k.clone(), v.clone());
562 }
563 }
564
565 result.extend(all_loaded);
567 result
568 }
569
570 pub fn load_one(&self, key: &K) -> Option<V> {
572 let result = self.load_many(std::slice::from_ref(key));
573 result.get(key).cloned()
574 }
575
576 pub fn clear_cache(&self) {
578 if let Ok(mut cache) = self.cache.write() {
579 cache.clear();
580 }
581 }
582
583 pub fn cache_size(&self) -> usize {
585 match self.cache.read() {
586 Ok(g) => g.len(),
587 Err(_) => 0,
588 }
589 }
590
591 pub fn batch_size(&self) -> usize {
593 self.batch_size
594 }
595}
596
597pub struct N1QueryDetector {
642 config: N1DetectionConfig,
644 counts: RwLock<HashMap<String, u64>>,
646 batch_counts: RwLock<HashMap<String, u64>>,
648 window_active: RwLock<bool>,
650 alerts: RwLock<Vec<N1Alert>>,
652}
653
654#[derive(Debug, Clone)]
656pub struct N1DetectionConfig {
657 pub threshold: u64,
659 pub enabled: bool,
661}
662
663impl Default for N1DetectionConfig {
664 fn default() -> Self {
665 Self {
666 threshold: 5,
667 enabled: true,
668 }
669 }
670}
671
672impl N1DetectionConfig {
673 pub fn new() -> Self {
675 Self::default()
676 }
677
678 pub fn with_threshold(mut self, threshold: u64) -> Self {
680 self.threshold = threshold.max(1);
681 self
682 }
683
684 pub fn with_enabled(mut self, enabled: bool) -> Self {
686 self.enabled = enabled;
687 self
688 }
689}
690
691#[derive(Debug, Clone, PartialEq, Eq)]
693pub struct N1Alert {
694 pub relation: String,
696 pub query_count: u64,
698 pub batch_count: u64,
700 pub threshold: u64,
702}
703
704impl N1Alert {
705 pub fn no_batch_used(&self) -> bool {
707 self.batch_count == 0
708 }
709
710 pub fn suggested_batch_size(&self) -> usize {
712 let n = self.query_count as usize;
713 if n <= 50 {
714 50
715 } else if n <= 100 {
716 100
717 } else if n <= 500 {
718 500
719 } else {
720 1000
721 }
722 }
723}
724
725impl N1QueryDetector {
726 pub fn new(config: N1DetectionConfig) -> Self {
728 Self {
729 config,
730 counts: RwLock::new(HashMap::new()),
731 batch_counts: RwLock::new(HashMap::new()),
732 window_active: RwLock::new(false),
733 alerts: RwLock::new(Vec::new()),
734 }
735 }
736
737 pub fn with_defaults() -> Self {
739 Self::new(N1DetectionConfig::default())
740 }
741
742 pub fn is_enabled(&self) -> bool {
744 self.config.enabled
745 }
746
747 pub fn threshold(&self) -> u64 {
749 self.config.threshold
750 }
751
752 pub fn start_window(&self) {
756 if !self.config.enabled {
757 return;
758 }
759 if let Ok(mut counts) = self.counts.write() {
760 *counts = HashMap::new();
761 }
762 if let Ok(mut batch_counts) = self.batch_counts.write() {
763 *batch_counts = HashMap::new();
764 }
765 if let Ok(mut alerts) = self.alerts.write() {
766 *alerts = Vec::new();
767 }
768 if let Ok(mut window_active) = self.window_active.write() {
769 *window_active = true;
770 }
771 }
772
773 pub fn end_window(&self) -> Vec<N1Alert> {
778 if !self.config.enabled {
779 return Vec::new();
780 }
781 if let Ok(mut window_active) = self.window_active.write() {
782 *window_active = false;
783 }
784
785 let new_alerts: Vec<N1Alert> = match (self.counts.read(), self.batch_counts.read()) {
787 (Ok(counts), Ok(batch_counts)) => {
788 let mut alerts: Vec<N1Alert> = counts
789 .iter()
790 .filter_map(|(rel, &cnt)| {
791 if cnt >= self.config.threshold {
792 Some(N1Alert {
793 relation: rel.clone(),
794 query_count: cnt,
795 batch_count: *batch_counts.get(rel).unwrap_or(&0),
796 threshold: self.config.threshold,
797 })
798 } else {
799 None
800 }
801 })
802 .collect();
803 alerts.sort_by(|a, b| a.relation.cmp(&b.relation));
805 alerts
806 }
807 _ => Vec::new(),
808 };
809
810 if let Ok(mut alerts) = self.alerts.write() {
811 *alerts = new_alerts.clone();
812 }
813 new_alerts
814 }
815
816 pub fn record_single_load(&self, relation: &str) {
818 if !self.config.enabled {
819 return;
820 }
821 {
822 let active = self.window_active.read().map(|g| *g).unwrap_or(false);
823 if !active {
824 return;
825 }
826 }
827 if let Ok(mut counts) = self.counts.write() {
828 *counts.entry(relation.to_string()).or_insert(0) += 1;
829 }
830 }
831
832 pub fn record_batch_load(&self, relation: &str, _keys_count: usize) {
837 if !self.config.enabled {
838 return;
839 }
840 {
841 let active = self.window_active.read().map(|g| *g).unwrap_or(false);
842 if !active {
843 return;
844 }
845 }
846 if let Ok(mut batch_counts) = self.batch_counts.write() {
847 *batch_counts.entry(relation.to_string()).or_insert(0) += 1;
848 }
849 }
850
851 pub fn alerts(&self) -> Vec<N1Alert> {
853 self.alerts.read().map(|g| g.clone()).unwrap_or_default()
854 }
855
856 pub fn current_count(&self, relation: &str) -> u64 {
858 self.counts
859 .read()
860 .map(|g| g.get(relation).copied().unwrap_or(0))
861 .unwrap_or(0)
862 }
863
864 pub fn current_batch_count(&self, relation: &str) -> u64 {
866 self.batch_counts
867 .read()
868 .map(|g| g.get(relation).copied().unwrap_or(0))
869 .unwrap_or(0)
870 }
871
872 pub fn is_window_active(&self) -> bool {
874 self.window_active.read().map(|g| *g).unwrap_or(false)
875 }
876
877 pub fn has_n_plus_one(&self) -> bool {
879 !self.alerts().is_empty()
880 }
881}
882
883impl Default for N1QueryDetector {
884 fn default() -> Self {
885 Self::with_defaults()
886 }
887}
888
889#[cfg(test)]
894mod tests {
895 use super::*;
896
897 #[test]
900 fn test_new_graph_is_empty() {
901 let g = EntityGraph::new();
902 assert!(g.is_empty());
903 assert_eq!(g.edge_count(), 0);
904 }
905
906 #[test]
907 fn test_add_edge() {
908 let mut g = EntityGraph::new();
909 g.add_edge("user", "posts");
910 assert_eq!(g.edge_count(), 1);
911 assert!(!g.is_empty());
912 }
913
914 #[test]
915 fn test_add_multiple_edges() {
916 let mut g = EntityGraph::new();
917 g.add_edge("user", "posts")
918 .add_edge("user", "profile")
919 .add_edge("user", "comments");
920 assert_eq!(g.edge_count(), 3);
921 }
922
923 #[test]
924 fn test_add_edge_with_sub_graph() {
925 let mut sub = EntityGraph::new();
926 sub.add_edge("comments", "author");
927
928 let mut g = EntityGraph::new();
929 g.add_edge_with_graph("user", "posts", sub);
930
931 assert_eq!(g.edge_count(), 1);
932 assert!(g.edges()[0].sub_graph.is_some());
933 assert_eq!(g.edges()[0].sub_graph.as_ref().unwrap().edge_count(), 1);
934 }
935
936 #[test]
937 fn test_relations_of() {
938 let mut g = EntityGraph::new();
939 g.add_edge("user", "posts")
940 .add_edge("user", "profile")
941 .add_edge("post", "comments");
942
943 let user_relations = g.relations_of("user");
944 assert_eq!(user_relations.len(), 2);
945 assert_eq!(user_relations[0].relation, "posts");
946 assert_eq!(user_relations[1].relation, "profile");
947
948 let post_relations = g.relations_of("post");
949 assert_eq!(post_relations.len(), 1);
950
951 let none = g.relations_of("nonexistent");
952 assert!(none.is_empty());
953 }
954
955 #[test]
956 fn test_all_relations() {
957 let mut g = EntityGraph::new();
958 g.add_edge("user", "posts")
959 .add_edge("user", "profile")
960 .add_edge("post", "comments");
961
962 let rels = g.all_relations();
963 assert_eq!(rels, vec!["comments", "posts", "profile"]);
964 }
965
966 #[test]
967 fn test_all_parent_fields() {
968 let mut g = EntityGraph::new();
969 g.add_edge("user", "posts")
970 .add_edge("user", "profile")
971 .add_edge("post", "comments");
972
973 let fields = g.all_parent_fields();
974 assert_eq!(fields, vec!["post", "user"]);
975 }
976
977 #[test]
978 fn test_all_relations_recursive() {
979 let mut sub = EntityGraph::new();
980 sub.add_edge("comments", "author")
981 .add_edge("comments", "likes");
982
983 let mut g = EntityGraph::new();
984 g.add_edge("user", "posts")
985 .add_edge_with_graph("user", "comments", sub);
986
987 let all = g.all_relations_recursive();
988 assert!(all.contains(&"posts".to_string()));
989 assert!(all.contains(&"comments".to_string()));
990 assert!(all.contains(&"author".to_string()));
991 assert!(all.contains(&"likes".to_string()));
992 assert_eq!(all.len(), 4);
993 }
994
995 #[test]
996 fn test_default_graph_is_empty() {
997 let g = EntityGraph::default();
998 assert!(g.is_empty());
999 }
1000
1001 #[test]
1004 fn test_strategy_name() {
1005 assert_eq!(BatchStrategy::In.name(), "in");
1006 assert_eq!(BatchStrategy::Join.name(), "join");
1007 assert_eq!(BatchStrategy::Subquery.name(), "subquery");
1008 }
1009
1010 #[test]
1011 fn test_strategy_default_is_in() {
1012 assert_eq!(BatchStrategy::default(), BatchStrategy::In);
1013 }
1014
1015 #[test]
1016 fn test_render_in_clause_empty() {
1017 let sql = BatchStrategy::render_in_clause("id", 0);
1018 assert_eq!(sql, "id IN ()");
1019 }
1020
1021 #[test]
1022 fn test_render_in_clause_single() {
1023 let sql = BatchStrategy::render_in_clause("id", 1);
1024 assert_eq!(sql, "id IN (?)");
1025 }
1026
1027 #[test]
1028 fn test_render_in_clause_multiple() {
1029 let sql = BatchStrategy::render_in_clause("user_id", 3);
1030 assert_eq!(sql, "user_id IN (?, ?, ?)");
1031 }
1032
1033 #[test]
1036 fn test_default_config() {
1037 let config = BatchSizeConfig::default();
1038 assert_eq!(config.size, 100);
1039 assert_eq!(config.strategy, BatchStrategy::In);
1040 }
1041
1042 #[test]
1043 fn test_with_size() {
1044 let config = BatchSizeConfig::with_size(50);
1045 assert_eq!(config.size, 50);
1046 assert_eq!(config.strategy, BatchStrategy::In);
1047 }
1048
1049 #[test]
1050 fn test_new_with_strategy() {
1051 let config = BatchSizeConfig::new(200, BatchStrategy::Join);
1052 assert_eq!(config.size, 200);
1053 assert_eq!(config.strategy, BatchStrategy::Join);
1054 }
1055
1056 #[test]
1057 fn test_batch_count_zero() {
1058 let config = BatchSizeConfig::with_size(100);
1059 assert_eq!(config.batch_count(0), 0);
1060 }
1061
1062 #[test]
1063 fn test_batch_count_exact_multiple() {
1064 let config = BatchSizeConfig::with_size(100);
1065 assert_eq!(config.batch_count(100), 1);
1066 assert_eq!(config.batch_count(200), 2);
1067 assert_eq!(config.batch_count(500), 5);
1068 }
1069
1070 #[test]
1071 fn test_batch_count_with_remainder() {
1072 let config = BatchSizeConfig::with_size(100);
1073 assert_eq!(config.batch_count(1), 1);
1074 assert_eq!(config.batch_count(99), 1);
1075 assert_eq!(config.batch_count(101), 2);
1076 assert_eq!(config.batch_count(150), 2);
1077 assert_eq!(config.batch_count(201), 3);
1078 }
1079
1080 #[test]
1081 fn test_batch_range() {
1082 let config = BatchSizeConfig::with_size(100);
1083
1084 assert_eq!(config.batch_range(0, 250), 0..100);
1085 assert_eq!(config.batch_range(1, 250), 100..200);
1086 assert_eq!(config.batch_range(2, 250), 200..250);
1087 }
1088
1089 #[test]
1090 fn test_batch_range_exact() {
1091 let config = BatchSizeConfig::with_size(100);
1092
1093 assert_eq!(config.batch_range(0, 100), 0..100);
1094 assert_eq!(config.batch_range(1, 100), 100..100); }
1096
1097 #[test]
1098 fn test_batch_range_small_batch() {
1099 let config = BatchSizeConfig::with_size(10);
1100
1101 assert_eq!(config.batch_range(0, 25), 0..10);
1102 assert_eq!(config.batch_range(1, 25), 10..20);
1103 assert_eq!(config.batch_range(2, 25), 20..25);
1104 }
1105
1106 fn make_loader() -> BatchLoader<i64, String> {
1109 let loader = Box::new(|ids: &[i64]| -> HashMap<i64, String> {
1110 ids.iter().map(|id| (*id, format!("user_{}", id))).collect()
1111 });
1112 BatchLoader::new(2, loader)
1113 }
1114
1115 #[test]
1116 fn test_batch_loader_load_many_single_batch() {
1117 let loader = make_loader();
1118 let result = loader.load_many(&[1, 2]);
1119 assert_eq!(result.len(), 2);
1120 assert_eq!(result.get(&1), Some(&"user_1".to_string()));
1121 assert_eq!(result.get(&2), Some(&"user_2".to_string()));
1122 }
1123
1124 #[test]
1125 fn test_batch_loader_load_many_multiple_batches() {
1126 let loader = make_loader();
1127 let result = loader.load_many(&[1, 2, 3, 4, 5]);
1129 assert_eq!(result.len(), 5);
1130 for id in 1..=5 {
1131 assert_eq!(
1132 result.get(&id),
1133 Some(&format!("user_{}", id)),
1134 "missing user {}",
1135 id
1136 );
1137 }
1138 }
1139
1140 #[test]
1141 fn test_batch_loader_load_one() {
1142 let loader = make_loader();
1143 let result = loader.load_one(&42);
1144 assert_eq!(result, Some("user_42".to_string()));
1145 }
1146
1147 #[test]
1148 fn test_batch_loader_load_one_missing() {
1149 let loader: BatchLoader<i64, String> =
1151 BatchLoader::new(10, Box::new(|_ids: &[i64]| HashMap::new()));
1152 let result = loader.load_one(&100);
1153 assert_eq!(result, None);
1154 }
1155
1156 #[test]
1157 fn test_batch_loader_caches_results() {
1158 let call_count = std::sync::Arc::new(std::sync::Mutex::new(0u32));
1159 let call_count_clone = call_count.clone();
1160
1161 let loader = Box::new(move |ids: &[i64]| -> HashMap<i64, String> {
1162 *call_count_clone.lock().unwrap() += 1;
1163 ids.iter().map(|id| (*id, format!("user_{}", id))).collect()
1164 });
1165
1166 let batch_loader = BatchLoader::new(100, loader);
1167
1168 batch_loader.load_many(&[1, 2, 3]);
1170 assert_eq!(*call_count.lock().unwrap(), 1);
1171
1172 batch_loader.load_many(&[1, 2, 3]);
1174 assert_eq!(*call_count.lock().unwrap(), 1); batch_loader.load_many(&[4, 5]);
1178 assert_eq!(*call_count.lock().unwrap(), 2);
1179 }
1180
1181 #[test]
1182 fn test_batch_loader_partial_cache_hit() {
1183 let call_count = std::sync::Arc::new(std::sync::Mutex::new(0u32));
1184 let call_count_clone = call_count.clone();
1185
1186 let loader = Box::new(move |ids: &[i64]| -> HashMap<i64, String> {
1187 *call_count_clone.lock().unwrap() += 1;
1188 ids.iter().map(|id| (*id, format!("user_{}", id))).collect()
1189 });
1190
1191 let batch_loader = BatchLoader::new(100, loader);
1192
1193 batch_loader.load_many(&[1, 2, 3]);
1195 assert_eq!(*call_count.lock().unwrap(), 1);
1196
1197 let result = batch_loader.load_many(&[1, 2, 3, 4, 5]);
1199 assert_eq!(result.len(), 5);
1200 assert_eq!(*call_count.lock().unwrap(), 2); assert_eq!(batch_loader.cache_size(), 5);
1204 }
1205
1206 #[test]
1207 fn test_batch_loader_clear_cache() {
1208 let loader = make_loader();
1209 loader.load_many(&[1, 2]);
1210 assert_eq!(loader.cache_size(), 2);
1211
1212 loader.clear_cache();
1213 assert_eq!(loader.cache_size(), 0);
1214 }
1215
1216 #[test]
1217 fn test_batch_loader_empty_input() {
1218 let loader = make_loader();
1219 let result = loader.load_many(&[]);
1220 assert!(result.is_empty());
1221 }
1222
1223 #[test]
1224 fn test_batch_loader_batch_size_attribute() {
1225 let loader = make_loader();
1226 assert_eq!(loader.batch_size(), 2);
1227 }
1228
1229 #[test]
1230 fn test_batch_loader_with_size_1() {
1231 let loader = BatchLoader::new(
1232 1,
1233 Box::new(|ids: &[i64]| ids.iter().map(|id| (*id, *id * 10)).collect()),
1234 );
1235 let result = loader.load_many(&[1, 2, 3]);
1236 assert_eq!(result.len(), 3);
1237 assert_eq!(result.get(&1), Some(&10));
1238 assert_eq!(result.get(&2), Some(&20));
1239 assert_eq!(result.get(&3), Some(&30));
1240 }
1241
1242 #[test]
1245 fn test_n1_config_default() {
1246 let cfg = N1DetectionConfig::default();
1247 assert_eq!(cfg.threshold, 5);
1248 assert!(cfg.enabled);
1249 }
1250
1251 #[test]
1252 fn test_n1_config_builder() {
1253 let cfg = N1DetectionConfig::new()
1254 .with_threshold(10)
1255 .with_enabled(false);
1256 assert_eq!(cfg.threshold, 10);
1257 assert!(!cfg.enabled);
1258
1259 let cfg2 = N1DetectionConfig::new().with_threshold(0);
1261 assert_eq!(cfg2.threshold, 1);
1262 }
1263
1264 #[test]
1265 fn test_n1_detector_default() {
1266 let det = N1QueryDetector::default();
1267 assert!(det.is_enabled());
1268 assert_eq!(det.threshold(), 5);
1269 assert!(!det.is_window_active());
1270 assert!(!det.has_n_plus_one());
1271 }
1272
1273 #[test]
1274 fn test_n1_detector_disabled_is_noop() {
1275 let det = N1QueryDetector::new(N1DetectionConfig::new().with_enabled(false));
1276 det.start_window();
1277 for _ in 0..100 {
1278 det.record_single_load("posts");
1279 }
1280 assert_eq!(det.current_count("posts"), 0);
1282 let alerts = det.end_window();
1283 assert!(alerts.is_empty());
1284 }
1285
1286 #[test]
1287 fn test_n1_detector_records_outside_window_ignored() {
1288 let det = N1QueryDetector::with_defaults();
1289 det.record_single_load("posts");
1291 assert_eq!(det.current_count("posts"), 0);
1292 }
1293
1294 #[test]
1295 fn test_n1_detector_below_threshold_no_alert() {
1296 let det = N1QueryDetector::with_defaults(); det.start_window();
1298 for _ in 0..4 {
1299 det.record_single_load("posts");
1300 }
1301 assert_eq!(det.current_count("posts"), 4);
1302 let alerts = det.end_window();
1303 assert!(alerts.is_empty(), "below threshold should not alert");
1304 assert!(!det.has_n_plus_one());
1305 }
1306
1307 #[test]
1308 fn test_n1_detector_at_threshold_triggers_alert() {
1309 let det = N1QueryDetector::with_defaults(); det.start_window();
1311 for _ in 0..5 {
1312 det.record_single_load("posts");
1313 }
1314 let alerts = det.end_window();
1315 assert_eq!(alerts.len(), 1);
1316 assert_eq!(alerts[0].relation, "posts");
1317 assert_eq!(alerts[0].query_count, 5);
1318 assert_eq!(alerts[0].threshold, 5);
1319 assert_eq!(alerts[0].batch_count, 0);
1320 assert!(alerts[0].no_batch_used());
1321 assert!(det.has_n_plus_one());
1322 }
1323
1324 #[test]
1325 fn test_n1_detector_above_threshold_triggers_alert() {
1326 let det = N1QueryDetector::with_defaults();
1327 det.start_window();
1328 for _ in 0..10 {
1329 det.record_single_load("posts");
1330 }
1331 let alerts = det.end_window();
1332 assert_eq!(alerts.len(), 1);
1333 assert_eq!(alerts[0].query_count, 10);
1334 assert!(alerts[0].suggested_batch_size() >= 50);
1336 }
1337
1338 #[test]
1339 fn test_n1_detector_multiple_relations() {
1340 let det = N1QueryDetector::with_defaults();
1341 det.start_window();
1342 for _ in 0..6 {
1343 det.record_single_load("posts");
1344 }
1345 for _ in 0..3 {
1346 det.record_single_load("comments"); }
1348 for _ in 0..8 {
1349 det.record_single_load("tags");
1350 }
1351 let alerts = det.end_window();
1352 assert_eq!(alerts.len(), 2);
1354 assert_eq!(alerts[0].relation, "posts");
1356 assert_eq!(alerts[0].query_count, 6);
1357 assert_eq!(alerts[1].relation, "tags");
1358 assert_eq!(alerts[1].query_count, 8);
1359 }
1360
1361 #[test]
1362 fn test_n1_detector_batch_load_recorded_separately() {
1363 let det = N1QueryDetector::with_defaults();
1364 det.start_window();
1365 for _ in 0..6 {
1367 det.record_single_load("posts");
1368 }
1369 det.record_batch_load("posts", 100);
1371 det.record_batch_load("posts", 50);
1372 let alerts = det.end_window();
1373 assert_eq!(alerts.len(), 1);
1374 assert_eq!(alerts[0].query_count, 6);
1375 assert_eq!(alerts[0].batch_count, 2);
1376 assert!(!alerts[0].no_batch_used());
1378 }
1379
1380 #[test]
1381 fn test_n1_detector_batch_only_does_not_trigger() {
1382 let det = N1QueryDetector::with_defaults();
1384 det.start_window();
1385 for _ in 0..100 {
1386 det.record_batch_load("posts", 50);
1387 }
1388 assert_eq!(det.current_batch_count("posts"), 100);
1389 assert_eq!(det.current_count("posts"), 0);
1390 let alerts = det.end_window();
1391 assert!(alerts.is_empty());
1392 }
1393
1394 #[test]
1395 fn test_n1_detector_start_window_resets() {
1396 let det = N1QueryDetector::with_defaults();
1397 det.start_window();
1398 for _ in 0..10 {
1399 det.record_single_load("posts");
1400 }
1401 let _ = det.end_window();
1402 assert_eq!(det.alerts().len(), 1);
1403
1404 det.start_window();
1406 assert_eq!(det.alerts().len(), 0);
1407 assert_eq!(det.current_count("posts"), 0);
1408 assert!(det.is_window_active());
1409 }
1410
1411 #[test]
1412 fn test_n1_detector_end_window_deactivates() {
1413 let det = N1QueryDetector::with_defaults();
1414 det.start_window();
1415 assert!(det.is_window_active());
1416 det.end_window();
1417 assert!(!det.is_window_active());
1418
1419 det.record_single_load("posts");
1421 assert_eq!(det.current_count("posts"), 0);
1422 }
1423
1424 #[test]
1425 fn test_n1_detector_custom_threshold() {
1426 let det = N1QueryDetector::new(N1DetectionConfig::new().with_threshold(100));
1427 det.start_window();
1428 for _ in 0..50 {
1429 det.record_single_load("posts");
1430 }
1431 let alerts = det.end_window();
1432 assert!(alerts.is_empty(), "below custom threshold should not alert");
1433
1434 det.start_window();
1435 for _ in 0..100 {
1436 det.record_single_load("posts");
1437 }
1438 let alerts = det.end_window();
1439 assert_eq!(alerts.len(), 1);
1440 assert_eq!(alerts[0].threshold, 100);
1441 assert_eq!(alerts[0].query_count, 100);
1442 }
1443
1444 #[test]
1445 fn test_n1_alert_suggested_batch_size() {
1446 let mk = |cnt: u64| N1Alert {
1447 relation: "x".into(),
1448 query_count: cnt,
1449 batch_count: 0,
1450 threshold: 5,
1451 };
1452 assert_eq!(mk(5).suggested_batch_size(), 50);
1453 assert_eq!(mk(50).suggested_batch_size(), 50);
1454 assert_eq!(mk(51).suggested_batch_size(), 100);
1455 assert_eq!(mk(100).suggested_batch_size(), 100);
1456 assert_eq!(mk(101).suggested_batch_size(), 500);
1457 assert_eq!(mk(500).suggested_batch_size(), 500);
1458 assert_eq!(mk(501).suggested_batch_size(), 1000);
1459 assert_eq!(mk(10000).suggested_batch_size(), 1000);
1460 }
1461
1462 #[test]
1463 fn test_n1_detector_real_n_plus_one_scenario() {
1464 let det = N1QueryDetector::with_defaults();
1466 det.start_window();
1467 let user_ids: Vec<i64> = (1..=20).collect();
1468 for _uid in &user_ids {
1469 det.record_single_load("posts");
1471 }
1472 let alerts = det.end_window();
1473 assert_eq!(alerts.len(), 1);
1474 assert_eq!(alerts[0].query_count, 20);
1475 assert!(alerts[0].no_batch_used());
1476
1477 det.start_window();
1479 det.record_batch_load("posts", 20); let alerts2 = det.end_window();
1481 assert!(alerts2.is_empty(), "batch loading should not trigger N+1");
1482 }
1483
1484 #[test]
1487 fn test_workflow_graph_and_batch_loader() {
1488 let mut graph = EntityGraph::new();
1490 graph.add_edge_with_graph("user", "posts", {
1491 let mut sub = EntityGraph::new();
1492 sub.add_edge("posts", "comments");
1493 sub
1494 });
1495 assert_eq!(graph.all_relations_recursive().len(), 2);
1496
1497 let user_loader = BatchLoader::new(
1499 50,
1500 Box::new(|ids: &[i64]| ids.iter().map(|id| (*id, format!("User#{}", id))).collect()),
1501 );
1502
1503 let user_ids: Vec<i64> = (1..=123).collect();
1505 let users = user_loader.load_many(&user_ids);
1506 assert_eq!(users.len(), 123);
1507 assert_eq!(user_loader.cache_size(), 123);
1508 }
1509
1510 #[test]
1511 fn test_n_plus_1_problem_solved() {
1512 let query_count = std::sync::Arc::new(std::sync::Mutex::new(0u32));
1517 let query_count_clone = query_count.clone();
1518
1519 let post_loader = BatchLoader::new(
1520 100,
1521 Box::new(move |user_ids: &[i64]| {
1522 *query_count_clone.lock().unwrap() += 1;
1523 user_ids
1525 .iter()
1526 .map(|uid| (*uid, format!("posts_for_user_{}", uid)))
1527 .collect()
1528 }),
1529 );
1530
1531 let user_ids: Vec<i64> = (1..=250).collect();
1533 let _posts = post_loader.load_many(&user_ids);
1534
1535 assert_eq!(*query_count.lock().unwrap(), 3);
1537 }
1538}