101 lines
3.6 KiB
Python
101 lines
3.6 KiB
Python
"""
|
||
轻量级数据库结构迁移
|
||
|
||
项目没有引入 Alembic，新增列通过幂等的 ALTER TABLE 完成：
|
||
先查询 information_schema 判断列是否已存在，不存在才执行 ALTER，
|
||
保证多次启动/执行也不会报错或重复加列。
|
||
"""
|
||
import logging
|
||
|
||
from sqlalchemy import text
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
# (表名, 列名, 建列语句)
|
||
CHAT_MESSAGE_COLUMNS = [
|
||
(
|
||
"chat_message",
|
||
"status",
|
||
"ALTER TABLE chat_message ADD COLUMN status VARCHAR(32) NOT NULL DEFAULT 'pending' "
|
||
"COMMENT '消息状态: pending/completed/interrupted/error'",
|
||
),
|
||
(
|
||
"chat_message",
|
||
"duration_ms",
|
||
"ALTER TABLE chat_message ADD COLUMN duration_ms INT DEFAULT NULL "
|
||
"COMMENT '生成耗时(毫秒)'",
|
||
),
|
||
(
|
||
"chat_message",
|
||
"thinking_log",
|
||
"ALTER TABLE chat_message ADD COLUMN thinking_log TEXT DEFAULT NULL "
|
||
"COMMENT '思考过程(JSON数组)'",
|
||
),
|
||
]
|
||
|
||
|
||
async def _column_exists(conn, table_name: str, column_name: str) -> bool:
|
||
result = await conn.execute(
|
||
text(
|
||
"SELECT COUNT(*) FROM information_schema.COLUMNS "
|
||
"WHERE TABLE_SCHEMA = DATABASE() "
|
||
"AND TABLE_NAME = :table_name AND COLUMN_NAME = :column_name"
|
||
),
|
||
{"table_name": table_name, "column_name": column_name},
|
||
)
|
||
return bool(result.scalar())
|
||
|
||
|
||
async def migrate_schema() -> None:
|
||
"""为存量数据库补齐新增列，并为历史助手消息回填状态。"""
|
||
from app.core.database import engine
|
||
|
||
added = []
|
||
async with engine.begin() as conn:
|
||
for table_name, column_name, ddl in CHAT_MESSAGE_COLUMNS:
|
||
try:
|
||
exists = await _column_exists(conn, table_name, column_name)
|
||
except Exception as exc: # noqa: BLE001
|
||
logger.warning("检查列 %s.%s 失败，跳过迁移: %s", table_name, column_name, exc)
|
||
return
|
||
if not exists:
|
||
await conn.execute(text(ddl))
|
||
added.append(f"{table_name}.{column_name}")
|
||
|
||
# 历史消息回填：旧的“中断”标记仅存在于 content 末尾，迁移后统一迁移到 status。
|
||
# 之后不再用内容比对判断中断，status 字段作为唯一依据。
|
||
await conn.execute(
|
||
text(
|
||
"UPDATE chat_message SET status = 'interrupted' "
|
||
"WHERE role = 'assistant' "
|
||
"AND (status IS NULL OR status = '' OR status = 'pending') "
|
||
"AND (content = 'interrupt' OR content LIKE '%\\n\\ninterrupt')"
|
||
)
|
||
)
|
||
# 清理历史“中断”标记文本：迁移后 status 是唯一依据，内容不再混入标记
|
||
await conn.execute(
|
||
text(
|
||
"UPDATE chat_message SET content = '' "
|
||
"WHERE status = 'interrupted' AND content = 'interrupt'"
|
||
)
|
||
)
|
||
await conn.execute(
|
||
text(
|
||
"UPDATE chat_message SET content = TRIM("
|
||
"LEFT(content, CHAR_LENGTH(content) - CHAR_LENGTH('\\n\\ninterrupt'))) "
|
||
"WHERE status = 'interrupted' AND content LIKE '%\\n\\ninterrupt'"
|
||
)
|
||
)
|
||
await conn.execute(
|
||
text(
|
||
"UPDATE chat_message SET status = 'completed' "
|
||
"WHERE (status IS NULL OR status = '' OR status = 'pending')"
|
||
)
|
||
)
|
||
|
||
if added:
|
||
logger.info("数据库迁移完成，新增列: %s", ", ".join(added))
|
||
else:
|
||
logger.info("数据库结构已是最新，无需迁移")
|