File size: 6,422 Bytes
32abc41 | 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 | """Resolve a query IR into the user-facing `DataUsed` record (KM-691).
Pure, deterministic id->name resolution against the catalog — NO LLM. Turns the
opaque IR the tools ran (column_id / table_id / source_id) into real names the
user can check against their own database, and splits the result set into columns
read straight from the data vs values the analysis computed.
Never raises on the caller's path: an unresolved id degrades to showing the raw id
(marked table='?'), never a crash — a trace slip must not break the user's answer.
"""
from __future__ import annotations
from typing import Any
from ..catalog.models import Catalog
from ..query.ir.models import AggSelect, ColumnSelect, QueryIR
from .schemas import (
ColumnRef,
DataUsed,
FilterRef,
JoinRef,
OrderByRef,
OutputColumn,
SourceRef,
TableRef,
)
class _CatalogIndex:
"""Flat id->name lookups built once from a Catalog (ids are globally unique)."""
def __init__(self, catalog: Catalog) -> None:
self.source: dict[str, tuple[str, str | None]] = {}
self.table: dict[str, str] = {}
# column_id -> (table_name, column_name, data_type, pii_flag)
self.col: dict[str, tuple[str, str, str | None, bool]] = {}
for s in catalog.sources:
self.source[s.source_id] = (s.name, s.source_type)
for t in s.tables:
self.table[t.table_id] = t.name
for c in t.columns:
self.col[c.column_id] = (t.name, c.name, c.data_type, c.pii_flag)
def col_name(self, column_id: str) -> str:
hit = self.col.get(column_id)
return hit[1] if hit else column_id
def col_qual(self, column_id: str) -> str:
"""`table.column` when resolvable, else the raw id (honest fallback)."""
hit = self.col.get(column_id)
return f"{hit[0]}.{hit[1]}" if hit else column_id
def _describe_filter(col: str, op: str, value: Any) -> str:
"""Plain-language rendering of one filter (fixed templates — no LLM)."""
if op == "between" and isinstance(value, list | tuple) and len(value) == 2:
return f"{col} is between {value[0]} and {value[1]}"
if op == "in":
return f"{col} is one of {value}"
if op == "not_in":
return f"{col} is not one of {value}"
if op == "is_null":
return f"{col} is empty"
if op == "is_not_null":
return f"{col} is not empty"
if op == "like":
return f"{col} matches {value}"
return f"{col} {op} {value}"
def resolve_data_used(
ir: QueryIR,
catalog: Catalog,
query: str | None = None,
rows_returned: int | None = None,
) -> DataUsed:
"""Resolve one `retrieve_data` IR into a `DataUsed`. Deterministic; no LLM."""
idx = _CatalogIndex(catalog)
src_name, src_type = idx.source.get(ir.source_id, (ir.source_id, None))
tables = [TableRef(id=ir.table_id, name=idx.table.get(ir.table_id, ir.table_id), role="base")]
joins: list[JoinRef] = []
for j in ir.joins:
ttid = j.target_table_id
tables.append(TableRef(id=ttid, name=idx.table.get(ttid, ttid), role="joined"))
left, right = idx.col_qual(j.left_column_id), idx.col_qual(j.right_column_id)
joins.append(JoinRef(type=j.type, condition=f"{left} = {right}"))
# roles[column_id] -> set of why-used tags, accumulated across every clause.
roles: dict[str, set[str]] = {}
def _tag(cid: str, role: str) -> None:
roles.setdefault(cid, set()).add(role)
output_columns: list[OutputColumn] = []
for s in ir.select:
if isinstance(s, ColumnSelect):
_tag(s.column_id, "selected")
output_columns.append(
OutputColumn(
name=s.alias or idx.col_name(s.column_id),
kind="column",
from_=idx.col_qual(s.column_id),
)
)
elif isinstance(s, AggSelect):
if s.column_id:
_tag(s.column_id, "aggregated")
formula = f"{s.fn.upper()}({idx.col_qual(s.column_id)})"
frm = idx.col_qual(s.column_id)
else:
formula = f"{s.fn.upper()}(*)" # count(*) carries no column
frm = None
output_columns.append(
OutputColumn(name=s.alias or formula, kind="computed", formula=formula, from_=frm)
)
filters: list[FilterRef] = []
for f in ir.filters:
_tag(f.column_id, "filtered")
col = idx.col_qual(f.column_id)
desc = _describe_filter(col, f.op, f.value)
filters.append(FilterRef(column=col, op=f.op, value=f.value, description=desc))
group_by: list[str] = []
for gid in ir.group_by:
_tag(gid, "grouped")
group_by.append(idx.col_qual(gid))
for j in ir.joins:
_tag(j.left_column_id, "joined")
_tag(j.right_column_id, "joined")
order_by: list[OrderByRef] = []
for o in ir.order_by:
# IR wart: order_by.column_id may hold a SELECT alias (a computed output),
# not a catalog column_id. Resolve as a real column when known, else treat
# it as a computed-output reference rather than inventing a name.
if o.column_id in idx.col:
_tag(o.column_id, "ordered")
order_by.append(OrderByRef(target=idx.col_qual(o.column_id), kind="column", dir=o.dir))
else:
order_by.append(OrderByRef(target=o.column_id, kind="computed", dir=o.dir))
columns_read: list[ColumnRef] = []
for cid, rs in roles.items():
hit = idx.col.get(cid)
if hit:
tname, cname, dtype, pii = hit
columns_read.append(
ColumnRef(
id=cid, name=cname, table=tname, data_type=dtype, pii=pii, roles=sorted(rs)
)
)
else:
# Unresolved id — keep it, mark it, never guess a name.
columns_read.append(ColumnRef(id=cid, name=cid, table="?", roles=sorted(rs)))
return DataUsed(
source=SourceRef(id=ir.source_id, name=src_name, type=src_type),
tables=tables,
joins=joins,
columns_read=columns_read,
output_columns=output_columns,
filters=filters,
group_by=group_by,
order_by=order_by,
limit=ir.limit,
rows_returned=rows_returned,
query=query,
)
|