feat(backup): restore a same-version backup in place, history kept, all or nothing (#5815)
#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>
This commit is contained in:
+4
-96
@@ -36,15 +36,14 @@ from sqlalchemy.ext.asyncio import create_async_engine
|
||||
from .contents import (ACTIVITY, ACTIVITY_EVERYWHERE, AREAS, EXCLUDED, HAND_SET_ONLY,
|
||||
STATE_SETTINGS)
|
||||
from .dump import BackupError, read_backup
|
||||
from .restore import _async_url
|
||||
from .pgdata import async_url as _async_url
|
||||
from .pgdata import copy_rows, restore_script, script_tables
|
||||
|
||||
CHECK_DEADLINE_SECONDS = 5 * 60
|
||||
# Names shown for a row, the first a table has; else its primary key.
|
||||
LABEL_COLUMNS = ("name", "username", "key", "title", "widget_key", "label", "mac")
|
||||
# A setting's two values are shown when they fit; longer ones are cut.
|
||||
SHOWN_CHARS = 120
|
||||
_TOKEN = re.compile(r"enc:v1:[A-Za-z0-9_\-=]+")
|
||||
_COPY = re.compile(r'^COPY ([^ ]+) \((.*)\) FROM stdin;$')
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -134,79 +133,6 @@ class CheckReport:
|
||||
return {"ok": self.ok, "changes": self.changes, "rows": self.rows, "summary": self.summary()}
|
||||
|
||||
|
||||
# ── COPY text ───────────────────────────────────────────────────────────────
|
||||
|
||||
_ESCAPES = {ord("b"): b"\b", ord("f"): b"\f", ord("n"): b"\n", ord("r"): b"\r",
|
||||
ord("t"): b"\t", ord("v"): b"\v"}
|
||||
|
||||
|
||||
def unescape(field: bytes) -> str | None:
|
||||
"""One field of COPY's text format: \\N is NULL, backslash escapes are
|
||||
undone (named, octal \\ooo and hex \\xhh)."""
|
||||
if field == b"\\N":
|
||||
return None
|
||||
if b"\\" not in field:
|
||||
return field.decode("utf-8", errors="replace")
|
||||
out = bytearray()
|
||||
i, n = 0, len(field)
|
||||
while i < n:
|
||||
c = field[i]
|
||||
if c != 0x5C or i + 1 == n:
|
||||
out.append(c)
|
||||
i += 1
|
||||
continue
|
||||
nxt = field[i + 1]
|
||||
if nxt in _ESCAPES:
|
||||
out += _ESCAPES[nxt]
|
||||
i += 2
|
||||
elif 0x30 <= nxt <= 0x37:
|
||||
j = i + 1
|
||||
while j < n and j < i + 4 and 0x30 <= field[j] <= 0x37:
|
||||
j += 1
|
||||
out.append(int(field[i + 1:j], 8) & 0xFF)
|
||||
i = j
|
||||
elif nxt == ord("x") and i + 2 < n and chr(field[i + 2]) in "0123456789abcdefABCDEF":
|
||||
j = i + 2
|
||||
while j < n and j < i + 4 and chr(field[j]) in "0123456789abcdefABCDEF":
|
||||
j += 1
|
||||
out.append(int(field[i + 2:j], 16))
|
||||
i = j
|
||||
else:
|
||||
out.append(nxt)
|
||||
i += 2
|
||||
return bytes(out).decode("utf-8", errors="replace")
|
||||
|
||||
|
||||
def copy_rows(data: bytes) -> list[list[str | None]]:
|
||||
"""COPY text rows (no header, no terminator) as lists of fields."""
|
||||
return [[unescape(f) for f in line.split(b"\t")]
|
||||
for line in data.split(b"\n") if line and line != b"\\."]
|
||||
|
||||
|
||||
def _ident(name: str) -> str:
|
||||
name = name.strip()
|
||||
return name[1:-1].replace('""', '"') if name.startswith('"') and name.endswith('"') else name
|
||||
|
||||
|
||||
def script_tables(script: bytes) -> dict[str, tuple[list[str], list[list[str | None]]]]:
|
||||
"""Each table's (columns, rows) from a pg_restore --data-only script."""
|
||||
out: dict[str, tuple[list[str], list[list[str | None]]]] = {}
|
||||
lines = iter(script.split(b"\n"))
|
||||
for line in lines:
|
||||
m = _COPY.match(line.decode("utf-8", errors="replace"))
|
||||
if not m:
|
||||
continue
|
||||
table = _ident(m.group(1).split(".")[-1])
|
||||
columns = [_ident(c) for c in m.group(2).split(",")]
|
||||
body = []
|
||||
for row in lines:
|
||||
if row == b"\\.":
|
||||
break
|
||||
body.append(row)
|
||||
out[table] = (columns, copy_rows(b"\n".join(body)))
|
||||
return out
|
||||
|
||||
|
||||
# ── Comparing ───────────────────────────────────────────────────────────────
|
||||
|
||||
def _fernet(key: str):
|
||||
@@ -355,26 +281,8 @@ def _labels(tables: dict[str, _Side]) -> dict[str, dict[str, str]]:
|
||||
# ── Reading both sides ──────────────────────────────────────────────────────
|
||||
|
||||
async def _backup_rows(dump: Path) -> dict[str, tuple[list[str], list[list[str | None]]]]:
|
||||
"""Every table's rows in the dump, as pg_restore would load them. Run
|
||||
without a database: pg_restore writes the script and nothing connects."""
|
||||
try:
|
||||
proc = await asyncio.create_subprocess_exec(
|
||||
"pg_restore", "--data-only", "--no-owner", "--no-privileges", "--file=-", str(dump),
|
||||
stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
|
||||
except FileNotFoundError:
|
||||
raise BackupError("pg_restore isn't installed here.") from None
|
||||
try:
|
||||
out, err = await asyncio.wait_for(proc.communicate(), CHECK_DEADLINE_SECONDS)
|
||||
except TimeoutError:
|
||||
proc.kill()
|
||||
await proc.wait()
|
||||
raise BackupError(f"Reading the backup didn't finish within "
|
||||
f"{CHECK_DEADLINE_SECONDS // 60} minutes.") from None
|
||||
if proc.returncode != 0:
|
||||
lines = [ln for ln in err.decode(errors="replace").splitlines() if ln.strip()]
|
||||
raise BackupError("The backup can't be read: " + (lines[-1] if lines else
|
||||
f"pg_restore exited {proc.returncode}."))
|
||||
return script_tables(out)
|
||||
"""Every table's rows in the dump, as pg_restore would load them."""
|
||||
return script_tables(await restore_script(dump))
|
||||
|
||||
|
||||
async def _here(database_url: str, tables: list[str]):
|
||||
|
||||
@@ -0,0 +1,131 @@
|
||||
"""A backup's rows as COPY text: read out of the dump by pg_restore, without a
|
||||
database, and parsed or handed on as they are. The check (check.py) parses
|
||||
them to compare; the in-place restore (restore.py) hands them to COPY
|
||||
unparsed, so what goes back in is byte for byte what the dump holds.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
from .dump import BackupError
|
||||
|
||||
SCRIPT_DEADLINE_SECONDS = 5 * 60
|
||||
_COPY = re.compile(r'^COPY ([^ ]+) \((.*)\) FROM stdin;$')
|
||||
_ESCAPES = {ord("b"): b"\b", ord("f"): b"\f", ord("n"): b"\n", ord("r"): b"\r",
|
||||
ord("t"): b"\t", ord("v"): b"\v"}
|
||||
_HEX = "0123456789abcdefABCDEF"
|
||||
|
||||
|
||||
def async_url(database_url: str) -> str:
|
||||
from sqlalchemy.engine import make_url
|
||||
url = make_url(database_url)
|
||||
return str(url.set(drivername="postgresql+asyncpg")) if url.drivername == "postgresql" else database_url
|
||||
|
||||
|
||||
async def restore_script(dump: Path) -> bytes:
|
||||
"""The dump's rows as pg_restore would load them: its --data-only script.
|
||||
Nothing connects; pg_restore only writes the script."""
|
||||
try:
|
||||
proc = await asyncio.create_subprocess_exec(
|
||||
"pg_restore", "--data-only", "--no-owner", "--no-privileges", "--file=-", str(dump),
|
||||
stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
|
||||
except FileNotFoundError:
|
||||
raise BackupError("pg_restore isn't installed here.") from None
|
||||
try:
|
||||
out, err = await asyncio.wait_for(proc.communicate(), SCRIPT_DEADLINE_SECONDS)
|
||||
except TimeoutError:
|
||||
proc.kill()
|
||||
await proc.wait()
|
||||
raise BackupError(f"Reading the backup didn't finish within "
|
||||
f"{SCRIPT_DEADLINE_SECONDS // 60} minutes.") from None
|
||||
if proc.returncode != 0:
|
||||
lines = [ln for ln in err.decode(errors="replace").splitlines() if ln.strip()]
|
||||
raise BackupError("The backup can't be read: " + (lines[-1] if lines else
|
||||
f"pg_restore exited {proc.returncode}."))
|
||||
return out
|
||||
|
||||
|
||||
def _ident(name: str) -> str:
|
||||
name = name.strip()
|
||||
return name[1:-1].replace('""', '"') if name.startswith('"') and name.endswith('"') else name
|
||||
|
||||
|
||||
def script_blocks(script: bytes) -> dict[str, tuple[list[str], bytes]]:
|
||||
"""Each table's (columns, COPY text) from a --data-only script, the text
|
||||
as it stands: one line per row, ready for COPY … FROM in text format."""
|
||||
out: dict[str, tuple[list[str], bytes]] = {}
|
||||
lines = iter(script.split(b"\n"))
|
||||
for line in lines:
|
||||
m = _COPY.match(line.decode("utf-8", errors="replace"))
|
||||
if not m:
|
||||
continue
|
||||
table = _ident(m.group(1).split(".")[-1])
|
||||
columns = [_ident(c) for c in m.group(2).split(",")]
|
||||
body = []
|
||||
for row in lines:
|
||||
if row == b"\\.":
|
||||
break
|
||||
body.append(row + b"\n")
|
||||
out[table] = (columns, b"".join(body))
|
||||
return out
|
||||
|
||||
|
||||
def unescape(field: bytes) -> str | None:
|
||||
"""One field of COPY's text format: \\N is NULL, backslash escapes are
|
||||
undone (named, octal \\ooo and hex \\xhh)."""
|
||||
if field == b"\\N":
|
||||
return None
|
||||
if b"\\" not in field:
|
||||
return field.decode("utf-8", errors="replace")
|
||||
out = bytearray()
|
||||
i, n = 0, len(field)
|
||||
while i < n:
|
||||
c = field[i]
|
||||
if c != 0x5C or i + 1 == n:
|
||||
out.append(c)
|
||||
i += 1
|
||||
continue
|
||||
nxt = field[i + 1]
|
||||
if nxt in _ESCAPES:
|
||||
out += _ESCAPES[nxt]
|
||||
i += 2
|
||||
elif 0x30 <= nxt <= 0x37:
|
||||
j = i + 1
|
||||
while j < n and j < i + 4 and 0x30 <= field[j] <= 0x37:
|
||||
j += 1
|
||||
out.append(int(field[i + 1:j], 8) & 0xFF)
|
||||
i = j
|
||||
elif nxt == ord("x") and i + 2 < n and chr(field[i + 2]) in _HEX:
|
||||
j = i + 2
|
||||
while j < n and j < i + 4 and chr(field[j]) in _HEX:
|
||||
j += 1
|
||||
out.append(int(field[i + 2:j], 16))
|
||||
i = j
|
||||
else:
|
||||
out.append(nxt)
|
||||
i += 2
|
||||
return bytes(out).decode("utf-8", errors="replace")
|
||||
|
||||
|
||||
def copy_rows(data: bytes) -> list[list[str | None]]:
|
||||
"""COPY text rows (no header, no terminator) as lists of fields."""
|
||||
return [[unescape(f) for f in line.split(b"\t")]
|
||||
for line in data.split(b"\n") if line and line != b"\\."]
|
||||
|
||||
|
||||
def script_tables(script: bytes) -> dict[str, tuple[list[str], list[list[str | None]]]]:
|
||||
"""Each table's (columns, rows) from a --data-only script."""
|
||||
return {t: (cols, copy_rows(raw)) for t, (cols, raw) in script_blocks(script).items()}
|
||||
|
||||
|
||||
def script_from(blocks: dict[str, tuple[list[str], bytes]]) -> bytes:
|
||||
"""A --data-only script holding `blocks`: the shape pg_restore writes,
|
||||
so a test can stand in for a dump."""
|
||||
out = []
|
||||
for table, (columns, raw) in blocks.items():
|
||||
out.append(f"COPY public.{table} ({', '.join(columns)}) FROM stdin;\n".encode())
|
||||
out.append(raw)
|
||||
out.append(b"\\.\n\n")
|
||||
return b"".join(out)
|
||||
+225
-7
@@ -1,4 +1,6 @@
|
||||
"""Bringing Steward back from a backup file into an empty database.
|
||||
"""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
|
||||
@@ -13,6 +15,7 @@ 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
|
||||
@@ -23,6 +26,7 @@ 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__)
|
||||
|
||||
@@ -147,12 +151,6 @@ async def _pg_restore(database_url: str, dump: Path, section: str) -> None:
|
||||
raise BackupError(f"Restoring {section}: " + (lines[-1] if lines else f"{name} exited {proc.returncode}."))
|
||||
|
||||
|
||||
def _async_url(database_url: str) -> str:
|
||||
from sqlalchemy.engine import make_url
|
||||
url = make_url(database_url)
|
||||
return str(url.set(drivername="postgresql+asyncpg")) if url.drivername == "postgresql" else database_url
|
||||
|
||||
|
||||
async def _tables(database_url: str) -> list[str]:
|
||||
engine = create_async_engine(_async_url(database_url))
|
||||
try:
|
||||
@@ -195,3 +193,223 @@ async def restore(path: Path, database_url: str, *, wipe: bool = False,
|
||||
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)
|
||||
|
||||
Reference in New Issue
Block a user