Archived
snowy-yang
9a34517454
为什么：同步层每次数据库操作都新建连接并执行 mkdir + PRAGMA journal_mode=WAL。实测拆解（各 100 次中位数）： mkdir(exist_ok) 190us sqlite3.connect 435us PRAGMA WAL(新连接首语句) ~1000us ← 大头 get_conn 合计 1641us → 501us 登录路径走同步层 4 次查询（身份解析/验证/位置/日志），每次请求 固定省 ~4.6ms；下载任务的 ledger/日志调用同样受益。 改动（对齐 core/db.py 既有模式——异步层早已做过同款优化并记录过 「并发 PRAGMA 争排他锁报 database is locked」的实测坑，同步层漏同步）： - _WAL_READY 按路径缓存：journal_mode 持久化在库文件头，设一次即可； _WAL_SETUP_LOCK 串行化首连切换（读路径不持 _DB_LOCK） - mkdir 移出 happy path：connect 失败才兜底补建（用 DB_PATH.parent， monkeypatch 场景下与 DATA_DIR 可能不同） - exists() 检查在 connect 之前（connect 会创建空文件，之后恒真）， 覆盖「删库重建后新文件回退 delete 模式」的边界 新增 tests/test_db_conn.py 6 例：WAL 首连生效、后续连接零 PRAGMA （set_trace_callback 实证）、换路径重设、目录缺失自动建、删库重建 恢复 WAL、8 线程并发首连无报错。 门禁：ruff 全过 · pyright 0/0 · pytest 382 passed（375→382）。
306 lines
12 KiB
Python
306 lines
12 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""路径常量、连接与建表。
|
||
|
||
（自 auth_system.py 拆出，行为与原实现一致。）
|
||
"""
|
||
|
||
import os
|
||
import secrets
|
||
import sqlite3
|
||
import sys
|
||
import threading
|
||
import time
|
||
from pathlib import Path
|
||
|
||
from .security import (
|
||
_hash_new,
|
||
_normalize_email,
|
||
)
|
||
|
||
# ============ 路径与常量 ============
|
||
# 必须与 core.db.DATA_DIR 用**同一条**表达式：两处指向同一个 data/site.db。
|
||
# 此前这里写死 Path('data') 而 core 侧认 SITE_DATA_DIR——一旦运维设置该变量，
|
||
# 站点库与设置/公告库会落到两个不同文件（脑裂），且不会报错。
|
||
DATA_DIR = Path(os.environ.get('SITE_DATA_DIR') or 'data')
|
||
|
||
DB_PATH = DATA_DIR / 'site.db'
|
||
|
||
BACKUP_DIR = DATA_DIR / 'backups'
|
||
|
||
SECRET_KEY_PATH = DATA_DIR / 'secret.key'
|
||
|
||
|
||
DEFAULT_ADMIN_USER = 'admin'
|
||
|
||
SESSION_KEY = 'site_uid'
|
||
|
||
DEFAULT_REGISTER_CREDITS = 3 # 注册直接赠送 3 次下载次数（无需领取；TRIAL_CREDITS 死常量已删）
|
||
|
||
|
||
_DB_LOCK = threading.RLock()
|
||
|
||
# journal_mode=WAL 持久化在库文件头里，按路径设一次即可。每连接重复执行
|
||
# 实测多付 ~1ms（新连接首条语句要建 WAL 索引读锁），且并发首连会争排他锁
|
||
# 报 database is locked（core/db.py 同款教训）。专用锁 + 二次复查防护。
|
||
_WAL_READY: set[str] = set()
|
||
_WAL_SETUP_LOCK = threading.Lock()
|
||
|
||
|
||
# ============ 数据库 ============
|
||
|
||
|
||
def get_conn() -> sqlite3.Connection:
|
||
key = str(DB_PATH)
|
||
# exists() 检查（~10us）必须在 connect 之前：connect 会创建空文件，
|
||
# 之后 exists 恒为 True，删库重建的场景就检测不到了
|
||
force_wal = not DB_PATH.exists()
|
||
# mkdir 移出 happy path（实测 190us/次）：connect 直接失败（目录缺失）
|
||
# 才补建重试。库文件被删的场景同样由此兜底重建。
|
||
try:
|
||
conn = sqlite3.connect(str(DB_PATH), timeout=15)
|
||
except sqlite3.OperationalError:
|
||
# DB_PATH.parent 而非 DATA_DIR：生产两者相同，测试 monkeypatch 只换
|
||
# DB_PATH 时兜底仍能命中正确目录
|
||
DB_PATH.parent.mkdir(parents=True, exist_ok=True)
|
||
conn = sqlite3.connect(str(DB_PATH), timeout=15)
|
||
conn.row_factory = sqlite3.Row
|
||
if key not in _WAL_READY or force_wal:
|
||
# 读路径不持 _DB_LOCK，多线程可能同时首连：锁内串行执行 journal_mode
|
||
# 切换（防并发争排他锁报 database is locked）。锁内不复查 DB_PATH.exists()——
|
||
# connect 已把文件建出来，exists 恒真；force_wal 是本线程的观测结论，
|
||
# 重复执行 PRAGMA 在已是 WAL 的库上幂等且廉价（~2us）
|
||
with _WAL_SETUP_LOCK:
|
||
conn.execute('PRAGMA journal_mode=WAL')
|
||
_WAL_READY.add(key)
|
||
return conn
|
||
|
||
|
||
def _q(sql: str, params: tuple = ()) -> list[sqlite3.Row]:
|
||
"""单语句读：**不持锁**。
|
||
|
||
WAL 模式下每个连接读到一致性快照，读与写互不阻塞；旧实现对读也加 _DB_LOCK，
|
||
等于把全部读串行化并被写挡住。此处仅放开单语句读——写事务（_x 与各
|
||
`with _DB_LOCK` 块）仍全程串行，故不引入写-写竞争。
|
||
"""
|
||
conn = get_conn()
|
||
try:
|
||
return conn.execute(sql, params).fetchall()
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
def _x(sql: str, params: tuple = ()) -> int:
|
||
"""执行写操作，返回 lastrowid / rowcount"""
|
||
with _DB_LOCK:
|
||
conn = get_conn()
|
||
try:
|
||
cur = conn.execute(sql, params)
|
||
conn.commit()
|
||
return cur.lastrowid if cur.lastrowid else cur.rowcount
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
_SCHEMA_SQL = '''
|
||
CREATE TABLE IF NOT EXISTS users (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
username TEXT UNIQUE NOT NULL,
|
||
password_hash TEXT NOT NULL,
|
||
salt TEXT NOT NULL,
|
||
email TEXT DEFAULT '',
|
||
role TEXT DEFAULT 'user',
|
||
duration_seconds INTEGER DEFAULT 0,
|
||
trial_device TEXT DEFAULT '',
|
||
register_ip TEXT DEFAULT '',
|
||
register_ua TEXT DEFAULT '',
|
||
netease_cookie TEXT DEFAULT '',
|
||
ann_seen TEXT DEFAULT '',
|
||
last_active INTEGER DEFAULT 0,
|
||
created_at INTEGER DEFAULT 0,
|
||
last_login INTEGER DEFAULT 0
|
||
);
|
||
CREATE TABLE IF NOT EXISTS email_codes (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
email TEXT NOT NULL,
|
||
code TEXT NOT NULL,
|
||
expires_at INTEGER NOT NULL,
|
||
created_at INTEGER DEFAULT 0,
|
||
purpose TEXT DEFAULT 'register'
|
||
);
|
||
CREATE TABLE IF NOT EXISTS cards (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
code TEXT UNIQUE NOT NULL,
|
||
minutes INTEGER NOT NULL,
|
||
created_by TEXT DEFAULT '',
|
||
created_at INTEGER DEFAULT 0,
|
||
used_by TEXT DEFAULT '',
|
||
used_at INTEGER DEFAULT 0
|
||
);
|
||
CREATE TABLE IF NOT EXISTS logs (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
type TEXT DEFAULT 'success',
|
||
username TEXT DEFAULT '',
|
||
ip TEXT DEFAULT '',
|
||
action TEXT DEFAULT '',
|
||
detail TEXT DEFAULT '',
|
||
time TEXT DEFAULT ''
|
||
);
|
||
'''
|
||
|
||
|
||
def _create_schema(conn: sqlite3.Connection) -> None:
|
||
"""建表 + 单行配置种子行（幂等：全部 IF NOT EXISTS / INSERT OR IGNORE）。"""
|
||
conn.executescript(_SCHEMA_SQL)
|
||
# 单行配置表
|
||
conn.execute('''CREATE TABLE IF NOT EXISTS smtp_config (
|
||
id INTEGER PRIMARY KEY CHECK (id = 1),
|
||
host TEXT DEFAULT '', port INTEGER DEFAULT 587,
|
||
sender TEXT DEFAULT '', password TEXT DEFAULT '',
|
||
updated_at INTEGER DEFAULT 0
|
||
)''')
|
||
conn.execute('INSERT OR IGNORE INTO smtp_config (id, port) VALUES (1, 587)')
|
||
# 键值设置表
|
||
conn.execute('''CREATE TABLE IF NOT EXISTS settings (
|
||
key TEXT PRIMARY KEY, value TEXT DEFAULT ''
|
||
)''')
|
||
conn.execute('''CREATE TABLE IF NOT EXISTS announcements (
|
||
id INTEGER PRIMARY KEY CHECK (id = 1),
|
||
title TEXT DEFAULT '', content TEXT DEFAULT '',
|
||
popup_mode TEXT DEFAULT 'off',
|
||
updated_at INTEGER DEFAULT 0
|
||
)''')
|
||
conn.execute('INSERT OR IGNORE INTO announcements (id) VALUES (1)')
|
||
# 计费账本（任务实扣/交付/退款/终态的持久化唯一记录，结算幂等）
|
||
conn.execute('''CREATE TABLE IF NOT EXISTS billing_ledger (
|
||
task_id TEXT PRIMARY KEY,
|
||
user_id INTEGER NOT NULL,
|
||
reserved INTEGER DEFAULT 0,
|
||
delivered INTEGER DEFAULT 0,
|
||
refunded INTEGER DEFAULT 0,
|
||
status TEXT DEFAULT 'open',
|
||
created_at INTEGER DEFAULT 0,
|
||
settled_at INTEGER DEFAULT 0
|
||
)''')
|
||
|
||
|
||
def _apply_column_migrations(conn: sqlite3.Connection) -> None:
|
||
"""轻量迁移：新增列。仅"列已存在"可忽略（其余迁移失败必须暴露，不能吞掉）"""
|
||
_already: set[str] = set()
|
||
for trow in conn.execute("SELECT name FROM sqlite_master WHERE type='table'").fetchall():
|
||
for crow in conn.execute(f"PRAGMA table_info({trow['name']})").fetchall():
|
||
_already.add(f"{trow['name']}.{crow['name']}")
|
||
for alter_sql, tbl_col in (
|
||
('ALTER TABLE users ADD COLUMN status INTEGER DEFAULT 1', 'users.status'),
|
||
('ALTER TABLE users ADD COLUMN login_seq INTEGER DEFAULT 0', 'users.login_seq'),
|
||
('ALTER TABLE users ADD COLUMN recent_provinces TEXT DEFAULT \'[]\'', 'users.recent_provinces'),
|
||
('ALTER TABLE users ADD COLUMN ban_reason TEXT DEFAULT \'\'', 'users.ban_reason'),
|
||
('ALTER TABLE cards ADD COLUMN note TEXT DEFAULT \'\'', 'cards.note'),
|
||
('ALTER TABLE cards ADD COLUMN expire_at INTEGER DEFAULT 0', 'cards.expire_at'),
|
||
('ALTER TABLE cards ADD COLUMN disabled INTEGER DEFAULT 0', 'cards.disabled'),
|
||
('ALTER TABLE email_codes ADD COLUMN purpose TEXT DEFAULT \'register\'', 'email_codes.purpose'),
|
||
):
|
||
if tbl_col in _already:
|
||
continue # 幂等：列已存在
|
||
try:
|
||
conn.execute(alter_sql)
|
||
except sqlite3.OperationalError as e:
|
||
if 'duplicate column' not in str(e).lower():
|
||
# 磁盘/锁/结构等真实故障：抛出，禁止静默带病初始化
|
||
raise
|
||
|
||
|
||
def _ensure_email_unique_index(conn: sqlite3.Connection) -> None:
|
||
"""非空邮箱唯一（部分索引，不约束历史空串）。
|
||
|
||
存量库若已有重复非空邮箱，索引创建失败仅告警不阻断启动
|
||
（查重逻辑在注册/建号路由另有兜底）。
|
||
"""
|
||
try:
|
||
conn.execute('CREATE UNIQUE INDEX IF NOT EXISTS idx_users_email_nonempty '
|
||
'ON users(email) WHERE email != \'\'')
|
||
except sqlite3.IntegrityError:
|
||
print('[警告] users.email 存在重复的非空邮箱，唯一索引未创建（请人工清理后重启）')
|
||
|
||
|
||
def _migrate_user_emails(conn: sqlite3.Connection) -> None:
|
||
"""users.email 口径统一为归一化（lower + 剥 Gmail +tag）——与
|
||
get_user_by_email/找回/邮箱登录的查询口径对齐。逐行迁移，归一化后与他
|
||
人撞唯一索引的行保留原值跳过（与"存量重复仅告警"同口径）。幂等：已
|
||
归一化的行零写入。
|
||
"""
|
||
for urow in conn.execute("SELECT id, email FROM users WHERE email != ''").fetchall():
|
||
fixed = _normalize_email(urow['email'])
|
||
if not fixed or fixed == urow['email']:
|
||
continue
|
||
try:
|
||
conn.execute('UPDATE users SET email = ? WHERE id = ?', (fixed, urow['id']))
|
||
except sqlite3.IntegrityError:
|
||
print(f"[警告] 用户 {urow['id']} 的邮箱归一化后与他人重复，保留原值")
|
||
|
||
|
||
def _ensure_initial_admin(conn: sqlite3.Connection) -> None:
|
||
"""初始管理员（新装不再使用公开默认口令，随机生成并写入日志一次）。
|
||
|
||
引导条件改为"不存在任何管理员角色"——按用户名查会在
|
||
管理员改名（如 admin→qfbsh）后误判缺失，重启即重建出一个新的 admin 账号。
|
||
"""
|
||
row = conn.execute("SELECT id FROM users WHERE role = 'admin' LIMIT 1").fetchone()
|
||
if row:
|
||
return
|
||
initial_pass = secrets.token_urlsafe(12)
|
||
conn.execute(
|
||
'INSERT INTO users (username, password_hash, salt, role, duration_seconds, created_at) '
|
||
'VALUES (?, ?, ?, ?, -1, ?)',
|
||
(DEFAULT_ADMIN_USER, _hash_new(initial_pass), '',
|
||
'admin', int(time.time())))
|
||
# ——非 ASCII 输出须防 GBK 重定向下 UnicodeEncodeError
|
||
try:
|
||
print(f'[初始化] 已创建管理员账号 admin，初始密码：{initial_pass}（请立即登录并修改，本提示不再显示）')
|
||
except UnicodeEncodeError:
|
||
print('[init] admin account created, initial password printed above failed to encode')
|
||
|
||
|
||
def init_db() -> None:
|
||
"""初始化表结构与默认数据"""
|
||
DATA_DIR.mkdir(parents=True, exist_ok=True)
|
||
with _DB_LOCK:
|
||
conn = get_conn()
|
||
try:
|
||
_create_schema(conn)
|
||
_apply_column_migrations(conn)
|
||
_ensure_email_unique_index(conn)
|
||
_migrate_user_emails(conn)
|
||
_ensure_initial_admin(conn)
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
def get_secret_key() -> str:
|
||
"""获取（或生成）持久化的会话签名密钥（原子写入 + 失败显式告警）"""
|
||
key = ''
|
||
try:
|
||
if SECRET_KEY_PATH.exists():
|
||
key = SECRET_KEY_PATH.read_text(encoding='utf-8').strip()
|
||
if key:
|
||
return key
|
||
except OSError:
|
||
pass
|
||
new_key = secrets.token_hex(32)
|
||
try:
|
||
SECRET_KEY_PATH.parent.mkdir(parents=True, exist_ok=True)
|
||
# 先写临时文件再原子替换，避免部分写入留下损坏密钥
|
||
tmp = SECRET_KEY_PATH.with_suffix('.key.tmp')
|
||
tmp.write_text(new_key, encoding='utf-8')
|
||
os.replace(str(tmp), str(SECRET_KEY_PATH))
|
||
# 并发初始化者以"替换后文件里的值"为准（最后写入者胜出但全员一致）
|
||
final = SECRET_KEY_PATH.read_text(encoding='utf-8').strip()
|
||
return final or new_key
|
||
except OSError as e:
|
||
# 写入失败不能装作成功：显著告警，返回临时值（本进程会话可用，重启后失效）
|
||
try:
|
||
print(f'[警告] 会话密钥持久化失败（{e}），本进程将使用临时密钥，重启后所有会话失效！', file=sys.stderr)
|
||
except UnicodeEncodeError:
|
||
print('[warn] secret key persistence failed; using temporary key for this process', file=sys.stderr)
|
||
return new_key
|