Rifqi Hafizuddin
[KM-691] Bound traceability payload + reject contradictory IR filter ranges
db6db73 | """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 typing import Any | |
| 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)})" | |
| ) | |
| self._reject_contradictory_ranges(ir) | |
| 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 _reject_contradictory_ranges(ir: QueryIR) -> None: | |
| """Reject disjoint BETWEEN ranges on the same column. | |
| All filters in an IR are ANDed, so two non-overlapping BETWEEN ranges on one | |
| column (e.g. Q1 2025 AND Q1 2026 on order_date) match no rows — a common LLM | |
| slip when it means a multi-period comparison. Rejecting lets the planner's | |
| re-prompt retry correct it (one spanning range + group-by the period, or one | |
| query per period). Best-effort: incomparable values are skipped, not raised. | |
| """ | |
| ranges: dict[str, list[tuple[Any, Any]]] = {} | |
| for f in ir.filters: | |
| if f.op == "between" and isinstance(f.value, list) and len(f.value) == 2: | |
| ranges.setdefault(f.column_id, []).append((f.value[0], f.value[1])) | |
| for col_id, rs in ranges.items(): | |
| if len(rs) < 2: | |
| continue | |
| try: | |
| # AND of ranges = intersection [max(lo), min(hi)]; empty when lo > hi. | |
| empty = max(lo for lo, _ in rs) > min(hi for _, hi in rs) | |
| except TypeError: | |
| continue # incomparable value types — skip (fail-open) | |
| if empty: | |
| raise IRValidationError( | |
| f"filters place {len(rs)} non-overlapping BETWEEN ranges on column " | |
| f"{col_id!r}; because all filters are ANDed this matches no rows. " | |
| "For a multi-period comparison use a single spanning range and group " | |
| "by the period, or emit one query per period." | |
| ) | |
| 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) | |
| 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 | |
| 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]})" | |
| ) | |
| 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]})" | |
| ) | |
| 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 | |