atuin-sqlite-to-postgresql.py

#!/usr/bin/env -S uv run --script
# /// script
# dependencies = ["psycopg2-binary"]
# ///
"""Migrate atuin server.db (SQLite) -> PostgreSQL."""

import sqlite3
import sys
import uuid
import psycopg2
import psycopg2.extras

SQLITE_PATH = "~/.local/share/atuin/server.db"
PG_DSN = "postgres://atuin:zepass@localhost:5432/atuin"
BATCH_SIZE = 5000


def to_uuid(val):
    if val is None:
        return None
    if isinstance(val, bytes):
        return str(uuid.UUID(bytes=val))
    return str(uuid.UUID(str(val)))


def migrate_users(sq, pg):
    cur = sq.execute("SELECT id, username, email, password, created_at, verified_at FROM users")
    rows = cur.fetchall()
    if not rows:
        print("users: 0 rows, skipping")
        return
    with pg.cursor() as pgc:
        psycopg2.extras.execute_values(
            pgc,
            """
            INSERT INTO users (id, username, email, password, created_at, verified_at)
            VALUES %s
            ON CONFLICT DO NOTHING
            """,
            rows,
        )
    pg.commit()
    print(f"users: migrated {len(rows)} row(s)")


def migrate_sessions(sq, pg):
    cur = sq.execute("SELECT id, user_id, token FROM sessions")
    rows = cur.fetchall()
    if not rows:
        print("sessions: 0 rows, skipping")
        return
    with pg.cursor() as pgc:
        psycopg2.extras.execute_values(
            pgc,
            "INSERT INTO sessions (id, user_id, token) VALUES %s ON CONFLICT DO NOTHING",
            rows,
        )
    pg.commit()
    print(f"sessions: migrated {len(rows)} row(s)")


def migrate_store(sq, pg):
    total = sq.execute("SELECT COUNT(*) FROM store").fetchone()[0]
    print(f"store: {total} rows to migrate...")

    offset = 0
    migrated = 0
    with pg.cursor() as pgc:
        while True:
            raw = sq.execute(
                "SELECT id, client_id, host, idx, timestamp, version, tag, data, cek, user_id, created_at "
                f"FROM store LIMIT {BATCH_SIZE} OFFSET {offset}"
            ).fetchall()
            if not raw:
                break
            rows = [
                (to_uuid(r[0]), to_uuid(r[1]), to_uuid(r[2]), r[3], r[4], r[5], r[6], r[7], r[8], r[9], r[10])
                for r in raw
            ]
            psycopg2.extras.execute_values(
                pgc,
                """
                INSERT INTO store (id, client_id, host, idx, timestamp, version, tag, data, cek, user_id, created_at)
                VALUES %s
                ON CONFLICT DO NOTHING
                """,
                rows,
            )
            pg.commit()
            migrated += len(rows)
            offset += BATCH_SIZE
            print(f"  store: {migrated}/{total}", end="\r", flush=True)

    print(f"\nstore: migrated {migrated} row(s)")


def main():
    print(f"Connecting to SQLite: {SQLITE_PATH}")
    sq = sqlite3.connect(SQLITE_PATH)

    print(f"Connecting to PostgreSQL...")
    pg = psycopg2.connect(PG_DSN)

    migrate_users(sq, pg)
    migrate_sessions(sq, pg)
    migrate_store(sq, pg)

    sq.close()
    pg.close()
    print("Done.")


if __name__ == "__main__":
    main()