monky-deployd/monky_deployd/compose.py

"""`docker compose` + a few `docker` facts. Everything shells out with a timeout and captured
output; nothing here ever sees a secret except the `.env` file compose reads by itself."""

from __future__ import annotations

import json
import logging
import os
import shutil
import subprocess
import time
from dataclasses import dataclass
from pathlib import Path

log = logging.getLogger("monky-deployd.compose")


class ComposeError(Exception):
    def __init__(self, what: str, rc: int, output: str):
        super().__init__(f"{what} failed (rc={rc})")
        self.what, self.rc, self.output = what, rc, output


@dataclass
class Container:
    name: str
    state: str
    health: str
    exit_code: int | None = None

    def as_dict(self) -> dict:
        return {"name": self.name, "state": self.state, "health": self.health}

    @property
    def ok(self) -> bool:
        if self.state == "running":
            return self.health in ("", "healthy", "none")
        if self.state == "exited":
            return (self.exit_code or 0) == 0
        return False


class Docker:
    def __init__(self, docker_bin: str = "docker", timeout_s: int = 600):
        self.bin = docker_bin
        self.timeout_s = timeout_s

    def available(self) -> bool:
        return shutil.which(self.bin) is not None

    def run(self, args: list[str], *, cwd: str | None = None, timeout: int | None = None, check: bool = True) -> str:
        cmd = [self.bin, *args]
        what = _what(cmd)
        log.debug("exec: %s", " ".join(cmd))
        try:
            p = subprocess.run(
                cmd,
                cwd=cwd,
                capture_output=True,
                text=True,
                timeout=timeout or self.timeout_s,
                env={**os.environ, "COMPOSE_INTERACTIVE_NO_CLI": "1"},
            )
        except FileNotFoundError as exc:
            raise ComposeError(what, 127, f"{self.bin} not found") from exc
        except subprocess.TimeoutExpired as exc:
            raise ComposeError(what, 124, f"timeout after {exc.timeout}s") from exc
        out = (p.stdout or "") + (("\n" + p.stderr) if p.stderr else "")
        if check and p.returncode != 0:
            raise ComposeError(what, p.returncode, out.strip())
        return out

    # -- facts ------------------------------------------------------------------------------------
    def version(self) -> str | None:
        try:
            return self.run(["version", "--format", "{{.Server.Version}}"], timeout=20).strip() or None
        except ComposeError:
            return None

    def compose_version(self) -> str | None:
        try:
            return self.run(["compose", "version", "--short"], timeout=20).strip() or None
        except ComposeError:
            return None

    def data_root(self) -> str:
        try:
            root = self.run(["info", "--format", "{{.DockerRootDir}}"], timeout=20).strip()
        except ComposeError:
            root = ""
        return root or "/var/lib/docker"

    def free_bytes(self, path: str | None = None) -> int | None:
        p = path or self.data_root()
        while p and not os.path.exists(p):
            p = os.path.dirname(p)
        try:
            st = os.statvfs(p or "/")
        except OSError:
            return None
        return st.f_bavail * st.f_frsize

    def image_prune(self) -> None:
        try:
            self.run(["image", "prune", "-f"], timeout=300)
        except ComposeError as exc:
            log.warning("image prune failed: %s", exc)


class Compose:
    def __init__(self, docker: Docker, project_dir: Path, project_name: str):
        self.docker = docker
        self.project_dir = Path(project_dir)
        self.project_name = project_name

    def _base(self) -> list[str]:
        return [
            "compose",
            "--project-name",
            self.project_name,
            "--project-directory",
            str(self.project_dir),
        ]

    def _with_files(self) -> list[str]:
        args = self._base()
        env_file = self.project_dir / ".env"
        if env_file.exists():
            args += ["--env-file", str(env_file)]
        return args

    def config_check(self) -> str:
        return self.docker.run([*self._with_files(), "config", "--quiet"], cwd=str(self.project_dir), timeout=120)

    def pull(self) -> str:
        return self.docker.run([*self._with_files(), "pull", "--quiet"], cwd=str(self.project_dir))

    def up(self) -> str:
        return self.docker.run(
            [*self._with_files(), "up", "-d", "--remove-orphans", "--quiet-pull"], cwd=str(self.project_dir)
        )

    def down(self, *, purge_volumes: bool) -> str:
        args = [*self._base(), "down", "--remove-orphans"]
        if purge_volumes:
            args.append("-v")
        return self.docker.run(args, cwd=str(self.project_dir) if self.project_dir.exists() else None, timeout=600)

    def ps(self) -> list[Container]:
        if not self.project_dir.exists():
            return []
        try:
            out = self.docker.run(
                [*self._base(), "ps", "-a", "--format", "json"], cwd=str(self.project_dir), timeout=60
            )
        except ComposeError as exc:
            log.warning("compose ps failed: %s", exc)
            return []
        return parse_ps(out)

    def logs_tail(self, lines: int = 80) -> str:
        try:
            return self.docker.run(
                [*self._base(), "logs", "--no-color", "--tail", str(lines)], cwd=str(self.project_dir), timeout=60
            )
        except ComposeError as exc:
            return exc.output

    def wait_healthy(self, timeout_s: int, poll_s: float = 5.0) -> tuple[bool, list[Container]]:
        deadline = time.monotonic() + timeout_s
        containers: list[Container] = []
        while True:
            containers = self.ps()
            if containers and all(c.ok for c in containers) and not any(c.health == "starting" for c in containers):
                return True, containers
            if any(c.state == "exited" and (c.exit_code or 0) != 0 for c in containers):
                return False, containers
            if time.monotonic() >= deadline:
                return False, containers
            time.sleep(poll_s)


def _what(cmd: list[str]) -> str:
    """`docker compose <sub>` / `docker <sub>` for error messages (compose global flags skipped)."""
    if len(cmd) > 1 and cmd[1] == "compose":
        rest = cmd[2:]
        while rest and rest[0].startswith("--"):
            rest = rest[2:]
        return "docker compose " + (rest[0] if rest else "")
    return " ".join(cmd[:2])


def parse_ps(out: str) -> list[Container]:
    """`compose ps --format json` is NDJSON on compose >= 2.21 and a JSON array before."""
    text = out.strip()
    if not text:
        return []
    rows: list[dict] = []
    if text.startswith("["):
        try:
            rows = json.loads(text)
        except json.JSONDecodeError:
            rows = []
    else:
        for line in text.splitlines():
            line = line.strip()
            if line.startswith("{"):
                try:
                    rows.append(json.loads(line))
                except json.JSONDecodeError:
                    continue
    out_list = []
    for r in rows:
        state = str(r.get("State") or r.get("state") or "").lower()
        health = str(r.get("Health") or r.get("health") or "").lower()
        code = r.get("ExitCode", r.get("exit_code"))
        try:
            code = int(code) if code is not None else None
        except (TypeError, ValueError):
            code = None
        out_list.append(
            Container(name=str(r.get("Name") or r.get("name") or "?"), state=state, health=health, exit_code=code)
        )
    return out_list