doris-sync/config/checker.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
初始化检测模块
"""

import sys
from typing import Dict, Any

from config.loader import get_config
from db.pool import get_mysql_conn, get_doris_conn


def init_check() -> bool:
    """
    初始化检测
    
    检测 MySQL 和 Doris 数据库连接是否正常
    
    Returns:
        所有连接是否正常
    """
    print("=" * 50)
    print("初始化检测")
    print("=" * 50)
    
    mysql_cfg: Dict[str, Any] = get_config("database.mysql", {})
    doris_cfg: Dict[str, Any] = get_config("database.doris", {})
    
    print(f"\nMySQL 配置:")
    print(f"  Host: {mysql_cfg.get('host', 'N/A')}")
    print(f"  Port: {mysql_cfg.get('port', 'N/A')}")
    print(f"  Database: {mysql_cfg.get('database', 'N/A')}")
    
    print(f"\nDoris 配置:")
    print(f"  Host: {doris_cfg.get('host', 'N/A')}")
    print(f"  Port: {doris_cfg.get('port', 'N/A')}")
    print(f"  Database: {doris_cfg.get('database', 'N/A')}")
    
    print("\n" + "-" * 50)
    
    # 检测 MySQL
    mysql_ok: bool = _check_mysql()
    
    # 检测 Doris
    doris_ok: bool = _check_doris()
    
    print("=" * 50)
    
    if mysql_ok and doris_ok:
        print("初始化检测完成，所有连接正常")
    else:
        print("初始化检测完成，部分连接失败")
        
    print("=" * 50)
    print()
    
    return mysql_ok and doris_ok


def _check_mysql() -> bool:
    """检测 MySQL 连接"""
    try:
        conn = get_mysql_conn()
        cursor = conn.cursor()
        cursor.execute("SELECT 1")
        result = cursor.fetchone()
        cursor.close()
        conn.close()

        if result and result.get("1") == 1:
            print("[OK] MySQL 连接成功")
            return True
        else:
            print("[警告] MySQL 连接异常")
            return False

    except Exception as e:
        print(f"[错误] MySQL 连接失败: {e}")
        return False


def _check_doris() -> bool:
    """检测 Doris 连接"""
    try:
        conn = get_doris_conn()
        cursor = conn.cursor()
        cursor.execute("SELECT 1")
        result = cursor.fetchone()
        cursor.close()
        conn.close()

        if result and result.get("1") == 1:
            print("[OK] Doris 连接成功")
            return True
        else:
            print("[警告] Doris 连接异常")
            return False

    except Exception as e:
        print(f"[错误] Doris 连接失败: {e}")
        return False


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