use std::path::{Path, PathBuf};
use clap::Args;
use sz_rust_core::orm::migration::{
FileMigrationResolver, Migration, MigrationContext, MigrationResolver, Migrator,
};
use sz_rust_core::orm::{Connection, ConnectionFactory, DbType};
use crate::error::CliError;
#[derive(Args, Debug)]
pub struct MigrateArgs {
#[arg(long)]
pub rollback: bool,
#[arg(short = 'p', long, default_value = "migrations")]
pub path: String,
#[arg(long, default_value = "postgres")]
pub db_type: String,
#[arg(long)]
pub show_sql: bool,
#[arg(long)]
pub url: Option<String>,
}
pub fn execute_migrate(args: &MigrateArgs) -> Result<(), CliError> {
let path = PathBuf::from(&args.path);
if !path.exists() {
return Err(CliError::Migration(format!(
"Migration directory not found: {}",
path.display()
)));
}
let db_type = DbType::from_str(&args.db_type)
.ok_or_else(|| CliError::Migration(format!("Unknown database type: {}", args.db_type)))?;
let migrations = resolve_migrations(&path, db_type)?;
if migrations.is_empty() {
println!("No migrations found in: {}", path.display());
return Ok(());
}
match &args.url {
None => execute_migrate_offline(args, &migrations),
Some(url) => execute_migrate_online(args, &migrations, url, db_type),
}
}
fn execute_migrate_offline(args: &MigrateArgs, migrations: &[Migration]) -> Result<(), CliError> {
if args.rollback {
println!("Rolling back last batch in: {}", args.path);
if let Some(last) = migrations.last() {
println!(" Would rollback: {} ({})", last.version, last.name);
if args.show_sql {
print_sql_block("SQL DOWN", &last.sql_down);
}
}
println!("Note: Actual rollback requires database connection (offline mode).");
} else {
println!("Running migrations in: {}", args.path);
for m in migrations {
println!(" Would apply: {} ({})", m.version, m.name);
if args.show_sql {
print_sql_block("SQL UP", &m.sql_up);
}
}
println!(
"Total: {} migration(s). Note: Actual execution requires database connection (offline mode).",
migrations.len()
);
}
Ok(())
}
fn execute_migrate_online(
args: &MigrateArgs,
migrations: &[Migration],
url: &str,
db_type: DbType,
) -> Result<(), CliError> {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.map_err(|e| CliError::Migration(format!("Failed to create tokio runtime: {}", e)))?;
rt.block_on(async move {
let mut conn = create_connection(url, db_type).await?;
if args.rollback {
let last = migrations
.last()
.ok_or_else(|| CliError::Migration("No migrations to rollback".to_string()))?;
println!("Rolling back: {} ({})", last.version, last.name);
if args.show_sql {
print_sql_block("SQL DOWN", &last.sql_down);
}
if !last.sql_down.is_empty() {
conn.execute(&last.sql_down)
.await
.map_err(|e| CliError::Migration(format!("Rollback failed: {}", e)))?;
}
delete_migration_record(&mut conn, &last.version, db_type).await?;
println!("Rollback completed: {} ({})", last.version, last.name);
} else {
ensure_migrations_table(&mut conn, db_type).await?;
let applied = fetch_applied_versions(&mut conn, db_type).await?;
let pending: Vec<&Migration> = migrations
.iter()
.filter(|m| !applied.contains(&m.version))
.collect();
if pending.is_empty() {
println!("No pending migrations. Database is up to date.");
return Ok(());
}
println!("Running {} pending migration(s):", pending.len());
let mut context = MigrationContext::default().with_db_type(db_type);
context.connection = Some(conn);
let mut migrator = Migrator::new(context);
for m in migrations {
let rebuilt = Migration::new(&m.version, &m.name, &m.sql_up, &m.sql_down);
if applied.contains(&m.version) {
migrator = migrator.add_migration(rebuilt.with_batch(1));
} else {
migrator = migrator.add_migration(rebuilt);
}
}
let applied_versions = migrator
.migrate()
.await
.map_err(|e| CliError::Migration(format!("Migration failed: {}", e)))?;
for v in &applied_versions {
println!(" Applied: {}", v);
}
println!("Migration completed: {} applied.", applied_versions.len());
}
Ok::<(), CliError>(())
})
}
pub fn execute_status(path: &str) -> Result<(), CliError> {
execute_status_full(path, "postgres", false, None)
}
pub fn execute_status_with(path: &str, db_type_str: &str, show_sql: bool) -> Result<(), CliError> {
execute_status_full(path, db_type_str, show_sql, None)
}
pub fn execute_status_full(
path: &str,
db_type_str: &str,
show_sql: bool,
url: Option<&str>,
) -> Result<(), CliError> {
let path_buf = PathBuf::from(path);
if !path_buf.exists() {
return Err(CliError::Migration(format!(
"Migration directory not found: {}",
path_buf.display()
)));
}
let db_type = DbType::from_str(db_type_str)
.ok_or_else(|| CliError::Migration(format!("Unknown database type: {}", db_type_str)))?;
let migrations = resolve_migrations(&path_buf, db_type)?;
if migrations.is_empty() {
println!("No migrations found in: {}", path_buf.display());
return Ok(());
}
let applied_versions = if let Some(url) = url {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.map_err(|e| CliError::Migration(format!("Failed to create tokio runtime: {}", e)))?;
rt.block_on(async move {
let mut conn = create_connection(url, db_type).await?;
ensure_migrations_table(&mut conn, db_type).await?;
fetch_applied_versions(&mut conn, db_type).await
})?
} else {
std::collections::HashSet::new()
};
println!(
"{:<15} {:<30} {:<20}",
"Version", "Migration Name", "Status"
);
println!("{}", "-".repeat(65));
for m in &migrations {
let status = if applied_versions.contains(&m.version) {
"Applied"
} else if url.is_some() {
"Pending"
} else {
"Pending*"
};
println!("{:<15} {:<30} {:<20}", m.version, m.name, status);
if show_sql {
print_sql_block("SQL UP", &m.sql_up);
print_sql_block("SQL DOWN", &m.sql_down);
}
}
println!();
if url.is_some() {
let applied = migrations
.iter()
.filter(|m| applied_versions.contains(&m.version))
.count();
println!(
"Total: {} migration(s), {} applied, {} pending.",
migrations.len(),
applied,
migrations.len() - applied
);
} else {
println!("* Status cannot be determined without database connection (offline mode).");
}
Ok(())
}
async fn create_connection(url: &str, db_type: DbType) -> Result<Box<dyn Connection>, CliError> {
use std::sync::Arc;
use sz_orm_sqlx::{
MySqlPoolHandle, PgPoolHandle, SqlitePoolHandle, SqlxMySqlConnectionFactory,
SqlxPgConnectionFactory, SqlxSqliteConnectionFactory,
};
match db_type {
DbType::PostgreSQL => {
let pool = PgPoolHandle::connect(url).await.map_err(|e| {
CliError::Migration(format!("PostgreSQL connect failed: {}", e))
})?;
let factory = SqlxPgConnectionFactory::new(Arc::new(pool));
let conn = factory.create().await.map_err(|e| {
CliError::Migration(format!("PostgreSQL acquire failed: {}", e))
})?;
Ok(conn)
}
DbType::MySQL => {
let pool = MySqlPoolHandle::connect(url).await.map_err(|e| {
CliError::Migration(format!("MySQL connect failed: {}", e))
})?;
let factory = SqlxMySqlConnectionFactory::new(Arc::new(pool));
let conn = factory.create().await.map_err(|e| {
CliError::Migration(format!("MySQL acquire failed: {}", e))
})?;
Ok(conn)
}
DbType::Sqlite => {
let pool = SqlitePoolHandle::connect(url).await.map_err(|e| {
CliError::Migration(format!("SQLite connect failed: {}", e))
})?;
let factory = SqlxSqliteConnectionFactory::new(Arc::new(pool));
let conn = factory.create().await.map_err(|e| {
CliError::Migration(format!("SQLite acquire failed: {}", e))
})?;
Ok(conn)
}
_ => Err(CliError::Migration(format!(
"Online migration not supported for db_type {:?}. Supported: PostgreSQL, MySQL, SQLite.",
db_type
))),
}
}
async fn ensure_migrations_table(
conn: &mut Box<dyn Connection>,
db_type: DbType,
) -> Result<(), CliError> {
let sql = match db_type {
DbType::PostgreSQL | DbType::Sqlite => {
"CREATE TABLE IF NOT EXISTS __migrations (
version VARCHAR(255) PRIMARY KEY,
name VARCHAR(255) NOT NULL,
batch INTEGER NOT NULL,
run_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
)"
}
DbType::MySQL => {
"CREATE TABLE IF NOT EXISTS __migrations (
version VARCHAR(255) PRIMARY KEY,
name VARCHAR(255) NOT NULL,
batch INT NOT NULL,
run_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
)"
}
_ => {
return Err(CliError::Migration(format!(
"Cannot ensure __migrations table for db_type {:?}",
db_type
)))
}
};
conn.execute(sql)
.await
.map_err(|e| CliError::Migration(format!("Failed to create __migrations table: {}", e)))?;
Ok(())
}
async fn fetch_applied_versions(
conn: &mut Box<dyn Connection>,
_db_type: DbType,
) -> Result<std::collections::HashSet<String>, CliError> {
let rows = conn
.query("SELECT version FROM __migrations")
.await
.map_err(|e| CliError::Migration(format!("Failed to query __migrations: {}", e)))?;
let mut versions = std::collections::HashSet::new();
for row in &rows {
if let Some(val) = row.get("version") {
use sz_rust_core::orm::Value;
match val {
Value::String(s) => versions.insert(s.clone()),
Value::I64(i) => versions.insert(i.to_string()),
Value::I32(i) => versions.insert(i.to_string()),
_ => false,
};
}
}
Ok(versions)
}
async fn delete_migration_record(
conn: &mut Box<dyn Connection>,
version: &str,
db_type: DbType,
) -> Result<(), CliError> {
let sql = match db_type {
DbType::PostgreSQL | DbType::Sqlite => {
format!("DELETE FROM __migrations WHERE version = '{}'", version)
}
DbType::MySQL => {
format!("DELETE FROM __migrations WHERE version = '{}'", version)
}
_ => return Ok(()),
};
conn.execute(&sql)
.await
.map_err(|e| CliError::Migration(format!("Failed to delete migration record: {}", e)))?;
Ok(())
}
fn resolve_migrations(path: &Path, db_type: DbType) -> Result<Vec<Migration>, CliError> {
let resolver = FileMigrationResolver::new(path.to_path_buf());
resolver
.resolve(db_type)
.map_err(|e| CliError::Migration(format!("Failed to resolve migrations: {}", e)))
}
fn print_sql_block(title: &str, sql: &str) {
if sql.is_empty() {
return;
}
println!(" --- {} ---", title);
for line in sql.lines() {
println!(" {}", line);
}
println!(" {}", "-".repeat(title.len() + 8));
}
#[cfg(test)]
mod tests {
use super::*;
use std::fs;
use std::io::Write;
fn create_test_migration(dir: &Path, version: &str, name: &str) {
let up_name = format!("{}_{}_up.sql", version, name);
let down_name = format!("{}_{}_down.sql", version, name);
let up_path = dir.join(up_name);
let down_path = dir.join(down_name);
let mut up_file = fs::File::create(&up_path).unwrap();
writeln!(up_file, "-- {} up", name).unwrap();
let mut down_file = fs::File::create(&down_path).unwrap();
writeln!(down_file, "-- {} down", name).unwrap();
}
#[test]
fn test_resolve_migrations_empty() {
let temp = tempfile::tempdir().unwrap();
let path = temp.path().to_path_buf();
let result = resolve_migrations(&path, DbType::PostgreSQL).unwrap();
assert!(result.is_empty());
}
#[test]
fn test_resolve_migrations_with_files() {
let temp = tempfile::tempdir().unwrap();
let path = temp.path().to_path_buf();
create_test_migration(&path, "001", "create_users");
create_test_migration(&path, "002", "add_index");
let result = resolve_migrations(&path, DbType::PostgreSQL).unwrap();
assert_eq!(result.len(), 2);
assert_eq!(result[0].version, "001");
assert_eq!(result[0].name, "create_users");
assert_eq!(result[1].version, "002");
assert_eq!(result[1].name, "add_index");
}
#[test]
fn test_resolve_migrations_returns_sql_content() {
let temp = tempfile::tempdir().unwrap();
let path = temp.path().to_path_buf();
let up_path = path.join("001_init_up.sql");
let down_path = path.join("001_init_down.sql");
fs::write(&up_path, "CREATE TABLE users (id INT);").unwrap();
fs::write(&down_path, "DROP TABLE users;").unwrap();
let result = resolve_migrations(&path, DbType::PostgreSQL).unwrap();
assert_eq!(result.len(), 1);
assert!(result[0].sql_up.contains("CREATE TABLE users"));
assert!(result[0].sql_down.contains("DROP TABLE users"));
}
#[test]
fn test_resolve_migrations_supports_multiple_db_types() {
let temp = tempfile::tempdir().unwrap();
let path = temp.path().to_path_buf();
create_test_migration(&path, "001", "init");
let mysql_result = resolve_migrations(&path, DbType::MySQL).unwrap();
let pg_result = resolve_migrations(&path, DbType::PostgreSQL).unwrap();
assert_eq!(mysql_result.len(), 1);
assert_eq!(pg_result.len(), 1);
}
#[test]
fn test_execute_status_nonexistent_dir() {
let result = execute_status("/nonexistent/path/migrations");
assert!(matches!(result, Err(CliError::Migration(_))));
}
#[test]
fn test_execute_status_empty_dir() {
let temp = tempfile::tempdir().unwrap();
let path = temp.path().to_str().unwrap();
let result = execute_status(path);
assert!(result.is_ok());
}
#[test]
fn test_execute_status_with_migrations() {
let temp = tempfile::tempdir().unwrap();
let path = temp.path().to_path_buf();
create_test_migration(&path, "001", "create_users");
let path_str = temp.path().to_str().unwrap();
let result = execute_status(path_str);
assert!(result.is_ok());
}
#[test]
fn test_execute_status_with_invalid_db_type() {
let temp = tempfile::tempdir().unwrap();
let path = temp.path().to_str().unwrap();
let result = execute_status_with(path, "invalid_db_type", false);
assert!(matches!(result, Err(CliError::Migration(_))));
}
#[test]
fn test_execute_status_with_show_sql() {
let temp = tempfile::tempdir().unwrap();
let path = temp.path().to_path_buf();
let up_path = path.join("001_init_up.sql");
let down_path = path.join("001_init_down.sql");
fs::write(&up_path, "CREATE TABLE users (id INT);").unwrap();
fs::write(&down_path, "DROP TABLE users;").unwrap();
let path_str = temp.path().to_str().unwrap();
let result = execute_status_with(path_str, "postgres", true);
assert!(result.is_ok());
}
#[test]
fn test_execute_migrate_nonexistent_dir() {
let args = MigrateArgs {
rollback: false,
path: "/nonexistent/migrations".to_string(),
db_type: "postgres".to_string(),
show_sql: false,
url: None,
};
let result = execute_migrate(&args);
assert!(matches!(result, Err(CliError::Migration(_))));
}
#[test]
fn test_execute_migrate_empty_dir() {
let temp = tempfile::tempdir().unwrap();
let args = MigrateArgs {
rollback: false,
path: temp.path().to_str().unwrap().to_string(),
db_type: "postgres".to_string(),
show_sql: false,
url: None,
};
let result = execute_migrate(&args);
assert!(result.is_ok());
}
#[test]
fn test_execute_migrate_with_files_offline() {
let temp = tempfile::tempdir().unwrap();
let path = temp.path().to_path_buf();
create_test_migration(&path, "001", "create_users");
let args = MigrateArgs {
rollback: false,
path: temp.path().to_str().unwrap().to_string(),
db_type: "postgres".to_string(),
show_sql: false,
url: None,
};
let result = execute_migrate(&args);
assert!(result.is_ok());
}
#[test]
fn test_execute_migrate_with_show_sql_offline() {
let temp = tempfile::tempdir().unwrap();
let path = temp.path().to_path_buf();
let up_path = path.join("001_init_up.sql");
let down_path = path.join("001_init_down.sql");
fs::write(&up_path, "CREATE TABLE users (id INT);").unwrap();
fs::write(&down_path, "DROP TABLE users;").unwrap();
let args = MigrateArgs {
rollback: false,
path: temp.path().to_str().unwrap().to_string(),
db_type: "postgres".to_string(),
show_sql: true,
url: None,
};
let result = execute_migrate(&args);
assert!(result.is_ok());
}
#[test]
fn test_execute_migrate_with_invalid_db_type() {
let temp = tempfile::tempdir().unwrap();
let args = MigrateArgs {
rollback: false,
path: temp.path().to_str().unwrap().to_string(),
db_type: "invalid_db_type".to_string(),
show_sql: false,
url: None,
};
let result = execute_migrate(&args);
assert!(matches!(result, Err(CliError::Migration(_))));
}
#[test]
fn test_execute_migrate_rollback_offline() {
let temp = tempfile::tempdir().unwrap();
let path = temp.path().to_path_buf();
create_test_migration(&path, "001", "create_users");
create_test_migration(&path, "002", "add_index");
let args = MigrateArgs {
rollback: true,
path: temp.path().to_str().unwrap().to_string(),
db_type: "postgres".to_string(),
show_sql: false,
url: None,
};
let result = execute_migrate(&args);
assert!(result.is_ok());
}
#[test]
fn test_execute_migrate_rollback_with_show_sql_offline() {
let temp = tempfile::tempdir().unwrap();
let path = temp.path().to_path_buf();
let up_path = path.join("001_init_up.sql");
let down_path = path.join("001_init_down.sql");
fs::write(&up_path, "CREATE TABLE users (id INT);").unwrap();
fs::write(&down_path, "DROP TABLE users;").unwrap();
let args = MigrateArgs {
rollback: true,
path: temp.path().to_str().unwrap().to_string(),
db_type: "postgres".to_string(),
show_sql: true,
url: None,
};
let result = execute_migrate(&args);
assert!(result.is_ok());
}
#[test]
fn test_print_sql_block_empty_sql() {
print_sql_block("SQL UP", "");
}
#[test]
fn test_print_sql_block_with_content() {
print_sql_block("SQL UP", "CREATE TABLE users (id INT);");
}
#[test]
fn test_execute_status_full_offline_no_url() {
let temp = tempfile::tempdir().unwrap();
let path = temp.path().to_path_buf();
create_test_migration(&path, "001", "init");
let path_str = temp.path().to_str().unwrap();
let result = execute_status_full(path_str, "postgres", false, None);
assert!(result.is_ok());
}
#[test]
fn test_execute_status_full_offline_with_show_sql() {
let temp = tempfile::tempdir().unwrap();
let path = temp.path().to_path_buf();
let up_path = path.join("001_init_up.sql");
let down_path = path.join("001_init_down.sql");
fs::write(&up_path, "CREATE TABLE t (id INT);").unwrap();
fs::write(&down_path, "DROP TABLE t;").unwrap();
let path_str = temp.path().to_str().unwrap();
let result = execute_status_full(path_str, "postgres", true, None);
assert!(result.is_ok());
}
#[test]
fn test_execute_status_full_invalid_db_type() {
let temp = tempfile::tempdir().unwrap();
let path_str = temp.path().to_str().unwrap();
let result = execute_status_full(path_str, "invalid_db", false, None);
assert!(matches!(result, Err(CliError::Migration(_))));
}
#[test]
fn test_execute_migrate_online_with_invalid_url_returns_error() {
let temp = tempfile::tempdir().unwrap();
let path = temp.path().to_path_buf();
create_test_migration(&path, "001", "init");
let args = MigrateArgs {
rollback: false,
path: temp.path().to_str().unwrap().to_string(),
db_type: "postgres".to_string(),
show_sql: false,
url: Some("postgres://invalid:invalid@127.0.0.1:1/invalid".to_string()),
};
let result = execute_migrate(&args);
assert!(result.is_err());
}
}