workbuddy-portal/workbuddy_portal/db.py

# -*- coding: utf-8 -*-
# SPDX-License-Identifier: MIT
# Copyright (c) 2026 Wang Chuanli

"""SQLite 访问层（多用户版）。

并发约定（重要）：
  * WAL 模式 —— 采集写入期间页面查询不会被 `database is locked` 挡住
  * 单写者 —— SQLite 同一时刻只允许一个写进程，所以采集必须串行
    （由 scheduler / CLI 共享的 collect.lock 保证）
  * busy_timeout=8s —— 偶发并发时等待而不是立刻报错
  * 每个线程独立连接（sqlite3 默认禁止跨线程复用连接）

多用户约定（改代码前务必先读）：
  * `user_id = 0` 在 settings / collect_runs / audit_log 里表示**实例级**；
    usage_records 里 0 是「历史遗留数据尚未归属」的兜底值，正常不会出现。
  * **settings 里只有两个键是个人级：`cookie` 与 `user_agent`**（即
    `config.USER_EDITABLE_KEYS`）。调度、采集参数、接口地址、注册策略全部
    是实例级（`config.GLOBAL_KEYS`），`set_setting()` 会把它们强制写到
    user_id=0 —— 因此「管理员改了但别人不生效」这类 bug 在结构上不存在。
  * `get_settings()` 会把 `ENCRYPTED_KEYS`（Cookie）**一律置空**；
    要拿明文只有 `get_secret()` 一条路。这样任何「顺手打印一下全部配置」
    的代码都不可能把凭证带出去。
  * `NO_FALLBACK_KEYS`（Cookie / User-Agent）**不参与实例级回退**：
    Cookie 是账号凭证，回落等于串号，是最严重的一类越权。
  * `slot:*` 调度簿记键是**个人级**（每个账号各自记「今天这个槽位跑过没」），
    虽然时刻本身是实例级的 —— 这两个千万别一起改。
"""
import os
import sqlite3
import threading
from datetime import datetime

from . import config, crypto

_local = threading.local()
_init_lock = threading.Lock()
_initialized = False

# 库结构版本。写在 PRAGMA user_version 里，用来判断是否需要迁移。
#   1 -> 单用户布局（settings 以 key 为主键，usage_records 以 request_id 为主键）
#   2 -> 多用户布局（见 schema.sql 顶部说明）
#   3 -> 调度与采集参数从个人级提升为实例级（普通用户只读；见 config.GLOBAL_KEYS）
DB_SCHEMA_VERSION = 3

# 这些键即使个人作用域没有值，也**不**回落到实例级
NO_FALLBACK_KEYS = {"cookie", "user_agent"}


class SecretUnreadable(Exception):
    """密文解不开 —— 通常是 data/instance.json 里的 cookie_key 被换过。"""


def now_str():
    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")


def connect(path=None):
    """新建一个连接（调用方负责关闭）。"""
    config.ensure_dirs()
    conn = sqlite3.connect(path or config.SQLITE_PATH, timeout=8.0,
                           isolation_level=None)   # autocommit，事务用显式 BEGIN
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA synchronous=NORMAL")
    conn.execute("PRAGMA busy_timeout=8000")
    conn.execute("PRAGMA foreign_keys=ON")
    return conn


def thread_conn():
    """按线程缓存的连接（采集线程、调度线程各自一份）。"""
    c = getattr(_local, "conn", None)
    if c is None:
        c = _local.conn = connect()
    return c


def close_thread_conn():
    c = getattr(_local, "conn", None)
    if c is not None:
        try:
            c.close()
        except sqlite3.Error:
            pass
        _local.conn = None


def _schema_sql():
    with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), "schema.sql"),
              "r", encoding="utf-8") as f:
        return f.read()


# ---------------- 初始化 / 迁移 ----------------
def _table_cols(conn, table):
    return {r["name"] for r in conn.execute("PRAGMA table_info(%s)" % table)}


def _has_table(conn, name):
    return bool(conn.execute(
        "SELECT 1 FROM sqlite_master WHERE type='table' AND name=?", (name,)).fetchone())


def _first_owner_uid(conn):
    """历史数据归谁：优先第一个管理员，其次第一个账号。"""
    row = conn.execute("SELECT id FROM users WHERE is_admin=1 ORDER BY id LIMIT 1").fetchone()
    if row:
        return row["id"]
    row = conn.execute("SELECT id FROM users ORDER BY id LIMIT 1").fetchone()
    return row["id"] if row else 0


