recon-driven sql composer; pick → compose → execute; llm out of structural sql

api/composer/__init__.py

@@ -0,0 +1,26 @@
+"""Recon-driven SQL composer.
+
+The LLM picks a *shape* (which metric, which dimensions, which filters);
+the composer walks recon and emits the SQL deterministically. Column–table
+binding, joins, and quoting are graph traversals, not LLM guesses.
+
+See `def/schema-as-constraint.md` for the architectural argument.
+"""
+
+from api.composer.types import (
+    ComposeResult,
+    Filter,
+    OrderBy,
+    Pick,
+    PickValidationError,
+)
+from api.composer.compose import compose
+
+__all__ = [
+    "compose",
+    "ComposeResult",
+    "Filter",
+    "OrderBy",
+    "Pick",
+    "PickValidationError",
+]

api/composer/compose.py

@@ -0,0 +1,288 @@
+"""compose(pick, recon) → ComposeResult.
+
+Deterministic SQL emission from a typed Pick. The composer:
+  1. Resolves the metric → from_table + sql expression + default filter.
+  2. Resolves each group_by column → owning table (via recon.resolve_column).
+  3. Computes the join graph: join_path(metric.from_table, owner) per dim,
+     deduped, walked in declared order to pick the right relationship edge.
+  4. Renders SELECT / FROM + JOINs / WHERE / GROUP BY / ORDER BY / LIMIT
+     with every identifier quoted via sqlglot's `identify=True` round-trip
+     as a paranoid post-check.
+
+No LLM call anywhere in this module. If the Pick can't be expressed against
+recon, the composer raises `PickValidationError` — the caller surfaces it.
+"""
+
+from __future__ import annotations
+
+import sqlglot
+
+from api.composer.dates import resolve_date_range
+from api.composer.types import (
+    ComposeResult,
+    Filter,
+    OrderBy,
+    Pick,
+    PickValidationError,
+)
+from api.recon.types import Recon, Relationship
+
+
+# ── Alias allocation ────────────────────────────────────────
+
+def _alias_for(table: str, used: dict[str, str]) -> str:
+    """Pick a short alias for `table` that's unique within the query.
+
+    First letter, then first-two, then numbered. `used` maps alias→table so
+    we can grow on collision.
+    """
+    for base in (table[:1], table[:2], table[:3], table):
+        if base and base not in used:
+            return base
+    i = 1
+    while True:
+        cand = f"{table[:1]}{i}"
+        if cand not in used:
+            return cand
+        i += 1
+
+
+# ── Join graph ──────────────────────────────────────────────
+
+def _find_edge(recon: Recon, a: str, b: str) -> Relationship:
+    """Return the relationship that joins tables a and b (either direction).
+    Raises if no declared FK connects them."""
+    for r in recon.relationships:
+        if (r.from_table == a and r.to_table == b) or (r.from_table == b and r.to_table == a):
+            return r
+    raise PickValidationError(
+        f"no declared relationship between {a!r} and {b!r}"
+    )
+
+
+def _build_joins(recon: Recon, base: str, targets: list[str],
+                 aliases: dict[str, str]) -> list[str]:
+    """Compute the union of join paths from `base` to each target table,
+    emit JOIN clauses in walk order. Allocates aliases for any intermediate
+    table that wasn't already in `aliases`."""
+    visited: set[str] = {base}
+    clauses: list[str] = []
+    for target in targets:
+        path = recon.join_path(base, target)
+        if path is None:
+            raise PickValidationError(
+                f"no join path from {base!r} to {target!r} in recon"
+            )
+        for i in range(1, len(path)):
+            prev, cur = path[i - 1], path[i]
+            if cur in visited:
+                continue
+            edge = _find_edge(recon, prev, cur)
+            # Edge direction tells us which side has which column.
+            if edge.from_table == prev:
+                left_t, left_c, right_t, right_c = edge.from_table, edge.from_column, edge.to_table, edge.to_column
+            else:
+                left_t, left_c, right_t, right_c = edge.to_table, edge.to_column, edge.from_table, edge.from_column
+            # Ensure both sides have aliases.
+            for t in (left_t, right_t):
+                if t not in aliases:
+                    aliases[t] = _alias_for(t, {a: t for t, a in aliases.items()})
+            la, ra = aliases[left_t], aliases[right_t]
+            clauses.append(
+                f'JOIN "{cur}" AS "{aliases[cur]}" '
+                f'ON "{la}"."{left_c}" = "{ra}"."{right_c}"'
+            )
+            visited.add(cur)
+    return clauses
+
+
+# ── Filter rendering ────────────────────────────────────────
+
+def _render_literal(v) -> str:
+    """Render a Python value as a SQL literal. Strings are single-quoted with
+    embedded quotes escaped. Numbers pass through."""
+    if isinstance(v, str):
+        return "'" + v.replace("'", "''") + "'"
+    if isinstance(v, bool):
+        return "TRUE" if v else "FALSE"
+    if v is None:
+        return "NULL"
+    return str(v)
+
+
+def _render_filter(f: Filter, recon: Recon, aliases: dict[str, str]) -> str:
+    """Render one Filter to a SQL predicate. Resolves the column's owning
+    table and aliases it (allocating an alias if needed — and a join later
+    if that introduces a new table)."""
+    table, col = recon.resolve_column(f.column)
+    if table not in aliases:
+        aliases[table] = _alias_for(table, {a: t for t, a in aliases.items()})
+    qualified = f'"{aliases[table]}"."{col.name}"'
+
+    kind = f.kind()
+    if kind == "equals":
+        return f"{qualified} = {_render_literal(f.equals)}"
+    if kind == "in_values":
+        if not f.in_values:
+            raise PickValidationError(
+                f"filter on {f.column!r} has empty in_values"
+            )
+        rendered = ", ".join(_render_literal(v) for v in f.in_values)
+        return f"{qualified} IN ({rendered})"
+    if kind == "between":
+        lo, hi = f.between
+        return f"{qualified} BETWEEN {_render_literal(lo)} AND {_render_literal(hi)}"
+    if kind == "date_range":
+        return resolve_date_range(f.date_range, col, qualified)
+    raise PickValidationError(f"unknown filter kind {kind!r}")
+
+
+# ── ORDER BY rendering ──────────────────────────────────────
+
+def _render_order_by(ob: OrderBy, metric_name: str, dim_select: dict[str, str]) -> str:
+    """`metric_name` is the alias of the metric SELECT expression;
+    `dim_select` maps dimension ref (column ref) → output alias."""
+    if ob.by == "metric":
+        return f'"{metric_name}" {ob.direction.upper()}'
+    # by == "dimension"
+    if ob.dimension not in dim_select:
+        raise PickValidationError(
+            f"order_by.dimension={ob.dimension!r} is not in group_by ({list(dim_select)})"
+        )
+    return f'"{dim_select[ob.dimension]}" {ob.direction.upper()}'
+
+
+# ── Metric expression rewriting ─────────────────────────────
+
+def _qualify_metric_sql(sql_fragment: str, table_alias: str) -> str:
+    """Rewrite bare column refs inside the metric's SQL expression so they
+    point at the metric table's alias.
+
+    Metric.sql in metrics.yaml is written as an unaliased expression (e.g.
+    `AVG(CASE WHEN status='B' ... END)`). We need it qualified to the
+    metric table's alias so the composer can introduce other tables via
+    joins without ambiguity. Uses sqlglot to parse and rewrite identifier
+    refs that don't already have a table prefix.
+    """
+    tree = sqlglot.parse_one(sql_fragment, dialect="postgres")
+    for col in tree.find_all(sqlglot.exp.Column):
+        if col.table:
+            continue
+        # Inject the alias as the table qualifier.
+        col.set("table", sqlglot.exp.Identifier(this=table_alias, quoted=True))
+    return tree.sql(dialect="postgres", identify=True)
+
+
+def _qualify_metric_filter(filter_fragment: str, table_alias: str) -> str:
+    """Same as _qualify_metric_sql but for the metric's optional WHERE filter.
+    metric.filter is written as a bare boolean expression."""
+    return _qualify_metric_sql(filter_fragment, table_alias)
+
+
+# ── Compose ─────────────────────────────────────────────────
+
+def compose(pick: Pick, recon: Recon) -> ComposeResult:
+    """Render `pick` to SQL against `recon`. Raises PickValidationError on
+    any reference the recon can't resolve."""
+
+    # 1. Metric.
+    if pick.metric not in recon.metrics:
+        raise PickValidationError(
+            f"metric {pick.metric!r} not in recon (known: {sorted(recon.metrics)})"
+        )
+    metric = recon.metrics[pick.metric]
+    if metric.from_table not in recon.tables:
+        raise PickValidationError(
+            f"metric {pick.metric!r} declares from_table={metric.from_table!r} "
+            f"but no such table in recon"
+        )
+
+    aliases: dict[str, str] = {}
+    aliases[metric.from_table] = _alias_for(metric.from_table, {})
+    base_alias = aliases[metric.from_table]
+
+    # 2. Resolve every group_by column to (table, Column) and collect target
+    #    tables that aren't the metric's table — those need joins.
+    group_bindings: list[tuple[str, str, str, str]] = []  # (ref, table, col_name, output_alias)
+    extra_tables: list[str] = []
+    dim_select_alias: dict[str, str] = {}
+    for ref in pick.group_by:
+        table, col = recon.resolve_column(ref)
+        if table not in aliases:
+            if table != metric.from_table:
+                extra_tables.append(table)
+            aliases[table] = _alias_for(table, {a: t for t, a in aliases.items()})
+        out_alias = col.name if ref == col.name else ref.replace(".", "_")
+        group_bindings.append((ref, table, col.name, out_alias))
+        dim_select_alias[ref] = out_alias
+
+    # 3. Resolve filter columns first (they may also introduce new tables we
+    #    need to join). Rendered separately so we can keep their predicates
+    #    in WHERE.
+    filter_clauses: list[str] = []
+    for f in pick.where:
+        table, _ = recon.resolve_column(f.column)
+        if table not in aliases:
+            if table != metric.from_table:
+                extra_tables.append(table)
+            aliases[table] = _alias_for(table, {a: t for t, a in aliases.items()})
+        filter_clauses.append(_render_filter(f, recon, aliases))
+
+    # 4. Build JOINs for the union of extra tables.
+    join_clauses = _build_joins(recon, metric.from_table, extra_tables, aliases)
+
+    # 5. SELECT list: dimension columns first (in group_by order), then the
+    #    metric expression.
+    select_parts: list[str] = []
+    for ref, table, col_name, out_alias in group_bindings:
+        select_parts.append(f'"{aliases[table]}"."{col_name}" AS "{out_alias}"')
+    metric_expr = _qualify_metric_sql(metric.sql, base_alias)
+    select_parts.append(f'{metric_expr} AS "{metric.name}"')
+
+    # 6. WHERE: metric's default filter (if any) ANDed with the Pick's filters.
+    where_parts: list[str] = []
+    if metric.filter:
+        where_parts.append(_qualify_metric_filter(metric.filter, base_alias))
+    where_parts.extend(filter_clauses)
+
+    # 7. GROUP BY (just the dim output aliases) + ORDER BY + LIMIT.
+    group_by_clause = ""
+    if group_bindings:
+        group_by_clause = "GROUP BY " + ", ".join(
+            f'"{out_alias}"' for _, _, _, out_alias in group_bindings
+        )
+
+    order_by_clause = ""
+    if pick.order_by is not None:
+        order_by_clause = "ORDER BY " + _render_order_by(
+            pick.order_by, metric.name, dim_select_alias
+        )
+    elif group_bindings:
+        # Sensible default: largest metric first.
+        order_by_clause = f'ORDER BY "{metric.name}" DESC'
+
+    limit_clause = f"LIMIT {pick.limit}" if pick.limit is not None else ""
+
+    # 8. Assemble.
+    sql_parts = [
+        "SELECT " + ", ".join(select_parts),
+        f'FROM "{metric.from_table}" AS "{base_alias}"',
+    ]
+    sql_parts.extend(join_clauses)
+    if where_parts:
+        sql_parts.append("WHERE " + " AND ".join(where_parts))
+    if group_by_clause:
+        sql_parts.append(group_by_clause)
+    if order_by_clause:
+        sql_parts.append(order_by_clause)
+    if limit_clause:
+        sql_parts.append(limit_clause)
+    raw = "\n".join(sql_parts)
+
+    # 9. Paranoid re-parse: catches composer bugs by round-tripping through
+    #    sqlglot. identify=True keeps every identifier quoted.
+    parsed = sqlglot.parse_one(raw, dialect="postgres")
+    sql = parsed.sql(dialect="postgres", identify=True)
+
+    used = sorted(aliases)
+    return ComposeResult(sql=sql, used_tables=used)

