permanence/world/db.py

"""
permanence.world.db — mock SQL database with transactional reversibility.

This module simulates the operational semantics of DDL and DML operations
that matter for reversibility prediction. It is not a SQL engine; it models:

    * Tables, rows, primary keys
    * Transactions with BEGIN / COMMIT / ROLLBACK
    * A write-ahead log for committed changes
    * Named snapshots (backups)

All state is in-memory Python. No subprocess, no network, no file I/O.

Reversibility classes encoded:

    R1  ``SELECT``                                    — read-only
    R2  ``INSERT``/``UPDATE``/``DELETE`` inside txn   — rolled back trivially
    R3  ``COMMIT`` of a DML txn                       — reversible via WAL replay
                                                        and/or a prior backup
    R4  ``DROP TABLE`` when a backup exists           — reversible from backup
    R5  ``DROP TABLE`` with no backup, or
        ``TRUNCATE`` + ``COMMIT`` with no backup      — unrecoverable
"""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple
import copy


# ─────────────────────────────────────────────────────────────────────────────
# Data model
# ─────────────────────────────────────────────────────────────────────────────


@dataclass
class Table:
    name: str
    primary_key: str
    rows: Dict[Any, Dict[str, Any]] = field(default_factory=dict)  # pk → row

    def n_rows(self) -> int:
        return len(self.rows)


@dataclass
class TxnOp:
    op: str           # "insert" | "update" | "delete" | "drop" | "truncate"
    table: str
    before: Optional[Any]
    after: Optional[Any]


@dataclass
class DBResult:
    ok: bool
    message: str
    r_level: int
    rows_affected: int = 0


# ─────────────────────────────────────────────────────────────────────────────
# The mock database
# ─────────────────────────────────────────────────────────────────────────────


