45 lines
1.7 KiB
Python
45 lines
1.7 KiB
Python
#!/usr/bin/env python3
|
||
"""数据库迁移脚本：为algorithm_services表添加技术分类和输出类型字段"""
|
||
|
||
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 = 'algorithm_services'
|
||
"""))
|
||
columns = [row[0] for row in result.fetchall()]
|
||
|
||
# 添加 tech_category 字段
|
||
if 'tech_category' not in columns:
|
||
conn.execute(text("ALTER TABLE algorithm_services ADD COLUMN tech_category VARCHAR(50) DEFAULT 'computer_vision'"))
|
||
print("✓ 已添加 tech_category 字段到 algorithm_services 表")
|
||
else:
|
||
print("✓ tech_category 字段已存在于 algorithm_services 表")
|
||
|
||
# 添加 output_type 字段
|
||
if 'output_type' not in columns:
|
||
conn.execute(text("ALTER TABLE algorithm_services ADD COLUMN output_type VARCHAR(50) DEFAULT 'image'"))
|
||
print("✓ 已添加 output_type 字段到 algorithm_services 表")
|
||
else:
|
||
print("✓ output_type 字段已存在于 algorithm_services 表")
|
||
|
||
conn.commit()
|
||
print("\n数据库迁移完成！")
|
||
|
||
except Exception as e:
|
||
print(f"数据库迁移失败: {e}")
|
||
sys.exit(1)
|
||
|
||
if __name__ == "__main__":
|
||
migrate() |