modbison/backend/app/services/world_store.py

"""World content storage service.

Implements the `world.md` §5.1 service interface over Postgres, per
`world_persistence_impl.md` §9 -- **minus** ``flush_dirty`` and
``rebuild_index``, which don't apply once there is no write-behind dirty
set and no derived index to rebuild (see that doc's §3.3 and §4.2).

This is the *only* place chunk content is written (world_persistence_impl.md
D2 / §5). Callers (the routers) own transaction boundaries: methods here
``flush()`` so later statements in the same request see earlier ones and so
integrity errors surface before a caller commits, but they never
``commit()`` -- that mirrors the existing convention in
``app/routers/worlds.py`` (``get_db`` docs: "Handlers are responsible for
committing or rolling back their own transactions").

Known gap this module had to resolve on its own (flagged, not silently
decided): `world_persistence_impl.md` §6/§7.2/§8.1 all reference a world's
monotonic ``content_revision`` ("bumped on every chunk write"), but no
column for it is defined anywhere in §4's schema. Rather than add an
undocumented column, ``content_revision()`` below derives it as
``MAX(id)`` over ``world_chunk_revisions`` scoped to the world: that table
is append-only (retention only nulls ``payload``, it never deletes rows),
so the value is genuinely monotonic and increments on every
create/update/delete/rollback.
"""

from __future__ import annotations

from dataclasses import dataclass

from sqlalchemy import func, select, update
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import load_only

from app.config import settings
from app.models.world_content import (
    MAIN_ROOM_NAME,
    MAIN_ROOM_SLUG,
    WorldChunk,
    WorldChunkRevision,
    WorldRoom,
)
from app.services import chunk_codec


class RevisionConflict(Exception):
    """Optimistic-concurrency conflict (world.md §5.4). ``current_revision``
    is 0 and ``current_checksum`` is None when the caller's ``base_revision``
    assumed a chunk exists but it doesn't."""

    def __init__(
        self,
        current_revision: int,
        current_checksum: str | None,
        coord: tuple[int, int, int] | None = None,
    ) -> None:
        self.current_revision = current_revision
        self.current_checksum = current_checksum
        self.coord = coord
        super().__init__(f"revision conflict: current_revision={current_revision}")


class RevisionNotAvailable(Exception):
    """Raised by ``rollback_chunk`` when the requested revision doesn't
    exist, or exists but its payload was pruned by the retention policy
    (world_persistence_impl.md §4.3)."""


@dataclass
class Bounds:
    cx_min: int | None = None
    cx_max: int | None = None
    cy_min: int | None = None
    cy_max: int | None = None
    cz: int | None = None


