oclaw/runtime/operations/scripts/migrate_assistant_sqlite_to_postgresql.py

"""Copy assistant data from SQLite (db_path file) into PostgreSQL (schema from Alembic / bootstrap).

**Prerequisite:** target PostgreSQL already has schema (``alembic upgrade head`` or
``svc/persistence/ddl/postgresql_bootstrap.sql``). This script copies **data only**.

Tables are copied in **foreign-key safe order** (from SQLite ``PRAGMA foreign_key_list``), and only
**columns present in both** SQLite and PostgreSQL are inserted.

By default the script aborts if **both** sides would supply rows for the same table: SQLite has at
least one row **and** PostgreSQL already has rows (typical after a gateway run seeded ``llm_profile``
etc.). Empty SQLite tables may coexist with PostgreSQL seed rows. Use ``--allow-non-empty`` to skip
that check (duplicates / primary-key violations are your risk).

**PostgreSQL URL** is taken from ``--pg-url`` if set; otherwise from the first non-empty environment
variable among ``AIA_ASSISTANT_DATABASE_URL``, ``OPS_ASSISTANT_DATABASE_URL``, ``AIA_ASSISTANT_PG_DSN``,
``OPS_ASSISTANT_PG_DSN``. Use ``--load-system-env`` to merge ``_local/system.env`` first (same as the
HTTP gateway).

**Open-source / headless device (Linux example)**::

    export AIA_ASSISTANT_DATABASE_URL='postgresql+psycopg://USER:PASS@HOST:5432/oclaw'
    export AIA_ASSISTANT_DB_PATH=/var/lib/oclaw/data/ai_ops.sqlite   # optional; default data/ai_ops.sqlite under repo
    cd /path/to/oclaw && PYTHONPATH=. python runtime/operations/scripts/migrate_assistant_sqlite_to_postgresql.py \\
        --load-system-env --sqlite-from-db-path --dry-run
    # then same without --dry-run

Or use the wrapper script ``runtime/operations/scripts/assistant_import_sqlite_to_postgresql.sh``.

**Explicit paths**::

    python runtime/operations/scripts/migrate_assistant_sqlite_to_postgresql.py \\
        --sqlite data/ai_ops.sqlite \\
        --pg-url postgresql+psycopg://postgres:pass@127.0.0.1:5432/oclaw

"""

from __future__ import annotations

import argparse
import os
import re
import sqlite3
import sys
from collections import defaultdict
from pathlib import Path
from typing import Any, Iterable

import psycopg
from psycopg.rows import dict_row

_REPO_ROOT = Path(__file__).resolve().parents[3]
if str(_REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(_REPO_ROOT))

from svc.persistence.pg_adapter import normalize_psycopg_conninfo


def _pg_row_first_value(row: Any) -> Any:
    if row is None:
        raise ValueError("expected a row")
    if isinstance(row, dict):
        return next(iter(row.values()))
    return row[0]


_IDENT = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*\Z")


def _require_ident(name: str) -> str:
    if not _IDENT.fullmatch(name):
        raise ValueError(f"invalid SQL identifier: {name!r}")
    return name


def _sqlite_user_tables(sl: sqlite3.Connection) -> list[str]:
    rows = sl.execute(
        """
        SELECT name FROM sqlite_master
        WHERE type='table' AND name NOT LIKE 'sqlite_%'
        ORDER BY name
        """
    ).fetchall()
    return [str(r[0]) for r in rows]


def _pg_public_tables(pg: psycopg.Connection) -> set[str]:
    with pg.cursor() as cur:
        cur.execute(
            """
            SELECT tablename FROM pg_catalog.pg_tables
            WHERE schemaname = 'public'
            """
        )
        return {str(_pg_row_first_value(r)) for r in cur.fetchall()}


def _fk_parents_for_table(sl: sqlite3.Connection, table: str) -> set[str]:
    t = _require_ident(table)
    rows = sl.execute(f"PRAGMA foreign_key_list({t})").fetchall()
    out: set[str] = set()
    for r in rows:
        # (id, seq, table, from, to, on_update, on_delete, match)
        ref = str(r[2])
        if ref:
            out.add(ref)
    return out


