krow_backend/scripts/export-local-data.py

#!/usr/bin/env python3
"""
Export this deployment's data as a replayable SQL file, for moving a whole
tenant from one database to another.

    make export-data                     # -> seed/exports/local-data.sql
    make export-data OUT=/tmp/krow.sql

This is NOT `make seed`. `seed/fixtures/seed.json` is the frontend's demo
dataset, generated from krow-demo and replayed by cmd/seed; it has never
reflected what is actually in a database. This reads the database itself, so
work done through the API — agents published, knowledge ingested, applications
screened — is carried across. The two are complementary and neither replaces
the other.

Why a generated SQL file rather than pg_dump straight into the target:

  - pg_dump orders tables by its own rules, not by foreign keys, so a
    --data-only restore into a populated schema fails on FK violations in a
    way that depends on table names. This sorts topologically from
    pg_constraint and dumps one table at a time in that order, so parents are
    always inserted before children.

  - Every statement is INSERT ... ON CONFLICT DO NOTHING. Re-running is safe
    and a row that already exists on the target is left exactly as it is.
    Nothing here can overwrite or delete data on the target — by construction,
    not by convention.

  - The file carries a guard preamble that refuses to apply against a schema
    older than the one it was taken from, because a restore into a half
    migrated database half-succeeds and that is worse than failing.

Apply it with a transaction, so a failure leaves nothing behind:

    psql "$TARGET_URL" --single-transaction -v ON_ERROR_STOP=1 \
      -f seed/exports/local-data.sql

The output contains real tenant data — password hashes, session tokens,
personal details. It is written under seed/exports/, which is gitignored. Do
not commit it; move it to the target host over scp or a secrets channel.
"""
import os
import subprocess
import sys
from collections import defaultdict

# Tables whose contents belong to the machine they were written on, not to the
# tenant. schema_migrations is owned by the migration tool and a stale row
# there would make the target lie about its own version; sessions are bound to
# cookies issued by one deployment and are worthless on another.
EXCLUDED = {"schema_migrations", "sessions"}

OUT = os.environ.get("OUT") or "seed/exports/local-data.sql"


def env(name: str, default: str = "") -> str:
    return os.environ.get(name, default).strip()


def psql_env() -> dict:
    """Connection environment for psql/pg_dump, from the same DATABASE_* vars
    the Go services read. No connection string is ever built by hand, so a
    password cannot end up in a process list."""
    e = dict(os.environ)
    e["PGHOST"] = env("DATABASE_HOST", "127.0.0.1")
    e["PGPORT"] = env("DATABASE_PORT", "5432")
    e["PGDATABASE"] = env("DATABASE_NAME")
    e["PGUSER"] = env("DATABASE_USER")
    if env("DATABASE_PASSWORD"):
        e["PGPASSWORD"] = env("DATABASE_PASSWORD")
    e["PGSSLMODE"] = env("DATABASE_SSLMODE", "disable")
    return e


def query(sql: str, e: dict) -> list[str]:
    out = subprocess.run(
        ["psql", "-Atq", "-c", sql], env=e, capture_output=True, text=True
    )
    if out.returncode != 0:
        sys.exit(f"export-local-data: query failed: {out.stderr.strip()}")
    return [line for line in out.stdout.splitlines() if line]


def topological_order(e: dict) -> list[str]:
    """Tables sorted so that every table appears after the tables it references.

    Self-references are ignored: a row pointing at another row of its own table
    is an ordering problem inside one INSERT batch, which ON CONFLICT DO
    NOTHING plus a second run resolves, not a table ordering problem.
    """
    tables = query(
        "select table_name from information_schema.tables "
        "where table_schema='public' and table_type='BASE TABLE' order by 1", e
    )
    tables = [t for t in tables if t not in EXCLUDED]
    known = set(tables)

    edges = query(
        "select distinct src.relname || '|' || tgt.relname "
        "from pg_constraint c "
        "join pg_class src on src.oid = c.conrelid "
        "join pg_class tgt on tgt.oid = c.confrelid "
        "join pg_namespace n on n.oid = src.relnamespace "
        "where c.contype = 'f' and n.nspname = 'public' "
        "and src.relname <> tgt.relname", e
    )

    deps: dict[str, set[str]] = defaultdict(set)   # table -> tables it needs first
    dependents: dict[str, set[str]] = defaultdict(set)
    for edge in edges:
        child, parent = edge.split("|", 1)
        if child in known and parent in known:
            deps[child].add(parent)
            dependents[parent].add(child)

    # Kahn's algorithm, alphabetical among ready tables so the output of two
    # runs against the same schema is byte-identical and reviewable in a diff.
    ready = sorted(t for t in tables if not deps[t])
    ordered: list[str] = []
    while ready:
        table = ready.pop(0)
        ordered.append(table)
        for child in sorted(dependents[table]):
            deps[child].discard(table)
            if not deps[child]:
                ready.append(child)
        ready.sort()

    if len(ordered) != len(tables):
        cyclic = sorted(set(tables) - set(ordered))
        sys.exit(
            "export-local-data: foreign keys form a cycle across "
            f"{', '.join(cyclic)}; no insert order can satisfy them. Make one "
            "of those constraints DEFERRABLE, or exclude a table here."
        )
    return ordered


