use std::path::Path;
use std::process::Command;
use diesel::connection::SimpleConnection as _;
use diesel::{Connection as _, PgConnection, RunQueryDsl as _, sql_query};
use crate::migrate;
pub mod backup;
pub mod s3;
pub mod scrub;
pub mod replica;
pub mod sqlite_snapshot;
pub mod retention;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum DbCommand {
Create,
Drop { force: bool },
Reset { force: bool },
}
#[derive(Debug)]
enum DbError {
NoUrl,
UnparsableUrl,
Connection { host: String, port: u16, db: String },
Sql(String),
ProductionRefused { profile: String },
}
impl std::fmt::Display for DbError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::NoUrl => write!(
f,
"No database URL found.\n Set database.primary_url (or database.url) in autumn.toml, \
or set AUTUMN_DATABASE__PRIMARY_URL / AUTUMN_DATABASE__URL / DATABASE_URL."
),
Self::UnparsableUrl => write!(
f,
"The resolved database URL could not be parsed or does not name a database."
),
Self::Connection { host, port, db } => write!(
f,
"Could not connect to the Postgres maintenance database for {db:?} at {host}:{port}.\n \
Is the server running and reachable?"
),
Self::Sql(message) => write!(f, "{message}"),
Self::ProductionRefused { profile } => write!(
f,
"Refusing to run a destructive database operation against the {profile:?} profile.\n \
Re-run with --force if you really mean it."
),
}
}
}
pub fn run(command: &DbCommand, profile: Option<&str>) {
eprintln!("\u{1F342} autumn db\n");
let result = match command {
DbCommand::Create => create(profile),
DbCommand::Drop { force } => drop(profile, *force),
DbCommand::Reset { force } => reset(profile, *force),
};
if let Err(e) = result {
eprintln!("\u{2717} {e}");
std::process::exit(1);
}
}
fn create(profile: Option<&str>) -> Result<(), DbError> {
let url = resolve_url(profile)?;
let MaintenanceTarget {
db_name,
maintenance_url,
host,
port,
} = maintenance_target(&url)?;
let mut conn = connect(&maintenance_url, &host, port, &db_name)?;
if database_exists(&mut conn, &db_name)? {
eprintln!(" \u{2139} Database {db_name:?} already exists \u{2014} nothing to do.");
return Ok(());
}
conn.batch_execute(&format!("CREATE DATABASE {}", quote_ident(&db_name)))
.map_err(|e| DbError::Sql(e.to_string()))?;
eprintln!(" \u{2713} Created database {db_name:?}.");
Ok(())
}
fn drop(profile: Option<&str>, force: bool) -> Result<(), DbError> {
let resolved = migrate::effective_profile(profile);
guard_destructive(&resolved, force)?;
let url = resolve_url(profile)?;
let MaintenanceTarget {
db_name,
maintenance_url,
host,
port,
} = maintenance_target(&url)?;
let mut conn = connect(&maintenance_url, &host, port, &db_name)?;
if !database_exists(&mut conn, &db_name)? {
eprintln!(" \u{2139} Database {db_name:?} does not exist \u{2014} nothing to do.");
return Ok(());
}
conn.batch_execute(&format!("DROP DATABASE {}", quote_ident(&db_name)))
.map_err(|e| DbError::Sql(e.to_string()))?;
eprintln!(" \u{2713} Dropped database {db_name:?}.");
Ok(())
}
fn reset(profile: Option<&str>, force: bool) -> Result<(), DbError> {
let resolved = migrate::effective_profile(profile);
guard_destructive(&resolved, force)?;
run_step("drop", &["db", "drop", "--force"], profile)?;
run_step("create", &["db", "create"], profile)?;
run_step("migrate", &["migrate"], profile)?;
if seed_binary_present() {
run_step("seed", &["seed"], profile)?;
} else {
eprintln!(" \u{2139} No src/bin/seed.rs found \u{2014} skipping the seed step.");
}
eprintln!(
"\n\u{2713} Database reset complete (drop \u{2192} create \u{2192} migrate \u{2192} seed)."
);
Ok(())
}
fn run_step(name: &str, args: &[&str], profile: Option<&str>) -> Result<(), DbError> {
eprintln!("\u{2500}\u{2500} reset: {name} \u{2500}\u{2500}");
let exe = std::env::current_exe()
.map_err(|e| DbError::Sql(format!("could not locate the autumn executable: {e}")))?;
let mut command = Command::new(exe);
command.args(args);
if let Some(profile) = profile {
command.args(["--profile", profile]);
}
let status = command
.status()
.map_err(|e| DbError::Sql(format!("failed to spawn `autumn {}`: {e}", args.join(" "))))?;
if status.success() {
Ok(())
} else {
Err(DbError::Sql(format!(
"reset failed at the {name:?} step (`autumn {}` exited {}).",
args.join(" "),
status
.code()
.map_or_else(|| "with a signal".to_owned(), |c| format!("with code {c}")),
)))
}
}
fn seed_binary_present() -> bool {
Path::new("src/bin/seed.rs").is_file()
}
fn resolve_url(profile: Option<&str>) -> Result<String, DbError> {
migrate::resolve_primary_url(profile).ok_or(DbError::NoUrl)
}
pub fn guard_destructive_public(profile: &str, force: bool) -> Result<(), ()> {
guard_destructive(profile, force).map_err(|_| ())
}
fn guard_destructive(profile: &str, force: bool) -> Result<(), DbError> {
if force || is_safe_destructive_profile(profile) {
Ok(())
} else {
Err(DbError::ProductionRefused {
profile: profile.to_owned(),
})
}
}
fn is_safe_destructive_profile(profile: &str) -> bool {
matches!(
profile.trim().to_ascii_lowercase().as_str(),
"dev" | "development" | "test"
)
}
struct MaintenanceTarget {
db_name: String,
maintenance_url: String,
host: String,
port: u16,
}
fn maintenance_target(url: &str) -> Result<MaintenanceTarget, DbError> {
let parsed = url::Url::parse(url).map_err(|_| DbError::UnparsableUrl)?;
let db_name = parsed
.path_segments()
.and_then(|mut segments| segments.next())
.map(decode_percent)
.filter(|name| !name.is_empty())
.ok_or(DbError::UnparsableUrl)?;
let host = parsed.host_str().unwrap_or("localhost").to_owned();
let port = parsed.port().unwrap_or(5432);
let mut maintenance = parsed;
maintenance.set_path("/postgres");
Ok(MaintenanceTarget {
db_name,
maintenance_url: maintenance.to_string(),
host,
port,
})
}
fn connect(
maintenance_url: &str,
host: &str,
port: u16,
db_name: &str,
) -> Result<PgConnection, DbError> {
PgConnection::establish(maintenance_url).map_err(|_| DbError::Connection {
host: host.to_owned(),
port,
db: db_name.to_owned(),
})
}
#[derive(diesel::QueryableByName)]
struct ExistsRow {
#[diesel(sql_type = diesel::sql_types::Bool)]
exists: bool,
}
fn database_exists(conn: &mut PgConnection, db_name: &str) -> Result<bool, DbError> {
let query = format!(
"SELECT EXISTS(SELECT 1 FROM pg_database WHERE datname = {}) AS exists",
quote_literal(db_name)
);
let row: ExistsRow = sql_query(query)
.get_result(conn)
.map_err(|e| DbError::Sql(e.to_string()))?;
Ok(row.exists)
}
fn quote_ident(name: &str) -> String {
format!("\"{}\"", name.replace('"', "\"\""))
}
pub fn quote_literal(value: &str) -> String {
format!("'{}'", value.replace('\'', "''"))
}
pub fn decode_percent(segment: &str) -> String {
let bytes = segment.as_bytes();
let mut out = Vec::with_capacity(bytes.len());
let mut i = 0;
while i < bytes.len() {
if bytes[i] == b'%' && i + 2 < bytes.len() {
let hi = (bytes[i + 1] as char).to_digit(16);
let lo = (bytes[i + 2] as char).to_digit(16);
if let (Some(hi), Some(lo)) = (hi, lo)
&& let Ok(byte) = u8::try_from(hi * 16 + lo)
{
out.push(byte);
i += 3;
continue;
}
}
out.push(bytes[i]);
i += 1;
}
String::from_utf8_lossy(&out).into_owned()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn maintenance_target_extracts_name_and_rewrites_to_postgres() {
let target = maintenance_target("postgres://user:pw@db.example.com:6543/my_app").unwrap();
assert_eq!(target.db_name, "my_app");
assert_eq!(target.host, "db.example.com");
assert_eq!(target.port, 6543);
assert!(
target.maintenance_url.ends_with("/postgres"),
"maintenance url should point at the postgres database: {}",
target.maintenance_url
);
}
#[test]
fn maintenance_target_defaults_port_and_handles_no_credentials() {
let target = maintenance_target("postgres://localhost/my_app").unwrap();
assert_eq!(target.db_name, "my_app");
assert_eq!(target.host, "localhost");
assert_eq!(target.port, 5432);
}
#[test]
fn maintenance_target_decodes_percent_encoded_name() {
let target = maintenance_target("postgres://localhost/my%20app").unwrap();
assert_eq!(target.db_name, "my app");
}
#[test]
fn maintenance_target_rejects_url_without_database_name() {
assert!(matches!(
maintenance_target("postgres://localhost/"),
Err(DbError::UnparsableUrl)
));
assert!(matches!(
maintenance_target("not a url"),
Err(DbError::UnparsableUrl)
));
}
#[test]
fn quote_ident_doubles_embedded_quotes() {
assert_eq!(quote_ident("my_app"), "\"my_app\"");
assert_eq!(quote_ident("we\"ird"), "\"we\"\"ird\"");
}
#[test]
fn quote_literal_doubles_embedded_single_quotes() {
assert_eq!(quote_literal("my_app"), "'my_app'");
assert_eq!(quote_literal("o'brien"), "'o''brien'");
}
#[test]
fn guard_allows_dev_and_test_without_force() {
for profile in ["dev", "development", "DEV", "test", "Test"] {
assert!(
guard_destructive(profile, false).is_ok(),
"{profile} should be allowed"
);
}
}
#[test]
fn guard_refuses_prod_and_custom_without_force() {
for profile in ["prod", "production", "staging", "anything-else"] {
assert!(
matches!(
guard_destructive(profile, false),
Err(DbError::ProductionRefused { .. })
),
"{profile} should be refused without --force"
);
assert!(
guard_destructive(profile, true).is_ok(),
"{profile} should pass with --force"
);
}
}
#[test]
fn errors_never_leak_credentials() {
let err = DbError::Connection {
host: "db.example.com".to_owned(),
port: 5432,
db: "my_app".to_owned(),
};
let rendered = err.to_string();
assert!(!rendered.contains("hunter2"));
assert!(rendered.contains("db.example.com"));
assert!(rendered.contains("my_app"));
let refused = DbError::ProductionRefused {
profile: "prod".to_owned(),
};
assert!(refused.to_string().contains("prod"));
assert!(!refused.to_string().contains("postgres://"));
}
}