mirror of
https://github.com/jxxghp/MoviePilot.git
synced 2026-11-12 16:51:21 +08:00
261 lines
8.5 KiB
Python
261 lines
8.5 KiB
Python
"""基于活动 SQLAlchemy 引擎的 SQLite 与 PostgreSQL 备份实现。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
import shutil
|
||
import sqlite3
|
||
import subprocess
|
||
from contextlib import closing
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
from typing import Callable, Mapping, Protocol, Sequence
|
||
|
||
from sqlalchemy.engine import Engine
|
||
|
||
|
||
@dataclass(frozen=True, slots=True)
|
||
class DatabaseBackupCheck:
|
||
"""数据库适配器返回的基础校验结果。"""
|
||
|
||
valid: bool
|
||
method: str
|
||
detail: str | None = None
|
||
|
||
|
||
class ProcessResult(Protocol):
|
||
"""数据库命令执行结果的最小合同。"""
|
||
|
||
returncode: int
|
||
stdout: str
|
||
stderr: str
|
||
|
||
|
||
class ProcessRunner(Protocol):
|
||
"""可替换的数据库命令执行边界。"""
|
||
|
||
def __call__(
|
||
self,
|
||
command: Sequence[str],
|
||
*,
|
||
env: Mapping[str, str],
|
||
capture_output: bool,
|
||
text: bool,
|
||
check: bool,
|
||
) -> ProcessResult:
|
||
"""执行命令并返回结果。"""
|
||
|
||
|
||
def verify_database_backup(
|
||
artifact: Path,
|
||
*,
|
||
db_type: str,
|
||
runner: ProcessRunner = subprocess.run,
|
||
tool_resolver: Callable[[str], str | None] = shutil.which,
|
||
pg_restore: str = "pg_restore",
|
||
) -> DatabaseBackupCheck:
|
||
"""在不访问活动数据库的前提下校验一个受管备份文件。"""
|
||
if db_type == "sqlite":
|
||
method = "PRAGMA integrity_check"
|
||
try:
|
||
# 正式备份不会再变化，immutable 可避免只读校验创建 WAL 旁路文件。
|
||
uri = f"{artifact.resolve().as_uri()}?mode=ro&immutable=1"
|
||
with closing(sqlite3.connect(uri, uri=True)) as connection:
|
||
rows = connection.execute("PRAGMA integrity_check").fetchall()
|
||
except sqlite3.Error as error:
|
||
return DatabaseBackupCheck(False, method, str(error))
|
||
valid = bool(rows) and all(row[0] == "ok" for row in rows)
|
||
detail = None if valid else "; ".join(str(row[0]) for row in rows)
|
||
return DatabaseBackupCheck(valid, method, detail)
|
||
|
||
if db_type == "postgresql":
|
||
method = "pg_restore --list"
|
||
executable = _require_tool(pg_restore, tool_resolver)
|
||
result = runner(
|
||
[executable, "--list", str(artifact)],
|
||
env=_postgres_environment(),
|
||
capture_output=True,
|
||
text=True,
|
||
check=False,
|
||
)
|
||
valid = result.returncode == 0 and bool(result.stdout.strip())
|
||
detail = None if valid else f"pg_restore 退出码 {result.returncode}"
|
||
return DatabaseBackupCheck(valid, method, detail)
|
||
|
||
raise ValueError(f"不支持的数据库备份类型：{db_type}")
|
||
|
||
|
||
def _require_tool(
|
||
executable: str,
|
||
tool_resolver: Callable[[str], str | None],
|
||
) -> str:
|
||
resolved = tool_resolver(executable)
|
||
if resolved is None:
|
||
raise RuntimeError(
|
||
f"未找到 {executable}，请安装与服务端同主版本或更高的 "
|
||
"PostgreSQL client 并加入 PATH"
|
||
)
|
||
return resolved
|
||
|
||
|
||
def _postgres_environment(
|
||
*,
|
||
password: str | None = None,
|
||
sslmode: str | None = None,
|
||
) -> dict[str, str]:
|
||
environment = dict(os.environ)
|
||
environment.pop("PGPASSWORD", None)
|
||
environment.pop("PGSSLMODE", None)
|
||
if password:
|
||
environment["PGPASSWORD"] = password
|
||
if sslmode:
|
||
environment["PGSSLMODE"] = sslmode
|
||
return environment
|
||
|
||
|
||
class SQLiteBackupBackend:
|
||
"""使用 SQLite 在线备份 API 管理活动文件数据库。"""
|
||
|
||
db_type = "sqlite"
|
||
suffix = ".db"
|
||
|
||
def __init__(self, engine: Engine) -> None:
|
||
self._engine = engine
|
||
database = engine.url.database
|
||
if not database or database == ":memory:":
|
||
raise ValueError("SQLite 内存数据库不支持文件备份")
|
||
self._database = Path(database)
|
||
|
||
def create(self, destination: Path) -> None:
|
||
"""从活动引擎指向的 SQLite 文件创建一致快照。"""
|
||
source = self._engine.raw_connection()
|
||
try:
|
||
with closing(sqlite3.connect(destination)) as target:
|
||
source.driver_connection.backup(target)
|
||
target.commit()
|
||
finally:
|
||
source.close()
|
||
|
||
def verify(self, artifact: Path) -> DatabaseBackupCheck:
|
||
"""通过 SQLite integrity_check 校验备份内容。"""
|
||
return verify_database_backup(artifact, db_type=self.db_type)
|
||
|
||
def restore(self, artifact: Path) -> None:
|
||
"""在 CLI 离线进程中原子替换活动 SQLite 文件。"""
|
||
temporary = self._database.with_name(f".{self._database.name}.restore")
|
||
self._engine.dispose()
|
||
try:
|
||
shutil.copy2(artifact, temporary)
|
||
temporary.chmod(0o600)
|
||
self._database.with_name(f"{self._database.name}-wal").unlink(missing_ok=True)
|
||
self._database.with_name(f"{self._database.name}-shm").unlink(missing_ok=True)
|
||
os.replace(temporary, self._database)
|
||
finally:
|
||
temporary.unlink(missing_ok=True)
|
||
|
||
|
||
class PostgreSQLBackupBackend:
|
||
"""使用 pg_dump 与 pg_restore 管理活动 PostgreSQL 数据库。"""
|
||
|
||
db_type = "postgresql"
|
||
suffix = ".dump"
|
||
|
||
def __init__(
|
||
self,
|
||
engine: Engine,
|
||
*,
|
||
runner: ProcessRunner = subprocess.run,
|
||
tool_resolver: Callable[[str], str | None] = shutil.which,
|
||
pg_dump: str = "pg_dump",
|
||
pg_restore: str = "pg_restore",
|
||
) -> None:
|
||
self._engine = engine
|
||
self._runner = runner
|
||
self._tool_resolver = tool_resolver
|
||
self._pg_dump = pg_dump
|
||
self._pg_restore = pg_restore
|
||
|
||
def create(self, destination: Path) -> None:
|
||
"""创建 PostgreSQL custom-format 在线备份。"""
|
||
command = [
|
||
self._require_tool(self._pg_dump),
|
||
"--format=custom",
|
||
"--no-owner",
|
||
"--no-acl",
|
||
"--file",
|
||
str(destination),
|
||
*self._connection_arguments(),
|
||
]
|
||
result = self._run(command, include_password=True)
|
||
if result.returncode != 0:
|
||
raise RuntimeError(f"pg_dump 执行失败，退出码 {result.returncode}")
|
||
if not destination.is_file() or destination.stat().st_size == 0:
|
||
raise RuntimeError("pg_dump 未生成有效的备份文件")
|
||
|
||
def verify(self, artifact: Path) -> DatabaseBackupCheck:
|
||
"""通过 pg_restore 目录读取校验 custom-format 归档。"""
|
||
return verify_database_backup(
|
||
artifact,
|
||
db_type=self.db_type,
|
||
runner=self._runner,
|
||
tool_resolver=self._tool_resolver,
|
||
pg_restore=self._pg_restore,
|
||
)
|
||
|
||
def restore(self, artifact: Path) -> None:
|
||
"""在 CLI 离线进程中覆盖当前 PostgreSQL 数据库内容。"""
|
||
command = [
|
||
self._require_tool(self._pg_restore),
|
||
"--clean",
|
||
"--if-exists",
|
||
"--no-owner",
|
||
"--no-acl",
|
||
"--single-transaction",
|
||
"--exit-on-error",
|
||
*self._connection_arguments(),
|
||
str(artifact),
|
||
]
|
||
result = self._run(command, include_password=True)
|
||
if result.returncode != 0:
|
||
raise RuntimeError(f"pg_restore 执行失败，退出码 {result.returncode}")
|
||
|
||
def _connection_arguments(self) -> list[str]:
|
||
url = self._engine.url
|
||
host = str(url.query.get("host") or url.host or "")
|
||
port = str(url.query.get("port") or url.port or "")
|
||
arguments = [
|
||
"--username",
|
||
str(url.username or ""),
|
||
"--dbname",
|
||
str(url.database or ""),
|
||
]
|
||
if host:
|
||
arguments.extend(["--host", host])
|
||
if port:
|
||
arguments.extend(["--port", port])
|
||
return arguments
|
||
|
||
def _run(self, command: Sequence[str], *, include_password: bool) -> ProcessResult:
|
||
return self._runner(
|
||
command,
|
||
env=self._environment(include_password=include_password),
|
||
capture_output=True,
|
||
text=True,
|
||
check=False,
|
||
)
|
||
|
||
def _require_tool(self, executable: str) -> str:
|
||
return _require_tool(executable, self._tool_resolver)
|
||
|
||
def _environment(self, *, include_password: bool) -> dict[str, str]:
|
||
password = (
|
||
str(self._engine.url.password)
|
||
if include_password and self._engine.url.password
|
||
else None
|
||
)
|
||
sslmode = self._engine.url.query.get("sslmode")
|
||
return _postgres_environment(
|
||
password=password,
|
||
sslmode=str(sslmode) if sslmode else None,
|
||
)
|