vent_agent_2.0/db/db_conn.py

"""数据库连接与方言适配。

两条路径（``DB_DIALECT``）：
- ``mysql``：pymysql + DictCursor + 线程本地连接（默认，生产/联调）
- ``sqlite``：单连接 + 互斥锁（本地开发与自动化测试用，无需任何数据库服务）

调用方只写 ``?`` 占位符，MySQL 适配器负责翻成 ``%s``；因此存储层 SQL 是
方言无关的公共子集，两条路径跑的是同一份逻辑。

读出的行统一做归一化：``datetime``/``Decimal`` 一律转成字符串/浮点，
调用方不必区分两种驱动返回的类型差异。
"""

from __future__ import annotations

import sqlite3
import threading
import time
from contextlib import contextmanager
from datetime import date, datetime
from decimal import Decimal
from pathlib import Path
from typing import Any, Iterator, Sequence

import config
from obs import get_logger

logger = get_logger(__name__)

DIALECT: str = config.DB_DIALECT

_schema_ready = False
_schema_lock = threading.Lock()


# ============================================================
# 时间与取值帮助函数
# ============================================================

def now() -> datetime:
    """统一的「当前时间」：本地时区的朴素 datetime。

    全仓只在这里取当前时间，将来要统一改 UTC 只需改这一处。
    """
    return datetime.now()


def fmt_dt(value: datetime | None = None) -> str:
    """把 datetime 格式化为 ``YYYY-MM-DD HH:MM:SS.ffffff``，供落库使用。

    必须带微秒：历史消息的排序完全依赖 ``created_at``，而一轮对话里
    user / assistant 两条消息往往落在同一秒。只到秒级时同秒记录只能靠
    随机 UUID 做次级排序，顺序会随机颠倒——冷启动补种图状态时这会直接把
    「用户说了什么、助手回了什么」的顺序喂反。
    """
    return (value or now()).strftime("%Y-%m-%d %H:%M:%S.%f")


def day_key(value: datetime | None = None) -> str:
    """用量归属日的键：``YYYYMMDD``，与 billing_usage.usage_day 对应。"""
    return (value or now()).strftime("%Y%m%d")


def new_id() -> str:
    """生成 32 位十六进制主键（对应 VARCHAR(64) 主键列）。"""
    import uuid

    return uuid.uuid4().hex


# ============================================================
# SQLite 适配
# ============================================================

class _SqliteAdapter:
    """SQLite 后端：单连接 + 锁。

    单连接是为了让 :memory: 与文件库行为一致，也避免多连接下的写锁竞争。
    代价是并发写串行化——对本地开发与测试足够；生产走 MySQL 分支。
    """

    def __init__(self, path: str) -> None:
        self._path = path
        self._lock = threading.RLock()
        self._conn: sqlite3.Connection | None = None
        self._transaction_depth = 0

    def _connection(self) -> sqlite3.Connection:
        if self._conn is None:
            if self._path != ":memory:":
                Path(self._path).parent.mkdir(parents=True, exist_ok=True)
            conn = sqlite3.connect(self._path, check_same_thread=False, timeout=30)
            conn.row_factory = sqlite3.Row
            conn.execute("PRAGMA journal_mode=WAL")
            conn.execute("PRAGMA foreign_keys=ON")
            self._conn = conn
        return self._conn

    def query(self, sql: str, params: Sequence[Any] = ()) -> list[dict[str, Any]]:
        with self._lock:
            cur = self._connection().execute(sql, tuple(params))
            return [dict(row) for row in cur.fetchall()]

    def execute(self, sql: str, params: Sequence[Any] = ()) -> int:
        with self._lock:
            conn = self._connection()
            try:
                cur = conn.execute(sql, tuple(params))
            except BaseException:
                if not self._transaction_depth:
                    conn.rollback()
                raise
            if not self._transaction_depth:
                conn.commit()
            return cur.rowcount

    def execute_many(self, sql: str, params_seq: Sequence[Sequence[Any]]) -> int:
        with self._lock:
            conn = self._connection()
            try:
                cur = conn.executemany(sql, [tuple(p) for p in params_seq])
            except BaseException:
                if not self._transaction_depth:
                    conn.rollback()
                raise
            if not self._transaction_depth:
                conn.commit()
            return cur.rowcount

    @contextmanager
    def transaction(self) -> Iterator[None]:
        with self._lock:
            conn = self._connection()
            depth = self._transaction_depth
            savepoint = f"nested_{depth}"
            conn.execute("BEGIN" if depth == 0 else f"SAVEPOINT {savepoint}")
            self._transaction_depth += 1
            try:
                yield
                if depth == 0:
                    conn.commit()
                else:
                    conn.execute(f"RELEASE SAVEPOINT {savepoint}")
            except BaseException:
                if depth == 0:
                    conn.rollback()
                else:
                    conn.execute(f"ROLLBACK TO SAVEPOINT {savepoint}")
                    conn.execute(f"RELEASE SAVEPOINT {savepoint}")
                raise
            finally:
                self._transaction_depth -= 1

    def schema_file(self) -> Path:
        return config.ROOT / "db" / "schema_sqlite.sql"

    def ping(self) -> None:
        self._connection().execute("SELECT 1").fetchone()

    def close(self) -> None:
        with self._lock:
            if self._conn is not None:
                self._conn.close()
                self._conn = None
            self._transaction_depth = 0


