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
172#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
178pub enum BatchStrategy {
179 #[default]
183 In,
184 Join,
188 Subquery,
192}
193
194impl BatchStrategy {
195 pub fn name(&self) -> &'static str {
197 match self {
198 BatchStrategy::In => "in",
199 BatchStrategy::Join => "join",
200 BatchStrategy::Subquery => "subquery",
201 }
202 }
203
204 pub fn render_in_clause(column: &str, placeholders: usize) -> String {
208 if placeholders == 0 {
209 return format!("{} IN ()", column);
210 }
211 let marks: Vec<&str> = vec!["?"; placeholders];
212 format!("{} IN ({})", column, marks.join(", "))
213 }
214}
215
216#[derive(Debug, Clone, Copy)]
224pub struct BatchSizeConfig {
225 pub size: usize,
227 pub strategy: BatchStrategy,
229}
230
231impl Default for BatchSizeConfig {
232 fn default() -> Self {
233 Self {
234 size: 100,
235 strategy: BatchStrategy::In,
236 }
237 }
238}
239
240impl BatchSizeConfig {
241 pub fn new(size: usize, strategy: BatchStrategy) -> Self {
243 Self { size, strategy }
244 }
245
246 pub fn with_size(size: usize) -> Self {
248 Self {
249 size,
250 strategy: BatchStrategy::In,
251 }
252 }
253
254 pub fn batch_count(&self, total: usize) -> usize {
269 if total == 0 {
270 0
271 } else {
272 total.div_ceil(self.size)
273 }
274 }
275
276 pub fn batch_range(&self, batch_index: usize, total: usize) -> std::ops::Range<usize> {
289 let start = batch_index * self.size;
290 let end = (start + self.size).min(total);
291 start..end
292 }
293}
294
295pub type BatchLoaderFn<K, V> = Box<dyn Fn(&[K]) -> HashMap<K, V> + Send + Sync>;
301
302pub struct BatchLoader<K, V>
326where
327 K: Hash + Eq + Clone + Send + Sync,
328 V: Clone + Send + Sync,
329{
330 batch_size: usize,
332 loader: BatchLoaderFn<K, V>,
334 cache: RwLock<HashMap<K, V>>,
336}
337
338impl<K, V> BatchLoader<K, V>
339where
340 K: Hash + Eq + Clone + Send + Sync,
341 V: Clone + Send + Sync,
342{
343 pub fn new(batch_size: usize, loader: BatchLoaderFn<K, V>) -> Self {
349 Self {
350 batch_size,
351 loader,
352 cache: RwLock::new(HashMap::new()),
353 }
354 }
355
356 pub fn load_many(&self, keys: &[K]) -> HashMap<K, V> {
362 let mut result: HashMap<K, V> = HashMap::new();
363
364 let cached = self
366 .cache
367 .read()
368 .expect("BatchLoader cache lock poisoned (read)");
369 let mut to_load: Vec<K> = Vec::new();
370 for k in keys {
371 if let Some(v) = cached.get(k) {
372 result.insert(k.clone(), v.clone());
373 } else {
374 to_load.push(k.clone());
375 }
376 }
377 drop(cached);
378
379 if to_load.is_empty() {
380 return result;
381 }
382
383 let batch_size = self.batch_size.max(1);
385 let mut all_loaded: HashMap<K, V> = HashMap::new();
386 for chunk in to_load.chunks(batch_size) {
387 let loaded = (self.loader)(chunk);
388 all_loaded.extend(loaded);
389 }
390
391 let mut cache = self
393 .cache
394 .write()
395 .expect("BatchLoader cache lock poisoned (write)");
396 for (k, v) in &all_loaded {
397 cache.insert(k.clone(), v.clone());
398 }
399 drop(cache);
400
401 result.extend(all_loaded);
403 result
404 }
405
406 pub fn load_one(&self, key: &K) -> Option<V> {
408 let result = self.load_many(std::slice::from_ref(key));
409 result.get(key).cloned()
410 }
411
412 pub fn clear_cache(&self) {
414 self.cache
415 .write()
416 .expect("BatchLoader cache lock poisoned (clear_cache)")
417 .clear();
418 }
419
420 pub fn cache_size(&self) -> usize {
422 self.cache
423 .read()
424 .expect("BatchLoader cache lock poisoned (cache_size)")
425 .len()
426 }
427
428 pub fn batch_size(&self) -> usize {
430 self.batch_size
431 }
432}
433
434#[cfg(test)]
439mod tests {
440 use super::*;
441
442 #[test]
445 fn test_new_graph_is_empty() {
446 let g = EntityGraph::new();
447 assert!(g.is_empty());
448 assert_eq!(g.edge_count(), 0);
449 }
450
451 #[test]
452 fn test_add_edge() {
453 let mut g = EntityGraph::new();
454 g.add_edge("user", "posts");
455 assert_eq!(g.edge_count(), 1);
456 assert!(!g.is_empty());
457 }
458
459 #[test]
460 fn test_add_multiple_edges() {
461 let mut g = EntityGraph::new();
462 g.add_edge("user", "posts")
463 .add_edge("user", "profile")
464 .add_edge("user", "comments");
465 assert_eq!(g.edge_count(), 3);
466 }
467
468 #[test]
469 fn test_add_edge_with_sub_graph() {
470 let mut sub = EntityGraph::new();
471 sub.add_edge("comments", "author");
472
473 let mut g = EntityGraph::new();
474 g.add_edge_with_graph("user", "posts", sub);
475
476 assert_eq!(g.edge_count(), 1);
477 assert!(g.edges()[0].sub_graph.is_some());
478 assert_eq!(g.edges()[0].sub_graph.as_ref().unwrap().edge_count(), 1);
479 }
480
481 #[test]
482 fn test_relations_of() {
483 let mut g = EntityGraph::new();
484 g.add_edge("user", "posts")
485 .add_edge("user", "profile")
486 .add_edge("post", "comments");
487
488 let user_relations = g.relations_of("user");
489 assert_eq!(user_relations.len(), 2);
490 assert_eq!(user_relations[0].relation, "posts");
491 assert_eq!(user_relations[1].relation, "profile");
492
493 let post_relations = g.relations_of("post");
494 assert_eq!(post_relations.len(), 1);
495
496 let none = g.relations_of("nonexistent");
497 assert!(none.is_empty());
498 }
499
500 #[test]
501 fn test_all_relations() {
502 let mut g = EntityGraph::new();
503 g.add_edge("user", "posts")
504 .add_edge("user", "profile")
505 .add_edge("post", "comments");
506
507 let rels = g.all_relations();
508 assert_eq!(rels, vec!["comments", "posts", "profile"]);
509 }
510
511 #[test]
512 fn test_all_parent_fields() {
513 let mut g = EntityGraph::new();
514 g.add_edge("user", "posts")
515 .add_edge("user", "profile")
516 .add_edge("post", "comments");
517
518 let fields = g.all_parent_fields();
519 assert_eq!(fields, vec!["post", "user"]);
520 }
521
522 #[test]
523 fn test_all_relations_recursive() {
524 let mut sub = EntityGraph::new();
525 sub.add_edge("comments", "author")
526 .add_edge("comments", "likes");
527
528 let mut g = EntityGraph::new();
529 g.add_edge("user", "posts")
530 .add_edge_with_graph("user", "comments", sub);
531
532 let all = g.all_relations_recursive();
533 assert!(all.contains(&"posts".to_string()));
534 assert!(all.contains(&"comments".to_string()));
535 assert!(all.contains(&"author".to_string()));
536 assert!(all.contains(&"likes".to_string()));
537 assert_eq!(all.len(), 4);
538 }
539
540 #[test]
541 fn test_default_graph_is_empty() {
542 let g = EntityGraph::default();
543 assert!(g.is_empty());
544 }
545
546 #[test]
549 fn test_strategy_name() {
550 assert_eq!(BatchStrategy::In.name(), "in");
551 assert_eq!(BatchStrategy::Join.name(), "join");
552 assert_eq!(BatchStrategy::Subquery.name(), "subquery");
553 }
554
555 #[test]
556 fn test_strategy_default_is_in() {
557 assert_eq!(BatchStrategy::default(), BatchStrategy::In);
558 }
559
560 #[test]
561 fn test_render_in_clause_empty() {
562 let sql = BatchStrategy::render_in_clause("id", 0);
563 assert_eq!(sql, "id IN ()");
564 }
565
566 #[test]
567 fn test_render_in_clause_single() {
568 let sql = BatchStrategy::render_in_clause("id", 1);
569 assert_eq!(sql, "id IN (?)");
570 }
571
572 #[test]
573 fn test_render_in_clause_multiple() {
574 let sql = BatchStrategy::render_in_clause("user_id", 3);
575 assert_eq!(sql, "user_id IN (?, ?, ?)");
576 }
577
578 #[test]
581 fn test_default_config() {
582 let config = BatchSizeConfig::default();
583 assert_eq!(config.size, 100);
584 assert_eq!(config.strategy, BatchStrategy::In);
585 }
586
587 #[test]
588 fn test_with_size() {
589 let config = BatchSizeConfig::with_size(50);
590 assert_eq!(config.size, 50);
591 assert_eq!(config.strategy, BatchStrategy::In);
592 }
593
594 #[test]
595 fn test_new_with_strategy() {
596 let config = BatchSizeConfig::new(200, BatchStrategy::Join);
597 assert_eq!(config.size, 200);
598 assert_eq!(config.strategy, BatchStrategy::Join);
599 }
600
601 #[test]
602 fn test_batch_count_zero() {
603 let config = BatchSizeConfig::with_size(100);
604 assert_eq!(config.batch_count(0), 0);
605 }
606
607 #[test]
608 fn test_batch_count_exact_multiple() {
609 let config = BatchSizeConfig::with_size(100);
610 assert_eq!(config.batch_count(100), 1);
611 assert_eq!(config.batch_count(200), 2);
612 assert_eq!(config.batch_count(500), 5);
613 }
614
615 #[test]
616 fn test_batch_count_with_remainder() {
617 let config = BatchSizeConfig::with_size(100);
618 assert_eq!(config.batch_count(1), 1);
619 assert_eq!(config.batch_count(99), 1);
620 assert_eq!(config.batch_count(101), 2);
621 assert_eq!(config.batch_count(150), 2);
622 assert_eq!(config.batch_count(201), 3);
623 }
624
625 #[test]
626 fn test_batch_range() {
627 let config = BatchSizeConfig::with_size(100);
628
629 assert_eq!(config.batch_range(0, 250), 0..100);
630 assert_eq!(config.batch_range(1, 250), 100..200);
631 assert_eq!(config.batch_range(2, 250), 200..250);
632 }
633
634 #[test]
635 fn test_batch_range_exact() {
636 let config = BatchSizeConfig::with_size(100);
637
638 assert_eq!(config.batch_range(0, 100), 0..100);
639 assert_eq!(config.batch_range(1, 100), 100..100); }
641
642 #[test]
643 fn test_batch_range_small_batch() {
644 let config = BatchSizeConfig::with_size(10);
645
646 assert_eq!(config.batch_range(0, 25), 0..10);
647 assert_eq!(config.batch_range(1, 25), 10..20);
648 assert_eq!(config.batch_range(2, 25), 20..25);
649 }
650
651 fn make_loader() -> BatchLoader<i64, String> {
654 let loader = Box::new(|ids: &[i64]| -> HashMap<i64, String> {
655 ids.iter().map(|id| (*id, format!("user_{}", id))).collect()
656 });
657 BatchLoader::new(2, loader)
658 }
659
660 #[test]
661 fn test_batch_loader_load_many_single_batch() {
662 let loader = make_loader();
663 let result = loader.load_many(&[1, 2]);
664 assert_eq!(result.len(), 2);
665 assert_eq!(result.get(&1), Some(&"user_1".to_string()));
666 assert_eq!(result.get(&2), Some(&"user_2".to_string()));
667 }
668
669 #[test]
670 fn test_batch_loader_load_many_multiple_batches() {
671 let loader = make_loader();
672 let result = loader.load_many(&[1, 2, 3, 4, 5]);
674 assert_eq!(result.len(), 5);
675 for id in 1..=5 {
676 assert_eq!(
677 result.get(&id),
678 Some(&format!("user_{}", id)),
679 "missing user {}",
680 id
681 );
682 }
683 }
684
685 #[test]
686 fn test_batch_loader_load_one() {
687 let loader = make_loader();
688 let result = loader.load_one(&42);
689 assert_eq!(result, Some("user_42".to_string()));
690 }
691
692 #[test]
693 fn test_batch_loader_load_one_missing() {
694 let loader: BatchLoader<i64, String> =
696 BatchLoader::new(10, Box::new(|_ids: &[i64]| HashMap::new()));
697 let result = loader.load_one(&100);
698 assert_eq!(result, None);
699 }
700
701 #[test]
702 fn test_batch_loader_caches_results() {
703 let call_count = std::sync::Arc::new(std::sync::Mutex::new(0u32));
704 let call_count_clone = call_count.clone();
705
706 let loader = Box::new(move |ids: &[i64]| -> HashMap<i64, String> {
707 *call_count_clone.lock().unwrap() += 1;
708 ids.iter().map(|id| (*id, format!("user_{}", id))).collect()
709 });
710
711 let batch_loader = BatchLoader::new(100, loader);
712
713 batch_loader.load_many(&[1, 2, 3]);
715 assert_eq!(*call_count.lock().unwrap(), 1);
716
717 batch_loader.load_many(&[1, 2, 3]);
719 assert_eq!(*call_count.lock().unwrap(), 1); batch_loader.load_many(&[4, 5]);
723 assert_eq!(*call_count.lock().unwrap(), 2);
724 }
725
726 #[test]
727 fn test_batch_loader_partial_cache_hit() {
728 let call_count = std::sync::Arc::new(std::sync::Mutex::new(0u32));
729 let call_count_clone = call_count.clone();
730
731 let loader = Box::new(move |ids: &[i64]| -> HashMap<i64, String> {
732 *call_count_clone.lock().unwrap() += 1;
733 ids.iter().map(|id| (*id, format!("user_{}", id))).collect()
734 });
735
736 let batch_loader = BatchLoader::new(100, loader);
737
738 batch_loader.load_many(&[1, 2, 3]);
740 assert_eq!(*call_count.lock().unwrap(), 1);
741
742 let result = batch_loader.load_many(&[1, 2, 3, 4, 5]);
744 assert_eq!(result.len(), 5);
745 assert_eq!(*call_count.lock().unwrap(), 2); assert_eq!(batch_loader.cache_size(), 5);
749 }
750
751 #[test]
752 fn test_batch_loader_clear_cache() {
753 let loader = make_loader();
754 loader.load_many(&[1, 2]);
755 assert_eq!(loader.cache_size(), 2);
756
757 loader.clear_cache();
758 assert_eq!(loader.cache_size(), 0);
759 }
760
761 #[test]
762 fn test_batch_loader_empty_input() {
763 let loader = make_loader();
764 let result = loader.load_many(&[]);
765 assert!(result.is_empty());
766 }
767
768 #[test]
769 fn test_batch_loader_batch_size_attribute() {
770 let loader = make_loader();
771 assert_eq!(loader.batch_size(), 2);
772 }
773
774 #[test]
775 fn test_batch_loader_with_size_1() {
776 let loader = BatchLoader::new(
777 1,
778 Box::new(|ids: &[i64]| ids.iter().map(|id| (*id, *id * 10)).collect()),
779 );
780 let result = loader.load_many(&[1, 2, 3]);
781 assert_eq!(result.len(), 3);
782 assert_eq!(result.get(&1), Some(&10));
783 assert_eq!(result.get(&2), Some(&20));
784 assert_eq!(result.get(&3), Some(&30));
785 }
786
787 #[test]
790 fn test_workflow_graph_and_batch_loader() {
791 let mut graph = EntityGraph::new();
793 graph.add_edge_with_graph("user", "posts", {
794 let mut sub = EntityGraph::new();
795 sub.add_edge("posts", "comments");
796 sub
797 });
798 assert_eq!(graph.all_relations_recursive().len(), 2);
799
800 let user_loader = BatchLoader::new(
802 50,
803 Box::new(|ids: &[i64]| ids.iter().map(|id| (*id, format!("User#{}", id))).collect()),
804 );
805
806 let user_ids: Vec<i64> = (1..=123).collect();
808 let users = user_loader.load_many(&user_ids);
809 assert_eq!(users.len(), 123);
810 assert_eq!(user_loader.cache_size(), 123);
811 }
812
813 #[test]
814 fn test_n_plus_1_problem_solved() {
815 let query_count = std::sync::Arc::new(std::sync::Mutex::new(0u32));
820 let query_count_clone = query_count.clone();
821
822 let post_loader = BatchLoader::new(
823 100,
824 Box::new(move |user_ids: &[i64]| {
825 *query_count_clone.lock().unwrap() += 1;
826 user_ids
828 .iter()
829 .map(|uid| (*uid, format!("posts_for_user_{}", uid)))
830 .collect()
831 }),
832 );
833
834 let user_ids: Vec<i64> = (1..=250).collect();
836 let _posts = post_loader.load_many(&user_ids);
837
838 assert_eq!(*query_count.lock().unwrap(), 3);
840 }
841}