import json from datetime import datetime, timezone from typing import AsyncGenerator, Literal from fastapi import APIRouter, Depends, Query from fastapi.responses import StreamingResponse from app.auth import verify_token from app.ops_runner import _BACKUP_TIMEOUT, stream_ops_host router = APIRouter() def _sse_line(payload: dict) -> str: """Format a dict as a single SSE data line.""" return f"data: {json.dumps(payload)}\n\n" async def _restore_generator( project: str, env: str, source: str, dry_run: bool, name: str | None = None, mode: str = "full", ) -> AsyncGenerator[str, None]: """Async generator that drives the restore workflow and yields SSE events. Runs on the host via nsenter because ops restore delegates to project CLIs that use host Python venvs incompatible with the container's Python. """ base_args = ["restore", project, env] # Pass the backup file path to avoid interactive selection prompt if name: backup_path = f"/opt/data/backups/{project}/{env}/{name}" base_args.append(backup_path) if dry_run: base_args.append("--dry-run") # Granular restore mode if mode == "db": base_args.append("--db-only") elif mode == "wp": base_args.append("--wp-only") if source == "offsite": # ops offsite restore download_args = ["offsite", "restore", project, env] yield _sse_line({"line": f"Downloading {project}/{env} from offsite...", "timestamp": _now()}) download_ok = True async for line in stream_ops_host(download_args, timeout=_BACKUP_TIMEOUT): yield _sse_line({"line": line, "timestamp": _now()}) if line.startswith("[error]"): download_ok = False if not download_ok: yield _sse_line({"done": True, "success": False}) return yield _sse_line({"line": "Download complete. Starting restore...", "timestamp": _now()}) success = True async for line in stream_ops_host(base_args, timeout=_BACKUP_TIMEOUT): yield _sse_line({"line": line, "timestamp": _now()}) if line.startswith("[error]"): success = False yield _sse_line({"done": True, "success": success}) def _now() -> str: return datetime.now(timezone.utc).isoformat() @router.get("/{project}/{env}", summary="Restore a backup with real-time output") async def restore_backup( project: str, env: str, source: Literal["local", "offsite"] = Query(default="local"), dry_run: bool = Query(default=False, alias="dry_run"), name: str | None = Query(default=None), mode: Literal["full", "db", "wp"] = Query(default="full"), _: str = Depends(verify_token), ) -> StreamingResponse: """ Restore a backup for the given project/env. Uses Server-Sent Events (SSE) to stream real-time progress. Runs on the host via nsenter for Python venv compatibility. Modes: full (default), db (database only), wp (wp-content only). """ return StreamingResponse( _restore_generator(project, env, source, dry_run, name, mode), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "X-Accel-Buffering": "no", }, )