"""Групповые операции по нескольким серверам с 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