yggdrasil/api/database/backup/
commands.rs

1//! PostgreSQL 命令构造：通过 PG 环境变量传递连接，备份去除属主和 ACL。
2
3use std::path::Path;
4
5/// 子进程不接收完整数据库 URI；连接参数拆入 PG 环境变量。
6fn pg_command_with_db_config(
7    program: &str,
8    db_url: &str,
9) -> std::io::Result<std::process::Command> {
10    use std::str::FromStr;
11    use tokio_postgres::config::{Host, SslMode};
12
13    let config = tokio_postgres::Config::from_str(db_url).map_err(|error| {
14        std::io::Error::new(
15            std::io::ErrorKind::InvalidInput,
16            format!("DATABASE_URL 解析失败: {error}"),
17        )
18    })?;
19    let mut command = std::process::Command::new(program);
20    // pg_dump/psql 不读取 DATABASE_URL；移除继承的敏感环境变量，只保留
21    // 拆出的 PG* 连接参数，避免子进程环境再次暴露完整 URI。
22    command.env_remove("DATABASE_URL");
23
24    let hosts: Vec<String> = config
25        .get_hosts()
26        .iter()
27        .map(|host| match host {
28            Host::Tcp(host) => host.clone(),
29            Host::Unix(path) => path.to_string_lossy().into_owned(),
30        })
31        .collect();
32    if !hosts.is_empty() {
33        command.env("PGHOST", hosts.join(","));
34    }
35    if !config.get_hostaddrs().is_empty() {
36        command.env(
37            "PGHOSTADDR",
38            config
39                .get_hostaddrs()
40                .iter()
41                .map(ToString::to_string)
42                .collect::<Vec<_>>()
43                .join(","),
44        );
45    }
46    if !config.get_ports().is_empty() {
47        command.env(
48            "PGPORT",
49            config
50                .get_ports()
51                .iter()
52                .map(ToString::to_string)
53                .collect::<Vec<_>>()
54                .join(","),
55        );
56    }
57    if let Some(user) = config.get_user() {
58        command.env("PGUSER", user);
59    }
60    if let Some(password) = config.get_password() {
61        command.env("PGPASSWORD", String::from_utf8_lossy(password).into_owned());
62    }
63    if let Some(dbname) = config.get_dbname() {
64        command.env("PGDATABASE", dbname);
65    }
66    if let Some(options) = config.get_options() {
67        command.env("PGOPTIONS", options);
68    }
69    if let Some(application_name) = config.get_application_name() {
70        command.env("PGAPPNAME", application_name);
71    }
72    let ssl_mode = match config.get_ssl_mode() {
73        SslMode::Disable => "disable",
74        SslMode::Prefer => "prefer",
75        SslMode::Require => "require",
76        _ => "prefer",
77    };
78    command.env("PGSSLMODE", ssl_mode);
79    if let Some(timeout) = config.get_connect_timeout() {
80        command.env("PGCONNECT_TIMEOUT", timeout.as_secs().max(1).to_string());
81    }
82    Ok(command)
83}
84
85pub(super) fn pg_dump_command(db_url: &str) -> std::io::Result<std::process::Command> {
86    let mut command = pg_command_with_db_config("pg_dump", db_url)?;
87    command.args(["--clean", "--if-exists", "--no-owner", "--no-privileges"]);
88    Ok(command)
89}
90pub(super) fn psql_restore_command(
91    db_url: &str,
92    path: &Path,
93) -> std::io::Result<std::process::Command> {
94    let mut command = pg_command_with_db_config("psql", db_url)?;
95    command
96        .args(["--single-transaction", "-v", "ON_ERROR_STOP=1", "-f"])
97        .arg(path)
98        .stdout(std::process::Stdio::null())
99        .stderr(std::process::Stdio::piped());
100    Ok(command)
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106
107    #[test]
108    fn pg_dump_is_owner_acl_neutral_without_db_url_in_args() {
109        let db_url = "postgres://user:secret@example.com:5432/dbname";
110        let command = pg_dump_command(db_url).expect("valid database URL");
111        let args: Vec<_> = command
112            .get_args()
113            .map(|arg| arg.to_string_lossy().into_owned())
114            .collect();
115        assert!(args.iter().any(|arg| arg == "--no-owner"), "{args:?}");
116        assert!(args.iter().any(|arg| arg == "--no-privileges"), "{args:?}");
117        assert!(!args.iter().any(|arg| arg.contains("secret")), "{args:?}");
118        assert!(command
119            .get_envs()
120            .any(|(key, value)| key == "PGPASSWORD" && value.is_some()));
121    }
122
123    #[test]
124    fn psql_restore_is_atomic_without_db_url_in_args() {
125        let db_url = "postgres://user:secret@example.com:5432/dbname";
126        let command =
127            psql_restore_command(db_url, Path::new("backup.sql")).expect("valid database URL");
128        let args: Vec<_> = command
129            .get_args()
130            .map(|arg| arg.to_string_lossy().into_owned())
131            .collect();
132        assert!(
133            args.iter().any(|arg| arg == "--single-transaction"),
134            "{args:?}"
135        );
136        assert!(!args.iter().any(|arg| arg.contains("secret")), "{args:?}");
137    }
138}
yggdrasil/api/database/backup/commands.rs

yggdrasil/api/database/backup/
commands.rs