# ============================================================
# MySQL 适配
# ============================================================

class _MySqlAdapter:
    """MySQL 后端：线程本地连接 + 空闲自动 ping。

    线程本地是因为 pymysql 的 Connection 不是线程安全的；空闲 ping 是为了
    绕开服务端 ``wait_timeout`` 主动断开导致的 "MySQL server has gone away"。
    """

    def __init__(self) -> None:
        self._local = threading.local()
        self._last_used: dict[int, float] = {}
        self._connections: dict[int, Any] = {}
        self._connections_lock = threading.Lock()

    def _connect(self, database: str | None):
        """建一个 pymysql 连接。``database=None`` 表示不选库。"""
        import pymysql
        from pymysql.cursors import DictCursor

        return pymysql.connect(
            host=config.DB_HOST,
            port=config.DB_PORT,
            user=config.DB_USER,
            password=config.DB_PASSWORD,
            database=database,
            charset="utf8mb4",
            cursorclass=DictCursor,
            autocommit=True,
            connect_timeout=10,
            read_timeout=60,
            write_timeout=60,
        )

    def _ensure_database(self) -> None:
        """库不存在则建库。

        骨架不该要求使用者先手工跑一次 CREATE DATABASE——那是最容易被漏掉、
        且报错信息（"Unknown database"）不提示该怎么解决的一步。
        """
        conn = self._connect(None)
        try:
            with conn.cursor() as cur:
                cur.execute(
                    f"CREATE DATABASE IF NOT EXISTS `{config.DB_NAME}` "
                    f"DEFAULT CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci"
                )
            logger.info("已确保数据库存在：%s", config.DB_NAME)
        finally:
            conn.close()

    def _connection(self):
        conn = getattr(self._local, "conn", None)
        if conn is None:
            try:
                conn = self._connect(config.DB_NAME)
            except Exception as exc:
                # 1049 = Unknown database：自动建库后重试一次
                if getattr(exc, "args", [None])[0] == 1049:
                    self._ensure_database()
                    conn = self._connect(config.DB_NAME)
                else:
                    raise
            self._local.conn = conn
            with self._connections_lock:
                self._connections[id(conn)] = conn
        else:
            ident = id(conn)
            if not getattr(self._local, "transaction_depth", 0) and time.time() - self._last_used.get(ident, 0) > config.DB_PING_IDLE_SECONDS:
                try:
                    conn.ping(reconnect=True)
                except Exception:
                    with self._connections_lock:
                        self._connections.pop(id(conn), None)
                    self._last_used.pop(id(conn), None)
                    try:
                        conn.close()
                    except Exception:
                        pass
                    self._local.conn = None
                    return self._connection()
        self._last_used[id(conn)] = time.time()
        return conn

    @staticmethod
    def _translate(sql: str) -> str:
        """把调用方写的 ``?`` 占位符翻成 pymysql 的 ``%s``。

        同时把 SQL 里字面的 ``%`` 转义成 ``%%``（pymysql 用 % 做参数格式化，
        不转义会把它误当占位符）。参数值本身不参与翻译。
        """
        parts = sql.split("?")
        escaped = [p.replace("%", "%%") for p in parts]
        return "%s".join(escaped)

    def query(self, sql: str, params: Sequence[Any] = ()) -> list[dict[str, Any]]:
        conn = self._connection()
        with conn.cursor() as cur:
            cur.execute(self._translate(sql), tuple(params))
            return list(cur.fetchall())

    def execute(self, sql: str, params: Sequence[Any] = ()) -> int:
        conn = self._connection()
        with conn.cursor() as cur:
            cur.execute(self._translate(sql), tuple(params))
            return cur.rowcount

    def execute_many(self, sql: str, params_seq: Sequence[Sequence[Any]]) -> int:
        conn = self._connection()
        with conn.cursor() as cur:
            cur.executemany(self._translate(sql), [tuple(p) for p in params_seq])
            return cur.rowcount

    @contextmanager
    def transaction(self) -> Iterator[None]:
        conn = self._connection()
        depth = getattr(self._local, "transaction_depth", 0)
        savepoint = f"nested_{depth}"
        if depth == 0:
            conn.begin()
        else:
            with conn.cursor() as cur:
                cur.execute(f"SAVEPOINT {savepoint}")
        self._local.transaction_depth = depth + 1
        try:
            yield
            if depth == 0:
                conn.commit()
            else:
                with conn.cursor() as cur:
                    cur.execute(f"RELEASE SAVEPOINT {savepoint}")
        except BaseException:
            if depth == 0:
                conn.rollback()
            else:
                with conn.cursor() as cur:
                    cur.execute(f"ROLLBACK TO SAVEPOINT {savepoint}")
                    cur.execute(f"RELEASE SAVEPOINT {savepoint}")
            raise
        finally:
            self._local.transaction_depth = depth

    def schema_file(self) -> Path:
        return config.ROOT / "db" / "mysql_schema.sql"

    def ping(self) -> None:
        self._connection().ping(reconnect=True)

    def close(self) -> None:
        """在所有数据库任务结束后关闭各线程连接，下一轮生命周期重新建立。"""
        with self._connections_lock:
            connections = list(self._connections.values())
            self._connections.clear()
            self._local = threading.local()
            self._last_used.clear()
        for conn in connections:
            try:
                conn.close()
            except Exception as exc:
                logger.warning("关闭 MySQL 连接失败：%s", type(exc).__name__)


