cleveragents-core/tools/controller/deploy/validate_db_mode.py

#!/usr/bin/env python3
"""Verify a SQLite db belongs to the launch mode (fork vs prod).

Called from ``tools/run-controller-state-machine-pipeline.sh`` once per
db (the controller db and the telemetry cache db) right after the
env is loaded and BEFORE the controller / OpenCode / telemetry start.

What it does
------------
- File missing → OK (will be created fresh by the writer).
- File exists with a ``_mode_marker`` table:
    * marker mode == requested mode → OK
    * marker mode != requested mode → **move file aside** to
      ``<file>.<prior-mode>.bak.<UTC-timestamp>`` and return. The writer
      will recreate the file fresh on first use.
- File exists without a ``_mode_marker``:
    * has no user tables (empty schema) → stamp and return.
    * has data (pre-marker era — e.g. the legacy ``tools/.cache/forgejo.sqlite``
      from the dispatcher era) →
        - default behavior: move aside as ``<file>.unmarked.bak.<ts>``.
          Safest assumption when we can't prove it belongs.
        - with ``--adopt``: stamp the existing data as the current mode
          (operator vouches that it belongs). Useful on first launch
          after this check is rolled out, to grandfather in already-good
          fork or prod data without losing history.

The script never deletes anything. Backups stay in place until the
operator removes them.

CLI
---

Pass either an absolute file path or a SQLAlchemy URL — non-SQLite URLs
(e.g. postgresql://...) are reported and skipped:

    validate_db_mode.py --mode prod --db /var/lib/cleveragents/db/controller.db
    validate_db_mode.py --mode prod --sqlalchemy-url "$CLEVERAGENTS_DB_URL"
    validate_db_mode.py --mode prod --db tools/.cache/forgejo.sqlite --label cache

The script exits 0 in every benign db-state case (missing, matching,
freshly stamped, backed-up-and-cleared). Exit 1 is reserved for argparse
errors (caller passed something wrong) — a real db inspection never
returns 1, so the launcher's ``|| die`` only fires when the script
itself was invoked incorrectly.
"""

from __future__ import annotations

import argparse
import datetime as dt
import sqlite3
import sys
from pathlib import Path
from urllib.parse import urlparse

# Schema kept here so the launcher and any future caller agree.
_MARKER_DDL = (
    "CREATE TABLE IF NOT EXISTS _mode_marker ("
    "  mode TEXT PRIMARY KEY,"
    "  stamped_at TEXT NOT NULL"
    ")"
)


def _resolve_path(db_arg: str | None, url_arg: str | None) -> Path | None:
    """Return the SQLite file path, or None when the URL points at a
    non-SQLite backend (postgres etc. — nothing for us to do)."""
    if db_arg:
        return Path(db_arg).expanduser()
    if not url_arg:
        return None
    parsed = urlparse(url_arg)
    if not parsed.scheme.startswith("sqlite"):
        return None
    # SQLAlchemy sqlite URLs:
    #   sqlite:///relative/path  →  urlparse.path = '/relative/path'  (the first
    #                               slash is the URL-path leading slash; the rest
    #                               is the "relative path" — confusingly absolute-
    #                               looking from urlparse)
    #   sqlite:////abs/path     →  urlparse.path = '//abs/path'  (double-leading
    #                               slash flags absolute on POSIX)
    raw = parsed.path
    if raw.startswith("//"):
        return Path(raw[1:]).expanduser()  # absolute → /abs/path
    if raw.startswith("/"):
        return Path(raw[1:]).expanduser()  # SQLAlchemy "relative" — strip leading /
    return Path(raw).expanduser()


def _log(label: str, level: str, msg: str) -> None:
    print(f"  [{label}] {level}: {msg}", file=sys.stderr)


