debian-manager/app/core/scheduler.py

"""Планировщик: запуск задач по cron-расписанию внутри приложения."""
from __future__ import annotations

import asyncio
from datetime import datetime, timezone
from typing import Optional

from croniter import croniter
from sqlmodel import select

from app.db.database import get_session
from app.db.models import ScheduledTask, Server

CHECK_INTERVAL = 30  # секунд между проверками

_next_run: dict[int, datetime] = {}
_loop_task: Optional[asyncio.Task] = None


# ---------- CRUD ----------

def list_tasks() -> list[ScheduledTask]:
    with get_session() as s:
        return list(s.exec(select(ScheduledTask).order_by(ScheduledTask.id)).all())


def add_task(**fields) -> ScheduledTask:
    with get_session() as s:
        t = ScheduledTask(**fields)
        s.add(t)
        s.commit()
        s.refresh(t)
        return t


def update_task(task_id: int, **fields) -> None:
    with get_session() as s:
        t = s.get(ScheduledTask, task_id)
        if t:
            for k, v in fields.items():
                setattr(t, k, v)
            s.add(t)
            s.commit()
    _next_run.pop(task_id, None)  # пересчитать расписание


def delete_task(task_id: int) -> None:
    with get_session() as s:
        t = s.get(ScheduledTask, task_id)
        if t:
            s.delete(t)
            s.commit()
    _next_run.pop(task_id, None)


# ---------- выполнение ----------

def _target_servers(target: str) -> list[Server]:
    with get_session() as s:
        servers = list(s.exec(select(Server)).all())
    if target == "all":
        return servers
    return [srv for srv in servers if target in (srv.tags or "").split(",")
            or target.strip() in [t.strip() for t in (srv.tags or "").split(",")]]


async def execute_task(task_id: int) -> str:
    """Выполнить задачу немедленно. Возвращает статус."""
    from app.core import servers as servers_core
    from app.core.group_ops import op_apt_update, op_upgrade, op_shell, run_group
    from app.core.tasks import tasks as tm

    with get_session() as s:
        t = s.get(ScheduledTask, task_id)
        if not t:
            return "error: задача не найдена"
        action, target, params, name = t.action, t.target, t.params, t.name
        t.last_run = datetime.now(timezone.utc)
        t.last_status = "running"
        s.add(t)
        s.commit()

    status = "success"
    log_lines: list[str] = []
    try:
        if action == "cleanup_metrics":
            from app.db.models import MetricSnapshot
            from datetime import timedelta
            cutoff = datetime.now(timezone.utc) - timedelta(days=7)
            with get_session() as s:
                old = s.exec(select(MetricSnapshot).where(MetricSnapshot.ts < cutoff)).all()
                for m in old:
                    s.delete(m)
                s.commit()
                log_lines.append(f"удалено {len(old)} снапшотов")
        elif action == "report":
            from app.core import reports
            path = await reports.generate()
            log_lines.append(f"отчёт: {path}")
        else:
            servers = _target_servers(target)
            if not servers:
                raise RuntimeError(f"нет серверов для target='{target}'")
            if action == "apt_update":
                op = op_apt_update()
            elif action == "upgrade":
                op = op_upgrade(False)
            elif action == "dist_upgrade":
                op = op_upgrade(True)
            elif action == "shell":
                op = op_shell(params)
            elif action == "backup":
                from app.core.backup import run_backup
                for srv in servers:
                    rec = await run_backup(srv, params or "/etc")
                    log_lines.append(f"{srv.name}: {rec.status}")
                op = None
            elif action == "playbook":
                from app.core import playbooks as pb
                from app.core.ssh import SSHClient
                data = pb.load(params)
                for srv in servers:
                    async with SSHClient(srv) as c:
                        await pb.run_playbook(c, srv, data, log_lines.append)
                op = None
            else:
                raise RuntimeError(f"неизвестное действие: {action}")

            if op is not None:
                res = await run_group(servers, op)
                log_lines.append(f"ok={res.ok_count} err={res.err_count}")
                for r in res.results.values():
                    if r.status == "error":
                        status = "error"
                        log_lines.append(f"{r.name}: {r.log[-300:]}")
    except Exception as e:
        status = "error"
        log_lines.append(f"[ОШИБКА] {e}")

    with get_session() as s:
        t = s.get(ScheduledTask, task_id)
        if t:
            t.last_status = status
            s.add(t)
            s.commit()
    # Telegram-уведомление
    try:
        from app.core import notify
        await notify.task_done(f"планировщик", name, status)
    except Exception:
        pass
    return status


# ---------- цикл ----------

async def _loop():
    while True:
        try:
            now = datetime.now(timezone.utc)
            for t in list_tasks():
                if not t.enabled:
                    continue
                nxt = _next_run.get(t.id)
                if nxt is None:
                    nxt = croniter(t.cron, now).get_next(datetime)
                    _next_run[t.id] = nxt
                if now >= nxt:
                    _next_run[t.id] = croniter(t.cron, now).get_next(datetime)
                    asyncio.ensure_future(execute_task(t.id))
        except Exception:
            pass
        await asyncio.sleep(CHECK_INTERVAL)


def start() -> None:
    global _loop_task
    if _loop_task is None or _loop_task.done():
        _loop_task = asyncio.ensure_future(_loop())


# ---------- пресеты ----------

PRESETS = [
    ("Ежедневный apt update (03:00)", "0 3 * * *", "apt_update", "all", ""),
    ("Еженедельный upgrade (вс 04:00)", "0 4 * * 0", "upgrade", "all", ""),
    ("Ежедневный бэкап /etc (02:00)", "0 2 * * *", "backup", "all", "/etc"),
    ("Бэкап /etc и /home (02:30)", "30 2 * * *", "backup", "all", "/etc /home"),
    ("Очистка старых метрик (ежечасно)", "7 * * * *", "cleanup_metrics", "all", ""),
    ("Ежедневный отчёт (08:00)", "0 8 * * *", "report", "all", ""),
    ("Docker prune (вс 05:00)", "0 5 * * 0", "shell", "all", "docker system prune -f"),
    ("Проверка диска (каждые 6 ч)", "13 */6 * * *", "shell", "all", "df -h"),
]