"""migrate_sqlite_to_postgres.py -------------------------------- Safely migrate all data from SQLite to PostgreSQL without data loss. Supports both the current FastAPI schema (users, api_keys, documents, chat_messages, shared_messages) and the extended schema that includes workspaces, workspace_invitations, workspace_members, chat_sessions, and drive_connections. Also supports the older legacy ``instance/users.db`` schema where the users table is named ``user`` (singular). Dependencies ------------ pip install sqlalchemy psycopg[binary] # or: pip install sqlalchemy psycopg2-binary Tables migrated (FK-safe order) --------------------------------- users → api_keys → workspaces → workspace_invitations → workspace_members → chat_sessions → documents → chat_messages → drive_connections → shared_messages Usage ----- # Dry-run (reads SQLite, prints counts — no Postgres connection needed): python migrate_sqlite_to_postgres.py \\ --sqlite sqlite:///./data/app.db \\ --dry-run # Live migration (Postgres URL via CLI): python migrate_sqlite_to_postgres.py \\ --sqlite sqlite:///./data/app.db \\ --postgres postgresql://user:pass@localhost:5432/mydb # Live migration (Postgres URL via environment variable): export SUPABASE_DB_URL="postgres://user:pass@db.supabase.co:5432/postgres" python migrate_sqlite_to_postgres.py --sqlite sqlite:///./data/app.db # Migrate from the legacy instance/users.db path: python migrate_sqlite_to_postgres.py \\ --sqlite-path instance/users.db \\ --postgres postgresql://user:pass@localhost:5432/mydb # Migrate specific tables only: python migrate_sqlite_to_postgres.py \\ --sqlite sqlite:///./data/app.db \\ --postgres postgresql://user:pass@localhost:5432/mydb \\ --tables users documents chat_messages # Wipe Postgres tables before migrating (fresh start): python migrate_sqlite_to_postgres.py \\ --sqlite sqlite:///./data/app.db \\ --postgres postgresql://user:pass@localhost:5432/mydb \\ --truncate # Verbose / debug logging: python migrate_sqlite_to_postgres.py \\ --sqlite sqlite:///./data/app.db \\ --postgres postgresql://user:pass@localhost:5432/mydb \\ --verbose Resolves: #279 """ from __future__ import annotations import argparse import logging import os import sys import uuid from datetime import datetime, timezone from pathlib import Path from typing import Any, Dict, List, Optional from sqlalchemy import create_engine, inspect, text from sqlalchemy.exc import OperationalError # --------------------------------------------------------------------------- # Logging # --------------------------------------------------------------------------- logging.basicConfig( level=logging.INFO, format="%(asctime)s %(levelname)-8s %(message)s", datefmt="%Y-%m-%d %H:%M:%S", ) logger = logging.getLogger("migration") # --------------------------------------------------------------------------- # Migration order — respects FK dependencies # --------------------------------------------------------------------------- TABLE_ORDER = [ "users", "api_keys", "workspaces", "workspace_invitations", "workspace_members", "chat_sessions", "documents", "chat_messages", "drive_connections", "shared_messages", ] # Columns that hold UUID / GUID values and need string normalisation _UUID_COLUMNS = { "id", "user_id", "document_id", "session_id", "message_id", "workspace_id", "inviter_id", "created_by", } # Boolean columns stored as 0/1 integers in SQLite _BOOL_COLUMNS = { "is_admin", "is_verified", "is_active", "is_deleted", "enabled", } # --------------------------------------------------------------------------- # URL helpers # --------------------------------------------------------------------------- def normalize_postgres_url(url: str) -> str: """ Normalise common Postgres URL forms to use the psycopg (v3) driver. Supabase and many hosting providers hand out ``postgres://`` or plain ``postgresql://`` connection strings. SQLAlchemy requires a driver specifier such as ``postgresql+psycopg://`` for psycopg v3, or ``postgresql+psycopg2://`` for psycopg2. We prefer psycopg v3 when no driver is already specified. """ if url.startswith("postgres://"): url = "postgresql+psycopg://" + url[len("postgres://"):] elif url.startswith("postgresql://"): url = "postgresql+psycopg://" + url[len("postgresql://"):] # Already has a driver specifier (e.g. postgresql+psycopg2://) — leave it. return url def sqlite_url_from_path(path: str) -> str: return f"sqlite:///{Path(path).resolve().as_posix()}" # --------------------------------------------------------------------------- # Engine factory # --------------------------------------------------------------------------- def _make_engine(url: str, label: str): """Create a SQLAlchemy engine with sensible defaults.""" is_sqlite = url.startswith("sqlite") kwargs: Dict[str, Any] = {"echo": False, "future": True} if is_sqlite: kwargs["connect_args"] = {"check_same_thread": False} else: kwargs.update(pool_pre_ping=True, pool_size=5, max_overflow=10) logger.info("Connecting to %s: %s", label, url) return create_engine(url, **kwargs) # --------------------------------------------------------------------------- # Row coercion # --------------------------------------------------------------------------- def _normalise_uuid(value: Any) -> Optional[str]: """Return a UUID as a plain lowercase string, or None.""" if value is None: return None if isinstance(value, uuid.UUID): return str(value) try: return str(uuid.UUID(str(value))) except (ValueError, AttributeError): return str(value) def _coerce_row(row: Dict[str, Any]) -> Dict[str, Any]: """ Coerce a SQLite row dict so it is safe to insert into Postgres: * UUID columns → normalised lowercase str * Boolean integers → Python bool (SQLite stores True/False as 1/0) * Naive datetimes → UTC-aware datetimes """ result: Dict[str, Any] = {} for col, val in row.items(): if col in _UUID_COLUMNS: result[col] = _normalise_uuid(val) elif col in _BOOL_COLUMNS: result[col] = bool(val) if val is not None else None elif isinstance(val, datetime) and val.tzinfo is None: result[col] = val.replace(tzinfo=timezone.utc) else: result[col] = val return result # --------------------------------------------------------------------------- # Table helpers # --------------------------------------------------------------------------- def _table_exists(inspector, table: str) -> bool: return table in inspector.get_table_names() def _resolve_users_table(src_inspector) -> Optional[str]: """ Return the name of the users table in the source database. Supports both the current schema (``users``) and the legacy schema (``user``, singular) used by older ``instance/users.db`` databases. Returns None if neither is found. """ if _table_exists(src_inspector, "users"): return "users" if _table_exists(src_inspector, "user"): logger.info("Legacy 'user' table detected — will migrate as 'users'.") return "user" return None # --------------------------------------------------------------------------- # Per-table migration # --------------------------------------------------------------------------- def _dry_run_table(table: str, src_conn) -> Dict[str, int]: """Read SQLite row counts and log them — no Postgres connection used.""" src_inspector = inspect(src_conn) if not _table_exists(src_inspector, table): logger.warning(" [%s] not found in SQLite — skipping", table) return {"would_migrate": 0, "skipped": 0} count = src_conn.execute(text(f"SELECT COUNT(*) FROM \"{table}\"")).scalar() or 0 logger.info(" [%s] %d rows found (dry-run — nothing written)", table, count) return {"would_migrate": count, "skipped": 0} def _live_migrate_table( table: str, src_table_name: str, src_conn, dst_conn, truncate: bool, batch_size: int, ) -> Dict[str, int]: """ Migrate one table from SQLite → Postgres. ``table`` is the canonical destination table name (always the current schema name, e.g. ``"users"``). ``src_table_name`` may differ for legacy sources (e.g. ``"user"``). """ stats = {"inserted": 0, "skipped": 0} src_inspector = inspect(src_conn) dst_inspector = inspect(dst_conn) if not _table_exists(src_inspector, src_table_name): logger.warning(" [%s] not found in SQLite — skipping", table) return stats if not _table_exists(dst_inspector, table): logger.warning( " [%s] not found in Postgres — skipping. " "Run your init_postgres.sql first to create the schema.", table, ) return stats # Fetch all rows from source rows_result = src_conn.execute(text(f"SELECT * FROM \"{src_table_name}\"")) column_names: List[str] = list(rows_result.keys()) raw_rows = rows_result.fetchall() if not raw_rows: logger.info(" [%s] 0 rows — nothing to migrate", table) return stats logger.info(" [%s] %d rows to migrate", table, len(raw_rows)) if truncate: logger.info(" [%s] truncating Postgres table (CASCADE)", table) dst_conn.execute(text(f'TRUNCATE TABLE "{table}" RESTART IDENTITY CASCADE')) # Build idempotent INSERT — skips duplicates silently col_list = ", ".join(f'"{c}"' for c in column_names) placeholders = ", ".join(f":{c}" for c in column_names) insert_sql = text( f'INSERT INTO "{table}" ({col_list}) ' f"VALUES ({placeholders}) " f"ON CONFLICT DO NOTHING" ) batch: List[Dict[str, Any]] = [] for raw in raw_rows: row_dict = dict(zip(column_names, raw)) batch.append(_coerce_row(row_dict)) if len(batch) >= batch_size: result = dst_conn.execute(insert_sql, batch) stats["inserted"] += result.rowcount if result.rowcount >= 0 else len(batch) batch = [] if batch: result = dst_conn.execute(insert_sql, batch) stats["inserted"] += result.rowcount if result.rowcount >= 0 else len(batch) logger.info(" [%s] inserted %d rows", table, stats["inserted"]) return stats # --------------------------------------------------------------------------- # Orchestration # --------------------------------------------------------------------------- def run_migration( sqlite_url: str, postgres_url: Optional[str], dry_run: bool, truncate: bool, tables: Optional[List[str]], batch_size: int, ) -> bool: """ Orchestrate the full migration. Returns True on success. In dry-run mode, no Postgres connection is opened; only SQLite row counts are reported. """ # Build the ordered table list target = set(tables) if tables else set(TABLE_ORDER) ordered = [t for t in TABLE_ORDER if t in target] # Append any user-supplied extras that aren't in TABLE_ORDER for t in (tables or []): if t not in ordered: logger.warning( "Table '%s' is not in the known FK-safe order. " "It will be migrated last — ensure FK dependencies are satisfied.", t, ) ordered.append(t) mode = "DRY RUN" if dry_run else "LIVE" logger.info("=" * 60) logger.info("Migration mode : %s", mode) logger.info("Source : %s", sqlite_url) if not dry_run: logger.info("Destination : %s", postgres_url) logger.info("Tables : %s", ", ".join(ordered)) logger.info("Truncate first : %s", truncate and not dry_run) logger.info("Batch size : %d", batch_size) logger.info("=" * 60) src_engine = _make_engine(sqlite_url, "SQLite") if dry_run: overall_stats: Dict[str, Dict[str, int]] = {} with src_engine.connect() as src_conn: src_inspector = inspect(src_conn) users_src = _resolve_users_table(src_inspector) for table in ordered: logger.info("Checking table : %s", table) # For the legacy 'user' → 'users' mapping src_table_name = users_src if table == "users" and users_src else table if src_table_name is None: logger.error("No users/user table found in SQLite.") return False overall_stats[table] = _dry_run_table(src_table_name, src_conn) else: if not postgres_url: logger.error( "A Postgres URL is required for a live migration. " "Pass --postgres or set SUPABASE_DB_URL / DATABASE_URL." ) return False pg_url = normalize_postgres_url(postgres_url) dst_engine = _make_engine(pg_url, "PostgreSQL") # Verify Postgres connectivity before touching data try: with dst_engine.connect() as probe: probe.execute(text("SELECT 1")) except OperationalError as exc: logger.error("Cannot connect to Postgres: %s", exc) return False overall_stats = {} # Single Postgres transaction — full rollback on any error with src_engine.connect() as src_conn, dst_engine.begin() as dst_conn: src_inspector = inspect(src_conn) users_src = _resolve_users_table(src_inspector) for table in ordered: logger.info("Migrating table: %s", table) src_table_name = users_src if table == "users" and users_src else table if table == "users" and src_table_name is None: logger.error("No users/user table found in SQLite — aborting.") raise RuntimeError("Missing users table in source database.") try: overall_stats[table] = _live_migrate_table( table=table, src_table_name=src_table_name or table, src_conn=src_conn, dst_conn=dst_conn, truncate=truncate, batch_size=batch_size, ) except Exception as exc: logger.error( " [%s] FAILED — rolling back all changes: %s", table, exc, exc_info=True, ) raise # triggers dst_engine.begin() rollback # Print summary logger.info("=" * 60) logger.info("Summary") logger.info("-" * 60) total_migrated = 0 for table, stats in overall_stats.items(): if dry_run: logger.info( " %-30s would_migrate=%-6d skipped=%d", table, stats.get("would_migrate", 0), stats.get("skipped", 0), ) total_migrated += stats.get("would_migrate", 0) else: logger.info( " %-30s inserted=%-6d skipped=%d", table, stats.get("inserted", 0), stats.get("skipped", 0), ) total_migrated += stats.get("inserted", 0) logger.info("-" * 60) if dry_run: logger.info(" TOTAL would_migrate=%d", total_migrated) logger.info("Dry run complete — no data was written to Postgres.") else: logger.info(" TOTAL inserted=%d", total_migrated) logger.info("Migration complete.") logger.info("=" * 60) return True # --------------------------------------------------------------------------- # CLI # --------------------------------------------------------------------------- def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Migrate data from SQLite to PostgreSQL.", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=__doc__, ) src_group = parser.add_mutually_exclusive_group() src_group.add_argument( "--sqlite", metavar="URL", help="SQLAlchemy URL for the SQLite source, e.g. sqlite:///./data/app.db", ) src_group.add_argument( "--sqlite-path", metavar="PATH", default="instance/users.db", help="Path to the SQLite file (alternative to --sqlite). " "Defaults to instance/users.db.", ) parser.add_argument( "--postgres", metavar="URL", default=( os.getenv("SUPABASE_DB_URL") or os.getenv("POSTGRES_DATABASE_URL") or os.getenv("DATABASE_URL") ), help="Postgres destination URL (postgresql:// or postgres://). " "Also read from SUPABASE_DB_URL, POSTGRES_DATABASE_URL, or DATABASE_URL. " "Not required for --dry-run.", ) parser.add_argument( "--dry-run", action="store_true", default=False, help="Read from SQLite and report row counts — no Postgres connection made.", ) parser.add_argument( "--truncate", action="store_true", default=False, help="TRUNCATE each Postgres table (CASCADE) before inserting. " "Useful for a clean re-migration. Ignored during --dry-run.", ) parser.add_argument( "--tables", nargs="+", metavar="TABLE", help="Migrate only these tables (space-separated). " "Defaults to all tables in FK-safe order.", ) parser.add_argument( "--batch-size", type=int, default=500, metavar="N", help="Number of rows per INSERT batch (default: 500).", ) parser.add_argument( "--verbose", action="store_true", help="Enable DEBUG logging.", ) return parser.parse_args() def main() -> None: args = parse_args() if args.verbose: logging.getLogger().setLevel(logging.DEBUG) logger.setLevel(logging.DEBUG) # Resolve SQLite URL if args.sqlite: sqlite_url = args.sqlite if not sqlite_url.startswith("sqlite"): logger.error("--sqlite must be a sqlite:/// URL, got: %s", sqlite_url) sys.exit(1) else: sqlite_url = sqlite_url_from_path(args.sqlite_path) # Postgres URL is optional for dry-run only postgres_url: Optional[str] = args.postgres if not args.dry_run: if not postgres_url: logger.error( "Provide a Postgres URL via --postgres, SUPABASE_DB_URL, " "POSTGRES_DATABASE_URL, or DATABASE_URL." ) sys.exit(2) if postgres_url.startswith("sqlite"): logger.error( "--postgres must be a Postgres URL, not a SQLite URL: %s", postgres_url, ) sys.exit(2) try: success = run_migration( sqlite_url=sqlite_url, postgres_url=postgres_url, dry_run=args.dry_run, truncate=args.truncate, tables=args.tables, batch_size=args.batch_size, ) except Exception as exc: logger.error("Migration failed: %s", exc) sys.exit(1) sys.exit(0 if success else 1) if __name__ == "__main__": main()