fithia2/apps/pipeline/filing_fetcher/main.py

"""Filing Fetcher: download exhibit text and cache locally."""
from __future__ import annotations

import argparse
import asyncio
import datetime as dt
import uuid

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__)


async def fetch_exhibits(run_id: str) -> dict[str, int]:
    settings = get_settings()
    app_config = settings.get_app_config()
    exhibit_types = app_config.get("pipeline", {}).get("exhibit_types", ["EX-99.1"])

    stats = {"seen": 0, "written": 0, "skipped": 0, "errors": 0}

    batch_size = 100

    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()

            result = await session.execute(
                select(Document).where(Document.parsed_status == "pending")
            )
            docs = result.scalars().all()
            stats["seen"] = len(docs)

            for idx, doc in enumerate(docs, 1):
                if not doc.accession_no:
                    stats["skipped"] += 1
                    continue

                fetched_any = False
                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,
                        )
                        fetched_any = True
                        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

                        cache_path = str(exhibit_path(doc.accession_no, exhibit_type))

                        existing_cache = await session.execute(
                            select(ExhibitCache).where(
                                ExhibitCache.accession_no == doc.accession_no,
                                ExhibitCache.exhibit_type == exhibit_type,
                            )
                        )
                        if existing_cache.scalar_one_or_none() is None:
                            cache_row = ExhibitCache(
                                accession_no=doc.accession_no,
                                exhibit_type=exhibit_type,
                                content_hash=checksum,
                                cache_path=cache_path,
                            )
                            session.add(cache_row)

                        fetched_any = True
                        stats["written"] += 1
                        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),
                        )
                        stats["errors"] += 1

                # Extract item_numbers from SGML header if not already set
                if not doc.item_numbers:
                    try:
                        items = await asyncio.wait_for(
                            svc.get_filing_items(doc.accession_no),
                            timeout=15.0,
                        )
                        if items:
                            doc.item_numbers = items
                            logger.info(
                                "item_numbers_extracted",
                                accession_no=doc.accession_no,
                                items=items,
                            )
                    except Exception as exc:
                        logger.debug(
                            "item_numbers_extraction_failed",
                            accession_no=doc.accession_no,
                            error=str(exc),
                        )

                # Always advance to ready_for_parse (exhibit may not exist)
                doc.parsed_status = "ready_for_parse"
                doc.updated_at_utc = dt.datetime.now(tz=dt.UTC)

                # Commit in batches to preserve progress
                if idx % batch_size == 0:
                    await session.commit()
                    logger.info(
                        "batch_committed",
                        processed=idx,
                        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())
    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))


if __name__ == "__main__":
    main()