NeteaseUrl/auth_system/backup.py

# -*- coding: utf-8 -*-
"""数据库备份、恢复与自动备份线程。

（自 auth_system.py 拆出，行为与原实现一致；DB 备份改用 aiosqlite 在线备份
API，文件系统/复制等阻塞 IO 经 asyncio.to_thread 卸载。）
"""

import asyncio
import math
import os
import shutil
import sqlite3
import sys
import threading
import time
from pathlib import Path
from typing import Any

import aiosqlite

import auth_system.db as db  # 绝对导入子模块：不经过包属性，避免与包 __init__ 互指成环

from .settings import (
    get_setting,
    set_setting,
)

# ============ 数据库备份 ============


# 备份命名专用锁：占位命名不阻塞写请求，但同秒并发的两次备份仍需保证文件名唯一
_BACKUP_NAME_LOCK = threading.Lock()


def _reserve_backup_path() -> Path:
    """锁内挑选唯一文件名并创建占位文件（同步微秒级本地 IO，仅建一个空文件）。

    先占位（而不是备份完再命名）保证同秒并发下命名唯一：备份前的空文件与
    sqlite3.connect 建出的空库等价，随后在线备份写入其内容。
    """
    with _BACKUP_NAME_LOCK:
        base = 'backup_' + time.strftime('%Y%m%d_%H%M%S')
        name = base + '.db'
        k = 1
        while (db.BACKUP_DIR / name).exists():
            name = f'{base}-{k}.db'
            k += 1
        dest = db.BACKUP_DIR / name
        dest.touch()
        return dest


async def create_backup() -> str:
    """创建备份（唯一文件名——同秒冲突自动追加序号，不再互相覆盖）。

    WAL 模式下用独立只读连接走 SQLite 在线备份 API（aiosqlite Connection.backup）：
    拿到的是一致性快照，且不持有写锁——原实现备份期间阻塞全站写请求，
    大库时会卡住所有用户。"""
    db.BACKUP_DIR.mkdir(parents=True, exist_ok=True)
    dest_path = await asyncio.to_thread(_reserve_backup_path)
    name = dest_path.name
    try:
        conn = await db.connect()
        try:
            dest = await aiosqlite.connect(str(dest_path))
            try:
                await conn.backup(dest)
            finally:
                await dest.close()
        finally:
            await conn.close()
        # 任何备份产生后统一裁剪（手动/自动/上传同配额）
        try:
            await prune_backups(20, 5)
        except Exception:
            pass
        return name
    except Exception:
        # 备份失败：清掉占位文件，不留半个坏备份
        try:
            await asyncio.to_thread(dest_path.unlink, missing_ok=True)
        except OSError:
            pass
        raise


def _list_backups_sync() -> list[dict[str, Any]]:
    """备份列表（按文件修改时间倒序——按名字排序会让 upload_ 永远排前，误删新自动备份）"""
    if not db.BACKUP_DIR.exists():
        return []
    out = []
    for p in db.BACKUP_DIR.glob('*.db'):
        st = p.stat()
        out.append({'name': p.name, 'size': st.st_size,
                    'time': time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(st.st_mtime)),
                    '_mtime': st.st_mtime})
    out.sort(key=lambda x: x['_mtime'], reverse=True)
    for o in out:
        o.pop('_mtime', None)
    return out


async def list_backups() -> list[dict[str, Any]]:
    return await asyncio.to_thread(_list_backups_sync)


def _backup_path_sync(name: str) -> Path | None:
    name = Path(name).name  # 防目录穿越
    p = db.BACKUP_DIR / name
    return p if p.exists() and p.suffix == '.db' else None


async def backup_path(name: str) -> Path | None:
    return await asyncio.to_thread(_backup_path_sync, name)


def _backup_is_valid_db_sync(src: Path) -> bool:
    """深度校验备份文件：SQLite 头 + 能真正打开 + 含 users 表，三者缺一不可"""
    try:
        with open(src, 'rb') as f:
            if f.read(16) != b'SQLite format 3\x00'[:16]:
                return False
        probe = sqlite3.connect(f'file:{src}?mode=ro', uri=True)
        try:
            tables = {r[0] for r in probe.execute("SELECT name FROM sqlite_master WHERE type='table'")}
            return 'users' in tables
        finally:
            probe.close()
    except Exception:
        return False


async def _backup_is_valid_db(src: Path) -> bool:
    """同步 sqlite3 探测（打开只读库）经 to_thread 卸载，不阻塞事件循环。"""
    return await asyncio.to_thread(_backup_is_valid_db_sync, src)


# 恢复前静止钩子——由应用层注册（取消全部在途下载任务），避免 worker 按旧库账目结算
_before_restore_hook = None


def register_before_restore_hook(fn) -> None:
    """注册恢复前钩子（纯内存登记，无 IO，保持同步）。"""
    global _before_restore_hook
    _before_restore_hook = fn