def _drop_all_user_indexes(conn):
    """删掉本项目自建的全部索引（idx_*）。

    必须先删：`ALTER TABLE ... RENAME TO` 会**把索引一起带走**（名字仍指向
    改名后的表），于是后面 `CREATE INDEX IF NOT EXISTS` 会被当成「已存在」
    静默跳过，最终新表上一个索引都没有 —— 表面完全正常，只是慢几百倍。
    `sqlite_autoindex_*` 是主键/唯一约束的隐式索引，不能动，靠前缀过滤掉。
    """
    names = [r["name"] for r in conn.execute(
        "SELECT name FROM sqlite_master WHERE type='index' AND name LIKE 'idx_%'")]
    for n in names:
        conn.execute('DROP INDEX IF EXISTS "%s"' % n)
    return names


def _migrate(conn):
    """把老库升到 DB_SCHEMA_VERSION。幂等，返回迁移说明列表。

    顺序不能变：
      1. 删索引（否则改名会把索引名占住，新表建不出索引）
      2. 改名主键变了的表（settings / usage_records）
      3. ALTER 加列（collect_runs / audit_log / users）
      4. 跑 schema.sql —— 此时所有列都齐了，表与索引一次建全
      5. 回填数据、删掉 _v1_ 旧表
    """
    ver = conn.execute("PRAGMA user_version").fetchone()[0]
    if ver >= DB_SCHEMA_VERSION:
        return []
    done = []
    ts = now_str()
    owner = _first_owner_uid(conn)

    tset = _table_cols(conn, "settings")
    tur = _table_cols(conn, "usage_records")
    rebuild_settings = "user_id" not in tset
    rebuild_records = "user_id" not in tur

    _drop_all_user_indexes(conn)
    if rebuild_settings:
        conn.execute("ALTER TABLE settings RENAME TO _v1_settings")
    if rebuild_records:
        conn.execute("ALTER TABLE usage_records RENAME TO _v1_usage_records")

    # ---- 只加列的表用 ALTER，代价小得多。放在建表之前，好让 schema.sql
    #      里的 CREATE INDEX 一次就成功（索引引用了这些新列）----
    if "user_id" not in _table_cols(conn, "collect_runs"):
        conn.execute("ALTER TABLE collect_runs ADD COLUMN user_id INTEGER NOT NULL DEFAULT 0")
        conn.execute("UPDATE collect_runs SET user_id=?", (owner,))
        done.append("collect_runs 增加 user_id")
    if "user_id" not in _table_cols(conn, "audit_log"):
        conn.execute("ALTER TABLE audit_log ADD COLUMN user_id INTEGER NOT NULL DEFAULT 0")
        done.append("audit_log 增加 user_id")
    tusers = _table_cols(conn, "users")
    for col, ddl in (("email", "TEXT"),
                     ("status", "TEXT NOT NULL DEFAULT 'active'"),
                     ("register_ip", "TEXT"),
                     ("last_login_ip", "TEXT")):
        if col not in tusers:
            conn.execute("ALTER TABLE users ADD COLUMN %s %s" % (col, ddl))
            done.append("users 增加 %s" % col)
    conn.execute("UPDATE users SET status='active' WHERE status IS NULL OR status=''")

    conn.executescript(_schema_sql())          # 建出新表 + 全部索引

    if rebuild_settings:
        # 老布局里 cookie / user_agent 是实例级的 —— 留在实例级等于
        # 「所有人共用管理员的凭证」，必须归到 owner 名下，且之后永不回落。
        conn.execute("INSERT INTO settings(user_id,key,value,updated_at)"
                     " SELECT 0,key,value,updated_at FROM _v1_settings"
                     " WHERE key NOT IN ('cookie','user_agent')")
        conn.execute("INSERT INTO settings(user_id,key,value,updated_at)"
                     " SELECT ?,key,value,updated_at FROM _v1_settings"
                     " WHERE key IN ('cookie','user_agent') AND value IS NOT NULL AND value <> ''",
                     (owner,))
        conn.execute("DROP TABLE _v1_settings")
        done.append("settings 改为 (user_id,key) 复合主键：旧值归实例级，"
                    "Cookie/UA 已归属账号 #%d" % owner)

    if rebuild_records:
        conn.execute("INSERT INTO usage_records(user_id,request_id,ts,day,hour,model,client,"
                     "credits,prompt,first_seen,last_seen,cloud_ts)"
                     " SELECT ?,request_id,ts,day,hour,model,client,credits,prompt,"
                     "first_seen,last_seen,cloud_ts FROM _v1_usage_records", (owner,))
        conn.execute("DROP TABLE _v1_usage_records")
        done.append("usage_records 增加 user_id，主键改为 (user_id, request_id)，"
                    "历史数据归属账号 #%d" % owner)

    conn.execute("UPDATE users SET is_admin=1 WHERE id=?", (owner,))
    conn.execute("PRAGMA user_version=%d" % DB_SCHEMA_VERSION)
    if done:
        conn.execute("INSERT INTO audit_log(user_id,at,actor,action,detail,ip)"
                     " VALUES(0,?,?,?,?,?)",
                     (ts, "system", "schema_migrate", "；".join(done)[:500], "127.0.0.1"))
    return done


