mariano/nvi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
buenosairesam
2026-06-03 12:10:30 -03:00
parent cbc7df8c60
commit be09fcde2c
16 changed files with 546 additions and 273 deletions
Download Patch File Download Diff File Expand all files Collapse all files

320
api/composer/compose.py
View File

@@ -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())