def _build_adapter() -> _SqliteAdapter | _MySqlAdapter:
    if DIALECT == "sqlite":
        path = config.DB_SQLITE_PATH
        if path != ":memory:":
            path = str(config.resolve_path(path))
        return _SqliteAdapter(path)
    if DIALECT == "mysql":
        return _MySqlAdapter()
    raise ValueError(
        f"不支持的 DB_DIALECT={DIALECT!r}，可选值：mysql | sqlite"
    )


_adapter = _build_adapter()


def backend_name() -> str:
    """当前后端名，供 /health 与日志展示。"""
    return DIALECT


# ============================================================
# 对外 API
# ============================================================

def _normalize(row: dict[str, Any]) -> dict[str, Any]:
    """把驱动特有类型收敛成可 JSON 化的通用类型。"""
    out: dict[str, Any] = {}
    for key, value in row.items():
        if isinstance(value, datetime):
            # 与 fmt_dt 对齐到微秒，保证「写入再读出」是同一个字符串
            out[key] = value.strftime("%Y-%m-%d %H:%M:%S.%f")
        elif isinstance(value, date):
            out[key] = value.isoformat()
        elif isinstance(value, Decimal):
            out[key] = float(value)
        elif isinstance(value, bytes):
            out[key] = value.decode("utf-8", errors="replace")
        else:
            out[key] = value
    return out


