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

2026-06-03 11:01:02 -03:00
parent 61494362a3
commit 29c620b2c2
27 changed files with 1516 additions and 249 deletions
--- a/api/composer/compose.py
+++ b/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)