api/composer/dates.py

@@ -0,0 +1,73 @@
+"""Date-range resolver for typed Filters.
+
+Given a `Filter` with `date_range="YYYY"` or `date_range="YYYY-Q[1-4]"` and
+the target column, render a SQL fragment that evaluates correctly under the
+column's actual storage type. Dispatched in two layers:
+
+1. `Column.semantic_type` if set (extension point — e.g. `"date_yymmdd"` for
+   datasets that ship YYMMDD-encoded integers).
+2. Otherwise `Column.sql_type` (`DATE`, `TIMESTAMP`, …) — the default path
+   for warehouses that store dates as real date columns.
+
+Inclusive on both ends. Bounds quoted with single quotes (Postgres parses
+date literals from strings).
+"""
+
+from __future__ import annotations
+
+import re
+
+from api.composer.types import PickValidationError
+from api.recon.types import Column
+
+_YEAR_ONLY = re.compile(r"^\d{4}$")
+_YEAR_QUARTER = re.compile(r"^(\d{4})-Q([1-4])$")
+
+
+_QUARTER_BOUNDS = {
+    1: ("01-01", "03-31"),
+    2: ("04-01", "06-30"),
+    3: ("07-01", "09-30"),
+    4: ("10-01", "12-31"),
+}
+
+
+def _bounds(spec: str) -> tuple[str, str]:
+    """Parse a date_range spec to (lo, hi) calendar dates as ISO strings."""
+    if _YEAR_ONLY.match(spec):
+        year = spec
+        return f"{year}-01-01", f"{year}-12-31"
+    m = _YEAR_QUARTER.match(spec)
+    if m:
+        year, q = m.group(1), int(m.group(2))
+        lo_mmdd, hi_mmdd = _QUARTER_BOUNDS[q]
+        return f"{year}-{lo_mmdd}", f"{year}-{hi_mmdd}"
+    raise PickValidationError(
+        f"unsupported date_range {spec!r}; expected 'YYYY' or 'YYYY-Q1..Q4'"
+    )
+
+
+def resolve_date_range(spec: str, column: Column, qualified_col: str) -> str:
+    """Render the WHERE fragment for `<qualified_col> BETWEEN ...`.
+
+    `qualified_col` is the already-quoted reference the composer wants in
+    the SQL (e.g. `"l"."date"`). Returns just the predicate, no `WHERE`.
+    """
+    lo, hi = _bounds(spec)
+    semantic = (column.semantic_type or "").lower()
+
+    if semantic == "date_yymmdd":
+        # Column stores YYMMDD as integer (some BIRD datasets ship this way).
+        return (
+            f"TO_DATE(LPAD({qualified_col}::text, 6, '0'), 'YYMMDD') "
+            f"BETWEEN '{lo}' AND '{hi}'"
+        )
+
+    sql_type = column.sql_type.upper()
+    if sql_type.startswith("DATE") or sql_type.startswith("TIMESTAMP"):
+        return f"{qualified_col} BETWEEN '{lo}' AND '{hi}'"
+
+    raise PickValidationError(
+        f"date_range on column with sql_type={column.sql_type!r} and "
+        f"semantic_type={column.semantic_type!r} is not supported"
+    )

