CLI command to generate a new migration

sea-orm-cli/Cargo.toml

@@ -45,6 +45,8 @@ sqlx = { version = "^0.5", default-features = false, features = [ "mysql", "post
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 tracing = { version = "0.1" }
 url = "^2.2"
+chrono = "0.4"
+regex = "1"
 
 [dev-dependencies]
 smol = "1.2.5"

sea-orm-cli/src/cli.rs

@@ -94,6 +94,17 @@ pub fn build_cli() -> App<'static, 'static> {
                 .about("Initialize migration directory")
                 .arg(arg_migration_dir.clone()),
         )
+        .subcommand(
+            SubCommand::with_name("generate")
+                .about("Generate a new, empty migration")
+                .arg(
+                    Arg::with_name("MIGRATION_NAME")
+                        .help("Name of the new migation")
+                        .required(true)
+                        .takes_value(true),
+                )
+                .arg(arg_migration_dir.clone()),
+        )
         .arg(arg_migration_dir.clone());
     for subcommand in sea_schema::migration::get_subcommands() {
         migrate_subcommands =

sea-orm-cli/src/commands.rs

@@ -1,4 +1,6 @@
+use chrono::Local;
 use clap::ArgMatches;
+use regex::Regex;
 use sea_orm_codegen::{EntityTransformer, OutputFile, WithSerde};
 use std::{error::Error, fmt::Display, fs, io::Write, path::Path, process::Command, str::FromStr};
 use url::Url;
@@ -220,6 +222,22 @@ pub fn run_migrate_command(matches: &ArgMatches<'_>) -> Result<(), Box<dyn Error
         println!("Done!");
         // Early exit!
         return Ok(());
+    } else if let ("generate", Some(args)) = migrate_subcommand {
+        let migration_dir = args.value_of("MIGRATION_DIR").unwrap();
+        let migration_name = args.value_of("MIGRATION_NAME").unwrap();
+        println!("Generating new migration...");
+
+        // build new migration filename
+        let now = Local::now();
+        let migration_name = format!(
+            "m{}_{}",
+            now.format("%Y%m%d_%H%M%S").to_string(),
+            migration_name
+        );
+
+        create_new_migration(&migration_name, migration_dir)?;
+        update_migrator(&migration_name, migration_dir)?;
+        return Ok(());
     }
     let (subcommand, migration_dir, steps, verbose) = match migrate_subcommand {
         // Catch all command with pattern `migrate xxx`
@@ -262,6 +280,63 @@ pub fn run_migrate_command(matches: &ArgMatches<'_>) -> Result<(), Box<dyn Error
     Ok(())
 }
 
+fn create_new_migration(migration_name: &str, migration_dir: &str) -> Result<(), Box<dyn Error>> {
+    let migration_filepath = Path::new(migration_dir)
+        .join("src")
+        .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 migration_content =
+        migration_template.replace("m20220101_000001_create_table", &migration_name);
+    let mut migration_file = fs::File::create(migration_filepath)?;
+    migration_file.write_all(migration_content.as_bytes())?;
+    Ok(())
+}
+
+fn update_migrator(migration_name: &str, migration_dir: &str) -> Result<(), Box<dyn Error>> {
+    let migrator_filepath = Path::new(migration_dir).join("src").join("lib.rs");
+    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.clone().with_file_name("lib.rs.bkp");
+    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+(?P<name>m\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::<Vec<String>>()
+        .join("\n");
+    boxed_migrations.push_str("\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(())
+}
+
 pub fn handle_error<E>(error: E)
 where
     E: Display,
@@ -371,4 +446,58 @@ mod tests {
 
         smol::block_on(run_generate_command(matches.subcommand().1.unwrap())).unwrap();
     }
+    #[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();
+        let migration_content =
+            migration_content.replace(&migration_name, "m20220101_000001_create_table");
+        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<name>\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<name>\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();
+    }
 }