201 lines
6.8 KiB
Python
201 lines
6.8 KiB
Python
"""db.py —— 生信服务独立 SQLite（aiosqlite + WAL，无 ORM），三张表 + 幂等迁移。
|
||
|
||
FR-BIO-028：与 BFF 完全隔离（独立进程 + 独立 bioflow.db）。
|
||
表（实施计划 §六三表清单）：
|
||
- bio_jobs：id / session_id / title / status / stage_chain(JSON) / params(JSON) /
|
||
archive_id / auto_continue / progress / error / created_at / updated_at
|
||
- bio_stages：id / job_id(FK CASCADE) / stage / status / params(JSON) / manifest_path /
|
||
log_tail / attempt / started_at / finished_at / error
|
||
- bio_archives：id / job_id / name / description / stage_chain(JSON) / params(JSON) /
|
||
report_artifact / disk_path / size_bytes / reused_count / created_at
|
||
|
||
时间一律 ISO 8601 UTC（由服务侧生成，同 docgen 模式）。
|
||
模式沿用 docgen 同款：CREATE TABLE IF NOT EXISTS + PRAGMA table_info 幂等加列。
|
||
"""
|
||
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 bio_jobs (
|
||
id TEXT PRIMARY KEY,
|
||
session_id TEXT NOT NULL DEFAULT '',
|
||
title TEXT NOT NULL DEFAULT '生信任务',
|
||
status TEXT NOT NULL DEFAULT 'created',
|
||
stage_chain TEXT NOT NULL DEFAULT '[]',
|
||
params TEXT NOT NULL DEFAULT '{}',
|
||
archive_id TEXT,
|
||
auto_continue INTEGER NOT NULL DEFAULT 0,
|
||
progress INTEGER NOT NULL DEFAULT 0,
|
||
error TEXT,
|
||
created_at TEXT NOT NULL,
|
||
updated_at TEXT NOT NULL
|
||
);
|
||
|
||
CREATE INDEX IF NOT EXISTS idx_bio_jobs_session ON bio_jobs(session_id);
|
||
CREATE INDEX IF NOT EXISTS idx_bio_jobs_status ON bio_jobs(status);
|
||
CREATE INDEX IF NOT EXISTS idx_bio_jobs_updated ON bio_jobs(updated_at);
|
||
|
||
CREATE TABLE IF NOT EXISTS bio_stages (
|
||
id TEXT PRIMARY KEY,
|
||
job_id TEXT NOT NULL REFERENCES bio_jobs(id) ON DELETE CASCADE,
|
||
stage TEXT NOT NULL,
|
||
status TEXT NOT NULL DEFAULT 'waiting',
|
||
params TEXT NOT NULL DEFAULT '{}',
|
||
manifest_path TEXT,
|
||
log_tail TEXT NOT NULL DEFAULT '',
|
||
attempt INTEGER NOT NULL DEFAULT 0,
|
||
started_at TEXT,
|
||
finished_at TEXT,
|
||
error TEXT,
|
||
UNIQUE(job_id, stage)
|
||
);
|
||
|
||
CREATE INDEX IF NOT EXISTS idx_bio_stages_job ON bio_stages(job_id);
|
||
|
||
CREATE TABLE IF NOT EXISTS bio_archives (
|
||
id TEXT PRIMARY KEY,
|
||
job_id TEXT NOT NULL,
|
||
name TEXT NOT NULL DEFAULT '',
|
||
description TEXT NOT NULL DEFAULT '',
|
||
stage_chain TEXT NOT NULL DEFAULT '[]',
|
||
params TEXT NOT NULL DEFAULT '{}',
|
||
report_artifact TEXT,
|
||
disk_path TEXT,
|
||
size_bytes INTEGER NOT NULL DEFAULT 0,
|
||
reused_count INTEGER NOT NULL DEFAULT 0,
|
||
created_at TEXT NOT NULL
|
||
);
|
||
|
||
CREATE INDEX IF NOT EXISTS idx_bio_archives_name ON bio_archives(name);
|
||
CREATE INDEX IF NOT EXISTS idx_bio_archives_created ON bio_archives(created_at);
|
||
"""
|
||
|
||
# 增量列幂等迁移（老库已建表时 CREATE IF NOT EXISTS 不会加列）
|
||
_BIO_JOBS_MIGRATIONS = (
|
||
("archive_id", "ALTER TABLE bio_jobs ADD COLUMN archive_id TEXT"),
|
||
("error", "ALTER TABLE bio_jobs ADD COLUMN error TEXT"),
|
||
)
|
||
_BIO_ARCHIVES_MIGRATIONS = (
|
||
("report_artifact", "ALTER TABLE bio_archives ADD COLUMN report_artifact TEXT"),
|
||
)
|
||
|
||
|
||
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；写串行化，同 docgen 模式）。"""
|
||
|
||
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._migrate_columns()
|
||
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_columns(self) -> None:
|
||
"""三表增量列幂等迁移（同 docgen _migrate_docgen_columns 模式）。"""
|
||
for table, migrations in (
|
||
("bio_jobs", _BIO_JOBS_MIGRATIONS),
|
||
("bio_archives", _BIO_ARCHIVES_MIGRATIONS),
|
||
):
|
||
cur = await self._conn.execute(f"PRAGMA table_info({table})")
|
||
cols = {row[1] for row in await cur.fetchall()}
|
||
for name, ddl in 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()
|
||
|
||
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:
|
||
global db
|
||
if path is not None:
|
||
db = Database(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()
|