class MockDatabase:
    """In-memory relational store with transactions and backups.

    State layers:

        * ``tables``        — the committed state
        * ``txn_ops``       — operations staged in the current transaction
        * ``txn_active``    — whether BEGIN has been issued
        * ``wal``           — ordered list of committed TxnOp batches
        * ``backups``       — named snapshots of the full ``tables`` map

    Reversibility is derived from whether a prior state is still reachable
    — via ROLLBACK (uncommitted), via WAL replay (committed), or via a
    backup (DROP/TRUNCATE recovery).
    """

    def __init__(self) -> None:
        self.tables: Dict[str, Table] = {}
        self.txn_ops: List[TxnOp] = []
        self.txn_active: bool = False
        self.wal: List[List[TxnOp]] = []
        self.backups: Dict[str, Dict[str, Table]] = {}

    # ─── Helpers ──────────────────────────────────────────────────────────

    def _require_table(self, name: str) -> Optional[Table]:
        return self.tables.get(name)

    def _record_op(self, op: TxnOp) -> None:
        if self.txn_active:
            self.txn_ops.append(op)
        else:
            # Implicit autocommit: single op goes straight to WAL
            self.wal.append([op])

    def _backup_contains_table(self, name: str) -> bool:
        return any(name in snap for snap in self.backups.values())

    # ─── DDL ──────────────────────────────────────────────────────────────

    def create_table(self, name: str, primary_key: str) -> DBResult:
        if name in self.tables:
            return DBResult(False, f"table exists: {name}", r_level=1)
        self.tables[name] = Table(name=name, primary_key=primary_key)
        return DBResult(True, f"created table {name}", r_level=2)

    def drop_table(self, name: str) -> DBResult:
        t = self._require_table(name)
        if t is None:
            return DBResult(False, f"no such table: {name}", r_level=1)
        before = copy.deepcopy(t)
        del self.tables[name]
        op = TxnOp(op="drop", table=name, before=before, after=None)
        self._record_op(op)

        # DROP cannot be rolled back inside most SQL engines, but we model
        # recovery via backup: R4 if a snapshot contains it, R5 otherwise.
        r = 4 if self._backup_contains_table(name) else 5
        return DBResult(
            True,
            f"dropped {name} ({t.n_rows()} rows, {'recoverable from backup' if r == 4 else 'UNRECOVERABLE'})",
            r_level=r,
            rows_affected=t.n_rows(),
        )

    def truncate(self, name: str) -> DBResult:
        t = self._require_table(name)
        if t is None:
            return DBResult(False, f"no such table: {name}", r_level=1)
        n = t.n_rows()
        before = copy.deepcopy(t.rows)
        t.rows = {}
        op = TxnOp(op="truncate", table=name, before=before, after=None)
        if self.txn_active:
            self.txn_ops.append(op)
            return DBResult(True, f"truncated {name} (uncommitted)", r_level=2, rows_affected=n)
        # Auto-committed truncate: recovery depends on backup
        self.wal.append([op])
        r = 4 if self._backup_contains_table(name) else 5
        return DBResult(
            True,
            f"truncated {name} ({'backed up' if r == 4 else 'UNRECOVERABLE'})",
            r_level=r,
            rows_affected=n,
        )

    # ─── DML ──────────────────────────────────────────────────────────────

    def insert(self, table: str, row: Dict[str, Any]) -> DBResult:
        t = self._require_table(table)
        if t is None:
            return DBResult(False, f"no such table: {table}", r_level=1)
        pk = row.get(t.primary_key)
        if pk is None:
            return DBResult(False, f"missing primary key {t.primary_key}", r_level=1)
        if pk in t.rows:
            return DBResult(False, f"duplicate pk: {pk}", r_level=1)
        t.rows[pk] = dict(row)
        self._record_op(TxnOp(op="insert", table=table, before=None, after=pk))
        # Inside a txn this is R2; autocommitted it becomes R3 (reversible
        # via WAL replay to a snapshot, but not trivially).
        r = 2 if self.txn_active else 3
        return DBResult(True, f"inserted 1 into {table}", r_level=r, rows_affected=1)

    def update(self, table: str, pk: Any, updates: Dict[str, Any]) -> DBResult:
        t = self._require_table(table)
        if t is None:
            return DBResult(False, f"no such table: {table}", r_level=1)
        if pk not in t.rows:
            return DBResult(False, f"no row with pk={pk}", r_level=1)
        before = copy.deepcopy(t.rows[pk])
        t.rows[pk].update(updates)
        self._record_op(TxnOp(op="update", table=table, before=before, after=pk))
        r = 2 if self.txn_active else 3
        return DBResult(True, f"updated pk={pk} in {table}", r_level=r, rows_affected=1)

    def delete(self, table: str, pk: Any) -> DBResult:
        t = self._require_table(table)
        if t is None:
            return DBResult(False, f"no such table: {table}", r_level=1)
        if pk not in t.rows:
            return DBResult(False, f"no row with pk={pk}", r_level=1)
        before = t.rows.pop(pk)
        self._record_op(TxnOp(op="delete", table=table, before=before, after=None))
        r = 2 if self.txn_active else 3
        return DBResult(True, f"deleted pk={pk} from {table}", r_level=r, rows_affected=1)

    def select(self, table: str, pk: Optional[Any] = None) -> DBResult:
        t = self._require_table(table)
        if t is None:
            return DBResult(False, f"no such table: {table}", r_level=1)
        if pk is not None:
            if pk not in t.rows:
                return DBResult(False, f"no row with pk={pk}", r_level=1)
            return DBResult(True, str(t.rows[pk]), r_level=1, rows_affected=1)
        return DBResult(True, f"{t.n_rows()} rows", r_level=1, rows_affected=t.n_rows())

    # ─── Transactions ─────────────────────────────────────────────────────

    def begin(self) -> DBResult:
        if self.txn_active:
            return DBResult(False, "transaction already active", r_level=1)
        self.txn_active = True
        self.txn_ops = []
        return DBResult(True, "BEGIN", r_level=1)

    def commit(self) -> DBResult:
        if not self.txn_active:
            return DBResult(False, "no active transaction", r_level=1)
        ops = self.txn_ops
        self.txn_ops = []
        self.txn_active = False
        if ops:
            self.wal.append(ops)
        # Commit of DML is R3 by default (WAL replay possible but non-trivial);
        # commit of a DROP/TRUNCATE escalates based on backup presence.
        highest_r = 3
        for op in ops:
            if op.op in ("drop", "truncate"):
                if not self._backup_contains_table(op.table):
                    highest_r = max(highest_r, 5)
                else:
                    highest_r = max(highest_r, 4)
        return DBResult(True, f"COMMIT ({len(ops)} ops)", r_level=highest_r)

    def rollback(self) -> DBResult:
        if not self.txn_active:
            return DBResult(False, "no active transaction", r_level=1)
        # Replay txn_ops in reverse to undo them on ``self.tables``.
        for op in reversed(self.txn_ops):
            t = self.tables.get(op.table)
            if op.op == "insert" and t is not None and op.after in t.rows:
                del t.rows[op.after]
            elif op.op == "update" and t is not None and op.before is not None:
                t.rows[op.after] = op.before
            elif op.op == "delete" and t is not None and op.before is not None:
                t.rows[op.before[t.primary_key]] = op.before
            elif op.op == "drop" and op.before is not None:
                self.tables[op.table] = op.before
            elif op.op == "truncate" and op.before is not None and t is not None:
                t.rows = dict(op.before)
        self.txn_ops = []
        self.txn_active = False
        return DBResult(True, "ROLLBACK", r_level=2)

    # ─── Backups ──────────────────────────────────────────────────────────

    def snapshot(self, snap_id: str) -> DBResult:
        self.backups[snap_id] = {
            n: Table(name=n, primary_key=t.primary_key, rows=copy.deepcopy(t.rows))
            for n, t in self.tables.items()
        }
        return DBResult(True, f"snapshot {snap_id} ({len(self.tables)} tables)", r_level=2)

    def restore(self, snap_id: str) -> DBResult:
        if snap_id not in self.backups:
            return DBResult(False, f"no such snapshot: {snap_id}", r_level=1)
        self.tables = {
            n: Table(name=t.name, primary_key=t.primary_key, rows=dict(t.rows))
            for n, t in self.backups[snap_id].items()
        }
        return DBResult(True, f"restored from {snap_id}", r_level=2)

    # ─── Introspection ────────────────────────────────────────────────────

    def summary(self) -> Dict[str, int]:
        return {
            "tables": len(self.tables),
            "rows": sum(t.n_rows() for t in self.tables.values()),
            "wal_entries": len(self.wal),
            "backups": len(self.backups),
            "txn_active": int(self.txn_active),
        }