heygo/scripts/export-django-links.py

#!/usr/bin/env python3
"""
Standalone exporter for the old Django "Links" shortlinks database.

Reads a Django SQLite DB (the `links` app's `Link` model, plus `auth_user`
for owner emails and `django_content_type` for table discovery) and writes
one JSON object per line (JSONL) with the canonical shortlink fields:

    alias, link_type, target_url, content_markdown, description,
    tags, created_at, updated_at

This script is intentionally standalone: it imports nothing from the old
Django project and makes no modifications to it. It only opens the SQLite
file read-only and runs SELECTs.

Field mapping (Django Link model -> export):
    alias              -> alias
    link_type (LINK/CUSTOM) -> link_type ("redirect" / "custom")
    original_url       -> target_url
    text               -> content_markdown
    description        -> description
    tags (M2M -> Tag)  -> tags  (list of tag names)
    created_at         -> created_at
    updated_at         -> updated_at

Owner handling:
    The current Django Link model has no owner/user FK. To stay forward-
    compatible with older dumps that did carry an owner column, the script
    introspects the `links_link` table for an owner-like FK
    (`owner_id` / `user_id` / `created_by_id`) and, when present, joins it to
    `auth_user` so `--owner-email` can filter. When no such column exists,
    `--owner-email` prints a warning to stderr and exports all links.
"""
from __future__ import annotations

import argparse
import json
import os
import sqlite3
import sys
from typing import Any, Iterable

# Django LinkType choices -> export vocabulary.
LINK_TYPE_MAP = {"LINK": "redirect", "CUSTOM": "custom"}
DEFAULT_LINK_TYPE = "redirect"

# Columns on `links_link` that, if present, are treated as an owner FK into
# auth_user.id (checked in this order; first match wins).
OWNER_COLUMNS = ("owner_id", "user_id", "created_by_id")

# Columns we expect to read from the link table. Names match the Django
# model; we tolerate missing columns by falling back to None.
LINK_COLUMNS = (
    "id",
    "alias",
    "link_type",
    "original_url",
    "text",
    "description",
    "created_at",
    "updated_at",
)


def _stderr(msg: str) -> None:
    print(msg, file=sys.stderr)


def resolve_db_path(arg_db: str | None) -> str:
    """Resolve the DB path from --db, then $DJANGO_DB_PATH. Exit on missing."""
    db_path = arg_db or os.environ.get("DJANGO_DB_PATH")
    if not db_path:
        _stderr("error: no DB path provided. Use --db or set DJANGO_DB_PATH.")
        sys.exit(2)
    if not os.path.exists(db_path):
        _stderr(f"error: DB file not found: {db_path}")
        sys.exit(2)
    if os.path.isdir(db_path):
        _stderr(f"error: DB path is a directory, not a file: {db_path}")
        sys.exit(2)
    return db_path


def open_db(db_path: str) -> sqlite3.Connection:
    """Open the SQLite DB read-only so we can never mutate the old DB."""
    uri = f"file:{db_path}?mode=ro"
    try:
        conn = sqlite3.connect(uri, uri=True)
    except sqlite3.OperationalError:
        # Very old / locked DBs may reject read-only URI; fall back to a
        # plain read-only connection via immutability, never writing.
        conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA query_only = 1")  # belt-and-suspenders: no writes
    return conn


def table_columns(conn: sqlite3.Connection, table: str) -> list[str]:
    rows = conn.execute(f'PRAGMA table_info("{table}")').fetchall()
    return [r["name"] for r in rows]


def discover_link_table(conn: sqlite3.Connection) -> str:
    """Find the links_link table via django_content_type, else a schema scan."""
    try:
        ct = conn.execute(
            "SELECT app_label, model FROM django_content_type "
            "WHERE app_label = ? AND model = ?",
            ("links", "link"),
        ).fetchone()
    except sqlite3.OperationalError:
        ct = None
    if ct:
        return f"{ct['app_label']}_{ct['model']}"  # => "links_link"
    # Fallback: any table containing a unique-ish "alias" column.
    tables = conn.execute(
        "SELECT name FROM sqlite_master WHERE type='table'"
    ).fetchall()
    for t in tables:
        cols = table_columns(conn, t["name"])
        if "alias" in cols and "original_url" in cols:
            return t["name"]
    raise LookupError("could not find the links_link table in this DB")


