You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

247 lines
8.5 KiB
Python

"""Filing Fetcher: download exhibit text and cache locally."""
from __future__ import annotations
import argparse
import asyncio
import datetime as dt
import uuid
from dataclasses import dataclass
from sqlalchemy import select
from libs.common.config import get_settings
from libs.common.file_store import exists_exhibit, write_exhibit
from libs.common.ids import new_job_run_id
from libs.common.logging import bind_job_run_id, configure_logging, get_logger
from libs.db.models import Document, ExhibitCache, JobRun
from libs.db.session import get_session
from libs.oracle_client import FilingsService, make_oracle_client
from libs.oracle_client.exceptions import OracleNotFoundError
logger = get_logger(__name__)
@dataclass(slots=True)
class _FetchedExhibit:
exhibit_type: str
checksum: str
cache_path: str
@dataclass(slots=True)
class _FetchResult:
doc_id: str
accession_no: str | None
fetched: list[_FetchedExhibit]
item_numbers: list[str] | None
errors: int
async def _fetch_doc(
svc: FilingsService,
doc: Document,
exhibit_types: list[str],
) -> _FetchResult:
if not doc.accession_no:
return _FetchResult(
doc_id=str(doc.document_id),
accession_no=None,
fetched=[],
item_numbers=None,
errors=0,
)
fetched: list[_FetchedExhibit] = []
errors = 0
for exhibit_type in exhibit_types:
if exists_exhibit(doc.accession_no, exhibit_type):
logger.info(
"exhibit_already_cached",
accession_no=doc.accession_no,
exhibit_type=exhibit_type,
)
continue
try:
response = await asyncio.wait_for(
svc.get_exhibit(doc.accession_no, exhibit_type),
timeout=30.0,
)
checksum = write_exhibit(doc.accession_no, exhibit_type, response.content)
from libs.common.file_store import exhibit_path
fetched.append(
_FetchedExhibit(
exhibit_type=exhibit_type,
checksum=checksum,
cache_path=str(exhibit_path(doc.accession_no, exhibit_type)),
)
)
logger.info(
"exhibit_fetched",
accession_no=doc.accession_no,
exhibit_type=exhibit_type,
)
except OracleNotFoundError:
logger.warning(
"exhibit_not_found",
accession_no=doc.accession_no,
exhibit_type=exhibit_type,
)
except Exception as exc:
logger.error(
"exhibit_fetch_error",
accession_no=doc.accession_no,
exhibit_type=exhibit_type,
error=str(exc),
)
errors += 1
item_numbers: list[str] | None = None
if not doc.item_numbers:
try:
item_numbers = await asyncio.wait_for(
svc.get_filing_items(doc.accession_no),
timeout=15.0,
)
if item_numbers:
logger.info(
"item_numbers_extracted",
accession_no=doc.accession_no,
items=item_numbers,
)
except Exception as exc:
logger.debug(
"item_numbers_extraction_failed",
accession_no=doc.accession_no,
error=str(exc),
)
return _FetchResult(
doc_id=str(doc.document_id),
accession_no=doc.accession_no,
fetched=fetched,
item_numbers=item_numbers,
errors=errors,
)
async def fetch_exhibits(
run_id: str,
start_date: str | None = None,
end_date: str | None = None,
) -> dict[str, int]:
settings = get_settings()
app_config = settings.get_app_config()
exhibit_types = app_config.get("pipeline", {}).get("exhibit_types", ["EX-99.1"])
concurrency = max(1, int(app_config.get("pipeline", {}).get("fetcher_concurrency", 20)))
stats = {"seen": 0, "written": 0, "skipped": 0, "errors": 0}
batch_size = max(10, int(app_config.get("pipeline", {}).get("fetcher_batch_size", 25)))
async with make_oracle_client() as client:
svc = FilingsService(client)
async with get_session() as session:
job = JobRun(
job_run_id=uuid.UUID(run_id),
job_name="filing_fetcher",
source_name="oracle",
run_date=dt.date.today(),
status="running",
)
session.add(job)
await session.flush()
await session.commit()
stmt = select(Document).where(Document.parsed_status == "pending")
if start_date:
stmt = stmt.where(Document.filing_date >= dt.date.fromisoformat(start_date))
if end_date:
stmt = stmt.where(Document.filing_date <= dt.date.fromisoformat(end_date))
result = await session.execute(stmt.order_by(Document.filing_date, Document.document_id))
docs = result.scalars().all()
stats["seen"] = len(docs)
doc_map = {str(doc.document_id): doc for doc in docs}
for batch_start in range(0, len(docs), batch_size):
batch_docs = docs[batch_start: batch_start + batch_size]
for task_start in range(0, len(batch_docs), concurrency):
task_docs = batch_docs[task_start: task_start + concurrency]
results = await asyncio.gather(
*[_fetch_doc(svc, doc, exhibit_types) for doc in task_docs]
)
for result_row in results:
doc = doc_map[result_row.doc_id]
if not result_row.accession_no:
stats["skipped"] += 1
continue
for fetched in result_row.fetched:
existing_cache = await session.execute(
select(ExhibitCache).where(
ExhibitCache.accession_no == result_row.accession_no,
ExhibitCache.exhibit_type == fetched.exhibit_type,
)
)
if existing_cache.scalar_one_or_none() is None:
session.add(
ExhibitCache(
accession_no=result_row.accession_no,
exhibit_type=fetched.exhibit_type,
content_hash=fetched.checksum,
cache_path=fetched.cache_path,
)
)
stats["written"] += 1
if result_row.item_numbers:
doc.item_numbers = result_row.item_numbers
doc.parsed_status = "ready_for_parse"
doc.updated_at_utc = dt.datetime.now(tz=dt.UTC)
stats["errors"] += result_row.errors
processed = min(batch_start + len(batch_docs), len(docs))
await session.commit()
logger.info(
"batch_committed",
processed=processed,
total=stats["seen"],
written=stats["written"],
errors=stats["errors"],
)
job.status = "succeeded" if stats["errors"] == 0 else "partial"
job.finished_at_utc = dt.datetime.now(tz=dt.UTC)
job.records_seen = stats["seen"]
job.records_written = stats["written"]
job.records_skipped = stats["skipped"]
job.error_count = stats["errors"]
logger.info("filing_fetcher_done", **stats)
return stats
def main() -> None:
parser = argparse.ArgumentParser(description="Filing Fetcher")
parser.add_argument("--run-id", default=new_job_run_id())
parser.add_argument("--start-date", default=None, metavar="YYYY-MM-DD")
parser.add_argument("--end-date", default=None, metavar="YYYY-MM-DD")
args = parser.parse_args()
settings = get_settings()
configure_logging(settings.log_level)
bind_job_run_id(args.run_id)
asyncio.run(fetch_exhibits(args.run_id, start_date=args.start_date, end_date=args.end_date))
if __name__ == "__main__":
main()