src/traceability/resolve.py

"""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,
    )