Roundtable-plugins/ups/scheduler.py

# plugins/ups/scheduler.py
from __future__ import annotations
import asyncio
import logging
from datetime import datetime, timezone
from pathlib import Path
from typing import TYPE_CHECKING

from fabledscryer.core.scheduler import ScheduledTask

if TYPE_CHECKING:
    from quart import Quart

logger = logging.getLogger(__name__)

# In-memory shutdown state.
# Limitation: resets on app restart. If the app restarts while on battery,
# the countdown begins again from zero. Given typical UPS runtimes (minutes),
# this is acceptable. A future improvement could persist this in the DB.
_on_battery_since: datetime | None = None
_shutdown_triggered: bool = False


def make_poll_task(app: "Quart") -> ScheduledTask:
    cfg = app.config["PLUGINS"]["ups"]
    interval = int(cfg.get("poll_interval_seconds", 30))

    async def poll() -> None:
        await _do_poll(app)

    return ScheduledTask(
        name="ups_poll",
        coro_factory=poll,
        interval_seconds=interval,
        run_on_startup=True,
    )


async def _do_poll(app: "Quart") -> None:
    global _on_battery_since, _shutdown_triggered

    from .client import NutClient, NutError, parse_status
    from .models import UpsStatus
    from fabledscryer.core.alerts import record_metric

    cfg = app.config["PLUGINS"]["ups"]
    client = NutClient(
        host=cfg.get("nut_host", "localhost"),
        port=int(cfg.get("nut_port", 3493)),
        username=cfg.get("nut_username", ""),
        password=cfg.get("nut_password", ""),
    )

    try:
        raw_vars = await client.get_vars(cfg.get("ups_name", "ups"))
    except NutError as exc:
        logger.warning("UPS poll failed: %s", exc)
        return

    parsed = parse_status(raw_vars)
    scraped_at = datetime.now(timezone.utc)

    status = UpsStatus(
        scraped_at=scraped_at,
        raw_status=parsed["raw_status"],
        on_battery=parsed["on_battery"],
        low_battery=parsed["low_battery"],
        battery_charge_pct=parsed["battery_charge_pct"],
        battery_runtime_secs=parsed["battery_runtime_secs"],
        battery_voltage=parsed["battery_voltage"],
        load_pct=parsed["load_pct"],
        input_voltage=parsed["input_voltage"],
        output_voltage=parsed["output_voltage"],
        temperature=parsed["temperature"],
        model=parsed["model"],
        manufacturer=parsed["manufacturer"],
    )

    async with app.db_sessionmaker() as session:
        async with session.begin():
            session.add(status)

            if parsed["battery_charge_pct"] is not None:
                await record_metric(
                    session=session,
                    source_module="ups",
                    resource_name="ups",
                    metric_name="battery_charge_pct",
                    value=parsed["battery_charge_pct"],
                )
            await record_metric(
                session=session,
                source_module="ups",
                resource_name="ups",
                metric_name="on_battery",
                value=1.0 if parsed["on_battery"] else 0.0,
            )
            if parsed["battery_runtime_secs"] is not None:
                await record_metric(
                    session=session,
                    source_module="ups",
                    resource_name="ups",
                    metric_name="battery_runtime_secs",
                    value=float(parsed["battery_runtime_secs"]),
                )

    # ── Shutdown logic ────────────────────────────────────────────────────────
    shutdown_after = int(cfg.get("shutdown_after_seconds", 0))

    if parsed["on_battery"]:
        if _on_battery_since is None:
            _on_battery_since = scraped_at
            logger.warning("UPS switched to battery at %s", scraped_at.isoformat())

        elapsed = (scraped_at - _on_battery_since).total_seconds()
        logger.debug("UPS on battery for %.0f seconds (shutdown_after=%d)", elapsed, shutdown_after)

        if (
            not _shutdown_triggered
            and shutdown_after > 0
            and elapsed >= shutdown_after
        ):
            logger.warning(
                "UPS on battery for %.0fs — triggering shutdown playbook", elapsed
            )
            _shutdown_triggered = True
            asyncio.create_task(_run_shutdown(app, cfg))
    else:
        if _on_battery_since is not None:
            logger.info("UPS back on AC power after %.0fs on battery",
                        (scraped_at - _on_battery_since).total_seconds())
        _on_battery_since = None
        _shutdown_triggered = False


async def _run_shutdown(app: "Quart", cfg: dict) -> None:
    """Locate the configured Ansible playbook and run it via subprocess."""
    source_name = cfg.get("shutdown_source", "")
    playbook_rel = cfg.get("shutdown_playbook", "")
    inventory_name = cfg.get("shutdown_inventory", "hosts")

    if not source_name or not playbook_rel:
        logger.error(
            "UPS shutdown triggered but shutdown_source or shutdown_playbook not configured"
        )
        return

    # Resolve source path from Ansible config
    ansible_cfg = app.config.get("ANSIBLE", {})
    sources = ansible_cfg.get("sources", [])
    source = next((s for s in sources if s.get("name") == source_name), None)
    if not source:
        logger.error("UPS shutdown: Ansible source %r not found", source_name)
        return

    source_path = Path(source["path"])
    playbook_path = source_path / playbook_rel
    inventory_path = source_path / inventory_name

    if not playbook_path.exists():
        logger.error("UPS shutdown: playbook not found at %s", playbook_path)
        return
    if not inventory_path.exists():
        logger.error("UPS shutdown: inventory not found at %s", inventory_path)
        return

    logger.warning("UPS shutdown: running %s with inventory %s", playbook_path, inventory_path)

    try:
        proc = await asyncio.create_subprocess_exec(
            "ansible-playbook", str(playbook_path), "-i", str(inventory_path),
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.STDOUT,
            cwd=str(source_path),
        )
        assert proc.stdout is not None
        async for line in proc.stdout:
            logger.info("[ups-shutdown] %s", line.decode(errors="replace").rstrip())
        await proc.wait()
        if proc.returncode == 0:
            logger.warning("UPS shutdown playbook completed successfully")
        else:
            logger.error("UPS shutdown playbook exited with code %d", proc.returncode)
    except Exception:
        logger.exception("UPS shutdown playbook execution failed")