def _backup(db_path: Path, prior_mode: str) -> Path:
    ts = dt.datetime.now(dt.timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    backup = db_path.with_name(db_path.name + f".{prior_mode}.bak.{ts}")
    db_path.rename(backup)
    return backup


def _stamp(db_path: Path, mode: str) -> None:
    """Create the marker table (if absent) and write the current mode."""
    db_path.parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(str(db_path))
    try:
        conn.execute(_MARKER_DDL)
        conn.execute(
            "INSERT OR REPLACE INTO _mode_marker (mode, stamped_at) VALUES (?, ?)",
            (mode, dt.datetime.now(dt.timezone.utc).isoformat()),
        )
        conn.commit()
    finally:
        conn.close()


def _user_tables(conn: sqlite3.Connection) -> list[str]:
    """Return user-table names — exclude SQLite internals and our own
    marker so we can detect a "schema present but no data" situation."""
    rows = conn.execute(
        "SELECT name FROM sqlite_master "
        "WHERE type='table' "
        "  AND name NOT LIKE 'sqlite_%' "
        "  AND name NOT LIKE '\\_%' ESCAPE '\\'"
    ).fetchall()
    return [r[0] for r in rows]


def _has_any_data(conn: sqlite3.Connection, tables: list[str]) -> bool:
    for t in tables:
        try:
            cnt = conn.execute(f'SELECT COUNT(*) FROM "{t}"').fetchone()[0]
            if cnt:
                return True
        except sqlite3.DatabaseError:
            # Treat unreadable as data-bearing to be safe.
            return True
    return False


def validate(db_path: Path, mode: str, label: str, adopt: bool = False) -> int:
    if not db_path.exists():
        _log(label, "ok", f"{db_path} does not exist; will be created fresh in {mode} mode")
        # We can't stamp a non-existent file (the writer needs to create
        # it first with its own schema). Stamping happens on first
        # subsequent launch once data shows up.
        return 0

    try:
        conn = sqlite3.connect(str(db_path))
    except sqlite3.DatabaseError as exc:
        _log(label, "warn", f"{db_path} not openable as sqlite ({exc!r}); leaving alone")
        return 0

    try:
        marker_present = conn.execute(
            "SELECT name FROM sqlite_master "
            "WHERE type='table' AND name='_mode_marker'"
        ).fetchone() is not None

        if marker_present:
            row = conn.execute(
                "SELECT mode FROM _mode_marker LIMIT 1"
            ).fetchone()
            existing = row[0] if row else None
            if existing == mode:
                _log(label, "ok", f"{db_path} mode marker confirms {existing!r}")
                return 0
            conn.close()
            backup = _backup(db_path, existing or "unknown")
            _log(label, "WIPED",
                 f"{db_path} stamped as {existing!r} but launching {mode!r}; "
                 f"moved aside → {backup.name}")
            return 0

        # No marker. Decide based on whether there's any data.
        tables = _user_tables(conn)
        if not tables or not _has_any_data(conn, tables):
            conn.close()
            _stamp(db_path, mode)
            _log(label, "ok",
                 f"{db_path} had no marker and no data; stamped as {mode!r}")
            return 0

        # Has data without a marker — almost certainly pre-marker era
        # state from before this check existed.
        conn.close()
        if adopt:
            # Operator vouches this data belongs to the current mode.
            _stamp(db_path, mode)
            _log(label, "ok",
                 f"{db_path} had data but no marker — ADOPTED as {mode!r} "
                 "(operator vouched via --adopt)")
            return 0
        backup = _backup(db_path, "unmarked")
        _log(label, "WIPED",
             f"{db_path} had data but no mode marker (pre-marker state); "
             f"moved aside → {backup.name}  "
             f"(re-run with --adopt to grandfather it in instead)")
        return 0
    finally:
        try:
            conn.close()
        except sqlite3.Error:
            pass


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Verify a SQLite db belongs to the launch mode (fork vs prod).",
    )
    parser.add_argument(
        "--mode", choices=("fork", "prod"), required=True,
        help="Launch mode the controller is starting in.",
    )
    src = parser.add_mutually_exclusive_group(required=True)
    src.add_argument("--db", help="SQLite file path (may not exist yet).")
    src.add_argument(
        "--sqlalchemy-url",
        help="SQLAlchemy URL — only sqlite:// URLs are validated; "
             "anything else is reported and skipped.",
    )
    parser.add_argument(
        "--label", default="db",
        help="Short label used in log output (e.g. 'controller', 'cache').",
    )
    parser.add_argument(
        "--adopt", action="store_true",
        help="If the db has data without a mode marker, STAMP it as the "
             "current mode rather than backing it up. Use this on first "
             "rollout to grandfather in already-good data without losing "
             "history. NOT for normal launches — only when the operator "
             "knows the existing data belongs to the requested mode.",
    )
    args = parser.parse_args(argv)

    db_path = _resolve_path(args.db, args.sqlalchemy_url)
    if db_path is None:
        _log(args.label, "skip",
             f"non-sqlite backend ({args.sqlalchemy_url!r}) — mode marker "
             "check not implemented for this backend; relying on the URL "
             "itself to point at the right db")
        return 0

    return validate(db_path, args.mode, args.label, adopt=args.adopt)


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