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