289 lines
12 KiB
Python
289 lines
12 KiB
Python
"""compose(pick, recon) → ComposeResult.
|
|
|
|
Deterministic SQL emission from a typed Pick. The composer:
|
|
1. Resolves the metric → from_table + sql expression + default filter.
|
|
2. Resolves each group_by column → owning table (via recon.resolve_column).
|
|
3. Computes the join graph: join_path(metric.from_table, owner) per dim,
|
|
deduped, walked in declared order to pick the right relationship edge.
|
|
4. Renders SELECT / FROM + JOINs / WHERE / GROUP BY / ORDER BY / LIMIT
|
|
with every identifier quoted via sqlglot's `identify=True` round-trip
|
|
as a paranoid post-check.
|
|
|
|
No LLM call anywhere in this module. If the Pick can't be expressed against
|
|
recon, the composer raises `PickValidationError` — the caller surfaces it.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sqlglot
|
|
|
|
from api.composer.dates import resolve_date_range
|
|
from api.composer.types import (
|
|
ComposeResult,
|
|
Filter,
|
|
OrderBy,
|
|
Pick,
|
|
PickValidationError,
|
|
)
|
|
from api.recon.types import Recon, Relationship
|
|
|
|
|
|
# ── Alias allocation ────────────────────────────────────────
|
|
|
|
def _alias_for(table: str, used: dict[str, str]) -> str:
|
|
"""Pick a short alias for `table` that's unique within the query.
|
|
|
|
First letter, then first-two, then numbered. `used` maps alias→table so
|
|
we can grow on collision.
|
|
"""
|
|
for base in (table[:1], table[:2], table[:3], table):
|
|
if base and base not in used:
|
|
return base
|
|
i = 1
|
|
while True:
|
|
cand = f"{table[:1]}{i}"
|
|
if cand not in used:
|
|
return cand
|
|
i += 1
|
|
|
|
|
|
# ── Join graph ──────────────────────────────────────────────
|
|
|
|
def _find_edge(recon: Recon, a: str, b: str) -> Relationship:
|
|
"""Return the relationship that joins tables a and b (either direction).
|
|
Raises if no declared FK connects them."""
|
|
for r in recon.relationships:
|
|
if (r.from_table == a and r.to_table == b) or (r.from_table == b and r.to_table == a):
|
|
return r
|
|
raise PickValidationError(
|
|
f"no declared relationship between {a!r} and {b!r}"
|
|
)
|
|
|
|
|
|
def _build_joins(recon: Recon, base: str, targets: list[str],
|
|
aliases: dict[str, str]) -> list[str]:
|
|
"""Compute the union of join paths from `base` to each target table,
|
|
emit JOIN clauses in walk order. Allocates aliases for any intermediate
|
|
table that wasn't already in `aliases`."""
|
|
visited: set[str] = {base}
|
|
clauses: list[str] = []
|
|
for target in targets:
|
|
path = recon.join_path(base, target)
|
|
if path is None:
|
|
raise PickValidationError(
|
|
f"no join path from {base!r} to {target!r} in recon"
|
|
)
|
|
for i in range(1, len(path)):
|
|
prev, cur = path[i - 1], path[i]
|
|
if cur in visited:
|
|
continue
|
|
edge = _find_edge(recon, prev, cur)
|
|
# Edge direction tells us which side has which column.
|
|
if edge.from_table == prev:
|
|
left_t, left_c, right_t, right_c = edge.from_table, edge.from_column, edge.to_table, edge.to_column
|
|
else:
|
|
left_t, left_c, right_t, right_c = edge.to_table, edge.to_column, edge.from_table, edge.from_column
|
|
# Ensure both sides have aliases.
|
|
for t in (left_t, right_t):
|
|
if t not in aliases:
|
|
aliases[t] = _alias_for(t, {a: t for t, a in aliases.items()})
|
|
la, ra = aliases[left_t], aliases[right_t]
|
|
clauses.append(
|
|
f'JOIN "{cur}" AS "{aliases[cur]}" '
|
|
f'ON "{la}"."{left_c}" = "{ra}"."{right_c}"'
|
|
)
|
|
visited.add(cur)
|
|
return clauses
|
|
|
|
|
|
# ── Filter rendering ────────────────────────────────────────
|
|
|
|
def _render_literal(v) -> str:
|
|
"""Render a Python value as a SQL literal. Strings are single-quoted with
|
|
embedded quotes escaped. Numbers pass through."""
|
|
if isinstance(v, str):
|
|
return "'" + v.replace("'", "''") + "'"
|
|
if isinstance(v, bool):
|
|
return "TRUE" if v else "FALSE"
|
|
if v is None:
|
|
return "NULL"
|
|
return str(v)
|
|
|
|
|
|
def _render_filter(f: Filter, recon: Recon, aliases: dict[str, str]) -> str:
|
|
"""Render one Filter to a SQL predicate. Resolves the column's owning
|
|
table and aliases it (allocating an alias if needed — and a join later
|
|
if that introduces a new table)."""
|
|
table, col = recon.resolve_column(f.column)
|
|
if table not in aliases:
|
|
aliases[table] = _alias_for(table, {a: t for t, a in aliases.items()})
|
|
qualified = f'"{aliases[table]}"."{col.name}"'
|
|
|
|
kind = f.kind()
|
|
if kind == "equals":
|
|
return f"{qualified} = {_render_literal(f.equals)}"
|
|
if kind == "in_values":
|
|
if not f.in_values:
|
|
raise PickValidationError(
|
|
f"filter on {f.column!r} has empty in_values"
|
|
)
|
|
rendered = ", ".join(_render_literal(v) for v in f.in_values)
|
|
return f"{qualified} IN ({rendered})"
|
|
if kind == "between":
|
|
lo, hi = f.between
|
|
return f"{qualified} BETWEEN {_render_literal(lo)} AND {_render_literal(hi)}"
|
|
if kind == "date_range":
|
|
return resolve_date_range(f.date_range, col, qualified)
|
|
raise PickValidationError(f"unknown filter kind {kind!r}")
|
|
|
|
|
|
# ── ORDER BY rendering ──────────────────────────────────────
|
|
|
|
def _render_order_by(ob: OrderBy, metric_name: str, dim_select: dict[str, str]) -> str:
|
|
"""`metric_name` is the alias of the metric SELECT expression;
|
|
`dim_select` maps dimension ref (column ref) → output alias."""
|
|
if ob.by == "metric":
|
|
return f'"{metric_name}" {ob.direction.upper()}'
|
|
# by == "dimension"
|
|
if ob.dimension not in dim_select:
|
|
raise PickValidationError(
|
|
f"order_by.dimension={ob.dimension!r} is not in group_by ({list(dim_select)})"
|
|
)
|
|
return f'"{dim_select[ob.dimension]}" {ob.direction.upper()}'
|
|
|
|
|
|
# ── Metric expression rewriting ─────────────────────────────
|
|
|
|
def _qualify_metric_sql(sql_fragment: str, table_alias: str) -> str:
|
|
"""Rewrite bare column refs inside the metric's SQL expression so they
|
|
point at the metric table's alias.
|
|
|
|
Metric.sql in metrics.yaml is written as an unaliased expression (e.g.
|
|
`AVG(CASE WHEN status='B' ... END)`). We need it qualified to the
|
|
metric table's alias so the composer can introduce other tables via
|
|
joins without ambiguity. Uses sqlglot to parse and rewrite identifier
|
|
refs that don't already have a table prefix.
|
|
"""
|
|
tree = sqlglot.parse_one(sql_fragment, dialect="postgres")
|
|
for col in tree.find_all(sqlglot.exp.Column):
|
|
if col.table:
|
|
continue
|
|
# Inject the alias as the table qualifier.
|
|
col.set("table", sqlglot.exp.Identifier(this=table_alias, quoted=True))
|
|
return tree.sql(dialect="postgres", identify=True)
|
|
|
|
|
|
def _qualify_metric_filter(filter_fragment: str, table_alias: str) -> str:
|
|
"""Same as _qualify_metric_sql but for the metric's optional WHERE filter.
|
|
metric.filter is written as a bare boolean expression."""
|
|
return _qualify_metric_sql(filter_fragment, table_alias)
|
|
|
|
|
|
# ── Compose ─────────────────────────────────────────────────
|
|
|
|
def compose(pick: Pick, recon: Recon) -> ComposeResult:
|
|
"""Render `pick` to SQL against `recon`. Raises PickValidationError on
|
|
any reference the recon can't resolve."""
|
|
|
|
# 1. Metric.
|
|
if pick.metric not in recon.metrics:
|
|
raise PickValidationError(
|
|
f"metric {pick.metric!r} not in recon (known: {sorted(recon.metrics)})"
|
|
)
|
|
metric = recon.metrics[pick.metric]
|
|
if metric.from_table not in recon.tables:
|
|
raise PickValidationError(
|
|
f"metric {pick.metric!r} declares from_table={metric.from_table!r} "
|
|
f"but no such table in recon"
|
|
)
|
|
|
|
aliases: dict[str, str] = {}
|
|
aliases[metric.from_table] = _alias_for(metric.from_table, {})
|
|
base_alias = aliases[metric.from_table]
|
|
|
|
# 2. Resolve every group_by column to (table, Column) and collect target
|
|
# tables that aren't the metric's table — those need joins.
|
|
group_bindings: list[tuple[str, str, str, str]] = [] # (ref, table, col_name, output_alias)
|
|
extra_tables: list[str] = []
|
|
dim_select_alias: dict[str, str] = {}
|
|
for ref in pick.group_by:
|
|
table, col = recon.resolve_column(ref)
|
|
if table not in aliases:
|
|
if table != metric.from_table:
|
|
extra_tables.append(table)
|
|
aliases[table] = _alias_for(table, {a: t for t, a in aliases.items()})
|
|
out_alias = col.name if ref == col.name else ref.replace(".", "_")
|
|
group_bindings.append((ref, table, col.name, out_alias))
|
|
dim_select_alias[ref] = out_alias
|
|
|
|
# 3. Resolve filter columns first (they may also introduce new tables we
|
|
# need to join). Rendered separately so we can keep their predicates
|
|
# in WHERE.
|
|
filter_clauses: list[str] = []
|
|
for f in pick.where:
|
|
table, _ = recon.resolve_column(f.column)
|
|
if table not in aliases:
|
|
if table != metric.from_table:
|
|
extra_tables.append(table)
|
|
aliases[table] = _alias_for(table, {a: t for t, a in aliases.items()})
|
|
filter_clauses.append(_render_filter(f, recon, aliases))
|
|
|
|
# 4. Build JOINs for the union of extra tables.
|
|
join_clauses = _build_joins(recon, metric.from_table, extra_tables, aliases)
|
|
|
|
# 5. SELECT list: dimension columns first (in group_by order), then the
|
|
# metric expression.
|
|
select_parts: list[str] = []
|
|
for ref, table, col_name, out_alias in group_bindings:
|
|
select_parts.append(f'"{aliases[table]}"."{col_name}" AS "{out_alias}"')
|
|
metric_expr = _qualify_metric_sql(metric.sql, base_alias)
|
|
select_parts.append(f'{metric_expr} AS "{metric.name}"')
|
|
|
|
# 6. WHERE: metric's default filter (if any) ANDed with the Pick's filters.
|
|
where_parts: list[str] = []
|
|
if metric.filter:
|
|
where_parts.append(_qualify_metric_filter(metric.filter, base_alias))
|
|
where_parts.extend(filter_clauses)
|
|
|
|
# 7. GROUP BY (just the dim output aliases) + ORDER BY + LIMIT.
|
|
group_by_clause = ""
|
|
if group_bindings:
|
|
group_by_clause = "GROUP BY " + ", ".join(
|
|
f'"{out_alias}"' for _, _, _, out_alias in group_bindings
|
|
)
|
|
|
|
order_by_clause = ""
|
|
if pick.order_by is not None:
|
|
order_by_clause = "ORDER BY " + _render_order_by(
|
|
pick.order_by, metric.name, dim_select_alias
|
|
)
|
|
elif group_bindings:
|
|
# Sensible default: largest metric first.
|
|
order_by_clause = f'ORDER BY "{metric.name}" DESC'
|
|
|
|
limit_clause = f"LIMIT {pick.limit}" if pick.limit is not None else ""
|
|
|
|
# 8. Assemble.
|
|
sql_parts = [
|
|
"SELECT " + ", ".join(select_parts),
|
|
f'FROM "{metric.from_table}" AS "{base_alias}"',
|
|
]
|
|
sql_parts.extend(join_clauses)
|
|
if where_parts:
|
|
sql_parts.append("WHERE " + " AND ".join(where_parts))
|
|
if group_by_clause:
|
|
sql_parts.append(group_by_clause)
|
|
if order_by_clause:
|
|
sql_parts.append(order_by_clause)
|
|
if limit_clause:
|
|
sql_parts.append(limit_clause)
|
|
raw = "\n".join(sql_parts)
|
|
|
|
# 9. Paranoid re-parse: catches composer bugs by round-tripping through
|
|
# sqlglot. identify=True keeps every identifier quoted.
|
|
parsed = sqlglot.parse_one(raw, dialect="postgres")
|
|
sql = parsed.sql(dialect="postgres", identify=True)
|
|
|
|
used = sorted(aliases)
|
|
return ComposeResult(sql=sql, used_tables=used)
|