def _encrypt_legacy_secrets(conn):
    """把历史**明文**凭证就地加密。幂等，可反复执行。

    老版本把 Cookie 直接明文写进 settings 表。升级后即便功能正常，
    「库里躺着一段明文凭证」本身就是风险：备份文件、磁盘镜像、
    误提交、排障时的一次 dump 都会把它带出去。
    这里对 `config.ENCRYPTED_KEYS` 里所有**非 v1. 前缀**的值加密一次；
    已加密的值会因前缀判定被跳过，所以每次启动跑一遍是安全的。
    """
    if not config.ENCRYPTED_KEYS:
        return []
    key = config.encryption_key()
    done = []
    for r in conn.execute("SELECT user_id,key,value FROM settings").fetchall():
        if r["key"] not in config.ENCRYPTED_KEYS:
            continue
        raw = r["value"]
        if not raw or crypto.is_encrypted(raw):
            continue
        conn.execute("UPDATE settings SET value=?,updated_at=? WHERE user_id=? AND key=?",
                     (crypto.encrypt(raw, key), now_str(), r["user_id"], r["key"]))
        done.append("settings[uid=%s].%s" % (r["user_id"], r["key"]))
    return done


def _promote_personal_to_global(conn):
    """把「原本个人级、现在实例级」的配置收敛到 user_id=0。幂等，可反复执行。

    v1.3.0 把调度与采集参数从个人级升为实例级（普通用户只读）。升级时
    必须做两件事，否则会静默丢配置：

      1. **把第一个管理员的个人值提升到实例级** —— 管理员此前设的
         `schedule_times=08:00` 若不提升，读取路径会因为「全局键不再看个人
         作用域」而直接跳过它，表现成「设置莫名其妙回到默认值」。
      2. **删掉所有个人作用域里的全局键** —— 留着不会被读（get_settings 会
         跳过），但会让后面排障的人以为「这个键是个人级的」。

    对本来就属于实例级的键（api_base / allow_register 等）这是空操作。
    """
    owner = _first_owner_uid(conn)
    promoted, cleaned = [], 0
    for k in sorted(config.GLOBAL_KEYS):
        if owner:
            mine = conn.execute("SELECT value FROM settings WHERE user_id=? AND key=?",
                                (owner, k)).fetchone()
            if mine is not None and mine["value"] is not None:
                cur = conn.execute("SELECT value FROM settings WHERE user_id=0 AND key=?",
                                   (k,)).fetchone()
                if cur is None or cur["value"] != mine["value"]:
                    set_setting(conn, k, mine["value"], 0)      # 全局键强制落 uid=0
                    promoted.append(k)
        # DELETE 的 rowcount 对「没有匹配行」返回 0，所以这里天然幂等
        cur2 = conn.execute("DELETE FROM settings WHERE user_id<>0 AND key=?", (k,))
        cleaned += cur2.rowcount or 0
    return promoted, cleaned


