bober-bbq-bot/scripts/migrate_sqlite_to_postgres.py

"""One-off MANUAL migration: copy all data from the existing SQLite database
into a fresh PostgreSQL database.

This is a standalone tool for a human to run by hand, once, when actually
moving a specific client's data off SQLite (e.g. because they've outgrown
SQLite's single-writer lock with multiple concurrent admin users). It is
NOT run automatically by the app, NOT wired into run_web.py's startup path,
and has nothing to do with bober_bbq/migrate.py (that module only ever adds
missing *columns* to whatever database is already configured, on every
startup — it never moves data between databases).

It reuses the app's own SQLAlchemy models (bober_bbq.models, via
db.metadata) instead of hand-writing table definitions or introspecting
raw SQL — that way it automatically stays in sync with the schema, and
SQLAlchemy's per-column type handling (Boolean, Date, DateTime, ...)
converts values correctly between SQLite's and Postgres's native
representations for you.

Usage (run from the repo root):

    python scripts/migrate_sqlite_to_postgres.py \\
        --sqlite-path instance/bober_bbq.db \\
        --postgres-url postgresql+psycopg2://user:pass@localhost/bober_bbq

Both flags can instead come from env vars SQLITE_DB_PATH / POSTGRES_URL.
Requires psycopg2-binary (already in requirements.txt) and a target Postgres
database that already exists (e.g. `createdb bober_bbq`) — this script
creates the schema in it itself (a plain `db.create_all()`, safe to run
against an empty database).

IMPORTANT — safety / idempotency notes:
  - This script REFUSES to run if the target already has rows in any table,
    to avoid partial/duplicate data from an accidental second run. Pass
    --truncate to instead wipe the target's tables first and start clean —
    only ever do that against a scratch/throwaway Postgres database, never
    against one already serving real traffic.
  - It is therefore safe to *re-run* (each run either refuses, or with
    --truncate starts from a clean slate) but it never merges/upserts
    against a target that already has *different* data you want to keep.
    Treat it as "run once against a fresh Postgres DB", not as a sync tool.
  - After copying, it resets each table's Postgres auto-increment sequence
    to MAX(id)+1. This step matters: rows are copied with their original
    SQLite ids, so without a sequence reset, the very first new row the
    app inserts after cutover (e.g. the next order) would collide with a
    migrated row's id.

After migration: point DATABASE_URL at the Postgres URL in .env and restart
the app. bober_bbq/migrate.py's startup column-migrations are dialect-
agnostic and will run harmlessly against the already-current schema.
"""

from __future__ import annotations

import argparse
import os
import sys
from pathlib import Path

# Allow running as `python scripts/migrate_sqlite_to_postgres.py` from the
# repo root without installing the package.
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from sqlalchemy import Table, create_engine, func, insert, select, text  # noqa: E402
from sqlalchemy.engine import Engine  # noqa: E402

from bober_bbq import models  # noqa: E402,F401  (imported for side effect: registers all tables on db.metadata)
from bober_bbq.extensions import db  # noqa: E402


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "--sqlite-path",
        default=os.getenv("SQLITE_DB_PATH", ""),
        help="Path to the existing SQLite .db file (or env SQLITE_DB_PATH)",
    )
    parser.add_argument(
        "--postgres-url",
        default=os.getenv("POSTGRES_URL", ""),
        help="Target Postgres SQLAlchemy URL, e.g. "
        "postgresql+psycopg2://user:pass@localhost/bober_bbq (or env POSTGRES_URL)",
    )
    parser.add_argument(
        "--truncate",
        action="store_true",
        help="Delete any existing rows in the target's tables first, instead of refusing to "
        "run against a non-empty database. Only use this against a scratch/throwaway "
        "Postgres database you're prepared to wipe.",
    )
    parser.add_argument("--yes", action="store_true", help="Skip the interactive confirmation prompt")
    return parser.parse_args()


def _confirm(sqlite_path: Path, postgres_url: str) -> None:
    print(f"Source (SQLite):   {sqlite_path}")
    print(f"Target (Postgres): {postgres_url}")
    reply = input("Copy ALL data from source into target? [y/N] ").strip().lower()
    if reply != "y":
        print("Aborted.")
        sys.exit(1)


