Skip to main content

drizzle_cli/commands/
new.rs

1//! Interactive schema builder (`drizzle new`)
2//!
3//! Walks the user through an interactive wizard to define tables, columns,
4//! indexes, and foreign keys, then generates Rust schema code using the
5//! existing codegen pipeline (the same one `drizzle introspect` uses).
6//!
7//! Supports JSON import/export for CI-friendly, reproducible schema generation:
8//! - `drizzle new --json` reads a JSON schema definition from stdin
9//! - `drizzle new --json --from file.json` reads from a file
10//! - `drizzle new --export-json out.json` exports the schema as JSON
11//! - `drizzle new --schema-help` prints the expected JSON shape
12
13use std::borrow::Cow;
14use std::collections::HashSet;
15use std::path::PathBuf;
16
17use inquire::validator::Validation;
18use inquire::{Confirm, MultiSelect, Select, Text};
19use schemars::JsonSchema;
20use serde::{Deserialize, Serialize};
21
22use crate::config::{Config, Dialect};
23use crate::error::CliError;
24use crate::output;
25
26// ── Public API ──────────────────────────────────────────────────────────────
27
28#[derive(clap::Args, Debug)]
29pub struct NewOptions {
30    /// Override dialect (sqlite, postgresql, mysql)
31    #[arg(long)]
32    pub dialect: Option<Dialect>,
33
34    /// Override schema output path
35    #[arg(long)]
36    pub schema: Option<String>,
37
38    /// Read schema definition from JSON (stdin by default)
39    #[arg(long)]
40    pub json: bool,
41
42    /// Read JSON from a file instead of stdin (requires --json)
43    #[arg(long, requires = "json", value_name = "PATH")]
44    pub from: Option<PathBuf>,
45
46    /// Export the schema definition as JSON after building
47    #[arg(long = "export-json", value_name = "PATH")]
48    pub export_json: Option<PathBuf>,
49
50    /// Print the expected JSON schema shape and exit
51    #[arg(long = "schema-help")]
52    pub schema_help: bool,
53}
54
55/// Run the `new` command to scaffold a schema definition.
56///
57/// # Errors
58///
59/// Returns [`CliError`] if loading or validating the JSON schema definition
60/// fails, if interactive prompts are cancelled, or if writing the generated
61/// Rust schema files fails.
62pub fn run(config: Option<&Config>, options: &NewOptions) -> Result<(), CliError> {
63    // --schema-help: print annotated example and exit
64    if options.schema_help {
65        print_json_schema();
66        return Ok(());
67    }
68
69    // Build the schema definition from either JSON input or interactive prompts
70    let def = if options.json {
71        load_json(options.from.as_deref())?
72    } else {
73        collect_interactively(config, options)?
74    };
75
76    // Validate the schema definition
77    validate_schema(&def)?;
78
79    // Export JSON if requested
80    if let Some(ref export_path) = options.export_json {
81        export_to_json(&def, export_path)?;
82    }
83
84    // Determine output path (JSON definition's output_path, or CLI override)
85    let output_path = if let Some(ref s) = options.schema {
86        s.clone()
87    } else {
88        def.output_path.clone()
89    };
90
91    // Generate code
92    let code = match def.dialect {
93        Dialect::Sqlite | Dialect::Turso => generate_sqlite(
94            &def.tables,
95            &def.indexes,
96            &def.foreign_keys,
97            &def.schema_name,
98            def.casing,
99        ),
100        Dialect::Postgresql => generate_postgres(
101            &def.tables,
102            &def.indexes,
103            &def.foreign_keys,
104            &def.enums,
105            &def.schema_name,
106            def.casing,
107        ),
108        Dialect::Mysql => generate_mysql(
109            &def.tables,
110            &def.indexes,
111            &def.foreign_keys,
112            &def.enums,
113            &def.schema_name,
114            def.casing,
115        )?,
116    };
117
118    // Write output
119    let path = PathBuf::from(&output_path);
120    if let Some(parent) = path.parent() {
121        std::fs::create_dir_all(parent)
122            .map_err(|e| CliError::IoError(format!("Failed to create directory: {e}")))?;
123    }
124    std::fs::write(&path, &code)
125        .map_err(|e| CliError::IoError(format!("Failed to write schema file: {e}")))?;
126
127    // Print summary
128    println!();
129    println!("{}", output::success("Schema generated successfully!"));
130    println!();
131    println!(
132        "  Tables: {}",
133        def.tables
134            .iter()
135            .map(|t| t.name.as_str())
136            .collect::<Vec<_>>()
137            .join(", ")
138    );
139    if !def.indexes.is_empty() {
140        println!(
141            "  Indexes: {}",
142            def.indexes
143                .iter()
144                .map(|i| i.name.as_str())
145                .collect::<Vec<_>>()
146                .join(", ")
147        );
148    }
149    if !def.foreign_keys.is_empty() {
150        println!("  Foreign keys: {}", def.foreign_keys.len());
151    }
152    println!("  Output: {output_path}");
153    if let Some(ref export_path) = options.export_json {
154        println!("  JSON export: {}", export_path.display());
155    }
156    println!();
157    println!("Next steps:");
158    println!(
159        "  Run {} to generate your first migration",
160        output::heading("drizzle generate")
161    );
162
163    Ok(())
164}
165
166// ── Schema definition (top-level JSON document) ─────────────────────────────
167
168#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
169pub struct SchemaDefinition {
170    pub dialect: Dialect,
171    #[serde(default = "default_casing")]
172    pub casing: FieldCasing,
173    #[serde(default = "default_schema_name")]
174    pub schema_name: String,
175    #[serde(default = "default_output_path")]
176    pub output_path: String,
177    #[serde(default)]
178    pub enums: Vec<EnumDef>,
179    pub tables: Vec<TableDef>,
180    #[serde(default)]
181    pub indexes: Vec<IndexDef>,
182    #[serde(default)]
183    pub foreign_keys: Vec<ForeignKeyDef>,
184}
185
186const fn default_casing() -> FieldCasing {
187    FieldCasing::Snake
188}
189
190fn default_schema_name() -> String {
191    "AppSchema".to_string()
192}
193
194fn default_output_path() -> String {
195    "src/schema.rs".to_string()
196}
197
198fn default_fk_action() -> String {
199    "No Action".to_string()
200}
201
202// ── Intermediate structs ────────────────────────────────────────────────────
203
204#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
205pub struct EnumDef {
206    pub name: String,
207    pub variants: Vec<String>,
208}
209
210#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
211pub struct TableDef {
212    pub name: String,
213    pub columns: Vec<ColumnDef>,
214    /// `SQLite` only
215    #[serde(default)]
216    pub strict: bool,
217    /// `SQLite` only
218    #[serde(default)]
219    pub without_rowid: bool,
220    /// `PostgreSQL` only
221    #[serde(default = "default_pg_schema")]
222    pub pg_schema: String,
223}
224
225fn default_pg_schema() -> String {
226    "public".to_string()
227}
228
229/// Auto-generation strategy for a column value.
230///
231/// `autoincrement` is supported by SQLite (`AUTOINCREMENT`) and MySQL
232/// (`AUTO_INCREMENT`); `identity` is PostgreSQL-specific. They are
233/// mutually exclusive dialect variants, so they live in a single optional enum
234/// rather than two parallel booleans.
235#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
236#[serde(rename_all = "lowercase")]
237pub enum AutoGenKind {
238    /// SQLite `INTEGER PRIMARY KEY AUTOINCREMENT` or MySQL `AUTO_INCREMENT`.
239    Autoincrement,
240    /// `PostgreSQL` `GENERATED ALWAYS AS IDENTITY`.
241    Identity,
242}
243
244#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
245pub struct ColumnDef {
246    pub name: String,
247    /// The SQL type string the codegen expects
248    pub sql_type: String,
249    #[serde(default)]
250    pub not_null: bool,
251    #[serde(default)]
252    pub primary_key: bool,
253    #[serde(default)]
254    pub unique: bool,
255    #[serde(default, skip_serializing_if = "Option::is_none")]
256    pub default: Option<String>,
257    /// Auto-generation strategy (SQLite/MySQL autoincrement or PostgreSQL identity).
258    #[serde(default, skip_serializing_if = "Option::is_none")]
259    pub auto_gen: Option<AutoGenKind>,
260    /// PostgreSQL enum type or MySQL inline enum definition name.
261    #[serde(default, skip_serializing_if = "Option::is_none")]
262    pub enum_name: Option<String>,
263}
264
265impl ColumnDef {
266    #[must_use]
267    pub const fn is_autoincrement(&self) -> bool {
268        matches!(self.auto_gen, Some(AutoGenKind::Autoincrement))
269    }
270
271    #[must_use]
272    pub const fn is_identity(&self) -> bool {
273        matches!(self.auto_gen, Some(AutoGenKind::Identity))
274    }
275}
276
277#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
278pub struct IndexDef {
279    pub name: String,
280    pub table: String,
281    pub columns: Vec<String>,
282    #[serde(default)]
283    pub unique: bool,
284    /// PG schema
285    #[serde(default)]
286    pub pg_schema: String,
287}
288
289#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
290pub struct ForeignKeyDef {
291    pub name: String,
292    pub table: String,
293    pub columns: Vec<String>,
294    pub table_to: String,
295    pub columns_to: Vec<String>,
296    #[serde(default = "default_fk_action")]
297    pub on_delete: String,
298    #[serde(default = "default_fk_action")]
299    pub on_update: String,
300    /// PG schema
301    #[serde(default)]
302    pub pg_schema: String,
303    #[serde(default)]
304    pub pg_schema_to: String,
305}
306
307#[derive(Debug, Clone, Copy, Serialize, Deserialize, JsonSchema)]
308pub enum FieldCasing {
309    #[serde(rename = "snake_case")]
310    Snake,
311    #[serde(rename = "camelCase")]
312    Camel,
313}
314
315// ── JSON import/export ──────────────────────────────────────────────────────
316
317fn load_json(from: Option<&std::path::Path>) -> Result<SchemaDefinition, CliError> {
318    let content = if let Some(path) = from {
319        std::fs::read_to_string(path)
320            .map_err(|e| CliError::IoError(format!("Failed to read {}: {e}", path.display())))?
321    } else {
322        use std::io::Read;
323        let mut buf = String::new();
324        std::io::stdin()
325            .read_to_string(&mut buf)
326            .map_err(|e| CliError::IoError(format!("Failed to read stdin: {e}")))?;
327        buf
328    };
329    serde_json::from_str(&content)
330        .map_err(|e| CliError::Other(format!("Invalid JSON schema definition: {e}")))
331}
332
333fn export_to_json(def: &SchemaDefinition, path: &std::path::Path) -> Result<(), CliError> {
334    let json = serde_json::to_string_pretty(def)
335        .map_err(|e| CliError::Other(format!("Failed to serialize schema: {e}")))?;
336    if let Some(parent) = path.parent() {
337        std::fs::create_dir_all(parent)
338            .map_err(|e| CliError::IoError(format!("Failed to create directory: {e}")))?;
339    }
340    std::fs::write(path, json)
341        .map_err(|e| CliError::IoError(format!("Failed to write JSON: {e}")))?;
342    Ok(())
343}
344
345fn print_json_schema() {
346    let schema = schemars::schema_for!(SchemaDefinition);
347    let json = serde_json::to_string_pretty(&schema).expect("schema serialization cannot fail");
348    println!("{json}");
349    println!();
350    println!(
351        "Valid on_delete/on_update actions: \"No Action\", \"Cascade\", \"Set Null\", \"Set Default\", \"Restrict\""
352    );
353    println!();
354    println!("Tip: Run `drizzle new --export-json schema.json` to export an interactive");
355    println!(
356        "session as valid JSON, then edit and replay with `drizzle new --json --from schema.json`."
357    );
358}
359
360// ── Validation ──────────────────────────────────────────────────────────────
361
362const VALID_FK_ACTIONS: &[&str] = &[
363    "No Action",
364    "Cascade",
365    "Set Null",
366    "Set Default",
367    "Restrict",
368];
369
370fn validate_schema(def: &SchemaDefinition) -> Result<(), CliError> {
371    // Must have at least one table
372    if def.tables.is_empty() {
373        return Err(CliError::Other(
374            "Schema must have at least one table".into(),
375        ));
376    }
377
378    // Check table names are valid and unique, and per-table column/dialect rules
379    let mut table_names = HashSet::new();
380    for table in &def.tables {
381        if !is_valid_identifier(&table.name) {
382            return Err(CliError::Other(format!(
383                "Invalid table name: '{}'",
384                table.name
385            )));
386        }
387        if !table_names.insert(&table.name) {
388            return Err(CliError::Other(format!(
389                "Duplicate table name: '{}'",
390                table.name
391            )));
392        }
393        validate_table(table, def.dialect)?;
394    }
395
396    let enum_names = validate_enums(def)?;
397
398    // Validate enum references in columns
399    for table in &def.tables {
400        for col in &table.columns {
401            if let Some(ref en) = col.enum_name
402                && !enum_names.contains(en.as_str())
403            {
404                return Err(CliError::Other(format!(
405                    "Column '{}.{}' references unknown enum '{}'",
406                    table.name, col.name, en
407                )));
408            }
409        }
410    }
411
412    validate_indexes(def)?;
413    validate_foreign_keys(def)?;
414
415    Ok(())
416}
417
418/// Validate a single table: column names, dialect-specific column rules, and
419/// that the table itself has at least one column.
420fn validate_table(table: &TableDef, dialect: Dialect) -> Result<(), CliError> {
421    if table.columns.is_empty() {
422        return Err(CliError::Other(format!(
423            "Table '{}' must have at least one column",
424            table.name
425        )));
426    }
427
428    let mut col_names = HashSet::new();
429    for col in &table.columns {
430        if !is_valid_identifier(&col.name) {
431            return Err(CliError::Other(format!(
432                "Invalid column name '{}' in table '{}'",
433                col.name, table.name
434            )));
435        }
436        if !col_names.insert(&col.name) {
437            return Err(CliError::Other(format!(
438                "Duplicate column name '{}' in table '{}'",
439                col.name, table.name
440            )));
441        }
442    }
443
444    match dialect {
445        Dialect::Sqlite | Dialect::Turso => {
446            for col in &table.columns {
447                if col.is_identity() {
448                    return Err(CliError::Other(format!(
449                        "Column '{}.{}': 'identity' is only supported for PostgreSQL",
450                        table.name, col.name
451                    )));
452                }
453                if col.enum_name.is_some() {
454                    return Err(CliError::Other(format!(
455                        "Column '{}.{}': 'enum_name' is only supported for PostgreSQL",
456                        table.name, col.name
457                    )));
458                }
459            }
460        }
461        Dialect::Postgresql => {
462            if table.strict {
463                return Err(CliError::Other(format!(
464                    "Table '{}': 'strict' is only supported for SQLite",
465                    table.name
466                )));
467            }
468            if table.without_rowid {
469                return Err(CliError::Other(format!(
470                    "Table '{}': 'without_rowid' is only supported for SQLite",
471                    table.name
472                )));
473            }
474            for col in &table.columns {
475                if col.is_autoincrement() {
476                    return Err(CliError::Other(format!(
477                        "Column '{}.{}': 'autoincrement' is only supported for SQLite (use 'identity' for PostgreSQL)",
478                        table.name, col.name
479                    )));
480                }
481            }
482        }
483        Dialect::Mysql => {
484            if table.strict || table.without_rowid {
485                return Err(CliError::Other(format!(
486                    "Table '{}': SQLite table options are not supported for MySQL",
487                    table.name
488                )));
489            }
490            for col in &table.columns {
491                if col.is_identity() {
492                    return Err(CliError::Other(format!(
493                        "Column '{}.{}': 'identity' is only supported for PostgreSQL (use 'autoincrement' for MySQL)",
494                        table.name, col.name
495                    )));
496                }
497            }
498        }
499    }
500
501    Ok(())
502}
503
504/// Validate enum definitions and return the set of declared enum names.
505fn validate_enums(def: &SchemaDefinition) -> Result<HashSet<&str>, CliError> {
506    if matches!(def.dialect, Dialect::Sqlite | Dialect::Turso) && !def.enums.is_empty() {
507        return Err(CliError::Other(
508            "Enums are only supported for PostgreSQL and MySQL".into(),
509        ));
510    }
511    let mut enum_names = HashSet::new();
512    for e in &def.enums {
513        if !is_valid_identifier(&e.name) {
514            return Err(CliError::Other(format!("Invalid enum name: '{}'", e.name)));
515        }
516        if !enum_names.insert(e.name.as_str()) {
517            return Err(CliError::Other(format!(
518                "Duplicate enum name: '{}'",
519                e.name
520            )));
521        }
522        if e.variants.is_empty() {
523            return Err(CliError::Other(format!(
524                "Enum '{}' must have at least one variant",
525                e.name
526            )));
527        }
528    }
529    Ok(enum_names)
530}
531
532/// Validate that each index references a real table and real columns.
533fn validate_indexes(def: &SchemaDefinition) -> Result<(), CliError> {
534    for idx in &def.indexes {
535        let table = def.tables.iter().find(|t| t.name == idx.table);
536        let Some(table) = table else {
537            return Err(CliError::Other(format!(
538                "Index '{}' references unknown table '{}'",
539                idx.name, idx.table
540            )));
541        };
542        for col_name in &idx.columns {
543            if !table.columns.iter().any(|c| &c.name == col_name) {
544                return Err(CliError::Other(format!(
545                    "Index '{}' references unknown column '{}.{}'",
546                    idx.name, idx.table, col_name
547                )));
548            }
549        }
550    }
551    Ok(())
552}
553
554/// Validate that each foreign key references real tables/columns and uses a
555/// recognized on-delete / on-update action.
556fn validate_foreign_keys(def: &SchemaDefinition) -> Result<(), CliError> {
557    for fk in &def.foreign_keys {
558        // Source table
559        let src = def.tables.iter().find(|t| t.name == fk.table);
560        let Some(src) = src else {
561            return Err(CliError::Other(format!(
562                "Foreign key '{}' references unknown source table '{}'",
563                fk.name, fk.table
564            )));
565        };
566        for col_name in &fk.columns {
567            if !src.columns.iter().any(|c| &c.name == col_name) {
568                return Err(CliError::Other(format!(
569                    "Foreign key '{}' references unknown source column '{}.{}'",
570                    fk.name, fk.table, col_name
571                )));
572            }
573        }
574
575        // Target table
576        let tgt = def.tables.iter().find(|t| t.name == fk.table_to);
577        let Some(tgt) = tgt else {
578            return Err(CliError::Other(format!(
579                "Foreign key '{}' references unknown target table '{}'",
580                fk.name, fk.table_to
581            )));
582        };
583        for col_name in &fk.columns_to {
584            if !tgt.columns.iter().any(|c| &c.name == col_name) {
585                return Err(CliError::Other(format!(
586                    "Foreign key '{}' references unknown target column '{}.{}'",
587                    fk.name, fk.table_to, col_name
588                )));
589            }
590        }
591
592        // Validate FK actions
593        if !VALID_FK_ACTIONS.contains(&fk.on_delete.as_str()) {
594            return Err(CliError::Other(format!(
595                "Foreign key '{}': invalid on_delete action '{}'. Valid: {}",
596                fk.name,
597                fk.on_delete,
598                VALID_FK_ACTIONS.join(", ")
599            )));
600        }
601        if !VALID_FK_ACTIONS.contains(&fk.on_update.as_str()) {
602            return Err(CliError::Other(format!(
603                "Foreign key '{}': invalid on_update action '{}'. Valid: {}",
604                fk.name,
605                fk.on_update,
606                VALID_FK_ACTIONS.join(", ")
607            )));
608        }
609        if def.dialect == Dialect::Mysql
610            && (fk.on_delete == "Set Default" || fk.on_update == "Set Default")
611        {
612            return Err(CliError::Other(format!(
613                "Foreign key '{}': MySQL does not support SET DEFAULT",
614                fk.name
615            )));
616        }
617        if fk.columns.len() != fk.columns_to.len() {
618            return Err(CliError::Other(format!(
619                "Foreign key '{}': source and target column counts must match",
620                fk.name
621            )));
622        }
623    }
624    Ok(())
625}
626
627// ── Interactive collection ──────────────────────────────────────────────────
628
629fn collect_interactively(
630    config: Option<&Config>,
631    options: &NewOptions,
632) -> Result<SchemaDefinition, CliError> {
633    // Phase 1: Setup
634    let dialect = resolve_dialect(config, options.dialect)?;
635    let casing = prompt_casing()?;
636    let output_path = resolve_output_path(config, options.schema.clone())?;
637    let schema_name = prompt_schema_name()?;
638
639    // Phase 2: named PostgreSQL enums and inline MySQL enums.
640    let enums: Vec<EnumDef> = if matches!(dialect, Dialect::Postgresql | Dialect::Mysql) {
641        prompt_enums()?
642    } else {
643        Vec::new()
644    };
645
646    // Phase 3 & 4: Tables + Columns
647    let mut tables: Vec<TableDef> = Vec::new();
648    loop {
649        let table = prompt_table(dialect, &enums)?;
650        tables.push(table);
651        if !confirm("Add another table?", false)? {
652            break;
653        }
654    }
655
656    // Phase 5: Indexes
657    let indexes: Vec<IndexDef> = if confirm("Add indexes?", false)? {
658        prompt_indexes(&tables)?
659    } else {
660        Vec::new()
661    };
662
663    // Phase 6: Foreign Keys
664    let foreign_keys: Vec<ForeignKeyDef> =
665        if tables.len() > 1 && confirm("Add foreign keys?", false)? {
666            prompt_foreign_keys(&tables, dialect)?
667        } else {
668            Vec::new()
669        };
670
671    Ok(SchemaDefinition {
672        dialect,
673        casing,
674        schema_name,
675        output_path,
676        enums,
677        tables,
678        indexes,
679        foreign_keys,
680    })
681}
682
683// ── Phase 1: Setup prompts ──────────────────────────────────────────────────
684
685fn resolve_dialect(
686    config: Option<&Config>,
687    cli_dialect: Option<Dialect>,
688) -> Result<Dialect, CliError> {
689    if let Some(d) = cli_dialect {
690        return Ok(d);
691    }
692    if let Some(c) = config {
693        return Ok(c.dialect());
694    }
695    let options = vec!["SQLite", "PostgreSQL", "MySQL"];
696    let answer = Select::new("Select database dialect:", options)
697        .prompt()
698        .map_err(|e| CliError::Other(format!("Prompt cancelled: {e}")))?;
699    match answer {
700        "SQLite" => Ok(Dialect::Sqlite),
701        "PostgreSQL" => Ok(Dialect::Postgresql),
702        "MySQL" => Ok(Dialect::Mysql),
703        _ => unreachable!(),
704    }
705}
706
707fn prompt_casing() -> Result<FieldCasing, CliError> {
708    let options = vec!["snake_case (default)", "camelCase"];
709    let answer = Select::new("Select field casing:", options)
710        .prompt()
711        .map_err(|e| CliError::Other(format!("Prompt cancelled: {e}")))?;
712    match answer {
713        s if s.starts_with("snake") => Ok(FieldCasing::Snake),
714        s if s.starts_with("camel") => Ok(FieldCasing::Camel),
715        _ => Ok(FieldCasing::Snake),
716    }
717}
718
719fn resolve_output_path(
720    config: Option<&Config>,
721    cli_schema: Option<String>,
722) -> Result<String, CliError> {
723    if let Some(s) = cli_schema {
724        return Ok(s);
725    }
726    let default = config.map_or_else(
727        || "src/schema.rs".to_string(),
728        super::super::config::Config::schema_display,
729    );
730    Text::new("Schema output path:")
731        .with_default(&default)
732        .prompt()
733        .map_err(|e| CliError::Other(format!("Prompt cancelled: {e}")))
734}
735
736fn prompt_schema_name() -> Result<String, CliError> {
737    Text::new("Schema struct name:")
738        .with_default("AppSchema")
739        .prompt()
740        .map_err(|e| CliError::Other(format!("Prompt cancelled: {e}")))
741}
742
743// ── Phase 2: Enums (PostgreSQL only) ────────────────────────────────────────
744
745fn prompt_enums() -> Result<Vec<EnumDef>, CliError> {
746    let mut enums = Vec::new();
747    if !confirm("Define any enums?", false)? {
748        return Ok(enums);
749    }
750    loop {
751        let name = Text::new("Enum name:")
752            .with_validator(|s: &str| {
753                if s.is_empty() {
754                    Ok(Validation::Invalid("Name cannot be empty".into()))
755                } else if !is_valid_identifier(s) {
756                    Ok(Validation::Invalid(
757                        "Must be a valid Rust identifier".into(),
758                    ))
759                } else {
760                    Ok(Validation::Valid)
761                }
762            })
763            .prompt()
764            .map_err(|e| CliError::Other(format!("Prompt cancelled: {e}")))?;
765
766        let mut variants = Vec::new();
767        loop {
768            let variant = Text::new("  Enum variant (empty to finish):")
769                .prompt()
770                .map_err(|e| CliError::Other(format!("Prompt cancelled: {e}")))?;
771            if variant.is_empty() {
772                break;
773            }
774            variants.push(variant);
775        }
776        if variants.is_empty() {
777            println!("  Skipping enum with no variants.");
778        } else {
779            enums.push(EnumDef { name, variants });
780        }
781        if !confirm("Add another enum?", false)? {
782            break;
783        }
784    }
785    Ok(enums)
786}
787
788// ── Phase 3 & 4: Tables + Columns ──────────────────────────────────────────
789
790fn prompt_table(dialect: Dialect, enums: &[EnumDef]) -> Result<TableDef, CliError> {
791    let name = Text::new("Table name:")
792        .with_validator(|s: &str| {
793            if s.is_empty() {
794                Ok(Validation::Invalid("Name cannot be empty".into()))
795            } else if !is_valid_identifier(s) {
796                Ok(Validation::Invalid(
797                    "Must be a valid Rust identifier (letters, digits, underscores)".into(),
798                ))
799            } else {
800                Ok(Validation::Valid)
801            }
802        })
803        .prompt()
804        .map_err(|e| CliError::Other(format!("Prompt cancelled: {e}")))?;
805
806    let mut strict = false;
807    let mut without_rowid = false;
808    let mut pg_schema = "public".to_string();
809
810    match dialect {
811        Dialect::Sqlite | Dialect::Turso => {
812            let table_opts = vec!["strict", "without_rowid"];
813            let selected = MultiSelect::new("Table options (space to toggle):", table_opts)
814                .prompt()
815                .map_err(|e| CliError::Other(format!("Prompt cancelled: {e}")))?;
816            strict = selected.contains(&"strict");
817            without_rowid = selected.contains(&"without_rowid");
818        }
819        Dialect::Postgresql => {
820            pg_schema = Text::new("PostgreSQL schema:")
821                .with_default("public")
822                .prompt()
823                .map_err(|e| CliError::Other(format!("Prompt cancelled: {e}")))?;
824        }
825        Dialect::Mysql => {}
826    }
827
828    // Columns
829    let mut columns = Vec::new();
830    println!();
831    println!("  Define columns for '{name}':");
832    loop {
833        let col = prompt_column(dialect, enums)?;
834        columns.push(col);
835        if !confirm("  Add another column?", true)? {
836            break;
837        }
838    }
839
840    Ok(TableDef {
841        name,
842        columns,
843        strict,
844        without_rowid,
845        pg_schema,
846    })
847}
848
849fn prompt_column(dialect: Dialect, enums: &[EnumDef]) -> Result<ColumnDef, CliError> {
850    let col_name = Text::new("  Column name:")
851        .with_validator(|s: &str| {
852            if s.is_empty() {
853                Ok(Validation::Invalid("Name cannot be empty".into()))
854            } else if !is_valid_identifier(s) {
855                Ok(Validation::Invalid("Must be a valid identifier".into()))
856            } else {
857                Ok(Validation::Valid)
858            }
859        })
860        .prompt()
861        .map_err(|e| CliError::Other(format!("Prompt cancelled: {e}")))?;
862
863    let (sql_type, enum_name) = prompt_type(dialect, enums)?;
864
865    let nullable = confirm("  Nullable (Option<T>)?", false)?;
866
867    let constraint_opts = match dialect {
868        Dialect::Sqlite | Dialect::Turso => {
869            vec!["Primary Key", "Autoincrement", "Unique", "Default value"]
870        }
871        Dialect::Postgresql => {
872            vec![
873                "Primary Key",
874                "Identity (auto-increment)",
875                "Unique",
876                "Default value",
877            ]
878        }
879        Dialect::Mysql => {
880            vec!["Primary Key", "Autoincrement", "Unique", "Default value"]
881        }
882    };
883    let selected = MultiSelect::new("  Column constraints (space to toggle):", constraint_opts)
884        .prompt()
885        .map_err(|e| CliError::Other(format!("Prompt cancelled: {e}")))?;
886
887    let primary_key = selected.iter().any(|s| s.starts_with("Primary"));
888    let autoincrement = selected.iter().any(|s| s.starts_with("Autoincrement"));
889    let identity = selected.iter().any(|s| s.starts_with("Identity"));
890    let unique = selected.iter().any(|s| s.starts_with("Unique"));
891    let has_default = selected.iter().any(|s| s.starts_with("Default"));
892
893    let default = if has_default {
894        let val = Text::new("  Default value:")
895            .prompt()
896            .map_err(|e| CliError::Other(format!("Prompt cancelled: {e}")))?;
897        Some(val)
898    } else {
899        None
900    };
901
902    let auto_gen = if autoincrement {
903        Some(AutoGenKind::Autoincrement)
904    } else if identity {
905        Some(AutoGenKind::Identity)
906    } else {
907        None
908    };
909
910    Ok(ColumnDef {
911        name: col_name,
912        sql_type,
913        not_null: !nullable,
914        primary_key,
915        unique,
916        default,
917        auto_gen,
918        enum_name,
919    })
920}
921
922fn prompt_type(dialect: Dialect, enums: &[EnumDef]) -> Result<(String, Option<String>), CliError> {
923    let mut options: Vec<String> = match dialect {
924        Dialect::Sqlite | Dialect::Turso => {
925            vec![
926                "i32".into(),
927                "i64".into(),
928                "f64".into(),
929                "String".into(),
930                "bool".into(),
931                "Vec<u8>".into(),
932            ]
933        }
934        Dialect::Postgresql => {
935            vec![
936                "i16".into(),
937                "i32".into(),
938                "i64".into(),
939                "f32".into(),
940                "f64".into(),
941                "String".into(),
942                "bool".into(),
943                "Vec<u8>".into(),
944                "uuid::Uuid".into(),
945                "chrono::NaiveDate".into(),
946                "chrono::NaiveDateTime".into(),
947                "chrono::DateTime<chrono::Utc>".into(),
948                "serde_json::Value".into(),
949            ]
950        }
951        Dialect::Mysql => {
952            vec![
953                "i8".into(),
954                "i16".into(),
955                "i32".into(),
956                "i64".into(),
957                "u8".into(),
958                "u16".into(),
959                "u32".into(),
960                "u64".into(),
961                "f32".into(),
962                "f64".into(),
963                "String".into(),
964                "bool".into(),
965                "Vec<u8>".into(),
966                "chrono::NaiveDate".into(),
967                "chrono::NaiveDateTime".into(),
968                "serde_json::Value".into(),
969            ]
970        }
971    };
972
973    // Append user-defined enums as type choices
974    for e in enums {
975        options.push(format!("enum:{}", e.name));
976    }
977
978    let refs: Vec<&str> = options.iter().map(std::string::String::as_str).collect();
979    let chosen = Select::new("  Rust type:", refs)
980        .prompt()
981        .map_err(|e| CliError::Other(format!("Prompt cancelled: {e}")))?;
982
983    // Map user-friendly Rust type -> SQL type string for codegen
984    if let Some(enum_name) = chosen.strip_prefix("enum:") {
985        // For enum columns, the sql_type is the enum name itself
986        return Ok((enum_name.to_string(), Some(enum_name.to_string())));
987    }
988
989    let sql_type = match dialect {
990        Dialect::Sqlite | Dialect::Turso => match chosen {
991            "i32" | "i64" => "integer",
992            "f64" => "real",
993            "bool" => "boolean",
994            "Vec<u8>" => "blob",
995            _ => "text",
996        },
997        Dialect::Postgresql => match chosen {
998            "i16" => "int2",
999            "i32" => "int4",
1000            "i64" => "int8",
1001            "f32" => "float4",
1002            "f64" => "float8",
1003            "bool" => "bool",
1004            "Vec<u8>" => "bytea",
1005            "uuid::Uuid" => "uuid",
1006            "chrono::NaiveDate" => "date",
1007            "chrono::NaiveDateTime" => "timestamp",
1008            "chrono::DateTime<chrono::Utc>" => "timestamptz",
1009            "serde_json::Value" => "jsonb",
1010            _ => "text",
1011        },
1012        Dialect::Mysql => match chosen {
1013            "i8" => "tinyint",
1014            "u8" => "tinyint unsigned",
1015            "i16" => "smallint",
1016            "u16" => "smallint unsigned",
1017            "i32" => "int",
1018            "u32" => "int unsigned",
1019            "i64" => "bigint",
1020            "u64" => "bigint unsigned",
1021            "f32" => "float",
1022            "f64" => "double",
1023            "bool" => "boolean",
1024            "Vec<u8>" => "blob",
1025            "chrono::NaiveDate" => "date",
1026            "chrono::NaiveDateTime" => "datetime",
1027            "serde_json::Value" => "json",
1028            _ => "text",
1029        },
1030    };
1031
1032    Ok((sql_type.to_string(), None))
1033}
1034
1035// ── Phase 5: Indexes ────────────────────────────────────────────────────────
1036
1037fn prompt_indexes(tables: &[TableDef]) -> Result<Vec<IndexDef>, CliError> {
1038    let mut indexes = Vec::new();
1039    loop {
1040        let table_names: Vec<&str> = tables.iter().map(|t| t.name.as_str()).collect();
1041        let table_name = Select::new("Index on which table?", table_names)
1042            .prompt()
1043            .map_err(|e| CliError::Other(format!("Prompt cancelled: {e}")))?;
1044
1045        let table = tables.iter().find(|t| t.name == table_name).unwrap();
1046        let col_names: Vec<&str> = table.columns.iter().map(|c| c.name.as_str()).collect();
1047
1048        if col_names.is_empty() {
1049            println!("  Table has no columns, skipping.");
1050            if !confirm("Add another index?", false)? {
1051                break;
1052            }
1053            continue;
1054        }
1055
1056        let selected_cols = MultiSelect::new("Select columns for index:", col_names)
1057            .prompt()
1058            .map_err(|e| CliError::Other(format!("Prompt cancelled: {e}")))?;
1059
1060        if selected_cols.is_empty() {
1061            println!("  No columns selected, skipping.");
1062            if !confirm("Add another index?", false)? {
1063                break;
1064            }
1065            continue;
1066        }
1067
1068        let is_unique = confirm("  Unique index?", false)?;
1069
1070        let suggested_name = format!("{}_{}_idx", table_name, selected_cols.join("_"));
1071        let idx_name = Text::new("  Index name:")
1072            .with_default(&suggested_name)
1073            .prompt()
1074            .map_err(|e| CliError::Other(format!("Prompt cancelled: {e}")))?;
1075
1076        indexes.push(IndexDef {
1077            name: idx_name,
1078            table: table_name.to_string(),
1079            columns: selected_cols
1080                .into_iter()
1081                .map(std::string::ToString::to_string)
1082                .collect(),
1083            unique: is_unique,
1084            pg_schema: table.pg_schema.clone(),
1085        });
1086
1087        if !confirm("Add another index?", false)? {
1088            break;
1089        }
1090    }
1091    Ok(indexes)
1092}
1093
1094// ── Phase 6: Foreign Keys ───────────────────────────────────────────────────
1095
1096fn prompt_foreign_keys(
1097    tables: &[TableDef],
1098    dialect: Dialect,
1099) -> Result<Vec<ForeignKeyDef>, CliError> {
1100    let mut fks = Vec::new();
1101    let action_options = vec![
1102        "No Action",
1103        "Cascade",
1104        "Set Null",
1105        "Set Default",
1106        "Restrict",
1107    ];
1108    loop {
1109        let table_names: Vec<&str> = tables.iter().map(|t| t.name.as_str()).collect();
1110
1111        let src_table_name = Select::new("Source table:", table_names.clone())
1112            .prompt()
1113            .map_err(|e| CliError::Other(format!("Prompt cancelled: {e}")))?;
1114        let src_table = tables.iter().find(|t| t.name == src_table_name).unwrap();
1115        let src_col_names: Vec<&str> = src_table.columns.iter().map(|c| c.name.as_str()).collect();
1116
1117        let src_cols = MultiSelect::new("Source column(s):", src_col_names)
1118            .prompt()
1119            .map_err(|e| CliError::Other(format!("Prompt cancelled: {e}")))?;
1120
1121        let tgt_table_name = Select::new("Target (referenced) table:", table_names)
1122            .prompt()
1123            .map_err(|e| CliError::Other(format!("Prompt cancelled: {e}")))?;
1124        let tgt_table = tables.iter().find(|t| t.name == tgt_table_name).unwrap();
1125        let tgt_col_names: Vec<&str> = tgt_table.columns.iter().map(|c| c.name.as_str()).collect();
1126
1127        let tgt_cols = MultiSelect::new("Target column(s):", tgt_col_names)
1128            .prompt()
1129            .map_err(|e| CliError::Other(format!("Prompt cancelled: {e}")))?;
1130
1131        let on_delete = Select::new("ON DELETE action:", action_options.clone())
1132            .prompt()
1133            .map_err(|e| CliError::Other(format!("Prompt cancelled: {e}")))?;
1134
1135        let on_update = Select::new("ON UPDATE action:", action_options.clone())
1136            .prompt()
1137            .map_err(|e| CliError::Other(format!("Prompt cancelled: {e}")))?;
1138
1139        let fk_name = format!("{}_{}_fk", src_table_name, src_cols.join("_"));
1140
1141        let pg_schema_to = match dialect {
1142            Dialect::Postgresql => tgt_table.pg_schema.clone(),
1143            _ => String::new(),
1144        };
1145
1146        fks.push(ForeignKeyDef {
1147            name: fk_name,
1148            table: src_table_name.to_string(),
1149            columns: src_cols
1150                .into_iter()
1151                .map(std::string::ToString::to_string)
1152                .collect(),
1153            table_to: tgt_table_name.to_string(),
1154            columns_to: tgt_cols
1155                .into_iter()
1156                .map(std::string::ToString::to_string)
1157                .collect(),
1158            on_delete: on_delete.to_string(),
1159            on_update: on_update.to_string(),
1160            pg_schema: src_table.pg_schema.clone(),
1161            pg_schema_to,
1162        });
1163
1164        if !confirm("Add another foreign key?", false)? {
1165            break;
1166        }
1167    }
1168    Ok(fks)
1169}
1170
1171// ── Phase 7: Code generation ────────────────────────────────────────────────
1172
1173fn generate_sqlite(
1174    tables: &[TableDef],
1175    indexes: &[IndexDef],
1176    fks: &[ForeignKeyDef],
1177    schema_name: &str,
1178    casing: FieldCasing,
1179) -> String {
1180    use drizzle_migrations::sqlite::codegen;
1181    use drizzle_migrations::sqlite::collection::SQLiteDDL;
1182    use drizzle_types::sqlite::ddl::{
1183        Column, ForeignKey, Index, IndexColumn, PrimaryKey, Table, UniqueConstraint,
1184    };
1185
1186    let mut ddl = SQLiteDDL::new();
1187
1188    for (table_idx, table) in tables.iter().enumerate() {
1189        let mut t = Table::new(table.name.clone());
1190        if table.strict {
1191            t = t.strict();
1192        }
1193        if table.without_rowid {
1194            t = t.without_rowid();
1195        }
1196        ddl.tables.push(t);
1197
1198        let mut pk_cols: Vec<String> = Vec::new();
1199        let mut unique_cols: Vec<String> = Vec::new();
1200
1201        for (col_idx, col) in table.columns.iter().enumerate() {
1202            let mut column =
1203                Column::new(table.name.clone(), col.name.clone(), col.sql_type.clone());
1204            if col.not_null {
1205                column = column.not_null();
1206            }
1207            if col.is_autoincrement() {
1208                column = column.autoincrement();
1209            }
1210            if let Some(ref default) = col.default {
1211                column = column.default_value(default.clone());
1212            }
1213            // Set ordinal position to preserve order
1214            column.ordinal_position = Some(
1215                i32::try_from(col_idx)
1216                    .ok()
1217                    .and_then(|i| i.checked_add(1))
1218                    .unwrap_or(i32::MAX),
1219            );
1220            ddl.columns.push(column);
1221
1222            if col.primary_key {
1223                pk_cols.push(col.name.clone());
1224            }
1225            if col.unique {
1226                unique_cols.push(col.name.clone());
1227            }
1228        }
1229
1230        if !pk_cols.is_empty() {
1231            ddl.pks.push(PrimaryKey::from_strings(
1232                table.name.clone(),
1233                format!("{}_pk", table.name),
1234                pk_cols,
1235            ));
1236        }
1237        for uc in unique_cols {
1238            ddl.uniques.push(UniqueConstraint::from_strings(
1239                table.name.clone(),
1240                format!("{}_{}_unique", table.name, uc),
1241                vec![uc],
1242            ));
1243        }
1244
1245        // Drop the table_idx binding explicitly
1246        let _ = table_idx;
1247    }
1248
1249    // Add indexes
1250    for idx in indexes {
1251        let columns: Vec<IndexColumn> = idx
1252            .columns
1253            .iter()
1254            .map(|c| IndexColumn::new(c.clone()))
1255            .collect();
1256        let mut index = Index::new(idx.table.clone(), idx.name.clone(), columns);
1257        if idx.unique {
1258            index = index.unique();
1259        }
1260        ddl.indexes.push(index);
1261    }
1262
1263    // Add foreign keys
1264    for fk in fks {
1265        let mut foreign_key = ForeignKey::from_strings(
1266            fk.table.clone(),
1267            fk.name.clone(),
1268            fk.columns.clone(),
1269            fk.table_to.clone(),
1270            fk.columns_to.clone(),
1271        );
1272        if fk.on_delete != "No Action" {
1273            foreign_key = foreign_key.on_delete(fk.on_delete.to_uppercase());
1274        }
1275        if fk.on_update != "No Action" {
1276            foreign_key = foreign_key.on_update(fk.on_update.to_uppercase());
1277        }
1278        ddl.fks.push(foreign_key);
1279    }
1280
1281    let field_casing = match casing {
1282        FieldCasing::Snake => codegen::FieldCasing::Snake,
1283        FieldCasing::Camel => codegen::FieldCasing::Camel,
1284    };
1285
1286    let options = codegen::CodegenOptions {
1287        module_doc: Some("Generated by `drizzle new`".to_string()),
1288        include_schema: true,
1289        schema_name: schema_name.to_string(),
1290        use_pub: true,
1291        field_casing,
1292    };
1293
1294    codegen::generate_rust_schema(&ddl, &options).code
1295}
1296
1297fn generate_mysql(
1298    tables: &[TableDef],
1299    indexes: &[IndexDef],
1300    fks: &[ForeignKeyDef],
1301    enums: &[EnumDef],
1302    schema_name: &str,
1303    casing: FieldCasing,
1304) -> Result<String, CliError> {
1305    use drizzle_migrations::mysql::codegen;
1306    use drizzle_migrations::mysql::collection::MySQLDDL;
1307    use drizzle_migrations::mysql::ddl::{
1308        Column, ForeignKey, Index, IndexColumn, InlineEnum, InlineType, PrimaryKey, Table,
1309        UniqueConstraint,
1310    };
1311
1312    let enum_values = enums
1313        .iter()
1314        .map(|definition| (definition.name.as_str(), definition.variants.as_slice()))
1315        .collect::<std::collections::HashMap<_, _>>();
1316    let mut ddl = MySQLDDL::new();
1317
1318    for table in tables {
1319        ddl.tables.push(Table::new(table.name.clone()));
1320        let mut primary_columns = Vec::new();
1321
1322        for column_definition in &table.columns {
1323            let mut column = Column::new(
1324                table.name.clone(),
1325                column_definition.name.clone(),
1326                column_definition.sql_type.clone(),
1327            );
1328            column.not_null = column_definition.not_null;
1329            column.autoincrement = column_definition.is_autoincrement();
1330            column.default = column_definition
1331                .default
1332                .as_ref()
1333                .map(|value| Cow::Owned(value.clone()));
1334            if let Some(enum_name) = column_definition.enum_name.as_deref()
1335                && let Some(values) = enum_values.get(enum_name)
1336            {
1337                column.inline_type =
1338                    Some(InlineType::Enum(InlineEnum::new(values.iter().cloned())));
1339            }
1340            ddl.columns.push(column);
1341
1342            if column_definition.primary_key {
1343                primary_columns.push(column_definition.name.clone());
1344            }
1345            if column_definition.unique {
1346                ddl.uniques.push(UniqueConstraint::new(
1347                    table.name.clone(),
1348                    format!("{}_{}_unique", table.name, column_definition.name),
1349                    [column_definition.name.clone()],
1350                ));
1351            }
1352        }
1353
1354        if !primary_columns.is_empty() {
1355            ddl.pks
1356                .push(PrimaryKey::new(table.name.clone(), primary_columns));
1357        }
1358    }
1359
1360    for definition in indexes {
1361        let columns = definition
1362            .columns
1363            .iter()
1364            .cloned()
1365            .map(IndexColumn::column)
1366            .collect();
1367        let mut index = Index::new(definition.table.clone(), definition.name.clone(), columns);
1368        index.unique = definition.unique;
1369        ddl.indexes.push(index);
1370    }
1371
1372    for definition in fks {
1373        let mut foreign_key = ForeignKey::new(
1374            definition.table.clone(),
1375            definition.name.clone(),
1376            definition.columns.clone(),
1377            definition.table_to.clone(),
1378            definition.columns_to.clone(),
1379        );
1380        foreign_key.on_delete = mysql_referential_action(&definition.on_delete);
1381        foreign_key.on_update = mysql_referential_action(&definition.on_update);
1382        ddl.fks.push(foreign_key);
1383    }
1384
1385    let options = codegen::CodegenOptions {
1386        module_doc: Some("Generated by `drizzle new`".to_string()),
1387        include_schema: true,
1388        schema_name: schema_name.to_string(),
1389        use_pub: true,
1390        field_casing: match casing {
1391            FieldCasing::Snake => codegen::FieldCasing::Snake,
1392            FieldCasing::Camel => codegen::FieldCasing::Camel,
1393        },
1394    };
1395    codegen::generate_rust_schema(&ddl, &options)
1396        .map(|generated| generated.code)
1397        .map_err(|error| CliError::Other(format!("Cannot generate MySQL schema: {error}")))
1398}
1399
1400fn mysql_referential_action(
1401    action: &str,
1402) -> Option<drizzle_migrations::mysql::ddl::ReferentialAction> {
1403    use drizzle_migrations::mysql::ddl::ReferentialAction;
1404
1405    match action {
1406        "Cascade" => Some(ReferentialAction::Cascade),
1407        "Set Null" => Some(ReferentialAction::SetNull),
1408        "Restrict" => Some(ReferentialAction::Restrict),
1409        "No Action" => Some(ReferentialAction::NoAction),
1410        // MySQL doesn't support SET DEFAULT, and validation rejects it for
1411        // MySQL before generation.
1412        _ => None,
1413    }
1414}
1415
1416fn generate_postgres(
1417    tables: &[TableDef],
1418    indexes: &[IndexDef],
1419    fks: &[ForeignKeyDef],
1420    enums: &[EnumDef],
1421    schema_name: &str,
1422    casing: FieldCasing,
1423) -> String {
1424    use drizzle_migrations::postgres::codegen;
1425    use drizzle_migrations::postgres::collection::PostgresDDL;
1426    use drizzle_types::postgres::ddl::{Enum, Table};
1427
1428    let mut ddl = PostgresDDL::new();
1429
1430    // Add enums
1431    for e in enums {
1432        let values: Vec<Cow<'static, str>> =
1433            e.variants.iter().map(|v| Cow::Owned(v.clone())).collect();
1434        ddl.enums.push(Enum::new(
1435            "public",
1436            Cow::<str>::Owned(e.name.clone()),
1437            Cow::<[Cow<'static, str>]>::Owned(values),
1438        ));
1439    }
1440
1441    // Add tables + columns
1442    for table in tables {
1443        ddl.tables
1444            .push(Table::new(table.pg_schema.clone(), table.name.clone()));
1445        add_postgres_table_columns(&mut ddl, table);
1446    }
1447
1448    add_postgres_indexes(&mut ddl, indexes);
1449    add_postgres_foreign_keys(&mut ddl, fks);
1450
1451    let field_casing = match casing {
1452        FieldCasing::Snake => codegen::FieldCasing::Snake,
1453        FieldCasing::Camel => codegen::FieldCasing::Camel,
1454    };
1455
1456    let options = codegen::CodegenOptions {
1457        module_doc: Some("Generated by `drizzle new`".to_string()),
1458        include_schema: true,
1459        schema_name: schema_name.to_string(),
1460        use_pub: true,
1461        field_casing,
1462    };
1463
1464    codegen::generate_rust_schema(&ddl, &options).code
1465}
1466
1467/// Populate `ddl.columns`, `ddl.pks`, `ddl.uniques` for a single postgres table.
1468fn add_postgres_table_columns(
1469    ddl: &mut drizzle_migrations::postgres::collection::PostgresDDL,
1470    table: &TableDef,
1471) {
1472    use drizzle_types::postgres::ddl::{Column, PrimaryKey, UniqueConstraint};
1473
1474    let mut pk_cols: Vec<String> = Vec::new();
1475    let mut unique_cols: Vec<String> = Vec::new();
1476
1477    for (col_idx, col) in table.columns.iter().enumerate() {
1478        let mut column = Column::new(
1479            table.pg_schema.clone(),
1480            table.name.clone(),
1481            col.name.clone(),
1482            col.sql_type.clone(),
1483        );
1484        if col.not_null {
1485            column = column.not_null();
1486        }
1487        if let Some(ref default) = col.default {
1488            column = column.default_value(default.clone());
1489        }
1490        if col.is_identity() {
1491            use drizzle_types::postgres::ddl::Identity;
1492            let seq_name = format!("{}_{}_seq", table.name, col.name);
1493            column.identity = Some(Identity::always(seq_name));
1494        }
1495        if col.enum_name.is_some() {
1496            // Set type_schema so codegen can find it in the enum_map
1497            column.type_schema = Some(Cow::Owned(table.pg_schema.clone()));
1498        }
1499        column.ordinal_position = Some(
1500            i32::try_from(col_idx)
1501                .ok()
1502                .and_then(|i| i.checked_add(1))
1503                .unwrap_or(i32::MAX),
1504        );
1505        ddl.columns.push(column);
1506
1507        if col.primary_key {
1508            pk_cols.push(col.name.clone());
1509        }
1510        if col.unique {
1511            unique_cols.push(col.name.clone());
1512        }
1513    }
1514
1515    if !pk_cols.is_empty() {
1516        ddl.pks.push(PrimaryKey::from_strings(
1517            table.pg_schema.clone(),
1518            table.name.clone(),
1519            format!("{}_pk", table.name),
1520            pk_cols,
1521        ));
1522    }
1523    for uc in unique_cols {
1524        ddl.uniques.push(UniqueConstraint::from_strings(
1525            table.pg_schema.clone(),
1526            table.name.clone(),
1527            format!("{}_{}_unique", table.name, uc),
1528            vec![uc],
1529        ));
1530    }
1531}
1532
1533/// Append index definitions to the postgres DDL collection.
1534fn add_postgres_indexes(
1535    ddl: &mut drizzle_migrations::postgres::collection::PostgresDDL,
1536    indexes: &[IndexDef],
1537) {
1538    use drizzle_types::postgres::ddl::{Index, IndexColumn};
1539    for idx in indexes {
1540        let columns: Vec<IndexColumn> = idx
1541            .columns
1542            .iter()
1543            .map(|c| IndexColumn::new(c.clone()))
1544            .collect();
1545        let mut index = Index::new(
1546            idx.pg_schema.clone(),
1547            idx.table.clone(),
1548            idx.name.clone(),
1549            columns,
1550        );
1551        if idx.unique {
1552            index = index.unique();
1553        }
1554        ddl.indexes.push(index);
1555    }
1556}
1557
1558/// Append foreign-key definitions to the postgres DDL collection.
1559fn add_postgres_foreign_keys(
1560    ddl: &mut drizzle_migrations::postgres::collection::PostgresDDL,
1561    fks: &[ForeignKeyDef],
1562) {
1563    use drizzle_types::postgres::ddl::ForeignKey;
1564    for fk in fks {
1565        let mut foreign_key = ForeignKey::from_strings(
1566            fk.pg_schema.clone(),
1567            fk.table.clone(),
1568            fk.name.clone(),
1569            fk.columns.clone(),
1570            fk.pg_schema_to.clone(),
1571            fk.table_to.clone(),
1572            fk.columns_to.clone(),
1573        );
1574        if fk.on_delete != "No Action" {
1575            foreign_key = foreign_key.on_delete(fk.on_delete.to_uppercase());
1576        }
1577        if fk.on_update != "No Action" {
1578            foreign_key = foreign_key.on_update(fk.on_update.to_uppercase());
1579        }
1580        ddl.fks.push(foreign_key);
1581    }
1582}
1583
1584// ── Utility helpers ─────────────────────────────────────────────────────────
1585
1586fn confirm(message: &str, default: bool) -> Result<bool, CliError> {
1587    Confirm::new(message)
1588        .with_default(default)
1589        .prompt()
1590        .map_err(|e| CliError::Other(format!("Prompt cancelled: {e}")))
1591}
1592
1593fn is_valid_identifier(s: &str) -> bool {
1594    if s.is_empty() {
1595        return false;
1596    }
1597    let mut chars = s.chars();
1598    let first = chars.next().unwrap();
1599    if !first.is_ascii_alphabetic() && first != '_' {
1600        return false;
1601    }
1602    chars.all(|c| c.is_ascii_alphanumeric() || c == '_')
1603}
1604
1605// ── Tests ───────────────────────────────────────────────────────────────────
1606
1607#[cfg(test)]
1608mod tests {
1609    use super::*;
1610
1611    fn minimal_sqlite_def() -> SchemaDefinition {
1612        SchemaDefinition {
1613            dialect: Dialect::Sqlite,
1614            casing: FieldCasing::Snake,
1615            schema_name: "TestSchema".into(),
1616            output_path: "src/schema.rs".into(),
1617            enums: vec![],
1618            tables: vec![TableDef {
1619                name: "users".into(),
1620                columns: vec![ColumnDef {
1621                    name: "id".into(),
1622                    sql_type: "integer".into(),
1623                    not_null: true,
1624                    primary_key: true,
1625                    unique: false,
1626                    default: None,
1627                    auto_gen: None,
1628                    enum_name: None,
1629                }],
1630                strict: false,
1631                without_rowid: false,
1632                pg_schema: String::new(),
1633            }],
1634            indexes: vec![],
1635            foreign_keys: vec![],
1636        }
1637    }
1638
1639    fn minimal_mysql_def() -> SchemaDefinition {
1640        let mut definition = minimal_sqlite_def();
1641        definition.dialect = Dialect::Mysql;
1642        definition.enums = vec![EnumDef {
1643            name: "status".into(),
1644            variants: vec!["active".into(), "disabled".into()],
1645        }];
1646        definition.tables[0].columns[0].sql_type = "bigint unsigned".into();
1647        definition.tables[0].columns[0].auto_gen = Some(AutoGenKind::Autoincrement);
1648        definition.tables[0].columns.push(ColumnDef {
1649            name: "status".into(),
1650            sql_type: "status".into(),
1651            not_null: true,
1652            primary_key: false,
1653            unique: true,
1654            default: Some("'active'".into()),
1655            auto_gen: None,
1656            enum_name: Some("status".into()),
1657        });
1658        definition
1659    }
1660
1661    #[test]
1662    fn validate_minimal_schema() {
1663        let def = minimal_sqlite_def();
1664        assert!(validate_schema(&def).is_ok());
1665    }
1666
1667    #[test]
1668    fn mysql_generation_uses_mysql_macros_and_dialect_features() {
1669        let definition = minimal_mysql_def();
1670        validate_schema(&definition).expect("valid MySQL definition");
1671        let code = generate_mysql(
1672            &definition.tables,
1673            &definition.indexes,
1674            &definition.foreign_keys,
1675            &definition.enums,
1676            &definition.schema_name,
1677            definition.casing,
1678        )
1679        .expect("valid MySQL definition");
1680
1681        assert!(code.contains("MySQLTable"), "{code}");
1682        assert!(code.contains("MySQLEnum"), "{code}");
1683        assert!(code.contains("AUTO_INCREMENT"), "{code}");
1684        assert!(code.contains("UNIQUE"), "{code}");
1685        assert!(!code.contains("SQLiteTable"), "{code}");
1686        assert!(!code.contains("PostgresTable"), "{code}");
1687    }
1688
1689    #[test]
1690    fn mysql_rejects_set_default_foreign_keys() {
1691        let mut definition = minimal_mysql_def();
1692        definition.tables.push(TableDef {
1693            name: "posts".into(),
1694            columns: vec![ColumnDef {
1695                name: "user_id".into(),
1696                sql_type: "bigint unsigned".into(),
1697                not_null: true,
1698                primary_key: false,
1699                unique: false,
1700                default: None,
1701                auto_gen: None,
1702                enum_name: None,
1703            }],
1704            strict: false,
1705            without_rowid: false,
1706            pg_schema: String::new(),
1707        });
1708        definition.foreign_keys.push(ForeignKeyDef {
1709            name: "posts_user_fk".into(),
1710            table: "posts".into(),
1711            columns: vec!["user_id".into()],
1712            table_to: "users".into(),
1713            columns_to: vec!["id".into()],
1714            on_delete: "Set Default".into(),
1715            on_update: "No Action".into(),
1716            pg_schema: String::new(),
1717            pg_schema_to: String::new(),
1718        });
1719
1720        let error = validate_schema(&definition).expect_err("SET DEFAULT is invalid in MySQL");
1721        assert!(error.to_string().contains("does not support SET DEFAULT"));
1722    }
1723
1724    #[test]
1725    fn validate_rejects_empty_tables() {
1726        let mut def = minimal_sqlite_def();
1727        def.tables.clear();
1728        let err = validate_schema(&def).unwrap_err();
1729        assert!(err.to_string().contains("at least one table"));
1730    }
1731
1732    #[test]
1733    fn validate_rejects_duplicate_table_names() {
1734        let mut def = minimal_sqlite_def();
1735        def.tables.push(def.tables[0].clone());
1736        let err = validate_schema(&def).unwrap_err();
1737        assert!(err.to_string().contains("Duplicate table name"));
1738    }
1739
1740    #[test]
1741    fn validate_rejects_empty_columns() {
1742        let mut def = minimal_sqlite_def();
1743        def.tables[0].columns.clear();
1744        let err = validate_schema(&def).unwrap_err();
1745        assert!(err.to_string().contains("at least one column"));
1746    }
1747
1748    #[test]
1749    fn validate_rejects_duplicate_column_names() {
1750        let mut def = minimal_sqlite_def();
1751        let dup = def.tables[0].columns[0].clone();
1752        def.tables[0].columns.push(dup);
1753        let err = validate_schema(&def).unwrap_err();
1754        assert!(err.to_string().contains("Duplicate column name"));
1755    }
1756
1757    #[test]
1758    fn validate_rejects_identity_on_sqlite() {
1759        let mut def = minimal_sqlite_def();
1760        def.tables[0].columns[0].auto_gen = Some(AutoGenKind::Identity);
1761        let err = validate_schema(&def).unwrap_err();
1762        assert!(err.to_string().contains("identity"));
1763        assert!(err.to_string().contains("PostgreSQL"));
1764    }
1765
1766    #[test]
1767    fn validate_rejects_autoincrement_on_postgres() {
1768        let mut def = minimal_sqlite_def();
1769        def.dialect = Dialect::Postgresql;
1770        def.tables[0].columns[0].auto_gen = Some(AutoGenKind::Autoincrement);
1771        let err = validate_schema(&def).unwrap_err();
1772        assert!(err.to_string().contains("autoincrement"));
1773        assert!(err.to_string().contains("SQLite"));
1774    }
1775
1776    #[test]
1777    fn validate_rejects_strict_on_postgres() {
1778        let mut def = minimal_sqlite_def();
1779        def.dialect = Dialect::Postgresql;
1780        def.tables[0].strict = true;
1781        let err = validate_schema(&def).unwrap_err();
1782        assert!(err.to_string().contains("strict"));
1783        assert!(err.to_string().contains("SQLite"));
1784    }
1785
1786    #[test]
1787    fn validate_rejects_enums_on_sqlite() {
1788        let mut def = minimal_sqlite_def();
1789        def.enums.push(EnumDef {
1790            name: "status".into(),
1791            variants: vec!["active".into()],
1792        });
1793        let err = validate_schema(&def).unwrap_err();
1794        assert!(err.to_string().contains("Enums"));
1795        assert!(err.to_string().contains("PostgreSQL"));
1796    }
1797
1798    #[test]
1799    fn validate_rejects_unknown_enum_reference() {
1800        let mut def = minimal_sqlite_def();
1801        def.dialect = Dialect::Postgresql;
1802        def.tables[0].columns[0].enum_name = Some("nonexistent".into());
1803        let err = validate_schema(&def).unwrap_err();
1804        assert!(err.to_string().contains("unknown enum"));
1805    }
1806
1807    #[test]
1808    fn validate_rejects_bad_fk_table_ref() {
1809        let mut def = minimal_sqlite_def();
1810        def.foreign_keys.push(ForeignKeyDef {
1811            name: "test_fk".into(),
1812            table: "nonexistent".into(),
1813            columns: vec!["id".into()],
1814            table_to: "users".into(),
1815            columns_to: vec!["id".into()],
1816            on_delete: "No Action".into(),
1817            on_update: "No Action".into(),
1818            pg_schema: String::new(),
1819            pg_schema_to: String::new(),
1820        });
1821        let err = validate_schema(&def).unwrap_err();
1822        assert!(err.to_string().contains("unknown source table"));
1823    }
1824
1825    #[test]
1826    fn validate_rejects_bad_fk_action() {
1827        let mut def = minimal_sqlite_def();
1828        def.tables.push(TableDef {
1829            name: "posts".into(),
1830            columns: vec![ColumnDef {
1831                name: "user_id".into(),
1832                sql_type: "integer".into(),
1833                not_null: true,
1834                primary_key: false,
1835                unique: false,
1836                default: None,
1837                auto_gen: None,
1838                enum_name: None,
1839            }],
1840            strict: false,
1841            without_rowid: false,
1842            pg_schema: String::new(),
1843        });
1844        def.foreign_keys.push(ForeignKeyDef {
1845            name: "posts_user_id_fk".into(),
1846            table: "posts".into(),
1847            columns: vec!["user_id".into()],
1848            table_to: "users".into(),
1849            columns_to: vec!["id".into()],
1850            on_delete: "INVALID".into(),
1851            on_update: "No Action".into(),
1852            pg_schema: String::new(),
1853            pg_schema_to: String::new(),
1854        });
1855        let err = validate_schema(&def).unwrap_err();
1856        assert!(err.to_string().contains("invalid on_delete"));
1857    }
1858
1859    #[test]
1860    fn validate_rejects_bad_index_column_ref() {
1861        let mut def = minimal_sqlite_def();
1862        def.indexes.push(IndexDef {
1863            name: "test_idx".into(),
1864            table: "users".into(),
1865            columns: vec!["nonexistent".into()],
1866            unique: false,
1867            pg_schema: String::new(),
1868        });
1869        let err = validate_schema(&def).unwrap_err();
1870        assert!(err.to_string().contains("unknown column"));
1871    }
1872
1873    #[test]
1874    fn json_round_trip() {
1875        let def = minimal_sqlite_def();
1876        let json = serde_json::to_string_pretty(&def).unwrap();
1877        let parsed: SchemaDefinition = serde_json::from_str(&json).unwrap();
1878        assert_eq!(parsed.dialect, def.dialect);
1879        assert_eq!(parsed.tables.len(), 1);
1880        assert_eq!(parsed.tables[0].name, "users");
1881        assert_eq!(parsed.tables[0].columns[0].name, "id");
1882    }
1883
1884    #[test]
1885    fn json_defaults_applied() {
1886        let json = r#"{
1887            "dialect": "sqlite",
1888            "tables": [{
1889                "name": "items",
1890                "columns": [{"name": "id", "sql_type": "integer"}]
1891            }]
1892        }"#;
1893        let def: SchemaDefinition = serde_json::from_str(json).unwrap();
1894        assert_eq!(def.schema_name, "AppSchema");
1895        assert_eq!(def.output_path, "src/schema.rs");
1896        assert!(def.enums.is_empty());
1897        assert!(def.indexes.is_empty());
1898        assert!(def.foreign_keys.is_empty());
1899        assert!(!def.tables[0].columns[0].not_null);
1900        assert!(!def.tables[0].columns[0].primary_key);
1901    }
1902
1903    #[test]
1904    fn json_fk_action_defaults() {
1905        let json = r#"{
1906            "name": "test_fk",
1907            "table": "a",
1908            "columns": ["x"],
1909            "table_to": "b",
1910            "columns_to": ["y"]
1911        }"#;
1912        let fk: ForeignKeyDef = serde_json::from_str(json).unwrap();
1913        assert_eq!(fk.on_delete, "No Action");
1914        assert_eq!(fk.on_update, "No Action");
1915    }
1916}