1use super::types::ColumnType;
18use std::collections::HashMap;
19
20#[derive(Debug, Clone, Default)]
22pub struct Schema {
23 pub tables: HashMap<String, Table>,
24 pub indexes: Vec<Index>,
25 pub migrations: Vec<MigrationHint>,
26}
27
28#[derive(Debug, Clone)]
29pub struct Table {
30 pub name: String,
31 pub columns: Vec<Column>,
32}
33
34#[derive(Debug, Clone)]
36pub struct Column {
37 pub name: String,
38 pub data_type: ColumnType,
39 pub nullable: bool,
40 pub primary_key: bool,
41 pub unique: bool,
42 pub default: Option<String>,
43 pub foreign_key: Option<ForeignKey>,
44 pub check: Option<CheckConstraint>,
46 pub generated: Option<Generated>,
48}
49
50#[derive(Debug, Clone)]
52pub struct ForeignKey {
53 pub table: String,
54 pub column: String,
55 pub on_delete: FkAction,
56 pub on_update: FkAction,
57 pub deferrable: Deferrable,
59}
60
61#[derive(Debug, Clone, Default, PartialEq)]
63pub enum FkAction {
64 #[default]
65 NoAction,
66 Cascade,
67 SetNull,
68 SetDefault,
69 Restrict,
70}
71
72#[derive(Debug, Clone)]
73pub struct Index {
74 pub name: String,
75 pub table: String,
76 pub columns: Vec<String>,
77 pub unique: bool,
78 pub method: IndexMethod,
80 pub where_clause: Option<CheckExpr>,
82 pub include: Vec<String>,
84 pub concurrently: bool,
86}
87
88#[derive(Debug, Clone)]
89pub enum MigrationHint {
90 Rename { from: String, to: String },
92 Transform { expression: String, target: String },
94 Drop { target: String, confirmed: bool },
96}
97
98#[derive(Debug, Clone)]
104pub enum CheckExpr {
105 GreaterThan { column: String, value: i64 },
107 GreaterOrEqual { column: String, value: i64 },
109 LessThan { column: String, value: i64 },
111 LessOrEqual { column: String, value: i64 },
113 Between { column: String, low: i64, high: i64 },
114 In { column: String, values: Vec<String> },
115 Regex { column: String, pattern: String },
117 MaxLength { column: String, max: usize },
119 MinLength { column: String, min: usize },
121 NotNull { column: String },
122 And(Box<CheckExpr>, Box<CheckExpr>),
123 Or(Box<CheckExpr>, Box<CheckExpr>),
124 Not(Box<CheckExpr>),
125}
126
127#[derive(Debug, Clone)]
129pub struct CheckConstraint {
130 pub expr: CheckExpr,
131 pub name: Option<String>,
132}
133
134#[derive(Debug, Clone, Default, PartialEq)]
140pub enum Deferrable {
141 #[default]
142 NotDeferrable,
143 Deferrable,
144 InitiallyDeferred,
145 InitiallyImmediate,
146}
147
148#[derive(Debug, Clone)]
154pub enum Generated {
155 AlwaysStored(String),
157 AlwaysIdentity,
159 ByDefaultIdentity,
161}
162
163#[derive(Debug, Clone, Default, PartialEq)]
169pub enum IndexMethod {
170 #[default]
171 BTree,
172 Hash,
173 Gin,
174 Gist,
175 Brin,
176 SpGist,
177}
178
179impl Schema {
180 pub fn new() -> Self {
181 Self::default()
182 }
183
184 pub fn add_table(&mut self, table: Table) {
185 self.tables.insert(table.name.clone(), table);
186 }
187
188 pub fn add_index(&mut self, index: Index) {
189 self.indexes.push(index);
190 }
191
192 pub fn add_hint(&mut self, hint: MigrationHint) {
193 self.migrations.push(hint);
194 }
195
196 pub fn validate(&self) -> Result<(), Vec<String>> {
199 let mut errors = Vec::new();
200
201 for table in self.tables.values() {
202 for col in &table.columns {
203 if let Some(ref fk) = col.foreign_key {
204 if !self.tables.contains_key(&fk.table) {
205 errors.push(format!(
206 "FK error: {}.{} references non-existent table '{}'",
207 table.name, col.name, fk.table
208 ));
209 } else {
210 let ref_table = &self.tables[&fk.table];
211 if !ref_table.columns.iter().any(|c| c.name == fk.column) {
212 errors.push(format!(
213 "FK error: {}.{} references non-existent column '{}.{}'",
214 table.name, col.name, fk.table, fk.column
215 ));
216 }
217 }
218 }
219 }
220 }
221
222 if errors.is_empty() {
223 Ok(())
224 } else {
225 Err(errors)
226 }
227 }
228}
229
230impl Table {
231 pub fn new(name: impl Into<String>) -> Self {
232 Self {
233 name: name.into(),
234 columns: Vec::new(),
235 }
236 }
237
238 pub fn column(mut self, col: Column) -> Self {
239 self.columns.push(col);
240 self
241 }
242}
243
244impl Column {
245 pub fn new(name: impl Into<String>, data_type: ColumnType) -> Self {
247 Self {
248 name: name.into(),
249 data_type,
250 nullable: true,
251 primary_key: false,
252 unique: false,
253 default: None,
254 foreign_key: None,
255 check: None,
256 generated: None,
257 }
258 }
259
260 pub fn not_null(mut self) -> Self {
261 self.nullable = false;
262 self
263 }
264
265 pub fn primary_key(mut self) -> Self {
270 if !self.data_type.can_be_primary_key() {
271 panic!(
272 "Column '{}' of type {} cannot be a primary key. \
273 Valid PK types: UUID, SERIAL, BIGSERIAL, INT, BIGINT",
274 self.name,
275 self.data_type.name()
276 );
277 }
278 self.primary_key = true;
279 self.nullable = false;
280 self
281 }
282
283 pub fn unique(mut self) -> Self {
287 if !self.data_type.supports_indexing() {
288 panic!(
289 "Column '{}' of type {} cannot have UNIQUE constraint. \
290 JSONB and BYTEA types do not support standard indexing.",
291 self.name,
292 self.data_type.name()
293 );
294 }
295 self.unique = true;
296 self
297 }
298
299 pub fn default(mut self, val: impl Into<String>) -> Self {
300 self.default = Some(val.into());
301 self
302 }
303
304 pub fn references(mut self, table: &str, column: &str) -> Self {
313 self.foreign_key = Some(ForeignKey {
314 table: table.to_string(),
315 column: column.to_string(),
316 on_delete: FkAction::default(),
317 on_update: FkAction::default(),
318 deferrable: Deferrable::default(),
319 });
320 self
321 }
322
323 pub fn on_delete(mut self, action: FkAction) -> Self {
325 if let Some(ref mut fk) = self.foreign_key {
326 fk.on_delete = action;
327 }
328 self
329 }
330
331 pub fn on_update(mut self, action: FkAction) -> Self {
333 if let Some(ref mut fk) = self.foreign_key {
334 fk.on_update = action;
335 }
336 self
337 }
338
339 pub fn check(mut self, expr: CheckExpr) -> Self {
343 self.check = Some(CheckConstraint { expr, name: None });
344 self
345 }
346
347 pub fn check_named(mut self, name: impl Into<String>, expr: CheckExpr) -> Self {
349 self.check = Some(CheckConstraint {
350 expr,
351 name: Some(name.into()),
352 });
353 self
354 }
355
356 pub fn deferrable(mut self) -> Self {
360 if let Some(ref mut fk) = self.foreign_key {
361 fk.deferrable = Deferrable::Deferrable;
362 }
363 self
364 }
365
366 pub fn initially_deferred(mut self) -> Self {
368 if let Some(ref mut fk) = self.foreign_key {
369 fk.deferrable = Deferrable::InitiallyDeferred;
370 }
371 self
372 }
373
374 pub fn initially_immediate(mut self) -> Self {
376 if let Some(ref mut fk) = self.foreign_key {
377 fk.deferrable = Deferrable::InitiallyImmediate;
378 }
379 self
380 }
381
382 pub fn generated_stored(mut self, expr: impl Into<String>) -> Self {
386 self.generated = Some(Generated::AlwaysStored(expr.into()));
387 self
388 }
389
390 pub fn generated_identity(mut self) -> Self {
392 self.generated = Some(Generated::AlwaysIdentity);
393 self
394 }
395
396 pub fn generated_by_default(mut self) -> Self {
398 self.generated = Some(Generated::ByDefaultIdentity);
399 self
400 }
401}
402
403impl Index {
404 pub fn new(name: impl Into<String>, table: impl Into<String>, columns: Vec<String>) -> Self {
405 Self {
406 name: name.into(),
407 table: table.into(),
408 columns,
409 unique: false,
410 method: IndexMethod::default(),
411 where_clause: None,
412 include: Vec::new(),
413 concurrently: false,
414 }
415 }
416
417 pub fn unique(mut self) -> Self {
418 self.unique = true;
419 self
420 }
421
422 pub fn using(mut self, method: IndexMethod) -> Self {
426 self.method = method;
427 self
428 }
429
430 pub fn partial(mut self, expr: CheckExpr) -> Self {
432 self.where_clause = Some(expr);
433 self
434 }
435
436 pub fn include(mut self, cols: Vec<String>) -> Self {
438 self.include = cols;
439 self
440 }
441
442 pub fn concurrently(mut self) -> Self {
444 self.concurrently = true;
445 self
446 }
447}
448
449pub fn to_qail_string(schema: &Schema) -> String {
451 let mut output = String::new();
452 output.push_str("# QAIL Schema\n\n");
453
454 for table in schema.tables.values() {
455 output.push_str(&format!("table {} {{\n", table.name));
456 for col in &table.columns {
457 let mut constraints: Vec<String> = Vec::new();
458 if col.primary_key {
459 constraints.push("primary_key".to_string());
460 }
461 if !col.nullable && !col.primary_key {
462 constraints.push("not_null".to_string());
463 }
464 if col.unique {
465 constraints.push("unique".to_string());
466 }
467 if let Some(def) = &col.default {
468 constraints.push(format!("default {}", def));
469 }
470 if let Some(ref fk) = col.foreign_key {
471 constraints.push(format!("references {}({})", fk.table, fk.column));
472 }
473
474 let constraint_str = if constraints.is_empty() {
475 String::new()
476 } else {
477 format!(" {}", constraints.join(" "))
478 };
479
480 output.push_str(&format!(
481 " {} {}{}\n",
482 col.name,
483 col.data_type.to_pg_type(),
484 constraint_str
485 ));
486 }
487 output.push_str("}\n\n");
488 }
489
490 for idx in &schema.indexes {
491 let unique = if idx.unique { "unique " } else { "" };
492 output.push_str(&format!(
493 "{}index {} on {} ({})\n",
494 unique,
495 idx.name,
496 idx.table,
497 idx.columns.join(", ")
498 ));
499 }
500
501 for hint in &schema.migrations {
502 match hint {
503 MigrationHint::Rename { from, to } => {
504 output.push_str(&format!("rename {} -> {}\n", from, to));
505 }
506 MigrationHint::Transform { expression, target } => {
507 output.push_str(&format!("transform {} -> {}\n", expression, target));
508 }
509 MigrationHint::Drop { target, confirmed } => {
510 let confirm = if *confirmed { " confirm" } else { "" };
511 output.push_str(&format!("drop {}{}\n", target, confirm));
512 }
513 }
514 }
515
516 output
517}
518
519#[cfg(test)]
520mod tests {
521 use super::*;
522
523 #[test]
524 fn test_schema_builder() {
525 let mut schema = Schema::new();
526
527 let users = Table::new("users")
528 .column(Column::new("id", ColumnType::Serial).primary_key())
529 .column(Column::new("name", ColumnType::Text).not_null())
530 .column(Column::new("email", ColumnType::Text).unique());
531
532 schema.add_table(users);
533 schema.add_index(Index::new("idx_users_email", "users", vec!["email".into()]).unique());
534
535 let output = to_qail_string(&schema);
536 assert!(output.contains("table users"));
537 assert!(output.contains("id SERIAL primary_key"));
538 assert!(output.contains("unique index idx_users_email"));
539 }
540
541 #[test]
542 fn test_migration_hints() {
543 let mut schema = Schema::new();
544 schema.add_hint(MigrationHint::Rename {
545 from: "users.username".into(),
546 to: "users.name".into(),
547 });
548
549 let output = to_qail_string(&schema);
550 assert!(output.contains("rename users.username -> users.name"));
551 }
552
553 #[test]
554 #[should_panic(expected = "cannot be a primary key")]
555 fn test_invalid_primary_key_type() {
556 Column::new("data", ColumnType::Text).primary_key();
558 }
559
560 #[test]
561 #[should_panic(expected = "cannot have UNIQUE")]
562 fn test_invalid_unique_type() {
563 Column::new("data", ColumnType::Jsonb).unique();
565 }
566
567 #[test]
568 fn test_foreign_key_valid() {
569 let mut schema = Schema::new();
570
571 schema.add_table(
572 Table::new("users").column(Column::new("id", ColumnType::Uuid).primary_key()),
573 );
574
575 schema.add_table(
576 Table::new("posts")
577 .column(Column::new("id", ColumnType::Uuid).primary_key())
578 .column(
579 Column::new("user_id", ColumnType::Uuid)
580 .references("users", "id")
581 .on_delete(FkAction::Cascade),
582 ),
583 );
584
585 assert!(schema.validate().is_ok());
587 }
588
589 #[test]
590 fn test_foreign_key_invalid_table() {
591 let mut schema = Schema::new();
592
593 schema.add_table(
594 Table::new("posts")
595 .column(Column::new("id", ColumnType::Uuid).primary_key())
596 .column(Column::new("user_id", ColumnType::Uuid).references("nonexistent", "id")),
597 );
598
599 let result = schema.validate();
601 assert!(result.is_err());
602 assert!(result.unwrap_err()[0].contains("non-existent table"));
603 }
604
605 #[test]
606 fn test_foreign_key_invalid_column() {
607 let mut schema = Schema::new();
608
609 schema.add_table(
610 Table::new("users").column(Column::new("id", ColumnType::Uuid).primary_key()),
611 );
612
613 schema.add_table(
614 Table::new("posts")
615 .column(Column::new("id", ColumnType::Uuid).primary_key())
616 .column(
617 Column::new("user_id", ColumnType::Uuid).references("users", "wrong_column"),
618 ),
619 );
620
621 let result = schema.validate();
623 assert!(result.is_err());
624 assert!(result.unwrap_err()[0].contains("non-existent column"));
625 }
626}