sz_orm_sharding/
routing.rs

1//! # 分片键提取器
2//!
3//! 提供 `ShardKeyExtractor` trait 与若干实现，用于从业务数据对象中提取分片键，
4//! 配合 [`ShardingRouter::route_by_data`](crate::ShardingRouter::route_by_data) 使用。
5//!
6//! ## 提供的提取器
7//!
8//! - [`FieldExtractor`]：通过闭包从数据中提取字段（闭包捕获字段访问逻辑）
9//! - [`CompositeKeyExtractor`]：组合多个 extractor，结果以 `:` 拼接
10
11use crate::ShardingError;
12use std::any::Any;
13
14/// 分片键提取器 trait
15///
16/// 实现者负责从 `data: &dyn Any` 中提取分片键字符串。
17/// 必须是 `Send + Sync` 以便在跨线程场景（如 `ScatterGather`）中使用。
18pub trait ShardKeyExtractor: Send + Sync {
19    /// 从数据中提取分片键
20    ///
21    /// # Errors
22    ///
23    /// 提取失败时返回 [`ShardingError`]。
24    fn extract(&self, data: &dyn Any) -> Result<String, ShardingError>;
25}
26
27/// 字段提取器：通过闭包从数据中提取字段
28///
29/// 闭包签名 `Fn() -> String`，由调用方在闭包内捕获对数据源的访问方式。
30/// `extract` 收到的 `data: &dyn Any` 参数被忽略（由闭包自行决定如何取值），
31/// 这种设计让用户可以灵活地把外部状态（如 `Arc<T>`、数据库行等）捕获进闭包。
32pub struct FieldExtractor {
33    extractor: Box<dyn Fn() -> String + Send + Sync>,
34}
35
36impl FieldExtractor {
37    /// 创建字段提取器
38    ///
39    /// # 参数
40    ///
41    /// - `f`: 闭包，捕获字段访问逻辑，返回字段值字符串
42    pub fn new<F>(f: F) -> Self
43    where
44        F: Fn() -> String + Send + Sync + 'static,
45    {
46        Self {
47            extractor: Box::new(f),
48        }
49    }
50}
51
52impl ShardKeyExtractor for FieldExtractor {
53    fn extract(&self, _data: &dyn Any) -> Result<String, ShardingError> {
54        Ok((self.extractor)())
55    }
56}
57
58impl std::fmt::Debug for FieldExtractor {
59    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
60        f.debug_struct("FieldExtractor")
61            .field("extractor", &"<closure>")
62            .finish()
63    }
64}
65
66/// 复合键提取器：组合多个 extractor，结果以 `:` 拼接
67///
68/// # 示例
69///
70/// ```rust,ignore
71/// use sz_orm_sharding::routing::{CompositeKeyExtractor, FieldExtractor};
72///
73/// let ext = CompositeKeyExtractor::new()
74///     .with(FieldExtractor::new(|| "cn".to_string()))
75///     .with(FieldExtractor::new(|| "user:123".to_string()));
76/// // extract 后得到 "cn:user:123"
77/// ```
78pub struct CompositeKeyExtractor {
79    extractors: Vec<Box<dyn ShardKeyExtractor>>,
80}
81
82impl CompositeKeyExtractor {
83    /// 创建空的复合提取器
84    pub fn new() -> Self {
85        Self {
86            extractors: Vec::new(),
87        }
88    }
89
90    /// 添加一个子提取器（链式 API）
91    ///
92    /// 命名为 `with` 而非 `add`，避免与 `std::ops::Add::add` trait 方法混淆。
93    pub fn with<E>(mut self, extractor: E) -> Self
94    where
95        E: ShardKeyExtractor + 'static,
96    {
97        self.extractors.push(Box::new(extractor));
98        self
99    }
100
101    /// 返回子提取器数量
102    pub fn len(&self) -> usize {
103        self.extractors.len()
104    }
105
106    /// 是否为空
107    pub fn is_empty(&self) -> bool {
108        self.extractors.is_empty()
109    }
110}
111
112impl Default for CompositeKeyExtractor {
113    fn default() -> Self {
114        Self::new()
115    }
116}
117
118impl ShardKeyExtractor for CompositeKeyExtractor {
119    fn extract(&self, data: &dyn Any) -> Result<String, ShardingError> {
120        let mut parts: Vec<String> = Vec::with_capacity(self.extractors.len());
121        for ext in &self.extractors {
122            parts.push(ext.extract(data)?);
123        }
124        Ok(parts.join(":"))
125    }
126}
127
128impl std::fmt::Debug for CompositeKeyExtractor {
129    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
130        f.debug_struct("CompositeKeyExtractor")
131            .field("count", &self.extractors.len())
132            .finish()
133    }
134}
135
136#[cfg(test)]
137mod tests {
138    use super::*;
139    use crate::{ShardingRouter, ShardingStrategy};
140    use std::collections::HashMap;
141
142    #[test]
143    fn test_field_extractor_basic() {
144        let ext = FieldExtractor::new(|| "key:123".to_string());
145        let key = ext.extract(&"ignored").unwrap();
146        assert_eq!(key, "key:123");
147    }
148
149    #[test]
150    fn test_field_extractor_deterministic() {
151        let ext = FieldExtractor::new(|| "stable_key".to_string());
152        let k1 = ext.extract(&42i32).unwrap();
153        let k2 = ext.extract(&"other").unwrap();
154        assert_eq!(k1, k2);
155        assert_eq!(k1, "stable_key");
156    }
157
158    #[test]
159    fn test_composite_key_extractor_combines() {
160        let ext = CompositeKeyExtractor::new()
161            .with(FieldExtractor::new(|| "cn".to_string()))
162            .with(FieldExtractor::new(|| "user:123".to_string()));
163        let key = ext.extract(&"ignored").unwrap();
164        assert_eq!(key, "cn:user:123");
165    }
166
167    #[test]
168    fn test_composite_key_extractor_three_parts() {
169        let ext = CompositeKeyExtractor::new()
170            .with(FieldExtractor::new(|| "2026-07-18".to_string()))
171            .with(FieldExtractor::new(|| "cn".to_string()))
172            .with(FieldExtractor::new(|| "user:42".to_string()));
173        let key = ext.extract(&"x").unwrap();
174        assert_eq!(key, "2026-07-18:cn:user:42");
175    }
176
177    #[test]
178    fn test_composite_key_extractor_empty() {
179        let ext = CompositeKeyExtractor::new();
180        assert!(ext.is_empty());
181        assert_eq!(ext.len(), 0);
182        let key = ext.extract(&"ignored").unwrap();
183        assert_eq!(key, "");
184    }
185
186    #[test]
187    fn test_route_by_data_with_field_extractor_enum() {
188        let mut mapping = HashMap::new();
189        mapping.insert("key:123".to_string(), "shard_a".to_string());
190        let router = ShardingRouter::new_enum(mapping, None);
191        let ext = FieldExtractor::new(|| "key:123".to_string());
192        let result = router.route_by_data(&"data", &ext).unwrap();
193        assert_eq!(result, "shard_a");
194    }
195
196    #[test]
197    fn test_route_by_data_no_match_errors() {
198        let router = ShardingRouter::new_enum(HashMap::new(), None);
199        let ext = FieldExtractor::new(|| "missing".to_string());
200        let result = router.route_by_data(&"data", &ext);
201        assert!(matches!(result, Err(ShardingError::NoMappingForKey(_))));
202    }
203
204    #[test]
205    fn test_route_by_data_with_hash_strategy() {
206        let router = ShardingRouter::new(ShardingStrategy::Hash, vec!["s0", "s1"]);
207        let ext = FieldExtractor::new(|| "user:42".to_string());
208        let result = router.route_by_data(&"data", &ext).unwrap();
209        assert!(result == "s0" || result == "s1");
210    }
211
212    #[test]
213    fn test_route_by_data_with_composite_extractor() {
214        // 二级 key 形如 "cn:user:123"
215        let mut mapping = HashMap::new();
216        mapping.insert("cn:user:123".to_string(), "shard_x".to_string());
217        let router = ShardingRouter::new_enum(mapping, None);
218        let ext = CompositeKeyExtractor::new()
219            .with(FieldExtractor::new(|| "cn".to_string()))
220            .with(FieldExtractor::new(|| "user:123".to_string()));
221        let result = router.route_by_data(&"data", &ext).unwrap();
222        assert_eq!(result, "shard_x");
223    }
224
225    #[test]
226    fn test_field_extractor_debug() {
227        let ext = FieldExtractor::new(|| "k".to_string());
228        let s = format!("{:?}", ext);
229        assert!(s.contains("FieldExtractor"));
230    }
231
232    #[test]
233    fn test_composite_key_extractor_debug() {
234        let ext = CompositeKeyExtractor::new().with(FieldExtractor::new(|| "a".to_string()));
235        let s = format!("{:?}", ext);
236        assert!(s.contains("CompositeKeyExtractor"));
237        assert!(s.contains("count"));
238    }
239}
sz_orm_sharding/routing.rs

sz_orm_sharding/
routing.rs