def dump_table(table: str, e: dict) -> str:
    r"""One table's rows as INSERT ... ON CONFLICT DO NOTHING.

    pg_dump is invoked per table because it reorders a multi-table dump
    internally and would undo the topological sort. That leaves a preamble and
    an epilogue around each table's statements, and both have to go: pg_dump 18
    wraps its output in the psql meta-commands `\restrict` and `\unrestrict`,
    so keeping the closing one without its opener makes psql fail with
    "not currently in restricted mode" and — under --single-transaction — roll
    the whole import back having inserted nothing.

    Neither end is trimmed line by line, because a text column may legitimately
    contain a line that looks like a comment or a meta-command. Instead the
    slice runs from the first line that starts an INSERT to the last line that
    ends a statement. Every generated statement ends in a semicolon and nothing
    in pg_dump's epilogue does, so this cannot cut a row's own text short.
    """
    out = subprocess.run(
        ["pg_dump", "--data-only", "--inserts", "--on-conflict-do-nothing",
         "--no-owner", "--no-privileges", "--table", f"public.{table}"],
        env=e, capture_output=True, text=True,
    )
    if out.returncode != 0:
        sys.exit(f"export-local-data: pg_dump {table} failed: {out.stderr.strip()}")

    lines = out.stdout.splitlines()
    first = next((i for i, l in enumerate(lines) if l.startswith("INSERT INTO ")), None)
    if first is None:
        return ""  # table is empty
    last = max(i for i, l in enumerate(lines) if l.rstrip().endswith(";"))
    return "\n".join(lines[first:last + 1]).rstrip() + "\n"


def main() -> None:
    e = psql_env()
    if not e["PGDATABASE"] or not e["PGUSER"]:
        sys.exit("export-local-data: DATABASE_NAME and DATABASE_USER are required "
                 "(source them from .env, or use `make export-data`)")

    version = (query("select version from schema_migrations", e) or ["0"])[0]
    dirty = (query("select dirty from schema_migrations", e) or ["f"])[0]
    if dirty == "t":
        sys.exit(f"export-local-data: source schema is dirty at version {version}; "
                 "fix it with `make migrate-force VERSION=…` before exporting")

    tables = topological_order(e)
    os.makedirs(os.path.dirname(OUT) or ".", exist_ok=True)

    counts: dict[str, int] = {}
    with open(OUT, "w") as f:
        f.write(
            f"-- Krow data export from {e['PGDATABASE']} at schema version {version}.\n"
            "-- Generated by scripts/export-local-data.py — do not edit by hand.\n"
            "--\n"
            "-- Every statement is ON CONFLICT DO NOTHING: applying this can create\n"
            "-- rows on the target but can never modify or delete one. Tables are\n"
            "-- ordered so parents precede children.\n"
            "--\n"
            "--   psql \"$TARGET_URL\" --single-transaction -v ON_ERROR_STOP=1 \\\n"
            f"--     -f {OUT}\n\n"
            "\\set ON_ERROR_STOP on\n\n"
            "DO $krow_guard$\n"
            "DECLARE target_version bigint;\n"
            "BEGIN\n"
            "  SELECT version INTO target_version FROM public.schema_migrations;\n"
            f"  IF target_version IS NULL OR target_version < {version} THEN\n"
            "    RAISE EXCEPTION 'target schema is at version %, but this export was "
            f"taken at {version}; run migrations on the target first', target_version;\n"
            "  END IF;\n"
            "END\n"
            "$krow_guard$;\n\n"
        )
        for table in tables:
            body = dump_table(table, e)
            counts[table] = body.count("INSERT INTO ")
            f.write(f"-- {table} ({counts[table]} rows)\n")
            f.write(body if body else "-- (empty)\n")
            f.write("\n")

    total = sum(counts.values())
    print(f"exported {e['PGDATABASE']} (schema version {version}) -> {OUT}")
    for table in tables:
        if counts[table]:
            print(f"  {table:<22} {counts[table]:>5}")
    skipped = [t for t in tables if not counts[t]]
    if skipped:
        print(f"  (empty: {', '.join(skipped)})")
    print(f"  {'TOTAL':<22} {total:>5} rows")
    print("\nThis file contains real tenant data. It is gitignored; move it to the "
          "target host over scp, not through git.")


if __name__ == "__main__":
    main()