refactor composer to use sqlalchemy, avoid string concatenations and follow conventions
This commit is contained in:
18
Makefile
18
Makefile
@@ -1,4 +1,5 @@
|
|||||||
.PHONY: kind tilt-up tilt-down seed seed-fetch seed-push recon evals test \
|
.PHONY: kind tilt-up tilt-down seed seed-fetch seed-push recon evals test \
|
||||||
|
run-log run-logs \
|
||||||
compose-up compose-down compose-clean compose-seed compose-recon compose-evals \
|
compose-up compose-down compose-clean compose-seed compose-recon compose-evals \
|
||||||
docs-graphs
|
docs-graphs
|
||||||
|
|
||||||
@@ -51,6 +52,23 @@ recon:
|
|||||||
evals:
|
evals:
|
||||||
kubectl $(KCTX) $(KNS) exec deploy/api -- uv run python -m api.evals.run_evals
|
kubectl $(KCTX) $(KNS) exec deploy/api -- uv run python -m api.evals.run_evals
|
||||||
|
|
||||||
|
# Pull a run's JSONL transcript from the api pod to the host. The api writes
|
||||||
|
# every published event into /app/.data/logs/<RUN>.jsonl inside the pod;
|
||||||
|
# this target copies it next to .data/logs/ on the host so it's easy to
|
||||||
|
# share / grep / paste a path to.
|
||||||
|
# Usage: make run-log RUN=<run_id>
|
||||||
|
run-log:
|
||||||
|
@[ -n "$(RUN)" ] || { echo "usage: make run-log RUN=<run_id>" >&2; exit 1; }
|
||||||
|
@mkdir -p .data/logs
|
||||||
|
@POD=$$(kubectl $(KCTX) $(KNS) get pod -l app=api -o jsonpath='{.items[0].metadata.name}'); \
|
||||||
|
kubectl $(KCTX) $(KNS) cp $$POD:/app/.data/logs/$(RUN).jsonl .data/logs/$(RUN).jsonl \
|
||||||
|
&& echo "→ .data/logs/$(RUN).jsonl"
|
||||||
|
|
||||||
|
# List the runs the api currently has log files for (most recent first).
|
||||||
|
run-logs:
|
||||||
|
@POD=$$(kubectl $(KCTX) $(KNS) get pod -l app=api -o jsonpath='{.items[0].metadata.name}'); \
|
||||||
|
kubectl $(KCTX) $(KNS) exec $$POD -- sh -c 'ls -1t /app/.data/logs 2>/dev/null || echo "(no logs yet)"'
|
||||||
|
|
||||||
# Unit tests run on the host venv (no postgres/llm calls needed).
|
# Unit tests run on the host venv (no postgres/llm calls needed).
|
||||||
# Ensures dev extras are installed first; uv is idempotent on no-ops.
|
# Ensures dev extras are installed first; uv is idempotent on no-ops.
|
||||||
test:
|
test:
|
||||||
|
|||||||
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.
|
"""compose(pick, recon) → ComposeResult.
|
||||||
|
|
||||||
Deterministic SQL emission from a typed Pick. The composer:
|
Deterministic SQL emission from a typed Pick. The composer assembles a single
|
||||||
1. Resolves the metric → from_table + sql expression + default filter.
|
sqlglot expression tree and renders it once — there is no hand-built SQL string
|
||||||
2. Resolves each group_by column → owning table (via recon.resolve_column).
|
and no after-the-fact re-parse; the tree *is* the structure, and
|
||||||
3. Computes the join graph: join_path(metric.from_table, owner) per dim,
|
`.sql(dialect="postgres", identify=True)` (every identifier quoted) is the one
|
||||||
deduped, walked in declared order to pick the right relationship edge.
|
place a string is produced.
|
||||||
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
|
The work is split across focused modules, each returning `exp` nodes:
|
||||||
recon, the composer raises `PickValidationError` — the caller surfaces it.
|
|
||||||
|
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
|
from __future__ import annotations
|
||||||
|
|
||||||
import sqlglot
|
from sqlglot import exp
|
||||||
|
|
||||||
from api.composer.dates import resolve_date_range
|
from api.composer.aliases import AliasMap
|
||||||
from api.composer.types import (
|
from api.composer.encodings import effective_encodings
|
||||||
ComposeResult,
|
from api.composer.filters import render_filter
|
||||||
Filter,
|
from api.composer.joins import build_joins
|
||||||
OrderBy,
|
from api.composer.metrics import qualify_metric_expr
|
||||||
Pick,
|
from api.composer.order import default_order_by, render_order_by
|
||||||
PickValidationError,
|
from api.composer.types import ComposeResult, Pick, PickValidationError
|
||||||
)
|
from api.recon.types import Recon
|
||||||
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:
|
def compose(pick: Pick, recon: Recon) -> ComposeResult:
|
||||||
"""Render `pick` to SQL against `recon`. Raises PickValidationError on
|
"""Render `pick` to SQL against `recon`. Raises PickValidationError on any
|
||||||
any reference the recon can't resolve."""
|
reference the recon can't resolve."""
|
||||||
|
|
||||||
# 1. Metric.
|
# 1. Metric → base table.
|
||||||
if pick.metric not in recon.metrics:
|
if pick.metric not in recon.metrics:
|
||||||
raise PickValidationError(
|
raise PickValidationError(
|
||||||
f"metric {pick.metric!r} not in recon (known: {sorted(recon.metrics)})"
|
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"
|
f"but no such table in recon"
|
||||||
)
|
)
|
||||||
|
|
||||||
aliases: dict[str, str] = {}
|
aliases = AliasMap()
|
||||||
aliases[metric.from_table] = _alias_for(metric.from_table, {})
|
base_alias = aliases.alias(metric.from_table)
|
||||||
base_alias = aliases[metric.from_table]
|
encodings = effective_encodings(recon)
|
||||||
|
|
||||||
# 2. Resolve every group_by column to (table, Column) and collect target
|
# 2. Resolve group_by columns → owning table + output alias. Tables other
|
||||||
# tables that aren't the metric's table — those need joins.
|
# than the metric's need joins.
|
||||||
group_bindings: list[tuple[str, str, str, str]] = [] # (ref, table, col_name, output_alias)
|
group_bindings: list[tuple[str, str, str]] = [] # (table, col_name, out_alias)
|
||||||
extra_tables: list[str] = []
|
extra_tables: list[str] = []
|
||||||
dim_select_alias: dict[str, str] = {}
|
dim_select_alias: dict[str, str] = {}
|
||||||
for ref in pick.group_by:
|
for ref in pick.group_by:
|
||||||
table, col = recon.resolve_column(ref)
|
table, col = recon.resolve_column(ref)
|
||||||
if table not in aliases:
|
if not aliases.has(table):
|
||||||
if table != metric.from_table:
|
if table != metric.from_table:
|
||||||
extra_tables.append(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(".", "_")
|
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
|
dim_select_alias[ref] = out_alias
|
||||||
|
|
||||||
# 3. Resolve filter columns first (they may also introduce new tables we
|
# 3. Resolve filters → predicate nodes (may also introduce new tables).
|
||||||
# need to join). Rendered separately so we can keep their predicates
|
filter_nodes: list[exp.Expression] = []
|
||||||
# in WHERE.
|
|
||||||
filter_clauses: list[str] = []
|
|
||||||
for f in pick.where:
|
for f in pick.where:
|
||||||
table, _ = recon.resolve_column(f.column)
|
table, _ = recon.resolve_column(f.column)
|
||||||
if table not in aliases:
|
if not aliases.has(table):
|
||||||
if table != metric.from_table:
|
if table != metric.from_table:
|
||||||
extra_tables.append(table)
|
extra_tables.append(table)
|
||||||
aliases[table] = _alias_for(table, {a: t for t, a in aliases.items()})
|
aliases.alias(table)
|
||||||
filter_clauses.append(_render_filter(f, recon, aliases))
|
filter_nodes.append(render_filter(f, recon, aliases, encodings))
|
||||||
|
|
||||||
# 4. Build JOINs for the union of extra tables.
|
# 4. JOINs for the union of extra tables.
|
||||||
join_clauses = _build_joins(recon, metric.from_table, extra_tables, aliases)
|
join_specs = build_joins(recon, metric.from_table, extra_tables, aliases)
|
||||||
|
|
||||||
# 5. SELECT list: dimension columns first (in group_by order), then the
|
# 5. SELECT list: dimensions (in group_by order), then the metric.
|
||||||
# metric expression.
|
select_nodes: list[exp.Expression] = [
|
||||||
select_parts: list[str] = []
|
exp.alias_(aliases.col(table, col_name), out_alias, quoted=True)
|
||||||
for ref, table, col_name, out_alias in group_bindings:
|
for 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_nodes.append(
|
||||||
select_parts.append(f'{metric_expr} AS "{metric.name}"')
|
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.
|
# 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:
|
if metric.filter:
|
||||||
where_parts.append(_qualify_metric_filter(metric.filter, base_alias))
|
where_nodes.append(qualify_metric_expr(metric.filter, base_alias))
|
||||||
where_parts.extend(filter_clauses)
|
where_nodes.extend(filter_nodes)
|
||||||
|
|
||||||
# 7. GROUP BY (just the dim output aliases) + ORDER BY + LIMIT.
|
# 7. ORDER BY: explicit, else largest-metric-first for grouped queries.
|
||||||
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:
|
if pick.order_by is not None:
|
||||||
order_by_clause = "ORDER BY " + _render_order_by(
|
order_node = render_order_by(pick.order_by, metric.name, dim_select_alias)
|
||||||
pick.order_by, metric.name, dim_select_alias
|
|
||||||
)
|
|
||||||
elif group_bindings:
|
elif group_bindings:
|
||||||
# Sensible default: largest metric first.
|
order_node = default_order_by(metric.name)
|
||||||
order_by_clause = f'ORDER BY "{metric.name}" DESC'
|
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 = sel.sql(dialect="postgres", identify=True)
|
||||||
sql_parts = [
|
return ComposeResult(sql=sql, used_tables=aliases.tables())
|
||||||
"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)
|
|
||||||
|
|||||||
@@ -1,22 +1,29 @@
|
|||||||
"""Date-range resolver for typed Filters.
|
"""Date-range resolver for typed Filters.
|
||||||
|
|
||||||
Given a `Filter` with `date_range="YYYY"` or `date_range="YYYY-Q[1-4]"` and
|
Given a `Filter` with `date_range="YYYY"` or `date_range="YYYY-Q[1-4]"`, the
|
||||||
the target column, render a SQL fragment that evaluates correctly under the
|
target column, and the active encoding map, build a SQL predicate node that
|
||||||
column's actual storage type. Dispatched in two layers:
|
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
|
1. `Column.semantic_type` if it names an encoding in the effective map — coerce
|
||||||
datasets that ship YYMMDD-encoded integers).
|
the raw column to a DATE via the encoding's `as_date` template, then compare.
|
||||||
2. Otherwise `Column.sql_type` (`DATE`, `TIMESTAMP`, …) — the default path
|
This is the extension point: BIRD's `date_yymmdd` ships as a default
|
||||||
for warehouses that store dates as real date columns.
|
(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
|
Returns an `exp.Between` node (inclusive on both ends) so it composes into the
|
||||||
date literals from strings).
|
query tree the composer assembles; the composer is the single place a SQL
|
||||||
|
string is produced.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import re
|
import re
|
||||||
|
|
||||||
|
import sqlglot
|
||||||
|
from sqlglot import exp
|
||||||
|
|
||||||
from api.composer.types import PickValidationError
|
from api.composer.types import PickValidationError
|
||||||
from api.recon.types import Column
|
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:
|
def resolve_date_range(
|
||||||
"""Render the WHERE fragment for `<qualified_col> BETWEEN ...`.
|
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
|
`col_expr` is the (already alias-qualified) column reference node;
|
||||||
the SQL (e.g. `"l"."date"`). Returns just the predicate, no `WHERE`.
|
`encodings` is the effective encoding map (see `api.composer.encodings`).
|
||||||
"""
|
"""
|
||||||
lo, hi = _bounds(spec)
|
lo, hi = _bounds(spec)
|
||||||
semantic = (column.semantic_type or "").lower()
|
bounds = (exp.convert(lo), exp.convert(hi))
|
||||||
|
|
||||||
if semantic == "date_yymmdd":
|
semantic = (column.semantic_type or "").lower()
|
||||||
# Column stores YYMMDD as integer (some BIRD datasets ship this way).
|
enc = encodings.get(semantic) if semantic else None
|
||||||
return (
|
if enc and enc.get("as_date"):
|
||||||
f"TO_DATE(LPAD({qualified_col}::text, 6, '0'), 'YYMMDD') "
|
# Coerce the raw column to a DATE via the dataset/default template.
|
||||||
f"BETWEEN '{lo}' AND '{hi}'"
|
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()
|
sql_type = column.sql_type.upper()
|
||||||
if sql_type.startswith("DATE") or sql_type.startswith("TIMESTAMP"):
|
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(
|
raise PickValidationError(
|
||||||
f"date_range on column with sql_type={column.sql_type!r} and "
|
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.
|
# tables.<name>: one-line description.
|
||||||
# columns.<table>.<col>: one-line description. Sparse — only the ones
|
# columns.<table>.<col>: one-line description. Sparse — only the ones
|
||||||
# worth describing. Undescribed columns just have
|
# 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).
|
# 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
|
schema: financial
|
||||||
|
|
||||||
|
|||||||
13
api/main.py
13
api/main.py
@@ -5,6 +5,7 @@ Endpoints:
|
|||||||
POST /ask — submit a question, returns run_id
|
POST /ask — submit a question, returns run_id
|
||||||
GET /runs/{run_id}/stream — SSE stream of run events
|
GET /runs/{run_id}/stream — SSE stream of run events
|
||||||
GET /runs/{run_id} — final state snapshot
|
GET /runs/{run_id} — final state snapshot
|
||||||
|
GET /runs/{run_id}/log — JSONL transcript (every published event)
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -17,7 +18,7 @@ from datetime import datetime, timezone
|
|||||||
|
|
||||||
from fastapi import FastAPI, HTTPException
|
from fastapi import FastAPI, HTTPException
|
||||||
from fastapi.middleware.cors import CORSMiddleware
|
from fastapi.middleware.cors import CORSMiddleware
|
||||||
from fastapi.responses import StreamingResponse
|
from fastapi.responses import FileResponse, StreamingResponse
|
||||||
from pydantic import BaseModel
|
from pydantic import BaseModel
|
||||||
|
|
||||||
from api.runtime import events
|
from api.runtime import events
|
||||||
@@ -116,3 +117,13 @@ async def stream_run(run_id: str):
|
|||||||
if run_id not in runs:
|
if run_id not in runs:
|
||||||
raise HTTPException(404, detail=f"Run {run_id} not found")
|
raise HTTPException(404, detail=f"Run {run_id} not found")
|
||||||
return StreamingResponse(events.stream(run_id), media_type="text/event-stream")
|
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 [])]
|
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]] = {}
|
column_to_tables: dict[str, list[str]] = {}
|
||||||
for t in tables.values():
|
for t in tables.values():
|
||||||
for c in t.columns:
|
for c in t.columns:
|
||||||
@@ -150,6 +155,7 @@ def merge_into_recon(dataset: str, extracted: dict[str, Any]) -> Recon:
|
|||||||
metrics=metrics,
|
metrics=metrics,
|
||||||
column_to_tables=column_to_tables,
|
column_to_tables=column_to_tables,
|
||||||
relationships=relationships,
|
relationships=relationships,
|
||||||
|
encodings=encodings,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -151,6 +151,12 @@ class Recon:
|
|||||||
metrics: dict[str, Metric]
|
metrics: dict[str, Metric]
|
||||||
column_to_tables: dict[str, list[str]] = field(default_factory=dict)
|
column_to_tables: dict[str, list[str]] = field(default_factory=dict)
|
||||||
relationships: list[Relationship] = field(default_factory=list)
|
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 ──
|
# ── 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()},
|
"metrics": {n: m.to_dict() for n, m in self.metrics.items()},
|
||||||
"column_to_tables": self.column_to_tables,
|
"column_to_tables": self.column_to_tables,
|
||||||
"relationships": [r.to_dict() for r in self.relationships],
|
"relationships": [r.to_dict() for r in self.relationships],
|
||||||
|
"encodings": self.encodings,
|
||||||
}
|
}
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
@@ -309,4 +316,5 @@ class Recon:
|
|||||||
metrics={n: Metric.from_dict(m) for n, m in d.get("metrics", {}).items()},
|
metrics={n: Metric.from_dict(m) for n, m in d.get("metrics", {}).items()},
|
||||||
column_to_tables=d.get("column_to_tables", {}),
|
column_to_tables=d.get("column_to_tables", {}),
|
||||||
relationships=[Relationship(**r) for r in d.get("relationships", [])],
|
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 asyncio
|
||||||
import json
|
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.analyses.types import Finding
|
||||||
from api.runtime.context import current_analysis, current_run_id
|
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.
|
# run_id -> queue of events. Events are plain dicts.
|
||||||
_queues: dict[str, asyncio.Queue[dict[str, Any] | None]] = {}
|
_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 ──
|
# ── Queue management ──
|
||||||
|
|
||||||
def open_run(run_id: str) -> asyncio.Queue:
|
def open_run(run_id: str) -> asyncio.Queue:
|
||||||
q: asyncio.Queue = asyncio.Queue()
|
q: asyncio.Queue = asyncio.Queue()
|
||||||
_queues[run_id] = q
|
_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
|
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:
|
async def publish(run_id: str, event: dict[str, Any]) -> None:
|
||||||
|
_write_log(run_id, event)
|
||||||
q = _queues.get(run_id)
|
q = _queues.get(run_id)
|
||||||
if q is not None:
|
if q is not None:
|
||||||
await q.put(event)
|
await q.put(event)
|
||||||
@@ -54,6 +89,12 @@ async def close(run_id: str) -> None:
|
|||||||
|
|
||||||
def drop(run_id: str) -> None:
|
def drop(run_id: str) -> None:
|
||||||
_queues.pop(run_id, 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 ──
|
# ── SSE formatting + streaming ──
|
||||||
|
|||||||
@@ -1,10 +1,25 @@
|
|||||||
"""Date-range resolver tests — bounds parsing + dispatch on column type."""
|
"""Date-range resolver tests — bounds parsing + dispatch on column type/encoding.
|
||||||
|
|
||||||
|
`resolve_date_range` now returns an `exp` node and takes the effective encoding
|
||||||
|
map (built-in defaults overlaid with any dataset overrides). Tests render the
|
||||||
|
node to SQL to assert on the predicate.
|
||||||
|
"""
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
from sqlglot import exp
|
||||||
|
|
||||||
from api.composer.dates import _bounds, resolve_date_range
|
from api.composer.dates import _bounds, resolve_date_range
|
||||||
|
from api.composer.encodings import DEFAULT_ENCODINGS, effective_encodings
|
||||||
from api.composer.types import PickValidationError
|
from api.composer.types import PickValidationError
|
||||||
from api.recon.types import Column
|
from api.recon.types import Column, Recon
|
||||||
|
|
||||||
|
|
||||||
|
def _col(name: str = "date", table: str = "l") -> exp.Column:
|
||||||
|
return exp.column(name, table=table, quoted=True)
|
||||||
|
|
||||||
|
|
||||||
|
def _sql(node: exp.Expression) -> str:
|
||||||
|
return node.sql(dialect="postgres", identify=True)
|
||||||
|
|
||||||
|
|
||||||
def test_bounds_year():
|
def test_bounds_year():
|
||||||
@@ -26,25 +41,56 @@ def test_bounds_unsupported_raises():
|
|||||||
|
|
||||||
def test_resolve_real_date_column():
|
def test_resolve_real_date_column():
|
||||||
col = Column(name="date", sql_type="DATE", nullable=False)
|
col = Column(name="date", sql_type="DATE", nullable=False)
|
||||||
out = resolve_date_range("1996", col, '"l"."date"')
|
out = _sql(resolve_date_range("1996", col, _col(), {}))
|
||||||
assert out == "\"l\".\"date\" BETWEEN '1996-01-01' AND '1996-12-31'"
|
assert out == "\"l\".\"date\" BETWEEN '1996-01-01' AND '1996-12-31'"
|
||||||
|
|
||||||
|
|
||||||
def test_resolve_timestamp_column():
|
def test_resolve_timestamp_column():
|
||||||
col = Column(name="created_at", sql_type="TIMESTAMP WITHOUT TIME ZONE", nullable=True)
|
col = Column(name="created_at", sql_type="TIMESTAMP WITHOUT TIME ZONE", nullable=True)
|
||||||
out = resolve_date_range("1996-Q2", col, '"x"."created_at"')
|
out = _sql(resolve_date_range("1996-Q2", col, _col("created_at", "x"), {}))
|
||||||
assert out == "\"x\".\"created_at\" BETWEEN '1996-04-01' AND '1996-06-30'"
|
assert out == "\"x\".\"created_at\" BETWEEN '1996-04-01' AND '1996-06-30'"
|
||||||
|
|
||||||
|
|
||||||
def test_resolve_yymmdd_semantic_type():
|
def test_resolve_yymmdd_via_default_encoding():
|
||||||
col = Column(name="date", sql_type="INTEGER", nullable=False,
|
col = Column(name="date", sql_type="INTEGER", nullable=False,
|
||||||
semantic_type="date_yymmdd")
|
semantic_type="date_yymmdd")
|
||||||
out = resolve_date_range("1996", col, '"l"."date"')
|
out = _sql(resolve_date_range("1996", col, _col(), DEFAULT_ENCODINGS))
|
||||||
assert "TO_DATE(LPAD(" in out
|
assert "TO_DATE(LPAD(" in out
|
||||||
assert "BETWEEN '1996-01-01' AND '1996-12-31'" in out
|
assert "BETWEEN '1996-01-01' AND '1996-12-31'" in out
|
||||||
|
|
||||||
|
|
||||||
|
def test_dataset_encoding_overrides_default():
|
||||||
|
# A dataset can override a built-in encoding with its own coercion SQL.
|
||||||
|
col = Column(name="d", sql_type="INTEGER", nullable=False,
|
||||||
|
semantic_type="date_yymmdd")
|
||||||
|
override = {"date_yymmdd": {"as_date": "MAKE_DATE(1900 + {col} / 10000, 1, 1)"}}
|
||||||
|
out = _sql(resolve_date_range("1996", col, _col("d"), override))
|
||||||
|
assert "MAKE_DATE(" in out
|
||||||
|
assert "TO_DATE" not in out
|
||||||
|
|
||||||
|
|
||||||
|
def test_resolve_semantic_type_without_encoding_raises():
|
||||||
|
# Declared semantic_type but no matching encoding and not a native date type.
|
||||||
|
col = Column(name="d", sql_type="INTEGER", nullable=False,
|
||||||
|
semantic_type="date_yymmdd")
|
||||||
|
with pytest.raises(PickValidationError, match="not supported"):
|
||||||
|
resolve_date_range("1996", col, _col("d"), {})
|
||||||
|
|
||||||
|
|
||||||
def test_resolve_unsupported_column_type_raises():
|
def test_resolve_unsupported_column_type_raises():
|
||||||
col = Column(name="x", sql_type="TEXT", nullable=True)
|
col = Column(name="x", sql_type="TEXT", nullable=True)
|
||||||
with pytest.raises(PickValidationError, match="not supported"):
|
with pytest.raises(PickValidationError, match="not supported"):
|
||||||
resolve_date_range("1996", col, '"x"."x"')
|
resolve_date_range("1996", col, _col("x", "x"), {})
|
||||||
|
|
||||||
|
|
||||||
|
def test_effective_encodings_merges_defaults_with_dataset():
|
||||||
|
recon = Recon(schema="t", tables={}, metrics={},
|
||||||
|
encodings={"my_epoch": {"as_date": "TO_TIMESTAMP({col})::date"}})
|
||||||
|
eff = effective_encodings(recon)
|
||||||
|
assert "date_yymmdd" in eff # default survives
|
||||||
|
assert eff["my_epoch"]["as_date"].startswith("TO_TIMESTAMP") # dataset added
|
||||||
|
|
||||||
|
# Dataset entry wins on key collision.
|
||||||
|
recon2 = Recon(schema="t", tables={}, metrics={},
|
||||||
|
encodings={"date_yymmdd": {"as_date": "CUSTOM({col})"}})
|
||||||
|
assert effective_encodings(recon2)["date_yymmdd"]["as_date"] == "CUSTOM({col})"
|
||||||
|
|||||||
@@ -141,7 +141,6 @@ function nodeClasses(n: GraphNode) {
|
|||||||
<div class="trace-pane">
|
<div class="trace-pane">
|
||||||
<Panel title="Plan" :status="statusOf(status)">
|
<Panel title="Plan" :status="statusOf(status)">
|
||||||
<div v-if="rationale" class="rationale">{{ rationale }}</div>
|
<div v-if="rationale" class="rationale">{{ rationale }}</div>
|
||||||
<div v-else class="muted small">Waiting for the planner…</div>
|
|
||||||
|
|
||||||
<div class="graph">
|
<div class="graph">
|
||||||
<div v-for="(n, idx) in graphNodes" :key="n.id" class="graph-row">
|
<div v-for="(n, idx) in graphNodes" :key="n.id" class="graph-row">
|
||||||
|
|||||||
Reference in New Issue
Block a user