def init_db(conn=None, create_admin=True, admin_user="admin", admin_password=None):
    """建表 / 迁移 / 灌默认配置。可重复执行（幂等）。返回迁移说明列表。"""
    global _initialized
    own = conn is None
    conn = conn or connect()
    try:
        if not _has_table(conn, "users"):
            # 全新库：schema.sql 一次到位（避免走迁移路径去 ALTER 不存在的表）
            conn.executescript(_schema_sql())
            conn.execute("PRAGMA user_version=%d" % DB_SCHEMA_VERSION)
            migrated = []
        else:
            migrated = _migrate(conn)
        # 无论如何再跑一次：幂等补齐（例如后续版本新增了表/索引，
        # 而老库的 user_version 已经是最新，就不会走 _migrate 了）
        conn.executescript(_schema_sql())
        ts = now_str()
        # 默认配置灌在**实例级**（user_id=0）。个人作用域不预置行，
        # 读取时按「个人 -> 实例 -> DEFAULTS」三级回落，语义更清楚。
        #
        # 凭证类键（cookie / user_agent）**刻意不灌**：它们是个人级的，
        # 实例级存一份没有任何读取路径会用到（NO_FALLBACK_KEYS 挡住了回落），
        # 只会让「读一下实例配置看看」的人拿到一个不该存在的凭证位。
        for k, v in config.DEFAULTS.items():
            if k in config.USER_EDITABLE_KEYS:
                continue
            conn.execute("INSERT OR IGNORE INTO settings(user_id,key,value,updated_at)"
                         " VALUES(0,?,?,?)", (k, v, ts))
        # 顺手清掉历史遗留在实例级的凭证行（老版本曾把它们当实例级配置存过）
        for k in sorted(config.USER_EDITABLE_KEYS):
            conn.execute("DELETE FROM settings WHERE user_id=0 AND key=?", (k,))
        # 顺手把历史明文凭证加密（幂等；新库无事可做）
        enc = _encrypt_legacy_secrets(conn)
        if enc:
            migrated.append("明文凭证已加密：%s" % ", ".join(enc))
            conn.execute("INSERT INTO audit_log(user_id,at,actor,action,detail,ip)"
                         " VALUES(0,?,?,?,?,?)",
                         (ts, "system", "encrypt_secrets",
                          "明文凭证已加密：%s" % ", ".join(enc)[:400], "127.0.0.1"))
        # 调度/采集参数在 v1.3.0 升为实例级：把管理员那份提升上去并清掉个人残留
        promoted, cleaned = _promote_personal_to_global(conn)
        if promoted or cleaned:
            note = "配置作用域收敛：提升 %s 到实例级%s" % (
                ", ".join(promoted) if promoted else "（无）",
                "，清理 %d 条个人级残留" % cleaned if cleaned else "")
            migrated.append(note)
            conn.execute("INSERT INTO audit_log(user_id,at,actor,action,detail,ip)"
                         " VALUES(0,?,?,?,?,?)",
                         (ts, "system", "promote_global_settings", note[:400], "127.0.0.1"))
        if create_admin:
            n = conn.execute("SELECT COUNT(*) FROM users").fetchone()[0]
            if n == 0:
                from .security import hash_password
                pwd = admin_password or "admin123"
                conn.execute(
                    "INSERT INTO users(username,password_hash,display_name,is_admin,status,"
                    " created_at) VALUES(?,?,?,1,'active',?)",
                    (admin_user, hash_password(pwd), "管理员", ts))
        _initialized = True
        return migrated
    finally:
        if own:
            conn.close()


# ---------------- 配置读写（按作用域） ----------------
def get_setting(conn, key, default=None, uid=0):
    """取单个配置。

    **加密键一律返回空串**：想拿 Cookie 明文只能用 get_secret()，
    避免任何「顺手读一下配置」的代码把凭证带进日志或响应体。
    """
    if key in config.ENCRYPTED_KEYS:
        return ""
    uid = 0 if config.is_global_key(key) else (uid or 0)
    row = conn.execute("SELECT value FROM settings WHERE user_id=? AND key=?",
                       (uid, key)).fetchone()
    if row is not None and row["value"] is not None:
        return row["value"]
    if uid and key not in NO_FALLBACK_KEYS:
        row = conn.execute("SELECT value FROM settings WHERE user_id=0 AND key=?",
                           (key,)).fetchone()
        if row is not None and row["value"] is not None:
            return row["value"]
    return config.DEFAULTS.get(key, default)


def get_settings(conn, keys=None, uid=0):
    """取该账号的**有效配置**（DEFAULTS -> 实例级 -> 个人级 三级合并）。

    Cookie 等加密键固定为空串，页面/接口可以直接整体回传。
    """
    out = dict(config.DEFAULTS)
    for r in conn.execute("SELECT key,value FROM settings WHERE user_id=0"):
        out[r["key"]] = r["value"]
    if uid:
        for r in conn.execute("SELECT key,value FROM settings WHERE user_id=?", (uid,)):
            if config.is_global_key(r["key"]):
                continue            # 个人作用域里不该有全局键，有也不认
            out[r["key"]] = r["value"]
    for k in config.ENCRYPTED_KEYS:
        out[k] = ""
    if keys:
        return {k: out.get(k) for k in keys}
    return out


