DataOps-platform/app/core/system/health.py

"""
系统健康检查模块
提供系统各组件健康状态检查和系统信息获取功能
"""

import logging
import platform
import psutil
import os
import socket
from datetime import datetime
from app.config.config import Config
from app.services.neo4j_driver import neo4j_driver

logger = logging.getLogger(__name__)

def check_neo4j_connection():
    """
    检查Neo4j数据库连接状态
    
    Returns:
        bool: 连接成功返回True，失败返回False
    """
    try:
        with neo4j_driver.get_session() as session:
            # 执行简单查询确认连接
            session.run("RETURN 1")
            return True
    except Exception as e:
        logger.error(f"Neo4j数据库连接失败: {str(e)}")
        return False

def check_system_health():
    """
    检查系统整体健康状态
    包括关键依赖组件的连接状态
    
    Returns:
        dict: 包含各组件健康状态的字典
    """
    # 检查Neo4j连接
    neo4j_status = check_neo4j_connection()
    
    # 可以添加其他组件的健康检查
    # 例如MySQL、Redis、MinIO等
    
    # 构造健康状态信息
    health_status = {
        "service": "DataOps-platform",
        "status": "UP" if neo4j_status else "DEGRADED",
        "version": "1.0.0",  # 可以从配置或版本文件中读取
        "time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "dependencies": {
            "neo4j": {
                "status": "UP" if neo4j_status else "DOWN",
                "details": {
                    "url": Config.NEO4J_URI,
                    "encrypted": Config.NEO4J_ENCRYPTED
                }
            }
            # 可以添加其他依赖的状态
        }
    }
    
    return health_status

def get_system_info():
    """
    获取系统运行环境信息
    包括操作系统、Python版本、CPU使用率、内存使用情况等
    
    Returns:
        dict: 包含系统信息的字典
    """
    try:
        # 获取基本系统信息
        sys_info = {
            "os": {
                "name": platform.system(),
                "version": platform.version(),
                "platform": platform.platform(),
            },
            "python": {
                "version": platform.python_version(),
                "implementation": platform.python_implementation(),
            },
            "network": {
                "hostname": socket.gethostname(),
                "ip": socket.gethostbyname(socket.gethostname()),
            },
            "resources": {
                "cpu": {
                    "cores": psutil.cpu_count(logical=False),
                    "logical_cores": psutil.cpu_count(logical=True),
                    "usage_percent": psutil.cpu_percent(interval=0.1),
                },
                "memory": {
                    "total": _format_bytes(psutil.virtual_memory().total),
                    "available": _format_bytes(psutil.virtual_memory().available),
                    "used": _format_bytes(psutil.virtual_memory().used),
                    "percent": psutil.virtual_memory().percent,
                },
                "disk": {
                    "total": _format_bytes(psutil.disk_usage("/").total),
                    "used": _format_bytes(psutil.disk_usage("/").used),
                    "free": _format_bytes(psutil.disk_usage("/").free),
                    "percent": psutil.disk_usage("/").percent,
                },
            },
            "application": {
                "environment": Config.ENVIRONMENT,
                "debug_mode": Config.DEBUG,
                "port": Config.PORT,
                "platform": Config.PLATFORM,
                "upload_folder": Config.UPLOAD_FOLDER,
                "bucket_name": Config.BUCKET_NAME,
                "prefix": Config.PREFIX,
                # 不返回敏感信息如密码、密钥等
            }
        }
        
        return sys_info
    except Exception as e:
        logger.error(f"获取系统信息失败: {str(e)}")
        return {"error": str(e)}

def _format_bytes(bytes_value):
    """
    将字节数格式化为易读形式
    
    Args:
        bytes_value: 字节数
        
    Returns:
        str: 格式化后的字符串，如"1.23 GB"
    """
    for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
        if bytes_value < 1024 or unit == 'TB':
            return f"{bytes_value:.2f} {unit}"
        bytes_value /= 1024