67 lines
2.6 KiB
Python
67 lines
2.6 KiB
Python
"""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
|