Rifqi Hafizuddin
Claude Fable 5.2
[NOTICKET] fix(planner): query-shape guards for bare measures and grouping by the measure
7d7f748 | """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"}) | |
| _NUMERIC_TYPES = frozenset({"int", "decimal"}) | |
| _PERIOD_NAMES = frozenset({"year", "month", "week", "quarter"}) | |
| def _is_measure(col: Column) -> bool: | |
| """Role of a column for the bare-select message: the dictionary's `classify` when | |
| set, else a number that is neither a period nor a key.""" | |
| role = (col.classify or "").strip().lower() | |
| if role in ("measure", "dimension"): | |
| return role == "measure" | |
| name = col.name.lower() | |
| return ( | |
| col.data_type in _NUMERIC_TYPES | |
| and name not in _PERIOD_NAMES | |
| and name != "id" | |
| and not name.endswith("_id") | |
| and not col.is_key | |
| ) | |
| 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 getattr(item, "computed", None) is not None: | |
| self._validate_computed(item, columns_by_id, where) | |
| 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}]") | |
| # A query mixing bare columns with aggregates must not select a bare column | |
| # that isn't in group_by — the database rejects it only at execution ("must | |
| # appear in the GROUP BY clause"), which is past the planner's corrective-retry | |
| # window. Catching it here turns a failed turn into a self-correcting re-prompt. | |
| # | |
| # Extended 2026-07-23 (F-17): the check used to be gated on `if ir.group_by`, | |
| # so a mixed select with group_by=[] passed. On a DB source Postgres then failed | |
| # loudly, but on a TABULAR source the pandas compiler silently DROPPED the bare | |
| # column while `output_columns` (built from the select list, not the result) | |
| # still advertised it — the user got a real-looking table whose first column was | |
| # entirely `—`. Wrong answers presented as correct is the worst outcome for this | |
| # product, so the check now fires whenever ANY aggregate is present. No false | |
| # positives are possible: a mixed select with no group_by is invalid SQL in | |
| # every dialect. | |
| has_agg = any(item.kind == "agg" for item in ir.select) | |
| if ir.group_by or has_agg: | |
| grouped = set(ir.group_by) | |
| bare = [ | |
| (i, item.column_id) for i, item in enumerate(ir.select) | |
| if item.kind == "column" and item.column_id not in grouped | |
| ] | |
| if bare: | |
| # Every offender, by NAME and by role (2026-11-10). The message used | |
| # to name the first one only and offer "add it to group_by" for any | |
| # column: a plan with a label and three measures bare was told about | |
| # the label, and its retry grouped by the measure itself — one row per | |
| # value, counted as units (PLANNER_SHAPE_FIXES_HANDOFF §1). | |
| context = ( | |
| "while group_by is present" | |
| if ir.group_by | |
| else "alongside an aggregate with no group_by" | |
| ) | |
| measures, labels = [], [] | |
| for _, col_id in bare: | |
| col = columns_by_id[col_id] | |
| (measures if _is_measure(col) else labels).append( | |
| f"{col.name!r} ({col_id})" | |
| ) | |
| first_i, first_id = bare[0] | |
| parts = [ | |
| f"select[{first_i}].column_id {first_id!r} is selected bare {context}." | |
| ] | |
| if measures: | |
| parts.append( | |
| f"Measure(s) {', '.join(measures)}: wrap each in an aggregate, " | |
| 'e.g. {"kind": "agg", "fn": "avg", "column_id": "<id>"} — the ' | |
| "catalog's `aggregate with`, else avg for rates and scores, sum " | |
| "for totals. Never add a measure to group_by: it keeps one row " | |
| "per value." | |
| ) | |
| if labels: | |
| parts.append( | |
| f"Column(s) {', '.join(labels)}: add each to group_by (this " | |
| "splits a group when the column varies inside it) or drop it " | |
| "from select." | |
| ) | |
| raise IRValidationError(" ".join(parts)) | |
| 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_computed( | |
| self, item: Any, columns_by_id: dict[str, Column], where: str | |
| ) -> None: | |
| """A computed ratio rides only on `avg` of a column, over numeric terms in | |
| the tables of this query (2026-11-09, RATIO_MEASURES_REQUIREMENTS M1).""" | |
| if item.fn != "avg" or item.column_id is None: | |
| raise IRValidationError( | |
| f"{where}.computed is only allowed on fn='avg' of a column" | |
| ) | |
| for side in ("numerator", "denominator"): | |
| for k, term in enumerate(getattr(item.computed, side)): | |
| col = self._require_column( | |
| columns_by_id, term.column_id, f"{where}.computed.{side}[{k}]" | |
| ) | |
| if col.data_type not in _NUMERIC_TYPES: | |
| raise IRValidationError( | |
| f"{where}.computed.{side}[{k}]: column {col.name!r} is " | |
| f"{col.data_type!r}, not numeric" | |
| ) | |
| 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 | |