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

"""db.py —— paper 服务独立 SQLite（aiosqlite + WAL，无 ORM），三张表。

- papers：论文主表（brief/outline/final_markdown/references/disclaimer）
- paper_sections：章节表（逐节 markdown + 状态）
- citations：结构化引用（temp_n 生成期全局递增，final_n complete 按出现序重排写回）

模式沿用 bioflow/docgen 同款：CREATE TABLE IF NOT EXISTS + aiosqlite 单连接 + WAL + 写锁。
"""
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 papers (
    id              TEXT PRIMARY KEY,
    session_id      TEXT NOT NULL DEFAULT '',
    title           TEXT NOT NULL DEFAULT '',
    paper_type      TEXT NOT NULL DEFAULT '',
    language        TEXT NOT NULL DEFAULT '',
    format_req      TEXT NOT NULL DEFAULT '',
    style           TEXT NOT NULL DEFAULT '',
    status          TEXT NOT NULL DEFAULT 'outline',
    brief_json      TEXT NOT NULL DEFAULT '{}',
    outline_json    TEXT NOT NULL DEFAULT '{}',
    final_markdown  TEXT NOT NULL DEFAULT '',
    references_json TEXT NOT NULL DEFAULT '[]',
    disclaimer      TEXT NOT NULL DEFAULT '',
    created_at      TEXT NOT NULL,
    updated_at      TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_papers_session ON papers(session_id);
CREATE INDEX IF NOT EXISTS idx_papers_status ON papers(status);

CREATE TABLE IF NOT EXISTS paper_sections (
    id          TEXT PRIMARY KEY,
    paper_id    TEXT NOT NULL REFERENCES papers(id) ON DELETE CASCADE,
    section_id  TEXT NOT NULL,
    title       TEXT NOT NULL DEFAULT '',
    key_points  TEXT NOT NULL DEFAULT '[]',
    status      TEXT NOT NULL DEFAULT 'pending',
    markdown    TEXT NOT NULL DEFAULT '',
    word_count  INTEGER NOT NULL DEFAULT 0,
    updated_at  TEXT NOT NULL,
    UNIQUE(paper_id, section_id)
);
CREATE INDEX IF NOT EXISTS idx_paper_sections_paper ON paper_sections(paper_id);

CREATE TABLE IF NOT EXISTS citations (
    id             INTEGER PRIMARY KEY AUTOINCREMENT,
    paper_id       TEXT NOT NULL REFERENCES papers(id) ON DELETE CASCADE,
    section_id     TEXT NOT NULL DEFAULT '',
    temp_n         INTEGER,
    final_n        INTEGER,
    kind           TEXT NOT NULL DEFAULT 'title',
    ref_text       TEXT NOT NULL DEFAULT '',
    source_doc_id  TEXT NOT NULL DEFAULT '',
    snippet        TEXT NOT NULL DEFAULT '',
    created_at     TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_citations_paper ON citations(paper_id);
"""


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


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


class Database:
    """aiosqlite 单连接封装（WAL；写串行化，同 bioflow 模式）。"""

    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 = Path(get_settings().db_path_abs)
        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()

    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:
    # 原地改单例路径后连接（不做 db = Database(path) 重赋值）：
    # 服务模块以 `from ...db import db` 持有了该单例引用，重赋值会让它们拿到旧对象。
    if path is not None:
        db._path = Path(path)
    await db.connect()
    return db


async def get_db() -> AsyncIterator[Database]:
    yield db


async def close_db() -> None:
    await db.close()


def monotonic() -> float:
    return time.monotonic()