diff --git a/steward/backup/check.py b/steward/backup/check.py index 118e6ba..e3357e4 100644 --- a/steward/backup/check.py +++ b/steward/backup/check.py @@ -35,7 +35,7 @@ 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 .dump import read_backup from .pgdata import async_url as _async_url from .pgdata import copy_rows, restore_script, script_tables diff --git a/steward/backup/restore.py b/steward/backup/restore.py index 60f094f..18905eb 100644 --- a/steward/backup/restore.py +++ b/steward/backup/restore.py @@ -321,17 +321,20 @@ async def restore_in_place(path: Path, database_url: str, *, 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()) + # 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) diff --git a/tests/integration/test_backup_restore_in_place.py b/tests/integration/test_backup_restore_in_place.py index 744bda5..27bdae7 100644 --- a/tests/integration/test_backup_restore_in_place.py +++ b/tests/integration/test_backup_restore_in_place.py @@ -47,6 +47,17 @@ async def _snapshot(app) -> bytes: return script_from(blocks) +async def _setting(app, key): + """A setting's stored value (get_all_settings knows only Steward's own keys).""" + import json + + from sqlalchemy import select + from steward.models.settings import AppSetting + async with app.db_sessionmaker() as s: + raw = (await s.execute(select(AppSetting.value_json).where(AppSetting.key == key))).scalar_one() + return json.loads(raw) + + async def _backup_file(app, tmp_path, heads=None): from steward.backup import dump, job (tmp_path / "steward.dump").write_bytes(b"PGDMP") @@ -61,7 +72,7 @@ def test_a_backup_goes_back_in_place_and_history_stays(app, monkeypatch, tmp_pat from sqlalchemy import func, select from steward.backup import pgdata from steward.backup.restore import restore_in_place - from steward.core.settings import get_all_settings, set_setting + from steward.core.settings import set_setting from steward.models.alerts import AlertEvent, AlertOperator, AlertRule from steward.models.audit import AuditEvent from steward.models.hosts import Host @@ -110,8 +121,8 @@ def test_a_backup_goes_back_in_place_and_history_stays(app, monkeypatch, tmp_pat done = await restore_in_place(path, app.config["DATABASE_URL"]) + settings = {k: await _setting(app, k) for k in (setting, "backup.last_attempt_at")} async with app.db_sessionmaker() as s: - settings = await get_all_settings(s) hosts = set((await s.execute(select(Host.id).where( Host.id.in_([ids["kept_host"], ids["new_host"]])))).scalars()) events = dict((await s.execute(select(AlertEvent.rule_id, func.count()).where( @@ -135,7 +146,7 @@ def test_a_restore_that_fails_part_way_changes_nothing(app, monkeypatch, tmp_pat from steward.backup import pgdata from steward.backup.dump import BackupError from steward.backup.restore import restore_in_place - from steward.core.settings import get_all_settings, set_setting + from steward.core.settings import set_setting tag = uuid.uuid4().hex[:8] setting = f"test.undone_{tag}" @@ -158,8 +169,7 @@ def test_a_restore_that_fails_part_way_changes_nothing(app, monkeypatch, tmp_pat await set_setting(s, setting, "changed") with pytest.raises(BackupError, match="undone, nothing changed"): await restore_in_place(path, app.config["DATABASE_URL"]) - async with app.db_sessionmaker() as s: - return (await get_all_settings(s))[setting] + return await _setting(app, setting) assert asyncio.run(_go()) == "changed"