fithia2/apps/pipeline/cli.py

"""CLI for running the data pipeline.

Usage:
    fithia2 pipeline run              # run all 5 steps
    fithia2 pipeline run --step poller|fetcher|parser|features|labels
    fithia2 pipeline run --start-date 2026-03-15
"""
from __future__ import annotations

import argparse
import asyncio
import sys

from rich.console import Console

_console = Console(width=120)

STEPS = ["poller", "fetcher", "parser", "features", "labels"]


def _run_id() -> str:
    from libs.common.ids import new_job_run_id
    return new_job_run_id()


def _configure() -> None:
    import logging
    try:
        from dotenv import load_dotenv
        load_dotenv()
    except ImportError:
        pass
    from libs.common.config import get_settings
    from libs.common.logging import configure_logging
    configure_logging(get_settings().log_level)
    # Suppress noisy HTTP request logs from httpx
    logging.getLogger("httpx").setLevel(logging.WARNING)
    logging.getLogger("httpcore").setLevel(logging.WARNING)


async def _run_steps_async(
    steps: list[str],
    start_date: str | None,
    end_date: str | None,
    continue_on_error: bool = False,
) -> dict[str, dict | str]:
    """Run all pipeline steps in a single event loop to avoid asyncpg loop conflicts."""
    from apps.pipeline.filing_poller.main import poll_filings
    from apps.pipeline.filing_fetcher.main import fetch_exhibits
    from apps.pipeline.event_parser.main import run_event_parser
    from apps.pipeline.feature_builder.main import run_feature_builder
    from apps.pipeline.label_generator.main import run_label_generator

    results: dict[str, dict | str] = {}
    for step in steps:
        try:
            if step == "poller":
                results[step] = await poll_filings(_run_id(), start_date=start_date, end_date=end_date)
            elif step == "fetcher":
                results[step] = await fetch_exhibits(_run_id(), start_date=start_date, end_date=end_date)
            elif step == "parser":
                results[step] = await run_event_parser(_run_id(), start_date=start_date, end_date=end_date)
            elif step == "features":
                results[step] = await run_feature_builder(_run_id(), start_date=start_date, end_date=end_date)
            elif step == "labels":
                results[step] = await run_label_generator(_run_id(), start_date=start_date, end_date=end_date)
        except Exception as exc:
            results[step] = f"ERROR: {exc}"
            if not continue_on_error:
                raise
    return results


_STEP_NAMES = {
    "poller": "Filing Poller",
    "fetcher": "Filing Fetcher",
    "parser": "Event Parser",
    "features": "Feature Builder",
    "labels": "Label Generator",
}


def cmd_run(args: argparse.Namespace) -> None:
    _configure()

    steps = [args.step] if args.step else STEPS

    for step in steps:
        _console.print(f"\n[bold cyan]▶ {_STEP_NAMES[step]}[/]")

    try:
        all_results = asyncio.run(
            _run_steps_async(steps, args.start_date, args.end_date, args.continue_on_error)
        )
    except Exception as exc:
        _console.print(f"[red]FAILED: {exc}[/]")
        sys.exit(1)

    _console.print()
    failed = False
    for step in steps:
        result = all_results.get(step, {})
        if isinstance(result, str) and result.startswith("ERROR:"):
            _console.print(f"  [bold cyan]{_STEP_NAMES[step]}:[/] [red]{result}[/]")
            failed = True
        else:
            parts = [f"{k}={v}" for k, v in (result.items() if isinstance(result, dict) else [])]
            _console.print(f"  [bold cyan]{_STEP_NAMES[step]}:[/] [green]done[/]  {', '.join(parts)}")

    if failed:
        sys.exit(1)

    _console.print("\n[bold green]Pipeline complete.[/]")


def main() -> None:
    parser = argparse.ArgumentParser(description="fithia2 Pipeline Runner")
    sub = parser.add_subparsers(dest="command", required=True)

    p = sub.add_parser("run", help="Run pipeline steps")
    p.add_argument(
        "--step", "-s",
        choices=STEPS,
        default=None,
        help="Run a single step (default: all steps in order)",
    )
    p.add_argument(
        "--start-date",
        default=None,
        metavar="YYYY-MM-DD",
        help="Start date for filing poller (default: 7 days ago)",
    )
    p.add_argument(
        "--end-date",
        default=None,
        metavar="YYYY-MM-DD",
        help="End date for filing poller (default: today)",
    )
    p.add_argument(
        "--continue-on-error",
        action="store_true",
        help="Continue to next step even if a step fails",
    )

    args = parser.parse_args()
    if args.command == "run":
        cmd_run(args)


if __name__ == "__main__":
    main()