mirror of
https://github.com/Portabase/agent.git
synced 2026-09-10 01:57:10 +00:00
Merge pull request #88 from Portabase/fix/postgres-provider
fix: postgres-provider
This commit is contained in:
+1
-1
@@ -21,7 +21,7 @@ services:
|
||||
LOG: debug
|
||||
TZ: "Europe/Paris"
|
||||
# TMPDIR: /scratch
|
||||
EDGE_KEY: "eyJzZXJ2ZXJVcmwiOiJodHRwOi8vbG9jYWxob3N0Ojg4ODciLCJhZ2VudElkIjoiNTljYzRjYTUtOTAyNy00ZThiLTk1NDktMjAzOTI3ZDVjNmUyIiwibWFzdGVyS2V5QjY0IjoiMUh0djdtWCtYVkJxL0IzUEV2WDlZZjlQeUdVZW5oRHlXemo5THRqNW90WT0ifQ=="
|
||||
EDGE_KEY: "eyJzZXJ2ZXJVcmwiOiJodHRwOi8vbG9jYWxob3N0Ojg4ODciLCJhZ2VudElkIjoiNmM4NWE3ODQtODRkMi00YzUyLTgzYmUtZTc2MDZkZjg2YjM5IiwibWFzdGVyS2V5QjY0IjoiMUh0djdtWCtYVkJxL0IzUEV2WDlZZjlQeUdVZW5oRHlXemo5THRqNW90WT0ifQ=="
|
||||
#CHUNK_SIZE_MB: "1"
|
||||
#POOLING: 1
|
||||
#DATABASES_CONFIG_FILE: "config.toml"
|
||||
|
||||
@@ -0,0 +1,25 @@
|
||||
use crate::services::config::DatabaseConfig;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum RestoreCleanMode {
|
||||
None,
|
||||
Clean,
|
||||
DropSchemas,
|
||||
DropDatabase,
|
||||
}
|
||||
|
||||
impl RestoreCleanMode {
|
||||
pub fn from_config(cfg: &DatabaseConfig) -> (Self, Option<String>) {
|
||||
match cfg.options.get("clean_mode").and_then(|v| v.as_str()) {
|
||||
None | Some("clean") => (Self::Clean, None),
|
||||
Some("none") => (Self::None, None),
|
||||
Some("drop_schemas") => (Self::DropSchemas, None),
|
||||
Some("drop_database") => (Self::DropDatabase, None),
|
||||
Some(other) => (Self::Clean, Some(other.to_string())),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn uses_pg_restore_clean(self) -> bool {
|
||||
matches!(self, Self::Clean)
|
||||
}
|
||||
}
|
||||
@@ -15,10 +15,11 @@ pub async fn run(
|
||||
env: HashMap<String, String>,
|
||||
logger: Arc<JobLogger>,
|
||||
) -> Result<()> {
|
||||
let handle = tokio::runtime::Handle::current();
|
||||
tokio::task::spawn_blocking(move || -> Result<()> {
|
||||
logger.log("info", format!("Starting cluster restore for {}", cfg.name));
|
||||
|
||||
let version = match futures::executor::block_on(server_version(&cfg)) {
|
||||
let version = match handle.block_on(server_version(&cfg)) {
|
||||
Ok(v) => v,
|
||||
Err(e) => {
|
||||
logger.log("error", format!("Failed to get server version for {}: {:?}", cfg.name, e));
|
||||
@@ -26,7 +27,7 @@ pub async fn run(
|
||||
}
|
||||
};
|
||||
|
||||
match futures::executor::block_on(is_superuser(&cfg)) {
|
||||
match handle.block_on(is_superuser(&cfg)) {
|
||||
Ok(true) => {}
|
||||
Ok(false) => {
|
||||
logger.log("error", format!("postgresql-cluster restore requires a superuser role for {}", cfg.name));
|
||||
@@ -40,7 +41,7 @@ pub async fn run(
|
||||
|
||||
let psql = select_pg_path(&version).join(psql_binary_name());
|
||||
|
||||
if let Err(e) = futures::executor::block_on(terminate_all_connections(&cfg)) {
|
||||
if let Err(e) = handle.block_on(terminate_all_connections(&cfg)) {
|
||||
logger.log("error", format!("Failed to terminate connections for cluster {}: {:?}", cfg.name, e));
|
||||
return Err(e.into());
|
||||
}
|
||||
|
||||
@@ -33,6 +33,14 @@ pub async fn server_version(cfg: &DatabaseConfig) -> Result<String> {
|
||||
Ok(version)
|
||||
}
|
||||
|
||||
pub async fn server_version_major(cfg: &DatabaseConfig) -> Result<u32> {
|
||||
let v = server_version(cfg).await?;
|
||||
Ok(v.split(['.', ' '])
|
||||
.next()
|
||||
.and_then(|s| s.parse().ok())
|
||||
.unwrap_or(17))
|
||||
}
|
||||
|
||||
pub async fn is_superuser(cfg: &DatabaseConfig) -> Result<bool> {
|
||||
let client = connect(cfg).await?;
|
||||
let is_super: bool = client
|
||||
@@ -43,6 +51,19 @@ pub async fn is_superuser(cfg: &DatabaseConfig) -> Result<bool> {
|
||||
Ok(is_super)
|
||||
}
|
||||
|
||||
pub async fn can_drop_database(cfg: &DatabaseConfig) -> Result<bool> {
|
||||
let client = connect(cfg).await?;
|
||||
let row = client
|
||||
.query_one(
|
||||
"SELECT r.rolsuper OR (r.rolcreatedb AND pg_catalog.pg_has_role(current_user, d.datdba, 'USAGE')) \
|
||||
FROM pg_roles r, pg_database d \
|
||||
WHERE r.rolname = current_user AND d.datname = current_database()",
|
||||
&[],
|
||||
)
|
||||
.await?;
|
||||
Ok(row.get(0))
|
||||
}
|
||||
|
||||
|
||||
pub fn select_pg_path(version: &str) -> std::path::PathBuf {
|
||||
select_pg_path_with(version, &CONFIG.pg_bin_dir)
|
||||
@@ -114,6 +135,22 @@ pub(crate) fn psql_binary_name() -> &'static str {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn pg_restore_binary_name() -> &'static str {
|
||||
if cfg!(target_os = "windows") {
|
||||
"pg_restore.exe"
|
||||
} else {
|
||||
"pg_restore"
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn quote_ident(s: &str) -> String {
|
||||
format!("\"{}\"", s.replace('"', "\"\""))
|
||||
}
|
||||
|
||||
pub(crate) fn quote_literal(s: &str) -> String {
|
||||
format!("'{}'", s.replace('\'', "''"))
|
||||
}
|
||||
|
||||
pub(crate) fn pg_dump_exists_in(dir: &std::path::Path) -> bool {
|
||||
dir.join(pg_dump_binary_name()).is_file()
|
||||
}
|
||||
@@ -165,6 +202,107 @@ pub async fn terminate_all_connections(cfg: &DatabaseConfig) -> Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn drop_and_recreate_database(cfg: &DatabaseConfig) -> Result<()> {
|
||||
let mut admin_cfg = cfg.clone();
|
||||
admin_cfg.database = "postgres".to_string();
|
||||
let admin = connect(&admin_cfg).await?;
|
||||
|
||||
let row = admin
|
||||
.query_opt(
|
||||
r#"
|
||||
SELECT pg_encoding_to_char(encoding), datcollate, datctype,
|
||||
pg_get_userbyid(datdba), datistemplate
|
||||
FROM pg_database WHERE datname = $1
|
||||
"#,
|
||||
&[&cfg.database],
|
||||
)
|
||||
.await?;
|
||||
|
||||
let (encoding, collate, ctype, owner) = match &row {
|
||||
Some(r) => (
|
||||
r.get::<_, String>(0),
|
||||
r.get::<_, String>(1),
|
||||
r.get::<_, String>(2),
|
||||
r.get::<_, String>(3),
|
||||
),
|
||||
None => ("UTF8".into(), "C".into(), "C".into(), cfg.username.clone()),
|
||||
};
|
||||
|
||||
if let Some(r) = &row {
|
||||
if r.get::<_, bool>(4) {
|
||||
anyhow::bail!("Refusing to drop template database {}", cfg.database);
|
||||
}
|
||||
}
|
||||
|
||||
let db = quote_ident(&cfg.database);
|
||||
|
||||
if let Err(e) = admin
|
||||
.batch_execute(&format!("ALTER DATABASE {db} WITH ALLOW_CONNECTIONS false"))
|
||||
.await
|
||||
{
|
||||
tracing::warn!("ALLOW_CONNECTIONS false failed for {}: {e}", cfg.database);
|
||||
}
|
||||
|
||||
let major = server_version_major(&admin_cfg).await?;
|
||||
let drop_stmt = if major >= 13 {
|
||||
format!("DROP DATABASE IF EXISTS {db} WITH (FORCE)")
|
||||
} else {
|
||||
format!("DROP DATABASE IF EXISTS {db}")
|
||||
};
|
||||
|
||||
let mut last_err = None;
|
||||
let mut dropped = false;
|
||||
for _ in 0..3 {
|
||||
let _ = terminate_connections(cfg).await;
|
||||
match admin.batch_execute(&drop_stmt).await {
|
||||
Ok(()) => {
|
||||
dropped = true;
|
||||
break;
|
||||
}
|
||||
Err(e) => {
|
||||
last_err = Some(e);
|
||||
tokio::time::sleep(std::time::Duration::from_millis(500)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !dropped {
|
||||
let _ = admin
|
||||
.batch_execute(&format!("ALTER DATABASE {db} WITH ALLOW_CONNECTIONS true"))
|
||||
.await;
|
||||
return Err(last_err
|
||||
.map(anyhow::Error::from)
|
||||
.unwrap_or_else(|| anyhow::anyhow!("DROP DATABASE {} failed", cfg.database)));
|
||||
}
|
||||
|
||||
admin
|
||||
.batch_execute(&format!(
|
||||
"CREATE DATABASE {db} OWNER {} TEMPLATE template0 ENCODING {} LC_COLLATE {} LC_CTYPE {}",
|
||||
quote_ident(&owner),
|
||||
quote_literal(&encoding),
|
||||
quote_literal(&collate),
|
||||
quote_literal(&ctype),
|
||||
))
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn sniff_format(restore_file: &Path) -> Result<PostgresDumpFormat> {
|
||||
use std::io::Read;
|
||||
let mut f = std::fs::File::open(restore_file)?;
|
||||
let mut magic = [0u8; 5];
|
||||
let n = f.read(&mut magic)?;
|
||||
let head = &magic[..n];
|
||||
if head.starts_with(b"PGDMP") {
|
||||
Ok(PostgresDumpFormat::Fc)
|
||||
} else if head.starts_with(&[0x1f, 0x8b]) {
|
||||
Ok(PostgresDumpFormat::Fd)
|
||||
} else {
|
||||
anyhow::bail!("Unrecognized dump format for {:?}", restore_file)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn detect_format_from_file(restore_file: &Path) -> PostgresDumpFormat {
|
||||
match restore_file.extension().and_then(|e| e.to_str()) {
|
||||
Some("dump") => PostgresDumpFormat::Fc,
|
||||
@@ -174,6 +312,44 @@ pub fn detect_format_from_file(restore_file: &Path) -> PostgresDumpFormat {
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn drop_all_schemas(cfg: &DatabaseConfig) -> Result<Vec<String>> {
|
||||
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<String> = 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 {:?} - {:?}",
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
#[derive(Clone, Copy)]
|
||||
#[derive(Clone, Copy, PartialEq, Debug)]
|
||||
pub enum PostgresDumpFormat {
|
||||
Fc,
|
||||
Fd,
|
||||
|
||||
@@ -1,9 +1,10 @@
|
||||
pub mod backup;
|
||||
pub(crate) mod cluster;
|
||||
pub(crate) mod clean_mode;
|
||||
pub(crate) mod connection;
|
||||
pub mod database;
|
||||
mod format;
|
||||
pub(crate) mod format;
|
||||
mod ping;
|
||||
mod restore;
|
||||
pub(crate) mod restore;
|
||||
|
||||
pub use connection::{detect_format_from_file, detect_format_from_size};
|
||||
|
||||
@@ -1,214 +0,0 @@
|
||||
use anyhow::Result;
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
use std::process::Command;
|
||||
use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
|
||||
use super::connection::{select_pg_path, server_version, terminate_connections};
|
||||
use super::format::PostgresDumpFormat;
|
||||
use crate::services::backup::logger::JobLogger;
|
||||
use crate::services::config::DatabaseConfig;
|
||||
|
||||
pub async fn run(
|
||||
cfg: DatabaseConfig,
|
||||
format: PostgresDumpFormat,
|
||||
restore_file: PathBuf,
|
||||
env: HashMap<String, String>,
|
||||
logger: Arc<JobLogger>,
|
||||
) -> Result<()> {
|
||||
tokio::task::spawn_blocking(move || -> Result<()> {
|
||||
logger.log("info", format!("Starting restore for database {}", cfg.name));
|
||||
|
||||
let version = match futures::executor::block_on(server_version(&cfg)) {
|
||||
Ok(v) => {
|
||||
logger.log("debug", format!("Postgres version detected: {}", v));
|
||||
v
|
||||
}
|
||||
Err(e) => {
|
||||
logger.log("error", format!("Failed to get server version for {}: {:?}", cfg.name, e));
|
||||
return Err(e.into());
|
||||
}
|
||||
};
|
||||
|
||||
let pg_restore = select_pg_path(&version).join("pg_restore");
|
||||
|
||||
logger.log("debug", format!("Using pg_restore at {:?}", pg_restore));
|
||||
|
||||
if let Err(e) = futures::executor::block_on(terminate_connections(&cfg)) {
|
||||
logger.log("error", format!("Failed to terminate connections for {}: {:?}", cfg.name, e));
|
||||
return Err(e.into());
|
||||
}
|
||||
logger.log("info", format!("Connections terminated for database {}", cfg.name));
|
||||
|
||||
let keep_ownership = cfg.options
|
||||
.get("keep_ownership")
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(false);
|
||||
|
||||
if keep_ownership {
|
||||
logger.log("info", format!("Restoring ownership and privileges for {}", cfg.name));
|
||||
} else {
|
||||
logger.log("info", format!("Stripping ownership and privileges for {} (--no-owner --no-privileges)", cfg.name));
|
||||
}
|
||||
|
||||
match format {
|
||||
PostgresDumpFormat::Fc => {
|
||||
logger.log("info", format!("Running FC restore for {}", cfg.name));
|
||||
let start = Instant::now();
|
||||
let mut cmd = Command::new(&pg_restore);
|
||||
if !keep_ownership {
|
||||
cmd.arg("--no-owner").arg("--no-privileges");
|
||||
}
|
||||
let output = cmd
|
||||
.arg("--clean")
|
||||
.arg("--if-exists")
|
||||
// .arg("--create")
|
||||
.arg("--host").arg(&cfg.host)
|
||||
.arg("--port").arg(cfg.port.to_string())
|
||||
.arg("--username").arg(&cfg.username)
|
||||
.arg("--dbname").arg(&cfg.database)
|
||||
.arg("-v")
|
||||
.arg(&restore_file)
|
||||
.envs(env)
|
||||
.output();
|
||||
|
||||
let duration_ms = start.elapsed().as_millis() as f64;
|
||||
|
||||
match output {
|
||||
Ok(o) => {
|
||||
let stderr = String::from_utf8_lossy(&o.stderr).to_string();
|
||||
let stdout = String::from_utf8_lossy(&o.stdout).to_string();
|
||||
let combined = format!("{}{}", stdout, stderr);
|
||||
let exit_code = o.status.code().unwrap_or(-1);
|
||||
|
||||
if o.status.success() {
|
||||
logger.log_command("pg_restore", if combined.is_empty() { None } else { Some(combined) }, Some(0), Some(duration_ms));
|
||||
logger.log("info", format!("FC restore completed successfully for {}", cfg.name))
|
||||
} else {
|
||||
logger.log_command("pg_restore", if combined.is_empty() { None } else { Some(combined) }, Some(exit_code), Some(duration_ms));
|
||||
logger.log("error", format!("FC restore failed with status {:?} for {}", o.status, cfg.name));
|
||||
anyhow::bail!("Postgres restore failed for {}", cfg.name);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
logger.log_command("pg_restore", Some(e.to_string()), Some(-1), Some(duration_ms));
|
||||
logger.log("error", format!("Error executing pg_restore for {}: {:?}", cfg.name, e));
|
||||
return Err(e.into());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
PostgresDumpFormat::Fd => {
|
||||
logger.log("info", format!("Running FD restore for {}", cfg.name));
|
||||
|
||||
let tar_gz = match std::fs::File::open(&restore_file) {
|
||||
Ok(f) => f,
|
||||
Err(e) => {
|
||||
logger.log("error", format!(
|
||||
"Failed to open restore file {:?} for {}: {:?}",
|
||||
restore_file, cfg.name, e
|
||||
));
|
||||
return Err(e.into());
|
||||
}
|
||||
};
|
||||
|
||||
logger.log("info", format!("tar_gz {:?}", tar_gz));
|
||||
|
||||
let dec = flate2::read::GzDecoder::new(tar_gz);
|
||||
let mut archive = tar::Archive::new(dec);
|
||||
|
||||
let tmp_dir = match tempfile::TempDir::new() {
|
||||
Ok(d) => d,
|
||||
Err(e) => {
|
||||
logger.log("error", format!(
|
||||
"Failed to create temporary directory for FD restore of {}: {:?}",
|
||||
cfg.name, e
|
||||
));
|
||||
return Err(e.into());
|
||||
}
|
||||
};
|
||||
|
||||
if let Err(e) = archive.unpack(tmp_dir.path()) {
|
||||
logger.log("error", format!("Failed to unpack FD archive for {}: {:?}", cfg.name, e));
|
||||
return Err(e.into());
|
||||
}
|
||||
|
||||
logger.log("debug", format!("Listing contents of temp dir: {}", tmp_dir.path().display()));
|
||||
|
||||
for entry in std::fs::read_dir(tmp_dir.path())? {
|
||||
if let Ok(entry) = entry {
|
||||
let path = entry.path();
|
||||
let file_type = entry.file_type()?;
|
||||
logger.log("debug", format!(
|
||||
" - {} | is_dir: {} | is_file: {}",
|
||||
path.display(),
|
||||
file_type.is_dir(),
|
||||
file_type.is_file()
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
let dump_dir = if tmp_dir.path().join("toc.dat").exists() {
|
||||
tmp_dir.path().to_path_buf()
|
||||
} else {
|
||||
std::fs::read_dir(tmp_dir.path())?
|
||||
.filter_map(|e| e.ok())
|
||||
.find(|entry| entry.path().join("toc.dat").exists())
|
||||
.map(|e| e.path())
|
||||
.ok_or_else(|| anyhow::anyhow!("Invalid FD archive: toc.dat not found"))?
|
||||
};
|
||||
|
||||
let start = Instant::now();
|
||||
let mut cmd = Command::new(&pg_restore);
|
||||
if !keep_ownership {
|
||||
cmd.arg("--no-owner").arg("--no-privileges");
|
||||
}
|
||||
let output = cmd
|
||||
.arg("--clean")
|
||||
.arg("--if-exists")
|
||||
// .arg("--create")
|
||||
.arg("--host").arg(&cfg.host)
|
||||
.arg("--port").arg(cfg.port.to_string())
|
||||
.arg("--username").arg(&cfg.username)
|
||||
.arg("--dbname").arg(&cfg.database)
|
||||
.arg("-v")
|
||||
.arg("-j")
|
||||
.arg("4")
|
||||
.arg(dump_dir)
|
||||
.envs(env)
|
||||
.output();
|
||||
|
||||
let duration_ms = start.elapsed().as_millis() as f64;
|
||||
|
||||
match output {
|
||||
Ok(o) => {
|
||||
let stderr = String::from_utf8_lossy(&o.stderr).to_string();
|
||||
let stdout = String::from_utf8_lossy(&o.stdout).to_string();
|
||||
let combined = format!("{}{}", stdout, stderr);
|
||||
let exit_code = o.status.code().unwrap_or(-1);
|
||||
|
||||
if o.status.success() {
|
||||
logger.log_command("pg_restore", if combined.is_empty() { None } else { Some(combined) }, Some(0), Some(duration_ms));
|
||||
logger.log("info", format!("FD restore completed successfully for {}", cfg.name))
|
||||
} else {
|
||||
logger.log_command("pg_restore", if combined.is_empty() { None } else { Some(combined) }, Some(exit_code), Some(duration_ms));
|
||||
logger.log("error", format!("FD restore failed with status {:?} for {}", o.status, cfg.name));
|
||||
anyhow::bail!("Postgres FD restore failed for {}", cfg.name);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
logger.log_command("pg_restore", Some(e.to_string()), Some(-1), Some(duration_ms));
|
||||
logger.log("error", format!("Error executing pg_restore for {}: {:?}", cfg.name, e));
|
||||
return Err(e.into());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
logger.log("info", format!("Restore finished for database {}", cfg.name));
|
||||
|
||||
Ok(())
|
||||
})
|
||||
.await?
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
use anyhow::Result;
|
||||
use std::process::Command;
|
||||
use std::time::Instant;
|
||||
|
||||
use crate::services::backup::logger::JobLogger;
|
||||
use crate::services::config::DatabaseConfig;
|
||||
|
||||
pub(crate) fn run_pg_restore(
|
||||
mut cmd: Command,
|
||||
logger: &JobLogger,
|
||||
cfg: &DatabaseConfig,
|
||||
) -> Result<()> {
|
||||
let start = Instant::now();
|
||||
let output = cmd.output();
|
||||
let duration_ms = start.elapsed().as_millis() as f64;
|
||||
|
||||
match output {
|
||||
Ok(o) => {
|
||||
let stderr = String::from_utf8_lossy(&o.stderr).to_string();
|
||||
let stdout = String::from_utf8_lossy(&o.stdout).to_string();
|
||||
let combined = format!("{}{}", stdout, stderr);
|
||||
let exit_code = o.status.code().unwrap_or(-1);
|
||||
let payload = if combined.is_empty() { None } else { Some(combined) };
|
||||
|
||||
if o.status.success() {
|
||||
logger.log_command("pg_restore", payload, Some(0), Some(duration_ms));
|
||||
logger.log("info", format!("Restore completed successfully for {}", cfg.name));
|
||||
Ok(())
|
||||
} else {
|
||||
logger.log_command("pg_restore", payload, Some(exit_code), Some(duration_ms));
|
||||
logger.log("error", format!("Restore failed with status {:?} for {}", o.status, cfg.name));
|
||||
anyhow::bail!("Postgres restore failed for {}", cfg.name);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
logger.log_command("pg_restore", Some(e.to_string()), Some(-1), Some(duration_ms));
|
||||
logger.log("error", format!("Error executing pg_restore for {}: {:?}", cfg.name, e));
|
||||
Err(e.into())
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
mod command;
|
||||
mod prepare;
|
||||
mod run;
|
||||
mod toc;
|
||||
|
||||
pub use run::run;
|
||||
pub(crate) use command::run_pg_restore;
|
||||
pub(crate) use prepare::prepare_archive;
|
||||
pub(crate) use toc::toc_creates_public_schema;
|
||||
@@ -0,0 +1,67 @@
|
||||
use anyhow::Result;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::process::Command;
|
||||
|
||||
use crate::domain::postgres::connection::sniff_format;
|
||||
use crate::domain::postgres::format::PostgresDumpFormat;
|
||||
use crate::services::backup::logger::JobLogger;
|
||||
|
||||
pub(crate) struct PreparedArchive {
|
||||
path: PathBuf,
|
||||
_tmp: Option<tempfile::TempDir>,
|
||||
toc: String,
|
||||
}
|
||||
|
||||
impl PreparedArchive {
|
||||
pub(crate) fn path(&self) -> &Path {
|
||||
&self.path
|
||||
}
|
||||
pub(crate) fn toc(&self) -> &str {
|
||||
&self.toc
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn prepare_archive(
|
||||
format: PostgresDumpFormat,
|
||||
restore_file: &Path,
|
||||
pg_restore: &Path,
|
||||
logger: &JobLogger,
|
||||
) -> Result<PreparedArchive> {
|
||||
let sniffed = sniff_format(restore_file)?;
|
||||
if sniffed != format {
|
||||
logger.log("warn", format!("Declared format {:?} != sniffed {:?}; using sniffed", format, sniffed));
|
||||
}
|
||||
let format = sniffed;
|
||||
|
||||
let (path, tmp) = match format {
|
||||
PostgresDumpFormat::Fc => (restore_file.to_path_buf(), None),
|
||||
PostgresDumpFormat::Fd => {
|
||||
let tar_gz = std::fs::File::open(restore_file)?;
|
||||
let dec = flate2::read::GzDecoder::new(tar_gz);
|
||||
let mut archive = tar::Archive::new(dec);
|
||||
let tmp_dir = tempfile::TempDir::new()?;
|
||||
archive.unpack(tmp_dir.path())?;
|
||||
|
||||
let dump_dir = if tmp_dir.path().join("toc.dat").exists() {
|
||||
tmp_dir.path().to_path_buf()
|
||||
} else {
|
||||
std::fs::read_dir(tmp_dir.path())?
|
||||
.filter_map(|e| e.ok())
|
||||
.find(|entry| entry.path().join("toc.dat").exists())
|
||||
.map(|e| e.path())
|
||||
.ok_or_else(|| anyhow::anyhow!("Invalid FD archive: toc.dat not found"))?
|
||||
};
|
||||
(dump_dir, Some(tmp_dir))
|
||||
}
|
||||
};
|
||||
|
||||
let toc_out = Command::new(pg_restore).arg("-l").arg(&path).output()?;
|
||||
if !toc_out.status.success() {
|
||||
let stderr = String::from_utf8_lossy(&toc_out.stderr).to_string();
|
||||
logger.log("error", format!("pg_restore -l failed: {}", stderr));
|
||||
anyhow::bail!("Archive validation failed (pg_restore -l): {}", stderr);
|
||||
}
|
||||
let toc = String::from_utf8_lossy(&toc_out.stdout).to_string();
|
||||
|
||||
Ok(PreparedArchive { path, _tmp: tmp, toc })
|
||||
}
|
||||
@@ -0,0 +1,108 @@
|
||||
use anyhow::Result;
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
use std::process::Command;
|
||||
use std::sync::Arc;
|
||||
|
||||
use super::{prepare_archive, run_pg_restore, toc_creates_public_schema};
|
||||
use crate::domain::postgres::clean_mode::RestoreCleanMode;
|
||||
use crate::domain::postgres::connection::{
|
||||
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 crate::domain::postgres::format::PostgresDumpFormat;
|
||||
use crate::services::backup::logger::JobLogger;
|
||||
use crate::services::config::DatabaseConfig;
|
||||
|
||||
pub async fn run(
|
||||
cfg: DatabaseConfig,
|
||||
format: PostgresDumpFormat,
|
||||
restore_file: PathBuf,
|
||||
env: HashMap<String, String>,
|
||||
logger: Arc<JobLogger>,
|
||||
) -> Result<()> {
|
||||
let handle = tokio::runtime::Handle::current();
|
||||
tokio::task::spawn_blocking(move || -> Result<()> {
|
||||
logger.log("info", format!("Starting restore for database {}", cfg.name));
|
||||
|
||||
let version = match handle.block_on(server_version(&cfg)) {
|
||||
Ok(v) => {
|
||||
logger.log("debug", format!("Postgres version detected: {}", v));
|
||||
v
|
||||
}
|
||||
Err(e) => {
|
||||
logger.log("error", format!("Failed to get server version for {}: {:?}", cfg.name, e));
|
||||
return Err(e.into());
|
||||
}
|
||||
};
|
||||
|
||||
let pg_restore = select_pg_path(&version).join(pg_restore_binary_name());
|
||||
|
||||
logger.log("debug", format!("Using pg_restore at {:?}", pg_restore));
|
||||
|
||||
let keep_ownership = cfg.options
|
||||
.get("keep_ownership")
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(false);
|
||||
|
||||
if keep_ownership {
|
||||
logger.log("info", format!("Restoring ownership and privileges for {}", cfg.name));
|
||||
} else {
|
||||
logger.log("info", format!("Stripping ownership and privileges for {} (--no-owner --no-privileges)", cfg.name));
|
||||
}
|
||||
|
||||
let (mode, bad_value) = RestoreCleanMode::from_config(&cfg);
|
||||
if let Some(v) = bad_value {
|
||||
logger.log("warn", format!("Unknown clean_mode '{}' for {}, falling back to 'clean'", v, cfg.name));
|
||||
}
|
||||
|
||||
let prepared = prepare_archive(format, &restore_file, &pg_restore, &logger)?;
|
||||
|
||||
match mode {
|
||||
RestoreCleanMode::DropSchemas => {
|
||||
handle.block_on(terminate_connections(&cfg))?;
|
||||
let owner = cfg.username.clone();
|
||||
let dropped = handle.block_on(drop_all_schemas(&cfg))?;
|
||||
logger.log("warn", format!("clean_mode=drop_schemas dropped schemas {:?} in {}", dropped, cfg.database));
|
||||
if !toc_creates_public_schema(prepared.toc()) {
|
||||
handle.block_on(recreate_public_schema(&cfg, &owner))?;
|
||||
}
|
||||
}
|
||||
RestoreCleanMode::DropDatabase => {
|
||||
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))?;
|
||||
}
|
||||
}
|
||||
|
||||
let mut cmd = Command::new(&pg_restore);
|
||||
if !keep_ownership {
|
||||
cmd.args(["--no-owner", "--no-privileges"]);
|
||||
}
|
||||
if mode.uses_pg_restore_clean() {
|
||||
cmd.args(["--clean", "--if-exists"]);
|
||||
}
|
||||
cmd.arg("--host").arg(&cfg.host)
|
||||
.arg("--port").arg(cfg.port.to_string())
|
||||
.arg("--username").arg(&cfg.username)
|
||||
.arg("--dbname").arg(&cfg.database)
|
||||
.arg("-v");
|
||||
if matches!(format, PostgresDumpFormat::Fd) {
|
||||
cmd.arg("-j").arg("4");
|
||||
}
|
||||
cmd.arg(prepared.path()).envs(env);
|
||||
|
||||
run_pg_restore(cmd, &logger, &cfg)?;
|
||||
logger.log("info", format!("Restore finished for database {}", cfg.name));
|
||||
Ok(())
|
||||
})
|
||||
.await?
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
pub(crate) fn toc_creates_public_schema(toc: &str) -> bool {
|
||||
toc.lines().any(|l| {
|
||||
l.split(" SCHEMA - ")
|
||||
.nth(1)
|
||||
.and_then(|rest| rest.split_whitespace().next())
|
||||
== Some("public")
|
||||
})
|
||||
}
|
||||
+700
-14
@@ -1,9 +1,14 @@
|
||||
use crate::domain::factory::DatabaseFactory;
|
||||
use crate::domain::postgres::connection::{pg_restore_binary_name, select_pg_path, server_version};
|
||||
use crate::domain::postgres::format::PostgresDumpFormat;
|
||||
use crate::domain::postgres::restore::prepare_archive;
|
||||
use crate::services::backup::logger::JobLogger;
|
||||
use crate::services::config::{DatabaseConfig, DbType};
|
||||
use crate::tests::init_tracing_for_test;
|
||||
use crate::utils::compress::{compress_to_tar_gz_large, decompress_large_tar_gz};
|
||||
use oauth2::url;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use tempfile::TempDir;
|
||||
use testcontainers::runners::AsyncRunner;
|
||||
use testcontainers::{ContainerAsync, ImageExt};
|
||||
@@ -63,7 +68,6 @@ async fn postgres_ping_test() {
|
||||
async fn is_superuser_detects_superuser_role() {
|
||||
init_tracing_for_test();
|
||||
|
||||
// The testcontainer's POSTGRES_USER ("testuser") is the bootstrap superuser.
|
||||
let (_container, config) = create_config().await;
|
||||
|
||||
let is_super = crate::domain::postgres::connection::is_superuser(&config)
|
||||
@@ -73,6 +77,31 @@ async fn is_superuser_detects_superuser_role() {
|
||||
assert!(is_super);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn can_drop_database_false_for_unprivileged_role() {
|
||||
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 lowpriv; CREATE ROLE lowpriv LOGIN PASSWORD 'x';")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let mut low = admin.clone();
|
||||
low.username = "lowpriv".into();
|
||||
low.password = "x".into();
|
||||
|
||||
assert_eq!(
|
||||
crate::domain::postgres::connection::can_drop_database(&low)
|
||||
.await
|
||||
.unwrap(),
|
||||
false
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn postgres_backup_restore_test() {
|
||||
init_tracing_for_test();
|
||||
@@ -170,17 +199,585 @@ async fn postgres_password_with_slash_test() {
|
||||
assert_eq!(reachable, true);
|
||||
}
|
||||
|
||||
fn pg_dump_env(config: &DatabaseConfig) -> std::collections::HashMap<String, String> {
|
||||
let mut env = std::env::vars().collect::<std::collections::HashMap<_, _>>();
|
||||
env.insert("PGPASSWORD".to_string(), config.password.clone());
|
||||
env
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn prepare_archive_fd_locates_toc_dir() {
|
||||
init_tracing_for_test();
|
||||
|
||||
let (_container, config) = create_config().await;
|
||||
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
|
||||
let backup_path = crate::domain::postgres::backup::run(
|
||||
config.clone(),
|
||||
PostgresDumpFormat::Fd,
|
||||
temp_dir.path().to_path_buf(),
|
||||
pg_dump_env(&config),
|
||||
Arc::new(JobLogger::new()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let compression = compress_to_tar_gz_large(&backup_path, Arc::new(JobLogger::new()))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(compression.compressed_path.is_file());
|
||||
|
||||
let version = server_version(&config).await.unwrap();
|
||||
let pg_restore = select_pg_path(&version).join(pg_restore_binary_name());
|
||||
|
||||
let logger = JobLogger::new();
|
||||
|
||||
let prepared = prepare_archive(
|
||||
PostgresDumpFormat::Fd,
|
||||
&compression.compressed_path,
|
||||
&pg_restore,
|
||||
&logger,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert!(prepared.path().join("toc.dat").exists());
|
||||
assert!(!prepared.toc().is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn prepare_archive_fc_returns_file_path_unchanged() {
|
||||
init_tracing_for_test();
|
||||
|
||||
let (_container, config) = create_config().await;
|
||||
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
|
||||
let backup_path = 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();
|
||||
|
||||
assert!(backup_path.is_file());
|
||||
|
||||
let version = server_version(&config).await.unwrap();
|
||||
let pg_restore = select_pg_path(&version).join(pg_restore_binary_name());
|
||||
|
||||
let logger = JobLogger::new();
|
||||
|
||||
let prepared = prepare_archive(PostgresDumpFormat::Fc, &backup_path, &pg_restore, &logger).unwrap();
|
||||
|
||||
assert_eq!(prepared.path(), backup_path.as_path());
|
||||
assert!(!prepared.toc().is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn restore_run_unified_fc_roundtrip() {
|
||||
init_tracing_for_test();
|
||||
|
||||
let (_container, config) = create_config().await;
|
||||
|
||||
let client = crate::domain::postgres::connection::connect(&config)
|
||||
.await
|
||||
.unwrap();
|
||||
client.execute("CREATE TABLE 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();
|
||||
|
||||
assert!(dump_file.is_file());
|
||||
|
||||
let format = crate::domain::postgres::connection::detect_format_from_file(&dump_file);
|
||||
|
||||
let result = crate::domain::postgres::restore::run(
|
||||
config.clone(),
|
||||
format,
|
||||
dump_file,
|
||||
pg_dump_env(&config),
|
||||
Arc::new(JobLogger::new()),
|
||||
)
|
||||
.await;
|
||||
|
||||
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);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn restore_drop_schemas_removes_extra_objects() {
|
||||
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 client = crate::domain::postgres::connection::connect(&config)
|
||||
.await
|
||||
.unwrap();
|
||||
client
|
||||
.batch_execute("CREATE TABLE orphan_only_here(id int);")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let mut cfg = config.clone();
|
||||
cfg.options
|
||||
.insert("clean_mode".into(), serde_json::json!("drop_schemas"));
|
||||
|
||||
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 client = crate::domain::postgres::connection::connect(&config)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let n: i64 = client
|
||||
.query_one(
|
||||
"SELECT count(*) FROM information_schema.tables WHERE table_name = 'orphan_only_here'",
|
||||
&[],
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.get(0);
|
||||
assert_eq!(n, 0);
|
||||
|
||||
let n: i64 = client
|
||||
.query_one(
|
||||
"SELECT count(*) FROM information_schema.tables WHERE table_name = 'base_t'",
|
||||
&[],
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.get(0);
|
||||
assert_eq!(n, 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn restore_clean_leaves_divergent_object() {
|
||||
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 client = crate::domain::postgres::connection::connect(&config)
|
||||
.await
|
||||
.unwrap();
|
||||
client
|
||||
.batch_execute("CREATE TABLE survives_clean(id int);")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let mut cfg = config.clone();
|
||||
cfg.options
|
||||
.insert("clean_mode".into(), serde_json::json!("clean"));
|
||||
|
||||
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 client = crate::domain::postgres::connection::connect(&config)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let n: i64 = client
|
||||
.query_one(
|
||||
"SELECT count(*) FROM information_schema.tables WHERE table_name = 'survives_clean'",
|
||||
&[],
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.get(0);
|
||||
assert_eq!(n, 1, "clean mode is not a reset; divergent object survives");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn restore_unknown_clean_mode_falls_back() {
|
||||
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!("wat"));
|
||||
|
||||
let format = crate::domain::postgres::connection::detect_format_from_file(&dump_file);
|
||||
|
||||
let logger = Arc::new(JobLogger::new());
|
||||
|
||||
let result = crate::domain::postgres::restore::run(
|
||||
cfg.clone(),
|
||||
format,
|
||||
dump_file,
|
||||
pg_dump_env(&cfg),
|
||||
logger.clone(),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert!(result.is_ok(), "restore::run failed: {:?}", result);
|
||||
|
||||
let entries = Arc::try_unwrap(logger)
|
||||
.expect("logger should have a single owner after run() completes")
|
||||
.into_entries();
|
||||
assert!(entries
|
||||
.iter()
|
||||
.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);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sniff_format_detects_custom_and_gzip() {
|
||||
use crate::domain::postgres::connection::sniff_format;
|
||||
|
||||
let dir = TempDir::new().unwrap();
|
||||
let fc = dir.path().join("a.dump");
|
||||
std::fs::write(&fc, b"PGDMP\x01\x0e\x00").unwrap();
|
||||
assert_eq!(sniff_format(&fc).unwrap(), PostgresDumpFormat::Fc);
|
||||
|
||||
let fd = dir.path().join("b.gz");
|
||||
std::fs::write(&fd, [0x1f, 0x8b, 0x08, 0x00]).unwrap();
|
||||
assert_eq!(sniff_format(&fd).unwrap(), PostgresDumpFormat::Fd);
|
||||
|
||||
let bad = dir.path().join("c.bin");
|
||||
std::fs::write(&bad, b"not a dump").unwrap();
|
||||
assert!(sniff_format(&bad).is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn corrupt_archive_leaves_database_untouched() {
|
||||
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 must_survive(id int);")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let mut cfg = config.clone();
|
||||
cfg.options
|
||||
.insert("clean_mode".into(), serde_json::json!("drop_schemas"));
|
||||
|
||||
let dir = TempDir::new().unwrap();
|
||||
let broken_file = dir.path().join("broken.tar.gz");
|
||||
std::fs::write(&broken_file, [0x1f, 0x8b, 0x08, 0x00, 0xde, 0xad, 0xbe, 0xef]).unwrap();
|
||||
|
||||
let result = crate::domain::postgres::restore::run(
|
||||
cfg.clone(),
|
||||
PostgresDumpFormat::Fd,
|
||||
broken_file,
|
||||
pg_dump_env(&cfg),
|
||||
Arc::new(JobLogger::new()),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert!(result.is_err(), "corrupt archive must be rejected before any destructive step");
|
||||
|
||||
let client = crate::domain::postgres::connection::connect(&config)
|
||||
.await
|
||||
.unwrap();
|
||||
let n: i64 = client
|
||||
.query_one(
|
||||
"SELECT count(*) FROM information_schema.tables WHERE table_name = 'must_survive'",
|
||||
&[],
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.get(0);
|
||||
assert_eq!(n, 1, "corrupt archive must never trigger the schema drop");
|
||||
}
|
||||
|
||||
mod select_pg_path_tests {
|
||||
use crate::domain::postgres::connection::{
|
||||
pg_dump_binary_name, pg_dump_exists_in, pg_dumpall_binary_name, psql_binary_name,
|
||||
select_pg_path_with,
|
||||
pg_dump_binary_name, pg_dump_exists_in, pg_dumpall_binary_name, pg_restore_binary_name,
|
||||
psql_binary_name, select_pg_path_with,
|
||||
};
|
||||
|
||||
// `select_pg_path_with` takes the `PG_BIN_DIR` override as a plain
|
||||
// argument, so these tests never touch process-global env state or the
|
||||
// cached `CONFIG`. They stay deterministic regardless of whether — or at
|
||||
// which version — a real PostgreSQL install exists on the host.
|
||||
|
||||
#[test]
|
||||
fn respects_pg_bin_dir_override() {
|
||||
let custom = if cfg!(target_os = "windows") {
|
||||
@@ -194,8 +791,6 @@ mod select_pg_path_tests {
|
||||
|
||||
#[test]
|
||||
fn pg_bin_dir_override_ignores_requested_version() {
|
||||
// The override is taken as-is, regardless of which version was
|
||||
// requested — this documents/locks in that behavior.
|
||||
let custom = if cfg!(target_os = "windows") {
|
||||
r"C:\custom\pg\bin"
|
||||
} else {
|
||||
@@ -207,10 +802,6 @@ mod select_pg_path_tests {
|
||||
|
||||
#[test]
|
||||
fn empty_pg_bin_dir_falls_through_to_detection() {
|
||||
// An empty override means "unset" (matches `CONFIG.pg_bin_dir` when
|
||||
// `PG_BIN_DIR` is absent). It must not be returned as a literal empty
|
||||
// path — resolution falls through to platform defaults / PATH lookup
|
||||
// and yields a non-empty path.
|
||||
let path = select_pg_path_with("17", "");
|
||||
assert_ne!(path, std::path::PathBuf::from(""));
|
||||
}
|
||||
@@ -250,4 +841,99 @@ mod select_pg_path_tests {
|
||||
assert_eq!(name, "psql");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pg_restore_binary_name_is_platform_correct() {
|
||||
let name = pg_restore_binary_name();
|
||||
if cfg!(target_os = "windows") {
|
||||
assert_eq!(name, "pg_restore.exe");
|
||||
} else {
|
||||
assert_eq!(name, "pg_restore");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
mod quoting_tests {
|
||||
use crate::domain::postgres::connection::{quote_ident, quote_literal};
|
||||
|
||||
#[test]
|
||||
fn quote_ident_escapes_double_quotes() {
|
||||
assert_eq!(quote_ident("devdb"), "\"devdb\"");
|
||||
assert_eq!(quote_ident("a\"b"), "\"a\"\"b\"");
|
||||
assert_eq!(quote_ident("drop\"; --"), "\"drop\"\"; --\"");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn quote_literal_escapes_single_quotes() {
|
||||
assert_eq!(quote_literal("UTF8"), "'UTF8'");
|
||||
assert_eq!(quote_literal("O'Brien"), "'O''Brien'");
|
||||
}
|
||||
}
|
||||
|
||||
mod clean_mode_tests {
|
||||
use crate::domain::postgres::clean_mode::RestoreCleanMode as M;
|
||||
use crate::services::config::{DatabaseConfig, DbType};
|
||||
|
||||
fn cfg_with(clean_mode: Option<&str>) -> DatabaseConfig {
|
||||
let mut options = std::collections::HashMap::new();
|
||||
if let Some(v) = clean_mode {
|
||||
options.insert("clean_mode".to_string(), serde_json::json!(v));
|
||||
}
|
||||
DatabaseConfig {
|
||||
name: "t".into(),
|
||||
database: "testdb".into(),
|
||||
db_type: DbType::Postgresql,
|
||||
username: "testuser".into(),
|
||||
password: "changeme".into(),
|
||||
port: 5432,
|
||||
host: "localhost".into(),
|
||||
generated_id: "00000000-0000-0000-0000-000000000000".into(),
|
||||
path: "".into(),
|
||||
max_packet_size: "".into(),
|
||||
volume_name: "".into(),
|
||||
container_name: None,
|
||||
options,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn clean_mode_parsing() {
|
||||
assert_eq!(M::from_config(&cfg_with(None)), (M::Clean, None));
|
||||
assert_eq!(M::from_config(&cfg_with(Some("clean"))), (M::Clean, None));
|
||||
assert_eq!(M::from_config(&cfg_with(Some("none"))), (M::None, None));
|
||||
assert_eq!(
|
||||
M::from_config(&cfg_with(Some("drop_schemas"))),
|
||||
(M::DropSchemas, None)
|
||||
);
|
||||
assert_eq!(
|
||||
M::from_config(&cfg_with(Some("drop_database"))),
|
||||
(M::DropDatabase, None)
|
||||
);
|
||||
assert_eq!(
|
||||
M::from_config(&cfg_with(Some("bogus"))),
|
||||
(M::Clean, Some("bogus".to_string()))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn uses_pg_restore_clean_behavior() {
|
||||
assert!(M::Clean.uses_pg_restore_clean());
|
||||
assert!(!M::DropSchemas.uses_pg_restore_clean());
|
||||
assert!(!M::None.uses_pg_restore_clean());
|
||||
assert!(!M::DropDatabase.uses_pg_restore_clean());
|
||||
}
|
||||
}
|
||||
|
||||
mod toc_tests {
|
||||
use crate::domain::postgres::restore::toc_creates_public_schema;
|
||||
|
||||
#[test]
|
||||
fn toc_public_schema_detection() {
|
||||
let with = "215; 2615 2200 SCHEMA - public pg_database_owner";
|
||||
let without_table = "200; 1259 12346 TABLE devschema users devuser";
|
||||
let without_similar_schema = "216; 2615 2201 SCHEMA - publicish someowner";
|
||||
assert!(toc_creates_public_schema(with));
|
||||
assert!(!toc_creates_public_schema(without_table));
|
||||
assert!(!toc_creates_public_schema(without_similar_schema));
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user