src/query/ir/validator.py

"""IRValidator — checks a QueryIR against a user's catalog.

See ARCHITECTURE.md §7 for the validation rules. On failure, the planner
is re-prompted with the error context (max 3 retries) — error messages
must therefore be specific enough that the LLM can self-correct.
"""

from ...catalog.models import Catalog, Column, Source, Table
from .models import QueryIR
from .operators import (
    ALLOWED_AGG_FNS,
    ALLOWED_FILTER_OPS,
    LIMIT_HARD_CAP,
    TYPE_COMPATIBILITY,
)

_NULLARY_FILTER_OPS = frozenset({"is_null", "is_not_null"})


class IRValidationError(Exception):
    pass


class IRValidator:
    """Reject IRs that reference unknown sources/tables/columns or use disallowed ops.

    Rules:
    - source_id exists in catalog for this user
    - table_id belongs to that source
    - every column_id exists in that table
    - every agg.fn and filter.op is whitelisted (see operators.py)
    - value_type consistent with column.data_type (TYPE_COMPATIBILITY)
    - limit positive int, ≤ LIMIT_HARD_CAP
    """

    def validate(self, ir: QueryIR, catalog: Catalog) -> None:
        source = self._find_source(catalog, ir.source_id)
        base_table = self._find_table(source, ir.table_id)
        # Columns available to select/filter/group/order — grows as joins are added.
        columns_by_id: dict[str, Column] = {c.column_id: c for c in base_table.columns}

        if ir.joins:
            self._validate_joins(ir, source, columns_by_id)

        select_aliases: set[str] = set()
        for i, item in enumerate(ir.select):
            where = f"select[{i}]"
            if item.kind == "column":
                self._require_column(columns_by_id, item.column_id, where)
            else:  # "agg"
                if item.fn not in ALLOWED_AGG_FNS:
                    raise IRValidationError(
                        f"{where}.fn: must be in {sorted(ALLOWED_AGG_FNS)}, "
                        f"got {item.fn!r}"
                    )
                if item.column_id is not None:
                    self._require_column(columns_by_id, item.column_id, where)
                elif item.fn != "count":
                    raise IRValidationError(
                        f"{where}.fn={item.fn!r} requires a column_id "
                        "(only 'count' may omit it for COUNT(*))"
                    )
            if item.alias:
                select_aliases.add(item.alias)

        for i, f in enumerate(ir.filters):
            where = f"filters[{i}]"
            col = self._require_column(columns_by_id, f.column_id, where)
            if f.op not in ALLOWED_FILTER_OPS:
                raise IRValidationError(
                    f"{where}.op: must be in {sorted(ALLOWED_FILTER_OPS)}, "
                    f"got {f.op!r}"
                )
            if f.op not in _NULLARY_FILTER_OPS:
                allowed = TYPE_COMPATIBILITY.get(col.data_type, frozenset())
                if f.value_type not in allowed:
                    raise IRValidationError(
                        f"{where}: value_type {f.value_type!r} incompatible with "
                        f"column.data_type {col.data_type!r} "
                        f"(allowed: {sorted(allowed)})"
                    )

        for i, col_id in enumerate(ir.group_by):
            self._require_column(columns_by_id, col_id, f"group_by[{i}]")

        for i, ob in enumerate(ir.order_by):
            if ob.column_id not in columns_by_id and ob.column_id not in select_aliases:
                raise IRValidationError(
                    f"order_by[{i}].column_id: {ob.column_id!r} not found in table "
                    f"{ir.table_id!r} columns or select aliases "
                    f"(known columns: {sorted(columns_by_id.keys())}, "
                    f"aliases: {sorted(select_aliases)})"
                )

        if ir.limit is not None:
            if ir.limit <= 0:
                raise IRValidationError(f"limit must be positive, got {ir.limit}")
            if ir.limit > LIMIT_HARD_CAP:
                raise IRValidationError(
                    f"limit {ir.limit} exceeds hard cap {LIMIT_HARD_CAP}"
                )

    def _validate_joins(
        self, ir: QueryIR, source: Source, columns_by_id: dict[str, Column]
    ) -> None:
        """Validate joins and grow `columns_by_id` with each joined table's columns.

        Joins are DB-only in v1; each join's column pair must be a declared foreign
        key, the left column must already be in the query, the right in the target.
        """
        if source.source_type != "schema":
            raise IRValidationError(
                f"joins are only supported on database sources in v1; source "
                f"{ir.source_id!r} is {source.source_type!r}"
            )
        fk_pairs = self._fk_pairs(source)
        for i, j in enumerate(ir.joins):
            where = f"joins[{i}]"
            target = self._find_table(source, j.target_table_id)
            target_cols = {c.column_id: c for c in target.columns}
            if j.left_column_id not in columns_by_id:
                raise IRValidationError(
                    f"{where}.left_column_id {j.left_column_id!r} is not in the query "
                    "so far (base table or an earlier join)"
                )
            if j.right_column_id not in target_cols:
                raise IRValidationError(
                    f"{where}.right_column_id {j.right_column_id!r} not in target table "
                    f"{j.target_table_id!r}"
                )
            if frozenset((j.left_column_id, j.right_column_id)) not in fk_pairs:
                raise IRValidationError(
                    f"{where}: ({j.left_column_id!r}, {j.right_column_id!r}) is not a "
                    f"declared foreign key in source {ir.source_id!r} — only FK-backed "
                    "joins are allowed"
                )
            columns_by_id.update(target_cols)

    @staticmethod
    def _fk_pairs(source: Source) -> set[frozenset[str]]:
        """All declared FK column pairs in the source, as unordered {col, col} sets."""
        pairs: set[frozenset[str]] = set()
        for t in source.tables:
            for fk in t.foreign_keys:
                pairs.add(frozenset((fk.column_id, fk.target_column_id)))
        return pairs

    @staticmethod
    def _find_source(catalog: Catalog, source_id: str) -> Source:
        for s in catalog.sources:
            if s.source_id == source_id:
                return s
        raise IRValidationError(
            f"source_id {source_id!r} not in catalog "
            f"(known: {[s.source_id for s in catalog.sources]})"
        )

    @staticmethod
    def _find_table(source: Source, table_id: str) -> Table:
        for t in source.tables:
            if t.table_id == table_id:
                return t
        raise IRValidationError(
            f"table_id {table_id!r} not in source {source.source_id!r} "
            f"(known: {[t.table_id for t in source.tables]})"
        )

    @staticmethod
    def _require_column(
        columns_by_id: dict[str, Column], col_id: str, where: str
    ) -> Column:
        col = columns_by_id.get(col_id)
        if col is None:
            raise IRValidationError(
                f"{where}.column_id: {col_id!r} not in table "
                f"(known: {sorted(columns_by_id.keys())})"
            )
        return col