use console::style;
use sea_orm::{ConnectionTrait, Database, DbBackend, Statement};
use std::env;
use std::fs;
use std::path::Path;
use std::process::Command;
use crate::templates;
use crate::templates::{ColumnInfo, TableInfo};
pub fn run(skip_migrations: bool, regenerate_models: bool) {
if !Path::new("src/models").exists() && !Path::new("src/migrations").exists() {
eprintln!(
"{} Not in a Ferro project directory",
style("Error:").red().bold()
);
std::process::exit(1);
}
if !skip_migrations {
run_migrations();
}
generate_entities(regenerate_models);
}
fn run_migrations() {
if !Path::new("src/migrations").exists() {
println!(
"{} No migrations directory found, skipping migrations",
style("Info:").yellow()
);
return;
}
if !Path::new("src/bin/migrate.rs").exists() {
println!(
"{} Migration binary not found, skipping migrations",
style("Info:").yellow()
);
return;
}
println!("{} Running pending migrations...", style("→").cyan());
let status = Command::new("cargo")
.args(["run", "--bin", "migrate", "--", "up"])
.status()
.expect("Failed to execute cargo command");
if !status.success() {
eprintln!("{} Migration failed", style("Error:").red().bold());
std::process::exit(1);
}
println!("{} Migrations complete", style("✓").green());
}
fn generate_entities(regenerate_models: bool) {
dotenvy::dotenv().ok();
let database_url = match env::var("DATABASE_URL") {
Ok(url) => url,
Err(_) => {
eprintln!(
"{} DATABASE_URL not set in .env",
style("Error:").red().bold()
);
std::process::exit(1);
}
};
println!("{} Discovering database schema...", style("→").cyan());
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
discover_and_generate(&database_url, regenerate_models).await;
});
}
async fn discover_and_generate(database_url: &str, regenerate_models: bool) {
let is_sqlite = database_url.starts_with("sqlite");
let db = match Database::connect(database_url).await {
Ok(db) => db,
Err(e) => {
eprintln!(
"{} Failed to connect to database: {}",
style("Error:").red().bold(),
e
);
std::process::exit(1);
}
};
let tables = if is_sqlite {
discover_sqlite_tables(&db).await
} else {
discover_postgres_tables(&db).await
};
let tables: Vec<_> = tables
.into_iter()
.filter(|t| t.name != "seaql_migrations" && !t.name.starts_with("_"))
.collect();
if tables.is_empty() {
println!("{} No tables found in database", style("Info:").yellow());
return;
}
println!(
"{} Found {} table(s): {}",
style("✓").green(),
tables.len(),
tables
.iter()
.map(|t| t.name.as_str())
.collect::<Vec<_>>()
.join(", ")
);
let models_dir = Path::new("src/models");
if !models_dir.exists() {
fs::create_dir_all(models_dir).expect("Failed to create models directory");
println!("{} Created src/models directory", style("✓").green());
}
let entities_dir = models_dir.join("entities");
if !entities_dir.exists() {
fs::create_dir_all(&entities_dir).expect("Failed to create entities directory");
println!(
"{} Created src/models/entities directory",
style("✓").green()
);
}
for table in &tables {
generate_entity_file(table, &entities_dir);
if regenerate_models {
generate_user_file(table, models_dir);
} else {
generate_user_file_if_not_exists(table, models_dir);
}
}
update_entities_mod(&tables, &entities_dir);
update_models_mod(&tables, models_dir);
println!();
println!(
"{} Entity files generated successfully!",
style("✓").green().bold()
);
println!();
println!("Generated files:");
for table in &tables {
println!(
" {} src/models/entities/{}.rs (auto-generated)",
style("•").dim(),
table.name
);
println!(
" {} src/models/{}.rs (user customizations)",
style("•").dim(),
table.name
);
}
}
async fn discover_sqlite_tables(db: &sea_orm::DatabaseConnection) -> Vec<TableInfo> {
let mut tables = Vec::new();
let table_names: Vec<String> = db
.query_all(Statement::from_string(
DbBackend::Sqlite,
"SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'",
))
.await
.unwrap_or_default()
.iter()
.filter_map(|row| row.try_get_by_index::<String>(0).ok())
.collect();
for table_name in table_names {
let columns = discover_sqlite_columns(db, &table_name).await;
tables.push(TableInfo {
name: table_name,
columns,
});
}
tables
}
async fn discover_sqlite_columns(
db: &sea_orm::DatabaseConnection,
table_name: &str,
) -> Vec<ColumnInfo> {
let query = format!("PRAGMA table_info({table_name})");
let rows = db
.query_all(Statement::from_string(DbBackend::Sqlite, query))
.await
.unwrap_or_default();
rows.iter()
.filter_map(|row| {
let name: String = row.try_get_by_index(1).ok()?;
let col_type: String = row.try_get_by_index(2).ok()?;
let notnull: i32 = row.try_get_by_index(3).ok()?;
let pk: i32 = row.try_get_by_index(5).ok()?;
Some(ColumnInfo {
name,
col_type,
is_nullable: notnull == 0,
is_primary_key: pk > 0,
})
})
.collect()
}
async fn discover_postgres_tables(db: &sea_orm::DatabaseConnection) -> Vec<TableInfo> {
let mut tables = Vec::new();
let table_names: Vec<String> = db
.query_all(Statement::from_string(
DbBackend::Postgres,
"SELECT table_name FROM information_schema.tables WHERE table_schema = 'public' AND table_type = 'BASE TABLE'",
))
.await
.unwrap_or_default()
.iter()
.filter_map(|row| row.try_get_by_index::<String>(0).ok())
.collect();
for table_name in table_names {
let columns = discover_postgres_columns(db, &table_name).await;
tables.push(TableInfo {
name: table_name,
columns,
});
}
tables
}
async fn discover_postgres_columns(
db: &sea_orm::DatabaseConnection,
table_name: &str,
) -> Vec<ColumnInfo> {
let query = format!(
r#"
SELECT
c.column_name,
c.data_type,
c.is_nullable,
CASE WHEN pk.column_name IS NOT NULL THEN true ELSE false END as is_pk
FROM information_schema.columns c
LEFT JOIN (
SELECT ku.column_name
FROM information_schema.table_constraints tc
JOIN information_schema.key_column_usage ku
ON tc.constraint_name = ku.constraint_name
WHERE tc.constraint_type = 'PRIMARY KEY'
AND tc.table_name = '{table_name}'
) pk ON c.column_name = pk.column_name
WHERE c.table_name = '{table_name}'
ORDER BY c.ordinal_position
"#
);
let rows = db
.query_all(Statement::from_string(DbBackend::Postgres, query))
.await
.unwrap_or_default();
rows.iter()
.filter_map(|row| {
let name: String = row.try_get_by_index(0).ok()?;
let col_type: String = row.try_get_by_index(1).ok()?;
let is_nullable_str: String = row.try_get_by_index(2).ok()?;
let is_pk: bool = row.try_get_by_index(3).ok()?;
Some(ColumnInfo {
name,
col_type,
is_nullable: is_nullable_str == "YES",
is_primary_key: is_pk,
})
})
.collect()
}
fn generate_entity_file(table: &TableInfo, entities_dir: &Path) {
let entity_file = entities_dir.join(format!("{}.rs", table.name));
let content = templates::entity_template(&table.name, &table.columns);
fs::write(&entity_file, content).expect("Failed to write entity file");
println!(
"{} Generated src/models/entities/{}.rs",
style("✓").green(),
table.name
);
}
fn generate_user_file_if_not_exists(table: &TableInfo, models_dir: &Path) {
let user_file = models_dir.join(format!("{}.rs", table.name));
if user_file.exists() {
println!(
"{} Skipped src/models/{}.rs (already exists)",
style("•").dim(),
table.name
);
return;
}
let struct_name = to_pascal_case(&singularize(&table.name));
let content = templates::user_model_template(&table.name, &struct_name, &table.columns);
fs::write(&user_file, content).expect("Failed to write user model file");
println!(
"{} Created src/models/{}.rs",
style("✓").green(),
table.name
);
}
fn generate_user_file(table: &TableInfo, models_dir: &Path) {
let user_file = models_dir.join(format!("{}.rs", table.name));
let struct_name = to_pascal_case(&singularize(&table.name));
let content = templates::user_model_template(&table.name, &struct_name, &table.columns);
fs::write(&user_file, content).expect("Failed to write user model file");
println!(
"{} Regenerated src/models/{}.rs",
style("✓").green(),
table.name
);
}
fn update_entities_mod(tables: &[TableInfo], entities_dir: &Path) {
let mod_file = entities_dir.join("mod.rs");
let content = templates::entities_mod_template(tables);
fs::write(&mod_file, content).expect("Failed to write entities/mod.rs");
println!("{} Updated src/models/entities/mod.rs", style("✓").green());
}
fn update_models_mod(tables: &[TableInfo], models_dir: &Path) {
let mod_file = models_dir.join("mod.rs");
let existing_content = if mod_file.exists() {
fs::read_to_string(&mod_file).unwrap_or_default()
} else {
"//! Application models\n\n".to_string()
};
let mut lines: Vec<String> = existing_content.lines().map(String::from).collect();
let has_entities_mod = lines.iter().any(|l| {
let trimmed = l.trim();
trimmed == "pub mod entities;" || trimmed == "mod entities;"
});
let mut insert_idx = 0;
for (i, line) in lines.iter().enumerate() {
if line.starts_with("//!") || line.is_empty() {
insert_idx = i + 1;
} else {
break;
}
}
if !has_entities_mod {
lines.insert(insert_idx, "pub mod entities;".to_string());
insert_idx += 1;
}
for table in tables {
let mod_decl = format!("pub mod {};", table.name);
let alt_mod_decl = format!("mod {};", table.name);
if !lines
.iter()
.any(|l| l.trim() == mod_decl || l.trim() == alt_mod_decl)
{
let mut last_mod_idx = insert_idx;
for (i, line) in lines.iter().enumerate() {
if line.trim().starts_with("pub mod ") || line.trim().starts_with("mod ") {
last_mod_idx = i + 1;
}
}
lines.insert(last_mod_idx, mod_decl);
}
}
let content = lines.join("\n");
fs::write(&mod_file, content).expect("Failed to write models/mod.rs");
println!("{} Updated src/models/mod.rs", style("✓").green());
}
fn to_pascal_case(s: &str) -> String {
let mut result = String::new();
let mut capitalize_next = true;
for c in s.chars() {
if c == '_' || c == '-' || c == ' ' {
capitalize_next = true;
} else if capitalize_next {
result.push(c.to_uppercase().next().unwrap());
capitalize_next = false;
} else {
result.push(c);
}
}
result
}
fn singularize(word: &str) -> String {
if let Some(stem) = word.strip_suffix("ies") {
format!("{stem}y")
} else if let Some(stem) = word.strip_suffix("es") {
if word.ends_with("ses") || word.ends_with("xes") {
word.to_string()
} else {
stem.to_string()
}
} else if let Some(stem) = word.strip_suffix('s') {
if word.ends_with("ss") || word.ends_with("us") {
word.to_string()
} else {
stem.to_string()
}
} else {
word.to_string()
}
}