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

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

import logging
import platform
import socket

import psutil
from flask import current_app

from app import db
from app.core.events import outbox_health
from app.services.db_healthcheck import check_database_connection
from app.services.neo4j_driver import Neo4jDriver

logger = logging.getLogger(__name__)


def data_source_pool_health():
    """Aggregate existing Worker pools without creating the runtime."""
    from app.core.data_source.runtime import peek_data_source_manager

    manager = peek_data_source_manager()
    statuses = manager.snapshot() if manager is not None else []
    return {
        "active_pool_count": sum(
            1 for status in statuses if not status.draining
        ),
        "degraded_pool_count": sum(
            1 for status in statuses if status.pool_state == "degraded"
        ),
        "open_circuit_count": sum(
            1 for status in statuses if status.pool_state == "open"
        ),
        "checked_out_total": sum(
            int(status.checked_out) for status in statuses
        ),
        "pool_timeout_total": sum(
            int(status.pool_timeout_total) for status in statuses
        ),
    }


def check_neo4j_connection():
    """
    检查Neo4j数据库连接状态

    Returns:
        bool: 连接成功返回True，失败返回False
    """
    try:
        with Neo4jDriver().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():
    """检查系统各个组件的健康状态"""
    health_status = {
        "database": check_database_connection(),
        "neo4j": Neo4jDriver().verify_connectivity(),
        "environment": current_app.config["FLASK_ENV"],
        "platform": current_app.config["PLATFORM"],
        "datasource_pools": data_source_pool_health(),
    }
    try:
        health_status["outbox"] = outbox_health(db.session)
    except Exception as exc:
        logger.warning("Outbox health unavailable: %s", exc)
        health_status["outbox"] = {"status": "unavailable"}

    # 检查所有组件是否都正常
    all_healthy = all([health_status["database"], health_status["neo4j"]])

    return {
        "status": "healthy" if all_healthy else "unhealthy",
        "components": 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": current_app.config["FLASK_ENV"],
                "debug_mode": current_app.config["DEBUG"],
                "port": current_app.config["PORT"],
                "platform": current_app.config["PLATFORM"],
                "bucket_name": current_app.config["BUCKET_NAME"],
                "prefix": current_app.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