async def restore_backup(name: str) -> None:
    """恢复备份（先验证、留回滚、后替换、失败自动回滚，源备份文件保留不移动）
    恢复前先通过钩子取消全部在途任务（worker 不再引用旧库用户结算）

    整个替换过程持**当前事件循环的写锁**（经 db._write_lock()）串行化；新库
    结构初始化调用 db._init_db_locked（不自行取锁）以避免 asyncio.Lock 非重入
    造成的自锁死。
    """
    src = await backup_path(name)
    if not src:
        raise ValueError('备份文件不存在')
    if not await _backup_is_valid_db(src):
        raise ValueError('不是有效的站点数据库备份（校验失败，未做任何修改）')
    if _before_restore_hook is not None:
        try:
            _before_restore_hook()
        except Exception:
            pass
    safety = db.BACKUP_DIR / ('pre_restore_' + time.strftime('%Y%m%d_%H%M%S') + '.db')
    async with db._write_lock():
        # 0) 先尽力清 WAL 并合并进主库（被占用导致 checkpoint 忙时明确失败，不假装成功）；
        #    必须在留安全副本之前做——否则副本缺 WAL 里未合并的最近提交，回滚点不完整
        conn = await db.connect()
        try:
            row = await (await conn.execute('PRAGMA wal_checkpoint(TRUNCATE)')).fetchone()
            busy = bool(row and row[0] == 1)  # 1 = busy
            if busy:
                raise RuntimeError('数据库当前有其他读写在途（checkpoint 忙），请稍后重试恢复')
            # 1) checkpoint 后留安全副本（包含全部已提交数据，可回滚）
            await asyncio.to_thread(shutil.copy2, str(db.DB_PATH), str(safety))
            # 2) 源备份复制到临时文件（不动源备份本身）
            tmp = db.DB_PATH.with_suffix('.db.restore_tmp')
            await asyncio.to_thread(shutil.copy2, str(src), str(tmp))
        finally:
            # 先关闭旧库连接：Windows 下替换被打开的库文件会失败
            await conn.close()
        # 3) 原子替换 + 立即校验
        await asyncio.to_thread(os.replace, str(tmp), str(db.DB_PATH))
        newconn = await db.connect()
        try:
            await db._init_db_locked(newconn)
        except Exception:
            await newconn.close()
            # 替换后初始化失败：回滚安全副本，恢复原库
            await asyncio.to_thread(shutil.copy2, str(safety), str(db.DB_PATH))
            rc = await db.connect()
            try:
                await db._init_db_locked(rc)
            finally:
                await rc.close()
            raise RuntimeError('恢复后数据库初始化失败，已自动回滚到恢复前状态') from None
        finally:
            await newconn.close()
    # 裁剪移出写锁——glob/unlink 的 I/O 不阻塞全站写
    try:
        await prune_backups(20, 5)
    except Exception:
        pass


async def delete_backup(name: str) -> bool:
    p = await backup_path(name)
    # 允许删除 pre_restore_ 安全副本（配额裁剪已保证至少保留
    # 最近 5 份；原实现一刀切拒删导致恢复频繁时无界累积且管理员手动也删不掉）
    if p:
        await asyncio.to_thread(p.unlink)
        return True
    return False


# 自动备份线程
_backup_thread_started = False

_BACKUP_ERR_TS = {'ts': 0.0, 'n': 0}   # 备份失败限频记录（失败不能零痕迹）


async def prune_backups(keep_normal: int = 20, keep_pre_restore: int = 5) -> None:
    """备份配额裁剪（从自动备份循环抽出为独立函数）——
    非pre_restore 保留最近 N 份、pre_restore_ 安全副本保留最近 M 份。
    在 create_backup/上传备份/恢复备份后统一调用（原实现只在自动备份循环里
    裁剪，手动备份/上传/恢复会让 data/backups 无界累积）。"""
    baks = [b for b in await list_backups() if not b['name'].startswith('pre_restore_')]
    for b in baks[keep_normal:]:
        try:
            await delete_backup(b['name'])
        except Exception:
            pass
    pre = [b for b in await list_backups() if b['name'].startswith('pre_restore_')]
    for b in pre[keep_pre_restore:]:
        try:
            await asyncio.to_thread((db.BACKUP_DIR / b['name']).unlink)
        except OSError:
            pass


async def _auto_backup_tick() -> None:
    """自动备份单次检查（在独立事件循环里运行，见 _auto_backup_loop）。"""
    interval = float(await get_setting('backup_interval_hours', '0') or 0)
    # 非有限值（NaN/inf）或负数一律视为关闭
    if not math.isfinite(interval):
        interval = 0
        await set_setting('backup_interval_hours', '0')
    if interval > 0:
        last = float(await get_setting('last_backup_ts', '0') or 0)
        if not math.isfinite(last):
            last = 0
        if time.time() - last >= interval * 3600:
            await create_backup()
            await set_setting('last_backup_ts', str(time.time()))
            # 仅保留最近 20 份（按修改时间，不按文件名）
            # pre_restore_* 安全副本不计入 20 份配额，仅各自保留最近 5 份
            await prune_backups(20, 5)


def _auto_backup_loop() -> None:
    """自动备份后台**线程**主循环（线程无法 await，每次 tick 用 asyncio.run 起独立 loop）。"""
    while True:
        try:
            asyncio.run(_auto_backup_tick())
        except Exception:
            # 定时备份失败绝不能静默（磁盘满/权限问题时 owner 以为有备份、
            # 出事时没有）——限频告警，每 10 分钟最多一条（对齐监控线程口径）
            import traceback as _tb
            _now = time.time()
            _BACKUP_ERR_TS['n'] += 1
            if _now - _BACKUP_ERR_TS['ts'] > 600:
                print(f'[警告] 自动备份失败（近 10 分钟累计 {_BACKUP_ERR_TS["n"]} 次）: '
                      f'{_tb.format_exc()[-400:]}', file=sys.stderr)
                _BACKUP_ERR_TS['ts'] = _now
                _BACKUP_ERR_TS['n'] = 0
        time.sleep(600)


def start_auto_backup() -> None:
    """启动自动备份后台线程（纯线程启动，无 IO，保持同步）。"""
    global _backup_thread_started
    if _backup_thread_started:
        return
    t = threading.Thread(target=_auto_backup_loop, daemon=True)
    t.start()
    # 仅在真实启动成功后置位，失败可重试
    _backup_thread_started = True