sz_orm_core/seeding/
faker.rs

1//! FakerGenerator — faker 数据生成器
2//!
3//! 按字段类型生成随机/语义化假数据，支持字段语义自定义生成器。
4
5use super::{FieldDef, FieldGenerator, FieldType, ModelDef, Record};
6use rand::rngs::StdRng;
7use rand::Rng;
8use rand::SeedableRng;
9use serde_json::{json, Value};
10use std::collections::HashMap;
11
12/// Faker 数据生成器
13pub struct FakerGenerator {
14    field_generators: HashMap<String, Box<dyn FieldGenerator>>,
15    rng: StdRng,
16}
17
18impl FakerGenerator {
19    /// 创建新的 FakerGenerator，注册内置字段生成器
20    pub fn new() -> Self {
21        let mut gen = Self {
22            field_generators: HashMap::new(),
23            rng: StdRng::from_entropy(),
24        };
25        gen.register_defaults();
26        gen
27    }
28
29    /// 创建带种子 rng 的 FakerGenerator（可复现）
30    pub fn with_seed(seed: u64) -> Self {
31        let mut gen = Self {
32            field_generators: HashMap::new(),
33            rng: StdRng::seed_from_u64(seed),
34        };
35        gen.register_defaults();
36        gen
37    }
38
39    fn register_defaults(&mut self) {
40        self.register("name", Box::new(NameGenerator));
41        self.register("email", Box::new(EmailGenerator));
42        self.register("address", Box::new(AddressGenerator));
43        self.register("phone", Box::new(PhoneGenerator));
44        self.register("uuid", Box::new(UuidGenerator));
45        self.register("date", Box::new(DateGenerator));
46        self.register("boolean", Box::new(BooleanGenerator));
47    }
48
49    /// 注册字段语义自定义生成器
50    pub fn register(&mut self, semantic: &str, generator: Box<dyn FieldGenerator>) {
51        self.field_generators
52            .insert(semantic.to_string(), generator);
53    }
54
55    /// 按字段类型推断默认生成器
56    pub fn infer_generator(field_type: &FieldType) -> Box<dyn FieldGenerator> {
57        match field_type {
58            FieldType::String => Box::new(NameGenerator),
59            FieldType::I32 | FieldType::U32 => Box::new(NumberGenerator::new(1, 100)),
60            FieldType::I64 | FieldType::U64 => Box::new(NumberGenerator::new(1, 10000)),
61            FieldType::F64 => Box::new(FloatGenerator::new(0.0, 100.0)),
62            FieldType::Boolean => Box::new(BooleanGenerator),
63            FieldType::Uuid => Box::new(UuidGenerator),
64            FieldType::DateTime => Box::new(DateGenerator),
65            FieldType::Json => Box::new(JsonGenerator),
66            FieldType::Enum(variants) => Box::new(EnumGenerator::new(variants.clone())),
67        }
68    }
69
70    /// 批量生成记录
71    pub fn generate_batch(&mut self, model: &ModelDef, count: usize) -> Vec<Record> {
72        let mut records = Vec::with_capacity(count);
73        for _ in 0..count {
74            let mut record = serde_json::Map::new();
75            for field in &model.fields {
76                let value = self.generate_field(field);
77                record.insert(field.name.clone(), value);
78            }
79            records.push(record);
80        }
81        records
82    }
83
84    fn generate_field(&mut self, field: &FieldDef) -> Value {
85        if field.nullable && self.rng.gen_bool(0.1) {
86            return Value::Null;
87        }
88        let semantic_key = format!("{}.{}", "", field.name);
89        if let Some(gen) = self.field_generators.get(&field.name) {
90            return gen.generate(&mut self.rng);
91        }
92        if let Some(gen) = self.field_generators.get(&semantic_key) {
93            return gen.generate(&mut self.rng);
94        }
95        let gen = Self::infer_generator(&field.field_type);
96        gen.generate(&mut self.rng)
97    }
98}
99
100impl Default for FakerGenerator {
101    fn default() -> Self {
102        Self::new()
103    }
104}
105
106// ===== 内置字段生成器 =====
107
108/// 姓名生成器（中文常见姓名）
109pub struct NameGenerator;
110impl FieldGenerator for NameGenerator {
111    fn generate(&self, rng: &mut StdRng) -> Value {
112        const NAMES: &[&str] = &[
113            "张伟", "王芳", "李娜", "刘洋", "陈静", "杨磊", "赵敏", "黄强", "周丽", "吴杰", "徐涛",
114            "孙艳", "马超", "朱琳", "胡军", "郭雪", "何明", "高飞", "林秀", "罗刚",
115        ];
116        Value::String(NAMES[rng.gen_range(0..NAMES.len())].to_string())
117    }
118}
119
120/// 邮箱生成器
121pub struct EmailGenerator;
122impl FieldGenerator for EmailGenerator {
123    fn generate(&self, rng: &mut StdRng) -> Value {
124        const DOMAINS: &[&str] = &[
125            "gmail.com",
126            "qq.com",
127            "163.com",
128            "outlook.com",
129            "huawei.com",
130        ];
131        let id: u64 = rng.gen_range(1000..99999);
132        let domain = DOMAINS[rng.gen_range(0..DOMAINS.len())];
133        Value::String(format!("user{}@{}", id, domain))
134    }
135}
136
137/// 地址生成器
138pub struct AddressGenerator;
139impl FieldGenerator for AddressGenerator {
140    fn generate(&self, rng: &mut StdRng) -> Value {
141        const CITIES: &[&str] = &[
142            "北京", "上海", "广州", "深圳", "杭州", "成都", "武汉", "西安",
143        ];
144        let city = CITIES[rng.gen_range(0..CITIES.len())];
145        let street: u32 = rng.gen_range(1..999);
146        Value::String(format!("{}市XX路{}号", city, street))
147    }
148}
149
150/// 手机号生成器
151pub struct PhoneGenerator;
152impl FieldGenerator for PhoneGenerator {
153    fn generate(&self, rng: &mut StdRng) -> Value {
154        let prefix: u32 = rng.gen_range(130..190);
155        let suffix: u64 = rng.gen_range(10000000..99999999);
156        Value::String(format!("{}{:08}", prefix, suffix))
157    }
158}
159
160/// UUID 生成器
161pub struct UuidGenerator;
162impl FieldGenerator for UuidGenerator {
163    fn generate(&self, rng: &mut StdRng) -> Value {
164        let bytes: [u8; 16] = rng.gen();
165        format!(
166            "{:02x}{:02x}{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}{:02x}{:02x}{:02x}{:02x}",
167            bytes[0], bytes[1], bytes[2], bytes[3],
168            bytes[4], bytes[5], bytes[6], bytes[7],
169            bytes[8], bytes[9], bytes[10], bytes[11],
170            bytes[12], bytes[13], bytes[14], bytes[15]
171        )
172        .into()
173    }
174}
175
176/// 日期生成器（YYYY-MM-DD 格式）
177pub struct DateGenerator;
178impl FieldGenerator for DateGenerator {
179    fn generate(&self, rng: &mut StdRng) -> Value {
180        let year: u32 = rng.gen_range(1970..2100);
181        let month: u32 = rng.gen_range(1..13);
182        let day: u32 = rng.gen_range(1..29);
183        Value::String(format!("{:04}-{:02}-{:02}", year, month, day))
184    }
185}
186
187/// 整数生成器（指定范围）
188pub struct NumberGenerator {
189    min: i64,
190    max: i64,
191}
192impl NumberGenerator {
193    /// 创建指定范围的整数生成器
194    pub fn new(min: i64, max: i64) -> Self {
195        Self { min, max }
196    }
197}
198impl FieldGenerator for NumberGenerator {
199    fn generate(&self, rng: &mut StdRng) -> Value {
200        json!(rng.gen_range(self.min..self.max))
201    }
202}
203
204/// 浮点数生成器（指定范围）
205pub struct FloatGenerator {
206    min: f64,
207    max: f64,
208}
209impl FloatGenerator {
210    /// 创建指定范围的浮点数生成器
211    pub fn new(min: f64, max: f64) -> Self {
212        Self { min, max }
213    }
214}
215impl FieldGenerator for FloatGenerator {
216    fn generate(&self, rng: &mut StdRng) -> Value {
217        let val = rng.gen_range(self.min..self.max);
218        json!(val)
219    }
220}
221
222/// 布尔生成器
223pub struct BooleanGenerator;
224impl FieldGenerator for BooleanGenerator {
225    fn generate(&self, rng: &mut StdRng) -> Value {
226        Value::Bool(rng.gen_bool(0.5))
227    }
228}
229
230/// 枚举生成器（从给定变体中随机选择）
231pub struct EnumGenerator {
232    variants: Vec<String>,
233}
234impl EnumGenerator {
235    /// 创建枚举生成器
236    pub fn new(variants: Vec<String>) -> Self {
237        Self { variants }
238    }
239}
240impl FieldGenerator for EnumGenerator {
241    fn generate(&self, rng: &mut StdRng) -> Value {
242        if self.variants.is_empty() {
243            return Value::Null;
244        }
245        Value::String(self.variants[rng.gen_range(0..self.variants.len())].clone())
246    }
247}
248
249/// JSON 生成器
250pub struct JsonGenerator;
251impl FieldGenerator for JsonGenerator {
252    fn generate(&self, rng: &mut StdRng) -> Value {
253        json!({
254            "id": rng.gen::<u32>(),
255            "active": rng.gen_bool(0.5),
256        })
257    }
258}
259
260#[cfg(test)]
261mod tests {
262    use super::*;
263    use crate::seeding::FieldDef;
264
265    fn user_model() -> ModelDef {
266        ModelDef {
267            table: "users".to_string(),
268            fields: vec![
269                FieldDef {
270                    name: "name".to_string(),
271                    field_type: FieldType::String,
272                    nullable: false,
273                },
274                FieldDef {
275                    name: "email".to_string(),
276                    field_type: FieldType::String,
277                    nullable: false,
278                },
279                FieldDef {
280                    name: "age".to_string(),
281                    field_type: FieldType::U32,
282                    nullable: false,
283                },
284            ],
285        }
286    }
287
288    #[test]
289    fn test_generate_batch_basic() {
290        let mut faker = FakerGenerator::with_seed(42);
291        let model = user_model();
292        let records = faker.generate_batch(&model, 10);
293        assert_eq!(records.len(), 10);
294        for record in &records {
295            assert!(record.contains_key("name"));
296            assert!(record.contains_key("email"));
297            assert!(record.contains_key("age"));
298            let name = record["name"].as_str().unwrap();
299            assert!(!name.is_empty());
300        }
301    }
302
303    #[test]
304    fn test_email_generator_format() {
305        let mut faker = FakerGenerator::with_seed(42);
306        let model = ModelDef {
307            table: "test".to_string(),
308            fields: vec![FieldDef {
309                name: "email".to_string(),
310                field_type: FieldType::String,
311                nullable: false,
312            }],
313        };
314        let records = faker.generate_batch(&model, 5);
315        for record in &records {
316            let email = record["email"].as_str().unwrap();
317            assert!(email.contains('@'));
318        }
319    }
320
321    #[test]
322    fn test_register_custom_generator() {
323        let mut faker = FakerGenerator::with_seed(42);
324        faker.register("custom_field", Box::new(NameGenerator));
325        let model = ModelDef {
326            table: "test".to_string(),
327            fields: vec![FieldDef {
328                name: "custom_field".to_string(),
329                field_type: FieldType::String,
330                nullable: false,
331            }],
332        };
333        let records = faker.generate_batch(&model, 3);
334        assert_eq!(records.len(), 3);
335        for record in &records {
336            assert!(record["custom_field"].is_string());
337        }
338    }
339
340    #[test]
341    fn test_nullable_field() {
342        let mut faker = FakerGenerator::with_seed(42);
343        let model = ModelDef {
344            table: "test".to_string(),
345            fields: vec![FieldDef {
346                name: "opt".to_string(),
347                field_type: FieldType::String,
348                nullable: true,
349            }],
350        };
351        let records = faker.generate_batch(&model, 100);
352        let null_count = records.iter().filter(|r| r["opt"].is_null()).count();
353        assert!(null_count > 0, "nullable field should produce some nulls");
354    }
355
356    #[test]
357    fn test_uuid_generator_format() {
358        let mut rng = StdRng::seed_from_u64(42);
359        let val = UuidGenerator.generate(&mut rng);
360        let uuid_str = val.as_str().unwrap();
361        assert_eq!(uuid_str.len(), 36);
362        assert_eq!(uuid_str.chars().nth(8), Some('-'));
363        assert_eq!(uuid_str.chars().nth(13), Some('-'));
364        assert_eq!(uuid_str.chars().nth(18), Some('-'));
365        assert_eq!(uuid_str.chars().nth(23), Some('-'));
366    }
367
368    #[test]
369    fn test_enum_generator() {
370        let mut rng = StdRng::seed_from_u64(42);
371        let gen = EnumGenerator::new(vec!["active".into(), "inactive".into(), "banned".into()]);
372        let val = gen.generate(&mut rng);
373        let s = val.as_str().unwrap();
374        assert!(["active", "inactive", "banned"].contains(&s));
375    }
376
377    #[test]
378    fn test_infer_generator() {
379        let gen = FakerGenerator::infer_generator(&FieldType::Boolean);
380        let mut rng = StdRng::seed_from_u64(42);
381        let val = gen.generate(&mut rng);
382        assert!(val.is_boolean());
383    }
384}
sz_orm_core/seeding/faker.rs

sz_orm_core/seeding/
faker.rs
