Create-knowledge-find-decode/scripts/run_acquisition.py

#!/usr/bin/env python3
"""Run formal acquisition against a cloud query batch."""
from __future__ import annotations

import argparse
import json
from uuid import UUID

from acquisition.repositories.postgres import PostgresAcquisitionRepository
from acquisition.runner import DEFAULT_PLATFORMS, run_batch
from core.config import CreationDbConfig, Settings
from core.db_session import transaction


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--batch-id", required=True, help="Formal query batch UUID")
    parser.add_argument(
        "--platform",
        action="append",
        choices=DEFAULT_PLATFORMS,
        help="Platform to run. Repeat to run multiple platforms. Default: all.",
    )
    parser.add_argument("--search-limit", type=int, default=10)
    parser.add_argument("--display-limit", type=int, default=5)
    parser.add_argument("--no-classify", action="store_true")
    parser.add_argument("--no-resume", action="store_true")
    parser.add_argument("--no-skip-done", action="store_true")
    parser.add_argument("--run-key")
    parser.add_argument("--env-file", default=".env")
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)
    settings = Settings.from_env(args.env_file)
    db_config = CreationDbConfig.from_env(args.env_file)
    platforms = tuple(args.platform or DEFAULT_PLATFORMS)
    with transaction(db_config) as conn:
        repo = PostgresAcquisitionRepository(conn)
        result = run_batch(
            repo,
            batch_id=UUID(args.batch_id),
            settings=settings,
            platforms=platforms,
            search_limit=args.search_limit,
            display_limit=args.display_limit,
            classify=not args.no_classify,
            resume=not args.no_resume,
            skip_done=not args.no_skip_done,
            run_key=args.run_key,
        )
    print(json.dumps(result.__dict__, ensure_ascii=False, default=str, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())