refactor composer to use sqlalchemy, avoid string concatenations and follow conventions

This commit is contained in:
2026-06-03 12:10:30 -03:00
parent cbc7df8c60
commit be09fcde2c
16 changed files with 546 additions and 273 deletions

View File

@@ -1,4 +1,5 @@
.PHONY: kind tilt-up tilt-down seed seed-fetch seed-push recon evals test \ .PHONY: kind tilt-up tilt-down seed seed-fetch seed-push recon evals test \
run-log run-logs \
compose-up compose-down compose-clean compose-seed compose-recon compose-evals \ compose-up compose-down compose-clean compose-seed compose-recon compose-evals \
docs-graphs docs-graphs
@@ -51,6 +52,23 @@ recon:
evals: evals:
kubectl $(KCTX) $(KNS) exec deploy/api -- uv run python -m api.evals.run_evals kubectl $(KCTX) $(KNS) exec deploy/api -- uv run python -m api.evals.run_evals
# Pull a run's JSONL transcript from the api pod to the host. The api writes
# every published event into /app/.data/logs/<RUN>.jsonl inside the pod;
# this target copies it next to .data/logs/ on the host so it's easy to
# share / grep / paste a path to.
# Usage: make run-log RUN=<run_id>
run-log:
@[ -n "$(RUN)" ] || { echo "usage: make run-log RUN=<run_id>" >&2; exit 1; }
@mkdir -p .data/logs
@POD=$$(kubectl $(KCTX) $(KNS) get pod -l app=api -o jsonpath='{.items[0].metadata.name}'); \
kubectl $(KCTX) $(KNS) cp $$POD:/app/.data/logs/$(RUN).jsonl .data/logs/$(RUN).jsonl \
&& echo "→ .data/logs/$(RUN).jsonl"
# List the runs the api currently has log files for (most recent first).
run-logs:
@POD=$$(kubectl $(KCTX) $(KNS) get pod -l app=api -o jsonpath='{.items[0].metadata.name}'); \
kubectl $(KCTX) $(KNS) exec $$POD -- sh -c 'ls -1t /app/.data/logs 2>/dev/null || echo "(no logs yet)"'
# Unit tests run on the host venv (no postgres/llm calls needed). # Unit tests run on the host venv (no postgres/llm calls needed).
# Ensures dev extras are installed first; uv is idempotent on no-ops. # Ensures dev extras are installed first; uv is idempotent on no-ops.
test: test:

56
api/composer/aliases.py Normal file
View File

@@ -0,0 +1,56 @@
"""Alias allocation for the composer.
Every table referenced in a composed query gets a short, query-unique alias.
`AliasMap` owns that allocation so the rest of the composer never threads a
bare dict around: ask it for a table's alias (allocating on first use) or for a
qualified `exp.Column`, and it remembers the assignment.
Allocation strategy (stable — golden SQL depends on it): first letter, then the
first two/three letters, then the full name, then a numbered `<first><n>` on
collision. So `loan→l`, `account→a`, `district→d`.
"""
from __future__ import annotations
from sqlglot import exp
def _alloc(table: str, used: set[str]) -> str:
"""Pick an alias for `table` not already in `used`."""
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
class AliasMap:
"""Maps table name → alias, allocating on first request."""
def __init__(self) -> None:
self._by_table: dict[str, str] = {}
def has(self, table: str) -> bool:
return table in self._by_table
def alias(self, table: str) -> str:
"""Return `table`'s alias, allocating a fresh unique one if needed."""
if table not in self._by_table:
self._by_table[table] = _alloc(table, set(self._by_table.values()))
return self._by_table[table]
def col(self, table: str, column: str) -> exp.Column:
"""A quoted, alias-qualified column reference for the final tree."""
return exp.column(column, table=self.alias(table), quoted=True)
def aliased_table(self, table: str) -> exp.Expression:
"""`"<table>" AS "<alias>"` for FROM / JOIN."""
return exp.alias_(exp.to_table(table, quoted=True), self.alias(table), quoted=True)
def tables(self) -> list[str]:
"""Sorted names of every table that has been referenced."""
return sorted(self._by_table)

View File

