diff --git a/src/domain/postgres/connection.rs b/src/domain/postgres/connection.rs index cae7010..a7ca406 100644 --- a/src/domain/postgres/connection.rs +++ b/src/domain/postgres/connection.rs @@ -190,6 +190,44 @@ pub fn detect_format_from_file(restore_file: &Path) -> PostgresDumpFormat { } } +pub async fn drop_all_schemas(cfg: &DatabaseConfig) -> Result> { + let client = connect(cfg).await?; + let rows = client + .query( + r#" + SELECT nspname FROM pg_namespace + WHERE nspname NOT IN ('pg_catalog', 'information_schema', 'pg_toast') + AND nspname NOT LIKE 'pg\_temp\_%' + AND nspname NOT LIKE 'pg\_toast\_temp\_%' + ORDER BY nspname + "#, + &[], + ) + .await?; + let schemas: Vec = rows.iter().map(|r| r.get::<_, String>(0)).collect(); + for s in &schemas { + client + .batch_execute(&format!("DROP SCHEMA IF EXISTS {} CASCADE", quote_ident(s))) + .await?; + } + client + .batch_execute("SELECT lo_unlink(oid) FROM pg_largeobject_metadata") + .await + .ok(); + Ok(schemas) +} + +pub async fn recreate_public_schema(cfg: &DatabaseConfig, owner: &str) -> Result<()> { + let client = connect(cfg).await?; + client + .batch_execute(&format!( + "CREATE SCHEMA IF NOT EXISTS public AUTHORIZATION {}; GRANT USAGE ON SCHEMA public TO PUBLIC;", + quote_ident(owner) + )) + .await?; + Ok(()) +} + pub async fn detect_format_from_size(cfg: &DatabaseConfig) -> PostgresDumpFormat { info!( "Detecting database format {:?} - {:?}", diff --git a/src/tests/domain/postgres.rs b/src/tests/domain/postgres.rs index 49b719b..500a5db 100644 --- a/src/tests/domain/postgres.rs +++ b/src/tests/domain/postgres.rs @@ -292,6 +292,39 @@ async fn restore_run_unified_fc_roundtrip() { assert!(result.is_ok(), "restore::run failed: {:?}", result); } +#[tokio::test] +async fn drop_all_schemas_removes_user_schema() { + init_tracing_for_test(); + + let (_container, config) = create_config().await; + + let client = crate::domain::postgres::connection::connect(&config) + .await + .unwrap(); + client + .batch_execute("CREATE SCHEMA IF NOT EXISTS extra_ns; CREATE TABLE IF NOT EXISTS extra_ns.t(id int);") + .await + .unwrap(); + + let dropped = crate::domain::postgres::connection::drop_all_schemas(&config) + .await + .unwrap(); + assert!(dropped.iter().any(|s| s == "extra_ns")); + + let client = crate::domain::postgres::connection::connect(&config) + .await + .unwrap(); + let row = client + .query_one( + "SELECT count(*) FROM pg_namespace WHERE nspname = 'extra_ns'", + &[], + ) + .await + .unwrap(); + let n: i64 = row.get(0); + assert_eq!(n, 0); +} + 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,