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")
|
||||
Reference in New Issue
Block a user