class WorldStore:
    """Owns chunk/room read/write operations over the async session
    handed to it. One instance per request is the expected usage
    (``WorldStore(db)`` where ``db`` is the ``DbDep`` session)."""

    def __init__(self, db: AsyncSession) -> None:
        self.db = db

    # ── rooms ────────────────────────────────────────────────────────────

    async def get_room(self, world_id: int, slug: str) -> WorldRoom | None:
        result = await self.db.execute(
            select(WorldRoom).where(WorldRoom.world_id == world_id, WorldRoom.slug == slug)
        )
        return result.scalar_one_or_none()

    async def list_rooms(self, world_id: int) -> list[WorldRoom]:
        result = await self.db.execute(
            select(WorldRoom).where(WorldRoom.world_id == world_id).order_by(WorldRoom.id)
        )
        return list(result.scalars().all())

    async def create_room(
        self,
        world_id: int,
        slug: str,
        name: str,
        *,
        grid: dict | None = None,
        bounds: dict | None = None,
        spawn: dict | None = None,
    ) -> WorldRoom:
        room = WorldRoom(world_id=world_id, slug=slug, name=name)
        if grid is not None:
            room.grid = grid
        if bounds is not None:
            room.bounds = bounds
        if spawn is not None:
            room.spawn = spawn
        self.db.add(room)
        await self.db.flush()
        await self.db.refresh(room)
        return room

    async def get_or_create_main_room(self, world_id: int) -> WorldRoom:
        """world_persistence_impl.md §4.1: every world gets a ``main`` room.
        Called from ``create_world`` for new worlds, and lazily here for
        worlds that predate this migration."""
        room = await self.get_room(world_id, MAIN_ROOM_SLUG)
        if room is not None:
            return room
        try:
            return await self.create_room(world_id, MAIN_ROOM_SLUG, MAIN_ROOM_NAME)
        except IntegrityError:
            # Lost a race with a concurrent creator; the row now exists.
            await self.db.rollback()
            room = await self.get_room(world_id, MAIN_ROOM_SLUG)
            if room is None:
                raise
            return room

    async def count_chunks(self, room_id: int) -> int:
        result = await self.db.execute(
            select(func.count()).select_from(WorldChunk).where(WorldChunk.room_id == room_id)
        )
        return int(result.scalar_one())

    async def content_revision(self, world_id: int) -> int:
        """See module docstring for why this is derived rather than stored."""
        result = await self.db.execute(
            select(func.max(WorldChunkRevision.id)).where(WorldChunkRevision.world_id == world_id)
        )
        return int(result.scalar() or 0)

    # ── chunk reads ──────────────────────────────────────────────────────

    async def get_chunk(self, key: str) -> WorldChunk | None:
        result = await self.db.execute(select(WorldChunk).where(WorldChunk.chunk_key == key))
        return result.scalar_one_or_none()

    async def get_chunk_by_coord(self, room_id: int, cx: int, cy: int, cz: int) -> WorldChunk | None:
        result = await self.db.execute(
            select(WorldChunk).where(
                WorldChunk.room_id == room_id, WorldChunk.cx == cx,
                WorldChunk.cy == cy, WorldChunk.cz == cz,
            )
        )
        return result.scalar_one_or_none()

    async def get_chunk_binary(self, key: str) -> bytes | None:
        chunk = await self.get_chunk(key)
        return chunk.payload if chunk is not None else None

    async def get_chunk_flags(self, key: str) -> dict | None:
        chunk = await self.get_chunk(key)
        return chunk.content_flags if chunk is not None else None

    async def list_chunks(
        self,
        world_id: int,
        room_id: int | None = None,
        bounds: Bounds | None = None,
    ) -> list[WorldChunk]:
        """Index query -- world.md §5.1's ``list_chunks``. Payload bytes are
        deliberately not loaded (world_persistence_impl.md §8.1: "Payloads
        are not included")."""
        query = select(WorldChunk).options(
            load_only(
                WorldChunk.id, WorldChunk.world_id, WorldChunk.room_id, WorldChunk.chunk_key,
                WorldChunk.cx, WorldChunk.cy, WorldChunk.cz, WorldChunk.revision,
                WorldChunk.checksum, WorldChunk.payload_len, WorldChunk.is_mature,
                WorldChunk.updated_at,
            )
        ).where(WorldChunk.world_id == world_id)
        if room_id is not None:
            query = query.where(WorldChunk.room_id == room_id)
        if bounds is not None:
            if bounds.cx_min is not None:
                query = query.where(WorldChunk.cx >= bounds.cx_min)
            if bounds.cx_max is not None:
                query = query.where(WorldChunk.cx <= bounds.cx_max)
            if bounds.cy_min is not None:
                query = query.where(WorldChunk.cy >= bounds.cy_min)
            if bounds.cy_max is not None:
                query = query.where(WorldChunk.cy <= bounds.cy_max)
            if bounds.cz is not None:
                query = query.where(WorldChunk.cz == bounds.cz)
        query = query.order_by(WorldChunk.cx, WorldChunk.cy, WorldChunk.cz)
        result = await self.db.execute(query)
        return list(result.scalars().all())

    # ── chunk writes ─────────────────────────────────────────────────────

    async def put_chunk(
        self,
        *,
        room: WorldRoom,
        cx: int,
        cy: int,
        cz: int,
        payload_bytes: bytes,
        base_revision: int,
        actor_id: int | None,
    ) -> WorldChunk:
        """Full-chunk write with optimistic concurrency (world.md §5.4,
        world_persistence_impl.md §5.4/§8.2). Validates the payload
        server-side (§5.5) before ever touching a row. ``base_revision=0``
        means "create; 409 if it already exists"."""
        key = chunk_codec.build_key(room.world_id, room.slug, cx, cy, cz)
        grid_size = tuple(room.grid.get("size", [32, 32, 1]))

        # Raises ChunkCodecError / ChunkValidationError -- the router maps
        # those to 422. Nothing below trusts the request body's own claims
        # about content_flags/is_mature; both come from `decoded`.
        decoded = chunk_codec.decode_and_validate(
            payload_bytes,
            expected_key=key,
            grid_size=grid_size,
            max_objects=settings.WORLD_CHUNK_MAX_OBJECTS,
        )
        body_checksum = chunk_codec.checksum(payload_bytes[chunk_codec.HEADER_SIZE:])

        result = await self.db.execute(
            select(WorldChunk)
            .where(
                WorldChunk.room_id == room.id, WorldChunk.cx == cx,
                WorldChunk.cy == cy, WorldChunk.cz == cz,
            )
            .with_for_update()
        )
        row = result.scalar_one_or_none()

        if base_revision == 0:
            if row is not None:
                raise RevisionConflict(row.revision, row.checksum, (cx, cy, cz))
            row = WorldChunk(
                world_id=room.world_id, room_id=room.id, chunk_key=key,
                cx=cx, cy=cy, cz=cz, revision=1,
            )
            self.db.add(row)
            action = "create"
        else:
            if row is None or row.revision != base_revision:
                raise RevisionConflict(
                    row.revision if row else 0, row.checksum if row else None, (cx, cy, cz)
                )
            row.revision = row.revision + 1
            action = "update"

        row.format_version = decoded.version
        row.payload = payload_bytes
        row.payload_len = len(payload_bytes)
        row.checksum = body_checksum
        row.content_flags = decoded.flags.to_dict()
        row.is_mature = decoded.flags.mature
        row.updated_by_id = actor_id

        try:
            await self.db.flush()
        except IntegrityError:
            # Lost a create race to a concurrent writer between our SELECT
            # and this flush. Roll back and report the real current state.
            await self.db.rollback()
            existing = await self.get_chunk_by_coord(room.id, cx, cy, cz)
            raise RevisionConflict(
                existing.revision if existing else 0,
                existing.checksum if existing else None,
                (cx, cy, cz),
            )

        await self._record_revision(key, room.world_id, row.revision, action, actor_id, body_checksum, payload_bytes)
        await self.db.refresh(row)
        return row

    async def put_chunks_batch(
        self,
        *,
        room: WorldRoom,
        items: list[dict],
        actor_id: int | None,
    ) -> list[WorldChunk]:
        """Batch write (world_persistence_impl.md §8.1). One transaction,
        all-or-nothing: none of this ``commit()``s, so a ``RevisionConflict``
        raised partway leaves nothing persisted once the caller's request
        handler lets the exception (or the HTTPException it becomes)
        propagate through ``get_db``, which rolls back on any exception."""
        rows: list[WorldChunk] = []
        for item in items:
            row = await self.put_chunk(
                room=room, cx=item["cx"], cy=item["cy"], cz=item["cz"],
                payload_bytes=item["payload_bytes"], base_revision=item["base_revision"],
                actor_id=actor_id,
            )
            rows.append(row)
        return rows

    async def delete_chunk(self, key: str, actor_id: int | None) -> None:
        row = await self.get_chunk(key)
        if row is None:
            return
        await self._record_revision(key, row.world_id, row.revision, "delete", actor_id, row.checksum, None)
        await self.db.delete(row)
        await self.db.flush()

    async def rollback_chunk(self, key: str, revision: int, actor_id: int | None) -> WorldChunk:
        """Restore a chunk to a previously-recorded revision's payload
        (world.md §5.1). Raises ``RevisionNotAvailable`` if that revision
        was never recorded, its payload was pruned by retention, or the
        chunk has since been deleted outright (nothing to update in place)."""
        rev_result = await self.db.execute(
            select(WorldChunkRevision)
            .where(WorldChunkRevision.chunk_key == key, WorldChunkRevision.revision == revision)
            .order_by(WorldChunkRevision.id.desc())
            .limit(1)
        )
        rev_row = rev_result.scalar_one_or_none()
        if rev_row is None:
            raise RevisionNotAvailable(f"no revision {revision} recorded for {key!r}")
        if rev_row.payload is None:
            raise RevisionNotAvailable(f"revision {revision} of {key!r} was pruned; payload unavailable")

        chunk = await self.get_chunk(key)
        if chunk is None:
            raise RevisionNotAvailable(f"chunk {key!r} no longer exists; cannot roll back in place")

        room_result = await self.db.execute(select(WorldRoom).where(WorldRoom.id == chunk.room_id))
        room = room_result.scalar_one()
        grid_size = tuple(room.grid.get("size", [32, 32, 1]))
        decoded = chunk_codec.decode_and_validate(
            rev_row.payload, expected_key=key, grid_size=grid_size,
            max_objects=settings.WORLD_CHUNK_MAX_OBJECTS,
        )
        body_checksum = chunk_codec.checksum(rev_row.payload[chunk_codec.HEADER_SIZE:])

        chunk.revision = chunk.revision + 1
        chunk.payload = rev_row.payload
        chunk.payload_len = len(rev_row.payload)
        chunk.checksum = body_checksum
        chunk.format_version = decoded.version
        chunk.content_flags = decoded.flags.to_dict()
        chunk.is_mature = decoded.flags.mature
        chunk.updated_by_id = actor_id

        await self.db.flush()
        # Tagged "update" (not a fourth action value) to keep the
        # documented 'create'|'update'|'delete' action enum intact.
        await self._record_revision(
            key, chunk.world_id, chunk.revision, "update", actor_id, body_checksum, rev_row.payload
        )
        await self.db.refresh(chunk)
        return chunk

    async def list_revisions(self, chunk_key: str, limit: int = 50) -> list[WorldChunkRevision]:
        """Not in world.md §5.1's interface by name, but required to
        implement the revision-history endpoint world_persistence_impl.md
        §8.1 lists (``GET .../revisions``)."""
        result = await self.db.execute(
            select(WorldChunkRevision)
            .where(WorldChunkRevision.chunk_key == chunk_key)
            .order_by(WorldChunkRevision.revision.desc(), WorldChunkRevision.id.desc())
            .limit(limit)
        )
        return list(result.scalars().all())

    async def migrate(self) -> None:
        """No-op. `world.md` §4.8's hand-rolled SQLite migration runner is
        superseded by Alembic (`backend/alembic/`) now that content lives
        in Postgres (world_persistence_impl.md §3.3). Kept only so this
        class's shape still matches the §5.1 interface it's implementing
        a subset of."""
        return None

    # ── internals ────────────────────────────────────────────────────────

    async def _record_revision(
        self,
        chunk_key: str,
        world_id: int,
        revision: int,
        action: str,
        actor_id: int | None,
        checksum_hex: str,
        payload_bytes: bytes | None,
    ) -> None:
        """Append an audit row and enforce the retention policy
        (world_persistence_impl.md §4.3, required): keep the newest
        ``WORLD_CHUNK_REVISION_KEEP`` payloads per ``chunk_key``, null out
        ``payload`` on older rows, in the same transaction as the write."""
        rev = WorldChunkRevision(
            chunk_key=chunk_key, world_id=world_id, revision=revision, action=action,
            actor_id=actor_id, checksum=checksum_hex, payload=payload_bytes,
        )
        self.db.add(rev)
        await self.db.flush()

        keep = settings.WORLD_CHUNK_REVISION_KEEP
        keep_ids_result = await self.db.execute(
            select(WorldChunkRevision.id)
            .where(WorldChunkRevision.chunk_key == chunk_key)
            .order_by(WorldChunkRevision.id.desc())
            .limit(keep)
        )
        keep_ids = [row[0] for row in keep_ids_result.all()]
        if keep_ids:
            await self.db.execute(
                update(WorldChunkRevision)
                .where(
                    WorldChunkRevision.chunk_key == chunk_key,
                    WorldChunkRevision.id.notin_(keep_ids),
                    WorldChunkRevision.payload.isnot(None),
                )
                .values(payload=None)
            )