web-petting/scripts/migrate_sqlite_to_postgres.py

#!/usr/bin/env python3
"""
One-off data migration from the old web-petting SQLite database to PostgreSQL.

The script intentionally lives outside the Rust application and does not affect
the app dependency graph. It migrates user-owned application data only and skips
Cot technical tables such as cot__migrations and cot__session.

Prerequisites:
  1. Run the new app migrations against PostgreSQL first, so the target tables
     already exist.
  2. Install one PostgreSQL Python driver in the environment that runs this
     script:
       python -m pip install "psycopg[binary]"
     psycopg2 is also supported if it is already installed.

Usage:
  python scripts/migrate_sqlite_to_postgres.py --sqlite db.sqlite3 --postgres-url postgresql://...
  python scripts/migrate_sqlite_to_postgres.py --dry-run
"""

from __future__ import annotations

import argparse
import os
import sqlite3
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable


Transform = Callable[[Any], Any]


@dataclass(frozen=True)
class TableSpec:
    name: str
    columns: tuple[str, ...]
    transforms: dict[str, Transform] | None = None


def optional_bool(value: Any) -> bool | None:
    if value is None:
        return None
    if value in (0, "0", False):
        return False
    if value in (1, "1", True):
        return True
    raise ValueError(f"cannot convert {value!r} to bool")


TABLES: tuple[TableSpec, ...] = (
    TableSpec(
        "web_petting__user",
        (
            "id",
            "login",
            "password_hash",
            "display_name",
            "telegram_chat_id",
            "telegram_notifications",
            "status",
            "created_at",
            "updated_at",
        ),
        transforms={"telegram_notifications": optional_bool},
    ),
    TableSpec(
        "web_petting__setting",
        ("id", "key", "value", "updated_at"),
    ),
    TableSpec(
        "web_petting__client",
        (
            "id",
            "name",
            "phone",
            "email",
            "address",
            "notes",
            "media_token",
            "color",
            "status",
            "created_at",
            "updated_at",
        ),
    ),
    TableSpec(
        "web_petting__testimonial",
        ("id", "text", "author_note", "image_path", "status", "sort_order", "created_at"),
    ),
    TableSpec(
        "web_petting__visit",
        (
            "id",
            "client_id",
            "user_id",
            "visit_date",
            "time_start",
            "time_end",
            "notes",
            "public_notes",
            "client_feedback",
            "status",
            "created_at",
            "updated_at",
        ),
    ),
    TableSpec(
        "web_petting__media",
        ("id", "client_id", "visit_id", "file_path", "file_type", "caption", "status", "created_at"),
    ),
    TableSpec(
        "web_petting__lead",
        ("id", "name", "phone", "email", "comment", "status", "client_id", "created_at", "updated_at"),
    ),
)


FK_CHECKS: tuple[tuple[str, str, str, str], ...] = (
    ("web_petting__lead", "client_id", "web_petting__client", "id"),
    ("web_petting__visit", "client_id", "web_petting__client", "id"),
    ("web_petting__visit", "user_id", "web_petting__user", "id"),
    ("web_petting__media", "client_id", "web_petting__client", "id"),
    ("web_petting__media", "visit_id", "web_petting__visit", "id"),
)


def quote_ident(identifier: str) -> str:
    return '"' + identifier.replace('"', '""') + '"'


def pg_table_name(table: str, schema: str | None) -> str:
    if schema:
        return f"{quote_ident(schema)}.{quote_ident(table)}"
    return quote_ident(table)


def sequence_relation_name(table: str, schema: str | None) -> str:
    if schema:
        return f"{schema}.{table}"
    return table


def sqlite_table_exists(conn: sqlite3.Connection, table: str) -> bool:
    row = conn.execute(
        "select 1 from sqlite_schema where type = 'table' and name = ?",
        (table,),
    ).fetchone()
    return row is not None


def sqlite_columns(conn: sqlite3.Connection, table: str) -> set[str]:
    return {row["name"] for row in conn.execute(f"pragma table_info({quote_ident(table)})")}


