275 lines
11 KiB
Python
275 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import secrets
|
|
import shutil
|
|
from datetime import timedelta
|
|
from pathlib import Path
|
|
|
|
from fastapi import APIRouter, Depends, File, Query, UploadFile
|
|
from fastapi.responses import FileResponse
|
|
from sqlalchemy import func, select, update
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
from starlette.background import BackgroundTask
|
|
|
|
from backend.auth import current_user, hash_token
|
|
from backend.config import get_settings
|
|
from backend.db import get_db
|
|
from backend.models import BackupPreflight, Countdown, User, new_id, utcnow
|
|
|
|
from .archive import MAX_ARCHIVE_BYTES as ARCHIVE_DEFAULT_LIMIT
|
|
from .archive import backup_error, parse_archive_path
|
|
from .schemas import RestoreRequest
|
|
from .service import BackupRepairPending, export_v2, restore_v2, validate_archive
|
|
|
|
router = APIRouter(prefix="/api/v1/backup", tags=["backup"])
|
|
MAX_ARCHIVE_BYTES = ARCHIVE_DEFAULT_LIMIT
|
|
MAX_PENDING_PREFLIGHTS_PER_USER = 3
|
|
_READ_CHUNK = 1024 * 1024
|
|
|
|
|
|
def _staging_root() -> Path:
|
|
root = Path(get_settings().backup_staging_dir).resolve()
|
|
root.mkdir(parents=True, exist_ok=True)
|
|
return root
|
|
|
|
|
|
def _safe_staging_path(value: str) -> Path:
|
|
root = _staging_root()
|
|
path = Path(value).resolve()
|
|
if path.parent != root:
|
|
raise backup_error("backup_preflight_invalid", "预检暂存无效", 409)
|
|
return path
|
|
|
|
|
|
def _path_size(path: str | None) -> int:
|
|
if not path:
|
|
return 0
|
|
root = Path(path)
|
|
if not root.exists():
|
|
return 0
|
|
if root.is_file():
|
|
return root.stat().st_size
|
|
return sum(item.stat().st_size for item in root.rglob("*") if item.is_file())
|
|
|
|
|
|
async def _prune(db: AsyncSession) -> None:
|
|
now = utcnow()
|
|
rows = list((await db.scalars(select(BackupPreflight).where(
|
|
BackupPreflight.expires_at <= now,
|
|
BackupPreflight.status.in_((
|
|
"reserved", "pending", "failed", "consuming", "cleanup_pending",
|
|
"repair_pending",
|
|
)),
|
|
).with_for_update())).all())
|
|
from .storage import remove_quarantine, restore_quarantine_dir
|
|
|
|
attachment_root = Path(get_settings().attachment_dir).resolve()
|
|
for row in rows:
|
|
try:
|
|
if row.status == "repair_pending":
|
|
if row.cleanup_path:
|
|
restore_quarantine_dir(attachment_root, Path(row.cleanup_path))
|
|
row.cleanup_path = None
|
|
_safe_staging_path(row.staging_path).unlink(missing_ok=True)
|
|
row.status = "failed"
|
|
continue
|
|
if row.cleanup_path:
|
|
remove_quarantine(Path(row.cleanup_path))
|
|
row.cleanup_path = None
|
|
_safe_staging_path(row.staging_path).unlink(missing_ok=True)
|
|
except (OSError, ValueError):
|
|
continue
|
|
if row.status == "cleanup_pending":
|
|
row.status = "consumed"
|
|
row.consumed_at = row.consumed_at or now
|
|
else:
|
|
await db.delete(row)
|
|
if rows:
|
|
await db.commit()
|
|
|
|
|
|
@router.get("/export.zip")
|
|
async def export_zip(user: User = Depends(current_user), db: AsyncSession = Depends(get_db)):
|
|
path = await export_v2(db, user)
|
|
await db.commit()
|
|
return FileResponse(
|
|
path,
|
|
media_type="application/zip",
|
|
filename="dodo-backup-v2.zip",
|
|
background=BackgroundTask(path.unlink, missing_ok=True),
|
|
)
|
|
|
|
|
|
@router.post("/preflight")
|
|
async def preflight(
|
|
mode: str = Query(pattern="^(merge|replace)$"),
|
|
file: UploadFile = File(...),
|
|
user: User = Depends(current_user),
|
|
db: AsyncSession = Depends(get_db),
|
|
):
|
|
await _prune(db)
|
|
settings = get_settings()
|
|
configured_limit = settings.backup_max_archive_mb * 1024 * 1024
|
|
limit = min(MAX_ARCHIVE_BYTES, configured_limit)
|
|
max_pending = min(MAX_PENDING_PREFLIGHTS_PER_USER, settings.backup_max_pending_per_user)
|
|
max_staged = settings.backup_max_staged_mb_per_user * 1024 * 1024
|
|
token = secrets.token_urlsafe(32)
|
|
staging = _staging_root() / f"{new_id()}.zip"
|
|
reservation = BackupPreflight(
|
|
token_hash=hash_token(token), user_id=user.id, backup_id=new_id(),
|
|
archive_sha256="0" * 64, archive_size=limit, staging_path=str(staging), mode=mode,
|
|
status="reserved", expires_at=utcnow() + timedelta(seconds=settings.backup_preflight_ttl_seconds),
|
|
)
|
|
# A user-row write serializes quota decisions across workers on PostgreSQL;
|
|
# SQLite serializes writers at the database level.
|
|
await db.execute(update(User).where(User.id == user.id).values(username=User.username))
|
|
pending_count, pending_bytes = (await db.execute(select(
|
|
func.count(BackupPreflight.id), func.coalesce(func.sum(BackupPreflight.archive_size), 0)
|
|
).where(
|
|
BackupPreflight.user_id == user.id,
|
|
BackupPreflight.status.in_((
|
|
"reserved", "pending", "failed", "consuming", "cleanup_pending",
|
|
"repair_pending",
|
|
)),
|
|
))).one()
|
|
pending_rows = list((await db.scalars(select(BackupPreflight).where(
|
|
BackupPreflight.user_id == user.id,
|
|
BackupPreflight.status.in_((
|
|
"reserved", "pending", "failed", "consuming", "cleanup_pending",
|
|
"repair_pending",
|
|
)),
|
|
))).all())
|
|
pending_bytes += sum(_path_size(item.cleanup_path) for item in pending_rows)
|
|
if pending_count >= max_pending or pending_bytes + limit > max_staged:
|
|
await db.rollback()
|
|
raise backup_error("backup_preflight_quota", "待处理预检配额已达上限", 429)
|
|
db.add(reservation)
|
|
await db.flush()
|
|
reservation_id = reservation.id
|
|
await db.commit()
|
|
|
|
size = 0
|
|
digest = hashlib.sha256()
|
|
archive = None
|
|
try:
|
|
with staging.open("xb") as output:
|
|
while chunk := await file.read(_READ_CHUNK):
|
|
size += len(chunk)
|
|
if size > limit:
|
|
raise backup_error("backup_size_invalid", "备份文件大小无效")
|
|
digest.update(chunk)
|
|
output.write(chunk)
|
|
archive = parse_archive_path(staging, max_archive_bytes=limit)
|
|
validate_archive(archive)
|
|
if mode == "merge":
|
|
incoming_pin_ids = {
|
|
str(item["id"])
|
|
for item in archive.entities["countdowns"]
|
|
if item.get("pinned") is True and item.get("archived_at") is None
|
|
}
|
|
existing_pin_ids = set((await db.scalars(select(Countdown.id).where(
|
|
Countdown.user_id == user.id,
|
|
Countdown.pinned.is_(True),
|
|
Countdown.archived_at.is_(None),
|
|
))).all())
|
|
if (
|
|
incoming_pin_ids
|
|
and existing_pin_ids
|
|
and incoming_pin_ids != {str(item) for item in existing_pin_ids}
|
|
):
|
|
raise backup_error("backup_constraint_invalid", "合并恢复会产生多个置顶倒数日")
|
|
shutil.rmtree(archive.staging_dir, ignore_errors=True)
|
|
if archive.archive_sha256 != digest.hexdigest():
|
|
raise backup_error("backup_checksum_mismatch", "备份校验和不匹配")
|
|
reservation.backup_id = archive.backup_id
|
|
reservation.archive_sha256 = archive.archive_sha256
|
|
reservation.archive_size = size
|
|
reservation.status = "pending"
|
|
await db.commit()
|
|
except Exception:
|
|
await db.rollback()
|
|
if archive is not None:
|
|
shutil.rmtree(archive.staging_dir, ignore_errors=True)
|
|
failed = await db.get(BackupPreflight, reservation_id)
|
|
if failed is not None:
|
|
await db.delete(failed)
|
|
await db.commit()
|
|
staging.unlink(missing_ok=True)
|
|
raise
|
|
return {"valid": True, "preflight_token": token, "backup_id": archive.backup_id,
|
|
"archive_sha256": archive.archive_sha256,
|
|
"entities": {name: len(rows) for name, rows in archive.entities.items()}}
|
|
|
|
|
|
@router.post("/restore")
|
|
async def restore(
|
|
payload: RestoreRequest,
|
|
user: User = Depends(current_user),
|
|
db: AsyncSession = Depends(get_db),
|
|
):
|
|
await _prune(db)
|
|
token_hash = hash_token(payload.preflight_token)
|
|
row = await db.scalar(select(BackupPreflight).where(
|
|
BackupPreflight.token_hash == token_hash, BackupPreflight.user_id == user.id
|
|
).with_for_update())
|
|
if row is None or row.mode != payload.mode or row.expires_at <= utcnow():
|
|
raise backup_error("backup_preflight_invalid", "预检令牌无效或已过期", 409)
|
|
if row.status == "repair_pending" and row.cleanup_path:
|
|
from .storage import restore_quarantine_dir
|
|
root = Path(get_settings().attachment_dir).resolve()
|
|
try:
|
|
restore_quarantine_dir(root, Path(row.cleanup_path))
|
|
except OSError as exc:
|
|
raise backup_error("backup_repair_pending", "原附件复原尚未完成，请稍后重试", 503) from exc
|
|
row.status = "failed"
|
|
row.cleanup_path = None
|
|
await db.commit()
|
|
raise backup_error("backup_restore_retry", "附件已复原，请重新预检后重试", 409)
|
|
if row.status == "cleanup_pending":
|
|
from .storage import remove_quarantine
|
|
try:
|
|
if row.cleanup_path:
|
|
remove_quarantine(Path(row.cleanup_path))
|
|
_safe_staging_path(row.staging_path).unlink(missing_ok=True)
|
|
except OSError as exc:
|
|
raise backup_error("backup_cleanup_pending", "清理尚未完成，请稍后重试", 503) from exc
|
|
row.status = "consumed"
|
|
row.cleanup_path = None
|
|
await db.commit()
|
|
return {"restored": 0, "mode": payload.mode, "already_imported": False, "cleanup_retried": True}
|
|
claimed = await db.execute(update(BackupPreflight).where(
|
|
BackupPreflight.id == row.id, BackupPreflight.status.in_(("pending", "failed"))
|
|
).values(status="consuming", consumed_at=utcnow()))
|
|
if claimed.rowcount != 1:
|
|
await db.rollback()
|
|
raise backup_error("backup_preflight_invalid", "预检令牌已使用", 409)
|
|
await db.commit()
|
|
staging = _safe_staging_path(row.staging_path)
|
|
archive = parse_archive_path(staging)
|
|
try:
|
|
if archive.archive_sha256 != row.archive_sha256 or archive.backup_id != row.backup_id:
|
|
row.status = "failed"
|
|
await db.commit()
|
|
raise backup_error("backup_preflight_invalid", "预检暂存已改变", 409)
|
|
result = await restore_v2(db, user, archive, payload.mode, operation=row)
|
|
except BackupRepairPending:
|
|
raise
|
|
except Exception:
|
|
await db.refresh(row)
|
|
if row.status not in {"cleanup_pending", "repair_pending"}:
|
|
row.status = "failed"
|
|
await db.commit()
|
|
raise
|
|
finally:
|
|
shutil.rmtree(archive.staging_dir, ignore_errors=True)
|
|
row.status = "cleanup_pending"
|
|
await db.commit()
|
|
try:
|
|
staging.unlink(missing_ok=True)
|
|
except OSError as exc:
|
|
raise backup_error("backup_cleanup_pending", "数据已恢复，但暂存清理未完成；请使用同一令牌重试", 500) from exc
|
|
row.status = "consumed"
|
|
await db.commit()
|
|
return result
|