forked from xiaohei/taiji-AI-PAD
102 lines
3.0 KiB
Python
102 lines
3.0 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
迁移脚本：添加 LiteLLM 集成相关表和字段
|
|
|
|
运行方式：
|
|
cd services/mcp-server
|
|
python migrations/run_011_migration.py
|
|
"""
|
|
|
|
import asyncio
|
|
import os
|
|
import sys
|
|
|
|
# 添加父目录到路径
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from sqlalchemy import text
|
|
from database import engine
|
|
|
|
|
|
async def run_migration():
|
|
"""执行迁移"""
|
|
|
|
# 读取 SQL 文件
|
|
sql_file = os.path.join(os.path.dirname(__file__), "011_add_litellm_integration.sql")
|
|
|
|
with open(sql_file, "r", encoding="utf-8") as f:
|
|
sql_content = f.read()
|
|
|
|
# 分割 SQL 语句（按分号分割，但忽略函数体内的分号）
|
|
statements = []
|
|
current_statement = []
|
|
in_function = False
|
|
|
|
for line in sql_content.split("\n"):
|
|
stripped = line.strip()
|
|
|
|
# 跳过注释
|
|
if stripped.startswith("--"):
|
|
continue
|
|
|
|
# 检测函数开始
|
|
if "AS $$" in line or "AS $" in line:
|
|
in_function = True
|
|
|
|
# 检测函数结束
|
|
if in_function and ("$$ language" in line.lower() or "$$ LANGUAGE" in line):
|
|
in_function = False
|
|
|
|
current_statement.append(line)
|
|
|
|
# 如果不在函数内且行以分号结尾，则完成一条语句
|
|
if not in_function and stripped.endswith(";"):
|
|
statement = "\n".join(current_statement).strip()
|
|
if statement and not statement.startswith("--"):
|
|
statements.append(statement)
|
|
current_statement = []
|
|
|
|
# 处理最后一条语句
|
|
if current_statement:
|
|
statement = "\n".join(current_statement).strip()
|
|
if statement and not statement.startswith("--"):
|
|
statements.append(statement)
|
|
|
|
print("=" * 60)
|
|
print("LiteLLM 集成迁移脚本")
|
|
print("=" * 60)
|
|
print(f"共 {len(statements)} 条 SQL 语句待执行")
|
|
print()
|
|
|
|
async with engine.begin() as conn:
|
|
for i, statement in enumerate(statements, 1):
|
|
# 显示语句摘要
|
|
first_line = statement.split("\n")[0][:60]
|
|
print(f"[{i}/{len(statements)}] 执行: {first_line}...")
|
|
|
|
try:
|
|
await conn.execute(text(statement))
|
|
print(f" ✓ 成功")
|
|
except Exception as e:
|
|
error_msg = str(e)
|
|
# 忽略 "already exists" 类型的错误
|
|
if "already exists" in error_msg.lower():
|
|
print(f" ⚠ 已存在，跳过")
|
|
else:
|
|
print(f" ✗ 失败: {error_msg}")
|
|
raise
|
|
|
|
print()
|
|
print("=" * 60)
|
|
print("迁移完成！")
|
|
print("=" * 60)
|
|
print()
|
|
print("新增内容：")
|
|
print(" - channels 表添加 litellm_team_id 字段")
|
|
print(" - 创建 tenant_model_keys 表")
|
|
print(" - 添加相关索引和触发器")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(run_migration())
|