fithia2/apps/tools/fix_null_symbols.py

"""Fix documents and events with NULL symbol_id / issuer_id.

The filing_poller stores ticker in document_id as TICKER::{ticker}.
When issuer_sync hadn't run before filing_poller, the lookup maps were
empty and symbol_id / issuer_id were inserted as NULL.

This script:
  1. Extracts ticker from document_id for NULL-symbol docs
  2. Looks up correct symbol_id / issuer_id from symbol_master / issuer_master
  3. Updates documents and events

Usage:
    python -m apps.tools.fix_null_symbols [--dry-run]
"""
from __future__ import annotations

import argparse
import asyncio

from sqlalchemy import select, text, update

from libs.common.config import get_settings
from libs.common.logging import configure_logging, get_logger
from libs.db.models import Document, Event, IssuerMaster, SymbolMaster
from libs.db.session import get_session

logger = get_logger(__name__)


async def fix_null_symbols(dry_run: bool = False) -> dict[str, int]:
    stats = {"docs_total_null": 0, "docs_fixed": 0, "docs_no_match": 0,
             "events_fixed": 0, "events_no_match": 0}

    async with get_session() as session:
        # 1. Build ticker lookup maps from current DB
        issuer_rows = await session.execute(select(IssuerMaster))
        ticker_to_issuer: dict[str, str] = {}
        for im in issuer_rows.scalars().all():
            if im.ticker:
                ticker_to_issuer[im.ticker.upper()] = im.issuer_id

        symbol_rows = await session.execute(
            select(SymbolMaster).where(SymbolMaster.is_primary == True)  # noqa: E712
        )
        ticker_to_symbol: dict[str, str] = {}
        for sm in symbol_rows.scalars().all():
            if sm.ticker:
                ticker_to_symbol[sm.ticker.upper()] = sm.symbol_id

        print(f"Lookup maps: {len(ticker_to_issuer)} issuers, {len(ticker_to_symbol)} symbols")

        # 2. Find documents with NULL symbol_id
        null_docs = await session.execute(
            select(Document).where(Document.symbol_id.is_(None))
        )
        docs = null_docs.scalars().all()
        stats["docs_total_null"] = len(docs)
        print(f"Documents with NULL symbol_id: {len(docs)}")

        # 3. Extract ticker from document_id and update
        no_match_tickers: set[str] = set()
        for doc in docs:
            # document_id format: DOC::sec::TICKER::{ticker}::{date}::{accession}
            parts = doc.document_id.split("::")
            if len(parts) >= 4 and parts[2] == "TICKER":
                ticker = parts[3].upper()
            else:
                logger.warning("unexpected_doc_id_format", document_id=doc.document_id)
                stats["docs_no_match"] += 1
                continue

            sym_id = ticker_to_symbol.get(ticker)
            iss_id = ticker_to_issuer.get(ticker)

            if sym_id:
                doc.symbol_id = sym_id
                doc.issuer_id = iss_id  # may still be None if issuer doesn't exist
                stats["docs_fixed"] += 1
            else:
                no_match_tickers.add(ticker)
                stats["docs_no_match"] += 1

        if no_match_tickers:
            print(f"No symbol_master match for {len(no_match_tickers)} tickers: "
                  f"{sorted(no_match_tickers)[:20]}{'...' if len(no_match_tickers) > 20 else ''}")

        # 4. Fix events with NULL symbol_id (join via primary_document_id)
        null_events = await session.execute(
            select(Event).where(Event.symbol_id.is_(None))
        )
        events = null_events.scalars().all()
        print(f"Events with NULL symbol_id: {len(events)}")

        for evt in events:
            # Look up the parent document
            doc_result = await session.execute(
                select(Document).where(Document.document_id == evt.primary_document_id)
            )
            parent_doc = doc_result.scalar_one_or_none()
            if parent_doc and parent_doc.symbol_id:
                evt.symbol_id = parent_doc.symbol_id
                evt.issuer_id = parent_doc.issuer_id
                stats["events_fixed"] += 1
            else:
                stats["events_no_match"] += 1

        # 5. Commit or rollback
        if dry_run:
            print("\n[DRY RUN] Rolling back — no changes saved.")
            await session.rollback()
        else:
            await session.flush()
            print("\nChanges committed.")

    # Report
    print(f"\n{'='*50}")
    print("Fix NULL Symbols Report")
    print(f"{'='*50}")
    print(f"  Documents with NULL symbol_id: {stats['docs_total_null']}")
    print(f"  Documents fixed:               {stats['docs_fixed']}")
    print(f"  Documents no match:            {stats['docs_no_match']}")
    print(f"  Events fixed:                  {stats['events_fixed']}")
    print(f"  Events no match:               {stats['events_no_match']}")
    print(f"{'='*50}")

    return stats


def main() -> None:
    parser = argparse.ArgumentParser(description="Fix NULL symbol_id in documents and events")
    parser.add_argument("--dry-run", action="store_true", help="Preview changes without saving")
    args = parser.parse_args()

    settings = get_settings()
    configure_logging(settings.log_level)

    asyncio.run(fix_null_symbols(dry_run=args.dry_run))


if __name__ == "__main__":
    main()