debian-manager/app/core/group_ops.py

"""Групповые операции по нескольким серверам с rolling-режимом."""
from __future__ import annotations

import asyncio
from dataclasses import dataclass, field
from typing import Callable, Coroutine, Dict, List, Optional

from app.core.ssh import SSHClient, check_server, wait_online
from app.db.models import Server

LogCallback = Callable[[str], None]


@dataclass
class ServerResult:
    server_id: int
    name: str
    status: str = "pending"   # pending/running/success/error
    log: str = ""


@dataclass
class GroupResult:
    results: Dict[int, ServerResult] = field(default_factory=dict)

    @property
    def ok_count(self) -> int:
        return sum(1 for r in self.results.values() if r.status == "success")

    @property
    def err_count(self) -> int:
        return sum(1 for r in self.results.values() if r.status == "error")


async def run_group(
    servers: List[Server],
    operation: Callable[[SSHClient, Server, LogCallback], Coroutine],
    concurrency: int = 5,
    fail_fast: bool = False,
    rolling: bool = False,
    on_status: Optional[Callable[[GroupResult], None]] = None,
) -> GroupResult:
    """Выполнить operation(c, server, log) на группе серверов."""
    result = GroupResult(results={
        s.id: ServerResult(server_id=s.id, name=s.name) for s in servers
    })
    sem = asyncio.Semaphore(1 if rolling else concurrency)
    stop_flag = asyncio.Event()

    async def worker(srv: Server):
        if stop_flag.is_set():
            result.results[srv.id].status = "error"
            result.results[srv.id].log = "пропущено (fail-fast)"
            return
        async with sem:
            r = result.results[srv.id]
            r.status = "running"
            if on_status:
                on_status(result)

            def log(line: str) -> None:
                r.log += line + "\n"
                if on_status:
                    on_status(result)

            try:
                async with SSHClient(srv) as c:
                    await operation(c, srv, log)
                r.status = "success"
            except Exception as e:
                r.status = "error"
                r.log += f"[ОШИБКА] {e}\n"
                if fail_fast:
                    stop_flag.set()
            if on_status:
                on_status(result)

    if rolling:
        for srv in servers:
            await worker(srv)
            if stop_flag.is_set():
                break
    else:
        await asyncio.gather(*(worker(s) for s in servers))
    return result


# ---------- готовые операции ----------

def op_apt_update():
    async def op(c: SSHClient, srv: Server, log: LogCallback):
        log("$ apt-get update")
        r = await c.sudo("DEBIAN_FRONTEND=noninteractive apt-get update", timeout=600)
        log(r.stdout[-2000:] or "(пустой вывод)")
        if not r.ok:
            raise RuntimeError(r.stderr[-500:])
    return op


def op_upgrade(dist: bool = False):
    async def op(c: SSHClient, srv: Server, log: LogCallback):
        kind = "dist-upgrade" if dist else "upgrade"
        log(f"$ apt-get {kind} -y")
        async for line in c.run_stream(
            f"DEBIAN_FRONTEND=noninteractive apt-get {kind} -y", use_sudo=True
        ):
            log(line)
    return op


def op_reboot(wait: bool = True):
    async def op(c: SSHClient, srv: Server, log: LogCallback):
        log("$ reboot")
        await c.sudo("(sleep 1; reboot) &", timeout=10)
        if wait:
            log("ожидание возврата сервера...")
            await asyncio.sleep(10)
            ok = await wait_online(srv, timeout_s=300)
            if not ok:
                raise RuntimeError("сервер не вернулся за 5 минут")
            log("сервер снова онлайн")
    return op


def op_install_package(names: List[str]):
    async def op(c: SSHClient, srv: Server, log: LogCallback):
        log(f"$ apt-get install -y {' '.join(names)}")
        async for line in c.run_stream(
            "DEBIAN_FRONTEND=noninteractive apt-get install -y " + " ".join(names),
            use_sudo=True,
        ):
            log(line)
    return op


def op_shell(cmd: str):
    async def op(c: SSHClient, srv: Server, log: LogCallback):
        log(f"$ {cmd}")
        r = await c.sudo(cmd, timeout=600)
        log(r.stdout[-3000:])
        if not r.ok:
            raise RuntimeError(r.stderr[-500:] or f"rc={r.returncode}")
    return op