forked from xiaohei/taiji-AI-PAD
145 lines
4.5 KiB
Python
145 lines
4.5 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Migration 016: Add Agent Access Info Fields
|
|
|
|
为 agent_billing_records 表添加访问信息字段（域名、外网IP等）
|
|
|
|
Usage:
|
|
python migrations/run_016_migration.py
|
|
"""
|
|
|
|
import asyncio
|
|
import os
|
|
import sys
|
|
|
|
# 添加项目根目录到 Python 路径
|
|
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():
|
|
"""运行迁移"""
|
|
print("=" * 60)
|
|
print("Migration 016: Add Agent Access Info Fields")
|
|
print("=" * 60)
|
|
|
|
# 读取 SQL 文件
|
|
sql_file = os.path.join(os.path.dirname(__file__), "016_add_agent_access_info_fields.sql")
|
|
with open(sql_file, "r", encoding="utf-8") as f:
|
|
sql_content = f.read()
|
|
|
|
# 分割 SQL 语句
|
|
statements = []
|
|
current_stmt = []
|
|
in_do_block = False
|
|
|
|
for line in sql_content.split("\n"):
|
|
line_stripped = line.strip()
|
|
|
|
# 跳过注释和空行（但保留 DO 块中的内容）
|
|
if not in_do_block:
|
|
if line_stripped.startswith("--") or not line_stripped:
|
|
continue
|
|
|
|
# 检测 DO 块开始
|
|
if line_stripped.upper().startswith("DO $$"):
|
|
in_do_block = True
|
|
current_stmt.append(line)
|
|
continue
|
|
|
|
# 检测 DO 块结束
|
|
if in_do_block and line_stripped == "END $$;":
|
|
current_stmt.append(line)
|
|
statements.append("\n".join(current_stmt))
|
|
current_stmt = []
|
|
in_do_block = False
|
|
continue
|
|
|
|
if in_do_block:
|
|
current_stmt.append(line)
|
|
continue
|
|
|
|
# 普通语句处理
|
|
current_stmt.append(line)
|
|
if line_stripped.endswith(";"):
|
|
stmt = "\n".join(current_stmt)
|
|
if stmt.strip():
|
|
statements.append(stmt)
|
|
current_stmt = []
|
|
|
|
# 执行迁移
|
|
async with engine.begin() as conn:
|
|
for i, stmt in enumerate(statements, 1):
|
|
try:
|
|
# 打印语句摘要
|
|
stmt_preview = stmt.strip()[:80].replace("\n", " ")
|
|
if len(stmt.strip()) > 80:
|
|
stmt_preview += "..."
|
|
print(f"\n[{i}/{len(statements)}] Executing: {stmt_preview}")
|
|
|
|
await conn.execute(text(stmt))
|
|
print(f" ✓ Success")
|
|
except Exception as e:
|
|
error_msg = str(e)
|
|
# 忽略 "column already exists" 错误
|
|
if "already exists" in error_msg.lower():
|
|
print(f" ⚠ Skipped (already exists)")
|
|
else:
|
|
print(f" ✗ Error: {error_msg}")
|
|
raise
|
|
|
|
print("\n" + "=" * 60)
|
|
print("Migration 016 completed successfully!")
|
|
print("=" * 60)
|
|
|
|
|
|
async def verify_migration():
|
|
"""验证迁移结果"""
|
|
print("\nVerifying migration...")
|
|
|
|
async with engine.begin() as conn:
|
|
# 检查新字段是否存在
|
|
result = await conn.execute(text("""
|
|
SELECT column_name, data_type, character_maximum_length
|
|
FROM information_schema.columns
|
|
WHERE table_name = 'agent_billing_records'
|
|
AND column_name IN ('external_ip', 'domain', 'domain_url', 'access_url', 'service_port', 'namespace')
|
|
ORDER BY column_name
|
|
"""))
|
|
columns = result.fetchall()
|
|
|
|
print(f"\nNew columns in agent_billing_records table:")
|
|
print("-" * 50)
|
|
for col in columns:
|
|
col_name, data_type, max_len = col
|
|
type_info = f"{data_type}({max_len})" if max_len else data_type
|
|
print(f" ✓ {col_name}: {type_info}")
|
|
|
|
expected_columns = {'external_ip', 'domain', 'domain_url', 'access_url', 'service_port', 'namespace'}
|
|
found_columns = {col[0] for col in columns}
|
|
|
|
if found_columns == expected_columns:
|
|
print(f"\n✓ All {len(expected_columns)} columns verified successfully!")
|
|
else:
|
|
missing = expected_columns - found_columns
|
|
if missing:
|
|
print(f"\n✗ Missing columns: {missing}")
|
|
return False
|
|
|
|
return True
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
asyncio.run(run_migration())
|
|
asyncio.run(verify_migration())
|
|
except KeyboardInterrupt:
|
|
print("\nMigration cancelled.")
|
|
sys.exit(1)
|
|
except Exception as e:
|
|
print(f"\nMigration failed: {e}")
|
|
sys.exit(1)
|
|
|