def _topological_sort(nodes: list[str], parents: dict[str, set[str]]) -> list[str]:
    """``parents[t]`` = tables that must be copied *before* ``t`` (referenced by FK)."""
    node_set = list(nodes)
    seen = set(node_set)
    if len(seen) != len(node_set):
        raise ValueError("duplicate table in migration list")

    children: dict[str, list[str]] = defaultdict(list)
    indegree: dict[str, int] = {}
    for t in node_set:
        ps = parents.get(t, set()) & seen
        indegree[t] = len(ps)
        for p in ps:
            children[p].append(t)
    for ch in children.values():
        ch.sort()

    queue = sorted([t for t in node_set if indegree[t] == 0])
    out: list[str] = []
    while queue:
        n = queue.pop(0)
        out.append(n)
        for c in children[n]:
            indegree[c] -= 1
            if indegree[c] == 0:
                queue.append(c)
                queue.sort()
    if len(out) != len(seen):
        remain = seen - set(out)
        raise SystemExit(
            "Cannot derive a foreign-key-safe copy order (cycle or unresolved FK). "
            f"Remaining tables: {sorted(remain)}"
        )
    return out


def _migration_order(sl: sqlite3.Connection, tables: list[str]) -> list[str]:
    node_set = list(tables)
    parents = {t: _fk_parents_for_table(sl, t) & set(node_set) for t in node_set}
    return _topological_sort(node_set, parents)


def _sqlite_columns(sl: sqlite3.Connection, table: str) -> list[str]:
    t = _require_ident(table)
    rows = sl.execute(f"PRAGMA table_info({t})").fetchall()
    # cid, name, type, notnull, dflt_value, pk
    return [str(r[1]) for r in rows]


def _pg_columns(pg: psycopg.Connection, table: str) -> set[str]:
    t = _require_ident(table)
    with pg.cursor() as cur:
        cur.execute(
            """
            SELECT column_name FROM information_schema.columns
            WHERE table_schema = 'public' AND table_name = %s
            """,
            (t,),
        )
        return {str(_pg_row_first_value(r)) for r in cur.fetchall()}


def _common_columns(sl: sqlite3.Connection, pg: psycopg.Connection, table: str) -> list[str]:
    sc = _sqlite_columns(sl, table)
    pc = _pg_columns(pg, table)
    return [c for c in sc if c in pc and _IDENT.fullmatch(c)]


def _sqlite_pg_nonempty_import_conflict_msg(table: str, sqlite_n: int, pg_n: int) -> str | None:
    """Return an error message if we cannot safely INSERT SQLite rows into PG for this table."""
    if sqlite_n <= 0 or pg_n <= 0:
        return None
    return (
        f"Refusing to import: table {table!r} has {sqlite_n} row(s) in SQLite but PostgreSQL "
        f"already has {pg_n} row(s). That usually means data or seeds already exist in PG "
        "(for example the gateway was started once). Use a fresh schema after `alembic upgrade head`, "
        "truncate this table on PostgreSQL, or pass --allow-non-empty knowing duplicates may cause "
        "primary-key violations."
    )


def _assert_pg_tables_compatible_for_copy(
    sl: sqlite3.Connection,
    pg: psycopg.Connection,
    tables: Iterable[str],
) -> None:
    """Abort when SQLite would copy rows into a PG table that is not empty."""
    for t in sorted(set(tables)):
        _require_ident(t)
        sqlite_n = int(sl.execute(f"SELECT COUNT(*) FROM {_require_ident(t)}").fetchone()[0])
        with pg.cursor() as cur:
            cur.execute(f'SELECT COUNT(*) AS n FROM "{_require_ident(t)}"')
            pg_n = int(_pg_row_first_value(cur.fetchone()))
        msg = _sqlite_pg_nonempty_import_conflict_msg(t, sqlite_n, pg_n)
        if msg:
            raise SystemExit(msg)


