209 lines
6.9 KiB
Python
209 lines
6.9 KiB
Python
"""Recon build — two-stage:
|
|
|
|
1. **DDL extraction** (auto, from Postgres): reads every table's columns +
|
|
types + nullability via SQLAlchemy Inspector and writes
|
|
`api/datasets/<name>/extracted_schema.json`. This is the structural
|
|
source of truth. Re-run whenever the warehouse changes.
|
|
|
|
2. **Augmentation merge** (human YAML): reads `schema_docs.yaml` (sparse —
|
|
table descriptions, column descriptions, relationships) and `metrics.yaml`,
|
|
merges them onto the extracted schema, and writes
|
|
`api/datasets/<name>/recon.json` — the artefact the runtime reads.
|
|
|
|
Run via `make recon` or `python -m api.recon.build`. Auto-triggered by
|
|
load_recon() on first call if recon.json is missing.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import logging
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import yaml
|
|
from sqlalchemy import inspect
|
|
|
|
from api.datasets import DATASETS_DIR
|
|
from api.recon.types import Column, Metric, Recon, Relationship, Table
|
|
from api.tools.db import get_engine
|
|
|
|
logger = logging.getLogger("nvi.recon.build")
|
|
|
|
|
|
# ── Stage 1: DDL extraction ──
|
|
|
|
def extract_ddl(dataset: str) -> dict[str, Any]:
|
|
"""Introspect Postgres for `<dataset>` schema → structural dict.
|
|
|
|
Shape:
|
|
{
|
|
"schema": "<name>",
|
|
"tables": {
|
|
"<table>": {
|
|
"columns": [
|
|
{"name": ..., "sql_type": ..., "nullable": ...},
|
|
...
|
|
]
|
|
},
|
|
...
|
|
}
|
|
}
|
|
"""
|
|
insp = inspect(get_engine())
|
|
tables: dict[str, dict[str, Any]] = {}
|
|
for name in insp.get_table_names(schema=dataset):
|
|
cols = []
|
|
for c in insp.get_columns(name, schema=dataset):
|
|
cols.append({
|
|
"name": c["name"],
|
|
"sql_type": str(c["type"]).upper(),
|
|
"nullable": bool(c["nullable"]),
|
|
})
|
|
tables[name] = {"columns": cols}
|
|
return {"schema": dataset, "tables": tables}
|
|
|
|
|
|
def write_extracted_schema(dataset: str, extracted: dict[str, Any]) -> Path:
|
|
out = DATASETS_DIR / dataset / "extracted_schema.json"
|
|
out.write_text(json.dumps(extracted, indent=2) + "\n", encoding="utf-8")
|
|
return out
|
|
|
|
|
|
# ── Stage 2: augmentation merge ──
|
|
|
|
def _read_yaml(path: Path) -> dict[str, Any]:
|
|
if not path.exists():
|
|
return {}
|
|
return yaml.safe_load(path.read_text()) or {}
|
|
|
|
|
|
def _parse_column_specs(raw: dict[str, Any]) -> dict[str, dict[str, Any]]:
|
|
"""Accept either the flat form `{table.col: <spec>}` or the nested form
|
|
`{table: {col: <spec>}}`. Each `<spec>` is either a bare description
|
|
string OR a dict `{desc?: str, type?: str}` where `type` is a domain-
|
|
level semantic type (e.g. "date_yymmdd"). Returns flat form: each value
|
|
is normalised to `{desc, type}`."""
|
|
def norm(v: Any) -> dict[str, Any]:
|
|
if isinstance(v, dict):
|
|
return {"desc": v.get("desc") or v.get("description"), "type": v.get("type")}
|
|
return {"desc": v, "type": None}
|
|
|
|
out: dict[str, dict[str, Any]] = {}
|
|
for k, v in raw.items():
|
|
if isinstance(v, dict) and not ("desc" in v or "description" in v or "type" in v):
|
|
# Nested: {table: {col: spec}}.
|
|
for col, sub in v.items():
|
|
out[f"{k}.{col}"] = norm(sub)
|
|
else:
|
|
out[k] = norm(v)
|
|
return out
|
|
|
|
|
|
def merge_into_recon(dataset: str, extracted: dict[str, Any]) -> Recon:
|
|
"""Combine extracted DDL with human-authored augmentation YAML."""
|
|
schema_name = extracted.get("schema", dataset)
|
|
|
|
aug = _read_yaml(DATASETS_DIR / dataset / "schema_docs.yaml")
|
|
metrics_yaml = _read_yaml(DATASETS_DIR / dataset / "metrics.yaml")
|
|
|
|
table_descs: dict[str, str] = aug.get("tables", {}) or {}
|
|
col_specs = _parse_column_specs(aug.get("columns", {}) or {})
|
|
|
|
tables: dict[str, Table] = {}
|
|
for tname, t_data in extracted["tables"].items():
|
|
cols = []
|
|
for c in t_data["columns"]:
|
|
spec = col_specs.get(f"{tname}.{c['name']}", {})
|
|
cols.append(Column(
|
|
name=c["name"],
|
|
sql_type=c["sql_type"],
|
|
nullable=c["nullable"],
|
|
description=spec.get("desc"),
|
|
semantic_type=spec.get("type"),
|
|
))
|
|
tables[tname] = Table(
|
|
name=tname,
|
|
description=table_descs.get(tname),
|
|
columns=cols,
|
|
)
|
|
|
|
metrics: dict[str, Metric] = {
|
|
name: Metric.from_spec(name, spec)
|
|
for name, spec in (metrics_yaml.get("metrics", {}) or {}).items()
|
|
}
|
|
|
|
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]] = {}
|
|
for t in tables.values():
|
|
for c in t.columns:
|
|
column_to_tables.setdefault(c.name, []).append(t.name)
|
|
for k in column_to_tables:
|
|
column_to_tables[k] = sorted(column_to_tables[k])
|
|
|
|
return Recon(
|
|
schema=schema_name,
|
|
tables=tables,
|
|
metrics=metrics,
|
|
column_to_tables=column_to_tables,
|
|
relationships=relationships,
|
|
encodings=encodings,
|
|
)
|
|
|
|
|
|
def write_recon(dataset: str, recon: Recon) -> Path:
|
|
out = DATASETS_DIR / dataset / "recon.json"
|
|
out.write_text(json.dumps(recon.to_dict(), indent=2) + "\n", encoding="utf-8")
|
|
return out
|
|
|
|
|
|
# ── Public entry: full build ──
|
|
|
|
def build_recon(dataset: str) -> Recon:
|
|
"""End-to-end: extract DDL, write the extracted artefact, merge with
|
|
YAML, write the recon artefact. Returns the in-memory Recon."""
|
|
extracted = extract_ddl(dataset)
|
|
write_extracted_schema(dataset, extracted)
|
|
recon = merge_into_recon(dataset, extracted)
|
|
write_recon(dataset, recon)
|
|
return recon
|
|
|
|
|
|
def main() -> None:
|
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s %(message)s")
|
|
ap = argparse.ArgumentParser(description="Build the recon artefacts for one or all datasets.")
|
|
ap.add_argument("--dataset", help="Dataset name (default: all subdirs of api/datasets/).")
|
|
args = ap.parse_args()
|
|
|
|
if args.dataset:
|
|
targets = [args.dataset]
|
|
else:
|
|
targets = [
|
|
p.name for p in DATASETS_DIR.iterdir()
|
|
if p.is_dir() and not p.name.startswith("_") and not p.name.startswith(".")
|
|
]
|
|
|
|
if not targets:
|
|
print("no datasets to build", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
for name in targets:
|
|
recon = build_recon(name)
|
|
print(
|
|
f"recon[{name}]: {len(recon.tables)} tables, "
|
|
f"{len(recon.metrics)} metrics, {len(recon.relationships)} relationships "
|
|
f"→ api/datasets/{name}/{{extracted_schema,recon}}.json"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|