@@ -1,191 +1,48 @@
"""compose(pick, recon) → ComposeResult. """compose(pick, recon) → ComposeResult.
Deterministic SQL emission from a typed Pick. The composer: Deterministic SQL emission from a typed Pick. The composer assembles a single
1. Resolves the metric → from_table + sql expression + default filter. sqlglot expression tree and renders it once — there is no hand-built SQL string
2. Resolves each group_by column → owning table (via recon.resolve_column). and no after-the-fact re-parse; the tree *is* the structure, and
3. Computes the join graph: join_path(metric.from_table, owner) per dim, `.sql(dialect="postgres", identify=True)` (every identifier quoted) is the one
deduped, walked in declared order to pick the right relationship edge. place a string is produced.
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 The work is split across focused modules, each returning `exp` nodes:
recon, the composer raises `PickValidationError` — the caller surfaces it.
aliases.py table → alias allocation (AliasMap)
joins.py join-graph walk → JOIN specs
filters.py typed Filter → predicate node
metrics.py metric YAML fragment → alias-qualified expression
dates.py date_range → BETWEEN node (encoding-aware)
order.py OrderBy → Ordered node
encodings.py default + dataset storage-encoding library
No LLM call anywhere. The Pick is the constraint boundary: if it can't be
expressed against recon the composer raises `PickValidationError`, which the
caller surfaces — it never fabricates SQL. New capability is additive: a new
filter shape is a branch in filters.py; a new Pick `kind` is a builder here.
Widening what a Pick can express is a deliberate decision — see
`def/schema-as-constraint.md` before doing so.
""" """
from __future__ import annotations from __future__ import annotations
import sqlglot from sqlglot import exp
from api.composer.dates import resolve_date_range from api.composer.aliases import AliasMap
from api.composer.types import ( from api.composer.encodings import effective_encodings
ComposeResult, from api.composer.filters import render_filter
Filter, from api.composer.joins import build_joins
OrderBy, from api.composer.metrics import qualify_metric_expr
Pick, from api.composer.order import default_order_by, render_order_by
PickValidationError, from api.composer.types import ComposeResult, Pick, PickValidationError
) from api.recon.types import Recon
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: def compose(pick: Pick, recon: Recon) -> ComposeResult:
"""Render `pick` to SQL against `recon`. Raises PickValidationError on """Render `pick` to SQL against `recon`. Raises PickValidationError on any
any reference the recon can't resolve.""" reference the recon can't resolve."""
# 1. Metric. # 1. Metric → base table.
if pick.metric not in recon.metrics: if pick.metric not in recon.metrics:
raise PickValidationError( raise PickValidationError(
f"metric {pick.metric!r} not in recon (known: {sorted(recon.metrics)})" f"metric {pick.metric!r} not in recon (known: {sorted(recon.metrics)})"
@@ -197,92 +54,73 @@ def compose(pick: Pick, recon: Recon) -> ComposeResult:
f"but no such table in recon" f"but no such table in recon"
) )
aliases: dict[str, str] = {} aliases = AliasMap()
aliases[metric.from_table] = _alias_for(metric.from_table, {}) base_alias = aliases.alias(metric.from_table)
base_alias = aliases[metric.from_table] encodings = effective_encodings(recon)
# 2. Resolve every group_by column to (table, Column) and collect target # 2. Resolve group_by columns → owning table + output alias. Tables other
# tables that aren't the metric's table — those need joins. # than the metric's need joins.
group_bindings: list[tuple[str, str, str, str]] = [] # (ref, table, col_name, output_alias) group_bindings: list[tuple[str, str, str]] = [] # (table, col_name, out_alias)
extra_tables: list[str] = [] extra_tables: list[str] = []
dim_select_alias: dict[str, str] = {} dim_select_alias: dict[str, str] = {}
for ref in pick.group_by: for ref in pick.group_by:
table, col = recon.resolve_column(ref) table, col = recon.resolve_column(ref)
if table not in aliases: if not aliases.has(table):
if table != metric.from_table: if table != metric.from_table:
extra_tables.append(table) extra_tables.append(table)
aliases[table] = _alias_for(table, {a: t for t, a in aliases.items()}) aliases.alias(table)
out_alias = col.name if ref == col.name else ref.replace(".", "_") out_alias = col.name if ref == col.name else ref.replace(".", "_")
group_bindings.append((ref, table, col.name, out_alias)) group_bindings.append((table, col.name, out_alias))
dim_select_alias[ref] = out_alias dim_select_alias[ref] = out_alias
# 3. Resolve filter columns first (they may also introduce new tables we # 3. Resolve filters → predicate nodes (may also introduce new tables).
# need to join). Rendered separately so we can keep their predicates filter_nodes: list[exp.Expression] = []
# in WHERE.
filter_clauses: list[str] = []
for f in pick.where: for f in pick.where:
table, _ = recon.resolve_column(f.column) table, _ = recon.resolve_column(f.column)
if table not in aliases: if not aliases.has(table):
if table != metric.from_table: if table != metric.from_table:
extra_tables.append(table) extra_tables.append(table)
aliases[table] = _alias_for(table, {a: t for t, a in aliases.items()}) aliases.alias(table)
filter_clauses.append(_render_filter(f, recon, aliases)) filter_nodes.append(render_filter(f, recon, aliases, encodings))
# 4. Build JOINs for the union of extra tables. # 4. JOINs for the union of extra tables.
join_clauses = _build_joins(recon, metric.from_table, extra_tables, aliases) join_specs = build_joins(recon, metric.from_table, extra_tables, aliases)
# 5. SELECT list: dimension columns first (in group_by order), then the # 5. SELECT list: dimensions (in group_by order), then the metric.
# metric expression. select_nodes: list[exp.Expression] = [
select_parts: list[str] = [] exp.alias_(aliases.col(table, col_name), out_alias, quoted=True)
for ref, table, col_name, out_alias in group_bindings: for 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_nodes.append(
select_parts.append(f'{metric_expr} AS "{metric.name}"') exp.alias_(qualify_metric_expr(metric.sql, base_alias), metric.name, quoted=True)
)
# 6. WHERE: metric's default filter (if any) ANDed with the Pick's filters. # 6. WHERE: metric's default filter (if any) ANDed with the Pick's filters.
where_parts: list[str] = [] where_nodes: list[exp.Expression] = []
if metric.filter: if metric.filter:
where_parts.append(_qualify_metric_filter(metric.filter, base_alias)) where_nodes.append(qualify_metric_expr(metric.filter, base_alias))
where_parts.extend(filter_clauses) where_nodes.extend(filter_nodes)
# 7. GROUP BY (just the dim output aliases) + ORDER BY + LIMIT. # 7. ORDER BY: explicit, else largest-metric-first for grouped queries.
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: if pick.order_by is not None:
order_by_clause = "ORDER BY " + _render_order_by( order_node = render_order_by(pick.order_by, metric.name, dim_select_alias)
pick.order_by, metric.name, dim_select_alias
)
elif group_bindings: elif group_bindings:
# Sensible default: largest metric first. order_node = default_order_by(metric.name)
order_by_clause = f'ORDER BY "{metric.name}" DESC' else:
order_node = None
limit_clause = f"LIMIT {pick.limit}" if pick.limit is not None else "" # 8. Assemble one Select and render once.
sel = exp.select(*select_nodes).from_(aliases.aliased_table(metric.from_table))
for table_expr, on_expr in join_specs:
sel = sel.join(table_expr, on=on_expr, join_type="")
for node in where_nodes:
sel = sel.where(node)
if group_bindings:
sel = sel.group_by(*(exp.column(oa, quoted=True) for _, _, oa in group_bindings))
if order_node is not None:
sel = sel.order_by(order_node)
if pick.limit is not None:
sel = sel.limit(pick.limit)
# 8. Assemble. sql = sel.sql(dialect="postgres", identify=True)
sql_parts = [ return ComposeResult(sql=sql, used_tables=aliases.tables())
"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)

View File

@@ -1,22 +1,29 @@
"""Date-range resolver for typed Filters. """Date-range resolver for typed Filters.
Given a `Filter` with `date_range="YYYY"` or `date_range="YYYY-Q[1-4]"` and Given a `Filter` with `date_range="YYYY"` or `date_range="YYYY-Q[1-4]"`, the
the target column, render a SQL fragment that evaluates correctly under the target column, and the active encoding map, build a SQL predicate node that
column's actual storage type. Dispatched in two layers: evaluates correctly under the column's actual storage. Dispatched in two layers:
1. `Column.semantic_type` if set (extension point — e.g. `"date_yymmdd"` for 1. `Column.semantic_type` if it names an encoding in the effective map — coerce
datasets that ship YYMMDD-encoded integers). the raw column to a DATE via the encoding's `as_date` template, then compare.
2. Otherwise `Column.sql_type` (`DATE`, `TIMESTAMP`, …) — the default path This is the extension point: BIRD's `date_yymmdd` ships as a default
for warehouses that store dates as real date columns. (see `api.composer.encodings`); a dataset can add or override encodings in
its own `schema_docs.yaml`.
2. Otherwise `Column.sql_type` (`DATE`, `TIMESTAMP`, …) — compare the column
directly.
Inclusive on both ends. Bounds quoted with single quotes (Postgres parses Returns an `exp.Between` node (inclusive on both ends) so it composes into the
date literals from strings). query tree the composer assembles; the composer is the single place a SQL
string is produced.
""" """
from __future__ import annotations from __future__ import annotations
import re import re
import sqlglot
from sqlglot import exp
from api.composer.types import PickValidationError from api.composer.types import PickValidationError
from api.recon.types import Column from api.recon.types import Column
@@ -47,27 +54,37 @@ def _bounds(spec: str) -> tuple[str, str]:
) )
def resolve_date_range(spec: str, column: Column, qualified_col: str) -> str: def resolve_date_range(
"""Render the WHERE fragment for `<qualified_col> BETWEEN ...`. spec: str,
column: Column,
col_expr: exp.Expression,
encodings: dict[str, dict[str, str]],
) -> exp.Between:
"""Build the `<comparable> BETWEEN '<lo>' AND '<hi>'` predicate for `spec`.
`qualified_col` is the already-quoted reference the composer wants in `col_expr` is the (already alias-qualified) column reference node;
the SQL (e.g. `"l"."date"`). Returns just the predicate, no `WHERE`. `encodings` is the effective encoding map (see `api.composer.encodings`).
""" """
lo, hi = _bounds(spec) lo, hi = _bounds(spec)
semantic = (column.semantic_type or "").lower() bounds = (exp.convert(lo), exp.convert(hi))
if semantic == "date_yymmdd": semantic = (column.semantic_type or "").lower()
# Column stores YYMMDD as integer (some BIRD datasets ship this way). enc = encodings.get(semantic) if semantic else None
return ( if enc and enc.get("as_date"):
f"TO_DATE(LPAD({qualified_col}::text, 6, '0'), 'YYMMDD') " # Coerce the raw column to a DATE via the dataset/default template.
f"BETWEEN '{lo}' AND '{hi}'" template = enc["as_date"]
coerced = sqlglot.parse_one(
template.format(col=col_expr.sql(dialect="postgres", identify=True)),
dialect="postgres",
) )
return exp.Between(this=coerced, low=bounds[0], high=bounds[1])
sql_type = column.sql_type.upper() sql_type = column.sql_type.upper()
if sql_type.startswith("DATE") or sql_type.startswith("TIMESTAMP"): if sql_type.startswith("DATE") or sql_type.startswith("TIMESTAMP"):
return f"{qualified_col} BETWEEN '{lo}' AND '{hi}'" return exp.Between(this=col_expr, low=bounds[0], high=bounds[1])
raise PickValidationError( raise PickValidationError(
f"date_range on column with sql_type={column.sql_type!r} and " f"date_range on column with sql_type={column.sql_type!r} and "
f"semantic_type={column.semantic_type!r} is not supported " f"semantic_type={column.semantic_type!r} is not supported "
f"(no native date type and no matching encoding)"
) )

42
api/composer/encodings.py Normal file
View File

@@ -0,0 +1,42 @@
"""Storage-encoding library — how a column's raw storage maps to a value the
composer can compare against.
Some warehouses store logical types in surprising physical encodings (BIRD
ships YYMMDD-encoded *integer* "dates", for instance). The composer needs a way
to coerce such a column to a comparable value without baking any one dataset's
quirk into the generic SQL builder.
The model is **defaults with override**:
- `DEFAULT_ENCODINGS` ships the generic, reusable patterns here — keyed by the
`Column.semantic_type` a dataset declares in `schema_docs.yaml`.
- A dataset may add a novel encoding or override a default via the optional
`encodings:` section of its `schema_docs.yaml` (carried onto
`Recon.encodings`). The dataset always wins.
Each encoding is plain data: a map of role → SQL template with a `{col}`
placeholder for the (already-qualified) column reference. Today the only role
is `as_date` — the expression that yields a DATE for `date_range` filters; new
roles are additive.
"""
from __future__ import annotations
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from api.recon.types import Recon
DEFAULT_ENCODINGS: dict[str, dict[str, str]] = {
# BIRD-style YYMMDD stored as a 6-digit integer → real DATE.
"date_yymmdd": {"as_date": "TO_DATE(LPAD({col}::text, 6, '0'), 'YYMMDD')"},
# Add further reusable patterns here (e.g. epoch_seconds, julian_day) as
# they recur across datasets — never a single dataset's one-off quirk.
}
def effective_encodings(recon: "Recon") -> dict[str, dict[str, str]]:
"""The encoding map the composer should use for `recon`: built-in defaults
overlaid with the dataset's own `encodings` (dataset entries win)."""
return {**DEFAULT_ENCODINGS, **(recon.encodings or {})}

