telegram-scraper/health.py

import sqlite3
from pathlib import Path
from typing import Any, Dict, List, Optional

from app_state import StateStore, load_account


def health_payload(
    data_dir: Path,
    session_dir: Path,
    state_store: StateStore,
    continuous_snapshot: Dict[str, Any],
    job_queue_size: int,
    account_ids: Optional[List[str]] = None,
) -> Dict[str, Any]:
    """
    Return a health-check payload for the application.

    NOTE: This endpoint does NOT support message search or query parameters.
    Health concerns itself with filesystem, database connectivity, job queue,
    and account session state.  Message search is handled by the message API
    endpoints (/api/channels/*/messages, /api/accounts/*/channels/*/messages).
    """
    checks = {
        "data_dir": _dir_check(data_dir, writable=True),
        "session_dir": _dir_check(session_dir, writable=True),
        "state_file": _state_check(state_store),
        "sqlite": _sqlite_check(),
        "continuous": _continuous_check(continuous_snapshot),
        "job_queue": {"ok": True, "size": job_queue_size},
    }

    # Per-account health
    if account_ids:
        account_checks: Dict[str, Any] = {}
        for acc_id in account_ids:
            acc_dir = data_dir / "accounts" / acc_id
            session_file = session_dir / f"{acc_id}.session"
            account_checks[acc_id] = {
                "data_dir": _dir_check(acc_dir),
                "session_file": {
                    "ok": session_file.exists(),
                    "path": str(session_file),
                },
            }
        checks["accounts"] = {
            "ok": all(
                item.get("data_dir", {}).get("ok", False)
                for item in account_checks.values()
            ),
            "items": account_checks,
        }

    ok = all(item.get("ok", False) for item in checks.values())
    return {"ok": ok, "status": "ok" if ok else "degraded", "checks": checks}


def _dir_check(path: Path, writable: bool = False) -> Dict[str, Any]:
    path.mkdir(parents=True, exist_ok=True)
    ok = path.exists() and path.is_dir()
    payload: Dict[str, Any] = {"ok": ok, "path": str(path)}
    if writable:
        probe = path / ".healthcheck"
        try:
            probe.write_text("ok", encoding="utf-8")
            probe.unlink(missing_ok=True)
            payload["writable"] = True
        except OSError as exc:
            payload.update({"ok": False, "writable": False, "error": str(exc)})
    return payload


def _state_check(state_store: StateStore) -> Dict[str, Any]:
    try:
        state = state_store.load()
        data_dir = state_store.path.parent
        accounts = state.get("accounts") or []

        if accounts:
            # Multi-account mode: the global store no longer holds api
            # credentials / channels.  Aggregate those from each account's own
            # state file so the reported values are meaningful.
            has_api_credentials = False
            tracked_channels = 0
            for acc_id in accounts:
                acc = load_account(data_dir, acc_id)
                if acc.get("api_id") and acc.get("api_hash"):
                    has_api_credentials = True
                tracked_channels += len(acc.get("channels", {}) or {})
            return {
                "ok": True,
                "path": str(state_store.path),
                "has_api_credentials": has_api_credentials,
                "tracked_channels": tracked_channels,
            }

        # Legacy single-account semantics (no accounts list).
        return {
            "ok": True,
            "path": str(state_store.path),
            "has_api_credentials": bool(state.get("api_id") and state.get("api_hash")),
            "tracked_channels": len(state.get("channels", {})),
        }
    except Exception as exc:
        return {"ok": False, "path": str(state_store.path), "error": str(exc)}


def _sqlite_check() -> Dict[str, Any]:
    try:
        conn = sqlite3.connect(":memory:")
        conn.execute("SELECT 1")
        conn.close()
        return {"ok": True}
    except sqlite3.Error as exc:
        return {"ok": False, "error": str(exc)}


def _continuous_check(snapshot: Dict[str, Any]) -> Dict[str, Any]:
    running_accounts = snapshot.get("running_accounts")
    if running_accounts is None:
        running_accounts = snapshot
    return {
        "ok": True,
        "account_count": len(running_accounts),
        "accounts": {
            aid: {
                "running": bool((info.get("status") or info).get("running")),
                "enabled": bool((info.get("config") or {}).get("enabled", False)),
                "last_error": (info.get("status") or info).get("last_error"),
            }
            for aid, info in running_accounts.items()
            if isinstance(info, dict)
        },
    }