#[path = "lib.rs"]
mod lib;
use anyhow::{anyhow, Context, Result};
use clap::{Parser, Subcommand};
use lib::*;
use regex::Regex;
use std::fs;
use std::fs::{create_dir_all, File};
use std::io::Write;
use std::path::{Path, PathBuf};
use tokio_postgres::Client;
use tracing::{error, info};
use tracing_subscriber;
#[derive(Parser)]
#[command(
name = "migro",
about = "A migration tool for Rust projects",
long_about = "A lightweight CLI tool to manage database migrations, seeders, and custom SQL queries in Rust projects."
)]
struct Cli {
#[command(subcommand)]
command: Commands,
}
#[derive(Subcommand)]
enum Commands {
Init,
Create {
name: String,
#[arg(short = 'S', long, default_value = "false")]
seeder: bool,
#[arg(short = 'Q', long, default_value = "false")]
query: bool,
#[arg(short = 'M', long, default_value = "false")]
migration: bool,
},
Up {
#[arg(short = 'S', long, default_value = "false")]
seed: bool,
},
Down,
Seed,
}
#[tokio::main]
async fn main() -> Result<()> {
dotenv::dotenv().ok();
tracing_subscriber::fmt::init();
let cli = Cli::parse();
match cli.command {
Commands::Init => {
ensure_database_dir()?;
info!("Database directory created successfully");
Ok(())
},
Commands::Create {
name,
seeder,
query,
migration
} => {
if !migration && !seeder && !query {
error!("You must specify at least one option: --migration, --seeder, or --query");
return Err(anyhow!("You must specify at least one option: --migration, --seeder, or --query"));
}
let migration_dir = ensure_migration_dir(&name)?;
if migration {
create_migration(&migration_dir)?;
info!("Migration file '{}' created", name);
}
if seeder {
create_seeder(&migration_dir)?;
info!("Seeder file '{}' created", name);
}
if query {
create_query(&migration_dir)?;
info!("Query folder '{}' created", name);
}
Ok(())
},
Commands::Up { seed } => {
run_migrations(seed).await?;
info!("Migrations applied successfully (seed: {})", seed);
Ok(())
},
Commands::Down => {
run_migrations_down().await?;
info!("Down migrations applied successfully");
Ok(())
},
Commands::Seed => {
seed_db().await?;
info!("Seeders executed successfully");
Ok(())
},
}
}
pub(crate) fn create_migration(path: &PathBuf) -> Result<()> {
let mut up_file = File::create(path.join("up.sql"))?;
up_file.write_all(b"-- Migration Up\n")?;
let mut down_file = File::create(path.join("down.sql"))?;
down_file.write_all(b"-- Migration Down\n")?;
Ok(())
}
pub(crate) fn create_seeder(path: &PathBuf) -> Result<()> {
let seeder_path = path.join("seeder.sql");
let mut seeder_file = File::create(&seeder_path)?;
seeder_file.write_all(b"-- Seeder\n")?;
Ok(())
}
pub(crate) fn create_query(path: &PathBuf) -> Result<()> {
let query_path = path.join("queries");
if !query_path.exists() {
create_dir_all(&query_path)?;
}
Ok(())
}
async fn run_migrations_down() -> Result<()> {
let client = connect_db().await?;
create_migrations_table(&client).await?;
let mut executed_migrations = get_executed_migrations(&client).await?;
let migrations_dir = ensure_database_dir()?;
let available_migrations = get_available_migrations(&migrations_dir)?;
executed_migrations.reverse();
for migration_name in executed_migrations {
let migration = available_migrations.iter()
.find(|m| m.name == migration_name);
if let Some(migration) = migration {
let down_script = fs::read_to_string(migration.path.join("down.sql"))?;
client.batch_execute(&down_script).await?;
remove_migration(&client, &migration.name).await?;
}
}
Ok(())
}
fn get_next_number(database_dir: &Path) -> Result<u32> {
let re = Regex::new(r"^(\d{5})_")?;
let mut highest_number = 0;
if !database_dir.exists() {
return Ok(1);
}
for entry in fs::read_dir(database_dir)? {
let entry = entry?;
let file_name = entry.file_name();
let file_name_str = file_name.to_string_lossy();
if let Some(captures) = re.captures(&file_name_str) {
if let Some(num_str) = captures.get(1) {
if let Ok(num) = num_str.as_str().parse::<u32>() {
highest_number = highest_number.max(num);
}
}
}
}
Ok(highest_number + 1)
}
async fn seed_db() -> Result<()> {
let conn = connect_db().await?;
let database_dir = PathBuf::from("database");
let migrations = fs::read_dir(&database_dir)?
.filter_map(|entry| {
let entry = entry.ok()?;
let path = entry.path();
if path.is_dir() {
Some(path)
} else {
None
}
})
.collect::<Vec<_>>();
let mut seeders = migrations.iter()
.filter(|path| path.join("seeder.sql").exists())
.map(|path| path.join("seeder.sql"))
.collect::<Vec<_>>();
seeders.sort_by(|a, b| a.parent().unwrap().file_name().cmp(&b.parent().unwrap().file_name()));
for seeder in seeders {
let seeder_script = fs::read_to_string(seeder)?;
conn.batch_execute(&seeder_script).await?;
}
Ok(())
}
async fn remove_migration(client: &Client, migration_name: &str) -> Result<()> {
client
.execute(
"DELETE FROM _migrations WHERE name = $1",
&[&migration_name]
)
.await?;
Ok(())
}
pub(crate) fn ensure_migration_dir(name: &str) -> Result<PathBuf> {
let re = Regex::new(r"^[a-zA-Z0-9_]+$")?;
if !re.is_match(&name) {
return Err(anyhow!("Le nom de migration ne peut contenir que des lettres, chiffres et underscores"));
}
let database_dir = ensure_database_dir()?;
let next_number = get_next_number(&database_dir)?;
let migration_prefix = format!("{:05}", next_number);
let migration_name = format!("{}_{}", migration_prefix, name);
let migration_dir = database_dir.join(migration_name);
if !migration_dir.exists() {
create_dir_all(&migration_dir)
.with_context(|| format!("Can't create the directory {}", migration_dir.display()))?;
}
Ok(migration_dir)
}