petal/scripts/migrate_local_user.py

#!/usr/bin/env python3
"""Move the pre-auth `local` user's writing onto a real account.

Petal ran as a single hardcoded user (`users.id = 'local'`) for its whole life
before sign-in existed. Everything she has written is owned by that row. This
script re-points it at the account she now signs in as.

Why a script and not a startup migration: the destination is an OIDC subject
id, which is not knowable from inside the app — it belongs to the identity
provider. Running it deliberately, with the app stopped and a backup taken, is
also the only way to be sure nothing is writing to the database halfway through.

What moves: `documents`, `tags`, `vocab_words` and `images` carry `user_id`
directly. `document_versions`, `suggestions` and `document_tags` hang off their
parents and follow without being touched — which is exactly why this must be one
transaction with foreign keys off: re-pointing a parent while its children are
enforced would either fail or cascade.

Sessions belonging to the old identity are deleted rather than moved. A session
is proof that *someone signed in*, and nobody ever signed in as `local`.

Safety: dry-run unless --apply; refuses to run while anything else has the
database open; refuses if the destination already owns writing of its own; takes
a `VACUUM INTO` backup first; and verifies every row it expected to move
actually moved before it commits.

Usage:
    # look at what would happen
    python3 scripts/migrate_local_user.py data/petal.db --to <oidc-sub>

    # do it
    python3 scripts/migrate_local_user.py data/petal.db --to <oidc-sub> \\
        --email her@example.com --name "Her Name" --apply

The subject id comes from the identity provider. For authentik with the default
`hashed_user_id` sub mode it is the user's `uid`, which is stable and knowable
before she has ever logged in:

    docker exec authentik-server-1 ak shell -c \\
        "from authentik.core.models import User; print(User.objects.get(username='claire').uid)"
"""
import argparse
import os
import sqlite3
import sys
import time

# The tables that name an owner directly. Everything else in the schema reaches
# its owner through one of these.
OWNED_TABLES = ("documents", "tags", "vocab_words", "images")

LOCAL_USER = "local"


def die(msg: str) -> None:
    print(f"error: {msg}", file=sys.stderr)
    sys.exit(1)


def open_db(path: str) -> sqlite3.Connection:
    if not os.path.exists(path):
        die(f"no database at {path}")
    conn = sqlite3.connect(path, isolation_level=None)
    conn.row_factory = sqlite3.Row
    return conn


def table_exists(conn: sqlite3.Connection, name: str) -> bool:
    row = conn.execute(
        "SELECT 1 FROM sqlite_master WHERE type='table' AND name=?", (name,)
    ).fetchone()
    return row is not None


def counts(conn: sqlite3.Connection, user_id: str) -> dict[str, int]:
    out = {}
    for table in OWNED_TABLES:
        if not table_exists(conn, table):
            continue
        out[table] = conn.execute(
            f"SELECT COUNT(*) FROM {table} WHERE user_id = ?", (user_id,)
        ).fetchone()[0]
    return out


def require_app_stopped(path: str) -> None:
    """Fail unless nothing else has the database open.

    `BEGIN EXCLUSIVE` is not enough, and quietly so: in WAL mode it only
    conflicts with another *writer*, so an idle-but-running Petal sails straight
    past it — which is precisely the case this guard exists to catch. Taking
    `locking_mode = EXCLUSIVE` conflicts with any other connection at all,
    because it locks the shared-memory index every WAL reader must map.

    Checking for a `-wal` file would be no use either: WAL mode leaves one
    behind whether or not anything is running.

    This runs on its own connection, which is then closed. Setting locking_mode
    back to NORMAL does not release the lock straight away — SQLite drops it on
    the connection's next database access, which on a WAL database can leave the
    file locked against everything else this script is about to do. Closing is
    the only unambiguous way to let go of it.
    """
    probe = sqlite3.connect(path, isolation_level=None)
    try:
        probe.execute("PRAGMA locking_mode = EXCLUSIVE")
        probe.execute("BEGIN IMMEDIATE")
        probe.execute("COMMIT")
    except sqlite3.OperationalError as err:
        die(
            f"{path} is in use ({err}). Stop Petal first:\n"
            "    docker compose stop petal      # VPS\n"
            "    systemctl --user stop petal    # millenia"
        )
    finally:
        probe.close()