def query(sql: str, params: Sequence[Any] = ()) -> list[dict[str, Any]]:
    """执行查询，返回归一化后的行列表。"""
    return [_normalize(row) for row in _adapter.query(sql, params)]


def query_one(sql: str, params: Sequence[Any] = ()) -> dict[str, Any] | None:
    """执行查询并返回首行；无结果返回 None。"""
    rows = query(sql, params)
    return rows[0] if rows else None


def scalar(sql: str, params: Sequence[Any] = (), default: Any = None) -> Any:
    """取首行首列；无结果返回 default。"""
    row = query_one(sql, params)
    if not row:
        return default
    return next(iter(row.values()), default)


def execute(sql: str, params: Sequence[Any] = ()) -> int:
    """执行写入语句，返回受影响行数。"""
    return _adapter.execute(sql, params)


def execute_many(sql: str, params_seq: Sequence[Sequence[Any]]) -> int:
    """批量写入，返回受影响行数。"""
    return _adapter.execute_many(sql, params_seq)


@contextmanager
def transaction() -> Iterator[None]:
    """事务上下文。块内异常自动回滚并向上抛出。"""
    with _adapter.transaction():
        yield


# ============================================================
# 建表与健康检查
# ============================================================

def _split_statements(sql_text: str) -> list[str]:
    """把 .sql 文件切成可逐条执行的语句。

    先剥掉 ``--`` 行注释与空行，再按 ``;`` 切分。刻意不做通用 SQL 解析——
    本项目的 schema 文件里没有字符串内的分号，够用即可。
    """
    lines: list[str] = []
    for raw in sql_text.splitlines():
        line = raw.split("--", 1)[0].rstrip()
        if line.strip():
            lines.append(line)
    return [stmt.strip() for stmt in "\n".join(lines).split(";") if stmt.strip()]


def init_schema(force: bool = False) -> None:
    """执行建表脚本（幂等，全部 IF NOT EXISTS）。"""

    global _schema_ready
    with _schema_lock:
        if _schema_ready and not force:
            return
        path = _adapter.schema_file()
        statements = _split_statements(path.read_text(encoding="utf-8"))
        for stmt in statements:
            _adapter.execute(stmt, ())
        _schema_ready = True
        logger.info(
            "建表完成 dialect=%s 语句数=%d 文件=%s", DIALECT, len(statements), path.name
        )


def health() -> dict[str, Any]:
    """连通性探测，供 /health 使用。异常不抛出，转为 status=error。"""
    try:
        _adapter.ping()
        detail = (
            f"{config.DB_USER}@{config.DB_HOST}:{config.DB_PORT}/{config.DB_NAME}"
            if DIALECT == "mysql"
            else config.DB_SQLITE_PATH
        )
        return {"status": "ok", "dialect": DIALECT, "target": detail}
    except Exception as exc:  # pragma: no cover - 取决于外部环境
        return {"status": "error", "dialect": DIALECT, "error": f"{type(exc).__name__}: {exc}"}


def close() -> None:
    """关闭业务库连接；只能在执行任务和后台维护全部结束后调用。"""
    global _schema_ready
    with _schema_lock:
        try:
            _adapter.close()
        finally:
            _schema_ready = False


def reset_for_tests() -> None:
    """仅供测试：丢弃缓存连接与建表标记，让下一次调用重新初始化。"""
    global _adapter, _schema_ready
    close()
    _adapter = _build_adapter()
    _schema_ready = False