def _sqlite_row_counts(sl: sqlite3.Connection, tables: Iterable[str]) -> dict[str, int]:
    out: dict[str, int] = {}
    for t in tables:
        _require_ident(t)
        n = int(sl.execute(f"SELECT COUNT(*) FROM {_require_ident(t)}").fetchone()[0])
        out[t] = n
    return out


def _pg_row_counts(pg: psycopg.Connection, tables: Iterable[str]) -> dict[str, int]:
    out: dict[str, int] = {}
    with pg.cursor() as cur:
        for t in tables:
            _require_ident(t)
            cur.execute(f'SELECT COUNT(*) AS n FROM "{t}"')
            row = cur.fetchone()
            out[t] = int(_pg_row_first_value(row))
    return out


def _copy_table(
    *,
    sl: sqlite3.Connection,
    pg: psycopg.Connection,
    table: str,
    cols: list[str],
    dry_run: bool,
    batch: int,
) -> int:
    if not cols:
        return 0
    t = _require_ident(table)
    cur = sl.execute(f"SELECT {', '.join(_require_ident(c) for c in cols)} FROM {t}")
    rows = cur.fetchall()
    if not rows:
        return 0
    if dry_run:
        return len(rows)
    col_sql = ", ".join(f'"{_require_ident(c)}"' for c in cols)
    placeholders = ", ".join(["%s"] * len(cols))
    sql = f'INSERT INTO "{t}" ({col_sql}) VALUES ({placeholders})'
    tuples = [tuple(r[c] for c in cols) for r in rows]
    with pg.cursor() as pc:
        for i in range(0, len(tuples), max(1, batch)):
            chunk = tuples[i : i + max(1, batch)]
            pc.executemany(sql, chunk)
    return len(rows)


def _serial_columns(pg: psycopg.Connection) -> list[tuple[str, str]]:
    """Tables/columns backed by a PostgreSQL sequence (BIGSERIAL etc.), for post-import setval."""
    with pg.cursor() as cur:
        cur.execute(
            """
            SELECT table_name, column_name
            FROM information_schema.columns
            WHERE table_schema = 'public'
              AND column_default IS NOT NULL
              AND column_default LIKE 'nextval%'
            ORDER BY table_name, column_name
            """
        )
        rows = cur.fetchall()
    out: list[tuple[str, str]] = []
    for r in rows:
        if isinstance(r, dict):
            t = str(r["table_name"])
            c = str(r["column_name"])
        else:
            t = str(r[0])
            c = str(r[1])
        if t and c:
            out.append((_require_ident(t), _require_ident(c)))
    return out


def _sync_sequences(pg: psycopg.Connection) -> None:
    for t, col in _serial_columns(pg):
        with pg.cursor() as cur:
            cur.execute(f'SELECT COALESCE(MAX("{col}"), 1) AS mx FROM "{t}"')
            row = cur.fetchone()
            mx = int(_pg_row_first_value(row))
            try:
                cur.execute(
                    "SELECT setval(pg_get_serial_sequence(%s, %s), %s, true)",
                    (t, col, mx),
                )
            except Exception:
                pass


def _pg_url_from_environ() -> str:
    return (
        os.getenv("AIA_ASSISTANT_DATABASE_URL")
        or os.getenv("OPS_ASSISTANT_DATABASE_URL")
        or os.getenv("AIA_ASSISTANT_PG_DSN")
        or os.getenv("OPS_ASSISTANT_PG_DSN")
        or ""
    ).strip()