def _reject_nonempty_target(target: Engine, tables: list[Table], allow_truncate: bool) -> None:
    with target.connect() as conn:
        nonempty = []
        for table in tables:
            count = conn.execute(select(func.count()).select_from(table)).scalar_one()
            if count:
                nonempty.append((table.name, count))

    if not nonempty:
        return

    if not allow_truncate:
        print("Target database already has data:", file=sys.stderr)
        for name, count in nonempty:
            print(f"  - {name}: {count} row(s)", file=sys.stderr)
        print(
            "\nRefusing to copy into a non-empty database — this would violate primary-key/"
            "unique constraints and could mix in stale data. Re-run with --truncate to wipe "
            "these tables first (only against a scratch database!), or point --postgres-url "
            "at a fresh one.",
            file=sys.stderr,
        )
        sys.exit(1)

    print("--truncate given: wiping existing rows in the target's tables first...")
    with target.begin() as conn:
        # Children before parents, so FK constraints don't block the deletes.
        for table in reversed(tables):
            conn.execute(table.delete())


def _copy_all_rows(source: Engine, target: Engine, tables: list[Table]) -> int:
    total = 0
    with source.connect() as sconn, target.begin() as tconn:
        for table in tables:
            # Small-business scale data (single cafe) — loading a whole
            # table into memory at once is not a concern here.
            rows = [dict(row) for row in sconn.execute(select(table)).mappings().all()]
            if not rows:
                print(f"  {table.name}: 0 rows")
                continue
            tconn.execute(insert(table), rows)
            total += len(rows)
            print(f"  {table.name}: {len(rows)} row(s) copied")
    return total


def _reset_sequences(target: Engine, tables: list[Table]) -> None:
    """Rows were copied with their original SQLite ids, so any Postgres
    SERIAL/IDENTITY sequence backing a single-column integer primary key is
    still sitting at its start value. Advance each one to MAX(id) so the
    next app-generated insert doesn't collide with a migrated row. Tables
    with a natural (non-generated) primary key — Setting.key,
    TelegramUser.telegram_id — simply have no such sequence and are
    skipped automatically. No-ops entirely if --postgres-url didn't
    actually point at Postgres (pg_get_serial_sequence()/setval() are
    Postgres-only functions)."""
    if target.dialect.name != "postgresql":
        return
    with target.connect() as conn:
        for table in tables:
            pk_cols = list(table.primary_key.columns)
            if len(pk_cols) != 1 or not isinstance(pk_cols[0].type, (db.Integer, db.BigInteger)):
                continue
            pk = pk_cols[0]
            seq_name = conn.execute(
                text("SELECT pg_get_serial_sequence(:tbl, :col)"),
                {"tbl": table.name, "col": pk.name},
            ).scalar_one_or_none()
            if not seq_name:
                continue  # not a generated column (e.g. TelegramUser.telegram_id)
            max_id = conn.execute(select(func.max(pk))).scalar_one_or_none()
            conn.execute(
                text("SELECT setval(:seq, :val, :is_called)"),
                {"seq": seq_name, "val": max_id or 1, "is_called": max_id is not None},
            )
        conn.commit()


def main() -> None:
    args = parse_args()

    if not args.sqlite_path:
        print("Missing --sqlite-path (or env SQLITE_DB_PATH)", file=sys.stderr)
        sys.exit(2)
    if not args.postgres_url:
        print("Missing --postgres-url (or env POSTGRES_URL)", file=sys.stderr)
        sys.exit(2)
    if not args.postgres_url.startswith("postgresql"):
        print(f"--postgres-url doesn't look like a Postgres URL: {args.postgres_url}", file=sys.stderr)
        sys.exit(2)

    sqlite_path = Path(args.sqlite_path).resolve()
    if not sqlite_path.is_file():
        print(f"SQLite file not found: {sqlite_path}", file=sys.stderr)
        sys.exit(2)

    if not args.yes:
        _confirm(sqlite_path, args.postgres_url)

    source = create_engine(f"sqlite:///{sqlite_path}")
    target = create_engine(args.postgres_url)

    # sorted_tables is topologically sorted by foreign key so parents are
    # always copied (and, reversed, deleted) before/after their children —
    # exactly the FK-safe order this migration needs.
    tables = db.metadata.sorted_tables

    print("Creating schema on target (no-op for tables that already exist)...")
    db.metadata.create_all(target)

    _reject_nonempty_target(target, tables, args.truncate)

    print("Copying rows...")
    total_rows = _copy_all_rows(source, target, tables)

    print("Resetting Postgres auto-increment sequences...")
    _reset_sequences(target, tables)

    print(f"\nDone: {total_rows} row(s) copied across {len(tables)} tables.")
    print("Next: point DATABASE_URL at the Postgres URL in .env and restart the app.")


if __name__ == "__main__":
    main()