def set_setting(conn, key, value, uid=0):
    """写单个配置。全局键强制落到 user_id=0；加密键自动加密后落库。"""
    uid = 0 if config.is_global_key(key) else (uid or 0)
    text = "" if value is None else str(value)
    if key in config.ENCRYPTED_KEYS and text:
        text = crypto.encrypt(text, config.encryption_key())
    conn.execute("INSERT INTO settings(user_id,key,value,updated_at) VALUES(?,?,?,?) "
                 "ON CONFLICT(user_id,key) DO UPDATE SET value=excluded.value,"
                 " updated_at=excluded.updated_at", (uid, key, text, now_str()))


def set_settings(conn, pairs, uid=0):
    for k, v in pairs.items():
        set_setting(conn, k, v, uid)


def get_int(conn, key, default=0, uid=0):
    try:
        return int(float(get_setting(conn, key, default, uid)))
    except (TypeError, ValueError):
        return default


def get_float(conn, key, default=0.0, uid=0):
    try:
        return float(get_setting(conn, key, default, uid))
    except (TypeError, ValueError):
        return default


def get_bool(conn, key, default=False, uid=0):
    v = str(get_setting(conn, key, "1" if default else "0", uid)).strip().lower()
    return v in ("1", "true", "yes", "on", "是")


# ---------------- 凭证（加密存储） ----------------
def get_secret(conn, key, uid=0):
    """取凭证明文。仅在真正要用它对外发请求时调用。"""
    row = conn.execute("SELECT value FROM settings WHERE user_id=? AND key=?",
                       (uid or 0, key)).fetchone()
    if row is None or not row["value"]:
        return ""
    try:
        return crypto.decrypt(row["value"], config.encryption_key())
    except crypto.DecryptError as e:
        raise SecretUnreadable("%s 无法解密：%s" % (key, e))


def set_secret(conn, key, value, uid=0):
    set_setting(conn, key, value, uid)


def secret_state(conn, key, uid=0):
    """给界面用的凭证状态：只回「有没有 / 多少字符 / 尾部 4 位」，绝不含明文。"""
    row = conn.execute("SELECT value FROM settings WHERE user_id=? AND key=?",
                       (uid or 0, key)).fetchone()
    if row is None or not row["value"]:
        return {"set": False, "chars": 0, "tail": "", "broken": False, "at": ""}
    try:
        plain = crypto.decrypt(row["value"], config.encryption_key())
    except crypto.DecryptError:
        return {"set": True, "chars": 0, "tail": "", "broken": True, "at": ""}
    at = conn.execute("SELECT updated_at FROM settings WHERE user_id=? AND key=?",
                      (uid or 0, key)).fetchone()
    return {"set": bool(plain), "chars": len(plain),
            "tail": plain[-4:] if len(plain) >= 4 else "",
            "broken": False, "at": (at["updated_at"] if at else "") or ""}


# ---------------- 账号 ----------------
def user_by_id(conn, uid):
    return conn.execute("SELECT * FROM users WHERE id=?", (uid,)).fetchone()


def user_by_name(conn, username):
    return conn.execute("SELECT * FROM users WHERE username=?", (username,)).fetchone()


def active_users(conn):
    """启用状态的账号（调度器按人遍历）。"""
    return conn.execute("SELECT * FROM users WHERE status='active' ORDER BY id").fetchall()


def user_count(conn):
    return conn.execute("SELECT COUNT(*) FROM users").fetchone()[0]


# ---------------- 审计 ----------------
def audit(conn, action, actor=None, detail=None, ip=None, uid=0):
    conn.execute("INSERT INTO audit_log(user_id,at,actor,action,detail,ip)"
                 " VALUES(?,?,?,?,?,?)", (uid or 0, now_str(), actor, action, detail, ip))


# ---------------- Flask 集成 ----------------
def get_db():
    from flask import g
    if "db" not in g:
        g.db = connect()
    return g.db


def close_db(exc=None):
    from flask import g
    db = g.pop("db", None)
    if db is not None:
        try:
            db.close()
        except sqlite3.Error:
            pass