PlantAI-rag-new/sdk/app/db.py

"""db.py —— SQLite 持久层（aiosqlite + WAL，无 ORM）。

Schema 与 API.md §5 对齐：
- sessions: id / title / status / meta / created_at / updated_at
- messages: id / session_id(FK CASCADE) / role / kind / status / blocks(JSON)
            / raw(JSON) / text / created_at / updated_at
- doc_library: id / name / category / tags / description / as_template /
               enc_path / embed_sha / active / created_at（docgen P2，无 FK，
               不依赖 vstore 可用性；enc_path/embed_sha 为 vstore ingest 回填键）
- doc_artifacts: id / session_id / status / title / template_name / md_path /
                 docx_path / pdf_path / error / created_at（docgen P2；
                 导出文件按需生成，路径列可空）

启动时幂等迁移（CREATE TABLE IF NOT EXISTS）。时间一律 ISO 8601 UTC
（`2026-08-25T12:34:56Z`），由 BFF 侧统一生成（SQLite 不写 DEFAULT 时间戳）。

docgen 两表为增量追加；迁移脚本执行在 try/except 内（见 connect()），
docgen 建表失败只记日志、不阻断 BFF 启动（docgen 端点届时返回 INTERNAL 信封）。
"""
from __future__ import annotations

import asyncio
import json
import logging
import time
from collections.abc import AsyncIterator, Iterable
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import aiosqlite

from .config import get_settings

logger = logging.getLogger(__name__)

_SCHEMA = """
CREATE TABLE IF NOT EXISTS sessions (
    id          TEXT PRIMARY KEY,
    title       TEXT NOT NULL DEFAULT '新会话',
    status      TEXT NOT NULL DEFAULT 'idle',
    meta        TEXT NOT NULL DEFAULT '{}',
    created_at  TEXT NOT NULL,
    updated_at  TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS messages (
    id          TEXT PRIMARY KEY,
    session_id  TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
    role        TEXT NOT NULL,
    kind        TEXT NOT NULL DEFAULT 'qa',
    status      TEXT NOT NULL DEFAULT 'done',
    blocks      TEXT NOT NULL DEFAULT '[]',
    raw         TEXT,
    text        TEXT,
    created_at  TEXT NOT NULL,
    updated_at  TEXT NOT NULL
);

CREATE INDEX IF NOT EXISTS idx_messages_session ON messages(session_id);
CREATE INDEX IF NOT EXISTS idx_sessions_updated ON sessions(updated_at);
"""

# docgen P2 增量表（单独脚本：迁移失败只记日志，不阻断启动）
_DOCGEN_SCHEMA = """
CREATE TABLE IF NOT EXISTS doc_library (
    id           TEXT PRIMARY KEY,
    name         TEXT NOT NULL DEFAULT '',
    category     TEXT NOT NULL DEFAULT '',
    tags         TEXT NOT NULL DEFAULT '',
    description  TEXT NOT NULL DEFAULT '',
    as_template  INTEGER NOT NULL DEFAULT 0,
    enc_path     TEXT,
    embed_sha    TEXT,
    active       INTEGER NOT NULL DEFAULT 1,
    created_at   TEXT NOT NULL
);

CREATE INDEX IF NOT EXISTS idx_doc_library_active ON doc_library(active);
CREATE INDEX IF NOT EXISTS idx_doc_library_category ON doc_library(category);

CREATE TABLE IF NOT EXISTS doc_artifacts (
    id            TEXT PRIMARY KEY,
    session_id    TEXT NOT NULL DEFAULT '',
    status        TEXT NOT NULL DEFAULT 'running',
    title         TEXT NOT NULL DEFAULT '',
    template_name TEXT NOT NULL DEFAULT '',
    md_path       TEXT,
    docx_path     TEXT,
    pdf_path      TEXT,
    error         TEXT,
    created_at    TEXT NOT NULL
);

CREATE INDEX IF NOT EXISTS idx_doc_artifacts_session ON doc_artifacts(session_id);
CREATE INDEX IF NOT EXISTS idx_doc_artifacts_created ON doc_artifacts(created_at);
"""

