interactive-story/backend/tests/test_snapshot_compression.py

"""context_snapshot, stored compressed.

The column is 89% of the database and the free tier allows 512 MB. Reads were
solved by deferring it; this is about the storage ceiling. Postgres already
TOASTs it and only gets 1.7x, because pglz is tuned for fast decompression of
data a query might filter on — and nothing ever filters on an assembled
prompt.

Three things have to hold, and only the first is obvious:

* what goes in comes back out, exactly, including a snapshot written before
  the conversion and one that is NULL;
* the model still hands callers a dict, so no call site changes;
* migration 44 drops the original column, so the backfill is the one
  destructive step in this file — it must convert every row or abort.

    python -m pytest tests/test_snapshot_compression.py -v
"""
import json
import os
import random
import tempfile
import zlib

_tmp = tempfile.NamedTemporaryFile(suffix=".db", delete=False)
_tmp.close()
os.environ["AIDND_DB_PATH"] = _tmp.name
os.environ.pop("AIDND_DATABASE_URL", None)
os.environ.pop("DATABASE_URL", None)

import pytest
from sqlalchemy import text

from app import compression, migrations, models
from app.database import Base, SessionLocal, engine
from tests import schema_rewind
from tools.fakeprose import prose


def snapshot(seed: int, nbytes: int = 60_000) -> dict:
    rng = random.Random(seed)
    return {
        "system": prose(rng, nbytes // 3),
        "story": prose(rng, nbytes - nbytes // 3),
        "world_state": {"delta": {"player.hp": -3}},
    }


@pytest.fixture()
def db():
    Base.metadata.create_all(bind=engine)
    session = SessionLocal()
    try:
        yield session
    finally:
        session.close()
        Base.metadata.drop_all(bind=engine)


@pytest.fixture()
def adventure(db):
    user = models.User(is_guest=False, email="snap@example.com")
    db.add(user)
    db.flush()
    adv = models.Adventure(user_id=user.id, title="Cave", script_state={})
    db.add(adv)
    db.commit()
    return adv


# ------------------------------------------------------------------- packing

def test_pack_round_trips_exactly():
    value = snapshot(1)
    assert compression.unpack(compression.pack(value)) == value


def test_pack_handles_the_awkward_values():
    for value in ({}, {"a": None}, {"nested": {"deep": [1, 2, {"x": "é"}]}}):
        assert compression.unpack(compression.pack(value)) == value


def test_pack_actually_shrinks_prose():
    """The whole justification. If this ratio collapses, the migration is
    spending CPU for nothing."""
    value = snapshot(2, 200_000)
    raw = json.dumps(value, separators=(",", ":")).encode()
    packed = compression.pack(value)
    assert len(packed) < len(raw) / 3, (
        f"{len(raw):,} B compressed to {len(packed):,} B — under 3x"
    )


def test_unpack_rejects_nothing_it_wrote():
    packed = compression.pack({"a": "b"})
    assert compression.unpack(bytearray(packed)) == {"a": "b"}
    assert compression.unpack(memoryview(bytes(packed))) == {"a": "b"}


# --------------------------------------------------------------- the column

def test_the_column_stores_bytes_and_returns_a_dict(db, adventure):
    value = snapshot(3)
    action = models.Action(
        adventure_id=adventure.id, index=0, type="ai", text="t",
        context_snapshot=value,
    )
    db.add(action)
    db.commit()
    db.expire_all()

    assert action.context_snapshot == value

    stored = db.execute(
        text("SELECT context_snapshot FROM actions WHERE id = :id"),
        {"id": action.id},
    ).scalar()
    assert isinstance(stored, (bytes, bytearray, memoryview))
    assert json.loads(zlib.decompress(bytes(stored))) == value


def test_the_column_is_smaller_than_the_json_it_holds(db, adventure):
    value = snapshot(4, 200_000)
    action = models.Action(
        adventure_id=adventure.id, index=0, type="ai", text="t",
        context_snapshot=value,
    )
    db.add(action)
    db.commit()

    stored = db.execute(
        text("SELECT length(context_snapshot) FROM actions WHERE id = :id"),
        {"id": action.id},
    ).scalar()
    assert stored < len(json.dumps(value)) / 3


def test_null_stays_null(db, adventure):
    action = models.Action(
        adventure_id=adventure.id, index=0, type="do", text="t",
        context_snapshot=None,
    )
    db.add(action)
    db.commit()
    db.expire_all()
    assert action.context_snapshot is None


def test_an_unreadable_snapshot_reads_as_none_rather_than_raising(db, adventure):
    """One corrupt row must not 500 the turn that happens to load it. The
    snapshot is a debugging view; the story is the thing that matters."""
    action = models.Action(
        adventure_id=adventure.id, index=0, type="ai", text="t",
        context_snapshot={"a": "b"},
    )
    db.add(action)
    db.commit()
    db.execute(
        text("UPDATE actions SET context_snapshot = :junk WHERE id = :id"),
        {"junk": b"not zlib at all", "id": action.id},
    )
    db.commit()
    db.expire_all()

    assert db.get(models.Action, action.id).context_snapshot is None


# ------------------------------------------------------- the upgrade in full

def as_json_column(db, rows: dict[int, dict]) -> None:
    """The pre-43 schema: context_snapshot as a JSON column, populated."""
    db.execute(text("ALTER TABLE actions DROP COLUMN context_snapshot"))
    db.execute(text("ALTER TABLE actions ADD COLUMN context_snapshot JSON"))
    for action_id, value in rows.items():
        db.execute(
            text("UPDATE actions SET context_snapshot = :v WHERE id = :id"),
            {"v": json.dumps(value), "id": action_id},
        )
    db.commit()


def seed_pre_43(db, adventure, count: int = 4) -> dict[int, dict]:
    ids = []
    for i in range(count):
        action = models.Action(
            adventure_id=adventure.id, index=i, type="ai", text=f"t{i}"
        )
        db.add(action)
        db.flush()
        ids.append(action.id)
    # One action with no snapshot at all, which must survive as NULL.
    plain = models.Action(
        adventure_id=adventure.id, index=count, type="do", text="look"
    )
    db.add(plain)
    db.commit()

    expected = {action_id: snapshot(action_id) for action_id in ids}
    as_json_column(db, expected)
    db.commit()
    # Take the later migrations' columns back off too, not just the stamp:
    # replaying 43-45 also replays everything appended after them.
    schema_rewind.rewind_to(engine, migrations.SNAPSHOT_COMPRESS_VERSION - 1)
    return expected


def test_bootstrap_converts_every_snapshot(db, adventure):
    expected = seed_pre_43(db, adventure)
    db.close()

    migrations.bootstrap(engine)

    check = SessionLocal()
    try:
        for action_id, value in expected.items():
            assert check.get(models.Action, action_id).context_snapshot == value
        assert check.execute(
            text("SELECT count(*) FROM actions WHERE context_snapshot IS NULL")
        ).scalar() == 1
    finally:
        check.close()


def test_bootstrap_leaves_the_column_named_context_snapshot(db, adventure):
    seed_pre_43(db, adventure)
    db.close()

    migrations.bootstrap(engine)

    with engine.begin() as conn:
        columns = {row[1] for row in conn.execute(text("PRAGMA table_info(actions)"))}
    assert "context_snapshot" in columns
    assert "context_snapshot_z" not in columns, "the swap left the scratch column behind"


def test_the_backfill_aborts_rather_than_dropping_unconvertible_data(
    db, adventure, monkeypatch
):
    """Migration 44 destroys the original. If anything cannot be converted the
    whole run has to roll back with the column still there — the alternative is
    losing somebody's prompts to a bug in this file."""
    seed_pre_43(db, adventure)
    db.close()

    def broken_unpack(_blob):
        return {"not": "what went in"}

    monkeypatch.setattr(compression, "unpack", broken_unpack)

    with pytest.raises(RuntimeError, match="round trip"):
        migrations.bootstrap(engine)

    with engine.begin() as conn:
        columns = {row[1] for row in conn.execute(text("PRAGMA table_info(actions)"))}
        version = conn.execute(text("PRAGMA user_version")).scalar()
        surviving = conn.execute(
            text("SELECT count(*) FROM actions WHERE context_snapshot IS NOT NULL")
        ).scalar()

    assert "context_snapshot" in columns
    assert version < migrations.SNAPSHOT_COMPRESS_VERSION, \
        "the version advanced past a backfill that failed"
    assert surviving == 4, "the prompts did not survive the rollback"