Kansho/backend/sqlite_to_postgres.py

#!/usr/bin/env python3
"""One-shot SQLite → PostgreSQL copy for the laptop/home journal instance.

Probe on a copy, never on the only backup. Does not send data to providers.
Media files are copied separately into the destination media root.
"""
from __future__ import annotations

import argparse
import os
import shutil
import sqlite3
import sys
from pathlib import Path

from sql_compat import postgres_connect_kwargs

SCHEMA_PATH = Path(__file__).resolve().parent / "schema.sql"
DEFAULT_SQLITE = Path(__file__).resolve().parent / "data" / "kansho.sqlite"
DEFAULT_MEDIA = Path(__file__).resolve().parent / "data" / "media"

# Parent tables first — matches schema.sql CREATE order.
TABLES = [
    "schema_migrations",
    "tiers",
    "features",
    "tier_limits",
    "profiles",
    "sessions",
    "user_feature_restrictions",
    "user_feature_usage",
    "ai_prompts",
    "usage_sessions",
    "conversations",
    "messages",
    "threads",
    "conversation_threads",
    "spaces",
    "thread_spaces",
    "derived_records",
    "handoffs",
    "external_refs",
    "re_grounding_events",
    "identity_mappings",
    "identity_review_proposals",
    "journal_days",
    "journal_drafts",
    "journal_entries",
    "journal_entry_versions",
    "journal_draft_source_refs",
    "journal_entry_version_source_refs",
    "media_assets",
    "writing_profiles",
    "writing_profile_sources",
    "writing_profile_facets",
    "writing_profile_traits",
    "writing_profile_trait_refs",
    "writing_profile_suggestions",
    "interaction_profiles",
    "interaction_preferences",
    "interaction_suggestions",
    "writing_profile_evidence",
    "writing_profile_reviews",
    "writing_profile_versions",
    "provider_settings",
    "llm_profiles",
    "journal_generation_selection",
    "generation_guidelines",
    "app_settings",
    "debug_runs",
]


def _sqlite_path() -> Path:
    env = (os.environ.get("KANSHO_DB_PATH") or "").strip()
    if env:
        return Path(env)
    data_dir = (os.environ.get("KANSHO_DATA_DIR") or "").strip()
    if data_dir:
        return Path(data_dir) / "kansho.sqlite"
    return DEFAULT_SQLITE


def _media_root() -> Path:
    env = (os.environ.get("KANSHO_MEDIA_ROOT") or "").strip()
    if env:
        return Path(env)
    data_dir = (os.environ.get("KANSHO_DATA_DIR") or "").strip()
    if data_dir:
        return Path(data_dir) / "media"
    return DEFAULT_MEDIA


def _connect_pg():
    import psycopg

    kwargs = postgres_connect_kwargs()
    if "conninfo" in kwargs:
        return psycopg.connect(kwargs["conninfo"])
    return psycopg.connect(**kwargs)


def _pg_columns(conn, table: str) -> list[str]:
    rows = conn.execute(
        """
        SELECT column_name
        FROM information_schema.columns
        WHERE table_schema = 'public' AND table_name = %s
        ORDER BY ordinal_position
        """,
        (table,),
    ).fetchall()
    return [row[0] for row in rows]


def _sqlite_rows(sqlite_path: Path, table: str) -> list[dict]:
    conn = sqlite3.connect(str(sqlite_path))
    conn.row_factory = sqlite3.Row
    try:
        rows = conn.execute(f"SELECT * FROM {table}").fetchall()
        return [dict(row) for row in rows]
    except sqlite3.OperationalError:
        return []
    finally:
        conn.close()


def _copy_media(src: Path, dest: Path) -> int:
    if not src.is_dir():
        print(f"No media directory at {src}")
        return 0
    dest.mkdir(parents=True, exist_ok=True)
    count = 0
    for path in src.rglob("*"):
        if not path.is_file():
            continue
        relative = path.relative_to(src)
        target = dest / relative
        target.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(path, target)
        count += 1
    return count


