video_decode/examples/test_db_connection.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
测试数据库连接

验证数据库配置是否正确，连接是否正常
"""

import sys
from src.models import get_engine, init_db, DecodeVideo
from src.models.database import get_database_url
from src.utils.logger import get_logger

logger = get_logger(__name__)


def test_connection():
    """测试数据库连接"""
    print("=" * 60)
    print("数据库连接测试")
    print("=" * 60)
    
    try:
        # 1. 显示连接信息（隐藏密码）
        database_url = get_database_url()
        # 隐藏密码显示
        safe_url = database_url.split('@')[0].split(':')[0] + ':***@' + '@'.join(database_url.split('@')[1:])
        print(f"\n✓ 数据库连接URL: {safe_url}")
        
        # 2. 测试连接
        print("\n正在测试数据库连接...")
        engine = get_engine()
        
        # 尝试连接
        with engine.connect() as conn:
            result = conn.execute("SELECT 1 as test")
            row = result.fetchone()
            if row and row[0] == 1:
                print("✓ 数据库连接成功！")
            else:
                print("✗ 数据库连接测试失败")
                return False
        
        # 3. 测试表是否存在
        print("\n检查 decode_videos 表...")
        from sqlalchemy import inspect
        inspector = inspect(engine)
        tables = inspector.get_table_names()
        
        if 'decode_videos' in tables:
            print("✓ decode_videos 表已存在")
            
            # 显示表结构
            columns = inspector.get_columns('decode_videos')
            print(f"\n表结构 ({len(columns)} 个字段):")
            for col in columns:
                nullable = "NULL" if col['nullable'] else "NOT NULL"
                print(f"  - {col['name']}: {col['type']} ({nullable})")
            
            # 统计记录数
            from src.models import get_db
            db = next(get_db())
            try:
                count = db.query(DecodeVideo).count()
                print(f"\n✓ 当前记录数: {count}")
            except Exception as e:
                print(f"⚠ 查询记录数时出错: {e}")
            finally:
                db.close()
        else:
            print("⚠ decode_videos 表不存在")
            print("\n是否要创建表？(y/n): ", end="")
            choice = input().strip().lower()
            if choice == 'y':
                print("\n正在创建表...")
                init_db()
                print("✓ 表创建完成")
            else:
                print("跳过表创建")
        
        print("\n" + "=" * 60)
        print("✓ 所有测试通过！")
        print("=" * 60)
        return True
        
    except Exception as e:
        print(f"\n✗ 数据库连接失败: {e}")
        print("\n请检查以下配置：")
        print("  1. 数据库主机地址是否正确")
        print("  2. 数据库端口是否开放")
        print("  3. 用户名和密码是否正确")
        print("  4. 数据库名称是否正确")
        print("  5. 网络连接是否正常（外网地址需要公网访问权限）")
        import traceback
        traceback.print_exc()
        return False


if __name__ == "__main__":
    success = test_connection()
    sys.exit(0 if success else 1)