dataforge/customization/
mod.rs1use std::collections::HashMap;
6use serde_json::Value;
7use crate::error::{DataForgeError, Result};
8
9pub mod api;
10pub mod plugins;
11
12pub use api::*;
13pub use plugins::*;
14
15pub trait CustomGenerator: Send + Sync {
17 fn name(&self) -> &str;
19
20 fn description(&self) -> &str;
22
23 fn generate(&self) -> Result<Value>;
25
26 fn generate_with_params(&self, params: &HashMap<String, Value>) -> Result<Value> {
28 let _ = params;
30 self.generate()
31 }
32
33 fn validate_params(&self, params: &HashMap<String, Value>) -> Result<()> {
35 let _ = params;
37 Ok(())
38 }
39
40 fn param_schema(&self) -> HashMap<String, ParamType> {
42 HashMap::new()
43 }
44}
45
46#[derive(Debug, Clone, PartialEq)]
48pub enum ParamType {
49 String { max_length: Option<usize> },
50 Integer { min: Option<i64>, max: Option<i64> },
51 Float { min: Option<f64>, max: Option<f64> },
52 Boolean,
53 Array(Box<ParamType>),
54 Object(HashMap<String, ParamType>),
55}
56
57pub struct CustomGeneratorRegistry {
59 generators: HashMap<String, Box<dyn CustomGenerator>>,
60}
61
62impl CustomGeneratorRegistry {
63 pub fn new() -> Self {
65 Self {
66 generators: HashMap::new(),
67 }
68 }
69
70 pub fn register<G>(&mut self, generator: G) -> Result<()>
72 where
73 G: CustomGenerator + 'static,
74 {
75 let name = generator.name().to_string();
76
77 if self.generators.contains_key(&name) {
78 return Err(DataForgeError::validation(&format!(
79 "Generator '{}' is already registered", name
80 )));
81 }
82
83 self.generators.insert(name, Box::new(generator));
84 Ok(())
85 }
86
87 pub fn get(&self, name: &str) -> Option<&dyn CustomGenerator> {
89 self.generators.get(name).map(|g| g.as_ref())
90 }
91
92 pub fn unregister(&mut self, name: &str) -> bool {
94 self.generators.remove(name).is_some()
95 }
96
97 pub fn list_generators(&self) -> Vec<GeneratorInfo> {
99 self.generators.iter().map(|(name, generator)| {
100 GeneratorInfo {
101 name: name.clone(),
102 description: generator.description().to_string(),
103 param_schema: generator.param_schema(),
104 }
105 }).collect()
106 }
107
108 pub fn generate(&self, name: &str, params: Option<&HashMap<String, Value>>) -> Result<Value> {
110 let generator = self.generators.get(name)
111 .ok_or_else(|| DataForgeError::validation(&format!("Generator '{}' not found", name)))?;
112
113 if let Some(params) = params {
114 generator.validate_params(params)?;
115 generator.generate_with_params(params)
116 } else {
117 generator.generate()
118 }
119 }
120
121 pub fn count(&self) -> usize {
123 self.generators.len()
124 }
125}
126
127impl Default for CustomGeneratorRegistry {
128 fn default() -> Self {
129 Self::new()
130 }
131}
132
133#[derive(Debug, Clone)]
135pub struct GeneratorInfo {
136 pub name: String,
137 pub description: String,
138 pub param_schema: HashMap<String, ParamType>,
139}
140
141pub struct RandomStringGenerator {
143 name: String,
144}
145
146impl RandomStringGenerator {
147 pub fn new() -> Self {
148 Self {
149 name: "random_string".to_string(),
150 }
151 }
152}
153
154impl CustomGenerator for RandomStringGenerator {
155 fn name(&self) -> &str {
156 &self.name
157 }
158
159 fn description(&self) -> &str {
160 "Generates random strings with configurable length and character set"
161 }
162
163 fn generate(&self) -> Result<Value> {
164 use rand::Rng;
165 let mut rng = rand::thread_rng();
166 let length = rng.gen_range(5..=20);
167 let chars: Vec<char> = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789".chars().collect();
168
169 let random_string: String = (0..length)
170 .map(|_| chars[rng.gen_range(0..chars.len())])
171 .collect();
172
173 Ok(Value::String(random_string))
174 }
175
176 fn generate_with_params(&self, params: &HashMap<String, Value>) -> Result<Value> {
177 use rand::Rng;
178 let mut rng = rand::thread_rng();
179
180 let length = params.get("length")
181 .and_then(|v| v.as_u64())
182 .unwrap_or(10) as usize;
183
184 let charset = params.get("charset")
185 .and_then(|v| v.as_str())
186 .unwrap_or("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789");
187
188 let chars: Vec<char> = charset.chars().collect();
189 if chars.is_empty() {
190 return Err(DataForgeError::validation("Character set cannot be empty"));
191 }
192
193 let random_string: String = (0..length)
194 .map(|_| chars[rng.gen_range(0..chars.len())])
195 .collect();
196
197 Ok(Value::String(random_string))
198 }
199
200 fn validate_params(&self, params: &HashMap<String, Value>) -> Result<()> {
201 if let Some(length) = params.get("length") {
202 if let Some(len) = length.as_u64() {
203 if len == 0 || len > 1000 {
204 return Err(DataForgeError::validation("Length must be between 1 and 1000"));
205 }
206 } else {
207 return Err(DataForgeError::validation("Length must be a positive integer"));
208 }
209 }
210
211 if let Some(charset) = params.get("charset") {
212 if let Some(cs) = charset.as_str() {
213 if cs.is_empty() {
214 return Err(DataForgeError::validation("Character set cannot be empty"));
215 }
216 } else {
217 return Err(DataForgeError::validation("Character set must be a string"));
218 }
219 }
220
221 Ok(())
222 }
223
224 fn param_schema(&self) -> HashMap<String, ParamType> {
225 let mut schema = HashMap::new();
226 schema.insert("length".to_string(), ParamType::Integer { min: Some(1), max: Some(1000) });
227 schema.insert("charset".to_string(), ParamType::String { max_length: None });
228 schema
229 }
230}
231
232impl Default for RandomStringGenerator {
233 fn default() -> Self {
234 Self::new()
235 }
236}
237
238#[cfg(test)]
239mod tests {
240 use super::*;
241
242 #[test]
243 fn test_custom_generator_registry() {
244 let mut registry = CustomGeneratorRegistry::new();
245 let generator = RandomStringGenerator::new();
246
247 assert!(registry.register(generator).is_ok());
248 assert_eq!(registry.count(), 1);
249
250 let result = registry.generate("random_string", None);
251 assert!(result.is_ok());
252
253 if let Ok(Value::String(s)) = result {
254 assert!(!s.is_empty());
255 } else {
256 panic!("Expected string value");
257 }
258 }
259
260 #[test]
261 fn test_random_string_generator() {
262 let generator = RandomStringGenerator::new();
263
264 assert_eq!(generator.name(), "random_string");
265 assert!(!generator.description().is_empty());
266
267 let result = generator.generate();
268 assert!(result.is_ok());
269
270 if let Ok(Value::String(s)) = result {
271 assert!(s.len() >= 5 && s.len() <= 20);
272 }
273 }
274
275 #[test]
276 fn test_random_string_generator_with_params() {
277 let generator = RandomStringGenerator::new();
278 let mut params = HashMap::new();
279 params.insert("length".to_string(), Value::Number(serde_json::Number::from(5)));
280 params.insert("charset".to_string(), Value::String("abc".to_string()));
281
282 assert!(generator.validate_params(¶ms).is_ok());
283
284 let result = generator.generate_with_params(¶ms);
285 assert!(result.is_ok());
286
287 if let Ok(Value::String(s)) = result {
288 assert_eq!(s.len(), 5);
289 assert!(s.chars().all(|c| "abc".contains(c)));
290 }
291 }
292
293 #[test]
294 fn test_invalid_params() {
295 let generator = RandomStringGenerator::new();
296 let mut params = HashMap::new();
297 params.insert("length".to_string(), Value::Number(serde_json::Number::from(0)));
298
299 assert!(generator.validate_params(¶ms).is_err());
300
301 params.insert("length".to_string(), Value::Number(serde_json::Number::from(2000)));
302 assert!(generator.validate_params(¶ms).is_err());
303
304 params.insert("charset".to_string(), Value::String("".to_string()));
305 assert!(generator.validate_params(¶ms).is_err());
306 }
307
308 #[test]
309 fn test_duplicate_registration() {
310 let mut registry = CustomGeneratorRegistry::new();
311 let generator1 = RandomStringGenerator::new();
312 let generator2 = RandomStringGenerator::new();
313
314 assert!(registry.register(generator1).is_ok());
315 assert!(registry.register(generator2).is_err());
316 }
317
318 #[test]
319 fn test_list_generators() {
320 let mut registry = CustomGeneratorRegistry::new();
321 let generator = RandomStringGenerator::new();
322
323 registry.register(generator).unwrap();
324
325 let generators = registry.list_generators();
326 assert_eq!(generators.len(), 1);
327 assert_eq!(generators[0].name, "random_string");
328 assert!(!generators[0].description.is_empty());
329 assert!(!generators[0].param_schema.is_empty());
330 }
331}