313 lines
11 KiB
Python
313 lines
11 KiB
Python
"""Recon types — the dataset's knowledge graph as seen by the rest of the api.
|
|
|
|
The dataset-model types (Column, Table, Metric, SchemaContext) used to live
|
|
in `api/tools/types.py`. They're authored here now because they describe the
|
|
dataset, not the tools that consume it.
|
|
|
|
`Recon` is a SchemaContext extended with derived indexes that let Analyses
|
|
ask things like "which table owns column A2?" or "how do I join loan to
|
|
district?" without re-deriving from raw YAML each time.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
|
|
# ── Schema model ──
|
|
|
|
@dataclass
|
|
class Column:
|
|
name: str
|
|
sql_type: str
|
|
nullable: bool
|
|
description: str | None = None
|
|
# Domain-level type that the composer dispatches on for filter rendering.
|
|
# Examples: "date_yymmdd" (BIRD's int-encoded YYMMDD dates), "date" (real
|
|
# postgres date). None = treat literally with the column's sql_type.
|
|
semantic_type: str | None = None
|
|
|
|
@classmethod
|
|
def from_inspector(cls, info: dict[str, Any], description: str | None = None) -> "Column":
|
|
"""Build from a SQLAlchemy Inspector column dict."""
|
|
return cls(
|
|
name=info["name"],
|
|
sql_type=str(info["type"]).upper(),
|
|
nullable=bool(info["nullable"]),
|
|
description=description,
|
|
)
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
return {
|
|
"name": self.name,
|
|
"sql_type": self.sql_type,
|
|
"nullable": self.nullable,
|
|
"description": self.description,
|
|
"semantic_type": self.semantic_type,
|
|
}
|
|
|
|
@classmethod
|
|
def from_dict(cls, d: dict[str, Any]) -> "Column":
|
|
return cls(
|
|
name=d["name"], sql_type=d["sql_type"],
|
|
nullable=d["nullable"], description=d.get("description"),
|
|
semantic_type=d.get("semantic_type"),
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class Table:
|
|
name: str
|
|
description: str | None
|
|
columns: list[Column]
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
return {
|
|
"name": self.name,
|
|
"description": self.description,
|
|
"columns": [c.to_dict() for c in self.columns],
|
|
}
|
|
|
|
@classmethod
|
|
def from_dict(cls, d: dict[str, Any]) -> "Table":
|
|
return cls(
|
|
name=d["name"],
|
|
description=d.get("description"),
|
|
columns=[Column.from_dict(c) for c in d.get("columns", [])],
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class Metric:
|
|
name: str
|
|
description: str
|
|
sql: str
|
|
from_table: str
|
|
filter: str | None
|
|
unit: str | None
|
|
|
|
@classmethod
|
|
def from_spec(cls, name: str, spec: dict[str, Any]) -> "Metric":
|
|
return cls(
|
|
name=name,
|
|
description=spec.get("description", ""),
|
|
sql=spec["sql"],
|
|
from_table=spec["from_table"],
|
|
filter=spec.get("filter"),
|
|
unit=spec.get("unit"),
|
|
)
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
return {
|
|
"name": self.name,
|
|
"description": self.description,
|
|
"sql": self.sql,
|
|
"from_table": self.from_table,
|
|
"filter": self.filter,
|
|
"unit": self.unit,
|
|
}
|
|
|
|
@classmethod
|
|
def from_dict(cls, d: dict[str, Any]) -> "Metric":
|
|
return cls(
|
|
name=d["name"], description=d.get("description", ""),
|
|
sql=d["sql"], from_table=d["from_table"],
|
|
filter=d.get("filter"), unit=d.get("unit"),
|
|
)
|
|
|
|
|
|
# ── Relationships ──
|
|
|
|
@dataclass
|
|
class Relationship:
|
|
from_table: str
|
|
from_column: str
|
|
to_table: str
|
|
to_column: str
|
|
|
|
@classmethod
|
|
def parse(cls, raw: dict[str, str]) -> "Relationship":
|
|
"""Accept either `{from: 'a.b', to: 'c.d'}` or `{from_table, from_column, ...}`."""
|
|
if "from_table" in raw:
|
|
return cls(raw["from_table"], raw["from_column"], raw["to_table"], raw["to_column"])
|
|
ft, fc = raw["from"].split(".", 1)
|
|
tt, tc = raw["to"].split(".", 1)
|
|
return cls(ft, fc, tt, tc)
|
|
|
|
def to_dict(self) -> dict[str, str]:
|
|
return {
|
|
"from_table": self.from_table, "from_column": self.from_column,
|
|
"to_table": self.to_table, "to_column": self.to_column,
|
|
}
|
|
|
|
|
|
# ── Recon — top-level dataset knowledge bundle ──
|
|
|
|
@dataclass
|
|
class Recon:
|
|
schema: str
|
|
tables: dict[str, Table]
|
|
metrics: dict[str, Metric]
|
|
column_to_tables: dict[str, list[str]] = field(default_factory=dict)
|
|
relationships: list[Relationship] = field(default_factory=list)
|
|
|
|
# ── Read-side helpers used by Analyses + the composer ──
|
|
|
|
def table_names(self) -> list[str]:
|
|
return sorted(self.tables)
|
|
|
|
def metric_names(self) -> list[str]:
|
|
return sorted(self.metrics)
|
|
|
|
def owning_tables(self, column: str) -> list[str]:
|
|
"""Tables that have a column with this name (case-sensitive)."""
|
|
return self.column_to_tables.get(column, [])
|
|
|
|
def resolve_column(self, ref: str) -> tuple[str, "Column"]:
|
|
"""Resolve a `column` or `table.column` reference to its (table_name,
|
|
Column) pair. Raises ValueError if the column doesn't exist or is
|
|
ambiguous (lives on multiple tables) without an explicit `table.` prefix.
|
|
|
|
Used by the composer to bind every Pick.group_by entry to a single
|
|
owning table before SQL composition. Disambiguation is on the caller:
|
|
if you mean `account.district_id`, say so."""
|
|
if "." in ref:
|
|
table, col = ref.split(".", 1)
|
|
if table not in self.tables:
|
|
raise ValueError(f"unknown table {table!r} in column ref {ref!r}")
|
|
for c in self.tables[table].columns:
|
|
if c.name == col:
|
|
return table, c
|
|
raise ValueError(f"column {col!r} not found on table {table!r}")
|
|
owners = self.owning_tables(ref)
|
|
if not owners:
|
|
raise ValueError(f"no table has a column named {ref!r}")
|
|
if len(owners) > 1:
|
|
raise ValueError(
|
|
f"column {ref!r} is ambiguous (lives in: {', '.join(owners)}); "
|
|
f"qualify it as 'table.{ref}'"
|
|
)
|
|
table = owners[0]
|
|
for c in self.tables[table].columns:
|
|
if c.name == ref:
|
|
return table, c
|
|
# Index says it's there but the table doesn't — shouldn't happen.
|
|
raise ValueError(f"column {ref!r} indexed under {table!r} but not found")
|
|
|
|
def join_path(self, src: str, dst: str) -> list[str] | None:
|
|
"""Shortest sequence of tables from src to dst via declared relationships.
|
|
Returns None if no path exists. The path includes both endpoints."""
|
|
if src == dst:
|
|
return [src]
|
|
if src not in self.tables or dst not in self.tables:
|
|
return None
|
|
# Build undirected adjacency once.
|
|
adj: dict[str, set[str]] = {t: set() for t in self.tables}
|
|
for r in self.relationships:
|
|
adj.setdefault(r.from_table, set()).add(r.to_table)
|
|
adj.setdefault(r.to_table, set()).add(r.from_table)
|
|
# BFS.
|
|
from collections import deque
|
|
prev: dict[str, str | None] = {src: None}
|
|
q: deque[str] = deque([src])
|
|
while q:
|
|
cur = q.popleft()
|
|
if cur == dst:
|
|
# reconstruct
|
|
path: list[str] = []
|
|
node: str | None = cur
|
|
while node is not None:
|
|
path.append(node)
|
|
node = prev[node]
|
|
return list(reversed(path))
|
|
for nb in adj.get(cur, ()):
|
|
if nb not in prev:
|
|
prev[nb] = cur
|
|
q.append(nb)
|
|
return None
|
|
|
|
# ── Prompt-rendering helpers ──
|
|
|
|
def render_tables(self, names: list[str] | None = None) -> str:
|
|
"""CREATE TABLE-like rendering for prompt context."""
|
|
sel = [self.tables[n] for n in (names or self.table_names()) if n in self.tables]
|
|
out: list[str] = []
|
|
for t in sel:
|
|
header = f"-- {t.description}\n" if t.description else ""
|
|
cols: list[str] = []
|
|
for c in t.columns:
|
|
line = f' "{c.name}" {c.sql_type}'
|
|
if not c.nullable:
|
|
line += " NOT NULL"
|
|
if c.description:
|
|
line += f" -- {c.description}"
|
|
cols.append(line)
|
|
out.append(header + f'CREATE TABLE "{t.name}" (\n' + ",\n".join(cols) + "\n);")
|
|
return "\n\n".join(out)
|
|
|
|
def render_metrics(self) -> str:
|
|
if not self.metrics:
|
|
return "(no metrics defined)"
|
|
lines: list[str] = []
|
|
for m in self.metrics.values():
|
|
lines.append(
|
|
f"- {m.name} ({m.unit or 'unitless'}): {m.description}\n"
|
|
f" sql: {m.sql}\n"
|
|
f" from: {m.from_table}"
|
|
+ (f"\n filter: {m.filter}" if m.filter else "")
|
|
)
|
|
return "\n".join(lines)
|
|
|
|
def render_metrics_brief(self) -> str:
|
|
"""One line per metric — name, unit, description. No SQL, no table.
|
|
|
|
For the planner: it picks WHICH analyses to run; it doesn't need
|
|
implementation detail. Saves prompt tokens and cuts the noise the
|
|
model has to filter through."""
|
|
if not self.metrics:
|
|
return "(no metrics defined)"
|
|
return "\n".join(
|
|
f"- {m.name} [{m.unit or 'unitless'}]: {m.description}"
|
|
for m in self.metrics.values()
|
|
)
|
|
|
|
def render_tables_brief(self) -> str:
|
|
"""One line per table — name + description. No columns, no types."""
|
|
if not self.tables:
|
|
return "(no tables)"
|
|
lines: list[str] = []
|
|
for t in self.tables.values():
|
|
desc = t.description or "(no description)"
|
|
lines.append(f"- {t.name}: {desc}")
|
|
return "\n".join(lines)
|
|
|
|
def render_relationships(self) -> str:
|
|
if not self.relationships:
|
|
return "(no relationships declared)"
|
|
return "\n".join(
|
|
f" {r.from_table}.{r.from_column} → {r.to_table}.{r.to_column}"
|
|
for r in self.relationships
|
|
)
|
|
|
|
# ── Cache serialisation ──
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
return {
|
|
"schema": self.schema,
|
|
"tables": {n: t.to_dict() for n, t in self.tables.items()},
|
|
"metrics": {n: m.to_dict() for n, m in self.metrics.items()},
|
|
"column_to_tables": self.column_to_tables,
|
|
"relationships": [r.to_dict() for r in self.relationships],
|
|
}
|
|
|
|
@classmethod
|
|
def from_dict(cls, d: dict[str, Any]) -> "Recon":
|
|
return cls(
|
|
schema=d["schema"],
|
|
tables={n: Table.from_dict(t) for n, t in d.get("tables", {}).items()},
|
|
metrics={n: Metric.from_dict(m) for n, m in d.get("metrics", {}).items()},
|
|
column_to_tables=d.get("column_to_tables", {}),
|
|
relationships=[Relationship(**r) for r in d.get("relationships", [])],
|
|
)
|