Migrate a SQL users table

CREATE TABLE newusers (
    id        BIGINT PRIMARY KEY,   -- same primary key as users.id
    usertoken UUID   NOT NULL UNIQUE
);

START_OFFSET=2000000 python migrate.py

import csv
import os
import psycopg2
import psycopg2.extras
from databunkerpro import DatabunkerproAPI

CHUNK = 4000
START_OFFSET = int(os.getenv("START_OFFSET", "0"))   # resume point after an interruption

api = DatabunkerproAPI(
    os.environ["DATABUNKER_API_URL"],       # e.g. https://databunker.internal.example.com
    os.environ["DATABUNKER_API_TOKEN"],
    os.getenv("DATABUNKER_TENANT_NAME", ""),
)
conn = psycopg2.connect(os.environ["APP_DATABASE_URL"])


def count_users():
    with conn.cursor() as cur:
        cur.execute("SELECT count(*) FROM users")
        return cur.fetchone()[0]


def fetch_chunk(offset):
    with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
        cur.execute(
            "SELECT * FROM users ORDER BY id LIMIT %s OFFSET %s",
            (CHUNK, offset),
        )
        return cur.fetchall()


def to_profile(row):
    """The whole row becomes the vault profile. 'custom' carries the legacy id."""
    profile = {k: v for k, v in row.items() if k != "id" and v is not None}
    profile["custom"] = str(row["id"])
    return profile


def save_tokens(created):
    pairs = [(int(item["profile"]["custom"]), item["token"]) for item in created]
    with conn.cursor() as cur:
        psycopg2.extras.execute_values(
            cur,
            "INSERT INTO newusers (id, usertoken) VALUES %s ON CONFLICT (id) DO NOTHING",
            pairs,
        )
    conn.commit()
    return len(pairs)


def main():
    total = count_users()
    migrated, failed = 0, 0
    print(f"{total} users to migrate, starting at offset {START_OFFSET}")

    # A fresh run starts the error log; a resumed run appends to it.
    with open("migration-errors.csv", "a" if START_OFFSET else "w", newline="") as fh:
        errors = csv.writer(fh)
        if not START_OFFSET:
            errors.writerow(["id", "reason"])

        for offset in range(START_OFFSET, total, CHUNK):
            # Printed before the work, so an interrupted run tells you where to resume:
            # re-run with START_OFFSET set to this value.
            print(f"offset {offset}/{total}", flush=True)
            rows = fetch_chunk(offset)

            result = api.create_users_bulk([{"profile": to_profile(r)} for r in rows])
            if result.get("status") != "ok":
                # Whole request rejected — nothing was inserted. Stop and investigate.
                raise SystemExit(f"Batch failed: {result.get('message')}")

            # 'created' echoes the submitted profile, so match on it rather than order.
            migrated += save_tokens(result.get("created", []))

            for item in result.get("errors", []):
                errors.writerow([item["profile"].get("custom"), item["reason"]])
                failed += 1

            print(f"  migrated={migrated} failed={failed}", flush=True)

    print(f"\nDone. Migrated {migrated}, failed {failed}.")
    if failed:
        print("Review migration-errors.csv before continuing.")


if __name__ == "__main__":
    main()

# reconcile.py — recover tokens for rows missing from newusers
with conn.cursor() as cur:
    cur.execute("SELECT id FROM users WHERE id NOT IN (SELECT id FROM newusers)")
    missing = [r[0] for r in cur.fetchall()]

unlock = api.bulk_list_unlock()["unlockuuid"]
found = api.bulk_list_users(
    unlock, [{"mode": "custom", "identity": str(uid)} for uid in missing]
)

pairs = [(int(r["profile"]["custom"]), r["token"]) for r in found.get("rows", [])]

if pairs:
    with conn.cursor() as cur:
        psycopg2.extras.execute_values(
            cur,
            "INSERT INTO newusers (id, usertoken) VALUES %s ON CONFLICT (id) DO NOTHING",
            pairs,
        )
    conn.commit()

print(f"{len(missing)} missing, {len(pairs)} recovered")

BEGIN;
ALTER TABLE users    RENAME TO users_old;
ALTER TABLE newusers RENAME TO users;
COMMIT;

DROP TABLE users_old;