46
api/composer/filters.py Normal file
View File

@@ -0,0 +1,46 @@
"""Filter rendering — one typed `Filter` → one SQL predicate node.
Each filter shape maps to a sqlglot expression node (`exp.EQ`, `exp.In`,
`exp.Between`, or the date-range subtree). Literals go through `exp.convert`,
so quoting and escaping are sqlglot's job, not ours. The column is resolved to
its owning table and qualified via the shared `AliasMap` (allocating an alias,
and implying a join later, if the filter introduces a new table).
Adding a new filter shape is a new branch here that returns an `exp` node — no
change to the assembly in `compose.py`. The `Filter` type stays the constraint
boundary: there is no raw-SQL escape hatch.
"""
from __future__ import annotations
from sqlglot import exp
from api.composer.aliases import AliasMap
from api.composer.dates import resolve_date_range
from api.composer.types import Filter, PickValidationError
from api.recon.types import Recon
def render_filter(
f: Filter,
recon: Recon,
aliases: AliasMap,
encodings: dict[str, dict[str, str]],
) -> exp.Expression:
"""Render one `Filter` to a predicate node bound to the column's table."""
table, col = recon.resolve_column(f.column)
col_expr = aliases.col(table, col.name)
kind = f.kind()
if kind == "equals":
return exp.EQ(this=col_expr, expression=exp.convert(f.equals))
if kind == "in_values":
if not f.in_values:
raise PickValidationError(f"filter on {f.column!r} has empty in_values")
return exp.In(this=col_expr, expressions=[exp.convert(v) for v in f.in_values])
if kind == "between":
lo, hi = f.between
return exp.Between(this=col_expr, low=exp.convert(lo), high=exp.convert(hi))
if kind == "date_range":
return resolve_date_range(f.date_range, col, col_expr, encodings)
raise PickValidationError(f"unknown filter kind {kind!r}")

