FabledSteward/steward/ansible/executor.py

# steward/ansible/executor.py
from __future__ import annotations
import asyncio
import logging
import time
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from quart import Quart

logger = logging.getLogger(__name__)

_OUTPUT_CAP_BYTES = 1024 * 1024  # 1 MB DB cap
_FLUSH_LINES = 50
_FLUSH_SECS = 5.0


@dataclass
class _Done:
    status: str


# Per-run broadcast state (lives until process restarts)
_run_lines: dict[str, list[str]] = {}          # all output lines received so far
_run_done: dict[str, _Done] = {}               # set when run completes
_run_listeners: dict[str, list[asyncio.Queue]] = {}  # active SSE client queues


def register_listener(run_id: str) -> asyncio.Queue:
    """Create a listener queue for a run. Replays existing lines immediately."""
    q: asyncio.Queue = asyncio.Queue()
    for line in _run_lines.get(run_id, []):
        q.put_nowait(line)
    done = _run_done.get(run_id)
    if done:
        q.put_nowait(done)
    else:
        _run_listeners.setdefault(run_id, []).append(q)
    return q


def deregister_listener(run_id: str, q: asyncio.Queue) -> None:
    listeners = _run_listeners.get(run_id, [])
    if q in listeners:
        listeners.remove(q)


def _broadcast(run_id: str, item: str | _Done) -> None:
    if isinstance(item, str):
        _run_lines.setdefault(run_id, []).append(item)
    else:
        _run_done[run_id] = item
    for q in _run_listeners.get(run_id, []):
        q.put_nowait(item)
    if isinstance(item, _Done):
        _run_listeners.pop(run_id, None)


def build_ansible_command(
    playbook_path: str,
    inventory_path: str,
    params: dict | None = None,
) -> list[str]:
    """Build the ansible-playbook argv from a run's optional params.

    Pure (no IO). Everything is passed as argv — never shell-interpolated — so
    extra-var values and limits can contain arbitrary characters safely.
    params keys: extra_vars (list of "key=value"), limit, tags, check (bool).
    """
    cmd = ["ansible-playbook", playbook_path, "-i", inventory_path]
    params = params or {}
    for ev in params.get("extra_vars") or []:
        cmd += ["-e", ev]
    if params.get("limit"):
        cmd += ["--limit", params["limit"]]
    if params.get("tags"):
        cmd += ["--tags", params["tags"]]
    if params.get("check"):
        cmd += ["--check", "--diff"]
    return cmd


async def start_run(
    app: "Quart",
    run_id: str,
    playbook_path: str,
    inventory_path: str,
    source_path: str,
    params: dict | None = None,
) -> None:
    """Execute ansible-playbook as a subprocess and update the DB run row."""
    _run_lines[run_id] = []

    db_output = ""
    truncated = False
    lines_since_flush = 0
    last_flush_at = time.monotonic()

    async def _flush(final_status: str | None = None) -> None:
        nonlocal lines_since_flush, last_flush_at
        from sqlalchemy import update
        from steward.models.ansible import AnsibleRun, AnsibleRunStatus

        values: dict = {"output": db_output}
        if final_status is not None:
            values["status"] = AnsibleRunStatus(final_status)
            values["finished_at"] = datetime.now(timezone.utc)

        async with app.db_sessionmaker() as session:
            async with session.begin():
                await session.execute(
                    update(AnsibleRun).where(AnsibleRun.id == run_id).values(**values)
                )
        lines_since_flush = 0
        last_flush_at = time.monotonic()

    cmd = build_ansible_command(playbook_path, inventory_path, params)
    cwd = source_path if Path(source_path).exists() else None

    try:
        proc = await asyncio.create_subprocess_exec(
            *cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.STDOUT,
            cwd=cwd,
        )
        assert proc.stdout is not None

        async for raw_line in proc.stdout:
            line = raw_line.decode(errors="replace").rstrip("\n\r")
            _broadcast(run_id, line)

            if not truncated:
                candidate = (db_output + "\n" + line) if db_output else line
                if len(candidate.encode()) > _OUTPUT_CAP_BYTES:
                    db_output += "\n[output truncated]"
                    truncated = True
                else:
                    db_output = candidate

            lines_since_flush += 1
            elapsed = time.monotonic() - last_flush_at
            if lines_since_flush >= _FLUSH_LINES or elapsed >= _FLUSH_SECS:
                await _flush()

        await proc.wait()
        final_status = "success" if proc.returncode == 0 else "failed"

    except Exception:
        logger.exception("Ansible run %s failed with exception", run_id)
        final_status = "failed"

    await _flush(final_status)
    _broadcast(run_id, _Done(status=final_status))