def discover_tag_table(conn: sqlite3.Connection) -> str | None:
    """Find the links_tag table (tag names), if it exists."""
    try:
        ct = conn.execute(
            "SELECT app_label, model FROM django_content_type "
            "WHERE app_label = ? AND model = ?",
            ("links", "tag"),
        ).fetchone()
    except sqlite3.OperationalError:
        ct = None
    if ct:
        return f"{ct['app_label']}_{ct['model']}".replace(" ", "_")
    # Fallback: a table with a `name` column in the links app namespace.
    tables = conn.execute(
        "SELECT name FROM sqlite_master WHERE type='table' AND name LIKE 'links_%'"
    ).fetchall()
    for t in tables:
        cols = table_columns(conn, t["name"])
        if "name" in cols and "id" in cols and "alias" not in cols:
            return t["name"]
    return None


def discover_m2m_table(conn: sqlite3.Connection, link_table: str, tag_table: str) -> str | None:
    """Find the links_link -> links_tag M2M through table, if any."""
    # Django default name: links_link_tags
    candidates = conn.execute(
        "SELECT name FROM sqlite_master WHERE type='table' "
        "AND (name LIKE ? OR name LIKE ?)",
        (f"{link_table}_%", "%_tags"),
    ).fetchall()
    for t in candidates:
        name = t["name"]
        cols = table_columns(conn, name)
        # An M2M through table has FKs to both sides. Django names them
        # <link_table> (or link) and <tag_table> (or tag).
        has_link_fk = any(c in cols for c in ("link_id", link_table + "_id"))
        has_tag_fk = any(c in cols for c in ("tag_id", tag_table + "_id"))
        if has_link_fk and has_tag_fk:
            return name
    return None


def load_users(conn: sqlite3.Connection) -> dict[int, str]:
    """Map auth_user.id -> email (lowercased) for owner filtering."""
    users: dict[int, str] = {}
    try:
        rows = conn.execute("SELECT id, email FROM auth_user").fetchall()
    except sqlite3.OperationalError:
        try:
            rows = conn.execute("SELECT id, username FROM auth_user").fetchall()
            # treat username as the "email"-ish identifier when email missing
        except sqlite3.OperationalError:
            return users
    for r in rows:
        users[r["id"]] = (r["email"] if "email" in r.keys() else r["username"]).lower()
    return users


def detect_owner_column(conn: sqlite3.Connection, link_table: str) -> str | None:
    """Return the owner FK column name on the link table, if any."""
    cols = table_columns(conn, link_table)
    for c in OWNER_COLUMNS:
        if c in cols:
            return c
    return None


def link_table_has_link_id_column(conn: sqlite3.Connection, m2m_table: str, link_table: str) -> str:
    """Pick the FK column on the M2M table that points at the link row."""
    cols = table_columns(conn, m2m_table)
    for cand in (link_table + "_id", "link_id"):
        if cand in cols:
            return cand
    # last resort: the first <something>_id column
    for c in cols:
        if c.endswith("_id") and c != "tag_id" and not c.endswith("tag_id"):
            return c
    raise LookupError(f"could not find link FK on M2M table {m2m_table}")


def query_links(
    conn: sqlite3.Connection,
    link_table: str,
    tag_table: str | None,
    m2m_table: str | None,
    owner_col: str | None,
    owner_email: str | None,
    users: dict[int, str],
) -> Iterable[sqlite3.Row]:
    """Build and run the SELECT, applying the owner filter when possible."""
    if owner_email and owner_col is None:
        _stderr(
            "warning: --owner-email given but the links_link table has no "
            "owner/user FK column; owner filter cannot be applied. "
            "Exporting all links."
        )

    # Assemble column list defensively (some columns may not exist in older
    # DBs); we read only the ones present.
    cols = table_columns(conn, link_table)
    select_cols = [c for c in LINK_COLUMNS if c in cols]
    if owner_col and owner_col in cols and owner_col not in select_cols:
        select_cols.append(owner_col)

    base = f'SELECT {", ".join(select_cols)} FROM "{link_table}"'
    where_parts: list[str] = []
    params: list[Any] = []

    # Validate owner email against known users when an owner column exists.
    if owner_email and owner_col is not None:
        owner_email_lc = owner_email.strip().lower()
        if not users:
            _stderr(
                "warning: --owner-email given but auth_user table is empty "
                "or missing; no links will match the owner filter."
            )
            where_parts.append("1=0")
        else:
            owner_user_ids = [uid for uid, em in users.items() if em == owner_email_lc]
            if not owner_user_ids:
                _stderr(
                    f"warning: no auth_user row matches --owner-email "
                    f"{owner_email!r}; export will be empty."
                )
                where_parts.append("1=0")
            else:
                placeholders = ",".join("?" for _ in owner_user_ids)
                where_parts.append(f"{owner_col} IN ({placeholders})")
                params.extend(owner_user_ids)

    sql = base
    if where_parts:
        sql += " WHERE " + " AND ".join(where_parts)
    sql += " ORDER BY alias ASC"
    return conn.execute(sql, params)