def migrate(sqlite_path: Path, media_src: Path, media_dest: Path, replace: bool) -> dict[str, tuple[int, int]]:
    if not sqlite_path.is_file():
        raise SystemExit(f"SQLite not found: {sqlite_path}")

    pg = _connect_pg()
    stats: dict[str, tuple[int, int]] = {}
    try:
        existing = pg.execute("SELECT COUNT(*) FROM profiles").fetchone()[0]
        seeds = pg.execute("SELECT COUNT(*) FROM tiers").fetchone()[0]
        if existing and not replace:
            raise SystemExit(
                f"PostgreSQL already has {existing} profiles. Pass --replace after a probe copy, not on the only backup."
            )
        if seeds and not replace:
            raise SystemExit(
                "PostgreSQL already has platform seeds from startup. "
                "Pass --replace to load a SQLite snapshot onto this instance (no user profiles yet is OK)."
            )
        if replace:
            print(f"Replacing Postgres snapshot (profiles={existing}, seed_tiers={seeds})")
            for table in reversed(TABLES):
                pg.execute(f"TRUNCATE TABLE {table} CASCADE")

        for table in TABLES:
            rows = _sqlite_rows(sqlite_path, table)
            columns = _pg_columns(pg, table)
            if not columns:
                print(f"  skip {table} (missing in Postgres schema)")
                stats[table] = (len(rows), 0)
                continue
            if not rows:
                print(f"  {table}: empty")
                stats[table] = (0, 0)
                continue
            usable = [col for col in columns if col in rows[0]]
            placeholders = ", ".join(["%s"] * len(usable))
            col_sql = ", ".join(usable)
            sql = f"INSERT INTO {table} ({col_sql}) VALUES ({placeholders})"
            values = [tuple(row.get(col) for col in usable) for row in rows]
            with pg.cursor() as cur:
                cur.executemany(sql, values)
            count = pg.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0]
            print(f"  {table}: {len(rows)} → {count}")
            stats[table] = (len(rows), int(count))
        pg.execute(
            "INSERT INTO schema_migrations (id) VALUES (%s) ON CONFLICT (id) DO NOTHING",
            ("022_postgres_runtime_baseline",),
        )
        pg.commit()
    except Exception:
        pg.rollback()
        raise
    finally:
        pg.close()

    copied = _copy_media(media_src, media_dest)
    print(f"Media files copied: {copied} → {media_dest}")
    return stats


def verify(stats: dict[str, tuple[int, int]]) -> None:
    failed = False
    print("\nVerification")
    for table, (src, dest) in stats.items():
        if table == "schema_migrations" and dest == src + 1:
            print(f"  ok       {table:36} sqlite={src:5} postgres={dest:5} (incl. 022)")
            continue
        mark = "ok" if src == dest else "MISMATCH"
        if src != dest:
            failed = True
        print(f"  {mark:8} {table:36} sqlite={src:5} postgres={dest:5}")
    if failed:
        raise SystemExit("Row counts do not match")
    print("Row counts match")


def main() -> None:
    parser = argparse.ArgumentParser(description="Copy Kanshō SQLite data into PostgreSQL")
    parser.add_argument("--sqlite", type=Path, default=_sqlite_path())
    parser.add_argument("--media-src", type=Path, default=_media_root())
    parser.add_argument("--media-dest", type=Path, default=Path(os.environ.get("KANSHO_MEDIA_DEST") or "/app/data/media"))
    parser.add_argument("--confirm", action="store_true", help="Required. Refuse a silent copy.")
    parser.add_argument("--replace", action="store_true", help="Truncate Postgres tables first")
    args = parser.parse_args()
    if not args.confirm:
        raise SystemExit("Refusing to copy personal data without --confirm")
    print(f"SQLite source: {args.sqlite}")
    print(f"Media source:  {args.media_src}")
    print(f"Media dest:    {args.media_dest}")
    stats = migrate(args.sqlite, args.media_src, args.media_dest, args.replace)
    verify(stats)
    print("Import complete. Check entry/message counts and one media GET before Prod.")


if __name__ == "__main__":
    main()