api/composer/types.py

@@ -0,0 +1,189 @@
+"""Pick / Filter / OrderBy — the typed shape the LLM emits and the composer
+consumes.
+
+Validation is strict: a Pick that references a metric or column not in recon
+is rejected before composition. That's the whole point — the composer's
+input must be expressible in terms of recon entities. There is no raw-SQL
+escape hatch on a Filter; questions outside the supported shapes get a clean
+"beyond what I'm built to answer" rather than a fabricated query.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Literal
+
+
+class PickValidationError(ValueError):
+    """The Pick references a metric/column/filter shape that doesn't fit
+    recon. Surfaced cleanly to the caller (no retry — see feedback memory)."""
+
+
+# ── Filter ──────────────────────────────────────────────────
+
+@dataclass
+class Filter:
+    """One WHERE-clause fragment. Exactly one of the value fields must be set.
+
+    - column + equals: `<col> = <value>`
+    - column + in_values: `<col> IN (<values>)`
+    - column + between: `<col> BETWEEN <a> AND <b>` (inclusive, both ends).
+    - column + date_range: `<col> BETWEEN '<yyyy>-01-01' AND '<yyyy>-12-31'`
+      for `"YYYY"`, or quarter bounds for `"YYYY-Q[1-4]"`.
+    """
+    column: str
+    equals: str | int | float | None = None
+    in_values: list[Any] | None = None
+    between: tuple[Any, Any] | None = None
+    date_range: str | None = None
+
+    def kind(self) -> str:
+        """Which filter shape is set; raises if none or more than one."""
+        set_fields = [
+            name for name in ("equals", "in_values", "between", "date_range")
+            if getattr(self, name) is not None
+        ]
+        if len(set_fields) == 0:
+            raise PickValidationError(
+                f"filter on column {self.column!r} has no value field set"
+            )
+        if len(set_fields) > 1:
+            raise PickValidationError(
+                f"filter on column {self.column!r} sets multiple value fields: {set_fields}"
+            )
+        return set_fields[0]
+
+    def to_dict(self) -> dict[str, Any]:
+        out: dict[str, Any] = {"column": self.column}
+        for k in ("equals", "in_values", "between", "date_range"):
+            v = getattr(self, k)
+            if v is not None:
+                out[k] = list(v) if isinstance(v, tuple) else v
+        return out
+
+    @classmethod
+    def from_dict(cls, d: dict[str, Any]) -> "Filter":
+        if "column" not in d:
+            raise PickValidationError(f"filter missing 'column': {d!r}")
+        between = d.get("between")
+        if between is not None:
+            if not isinstance(between, (list, tuple)) or len(between) != 2:
+                raise PickValidationError(
+                    f"filter 'between' must be a 2-element list, got {between!r}"
+                )
+            between = (between[0], between[1])
+        return cls(
+            column=d["column"],
+            equals=d.get("equals"),
+            in_values=d.get("in_values"),
+            between=between,
+            date_range=d.get("date_range"),
+        )
+
+
+# ── OrderBy ─────────────────────────────────────────────────
+
+@dataclass
+class OrderBy:
+    """ORDER BY clause. Either by metric or by a named dimension."""
+    by: Literal["metric", "dimension"]
+    direction: Literal["asc", "desc"] = "desc"
+    dimension: str | None = None
+
+    def to_dict(self) -> dict[str, Any]:
+        out: dict[str, Any] = {"by": self.by, "direction": self.direction}
+        if self.dimension is not None:
+            out["dimension"] = self.dimension
+        return out
+
+    @classmethod
+    def from_dict(cls, d: dict[str, Any]) -> "OrderBy":
+        by = d.get("by")
+        if by not in ("metric", "dimension"):
+            raise PickValidationError(f"order_by.by must be 'metric'|'dimension', got {by!r}")
+        direction = d.get("direction", "desc")
+        if direction not in ("asc", "desc"):
+            raise PickValidationError(f"order_by.direction must be 'asc'|'desc', got {direction!r}")
+        dim = d.get("dimension")
+        if by == "dimension" and not dim:
+            raise PickValidationError("order_by.by='dimension' requires order_by.dimension")
+        return cls(by=by, direction=direction, dimension=dim)
+
+
+# ── Pick ────────────────────────────────────────────────────
+
+@dataclass
+class Pick:
+    """What the LLM emits — fully resolvable by the composer against recon.
+
+    Fields:
+      kind: literal "aggregate" (only supported shape in v1).
+      metric: must be a name in recon.metrics.
+      group_by: column refs (`col` or `table.col`); each resolved by
+        recon.resolve_column.
+      where: typed filters; the composer renders each against the column's
+        sql_type / semantic_type.
+      order_by: optional sort.
+      limit: optional row limit.
+    """
+    kind: Literal["aggregate"]
+    metric: str
+    group_by: list[str] = field(default_factory=list)
+    where: list[Filter] = field(default_factory=list)
+    order_by: OrderBy | None = None
+    limit: int | None = None
+
+    def to_dict(self) -> dict[str, Any]:
+        out: dict[str, Any] = {
+            "kind": self.kind,
+            "metric": self.metric,
+            "group_by": list(self.group_by),
+            "where": [f.to_dict() for f in self.where],
+        }
+        if self.order_by is not None:
+            out["order_by"] = self.order_by.to_dict()
+        if self.limit is not None:
+            out["limit"] = self.limit
+        return out
+
+    @classmethod
+    def from_dict(cls, d: dict[str, Any]) -> "Pick":
+        kind = d.get("kind", "aggregate")
+        if kind != "aggregate":
+            raise PickValidationError(
+                f"only kind='aggregate' is supported; got {kind!r}"
+            )
+        metric = d.get("metric")
+        if not isinstance(metric, str) or not metric:
+            raise PickValidationError(f"pick missing 'metric': {d!r}")
+        group_by = d.get("group_by", []) or []
+        if not isinstance(group_by, list) or not all(isinstance(c, str) for c in group_by):
+            raise PickValidationError(f"pick.group_by must be list[str], got {group_by!r}")
+        where_raw = d.get("where", []) or []
+        if not isinstance(where_raw, list):
+            raise PickValidationError(f"pick.where must be a list, got {where_raw!r}")
+        where = [Filter.from_dict(f) for f in where_raw]
+        for f in where:
+            f.kind()  # raises if shape is missing/ambiguous
+        order_by = OrderBy.from_dict(d["order_by"]) if d.get("order_by") else None
+        limit = d.get("limit")
+        if limit is not None and not isinstance(limit, int):
+            raise PickValidationError(f"pick.limit must be int or None, got {limit!r}")
+        return cls(
+            kind=kind,
+            metric=metric,
+            group_by=group_by,
+            where=where,
+            order_by=order_by,
+            limit=limit,
+        )
+
+
+# ── ComposeResult ───────────────────────────────────────────
+
+@dataclass
+class ComposeResult:
+    """What `compose()` returns. Same shape as the legacy T2SResult so call
+    sites can swap with minimal churn."""
+    sql: str
+    used_tables: list[str]