spb/uqo-research-db

1# ============================================================2# Project: EDGAR Research Database3# Description: SEC EDGAR Research Infrastructure4#5# Authors:6# - Manel Kammoun — manel.kammoun@uqo.ca7# - Djerrycharli Tandjambianda — djerrycharli.tandjambianda@uqo.ca8# - Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca9#10# Institution:11# Université du Québec en Outaouais (UQO)12#13# License:14# Research / Academic Use15# ============================================================1617"""Initialise the DuckDB schema and load companies + filings from the metadata.1819Companies come from ``data/metadata/universe.json`` (S&P 500 ∪ Russell 3000;20falls back to the legacy ``sp500_constituents.json``), enriched with the SEC21submissions headers saved by the bulk downloader.  Filings come from the22download log(s).  Re-runnable: everything is INSERT OR REPLACE and a filing23already marked ``parse_status = 'complete'`` keeps that status.24"""2526from __future__ import annotations2728import argparse29import json30import sys31from pathlib import Path3233import pandas as pd3435sys.path.insert(0, str(Path(__file__).resolve().parent.parent))3637from src.storage.database import DatabaseManager38from src.storage.parquet_manager import ParquetManager39from src.utils.constants import DATA_DIRS40from src.utils.logging import setup_logging4142logger = setup_logging("build_database")434445def parse_args() -> argparse.Namespace:46    parser = argparse.ArgumentParser(description="Initialize and populate the EDGAR research database")47    parser.add_argument("--reset", action="store_true", default=False, help="Drop and recreate all tables")48    parser.add_argument("--only-ciks-file", type=str, default=None,49                        help="restrict the filings load to these CIKs (one per line)")50    parser.add_argument("--shard-node", type=str, default=None, help="tag companies with the node name")51    parser.add_argument("--no-parquet", action="store_true", help="skip the companies/filings parquet export")52    return parser.parse_args()535455def load_download_log(paths: list[Path]) -> list[dict]:56    entries: list[dict] = []57    for log_path in paths:58        if not log_path.exists():59            continue60        with open(log_path) as f:61            for line in f:62                line = line.strip()63                if not line:64                    continue65                try:66                    entries.append(json.loads(line))67                except json.JSONDecodeError:68                    logger.warning("Skipping malformed line in %s", log_path)69    logger.info("Loaded %d download-log entries from %d file(s)", len(entries), len(paths))70    return entries717273def load_companies() -> list[dict]:74    meta = DATA_DIRS["metadata"]75    universe_path = meta / "universe.json"76    if universe_path.exists():77        with open(universe_path) as f:78            constituents = json.load(f)["constituents"]79        logger.info("Universe: %d companies from %s", len(constituents), universe_path)80    else:81        with open(meta / "sp500_constituents.json") as f:82            constituents = json.load(f).get("constituents", [])83        for c in constituents:84            c["in_sp500"], c["in_russell3000"] = True, False85        logger.info("Universe: %d companies from legacy sp500_constituents.json", len(constituents))8687    sub_dir = meta / "submissions"88    records = []89    for c in constituents:90        cik = str(c.get("cik", "")).zfill(10) if c.get("cik") else ""91        if not cik or not c.get("ticker"):92            continue93        header = {}94        hp = sub_dir / f"{cik}.json"95        if hp.exists():96            try:97                header = json.loads(hp.read_text())98            except json.JSONDecodeError:99                header = {}100        date_added = c.get("date_added") or None101        if date_added and not str(date_added)[:4].isdigit():102            date_added = None103        records.append({104            "cik": cik,105            "ticker": c["ticker"],106            "company_name": c.get("company_name") or header.get("name") or "",107            "sector": c.get("sector", ""),108            "sub_industry": c.get("sub_industry", "") or header.get("sic_description", "") or "",109            "sic_code": header.get("sic") or None,110            "state": header.get("state_of_business") or header.get("state_of_incorporation") or None,111            "date_added": str(date_added)[:10] if date_added else None,112            "in_sp500": bool(c.get("in_sp500", False)),113            "in_russell3000": bool(c.get("in_russell3000", False)),114            "sic_description": header.get("sic_description") or None,115            "state_of_incorporation": header.get("state_of_incorporation") or None,116            "exchange": ",".join(header.get("exchanges") or []) or None,117            "fiscal_year_end": header.get("fiscal_year_end") or None,118            "russell_weight_pct": float(c.get("russell_weight_pct", 0.0) or 0.0),119        })120    return records121122123def main() -> None:124    args = parse_args()125    db = DatabaseManager()126127    if args.reset:128        logger.warning("Resetting database: dropping all tables")129        for table in ["processing_log", "embeddings_metadata", "filing_chunks", "filing_sections", "filings", "companies"]:130            try:131                db.execute(f"DROP TABLE IF EXISTS {table} CASCADE")132            except Exception:  # noqa: BLE001133                db.execute(f"DROP TABLE IF EXISTS {table}")134        for view in ["filing_summary", "section_detail", "chunk_detail"]:135            db.execute(f"DROP VIEW IF EXISTS {view}")136137    db.initialize()138139    companies = load_companies()140    if args.shard_node:141        for c in companies:142            c["shard_node"] = args.shard_node143    if companies:144        db.insert_df("companies", pd.DataFrame.from_records(companies))145        logger.info("Loaded %d companies into database", len(companies))146    known_ciks = {c["cik"] for c in companies}147    company_names = {c["cik"]: c["company_name"] for c in companies}148149    only: set[str] | None = None150    if args.only_ciks_file:151        only = {l.strip().zfill(10) for l in Path(args.only_ciks_file).read_text().splitlines() if l.strip()}152153    meta = DATA_DIRS["metadata"]154    logs = [meta / "download_log.json"] + sorted(meta.glob("download_log.*.jsonl"))155    entries = load_download_log(logs)156157    existing = {}158    try:159        existing = dict(db.execute("SELECT accession_number, parse_status FROM filings").fetchall())160    except Exception:  # noqa: BLE001161        pass162163    filing_records = []164    seen: set[str] = set()165    for e in entries:166        accession = e.get("accession_number", "")167        key = e.get("key", "")168        cik = key.split("|")[0] if "|" in key else ""169        if not accession or accession in seen or not cik:170            continue171        if only is not None and cik not in only:172            continue173        if cik not in known_ciks:174            logger.debug("CIK %s not in universe, filing %s skipped", cik, accession)175            continue176        seen.add(accession)177        primary = e.get("primary_document") or (key.split("|")[2] if key.count("|") >= 2 else "")178        if primary:179            local_path = str(DATA_DIRS["raw"] / cik / e.get("form_type", "") / accession.replace("-", "") / primary)180        else:181            local_path = e.get("local_path", "")182        filing_records.append({183            "accession_number": accession,184            "cik": cik,185            "ticker": e.get("ticker", ""),186            "company_name": company_names.get(cik, ""),187            "form_type": e.get("form_type", ""),188            "filing_date": e.get("filing_date", ""),189            "report_date": (e.get("report_date") or None),190            "sec_url": e.get("filing_url") or None,191            "local_path": local_path,192            "primary_document": primary or None,193            "file_size_bytes": e.get("file_size_bytes") or None,194            "download_status": "complete",195            "parse_status": existing.get(accession, "pending"),196        })197198    if filing_records:199        df = pd.DataFrame.from_records(filing_records)200        df["filing_date"] = pd.to_datetime(df["filing_date"], errors="coerce").dt.date201        df["report_date"] = pd.to_datetime(df["report_date"], errors="coerce").dt.date202        df["fiscal_year"] = pd.to_datetime(df["filing_date"], errors="coerce").dt.year.astype("Int64")203        df = df[df["filing_date"].notna()]204        n = db.insert_df("filings", df)205        logger.info("Loaded %d filings into database", n)206207    if not args.no_parquet:208        parquet = ParquetManager()209        for table_name in ["companies", "filings"]:210            try:211                if db.execute(f"SELECT COUNT(*) FROM {table_name}").fetchone()[0] > 0:212                    parquet.export_to_parquet(db, table_name)213            except Exception:  # noqa: BLE001214                logger.warning("Could not export '%s' to parquet", table_name, exc_info=True)215216    stats = db.get_stats()217    logger.info("Database statistics:")218    for table, count in stats.items():219        logger.info("  %-25s %d rows", table, count)220    db.close()221    logger.info("Database build complete")222223224if __name__ == "__main__":225    main()226