refactor composer to use sqlalchemy, avoid string concatenations and follow conventions
This commit is contained in:
56
api/composer/aliases.py
Normal file
56
api/composer/aliases.py
Normal 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)
|
||||
@@ -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())
|
||||
|
||||
@@ -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 `<qualified_col> BETWEEN ...`.
|
||||
def resolve_date_range(
|
||||
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
|
||||
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)"
|
||||
)
|
||||
|
||||
42
api/composer/encodings.py
Normal file
42
api/composer/encodings.py
Normal 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
46
api/composer/filters.py
Normal 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
66
api/composer/joins.py
Normal 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
30
api/composer/metrics.py
Normal 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
38
api/composer/order.py
Normal 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")
|
||||
@@ -9,8 +9,19 @@
|
||||
# tables.<name>: one-line description.
|
||||
# columns.<table>.<col>: 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
|
||||
|
||||
|
||||
13
api/main.py
13
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)
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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 {},
|
||||
)
|
||||
|
||||
@@ -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: `<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 ──
|
||||
|
||||
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 ──
|
||||
|
||||
Reference in New Issue
Block a user