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# ============================================================1617from __future__ import annotations1819import uuid20from datetime import datetime, timezone21from pathlib import Path22from typing import Any2324import duckdb25import pandas as pd26import polars as pl2728from src.storage.schema import DatabaseSchema29from src.utils.constants import DUCKDB_PATH30from src.utils.logging import get_logger3132logger = get_logger("storage.database")333435class DatabaseManager:3637 def __init__(self, db_path: Path | str | None = None) -> None:38 self.db_path = Path(db_path) if db_path else DUCKDB_PATH39 self.db_path.parent.mkdir(parents=True, exist_ok=True)40 self.conn = duckdb.connect(str(self.db_path))41 logger.info("Connected to DuckDB at %s", self.db_path)4243 def __enter__(self) -> DatabaseManager:44 return self4546 def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:47 self.close()4849 def close(self) -> None:50 if self.conn:51 self.conn.close()52 self.conn = None53 logger.info("Closed DuckDB connection")5455 def initialize(self) -> None:56 for sql in DatabaseSchema.get_create_tables_sql():57 self.conn.execute(sql)58 for sql in DatabaseSchema.get_upgrade_sql():59 try:60 self.conn.execute(sql)61 except Exception: # noqa: BLE001 - column already present62 pass63 for sql in DatabaseSchema.get_create_indexes_sql():64 self.conn.execute(sql)65 for sql in DatabaseSchema.get_create_views_sql():66 self.conn.execute(sql)67 logger.info("Database schema initialized")6869 def execute(self, sql: str, params: list | tuple | None = None) -> duckdb.DuckDBPyRelation:70 if params:71 return self.conn.execute(sql, params)72 return self.conn.execute(sql)7374 def execute_many(self, sql: str, data: list) -> None:75 self.conn.executemany(sql, data)7677 def fetch_df(self, sql: str, params: list | tuple | None = None) -> pd.DataFrame:78 result = self.execute(sql, params)79 return result.fetchdf()8081 def fetch_polars(self, sql: str, params: list | tuple | None = None) -> pl.DataFrame:82 result = self.execute(sql, params)83 return result.pl()8485 # ---- CRUD: companies ----8687 def insert_company(self, company: dict) -> None:88 cols = list(company.keys())89 placeholders = ", ".join(["?"] * len(cols))90 col_names = ", ".join(cols)91 sql = f"INSERT OR REPLACE INTO companies ({col_names}) VALUES ({placeholders})"92 self.conn.execute(sql, list(company.values()))9394 def insert_companies(self, companies: list[dict]) -> None:95 if not companies:96 return97 cols = list(companies[0].keys())98 placeholders = ", ".join(["?"] * len(cols))99 col_names = ", ".join(cols)100 sql = f"INSERT OR REPLACE INTO companies ({col_names}) VALUES ({placeholders})"101 self.conn.executemany(sql, [list(c.values()) for c in companies])102 logger.info("Inserted %d companies", len(companies))103104 # ---- CRUD: filings ----105106 def insert_filing(self, filing: dict) -> None:107 cols = list(filing.keys())108 placeholders = ", ".join(["?"] * len(cols))109 col_names = ", ".join(cols)110 sql = f"INSERT OR REPLACE INTO filings ({col_names}) VALUES ({placeholders})"111 self.conn.execute(sql, list(filing.values()))112113 def insert_filings(self, filings: list[dict]) -> None:114 if not filings:115 return116 cols = list(filings[0].keys())117 placeholders = ", ".join(["?"] * len(cols))118 col_names = ", ".join(cols)119 sql = f"INSERT OR REPLACE INTO filings ({col_names}) VALUES ({placeholders})"120 self.conn.executemany(sql, [list(f.values()) for f in filings])121 logger.info("Inserted %d filings", len(filings))122123 # ---- CRUD: sections ----124125 def insert_sections(self, sections: list[dict]) -> None:126 if not sections:127 return128 cols = list(sections[0].keys())129 placeholders = ", ".join(["?"] * len(cols))130 col_names = ", ".join(cols)131 sql = f"INSERT OR REPLACE INTO filing_sections ({col_names}) VALUES ({placeholders})"132 self.conn.executemany(sql, [list(s.values()) for s in sections])133 logger.info("Inserted %d sections", len(sections))134135 # ---- CRUD: chunks ----136137 def insert_chunks(self, chunks: list[dict]) -> None:138 if not chunks:139 return140 cols = list(chunks[0].keys())141 placeholders = ", ".join(["?"] * len(cols))142 col_names = ", ".join(cols)143 sql = f"INSERT OR REPLACE INTO filing_chunks ({col_names}) VALUES ({placeholders})"144 self.conn.executemany(sql, [list(c.values()) for c in chunks])145 logger.info("Inserted %d chunks", len(chunks))146147 _PRIMARY_KEYS = {148 "companies": "cik", "filings": "accession_number", "filing_sections": "section_id",149 "filing_chunks": "chunk_id", "embeddings_metadata": "embedding_id",150 }151152 # columns that are indexed / part of a foreign key: DuckDB rewrites their UPDATE as delete+insert,153 # which fails on referenced rows, so they are never updated in place154 _FROZEN_COLS = {"filings": {"cik", "form_type", "filing_date", "fiscal_year"}}155156 def insert_df(self, table: str, df: pd.DataFrame, upsert: bool = True) -> int:157 """Bulk upsert from a DataFrame (orders of magnitude faster than executemany).158159 Existing rows are UPDATEd column by column and new rows INSERTed, which is160 the only form DuckDB accepts for rows referenced by foreign keys161 (companies <- filings <- sections <- chunks).162 """163 if df is None or df.empty:164 return 0165 cols = list(df.columns)166 col_list = ", ".join(cols)167 pk = self._PRIMARY_KEYS.get(table)168 self.conn.register("_df_in", df)169 try:170 if upsert and pk and pk in cols:171 frozen = self._FROZEN_COLS.get(table, set())172 upd_cols = [c for c in cols if c != pk and c not in frozen]173 if upd_cols:174 updates = ", ".join(f"{c} = d.{c}" for c in upd_cols)175 self.conn.execute(176 f"UPDATE {table} SET {updates} FROM _df_in d WHERE {table}.{pk} = d.{pk}"177 )178 self.conn.execute(179 f"INSERT INTO {table} ({col_list}) SELECT {col_list} FROM _df_in d "180 f"WHERE NOT EXISTS (SELECT 1 FROM {table} t WHERE t.{pk} = d.{pk})"181 )182 else:183 self.conn.execute(f"INSERT OR REPLACE INTO {table} ({col_list}) SELECT {col_list} FROM _df_in")184 finally:185 self.conn.unregister("_df_in")186 return len(df)187188 def insert_records_df(self, table: str, records: list[dict]) -> int:189 if not records:190 return 0191 return self.insert_df(table, pd.DataFrame.from_records(records))192193 # ---- Status updates ----194195 def update_filing_status(self, accession_number: str, field: str, status: str) -> None:196 allowed_fields = {"download_status", "parse_status"}197 if field not in allowed_fields:198 raise ValueError(f"field must be one of {allowed_fields}, got '{field}'")199 sql = f"UPDATE filings SET {field} = ?, updated_at = current_timestamp WHERE accession_number = ?"200 self.conn.execute(sql, [status, accession_number])201202 def log_operation(203 self,204 operation: str,205 entity_type: str,206 entity_id: str,207 status: str,208 message: str = "",209 ) -> None:210 log_id = str(uuid.uuid4())211 now = datetime.now(timezone.utc).isoformat()212 sql = """213 INSERT INTO processing_log (log_id, operation, entity_type, entity_id, status, message, started_at, completed_at)214 VALUES (?, ?, ?, ?, ?, ?, ?, ?)215 """216 self.conn.execute(sql, [log_id, operation, entity_type, entity_id, status, message, now, now])217218 # ---- Query methods ----219220 def get_filings(221 self,222 cik: str | None = None,223 ticker: str | None = None,224 form_type: str | None = None,225 start_date: str | None = None,226 end_date: str | None = None,227 fiscal_year: int | None = None,228 ) -> pd.DataFrame:229 conditions: list[str] = []230 params: list[Any] = []231232 if cik:233 conditions.append("cik = ?")234 params.append(cik)235 if ticker:236 conditions.append("ticker = ?")237 params.append(ticker)238 if form_type:239 conditions.append("form_type = ?")240 params.append(form_type)241 if start_date:242 conditions.append("filing_date >= ?")243 params.append(start_date)244 if end_date:245 conditions.append("filing_date <= ?")246 params.append(end_date)247 if fiscal_year is not None:248 conditions.append("fiscal_year = ?")249 params.append(fiscal_year)250251 where = f" WHERE {' AND '.join(conditions)}" if conditions else ""252 sql = f"SELECT * FROM filings{where} ORDER BY filing_date DESC"253 return self.fetch_df(sql, params if params else None)254255 def get_sections(self, accession_number: str) -> pd.DataFrame:256 sql = "SELECT * FROM filing_sections WHERE accession_number = ? ORDER BY section_order"257 return self.fetch_df(sql, [accession_number])258259 def get_chunks(260 self,261 accession_number: str | None = None,262 section_id: str | None = None,263 ) -> pd.DataFrame:264 conditions: list[str] = []265 params: list[Any] = []266267 if accession_number:268 conditions.append("accession_number = ?")269 params.append(accession_number)270 if section_id:271 conditions.append("section_id = ?")272 params.append(section_id)273274 where = f" WHERE {' AND '.join(conditions)}" if conditions else ""275 sql = f"SELECT * FROM filing_chunks{where} ORDER BY chunk_order"276 return self.fetch_df(sql, params if params else None)277278 def get_companies(self) -> pd.DataFrame:279 return self.fetch_df("SELECT * FROM companies ORDER BY ticker")280281 def get_filing_count(self) -> int:282 result = self.conn.execute("SELECT COUNT(*) FROM filings").fetchone()283 return result[0]284285 def get_stats(self) -> dict:286 tables = ["companies", "filings", "filing_sections", "filing_chunks", "embeddings_metadata", "processing_log"]287 stats: dict[str, int] = {}288 for table in tables:289 result = self.conn.execute(f"SELECT COUNT(*) FROM {table}").fetchone()290 stats[table] = result[0]291 return stats292