45 lines
1.6 KiB
Python
45 lines
1.6 KiB
Python
#!/usr/bin/env python3
|
||
"""数据库迁移脚本：添加技术分类和输出类型字段"""
|
||
|
||
import sys
|
||
import os
|
||
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
|
||
|
||
from sqlalchemy import text
|
||
from app.models.database import engine
|
||
|
||
def migrate():
|
||
"""执行数据库迁移"""
|
||
try:
|
||
with engine.connect() as conn:
|
||
# 检查字段是否已存在（PostgreSQL语法）
|
||
result = conn.execute(text("""
|
||
SELECT column_name
|
||
FROM information_schema.columns
|
||
WHERE table_name = 'algorithms'
|
||
"""))
|
||
columns = [row[0] for row in result.fetchall()]
|
||
|
||
# 添加 tech_category 字段
|
||
if 'tech_category' not in columns:
|
||
conn.execute(text("ALTER TABLE algorithms ADD COLUMN tech_category VARCHAR(50) DEFAULT 'computer_vision'"))
|
||
print("✓ 已添加 tech_category 字段")
|
||
else:
|
||
print("✓ tech_category 字段已存在")
|
||
|
||
# 添加 output_type 字段
|
||
if 'output_type' not in columns:
|
||
conn.execute(text("ALTER TABLE algorithms ADD COLUMN output_type VARCHAR(50) DEFAULT 'image'"))
|
||
print("✓ 已添加 output_type 字段")
|
||
else:
|
||
print("✓ output_type 字段已存在")
|
||
|
||
conn.commit()
|
||
print("\n数据库迁移完成！")
|
||
|
||
except Exception as e:
|
||
print(f"数据库迁移失败: {e}")
|
||
sys.exit(1)
|
||
|
||
if __name__ == "__main__":
|
||
migrate() |