def sqlite_count(conn: sqlite3.Connection, table: str) -> int:
    return int(conn.execute(f"select count(*) from {quote_ident(table)}").fetchone()[0])


def read_rows(conn: sqlite3.Connection, spec: TableSpec) -> list[tuple[Any, ...]]:
    columns_sql = ", ".join(quote_ident(column) for column in spec.columns)
    rows: list[tuple[Any, ...]] = []
    for row in conn.execute(f"select {columns_sql} from {quote_ident(spec.name)} order by id"):
        values = []
        for column in spec.columns:
            value = row[column]
            if spec.transforms and column in spec.transforms:
                value = spec.transforms[column](value)
            values.append(value)
        rows.append(tuple(values))
    return rows


def validate_sqlite(conn: sqlite3.Connection) -> None:
    for spec in TABLES:
        if not sqlite_table_exists(conn, spec.name):
            raise RuntimeError(f"SQLite table is missing: {spec.name}")
        existing_columns = sqlite_columns(conn, spec.name)
        missing_columns = [column for column in spec.columns if column not in existing_columns]
        if missing_columns:
            joined = ", ".join(missing_columns)
            raise RuntimeError(f"SQLite table {spec.name} is missing columns: {joined}")

    for source_table, source_column, target_table, target_column in FK_CHECKS:
        rows = conn.execute(
            f"""
            select s.id, s.{quote_ident(source_column)}
            from {quote_ident(source_table)} s
            left join {quote_ident(target_table)} t
              on s.{quote_ident(source_column)} = t.{quote_ident(target_column)}
            where s.{quote_ident(source_column)} is not null
              and t.{quote_ident(target_column)} is null
            order by s.id
            limit 20
            """
        ).fetchall()
        if rows:
            examples = ", ".join(
                f"id={row[0]} {source_column}={row[1]}" for row in rows
            )
            raise RuntimeError(
                f"SQLite FK check failed: {source_table}.{source_column} -> "
                f"{target_table}.{target_column}; examples: {examples}"
            )


def import_postgres_driver() -> Any:
    try:
        import psycopg  # type: ignore

        return psycopg
    except ModuleNotFoundError:
        pass

    try:
        import psycopg2  # type: ignore

        return psycopg2
    except ModuleNotFoundError as exc:
        raise RuntimeError(
            "No PostgreSQL Python driver found. Install one with: "
            'python -m pip install "psycopg[binary]"'
        ) from exc


def connect_postgres(postgres_url: str) -> Any:
    driver = import_postgres_driver()
    return driver.connect(postgres_url)


def postgres_table_exists(cur: Any, table: str, schema: str | None) -> bool:
    if schema:
        cur.execute(
            """
            select 1
            from information_schema.tables
            where table_schema = %s and table_name = %s
            """,
            (schema, table),
        )
    else:
        cur.execute("select to_regclass(%s)", (table,))
    row = cur.fetchone()
    if not row:
        return False
    return bool(row[0])


def postgres_count(cur: Any, table: str, schema: str | None) -> int:
    cur.execute(f"select count(*) from {pg_table_name(table, schema)}")
    return int(cur.fetchone()[0])


def ensure_postgres_ready(cur: Any, schema: str | None, allow_existing: bool) -> None:
    for spec in TABLES:
        if not postgres_table_exists(cur, spec.name, schema):
            raise RuntimeError(
                f"PostgreSQL table is missing: {spec.name}. "
                "Run the new app migrations before this script."
            )
        count = postgres_count(cur, spec.name, schema)
        if count and not allow_existing:
            raise RuntimeError(
                f"PostgreSQL table {spec.name} already contains {count} rows. "
                "Use --truncate-user-tables to replace user data."
            )


def truncate_user_tables(cur: Any, schema: str | None) -> None:
    tables = ", ".join(pg_table_name(spec.name, schema) for spec in TABLES)
    cur.execute(f"truncate table {tables} restart identity")


def insert_table(cur: Any, spec: TableSpec, rows: list[tuple[Any, ...]], schema: str | None) -> None:
    if not rows:
        return
    columns_sql = ", ".join(quote_ident(column) for column in spec.columns)
    placeholders = ", ".join(["%s"] * len(spec.columns))
    sql = f"insert into {pg_table_name(spec.name, schema)} ({columns_sql}) values ({placeholders})"
    cur.executemany(sql, rows)


