diff --git a/Makefile b/Makefile index 202ae52..d34a060 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,5 @@ .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 \ docs-graphs @@ -51,6 +52,23 @@ recon: 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/.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-log: + @[ -n "$(RUN)" ] || { echo "usage: make run-log RUN=" >&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). # Ensures dev extras are installed first; uv is idempotent on no-ops. test: diff --git a/api/composer/aliases.py b/api/composer/aliases.py new file mode 100644 index 0000000..46aff10 --- /dev/null +++ b/api/composer/aliases.py @@ -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 `` 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: + """`"" AS ""` 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) diff --git a/api/composer/compose.py b/api/composer/compose.py index d21b7b5..65557e2 100644 --- a/api/composer/compose.py +++ b/api/composer/compose.py @@ -1,191 +1,48 @@ """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. +Deterministic SQL emission from a typed Pick. The composer assembles a single +sqlglot expression tree and renders it once — there is no hand-built SQL string +and no after-the-fact re-parse; the tree *is* the structure, and +`.sql(dialect="postgres", identify=True)` (every identifier quoted) is the one +place a string is produced. -No LLM call anywhere in this module. If the Pick can't be expressed against -recon, the composer raises `PickValidationError` — the caller surfaces it. +The work is split across focused modules, each returning `exp` nodes: + + 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 -import sqlglot +from sqlglot import exp -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 +from api.composer.aliases import AliasMap +from api.composer.encodings import effective_encodings +from api.composer.filters import render_filter +from api.composer.joins import build_joins +from api.composer.metrics import qualify_metric_expr +from api.composer.order import default_order_by, render_order_by +from api.composer.types import ComposeResult, Pick, PickValidationError +from api.recon.types import Recon -# ── 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.""" + """Render `pick` to SQL against `recon`. Raises PickValidationError on any + reference the recon can't resolve.""" - # 1. Metric. + # 1. Metric → base table. if pick.metric not in recon.metrics: raise PickValidationError( 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" ) - aliases: dict[str, str] = {} - aliases[metric.from_table] = _alias_for(metric.from_table, {}) - base_alias = aliases[metric.from_table] + aliases = AliasMap() + base_alias = aliases.alias(metric.from_table) + encodings = effective_encodings(recon) - # 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) + # 2. Resolve group_by columns → owning table + output alias. Tables other + # than the metric's need joins. + group_bindings: list[tuple[str, str, str]] = [] # (table, col_name, out_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 not aliases.has(table): if table != metric.from_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(".", "_") - group_bindings.append((ref, table, col.name, out_alias)) + group_bindings.append((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] = [] + # 3. Resolve filters → predicate nodes (may also introduce new tables). + filter_nodes: list[exp.Expression] = [] for f in pick.where: table, _ = recon.resolve_column(f.column) - if table not in aliases: + if not aliases.has(table): 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)) + aliases.alias(table) + filter_nodes.append(render_filter(f, recon, aliases, encodings)) - # 4. Build JOINs for the union of extra tables. - join_clauses = _build_joins(recon, metric.from_table, extra_tables, aliases) + # 4. JOINs for the union of extra tables. + join_specs = 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}"') + # 5. SELECT list: dimensions (in group_by order), then the metric. + select_nodes: list[exp.Expression] = [ + exp.alias_(aliases.col(table, col_name), out_alias, quoted=True) + for table, col_name, out_alias in group_bindings + ] + select_nodes.append( + 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. - where_parts: list[str] = [] + where_nodes: list[exp.Expression] = [] if metric.filter: - where_parts.append(_qualify_metric_filter(metric.filter, base_alias)) - where_parts.extend(filter_clauses) + where_nodes.append(qualify_metric_expr(metric.filter, base_alias)) + where_nodes.extend(filter_nodes) - # 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 = "" + # 7. ORDER BY: explicit, else largest-metric-first for grouped queries. if pick.order_by is not None: - order_by_clause = "ORDER BY " + _render_order_by( - pick.order_by, metric.name, dim_select_alias - ) + order_node = 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' + order_node = default_order_by(metric.name) + 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_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) + sql = sel.sql(dialect="postgres", identify=True) + return ComposeResult(sql=sql, used_tables=aliases.tables()) diff --git a/api/composer/dates.py b/api/composer/dates.py index d870160..18256aa 100644 --- a/api/composer/dates.py +++ b/api/composer/dates.py @@ -1,22 +1,29 @@ """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: +Given a `Filter` with `date_range="YYYY"` or `date_range="YYYY-Q[1-4]"`, the +target column, and the active encoding map, build a SQL predicate node that +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 - 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. +1. `Column.semantic_type` if it names an encoding in the effective map — coerce + the raw column to a DATE via the encoding's `as_date` template, then compare. + This is the extension point: BIRD's `date_yymmdd` ships as a default + (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 -date literals from strings). +Returns an `exp.Between` node (inclusive on both ends) so it composes into the +query tree the composer assembles; the composer is the single place a SQL +string is produced. """ from __future__ import annotations import re +import sqlglot +from sqlglot import exp + from api.composer.types import PickValidationError 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: - """Render the WHERE fragment for ` BETWEEN ...`. +def resolve_date_range( + spec: str, + column: Column, + col_expr: exp.Expression, + encodings: dict[str, dict[str, str]], +) -> exp.Between: + """Build the ` BETWEEN '' AND ''` predicate for `spec`. - `qualified_col` is the already-quoted reference the composer wants in - the SQL (e.g. `"l"."date"`). Returns just the predicate, no `WHERE`. + `col_expr` is the (already alias-qualified) column reference node; + `encodings` is the effective encoding map (see `api.composer.encodings`). """ lo, hi = _bounds(spec) - semantic = (column.semantic_type or "").lower() + bounds = (exp.convert(lo), exp.convert(hi)) - 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}'" + semantic = (column.semantic_type or "").lower() + enc = encodings.get(semantic) if semantic else None + if enc and enc.get("as_date"): + # Coerce the raw column to a DATE via the dataset/default template. + 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() 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( 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)" ) diff --git a/api/composer/encodings.py b/api/composer/encodings.py new file mode 100644 index 0000000..54e5225 --- /dev/null +++ b/api/composer/encodings.py @@ -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 {})} diff --git a/api/composer/filters.py b/api/composer/filters.py new file mode 100644 index 0000000..e0be3de --- /dev/null +++ b/api/composer/filters.py @@ -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}") diff --git a/api/composer/joins.py b/api/composer/joins.py new file mode 100644 index 0000000..435f018 --- /dev/null +++ b/api/composer/joins.py @@ -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 diff --git a/api/composer/metrics.py b/api/composer/metrics.py new file mode 100644 index 0000000..a56ed6d --- /dev/null +++ b/api/composer/metrics.py @@ -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 diff --git a/api/composer/order.py b/api/composer/order.py new file mode 100644 index 0000000..2e81ff9 --- /dev/null +++ b/api/composer/order.py @@ -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") diff --git a/api/datasets/financial/schema_docs.yaml b/api/datasets/financial/schema_docs.yaml index 84f2c87..30eadef 100644 --- a/api/datasets/financial/schema_docs.yaml +++ b/api/datasets/financial/schema_docs.yaml @@ -9,8 +9,19 @@ # tables.: one-line description. # columns.
.: one-line description. Sparse — only the ones # 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). +# 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 diff --git a/api/main.py b/api/main.py index d9328a6..a1b8ce6 100644 --- a/api/main.py +++ b/api/main.py @@ -5,6 +5,7 @@ Endpoints: POST /ask — submit a question, returns run_id GET /runs/{run_id}/stream — SSE stream of run events GET /runs/{run_id} — final state snapshot + GET /runs/{run_id}/log — JSONL transcript (every published event) """ from __future__ import annotations @@ -17,7 +18,7 @@ from datetime import datetime, timezone from fastapi import FastAPI, HTTPException from fastapi.middleware.cors import CORSMiddleware -from fastapi.responses import StreamingResponse +from fastapi.responses import FileResponse, StreamingResponse from pydantic import BaseModel from api.runtime import events @@ -116,3 +117,13 @@ async def stream_run(run_id: str): if run_id not in runs: raise HTTPException(404, detail=f"Run {run_id} not found") 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) diff --git a/api/recon/build.py b/api/recon/build.py index 5aa287b..68ed35c 100644 --- a/api/recon/build.py +++ b/api/recon/build.py @@ -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 [])] + # 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]] = {} for t in tables.values(): for c in t.columns: @@ -150,6 +155,7 @@ def merge_into_recon(dataset: str, extracted: dict[str, Any]) -> Recon: metrics=metrics, column_to_tables=column_to_tables, relationships=relationships, + encodings=encodings, ) diff --git a/api/recon/types.py b/api/recon/types.py index e61556c..b84b082 100644 --- a/api/recon/types.py +++ b/api/recon/types.py @@ -151,6 +151,12 @@ class Recon: metrics: dict[str, Metric] column_to_tables: dict[str, list[str]] = field(default_factory=dict) 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 ── @@ -299,6 +305,7 @@ class Recon: "metrics": {n: m.to_dict() for n, m in self.metrics.items()}, "column_to_tables": self.column_to_tables, "relationships": [r.to_dict() for r in self.relationships], + "encodings": self.encodings, } @classmethod @@ -309,4 +316,5 @@ class Recon: metrics={n: Metric.from_dict(m) for n, m in d.get("metrics", {}).items()}, column_to_tables=d.get("column_to_tables", {}), relationships=[Relationship(**r) for r in d.get("relationships", [])], + encodings=d.get("encodings", {}) or {}, ) diff --git a/api/runtime/events.py b/api/runtime/events.py index c679e50..ab5e8dd 100644 --- a/api/runtime/events.py +++ b/api/runtime/events.py @@ -15,24 +15,59 @@ from __future__ import annotations import asyncio 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.runtime.context import current_analysis, current_run_id +logger = logging.getLogger("nvi.runtime.events") + # run_id -> queue of events. Events are plain dicts. _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: `/.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 ── def open_run(run_id: str) -> asyncio.Queue: q: asyncio.Queue = asyncio.Queue() _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 +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: + _write_log(run_id, event) q = _queues.get(run_id) if q is not None: await q.put(event) @@ -54,6 +89,12 @@ async def close(run_id: str) -> None: def drop(run_id: str) -> 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 ── diff --git a/tests/test_composer_dates.py b/tests/test_composer_dates.py index 310a498..639d2e1 100644 --- a/tests/test_composer_dates.py +++ b/tests/test_composer_dates.py @@ -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 +from sqlglot import exp 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.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(): @@ -26,25 +41,56 @@ def test_bounds_unsupported_raises(): def test_resolve_real_date_column(): 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'" def test_resolve_timestamp_column(): 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'" -def test_resolve_yymmdd_semantic_type(): +def test_resolve_yymmdd_via_default_encoding(): col = Column(name="date", sql_type="INTEGER", nullable=False, 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 "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(): col = Column(name="x", sql_type="TEXT", nullable=True) 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})" diff --git a/ui/app/src/pages/Ask.vue b/ui/app/src/pages/Ask.vue index 10123ae..dad5ca2 100644 --- a/ui/app/src/pages/Ask.vue +++ b/ui/app/src/pages/Ask.vue @@ -141,7 +141,6 @@ function nodeClasses(n: GraphNode) {
{{ rationale }}
-
Waiting for the planner…