From f24b1ec3e48968bb2f32e28c08192c749631fbaa Mon Sep 17 00:00:00 2001 From: charles-gauthereau Date: Thu, 23 Jul 2026 09:51:49 +0200 Subject: [PATCH] feat(postgres): wire clean_mode=drop_database with preflight gate --- src/domain/postgres/connection.rs | 2 +- src/domain/postgres/restore.rs | 13 ++- src/tests/domain/postgres.rs | 171 ++++++++++++++++++++++++++++++ 3 files changed, 182 insertions(+), 4 deletions(-) diff --git a/src/domain/postgres/connection.rs b/src/domain/postgres/connection.rs index f6e1fa2..27a1665 100644 --- a/src/domain/postgres/connection.rs +++ b/src/domain/postgres/connection.rs @@ -243,7 +243,7 @@ pub async fn drop_and_recreate_database(cfg: &DatabaseConfig) -> Result<()> { tracing::warn!("ALLOW_CONNECTIONS false failed for {}: {e}", cfg.database); } - let major = server_version_major(cfg).await?; + let major = server_version_major(&admin_cfg).await?; let drop_stmt = if major >= 13 { format!("DROP DATABASE IF EXISTS {db} WITH (FORCE)") } else { diff --git a/src/domain/postgres/restore.rs b/src/domain/postgres/restore.rs index f63134d..ed86240 100644 --- a/src/domain/postgres/restore.rs +++ b/src/domain/postgres/restore.rs @@ -7,8 +7,8 @@ use std::time::Instant; use super::clean_mode::RestoreCleanMode; use super::connection::{ - drop_all_schemas, pg_restore_binary_name, recreate_public_schema, select_pg_path, - server_version, terminate_connections, + can_drop_database, drop_all_schemas, drop_and_recreate_database, pg_restore_binary_name, + recreate_public_schema, select_pg_path, server_version, terminate_connections, }; use super::format::PostgresDumpFormat; use crate::services::backup::logger::JobLogger; @@ -163,7 +163,14 @@ pub async fn run( } } RestoreCleanMode::DropDatabase => { - anyhow::bail!("clean_mode=drop_database not yet available"); + if !handle.block_on(can_drop_database(&cfg))? { + anyhow::bail!( + "clean_mode=drop_database requires CREATEDB + ownership on {}; use clean_mode=drop_schemas instead", + cfg.database + ); + } + logger.log("warn", format!("clean_mode=drop_database DROPPING database {} before restore", cfg.database)); + handle.block_on(drop_and_recreate_database(&cfg))?; } RestoreCleanMode::Clean | RestoreCleanMode::None => { handle.block_on(terminate_connections(&cfg))?; diff --git a/src/tests/domain/postgres.rs b/src/tests/domain/postgres.rs index 2bb269f..1c227ac 100644 --- a/src/tests/domain/postgres.rs +++ b/src/tests/domain/postgres.rs @@ -537,6 +537,177 @@ async fn restore_unknown_clean_mode_falls_back() { .any(|e| e.message.contains("Unknown clean_mode 'wat'"))); } +#[tokio::test] +async fn drop_database_preflight_preserves_data_when_unprivileged() { + init_tracing_for_test(); + + let (_container, admin) = create_config().await; + + let a = crate::domain::postgres::connection::connect(&admin) + .await + .unwrap(); + a.batch_execute("DROP ROLE IF EXISTS lowpriv2; CREATE ROLE lowpriv2 LOGIN PASSWORD 'x';") + .await + .unwrap(); + a.batch_execute("CREATE TABLE IF NOT EXISTS keep_me(id int);") + .await + .unwrap(); + + let temp_dir = TempDir::new().unwrap(); + + let dump_file = crate::domain::postgres::backup::run( + admin.clone(), + PostgresDumpFormat::Fc, + temp_dir.path().to_path_buf(), + pg_dump_env(&admin), + Arc::new(JobLogger::new()), + ) + .await + .unwrap(); + + let mut low = admin.clone(); + low.username = "lowpriv2".into(); + low.password = "x".into(); + low.options + .insert("clean_mode".into(), serde_json::json!("drop_database")); + + let format = crate::domain::postgres::connection::detect_format_from_file(&dump_file); + + let res = crate::domain::postgres::restore::run( + low.clone(), + format, + dump_file, + pg_dump_env(&low), + Arc::new(JobLogger::new()), + ) + .await; + assert!(res.is_err(), "preflight must reject an unprivileged role"); + + let a = crate::domain::postgres::connection::connect(&admin) + .await + .unwrap(); + let n: i64 = a + .query_one( + "SELECT count(*) FROM information_schema.tables WHERE table_name = 'keep_me'", + &[], + ) + .await + .unwrap() + .get(0); + assert_eq!(n, 1, "preflight must fail before dropping anything"); +} + +#[tokio::test] +async fn drop_database_preserves_encoding_and_owner() { + init_tracing_for_test(); + + let (_container, config) = create_config().await; + + let client = crate::domain::postgres::connection::connect(&config) + .await + .unwrap(); + client + .batch_execute("CREATE TABLE base_t(id int);") + .await + .unwrap(); + + let temp_dir = TempDir::new().unwrap(); + + let dump_file = crate::domain::postgres::backup::run( + config.clone(), + PostgresDumpFormat::Fc, + temp_dir.path().to_path_buf(), + pg_dump_env(&config), + Arc::new(JobLogger::new()), + ) + .await + .unwrap(); + + let before = crate::domain::postgres::connection::connect(&config) + .await + .unwrap() + .query_one( + "SELECT pg_encoding_to_char(encoding), datcollate FROM pg_database WHERE datname = current_database()", + &[], + ) + .await + .unwrap(); + let enc0: String = before.get(0); + let coll0: String = before.get(1); + + let mut cfg = config.clone(); + cfg.options + .insert("clean_mode".into(), serde_json::json!("drop_database")); + + let format = crate::domain::postgres::connection::detect_format_from_file(&dump_file); + + let result = crate::domain::postgres::restore::run( + cfg.clone(), + format, + dump_file, + pg_dump_env(&cfg), + Arc::new(JobLogger::new()), + ) + .await; + assert!(result.is_ok(), "restore::run failed: {:?}", result); + + let after = crate::domain::postgres::connection::connect(&config) + .await + .unwrap() + .query_one( + "SELECT pg_encoding_to_char(encoding), datcollate FROM pg_database WHERE datname = current_database()", + &[], + ) + .await + .unwrap(); + let enc1: String = after.get(0); + let coll1: String = after.get(1); + assert_eq!(enc0, enc1); + assert_eq!(coll0, coll1); +} + +#[tokio::test] +async fn drop_database_force_wins_over_open_connection() { + init_tracing_for_test(); + + let (_container, config) = create_config().await; + + let temp_dir = TempDir::new().unwrap(); + + let dump_file = crate::domain::postgres::backup::run( + config.clone(), + PostgresDumpFormat::Fc, + temp_dir.path().to_path_buf(), + pg_dump_env(&config), + Arc::new(JobLogger::new()), + ) + .await + .unwrap(); + + let mut cfg = config.clone(); + cfg.options + .insert("clean_mode".into(), serde_json::json!("drop_database")); + + let squatter = crate::domain::postgres::connection::connect(&config) + .await + .unwrap(); + let _keep = tokio::spawn(async move { + let _ = squatter.query_one("SELECT pg_sleep(5)", &[]).await; + }); + + let format = crate::domain::postgres::connection::detect_format_from_file(&dump_file); + + let result = crate::domain::postgres::restore::run( + cfg.clone(), + format, + dump_file, + pg_dump_env(&cfg), + Arc::new(JobLogger::new()), + ) + .await; + assert!(result.is_ok(), "restore::run failed: {:?}", result); +} + mod select_pg_path_tests { use crate::domain::postgres::connection::{ pg_dump_binary_name, pg_dump_exists_in, pg_dumpall_binary_name, pg_restore_binary_name,