74 lines
2.4 KiB
Python
74 lines
2.4 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:
|
||
"""为存量数据库补齐新增列。
|
||
|
||
只做“新增列”这类非破坏性变更：不修改、不删除、不回填任何已有数据。
|
||
历史消息的中断状态由读取路径（get messages）按旧标记做只读兼容判断。
|
||
"""
|
||
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}")
|
||
|
||
if added:
|
||
logger.info("数据库迁移完成，新增列: %s", ", ".join(added))
|
||
else:
|
||
logger.info("数据库结构已是最新，无需迁移")
|