export DATABASE_URL=postgres://...
export DB_PARAM_GATEWAY_PASS=iejwi
cargo run -p your-service --bin dbmigrate -- up --app enableiot --env dev
cargo run -p your-service --bin dbmigrate -- status
cargo run -p your-service --bin dbmigrate -- down
cargo run -p your-service --bin dbmigrate -- reset
| [[bin]] | |
| name = "dbmigrate" | |
| path = "src/bin/dbmigrate.rs" | |
| [dependencies] | |
| anyhow = "1" | |
| tracing = "0.1" | |
| tracing-subscriber = "0.3" | |
| diesel = { version = "2", features = ["postgres"] } |
| use anyhow::{Context, Result}; | |
| use diesel::pg::PgConnection; | |
| use diesel::prelude::*; | |
| use diesel::sql_query; | |
| use diesel_migrations::{embed_migrations, EmbeddedMigrations, MigrationHarness}; | |
| pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!("./migrations"); | |
| #[derive(Debug, Clone, Copy)] | |
| pub enum MigrationAction { | |
| Up, | |
| DownOne, | |
| DownAll, | |
| Reset, | |
| } | |
| #[derive(Debug, Clone)] | |
| pub struct MigrationOptions<'a> { | |
| pub application: Option<&'a str>, | |
| pub environment: Option<&'a str>, | |
| pub use_advisory_lock: bool, | |
| pub transactional: bool, | |
| /// Auto-discover env vars and pass them into the DB session as settings. | |
| /// | |
| /// Example: DB_PARAM_GATEWAY_PASS=... becomes `db.param.gateway_pass`. | |
| pub env_prefixes: &'a [&'a str], | |
| /// Namespace used when creating the session setting keys. | |
| /// If None, we derive from prefix (DB_PARAM_ -> "db.param"). | |
| pub env_namespace_override: Option<&'a str>, | |
| /// Optional denylist (exact env var names) you do NOT want to pass through. | |
| pub env_denylist: &'a [&'a str], | |
| } | |
| impl<'a> Default for MigrationOptions<'a> { | |
| fn default() -> Self { | |
| Self { | |
| application: None, | |
| environment: None, | |
| use_advisory_lock: true, | |
| transactional: true, | |
| env_prefixes: &["DB_PARAM_"], | |
| env_namespace_override: None, | |
| env_denylist: &[], | |
| } | |
| } | |
| } | |
| pub fn migrate(conn: &mut PgConnection, action: MigrationAction, opts: MigrationOptions<'_>) -> Result<()> { | |
| // Optional, explicit session settings | |
| if let Some(app) = opts.application { | |
| set_setting(conn, "app.current_application", app) | |
| .context("failed setting app.current_application")?; | |
| } | |
| if let Some(env) = opts.environment { | |
| set_setting(conn, "app.environment", env) | |
| .context("failed setting app.environment")?; | |
| } | |
| // NEW: auto-pass env vars (DB_PARAM_* etc) into session settings. | |
| apply_env_settings(conn, &opts).context("failed applying env-derived session settings")?; | |
| if opts.use_advisory_lock { | |
| advisory_lock(conn, 0x45_4E_41_42_4C_45_49_4F_54_i64)?; | |
| } | |
| let run = |c: &mut PgConnection| -> Result<()> { | |
| match action { | |
| MigrationAction::Up => { | |
| let applied = c.run_pending_migrations(MIGRATIONS) | |
| .map_err(anyhow::Error::from)?; | |
| tracing::info!("migrations applied: {}", applied.len()); | |
| } | |
| MigrationAction::DownOne => { | |
| c.revert_last_migration(MIGRATIONS).map_err(anyhow::Error::from)?; | |
| tracing::info!("reverted last migration"); | |
| } | |
| MigrationAction::DownAll => { | |
| c.revert_all_migrations(MIGRATIONS).map_err(anyhow::Error::from)?; | |
| tracing::warn!("reverted ALL migrations"); | |
| } | |
| MigrationAction::Reset => { | |
| c.revert_all_migrations(MIGRATIONS).map_err(anyhow::Error::from)?; | |
| let applied = c.run_pending_migrations(MIGRATIONS) | |
| .map_err(anyhow::Error::from)?; | |
| tracing::warn!("reset complete; migrations applied: {}", applied.len()); | |
| } | |
| } | |
| Ok(()) | |
| }; | |
| if opts.transactional { | |
| conn.transaction::<_, diesel::result::Error, _>(|c| { | |
| run(c).map_err(|_| diesel::result::Error::RollbackTransaction) | |
| }).map_err(|_| anyhow::anyhow!("migration transaction rolled back"))?; | |
| Ok(()) | |
| } else { | |
| run(conn) | |
| } | |
| } | |
| /// Discover env vars matching prefixes and apply them via set_config(). | |
| /// | |
| /// Example: | |
| /// DB_PARAM_GATEWAY_PASS=iejwi | |
| /// becomes: | |
| /// set_config('db.param.gateway_pass', 'iejwi', true) | |
| fn apply_env_settings(conn: &mut PgConnection, opts: &MigrationOptions<'_>) -> Result<()> { | |
| // Fast lookup for denylist | |
| let deny = opts.env_denylist; | |
| for (k, v) in std::env::vars() { | |
| if deny.iter().any(|d| *d == k) { | |
| continue; | |
| } | |
| let (prefix, suffix) = match opts.env_prefixes.iter().find_map(|p| { | |
| k.strip_prefix(p).map(|s| (*p, s)) | |
| }) { | |
| Some(x) => x, | |
| None => continue, | |
| }; | |
| // Decide namespace | |
| // - override if provided | |
| // - else derive from prefix: "DB_PARAM_" -> "db.param" | |
| let namespace = opts.env_namespace_override.unwrap_or_else(|| derive_namespace(prefix)); | |
| // Convert suffix "GATEWAY_PASS" -> "gateway_pass" | |
| let leaf = suffix.to_ascii_lowercase(); | |
| // Final setting key: "db.param.gateway_pass" | |
| let setting_key = format!("{}.{}", namespace, leaf); | |
| set_setting(conn, &setting_key, &v) | |
| .with_context(|| format!("failed setting {} from env {}", setting_key, k))?; | |
| } | |
| Ok(()) | |
| } | |
| /// Turn an env prefix into a dotted namespace. | |
| /// Example: "DB_PARAM_" -> "db.param" | |
| fn derive_namespace(prefix: &str) -> &str { | |
| // Keep this simple and deterministic. | |
| // You can expand if you add more prefixes later. | |
| match prefix { | |
| "DB_PARAM_" => "db.param", | |
| "APP_PARAM_" => "app.param", | |
| _ => "env.param", | |
| } | |
| } | |
| fn set_setting(conn: &mut PgConnection, key: &str, value: &str) -> Result<()> { | |
| // Session-local setting | |
| sql_query("select set_config($1, $2, true);") | |
| .bind::<diesel::sql_types::Text, _>(key) | |
| .bind::<diesel::sql_types::Text, _>(value) | |
| .execute(conn) | |
| .context("set_config failed")?; | |
| Ok(()) | |
| } | |
| fn advisory_lock(conn: &mut PgConnection, key: i64) -> Result<()> { | |
| sql_query("select pg_advisory_lock($1);") | |
| .bind::<diesel::sql_types::BigInt, _>(key) | |
| .execute(conn) | |
| .context("pg_advisory_lock failed")?; | |
| Ok(()) | |
| } |
| use anyhow::{bail, Context, Result}; | |
| use diesel::pg::PgConnection; | |
| use diesel::prelude::*; | |
| use your_crate_name::db_migrations::{self, MigrationAction, MigrationOptions}; // <-- change crate path | |
| fn main() -> Result<()> { | |
| tracing_subscriber::fmt() | |
| .with_env_filter( | |
| std::env::var("RUST_LOG").unwrap_or_else(|_| "info".to_string()), | |
| ) | |
| .init(); | |
| let mut args = std::env::args().skip(1); | |
| let cmd = args.next().unwrap_or_else(|| "help".to_string()); | |
| if cmd == "help" || cmd == "--help" || cmd == "-h" { | |
| print_help(); | |
| return Ok(()); | |
| } | |
| // Allow DATABASE_URL from env (common), but also accept --database-url | |
| let mut database_url = std::env::var("DATABASE_URL").ok(); | |
| // Optional flags | |
| let mut app: Option<String> = None; | |
| let mut env_name: Option<String> = std::env::var("APP_ENV").ok(); | |
| let mut no_lock = false; | |
| let mut non_transactional = false; | |
| // Very simple flag parsing: supports: | |
| // --database-url <url> | |
| // --app <name> | |
| // --env <name> | |
| // --no-lock | |
| // --non-transactional | |
| while let Some(a) = args.next() { | |
| match a.as_str() { | |
| "--database-url" => { | |
| let v = args.next().context("--database-url requires a value")?; | |
| database_url = Some(v); | |
| } | |
| "--app" => { | |
| app = Some(args.next().context("--app requires a value")?); | |
| } | |
| "--env" => { | |
| env_name = Some(args.next().context("--env requires a value")?); | |
| } | |
| "--no-lock" => no_lock = true, | |
| "--non-transactional" => non_transactional = true, | |
| other => bail!("Unknown flag: {other}\nRun: dbmigrate help"), | |
| } | |
| } | |
| let database_url = database_url.context("DATABASE_URL not set (or pass --database-url)")?; | |
| let mut conn = PgConnection::establish(&database_url) | |
| .with_context(|| "failed to connect to DATABASE_URL")?; | |
| // Build options: auto-discover DB_PARAM_* and set them on the session. | |
| let opts = MigrationOptions { | |
| application: app.as_deref(), | |
| environment: env_name.as_deref(), | |
| use_advisory_lock: !no_lock, | |
| transactional: !non_transactional, | |
| // Auto discover: | |
| env_prefixes: &["DB_PARAM_"], | |
| env_namespace_override: None, | |
| env_denylist: &[], | |
| }; | |
| match cmd.as_str() { | |
| "up" => { | |
| db_migrations::migrate(&mut conn, MigrationAction::Up, opts)?; | |
| } | |
| "down" => { | |
| db_migrations::migrate(&mut conn, MigrationAction::DownOne, opts)?; | |
| } | |
| "down-all" => { | |
| db_migrations::migrate(&mut conn, MigrationAction::DownAll, opts)?; | |
| } | |
| "reset" | "clean" => { | |
| db_migrations::migrate(&mut conn, MigrationAction::Reset, opts)?; | |
| } | |
| "status" => { | |
| let applied = db_migrations::migration_status(&mut conn)?; | |
| if applied.is_empty() { | |
| println!("No migrations applied."); | |
| } else { | |
| println!("Applied migrations:"); | |
| for v in applied { | |
| println!(" - {v}"); | |
| } | |
| } | |
| } | |
| other => bail!("Unknown command: {other}\nRun: dbmigrate help"), | |
| } | |
| Ok(()) | |
| } | |
| fn print_help() { | |
| eprintln!( | |
| r#"dbmigrate — Diesel migration CLI | |
| Usage: | |
| dbmigrate <command> [flags] | |
| Commands: | |
| up Apply all pending migrations | |
| down Revert the last applied migration | |
| down-all Revert all applied migrations (dangerous) | |
| reset Revert all then re-apply (dangerous) | |
| clean Alias for reset | |
| status List applied migration versions | |
| Flags: | |
| --database-url <url> Override DATABASE_URL | |
| --app <name> Sets session setting: app.current_application | |
| --env <name> Sets session setting: app.environment (default: APP_ENV if set) | |
| --no-lock Disable pg_advisory_lock | |
| --non-transactional Don't wrap the action in one transaction | |
| Environment passthrough: | |
| Any environment variable starting with DB_PARAM_ is set on the DB session using set_config(). | |
| Example: | |
| DB_PARAM_GATEWAY_PASS=iejwi | |
| becomes: | |
| db.param.gateway_pass = 'iejwi' | |
| Readable in SQL via: | |
| current_setting('db.param.gateway_pass', true) | |
| "# | |
| ); | |
| } |