bvandeusenandClaude Opus 5.5
aab8581f20
#5815 "Restore engine: in place for a same-version backup". restore_in_place (steward/backup/restore.py) runs in one transaction: 1. It keeps aside what Steward records as it runs: ACTIVITY columns of rows that come back, and STATE_SETTINGS. 2. It drops every foreign key touching a replaced table. 3. It truncates the replaced tables and COPYs the backup's rows in unparsed, through asyncpg. 4. It puts the kept state back and moves sequences past their rows. 5. It puts right the rows whose parent didn't come back, by each key's own ON DELETE. A RESTRICT key that may be empty is cleared, so a run whose user is gone keeps its record. 6. It re-adds the keys, which checks every row. A failure anywhere rolls the whole restore back. A backup made at another schema is refused with NotSameSchema ("older": a full replace is the way; "newer": update Steward). The COPY-text reading moves into steward/backup/pgdata.py, shared by the check and the restore. The integration tests stand in for pg_restore with this database's own COPY output: - settings come back while the audit log and the history of what remains stay; - history of what didn't come back goes; - a failure part-way changes nothing; - a newer schema is refused. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
416 lines
20 KiB
Python
416 lines
20 KiB
Python
"""Bringing Steward back from a backup file: into an empty database
|
|
(`restore`, the CLI and first-run setup), or in place on a running Steward
|
|
(`restore_in_place`, below).
|
|
|
|
pg_restore runs in its three sections so the rows can be put right between
|
|
the data and the constraints: a kept table may point at a row of a table
|
|
whose rows were left out (an alert event naming the Ansible run it
|
|
started). Each such reference gets what its foreign key's ON DELETE says
|
|
would happen had that row been deleted: SET NULL clears it, anything else
|
|
removes the row. Then the constraints and indexes go on.
|
|
|
|
Migrations are not run here: the next start of Steward runs them, as it
|
|
always does, so a backup made by an older Steward comes up to date.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import io
|
|
import logging
|
|
import tempfile
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
from sqlalchemy import MetaData, text
|
|
from sqlalchemy.ext.asyncio import create_async_engine
|
|
|
|
from .contents import EXCLUDED
|
|
from .dump import BackupError, pg_env, read_backup
|
|
from .pgdata import async_url as _async_url
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
RESTORE_DEADLINE_SECONDS = 30 * 60
|
|
|
|
|
|
@dataclass
|
|
class Restored:
|
|
manifest: dict
|
|
secret_key: str
|
|
fixed: int # rows put right between data and constraints
|
|
|
|
|
|
def all_metadata(plugin_dirs: list[Path] | None = None) -> MetaData:
|
|
"""Steward's tables, every plugin's included, without an app: the
|
|
restore runs before anything else would load them. Plugins are found in
|
|
`plugin_dirs` (the bootstrap's, bundled first) and loaded by path under
|
|
the same module names the app gives them, so a plugin the app already
|
|
loaded isn't mapped twice. Without dirs (tests, a checkout): the
|
|
`plugins` package, as the tests import it."""
|
|
import importlib
|
|
import sys
|
|
|
|
import steward
|
|
import steward.models # noqa: F401 — registers core tables
|
|
from steward.core.plugin_manager import _import_plugin, resolve_plugin_path
|
|
from steward.models.base import Base
|
|
|
|
roots = [Path(d).resolve() for d in (plugin_dirs or [Path(steward.__file__).parent.parent / "plugins"])]
|
|
by_path = bool(plugin_dirs)
|
|
names = sorted({p.name for root in roots if root.is_dir() for p in root.iterdir()
|
|
if (p / "__init__.py").exists() and (p / "models.py").exists()})
|
|
for name in names:
|
|
if f"plugins.{name}.models" in sys.modules or f"_steward_plugin_{name}.models" in sys.modules:
|
|
continue # already mapped, under one name or the other
|
|
try:
|
|
if by_path:
|
|
package = _import_plugin(name, resolve_plugin_path(roots, name)).__name__
|
|
else:
|
|
package = f"plugins.{name}"
|
|
importlib.import_module(f"{package}.models")
|
|
except Exception: # a plugin that won't load: its tables are left as pg_restore put them
|
|
logger.warning("Restore: couldn't load plugin %r's tables", name, exc_info=True)
|
|
return Base.metadata
|
|
|
|
|
|
def dangling_fixes(metadata: MetaData) -> list[str]:
|
|
"""SQL putting right the rows of kept tables that point at a table whose
|
|
rows were left out, by each foreign key's own ON DELETE."""
|
|
out: list[str] = []
|
|
for table in sorted(metadata.tables.values(), key=lambda t: t.name):
|
|
if table.name in EXCLUDED:
|
|
continue
|
|
for fk in sorted(table.foreign_keys, key=lambda f: f.parent.name):
|
|
parent = fk.column.table.name
|
|
if parent not in EXCLUDED:
|
|
continue
|
|
child, col, pcol = table.name, fk.parent.name, fk.column.name
|
|
gone = (f'"{col}" IS NOT NULL AND NOT EXISTS '
|
|
f'(SELECT 1 FROM "{parent}" p WHERE p."{pcol}" = "{child}"."{col}")')
|
|
if (fk.ondelete or "").upper() == "SET NULL":
|
|
out.append(f'UPDATE "{child}" SET "{col}" = NULL WHERE {gone}')
|
|
else:
|
|
out.append(f'DELETE FROM "{child}" WHERE {gone}')
|
|
return out
|
|
|
|
|
|
# pg_restore 17 opens every script with this; a Postgres 16 server doesn't
|
|
# know the setting and refuses it, so the line is dropped on the way through.
|
|
_NEWER_THAN_SERVER = {b"SET transaction_timeout = 0;"}
|
|
# A COPY row is one line of the script; this bounds the longest one read.
|
|
_LINE_LIMIT = 256 * 1024 * 1024
|
|
|
|
|
|
async def _pg_restore(database_url: str, dump: Path, section: str) -> None:
|
|
"""One section of the dump into the database: pg_restore writes it as a
|
|
script, psql runs it and stops at the first error."""
|
|
env = pg_env(database_url)
|
|
try:
|
|
script = await asyncio.create_subprocess_exec(
|
|
"pg_restore", "--no-owner", "--no-privileges", f"--section={section}", "--file=-", str(dump),
|
|
env=env, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, limit=_LINE_LIMIT)
|
|
psql = await asyncio.create_subprocess_exec(
|
|
"psql", "-X", "-q", "-v", "ON_ERROR_STOP=1",
|
|
env=env, stdin=asyncio.subprocess.PIPE,
|
|
stdout=asyncio.subprocess.DEVNULL, stderr=asyncio.subprocess.PIPE)
|
|
except FileNotFoundError:
|
|
raise BackupError("pg_restore or psql isn't installed here.") from None
|
|
|
|
async def pump() -> None:
|
|
try:
|
|
while line := await script.stdout.readline():
|
|
if line.strip() in _NEWER_THAN_SERVER:
|
|
continue
|
|
psql.stdin.write(line)
|
|
await psql.stdin.drain()
|
|
except (BrokenPipeError, ConnectionResetError):
|
|
# psql stopped on an error (its stderr says which); pg_restore
|
|
# would block on a pipe nobody reads.
|
|
script.kill()
|
|
finally:
|
|
psql.stdin.close()
|
|
|
|
async def run() -> tuple[bytes, bytes]:
|
|
_, script_err, psql_err = await asyncio.gather(
|
|
pump(), script.stderr.read(), psql.stderr.read())
|
|
await asyncio.gather(script.wait(), psql.wait())
|
|
return script_err, psql_err
|
|
|
|
try:
|
|
script_err, psql_err = await asyncio.wait_for(run(), RESTORE_DEADLINE_SECONDS)
|
|
except TimeoutError:
|
|
for proc in (script, psql):
|
|
if proc.returncode is None:
|
|
proc.kill()
|
|
await proc.wait()
|
|
raise BackupError(f"Restoring {section} didn't finish within "
|
|
f"{RESTORE_DEADLINE_SECONDS // 60} minutes.") from None
|
|
for proc, err, name in ((psql, psql_err, "psql"), (script, script_err, "pg_restore")):
|
|
if proc.returncode != 0:
|
|
lines = [ln for ln in err.decode(errors="replace").splitlines() if ln.strip()]
|
|
raise BackupError(f"Restoring {section}: " + (lines[-1] if lines else f"{name} exited {proc.returncode}."))
|
|
|
|
|
|
async def _tables(database_url: str) -> list[str]:
|
|
engine = create_async_engine(_async_url(database_url))
|
|
try:
|
|
async with engine.connect() as conn:
|
|
return list((await conn.execute(text(
|
|
"SELECT tablename FROM pg_tables WHERE schemaname = 'public' ORDER BY tablename"))).scalars())
|
|
finally:
|
|
await engine.dispose()
|
|
|
|
|
|
async def _execute(database_url: str, statements: list[str]) -> int:
|
|
engine = create_async_engine(_async_url(database_url))
|
|
changed = 0
|
|
try:
|
|
async with engine.begin() as conn:
|
|
for sql in statements:
|
|
result = await conn.execute(text(sql))
|
|
changed += max(result.rowcount or 0, 0)
|
|
finally:
|
|
await engine.dispose()
|
|
return changed
|
|
|
|
|
|
async def restore(path: Path, database_url: str, *, wipe: bool = False,
|
|
metadata: MetaData | None = None, plugin_dirs: list[Path] | None = None) -> Restored:
|
|
"""Restore a backup file into the database at `database_url`. It must be
|
|
empty unless `wipe`, which drops everything in its public schema first."""
|
|
existing = await _tables(database_url)
|
|
if existing and not wipe:
|
|
raise BackupError(f"The database already has {len(existing)} tables "
|
|
f"({', '.join(existing[:3])}{'…' if len(existing) > 3 else ''}). "
|
|
"Restore into an empty database, or pass --wipe to replace what is there.")
|
|
if existing:
|
|
logger.warning("Restore: dropping %d tables in the public schema first", len(existing))
|
|
await _execute(database_url, ["DROP SCHEMA public CASCADE", "CREATE SCHEMA public"])
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
info, dump, key = await asyncio.to_thread(read_backup, path, Path(tmp))
|
|
await _pg_restore(database_url, dump, "pre-data")
|
|
await _pg_restore(database_url, dump, "data")
|
|
fixed = await _execute(database_url, dangling_fixes(metadata or all_metadata(plugin_dirs)))
|
|
await _pg_restore(database_url, dump, "post-data")
|
|
return Restored(info, key, fixed)
|
|
|
|
|
|
# ── In place, on a running Steward (#5815) ───────────────────────────────────
|
|
# A live restore replaces what the backup carries and keeps everything else:
|
|
# the audit log, alert history and metrics stay. It runs in one transaction,
|
|
# so a failure anywhere leaves the database as it was:
|
|
# 1. keep what Steward records as it runs (contents.ACTIVITY on rows that
|
|
# come back, contents.STATE_SETTINGS) aside in temporary tables;
|
|
# 2. drop every foreign key touching a replaced table;
|
|
# 3. empty the replaced tables and COPY the backup's rows in, unparsed;
|
|
# 4. put back what was kept aside, and move each sequence past its rows;
|
|
# 5. put right the rows whose parent didn't come back, by their foreign
|
|
# key: CASCADE removes the row, SET NULL clears the link, and a RESTRICT
|
|
# or NO ACTION link is cleared when it may be empty (a run whose user is
|
|
# gone keeps its record) and removed when it may not;
|
|
# 6. add the foreign keys back, which checks every row.
|
|
# Steward's database user owns its tables, so none of this needs a superuser.
|
|
|
|
|
|
class NotSameSchema(BackupError):
|
|
"""The backup was made at another schema: in place won't do."""
|
|
|
|
def __init__(self, schema: str, backup_heads: list[str], here_heads: list[str]):
|
|
self.schema = schema
|
|
if schema == "newer":
|
|
said = "It was made by a newer Steward than this one; update Steward first."
|
|
else:
|
|
said = ("It was made by an older Steward, so it can't go back in place; restoring it "
|
|
"replaces everything, history included.")
|
|
super().__init__(said)
|
|
self.backup_heads, self.here_heads = backup_heads, here_heads
|
|
|
|
|
|
@dataclass
|
|
class RestoredInPlace:
|
|
manifest: dict
|
|
secret_key: str
|
|
tables: list[str]
|
|
rows: int
|
|
fixed: int
|
|
|
|
|
|
@dataclass
|
|
class _ForeignKey:
|
|
name: str
|
|
child: str
|
|
parent: str
|
|
columns: list[str]
|
|
parent_columns: list[str]
|
|
on_delete: str # pg_constraint.confdeltype: c, n, a, r, d
|
|
nullable: bool
|
|
definition: str
|
|
|
|
|
|
def _q(name: str) -> str:
|
|
return '"' + name.replace('"', '""') + '"'
|
|
|
|
|
|
def fk_fix(fk: _ForeignKey) -> str | None:
|
|
"""The statement putting right `fk`'s rows whose parent is gone; None for
|
|
a key over several columns (re-adding it still checks every row)."""
|
|
if len(fk.columns) != 1:
|
|
return None
|
|
col, pcol = _q(fk.columns[0]), _q(fk.parent_columns[0])
|
|
gone = (f"x.{col} IS NOT NULL AND NOT EXISTS "
|
|
f"(SELECT 1 FROM public.{_q(fk.parent)} p WHERE p.{pcol} = x.{col})")
|
|
if fk.on_delete == "n" or (fk.on_delete in ("a", "r") and fk.nullable):
|
|
return f"UPDATE public.{_q(fk.child)} x SET {col} = NULL WHERE {gone}"
|
|
return f"DELETE FROM public.{_q(fk.child)} x WHERE {gone}"
|
|
|
|
|
|
async def _layout(conn, tables: list[str]):
|
|
"""What the transaction needs to know first: the foreign keys touching
|
|
`tables`, each table's primary key and columns, and its sequences."""
|
|
fks = []
|
|
for row in (await conn.execute(text(
|
|
"SELECT con.conname, cl.relname, ref.relname, con.confdeltype, "
|
|
"pg_get_constraintdef(con.oid), "
|
|
"ARRAY(SELECT a.attname FROM unnest(con.conkey) WITH ORDINALITY k(n, i) "
|
|
" JOIN pg_attribute a ON a.attrelid = con.conrelid AND a.attnum = k.n ORDER BY k.i), "
|
|
"ARRAY(SELECT a.attname FROM unnest(con.confkey) WITH ORDINALITY k(n, i) "
|
|
" JOIN pg_attribute a ON a.attrelid = con.confrelid AND a.attnum = k.n ORDER BY k.i), "
|
|
"NOT bool_or_notnull.v "
|
|
"FROM pg_constraint con "
|
|
"JOIN pg_class cl ON cl.oid = con.conrelid "
|
|
"JOIN pg_class ref ON ref.oid = con.confrelid "
|
|
"JOIN pg_namespace ns ON ns.oid = cl.relnamespace "
|
|
"CROSS JOIN LATERAL (SELECT bool_or(a.attnotnull) AS v FROM pg_attribute a "
|
|
" WHERE a.attrelid = con.conrelid AND a.attnum = ANY(con.conkey)) bool_or_notnull "
|
|
"WHERE con.contype = 'f' AND ns.nspname = 'public' "
|
|
"AND (cl.relname = ANY(:t) OR ref.relname = ANY(:t)) ORDER BY con.conname"),
|
|
{"t": tables})).all():
|
|
fks.append(_ForeignKey(row[0], row[1], row[2], list(row[5]), list(row[6]), row[3],
|
|
bool(row[7]), row[4]))
|
|
keys: dict[str, list[str]] = {}
|
|
for t, c in (await conn.execute(text(
|
|
"SELECT cl.relname, a.attname FROM pg_index i "
|
|
"JOIN pg_class cl ON cl.oid = i.indrelid "
|
|
"JOIN pg_namespace ns ON ns.oid = cl.relnamespace "
|
|
"JOIN pg_attribute a ON a.attrelid = cl.oid AND a.attnum = ANY(i.indkey) "
|
|
"WHERE i.indisprimary AND ns.nspname = 'public' AND cl.relname = ANY(:t) "
|
|
"ORDER BY cl.relname, a.attnum"), {"t": tables})).all():
|
|
keys.setdefault(t, []).append(c)
|
|
columns: dict[str, list[str]] = {}
|
|
sequences: list[tuple[str, str, str]] = []
|
|
for t, c, default in (await conn.execute(text(
|
|
"SELECT table_name, column_name, column_default FROM information_schema.columns "
|
|
"WHERE table_schema = 'public' AND table_name = ANY(:t) "
|
|
"ORDER BY table_name, ordinal_position"), {"t": tables})).all():
|
|
columns.setdefault(t, []).append(c)
|
|
if default and default.startswith("nextval("):
|
|
seq = (await conn.execute(text("SELECT pg_get_serial_sequence(:t, :c)"),
|
|
{"t": f"public.{_q(t)}", "c": c})).scalar()
|
|
if seq:
|
|
sequences.append((t, c, seq))
|
|
return fks, keys, columns, sequences
|
|
|
|
|
|
async def restore_in_place(path: Path, database_url: str, *,
|
|
plugin_dirs: list[Path] | None = None) -> RestoredInPlace:
|
|
"""Put the backup at `path` back into the running Steward's database,
|
|
replacing the tables it carries and keeping the rest (see above). Raises
|
|
NotSameSchema for a backup made at another schema, and BackupError for
|
|
one that can't be read or doesn't fit; the database is then untouched."""
|
|
from .check import known_revisions, schema_of
|
|
from .contents import ACTIVITY, STATE_SETTINGS
|
|
from .pgdata import restore_script, script_blocks
|
|
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
info, dump, key = await asyncio.to_thread(read_backup, path, Path(tmp))
|
|
blocks = script_blocks(await restore_script(dump))
|
|
engine = create_async_engine(_async_url(database_url))
|
|
try:
|
|
async with engine.connect() as conn:
|
|
heads = sorted((await conn.execute(text("SELECT version_num FROM alembic_version"))).scalars())
|
|
backup_heads = sorted(info.get("migration_heads") or [])
|
|
schema = schema_of(backup_heads, heads, known_revisions(plugin_dirs))
|
|
if schema != "same":
|
|
raise NotSameSchema(schema, backup_heads, heads)
|
|
present = set((await conn.execute(text(
|
|
"SELECT tablename FROM pg_tables WHERE schemaname = 'public'"))).scalars())
|
|
tables = sorted(t for t in blocks if t not in EXCLUDED and t in present)
|
|
fks, keys, columns, sequences = await _layout(conn, tables)
|
|
for t in tables:
|
|
extra = set(blocks[t][0]) - set(columns.get(t, []))
|
|
if extra:
|
|
raise BackupError(f"The backup's {t} has columns this Steward lacks: "
|
|
f"{', '.join(sorted(extra))}.")
|
|
|
|
rows = sum(raw.count(b"\n") for _, raw in (blocks[t] for t in tables))
|
|
try:
|
|
async with engine.begin() as conn:
|
|
# 1. What Steward records as it runs, kept aside.
|
|
if "app_settings" in tables:
|
|
await conn.execute(text(
|
|
"CREATE TEMP TABLE _keep_settings ON COMMIT DROP AS "
|
|
"SELECT * FROM public.app_settings WHERE key = ANY(:k)"),
|
|
{"k": sorted(STATE_SETTINGS)})
|
|
kept_aside: dict[str, list[str]] = {}
|
|
for t in tables:
|
|
act = [c for c in columns.get(t, []) if c in ACTIVITY.get(t, ())]
|
|
if act and keys.get(t):
|
|
kept_aside[t] = act
|
|
cols = ", ".join(_q(c) for c in keys[t] + act)
|
|
await conn.execute(text(
|
|
f"CREATE TEMP TABLE {_q('_keep_' + t)} ON COMMIT DROP AS "
|
|
f"SELECT {cols} FROM public.{_q(t)}"))
|
|
# 2. The foreign keys touching what is replaced.
|
|
for fk in fks:
|
|
await conn.execute(text(
|
|
f"ALTER TABLE public.{_q(fk.child)} DROP CONSTRAINT {_q(fk.name)}"))
|
|
# 3. The backup's rows, as the dump holds them.
|
|
await conn.execute(text("TRUNCATE " + ", ".join(f"public.{_q(t)}" for t in tables)))
|
|
raw = (await conn.get_raw_connection()).driver_connection
|
|
for t in tables:
|
|
cols, data = blocks[t]
|
|
if data:
|
|
await raw.copy_to_table(t, schema_name="public", source=io.BytesIO(data),
|
|
columns=cols, format="text")
|
|
# 4. Put back what was kept aside; sequences past their rows.
|
|
if "app_settings" in tables:
|
|
await conn.execute(text("DELETE FROM public.app_settings WHERE key = ANY(:k)"),
|
|
{"k": sorted(STATE_SETTINGS)})
|
|
await conn.execute(text("INSERT INTO public.app_settings SELECT * FROM _keep_settings"))
|
|
for t, act in kept_aside.items():
|
|
sets = ", ".join(f"{_q(c)} = k.{_q(c)}" for c in act)
|
|
match = " AND ".join(f"x.{_q(c)} = k.{_q(c)}" for c in keys[t])
|
|
await conn.execute(text(
|
|
f"UPDATE public.{_q(t)} x SET {sets} FROM {_q('_keep_' + t)} k WHERE {match}"))
|
|
for t, c, seq in sequences:
|
|
# `seq` is pg_get_serial_sequence's own answer, already quoted.
|
|
literal = "'" + seq.replace("'", "''") + "'"
|
|
await conn.execute(text(
|
|
f"SELECT setval({literal}, GREATEST(COALESCE((SELECT max({_q(c)}) FROM public.{_q(t)}), 0), "
|
|
f"(SELECT last_value FROM {seq}), 1))"))
|
|
# 5. Rows whose parent didn't come back. A removal can leave
|
|
# another row without its parent, so until nothing moves.
|
|
fixes = [sql for sql in (fk_fix(fk) for fk in fks) if sql]
|
|
fixed = 0
|
|
for _ in range(5):
|
|
moved = 0
|
|
for sql in fixes:
|
|
moved += max((await conn.execute(text(sql))).rowcount or 0, 0)
|
|
fixed += moved
|
|
if not moved:
|
|
break
|
|
# 6. The keys back on, each row checked.
|
|
for fk in fks:
|
|
await conn.execute(text(
|
|
f"ALTER TABLE public.{_q(fk.child)} ADD CONSTRAINT {_q(fk.name)} {fk.definition}"))
|
|
except BackupError:
|
|
raise
|
|
except Exception as exc:
|
|
said = getattr(exc, "orig", None) or exc
|
|
raise BackupError(f"The restore was undone, nothing changed: {said}") from exc
|
|
finally:
|
|
await engine.dispose()
|
|
logger.warning("Restore in place: replaced %d tables (%d rows) from %s; %d rows put right",
|
|
len(tables), rows, path.name, fixed)
|
|
return RestoredInPlace(info, key, tables, rows, fixed)
|