yggdrasil/api/database/backup/
creation.rs

1//! 数据库备份执行、进度更新与 uploads 素材打包。
2
3use super::commands::pg_dump_command;
4use super::files::{
5    backup_path, backup_stamp, create_new_backup_file, rotate_auto_backups, uploads_archive_name,
6    BACKUP_DIR, BACKUP_SIGNATURE,
7};
8use super::{BackupOrigin, BackupRunOutcome};
9use crate::api::database::tasks::{self, TaskKind, TaskStatus};
10use chrono::Utc;
11use std::io::Write;
12use std::path::Path;
13use std::sync::atomic::{AtomicBool, Ordering};
14
15/// 已有备份在跑时拒绝新备份的提示（手动与自动共用）。
16pub(super) const BACKUP_BUSY: &str = "已有备份正在进行，请等待完成后再试";
17
18/// 进程内备份单飞标记：同一时刻只允许一个备份任务写 backups/。
19static BACKUP_RUNNING: AtomicBool = AtomicBool::new(false);
20
21/// 持有期间占住备份单飞标记；drop（含任务 panic 展开）时释放。
22/// [`run_backup`] 按值接收，保证每个备份任务都先拿到它。
23pub(super) struct BackupGuard(());
24
25impl BackupGuard {
26    /// 已有备份在跑时返回 None。
27    pub(super) fn try_acquire() -> Option<Self> {
28        BACKUP_RUNNING
29            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
30            .ok()
31            .map(|_| Self(()))
32    }
33}
34
35impl Drop for BackupGuard {
36    fn drop(&mut self) {
37        BACKUP_RUNNING.store(false, Ordering::Release);
38    }
39}
40
41/// 定时任务入口：注册任务后执行自动备份，返回结果供调度器落库 last_run。
42/// 手动备份仍在跑时跳过本次，不注册任务。
43pub(crate) async fn run_auto_backup() -> Result<BackupRunOutcome, String> {
44    let guard = BackupGuard::try_acquire().ok_or_else(|| BACKUP_BUSY.to_string())?;
45    let task_id = uuid::Uuid::new_v4().to_string();
46    tasks::insert(task_id.clone(), TaskKind::Backup);
47    run_backup(&task_id, BackupOrigin::Auto, guard).await
48}
49
50/// 后台执行备份：pg_dump 优先，不可用回退纯 SQL；成功后按设置打包 uploads，
51/// 自动备份再做轮转。Ok/Err 均已同步进任务进度表（供前端轮询），返回值供
52/// 调度器落库 last_run。
53pub(super) async fn run_backup(
54    task_id: &str,
55    origin: BackupOrigin,
56    _guard: BackupGuard,
57) -> Result<BackupRunOutcome, String> {
58    let _ = std::fs::create_dir_all(BACKUP_DIR);
59    let stamp = backup_stamp(Utc::now());
60    let prefix = origin.file_prefix();
61
62    // 读取备份设置（include_uploads / retention_count）；读取失败用默认值——
63    // DB 不可达时 pg_dump 自身也会失败并按原路径上报。
64    let settings = match crate::db::pool::get_conn().await {
65        Ok(conn) => crate::api::settings::load_backup_settings(&conn)
66            .await
67            .unwrap_or_default(),
68        Err(_) => crate::models::settings::BackupSettings::default(),
69    };
70
71    // 探测 pg_dump（fork+exec+wait 仍是阻塞系统调用，移出 tokio worker 线程）
72    let pg_dump_ok = tokio::task::spawn_blocking(|| {
73        std::process::Command::new("pg_dump")
74            .arg("--version")
75            .output()
76            .is_ok()
77    })
78    .await
79    .unwrap_or(false);
80
81    let sql_result = if pg_dump_ok {
82        run_pg_dump_backup(task_id, prefix, &stamp).await
83    } else {
84        run_sql_fallback_backup(task_id, prefix, &stamp).await
85    };
86    let sql_filename = match sql_result {
87        Ok(f) => f,
88        // 子流程已把具体错误写进任务进度表，这里仅透传给调度器。
89        Err(e) => return Err(e),
90    };
91
92    // uploads 素材打包：失败不拖垮已成功的 DB 备份，降级为告警。
93    let mut uploads_filename = None;
94    let mut warning = None;
95    if settings.include_uploads {
96        tasks::update(
97            task_id,
98            "正在打包 uploads 素材",
99            92,
100            TaskStatus::Running,
101            None,
102            None,
103            None,
104        );
105        let tar_name = uploads_archive_name(&sql_filename);
106        let tar_path = backup_path(&tar_name);
107        match tokio::task::spawn_blocking(move || {
108            create_uploads_tarball(Path::new("uploads"), &tar_path)
109        })
110        .await
111        {
112            Ok(Ok(())) => uploads_filename = Some(tar_name),
113            Ok(Err(e)) => warning = Some(format!("uploads 打包失败: {e}")),
114            Err(e) => warning = Some(format!("uploads 打包任务panic: {e}")),
115        }
116        if warning.is_some() {
117            tracing::warn!("backup uploads tarball failed: {:?}", warning);
118        }
119    }
120
121    // 自动备份轮转：只删 auto_ 前缀的最旧者，手动备份永不触碰。
122    if origin == BackupOrigin::Auto {
123        rotate_auto_backups(settings.retention_count);
124    }
125
126    tasks::update(
127        task_id,
128        "完成",
129        100,
130        TaskStatus::Done,
131        warning.clone(),
132        None,
133        Some(sql_filename.clone()),
134    );
135    Ok(BackupRunOutcome {
136        sql_filename,
137        uploads_filename,
138        warning,
139    })
140}
141/// 把 uploads 目录打成 tar.gz（排除可重建的 .cache/ 与 VCS 占位 .gitkeep）。
142/// 阻塞 IO + CPU（gzip），调用方须放 spawn_blocking。
143/// uploads_dir 显式传入（而非写死相对路径）便于测试注入临时目录。
144fn create_uploads_tarball(uploads_dir: &Path, out_path: &Path) -> std::io::Result<()> {
145    let file = create_new_backup_file(out_path)?;
146    let gz = flate2::write::GzEncoder::new(file, flate2::Compression::default());
147    let mut builder = tar::Builder::new(gz);
148    if uploads_dir.is_dir() {
149        for entry in std::fs::read_dir(uploads_dir)? {
150            let entry = entry?;
151            let name = entry.file_name();
152            if name == ".cache" || name == ".gitkeep" {
153                continue;
154            }
155            let path = entry.path();
156            if path.is_dir() {
157                builder.append_dir_all(&name, &path)?;
158            } else if path.is_file() {
159                builder.append_path_with_name(&path, &name)?;
160            }
161        }
162    }
163    let gz = builder.into_inner()?;
164    gz.finish()?;
165    Ok(())
166}
167async fn run_pg_dump_backup(task_id: &str, prefix: &str, stamp: &str) -> Result<String, String> {
168    tasks::update(
169        task_id,
170        "正在用 pg_dump 导出",
171        10,
172        TaskStatus::Running,
173        None,
174        None,
175        None,
176    );
177    let filename = format!("{prefix}_{stamp}.sql");
178    let path = backup_path(&filename);
179    let db_url = match std::env::var("DATABASE_URL") {
180        Ok(u) if !u.is_empty() => u,
181        _ => {
182            let msg = "pg_dump 备份需要 DATABASE_URL".to_string();
183            tasks::update(
184                task_id,
185                "DATABASE_URL 未配置",
186                100,
187                TaskStatus::Failed,
188                None,
189                Some(msg.clone()),
190                None,
191            );
192            return Err(msg);
193        }
194    };
195
196    let mut header = String::new();
197    header.push_str(&format!("{}\n", BACKUP_SIGNATURE));
198    header.push_str(&format!("-- created_at: {}\n", Utc::now()));
199    header.push_str("-- mode: pg_dump\n");
200
201    // 独占新建后先写签名头，再把同一文件句柄交给 pg_dump 续写：
202    // 子进程继承的文件偏移就在签名头之后。
203    let mut stdout_file = match create_new_backup_file(&path) {
204        Ok(f) => f,
205        Err(e) => {
206            let msg = format!("无法写入备份目录: {e}");
207            tasks::update(
208                task_id,
209                "写入备份文件失败",
210                100,
211                TaskStatus::Failed,
212                None,
213                Some(msg.clone()),
214                None,
215            );
216            return Err(msg);
217        }
218    };
219    if let Err(e) = stdout_file.write_all(header.as_bytes()) {
220        let msg = format!("无法写入备份目录: {e}");
221        tasks::update(
222            task_id,
223            "写入备份文件失败",
224            100,
225            TaskStatus::Failed,
226            None,
227            Some(msg.clone()),
228            None,
229        );
230        return Err(msg);
231    }
232    // pg_dump 导出可持续数十秒到数分钟，整个子进程生命周期（spawn + wait_with_output）
233    // 移入 spawn_blocking，避免阻塞 tokio worker 线程。注意 stdout 重定向到备份文件，
234    // 故不能用 .output()（它会用 piped 覆盖 stdout 配置）；保留 spawn() + wait_with_output()
235    // 两段式。闭包返回 Result<Output, (bool, io::Error)>：true=启动(spawn)阶段失败，
236    // false=等待(wait)阶段失败，分别对应原有「启动失败」「执行失败」两条上报路径；
237    // 闭包 panic（JoinError）按「执行失败」处理。
238    let dump_result = tokio::task::spawn_blocking(
239        move || -> Result<std::process::Output, (bool, std::io::Error)> {
240            let mut command = pg_dump_command(&db_url).map_err(|error| (true, error))?;
241            command
242                .stdout(std::process::Stdio::from(stdout_file))
243                .stderr(std::process::Stdio::piped())
244                .spawn()
245                .map_err(|e| (true, e))?
246                .wait_with_output()
247                .map_err(|e| (false, e))
248        },
249    )
250    .await
251    .unwrap_or_else(|join_e| Err((false, std::io::Error::other(join_e.to_string()))));
252    match dump_result {
253        Ok(o) if o.status.success() => Ok(filename),
254        Ok(o) => {
255            let msg = String::from_utf8_lossy(&o.stderr).to_string();
256            tasks::update(
257                task_id,
258                "pg_dump 失败",
259                100,
260                TaskStatus::Failed,
261                None,
262                Some(msg.clone()),
263                None,
264            );
265            Err(msg)
266        }
267        Err((true, e)) => {
268            let msg = e.to_string();
269            tasks::update(
270                task_id,
271                "pg_dump 启动失败",
272                100,
273                TaskStatus::Failed,
274                None,
275                Some(msg.clone()),
276                None,
277            );
278            Err(msg)
279        }
280        Err((false, e)) => {
281            let msg = e.to_string();
282            tasks::update(
283                task_id,
284                "pg_dump 执行失败",
285                100,
286                TaskStatus::Failed,
287                None,
288                Some(msg.clone()),
289                None,
290            );
291            Err(msg)
292        }
293    }
294}
295
296/// 纯 SQL 回退：仅备份数据（不含 schema），按表计数精确进度。
297///
298/// 注意：该产物不是合法 SQL 脚本（COPY TO STDOUT 的 CSV 原文拼接），
299/// 不能经 psql 恢复——仅供极端环境下抢救数据。生产镜像自带 pg_dump，
300/// 正常不会走到这里。
301///
302/// 成功返回 SQL 文件名（最终「完成」进度由调用方统一上报）。
303async fn run_sql_fallback_backup(
304    task_id: &str,
305    prefix: &str,
306    stamp: &str,
307) -> Result<String, String> {
308    tasks::update(
309        task_id,
310        "pg_dump 不可用，使用纯 SQL 回退（仅数据）",
311        10,
312        TaskStatus::Running,
313        Some("仅备份数据，不含 schema/索引/触发器，且不可经 psql 恢复".to_string()),
314        None,
315        None,
316    );
317    let filename = format!("{prefix}_{stamp}_sqlfallback.sql");
318    let path = backup_path(&filename);
319
320    let client = match crate::db::pool::get_conn().await {
321        Ok(c) => c,
322        Err(e) => {
323            let msg = e.to_string();
324            tasks::update(
325                task_id,
326                "数据库连接失败",
327                100,
328                TaskStatus::Failed,
329                None,
330                Some(msg.clone()),
331                None,
332            );
333            return Err(msg);
334        }
335    };
336
337    // 取 public schema 下所有表名
338    let tables: Vec<String> = match client
339        .query(
340            "SELECT tablename FROM pg_tables WHERE schemaname = 'public' ORDER BY tablename",
341            &[],
342        )
343        .await
344    {
345        Ok(rows) => rows.into_iter().map(|r| r.get(0)).collect(),
346        Err(e) => {
347            let msg = e.to_string();
348            tasks::update(
349                task_id,
350                "读取表清单失败",
351                100,
352                TaskStatus::Failed,
353                None,
354                Some(msg.clone()),
355                None,
356            );
357            return Err(msg);
358        }
359    };
360    let total = tables.len().max(1);
361
362    let mut out = String::new();
363    out.push_str(&format!("{}\n", BACKUP_SIGNATURE));
364    out.push_str(&format!("-- created_at: {}\n", Utc::now()));
365    out.push_str("-- mode: sql-fallback\n\n");
366
367    for (i, table) in tables.iter().enumerate() {
368        out.push_str(&format!("\n-- table: {}\n", table));
369        let copy_stmt = format!("COPY \"{}\" TO STDOUT WITH CSV", table);
370        match client.copy_out(&copy_stmt).await {
371            Ok(stream) => {
372                use futures::StreamExt;
373                // CopyOutStream 是 !Unpin，必须 pin 才能调 next。
374                tokio::pin!(stream);
375                while let Some(chunk) = stream.next().await {
376                    if let Ok(bytes) = chunk {
377                        out.push_str(&String::from_utf8_lossy(&bytes));
378                    }
379                }
380            }
381            Err(e) => {
382                out.push_str(&format!("-- 导出失败: {}\n", e));
383            }
384        }
385        // 按表更新进度（用 u32 避免大 schema 下的截断/溢出）
386        tasks::update(
387            task_id,
388            &format!("导出表 {}/{}", i + 1, total),
389            (10 + (i + 1) as u32 * 90 / total as u32).min(99) as u8,
390            TaskStatus::Running,
391            None,
392            None,
393            None,
394        );
395    }
396
397    if let Err(e) = create_new_backup_file(&path).and_then(|mut f| f.write_all(out.as_bytes())) {
398        let msg = format!("无法写入备份目录: {e}");
399        tasks::update(
400            task_id,
401            "写入备份文件失败",
402            100,
403            TaskStatus::Failed,
404            None,
405            Some(msg.clone()),
406            None,
407        );
408        return Err(msg);
409    }
410    Ok(filename)
411}
412
413#[cfg(test)]
414mod tests {
415    use super::*;
416
417    #[test]
418    fn backup_guard_is_single_flight_until_dropped() {
419        let guard = BackupGuard::try_acquire().expect("空闲时应能拿到备份标记");
420        assert!(
421            BackupGuard::try_acquire().is_none(),
422            "已有备份在跑时必须拒绝第二个"
423        );
424        drop(guard);
425        assert!(
426            BackupGuard::try_acquire().is_some(),
427            "上一个备份结束后应能再次发起"
428        );
429    }
430
431    #[test]
432    fn tarball_excludes_cache_and_gitkeep() {
433        let nanos = std::time::SystemTime::now()
434            .duration_since(std::time::UNIX_EPOCH)
435            .expect("系统时间必然晚于 UNIX_EPOCH")
436            .as_nanos();
437        let dir = std::env::temp_dir().join(format!(
438            "yggdrasil_backup_tar_test_{}_{}",
439            nanos,
440            std::process::id()
441        ));
442        let uploads = dir.join("uploads");
443        std::fs::create_dir_all(uploads.join("2026")).expect("创建测试目录");
444        std::fs::create_dir_all(uploads.join(".cache")).expect("创建测试目录");
445        std::fs::write(uploads.join("2026/pic.webp"), b"img").expect("写测试文件");
446        std::fs::write(uploads.join(".cache/x.webp"), b"cache").expect("写测试文件");
447        std::fs::write(uploads.join(".gitkeep"), b"").expect("写测试文件");
448
449        let out = dir.join("out.tar.gz");
450        create_uploads_tarball(&uploads, &out).expect("打包应成功");
451
452        let file = std::fs::File::open(&out).expect("打开打包产物");
453        let gz = flate2::read::GzDecoder::new(file);
454        let mut archive = tar::Archive::new(gz);
455        let entries: Vec<String> = archive
456            .entries()
457            .expect("读取 tar 条目")
458            .map(|e| {
459                e.expect("tar 条目有效")
460                    .path()
461                    .expect("路径有效")
462                    .to_string_lossy()
463                    .to_string()
464            })
465            .collect();
466        assert!(
467            entries.iter().any(|p| p.contains("2026/pic.webp")),
468            "应包含素材文件: {entries:?}"
469        );
470        assert!(
471            !entries
472                .iter()
473                .any(|p| p.contains(".cache") || p.contains(".gitkeep")),
474            "应排除 .cache 与 .gitkeep: {entries:?}"
475        );
476        std::fs::remove_dir_all(&dir).expect("清理测试目录");
477    }
478}
yggdrasil/api/database/backup/creation.rs

yggdrasil/api/database/backup/
creation.rs
