spb/uqo-research-db
Public
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"""Parse downloaded filings into sections and chunks, in bounded batches.1819 python scripts/parse_filings.py --max-workers 26 --batch-size 1500 --no-json2021Only ``parse_status = 'pending'`` filings are processed, so the script can be22re-run after an interruption. Results of each batch are written to DuckDB23before the next batch starts (memory stays bounded).24"""2526from __future__ import annotations2728import argparse29import sys30import time31from pathlib import Path3233import pandas as pd3435sys.path.insert(0, str(Path(__file__).resolve().parent.parent))3637from src.parsing import FilingParser38from src.storage.database import DatabaseManager39from src.utils.config import get_parsing_config40from src.utils.logging import setup_logging4142logger = setup_logging("parse_filings")4344SECTION_COLS = ["section_id", "accession_number", "section_name", "item_number",45 "section_order", "text", "word_count", "char_count"]46CHUNK_COLS = ["chunk_id", "section_id", "accession_number", "chunk_text", "chunk_order",47 "word_count", "char_count", "start_char", "end_char"]484950def parse_args() -> argparse.Namespace:51 parser = argparse.ArgumentParser(description="Parse downloaded SEC filings into sections and chunks")52 parser.add_argument("--form-type", type=str, default=None, help="Filter by form type (e.g., 10-K)")53 parser.add_argument("--max-workers", type=int, default=4, help="Number of parallel workers (default: 4)")54 parser.add_argument("--batch-size", type=int, default=1000, help="Filings per batch (default: 1000)")55 parser.add_argument("--accession", type=str, default=None, help="Parse a specific filing by accession number")56 parser.add_argument("--no-json", action="store_true", help="Do not write per-filing JSON to data/processed")57 parser.add_argument("--limit", type=int, default=None, help="Stop after N filings (testing)")58 return parser.parse_args()596061def main() -> None:62 args = parse_args()63 db = DatabaseManager()64 db.initialize()6566 if args.accession:67 filing_df = db.fetch_df(68 "SELECT accession_number, form_type, local_path FROM filings WHERE accession_number = ?",69 [args.accession])70 else:71 conditions = ["download_status = 'complete'", "parse_status = 'pending'"]72 params: list = []73 if args.form_type:74 conditions.append("form_type = ?")75 params.append(args.form_type)76 where = " AND ".join(conditions)77 sql = (f"SELECT accession_number, form_type, local_path FROM filings WHERE {where} "78 f"ORDER BY filing_date")79 if args.limit:80 sql += f" LIMIT {int(args.limit)}"81 filing_df = db.fetch_df(sql, params if params else None)8283 if filing_df.empty:84 logger.info("No filings to parse")85 db.close()86 return87 logger.info("Found %d filings to parse", len(filing_df))8889 file_list, missing = [], []90 for row in filing_df.itertuples(index=False):91 lp = row.local_path or ""92 if not lp or not Path(lp).exists():93 missing.append(row.accession_number)94 continue95 file_list.append({"file_path": Path(lp), "form_type": row.form_type,96 "accession_number": row.accession_number})97 if missing:98 logger.warning("%d filings have no local file; marking parse_status='missing'", len(missing))99 for acc in missing:100 db.update_filing_status(acc, "parse_status", "missing")101102 config = dict(get_parsing_config())103 if args.no_json:104 config["save_json"] = False105 parser = FilingParser(config=config)106107 total = len(file_list)108 succeeded = failed = 0109 t0 = time.time()110 for start in range(0, total, args.batch_size):111 batch = file_list[start:start + args.batch_size]112 results = parser.parse_batch(batch, max_workers=args.max_workers)113114 sections, chunks, ok_acc, err_acc = [], [], [], []115 for r in results:116 if r.get("status") != "success":117 err_acc.append(r["accession_number"])118 continue119 sections.extend({k: s.get(k) for k in SECTION_COLS} for s in r["sections"])120 chunks.extend({k: c.get(k) for k in CHUNK_COLS} for c in r["chunks"])121 ok_acc.append(r["accession_number"])122 del results123124 try:125 db.conn.execute("BEGIN")126 if sections:127 db.insert_records_df("filing_sections", sections)128 if chunks:129 db.insert_records_df("filing_chunks", chunks)130 if ok_acc:131 db.conn.register("_ok", pd.DataFrame({"accession_number": ok_acc}))132 db.conn.execute("UPDATE filings SET parse_status='complete', updated_at=current_timestamp "133 "WHERE accession_number IN (SELECT accession_number FROM _ok)")134 db.conn.unregister("_ok")135 if err_acc:136 db.conn.register("_err", pd.DataFrame({"accession_number": err_acc}))137 db.conn.execute("UPDATE filings SET parse_status='error', updated_at=current_timestamp "138 "WHERE accession_number IN (SELECT accession_number FROM _err)")139 db.conn.unregister("_err")140 db.conn.execute("COMMIT")141 except Exception:142 db.conn.execute("ROLLBACK")143 logger.error("Failed to store batch starting at %d", start, exc_info=True)144 raise145 succeeded += len(ok_acc)146 failed += len(err_acc)147 el = time.time() - t0148 done = min(start + args.batch_size, total)149 logger.info("batch done %d/%d filings | +%d sections +%d chunks | ok=%d err=%d | %.0fs (%.1f filings/s)",150 done, total, len(sections), len(chunks), succeeded, failed, el, done / max(el, 1))151152 logger.info("Parsing complete: %d succeeded, %d failed", succeeded, failed)153 db.close()154155156if __name__ == "__main__":157 main()158