# ============================================================ # Project: EDGAR Research Database # Description: SEC EDGAR Research Infrastructure # # Authors: # - Manel Kammoun — manel.kammoun@uqo.ca # - Djerrycharli Tandjambianda — djerrycharli.tandjambianda@uqo.ca # - Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca # # Institution: # Université du Québec en Outaouais (UQO) # # License: # Research / Academic Use # ============================================================ """Initialise the DuckDB schema and load companies + filings from the metadata. Companies come from ``data/metadata/universe.json`` (S&P 500 ∪ Russell 3000; falls back to the legacy ``sp500_constituents.json``), enriched with the SEC submissions headers saved by the bulk downloader. Filings come from the download log(s). Re-runnable: everything is INSERT OR REPLACE and a filing already marked ``parse_status = 'complete'`` keeps that status. """ from __future__ import annotations import argparse import json import sys from pathlib import Path import pandas as pd sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) from src.storage.database import DatabaseManager from src.storage.parquet_manager import ParquetManager from src.utils.constants import DATA_DIRS from src.utils.logging import setup_logging logger = setup_logging("build_database") def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Initialize and populate the EDGAR research database") parser.add_argument("--reset", action="store_true", default=False, help="Drop and recreate all tables") parser.add_argument("--only-ciks-file", type=str, default=None, help="restrict the filings load to these CIKs (one per line)") parser.add_argument("--shard-node", type=str, default=None, help="tag companies with the node name") parser.add_argument("--no-parquet", action="store_true", help="skip the companies/filings parquet export") return parser.parse_args() def load_download_log(paths: list[Path]) -> list[dict]: entries: list[dict] = [] for log_path in paths: if not log_path.exists(): continue with open(log_path) as f: for line in f: line = line.strip() if not line: continue try: entries.append(json.loads(line)) except json.JSONDecodeError: logger.warning("Skipping malformed line in %s", log_path) logger.info("Loaded %d download-log entries from %d file(s)", len(entries), len(paths)) return entries def load_companies() -> list[dict]: meta = DATA_DIRS["metadata"] universe_path = meta / "universe.json" if universe_path.exists(): with open(universe_path) as f: constituents = json.load(f)["constituents"] logger.info("Universe: %d companies from %s", len(constituents), universe_path) else: with open(meta / "sp500_constituents.json") as f: constituents = json.load(f).get("constituents", []) for c in constituents: c["in_sp500"], c["in_russell3000"] = True, False logger.info("Universe: %d companies from legacy sp500_constituents.json", len(constituents)) sub_dir = meta / "submissions" records = [] for c in constituents: cik = str(c.get("cik", "")).zfill(10) if c.get("cik") else "" if not cik or not c.get("ticker"): continue header = {} hp = sub_dir / f"{cik}.json" if hp.exists(): try: header = json.loads(hp.read_text()) except json.JSONDecodeError: header = {} date_added = c.get("date_added") or None if date_added and not str(date_added)[:4].isdigit(): date_added = None records.append({ "cik": cik, "ticker": c["ticker"], "company_name": c.get("company_name") or header.get("name") or "", "sector": c.get("sector", ""), "sub_industry": c.get("sub_industry", "") or header.get("sic_description", "") or "", "sic_code": header.get("sic") or None, "state": header.get("state_of_business") or header.get("state_of_incorporation") or None, "date_added": str(date_added)[:10] if date_added else None, "in_sp500": bool(c.get("in_sp500", False)), "in_russell3000": bool(c.get("in_russell3000", False)), "sic_description": header.get("sic_description") or None, "state_of_incorporation": header.get("state_of_incorporation") or None, "exchange": ",".join(header.get("exchanges") or []) or None, "fiscal_year_end": header.get("fiscal_year_end") or None, "russell_weight_pct": float(c.get("russell_weight_pct", 0.0) or 0.0), }) return records def main() -> None: args = parse_args() db = DatabaseManager() if args.reset: logger.warning("Resetting database: dropping all tables") for table in ["processing_log", "embeddings_metadata", "filing_chunks", "filing_sections", "filings", "companies"]: try: db.execute(f"DROP TABLE IF EXISTS {table} CASCADE") except Exception: # noqa: BLE001 db.execute(f"DROP TABLE IF EXISTS {table}") for view in ["filing_summary", "section_detail", "chunk_detail"]: db.execute(f"DROP VIEW IF EXISTS {view}") db.initialize() companies = load_companies() if args.shard_node: for c in companies: c["shard_node"] = args.shard_node if companies: db.insert_df("companies", pd.DataFrame.from_records(companies)) logger.info("Loaded %d companies into database", len(companies)) known_ciks = {c["cik"] for c in companies} company_names = {c["cik"]: c["company_name"] for c in companies} only: set[str] | None = None if args.only_ciks_file: only = {l.strip().zfill(10) for l in Path(args.only_ciks_file).read_text().splitlines() if l.strip()} meta = DATA_DIRS["metadata"] logs = [meta / "download_log.json"] + sorted(meta.glob("download_log.*.jsonl")) entries = load_download_log(logs) existing = {} try: existing = dict(db.execute("SELECT accession_number, parse_status FROM filings").fetchall()) except Exception: # noqa: BLE001 pass filing_records = [] seen: set[str] = set() for e in entries: accession = e.get("accession_number", "") key = e.get("key", "") cik = key.split("|")[0] if "|" in key else "" if not accession or accession in seen or not cik: continue if only is not None and cik not in only: continue if cik not in known_ciks: logger.debug("CIK %s not in universe, filing %s skipped", cik, accession) continue seen.add(accession) primary = e.get("primary_document") or (key.split("|")[2] if key.count("|") >= 2 else "") if primary: local_path = str(DATA_DIRS["raw"] / cik / e.get("form_type", "") / accession.replace("-", "") / primary) else: local_path = e.get("local_path", "") filing_records.append({ "accession_number": accession, "cik": cik, "ticker": e.get("ticker", ""), "company_name": company_names.get(cik, ""), "form_type": e.get("form_type", ""), "filing_date": e.get("filing_date", ""), "report_date": (e.get("report_date") or None), "sec_url": e.get("filing_url") or None, "local_path": local_path, "primary_document": primary or None, "file_size_bytes": e.get("file_size_bytes") or None, "download_status": "complete", "parse_status": existing.get(accession, "pending"), }) if filing_records: df = pd.DataFrame.from_records(filing_records) df["filing_date"] = pd.to_datetime(df["filing_date"], errors="coerce").dt.date df["report_date"] = pd.to_datetime(df["report_date"], errors="coerce").dt.date df["fiscal_year"] = pd.to_datetime(df["filing_date"], errors="coerce").dt.year.astype("Int64") df = df[df["filing_date"].notna()] n = db.insert_df("filings", df) logger.info("Loaded %d filings into database", n) if not args.no_parquet: parquet = ParquetManager() for table_name in ["companies", "filings"]: try: if db.execute(f"SELECT COUNT(*) FROM {table_name}").fetchone()[0] > 0: parquet.export_to_parquet(db, table_name) except Exception: # noqa: BLE001 logger.warning("Could not export '%s' to parquet", table_name, exc_info=True) stats = db.get_stats() logger.info("Database statistics:") for table, count in stats.items(): logger.info(" %-25s %d rows", table, count) db.close() logger.info("Database build complete") if __name__ == "__main__": main()