#!/usr/bin/env python from __future__ import annotations import argparse import json import os import shlex import shutil import subprocess import sys import tempfile from collections.abc import Callable, Iterable from contextlib import contextmanager from pathlib import Path from types import SimpleNamespace from typing import Any from urllib.error import URLError META_ROOT = Path(__file__).resolve().parents[2] ROOT = Path(os.environ.get("GOVOPLAN_CORE_ROOT", META_ROOT.parent / "govoplan-core")).resolve() SRC = ROOT / "src" if str(SRC) not in sys.path: sys.path.insert(0, str(SRC)) os.environ.setdefault("APP_ENV", "test") os.environ.setdefault("DEV_BOOTSTRAP_ENABLED", "false") os.environ.setdefault("CELERY_ENABLED", "false") from govoplan_core.admin.models import SystemSettings from govoplan_core.core.maintenance import MaintenanceMode, save_maintenance_mode from govoplan_core.core.module_installer import ( cancel_module_installer_request, claim_next_module_installer_request, module_installer_daemon_status, module_installer_lock_status, queue_module_installer_request, read_module_installer_run, retry_module_installer_request, rollback_module_install_run, run_module_install_plan, supervise_module_install_plan, update_module_installer_daemon_status, update_module_installer_request, ) from govoplan_core.core.module_management import ( ModuleInstallPlan, ModuleInstallPlanItem, saved_desired_enabled_modules, saved_module_install_plan, save_desired_enabled_modules, save_module_install_plan, ) from govoplan_core.core.modules import MigrationRetirementPlan, MigrationSpec, ModuleManifest from govoplan_core.db.base import Base from govoplan_core.db.session import configure_database, get_database from govoplan_core.server.registry import available_module_manifests Scenario = Callable[[Path], dict[str, object]] def main() -> int: parser = argparse.ArgumentParser(description="Run safe module-installer rollback drills in temporary runtimes.") parser.add_argument("--scenario", action="append", choices=sorted(SCENARIOS), help="Scenario to run. Defaults to all scenarios.") parser.add_argument("--runtime-root", type=Path, help="Directory for temporary drill runtimes. Defaults to a new temp directory.") parser.add_argument("--keep-runtime", action="store_true", help="Keep the temporary drill runtime after completion.") parser.add_argument("--format", choices=("json", "text"), default="text", help="Output format.") args = parser.parse_args() runtime_root = args.runtime_root or Path(tempfile.mkdtemp(prefix="govoplan-installer-drill-")) runtime_root.mkdir(parents=True, exist_ok=True) selected = args.scenario or sorted(SCENARIOS) results: list[dict[str, object]] = [] ok = True try: for name in selected: scenario_root = runtime_root / name scenario_root.mkdir(parents=True, exist_ok=True) try: result = SCENARIOS[name](scenario_root) except Exception as exc: ok = False result = {"scenario": name, "ok": False, "error": str(exc), "root": str(scenario_root)} else: ok = ok and bool(result.get("ok")) results.append(result) finally: if not args.keep_runtime and args.runtime_root is None: shutil.rmtree(runtime_root, ignore_errors=True) payload = {"ok": ok, "runtime_root": str(runtime_root), "scenarios": results} if args.format == "json": print(json.dumps(payload, indent=2, sort_keys=True)) else: for result in results: status = "ok" if result.get("ok") else "FAILED" print(f"{status}: {result['scenario']}") if result.get("run_id"): print(f" run: {result['run_id']}") if result.get("record_path"): print(f" record: {result['record_path']}") if result.get("error"): print(f" error: {result['error']}") print(f"runtime: {runtime_root}") return 0 if ok else 1 def package_failure(root: Path) -> dict[str, object]: database_url, database = _sqlite_database(root) plan = _saved_install_plan(database, "package-failure-example") with database.session() as session, _patch_subprocess(_package_failure_run): result = supervise_module_install_plan( session=session, plan=plan, available=available_module_manifests(), current_enabled=("tenancy", "access"), desired_enabled=("tenancy", "access"), database_url=database_url, runtime_dir=root / "installer", ) restored_plan = saved_module_install_plan(session) record = _run_record(root, result.run_id) return _scenario_result( "package-failure", root, result, expected_status="rolled-back", checks={ "supervisor_status": _nested(record, "supervisor", "status") == "rolled-back", "plan_restored": tuple(item.status for item in restored_plan.items) == ("planned",), }, ) def migration_failure(root: Path) -> dict[str, object]: database_url, database = _sqlite_database(root) plan = _saved_install_plan(database, "migration-failure-example") with database.session() as session, _patch_subprocess(_migration_failure_run): result = supervise_module_install_plan( session=session, plan=plan, available=available_module_manifests(), current_enabled=("tenancy", "access"), desired_enabled=("tenancy", "access"), database_url=database_url, runtime_dir=root / "installer", migrate_database=True, ) record = _run_record(root, result.run_id) return _scenario_result( "migration-failure", root, result, expected_status="rolled-back", checks={ "sqlite_backup_recorded": isinstance(_nested(record, "snapshot", "database_backup"), dict), "supervisor_status": _nested(record, "supervisor", "status") == "rolled-back", }, ) def restart_failure(root: Path) -> dict[str, object]: database_url, database = _sqlite_database(root) plan = _saved_install_plan(database, "restart-failure-example") with database.session() as session, _patch_subprocess(_restart_failure_run): result = supervise_module_install_plan( session=session, plan=plan, available=available_module_manifests(), current_enabled=("tenancy", "access"), desired_enabled=("tenancy", "access"), database_url=database_url, runtime_dir=root / "installer", restart_command="restart-fails", ) record = _run_record(root, result.run_id) return _scenario_result( "restart-failure", root, result, expected_status="rolled-back", checks={"supervisor_status": _nested(record, "supervisor", "status") == "rolled-back"}, ) def health_timeout(root: Path) -> dict[str, object]: database_url, database = _sqlite_database(root) plan = _saved_install_plan(database, "health-timeout-example") with database.session() as session, _patch_subprocess(_success_run), _patch_urlopen(): result = supervise_module_install_plan( session=session, plan=plan, available=available_module_manifests(), current_enabled=("tenancy", "access"), desired_enabled=("tenancy", "access"), database_url=database_url, runtime_dir=root / "installer", restart_command="restart-ok", health_url="http://127.0.0.1:9/health", health_timeout_seconds=0.2, health_interval_seconds=0.05, ) record = _run_record(root, result.run_id) return _scenario_result( "health-timeout", root, result, expected_status="rolled-back", checks={ "health_failure_recorded": "127.0.0.1:9/health" in str(_nested(record, "supervisor", "failure_reason")), "supervisor_status": _nested(record, "supervisor", "status") == "rolled-back", }, ) def destructive_retirement_failure(root: Path) -> dict[str, object]: database_url, database = _sqlite_database(root) available = {"retirement-example": _retirement_failure_manifest()} with database.session() as session: save_maintenance_mode(session, MaintenanceMode(enabled=True)) plan = save_module_install_plan(session, [{ "module_id": "retirement-example", "action": "uninstall", "python_package": "govoplan-retirement-example", "destroy_data": True, }]) session.commit() with _patch_subprocess(_success_run): result = run_module_install_plan( session=session, plan=plan, available=available, current_enabled=("tenancy", "access"), desired_enabled=("tenancy", "access"), database_url=database_url, runtime_dir=root / "installer", ) record = _run_record(root, result.run_id) return _scenario_result( "destructive-retirement-failure", root, result, expected_status="rolled-back", checks={ "rollback_recorded": isinstance(record.get("destructive_retirement_rollback"), dict), "sqlite_backup_recorded": isinstance(_nested(record, "snapshot", "database_backup"), dict), }, ) def postgres_hooks(root: Path) -> dict[str, object]: _, database = _sqlite_database(root) backup_script = ( "import json, os, pathlib; " "pathlib.Path(os.environ['GOVOPLAN_DATABASE_BACKUP_PATH']).write_text('backup', encoding='utf-8'); " "pathlib.Path(os.environ['GOVOPLAN_DATABASE_BACKUP_METADATA']).write_text(json.dumps({'snapshot': 'external'}), encoding='utf-8')" ) check_script = ( "import os, pathlib; " "pathlib.Path(os.environ['GOVOPLAN_INSTALLER_RUN_DIR'], 'restore-check.marker').write_text(" "pathlib.Path(os.environ['GOVOPLAN_DATABASE_BACKUP_PATH']).read_text(encoding='utf-8'), encoding='utf-8')" ) restore_script = ( "import os, pathlib; " "pathlib.Path(os.environ['GOVOPLAN_INSTALLER_RUN_DIR'], 'restore.marker').write_text(" "os.environ.get('GOVOPLAN_DATABASE_URL', ''), encoding='utf-8')" ) with database.session() as session: save_maintenance_mode(session, MaintenanceMode(enabled=True)) plan = save_module_install_plan(session, [{ "module_id": "postgres-hook-example", "action": "install", "python_package": "govoplan-postgres-hook-example", "python_ref": "govoplan-postgres-hook-example==0.1.0", }]) session.commit() result = run_module_install_plan( session=session, plan=plan, available=available_module_manifests(), current_enabled=("tenancy", "access"), desired_enabled=("tenancy", "access"), database_url="postgresql://db.example.invalid/govoplan", runtime_dir=root / "installer", migrate_database=True, database_backup_command=f"{sys.executable} -c {shlex.quote(backup_script)}", database_restore_command=f"{sys.executable} -c {shlex.quote(restore_script)}", database_restore_check_command=f"{sys.executable} -c {shlex.quote(check_script)}", dry_run=True, ) with _patch_subprocess(_non_shell_success_run): rollback = rollback_module_install_run( run_id=result.run_id, runtime_dir=root / "installer", database_url="postgresql://db.example.invalid/govoplan", ) run_dir = result.record_path.parent record = _run_record(root, result.run_id) ok = ( result.status == "dry-run" and rollback.status == "rolled-back" and (run_dir / "restore-check.marker").read_text(encoding="utf-8") == "backup" and (run_dir / "restore.marker").read_text(encoding="utf-8") == "postgresql://db.example.invalid/govoplan" and isinstance(_nested(record, "snapshot", "database_backup"), dict) ) return { "scenario": "postgres-hooks", "ok": ok, "root": str(root), "run_id": result.run_id, "record_path": str(result.record_path), "rollback_status": rollback.status, } def queue_and_lock(root: Path) -> dict[str, object]: runtime_dir = root / "installer" status = update_module_installer_daemon_status( runtime_dir=runtime_dir, patch={"status": "polling", "current_request_id": None}, ) request = queue_module_installer_request( runtime_dir=runtime_dir, requested_by="drill", options={"migrate_database": True, "health_urls": ["http://127.0.0.1:8000/health"]}, ) claimed = claim_next_module_installer_request(runtime_dir=runtime_dir) assert claimed is not None failed = update_module_installer_request( runtime_dir=runtime_dir, request_id=str(claimed["request_id"]), patch={"status": "failed", "error": "simulated drill failure"}, ) retry = retry_module_installer_request(runtime_dir=runtime_dir, request_id=str(failed["request_id"]), requested_by="drill") cancelled = cancel_module_installer_request(runtime_dir=runtime_dir, request_id=str(retry["request_id"]), cancelled_by="drill") lock_path = runtime_dir / "install.lock" lock_path.write_text(json.dumps({"pid": 999999, "created_at": "drill"}), encoding="utf-8") lock = module_installer_lock_status(runtime_dir=runtime_dir) lock_path.unlink() return { "scenario": "queue-and-lock", "ok": bool(status["running"]) and request["status"] == "queued" and failed["status"] == "failed" and cancelled["status"] == "cancelled" and lock["locked"] is True and not module_installer_lock_status(runtime_dir=runtime_dir)["locked"], "root": str(root), "daemon_status_path": str(runtime_dir / "daemon.status.json"), } def _sqlite_database(root: Path) -> tuple[str, Any]: database_url = f"sqlite:///{root / 'drill.db'}" configure_database(database_url) database = get_database() Base.metadata.create_all(bind=database.engine, tables=[SystemSettings.__table__]) return database_url, database def _saved_install_plan(database: Any, module_id: str) -> ModuleInstallPlan: with database.session() as session: save_maintenance_mode(session, MaintenanceMode(enabled=True)) plan = save_module_install_plan(session, [{ "module_id": module_id, "action": "install", "python_package": f"govoplan-{module_id}", "python_ref": f"govoplan-{module_id}==0.1.0", }]) save_desired_enabled_modules(session, ("tenancy", "access")) session.commit() return plan def _retirement_failure_manifest() -> ModuleManifest: def destroy(_session: object, _module_id: str) -> None: raise RuntimeError("simulated destructive retirement failure") def provider(_session: object | None, _module_id: str) -> MigrationRetirementPlan: return MigrationRetirementPlan( supported=True, summary="Drill retirement provider supports destructive cleanup.", destroy_data_supported=True, destroy_data_summary="Drill destructive cleanup will fail during execution.", destroy_data_executor=destroy, ) return ModuleManifest( id="retirement-example", name="Retirement example", version="0.1.0", migration_spec=MigrationSpec( module_id="retirement-example", retirement_supported=True, retirement_provider=provider, ), ) def _scenario_result( scenario: str, root: Path, result: Any, *, expected_status: str, checks: dict[str, bool], ) -> dict[str, object]: return { "scenario": scenario, "ok": result.status == expected_status and all(checks.values()), "root": str(root), "run_id": result.run_id, "record_path": str(result.record_path), "status": result.status, "checks": checks, } def _run_record(root: Path, run_id: str) -> dict[str, object]: return read_module_installer_run(runtime_dir=root / "installer", run_id=run_id) def _nested(value: dict[str, object], *keys: str) -> object: current: object = value for key in keys: if not isinstance(current, dict): return None current = current.get(key) return current def _command_text(argv: object) -> str: if isinstance(argv, list | tuple): return " ".join(str(item) for item in argv) return str(argv) def _completed(return_code: int = 0, stdout: str = "", stderr: str = "") -> SimpleNamespace: return SimpleNamespace(returncode=return_code, stdout=stdout, stderr=stderr) def _success_run(argv: object, *_args: object, **_kwargs: object) -> SimpleNamespace: text = _command_text(argv) if "pip freeze" in text: return _completed(stdout="govoplan-core==0.0.0\n") return _completed() def _non_shell_success_run(argv: object, *_args: object, **kwargs: object) -> Any: if kwargs.get("shell"): return ORIGINAL_SUBPROCESS_RUN(argv, *_args, **kwargs) return _success_run(argv, *_args, **kwargs) def _package_failure_run(argv: object, *_args: object, **kwargs: object) -> SimpleNamespace: text = _command_text(argv) if "pip install" in text and "==" in text and "-r" not in text: return _completed(23, stderr="simulated package install failure") return _success_run(argv, *_args, **kwargs) def _migration_failure_run(argv: object, *_args: object, **kwargs: object) -> SimpleNamespace: text = _command_text(argv) if "govoplan_core.commands.init_db" in text: return _completed(24, stderr="simulated migration failure") return _success_run(argv, *_args, **kwargs) def _restart_failure_run(argv: object, *_args: object, **kwargs: object) -> SimpleNamespace: if kwargs.get("shell") and str(argv).strip() == "restart-fails": return _completed(25, stderr="simulated restart failure") return _success_run(argv, *_args, **kwargs) ORIGINAL_SUBPROCESS_RUN = subprocess.run @contextmanager def _patch_subprocess(fake_run: Callable[..., Any]) -> Iterable[None]: import govoplan_core.core.module_installer as installer original = installer.subprocess.run installer.subprocess.run = fake_run # type: ignore[assignment] try: yield finally: installer.subprocess.run = original # type: ignore[assignment] @contextmanager def _patch_urlopen() -> Iterable[None]: import govoplan_core.core.module_installer as installer original = installer.urllib.request.urlopen def fail(*_args: object, **_kwargs: object) -> None: raise URLError("simulated health timeout") installer.urllib.request.urlopen = fail # type: ignore[assignment] try: yield finally: installer.urllib.request.urlopen = original # type: ignore[assignment] SCENARIOS: dict[str, Scenario] = { "destructive-retirement-failure": destructive_retirement_failure, "health-timeout": health_timeout, "migration-failure": migration_failure, "package-failure": package_failure, "postgres-hooks": postgres_hooks, "queue-and-lock": queue_and_lock, "restart-failure": restart_failure, } if __name__ == "__main__": raise SystemExit(main())