File size: 10,042 Bytes
6bff5d9 f873f92 6bff5d9 0721bb4 6bff5d9 f873f92 6bff5d9 49b0848 6bff5d9 f873f92 0721bb4 6bff5d9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 | """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}]")
# A grouped query must not select bare columns that aren'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.
if ir.group_by:
grouped = set(ir.group_by)
for i, item in enumerate(ir.select):
if item.kind == "column" and item.column_id not in grouped:
raise IRValidationError(
f"select[{i}].column_id {item.column_id!r} is selected bare "
"while group_by is present — every selected column must "
"either appear in group_by or be wrapped in an aggregate "
f'(e.g. {{"kind": "agg", "fn": "sum", '
f'"column_id": {item.column_id!r}}})'
)
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}"
)
@staticmethod
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)
@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
|