148 lines
4.6 KiB
Python
148 lines
4.6 KiB
Python
"""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()
|