ModelForge/scripts/bootstrap.py

"""Bring a ModelForge installation from an empty database to a serving control plane.

    python scripts/bootstrap.py --database-url postgresql+psycopg://...

The whole point is that this is safe to run twice. Every step either creates what is missing or
confirms what is already there, and the report says which of the two happened, so an operator can
re-run it after a failure without wondering what state they are in.

It never invents an operator credential. Secrets are generated by the operator and supplied through
configuration; bootstrap verifies one is present and refuses to continue in production without it,
because a platform that mints its own admin secret has no way to tell you it did.
"""

from __future__ import annotations

import argparse
import json
import sys
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from alembic import command
from alembic.config import Config
from sqlalchemy import create_engine, inspect, text
from sqlalchemy.orm import Session

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "backend" / "src"))

from modelforge_api.domain.release import (  # noqa: E402
    PRODUCT_NAME,
    PRODUCT_VERSION,
    TARGET_SCHEMA_REVISION,
    Compatibility,
    schema_compatibility,
)
from modelforge_api.services.lifecycle import LifecycleService  # noqa: E402
from modelforge_api.services.manifest_registry import ManifestRegistry  # noqa: E402
from modelforge_api.services.migration_engine import MigrationEngineService  # noqa: E402
from modelforge_api.services.observability import ObservabilityService  # noqa: E402
from modelforge_api.services.project_registry import sync_project_registry  # noqa: E402
from modelforge_api.services.recovery import RecoveryService  # noqa: E402
from modelforge_api.services.registry import seed_candidate_registry  # noqa: E402
from modelforge_api.settings import Settings  # noqa: E402


@dataclass
class Step:
    name: str
    outcome: str
    detail: str
    seconds: float = 0.0


@dataclass
class BootstrapReport:
    steps: list[Step] = field(default_factory=list)
    started_at: float = field(default_factory=time.time)

    def record(self, name: str, outcome: str, detail: str, seconds: float = 0.0) -> None:
        self.steps.append(Step(name, outcome, detail, round(seconds, 3)))

    def as_dict(self) -> dict[str, Any]:
        return {
            "product": PRODUCT_NAME,
            "version": PRODUCT_VERSION,
            "total_seconds": round(time.time() - self.started_at, 3),
            "steps": [
                {
                    "name": step.name,
                    "outcome": step.outcome,
                    "detail": step.detail,
                    "seconds": step.seconds,
                }
                for step in self.steps
            ],
        }


def wait_for_database(url: str, report: BootstrapReport, timeout: float = 120.0) -> None:
    started = time.time()
    engine = create_engine(url, pool_pre_ping=True)
    last: str = ""
    while time.time() - started < timeout:
        try:
            with engine.connect() as connection:
                connection.exec_driver_sql("select 1")
            report.record(
                "database reachable", "OK", "accepted a connection", time.time() - started
            )
            engine.dispose()
            return
        except Exception as error:  # noqa: BLE001 - any failure means not ready yet
            last = type(error).__name__
            time.sleep(2)
    engine.dispose()
    raise SystemExit(f"the database never became reachable within {timeout:.0f}s (last: {last})")


def migrate(url: str, report: BootstrapReport) -> str:
    started = time.time()
    engine = create_engine(url)
    with engine.connect() as connection:
        had_schema = inspect(connection).has_table("alembic_version")
        before = (
            connection.execute(text("select version_num from alembic_version")).scalar_one_or_none()
            if had_schema
            else None
        )
    config = Config(str(ROOT / "backend" / "alembic.ini"))
    config.set_main_option("script_location", str(ROOT / "backend" / "alembic"))
    config.set_main_option("sqlalchemy.url", url)
    command.upgrade(config, "head")
    with engine.connect() as connection:
        after = connection.execute(text("select version_num from alembic_version")).scalar_one()
    engine.dispose()
    outcome = "ALREADY_CURRENT" if before == after else ("CREATED" if before is None else "UPGRADED")
    report.record("migrations", outcome, f"{before or '(empty)'} -> {after}", time.time() - started)
    if after != TARGET_SCHEMA_REVISION:
        raise SystemExit(
            f"migrations landed on {after}, but this release targets {TARGET_SCHEMA_REVISION}"
        )
    return str(after)


