verbose live UI + tool-level SSE events + Groq default + regression tests

ctrl/seed/load_bird.py

@@ -0,0 +1,147 @@
+"""Load a BIRD SQLite dataset into Postgres.
+
+Runs inside the api container (`make seed`). Reads which dataset to load
+from `settings.dataset` (env var `NVI_DATASET`, default `financial`) and
+expects the source SQLite at `ctrl/seed/data/<dataset>.sqlite` — produced
+on the host by `ctrl/seed/download.sh` and reaching the container via
+Tilt live_update sync (kind) or the docker-compose bind mount.
+
+DDL is built from SQLAlchemy `MetaData` + `Table` + `Column` objects, not
+string SQL. Bulk insert uses psycopg's COPY protocol via the engine's raw
+connection — there's no SQLAlchemy equivalent, and INSERT for ~1M rows
+would be orders of magnitude slower. Identifiers in the COPY statement
+are rendered through SA's dialect-aware identifier preparer.
+"""
+
+from __future__ import annotations
+
+import sqlite3
+import sys
+from pathlib import Path
+from typing import Any
+
+from sqlalchemy import (
+    BigInteger, Boolean, Column, Date, DateTime, Float, Integer, LargeBinary,
+    MetaData, Numeric, SmallInteger, Table, Text, text,
+)
+from sqlalchemy.engine import Engine
+
+from api.config import get_settings
+from api.tools.db import get_engine
+
+# SQLite is type-affinity, not strict — pick a SQLAlchemy type per affinity.
+# Anything unrecognised falls through to Text, which is always safe.
+TYPE_MAP: dict[str, Any] = {
+    "INT": BigInteger,
+    "INTEGER": BigInteger,
+    "TINYINT": SmallInteger,
+    "SMALLINT": SmallInteger,
+    "MEDIUMINT": Integer,
+    "BIGINT": BigInteger,
+    "REAL": Float,
+    "DOUBLE": Float,
+    "FLOAT": Float,
+    "NUMERIC": Numeric,
+    "DECIMAL": Numeric,
+    "BOOLEAN": Boolean,
+    "DATE": Date,
+    "DATETIME": DateTime,
+    "TIMESTAMP": DateTime,
+    "TEXT": Text,
+    "CHAR": Text,
+    "VARCHAR": Text,
+    "CLOB": Text,
+    "BLOB": LargeBinary,
+}
+
+
+def _sa_type(sqlite_type: str):
+    base = sqlite_type.upper().split("(")[0].strip()
+    return TYPE_MAP.get(base, Text)
+
+
+def _list_tables(src: sqlite3.Connection) -> list[str]:
+    return [
+        r[0] for r in src.execute(
+            "SELECT name FROM sqlite_master "
+            "WHERE type='table' AND name NOT LIKE 'sqlite_%' "
+            "ORDER BY name"
+        ).fetchall()
+    ]
+
+
+def _discover_table(src: sqlite3.Connection, metadata: MetaData,
+                    name: str, schema: str) -> Table:
+    cols = src.execute(f'PRAGMA table_info("{name}")').fetchall()
+    # PRAGMA returns (cid, name, type, notnull, dflt_value, pk).
+    sa_cols = [
+        Column(
+            c[1],
+            _sa_type(c[2]),
+            nullable=not c[3],
+            primary_key=bool(c[5]),
+        )
+        for c in cols
+    ]
+    return Table(name, metadata, *sa_cols, schema=schema)
+
+
+def _copy_table(engine: Engine, src: sqlite3.Connection, sa_table: Table) -> int:
+    """Bulk-load via psycopg COPY. Identifiers go through SA's dialect preparer
+    so we never hand-quote names."""
+    prep = engine.dialect.identifier_preparer
+    table_ident = prep.format_table(sa_table)                       # "<schema>"."<table>"
+    col_list = ", ".join(prep.quote(c.name) for c in sa_table.columns)
+    src_ident = '"' + sa_table.name.replace('"', '""') + '"'        # SQLite side
+
+    count = 0
+    raw = engine.raw_connection()
+    try:
+        cur = raw.cursor()
+        with cur.copy(f"COPY {table_ident} ({col_list}) FROM STDIN") as copy:
+            for row in src.execute(f"SELECT * FROM {src_ident}"):
+                copy.write_row(tuple(row))
+                count += 1
+                if count % 100_000 == 0:
+                    print(f"  {sa_table.name}: {count:,} rows...")
+        raw.commit()
+    finally:
+        raw.close()
+    return count
+
+
+def main() -> None:
+    dataset = get_settings().dataset
+    sqlite_path = Path(f"seed/data/{dataset}.sqlite")
+
+    if not sqlite_path.exists():
+        sys.exit(
+            f"{dataset} SQLite not found at {sqlite_path}.\n"
+            f"Run `bash ctrl/seed/download.sh` on the host first."
+        )
+
+    engine = get_engine()
+    src = sqlite3.connect(sqlite_path)
+    src.row_factory = sqlite3.Row
+
+    prep = engine.dialect.identifier_preparer
+    schema_ident = prep.quote_schema(dataset)
+    with engine.begin() as conn:
+        conn.execute(text(f"DROP SCHEMA IF EXISTS {schema_ident} CASCADE"))
+        conn.execute(text(f"CREATE SCHEMA {schema_ident}"))
+
+    metadata = MetaData(schema=dataset)
+    tables = [_discover_table(src, metadata, t, dataset) for t in _list_tables(src)]
+    metadata.create_all(engine)
+    print(f"created {len(tables)} tables in schema {dataset!r}")
+
+    for sa_table in tables:
+        count = _copy_table(engine, src, sa_table)
+        print(f"  loaded {sa_table.name}: {count:,} rows")
+
+    src.close()
+    print("done.")
+
+
+if __name__ == "__main__":
+    main()