mirror of
https://github.com/jxxghp/MoviePilot.git
synced 2026-11-01 21:07:01 +08:00
341 lines
11 KiB
Python
341 lines
11 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, Literal, Mapping, Protocol, Sequence
|
||
|
||
from sqlalchemy.engine import Engine
|
||
|
||
|
||
@dataclass(frozen=True, slots=True)
|
||
class DatabaseBackupCheck:
|
||
"""
|
||
数据库备份校验的结论与诊断信息。
|
||
|
||
:param valid: 校验是否通过；True 不代表已完成恢复演练
|
||
:param method: 实际使用的检查命令或方法名称
|
||
:param detail: 失败或异常说明；校验通过时通常为 None
|
||
"""
|
||
|
||
valid: bool
|
||
method: str
|
||
detail: str | None = None
|
||
|
||
|
||
class ProcessResult(Protocol):
|
||
"""
|
||
数据库客户端进程退出后可读取的最小结果接口。
|
||
|
||
:ivar returncode: 进程退出码，0 表示命令成功
|
||
:ivar stdout: 已捕获并解码的标准输出文本
|
||
:ivar stderr: 已捕获并解码的标准错误文本
|
||
"""
|
||
|
||
returncode: int
|
||
stdout: str
|
||
stderr: str
|
||
|
||
|
||
class ProcessRunner(Protocol):
|
||
"""
|
||
可替换的数据库命令执行边界。
|
||
"""
|
||
|
||
def __call__(
|
||
self,
|
||
command: Sequence[str],
|
||
/,
|
||
*,
|
||
env: Mapping[str, str],
|
||
capture_output: Literal[True],
|
||
text: Literal[True],
|
||
check: bool,
|
||
) -> ProcessResult:
|
||
"""
|
||
执行数据库客户端命令并等待退出，以文本模式捕获输出。
|
||
|
||
:param command: 可执行文件与独立参数组成的序列，不交给 shell 解析
|
||
:param env: 此次进程使用的完整环境变量，包含经隔离的数据库连接凭据
|
||
:param capture_output: 固定为 True，要求捕获标准输出与标准错误
|
||
:param text: 固定为 True，要求返回已解码的文本
|
||
:param check: 是否在非零退出码时抛出进程执行异常
|
||
:return: 退出码与两路输出；调用结束时子进程已经退出
|
||
:raises subprocess.CalledProcessError: check 为 True 且进程退出码非零
|
||
"""
|
||
|
||
|
||
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:
|
||
"""
|
||
在不连接活动数据库的前提下校验受管备份文件。
|
||
|
||
SQLite 运行完整性检查；PostgreSQL 只读取归档目录，不进行恢复演练。
|
||
|
||
:param artifact: 待检查的备份文件路径
|
||
:param db_type: 数据库类型，只接受 sqlite 或 postgresql
|
||
:param runner: 可替换的 PostgreSQL 客户端命令执行器
|
||
:param tool_resolver: 按名称定位数据库客户端可执行文件的回调
|
||
:param pg_restore: PostgreSQL 归档读取工具名称或路径
|
||
:return: 完整性或归档目录检查结果，含失败说明
|
||
:raises ValueError: 数据库类型不支持
|
||
:raises RuntimeError: PostgreSQL 客户端工具不可用
|
||
"""
|
||
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]:
|
||
"""
|
||
隔离继承的 PostgreSQL 凭证，仅传入本次操作所需的连接设置。
|
||
"""
|
||
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:
|
||
"""
|
||
绑定活动 SQLite 文件路径，拒绝无法生成文件快照的内存数据库。
|
||
"""
|
||
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 查询项中的主机和端口。
|
||
"""
|
||
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:
|
||
"""
|
||
通过当前解析器确认 PostgreSQL 客户端可用。
|
||
"""
|
||
return _require_tool(executable, self._tool_resolver)
|
||
|
||
def _environment(self, *, include_password: bool) -> dict[str, str]:
|
||
"""
|
||
从活动引擎提取当前命令需要的密码与 SSL 设置。
|
||
"""
|
||
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,
|
||
)
|