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("数据库结构已是最新，无需迁移")
|