""" 图像翻译模块数据库迁移脚本 功能： 1. 创建 image_translation 表 2. 插入模型配置（model 和 model_price） 3. 插入百宝箱应用数据（toolbox_apps） 4. 验证数据完整性 从 .env 文件读取数据库连接信息 """ import os import sys from pathlib import Path from dotenv import load_dotenv import psycopg2 from psycopg2 import sql # 添加项目根目录到路径 sys.path.insert(0, str(Path(__file__).parent.parent)) # 加载环境变量 env_path = Path(__file__).parent.parent / '.env' load_dotenv(dotenv_path=env_path) def get_db_connection(): """从环境变量获取数据库连接""" try: conn = psycopg2.connect( host=os.getenv("DB_HOST", "localhost"), port=int(os.getenv("DB_PORT", "5432")), database=os.getenv("DB_NAME", "aigcspace"), user=os.getenv("DB_USER", "postgres"), password=os.getenv("DB_PASSWORD", "") ) return conn except Exception as e: print(f"❌ 数据库连接失败: {e}") sys.exit(1) def execute_sql_file(cursor, file_path: Path): """执行SQL文件""" try: with open(file_path, 'r', encoding='utf-8') as f: sql_content = f.read() # 执行SQL cursor.execute(sql_content) print(f"✓ 执行成功: {file_path.name}") return True except Exception as e: print(f"❌ 执行失败: {file_path.name}") print(f" 错误信息: {e}") return False def insert_model_data(cursor): """插入模型配置数据""" try: # 检查模型是否已存在 cursor.execute(""" SELECT id FROM aigcspace.model WHERE title = 'qwen-mt-image' LIMIT 1; """) existing_model = cursor.fetchone() if existing_model: print(f"✓ 模型 qwen-mt-image 已存在 (ID: {existing_model[0]})") model_id = existing_model[0] else: # 插入模型记录 cursor.execute(""" INSERT INTO aigcspace.model ( title, description, provider, model_type, status, created_at, updated_at ) VALUES ( 'qwen-mt-image', '通义千问图像翻译模型，支持多语种图像文字翻译', 'Alibaba', 'image-translation', 'active', NOW(), NOW() ) RETURNING id; """) model_id = cursor.fetchone()[0] print(f"✓ 模型 qwen-mt-image 插入成功 (ID: {model_id})") # 检查价格配置是否已存在 cursor.execute(""" SELECT id FROM aigcspace.model_price WHERE model_id = %s LIMIT 1; """, (model_id,)) existing_price = cursor.fetchone() if existing_price: print(f"✓ 模型价格配置已存在 (ID: {existing_price[0]})") else: # 插入价格配置 cursor.execute(""" INSERT INTO aigcspace.model_price ( model_id, input_price, output_price, created_at, updated_at ) VALUES ( %s, 0.0000, 0.0200, NOW(), NOW() ); """, (model_id,)) print(f"✓ 模型价格配置插入成功（每张图片 0.02 元）") return True except Exception as e: print(f"❌ 模型数据插入失败: {e}") return False def insert_toolbox_app(cursor): """插入百宝箱应用数据""" try: # 检查应用是否已存在 cursor.execute(""" SELECT id FROM aigcspace.toolbox_apps WHERE app_key = 'qwen-image-translate' LIMIT 1; """) existing_app = cursor.fetchone() if existing_app: print(f"✓ 百宝箱应用 qwen-image-translate 已存在 (ID: {existing_app[0]})") return True # 插入应用记录 cursor.execute(""" INSERT INTO aigcspace.toolbox_apps ( app_key, name, description, category, icon, is_new, is_hot, status, sort_order, version, created_at, updated_at ) VALUES ( 'qwen-image-translate', '通义千问图像翻译', '支持多语种图像文字翻译，保持原图排版和风格', '图像处理', 'translate', true, true, 'active', 100, '1.0.0', NOW(), NOW() ) RETURNING id; """) app_id = cursor.fetchone()[0] print(f"✓ 百宝箱应用 qwen-image-translate 插入成功 (ID: {app_id})") return True except Exception as e: print(f"❌ 百宝箱应用插入失败: {e}") return False def verify_data(cursor): """验证数据是否正确插入""" try: print(f"\n📊 数据验证:") # 验证表是否创建 cursor.execute(""" SELECT EXISTS ( SELECT FROM information_schema.tables WHERE table_schema = 'aigcspace' AND table_name = 'image_translation' ); """) table_exists = cursor.fetchone()[0] print(f" ✓ image_translation 表: {'已创建' if table_exists else '❌ 未创建'}") # 验证模型配置 cursor.execute(""" SELECT m.id, m.title, mp.output_price FROM aigcspace.model m LEFT JOIN aigcspace.model_price mp ON m.id = mp.model_id WHERE m.title = 'qwen-mt-image'; """) model_data = cursor.fetchone() if model_data: print(f" ✓ 模型配置: qwen-mt-image (ID: {model_data[0]}, 价格: {model_data[2]}元/张)") else: print(f" ❌ 模型配置未找到") # 验证百宝箱应用 cursor.execute(""" SELECT id, name, category, is_new, version, status FROM aigcspace.toolbox_apps WHERE app_key = 'qwen-image-translate'; """) app_data = cursor.fetchone() if app_data: app_id, name, category, is_new, version, status = app_data new_tag = " [New]" if is_new else "" print(f" ✓ 百宝箱应用: {name}{new_tag} (版本: {version}, 状态: {status})") else: print(f" ❌ 百宝箱应用未找到") # 验证索引 cursor.execute(""" SELECT indexname FROM pg_indexes WHERE schemaname = 'aigcspace' AND tablename = 'image_translation'; """) indexes = cursor.fetchall() print(f" ✓ 索引数量: {len(indexes)}") for idx in indexes: print(f" - {idx[0]}") return True except Exception as e: print(f"❌ 数据验证失败: {e}") return False def main(): """主函数""" print("=" * 60) print("图像翻译模块数据库迁移脚本") print("=" * 60) # 获取迁移文件路径 migrations_dir = Path(__file__).parent.parent / 'migrations' migration_file = migrations_dir / 'add_image_translation_table.sql' # 检查文件是否存在 if not migration_file.exists(): print(f"❌ 迁移文件不存在: {migration_file}") sys.exit(1) print(f"\n📁 找到迁移文件:") print(f" - {migration_file.name}") # 连接数据库 print(f"\n🔌 连接数据库...") print(f" Host: {os.getenv('DB_HOST', 'localhost')}") print(f" Port: {os.getenv('DB_PORT', '5432')}") print(f" Database: {os.getenv('DB_NAME', 'aigcspace')}") print(f" User: {os.getenv('DB_USER', 'postgres')}") conn = get_db_connection() cursor = conn.cursor() print(f"✓ 数据库连接成功") # 执行迁移 print(f"\n🚀 开始执行迁移...") # 1. 执行SQL文件（创建表） print(f"\n[1/3] 创建数据表...") if not execute_sql_file(cursor, migration_file): conn.rollback() cursor.close() conn.close() print(f"\n❌ 迁移失败，已回滚所有更改") sys.exit(1) # 提交事务 conn.commit() print(f"\n✓ 所有迁移执行成功") # 验证数据 verify_data(cursor) # 关闭连接 cursor.close() conn.close() print(f"\n" + "=" * 60) print("✅ 图像翻译模块数据库迁移完成！") print("=" * 60) print(f"\n💡 下一步:") print(f" 1. 重启后端服务: cd backend && python main.py") print(f" 2. 访问百宝箱: http://localhost:5173/toolbox") print(f" 3. 点击「通义千问图像翻译」卡片的「立即使用」按钮") print(f" 4. 开始使用图像翻译功能") print() if __name__ == "__main__": try: main() except KeyboardInterrupt: print(f"\n\n⚠️ 用户中断执行") sys.exit(1) except Exception as e: print(f"\n\n❌ 执行过程中发生错误: {e}") import traceback traceback.print_exc() sys.exit(1)