1use crate::foreign_keys::ForeignKeyDefinition;
24use crate::types::Type;
25use serde::{Deserialize, Serialize};
26use std::collections::HashMap;
27use std::fs;
28use std::path::Path;
29
30#[derive(Debug, thiserror::Error)]
36pub enum SchemaError {
37 #[error("Failed to read schema file: {0}")]
39 IoError(#[from] std::io::Error),
40
41 #[error("Failed to parse YAML: {0}")]
43 YamlError(#[from] serde_yaml::Error),
44
45 #[error("Table not found: {0}")]
47 TableNotFound(String),
48
49 #[error("Field '{field}' not found in table '{table}'")]
51 FieldNotFound { table: String, field: String },
52
53 #[error("Column '{column}' not found in table '{table}'")]
55 ColumnNotFound { table: String, column: String },
56}
57
58#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
67pub struct ColumnDefinition {
68 pub name: String,
70
71 #[serde(rename = "type")]
73 pub column_type: Type,
74
75 #[serde(default)]
77 pub nullable: bool,
78}
79
80impl ColumnDefinition {
81 pub fn new(name: impl Into<String>, column_type: Type) -> Self {
83 Self {
84 name: name.into(),
85 column_type,
86 nullable: false,
87 }
88 }
89
90 pub fn nullable(name: impl Into<String>, column_type: Type) -> Self {
92 Self {
93 name: name.into(),
94 column_type,
95 nullable: true,
96 }
97 }
98}
99
100#[derive(Debug, Clone, Serialize, Deserialize)]
105pub struct TableDefinition {
106 pub name: String,
108
109 pub primary_key: ColumnDefinition,
111
112 pub columns: Vec<ColumnDefinition>,
114
115 #[serde(default)]
117 pub foreign_keys: Vec<ForeignKeyDefinition>,
118
119 #[serde(default)]
122 pub composite_primary_key: Option<Vec<String>>,
123}
124
125impl TableDefinition {
126 pub fn new(
128 name: impl Into<String>,
129 primary_key: ColumnDefinition,
130 columns: Vec<ColumnDefinition>,
131 ) -> Self {
132 Self {
133 name: name.into(),
134 primary_key,
135 columns,
136 foreign_keys: Vec::new(),
137 composite_primary_key: None,
138 }
139 }
140
141 pub fn get_column(&self, name: &str) -> Option<&ColumnDefinition> {
143 if self.primary_key.name == name {
144 return Some(&self.primary_key);
145 }
146 self.columns.iter().find(|c| c.name == name)
147 }
148
149 pub fn get_column_type(&self, name: &str) -> Option<&Type> {
151 self.get_column(name).map(|c| &c.column_type)
152 }
153
154 pub fn column_names(&self) -> Vec<&str> {
156 let mut names = vec![self.primary_key.name.as_str()];
157 names.extend(self.columns.iter().map(|c| c.name.as_str()));
158 names
159 }
160
161 pub fn primary_key_column_names(&self) -> Vec<&str> {
163 if let Some(ref cpk) = self.composite_primary_key {
164 cpk.iter().map(|s| s.as_str()).collect()
165 } else {
166 vec![self.primary_key.name.as_str()]
167 }
168 }
169}
170
171#[derive(Debug, Clone, Default, Serialize, Deserialize)]
176pub struct DatabaseSchema {
177 pub tables: Vec<TableDefinition>,
179
180 #[serde(skip)]
182 table_map: HashMap<String, usize>,
183}
184
185impl DatabaseSchema {
186 pub fn new(tables: Vec<TableDefinition>) -> Self {
188 let mut schema = Self {
189 tables,
190 table_map: HashMap::new(),
191 };
192 schema.build_table_map();
193 schema
194 }
195
196 fn build_table_map(&mut self) {
198 self.table_map = self
199 .tables
200 .iter()
201 .enumerate()
202 .map(|(idx, table)| (table.name.clone(), idx))
203 .collect();
204 }
205
206 pub fn get_table(&self, name: &str) -> Option<&TableDefinition> {
208 self.table_map
209 .get(name)
210 .and_then(|&idx| self.tables.get(idx))
211 }
212
213 pub fn get_table_mut(&mut self, name: &str) -> Option<&mut TableDefinition> {
215 let idx = *self.table_map.get(name)?;
216 self.tables.get_mut(idx)
217 }
218
219 pub fn get_column_type(&self, table: &str, column: &str) -> Result<&Type, SchemaError> {
221 let table_schema = self
222 .get_table(table)
223 .ok_or_else(|| SchemaError::TableNotFound(table.to_string()))?;
224
225 table_schema
226 .get_column_type(column)
227 .ok_or_else(|| SchemaError::ColumnNotFound {
228 table: table.to_string(),
229 column: column.to_string(),
230 })
231 }
232
233 pub fn table_names(&self) -> Vec<&str> {
235 self.tables.iter().map(|t| t.name.as_str()).collect()
236 }
237
238 pub fn add_table(&mut self, table: TableDefinition) {
240 let idx = self.tables.len();
241 self.table_map.insert(table.name.clone(), idx);
242 self.tables.push(table);
243 }
244}
245
246#[derive(Debug, Clone, Serialize, Deserialize)]
255#[serde(tag = "type", rename_all = "snake_case")]
256pub enum GeneratorConfig {
257 UuidV4,
259
260 Sequential {
262 #[serde(default)]
264 start: i64,
265 },
266
267 Pattern {
269 pattern: String,
271 },
272
273 IntRange {
275 min: i64,
277 max: i64,
279 },
280
281 FloatRange {
283 min: f64,
285 max: f64,
287 },
288
289 DecimalRange {
291 min: f64,
293 max: f64,
295 },
296
297 TimestampRange {
299 start: String,
301 end: String,
303 },
304
305 TimestampNow,
313
314 WeightedBool {
316 true_weight: f64,
318 },
319
320 OneOf {
322 values: Vec<serde_yaml::Value>,
324 },
325
326 SampleArray {
328 pool: Vec<String>,
330 #[serde(default)]
332 min_length: usize,
333 max_length: usize,
335 },
336
337 Static {
339 value: serde_yaml::Value,
341 },
342
343 Null,
345
346 DurationRange {
348 min_secs: u64,
350 max_secs: u64,
352 },
353}
354
355#[derive(Debug, Clone, Serialize, Deserialize)]
357pub struct GeneratorFieldDefinition {
358 pub name: String,
360
361 #[serde(rename = "type")]
363 pub field_type: Type,
364
365 pub generator: GeneratorConfig,
367
368 #[serde(default)]
370 pub nullable: bool,
371}
372
373impl GeneratorFieldDefinition {
374 pub fn to_column_definition(&self) -> ColumnDefinition {
376 ColumnDefinition {
377 name: self.name.clone(),
378 column_type: self.field_type.clone(),
379 nullable: self.nullable,
380 }
381 }
382}
383
384#[derive(Debug, Clone, Serialize, Deserialize)]
386pub struct GeneratorIDDefinition {
387 #[serde(rename = "type")]
389 pub id_type: Type,
390
391 pub generator: GeneratorConfig,
393}
394
395impl GeneratorIDDefinition {
396 pub fn to_column_definition(&self) -> ColumnDefinition {
399 ColumnDefinition {
400 name: "id".to_string(),
401 column_type: self.id_type.clone(),
402 nullable: false, }
404 }
405}
406
407#[derive(Debug, Clone, Serialize, Deserialize)]
409pub struct GeneratorTableDefinition {
410 pub name: String,
412
413 pub id: GeneratorIDDefinition,
415
416 pub fields: Vec<GeneratorFieldDefinition>,
418}
419
420impl GeneratorTableDefinition {
421 pub fn get_field(&self, name: &str) -> Option<&GeneratorFieldDefinition> {
423 self.fields.iter().find(|f| f.name == name)
424 }
425
426 pub fn get_field_type(&self, name: &str) -> Option<&Type> {
428 self.get_field(name).map(|f| &f.field_type)
429 }
430
431 pub fn field_names(&self) -> Vec<&str> {
433 self.fields.iter().map(|f| f.name.as_str()).collect()
434 }
435
436 pub fn to_table_definition(&self) -> TableDefinition {
438 TableDefinition {
439 name: self.name.clone(),
440 primary_key: self.id.to_column_definition(),
441 columns: self
442 .fields
443 .iter()
444 .map(|f| f.to_column_definition())
445 .collect(),
446 foreign_keys: Vec::new(),
447 composite_primary_key: None,
448 }
449 }
450}
451
452fn default_version() -> u32 {
453 1
454}
455
456#[derive(Debug, Clone, Serialize, Deserialize)]
462pub struct GeneratorSchema {
463 #[serde(default = "default_version")]
465 pub version: u32,
466
467 pub tables: Vec<GeneratorTableDefinition>,
469
470 #[serde(skip)]
472 table_map: HashMap<String, usize>,
473}
474
475impl GeneratorSchema {
476 pub fn from_file<P: AsRef<Path>>(path: P) -> Result<Self, SchemaError> {
478 let content = fs::read_to_string(path)?;
479 Self::from_yaml(&content)
480 }
481
482 pub fn from_yaml(yaml: &str) -> Result<Self, SchemaError> {
484 let mut schema: GeneratorSchema = serde_yaml::from_str(yaml)?;
485 schema.build_table_map();
486 Ok(schema)
487 }
488
489 fn build_table_map(&mut self) {
491 self.table_map = self
492 .tables
493 .iter()
494 .enumerate()
495 .map(|(idx, table)| (table.name.clone(), idx))
496 .collect();
497 }
498
499 pub fn get_table(&self, name: &str) -> Option<&GeneratorTableDefinition> {
501 self.table_map
502 .get(name)
503 .and_then(|&idx| self.tables.get(idx))
504 }
505
506 pub fn get_field_type(&self, table: &str, field: &str) -> Result<&Type, SchemaError> {
508 let table_schema = self
509 .get_table(table)
510 .ok_or_else(|| SchemaError::TableNotFound(table.to_string()))?;
511
512 table_schema
513 .get_field(field)
514 .map(|f| &f.field_type)
515 .ok_or_else(|| SchemaError::FieldNotFound {
516 table: table.to_string(),
517 field: field.to_string(),
518 })
519 }
520
521 pub fn table_names(&self) -> Vec<&str> {
523 self.tables.iter().map(|t| t.name.as_str()).collect()
524 }
525
526 pub fn to_database_schema(&self) -> DatabaseSchema {
528 DatabaseSchema::new(
529 self.tables
530 .iter()
531 .map(|t| t.to_table_definition())
532 .collect(),
533 )
534 }
535}
536
537pub type Schema = GeneratorSchema;
544
545pub type TableDefinitionWithGenerators = GeneratorTableDefinition;
548
549pub type FieldDefinition = GeneratorFieldDefinition;
552
553pub type IDDefinition = GeneratorIDDefinition;
556
557#[cfg(test)]
562mod tests {
563 use super::*;
564
565 #[test]
570 fn test_column_definition_serde() {
571 let col = ColumnDefinition::new("email", Type::VarChar { length: 255 });
572
573 let yaml = serde_yaml::to_string(&col).unwrap();
574 let parsed: ColumnDefinition = serde_yaml::from_str(&yaml).unwrap();
575
576 assert_eq!(col.name, parsed.name);
577 assert_eq!(col.column_type, parsed.column_type);
578 assert_eq!(col.nullable, parsed.nullable);
579 }
580
581 #[test]
582 fn test_table_definition_get_column() {
583 let table = TableDefinition::new(
584 "users",
585 ColumnDefinition::new("id", Type::Int64),
586 vec![
587 ColumnDefinition::new("email", Type::VarChar { length: 255 }),
588 ColumnDefinition::nullable("bio", Type::Text),
589 ],
590 );
591
592 let id_col = table.get_column("id").expect("id should exist");
594 assert_eq!(id_col.column_type, Type::Int64);
595
596 let email_col = table.get_column("email").expect("email should exist");
598 assert_eq!(email_col.column_type, Type::VarChar { length: 255 });
599
600 assert!(table.get_column("nonexistent").is_none());
602
603 let names = table.column_names();
605 assert!(names.contains(&"id"));
606 assert!(names.contains(&"email"));
607 assert!(names.contains(&"bio"));
608 }
609
610 #[test]
611 fn test_database_schema_get_table() {
612 let schema = DatabaseSchema::new(vec![
613 TableDefinition::new(
614 "users",
615 ColumnDefinition::new("id", Type::Uuid),
616 vec![ColumnDefinition::new("name", Type::Text)],
617 ),
618 TableDefinition::new(
619 "posts",
620 ColumnDefinition::new("id", Type::Int64),
621 vec![ColumnDefinition::new("title", Type::Text)],
622 ),
623 ]);
624
625 let users = schema.get_table("users").expect("users should exist");
627 assert_eq!(users.primary_key.column_type, Type::Uuid);
628
629 let posts = schema.get_table("posts").expect("posts should exist");
630 assert_eq!(posts.primary_key.column_type, Type::Int64);
631
632 assert!(schema.get_table("nonexistent").is_none());
634
635 let col_type = schema.get_column_type("users", "name").unwrap();
637 assert_eq!(col_type, &Type::Text);
638 }
639
640 const SAMPLE_SCHEMA: &str = r#"
645version: 1
646
647tables:
648 - name: users
649 id:
650 type: uuid
651 generator:
652 type: uuid_v4
653
654 fields:
655 - name: email
656 type:
657 type: var_char
658 length: 255
659 generator:
660 type: pattern
661 pattern: "user_{index}@example.com"
662
663 - name: age
664 type: int
665 generator:
666 type: int_range
667 min: 18
668 max: 80
669
670 - name: is_active
671 type:
672 type: tiny_int
673 width: 1
674 generator:
675 type: weighted_bool
676 true_weight: 0.8
677"#;
678
679 #[test]
680 fn test_parse_generator_schema() {
681 let schema = GeneratorSchema::from_yaml(SAMPLE_SCHEMA).unwrap();
682
683 assert_eq!(schema.version, 1);
684 assert_eq!(schema.tables.len(), 1);
685
686 let users = schema.get_table("users").unwrap();
687 assert_eq!(users.name, "users");
688 assert_eq!(users.id.id_type, Type::Uuid);
689 assert_eq!(users.fields.len(), 3);
690 }
691
692 #[test]
693 fn test_get_field_type() {
694 let schema = GeneratorSchema::from_yaml(SAMPLE_SCHEMA).unwrap();
695
696 let email_type = schema.get_field_type("users", "email").unwrap();
697 assert_eq!(email_type, &Type::VarChar { length: 255 });
698
699 let age_type = schema.get_field_type("users", "age").unwrap();
700 assert_eq!(age_type, &Type::Int32);
701 }
702
703 #[test]
704 fn test_table_not_found() {
705 let schema = GeneratorSchema::from_yaml(SAMPLE_SCHEMA).unwrap();
706
707 let result = schema.get_field_type("nonexistent", "field");
708 assert!(matches!(result, Err(SchemaError::TableNotFound(_))));
709 }
710
711 #[test]
712 fn test_field_not_found() {
713 let schema = GeneratorSchema::from_yaml(SAMPLE_SCHEMA).unwrap();
714
715 let result = schema.get_field_type("users", "nonexistent");
716 assert!(matches!(result, Err(SchemaError::FieldNotFound { .. })));
717 }
718
719 #[test]
720 fn test_generator_configs() {
721 let yaml = r#"
722version: 1
723tables:
724 - name: test
725 id:
726 type: int
727 generator:
728 type: sequential
729 start: 1
730 fields:
731 - name: score
732 type: double
733 generator:
734 type: float_range
735 min: 0.0
736 max: 100.0
737 - name: status
738 type: text
739 generator:
740 type: one_of
741 values: ["active", "inactive", "pending"]
742 - name: metadata
743 type: json
744 generator:
745 type: static
746 value: { "version": 1 }
747"#;
748
749 let schema = GeneratorSchema::from_yaml(yaml).unwrap();
750 let table = schema.get_table("test").unwrap();
751
752 assert!(matches!(
754 table.id.generator,
755 GeneratorConfig::Sequential { start: 1 }
756 ));
757
758 let score_field = table.get_field("score").unwrap();
760 assert!(matches!(
761 score_field.generator,
762 GeneratorConfig::FloatRange {
763 min: 0.0,
764 max: 100.0
765 }
766 ));
767 }
768
769 #[test]
770 fn test_generator_field_definition_flatten_serde() {
771 let field = GeneratorFieldDefinition {
773 name: "email".to_string(),
774 field_type: Type::VarChar { length: 255 },
775 generator: GeneratorConfig::Pattern {
776 pattern: "user_{index}@test.com".to_string(),
777 },
778 nullable: false,
779 };
780
781 let yaml = serde_yaml::to_string(&field).unwrap();
782 let parsed: GeneratorFieldDefinition = serde_yaml::from_str(&yaml).unwrap();
783
784 assert_eq!(field.name, parsed.name);
785 assert_eq!(field.field_type, parsed.field_type);
786 assert_eq!(field.nullable, parsed.nullable);
787
788 let col = field.to_column_definition();
790 assert_eq!(col.name, "email");
791 assert_eq!(col.column_type, Type::VarChar { length: 255 });
792 assert!(!col.nullable);
793 }
794
795 #[test]
796 fn test_generator_schema_yaml_compatibility() {
797 let schema = GeneratorSchema::from_yaml(SAMPLE_SCHEMA).unwrap();
799
800 let db_schema = schema.to_database_schema();
802
803 assert_eq!(db_schema.tables.len(), 1);
804 let users = db_schema.get_table("users").expect("users should exist");
805 assert_eq!(users.primary_key.column_type, Type::Uuid);
806 assert_eq!(users.columns.len(), 3);
807
808 let email_type = db_schema.get_column_type("users", "email").unwrap();
810 assert_eq!(email_type, &Type::VarChar { length: 255 });
811 }
812
813 #[test]
814 fn test_type_alias_compatibility() {
815 let schema: Schema = GeneratorSchema::from_yaml(SAMPLE_SCHEMA).unwrap();
817 let _table: &TableDefinitionWithGenerators = schema.get_table("users").unwrap();
818 let _field: &FieldDefinition = schema
819 .get_table("users")
820 .unwrap()
821 .get_field("email")
822 .unwrap();
823 let _id: &IDDefinition = &schema.get_table("users").unwrap().id;
824 }
825}