"""The staging refresh: production's dump replaces staging's data and the client's rows survive. Sources of truth: this module, `app/services/discussion.py` (`CLIENT_WRITTEN_TABLES`, the one list of tables people write on the site), `scripts/staging-refresh.py` (the host command that runs the halves in order), `docs/RUNBOOK.md` (Staging refresh), `tests/test_staging_refresh.py`. Why this exists: Adam comments on staging. A refresh that restores production's dump over staging replaces those tables wholesale, and a lost comment is not recoverable. So a refresh is three halves. `staging-refresh export` reads every client-written row into one JSON document BEFORE anything is touched (the host script keeps it under `backups/refresh/`); the dump is restored and migrated to head; `staging-refresh apply` puts the rows back. Apply never loses a row. Per table, in dependency order: a row absent from the restored database is inserted with its own id; a row already there and identical is left alone; a row whose NATURAL key exists with different content is decided by its timestamp (the newer wins, a tie keeps production's); a row whose primary key is taken by different content is inserted under a fresh id and reported as renumbered. Who-columns are remapped by username, because account ids need not agree between hosts; an author whose account is gone keeps the typed name with an empty link. `--check` runs the whole apply in one transaction and rolls it back, so the report is exact and nothing is written. python -m app.cli staging-refresh export > client-rows.json python -m app.cli staging-refresh apply --check --from client-rows.json python -m app.cli staging-refresh apply --from - < client-rows.json """ from __future__ import annotations import argparse import json import sys from dataclasses import dataclass, field from datetime import UTC, datetime from sqlalchemy import DateTime, Table, and_, func, insert, select, text, update from sqlalchemy.orm import Session from app.db import SessionLocal from app.models import Base from app.services.discussion import CLIENT_WRITTEN_TABLES FORMAT = 1 @dataclass(frozen=True) class TableSpec: """How one client-written table is put back. `natural`: the columns that identify a row across hosts (None: the primary key alone). `ts`: the column whose later value wins a natural-key conflict (None: production keeps its row). `fks`: column -> the table it points at; `accounts` remaps by username, any other name by the id map built while applying that table earlier in the order. """ pk: str = "id" natural: tuple[str, ...] | None = None ts: str | None = None fks: dict[str, str] = field(default_factory=dict) #: (the type column, {type value: the table whose id map the text `source_id` follows}). source_ref: tuple[str, dict[str, str]] | None = None #: One spec per table in CLIENT_WRITTEN_TABLES; the test pins the two sets equal. SPECS: dict[str, TableSpec] = { # Cards are imported by title on both hosts (services/discussion_import.py). "discussion_items": TableSpec(natural=("title",), ts="updated_at"), # `closing_comment_id` is a plain comment id, not mapped: a renumbered comment (rare) leaves # the closing word to be read from the thread's last comment by the resolver. "threads": TableSpec(natural=("subject_type", "subject_id"), ts="resolved_at", fks={"resolved_by_id": "accounts", "followup_by_id": "accounts", "archived_by_id": "accounts"}), "discussion_comments": TableSpec(fks={"thread_id": "threads", "item_id": "discussion_items", "author_id": "accounts"}), "feature_priorities": TableSpec(natural=("feature_key",), ts="updated_at", fks={"author_id": "accounts"}), "quote_selections": TableSpec(natural=("item_key",), ts="updated_at", fks={"author_id": "accounts"}), "quote_requests": TableSpec(fks={"author_id": "accounts"}), # Seeded by (owner, title) on both hosts (scripts/seed-todos.py); a completed row beats an open one. "client_todos": TableSpec(natural=("owner", "title"), ts="completed_at", fks={"completed_by_id": "accounts"}), "client_uploads": TableSpec(natural=("stored_name",), ts="created_at", fks={"todo_id": "client_todos", "uploaded_by_id": "accounts"}), "owner_item_states": TableSpec(pk="item_id", natural=("item_id",), ts="updated_at", fks={"acted_by_id": "accounts"}), "notifications": TableSpec(natural=("dedupe_key",), ts="created_at", fks={"recipient_id": "accounts", "actor_id": "accounts"}, source_ref=("source_type", {"comment": "discussion_comments", "thread": "threads"})), "account_preferences": TableSpec(pk="account_id", natural=("account_id",), ts="updated_at", fks={"account_id": "accounts"}), # A person's read stamp per thread (T12): the later stamp wins. "thread_reads": TableSpec(natural=("thread_id", "account_id"), ts="read_at", fks={"thread_id": "threads", "account_id": "accounts"}), "thread_asks": TableSpec(fks={"thread_id": "threads", "for_id": "accounts", "by_id": "accounts", "done_by_id": "accounts"}), "comment_acks": TableSpec(natural=("comment_id", "account_id"), ts="created_at", fks={"comment_id": "discussion_comments", "account_id": "accounts"}), "comment_flags": TableSpec(natural=("comment_id", "account_id"), ts="created_at", fks={"comment_id": "discussion_comments", "account_id": "accounts"}), } OUTCOMES = ("restored", "same", "updated", "kept", "renumbered", "unresolvable") def _table(name: str) -> Table: return Base.metadata.tables[name] def _aware(value): """SQLite hands naive datetimes back; treat them as UTC so both sides compare.""" if isinstance(value, datetime) and value.tzinfo is None: return value.replace(tzinfo=UTC) return value def _jsonable(value): if isinstance(value, datetime): return _aware(value).isoformat() return value def _typed(table: Table, row: dict) -> dict: """A JSON row back into column values: ISO strings become aware datetimes.""" out = {} for name, value in row.items(): if value is not None and isinstance(table.c[name].type, DateTime): value = _aware(datetime.fromisoformat(value)) out[name] = value return out def schema_version(db: Session) -> str | None: try: return db.execute(text("select version_num from alembic_version")).scalar() except Exception: # noqa: BLE001 - no alembic table (a bare test database) db.rollback() return None # --- export ---------------------------------------------------------------------------------- def export(db: Session) -> dict: """Every client-written row, plus the accounts map the who-columns are remapped through.""" tables: dict[str, list[dict]] = {} for name in CLIENT_WRITTEN_TABLES: t = _table(name) rows = db.execute(select(t).order_by(t.c[SPECS[name].pk])).mappings().all() tables[name] = [{k: _jsonable(v) for k, v in r.items()} for r in rows] accounts = _table("accounts") people = db.execute(select(accounts.c.id, accounts.c.username, accounts.c.display_name)).all() return { "format": FORMAT, "exported_at": datetime.now(UTC).isoformat(timespec="seconds"), "schema_version": schema_version(db), "accounts": {str(i): {"username": u, "display_name": d} for i, u, d in people}, "tables": tables, "counts": {name: len(rows) for name, rows in tables.items()}, } # --- apply ----------------------------------------------------------------------------------- @dataclass class Report: tables: dict[str, dict[str, int]] = field(default_factory=dict) notes: list[str] = field(default_factory=list) check: bool = False def count(self, table: str, outcome: str, n: int = 1) -> None: self.tables.setdefault(table, {o: 0 for o in ("exported", *OUTCOMES, "unlinked")}) self.tables[table][outcome] += n @property def unresolvable(self) -> int: return sum(t["unresolvable"] for t in self.tables.values()) def lines(self) -> list[str]: out = ["[check, nothing written]" if self.check else "[applied]"] for name, c in self.tables.items(): out.append( f"{name}: {c['exported']} exported; restored {c['restored']}, same {c['same']}, " f"updated {c['updated']}, kept {c['kept']}, renumbered {c['renumbered']}, " f"unresolvable {c['unresolvable']}" + (f", {c['unlinked']} author link(s) emptied" if c["unlinked"] else "") ) out.extend(self.notes) return out def _accounts_by_username(db: Session) -> dict[str, int]: accounts = _table("accounts") return {u: i for i, u in db.execute(select(accounts.c.id, accounts.c.username)).all() if u} def _content(row: dict, pk: str) -> dict: return {k: _aware(v) for k, v in row.items() if k != pk} def _find(db: Session, t: Table, where) -> dict | None: found = db.execute(select(t).where(where)).mappings().first() return dict(found) if found is not None else None def _find_by_content(db: Session, t: Table, content: dict) -> dict | None: """The same row under any id: what a second apply must recognise after a renumbering.""" conds = [t.c[k].is_(None) if v is None else t.c[k] == v for k, v in content.items()] return _find(db, t, and_(*conds)) def apply(db: Session, payload: dict, *, check: bool = False) -> Report: """Put the exported rows back; see the module docstring for the policy. Commits unless `check`, in which case everything is rolled back and the report says what would happen.""" report = Report(check=check) if payload.get("format") != FORMAT: raise ValueError(f"staging-refresh: export format {payload.get('format')!r} is not {FORMAT}") exported = payload.get("tables") or {} unknown = set(exported) - set(CLIENT_WRITTEN_TABLES) if unknown: raise ValueError(f"staging-refresh: the export carries tables this code does not know: {sorted(unknown)}") missing = [name for name in CLIENT_WRITTEN_TABLES if name not in exported] if missing: raise ValueError(f"staging-refresh: the export lacks {missing}; export again from the same code") here, there = schema_version(db), payload.get("schema_version") if here != there: report.notes.append(f"schema: the export came from {there or 'no alembic table'}; this database is at " f"{here or 'no alembic table'} (run alembic upgrade head before apply if it is behind)") people = _accounts_by_username(db) exported_people = payload.get("accounts") or {} id_maps: dict[str, dict] = {} # table -> old id -> new id for name in CLIENT_WRITTEN_TABLES: t, spec = _table(name), SPECS[name] rows = exported[name] report.count(name, "exported", len(rows)) id_map: dict = {} id_maps[name] = id_map # A renumbered row takes an id above both the table's and the export's highest, so it # can never land on an id a later export row still wants (a fresh serial would: the # next value after production's 214 is staging's 215). next_free = None if spec.pk == "id": table_max = db.execute(select(func.max(t.c.id))).scalar() or 0 next_free = max(table_max, max((r["id"] for r in rows), default=0)) + 1 for raw in rows: row = _typed(t, raw) old_id = row[spec.pk] # 1. remap every foreign key broken = False for col, target in spec.fks.items(): value = row.get(col) if value is None: continue if target == "accounts": username = (exported_people.get(str(value)) or {}).get("username") new = people.get(username) if username else None if new is None and not t.c[col].nullable: broken = True # a row that names a person who is not there cannot come back break if new is None: report.count(name, "unlinked") row[col] = new else: if value not in id_maps.get(target, {}): broken = True break row[col] = id_maps[target][value] if not broken and spec.source_ref: type_col, targets = spec.source_ref ref_table = targets.get(row.get(type_col) or "") if ref_table and row.get("source_id"): try: row["source_id"] = str(id_maps[ref_table][int(row["source_id"])]) except (KeyError, ValueError): broken, col, value, target = True, "source_id", row["source_id"], ref_table if broken: report.count(name, "unresolvable") report.notes.append(f"{name} {old_id}: its {col} {value} is not in the restored {target}; left in the export file") continue content = _content(row, spec.pk) # 2. by natural key existing = None if spec.natural: existing = _find(db, t, and_(*[t.c[c] == row[c] for c in spec.natural])) if existing is not None: id_map[old_id] = existing[spec.pk] if _content(existing, spec.pk) == content: report.count(name, "same") elif spec.ts and (_aware(content.get(spec.ts)) or datetime.min.replace(tzinfo=UTC)) > \ (_aware(existing.get(spec.ts)) or datetime.min.replace(tzinfo=UTC)): db.execute(update(t).where(t.c[spec.pk] == existing[spec.pk]).values(**content)) report.count(name, "updated") else: report.count(name, "kept") continue # 3. by primary key taken = _find(db, t, t.c[spec.pk] == old_id) if taken is not None: if _content(taken, spec.pk) == content: id_map[old_id] = old_id report.count(name, "same") continue elsewhere = _find_by_content(db, t, content) if elsewhere is not None: id_map[old_id] = elsewhere[spec.pk] report.count(name, "same") continue assert next_free is not None, name db.execute(insert(t).values(**{**content, spec.pk: next_free})) id_map[old_id] = next_free next_free += 1 report.count(name, "renumbered") report.notes.append(f"{name} {old_id}: its id was taken by a different row; now {id_map[old_id]}") continue db.execute(insert(t).values(**row)) id_map[old_id] = old_id report.count(name, "restored") db.flush() if check: db.rollback() return report db.commit() if db.bind is not None and db.bind.dialect.name == "postgresql": # Rows came back with their own ids; the sequences must move past them or the next # insert collides. setval is not transactional, hence after the commit. for name in CLIENT_WRITTEN_TABLES: pk = SPECS[name].pk if pk == "id": db.execute(text( f"select setval(pg_get_serial_sequence('{name}', 'id'), coalesce(max(id), 1)) from {name}" )) db.commit() return report # --- the CLI --------------------------------------------------------------------------------- def cmd_export(_: argparse.Namespace) -> int: with SessionLocal() as db: payload = export(db) json.dump(payload, sys.stdout, indent=1) sys.stdout.write("\n") print("exported " + ", ".join(f"{k}={v}" for k, v in payload["counts"].items()), file=sys.stderr) return 0 def cmd_apply(args: argparse.Namespace) -> int: if args.source == "-": payload = json.load(sys.stdin) else: with open(args.source, encoding="utf-8") as fh: payload = json.load(fh) with SessionLocal() as db: try: report = apply(db, payload, check=args.check) except ValueError as exc: db.rollback() print(exc) return 2 for line in report.lines(): print(line) return 1 if report.unresolvable else 0 def register(sub: argparse._SubParsersAction) -> None: refresh = sub.add_parser("staging-refresh", help="the client-written rows out before a restore, and back in after") rsub = refresh.add_subparsers(dest="refresh_command", required=True) exp = rsub.add_parser("export", help="every client-written row as one JSON document on stdout") exp.set_defaults(func=cmd_export) app_ = rsub.add_parser("apply", help="put an exported document's rows back (by natural key, never losing one)") app_.add_argument("--from", dest="source", required=True, metavar="FILE", help="the exported JSON, or - for stdin") app_.add_argument("--check", action="store_true", help="report what would change; write nothing") app_.set_defaults(func=cmd_apply)