# docgen 模板预览（2026-09-04）增量列：幂等 ALTER（PRAGMA table_info 判列存在）。
# html_path：生成时落盘的 html 文件相对路径（渲染读取按 md 现渲染，此列作 TTL 清理键）
# template_key：内置模板 key（official/report/...）
# template_css：生成时定稿的 style_hints 覆盖 css（渲染/导出免 vstore 往返）
_DOCGEN_MIGRATIONS = (
    ("html_path", "ALTER TABLE doc_artifacts ADD COLUMN html_path TEXT"),
    ("template_key", "ALTER TABLE doc_artifacts ADD COLUMN template_key TEXT NOT NULL DEFAULT ''"),
    ("template_css", "ALTER TABLE doc_artifacts ADD COLUMN template_css TEXT NOT NULL DEFAULT ''"),
    ("letterhead", "ALTER TABLE doc_artifacts ADD COLUMN letterhead TEXT NOT NULL DEFAULT ''"),
    ("file_name", "ALTER TABLE doc_artifacts ADD COLUMN file_name TEXT NOT NULL DEFAULT ''"),
)


def _now_iso() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def _db_path() -> Path:
    path = Path(get_settings().db_path)
    path.parent.mkdir(parents=True, exist_ok=True)
    return path


class Database:
    """aiosqlite 单连接封装。

    SQLite 写并发不高，BFF 用单写连接（WAL 模式下读写不互斥）；
    依赖注入经 get_db() 提供本实例。
    """

    def __init__(self, path: str | Path | None = None) -> None:
        self._path: Path | None = Path(path) if path is not None else None
        self._conn: aiosqlite.Connection | None = None
        self._lock = asyncio.Lock()

    @property
    def path(self) -> Path:
        if self._path is None:
            self._path = _db_path()
        return self._path

    async def connect(self) -> None:
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self._conn = await aiosqlite.connect(self.path)
        self._conn.row_factory = aiosqlite.Row
        await self._conn.execute("PRAGMA journal_mode=WAL")
        await self._conn.execute("PRAGMA foreign_keys=ON")
        await self._conn.execute("PRAGMA busy_timeout=5000")
        await self._conn.executescript(_SCHEMA)
        # docgen 增量表（P2）：单独执行；失败只记日志，不阻断 BFF 启动
        try:
            await self._conn.executescript(_DOCGEN_SCHEMA)
            await self._migrate_docgen_columns()
        except Exception:
            logger.exception("docgen 表迁移失败（docgen 端点将不可用，其余不受影响）")
        await self._conn.commit()

    async def close(self) -> None:
        if self._conn is not None:
            await self._conn.close()
            self._conn = None

    async def _migrate_docgen_columns(self) -> None:
        """doc_artifacts 增量列幂等迁移（老库已建表时 CREATE IF NOT EXISTS 不会加列）。"""
        cur = await self._conn.execute("PRAGMA table_info(doc_artifacts)")
        cols = {row[1] for row in await cur.fetchall()}
        for name, ddl in _DOCGEN_MIGRATIONS:
            if name not in cols:
                await self._conn.execute(ddl)

    @property
    def conn(self) -> aiosqlite.Connection:
        if self._conn is None:
            raise RuntimeError("Database 未初始化：先 await init_db()")
        return self._conn

    # ---- 通用查询（只读）----

    async def fetchall(self, sql: str, params: Iterable[Any] = ()) -> list[aiosqlite.Row]:
        async with self._conn.execute(sql, tuple(params)) as cursor:
            return list(await cursor.fetchall())

    async def fetchone(self, sql: str, params: Iterable[Any] = ()) -> aiosqlite.Row | None:
        async with self._conn.execute(sql, tuple(params)) as cursor:
            return await cursor.fetchone()

    # ---- 写操作（串行化，提交后返回 lastrowid）----

    async def execute(self, sql: str, params: Iterable[Any] = ()) -> int:
        async with self._lock:
            cur = await self._conn.execute(sql, tuple(params))
            await self._conn.commit()
            return cur.lastrowid


db = Database()


async def init_db(path: str | Path | None = None) -> Database:
    """app 启动钩子：创建并迁移数据库（幂等）。返回全局单例（或指定路径的实例）。"""
    global db
    if path is not None:
        db = Database(path)
    await db.connect()
    return db


async def get_db() -> AsyncIterator[Database]:
    """FastAPI 依赖：每个请求交出同一数据库实例。"""
    yield db


async def close_db() -> None:
    """app shutdown 钩子。"""
    await db.close()


# ---- JSON 编解码（落库用；失败时按契约置 null/[] 而不吞请求）----

def dumps_json(value: Any, default: Any = None) -> str:
    try:
        return json.dumps(value, ensure_ascii=False)
    except (TypeError, ValueError):
        return json.dumps(default, ensure_ascii=False)


def loads_json(text: str | None, default: Any = None) -> Any:
    if text is None:
        return default
    try:
        return json.loads(text)
    except (TypeError, ValueError):
        return default


def monotonic() -> float:
    """测试/内部计时用（time.monotonic 别名，避免各处直接 import time）。"""
    return time.monotonic()