66
api/composer/joins.py Normal file
View File

@@ -0,0 +1,66 @@
"""Join-graph walking.
The composer never lets the LLM choose joins. Given the metric's base table and
the set of other tables a Pick pulled in (via group_by / filters), `build_joins`
walks `recon.join_path` for each target, dedupes the visited tables, and emits a
JOIN spec per new edge — direction resolved from the declared relationship so
the ON clause names the right column on each side.
Each spec is `(aliased_table_expr, on_expr)`: the composer applies them with
`select.join(table, on=on, join_type="")`. Aliases for intermediate tables are
allocated through the shared `AliasMap` as they're encountered.
"""
from __future__ import annotations
from sqlglot import exp
from api.composer.aliases import AliasMap
from api.composer.types import PickValidationError
from api.recon.types import Recon, Relationship
def _find_edge(recon: Recon, a: str, b: str) -> Relationship:
"""The declared relationship joining tables `a` and `b` (either direction).
Raises if no 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: AliasMap,
) -> list[tuple[exp.Expression, exp.Expression]]:
"""Union of join paths from `base` to each target → list of
`(aliased_table, on_predicate)` specs in walk order."""
visited: set[str] = {base}
specs: list[tuple[exp.Expression, exp.Expression]] = []
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 owns 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,
)
on_expr = exp.EQ(
this=aliases.col(left_t, left_c),
expression=aliases.col(right_t, right_c),
)
specs.append((aliases.aliased_table(cur), on_expr))
visited.add(cur)
return specs

30
api/composer/metrics.py Normal file
View File

@@ -0,0 +1,30 @@
"""Metric-fragment qualification.
Metric expressions in `metrics.yaml` are written *unaliased* — e.g.
`AVG(CASE WHEN status = 'B' THEN 1.0 WHEN status = 'A' THEN 0.0 END)`, `A13`,
`SUM(amount)`. Before such a fragment can sit in a multi-table query it has to
be bound to the metric's own table alias, so the composer can introduce other
tables via joins without ambiguity.
We parse the fragment with sqlglot and inject the alias as the table qualifier
on every bare column reference, returning the resulting expression node (not a
string) so it splices directly into the query tree the composer assembles.
"""
from __future__ import annotations
import sqlglot
from sqlglot import exp
def qualify_metric_expr(sql_fragment: str, table_alias: str) -> exp.Expression:
"""Parse `sql_fragment` and qualify its bare column refs to `table_alias`.
Used for both the metric's SELECT expression and its optional WHERE filter
(both are written as bare expressions over the metric's table)."""
tree = sqlglot.parse_one(sql_fragment, dialect="postgres")
for col in tree.find_all(exp.Column):
if col.table:
continue
col.set("table", exp.Identifier(this=table_alias, quoted=True))
return tree

38
api/composer/order.py Normal file
View File

@@ -0,0 +1,38 @@
"""ORDER BY rendering.
A Pick orders by either the metric or a named dimension; both reference the
SELECT-list *output alias* (the metric's name, or the dimension's output alias),
not the underlying qualified column. Returns an `exp.Ordered` node.
"""
from __future__ import annotations
from sqlglot import exp
from api.composer.types import OrderBy, PickValidationError
def _ordered(alias: str, direction: str) -> exp.Ordered:
return exp.Ordered(this=exp.column(alias, quoted=True), desc=(direction == "desc"))
def render_order_by(
ob: OrderBy,
metric_name: str,
dim_select: dict[str, str],
) -> exp.Ordered:
"""`metric_name` is the metric SELECT alias; `dim_select` maps a dimension
ref → its output alias."""
if ob.by == "metric":
return _ordered(metric_name, ob.direction)
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 _ordered(dim_select[ob.dimension], ob.direction)
def default_order_by(metric_name: str) -> exp.Ordered:
"""Sensible default when a grouped query gives no explicit order: largest
metric value first."""
return _ordered(metric_name, "desc")

View File

@@ -9,8 +9,19 @@
# tables.<name>: one-line description. # tables.<name>: one-line description.
# columns.<table>.<col>: one-line description. Sparse — only the ones # columns.<table>.<col>: one-line description. Sparse — only the ones
# worth describing. Undescribed columns just have # worth describing. Undescribed columns just have
# no description in the recon. # no description in the recon. A column may also
# carry a domain type, e.g.
# loan.date: {desc: "...", type: date_yymmdd}
# relationships: declared FKs (BIRD ships SQLite without them). # relationships: declared FKs (BIRD ships SQLite without them).
# encodings: OPTIONAL. Storage-encoding overrides keyed by a
# column's `type`. Each maps a role → SQL template
# with a `{col}` placeholder. The composer ships
# defaults (e.g. date_yymmdd); add an entry here only
# to introduce a novel encoding or override a default:
# encodings:
# date_yymmdd:
# as_date: "TO_DATE(LPAD({col}::text, 6, '0'), 'YYMMDD')"
# `financial` stores real DATE columns, so it needs none.
schema: financial schema: financial

View File

@@ -5,6 +5,7 @@ Endpoints:
POST /ask — submit a question, returns run_id POST /ask — submit a question, returns run_id
GET /runs/{run_id}/stream — SSE stream of run events GET /runs/{run_id}/stream — SSE stream of run events
GET /runs/{run_id} — final state snapshot GET /runs/{run_id} — final state snapshot
GET /runs/{run_id}/log — JSONL transcript (every published event)
""" """
from __future__ import annotations from __future__ import annotations
@@ -17,7 +18,7 @@ from datetime import datetime, timezone
from fastapi import FastAPI, HTTPException from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse from fastapi.responses import FileResponse, StreamingResponse
from pydantic import BaseModel from pydantic import BaseModel
from api.runtime import events from api.runtime import events
@@ -116,3 +117,13 @@ async def stream_run(run_id: str):
if run_id not in runs: if run_id not in runs:
raise HTTPException(404, detail=f"Run {run_id} not found") raise HTTPException(404, detail=f"Run {run_id} not found")
return StreamingResponse(events.stream(run_id), media_type="text/event-stream") return StreamingResponse(events.stream(run_id), media_type="text/event-stream")
@app.get("/runs/{run_id}/log")
async def get_run_log(run_id: str):
"""Serve the JSONL transcript written by `events.publish`. Available even
for in-flight runs (the file is flushed after every write)."""
path = events.log_path(run_id)
if not path.exists():
raise HTTPException(404, detail=f"No log file for run {run_id}")
return FileResponse(path, media_type="application/x-ndjson", filename=path.name)

View File

@@ -137,6 +137,11 @@ def merge_into_recon(dataset: str, extracted: dict[str, Any]) -> Recon:
relationships = [Relationship.parse(r) for r in (aug.get("relationships", []) or [])] relationships = [Relationship.parse(r) for r in (aug.get("relationships", []) or [])]
# Storage-encoding overrides (optional). Keyed by Column.semantic_type; each
# value maps a role (e.g. "as_date") to a SQL template with a `{col}`
# placeholder. The composer merges these over its built-in defaults.
encodings: dict[str, dict[str, str]] = aug.get("encodings", {}) or {}
column_to_tables: dict[str, list[str]] = {} column_to_tables: dict[str, list[str]] = {}
for t in tables.values(): for t in tables.values():
for c in t.columns: for c in t.columns:
@@ -150,6 +155,7 @@ def merge_into_recon(dataset: str, extracted: dict[str, Any]) -> Recon:
metrics=metrics, metrics=metrics,
column_to_tables=column_to_tables, column_to_tables=column_to_tables,
relationships=relationships, relationships=relationships,
encodings=encodings,
) )