def collect_tags(
    conn: sqlite3.Connection,
    link_id: int,
    tag_table: str | None,
    m2m_table: str | None,
    link_fk_col: str,
) -> list[str]:
    if not tag_table or not m2m_table:
        return []
    rows = conn.execute(
        f'SELECT t."name" FROM "{tag_table}" t '
        f'JOIN "{m2m_table}" m ON m."tag_id" = t."id" '
        f'WHERE m."{link_fk_col}" = ? ORDER BY t."name" ASC',
        (link_id,),
    ).fetchall()
    return [r["name"] for r in rows]


def row_to_record(
    conn: sqlite3.Connection,
    row: sqlite3.Row,
    tag_table: str | None,
    m2m_table: str | None,
    link_fk_col: str,
) -> dict[str, Any]:
    link_id = row["id"] if "id" in row.keys() else None

    def g(col: str) -> Any:
        return row[col] if col in row.keys() else None

    raw_link_type = g("link_type")
    if raw_link_type in LINK_TYPE_MAP:
        link_type = LINK_TYPE_MAP[raw_link_type]
    elif not raw_link_type:
        link_type = DEFAULT_LINK_TYPE
    else:
        link_type = raw_link_type

    tags = collect_tags(conn, link_id, tag_table, m2m_table, link_fk_col) if link_id is not None else []

    return {
        "alias": g("alias"),
        "link_type": link_type,
        "target_url": g("original_url"),
        "content_markdown": g("text"),
        "description": g("description"),
        "tags": tags,
        "created_at": g("created_at"),
        "updated_at": g("updated_at"),
    }


def emit(records: Iterable[dict[str, Any]], out_path: str | None) -> int:
    """Write records as JSONL to out_path or stdout. Returns count written."""
    fh: Any = sys.stdout
    if out_path and out_path != "-":
        fh = open(out_path, "w", encoding="utf-8")
    count = 0
    for rec in records:
        fh.write(json.dumps(rec, ensure_ascii=False, default=str) + "\n")
        count += 1
    if fh is not sys.stdout:
        fh.close()
    return count


def build_arg_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        prog="export-django-links",
        description=(
            "Export shortlinks from an old Django 'links' app SQLite DB to "
            "JSONL. Standalone; does not modify the Django project."
        ),
    )
    p.add_argument(
        "--db",
        metavar="PATH",
        default=None,
        help="Path to the Django SQLite DB file (overrides $DJANGO_DB_PATH).",
    )
    p.add_argument(
        "--output",
        "-o",
        metavar="PATH",
        default=None,
        help=(
            "Output JSONL file path. Use '-' or omit for stdout (default)."
        ),
    )
    p.add_argument(
        "--owner-email",
        metavar="EMAIL",
        default=None,
        help=(
            "Only export links owned by this admin user (by auth_user email). "
            "Requires an owner/user FK column on the links table; if none "
            "exists, warns and exports all links."
        ),
    )
    return p


def main(argv: list[str] | None = None) -> int:
    args = build_arg_parser().parse_args(argv)
    db_path = resolve_db_path(args.db)
    conn = open_db(db_path)

    try:
        link_table = discover_link_table(conn)
    except LookupError as e:
        _stderr(f"error: {e}")
        conn.close()
        return 1

    tag_table = discover_tag_table(conn)
    m2m_table = (
        discover_m2m_table(conn, link_table, tag_table) if tag_table else None
    )

    users = load_users(conn)
    owner_col = detect_owner_column(conn, link_table)

    link_fk_col = "link_id"
    if m2m_table:
        try:
            link_fk_col = link_table_has_link_id_column(conn, m2m_table, link_table)
        except LookupError:
            link_fk_col = "link_id"

    rows = query_links(
        conn,
        link_table,
        tag_table,
        m2m_table,
        owner_col,
        args.owner_email,
        users,
    )

    records = (
        row_to_record(conn, r, tag_table, m2m_table, link_fk_col) for r in rows
    )
    count = emit(records, args.output)

    conn.close()

    if args.output and args.output != "-":
        _stderr(f"exported {count} link(s) to {args.output}")
    else:
        _stderr(f"exported {count} link(s) to stdout")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())