use chrono::{Local, Utc}; use regex::Regex; use std::{ error::Error, fmt::Display, fs, io::Write, path::{Path, PathBuf}, process::Command, }; #[cfg(feature = "cli")] use crate::MigrateSubcommands; #[cfg(feature = "cli")] pub fn run_migrate_command( command: Option, migration_dir: &str, database_schema: Option, database_url: Option, verbose: bool, ) -> Result<(), Box> { match command { Some(MigrateSubcommands::Init) => run_migrate_init(migration_dir)?, Some(MigrateSubcommands::Generate { migration_name, universal_time, }) => run_migrate_generate(migration_dir, &migration_name, universal_time)?, _ => { let (subcommand, migration_dir, steps, verbose) = match command { Some(MigrateSubcommands::Fresh) => ("fresh", migration_dir, None, verbose), Some(MigrateSubcommands::Refresh) => ("refresh", migration_dir, None, verbose), Some(MigrateSubcommands::Reset) => ("reset", migration_dir, None, verbose), Some(MigrateSubcommands::Status) => ("status", migration_dir, None, verbose), Some(MigrateSubcommands::Up { num }) => ("up", migration_dir, num, verbose), Some(MigrateSubcommands::Down { num }) => { ("down", migration_dir, Some(num), verbose) } _ => ("up", migration_dir, None, verbose), }; // Construct the `--manifest-path` let manifest_path = if migration_dir.ends_with('/') { format!("{}Cargo.toml", migration_dir) } else { format!("{}/Cargo.toml", migration_dir) }; // Construct the arguments that will be supplied to `cargo` command let mut args = vec!["run", "--manifest-path", &manifest_path, "--", subcommand]; let mut num: String = "".to_string(); if let Some(steps) = steps { num = steps.to_string(); } if !num.is_empty() { args.extend(["-n", &num]) } if let Some(database_url) = &database_url { args.extend(["-u", database_url]); } if let Some(database_schema) = &database_schema { args.extend(["-s", database_schema]); } if verbose { args.push("-v"); } // Run migrator CLI on user's behalf println!("Running `cargo {}`", args.join(" ")); Command::new("cargo").args(args).spawn()?.wait()?; } } Ok(()) } pub fn run_migrate_init(migration_dir: &str) -> Result<(), Box> { let migration_dir = match migration_dir.ends_with('/') { true => migration_dir.to_string(), false => format!("{}/", migration_dir), }; println!("Initializing migration directory..."); macro_rules! write_file { ($filename: literal) => { let fn_content = |content: String| content; write_file!($filename, $filename, fn_content); }; ($filename: literal, $template: literal, $fn_content: expr) => { let filepath = [&migration_dir, $filename].join(""); println!("Creating file `{}`", filepath); let path = Path::new(&filepath); let prefix = path.parent().unwrap(); fs::create_dir_all(prefix).unwrap(); let mut file = fs::File::create(path)?; let content = include_str!(concat!("../../template/migration/", $template)); let content = $fn_content(content.to_string()); file.write_all(content.as_bytes())?; }; } write_file!("src/lib.rs"); write_file!("src/m20220101_000001_create_table.rs"); write_file!("src/main.rs"); write_file!("Cargo.toml", "_Cargo.toml", |content: String| { let ver = format!( "{}.{}.0", env!("CARGO_PKG_VERSION_MAJOR"), env!("CARGO_PKG_VERSION_MINOR") ); content.replace("", &ver) }); write_file!("README.md"); println!("Done!"); Ok(()) } pub fn run_migrate_generate( migration_dir: &str, migration_name: &str, universal_time: bool, ) -> Result<(), Box> { // Make sure the migration name doesn't contain any characters that // are invalid module names in Rust. if migration_name.contains('-') { return Err(Box::new(MigrationCommandError::InvalidName( "Hyphen `-` cannot be used in migration name".to_string(), ))); } println!("Generating new migration..."); // build new migration filename const FMT: &str = "%Y%m%d_%H%M%S"; let formatted_now = if universal_time { Utc::now().format(FMT) } else { Local::now().format(FMT) }; let migration_name = format!("m{}_{}", formatted_now, migration_name); create_new_migration(&migration_name, migration_dir)?; update_migrator(&migration_name, migration_dir)?; Ok(()) } /// `get_full_migration_dir` looks for a `src` directory /// inside of `migration_dir` and appends that to the returned path if found. /// /// Otherwise, `migration_dir` can point directly to a directory containing the /// migrations. In that case, nothing is appended. /// /// This way, `src` doesn't need to be appended in the standard case where /// migrations are in their own crate. If the migrations are in a submodule /// of another crate, `migration_dir` can point directly to that module. fn get_full_migration_dir(migration_dir: &str) -> PathBuf { let without_src = Path::new(migration_dir).to_owned(); let with_src = without_src.join("src"); match () { _ if with_src.is_dir() => with_src, _ => without_src, } } fn create_new_migration(migration_name: &str, migration_dir: &str) -> Result<(), Box> { let migration_filepath = get_full_migration_dir(migration_dir).join(format!("{}.rs", &migration_name)); println!("Creating migration file `{}`", migration_filepath.display()); // TODO: make OS agnostic let migration_template = include_str!("../../template/migration/src/m20220101_000001_create_table.rs"); let mut migration_file = fs::File::create(migration_filepath)?; migration_file.write_all(migration_template.as_bytes())?; Ok(()) } /// `get_migrator_filepath` looks for a file `migration_dir/src/lib.rs` /// and returns that path if found. /// /// If `src` is not found, it will look directly in `migration_dir` for `lib.rs`. /// /// If `lib.rs` is not found, it will look for `mod.rs` instead, /// e.g. `migration_dir/mod.rs`. /// /// This way, `src` doesn't need to be appended in the standard case where /// migrations are in their own crate (with a file `lib.rs`). If the /// migrations are in a submodule of another crate (with a file `mod.rs`), /// `migration_dir` can point directly to that module. fn get_migrator_filepath(migration_dir: &str) -> PathBuf { let full_migration_dir = get_full_migration_dir(migration_dir); let with_lib = full_migration_dir.join("lib.rs"); match () { _ if with_lib.is_file() => with_lib, _ => full_migration_dir.join("mod.rs"), } } fn update_migrator(migration_name: &str, migration_dir: &str) -> Result<(), Box> { let migrator_filepath = get_migrator_filepath(migration_dir); println!( "Adding migration `{}` to `{}`", migration_name, migrator_filepath.display() ); let migrator_content = fs::read_to_string(&migrator_filepath)?; let mut updated_migrator_content = migrator_content.clone(); // create a backup of the migrator file in case something goes wrong let migrator_backup_filepath = migrator_filepath.with_extension("rs.bak"); fs::copy(&migrator_filepath, &migrator_backup_filepath)?; let mut migrator_file = fs::File::create(&migrator_filepath)?; // find existing mod declarations, add new line let mod_regex = Regex::new(r"mod\s+(?Pm\d{8}_\d{6}_\w+);")?; let mods: Vec<_> = mod_regex.captures_iter(&migrator_content).collect(); let mods_end = mods.last().unwrap().get(0).unwrap().end() + 1; updated_migrator_content.insert_str(mods_end, format!("mod {};\n", migration_name).as_str()); // build new vector from declared migration modules let mut migrations: Vec<&str> = mods .iter() .map(|cap| cap.name("name").unwrap().as_str()) .collect(); migrations.push(migration_name); let mut boxed_migrations = migrations .iter() .map(|migration| format!(" Box::new({}::Migration),", migration)) .collect::>() .join("\n"); boxed_migrations.push('\n'); let boxed_migrations = format!("vec![\n{} ]\n", boxed_migrations); let vec_regex = Regex::new(r"vec!\[[\s\S]+\]\n")?; let updated_migrator_content = vec_regex.replace(&updated_migrator_content, &boxed_migrations); migrator_file.write_all(updated_migrator_content.as_bytes())?; fs::remove_file(&migrator_backup_filepath)?; Ok(()) } #[derive(Debug)] enum MigrationCommandError { InvalidName(String), } impl Display for MigrationCommandError { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { MigrationCommandError::InvalidName(name) => { write!(f, "Invalid migration name: {}", name) } } } } impl Error for MigrationCommandError {} #[cfg(test)] mod tests { use super::*; #[test] fn test_create_new_migration() { let migration_name = "test_name"; let migration_dir = "/tmp/sea_orm_cli_test_new_migration/"; fs::create_dir_all(format!("{}src", migration_dir)).unwrap(); create_new_migration(migration_name, migration_dir).unwrap(); let migration_filepath = Path::new(migration_dir) .join("src") .join(format!("{}.rs", migration_name)); assert!(migration_filepath.exists()); let migration_content = fs::read_to_string(migration_filepath).unwrap(); assert_eq!( &migration_content, include_str!("../../template/migration/src/m20220101_000001_create_table.rs") ); fs::remove_dir_all("/tmp/sea_orm_cli_test_new_migration/").unwrap(); } #[test] fn test_update_migrator() { let migration_name = "test_name"; let migration_dir = "/tmp/sea_orm_cli_test_update_migrator/"; fs::create_dir_all(format!("{}src", migration_dir)).unwrap(); let migrator_filepath = Path::new(migration_dir).join("src").join("lib.rs"); fs::copy("./template/migration/src/lib.rs", &migrator_filepath).unwrap(); update_migrator(migration_name, migration_dir).unwrap(); assert!(&migrator_filepath.exists()); let migrator_content = fs::read_to_string(&migrator_filepath).unwrap(); let mod_regex = Regex::new(r"mod (?P\w+);").unwrap(); let migrations: Vec<&str> = mod_regex .captures_iter(&migrator_content) .map(|cap| cap.name("name").unwrap().as_str()) .collect(); assert_eq!(migrations.len(), 2); assert_eq!( *migrations.first().unwrap(), "m20220101_000001_create_table" ); assert_eq!(migrations.last().unwrap(), &migration_name); let boxed_regex = Regex::new(r"Box::new\((?P\S+)::Migration\)").unwrap(); let migrations: Vec<&str> = boxed_regex .captures_iter(&migrator_content) .map(|cap| cap.name("name").unwrap().as_str()) .collect(); assert_eq!(migrations.len(), 2); assert_eq!( *migrations.first().unwrap(), "m20220101_000001_create_table" ); assert_eq!(migrations.last().unwrap(), &migration_name); fs::remove_dir_all("/tmp/sea_orm_cli_test_update_migrator/").unwrap(); } }