def reset_sequence(cur: Any, table: str, schema: str | None) -> None:
    relation = sequence_relation_name(table, schema)
    cur.execute(
        """
        select setval(
            pg_get_serial_sequence(%s, 'id'),
            greatest(coalesce((select max(id) from %s), 0), 1),
            coalesce((select max(id) from %s), 0) > 0
        )
        """
        % ("%s", pg_table_name(table, schema), pg_table_name(table, schema)),
        (relation,),
    )


def migrate(sqlite_path: Path, postgres_url: str, schema: str | None, truncate: bool) -> None:
    sqlite_conn = sqlite3.connect(f"file:{sqlite_path}?mode=ro", uri=True)
    sqlite_conn.row_factory = sqlite3.Row
    try:
        validate_sqlite(sqlite_conn)
        source_rows = {spec.name: read_rows(sqlite_conn, spec) for spec in TABLES}
    finally:
        sqlite_conn.close()

    pg_conn = connect_postgres(postgres_url)
    try:
        with pg_conn:
            with pg_conn.cursor() as cur:
                ensure_postgres_ready(cur, schema, allow_existing=truncate)
                if truncate:
                    truncate_user_tables(cur, schema)
                for spec in TABLES:
                    insert_table(cur, spec, source_rows[spec.name], schema)
                    print(f"inserted {len(source_rows[spec.name])} rows into {spec.name}")
                for spec in TABLES:
                    reset_sequence(cur, spec.name, schema)
                for spec in TABLES:
                    target_count = postgres_count(cur, spec.name, schema)
                    source_count = len(source_rows[spec.name])
                    if target_count != source_count:
                        raise RuntimeError(
                            f"count mismatch for {spec.name}: source={source_count}, target={target_count}"
                        )
    finally:
        pg_conn.close()


def dry_run(sqlite_path: Path) -> None:
    conn = sqlite3.connect(f"file:{sqlite_path}?mode=ro", uri=True)
    conn.row_factory = sqlite3.Row
    try:
        validate_sqlite(conn)
        print("SQLite source is valid.")
        for spec in TABLES:
            print(f"{spec.name}: {sqlite_count(conn, spec.name)} rows")
        print("Cot technical tables are intentionally skipped.")
    finally:
        conn.close()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Migrate web-petting user data from SQLite to PostgreSQL.",
    )
    parser.add_argument(
        "--sqlite",
        default="db.sqlite3",
        type=Path,
        help="Path to the downloaded SQLite database (default: db.sqlite3).",
    )
    parser.add_argument(
        "--postgres-url",
        default=os.environ.get("WEB_PETTING_DATABASE_URL") or os.environ.get("DATABASE_URL"),
        help="PostgreSQL URL. Defaults to WEB_PETTING_DATABASE_URL or DATABASE_URL.",
    )
    parser.add_argument(
        "--schema",
        default="public",
        help="Target PostgreSQL schema (default: public). Use an empty string for current search_path.",
    )
    parser.add_argument(
        "--truncate-user-tables",
        action="store_true",
        help="Delete existing rows from migrated application tables before inserting.",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Validate and count SQLite rows without connecting to PostgreSQL.",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    sqlite_path = args.sqlite
    if not sqlite_path.exists():
        print(f"SQLite file not found: {sqlite_path}", file=sys.stderr)
        return 2

    schema = args.schema or None

    try:
        if args.dry_run:
            dry_run(sqlite_path)
            return 0

        if not args.postgres_url:
            print(
                "PostgreSQL URL is required. Pass --postgres-url or set WEB_PETTING_DATABASE_URL.",
                file=sys.stderr,
            )
            return 2

        migrate(sqlite_path, args.postgres_url, schema, args.truncate_user_tables)
        print("Migration complete.")
        return 0
    except Exception as exc:
        print(f"Migration failed: {exc}", file=sys.stderr)
        return 1


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