"""Registry of target databases the pipeline knows about. The default registry is populated from disk: any SQLite file under data/ that matches a known shape (Chinook, BIRD slices) is auto-registered. Postgres- backed databases are registered explicitly when the docker-compose stack is running. """ from __future__ import annotations from dataclasses import dataclass, field from pathlib import Path from nl_sql.db.connection import DatabaseSpec, sqlite_url_readonly from nl_sql.paths import under_root # Anchored to the repo root (not CWD) so scanning finds the data/ tree no matter # where the process was launched from — Streamlit/uvicorn/pytest don't agree on CWD. DATA_ROOT = under_root("data") @dataclass(slots=True) class DatabaseRegistry: specs: dict[str, DatabaseSpec] = field(default_factory=dict) def register(self, spec: DatabaseSpec) -> None: self.specs[spec.id] = spec def get(self, db_id: str) -> DatabaseSpec: if db_id not in self.specs: raise KeyError(f"database {db_id!r} not registered. Known: {sorted(self.specs)}") return self.specs[db_id] def ids(self) -> list[str]: return sorted(self.specs) def get_default_registry( data_root: Path = DATA_ROOT, *, pg_dsn: str = "", pg_db_id: str = "pg_codebase_community", pg_description: str = "", ) -> DatabaseRegistry: """Build a registry by scanning the data/ tree. Resolution order: - data/chinook/Chinook.sqlite → id="chinook" - data/bird_mini_dev/MINIDEV/dev_databases//.sqlite → id=f"bird_{db}" When ``pg_dsn`` is non-empty, a Postgres-backed database is also registered under ``pg_db_id`` (load it first with scripts/load_postgres.py, or — for a BIRD slice with Postgres gold — scripts/extract_pg_dump_slice.py). The DSN should point at the read-only role; the engine additionally forces read-only transactions (see db/connection.py). Registration order matters: Postgres is registered *last*, so passing an id that a SQLite scan also produces (e.g. ``pg_db_id="bird_codebase_community"``) deliberately re-points that database at Postgres. That is how the Postgres eval runs the same BIRD questions against a different engine. """ registry = DatabaseRegistry() chinook_path = data_root / "chinook" / "Chinook.sqlite" if chinook_path.exists(): registry.register( DatabaseSpec( id="chinook", dialect="sqlite", url=sqlite_url_readonly(chinook_path), description="Chinook music store — invoices, tracks, customers (smoke / sanity).", ) ) bird_dev_root = data_root / "bird_mini_dev" / "MINIDEV" / "dev_databases" if bird_dev_root.is_dir(): for db_dir in sorted(p for p in bird_dev_root.iterdir() if p.is_dir()): sqlite_file = db_dir / f"{db_dir.name}.sqlite" if sqlite_file.exists(): registry.register( DatabaseSpec( id=f"bird_{db_dir.name}", dialect="sqlite", url=sqlite_url_readonly(sqlite_file), description=f"BIRD Mini-Dev / {db_dir.name}.", ) ) if pg_dsn: registry.register( DatabaseSpec( id=pg_db_id, dialect="postgresql", url=pg_dsn, description=pg_description, ) ) return registry