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

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

import asyncio
import json
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

_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);
"""


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)
        await self._conn.commit()

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

    @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()