Maas-Model-platform/backend/scripts/test_db_connection.py

"""
数据库连接测试脚本

用于测试 .env 文件中的数据库配置是否正确
"""
import os
import sys
from pathlib import Path
from dotenv import load_dotenv
import psycopg2

# 添加项目根目录到路径
sys.path.insert(0, str(Path(__file__).parent.parent))

# 加载环境变量
env_path = Path(__file__).parent.parent / '.env'
load_dotenv(dotenv_path=env_path)


def test_connection():
    """测试数据库连接"""
    print("=" * 60)
    print("数据库连接测试")
    print("=" * 60)
    
    # 显示配置信息
    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')}")
    print(f"   Password: {'*' * len(os.getenv('DB_PASSWORD', ''))}")
    
    # 尝试连接
    print(f"\n🔌 正在连接数据库...")
    
    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", "")
        )
        
        print(f"✓ 数据库连接成功！")
        
        # 获取数据库版本
        cursor = conn.cursor()
        cursor.execute("SELECT version();")
        version = cursor.fetchone()[0]
        print(f"\n📊 数据库信息:")
        print(f"   版本: {version.split(',')[0]}")
        
        # 检查schema是否存在
        cursor.execute("""
            SELECT schema_name 
            FROM information_schema.schemata 
            WHERE schema_name = 'aigcspace';
        """)
        schema_exists = cursor.fetchone()
        
        if schema_exists:
            print(f"   ✓ aigcspace schema 存在")
            
            # 检查相关表是否存在
            cursor.execute("""
                SELECT table_name 
                FROM information_schema.tables 
                WHERE table_schema = 'aigcspace' 
                AND table_name IN ('model', 'model_price', 'toolbox_apps', 'image_translation')
                ORDER BY table_name;
            """)
            tables = cursor.fetchall()
            
            print(f"\n📋 相关表状态:")
            required_tables = ['model', 'model_price', 'toolbox_apps', 'image_translation']
            existing_tables = [t[0] for t in tables]
            
            for table in required_tables:
                if table in existing_tables:
                    print(f"   ✓ {table} - 已存在")
                else:
                    print(f"   ✗ {table} - 不存在")
        else:
            print(f"   ✗ aigcspace schema 不存在")
            print(f"\n⚠️  警告: 需要先创建 aigcspace schema")
        
        cursor.close()
        conn.close()
        
        print(f"\n" + "=" * 60)
        print("✅ 连接测试完成！")
        print("=" * 60)
        
        if not schema_exists:
            print(f"\n💡 下一步:")
            print(f"   1. 创建 schema: CREATE SCHEMA aigcspace;")
            print(f"   2. 执行相关表的迁移脚本")
        elif 'image_translation' not in existing_tables:
            print(f"\n💡 下一步:")
            print(f"   执行初始化脚本: python scripts/init_image_translation.py")
        else:
            print(f"\n💡 提示:")
            print(f"   image_translation 表已存在，无需重复初始化")
        
        print()
        
    except Exception as e:
        print(f"❌ 数据库连接失败！")
        print(f"\n错误信息: {e}")
        print(f"\n💡 请检查:")
        print(f"   1. 数据库服务是否运行")
        print(f"   2. .env 文件中的配置是否正确")
        print(f"   3. 网络连接是否正常")
        print(f"   4. 用户权限是否足够")
        print()
        sys.exit(1)


if __name__ == "__main__":
    try:
        test_connection()
    except KeyboardInterrupt:
        print(f"\n\n⚠️  用户中断执行")
        sys.exit(1)