Rifqi Hafizuddin
Claude Opus 4.8
[KM-652] feat(query): single-level FK joins in QueryIR (T4)
7d36567 | """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) | |
| 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 | |