def seed(url: str, settings: Settings, report: BootstrapReport) -> None:
    """Every seed is an ensure_*: running it twice must not produce a second copy of anything."""

    engine = create_engine(url)
    manifests = ManifestRegistry(settings.config_root)
    with Session(engine) as session:
        started = time.time()
        candidates = seed_candidate_registry(session, manifests)
        session.commit()
        report.record(
            "candidate registry",
            "SEEDED" if candidates else "ALREADY_PRESENT",
            f"{candidates} candidate(s) added",
            time.time() - started,
        )

        started = time.time()
        sync = sync_project_registry(session, manifests)
        session.commit()
        detail = f"{len(sync.unavailable_contracts)} binding(s) waiting for a contract"
        report.record("project registry", "SYNCED", detail, time.time() - started)

        for name, service in (
            ("lifecycle policies", LifecycleService(session)),
            ("migration policies", MigrationEngineService(session)),
            ("observability rules", ObservabilityService(session, settings)),
            (
                "recovery policies",
                RecoveryService(session, settings, "control_plane", "bootstrap"),
            ),
        ):
            started = time.time()
            service.ensure_defaults()
            session.commit()
            report.record(name, "ENSURED", "defaults present", time.time() - started)
    engine.dispose()


def verify(url: str, settings: Settings, report: BootstrapReport) -> None:
    engine = create_engine(url)
    with engine.connect() as connection:
        revision = connection.execute(text("select version_num from alembic_version")).scalar_one()
        tables = int(
            connection.execute(
                text(
                    "select count(*) from information_schema.tables where table_schema='public'"
                )
            ).scalar_one()
        )
        duplicates = connection.execute(
            text(
                "select count(*) from (select capability_contract_id from capability_deployments "
                "where status='stable' group by capability_contract_id having count(*) > 1) as d"
            )
        ).scalar_one()
    engine.dispose()
    compatibility = schema_compatibility(revision)
    report.record(
        "schema",
        "OK" if compatibility is Compatibility.COMPATIBLE else "INCOMPATIBLE",
        f"revision {revision}, {tables} tables",
    )
    report.record(
        "no duplicate stable identity",
        "OK" if duplicates == 0 else "VIOLATED",
        f"{duplicates} contract(s) with more than one stable deployment",
    )
    has_operator_key = settings.operator_api_key is not None
    report.record(
        "operator credential",
        "PRESENT" if has_operator_key else "MISSING",
        "supplied through configuration"
        if has_operator_key
        else "set MODELFORGE_OPERATOR_API_KEY before serving",
    )
    if not has_operator_key and settings.env == "production":
        raise SystemExit(
            "refusing to finish: production requires MODELFORGE_OPERATOR_API_KEY. ModelForge "
            "never mints its own admin secret, because it would have no way to tell you it did."
        )


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--database-url", default=None)
    parser.add_argument("--env-file", default=None)
    parser.add_argument("--report", default=None)
    parser.add_argument("--skip-seed", action="store_true")
    args = parser.parse_args(argv)

    settings = (
        Settings(_env_file=args.env_file)  # type: ignore[call-arg]
        if args.env_file
        else Settings()
    )
    url = args.database_url or settings.database_url

    report = BootstrapReport()
    print(f"{PRODUCT_NAME} {PRODUCT_VERSION} — bootstrap", flush=True)
    wait_for_database(url, report)
    migrate(url, report)
    if not args.skip_seed:
        seed(url, settings, report)
    verify(url, settings, report)

    for step in report.steps:
        print(f"  {step.outcome:16} {step.name:28} {step.detail} ({step.seconds:.2f}s)", flush=True)
    print(f"  bootstrap completed in {report.as_dict()['total_seconds']:.2f}s", flush=True)

    if args.report:
        Path(args.report).write_text(json.dumps(report.as_dict(), indent=2), encoding="utf-8")
        print(f"  report written to {args.report}", flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())