bvandeusenandClaude Opus 5.5
9e5f6e183a
- restore_in_place checks the backup's migration heads before pg_restore writes anything, so a backup from another schema is refused at once. - The integration tests read their own setting rows directly: get_all_settings knows only Steward's own keys. - check.py: drop an import left unused by the move to pgdata. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
419 lines
20 KiB
Python
419 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
|
|
|
|
engine = create_async_engine(_async_url(database_url))
|
|
try:
|
|
# The schema first: a backup that can't go in place is refused before
|
|
# anything else is read.
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
info, dump, key = await asyncio.to_thread(read_backup, path, Path(tmp))
|
|
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)
|
|
blocks = script_blocks(await restore_script(dump))
|
|
async with engine.connect() as conn:
|
|
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)
|