FabledSteward/tests/integration/test_backup_restore_in_place.py

"""Integration: a backup put back in place on a running Steward (#5815). It
replaces what the backup carries and keeps the rest: settings, hosts and
rules come back as they were; the audit log and the history of what is
still there stay; history of what didn't come back goes with it. A failure
part-way changes nothing. pg_restore isn't in this lane, so the backup's
rows are this database's own, read out as COPY text before the changes.
Requires STEWARD_DATABASE_URL."""
from __future__ import annotations

import asyncio
import io
import os
import uuid
from datetime import datetime, timezone

import pytest

pytestmark = pytest.mark.integration


@pytest.fixture
def app():
    if not os.environ.get("STEWARD_DATABASE_URL"):
        pytest.skip("needs Postgres")
    from steward.app import create_app
    return create_app(testing=False)


async def _snapshot(app) -> bytes:
    """The kept tables as they stand, as a --data-only script."""
    from sqlalchemy import text
    from steward.backup import contents
    from steward.backup.pgdata import script_from
    blocks = {}
    async with app._db_engine.connect() as conn:
        order: dict[str, list[str]] = {}
        for t, c in (await conn.execute(text(
                "SELECT table_name, column_name FROM information_schema.columns "
                "WHERE table_schema = 'public' AND table_name = ANY(:t) "
                "ORDER BY table_name, ordinal_position"), {"t": sorted(contents.KEPT)})).all():
            order.setdefault(t, []).append(c)
        raw = (await conn.get_raw_connection()).driver_connection
        for t, cols in order.items():
            buf = io.BytesIO()
            await raw.copy_from_table(t, columns=cols, output=buf, format="text")
            blocks[t] = (cols, buf.getvalue())
    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")
    path = tmp_path / "steward-20261015T030000Z.tar"
    heads = heads if heads is not None else await job.migration_heads(app.config["DATABASE_URL"])
    dump.pack(path, tmp_path / "steward.dump", app.config["SECRET_KEY"],
              dump.manifest(datetime.now(timezone.utc), heads))
    return path


def test_a_backup_goes_back_in_place_and_history_stays(app, monkeypatch, tmp_path):
    from sqlalchemy import func, select
    from steward.backup import pgdata
    from steward.backup.restore import restore_in_place
    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
    from steward.models.users import User, UserRole

    tag = uuid.uuid4().hex[:8]
    setting = f"test.in_place_{tag}"

    def rule(name, uid):
        return AlertRule(name=name, source_module="host_agent", resource_name=f"r-{tag}",
                         metric_name="cpu_pct", operator=AlertOperator.gt, threshold=90.0,
                         consecutive_failures_required=1, enabled=True, created_by=uid)

    async def _go():
        async with app.db_sessionmaker() as s:
            async with s.begin():
                user = User(username=f"ip-{tag}", email=f"ip-{tag}@test", password_hash="x",
                            role=UserRole.admin)
                s.add(user)
                await s.flush()
                await set_setting(s, setting, "before")
                kept_host = Host(name=f"kept-{tag}", address="10.7.0.1")
                kept_rule = rule(f"kept {tag}", user.id)
                s.add_all([kept_host, kept_rule])
                await s.flush()
                s.add(AlertEvent(rule_id=kept_rule.id, from_state="inactive", to_state="firing"))
        ids = {"user": user.id, "kept_host": kept_host.id, "kept_rule": kept_rule.id}
        script = await _snapshot(app)

        async def restore_script(_dump):
            return script
        monkeypatch.setattr(pgdata, "restore_script", restore_script)
        path = await _backup_file(app, tmp_path)

        async with app.db_sessionmaker() as s:
            async with s.begin():
                await set_setting(s, setting, "after")
                await set_setting(s, "backup.last_attempt_at", f"kept-{tag}")   # Steward's own record
                new_host = Host(name=f"new-{tag}", address="10.7.0.2")
                new_rule = rule(f"new {tag}", ids["user"])
                s.add_all([new_host, new_rule])
                await s.flush()
                s.add(AlertEvent(rule_id=new_rule.id, from_state="inactive", to_state="firing"))
                s.add(AuditEvent(username="test", action=f"test.{tag}"))
        ids |= {"new_host": new_host.id, "new_rule": new_rule.id}

        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:
            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(
                AlertEvent.rule_id.in_([ids["kept_rule"], ids["new_rule"]]))
                .group_by(AlertEvent.rule_id))).all())
            audits = (await s.execute(select(func.count()).select_from(AuditEvent).where(
                AuditEvent.action == f"test.{tag}"))).scalar_one()
        return ids, done, settings, hosts, events, audits

    ids, done, settings, hosts, events, audits = asyncio.run(_go())
    assert settings[setting] == "before"                       # settings came back
    assert settings["backup.last_attempt_at"] == f"kept-{tag}"  # Steward's own record stayed
    assert hosts == {ids["kept_host"]}                          # the new host is gone
    assert events == {ids["kept_rule"]: 1}                      # its rule's history went with it
    assert audits == 1                                          # the audit log stayed
    assert "app_settings" in done.tables and "audit_events" not in done.tables
    assert done.fixed >= 1


def test_a_restore_that_fails_part_way_changes_nothing(app, monkeypatch, tmp_path):
    from steward.backup import pgdata
    from steward.backup.dump import BackupError
    from steward.backup.restore import restore_in_place
    from steward.core.settings import set_setting

    tag = uuid.uuid4().hex[:8]
    setting = f"test.undone_{tag}"

    async def _go():
        async with app.db_sessionmaker() as s:
            async with s.begin():
                await set_setting(s, setting, "here")
        script = await _snapshot(app)
        # A monitor row with too few fields: COPY refuses it after the
        # settings have already been replaced in the same transaction.
        script += b"COPY public.monitors (id, name) FROM stdin;\nonly-one-field\n\\.\n"

        async def restore_script(_dump):
            return script
        monkeypatch.setattr(pgdata, "restore_script", restore_script)
        path = await _backup_file(app, tmp_path)
        async with app.db_sessionmaker() as s:
            async with s.begin():
                await set_setting(s, setting, "changed")
        with pytest.raises(BackupError, match="undone, nothing changed"):
            await restore_in_place(path, app.config["DATABASE_URL"])
        return await _setting(app, setting)

    assert asyncio.run(_go()) == "changed"


def test_a_backup_from_another_schema_is_refused_in_place(app, tmp_path):
    from steward.backup.restore import NotSameSchema, restore_in_place

    async def _go():
        path = await _backup_file(app, tmp_path, heads=["9999_from_the_future"])
        with pytest.raises(NotSameSchema) as caught:
            await restore_in_place(path, app.config["DATABASE_URL"])
        return caught.value

    refused = asyncio.run(_go())
    assert refused.schema == "newer" and "newer Steward" in str(refused)