303 lines
11 KiB
Python
303 lines
11 KiB
Python
"""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 异步任务（缺陷 #52，2026-09-10）：generate 202+task_id 轮询契约
|
||
CREATE TABLE IF NOT EXISTS docgen_tasks (
|
||
id TEXT PRIMARY KEY,
|
||
session_id TEXT NOT NULL DEFAULT '',
|
||
status TEXT NOT NULL DEFAULT 'queued',
|
||
artifact_id TEXT,
|
||
title TEXT NOT NULL DEFAULT '',
|
||
error TEXT,
|
||
progress TEXT NOT NULL DEFAULT '[]',
|
||
created_at TEXT NOT NULL,
|
||
updated_at TEXT NOT NULL
|
||
);
|
||
|
||
CREATE INDEX IF NOT EXISTS idx_docgen_tasks_session ON docgen_tasks(session_id);
|
||
CREATE INDEX IF NOT EXISTS idx_docgen_tasks_status ON docgen_tasks(status);
|
||
"""
|
||
|
||
# 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 ''"),
|
||
)
|
||
|
||
|
||
# 管理后台/鉴权增量表(users + auth_tokens;行为日志走文件,不建表)
|
||
_ADMIN_SCHEMA = """
|
||
CREATE TABLE IF NOT EXISTS users (
|
||
id TEXT PRIMARY KEY,
|
||
username TEXT NOT NULL UNIQUE,
|
||
password_hash TEXT NOT NULL,
|
||
salt TEXT NOT NULL,
|
||
role TEXT NOT NULL DEFAULT 'user',
|
||
authorizations TEXT NOT NULL DEFAULT '[]',
|
||
created_at TEXT NOT NULL
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS auth_tokens (
|
||
token TEXT PRIMARY KEY,
|
||
user_id TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
|
||
expires_at TEXT NOT NULL,
|
||
created_at TEXT NOT NULL,
|
||
revoked INTEGER NOT NULL DEFAULT 0
|
||
);
|
||
|
||
CREATE INDEX IF NOT EXISTS idx_auth_tokens_user ON auth_tokens(user_id);
|
||
|
||
CREATE TABLE IF NOT EXISTS settings (
|
||
key TEXT PRIMARY KEY,
|
||
value TEXT NOT NULL
|
||
);
|
||
"""
|
||
|
||
|
||
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()
|
||
# 启动清理（缺陷 #52）：running/queued 任务因服务重启中断 →
|
||
# 置 error（可读文案），避免前端永久挂「生成中」
|
||
await self._conn.execute(
|
||
"UPDATE docgen_tasks SET status = 'error', "
|
||
"error = '服务重启，任务中断，请重试', updated_at = ? "
|
||
"WHERE status IN ('running', 'queued')",
|
||
(_now_iso(),),
|
||
)
|
||
except Exception:
|
||
logger.exception("docgen 表迁移失败（docgen 端点将不可用，其余不受影响）")
|
||
# 管理后台/鉴权增量表(users + auth_tokens):单独执行;失败只记日志
|
||
try:
|
||
await self._conn.executescript(_ADMIN_SCHEMA)
|
||
except Exception:
|
||
logger.exception("admin 表迁移失败（管理后台将不可用，其余不受影响）")
|
||
# sessions 加 owner_id（用户级数据隔离，幂等 ALTER）
|
||
try:
|
||
cur = await self._conn.execute("PRAGMA table_info(sessions)")
|
||
cols = {row[1] for row in await cur.fetchall()}
|
||
if "owner_id" not in cols:
|
||
await self._conn.execute("ALTER TABLE sessions ADD COLUMN owner_id TEXT")
|
||
except Exception:
|
||
logger.exception("sessions owner_id 迁移失败")
|
||
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()
|