View File

@@ -151,6 +151,12 @@ class Recon:
metrics: dict[str, Metric] metrics: dict[str, Metric]
column_to_tables: dict[str, list[str]] = field(default_factory=dict) column_to_tables: dict[str, list[str]] = field(default_factory=dict)
relationships: list[Relationship] = field(default_factory=list) relationships: list[Relationship] = field(default_factory=list)
# Storage-encoding overrides keyed by Column.semantic_type. Each value is a
# map of role → SQL template with a `{col}` placeholder (e.g.
# {"date_yymmdd": {"as_date": "TO_DATE(LPAD({col}::text, 6, '0'), 'YYMMDD')"}}).
# The composer merges these over its DEFAULT_ENCODINGS — the dataset wins.
# Empty for datasets whose columns are stored in native types.
encodings: dict[str, dict[str, str]] = field(default_factory=dict)
# ── Read-side helpers used by Analyses + the composer ── # ── Read-side helpers used by Analyses + the composer ──
@@ -299,6 +305,7 @@ class Recon:
"metrics": {n: m.to_dict() for n, m in self.metrics.items()}, "metrics": {n: m.to_dict() for n, m in self.metrics.items()},
"column_to_tables": self.column_to_tables, "column_to_tables": self.column_to_tables,
"relationships": [r.to_dict() for r in self.relationships], "relationships": [r.to_dict() for r in self.relationships],
"encodings": self.encodings,
} }
@classmethod @classmethod
@@ -309,4 +316,5 @@ class Recon:
metrics={n: Metric.from_dict(m) for n, m in d.get("metrics", {}).items()}, metrics={n: Metric.from_dict(m) for n, m in d.get("metrics", {}).items()},
column_to_tables=d.get("column_to_tables", {}), column_to_tables=d.get("column_to_tables", {}),
relationships=[Relationship(**r) for r in d.get("relationships", [])], relationships=[Relationship(**r) for r in d.get("relationships", [])],
encodings=d.get("encodings", {}) or {},
) )

