app/database.py

"""SQLite 初始化与连接管理。

简化点：相比原 Netlify Blobs KV，这里用本地 SQLite（容器内 /data）。
HF Space 重建时数据可能丢失，启动时通过 HF Hub 仓库恢复配置（见 services/config_store.py）。
"""
from __future__ import annotations

import os
import sqlite3
import threading
from contextlib import contextmanager
from pathlib import Path
from typing import Iterator

from .config import get_settings

_lock = threading.Lock()
_conn: sqlite3.Connection | None = None


def _ensure_parent_dir(db_path: str) -> None:
    p = Path(db_path).expanduser()
    if p.parent and not p.parent.exists():
        try:
            p.parent.mkdir(parents=True, exist_ok=True)
        except Exception:
            # 容器内只读卷会失败，退回 /tmp
            pass


def _resolve_db_path() -> str:
    settings = get_settings()
    db_path = settings.db_path
    # HF 免费档无持久卷时，/data 不存在或不可写，回退到 /tmp
    try:
        _ensure_parent_dir(db_path)
        test_path = Path(db_path)
        test_path.touch(exist_ok=True)
        if not os.access(db_path, os.W_OK):
            raise PermissionError(db_path)
        return db_path
    except Exception:
        fallback = "/tmp/xtc.db"
        print(f"[database] db_path {db_path} not writable, fallback to {fallback}")
        return fallback


def init_db() -> sqlite3.Connection:
    global _conn
    with _lock:
        if _conn is not None:
            return _conn
        db_path = _resolve_db_path()
        conn = sqlite3.connect(db_path, check_same_thread=False, isolation_level=None)
        conn.row_factory = sqlite3.Row
        conn.execute("PRAGMA journal_mode=WAL")
        conn.execute("PRAGMA synchronous=NORMAL")
        conn.execute("PRAGMA foreign_keys=ON")
        _create_schema(conn)
        _conn = conn
        return _conn


