utils/sql_drills/sql_drills/_generate.py

"""Validate and save a generated drill; the reference SQL is never stored."""

from __future__ import annotations

from dataclasses import dataclass
import hashlib
import json
from typing import TYPE_CHECKING, Any

from sql_drills import _paths, _sandbox
from sql_drills._model import TOPICS, Table

if TYPE_CHECKING:
    from collections.abc import Mapping

    from sql_drills._normalize import Rows

KINDS = ("write", "predict")
PREDICT_HOWTO = (
    "Predict the exact output of the query before running it: one grid row per\n"
    "result row, `NULL` for SQL NULL, no header. If the query fails, type `ERROR`\n"
    "in the first cell."
)
_TEXT_FIELDS = ("title", "task_md", "schema_sql", "reference_sql")


class SpecError(ValueError):
    pass


@dataclass(frozen=True)
class DrillSpec:
    kind: str
    topic: str
    title: str
    task_md: str
    schema_sql: str
    reference_sql: str  # write: the solution; predict: the query to predict
    expect_error: bool = False

    def __post_init__(self) -> None:
        for name in _TEXT_FIELDS:
            value = getattr(self, name)
            if not isinstance(value, str) or not value.strip():
                msg = f"{name} must be a non-empty string"
                raise SpecError(msg)
        if self.kind not in KINDS:
            msg = f"kind must be one of {KINDS}"
            raise SpecError(msg)
        if self.topic not in TOPICS:
            msg = f"topic must be one of {TOPICS}"
            raise SpecError(msg)
        if self.expect_error and self.kind != "predict":
            msg = "expect_error only applies to a predict drill"
            raise SpecError(msg)

    @classmethod
    def from_mapping(cls, raw: Mapping[str, object]) -> DrillSpec:
        missing = [k for k in ("kind", "topic", *_TEXT_FIELDS) if k not in raw]
        if missing:
            msg = f"missing field: {', '.join(missing)}"
            raise SpecError(msg)
        known = ("kind", "topic", "expect_error", *_TEXT_FIELDS)
        fields: dict[str, Any] = {k: raw[k] for k in known if k in raw}
        return cls(**fields)

    def digest(self) -> str:
        blob = json.dumps(self.__dict__, sort_keys=True).encode()
        return hashlib.sha256(blob).hexdigest()[:8]


@dataclass(frozen=True)
class _Reference:
    ok: bool
    code: int | None
    error: str | None
    rows: Rows
    tables: dict[str, Table]


def _run_reference(spec: DrillSpec) -> _Reference:
    with _sandbox.sandbox(spec.schema_sql) as box:
        done = box.execute(spec.reference_sql)
        tables = box.tables()
    return _Reference(done.ok, done.code, done.error, done.rows, tables)


def _validated(spec: DrillSpec) -> _Reference:
    first = _run_reference(spec)
    if not first.ok and not spec.expect_error:
        msg = f"reference SQL failed: {first.error}"
        raise SpecError(msg)
    if first.ok and spec.expect_error:
        msg = "reference SQL did not fail but expect_error is set"
        raise SpecError(msg)
    if not first.ok and first.code is None:
        msg = f"no MySQL error code in: {first.error}"
        raise SpecError(msg)
    if first.ok and not first.rows:
        msg = "reference output is empty"
        raise SpecError(msg)
    again = _run_reference(spec)
    if (again.ok, again.code, again.rows) != (first.ok, first.code, first.rows):
        msg = "reference output is not deterministic"
        raise SpecError(msg)
    return first


def _write(spec: DrillSpec, ref: _Reference, drill_id: str) -> None:
    root = _paths.generated_dir()
    root.mkdir(parents=True, exist_ok=True)
    staging = root / f".tmp-{drill_id}"
    staging.mkdir(exist_ok=True)
    predict = spec.kind == "predict"
    info: dict[str, object] = {"kind": spec.kind, "topic": spec.topic}
    expected = "".join("\t".join(r) + "\n" for r in ref.rows)
    if not ref.ok:
        info |= {"error_code": ref.code, "error_message": ref.error}
        expected = f"ERROR {ref.code}\n"
    task = f"# {spec.title}\n\n{spec.task_md.strip()}\n"
    files = {
        "meta.json": json.dumps(info),
        "DOCS-task.md": task + (f"\n{PREDICT_HOWTO}\n" if predict else ""),
        "schema.sql": spec.schema_sql.strip() + "\n",
        "expected.txt": expected,
        "tables.json": json.dumps({n: t.__dict__ for n, t in ref.tables.items()}),
    }
    if predict:
        files["query.sql"] = spec.reference_sql.strip() + "\n"
    for name, text in files.items():
        (staging / name).write_text(text)
    staging.rename(root / drill_id)


def save_generated(spec: DrillSpec | Mapping[str, object]) -> str:
    """Return the id of the saved drill; raise ``SpecError`` if it is unusable."""
    parsed = spec if isinstance(spec, DrillSpec) else DrillSpec.from_mapping(spec)
    ref = _validated(parsed)
    drill_id = f"g-{parsed.topic}-{parsed.digest()}"
    if not (_paths.generated_dir() / drill_id).exists():
        _write(parsed, ref, drill_id)
    return drill_id