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(
242 &self,
243 adj: &mut std::collections::HashMap<String, Vec<String>>,
244 ) {
245 for edge in &self.edges {
246 adj.entry(edge.parent_field.clone())
247 .or_default()
248 .push(edge.relation.clone());
249 if let Some(sub) = &edge.sub_graph {
250 sub.collect_edges_recursive(adj);
251 }
252 }
253 }
254
255 pub fn detect_duplicate_edges(&self) -> Result<(), Vec<(String, String)>> {
261 let mut seen = std::collections::HashSet::new();
262 let mut duplicates = Vec::new();
263 for edge in &self.edges {
264 let key = (edge.parent_field.clone(), edge.relation.clone());
265 if !seen.insert(key.clone()) {
266 duplicates.push((edge.parent_field.clone(), edge.relation.clone()));
267 }
268 }
269 if duplicates.is_empty() {
270 Ok(())
271 } else {
272 Err(duplicates)
273 }
274 }
275
276 pub fn validate(&self) -> Result<(), String> {
280 if let Err(cycle) = self.detect_cycles() {
282 return Err(format!(
283 "EntityGraph 循环引用检测失败:{}",
284 cycle.join(" → ")
285 ));
286 }
287 if let Err(duplicates) = self.detect_duplicate_edges() {
289 let dup_str: Vec<String> = duplicates
290 .iter()
291 .map(|(p, r)| format!("({}->{})", p, r))
292 .collect();
293 return Err(format!(
294 "EntityGraph 重复边检测失败:{}",
295 dup_str.join(", ")
296 ));
297 }
298 Ok(())
299 }
300}
301
302fn dfs_cycle_detect(
307 node: &str,
308 adj: &std::collections::HashMap<String, Vec<String>>,
309 visited: &mut std::collections::HashSet<String>,
310 visiting: &mut std::collections::HashSet<String>,
311 path: &mut Vec<String>,
312) -> Result<(), Vec<String>> {
313 if visiting.contains(node) {
315 let cycle_start = path.iter().position(|n| n == node).unwrap_or(0);
316 let mut cycle = path[cycle_start..].to_vec();
317 cycle.push(node.to_string());
318 return Err(cycle);
319 }
320 if visited.contains(node) {
322 return Ok(());
323 }
324
325 visiting.insert(node.to_string());
327 path.push(node.to_string());
328
329 if let Some(neighbors) = adj.get(node) {
331 for neighbor in neighbors {
332 dfs_cycle_detect(neighbor, adj, visited, visiting, path)?;
333 }
334 }
335
336 visiting.remove(node);
338 visited.insert(node.to_string());
339 path.pop();
340 Ok(())
341}
342
343#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
349pub enum BatchStrategy {
350 #[default]
354 In,
355 Join,
359 Subquery,
363}
364
365impl BatchStrategy {
366 pub fn name(&self) -> &'static str {
368 match self {
369 BatchStrategy::In => "in",
370 BatchStrategy::Join => "join",
371 BatchStrategy::Subquery => "subquery",
372 }
373 }
374
375 pub fn render_in_clause(column: &str, placeholders: usize) -> String {
379 if placeholders == 0 {
380 return format!("{} IN ()", column);
381 }
382 let marks: Vec<&str> = vec!["?"; placeholders];
383 format!("{} IN ({})", column, marks.join(", "))
384 }
385}
386
387#[derive(Debug, Clone, Copy)]
395pub struct BatchSizeConfig {
396 pub size: usize,
398 pub strategy: BatchStrategy,
400}
401
402impl Default for BatchSizeConfig {
403 fn default() -> Self {
404 Self {
405 size: 100,
406 strategy: BatchStrategy::In,
407 }
408 }
409}
410
411impl BatchSizeConfig {
412 pub fn new(size: usize, strategy: BatchStrategy) -> Self {
414 Self { size, strategy }
415 }
416
417 pub fn with_size(size: usize) -> Self {
419 Self {
420 size,
421 strategy: BatchStrategy::In,
422 }
423 }
424
425 pub fn batch_count(&self, total: usize) -> usize {
440 if total == 0 {
441 0
442 } else {
443 total.div_ceil(self.size)
444 }
445 }
446
447 pub fn batch_range(&self, batch_index: usize, total: usize) -> std::ops::Range<usize> {
460 let start = batch_index * self.size;
461 let end = (start + self.size).min(total);
462 start..end
463 }
464}
465
466pub type BatchLoaderFn<K, V> = Box<dyn Fn(&[K]) -> HashMap<K, V> + Send + Sync>;
472
473pub struct BatchLoader<K, V>
497where
498 K: Hash + Eq + Clone + Send + Sync,
499 V: Clone + Send + Sync,
500{
501 batch_size: usize,
503 loader: BatchLoaderFn<K, V>,
505 cache: RwLock<HashMap<K, V>>,
507}
508
509impl<K, V> BatchLoader<K, V>
510where
511 K: Hash + Eq + Clone + Send + Sync,
512 V: Clone + Send + Sync,
513{
514 pub fn new(batch_size: usize, loader: BatchLoaderFn<K, V>) -> Self {
520 Self {
521 batch_size,
522 loader,
523 cache: RwLock::new(HashMap::new()),
524 }
525 }
526
527 pub fn load_many(&self, keys: &[K]) -> HashMap<K, V> {
533 let mut result: HashMap<K, V> = HashMap::new();
534
535 let mut to_load: Vec<K> = Vec::new();
537 if let Ok(cached) = self.cache.read() {
538 for k in keys {
539 if let Some(v) = cached.get(k) {
540 result.insert(k.clone(), v.clone());
541 } else {
542 to_load.push(k.clone());
543 }
544 }
545 } else {
546 to_load.extend(keys.iter().cloned());
547 }
548
549 if to_load.is_empty() {
550 return result;
551 }
552
553 let batch_size = self.batch_size.max(1);
555 let mut all_loaded: HashMap<K, V> = HashMap::new();
556 for chunk in to_load.chunks(batch_size) {
557 let loaded = (self.loader)(chunk);
558 all_loaded.extend(loaded);
559 }
560
561 if let Ok(mut cache) = self.cache.write() {
563 for (k, v) in &all_loaded {
564 cache.insert(k.clone(), v.clone());
565 }
566 }
567
568 result.extend(all_loaded);
570 result
571 }
572
573 pub fn load_one(&self, key: &K) -> Option<V> {
575 let result = self.load_many(std::slice::from_ref(key));
576 result.get(key).cloned()
577 }
578
579 pub fn clear_cache(&self) {
581 if let Ok(mut cache) = self.cache.write() {
582 cache.clear();
583 }
584 }
585
586 pub fn cache_size(&self) -> usize {
588 match self.cache.read() {
589 Ok(g) => g.len(),
590 Err(_) => 0,
591 }
592 }
593
594 pub fn batch_size(&self) -> usize {
596 self.batch_size
597 }
598}
599
600pub struct N1QueryDetector {
645 config: N1DetectionConfig,
647 counts: RwLock<HashMap<String, u64>>,
649 batch_counts: RwLock<HashMap<String, u64>>,
651 window_active: RwLock<bool>,
653 alerts: RwLock<Vec<N1Alert>>,
655}
656
657#[derive(Debug, Clone)]
659pub struct N1DetectionConfig {
660 pub threshold: u64,
662 pub enabled: bool,
664}
665
666impl Default for N1DetectionConfig {
667 fn default() -> Self {
668 Self {
669 threshold: 5,
670 enabled: true,
671 }
672 }
673}
674
675impl N1DetectionConfig {
676 pub fn new() -> Self {
678 Self::default()
679 }
680
681 pub fn with_threshold(mut self, threshold: u64) -> Self {
683 self.threshold = threshold.max(1);
684 self
685 }
686
687 pub fn with_enabled(mut self, enabled: bool) -> Self {
689 self.enabled = enabled;
690 self
691 }
692}
693
694#[derive(Debug, Clone, PartialEq, Eq)]
696pub struct N1Alert {
697 pub relation: String,
699 pub query_count: u64,
701 pub batch_count: u64,
703 pub threshold: u64,
705}
706
707impl N1Alert {
708 pub fn no_batch_used(&self) -> bool {
710 self.batch_count == 0
711 }
712
713 pub fn suggested_batch_size(&self) -> usize {
715 let n = self.query_count as usize;
716 if n <= 50 {
717 50
718 } else if n <= 100 {
719 100
720 } else if n <= 500 {
721 500
722 } else {
723 1000
724 }
725 }
726}
727
728impl N1QueryDetector {
729 pub fn new(config: N1DetectionConfig) -> Self {
731 Self {
732 config,
733 counts: RwLock::new(HashMap::new()),
734 batch_counts: RwLock::new(HashMap::new()),
735 window_active: RwLock::new(false),
736 alerts: RwLock::new(Vec::new()),
737 }
738 }
739
740 pub fn with_defaults() -> Self {
742 Self::new(N1DetectionConfig::default())
743 }
744
745 pub fn is_enabled(&self) -> bool {
747 self.config.enabled
748 }
749
750 pub fn threshold(&self) -> u64 {
752 self.config.threshold
753 }
754
755 pub fn start_window(&self) {
759 if !self.config.enabled {
760 return;
761 }
762 if let Ok(mut counts) = self.counts.write() {
763 *counts = HashMap::new();
764 }
765 if let Ok(mut batch_counts) = self.batch_counts.write() {
766 *batch_counts = HashMap::new();
767 }
768 if let Ok(mut alerts) = self.alerts.write() {
769 *alerts = Vec::new();
770 }
771 if let Ok(mut window_active) = self.window_active.write() {
772 *window_active = true;
773 }
774 }
775
776 pub fn end_window(&self) -> Vec<N1Alert> {
781 if !self.config.enabled {
782 return Vec::new();
783 }
784 if let Ok(mut window_active) = self.window_active.write() {
785 *window_active = false;
786 }
787
788 let new_alerts: Vec<N1Alert> =
790 match (self.counts.read(), self.batch_counts.read()) {
791 (Ok(counts), Ok(batch_counts)) => {
792 let mut alerts: Vec<N1Alert> = counts
793 .iter()
794 .filter_map(|(rel, &cnt)| {
795 if cnt >= self.config.threshold {
796 Some(N1Alert {
797 relation: rel.clone(),
798 query_count: cnt,
799 batch_count: *batch_counts.get(rel).unwrap_or(&0),
800 threshold: self.config.threshold,
801 })
802 } else {
803 None
804 }
805 })
806 .collect();
807 alerts.sort_by(|a, b| a.relation.cmp(&b.relation));
809 alerts
810 }
811 _ => Vec::new(),
812 };
813
814 if let Ok(mut alerts) = self.alerts.write() {
815 *alerts = new_alerts.clone();
816 }
817 new_alerts
818 }
819
820 pub fn record_single_load(&self, relation: &str) {
822 if !self.config.enabled {
823 return;
824 }
825 {
826 let active = self
827 .window_active
828 .read()
829 .map(|g| *g)
830 .unwrap_or(false);
831 if !active {
832 return;
833 }
834 }
835 if let Ok(mut counts) = self.counts.write() {
836 *counts.entry(relation.to_string()).or_insert(0) += 1;
837 }
838 }
839
840 pub fn record_batch_load(&self, relation: &str, _keys_count: usize) {
845 if !self.config.enabled {
846 return;
847 }
848 {
849 let active = self
850 .window_active
851 .read()
852 .map(|g| *g)
853 .unwrap_or(false);
854 if !active {
855 return;
856 }
857 }
858 if let Ok(mut batch_counts) = self.batch_counts.write() {
859 *batch_counts.entry(relation.to_string()).or_insert(0) += 1;
860 }
861 }
862
863 pub fn alerts(&self) -> Vec<N1Alert> {
865 self.alerts
866 .read()
867 .map(|g| g.clone())
868 .unwrap_or_default()
869 }
870
871 pub fn current_count(&self, relation: &str) -> u64 {
873 self.counts
874 .read()
875 .map(|g| g.get(relation).copied().unwrap_or(0))
876 .unwrap_or(0)
877 }
878
879 pub fn current_batch_count(&self, relation: &str) -> u64 {
881 self.batch_counts
882 .read()
883 .map(|g| g.get(relation).copied().unwrap_or(0))
884 .unwrap_or(0)
885 }
886
887 pub fn is_window_active(&self) -> bool {
889 self.window_active
890 .read()
891 .map(|g| *g)
892 .unwrap_or(false)
893 }
894
895 pub fn has_n_plus_one(&self) -> bool {
897 !self.alerts().is_empty()
898 }
899}
900
901impl Default for N1QueryDetector {
902 fn default() -> Self {
903 Self::with_defaults()
904 }
905}
906
907#[cfg(test)]
912mod tests {
913 use super::*;
914
915 #[test]
918 fn test_new_graph_is_empty() {
919 let g = EntityGraph::new();
920 assert!(g.is_empty());
921 assert_eq!(g.edge_count(), 0);
922 }
923
924 #[test]
925 fn test_add_edge() {
926 let mut g = EntityGraph::new();
927 g.add_edge("user", "posts");
928 assert_eq!(g.edge_count(), 1);
929 assert!(!g.is_empty());
930 }
931
932 #[test]
933 fn test_add_multiple_edges() {
934 let mut g = EntityGraph::new();
935 g.add_edge("user", "posts")
936 .add_edge("user", "profile")
937 .add_edge("user", "comments");
938 assert_eq!(g.edge_count(), 3);
939 }
940
941 #[test]
942 fn test_add_edge_with_sub_graph() {
943 let mut sub = EntityGraph::new();
944 sub.add_edge("comments", "author");
945
946 let mut g = EntityGraph::new();
947 g.add_edge_with_graph("user", "posts", sub);
948
949 assert_eq!(g.edge_count(), 1);
950 assert!(g.edges()[0].sub_graph.is_some());
951 assert_eq!(g.edges()[0].sub_graph.as_ref().unwrap().edge_count(), 1);
952 }
953
954 #[test]
955 fn test_relations_of() {
956 let mut g = EntityGraph::new();
957 g.add_edge("user", "posts")
958 .add_edge("user", "profile")
959 .add_edge("post", "comments");
960
961 let user_relations = g.relations_of("user");
962 assert_eq!(user_relations.len(), 2);
963 assert_eq!(user_relations[0].relation, "posts");
964 assert_eq!(user_relations[1].relation, "profile");
965
966 let post_relations = g.relations_of("post");
967 assert_eq!(post_relations.len(), 1);
968
969 let none = g.relations_of("nonexistent");
970 assert!(none.is_empty());
971 }
972
973 #[test]
974 fn test_all_relations() {
975 let mut g = EntityGraph::new();
976 g.add_edge("user", "posts")
977 .add_edge("user", "profile")
978 .add_edge("post", "comments");
979
980 let rels = g.all_relations();
981 assert_eq!(rels, vec!["comments", "posts", "profile"]);
982 }
983
984 #[test]
985 fn test_all_parent_fields() {
986 let mut g = EntityGraph::new();
987 g.add_edge("user", "posts")
988 .add_edge("user", "profile")
989 .add_edge("post", "comments");
990
991 let fields = g.all_parent_fields();
992 assert_eq!(fields, vec!["post", "user"]);
993 }
994
995 #[test]
996 fn test_all_relations_recursive() {
997 let mut sub = EntityGraph::new();
998 sub.add_edge("comments", "author")
999 .add_edge("comments", "likes");
1000
1001 let mut g = EntityGraph::new();
1002 g.add_edge("user", "posts")
1003 .add_edge_with_graph("user", "comments", sub);
1004
1005 let all = g.all_relations_recursive();
1006 assert!(all.contains(&"posts".to_string()));
1007 assert!(all.contains(&"comments".to_string()));
1008 assert!(all.contains(&"author".to_string()));
1009 assert!(all.contains(&"likes".to_string()));
1010 assert_eq!(all.len(), 4);
1011 }
1012
1013 #[test]
1014 fn test_default_graph_is_empty() {
1015 let g = EntityGraph::default();
1016 assert!(g.is_empty());
1017 }
1018
1019 #[test]
1022 fn test_strategy_name() {
1023 assert_eq!(BatchStrategy::In.name(), "in");
1024 assert_eq!(BatchStrategy::Join.name(), "join");
1025 assert_eq!(BatchStrategy::Subquery.name(), "subquery");
1026 }
1027
1028 #[test]
1029 fn test_strategy_default_is_in() {
1030 assert_eq!(BatchStrategy::default(), BatchStrategy::In);
1031 }
1032
1033 #[test]
1034 fn test_render_in_clause_empty() {
1035 let sql = BatchStrategy::render_in_clause("id", 0);
1036 assert_eq!(sql, "id IN ()");
1037 }
1038
1039 #[test]
1040 fn test_render_in_clause_single() {
1041 let sql = BatchStrategy::render_in_clause("id", 1);
1042 assert_eq!(sql, "id IN (?)");
1043 }
1044
1045 #[test]
1046 fn test_render_in_clause_multiple() {
1047 let sql = BatchStrategy::render_in_clause("user_id", 3);
1048 assert_eq!(sql, "user_id IN (?, ?, ?)");
1049 }
1050
1051 #[test]
1054 fn test_default_config() {
1055 let config = BatchSizeConfig::default();
1056 assert_eq!(config.size, 100);
1057 assert_eq!(config.strategy, BatchStrategy::In);
1058 }
1059
1060 #[test]
1061 fn test_with_size() {
1062 let config = BatchSizeConfig::with_size(50);
1063 assert_eq!(config.size, 50);
1064 assert_eq!(config.strategy, BatchStrategy::In);
1065 }
1066
1067 #[test]
1068 fn test_new_with_strategy() {
1069 let config = BatchSizeConfig::new(200, BatchStrategy::Join);
1070 assert_eq!(config.size, 200);
1071 assert_eq!(config.strategy, BatchStrategy::Join);
1072 }
1073
1074 #[test]
1075 fn test_batch_count_zero() {
1076 let config = BatchSizeConfig::with_size(100);
1077 assert_eq!(config.batch_count(0), 0);
1078 }
1079
1080 #[test]
1081 fn test_batch_count_exact_multiple() {
1082 let config = BatchSizeConfig::with_size(100);
1083 assert_eq!(config.batch_count(100), 1);
1084 assert_eq!(config.batch_count(200), 2);
1085 assert_eq!(config.batch_count(500), 5);
1086 }
1087
1088 #[test]
1089 fn test_batch_count_with_remainder() {
1090 let config = BatchSizeConfig::with_size(100);
1091 assert_eq!(config.batch_count(1), 1);
1092 assert_eq!(config.batch_count(99), 1);
1093 assert_eq!(config.batch_count(101), 2);
1094 assert_eq!(config.batch_count(150), 2);
1095 assert_eq!(config.batch_count(201), 3);
1096 }
1097
1098 #[test]
1099 fn test_batch_range() {
1100 let config = BatchSizeConfig::with_size(100);
1101
1102 assert_eq!(config.batch_range(0, 250), 0..100);
1103 assert_eq!(config.batch_range(1, 250), 100..200);
1104 assert_eq!(config.batch_range(2, 250), 200..250);
1105 }
1106
1107 #[test]
1108 fn test_batch_range_exact() {
1109 let config = BatchSizeConfig::with_size(100);
1110
1111 assert_eq!(config.batch_range(0, 100), 0..100);
1112 assert_eq!(config.batch_range(1, 100), 100..100); }
1114
1115 #[test]
1116 fn test_batch_range_small_batch() {
1117 let config = BatchSizeConfig::with_size(10);
1118
1119 assert_eq!(config.batch_range(0, 25), 0..10);
1120 assert_eq!(config.batch_range(1, 25), 10..20);
1121 assert_eq!(config.batch_range(2, 25), 20..25);
1122 }
1123
1124 fn make_loader() -> BatchLoader<i64, String> {
1127 let loader = Box::new(|ids: &[i64]| -> HashMap<i64, String> {
1128 ids.iter().map(|id| (*id, format!("user_{}", id))).collect()
1129 });
1130 BatchLoader::new(2, loader)
1131 }
1132
1133 #[test]
1134 fn test_batch_loader_load_many_single_batch() {
1135 let loader = make_loader();
1136 let result = loader.load_many(&[1, 2]);
1137 assert_eq!(result.len(), 2);
1138 assert_eq!(result.get(&1), Some(&"user_1".to_string()));
1139 assert_eq!(result.get(&2), Some(&"user_2".to_string()));
1140 }
1141
1142 #[test]
1143 fn test_batch_loader_load_many_multiple_batches() {
1144 let loader = make_loader();
1145 let result = loader.load_many(&[1, 2, 3, 4, 5]);
1147 assert_eq!(result.len(), 5);
1148 for id in 1..=5 {
1149 assert_eq!(
1150 result.get(&id),
1151 Some(&format!("user_{}", id)),
1152 "missing user {}",
1153 id
1154 );
1155 }
1156 }
1157
1158 #[test]
1159 fn test_batch_loader_load_one() {
1160 let loader = make_loader();
1161 let result = loader.load_one(&42);
1162 assert_eq!(result, Some("user_42".to_string()));
1163 }
1164
1165 #[test]
1166 fn test_batch_loader_load_one_missing() {
1167 let loader: BatchLoader<i64, String> =
1169 BatchLoader::new(10, Box::new(|_ids: &[i64]| HashMap::new()));
1170 let result = loader.load_one(&100);
1171 assert_eq!(result, None);
1172 }
1173
1174 #[test]
1175 fn test_batch_loader_caches_results() {
1176 let call_count = std::sync::Arc::new(std::sync::Mutex::new(0u32));
1177 let call_count_clone = call_count.clone();
1178
1179 let loader = Box::new(move |ids: &[i64]| -> HashMap<i64, String> {
1180 *call_count_clone.lock().unwrap() += 1;
1181 ids.iter().map(|id| (*id, format!("user_{}", id))).collect()
1182 });
1183
1184 let batch_loader = BatchLoader::new(100, loader);
1185
1186 batch_loader.load_many(&[1, 2, 3]);
1188 assert_eq!(*call_count.lock().unwrap(), 1);
1189
1190 batch_loader.load_many(&[1, 2, 3]);
1192 assert_eq!(*call_count.lock().unwrap(), 1); batch_loader.load_many(&[4, 5]);
1196 assert_eq!(*call_count.lock().unwrap(), 2);
1197 }
1198
1199 #[test]
1200 fn test_batch_loader_partial_cache_hit() {
1201 let call_count = std::sync::Arc::new(std::sync::Mutex::new(0u32));
1202 let call_count_clone = call_count.clone();
1203
1204 let loader = Box::new(move |ids: &[i64]| -> HashMap<i64, String> {
1205 *call_count_clone.lock().unwrap() += 1;
1206 ids.iter().map(|id| (*id, format!("user_{}", id))).collect()
1207 });
1208
1209 let batch_loader = BatchLoader::new(100, loader);
1210
1211 batch_loader.load_many(&[1, 2, 3]);
1213 assert_eq!(*call_count.lock().unwrap(), 1);
1214
1215 let result = batch_loader.load_many(&[1, 2, 3, 4, 5]);
1217 assert_eq!(result.len(), 5);
1218 assert_eq!(*call_count.lock().unwrap(), 2); assert_eq!(batch_loader.cache_size(), 5);
1222 }
1223
1224 #[test]
1225 fn test_batch_loader_clear_cache() {
1226 let loader = make_loader();
1227 loader.load_many(&[1, 2]);
1228 assert_eq!(loader.cache_size(), 2);
1229
1230 loader.clear_cache();
1231 assert_eq!(loader.cache_size(), 0);
1232 }
1233
1234 #[test]
1235 fn test_batch_loader_empty_input() {
1236 let loader = make_loader();
1237 let result = loader.load_many(&[]);
1238 assert!(result.is_empty());
1239 }
1240
1241 #[test]
1242 fn test_batch_loader_batch_size_attribute() {
1243 let loader = make_loader();
1244 assert_eq!(loader.batch_size(), 2);
1245 }
1246
1247 #[test]
1248 fn test_batch_loader_with_size_1() {
1249 let loader = BatchLoader::new(
1250 1,
1251 Box::new(|ids: &[i64]| ids.iter().map(|id| (*id, *id * 10)).collect()),
1252 );
1253 let result = loader.load_many(&[1, 2, 3]);
1254 assert_eq!(result.len(), 3);
1255 assert_eq!(result.get(&1), Some(&10));
1256 assert_eq!(result.get(&2), Some(&20));
1257 assert_eq!(result.get(&3), Some(&30));
1258 }
1259
1260 #[test]
1263 fn test_n1_config_default() {
1264 let cfg = N1DetectionConfig::default();
1265 assert_eq!(cfg.threshold, 5);
1266 assert!(cfg.enabled);
1267 }
1268
1269 #[test]
1270 fn test_n1_config_builder() {
1271 let cfg = N1DetectionConfig::new()
1272 .with_threshold(10)
1273 .with_enabled(false);
1274 assert_eq!(cfg.threshold, 10);
1275 assert!(!cfg.enabled);
1276
1277 let cfg2 = N1DetectionConfig::new().with_threshold(0);
1279 assert_eq!(cfg2.threshold, 1);
1280 }
1281
1282 #[test]
1283 fn test_n1_detector_default() {
1284 let det = N1QueryDetector::default();
1285 assert!(det.is_enabled());
1286 assert_eq!(det.threshold(), 5);
1287 assert!(!det.is_window_active());
1288 assert!(!det.has_n_plus_one());
1289 }
1290
1291 #[test]
1292 fn test_n1_detector_disabled_is_noop() {
1293 let det = N1QueryDetector::new(N1DetectionConfig::new().with_enabled(false));
1294 det.start_window();
1295 for _ in 0..100 {
1296 det.record_single_load("posts");
1297 }
1298 assert_eq!(det.current_count("posts"), 0);
1300 let alerts = det.end_window();
1301 assert!(alerts.is_empty());
1302 }
1303
1304 #[test]
1305 fn test_n1_detector_records_outside_window_ignored() {
1306 let det = N1QueryDetector::with_defaults();
1307 det.record_single_load("posts");
1309 assert_eq!(det.current_count("posts"), 0);
1310 }
1311
1312 #[test]
1313 fn test_n1_detector_below_threshold_no_alert() {
1314 let det = N1QueryDetector::with_defaults(); det.start_window();
1316 for _ in 0..4 {
1317 det.record_single_load("posts");
1318 }
1319 assert_eq!(det.current_count("posts"), 4);
1320 let alerts = det.end_window();
1321 assert!(alerts.is_empty(), "below threshold should not alert");
1322 assert!(!det.has_n_plus_one());
1323 }
1324
1325 #[test]
1326 fn test_n1_detector_at_threshold_triggers_alert() {
1327 let det = N1QueryDetector::with_defaults(); det.start_window();
1329 for _ in 0..5 {
1330 det.record_single_load("posts");
1331 }
1332 let alerts = det.end_window();
1333 assert_eq!(alerts.len(), 1);
1334 assert_eq!(alerts[0].relation, "posts");
1335 assert_eq!(alerts[0].query_count, 5);
1336 assert_eq!(alerts[0].threshold, 5);
1337 assert_eq!(alerts[0].batch_count, 0);
1338 assert!(alerts[0].no_batch_used());
1339 assert!(det.has_n_plus_one());
1340 }
1341
1342 #[test]
1343 fn test_n1_detector_above_threshold_triggers_alert() {
1344 let det = N1QueryDetector::with_defaults();
1345 det.start_window();
1346 for _ in 0..10 {
1347 det.record_single_load("posts");
1348 }
1349 let alerts = det.end_window();
1350 assert_eq!(alerts.len(), 1);
1351 assert_eq!(alerts[0].query_count, 10);
1352 assert!(alerts[0].suggested_batch_size() >= 50);
1354 }
1355
1356 #[test]
1357 fn test_n1_detector_multiple_relations() {
1358 let det = N1QueryDetector::with_defaults();
1359 det.start_window();
1360 for _ in 0..6 {
1361 det.record_single_load("posts");
1362 }
1363 for _ in 0..3 {
1364 det.record_single_load("comments"); }
1366 for _ in 0..8 {
1367 det.record_single_load("tags");
1368 }
1369 let alerts = det.end_window();
1370 assert_eq!(alerts.len(), 2);
1372 assert_eq!(alerts[0].relation, "posts");
1374 assert_eq!(alerts[0].query_count, 6);
1375 assert_eq!(alerts[1].relation, "tags");
1376 assert_eq!(alerts[1].query_count, 8);
1377 }
1378
1379 #[test]
1380 fn test_n1_detector_batch_load_recorded_separately() {
1381 let det = N1QueryDetector::with_defaults();
1382 det.start_window();
1383 for _ in 0..6 {
1385 det.record_single_load("posts");
1386 }
1387 det.record_batch_load("posts", 100);
1389 det.record_batch_load("posts", 50);
1390 let alerts = det.end_window();
1391 assert_eq!(alerts.len(), 1);
1392 assert_eq!(alerts[0].query_count, 6);
1393 assert_eq!(alerts[0].batch_count, 2);
1394 assert!(!alerts[0].no_batch_used());
1396 }
1397
1398 #[test]
1399 fn test_n1_detector_batch_only_does_not_trigger() {
1400 let det = N1QueryDetector::with_defaults();
1402 det.start_window();
1403 for _ in 0..100 {
1404 det.record_batch_load("posts", 50);
1405 }
1406 assert_eq!(det.current_batch_count("posts"), 100);
1407 assert_eq!(det.current_count("posts"), 0);
1408 let alerts = det.end_window();
1409 assert!(alerts.is_empty());
1410 }
1411
1412 #[test]
1413 fn test_n1_detector_start_window_resets() {
1414 let det = N1QueryDetector::with_defaults();
1415 det.start_window();
1416 for _ in 0..10 {
1417 det.record_single_load("posts");
1418 }
1419 let _ = det.end_window();
1420 assert_eq!(det.alerts().len(), 1);
1421
1422 det.start_window();
1424 assert_eq!(det.alerts().len(), 0);
1425 assert_eq!(det.current_count("posts"), 0);
1426 assert!(det.is_window_active());
1427 }
1428
1429 #[test]
1430 fn test_n1_detector_end_window_deactivates() {
1431 let det = N1QueryDetector::with_defaults();
1432 det.start_window();
1433 assert!(det.is_window_active());
1434 det.end_window();
1435 assert!(!det.is_window_active());
1436
1437 det.record_single_load("posts");
1439 assert_eq!(det.current_count("posts"), 0);
1440 }
1441
1442 #[test]
1443 fn test_n1_detector_custom_threshold() {
1444 let det = N1QueryDetector::new(N1DetectionConfig::new().with_threshold(100));
1445 det.start_window();
1446 for _ in 0..50 {
1447 det.record_single_load("posts");
1448 }
1449 let alerts = det.end_window();
1450 assert!(alerts.is_empty(), "below custom threshold should not alert");
1451
1452 det.start_window();
1453 for _ in 0..100 {
1454 det.record_single_load("posts");
1455 }
1456 let alerts = det.end_window();
1457 assert_eq!(alerts.len(), 1);
1458 assert_eq!(alerts[0].threshold, 100);
1459 assert_eq!(alerts[0].query_count, 100);
1460 }
1461
1462 #[test]
1463 fn test_n1_alert_suggested_batch_size() {
1464 let mk = |cnt: u64| N1Alert {
1465 relation: "x".into(),
1466 query_count: cnt,
1467 batch_count: 0,
1468 threshold: 5,
1469 };
1470 assert_eq!(mk(5).suggested_batch_size(), 50);
1471 assert_eq!(mk(50).suggested_batch_size(), 50);
1472 assert_eq!(mk(51).suggested_batch_size(), 100);
1473 assert_eq!(mk(100).suggested_batch_size(), 100);
1474 assert_eq!(mk(101).suggested_batch_size(), 500);
1475 assert_eq!(mk(500).suggested_batch_size(), 500);
1476 assert_eq!(mk(501).suggested_batch_size(), 1000);
1477 assert_eq!(mk(10000).suggested_batch_size(), 1000);
1478 }
1479
1480 #[test]
1481 fn test_n1_detector_real_n_plus_one_scenario() {
1482 let det = N1QueryDetector::with_defaults();
1484 det.start_window();
1485 let user_ids: Vec<i64> = (1..=20).collect();
1486 for _uid in &user_ids {
1487 det.record_single_load("posts");
1489 }
1490 let alerts = det.end_window();
1491 assert_eq!(alerts.len(), 1);
1492 assert_eq!(alerts[0].query_count, 20);
1493 assert!(alerts[0].no_batch_used());
1494
1495 det.start_window();
1497 det.record_batch_load("posts", 20); let alerts2 = det.end_window();
1499 assert!(alerts2.is_empty(), "batch loading should not trigger N+1");
1500 }
1501
1502 #[test]
1505 fn test_workflow_graph_and_batch_loader() {
1506 let mut graph = EntityGraph::new();
1508 graph.add_edge_with_graph("user", "posts", {
1509 let mut sub = EntityGraph::new();
1510 sub.add_edge("posts", "comments");
1511 sub
1512 });
1513 assert_eq!(graph.all_relations_recursive().len(), 2);
1514
1515 let user_loader = BatchLoader::new(
1517 50,
1518 Box::new(|ids: &[i64]| ids.iter().map(|id| (*id, format!("User#{}", id))).collect()),
1519 );
1520
1521 let user_ids: Vec<i64> = (1..=123).collect();
1523 let users = user_loader.load_many(&user_ids);
1524 assert_eq!(users.len(), 123);
1525 assert_eq!(user_loader.cache_size(), 123);
1526 }
1527
1528 #[test]
1529 fn test_n_plus_1_problem_solved() {
1530 let query_count = std::sync::Arc::new(std::sync::Mutex::new(0u32));
1535 let query_count_clone = query_count.clone();
1536
1537 let post_loader = BatchLoader::new(
1538 100,
1539 Box::new(move |user_ids: &[i64]| {
1540 *query_count_clone.lock().unwrap() += 1;
1541 user_ids
1543 .iter()
1544 .map(|uid| (*uid, format!("posts_for_user_{}", uid)))
1545 .collect()
1546 }),
1547 );
1548
1549 let user_ids: Vec<i64> = (1..=250).collect();
1551 let _posts = post_loader.load_many(&user_ids);
1552
1553 assert_eq!(*query_count.lock().unwrap(), 3);
1555 }
1556}