Skip to main content

sz_orm_core/
entity_graph.rs

1//! Entity Graph + @BatchSize 批量抓取
2//!
3//! 对应文档 6.8 节改进项 27(Entity Graph)+ 28(@BatchSize 批量抓取)。
4//!
5//! # 核心概念
6//!
7//! - **EntityGraph**:实体图,定义一组关联关系一起抓取,避免 N+1 查询
8//! - **BatchSizeConfig**:批量抓取配置(大小 + 策略)
9//! - **BatchLoader**:通用批量加载器,将 N 次单条查询合并为 ⌈N/batch_size⌉ 次批量查询
10//! - **BatchStrategy**:批量策略(IN / JOIN / SUBQUERY)
11//!
12//! # 设计灵感
13//!
14//! - Hibernate `@NamedEntityGraph` / `@BatchSize`
15//! - Doctrine `FetchMode::EAGER` / partial
16//! - Django `select_related` / `prefetch_related`
17//! - Sequelize `include` + `separate: true`
18//!
19//! # 使用示例
20//!
21//! ```
22//! use sz_orm_core::entity_graph::{EntityGraph, BatchLoader, BatchSizeConfig, BatchStrategy};
23//! use std::collections::HashMap;
24//!
25//! // 1. 定义 EntityGraph:抓取 user.posts.comments
26//! let mut graph = EntityGraph::new();
27//! graph.add_edge("user", "posts");
28//! graph.add_edge_with_graph("posts", "comments", {
29//!     let mut sub = EntityGraph::new();
30//!     sub.add_edge("comments", "author");
31//!     sub
32//! });
33//!
34//! // 2. 使用 BatchLoader 批量加载用户
35//! fn load_users(ids: &[i64]) -> HashMap<i64, String> {
36//!     ids.iter().map(|id| (*id, format!("user_{}", id))).collect()
37//! }
38//!
39//! let loader = BatchLoader::new(100, Box::new(load_users));
40//! let users = loader.load_many(&[1, 2, 3, 4, 5]);
41//! assert_eq!(users.len(), 5);
42//! assert_eq!(users.get(&3), Some(&"user_3".to_string()));
43//! ```
44
45use std::collections::HashMap;
46use std::hash::Hash;
47use std::sync::RwLock;
48
49// ============================================================================
50// EntityGraph — 实体图
51// ============================================================================
52
53/// 实体图边(关联关系)
54#[derive(Debug, Clone)]
55pub struct GraphEdge {
56    /// 父字段名
57    pub parent_field: String,
58    /// 关联名(如 "posts"、"comments")
59    pub relation: String,
60    /// 嵌套子图(用于递归抓取)
61    pub sub_graph: Option<Box<EntityGraph>>,
62}
63
64/// 实体图
65///
66/// 描述一组关联关系的抓取计划。
67///
68/// # 示例
69///
70/// ```
71/// use sz_orm_core::entity_graph::EntityGraph;
72///
73/// let mut graph = EntityGraph::new();
74/// graph.add_edge("user", "profile");
75/// graph.add_edge("user", "posts");
76/// ```
77#[derive(Debug, Clone, Default)]
78pub struct EntityGraph {
79    /// 边列表
80    edges: Vec<GraphEdge>,
81}
82
83impl EntityGraph {
84    /// 创建空实体图
85    pub fn new() -> Self {
86        Self { edges: Vec::new() }
87    }
88
89    /// 添加一条边
90    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    /// 添加一条带子图的边(嵌套抓取)
104    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    /// 返回所有边
119    pub fn edges(&self) -> &[GraphEdge] {
120        &self.edges
121    }
122
123    /// 返回边的数量
124    pub fn edge_count(&self) -> usize {
125        self.edges.len()
126    }
127
128    /// 查询某个父字段的所有关联
129    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    /// 收集图中所有关联名(去重)
137    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    /// 收集图中所有父字段(去重)
145    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    /// 是否为空图
153    pub fn is_empty(&self) -> bool {
154        self.edges.is_empty()
155    }
156
157    /// 递归收集图中所有关联(含子图)
158    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// ============================================================================
173// BatchStrategy — 批量策略
174// ============================================================================
175
176/// 批量抓取策略
177#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
178pub enum BatchStrategy {
179    /// 使用 `WHERE id IN (?, ?, ...)` 子句批量加载
180    ///
181    /// 适用场景:关联数量较少、目标表无索引时的备选方案
182    #[default]
183    In,
184    /// 使用 `LEFT JOIN` 一次性加载所有关联
185    ///
186    /// 适用场景:关联数量较少、需要原子性读取
187    Join,
188    /// 使用 `WHERE id IN (SELECT ... FROM ...)` 子查询批量加载
189    ///
190    /// 适用场景:子查询可被数据库优化器优化时
191    Subquery,
192}
193
194impl BatchStrategy {
195    /// 策略名称
196    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    /// 生成 IN 子句的 SQL 片段
205    ///
206    /// 返回形如 `"id IN (?, ?, ?)"` 的字符串(占位符数量与 values 一致)。
207    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// ============================================================================
217// BatchSizeConfig — 批量大小配置
218// ============================================================================
219
220/// 批量大小配置
221///
222/// 对应 Hibernate `@BatchSize(size = 100)` 注解。
223#[derive(Debug, Clone, Copy)]
224pub struct BatchSizeConfig {
225    /// 每批数量
226    pub size: usize,
227    /// 抓取策略
228    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    /// 创建配置
242    pub fn new(size: usize, strategy: BatchStrategy) -> Self {
243        Self { size, strategy }
244    }
245
246    /// 创建默认策略的配置(IN)
247    pub fn with_size(size: usize) -> Self {
248        Self {
249            size,
250            strategy: BatchStrategy::In,
251        }
252    }
253
254    /// 计算给定总数需要分多少批
255    ///
256    /// # 示例
257    ///
258    /// ```
259    /// use sz_orm_core::entity_graph::BatchSizeConfig;
260    ///
261    /// let config = BatchSizeConfig::with_size(100);
262    /// assert_eq!(config.batch_count(0), 0);
263    /// assert_eq!(config.batch_count(1), 1);
264    /// assert_eq!(config.batch_count(100), 1);
265    /// assert_eq!(config.batch_count(101), 2);
266    /// assert_eq!(config.batch_count(250), 3);
267    /// ```
268    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    /// 返回第 `batch_index` 批的范围(start..end,end 不超过 total)
277    ///
278    /// # 示例
279    ///
280    /// ```
281    /// use sz_orm_core::entity_graph::BatchSizeConfig;
282    ///
283    /// let config = BatchSizeConfig::with_size(100);
284    /// assert_eq!(config.batch_range(0, 250), 0..100);
285    /// assert_eq!(config.batch_range(1, 250), 100..200);
286    /// assert_eq!(config.batch_range(2, 250), 200..250);
287    /// ```
288    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
295// ============================================================================
296// BatchLoader — 通用批量加载器
297// ============================================================================
298
299/// 批量加载函数类型
300pub type BatchLoaderFn<K, V> = Box<dyn Fn(&[K]) -> HashMap<K, V> + Send + Sync>;
301
302/// 批量加载器
303///
304/// 将 N 个单条加载请求合并为 ⌈N/batch_size⌉ 次批量加载,避免 N+1 查询问题。
305///
306/// # 泛型参数
307///
308/// - `K`:主键类型(必须实现 `Hash + Eq + Clone`)
309/// - `V`:值类型
310///
311/// # 示例
312///
313/// ```
314/// use sz_orm_core::entity_graph::BatchLoader;
315/// use std::collections::HashMap;
316///
317/// fn load_users(ids: &[i64]) -> HashMap<i64, String> {
318///     ids.iter().map(|id| (*id, format!("user_{}", id))).collect()
319/// }
320///
321/// let loader = BatchLoader::new(100, Box::new(load_users));
322/// let users = loader.load_many(&[1, 2, 3]);
323/// assert_eq!(users.len(), 3);
324/// ```
325pub struct BatchLoader<K, V>
326where
327    K: Hash + Eq + Clone + Send + Sync,
328    V: Clone + Send + Sync,
329{
330    /// 每批数量
331    batch_size: usize,
332    /// 实际加载函数
333    loader: BatchLoaderFn<K, V>,
334    /// 缓存(避免重复加载相同的 key)
335    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    /// 创建批量加载器
344    ///
345    /// # 参数
346    /// - `batch_size`:每批数量
347    /// - `loader`:实际加载函数,接收一批 key,返回 key→value 的 HashMap
348    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    /// 批量加载多个 key
357    ///
358    /// - 自动跳过缓存中已有的 key
359    /// - 按 batch_size 分批调用 loader
360    /// - 返回所有 key 对应的 value(包含缓存与新加载的)
361    pub fn load_many(&self, keys: &[K]) -> HashMap<K, V> {
362        let mut result: HashMap<K, V> = HashMap::new();
363
364        // 1. 从缓存读取
365        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        // 2. 分批加载
384        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        // 3. 写入缓存
392        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        // 4. 合并结果
402        result.extend(all_loaded);
403        result
404    }
405
406    /// 加载单个 key(便捷方法)
407    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    /// 清空缓存
413    pub fn clear_cache(&self) {
414        self.cache
415            .write()
416            .expect("BatchLoader cache lock poisoned (clear_cache)")
417            .clear();
418    }
419
420    /// 返回当前缓存大小
421    pub fn cache_size(&self) -> usize {
422        self.cache
423            .read()
424            .expect("BatchLoader cache lock poisoned (cache_size)")
425            .len()
426    }
427
428    /// 返回 batch_size
429    pub fn batch_size(&self) -> usize {
430        self.batch_size
431    }
432}
433
434// ============================================================================
435// 单元测试
436// ============================================================================
437
438#[cfg(test)]
439mod tests {
440    use super::*;
441
442    // ===== EntityGraph 测试 =====
443
444    #[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    // ===== BatchStrategy 测试 =====
547
548    #[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    // ===== BatchSizeConfig 测试 =====
579
580    #[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); // 空范围
640    }
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    // ===== BatchLoader 测试 =====
652
653    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        // batch_size=2, 5 keys → 3 batches
673        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        // loader 返回的 map 没有 key 100
695        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        // 第一次加载
714        batch_loader.load_many(&[1, 2, 3]);
715        assert_eq!(*call_count.lock().unwrap(), 1);
716
717        // 第二次加载相同 key,应命中缓存
718        batch_loader.load_many(&[1, 2, 3]);
719        assert_eq!(*call_count.lock().unwrap(), 1); // 未增加
720
721        // 加载新 key,应触发新的 loader 调用
722        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        // 加载 1, 2, 3
739        batch_loader.load_many(&[1, 2, 3]);
740        assert_eq!(*call_count.lock().unwrap(), 1);
741
742        // 加载 1, 2, 3, 4, 5(前 3 个命中缓存)
743        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); // 只为 4, 5 调用一次
746
747        // 缓存大小应为 5
748        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    // ===== 集成场景测试 =====
788
789    #[test]
790    fn test_workflow_graph_and_batch_loader() {
791        // 模拟 User → Posts → Comments 的批量加载场景
792        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        // 模拟批量加载用户
801        let user_loader = BatchLoader::new(
802            50,
803            Box::new(|ids: &[i64]| ids.iter().map(|id| (*id, format!("User#{}", id))).collect()),
804        );
805
806        // 加载 123 个用户(应分 3 批)
807        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        // 经典 N+1 问题演示:
816        // - 错误做法:N 个用户各发 1 次查询加载 posts → N+1 次查询
817        // - 正确做法:用 BatchLoader 一次批量加载 → ⌈N/batch⌉+1 次查询
818
819        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_id 返回 posts
827                user_ids
828                    .iter()
829                    .map(|uid| (*uid, format!("posts_for_user_{}", uid)))
830                    .collect()
831            }),
832        );
833
834        // 250 个用户
835        let user_ids: Vec<i64> = (1..=250).collect();
836        let _posts = post_loader.load_many(&user_ids);
837
838        // 应分 3 批(100+100+50),调用 loader 3 次
839        assert_eq!(*query_count.lock().unwrap(), 3);
840    }
841}