|
| 1 | +"""Mirror ALL prod app data into the staging branch (spec D8). |
| 2 | +
|
| 3 | +Truncate-and-COPY every public table except the exclusions, in FK-topological |
| 4 | +order, then apply the quiesce transform. Known FK back-edges (the same two |
| 5 | +preview_seed declares) are nulled during load and patched afterward. |
| 6 | +Self-enforcing: unclassified tables or unexpected FK cycles abort the run. |
| 7 | +""" |
| 8 | + |
| 9 | +import asyncio |
| 10 | +import os |
| 11 | +import re |
| 12 | +import shutil |
| 13 | +import sys |
| 14 | +import tempfile |
| 15 | +import time |
| 16 | + |
| 17 | +import asyncpg |
| 18 | + |
| 19 | +EXCLUDE = { |
| 20 | + # Alembic pointers: staging runs AHEAD of prod — never overwrite. |
| 21 | + "alembic_version_oddish", "alembic_version_backend", "alembic_version", |
| 22 | + # Runtime/queue state: rebuilds itself. |
| 23 | + "queue_slots", "trial_events", "_preview_seed_state", |
| 24 | + # Prod idempotency claims must not dedupe staging submissions. |
| 25 | + "submission_idempotency", |
| 26 | + # Prod API keys must never authenticate against staging — they would blur |
| 27 | + # cost/usage attribution across environments. Staging keys are minted per |
| 28 | + # person from the staging dashboard. (Purged below as well, since an |
| 29 | + # excluded table is not truncated and may hold rows from an earlier run.) |
| 30 | + "api_keys", |
| 31 | +} |
| 32 | +# (task, column) back-edges nulled during COPY and patched after — mirrors |
| 33 | +# preview_seed._BACKEDGES + _LINKAGE_COLUMNS. |
| 34 | +BACKEDGES = {("tasks", "current_version_id"), ("trials", "superseded_by_trial_id")} |
| 35 | + |
| 36 | +QUIESCE = [ |
| 37 | + "UPDATE worker_jobs SET status='CANCELLED' WHERE status::text NOT IN ('SUCCESS','FAILED','CANCELLED')", |
| 38 | + "UPDATE trials SET status='FAILED' WHERE status::text NOT IN ('SUCCESS','FAILED')", |
| 39 | + "UPDATE tasks SET status='FAILED' WHERE status::text NOT IN ('COMPLETED','FAILED')", |
| 40 | +] |
| 41 | + |
| 42 | +# Destination-only hygiene applied after the load. Separate from QUIESCE |
| 43 | +# (which neutralises copied in-flight work) because this enforces an |
| 44 | +# environment boundary rather than fixing job state. |
| 45 | +PURGE = [ |
| 46 | + "DELETE FROM api_keys", |
| 47 | +] |
| 48 | + |
| 49 | + |
| 50 | +def _plain(url: str) -> str: |
| 51 | + return url.replace("postgresql+asyncpg://", "postgresql://", 1) |
| 52 | + |
| 53 | + |
| 54 | +def _session_pooler(url: str) -> str: |
| 55 | + # Prod's app URL rides the transaction pooler (6543), which reaps |
| 56 | + # long-lived transactions — the multi-minute snapshot COPYs must use |
| 57 | + # the session pooler (5432) on the same host. No-op for URLs already |
| 58 | + # on 5432 or direct connections. |
| 59 | + return url.replace(":6543/", ":5432/", 1) |
| 60 | + |
| 61 | + |
| 62 | +async def _tables(conn) -> set[str]: |
| 63 | + rows = await conn.fetch( |
| 64 | + "SELECT tablename FROM pg_tables WHERE schemaname='public'" |
| 65 | + ) |
| 66 | + return {r["tablename"] for r in rows} |
| 67 | + |
| 68 | + |
| 69 | +async def _topo_order(conn, include: set[str]) -> list[str]: |
| 70 | + fks = await conn.fetch(""" |
| 71 | + SELECT tc.table_name AS child, ccu.table_name AS parent, |
| 72 | + kcu.column_name AS child_col |
| 73 | + FROM information_schema.table_constraints tc |
| 74 | + JOIN information_schema.key_column_usage kcu |
| 75 | + ON tc.constraint_name = kcu.constraint_name AND tc.table_schema = kcu.table_schema |
| 76 | + JOIN information_schema.constraint_column_usage ccu |
| 77 | + ON tc.constraint_name = ccu.constraint_name AND tc.table_schema = ccu.table_schema |
| 78 | + WHERE tc.constraint_type='FOREIGN KEY' AND tc.table_schema='public' |
| 79 | + """) |
| 80 | + deps: dict[str, set[str]] = {t: set() for t in include} |
| 81 | + for r in fks: |
| 82 | + c, p = r["child"], r["parent"] |
| 83 | + if c in include and p in include and c != p: |
| 84 | + if (c, r["child_col"]) in BACKEDGES: |
| 85 | + continue # broken edge: loaded as NULL, patched later |
| 86 | + deps[c].add(p) |
| 87 | + order, done = [], set() |
| 88 | + while deps: |
| 89 | + ready = sorted(t for t, ds in deps.items() if ds <= done) |
| 90 | + if not ready: |
| 91 | + sys.exit(f"FK cycle not covered by BACKEDGES: {sorted(deps)}") |
| 92 | + for t in ready: |
| 93 | + order.append(t); done.add(t); deps.pop(t) |
| 94 | + return order |
| 95 | + |
| 96 | + |
| 97 | +async def _connect_with_retry(url: str, attempts: int = 12, delay: float = 15.0): |
| 98 | + # The staging branch restarts on compute/disk changes, and its pooler can |
| 99 | + # report healthy while still refusing connections for a short window |
| 100 | + # (observed: econnrefused seconds after ACTIVE_HEALTHY). Retry connects. |
| 101 | + last: Exception | None = None |
| 102 | + for i in range(1, attempts + 1): |
| 103 | + try: |
| 104 | + return await asyncpg.connect(url, statement_cache_size=0) |
| 105 | + except (asyncpg.PostgresConnectionError, ConnectionError, OSError) as exc: |
| 106 | + last = exc |
| 107 | + print(f"connect attempt {i}/{attempts} failed: {exc!r}; retrying in {delay:.0f}s", |
| 108 | + flush=True) |
| 109 | + await asyncio.sleep(delay) |
| 110 | + raise last # type: ignore[misc] |
| 111 | + |
| 112 | + |
| 113 | +async def _src_connect(url: str): |
| 114 | + return await _connect_with_retry(url) |
| 115 | + |
| 116 | + |
| 117 | +def _pick_spool_dir() -> str: |
| 118 | + # The per-table spool (src -> tempfile -> dst) can run tens of GB for the |
| 119 | + # largest tables, and this job runs inside a container (ghcr.io ci-base) |
| 120 | + # on a GitHub-hosted runner, where neither the container's default temp |
| 121 | + # dir nor a host-mounted volume is reliably the roomier option (and /mnt |
| 122 | + # may not even be mounted in-container). Measure the real candidates and |
| 123 | + # spool wherever there's the most free space. |
| 124 | + candidates = ["/__w", os.environ.get("RUNNER_TEMP"), tempfile.gettempdir()] |
| 125 | + best_dir, best_free = tempfile.gettempdir(), -1 |
| 126 | + for c in candidates: |
| 127 | + if not c or not os.path.isdir(c) or not os.access(c, os.W_OK): |
| 128 | + continue |
| 129 | + try: |
| 130 | + free = shutil.disk_usage(c).free |
| 131 | + except OSError: |
| 132 | + continue |
| 133 | + if free > best_free: |
| 134 | + best_dir, best_free = c, free |
| 135 | + return best_dir |
| 136 | + |
| 137 | + |
| 138 | +async def _run_mirror() -> None: |
| 139 | + spool_dir = _pick_spool_dir() |
| 140 | + print(f"spool dir: {spool_dir} ({shutil.disk_usage(spool_dir).free / 1e9:.2f} GB free)", |
| 141 | + flush=True) |
| 142 | + src_url = _session_pooler(_plain(os.environ["SOURCE_DB_URL"])) |
| 143 | + holder = await _src_connect(src_url) |
| 144 | + dst = await _connect_with_retry(_plain(os.environ["TARGET_DB_URL"])) |
| 145 | + try: |
| 146 | + src_tables, dst_tables = await _tables(holder), await _tables(dst) |
| 147 | + include = (src_tables & dst_tables) - EXCLUDE |
| 148 | + unclassified = (src_tables - dst_tables - EXCLUDE) |
| 149 | + if unclassified: |
| 150 | + sys.exit(f"Tables on prod but not staging (run migrations first?): {sorted(unclassified)}") |
| 151 | + order = await _topo_order(holder, include) |
| 152 | + backedge_cols = {t: [c for (bt, c) in BACKEDGES if bt == t] for t in include} |
| 153 | + |
| 154 | + # The holder pins ONE repeatable-read snapshot for the whole run and is kept |
| 155 | + # warm by a keepalive ping; every table then reads through its own |
| 156 | + # short-lived connection importing that snapshot, so no source connection |
| 157 | + # ever idles longer than its own active COPY (the observed kill mode: |
| 158 | + # the single source socket sat idle through 30+ minute destination uploads). |
| 159 | + async with holder.transaction(isolation="repeatable_read", readonly=True): |
| 160 | + await holder.execute("SET LOCAL statement_timeout = 0") |
| 161 | + await holder.execute("SET LOCAL idle_in_transaction_session_timeout = 0") |
| 162 | + snapshot_id = await holder.fetchval("SELECT pg_export_snapshot()") |
| 163 | + assert re.fullmatch(r"[0-9A-Fa-f\-]+", snapshot_id) |
| 164 | + holder_lock = asyncio.Lock() |
| 165 | + stop_ping = asyncio.Event() |
| 166 | + |
| 167 | + async def _keepalive() -> None: |
| 168 | + while not stop_ping.is_set(): |
| 169 | + try: |
| 170 | + await asyncio.wait_for(stop_ping.wait(), timeout=60) |
| 171 | + except asyncio.TimeoutError: |
| 172 | + async with holder_lock: |
| 173 | + await holder.execute("SELECT 1") |
| 174 | + |
| 175 | + ping_task = asyncio.create_task(_keepalive()) |
| 176 | + |
| 177 | + async def _copy_table_from_src(t: str, select_cols: str, cols: list[str]) -> None: |
| 178 | + with tempfile.TemporaryFile(dir=spool_dir) as spool: |
| 179 | + conn = await _src_connect(src_url) |
| 180 | + try: |
| 181 | + async with conn.transaction(isolation="repeatable_read", readonly=True): |
| 182 | + await conn.execute("SET LOCAL statement_timeout = 0") |
| 183 | + await conn.execute(f"SET TRANSACTION SNAPSHOT '{snapshot_id}'") |
| 184 | + await conn.copy_from_query( |
| 185 | + f'SELECT {select_cols} FROM "{t}"', output=spool) |
| 186 | + finally: |
| 187 | + await conn.close() |
| 188 | + # Source is fully closed before the (potentially very long) |
| 189 | + # destination upload starts — no source connection ever idles |
| 190 | + # through it (the observed kill mode on both pooler modes). |
| 191 | + spool.seek(0) |
| 192 | + await dst.copy_to_table(t, source=spool, columns=cols) |
| 193 | + |
| 194 | + try: |
| 195 | + async with dst.transaction(): |
| 196 | + await dst.execute("SET LOCAL statement_timeout = 0") |
| 197 | + await dst.execute("SET LOCAL synchronous_commit = off") |
| 198 | + for t in reversed(order): |
| 199 | + await dst.execute(f'TRUNCATE TABLE "{t}" CASCADE') |
| 200 | + for t in order: |
| 201 | + started = time.monotonic() |
| 202 | + async with holder_lock: |
| 203 | + cols = [r["column_name"] for r in await holder.fetch( |
| 204 | + "SELECT column_name FROM information_schema.columns " |
| 205 | + "WHERE table_schema='public' AND table_name=$1 ORDER BY ordinal_position", t)] |
| 206 | + select_cols = ", ".join( |
| 207 | + f'NULL AS "{c}"' if c in backedge_cols.get(t, []) else f'"{c}"' for c in cols) |
| 208 | + attempts = 0 |
| 209 | + while True: |
| 210 | + attempts += 1 |
| 211 | + try: |
| 212 | + async with dst.transaction(): # savepoint |
| 213 | + # CASCADE, not plain TRUNCATE: Postgres refuses to |
| 214 | + # truncate a table with an existing FK reference |
| 215 | + # from another table regardless of whether that |
| 216 | + # DELETE, not TRUNCATE: plain TRUNCATE errors |
| 217 | + # structurally when any FK references t, and |
| 218 | + # TRUNCATE ... CASCADE follows EVERY FK edge — |
| 219 | + # including the back-edges dropped from the topo |
| 220 | + # order (tasks.current_version_id -> task_versions), |
| 221 | + # so a task_versions retry would silently empty the |
| 222 | + # already-loaded tasks table. DELETE is safe: at |
| 223 | + # retry time no loaded row references t's rows |
| 224 | + # (children load later; back-edge columns are still |
| 225 | + # NULL until the patch phase). |
| 226 | + await dst.execute(f'DELETE FROM "{t}"') |
| 227 | + await _copy_table_from_src(t, select_cols, cols) |
| 228 | + break |
| 229 | + except (asyncpg.PostgresConnectionError, |
| 230 | + asyncpg.InterfaceError, |
| 231 | + ConnectionError, OSError) as exc: |
| 232 | + if attempts >= 2: |
| 233 | + raise |
| 234 | + print(f"retrying {t} after transient failure: {exc!r}", flush=True) |
| 235 | + await asyncio.sleep(5) |
| 236 | + n = await dst.fetchval(f'SELECT count(*) FROM "{t}"') |
| 237 | + print(f"mirrored {t}: {n} rows in {time.monotonic() - started:.0f}s", flush=True) |
| 238 | + for (t, col) in sorted(BACKEDGES): |
| 239 | + if t not in include: |
| 240 | + continue |
| 241 | + pk = "id" |
| 242 | + conn = await _src_connect(src_url) |
| 243 | + try: |
| 244 | + async with conn.transaction(isolation="repeatable_read", readonly=True): |
| 245 | + await conn.execute(f"SET TRANSACTION SNAPSHOT '{snapshot_id}'") |
| 246 | + rows = await conn.fetch( |
| 247 | + f'SELECT "{pk}", "{col}" FROM "{t}" WHERE "{col}" IS NOT NULL') |
| 248 | + finally: |
| 249 | + await conn.close() |
| 250 | + # Set-based patch: row-by-row UPDATEs over the pooler ran |
| 251 | + # for hours on millions of rows (observed: 4h silent). |
| 252 | + # Stage the pairs in a typed temp table and join once. |
| 253 | + started = time.monotonic() |
| 254 | + if rows: |
| 255 | + tmp = f"_patch_{t}_{col}" |
| 256 | + await dst.execute( |
| 257 | + f'CREATE TEMP TABLE "{tmp}" ON COMMIT DROP AS ' |
| 258 | + f'SELECT "{pk}", "{col}" FROM "{t}" WITH NO DATA') |
| 259 | + await dst.copy_records_to_table( |
| 260 | + tmp, records=[(r[pk], r[col]) for r in rows]) |
| 261 | + await dst.execute( |
| 262 | + f'UPDATE "{t}" AS tgt SET "{col}" = p."{col}" ' |
| 263 | + f'FROM "{tmp}" AS p WHERE tgt."{pk}" = p."{pk}"') |
| 264 | + print(f"patched {t}.{col}: {len(rows)} rows in " |
| 265 | + f"{time.monotonic() - started:.0f}s", flush=True) |
| 266 | + for stmt in QUIESCE: |
| 267 | + tag = await dst.execute(stmt) |
| 268 | + print(f"quiesce: {stmt.split(' WHERE')[0]} -> {tag}", flush=True) |
| 269 | + for stmt in PURGE: |
| 270 | + tag = await dst.execute(stmt) |
| 271 | + print(f"purge: {stmt} -> {tag}", flush=True) |
| 272 | + finally: |
| 273 | + stop_ping.set() |
| 274 | + ping_task.cancel() |
| 275 | + try: |
| 276 | + await ping_task |
| 277 | + except (asyncio.CancelledError, Exception): |
| 278 | + pass |
| 279 | + live = await dst.fetchval( |
| 280 | + "SELECT count(*) FROM worker_jobs WHERE status::text NOT IN ('SUCCESS','FAILED','CANCELLED')") |
| 281 | + assert live == 0, f"{live} non-terminal worker_jobs after quiesce" |
| 282 | + print("mirror complete") |
| 283 | + finally: |
| 284 | + # A failed attempt must not leak connections into the next retry: close |
| 285 | + # both ends regardless of how _run_mirror exited, tolerating a |
| 286 | + # connection that is already dead/closed (e.g. the source drop that |
| 287 | + # motivated the outer retry in main()). |
| 288 | + try: |
| 289 | + await holder.close() |
| 290 | + except Exception: |
| 291 | + pass |
| 292 | + try: |
| 293 | + await dst.close() |
| 294 | + except Exception: |
| 295 | + pass |
| 296 | + |
| 297 | + |
| 298 | +async def main() -> None: |
| 299 | + attempts = 3 |
| 300 | + for attempt in range(1, attempts + 1): |
| 301 | + try: |
| 302 | + await _run_mirror() |
| 303 | + return |
| 304 | + except (asyncpg.PostgresConnectionError, asyncpg.InterfaceError, |
| 305 | + ConnectionError, OSError) as exc: |
| 306 | + if attempt >= attempts: |
| 307 | + raise |
| 308 | + print(f"mirror attempt {attempt}/{attempts} failed with a connection error " |
| 309 | + f"({exc!r}); restarting the whole mirror with fresh connections " |
| 310 | + f"and a fresh snapshot", flush=True) |
| 311 | + await asyncio.sleep(30) |
| 312 | + |
| 313 | + |
| 314 | +if __name__ == "__main__": |
| 315 | + asyncio.run(main()) |
0 commit comments