View File

@@ -15,24 +15,59 @@ from __future__ import annotations
import asyncio import asyncio
import json import json
from typing import Any import logging
import time
from pathlib import Path
from typing import Any, IO
from api.analyses.types import Finding from api.analyses.types import Finding
from api.runtime.context import current_analysis, current_run_id from api.runtime.context import current_analysis, current_run_id
logger = logging.getLogger("nvi.runtime.events")
# run_id -> queue of events. Events are plain dicts. # run_id -> queue of events. Events are plain dicts.
_queues: dict[str, asyncio.Queue[dict[str, Any] | None]] = {} _queues: dict[str, asyncio.Queue[dict[str, Any] | None]] = {}
# run_id -> append-mode JSONL log file. Mirrors every published event so the
# whole run is replayable from disk. Path: `<LOG_DIR>/<run_id>.jsonl`.
_log_files: dict[str, IO[str]] = {}
LOG_DIR = Path(".data/logs")
def log_path(run_id: str) -> Path:
"""Where this run's JSONL log lives. Resolved relative to cwd of the api
process — `/app/.data/logs/` in the pod."""
return LOG_DIR / f"{run_id}.jsonl"
# ── Queue management ── # ── Queue management ──
def open_run(run_id: str) -> asyncio.Queue: def open_run(run_id: str) -> asyncio.Queue:
q: asyncio.Queue = asyncio.Queue() q: asyncio.Queue = asyncio.Queue()
_queues[run_id] = q _queues[run_id] = q
try:
LOG_DIR.mkdir(parents=True, exist_ok=True)
_log_files[run_id] = log_path(run_id).open("a", encoding="utf-8")
except OSError as e:
# Filesystem issues shouldn't kill the run; just lose the on-disk log.
logger.warning("could not open log file for run %s: %s", run_id, e)
return q return q
def _write_log(run_id: str, event: dict[str, Any]) -> None:
f = _log_files.get(run_id)
if f is None:
return
try:
record = {"t": time.time(), **event}
f.write(json.dumps(record, default=str) + "\n")
f.flush() # so `kubectl exec ... -- tail -f` shows live progress
except (OSError, ValueError) as e:
logger.warning("log write failed for run %s: %s", run_id, e)
async def publish(run_id: str, event: dict[str, Any]) -> None: async def publish(run_id: str, event: dict[str, Any]) -> None:
_write_log(run_id, event)
q = _queues.get(run_id) q = _queues.get(run_id)
if q is not None: if q is not None:
await q.put(event) await q.put(event)
@@ -54,6 +89,12 @@ async def close(run_id: str) -> None:
def drop(run_id: str) -> None: def drop(run_id: str) -> None:
_queues.pop(run_id, None) _queues.pop(run_id, None)
f = _log_files.pop(run_id, None)
if f is not None:
try:
f.close()
except OSError:
pass
# ── SSE formatting + streaming ── # ── SSE formatting + streaming ──