def _resolve_sqlite_path(arg: str | None, use_db_path: bool) -> Path:
    if use_db_path:
        from svc.config.paths import db_path

        return Path(db_path()).expanduser().resolve()
    if not arg:
        raise SystemExit("Either pass --sqlite PATH or --sqlite-from-db-path")
    return Path(arg).expanduser().resolve()


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--sqlite", default=None, help="Path to source SQLite assistant DB")
    ap.add_argument(
        "--sqlite-from-db-path",
        action="store_true",
        help="Use db_path() from env (AIA_ASSISTANT_DB_PATH / default) as SQLite source",
    )
    ap.add_argument(
        "--load-system-env",
        action="store_true",
        help="Merge _local/system.env into the process (for DB_PATH / DATABASE_URL on devices)",
    )
    ap.add_argument(
        "--pg-url",
        default=None,
        help="Target PostgreSQL URL; if omitted, use AIA_ASSISTANT_DATABASE_URL (or OPS_* / *_PG_DSN)",
    )
    ap.add_argument("--dry-run", action="store_true", help="Count rows only; do not write to PG")
    ap.add_argument(
        "--allow-non-empty",
        action="store_true",
        help="Skip the SQLite+PG nonempty conflict check (risk of duplicate keys / bad counts)",
    )
    ap.add_argument(
        "--batch",
        type=int,
        default=500,
        help="Rows per executemany batch (default 500)",
    )
    args = ap.parse_args()
    if bool(args.load_system_env):
        from svc.config.bootstrap_env import load_system_env

        load_system_env(force=True)

    sqlite_path = _resolve_sqlite_path(args.sqlite, bool(args.sqlite_from_db_path))
    if not sqlite_path.is_file():
        raise SystemExit(f"sqlite file not found: {sqlite_path}")

    raw_pg = (args.pg_url or "").strip() or _pg_url_from_environ()
    if not raw_pg:
        raise SystemExit(
            "No PostgreSQL URL: pass --pg-url or set one of "
            "AIA_ASSISTANT_DATABASE_URL, OPS_ASSISTANT_DATABASE_URL, "
            "AIA_ASSISTANT_PG_DSN, OPS_ASSISTANT_PG_DSN (use --load-system-env to read _local/system.env)."
        )
    pg_url = normalize_psycopg_conninfo(raw_pg)
    sl = sqlite3.connect(str(sqlite_path))
    sl.row_factory = sqlite3.Row
    try:
        sl.execute("PRAGMA foreign_keys=ON")
    except Exception:
        pass
    pg = psycopg.connect(pg_url, row_factory=dict_row, autocommit=False)
    try:
        sqlite_tables = _sqlite_user_tables(sl)
        pg_tables = _pg_public_tables(pg)
        common = [t for t in sqlite_tables if t in pg_tables]
        skipped_sqlite = [t for t in sqlite_tables if t not in pg_tables]
        if skipped_sqlite:
            print("skip (not in PG public schema):", ", ".join(skipped_sqlite))

        if not common:
            raise SystemExit("No common tables between SQLite and PostgreSQL; nothing to copy.")

        order = _migration_order(sl, common)
        if not args.dry_run and not args.allow_non_empty:
            _assert_pg_tables_compatible_for_copy(sl, pg, common)

        src_counts = _sqlite_row_counts(sl, order)
        total = 0
        copied: list[str] = []
        for t in order:
            cols = _common_columns(sl, pg, t)
            if not cols:
                n0 = src_counts.get(t, 0)
                if n0 > 0:
                    print(f"{t}: SKIP (no common columns; sqlite has {n0} rows — schema drift)")
                continue
            n = _copy_table(sl=sl, pg=pg, table=t, cols=cols, dry_run=bool(args.dry_run), batch=int(args.batch))
            print(f"{t}: {n} rows ({len(cols)} columns)")
            total += n
            copied.append(t)
        if not args.dry_run:
            pg.commit()
            _sync_sequences(pg)
            pg.commit()
            verify = _pg_row_counts(pg, order)
            bad = [t for t in copied if verify.get(t, 0) != src_counts.get(t, 0)]
            if bad:
                print("WARNING: row count mismatch PG vs SQLite for:", ", ".join(bad))
                for t in bad:
                    print(f"  {t}: sqlite={src_counts.get(t, 0)} pg={verify.get(t, 0)}")
            else:
                print("verify: row counts match SQLite for all copied tables")
        print("total rows (copied or dry-run counted):", total)
    except BaseException:
        pg.rollback()
        raise
    finally:
        sl.close()
        pg.close()


if __name__ == "__main__":
    main()