pobsync/src/pobsync_backend/management/commands/run_pobsync_scheduler.py

from __future__ import annotations

import time
from pathlib import Path
from typing import Any

from django.conf import settings
from django.core.management import call_command
from django.core.management.base import BaseCommand
from django.db import transaction
from django.utils import timezone

from pobsync_backend.models import BackupRun, ScheduleConfig
from pobsync_backend.scheduler import due_key, is_due


class Command(BaseCommand):
    help = "Run due pobsync schedules from the Django database."

    def add_arguments(self, parser) -> None:
        parser.add_argument("--prefix", default=settings.POBSYNC_HOME, help="Runtime state root")
        parser.add_argument("--once", action="store_true", help="Check once and exit")
        parser.add_argument("--loop", action="store_true", help="Keep checking schedules")
        parser.add_argument("--interval", type=int, default=60, help="Loop interval in seconds")
        parser.add_argument("--dry-run", action="store_true", help="Pass --dry-run to backup runs")

    def handle(self, *args: Any, **options: Any) -> None:
        if not options["once"] and not options["loop"]:
            options["once"] = True

        prefix = Path(options["prefix"])
        while True:
            count = self._run_due(prefix=prefix, dry_run=bool(options["dry_run"]))
            self.stdout.write(f"Ran {count} due schedule(s).")
            if options["once"]:
                return
            time.sleep(max(1, int(options["interval"])))

    def _run_due(self, *, prefix: Path, dry_run: bool) -> int:
        now = timezone.localtime(timezone.now())
        current_due_key = due_key(now)
        ran = 0

        schedules = (
            ScheduleConfig.objects.select_related("host")
            .filter(enabled=True, host__enabled=True)
            .order_by("host__host")
        )
        for schedule in schedules:
            if schedule.last_due_key == current_due_key:
                continue
            if not is_due(schedule.cron_expr, now):
                continue

            schedule_started_at = timezone.now()
            with transaction.atomic():
                locked = ScheduleConfig.objects.select_for_update().get(pk=schedule.pk)
                if locked.last_due_key == current_due_key:
                    continue
                locked.last_due_key = current_due_key
                locked.last_started_at = schedule_started_at
                locked.last_status = "running"
                locked.save(update_fields=["last_due_key", "last_started_at", "last_status", "updated_at"])

            status = "success"
            try:
                call_command(
                    "run_pobsync_backup",
                    schedule.host.host,
                    prefix=str(prefix),
                    dry_run=dry_run,
                    prune=schedule.prune,
                    prune_max_delete=schedule.prune_max_delete,
                    prune_protect_bases=schedule.prune_protect_bases,
                )
                status = _latest_scheduled_run_status(host_id=schedule.host_id, started_at=schedule_started_at) or status
            except Exception as exc:
                status = "failed"
                self.stderr.write(f"{schedule.host.host}: {type(exc).__name__}: {exc}")
            finally:
                ScheduleConfig.objects.filter(pk=schedule.pk).update(
                    last_finished_at=timezone.now(),
                    last_status=status,
                )
            ran += 1

        return ran


def _latest_scheduled_run_status(*, host_id: int, started_at) -> str | None:
    run = (
        BackupRun.objects.filter(
            host_id=host_id,
            run_type=BackupRun.RunType.SCHEDULED,
            created_at__gte=started_at,
        )
        .order_by("-created_at", "-id")
        .first()
    )
    return run.status if run is not None else None