rbox/mywebdav/monitoring/health.py

import asyncio
import time
import os
from datetime import datetime
from typing import Dict, Any
from fastapi import APIRouter, Response
from pydantic import BaseModel

router = APIRouter(prefix="/health", tags=["health"])


class HealthCheck(BaseModel):
    status: str
    timestamp: str
    checks: Dict[str, Any]


class ReadinessCheck(BaseModel):
    ready: bool


class LivenessCheck(BaseModel):
    alive: bool


async def check_database() -> Dict[str, Any]:
    try:
        from ..database import get_user_db_manager
        db_manager = get_user_db_manager()
        async with db_manager.get_master_connection() as conn:
            cursor = await conn.execute("SELECT 1")
            await cursor.fetchone()
        return {"ok": True, "message": "Connected"}
    except Exception as e:
        return {"ok": False, "message": str(e)}


async def check_cache() -> Dict[str, Any]:
    try:
        from ..cache import get_cache
        cache = get_cache()
        stats = cache.get_stats()
        return {"ok": True, "stats": stats}
    except Exception as e:
        return {"ok": False, "message": str(e)}


async def check_locks() -> Dict[str, Any]:
    try:
        from ..concurrency import get_lock_manager
        lock_manager = get_lock_manager()
        stats = await lock_manager.get_stats()
        return {"ok": True, "stats": stats}
    except Exception as e:
        return {"ok": False, "message": str(e)}


async def check_task_queue() -> Dict[str, Any]:
    try:
        from ..workers import get_task_queue
        queue = get_task_queue()
        stats = await queue.get_stats()
        return {"ok": True, "stats": stats}
    except Exception as e:
        return {"ok": False, "message": str(e)}


async def check_storage() -> Dict[str, Any]:
    try:
        from ..settings import settings
        storage_path = settings.STORAGE_PATH
        if os.path.exists(storage_path):
            stat = os.statvfs(storage_path)
            free_bytes = stat.f_bavail * stat.f_frsize
            total_bytes = stat.f_blocks * stat.f_frsize
            used_percent = ((total_bytes - free_bytes) / total_bytes) * 100
            return {
                "ok": True,
                "free_gb": round(free_bytes / (1024**3), 2),
                "total_gb": round(total_bytes / (1024**3), 2),
                "used_percent": round(used_percent, 2)
            }
        return {"ok": False, "message": "Storage path does not exist"}
    except Exception as e:
        return {"ok": False, "message": str(e)}


@router.get("", response_model=HealthCheck)
async def health_check():
    checks = {}
    check_funcs = {
        "database": check_database,
        "cache": check_cache,
        "locks": check_locks,
        "task_queue": check_task_queue,
        "storage": check_storage,
    }

    results = await asyncio.gather(
        *[func() for func in check_funcs.values()],
        return_exceptions=True
    )

    for name, result in zip(check_funcs.keys(), results):
        if isinstance(result, Exception):
            checks[name] = {"ok": False, "message": str(result)}
        else:
            checks[name] = result

    all_ok = all(check.get("ok", False) for check in checks.values())
    status = "healthy" if all_ok else "degraded"

    return HealthCheck(
        status=status,
        timestamp=datetime.utcnow().isoformat(),
        checks=checks
    )


@router.get("/ready", response_model=ReadinessCheck)
async def readiness_check():
    try:
        db_check = await check_database()
        return ReadinessCheck(ready=db_check.get("ok", False))
    except Exception:
        return ReadinessCheck(ready=False)


@router.get("/live", response_model=LivenessCheck)
async def liveness_check():
    return LivenessCheck(alive=True)


@router.get("/metrics")
async def metrics():
    metrics_data = []

    try:
        from ..cache import get_cache
        cache = get_cache()
        stats = cache.get_stats()
        metrics_data.append(f'cache_hits_total {stats.get("hits", 0)}')
        metrics_data.append(f'cache_misses_total {stats.get("misses", 0)}')
        metrics_data.append(f'cache_hit_rate_percent {stats.get("hit_rate_percent", 0)}')
    except Exception:
        pass

    try:
        from ..concurrency import get_lock_manager
        lock_manager = get_lock_manager()
        stats = await lock_manager.get_stats()
        metrics_data.append(f'locks_total {stats.get("total_locks", 0)}')
        metrics_data.append(f'locks_active {stats.get("active_locks", 0)}')
    except Exception:
        pass

    try:
        from ..workers import get_task_queue
        queue = get_task_queue()
        stats = await queue.get_stats()
        metrics_data.append(f'tasks_enqueued_total {stats.get("enqueued", 0)}')
        metrics_data.append(f'tasks_completed_total {stats.get("completed", 0)}')
        metrics_data.append(f'tasks_failed_total {stats.get("failed", 0)}')
        metrics_data.append(f'tasks_pending {stats.get("pending_tasks", 0)}')
    except Exception:
        pass

    try:
        storage_check = await check_storage()
        if storage_check.get("ok"):
            metrics_data.append(f'storage_free_gb {storage_check.get("free_gb", 0)}')
            metrics_data.append(f'storage_used_percent {storage_check.get("used_percent", 0)}')
    except Exception:
        pass

    return Response(
        content="\n".join(metrics_data),
        media_type="text/plain"
    )