def _create_schema(conn: sqlite3.Connection) -> None:
    conn.executescript(
        """
        CREATE TABLE IF NOT EXISTS kv (
            key   TEXT PRIMARY KEY,
            value TEXT NOT NULL,
            updated_at INTEGER NOT NULL
        );

        CREATE TABLE IF NOT EXISTS access_tokens (
            token       TEXT PRIMARY KEY,
            issued_at   INTEGER NOT NULL,
            expires_at  INTEGER NOT NULL,
            issued_by   TEXT,
            revoked     INTEGER NOT NULL DEFAULT 0
        );
        CREATE INDEX IF NOT EXISTS idx_access_tokens_expires ON access_tokens(expires_at);

        CREATE TABLE IF NOT EXISTS pseudo_sessions (
            id          TEXT PRIMARY KEY,
            created_at  INTEGER NOT NULL,
            updated_at  INTEGER NOT NULL,
            done        INTEGER NOT NULL DEFAULT 0,
            error       TEXT,
            thought     TEXT NOT NULL DEFAULT '',
            text        TEXT NOT NULL DEFAULT '',
            seq         INTEGER NOT NULL DEFAULT 0,
            payload     TEXT
        );
        CREATE INDEX IF NOT EXISTS idx_pseudo_sessions_updated ON pseudo_sessions(updated_at);

        CREATE TABLE IF NOT EXISTS pseudo_events (
            seq         INTEGER NOT NULL,
            session_id  TEXT NOT NULL,
            at          TEXT NOT NULL,
            type        TEXT NOT NULL,
            delta       TEXT,
            finish_reason TEXT,
            extra       TEXT,
            PRIMARY KEY (session_id, seq)
        );
        CREATE INDEX IF NOT EXISTS idx_pseudo_events_session ON pseudo_events(session_id);

        CREATE TABLE IF NOT EXISTS rr_counter (
            provider_id TEXT PRIMARY KEY,
            counter     INTEGER NOT NULL DEFAULT 0
        );

        CREATE TABLE IF NOT EXISTS app_logs (
            id          INTEGER PRIMARY KEY AUTOINCREMENT,
            ts          INTEGER NOT NULL,
            level       TEXT NOT NULL,
            module      TEXT NOT NULL,
            event       TEXT NOT NULL,
            trace_id    TEXT,
            data        TEXT
        );
        CREATE INDEX IF NOT EXISTS idx_app_logs_ts ON app_logs(ts);
        CREATE INDEX IF NOT EXISTS idx_app_logs_level ON app_logs(level);

        CREATE TABLE IF NOT EXISTS usage_log (
            id          INTEGER PRIMARY KEY AUTOINCREMENT,
            ts          INTEGER NOT NULL,
            access_key  TEXT,
            provider    TEXT,
            model       TEXT,
            prompt_tokens INTEGER,
            completion_tokens INTEGER,
            total_tokens INTEGER,
            ok          INTEGER NOT NULL DEFAULT 0,
            error_code  TEXT
        );
        CREATE INDEX IF NOT EXISTS idx_usage_log_ts ON usage_log(ts);
        CREATE INDEX IF NOT EXISTS idx_usage_log_key ON usage_log(access_key);

        CREATE TABLE IF NOT EXISTS audit_log (
            id          INTEGER PRIMARY KEY AUTOINCREMENT,
            ts          INTEGER NOT NULL,
            action      TEXT NOT NULL,
            actor       TEXT,
            target      TEXT,
            detail      TEXT
        );
        CREATE INDEX IF NOT EXISTS idx_audit_log_ts ON audit_log(ts);

        CREATE TABLE IF NOT EXISTS file_meta (
            key         TEXT PRIMARY KEY,
            namespace   TEXT NOT NULL,
            filename    TEXT NOT NULL,
            mime        TEXT,
            size        INTEGER NOT NULL,
            sha256      TEXT,
            uploaded_at INTEGER NOT NULL,
            uploaded_by TEXT,
            access_key  TEXT,
            refs        TEXT
        );
        CREATE INDEX IF NOT EXISTS idx_file_meta_ns ON file_meta(namespace);
        CREATE INDEX IF NOT EXISTS idx_file_meta_owner ON file_meta(access_key);
        CREATE INDEX IF NOT EXISTS idx_file_meta_uploaded ON file_meta(uploaded_at);

        CREATE TABLE IF NOT EXISTS sessions (
            id          TEXT PRIMARY KEY,
            title       TEXT,
            provider    TEXT,
            model       TEXT,
            access_key  TEXT,
            created_at  INTEGER NOT NULL,
            updated_at  INTEGER NOT NULL
        );
        CREATE INDEX IF NOT EXISTS idx_sessions_owner ON sessions(access_key);
        CREATE INDEX IF NOT EXISTS idx_sessions_updated ON sessions(updated_at);

        CREATE TABLE IF NOT EXISTS session_messages (
            id          INTEGER PRIMARY KEY AUTOINCREMENT,
            session_id  TEXT NOT NULL,
            seq         INTEGER NOT NULL,
            role        TEXT NOT NULL,
            content     TEXT NOT NULL,
            thought     TEXT,
            provider    TEXT,
            model       TEXT,
            ts          INTEGER NOT NULL,
            file_keys   TEXT,
            UNIQUE(session_id, seq)
        );
        CREATE INDEX IF NOT EXISTS idx_session_messages_session ON session_messages(session_id);

        CREATE TABLE IF NOT EXISTS rate_limit (
            bucket      TEXT PRIMARY KEY,
            count       INTEGER NOT NULL,
            window_start INTEGER NOT NULL
        );

        CREATE TABLE IF NOT EXISTS webhook_config (
            id          INTEGER PRIMARY KEY AUTOINCREMENT,
            name        TEXT NOT NULL,
            url         TEXT NOT NULL,
            events      TEXT NOT NULL,
            enabled     INTEGER NOT NULL DEFAULT 1,
            created_at  INTEGER NOT NULL
        );

        CREATE TABLE IF NOT EXISTS user_accounts (
            id            INTEGER PRIMARY KEY AUTOINCREMENT,
            username      TEXT NOT NULL UNIQUE,
            password_hash TEXT NOT NULL,
            created_at    INTEGER NOT NULL
        );
        CREATE INDEX IF NOT EXISTS idx_user_accounts_username ON user_accounts(username);

        CREATE TABLE IF NOT EXISTS request_log (
            id              INTEGER PRIMARY KEY AUTOINCREMENT,
            ts              INTEGER NOT NULL,
            method          TEXT NOT NULL,
            path            TEXT NOT NULL,
            query           TEXT,
            status_code     INTEGER,
            elapsed_ms      INTEGER,
            access_key      TEXT,
            client_ip       TEXT,
            user_agent      TEXT,
            provider        TEXT,
            model           TEXT,
            stream          INTEGER NOT NULL DEFAULT 0,
            ok              INTEGER NOT NULL DEFAULT 0,
            error_code      TEXT,
            error_message   TEXT,
            request_headers TEXT,
            request_body    TEXT,
            response_body   TEXT,
            response_headers TEXT,
            device_id       TEXT
        );
        CREATE INDEX IF NOT EXISTS idx_request_log_ts ON request_log(ts);
        CREATE INDEX IF NOT EXISTS idx_request_log_path ON request_log(path);
        CREATE INDEX IF NOT EXISTS idx_request_log_status ON request_log(status_code);
        CREATE INDEX IF NOT EXISTS idx_request_log_key ON request_log(access_key);
        CREATE INDEX IF NOT EXISTS idx_request_log_device ON request_log(device_id);
        """
    )
    # 兼容已存在的库：若旧表无 device_id 列则补加（ALTER TABLE 幂等检查）
    _ensure_column(conn, "request_log", "device_id", "TEXT")
    # file_meta 增加 alias 列：原始文件名（可能是中文），filename 改存随机英文名
    _ensure_column(conn, "file_meta", "alias", "TEXT")
    # request_log 容量上限从 kv 读取，默认 500（直接用当前 conn，避免递归 init_db）
    _ensure_request_log_limit_default(conn)


def _ensure_request_log_limit_default(conn: sqlite3.Connection) -> None:
    """首次启动写入默认 request_log 上限。"""
    import time as _time
    try:
        row = conn.execute(
            "SELECT value FROM kv WHERE key = 'request_log_limit'"
        ).fetchone()
        if not row:
            conn.execute(
                "INSERT INTO kv(key, value, updated_at) VALUES(?,?,?)",
                ("request_log_limit", "500", int(_time.time())),
            )
    except Exception as e:
        print(f"[database] request_log_limit init failed: {e}")


def _ensure_column(conn: sqlite3.Connection, table: str, column: str, ddl_type: str) -> None:
    """幂等加列：若旧表缺少某列则 ALTER TABLE 补加。仅用于轻量迁移。"""
    try:
        cols = {row["name"] for row in conn.execute(f"PRAGMA table_info({table})")}
    except Exception:
        return
    if column not in cols:
        try:
            conn.execute(f"ALTER TABLE {table} ADD COLUMN {column} {ddl_type}")
        except Exception as e:
            print(f"[database] alter {table}.{column} failed: {e}")


@contextmanager
def get_conn() -> Iterator[sqlite3.Connection]:
    conn = init_db()
    yield conn


def close_db() -> None:
    global _conn
    with _lock:
        if _conn is not None:
            _conn.close()
            _conn = None