View File

@@ -1,10 +1,25 @@
"""Date-range resolver tests — bounds parsing + dispatch on column type.""" """Date-range resolver tests — bounds parsing + dispatch on column type/encoding.
`resolve_date_range` now returns an `exp` node and takes the effective encoding
map (built-in defaults overlaid with any dataset overrides). Tests render the
node to SQL to assert on the predicate.
"""
import pytest import pytest
from sqlglot import exp
from api.composer.dates import _bounds, resolve_date_range from api.composer.dates import _bounds, resolve_date_range
from api.composer.encodings import DEFAULT_ENCODINGS, effective_encodings
from api.composer.types import PickValidationError from api.composer.types import PickValidationError
from api.recon.types import Column from api.recon.types import Column, Recon
def _col(name: str = "date", table: str = "l") -> exp.Column:
return exp.column(name, table=table, quoted=True)
def _sql(node: exp.Expression) -> str:
return node.sql(dialect="postgres", identify=True)
def test_bounds_year(): def test_bounds_year():
@@ -26,25 +41,56 @@ def test_bounds_unsupported_raises():
def test_resolve_real_date_column(): def test_resolve_real_date_column():
col = Column(name="date", sql_type="DATE", nullable=False) col = Column(name="date", sql_type="DATE", nullable=False)
out = resolve_date_range("1996", col, '"l"."date"') out = _sql(resolve_date_range("1996", col, _col(), {}))
assert out == "\"l\".\"date\" BETWEEN '1996-01-01' AND '1996-12-31'" assert out == "\"l\".\"date\" BETWEEN '1996-01-01' AND '1996-12-31'"
def test_resolve_timestamp_column(): def test_resolve_timestamp_column():
col = Column(name="created_at", sql_type="TIMESTAMP WITHOUT TIME ZONE", nullable=True) col = Column(name="created_at", sql_type="TIMESTAMP WITHOUT TIME ZONE", nullable=True)
out = resolve_date_range("1996-Q2", col, '"x"."created_at"') out = _sql(resolve_date_range("1996-Q2", col, _col("created_at", "x"), {}))
assert out == "\"x\".\"created_at\" BETWEEN '1996-04-01' AND '1996-06-30'" assert out == "\"x\".\"created_at\" BETWEEN '1996-04-01' AND '1996-06-30'"
def test_resolve_yymmdd_semantic_type(): def test_resolve_yymmdd_via_default_encoding():
col = Column(name="date", sql_type="INTEGER", nullable=False, col = Column(name="date", sql_type="INTEGER", nullable=False,
semantic_type="date_yymmdd") semantic_type="date_yymmdd")
out = resolve_date_range("1996", col, '"l"."date"') out = _sql(resolve_date_range("1996", col, _col(), DEFAULT_ENCODINGS))
assert "TO_DATE(LPAD(" in out assert "TO_DATE(LPAD(" in out
assert "BETWEEN '1996-01-01' AND '1996-12-31'" in out assert "BETWEEN '1996-01-01' AND '1996-12-31'" in out
def test_dataset_encoding_overrides_default():
# A dataset can override a built-in encoding with its own coercion SQL.
col = Column(name="d", sql_type="INTEGER", nullable=False,
semantic_type="date_yymmdd")
override = {"date_yymmdd": {"as_date": "MAKE_DATE(1900 + {col} / 10000, 1, 1)"}}
out = _sql(resolve_date_range("1996", col, _col("d"), override))
assert "MAKE_DATE(" in out
assert "TO_DATE" not in out
def test_resolve_semantic_type_without_encoding_raises():
# Declared semantic_type but no matching encoding and not a native date type.
col = Column(name="d", sql_type="INTEGER", nullable=False,
semantic_type="date_yymmdd")
with pytest.raises(PickValidationError, match="not supported"):
resolve_date_range("1996", col, _col("d"), {})
def test_resolve_unsupported_column_type_raises(): def test_resolve_unsupported_column_type_raises():
col = Column(name="x", sql_type="TEXT", nullable=True) col = Column(name="x", sql_type="TEXT", nullable=True)
with pytest.raises(PickValidationError, match="not supported"): with pytest.raises(PickValidationError, match="not supported"):
resolve_date_range("1996", col, '"x"."x"') resolve_date_range("1996", col, _col("x", "x"), {})
def test_effective_encodings_merges_defaults_with_dataset():
recon = Recon(schema="t", tables={}, metrics={},
encodings={"my_epoch": {"as_date": "TO_TIMESTAMP({col})::date"}})
eff = effective_encodings(recon)
assert "date_yymmdd" in eff # default survives
assert eff["my_epoch"]["as_date"].startswith("TO_TIMESTAMP") # dataset added
# Dataset entry wins on key collision.
recon2 = Recon(schema="t", tables={}, metrics={},
encodings={"date_yymmdd": {"as_date": "CUSTOM({col})"}})
assert effective_encodings(recon2)["date_yymmdd"]["as_date"] == "CUSTOM({col})"

View File

@@ -141,7 +141,6 @@ function nodeClasses(n: GraphNode) {
<div class="trace-pane"> <div class="trace-pane">
<Panel title="Plan" :status="statusOf(status)"> <Panel title="Plan" :status="statusOf(status)">
<div v-if="rationale" class="rationale">{{ rationale }}</div> <div v-if="rationale" class="rationale">{{ rationale }}</div>
<div v-else class="muted small">Waiting for the planner</div>
<div class="graph"> <div class="graph">
<div v-for="(n, idx) in graphNodes" :key="n.id" class="graph-row"> <div v-for="(n, idx) in graphNodes" :key="n.id" class="graph-row">