db.py

"""SQLite access layer.

- Opens the database in read-only URI mode so the agent CANNOT write,
  even if it bypasses the validator.
- Provides schema introspection used by the agent's `get_schema` tool.
"""
from __future__ import annotations

import sqlite3
from pathlib import Path
from typing import Any


class ReadOnlyDB:
    def __init__(self, db_path: str) -> None:
        if not Path(db_path).exists():
            raise FileNotFoundError(
                f"Database not found at {db_path}. Run `python -m src.seed` first."
            )
        # mode=ro -> SQLite enforces read-only at the driver level.
        uri = f"file:{db_path}?mode=ro"
        self._conn = sqlite3.connect(uri, uri=True, check_same_thread=False)
        self._conn.row_factory = sqlite3.Row

    # ---------- Schema introspection ----------

    def list_tables(self) -> list[str]:
        rows = self._conn.execute(
            "SELECT name FROM sqlite_master WHERE type='table' "
            "AND name NOT LIKE 'sqlite_%' ORDER BY name"
        ).fetchall()
        return [r["name"] for r in rows]

    def describe_table(self, table: str) -> dict[str, Any]:
        if table not in self.list_tables():
            raise ValueError(f"Unknown table: {table}")

        cols = self._conn.execute(f"PRAGMA table_info({table})").fetchall()
        fks = self._conn.execute(f"PRAGMA foreign_key_list({table})").fetchall()
        sample = self._conn.execute(f"SELECT * FROM {table} LIMIT 3").fetchall()

        return {
            "table": table,
            "columns": [
                {
                    "name": c["name"],
                    "type": c["type"],
                    "not_null": bool(c["notnull"]),
                    "pk": bool(c["pk"]),
                }
                for c in cols
            ],
            "foreign_keys": [
                {"column": f["from"], "references": f"{f['table']}.{f['to']}"}
                for f in fks
            ],
            "sample_rows": [dict(r) for r in sample],
        }

    def full_schema(self) -> list[dict[str, Any]]:
        return [self.describe_table(t) for t in self.list_tables()]

    # ---------- Execution ----------

    def run_select(self, sql: str, max_rows: int) -> dict[str, Any]:
        """Execute a SELECT-only SQL and return columns + rows + truncation flag."""
        cur = self._conn.execute(sql)
        # SQLite cursor.description gives column metadata for SELECT only.
        cols = [d[0] for d in cur.description] if cur.description else []
        all_rows = cur.fetchall()
        truncated = len(all_rows) > max_rows
        rows = [list(r) for r in all_rows[:max_rows]]
        return {
            "columns": cols,
            "rows": rows,
            "row_count_returned": len(rows),
            "row_count_total": len(all_rows),
            "truncated": truncated,
        }

    def close(self) -> None:
        self._conn.close()