interactive-story/backend/tests/test_snapshot_compression.py

"""context_snapshot, stored compressed.

The column makes up 89% of the database, and the free tier allows only
512 MB. Deferring the column already solved the read cost, so this is
about the storage ceiling. Postgres already TOASTs the column and gets
only a 1.7x ratio, because pglz favors fast decompression for data a query
might filter on. Nothing ever filters on an assembled prompt.

Three things must 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 return a 500 error for the turn that loads
    it. The snapshot is a debugging view. The story is what actually
    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 column. If anything fails to
    convert, the whole run must roll back with the column still there.
    Otherwise a bug in this file could lose someone's prompts."""
    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"