def backup(path: str) -> str:
    """Snapshot the database with VACUUM INTO — one coherent file including the
    WAL, taken without a write lock, and it refuses to overwrite."""
    dest = f"{path}.pre-migrate-{time.strftime('%Y%m%d-%H%M%S')}"
    conn = sqlite3.connect(path)
    try:
        conn.execute("VACUUM INTO ?", (dest,))
    finally:
        conn.close()
    return dest


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("database", help="path to petal.db")
    ap.add_argument("--to", required=True, metavar="SUB",
                    help="OIDC subject id of the destination account")
    ap.add_argument("--from", dest="source", default=LOCAL_USER,
                    help=f"account to move from (default: {LOCAL_USER})")
    ap.add_argument("--email", default="", help="email for the destination account, if it doesn't exist yet")
    ap.add_argument("--name", default="", help="display name for the destination account")
    ap.add_argument("--pair-lang", default="zh", help="language pair for a newly created account (default: zh)")
    ap.add_argument("--apply", action="store_true", help="actually commit (default: dry run)")
    ap.add_argument("--no-backup", action="store_true", help="skip the pre-migration snapshot")
    args = ap.parse_args()

    if args.to == args.source:
        die("source and destination are the same account")

    require_app_stopped(args.database)
    conn = open_db(args.database)

    src_counts = counts(conn, args.source)
    if not any(src_counts.values()):
        die(f"account {args.source!r} owns nothing in this database — wrong file, or already migrated?")

    dst_counts = counts(conn, args.to)
    if any(dst_counts.values()):
        die(
            f"account {args.to!r} already owns writing here "
            f"({', '.join(f'{k}={v}' for k, v in dst_counts.items() if v)}). "
            "Refusing to merge two accounts — that is not something this script can undo."
        )

    dst = conn.execute("SELECT id, email, display_name FROM users WHERE id = ?", (args.to,)).fetchone()

    print(f"database:    {args.database}")
    print(f"moving from: {args.source}")
    print(f"moving to:   {args.to}" + ("" if dst else "  (will be created)"))
    for table, n in src_counts.items():
        print(f"  {table:<14} {n}")
    sessions = 0
    if table_exists(conn, "sessions"):
        sessions = conn.execute(
            "SELECT COUNT(*) FROM sessions WHERE user_id = ?", (args.source,)
        ).fetchone()[0]
        print(f"  {'sessions':<14} {sessions}  (deleted, not moved)")

    if not args.apply:
        print("\ndry run — nothing changed. Re-run with --apply to commit.")
        return

    if not args.no_backup:
        dest = backup(args.database)
        print(f"\nbackup:      {dest}")

    # Foreign keys off for the duration: children reference the parents being
    # re-pointed, and this is one atomic swap of an identity, not a data change.
    conn.execute("PRAGMA foreign_keys = OFF")
    conn.execute("BEGIN EXCLUSIVE")
    try:
        if not dst:
            conn.execute(
                "INSERT INTO users (id, email, display_name, pair_lang) VALUES (?, ?, ?, ?)",
                (args.to, args.email, args.name or args.email, args.pair_lang),
            )

        for table in src_counts:
            conn.execute(
                f"UPDATE {table} SET user_id = ? WHERE user_id = ?", (args.to, args.source)
            )
        if table_exists(conn, "sessions"):
            conn.execute("DELETE FROM sessions WHERE user_id = ?", (args.source,))

        # Verify before committing: every row that was the source's is now the
        # destination's, and the source owns nothing.
        moved = counts(conn, args.to)
        left = counts(conn, args.source)
        if moved != src_counts or any(left.values()):
            raise RuntimeError(
                f"row counts do not match after the move (expected {src_counts}, "
                f"got {moved}, source still holds {left})"
            )

        conn.execute("DELETE FROM users WHERE id = ?", (args.source,))
        conn.execute("COMMIT")
    except Exception as err:  # noqa: BLE001 — any failure must roll the whole thing back
        conn.execute("ROLLBACK")
        die(f"migration rolled back: {err}")
    finally:
        conn.execute("PRAGMA foreign_keys = ON")

    print("\nmigrated. Start Petal and sign in as the destination account.")
    for table, n in src_counts.items():
        print(f"  {table:<14} {n}")


if __name__ == "__main__":
    main()